Skip to content

Commit

Permalink
PEP8 formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@VChristiaens
VChristiaens committed Mar 20, 2024
1 parent 8e67ef3 commit 472c205
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 68 deletions.
112 changes: 59 additions & 53 deletions vip_hci/preproc/skysubtraction.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
| *Astronomy & Astrophysics, Volume 679, p. 8*
| `https://arxiv.org/abs/2308.16912
<https://arxiv.org/abs/2308.16912>`_
"""

__author__ = 'Sandrine Juillard'
Expand All @@ -35,21 +35,21 @@

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
----------
sci_cube : numpy ndarray
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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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]:
Expand All @@ -123,22 +125,21 @@ 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)
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])
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)
Expand All @@ -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)
Expand All @@ -165,58 +167,61 @@ 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]):
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_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)
for i in range(sci_cube.shape[0]):
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)
Expand All @@ -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:
Expand Down
32 changes: 17 additions & 15 deletions vip_hci/psfsub/pca_fullfr.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 = []
Expand All @@ -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,
Expand All @@ -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,
Expand Down

0 comments on commit 472c205

Please sign in to comment.