"""
Functions for combining the residuals of the PSF subtraction.
"""
from typing import Optional
import numpy as np
from typeguard import typechecked
from scipy.ndimage import rotate
[docs]
@typechecked
def combine_residuals(
method: str,
res_rot: np.ndarray,
residuals: Optional[np.ndarray] = None,
angles: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Wavelength wrapper for the combine_residual function. Produces an array with either 1
or number of wavelengths sized array.
Parameters
----------
method : str
Method used for combining the residuals ('mean', 'median', 'weighted', or 'clipped').
res_rot : np.ndarray
Derotated residuals of the PSF subtraction (3D).
residuals : np.ndarray, None
Non-derotated residuals of the PSF subtraction (3D). Only required for the noise-weighted
residuals.
angles : np.ndarray, None
Derotation angles (deg). Only required for the noise-weighted residuals.
Returns
-------
np.ndarray
Collapsed residuals (3D).
"""
if res_rot.ndim == 3:
output = _residuals(
method=method,
res_rot=np.asarray(res_rot),
residuals=residuals,
angles=angles,
)
if res_rot.ndim == 4:
output = np.zeros((res_rot.shape[0], res_rot.shape[2], res_rot.shape[3]))
for i in range(res_rot.shape[0]):
if residuals is None:
output[i,] = _residuals(
method=method,
res_rot=res_rot[i,],
residuals=residuals,
angles=angles,
)[0]
else:
output[i,] = _residuals(
method=method,
res_rot=res_rot[i,],
residuals=residuals[i,],
angles=angles,
)[0]
return output
@typechecked
def _residuals(
method: str,
res_rot: np.ndarray,
residuals: Optional[np.ndarray] = None,
angles: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Function for combining the derotated residuals of the PSF subtraction.
Parameters
----------
method : str
Method used for combining the residuals ('mean', 'median', 'weighted', or 'clipped').
res_rot : np.ndarray
Derotated residuals of the PSF subtraction (3D).
residuals : np.ndarray, None
Non-derotated residuals of the PSF subtraction (3D). Only required for the noise-weighted
residuals.
angles : np.ndarray, None
Derotation angles (deg). Only required for the noise-weighted residuals.
Returns
-------
np.ndarray
Combined residuals (3D).
"""
if method == "mean":
stack = np.mean(res_rot, axis=0)
elif method == "median":
stack = np.median(res_rot, axis=0)
elif method == "weighted":
tmp_res_var = np.var(residuals, axis=0)
res_repeat = np.repeat(
tmp_res_var[np.newaxis, :, :], repeats=residuals.shape[0], axis=0
)
res_var = np.zeros(res_repeat.shape)
for j, angle in enumerate(angles):
# scipy.ndimage.rotate rotates in clockwise direction for positive angles
res_var[j,] = rotate(input=res_repeat[j,], angle=angle, reshape=False)
weight1 = np.divide(
res_rot,
res_var,
out=np.zeros_like(res_var),
where=(np.abs(res_var) > 1e-100) & (res_var != np.nan),
)
weight2 = np.divide(
1.0,
res_var,
out=np.zeros_like(res_var),
where=(np.abs(res_var) > 1e-100) & (res_var != np.nan),
)
sum1 = np.sum(weight1, axis=0)
sum2 = np.sum(weight2, axis=0)
stack = np.divide(
sum1,
sum2,
out=np.zeros_like(sum2),
where=(np.abs(sum2) > 1e-100) & (sum2 != np.nan),
)
elif method == "clipped":
stack = np.zeros(res_rot.shape[-2:])
for i in range(stack.shape[0]):
for j in range(stack.shape[1]):
pix_line = res_rot[:, i, j]
if np.var(pix_line) > 0.0:
no_mean = pix_line - np.mean(pix_line)
part1 = no_mean.compress(
(no_mean < 3.0 * np.sqrt(np.var(no_mean))).flat
)
part2 = part1.compress(
(part1 > -3.0 * np.sqrt(np.var(no_mean))).flat
)
stack[i, j] = np.mean(pix_line) + np.mean(part2)
return stack[np.newaxis, ...]
