diff --git a/tests/pre_3_10/test_preproc_recentering.py b/tests/pre_3_10/test_preproc_recentering.py index f8e24983..25b95436 100644 --- a/tests/pre_3_10/test_preproc_recentering.py +++ b/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( [ diff --git a/vip_hci/fm/fakecomp.py b/vip_hci/fm/fakecomp.py index e04af5fc..e530810e 100644 --- a/vip_hci/fm/fakecomp.py +++ b/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])) diff --git a/vip_hci/metrics/detection.py b/vip_hci/metrics/detection.py index bc9e4b06..7e187295 100644 --- a/vip_hci/metrics/detection.py +++ b/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 ---------- diff --git a/vip_hci/metrics/stim.py b/vip_hci/metrics/stim.py index 776830b6..8d24308d 100644 --- a/vip_hci/metrics/stim.py +++ b/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 `_ - + """ __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 diff --git a/vip_hci/objects/postproc.py b/vip_hci/objects/postproc.py index 2d046c40..59dc755f 100644 --- a/vip_hci/objects/postproc.py +++ b/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) diff --git a/vip_hci/preproc/recentering.py b/vip_hci/preproc/recentering.py index bdc4d333..13c44d43 100644 --- a/vip_hci/preproc/recentering.py +++ b/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 diff --git a/vip_hci/preproc/skysubtraction.py b/vip_hci/preproc/skysubtraction.py index 4c62873d..a335663a 100644 --- a/vip_hci/preproc/skysubtraction.py +++ b/vip_hci/preproc/skysubtraction.py @@ -23,7 +23,7 @@ | *Astronomy & Astrophysics, Volume 679, p. 8* | `https://arxiv.org/abs/2308.16912 `_ - + """ __author__ = 'Sandrine Juillard' @@ -35,8 +35,8 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, full_output=False): - """ PCA-based sky subtraction as explained in [REN23]_. - (see also [GOM17]_ and [HUN18]_) + """PCA-based sky subtraction as explained in [REN23]_ (see also\ + [GOM17]_ and [HUN18]_). Parameters ---------- @@ -44,12 +44,12 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, 3d array of science frames. sky_cube : numpy ndarray 3d array of sky frames. - masks : tuple of two numpy ndarray or one signle numpy ndarray - Mask indicating the boat and anchor regions for the analysis. + masks : tuple of two numpy ndarray or one signle numpy ndarray + Mask indicating the boat and anchor regions for the analysis. If two masks are provided, they will be assigned to mask_anchor and - mask_boat in that order. - If only one mask is provided, it will be used as the anchor, and the - boat images will not be masked (i.e., full frames used). + mask_boat in that order. + If only one mask is provided, it will be used as the anchor, and the + boat images will not be masked (i.e., full frames used). ref_cube : numpy ndarray or None, opt Reference cube. ncomp : int, opt @@ -59,9 +59,10 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, derotated residual cube. Notes - ---------- - Masks can be created with the function `vip_hci.var.create_ringed_spider_mask` - or `get_annulus_segments` (see Usage Exemple below) + ----- + Masks can be created with the function + ``vip_hci.var.create_ringed_spider_mask`` or + ``vip_hci.var.get_annulus_segments`` (see Usage Example below). Returns @@ -77,29 +78,30 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, - anchor principal components, and - reconstructed cube. - Usage Exemple - ------- + Usage Example + ------------- You can create the masks using `get_annulus_segments` from `vip_hci.var`. - + .. code-block:: python from vip_hci.var import get_annulus_segments - - The function must be used as follows, where `ring_out`, `ring_in`, and + + The function must be used as follows, where `ring_out`, `ring_in`, and `coro` define the radius of the different annulus masks. They must have the same shape as a frame of the science cube. - + .. code-block:: python ones = np.ones(cube[0].shape) boat = get_annulus_segments(ones,coro,ring_out-coro, mode="mask")[0] - anchor = get_annulus_segments(ones,ring_in,ring_out-ring_in, mode="mask")[0] + anchor = get_annulus_segments(ones,ring_in,ring_out-ring_in, + mode="mask")[0] + - Masks should be provided as 'mask_rdi' argument when using PCA. - + .. code-block:: python res = pca(cube, angles, ref, mask_rdi=(boat, anchor), ncomp=2) @@ -114,7 +116,7 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, if sci_cube.shape[1] != sky_cube.shape[1] or sci_cube.shape[2] != \ sky_cube.shape[2]: raise TypeError('Science and Sky frames sizes do not match') - + if ref_cube is not None: if sci_cube.shape[1] != ref_cube.shape[1] or sci_cube.shape[2] != \ ref_cube.shape[2]: @@ -123,13 +125,12 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, # If only one mask is provided, the second mask is generated mask_anchor = masks mask_boat = np.ones(masks.shape) - elif len(masks)!=2: + elif len(masks) != 2: raise TypeError('Science and Reference frames sizes do not match') - else : + else: mask_anchor, mask_boat = masks - - ## -- Generate boat and anchor matrixes - + + # -- Generate boat and anchor matrices # Masking the sky cube with anchor sky_cube_masked = np.zeros_like(sky_cube) ind_masked = np.where(mask_anchor == 0) @@ -137,8 +138,8 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, masked_image = np.copy(sky_cube[i]) masked_image[ind_masked] = 0 sky_cube_masked[i] = masked_image - sky_anchor = sky_cube_masked.reshape(sky_cube.shape[0], - sky_cube.shape[1]*sky_cube.shape[2]) + sky_anchor = sky_cube_masked.reshape(sky_cube.shape[0], + sky_cube.shape[1]*sky_cube.shape[2]) # Masking the science cube with anchor sci_cube_anchor = np.zeros_like(sci_cube) @@ -147,7 +148,8 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, masked_image = np.copy(sci_cube[i]) masked_image[ind_masked] = 0 sci_cube_anchor[i] = masked_image - Msci_masked_anchor = prepare_matrix(sci_cube_anchor, scaling=None, verbose=False) + Msci_masked_anchor = prepare_matrix(sci_cube_anchor, scaling=None, + verbose=False) # Masking the science cube with boat sci_cube_boat = np.zeros_like(sci_cube) @@ -165,26 +167,28 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, masked_image = np.copy(sky_cube[i]) masked_image[ind_masked] = 0 sky_cube_boat[i] = masked_image - sky_boat = sky_cube_boat.reshape(sky_cube.shape[0], - sky_cube.shape[1]*sky_cube.shape[2]) + sky_boat = sky_cube_boat.reshape(sky_cube.shape[0], + sky_cube.shape[1]*sky_cube.shape[2]) - ## -- Generate eigenvectors of R(a)T R(a) + # -- Generate eigenvectors of R(a)T R(a) sky_kl = np.dot(sky_anchor, sky_anchor.T) Msky_kl = prepare_matrix(sky_kl, scaling=None, verbose=False) sky_pcs = svd_wrapper(Msky_kl, 'lapack', sky_kl.shape[0], False) sky_pcs_kl = sky_pcs.reshape(sky_kl.shape[0], sky_kl.shape[1]) - - ## -- Generate Kl and Dikl transform - - sky_pc_anchor = np.dot(sky_pcs_kl,sky_anchor) - sky_pcs_anchor_cube = sky_pc_anchor.reshape(sky_cube.shape[0], - sky_cube.shape[1], sky_cube.shape[2]) - - sky_pcs_boat_cube = np.dot(sky_pcs_kl,sky_boat).reshape(sky_cube.shape[0], - sky_cube.shape[1], sky_cube.shape[2]) - - ## -- Generate Kl projection to get coeff + + # -- Generate Kl and Dikl transform + + sky_pc_anchor = np.dot(sky_pcs_kl, sky_anchor) + sky_pcs_anchor_cube = sky_pc_anchor.reshape(sky_cube.shape[0], + sky_cube.shape[1], + sky_cube.shape[2]) + + sky_pcs_boat_cube = np.dot(sky_pcs_kl, sky_boat).reshape(sky_cube.shape[0], + sky_cube.shape[1], + sky_cube.shape[2]) + + # -- Generate Kl projection to get coeff transf_sci = np.zeros((sky_cube.shape[0], Msci_masked_anchor.shape[0])) for i in range(Msci_masked_anchor.shape[0]): @@ -192,23 +196,23 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, Msky_pcs_anchor = prepare_matrix(sky_pcs_anchor_cube, scaling=None, verbose=False) - + mat_inv = np.linalg.inv(np.dot(Msky_pcs_anchor, Msky_pcs_anchor.T)) transf_sci_scaled = np.dot(mat_inv, transf_sci) + # -- Subtraction Dikl projection using anchor coeff to sci cube - ## -- Subtraction Dikl projection using anchor coeff to sci cube - sci_cube_skysub = np.zeros_like(sci_cube) sky_opt = sci_cube.copy() for i in range(Msci_masked.shape[0]): - sky_opt[i] = np.array([np.sum( - transf_sci_scaled[j, i] * sky_pcs_boat_cube[j] for j in range(ncomp))]) + tmp_sky = [np.sum(transf_sci_scaled[j, i]*sky_pcs_boat_cube[j] + for j in range(ncomp))] + sky_opt[i] = np.array(tmp_sky) sci_cube_skysub[i] = sci_cube_boat[i] - sky_opt[i] - ## -- Processing the reference cube (if any) + # -- Processing the reference cube (if any) if ref_cube is not None: - + # Masking the ref cube with anchor ref_cube_anchor = np.zeros_like(sci_cube) ind_masked = np.where(mask_anchor == 0) @@ -216,7 +220,8 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, masked_image = np.copy(sci_cube[i]) masked_image[ind_masked] = 0 ref_cube_anchor[i] = masked_image - Mref_masked_anchor = prepare_matrix(ref_cube_anchor, scaling=None, verbose=False) + Mref_masked_anchor = prepare_matrix(ref_cube_anchor, scaling=None, + verbose=False) # Masking the ref cube with boat ref_cube_boat = np.zeros_like(sci_cube) @@ -235,8 +240,9 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, masks, ref_cube=None, ncomp=2, ref_cube_skysub = np.zeros_like(ref_cube) for i in range(Mref_masked.shape[0]): - sky_opt = np.array([np.sum(transf_ref_scaled[j, i] * sky_pcs_boat_cube[j] - for j in range(ncomp))]) + tmp_sky = [np.sum(transf_ref_scaled[j, i]*sky_pcs_boat_cube[j] + for j in range(ncomp))] + sky_opt = np.array(tmp_sky) ref_cube_skysub[i] = ref_cube_boat[i] - sky_opt if full_output: diff --git a/vip_hci/preproc/subsampling.py b/vip_hci/preproc/subsampling.py index 9b15c090..da66ccf8 100644 --- a/vip_hci/preproc/subsampling.py +++ b/vip_hci/preproc/subsampling.py @@ -21,9 +21,14 @@ def cube_collapse(cube, mode='median', n=50, w=None): - """ Collapses a cube into a frame (3D array -> 2D array) depending on the - parameter ``mode``. It's possible to perform a trimmed mean combination of - the frames based on description in [BRA13]_. + """Collapse a 3D or 4D cube into a 2D frame or 3D cube, respectively. + + The ``mode`` parameter determines how the collapse should be done. It is + possible to perform a trimmed mean combination of the frames, as in + [BRA13]_. In case of a 4D input cube, it is assumed to be an IFS dataset + with the zero-th axis being the spectral dimension, and the first axis the + temporal dimension. + Parameters ---------- @@ -47,7 +52,12 @@ def cube_collapse(cube, mode='median', n=50, w=None): Output array, cube combined. """ arr = cube - if arr.ndim != 3: + if arr.ndim == 3: + ax = 0 + elif arr.ndim == 4: + nch = arr.shape[0] + ax = 1 + else: raise TypeError('The input array is not a cube or 3d array.') if mode == 'wmean': @@ -60,27 +70,39 @@ def cube_collapse(cube, mode='median', n=50, w=None): w = np.array(w) if mode == 'mean': - frame = np.nanmean(arr, axis=0) + frame = np.nanmean(arr, axis=ax) elif mode == 'median': - frame = np.nanmedian(arr, axis=0) + frame = np.nanmedian(arr, axis=ax) elif mode == 'sum': - frame = np.nansum(arr, axis=0) + frame = np.nansum(arr, axis=ax) elif mode == 'max': - frame = np.nanmax(arr, axis=0) + frame = np.nanmax(arr, axis=ax) elif mode == 'trimmean': - N = arr.shape[0] + N = arr.shape[ax] k = (N - n)//2 if N % 2 != n % 2: n += 1 - frame = np.empty_like(arr[0]) - for index, _ in np.ndenumerate(arr[0]): - sort = np.sort(arr[:, index[0], index[1]]) - frame[index] = np.nanmean(sort[k:k+n]) + if ax == 0: + frame = np.empty_like(arr[0]) + for index, _ in np.ndenumerate(arr[0]): + sort = np.sort(arr[:, index[0], index[1]]) + frame[index] = np.nanmean(sort[k:k+n]) + else: + frame = np.empty_like(arr[:, 0]) + for j in range(nch): + for index, _ in np.ndenumerate(arr[:, 0]): + sort = np.sort(arr[j, :, index[0], index[1]]) + frame[j][index] = np.nanmean(sort[k:k+n]) elif mode == 'wmean': arr[np.where(np.isnan(arr))] = 0 # to avoid product with nan - frame = np.inner(w, np.moveaxis(arr, 0, -1)) + if ax == 0: + frame = np.inner(w, np.moveaxis(arr, 0, -1)) + else: + frame = np.empty_like(arr[:, 0]) + for j in range(nch): + frame[j] = np.inner(w, np.moveaxis(arr[j], 0, -1)) elif mode == 'absmean': - frame = np.mean(np.abs(arr), axis=0) + frame = np.nanmean(np.abs(arr), axis=ax) else: raise TypeError("mode not recognized") @@ -170,8 +192,9 @@ def cube_subsample(array, n, mode="mean", w=None, parallactic=None, def cube_subsample_trimmean(arr, n, m): - """Performs a trimmed mean combination every m frames in a cube. Based on - description in Brandt+ 2012. + """Perform a trimmed mean combination every m frames in a cube. + + Details in [BRA13]_. Parameters ---------- diff --git a/vip_hci/psfsub/loci.py b/vip_hci/psfsub/loci.py index 7f66d3c0..ce13c399 100644 --- a/vip_hci/psfsub/loci.py +++ b/vip_hci/psfsub/loci.py @@ -411,7 +411,8 @@ def _leastsq_adi( # indices indices = get_annulus_segments( - cube[0], inner_radius=inner_radius_ann, width=asize, nsegm=n_segments_ann + cube[0], inner_radius=inner_radius_ann, width=asize, + nsegm=n_segments_ann ) ind_opt = get_annulus_segments( cube[0], @@ -470,7 +471,8 @@ def _leastsq_adi( return frame_der_median -def _leastsq_patch(ayxyx, pa_thresholds, angles, metric, dist_threshold, solver, tol): +def _leastsq_patch(ayxyx, pa_thresholds, angles, metric, dist_threshold, solver, + tol): """Helper function for _leastsq_ann. Parameters @@ -516,7 +518,10 @@ def _leastsq_patch(ayxyx, pa_thresholds, angles, metric, dist_threshold, solver, A = values_opt[ind_ref] b = values_opt[i] if solver == "lstsq": - coef = sp.linalg.lstsq(A.T, b, cond=tol)[0] # SVD method + try: + coef = sp.linalg.lstsq(A.T, b, cond=tol)[0] # SVD method + except: + coef = sp.optimize.nnls(A.T, b)[0] # if SVD does not work elif solver == "nnls": coef = sp.optimize.nnls(A.T, b)[0] elif solver == "lsq": # TODO diff --git a/vip_hci/psfsub/pca_fullfr.py b/vip_hci/psfsub/pca_fullfr.py index f3278e6e..573d0921 100644 --- a/vip_hci/psfsub/pca_fullfr.py +++ b/vip_hci/psfsub/pca_fullfr.py @@ -280,12 +280,12 @@ def pca(*all_args: List, **all_kwargs: dict): If a tuple, it should contain the first and last channels where the mSDI residual channels will be collapsed (by default collapses all channels). mask_rdi: tuple of two numpy array or one signle 2d numpy array, opt - If provided, binary mask(s) will be used either in RDI mode or in - ADI+mSDI (2 steps) mode. They will be used as boat and anchor masks + If provided, binary mask(s) will be used either in RDI mode or in + ADI+mSDI (2 steps) mode. They will be used as boat and anchor masks following the procedure described in [REN23], which is useful to avoid - self-subtraction in the presence of a bright disc signal. If only one - mask is provided, the boat images will be unmasked - (i.e., full frames will be used). + self-subtraction in the presence of a bright disc signal. If only one + mask is provided, the boat images will be unmasked + (i.e., full frames will be used). check_memory : bool, optional If True, it checks that the input cube is smaller than the available system memory. @@ -849,11 +849,11 @@ def _adi_rdi_pca( recon = reshape_matrix(reconstructed, y, x) else: residuals_cube = residuals_result - + # A rotation threshold is applied else: if delta_rot is None or fwhm is None: - msg = "Delta_rot or fwhm parameters missing. Needed for the" + msg = "Delta_rot or fwhm parameters missing. Needed for" msg += "PA-based rejection of frames from the library" raise TypeError(msg) nfrslib = [] @@ -863,10 +863,10 @@ def _adi_rdi_pca( x1, y1 = source_xy ann_center = dist(yc, xc, y1, x1) pa_thr = _compute_pa_thresh(ann_center, fwhm, delta_rot) - + for frame in range(n): ind = _find_indices_adi(angle_list, frame, pa_thr) - + res_result = _project_subtract( cube, cube_ref, @@ -886,25 +886,27 @@ def _adi_rdi_pca( nfrslib.append(res_result[0]) residual_frame = res_result[1] recon_frame = res_result[2] - residuals_cube[frame] = residual_frame.reshape((y, x)) + residuals_cube[frame] = residual_frame.reshape((y, + x)) recon_cube[frame] = recon_frame.reshape((y, x)) else: nfrslib.append(res_result[0]) residual_frame = res_result[1] - residuals_cube[frame] = residual_frame.reshape((y, x)) - - # number of frames in library printed for each annular quadrant + residuals_cube[frame] = residual_frame.reshape((y, + x)) + + # number of frames in library printed for each ann. quadrant if verbose: descriptive_stats(nfrslib, verbose=verbose, label="Size LIB: ") - else: + else: residuals_result = cube_subtract_sky_pca( cube, cube_ref, mask_rdi, ncomp=ncomp, full_output=True ) residuals_cube = residuals_result[0] pcs = residuals_result[2] recon = residuals_result[-1] - + residuals_cube_ = cube_derotate( residuals_cube, angle_list,