From 3a9bbf662ad89b343ae3a5994ae8aae24544d885 Mon Sep 17 00:00:00 2001 From: Sand-jrd Date: Wed, 20 Mar 2024 18:38:53 +0100 Subject: [PATCH] Replacement of sci_cube_skysub with DIKL Replace the content of sci_cube_skysub with the DIKL procedure as explained in [REN23]. In pca_fullfr, the mask_rdi has been added to _adi_rdi_pca and is used along with sci_cube_skysub when it is provided. If not, it follows the procedure in place. The argument mask_rdi in docstring has been updated accordingly. --- vip_hci/preproc/skysubtraction.py | 201 +++++++++++++++++++++++------- vip_hci/psfsub/pca_fullfr.py | 147 ++++++++++++---------- 2 files changed, 235 insertions(+), 113 deletions(-) diff --git a/vip_hci/preproc/skysubtraction.py b/vip_hci/preproc/skysubtraction.py index dd4249d5..4c62873d 100644 --- a/vip_hci/preproc/skysubtraction.py +++ b/vip_hci/preproc/skysubtraction.py @@ -17,18 +17,26 @@ | `https://arxiv.org/abs/1706.10069 `_ +.. [REN23] + | Ren 2023 + | **Karhunen-Loève data imputation in high-contrast imaging** + | *Astronomy & Astrophysics, Volume 679, p. 8* + | `https://arxiv.org/abs/2308.16912 + `_ + """ -__author__ = 'Carlos Alberto Gomez Gonzalez' +__author__ = 'Sandrine Juillard' __all__ = ['cube_subtract_sky_pca'] import numpy as np from ..var import prepare_matrix -def cube_subtract_sky_pca(sci_cube, sky_cube, mask, ref_cube=None, ncomp=2, +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 [GOM17]_ and [HUN18]_. + """ PCA-based sky subtraction as explained in [REN23]_. + (see also [GOM17]_ and [HUN18]_) Parameters ---------- @@ -36,9 +44,12 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, mask, ref_cube=None, ncomp=2, 3d array of science frames. sky_cube : numpy ndarray 3d array of sky frames. - mask : numpy ndarray - Mask indicating the region for the analysis. Can be created with the - function vip_hci.var.create_ringed_spider_mask. + 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). ref_cube : numpy ndarray or None, opt Reference cube. ncomp : int, opt @@ -47,6 +58,12 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, mask, ref_cube=None, ncomp=2, Whether to also output pcs, reconstructed cube, residuals cube and 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) + + Returns ------- sci_cube_skysub : numpy ndarray @@ -56,10 +73,38 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, mask, ref_cube=None, ncomp=2, If full_output is set to True, returns (in the following order): - sky-subtracted science cube, - sky-subtracted reference cube (if any provided), - - principal components (non-masked), - - principal components (masked), and + - boat principal components, + - anchor principal components, and - reconstructed cube. + Usage Exemple + ------- + + 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 + `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] + + + 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) + + """ try: from ..psfsub.svd import svd_wrapper @@ -69,79 +114,139 @@ def cube_subtract_sky_pca(sci_cube, sky_cube, mask, 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]: raise TypeError('Science and Reference frames sizes do not match') + if type(masks) not in (list, tuple): + # If only one mask is provided, the second mask is generated + mask_anchor = masks + mask_boat = np.ones(masks.shape) + elif len(masks)!=2: + raise TypeError('Science and Reference frames sizes do not match') + else : + mask_anchor, mask_boat = masks + + ## -- Generate boat and anchor matrixes + + # Masking the sky cube with anchor + sky_cube_masked = np.zeros_like(sky_cube) + ind_masked = np.where(mask_anchor == 0) + for i in range(sky_cube.shape[0]): + 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]) - # Getting the EVs from the sky cube - Msky = prepare_matrix(sky_cube, scaling=None, verbose=False) - sky_pcs = svd_wrapper(Msky, 'lapack', sky_cube.shape[0], False) - sky_pcs_cube = sky_pcs.reshape(sky_cube.shape[0], sky_cube.shape[1], - sky_cube.shape[2]) - - # Masking the science cube - sci_cube_masked = np.zeros_like(sci_cube) - ind_masked = np.where(mask == 0) + # Masking the science cube with anchor + sci_cube_anchor = np.zeros_like(sci_cube) + ind_masked = np.where(mask_anchor == 0) for i in range(sci_cube.shape[0]): masked_image = np.copy(sci_cube[i]) masked_image[ind_masked] = 0 - sci_cube_masked[i] = masked_image - Msci_masked = prepare_matrix(sci_cube_masked, scaling=None, verbose=False) + sci_cube_anchor[i] = masked_image + Msci_masked_anchor = prepare_matrix(sci_cube_anchor, scaling=None, verbose=False) - # Masking the PCs learned from the skies - sky_pcs_cube_masked = np.zeros_like(sky_pcs_cube) - for i in range(sky_pcs_cube.shape[0]): - masked_image = np.copy(sky_pcs_cube[i]) + # Masking the science cube with boat + sci_cube_boat = np.zeros_like(sci_cube) + ind_masked = np.where(mask_boat == 0) + for i in range(sci_cube.shape[0]): + masked_image = np.copy(sci_cube[i]) masked_image[ind_masked] = 0 - sky_pcs_cube_masked[i] = masked_image - - # Project the masked frames onto the sky PCs to get the coefficients - transf_sci = np.zeros((sky_cube.shape[0], Msci_masked.shape[0])) - for i in range(Msci_masked.shape[0]): - transf_sci[:, i] = np.inner(sky_pcs, Msci_masked[i].T) - - Msky_pcs_masked = prepare_matrix(sky_pcs_cube_masked, scaling=None, + sci_cube_boat[i] = masked_image + Msci_masked = prepare_matrix(sci_cube_boat, scaling=None, verbose=False) + + # Masking the sky cube with boat + sky_cube_boat = np.zeros_like(sky_cube) + ind_masked = np.where(mask_boat == 0) + for i in range(sky_cube.shape[0]): + 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]) + + ## -- 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 + + transf_sci = np.zeros((sky_cube.shape[0], Msci_masked_anchor.shape[0])) + for i in range(Msci_masked_anchor.shape[0]): + transf_sci[:, i] = np.inner(sky_pc_anchor, Msci_masked_anchor[i].T) + + Msky_pcs_anchor = prepare_matrix(sky_pcs_anchor_cube, scaling=None, verbose=False) - mat_inv = np.linalg.inv(np.dot(Msky_pcs_masked, Msky_pcs_masked.T)) + + mat_inv = np.linalg.inv(np.dot(Msky_pcs_anchor, Msky_pcs_anchor.T)) transf_sci_scaled = np.dot(mat_inv, transf_sci) - # Obtaining the optimized sky and subtraction + + ## -- 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_cube[j] for j in range(ncomp))]) - sci_cube_skysub[i] = sci_cube[i] - sky_opt[i] + transf_sci_scaled[j, i] * sky_pcs_boat_cube[j] for j in range(ncomp))]) + 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: - ref_cube_masked = np.zeros_like(ref_cube) - for i in range(ref_cube.shape[0]): - masked_image = np.copy(ref_cube[i]) + + # Masking the ref cube with anchor + ref_cube_anchor = np.zeros_like(sci_cube) + ind_masked = np.where(mask_anchor == 0) + for i in range(sci_cube.shape[0]): + masked_image = np.copy(sci_cube[i]) masked_image[ind_masked] = 0 - ref_cube_masked[i] = masked_image - Mref_masked = prepare_matrix(ref_cube_masked, scaling=None, - verbose=False) + ref_cube_anchor[i] = masked_image + 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) + ind_masked = np.where(mask_boat == 0) + for i in range(sci_cube.shape[0]): + masked_image = np.copy(sci_cube[i]) + masked_image[ind_masked] = 0 + ref_cube_boat[i] = masked_image + Mref_masked = prepare_matrix(ref_cube_boat, scaling=None, verbose=False) + transf_ref = np.zeros((sky_cube.shape[0], Mref_masked.shape[0])) for i in range(Mref_masked.shape[0]): - transf_ref[:, i] = np.inner(sky_pcs, Mref_masked[i].T) + transf_ref[:, i] = np.inner(sky_pc_anchor, Mref_masked_anchor[i].T) transf_ref_scaled = np.dot(mat_inv, transf_ref) 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_cube[j] + sky_opt = np.array([np.sum(transf_ref_scaled[j, i] * sky_pcs_boat_cube[j] for j in range(ncomp))]) - ref_cube_skysub[i] = ref_cube[i] - sky_opt + ref_cube_skysub[i] = ref_cube_boat[i] - sky_opt if full_output: - return (sci_cube_skysub, ref_cube_skysub, sky_pcs_cube, - sky_pcs_cube_masked, sky_opt) + return (sci_cube_skysub, ref_cube_skysub, sky_pcs_anchor_cube, + sky_pcs_boat_cube, sky_opt) else: return sci_cube_skysub, ref_cube_skysub else: if full_output: - return (sci_cube_skysub, sky_pcs_cube, sky_pcs_cube_masked, sky_opt) + return (sci_cube_skysub, sky_pcs_anchor_cube, sky_pcs_boat_cube, + sky_opt) else: return sci_cube_skysub diff --git a/vip_hci/psfsub/pca_fullfr.py b/vip_hci/psfsub/pca_fullfr.py index cfe13c69..f3278e6e 100644 --- a/vip_hci/psfsub/pca_fullfr.py +++ b/vip_hci/psfsub/pca_fullfr.py @@ -37,6 +37,13 @@ | `https://arxiv.org/abs/0909.4061 `_ +.. [REN23] + | Ren 2023 + | **Karhunen-Loève data imputation in high-contrast imaging** + | *Astronomy & Astrophysics, Volume 679, p. 8* + | `https://arxiv.org/abs/2308.16912 + `_ + """ __author__ = "C.A. Gomez Gonzalez, V. Christiaens, T. Bédrine" @@ -272,11 +279,13 @@ def pca(*all_args: List, **all_kwargs: dict): ifs_collapse_range: str 'all' or tuple of 2 int 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: 2d numpy array, opt - If provided, this binary mask will be used either in RDI mode or in - ADI+mSDI (2 steps) mode. The projection coefficients for the principal - components will be found considering the area covered by the mask - (useful to avoid self-subtraction in presence of bright disc signal) + 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 + 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). check_memory : bool, optional If True, it checks that the input cube is smaller than the available system memory. @@ -761,6 +770,7 @@ def _adi_rdi_pca( nproc, full_output, weights=None, + mask_rdi=None, cube_sig=None, left_eigv=False, min_frames_pca=10, @@ -815,49 +825,9 @@ def _adi_rdi_pca( "Number of PCs too high (max PCs={}), using {} PCs " "instead.".format(n, ncomp) ) - - if source_xy is None: - residuals_result = _project_subtract( - cube, - cube_ref, - ncomp, - scaling, - mask_center_px, - svd_mode, - verbose, - full_output, - cube_sig=cube_sig, - left_eigv=left_eigv, - ) - if verbose: - timing(start_time) - if full_output: - residuals_cube = residuals_result[0] - reconstructed = residuals_result[1] - V = residuals_result[2] - pcs = reshape_matrix(V, y, x) if not left_eigv else V.T - 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 += "PA-based rejection of frames from the library" - raise TypeError(msg) - nfrslib = [] - residuals_cube = np.zeros_like(cube) - recon_cube = np.zeros_like(cube) - yc, xc = frame_center(cube[0], False) - 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( + if mask_rdi is None: + if source_xy is None: + residuals_result = _project_subtract( cube, cube_ref, ncomp, @@ -866,28 +836,75 @@ def _adi_rdi_pca( svd_mode, verbose, full_output, - ind, - frame, cube_sig=cube_sig, left_eigv=left_eigv, - min_frames_pca=min_frames_pca, ) + if verbose: + timing(start_time) if full_output: - nfrslib.append(res_result[0]) - residual_frame = res_result[1] - recon_frame = res_result[2] - residuals_cube[frame] = residual_frame.reshape((y, x)) - recon_cube[frame] = recon_frame.reshape((y, x)) + residuals_cube = residuals_result[0] + reconstructed = residuals_result[1] + V = residuals_result[2] + pcs = reshape_matrix(V, y, x) if not left_eigv else V.T + recon = reshape_matrix(reconstructed, 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 - if verbose: - descriptive_stats(nfrslib, verbose=verbose, - label="Size LIB: ") - + 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 += "PA-based rejection of frames from the library" + raise TypeError(msg) + nfrslib = [] + residuals_cube = np.zeros_like(cube) + recon_cube = np.zeros_like(cube) + yc, xc = frame_center(cube[0], False) + 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, + ncomp, + scaling, + mask_center_px, + svd_mode, + verbose, + full_output, + ind, + frame, + cube_sig=cube_sig, + left_eigv=left_eigv, + min_frames_pca=min_frames_pca, + ) + if full_output: + nfrslib.append(res_result[0]) + residual_frame = res_result[1] + recon_frame = res_result[2] + 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 + if verbose: + descriptive_stats(nfrslib, verbose=verbose, + label="Size LIB: ") + 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,