Source code for pynpoint.util.mcmc

"""
Functions for MCMC sampling.
"""

import math

from typing import Optional, Tuple

import numpy as np

from typeguard import typechecked

from pynpoint.util.analysis import fake_planet, merit_function
from pynpoint.util.psf import pca_psf_subtraction
from pynpoint.util.residuals import combine_residuals


[docs]@typechecked def lnprob(param: np.ndarray, bounds: Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]], images: np.ndarray, psf: np.ndarray, mask: np.ndarray, parang: np.ndarray, psf_scaling: float, pixscale: float, pca_number: int, extra_rot: float, aperture: Tuple[int, int, float], indices: np.ndarray, merit: str, residuals: str, var_noise: Optional[float]) -> float: """ Function for the log posterior function. Should be placed at the highest level of the Python module to be pickable for the multiprocessing. Parameters ---------- param : numpy.ndarray The separation (arcsec), angle (deg), and contrast (mag). The angle is measured in counterclockwise direction with respect to the positive y-axis. bounds : tuple(tuple(float, float), tuple(float, float), tuple(float, float)) The boundaries of the separation (arcsec), angle (deg), and contrast (mag). Each set of boundaries is specified as a tuple. images : numpy.ndarray Stack with images. psf : numpy.ndarray PSF template, either a single image (2D) or a cube (3D) with the dimensions equal to *images*. mask : numpy.ndarray Array with the circular mask (zeros) of the central and outer regions. parang : numpy.ndarray Array with the angles for derotation. psf_scaling : float Additional scaling factor of the planet flux (e.g., to correct for a neutral density filter). Should be negative in order to inject negative fake planets. pixscale : float Additional scaling factor of the planet flux (e.g., to correct for a neutral density filter). Should be negative in order to inject negative fake planets. pca_number : int Number of principal components used for the PSF subtraction. extra_rot : float Additional rotation angle of the images (deg). aperture : tuple(int, int, float) Position (y, x) of the aperture center (pix) and aperture radius (pix). indices : numpy.ndarray Non-masked image indices. merit : str Figure of merit that is used for the likelihood function ('gaussian' or 'poisson'). Pixels are assumed to be independent measurements which are expected to be equal to zero in case the best-fit negative PSF template is injected. With 'gaussian', the variance is estimated from the pixel values within an annulus at the separation of the aperture (but excluding the pixels within the aperture). With 'poisson', a Poisson distribution is assumed for the variance of each pixel value. residuals : str Method used for combining the residuals ('mean', 'median', 'weighted', or 'clipped'). var_noise : float, None Variance of the noise which is required when `merit` is set to 'gaussian' or 'hessian'. Returns ------- float Log posterior probability. """ @typechecked def _lnprior() -> float: """ Internal function for the log prior function. Returns ------- float Log prior. """ if bounds[0][0] <= param[0] <= bounds[0][1] and \ bounds[1][0] <= param[1] <= bounds[1][1] and \ bounds[2][0] <= param[2] <= bounds[2][1]: ln_prior = 0. else: ln_prior = -np.inf return ln_prior @typechecked def _lnlike() -> float: """ Internal function for the log likelihood function. Returns ------- float Log likelihood. """ sep, ang, mag = param fake = fake_planet(images=images, psf=psf, parang=parang-extra_rot, position=(sep/pixscale, ang), magnitude=mag, psf_scaling=psf_scaling) im_res_rot, im_res_derot = pca_psf_subtraction(images=fake*mask, angles=-1.*parang+extra_rot, pca_number=pca_number, indices=indices) res_stack = combine_residuals(method=residuals, res_rot=im_res_derot, residuals=im_res_rot, angles=parang) chi_square = merit_function(residuals=res_stack[0, ], merit=merit, aperture=aperture, sigma=0., var_noise=var_noise) return -0.5*chi_square ln_prior = _lnprior() if math.isinf(ln_prior): ln_prob = -np.inf else: ln_prob = ln_prior + _lnlike() return ln_prob