Source code for vip_hci.fm.utils_negfc

#! /usr/bin/env python
"""
Module with post-processing related functions called from within the NEGFC
algorithm.
"""

__author__ = 'Carlos Alberto Gomez Gonzalez'
__all__ = ['cube_planet_free',
           'find_nearest']

import numpy as np
from ..fm import cube_inject_companions


[docs] def cube_planet_free(planet_parameter, cube, angs, psfn, imlib='vip-fft', interpolation='lanczos4', transmission=None): """ Return a cube in which we have injected negative fake companion at the position/flux given by planet_parameter. Parameters ---------- planet_parameter: numpy.array or list or tuple The (r, theta, flux) for all known companions. For a 4d cube r, theta and flux must all be 1d arrays with length equal to cube.shape[0]; i.e. planet_parameter should have shape: (n_pl,3,n_ch). cube: numpy ndarray The cube of fits images expressed as a numpy.array. angs: numpy ndarray The parallactic angle fits image expressed as a numpy.array. psfn: 2d or 3d numpy ndarray The normalized psf expressed as a numpy ndarray. Can be 3d for a 4d (spectral+ADI) input cube. imlib : str, optional See the documentation of the ``vip_hci.preproc.frame_rotate`` function. interpolation : str, optional See the documentation of the ``vip_hci.preproc.frame_rotate`` function. transmission: numpy array, optional Radial transmission of the coronagraph, if any. Array with either 2 x n_rad, 1+n_ch x n_rad columns. The first column should contain the radial separation in pixels, while the other column(s) are the corresponding off-axis transmission (between 0 and 1), for either all, or each spectral channel (only relevant for a 4D input cube). Returns ------- cpf : numpy.array The cube with negative companions injected at the position given in planet_parameter. """ cpf = np.zeros_like(cube) # unify planet_parameter format planet_parameter = np.array(planet_parameter) cond1 = cube.ndim == 3 and planet_parameter.ndim < 2 cond2 = cube.ndim == 4 and planet_parameter.ndim < 3 if cond1 or cond2: planet_parameter = planet_parameter[np.newaxis, :] if cube.ndim == 4: if planet_parameter.shape[2] != cube.shape[0]: raise TypeError("Input planet parameter with wrong dimensions.") for i in range(planet_parameter.shape[0]): if i == 0: cube_temp = cube else: cube_temp = cpf if cube.ndim == 4: for j in range(cube.shape[0]): flevel = -planet_parameter[i, 2, j] r = planet_parameter[i, 0, j] theta = planet_parameter[i, 1, j] cpf[j] = cube_inject_companions(cube_temp[j], psfn[j], angs, flevel=flevel, rad_dists=[r], n_branches=1, theta=theta, imlib=imlib, interpolation=interpolation, verbose=False, transmission=transmission) else: cpf = cube_inject_companions(cube_temp, psfn, angs, n_branches=1, flevel=-planet_parameter[i, 2], rad_dists=[planet_parameter[i, 0]], theta=planet_parameter[i, 1], imlib=imlib, verbose=False, interpolation=interpolation, transmission=transmission) return cpf
[docs] def find_nearest(array, value, output='index', constraint=None, n=1): """ Function to find the indices, and optionally the values, of an array's n closest elements to a certain value. By default, only returns the index/indices. Possible constraints: 'ceil', 'floor', None ("ceil" will return the closest element with a value greater than 'value', "floor" the opposite). Parameters ---------- array: 1d numpy array or list Array in which to check the closest element to value. value: float Value for which the algorithm searches for the n closest elements in the array. output: str, opt {'index','value','both' } Set what is returned constraint: str, opt {None, 'ceil', 'floor', 'ceil=', 'floor='} If not None, will check for the closest element larger (or equal) than value if set to 'ceil' ('ceil='), or closest element smaller (or equal) than value if set to 'floor' ('floor='). n: int, opt Number of elements to be returned, sorted by proximity to the values. Default: only the closest value is returned. Returns ------- [output='index']: index/indices of the closest n value(s) in the array; [output='value']: the closest n value(s) in the array, [output='both']: closest value(s) and index/-ices, respectively. """ if isinstance(array, np.ndarray): pass elif isinstance(array, list): array = np.array(array) else: raise ValueError("Input type for array should be np.ndarray or list.") if constraint is None: fm = np.absolute(array-value) idx = fm.argsort()[:n] elif 'floor' in constraint or 'ceil' in constraint: indices = np.arange(len(array), dtype=np.int32) if 'floor' in constraint: fm = -(array-value) else: fm = array-value if '=' in constraint: crop_indices = indices[np.where(fm >= 0)] fm = fm[np.where(fm >= 0)] else: crop_indices = indices[np.where(fm > 0)] fm = fm[np.where(fm > 0)] idx = fm.argsort()[:n] idx = crop_indices[idx] if len(idx) == 0: msg = "No indices match the constraint ({} w.r.t {:.2f})" print(msg.format(constraint, value)) raise ValueError("No indices match the constraint") else: raise ValueError("Constraint not recognised") if n == 1: idx = idx[0] if output == 'index': return idx elif output == 'value': return array[idx] else: return array[idx], idx