Adapted cube_collapse function to also accept an input 4D cube

tests/pre_3_10/test_preproc_recentering.py

@@ -604,7 +604,7 @@ def test_satspots(debug=False):
     cube, spotcoords = create_cube_with_satspots(n_frames=n_frames)
 
     # ===== shift
-    shift_magnitude = 1
+    shift_magnitude = 0.8
     randax = seed.uniform(-shift_magnitude, 0, size=n_frames)
     randay = seed.uniform(0, shift_magnitude, size=n_frames)
 
@@ -645,7 +645,8 @@ def test_radon(debug=False):
     fit_res = vip.var.fit_2dgaussian(med_fr, crop=True, cropsize=13,
                                      cent=(144, 147), debug=False,
                                      full_output=True)
-    med_y, med_x = float(fit_res['centroid_y']), float(fit_res['centroid_x'])
+    med_y = float(fit_res['centroid_y'][0])
+    med_x = float(fit_res['centroid_x'][0])
     # remove BKG star
     fit_flux = np.array(
         [

vip_hci/fm/fakecomp.py

@@ -1,7 +1,5 @@
 #! /usr/bin/env python
-"""
-Module with fake companion injection functions.
-"""
+"""Module with fake companion injection functions."""
 
 __author__ = 'Carlos Alberto Gomez Gonzalez, Valentin Christiaens'
 __all__ = ['collapse_psf_cube',
@@ -38,7 +36,8 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists,
                            interpolation='lanczos4', transmission=None,
                            radial_gradient=False, full_output=False,
                            verbose=False, nproc=1):
-    """ Injects fake companions in branches, at given radial distances.
+    """Inject fake companions in branches and given radial distances in an ADI\
+    cube.
 
     Parameters
     ----------
@@ -165,14 +164,6 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc,
                                                  -(ang*180/np.pi-angle_list[0]),
                                                  imlib=imlib_rot,
                                                  interpolation=interpolation)
-
-                        # shift_y = rad * np.sin(ang - np.deg2rad(angle_list[0]))
-                        # shift_x = rad * np.cos(ang - np.deg2rad(angle_list[0]))
-                        # dsy = shift_y-int(shift_y)
-                        # dsx = shift_x-int(shift_x)
-                        # fc_fr_ang = frame_shift(psf_trans, dsy, dsx, imlib_sh,
-                        #                         interpolation,
-                        #                         border_mode='constant')
                     else:
                         fc_fr_rad = interp_trans(rad)*fc_fr
                 if nproc == 1:
@@ -247,7 +238,7 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc,
         if transmission.ndim != 2:
             raise ValueError("transmission should be a 2D ndarray")
         elif t_nz != 2 and t_nz != 1+array.shape[0]:
-            msg = "transmission dimensions should be either (2,N) or (n_wave+1, N)"
+            msg = "transmission dimensions should be (2,N) or (n_wave+1, N)"
             raise ValueError(msg)
         # if transmission doesn't have right format for interpolation, adapt it
         diag = np.sqrt(2)*array.shape[-1]
@@ -324,10 +315,7 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc,
 def _frame_shift_fcp(fc_fr_rad, array, rad, ang, derot_ang, flevel, size_fc,
                      imlib_sh, imlib_rot, interpolation, transmission,
                      radial_gradient):
-    """
-    Specific cube shift algorithm for injection of fake companions
-    """
-
+    """Specific cube shift algorithm to inject fake companions."""
     ceny, cenx = frame_center(array)
     sizey = array.shape[-2]
     sizex = array.shape[-1]
@@ -480,9 +468,9 @@ def generate_cube_copies_with_injections(array, psf_template, angle_list, plsc,
 def frame_inject_companion(array, array_fc, pos_y, pos_x, flux,
                            imlib='vip-fft', interpolation='lanczos4'):
     """
-    Injects a fake companion in a single frame (it could be a single
-    multi-wavelength frame) at given coordinates, or in a cube (at the same
-    coordinates, flux and with same fake companion image throughout the cube).
+    Inject a fake companion in a single frame (can be a single multi-wavelength\
+    frame) at given coordinates, or in a cube (at the same coordinates, flux\
+    and with same fake companion image throughout the cube).
 
     Parameters
     ----------
@@ -539,7 +527,7 @@ def frame_inject_companion(array, array_fc, pos_y, pos_x, flux,
 
 
 def collapse_psf_cube(array, size, fwhm=4, verbose=True, collapse='mean'):
-    """ Creates a 2d PSF template from a cube of non-saturated off-axis frames
+    """Create a 2d PSF template from a cube of non-saturated off-axis frames\
     of the star by taking the mean and normalizing the PSF flux.
 
     Parameters
@@ -560,6 +548,7 @@ def collapse_psf_cube(array, size, fwhm=4, verbose=True, collapse='mean'):
     -------
     psf_normd : numpy ndarray
         Normalized PSF.
+
     """
     if array.ndim != 3 and array.ndim != 4:
         raise TypeError('Array is not a cube, 3d or 4d array.')
@@ -584,9 +573,11 @@ def normalize_psf(array, fwhm='fit', size=None, threshold=None, mask_core=None,
                   model='gauss', imlib='vip-fft', interpolation='lanczos4',
                   force_odd=True, correct_outliers=True, full_output=False,
                   verbose=True, debug=False):
-    """ Normalizes a PSF (2d or 3d array), to have the flux in a 1xFWHM
-    aperture equal to one. It also allows to crop the array and center the PSF
-    at the center of the array(s).
+    """Normalize a PSF (2d or 3d array), to have the flux in a 1xFWHM aperture\
+    equal to one.
+
+    It also allows cropping of the array. Automatic recentering of the PSF is
+    done internally - as long as it is already roughly centered within ~2px.
 
     Parameters
     ----------
@@ -599,8 +590,8 @@ def normalize_psf(array, fwhm='fit', size=None, threshold=None, mask_core=None,
         estimate the FWHM in 2D or 3D PSF arrays.
     size : int or None, optional
         If int it will correspond to the size of the centered sub-image to be
-        cropped form the PSF array. The PSF is assumed to be roughly centered wrt
-        the array.
+        cropped form the PSF array. The PSF is assumed to be roughly centered
+        with respect to the array.
     threshold : None or float, optional
         Sets to zero values smaller than threshold (in the normalized image).
         This can be used to only leave the core of the PSF.
@@ -646,7 +637,7 @@ def normalize_psf(array, fwhm='fit', size=None, threshold=None, mask_core=None,
 
     """
     def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
-        """ 2d case """
+        """Normalize PSF in the 2d case."""
         # we check if the psf is centered and fix it if needed
         cy, cx = frame_center(psf, verbose=False)
         xcom, ycom = cen_com(psf)
@@ -766,8 +757,9 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
             if fwhm != 'fit':
                 fwhm = [fwhm] * array.shape[0]
             else:
-                fits_vect = [fit_2d(array[i], full_output=True, debug=debug) for i
-                             in range(n)]
+                fits_vect = [fit_2d(array[i],
+                                    full_output=True,
+                                    debug=debug) for i in range(n)]
                 if model == 'gauss':
                     fwhmx = [fits_vect[i]['fwhm_x'] for i in range(n)]
                     fwhmy = [fits_vect[i]['fwhm_y'] for i in range(n)]
@@ -792,8 +784,8 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
                                 fwhm[f] = np.nanmean(np.array([fwhm[f-1],
                                                                fwhm[f+1]]))
                             elif np.isnan(fwhm[f]):
-                                msg = "2D fit failed for first or last channel. "
-                                msg += "Try other parameters?"
+                                msg = "2D fit failed for first or last channel."
+                                msg += " Try other parameters?"
                                 raise ValueError(msg)
         elif len(fwhm) != array.shape[0]:
             msg = "If fwhm is a list/1darray it should have a length of {}"
@@ -817,3 +809,9 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
             return array_out, fwhm_flux, fwhm
         else:
             return array_out
+
+    else:
+        msg = "Input psf should be 2D or 3D. If of higher dimension, either use"
+        msg += " ``vip_hci.preproc.cube_collapse`` first, or loop on the "
+        msg += "temporal axis."
+        raise ValueError(msg.format(array.shape[0]))

vip_hci/metrics/detection.py

@@ -26,11 +26,13 @@
 def detection(array, fwhm=4, psf=None, mode='lpeaks', bkg_sigma=5,
               matched_filter=False, mask=True, snr_thresh=5, nproc=1, plot=True,
               debug=False, full_output=False, verbose=True, **kwargs):
-    """ Finds blobs in a 2d array. The algorithm is designed for automatically
-    finding planets in post-processed high contrast final frames. Blob can be
-    defined as a region of an image in which some properties are constant or
-    vary within a prescribed range of values. See ``Notes`` below to read about
-    the algorithm details.
+    """Automatically find point-like sources in a 2d array.
+
+    The algorithm is designed for automatically finding planets in
+    post-processed high contrast final frames. Blob can be defined as a region
+    of an image in which some properties are constant or vary within a
+    prescribed range of values. See ``Notes`` below to read about the algorithm
+    details.
 
     Parameters
     ----------

vip_hci/metrics/stim.py

@@ -1,27 +1,28 @@
 #! /usr/bin/env python
-
 """
 Implementation of the STIM map from [PAI19]
 
 .. [PAI19]
    | Pairet et al. 2019
-   | **STIM map: detection map for exoplanets imaging beyond asymptotic Gaussian residual speckle noise**
+   | **STIM map: detection map for exoplanets imaging beyond asymptotic Gaussian
+   residual speckle noise**
    | *MNRAS, 487, 2262*
    | `http://doi.org/10.1093/mnras/stz1350
      <http://doi.org/10.1093/mnras/stz1350>`_
-   
+
 """
 __author__ = 'Benoit Pairet'
 __all__ = ['stim_map',
-           'inverse_stim_map']
+           'inverse_stim_map',
+           'normalized_stim_map']
 
 import numpy as np
 from ..preproc import cube_derotate
-from ..var import get_circle
+from ..var import get_circle, mask_circle
 
 
 def stim_map(cube_der):
-    """ Computes the STIM detection map as in [PAI19]_.
+    """Compute the STIM detection map as in [PAI19]_.
 
     Parameters
     ----------
@@ -33,8 +34,8 @@ def stim_map(cube_der):
     -------
     detection_map : 2d ndarray
         STIM detection map.
-    """
 
+    """
     t, n, _ = cube_der.shape
     mu = np.mean(cube_der, axis=0)
     sigma = np.sqrt(np.var(cube_der, axis=0))
@@ -44,7 +45,7 @@ def stim_map(cube_der):
 
 
 def inverse_stim_map(cube, angle_list, **rot_options):
-    """ Computes the inverse STIM detection map as in [PAI19]_.
+    """Compute the inverse STIM detection map as in [PAI19]_.
 
     Parameters
     ----------
@@ -63,9 +64,55 @@ def inverse_stim_map(cube, angle_list, **rot_options):
     -------
     inv_stim_map : 2d ndarray
         Inverse STIM detection map.
-    """
 
+    """
     t, n, _ = cube.shape
     cube_inv_der = cube_derotate(cube, -angle_list, **rot_options)
     inv_stim_map = stim_map(cube_inv_der)
     return inv_stim_map
+
+
+def normalized_stim_map(cube, angle_list, mask=None, **rot_options):
+    """Compute the normalized STIM detection map as in [PAI19]_.
+
+    Parameters
+    ----------
+    cube : 3d numpy ndarray
+        Non de-rotated residuals from reduction algorithm, eg.
+        ``residuals_cube`` output from ``vip_hci.psfsub.pca``.
+    angle_list : numpy ndarray, 1d
+        Corresponding parallactic angle for each frame.
+    mask : int, float, numpy ndarray 2d or None
+        Mask informing where the maximum value in the inverse STIM map should
+        be calculated. If an integer or float, a circular mask with that radius
+        masking the central part of the image will be used. If a 2D array, it
+        should be a binary mask (ones in the areas that can be used).
+    rot_options: dictionary, optional
+        Dictionary with optional keyword values for "nproc", "imlib",
+        "interpolation, "border_mode", "mask_val",  "edge_blend",
+        "interp_zeros", "ker" (see documentation of
+        ``vip_hci.preproc.frame_rotate``)
+
+    Returns
+    -------
+    normalized STIM map : 2d ndarray
+        STIM detection map.
+
+    """
+    inv_map = inverse_stim_map(cube, angle_list, **rot_options)
+
+    if mask is not None:
+        if np.isscalar(mask):
+            inv_map = mask_circle(inv_map, mask)
+        else:
+            inv_map *= mask
+
+    max_inv = np.nanmax(inv_map)
+    if max_inv <= 0:
+        msg = "The normalization value is found to be {}".format(max_inv)
+        raise ValueError(msg)
+
+    cube_der = cube_derotate(cube, angle_list, **rot_options)
+    det_map = stim_map(cube_der)
+
+    return det_map/max_inv

vip_hci/objects/postproc.py

@@ -474,7 +474,7 @@ def get_params_from_results(self, session_id: int) -> None:
         print_algo_params(res[session_id].parameters)
 
     # TODO : identify the problem around the element `_repr_html_`
-    def _get_calculations(self) -> dict:
+    def _get_calculations(self, debug=False) -> dict:
         """
         Get a list of all attributes which are *calculated*.
 
@@ -494,33 +494,36 @@ def _get_calculations(self) -> dict:
         for element in dir(self):
             # BLACKMAGIC : _repr_html_ must be skipped
             """
-            `_repr_html_` is an element of the directory of the PostProc object which
-            causes the search of calculated attributes to overflow, looping indefinitely
-            and never reaching the actual elements containing those said attributes.
-            It will be skipped until the issue has been properly identified and fixed.
-            You can uncomment the block below to observe how the directory loops after
-            reaching that element - acknowledging you are not skipping it.
+            `_repr_html_` is an element of the directory of the PostProc object
+            which causes the search of calculated attributes to overflow,
+            looping indefinitely and never reaching the actual elements
+            containing those said attributes. It will be skipped until the issue
+            has been properly identified and fixed. You can set debug=True to
+            observe how the directory loops after reaching that element -
+            acknowledging you are not skipping it.
             """
             if element not in PROBLEMATIC_ATTRIBUTE_NAMES:
                 try:
-                    # print(
-                    #     "directory element : ",
-                    #     element,
-                    #     ", calculations list : ",
-                    #     calculations,
-                    # )
+                    if debug:
+                        print(
+                            "directory element : ",
+                            element,
+                            ", calculations list : ",
+                            calculations,
+                        )
                     for k in getattr(getattr(self, element), "_calculates"):
                         calculations[k] = element
                 except AttributeError:
                     pass
             # below can be commented after debug
             else:
-                print(
-                    "directory element SKIPPED: ",
-                    element,
-                    ", calculations list : ",
-                    calculations,
-                )
+                if debug:
+                    print(
+                        "directory element SKIPPED: ",
+                        element,
+                        ", calculations list : ",
+                        calculations,
+                    )
 
         return calculations
 
@@ -550,9 +553,9 @@ def __getattr__(self, attr: str) -> NoReturn:
         """
         calculations = self._get_calculations()
         if attr in calculations:
-            raise AttributeError(
-                f"The {attr} was not calculated yet. Call {calculations[attr]} first."
-            )
+            msg = f"The {attr} was not calculated yet. "
+            msg += f"Call {calculations[attr]} first."
+            raise AttributeError(msg)
         # this raises a regular AttributeError:
         return self.__getattribute__(attr)
 

vip_hci/preproc/recentering.py

@@ -844,13 +844,13 @@ def _center_radon(array, cropsize=None, hsize=1., step=0.1,
                     raise ValueError(msg)
                 sinogram = radon(frame, theta=theta, circle=True)
                 plot_frames((frame, sinogram))
-                print(np.sum(np.abs(sinogram[int(cent), :])))
+                # print(np.sum(np.abs(sinogram[int(cent), :])))
             else:
                 theta = np.linspace(start=0, stop=360, num=int(cent*2),
                                     endpoint=False)
                 sinogram = radon(frame, theta=theta, circle=True)
                 plot_frames((frame, sinogram))
-                print(np.sum(np.abs(sinogram[int(cent), :])))
+                # print(np.sum(np.abs(sinogram[int(cent), :])))
 
         if nproc is None:
             nproc = cpu_count() // 2    # Hyper-threading doubles the # of cores