"""
Module to support frame combination
"""
import warnings
import numpy as np
import corgidrp.data as data
from pyklip.klip import rotate
import corgidrp
[docs]
def combine_images(data_subset, err_subset, dq_subset, collapse, num_frames_scaling, other_hdus=None):
"""
Combines several images together
Args:
data_subset (np.array): 3-D array of N 2-D images
err_subset (np.array): 4-D array of N 3-D error maps
dq_subset (np.array): 3-D array of N 2-D DQ maps
collapse (str): "mean" or "median".
num_frames_scaling (bool): Multiply by number of frames in sequence in order to ~conserve photons
other_hdus (list of HDULists, optional): list of other HDULists to be combined in the same way
Returns:
tuple: combined images
np.array: 2-D array of combined images
np.array: 3-D array of combined error map
np.array: 2-D array of combined DQ maps
list of np.array: list of the combined data of other HUDs
"""
tot_frames = data_subset.shape[0]
# mask bad pixels
bad = np.where(dq_subset > 0)
data_subset[bad] = np.nan
err_subset[bad[0],:,bad[1],bad[2]] = np.nan
# track the number of good values that go into the combination
n_samples = np.ones(data_subset.shape)
n_samples[bad] = 0
n_samples = np.sum(n_samples, axis=0)
if collapse.lower() == "mean":
with warnings.catch_warnings():
# prevent RuntimeWarning: Mean of empty slice
warnings.filterwarnings('ignore', category=RuntimeWarning)
data_collapse = np.nanmean(data_subset, axis=0)
err_collapse = np.sqrt(np.nanmean(err_subset**2, axis=0)) /np.sqrt(n_samples) # correct assuming standard error propagation
elif collapse.lower() == "median":
with warnings.catch_warnings():
# prevent RuntimeWarning: Mean of empty slice
warnings.filterwarnings('ignore', category=RuntimeWarning)
data_collapse = np.nanmedian(data_subset, axis=0)
err_collapse = np.sqrt(np.nanmean(err_subset**2, axis=0)) /np.sqrt(n_samples) * np.sqrt(np.pi/2) # inflate median error
if num_frames_scaling:
# scale up by the number of frames
data_collapse *= tot_frames
err_collapse *= tot_frames
# dq collpase: keep all flags on
dq_collapse = np.bitwise_or.reduce(dq_subset, axis=0)
# except for those pixels that have been replaced with good values
dq_collapse[np.where((dq_collapse > 0) & (~np.isnan(data_collapse)))] = 0
# other hdus
if other_hdus is not None:
combined_hdus = [[] for _ in range(len(other_hdus[0]))]
# iterate over each hdulist and append the data
for hdul in other_hdus:
for i, hdu in enumerate(hdul):
combined_hdus[i].append(np.copy(hdu.data))
# now combine each hdu data
for i in range(len(combined_hdus)):
# TODO: not implemented how to take means of anything beyond np.arrays (e.g., np.recarray)
try:
if collapse.lower() == "mean":
combined_hdus[i] = np.nanmean(np.array(combined_hdus[i]), axis=0)
elif collapse.lower() == "median":
combined_hdus[i] = np.nanmedian(np.array(combined_hdus[i]), axis=0)
# nothing here makes sense to scale by number of frames
# if num_frames_scaling:
# combined_hdus[i] *= tot_frames
except:
combined_hdus[i] = combined_hdus[i][0] # just take the first one if cannot combine
else:
combined_hdus = None
return data_collapse, err_collapse, dq_collapse, combined_hdus
[docs]
def combine_subexposures(input_dataset, num_frames_per_group=None, collapse="mean", num_frames_scaling=True,
combine_other_hdus=False):
"""
Combines a sequence of exposures assuming a constant nubmer of frames per group.
The length of the dataset must be divisible by the number of frames per group.
The combination is done with either the mean or median, but the collapsed image can be scaled
in order to ~conserve the total number of photons in the input dataset (this essentially turns a
median into a sum)
Args:
input_dataset (corgidrp.data.Dataset): input data.
num_frames_per_group (int): number of subexposures per group. If None, combines all images together
collapse (str): "mean" or "median". (default: mean)
num_frames_scaling (bool): Multiply by number of frames in sequence in order to ~conserve photons (default: True)
combine_other_hdus (bool): Whether to combine other HDUs in the same way as the main data, err, DQ.
Otherwise, uses the HDUs from the first frame in a subset (default: False)
Returns:
corgidrp.data.Dataset: dataset after combination of every "num_frames_per_group" frames together
"""
if num_frames_per_group is None:
num_frames_per_group = len(input_dataset)
if len(input_dataset) % num_frames_per_group != 0:
raise ValueError("Input dataset of length {0} cannot be grouped in sets of {1}".format(len(input_dataset), num_frames_per_group))
if collapse.lower() not in ["mean", "median"]:
raise ValueError("combine_subexposures can only collapse with mean or median")
num_groups = len(input_dataset) // num_frames_per_group
new_dataset = []
for i in range(num_groups):
data_subset = np.copy(input_dataset.all_data[num_frames_per_group*i:num_frames_per_group*(i+1)])
err_subset = np.copy(input_dataset.all_err[num_frames_per_group*i:num_frames_per_group*(i+1)])
dq_subset = np.copy(input_dataset.all_dq[num_frames_per_group*i:num_frames_per_group*(i+1)])
if combine_other_hdus:
other_hdus = [input_dataset[j].hdu_list for j in range(num_frames_per_group*i, num_frames_per_group*(i+1))]
else:
other_hdus = None
data_collapse, err_collapse, dq_collapse, combined_hdus = combine_images(data_subset, err_subset, dq_subset, collapse=collapse,
num_frames_scaling=num_frames_scaling, other_hdus=other_hdus)
# grab the headers from the first frame in this sub sequence
pri_hdr = input_dataset[num_frames_per_group*i].pri_hdr.copy()
ext_hdr = input_dataset[num_frames_per_group*i].ext_hdr.copy()
ext_hdr["NUM_FR"] = num_frames_per_group
err_hdr = input_dataset[num_frames_per_group*i].err_hdr.copy()
dq_hdr = input_dataset[num_frames_per_group*i].dq_hdr.copy()
hdulist = input_dataset[num_frames_per_group*i].hdu_list.copy()
# update other hdus if needed
if combine_other_hdus:
for j, hdu in enumerate(hdulist):
hdu.data = combined_hdus[j]
hdu.header["NUM_FR"] = num_frames_per_group
hdu.header['HISTORY'] = "Combined {0} frames by {1}".format(num_frames_per_group, collapse)
new_image = data.Image(data_collapse, pri_hdr=pri_hdr, ext_hdr=ext_hdr, err=err_collapse, dq=dq_collapse, err_hdr=err_hdr,
dq_hdr=dq_hdr, input_hdulist=hdulist)
# always take the last filename in the group for the combined frame
last_idx_in_group = num_frames_per_group*(i+1) - 1
new_image.filename = input_dataset[last_idx_in_group].filename
new_image._record_parent_filenames(input_dataset[num_frames_per_group*i:num_frames_per_group*(i+1)])
new_dataset.append(new_image)
new_dataset = data.Dataset(new_dataset)
drpnfile = new_dataset[0].ext_hdr['DRPNFILE']
# Here we change header keywords only for the combined non-coronagraphic imaging datasets
if (input_dataset[0].ext_hdr['DPAMNAME'] == 'IMAGING' and input_dataset[0].ext_hdr['LSAMNAME'] == 'OPEN') and input_dataset[0].ext_hdr['DATALVL'] == 'L3':
# average/delete header keywords as L4 involves combination of multiple frames
pri_hdr_comb, ext_hdr_comb, _, _ = corgidrp.check.merge_headers(input_dataset,
last_frame_keywords=['VISITID', 'MJDEND', 'SCTEND'],
first_frame_keywords=['MJDSRT','SCTSRT','CD1_1', 'CD1_2', 'CD2_1', 'CD2_2', 'CRPIX1', 'CRPIX2','NORTHANG'],
deleted_keywords=['CDELT1','CDELT2','FILE0'] + corgidrp.check.deleted_keywords_default, #we re-add FILE0 below
invalid_keywords=[
#Primary header keywords
'FILETIME', 'PA_V3', 'PA_APER','SVB_1', 'SVB_2', 'SVB_3',
'ROLL', 'PITCH', 'YAW', 'WBJ_1', 'WBJ_2', 'WBJ_3',
#Extension header keywords
'DATETIME', 'FTIMEUTC','DATATYPE'],
averaged_keywords=['EXCAMT','NOVEREXP','PROXET',
'FCMPOS','FSMSG1', 'FSMSG2', 'FSMSG3', 'FSMX', 'FSMY',
'SB_FP_DX', 'SB_FP_DY', 'SB_FS_DX', 'SB_FS_DY',
'Z2AVG', 'Z3AVG', 'Z4AVG', 'Z5AVG', 'Z6AVG', 'Z7AVG', 'Z8AVG', 'Z9AVG',
'Z10AVG', 'Z11AVG', 'Z12AVG', 'Z13AVG', 'Z14AVG',
'Z2RES', 'Z3RES', 'Z4RES', 'Z5RES', 'Z6RES', 'Z7RES', 'Z8RES', 'Z9RES',
'Z10RES', 'Z11RES',
'Z2VAR', 'Z3VAR'])
# incorporate modified headers in L4 dataset
for img in new_dataset:
img.pri_hdr = pri_hdr_comb
img.ext_hdr = ext_hdr_comb
img.ext_hdr['NUM_FR'] = num_frames_per_group
img.ext_hdr['DRPNFILE'] = drpnfile
img._record_parent_filenames(input_dataset)
new_dataset.update_after_processing_step("Combine_subexposures: combined every {0} frames by {1}".format(num_frames_per_group, collapse))
return new_dataset
def _append_unique(default_keywords, input_keywords):
"""
Append input keywords after default keywords while preserving order and
removing duplicates.
"""
combined_keywords = []
for keyword in default_keywords + list(input_keywords):
if keyword not in combined_keywords:
combined_keywords.append(keyword)
return combined_keywords
[docs]
def combine_frames_per_visit(input_dataset, collapse="mean", num_frames_per_group=None, num_frames_scaling=False,
combine_other_hdus=False, max_combined=100, pri_split_keywords=['VISITID'],
ext_split_keywords=['DPAMNAME']):
"""
Combines frames in the input dataset by calling combine_subexposures on each
unique split of VISITID and polarization state (DPAMNAME). The combination is done
with either the mean or median, but the collapsed image can be scaled in order to ~conserve
the total number of photons in the input dataset (this essentially turns a median into a sum).
Args:
input_dataset (corgidrp.data.Dataset): a dataset of Images (L3-level)
collapse (str): "mean" or "median". (default: mean)
num_frames_per_group (int or None): number of consecutive frames to combine
within each VISITID/DPAMNAME split subset.
If None, combine all frames in each subset
into one output frame. If the requested
grouping would create more than ``max_combined`` output
frames in a split subset, it is overridden
with the smallest group size that keeps the
per-subset output count at or below
``max_combined``; trailing remainder frames
are left out.
(default: None)
num_frames_scaling (bool): Multiply by number of frames in sequence in order to ~conserve photons (default: False)
combine_other_hdus (bool): Whether to combine other HDUs in the same way as the main data, err, DQ.
Otherwise, uses the HDUs from the first frame in a subset (default: False)
max_combined (int): maximum number of output frames allowed per split subset
before ``num_frames_per_group`` is automatically increased.
(default: 100)
pri_split_keywords (list of str): Additional primary header keywords to use
for splitting the dataset into subsets. The default split
keyword ['VISITID'] is always included, and input values
are appended. (default: ['VISITID'])
ext_split_keywords (list of str): Additional extension header keywords to use
for splitting the dataset into subsets. The default split
keyword ['DPAMNAME'] is always included, and input values
are appended. (default: ['DPAMNAME'])
Returns:
corgidrp.data.Dataset: dataset of combined frames. If
``num_frames_per_group`` is None, there is one output frame per split
subset, so ``len(output_dataset)`` equals the number of unique VISITID
groups or unique ``(VISITID, DPAMNAME)`` groups when ``DPAMNAME`` is
present for all frames. If ``num_frames_per_group`` is set, each split
subset is further combined in sequential groups of
``num_frames_per_group``. If a requested grouping would create more than
``max_combined`` output frames in a split subset, the grouping is
increased so that the output count stays at or below ``max_combined``
and trailing remainder frames are left out. Each output frame keeps the
same internal data cube shape as an input frame; only the dataset/frame
axis is collapsed.
"""
if num_frames_per_group is not None and num_frames_per_group <= 0:
raise ValueError("num_frames_per_group must be a positive integer or None")
if max_combined <= 0:
raise ValueError("max_combined must be a positive integer")
default_pri_split_keywords = ['VISITID']
default_ext_split_keywords = ['DPAMNAME']
if pri_split_keywords is None:
pri_split_keywords = []
if ext_split_keywords is None:
ext_split_keywords = []
pri_split_keywords = _append_unique(default_pri_split_keywords, pri_split_keywords)
ext_split_keywords = _append_unique(default_ext_split_keywords, ext_split_keywords)
split_keywords = list(pri_split_keywords) + list(ext_split_keywords)
split_description = "/".join(split_keywords) if split_keywords else "input"
# Combine each split subset independently.
split_datasets, _ = input_dataset.split_dataset(prihdr_keywords=pri_split_keywords, exthdr_keywords=ext_split_keywords)
combined_frames = []
any_auto_override = False
for split_dataset in split_datasets:
effective_num_frames_per_group = num_frames_per_group
if num_frames_per_group is not None:
requested_output_frames = len(split_dataset) / num_frames_per_group
if requested_output_frames > max_combined:
effective_num_frames_per_group = -(-len(split_dataset) // max_combined)
any_auto_override = True
combined_split = combine_subexposures(
split_dataset,
num_frames_per_group=effective_num_frames_per_group,
collapse=collapse,
num_frames_scaling=num_frames_scaling,
combine_other_hdus=combine_other_hdus,
)
combined_frames.extend(combined_split.frames.tolist())
combined_dataset = data.Dataset(combined_frames)
if num_frames_per_group is None:
history_msg = (
f"Combined frames by applying {collapse} within each {split_description} subset, "
f"result is a dataset with {len(combined_dataset)} frame(s)"
)
else:
history_msg = (
f"Combined frames by applying {collapse} in sequential groups of {num_frames_per_group} "
f"within each {split_description} subset, result is a dataset with "
f"{len(combined_dataset)} frame(s)"
)
if any_auto_override:
history_msg += f" Automatically increased some group sizes to keep per-subset output frame counts at or below {max_combined}."
combined_dataset.update_after_processing_step(history_msg)
return combined_dataset
[docs]
def derotate_arr(data_arr,northang_deg, xcen,ycen,new_center=None,astr_hdr=None,
is_dq=False,dq_round_threshold=0.05):
"""Derotates an array based on the provided NORTHANG angle, about the provided
center. Treats DQ arrays specially, converting to float to do the rotation,
and converting back to np.int64 afterwards. DQ output becomes only zeros and
ones, so detailed DQ flag information is not preserved.
Args:
data_arr (np.array): an array with 2-4 dimensions
northang_deg (float): angle (measured counter-clockwise) of the detector y axis from
celestial north (degrees). Calculated from the northangle of the astrometric cal frame and
the PA_APER offset between the astrom cal frame and the science frame.
xcen (float): x-coordinate of center about which to rotate
ycen (float): y-coordinate of center about which to rotate
new_center (tuple, optional): tuple of x- and y- coordinate of the new center to shift to.
astr_hdr (astropy.fits.Header, optional): WCS header which will be updated. Defaults to None.
is_dq (bool, optional): Flag to determine if this is a DQ array. Defaults to False.
dq_round_threshold (float, optional): value between 0-1 which determines the
threshold for spreading dq values to neighboring pixels after derotation.
Returns:
np.array: The derotated array.
"""
# Temporarily convert dq to floats
if is_dq:
data_arr = data_arr.astype(np.float32)
if data_arr.ndim == 2:
derotated_arr = rotate(data_arr,northang_deg,(xcen,ycen),
new_center=new_center,
astr_hdr=astr_hdr) # astr_hdr is corrected at above lines
elif data_arr.ndim == 3:
derotated_arr = []
for i,im in enumerate(data_arr):
derotated_im = rotate(im,northang_deg,(xcen,ycen),
new_center=new_center,
astr_hdr=astr_hdr if (i==0) else None) # astr_hdr is corrected only once
derotated_arr.append(derotated_im)
derotated_arr = np.array(derotated_arr)
elif data_arr.ndim == 4:
derotated_arr = []
for s,set in enumerate(data_arr):
derotated_set = []
for i,im in enumerate(set):
derotated_im = rotate(im,northang_deg,(xcen,ycen),
new_center=new_center,
astr_hdr=astr_hdr if (i==0 and s==0) else None) # astr_hdr is corrected only once
derotated_set.append(derotated_im)
derotated_arr.append(derotated_set)
derotated_arr = np.array(derotated_arr)
else:
raise ValueError('derotate_arr() not configured for data with >4 dimensions')
# convert dq_array back to ints
if is_dq:
derotated_arr[np.isnan(derotated_arr)] = 1 # assign nans to 1
derotated_arr_int = (derotated_arr>dq_round_threshold).astype(np.int64)
# import matplotlib.pyplot as plt
# plt.imshow(derotated_arr_int,origin='lower')
# plt.colorbar()
# plt.title(f'round_threshold: {round_threshold}')
# plt.show()
return derotated_arr_int
return derotated_arr
[docs]
def prop_err_dq(sci_dataset,ref_dataset,mode,dq_thresh=1,new_center=None):
"""Applies logic to propagate the dq arrays and error arrays
in a dataset through PSF subtraction.
Args:
sci_dataset (corgidrp.data.Dataset): The input science dataset.
ref_dataset (corgidrp.data.Dataset): The input reference dataset (or None if ADI only).
mode (str): The PSF subtraction mode, e.g. "ADI", "RDI", "ADI+RDI".
dq_thresh (int): Minimum dq flag value to be considered a bad pixel. Defaults to 1.
new_center (tuple): New center (xy) to align all frames. Defaults to pixel closest to array center.
Returns:
tuple of np.array: the dq array and err array which should apply to the PSF subtraction output dataset.
"""
# Assign master output dq & error (before derotation)
# dq shape = (n_rolls, n_wls(optional), y, x)
sci_input_dqs = sci_dataset.all_dq >= dq_thresh
sci_input_errs = np.full_like(sci_dataset.all_err,np.nan) # Set errors to np.nan for now
if new_center is None:
new_center = [int(sci_dataset.all_data.shape[-1]//2), int(sci_dataset.all_data.shape[-2]//2)]
# Align frames
aligned_sci_dq_arr = []
aligned_sci_err_arr = []
for i,frame in enumerate(sci_dataset):
xcen, ycen = frame.ext_hdr['STARLOCX'], frame.ext_hdr['STARLOCY']
frame.ext_hdr['STARLOCX'], frame.ext_hdr['STARLOCY'] = new_center
aligned_sci_dq = derotate_arr(sci_input_dqs[i],0, xcen,ycen,
new_center=new_center,is_dq=True)
aligned_sci_err = derotate_arr(sci_input_errs[i],0, xcen,ycen,
new_center=new_center)
aligned_sci_dq_arr.append(aligned_sci_dq)
aligned_sci_err_arr.append(aligned_sci_err)
aligned_sci_dq_arr = np.array(aligned_sci_dq_arr)
aligned_sci_err_arr = np.array(aligned_sci_err_arr)
if "RDI" in mode:
ref_input_dqs = ref_dataset.all_dq >= dq_thresh
ref_input_errs = np.full_like(ref_dataset.all_err,np.nan) # Set errors to np.nan for now
aligned_ref_dq_arr = []
aligned_ref_err_arr = []
for i,frame in enumerate(ref_dataset):
xcen, ycen = frame.ext_hdr['STARLOCX'], frame.ext_hdr['STARLOCY']
frame.ext_hdr['STARLOCX'], frame.ext_hdr['STARLOCY'] = new_center
aligned_ref_dq = derotate_arr(ref_input_dqs[i],0, xcen,ycen,
new_center=new_center,is_dq=True)
aligned_ref_err = derotate_arr(ref_input_errs[i],0, xcen,ycen,
new_center=new_center)
aligned_ref_dq_arr.append(aligned_ref_dq)
aligned_ref_err_arr.append(aligned_ref_err)
aligned_ref_dq_arr = np.array(aligned_ref_dq_arr)
aligned_ref_err_arr = np.array(aligned_ref_err_arr)
# If doing ADI, flag pixels that are bad in all science frames
if 'ADI' in mode:
aligned_sci_dq_arr[:] = np.all(aligned_sci_dq_arr,axis=0)
# If using references, flag pixels that are bad in all the ref frames
if 'RDI' in mode:
ref_output_dqs_flat = np.all(aligned_ref_dq_arr,axis=0,keepdims=True)
aligned_sci_dq_arr = np.logical_or(aligned_sci_dq_arr,ref_output_dqs_flat)
# Derotate dq & error
derotated_dq_arr = []
derotated_err_arr = []
for i,frame in enumerate(sci_dataset):
northang_deg = frame.ext_hdr['NORTHANG']
xcen, ycen = frame.ext_hdr['STARLOCX'], frame.ext_hdr['STARLOCY']
derotated_dq = derotate_arr(aligned_sci_dq_arr[i],northang_deg, xcen,ycen,is_dq=True)
derotated_err = derotate_arr(aligned_sci_err_arr[i],northang_deg, xcen,ycen)
derotated_dq_arr.append(derotated_dq)
derotated_err_arr.append(derotated_err)
# Collapse dq & error
dq_out_collapsed = np.where(np.all(derotated_dq_arr,axis=0),1,0)
err_out_collapsed = np.sqrt(np.sum(np.array(derotated_err_arr)**2,axis=0))
return dq_out_collapsed, err_out_collapsed
