New options for NEGFC + minor bug fixes

tests/pre_3_10/test_metrics_contrcurve.py

@@ -8,7 +8,7 @@
 from tests.helpers import fixture
 from tests.helpers import np
 from vip_hci.config import VLT_NACO
-from vip_hci.fm.utils_negfc import cube_planet_free
+from vip_hci.fm import cube_planet_free
 from vip_hci.fm.utils_negfc import find_nearest
 from vip_hci.metrics import contrast_curve
 from vip_hci.preproc import frame_crop

vip_hci/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.6.2"
+__version__ = "1.6.4"
 
 from . import preproc
 from . import config

vip_hci/fm/fakecomp.py

@@ -6,7 +6,8 @@
            'normalize_psf',
            'cube_inject_companions',
            'generate_cube_copies_with_injections',
-           'frame_inject_companion']
+           'frame_inject_companion',
+           'cube_planet_free']
 
 from multiprocessing import cpu_count
 import numpy as np
@@ -810,3 +811,85 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
         msg += " ``vip_hci.preproc.cube_collapse`` first, or loop on the "
         msg += "temporal axis."
         raise ValueError(msg.format(array.shape[0]))
+
+
+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 3D cube, this
+        parameter must have a shape (n_pl,3) or (3,) -- the latter case assumes
+        a single planet in the data. 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