class ScaleTemplates(MultiInputCommand):
"""Compute a scale factor fitting the radial profile of
a template PSF to frames from a data cube"""
template_path : str = xconf.field(help="Path to FITS file with templates in extensions")
saturation_threshold : Optional[float] = xconf.field(default=None, help="Value in counts above which pixels should be considered saturated and ignored for scaling purposes")
def main(self):
import dask
from ..core import LazyPipelineCollection
from ..tasks import iofits, improc
from .. import pipelines, utils
destination = self.destination
dest_fs = utils.get_fs(destination)
dest_fs.makedirs(destination, exist_ok=True)
output_filepath = utils.join(destination, utils.basename("template_scale_factors.fits"))
self.quit_if_outputs_exist([output_filepath])
all_inputs = self.get_all_inputs()
template_hdul = iofits.load_fits_from_path(self.template_path)
if len(all_inputs) == 1:
data_hdul = iofits.load_fits_from_path(all_inputs[0])
factors_hdul = pipelines.compute_scale_factors(data_hdul, template_hdul, self.saturation_threshold)
else:
coll = LazyPipelineCollection(all_inputs).map(iofits.load_fits_from_path)
factors_hdul = pipelines.compute_scale_factors(coll, template_hdul, self.saturation_threshold)
factors_hdul, = dask.compute(factors_hdul)
iofits.write_fits(factors_hdul, output_filepath)
class MultiInputCommand(InputCommand):
input: str = xconf.field(
help=
"Input file, directory, or wildcard pattern matching multiple files")
sample_every_n: int = xconf.field(
default=1,
help="Take every Nth file from inputs (for speed of debugging)")
file_extensions: list[str] = xconf.field(
default=DEFAULT_EXTENSIONS,
help="File extensions to match in the input (when given a directory)")
def get_all_inputs(self):
src_fs = utils.get_fs(self.input)
if '*' in self.input:
# handle globbing
all_inputs = src_fs.glob(self.input)
else:
# handle directory
if src_fs.isdir(self.input):
all_inputs = []
for extension in self.file_extensions:
glob_result = src_fs.glob(
utils.join(self.input, f"*{extension}"))
# returned paths from glob won't have protocol string or host
# so take the basenames of the files and we stick the other
# part back on from `entry`
all_inputs.extend([
utils.join(self.input, utils.basename(x))
for x in glob_result
])
# handle single file
else:
all_inputs = [self.input]
return list(sorted(all_inputs))[::self.sample_every_n]
class PixelRotationRangeConfig:
delta_px: float = xconf.field(
default=0,
help="Maximum difference between target frame value and matching frames"
)
r_px: float = xconf.field(
default=None, help="Radius at which to calculate motion in pixels")
class UpdateHeaders(MultiInputCommand):
"""Update FITS headers for file (See https://archive.stsci.edu/fits/fits_standard/node40.html for examples)"""
keywords: dict[str, types.FITS_KW_VAL] = xconf.field(
default_factory=dict,
help="Header keyword values to set in extension 0")
extensions: dict[typing.Union[str, int], dict[
str, types.FITS_KW_VAL]] = xconf.field(
default_factory=dict,
help=
"Mapping of extension names to a table of KEYWORD=value pairs to set"
)
def main(self):
import fsspec.spec
from .. import utils
from .. import pipelines
from ..ref import clio
from ..tasks import iofits, sky_model, improc
destination = self.destination
dest_fs = utils.get_fs(destination)
assert isinstance(dest_fs, fsspec.spec.AbstractFileSystem)
log.debug(f"calling makedirs on {dest_fs} at {destination}")
dest_fs.makedirs(destination, exist_ok=True)
all_inputs = self.get_all_inputs()
output_filepaths = [
utils.join(destination, utils.basename(filepath))
for filepath in all_inputs
]
self.quit_if_outputs_exist(output_filepaths)
coll = LazyPipelineCollection(all_inputs).map(
iofits.load_fits_from_path)
any_updates = False
header_updates = {}
if len(self.keywords):
header_updates = {0: self.keywords}
any_updates = True
if 'PRIMARY' in self.extensions and 0 in header_updates: # unlikely edge case where ext 0 is referenced by name too
primary_updates = self.extensions.pop('PRIMARY')
any_updates = any_updates or len(primary_updates) > 0
header_updates[0].update(primary_updates)
for key in self.extensions:
any_updates = any_updates or len(self.extensions[key]) > 0
header_updates[key] = self.extensions[key]
if not any_updates:
log.error(
"Supplied configuration would not make any modifications, exiting"
)
sys.exit(0)
output_coll = coll.map(
lambda x, hdrs: x.updated_copy(new_headers_for_exts=hdrs),
header_updates)
result = output_coll.zip_map(iofits.write_fits,
output_filepaths,
overwrite=True).compute()
log.info(result)
class ExcludeRangeConfig:
angle: Union[AngleRangeConfig, PixelRotationRangeConfig] = xconf.field(
default=AngleRangeConfig(),
help="Apply exclusion to derotation angles")
nearest_n_frames: int = xconf.field(
default=0,
help=
"Number of additional temporally-adjacent frames on either side of the target frame to exclude from the sequence when computing the KLIP eigenimages"
)
class VappAdiObservation(AdiObservation):
vapp_left_ext : types.FITS_EXT = xconf.field(help=utils.unwrap(
"""FITS extension to combine across dataset containing
complementary gvAPP-180 PSF where the dark hole is left of +Y"""
))
vapp_right_ext : types.FITS_EXT = xconf.field(help=utils.unwrap(
"""FITS extension to combine across dataset containing
complementary gvAPP-180 PSF where the dark hole is right of +Y"""
))
class KlipInputConfig:
sci_arr: FitsConfig
signal_arr: Optional[FitsConfig]
estimation_mask: Optional[FitsConfig]
mask_min_r_px: float = xconf.field(
default=0,
help="Apply radial mask excluding pixels < mask_min_r_px from center")
mask_max_r_px: Optional[float] = xconf.field(
default=None,
help="Apply radial mask excluding pixels > mask_max_r_px from center")
class SearchConfig:
min_r_px: float = xconf.field(
default=None,
help="Limit blind search to pixels more than this radius from center")
max_r_px: float = xconf.field(
default=None,
help="Limit blind search to pixels less than this radius from center")
snr_threshold: float = xconf.field(
default=5.0,
help="Threshold above which peaks of interest should be reported")
class CutoutConfig:
search_box: typing.Union[BoxFromCenter, BoxFromOrigin, Box] = xconf.field(
default=Box(),
help="Search box to find the PSF to cross-correlate with the template")
template: typing.Union[base.FitsConfig, GaussianTemplate] = xconf.field(
default=GaussianTemplate(),
help=utils.unwrap("""
Template cross-correlated with the search region to align images to a common grid, either given as a FITS image
or specified as a centered 2D Gaussian with given FWHM
"""))
class FitsConfig:
path: str = xconf.field(
help="Path from which to load the containing FITS file")
ext: typing.Union[int,
str] = xconf.field(default=0,
help="Extension from which to load")
def load(self):
from ..tasks import iofits
hdul = iofits.load_fits_from_path(self.path)
return hdul[self.ext].data
class ComponentConfig(FitsConfig):
derotate_by: Optional[FitsTableColumnConfig] = xconf.field(
default=None,
help="Metadata table column with derotation angle in degrees")
mask: Optional[FitsConfig] = xconf.field(default=None)
destination_idx: int = xconf.field(default=0)
amplitude_scale: float = xconf.field(
default=1,
help=
"Factor by which pixels in this component are multiplied before adding to the stack"
)
class FitsTableColumnConfig(FitsConfig):
table_column: str = xconf.field(
help="Path from which to load the containing FITS file")
ext: typing.Union[int,
str] = xconf.field(default="OBSTABLE",
help="Extension from which to load")
def load(self):
from ..tasks import iofits
hdul = iofits.load_fits_from_path(self.path)
return hdul[self.ext].data[self.table_column]
class KlipInputConfig:
sci_arr: FitsConfig
estimation_mask: FitsConfig
combination_mask: Optional[FitsConfig]
initial_decomposition_path: str = None
mask_min_r_px: Union[int, float] = xconf.field(
default=None,
help="Apply radial mask excluding pixels < mask_min_r_px from center")
mask_max_r_px: Union[int, float] = xconf.field(
default=None,
help="Apply radial mask excluding pixels > mask_max_r_px from center")
class LocalRayConfig(CommonRayConfig):
cpus: Optional[int] = xconf.field(
default=None,
help="CPUs available to built-in Ray cluster (default is auto-detected)"
)
gpus: Optional[int] = xconf.field(
default=None,
help="GPUs available to built-in Ray cluster (default is auto-detected)"
)
resources: Optional[dict[str, float]] = xconf.field(
default_factory=dict,
help="Node resources available when running in standalone mode")
class FileConfig:
path: str = xconf.field(help="File path")
def open(self, mode='r'):
from ..utils import get_fs
fs = get_fs(self.path)
return fs.open(self.path, mode)
class DaskConfig:
distributed: bool = xconf.field(
default=False,
help=
"Whether to execute Dask workflows with a cluster to parallelize work")
log_level: str = xconf.field(
default='WARN',
help="What level of logs to show from the scheduler and workers")
port: int = xconf.field(default=8786,
help="Port to contact dask-scheduler")
host: typing.Optional[str] = xconf.field(
default=None, help="Hostname of running dask-scheduler")
n_processes: int = xconf.field(
default=None,
help="Override automatically selected number of processes to spawn")
n_threads: int = xconf.field(
default=None,
help=
"Override automatically selected number of threads per process to spawn"
)
synchronous: bool = xconf.field(
default=True,
help=
"Whether to disable Dask's parallelism in favor of lower levels (ignored when distributed=True)"
)
class GetReferenceData(xconf.Command):
exclude : list[str] = xconf.field(default_factory=list, help='Dotted module path for resources to exclude (e.g. doodads.ref.mko_filters)')
def main(self):
import matplotlib
matplotlib.use('Agg')
log.info(f"Excluded modules: {self.exclude}")
selected_resources = REMOTE_RESOURCES.filter(exclude=self.exclude)
log.info(f"Processing registered resources from {pformat(selected_resources.resources)}")
selected_resources.download_and_convert()
class CompanionConfig(MeasurementConfig):
scale: float = xconf.field(help=utils.unwrap(
"""Scale factor multiplied by template (and optional template
per-frame scale factor) to give companion image,
i.e., contrast ratio. Can be negative or zero."""))
def to_companionspec(self):
from xpipeline.tasks.characterization import CompanionSpec
return CompanionSpec(self.r_px, self.pa_deg, self.scale)
class SamplingConfig:
n_radii: int = xconf.field(
help="Number of steps in radius at which to probe contrast")
spacing_px: float = xconf.field(
help=
"Spacing in pixels between contrast probes along circle (sets number of probes at radius by 2 * pi * r / spacing)"
)
scales: list[float] = xconf.field(
default_factory=lambda: [0.0],
help="Probe contrast levels (C = companion / host)")
iwa_px: float = xconf.field(help="Inner working angle (px)")
owa_px: float = xconf.field(help="Outer working angle (px)")
def __post_init__(self):
# to make use of this for detection, we must also apply the matched
# filter in the no-injection case for each combination of parameters
self.scales = [float(s) for s in self.scales]
if 0.0 not in self.scales:
self.scales.insert(0, 0.0)
class TemplateConfig:
path: str = xconf.field(
help=utils.unwrap("""Path to FITS image of template PSF, scaled to the
average amplitude of the host star signal such that
multiplying by the contrast gives an appropriately
scaled planet PSF"""))
ext: typing.Optional[types.FITS_EXT] = xconf.field(default=None,
help=utils.unwrap("""
Extension containing the template data (default: same as template name)
"""))
scale_factors_path: typing.Optional[str] = xconf.field(help=utils.unwrap(
"""Path to FITS file with extensions for each data extension
containing 1D arrays of scale factors that match template PSF
intensity to real PSF intensity per-frame"""))
scale_factors_ext: typing.Optional[types.FITS_EXT] = xconf.field(
default=None,
help=utils.unwrap("""
Extension containing the per-frame scale factors by which
the template data is multiplied before applying the
companion scale factor (default: same as template name)
"""))
class BaseCommand(xconf.Command):
destination: str = xconf.field(default=".", help="Output directory")
random_state: int = xconf.field(
default=0,
help="Initialize NumPy's random number generator with this seed")
cpus: int = xconf.field(default=utils.available_cpus(),
help="Number of CPUs free for use")
def get_dest_fs(self):
dest_fs = utils.get_fs(self.destination)
log.debug(f"calling makedirs on {dest_fs} at {self.destination}")
dest_fs.makedirs(self.destination, exist_ok=True)
return dest_fs
def check_for_outputs(self, output_paths):
dest_fs = self.get_dest_fs()
for op in output_paths:
if dest_fs.exists(op):
return True
return False
def quit_if_outputs_exist(self, output_paths):
if self.check_for_outputs(output_paths):
log.info(f"Outputs exist at {output_paths}")
log.info(f"Remove to re-run")
sys.exit(0)
def get_output_paths(self, *output_paths):
output_paths = [
utils.join(self.destination, op) for op in output_paths
]
return output_paths
def __post_init__(self):
numpy.random.seed(self.random_state)
log.debug(f"Set NumPy random seed to {self.random_state}")
class SummarizeGrid(InputCommand):
ext : str = xconf.field(default="grid", help="FITS binary table extension with calibration grid")
columns : GridColumnsConfig = xconf.field(default=GridColumnsConfig(), help="Grid column names")
arcsec_per_pixel : float = xconf.field(default=clio.CLIO2_PIXEL_SCALE.value)
wavelength_um : float = xconf.field(help="Wavelength in microns")
primary_diameter_m : float = xconf.field(default=magellan.PRIMARY_MIRROR_DIAMETER.to(u.m).value)
coverage_mask : FitsConfig = xconf.field(help="Mask image with 1s where pixels have observation coverage and 0 elsewhere")
def main(self):
import pandas as pd
from ..tasks import iofits, improc
output_filepath, = self.get_output_paths("summarize_grid.fits")
self.quit_if_outputs_exist([output_filepath])
coverage_mask = self.coverage_mask.load() == 1
yc, xc = improc.arr_center(coverage_mask)
grid_hdul = iofits.load_fits_from_path(self.input)
grid_tbl = grid_hdul[self.ext].data
log.info(f"Loaded {len(grid_tbl)} points for evaluation")
grid_df = pd.DataFrame(grid_tbl)
import pandas as pd
limits_df, detections_df = characterization.summarize_grid(
grid_df,
r_px_colname=self.columns.r_px,
pa_deg_colname=self.columns.pa_deg,
snr_colname=self.columns.snr,
injected_scale_colname=self.columns.injected_scale,
hyperparameter_colnames=self.columns.hyperparameters,
)
limits_df['delta_mag_contrast_limit_5sigma'] = characterization.contrast_to_deltamag(limits_df['contrast_limit_5sigma'].to_numpy())
for df in [limits_df, detections_df]:
df['r_as'] = (limits_df['r_px'].to_numpy() * u.pix * self.arcsec_per_pixel).value
df['r_lambda_over_d'] = characterization.arcsec_to_lambda_over_d(
df['r_as'].to_numpy() * u.arcsec,
self.wavelength_um * u.um,
d=self.primary_diameter_m * u.m
)
log.info(f"Sampled {len(limits_df)} locations for contrast limits and detection")
lim_xs, lim_ys = characterization.r_pa_to_x_y(limits_df[self.columns.r_px], limits_df[self.columns.pa_deg], xc, yc)
contrast_lim_map = characterization.points_to_map(lim_xs, lim_ys, limits_df['contrast_limit_5sigma'], coverage_mask)
det_xs, det_ys = characterization.r_pa_to_x_y(detections_df[self.columns.r_px], detections_df[self.columns.pa_deg], xc, yc)
detection_map = characterization.points_to_map(det_xs, det_ys, detections_df['snr'], coverage_mask)
hdus = [iofits.DaskHDU(None, kind="primary")]
hdus.append(iofits.DaskHDU(contrast_lim_map, name="limits_5sigma_contrast_map"))
hdus.append(iofits.DaskHDU(limits_df.to_records(), kind="bintable", name="limits"))
hdus.append(iofits.DaskHDU(detection_map, name="detection_snr_map"))
hdus.append(iofits.DaskHDU(detections_df.to_records(), kind="bintable", name="detection"))
iofits.write_fits(iofits.DaskHDUList(hdus), output_filepath)
class ComputeSkyModel(MultiInputCommand):
"""Compute sky model eigenimages"""
n_components : int = xconf.field(default=6, help="Number of PCA components to calculate")
mask_dilate_iters : int = xconf.field(default=6, help="Number of times to grow mask regions before selecting cross-validation pixels")
test_fraction : float = xconf.field(default=0.25, help="Fraction of inputs to reserve for cross-validation")
excluded_regions : Optional[FileConfig] = xconf.field(default_factory=list, help="Regions presumed illuminated to be excluded from background estimation, stored as DS9 region file (reg format)")
ext : FITS_EXT = xconf.field(default='SCI', help="Extension containing science data")
dq_ext : FITS_EXT = xconf.field(default='DQ', help="Extension containing data quality metadata")
def main(self):
import dask
from .. import utils
from .. import pipelines
from ..core import LazyPipelineCollection
from ..tasks import iofits, regions
# outputs
destination = self.destination
dest_fs = utils.get_fs(destination)
log.debug(f"calling makedirs on {dest_fs} at {self.destination}")
dest_fs.makedirs(destination, exist_ok=True)
model_fn = utils.join(self.destination, "sky_model.fits")
self.quit_if_outputs_exist([model_fn])
# execute
inputs_coll = LazyPipelineCollection(self.get_all_inputs()).map(iofits.load_fits_from_path)
one_input_hdul = dask.compute(inputs_coll.items[0])[0]
if self.ext not in one_input_hdul or self.dq_ext not in one_input_hdul:
raise RuntimeError(f"Looking for {self.ext=} and {self.dq_ext=} in first input failed, check inputs?")
plane_shape = one_input_hdul[0].data.shape
if isinstance(self.excluded_regions, FileConfig):
with self.excluded_regions.open() as fh:
excluded_regions = regions.load_file(fh)
excluded_pixels_mask = regions.make_mask(excluded_regions, plane_shape)
d_sky_model = pipelines.compute_sky_model(
inputs_coll,
plane_shape,
self.test_fraction,
self.random_state,
self.n_components,
self.mask_dilate_iters,
excluded_pixels_mask=excluded_pixels_mask,
)
the_sky_model = dask.compute(d_sky_model)[0]
hdul = the_sky_model.to_hdulist()
hdul.writeto(model_fn, overwrite=True)
log.info(f"Sky model written to {model_fn}")
class ClioCalibrate(MultiInputCommand):
"""Apply bad pixel map, linearity correction, and saturation flags"""
badpix_path: str = xconf.field(
help=
"Path to full detector bad pixel map FITS file (1 where pixel is bad)")
def main(self):
from .. import utils
from .. import pipelines
from ..ref import clio
from ..tasks import iofits
destination = self.destination
dest_fs = utils.get_fs(destination)
assert isinstance(dest_fs, fsspec.spec.AbstractFileSystem)
log.debug(f"calling makedirs on {dest_fs} at {destination}")
dest_fs.makedirs(destination, exist_ok=True)
# infer planes per cube
all_inputs = self.get_all_inputs()
hdul = iofits.load_fits_from_path(all_inputs[0])
plane_shape = hdul[0].data.shape
n_output_files = len(all_inputs)
output_filepaths = [
utils.join(destination, f"{self.name}_{i:04}.fits")
for i in range(n_output_files)
]
self.quit_if_outputs_exist(output_filepaths)
coll = LazyPipelineCollection(all_inputs).map(
iofits.load_fits_from_path)
badpix_path = self.badpix_path
full_badpix_arr = iofits.load_fits_from_path(badpix_path)[0].data
badpix_arr = clio.badpix_for_shape(full_badpix_arr, plane_shape)
output_coll = pipelines.clio_badpix_linearity(coll, badpix_arr,
plane_shape)
result = output_coll.zip_map(iofits.write_fits,
output_filepaths,
overwrite=True).compute()
log.info(result)
class FitsTable(xconf.Command):
input : str = xconf.field(help="FITS table file")
ext : typing.Optional[str] = xconf.field(default=None, help="FITS extension to read containing a binary table (default: first one found)")
rows : int = xconf.field(default=10)
all : bool = xconf.field(default=False, help="Whether to ignore the rows setting and print all")
sort_on : typing.Optional[str] = xconf.field(default=None, help="Column name on which to sort")
reverse : bool = xconf.field(default=False, help="Reverse order (before truncating at rows)")
def main(self):
from astropy.io import fits
import numpy as np
with open(self.input, 'rb') as fh:
hdul = fits.open(fh)
if self.ext is not None:
hdu = hdul[self.ext]
else:
for idx, hdu in enumerate(hdul):
if isinstance(hdu, fits.BinTableHDU):
break
if not isinstance(hdu, fits.BinTableHDU):
raise ValueError("Not a BinTable extension")
tbl = hdu.data
fields = tbl.dtype.fields
if self.sort_on is not None and self.sort_on in fields:
sorter = np.argsort(tbl[self.sort_on])
tbl = tbl[sorter]
if self.reverse:
tbl = tbl[::-1]
if not self.all:
tbl = tbl[:self.rows]
cols = []
colfmts = []
for fld in fields:
width = max(len(fld), 9)
cols.append(f"{fld:9}")
if np.issubdtype(tbl[fld].dtype, np.number):
colfmts.append("{:< " + str(width) + ".3g}")
else:
colfmts.append("{:<" + str(width) + "}")
print(' '.join(cols))
for row in tbl:
cols = []
for fld, colfmt in zip(fields, colfmts):
cols.append(colfmt.format(row[fld]))
print(' '.join(cols))
class SubtractStarlight(BaseCommand):
"Subtract starlight with KLIP"
inputs: list[KlipInputConfig] = xconf.field(
help="Input data to simultaneously reduce")
destination: str = xconf.field(help="Destination path to save results")
obstable: FitsConfig = xconf.field(help="Metadata table in FITS")
k_klip: int = xconf.field(
default=10,
help="Number of modes to subtract in starlight subtraction")
exclude_ranges: ExcludeRangeConfig = xconf.field(
default=ExcludeRangeConfig(),
help="How to exclude frames from reference sample")
exclude_frames: Optional[ExcludeFramesConfig] = xconf.field(
default=None, help="How to select frames to drop from the inputs")
strategy: constants.KlipStrategy = xconf.field(
default=constants.KlipStrategy.DOWNDATE_SVD,
help="Implementation of KLIP to use")
reuse_eigenimages: bool = xconf.field(
default=False,
help=
"Apply KLIP without adjusting the eigenimages at each step (much faster, less powerful)"
)
initial_decomposition_only: bool = xconf.field(
default=False, help="Whether to output initial decomposition and exit")
initial_decomposition_path: Optional[str] = xconf.field(
default=None, help="Initial decomposition as FITS file")
def _exclude_frames_mask(self, config: ExcludeFramesConfig):
values = config.values.load()
if config.op is constants.CompareOperation.GT:
return values > config.compare_to
elif config.op is constants.CompareOperation.GE:
return values >= config.compare_to
elif config.op is constants.CompareOperation.EQ:
return values == config.compare_to
elif config.op is constants.CompareOperation.LE:
return values <= config.compare_to
elif config.op is constants.CompareOperation.LT:
return values < config.compare_to
elif config.op is constants.CompareOperation.NE:
return values != config.compare_to
else:
raise ValueError(f"{config.op=} unknown")
def _load_obstable(self):
from ..tasks import iofits
obstable_hdul = iofits.load_fits_from_path(self.obstable.path)
obstable = obstable_hdul[self.obstable.ext].data
return obstable
def _load_initial_decomposition(self, path):
from ..tasks import iofits, starlight_subtraction
decomp_hdul = iofits.load_fits_from_path(path)
return starlight_subtraction.InitialDecomposition(
mtx_u0=decomp_hdul['MTX_U0'].data,
diag_s0=decomp_hdul['DIAG_S0'].data,
mtx_v0=decomp_hdul['MTX_V0'].data,
)
def _save_initial_decomposition(self, initial_decomposition,
output_initial_decomp_fn):
from ..tasks import iofits
hdus = [
iofits.DaskHDU(data=None, kind="primary"),
iofits.DaskHDU(initial_decomposition.mtx_u0, name="MTX_U0"),
iofits.DaskHDU(initial_decomposition.diag_s0, name="DIAG_S0"),
iofits.DaskHDU(initial_decomposition.mtx_v0, name="MTX_V0"),
]
iofits.write_fits(iofits.DaskHDUList(hdus), output_initial_decomp_fn)
def _make_mask(self, klip_input_cfg: KlipInputConfig, sci_arr: np.ndarray):
from ..tasks import iofits, improc
# mask file(s)
if klip_input_cfg.estimation_mask is not None:
estimation_mask_hdul = iofits.load_fits_from_path(
klip_input_cfg.estimation_mask.path)
estimation_mask = estimation_mask_hdul[
klip_input_cfg.estimation_mask.ext].data
else:
estimation_mask = np.ones(sci_arr.shape[1:])
# coerce to bools
estimation_mask = estimation_mask == 1
if klip_input_cfg.mask_min_r_px == 0 and klip_input_cfg.mask_max_r_px is None:
return estimation_mask
else:
ctr = improc.arr_center(sci_arr.shape[1:])
frame_shape = sci_arr.shape[1:]
if klip_input_cfg.mask_min_r_px is None:
min_r = 0
else:
min_r = klip_input_cfg.mask_min_r_px
annular_mask = improc.mask_arc(
ctr,
frame_shape,
from_radius=min_r,
to_radius=klip_input_cfg.mask_max_r_px)
estimation_mask &= annular_mask
return estimation_mask
def _assemble_klip_input(self, klip_input_cfg: KlipInputConfig,
obstable: np.ndarray):
from ..tasks import iofits, starlight_subtraction
input_hdul = iofits.load_fits_from_path(klip_input_cfg.sci_arr.path)
sci_arr = input_hdul[klip_input_cfg.sci_arr.ext].data
if klip_input_cfg.signal_arr is not None:
if klip_input_cfg.signal_arr.path != klip_input_cfg.sci_arr.path:
signal_hdul = iofits.load_fits_from_path(
klip_input_cfg.sci_arr.path)
else:
signal_hdul = input_hdul
signal_arr = signal_hdul[klip_input_cfg.signal_arr.ext].data
else:
signal_arr = None
estimation_mask = self._make_mask(klip_input_cfg, sci_arr)
return starlight_subtraction.KlipInput(sci_arr, obstable,
estimation_mask, signal_arr)
def main(self):
from ..tasks import iofits
output_subtracted_fn = utils.join(self.destination,
"starlight_subtracted.fits")
output_initial_decomp_fn = utils.join(self.destination,
"initial_decomposition.fits")
destination = self.destination
dest_fs = utils.get_fs(destination)
dest_fs.makedirs(destination, exist_ok=True)
if self.initial_decomposition_only:
if dest_fs.exists(output_initial_decomp_fn):
raise FileExistsError(
f"Output exists at {output_initial_decomp_fn}")
else:
if dest_fs.exists(output_subtracted_fn):
raise FileExistsError(
f"Output exists at {output_subtracted_fn}")
obstable = self._load_obstable()
klip_inputs = []
for inputcfg in self.inputs:
klip_inputs.append(self._assemble_klip_input(inputcfg, obstable))
if self.initial_decomposition_path is not None:
initial_decomposition = self._load_initial_decomposition(
self.initial_decomposition_path)
else:
initial_decomposition = None
klip_params = self._assemble_klip_params(obstable,
initial_decomposition)
import time
start = time.perf_counter()
from ..pipelines import klip_multi
result = klip_multi(klip_inputs, klip_params)
if klip_params.initial_decomposition_only:
self._save_initial_decomposition(result, output_initial_decomp_fn)
return 0
else:
outcubes, outmeans = result
elapsed = time.perf_counter() - start
log.info(f"Computed in {elapsed} sec")
hdus = [iofits.DaskHDU(data=None, kind="primary")]
for idx in range(len(outcubes)):
hdus.append(iofits.DaskHDU(outcubes[idx], name=f"SCI_{idx}"))
hdus.append(iofits.DaskHDU(outmeans[idx], name=f"MEAN_{idx}"))
iofits.write_fits(iofits.DaskHDUList(hdus), output_subtracted_fn)
def _make_exclusions(self, exclude: ExcludeRangeConfig, obstable):
import numpy as np
from ..tasks import starlight_subtraction
exclusions = []
derotation_angles = obstable[exclude.angle_deg_col]
if exclude.nearest_n_frames > 0:
indices = np.arange(derotation_angles.shape[0])
exc = starlight_subtraction.ExclusionValues(
exclude_within_delta=exclude.nearest_n_frames,
values=indices,
num_excluded_max=2 * exclude.nearest_n_frames + 1)
exclusions.append(exc)
if isinstance(exclude.angle,
PixelRotationRangeConfig) and exclude.angle.delta_px > 0:
exc = starlight_subtraction.ExclusionValues(
exclude_within_delta=exclude.angle.delta_px,
values=exclude.angle.r_px *
np.unwrap(np.deg2rad(derotation_angles)))
exclusions.append(exc)
elif isinstance(exclude.angle,
AngleRangeConfig) and exclude.angle.delta_deg > 0:
exc = starlight_subtraction.ExclusionValues(
exclude_within_delta=exclude.angle.delta_deg,
values=derotation_angles)
exclusions.append(exc)
else:
pass # not an error to have delta of zero, just don't exclude based on rotation
return exclusions
def _assemble_klip_params(self, obstable, initial_decomposition):
import numpy as np
from ..tasks import starlight_subtraction
exclusions = self._make_exclusions(self.exclude_ranges, obstable)
klip_params = starlight_subtraction.KlipParams(
self.k_klip,
exclusions,
decomposer=starlight_subtraction.DEFAULT_DECOMPOSERS[
self.strategy],
strategy=self.strategy,
reuse=self.reuse_eigenimages,
initial_decomposition_only=self.initial_decomposition_only,
initial_decomposition=initial_decomposition)
log.debug(klip_params)
return klip_params
class ExcludeFramesConfig:
values: Union[FitsConfig, FitsTableColumnConfig] = xconf.field(help="")
op: constants.CompareOperation = xconf.field(help="")
compare_to: Union[float, str] = xconf.field(help="")
class PerTaskConfig:
generate: float = xconf.field(help="amount required in model generation")
decompose: float = xconf.field(help="amount required in basis computation")
evaluate: float = xconf.field(help="amount required in fitting")
class VappTrap(InputCommand):
checkpoint: str = xconf.field(
default=None,
help=
"Save checkpoints to this path, and/or resume grid from this checkpoint (no verification of parameters used to generate the grid is performed)"
)
checkpoint_every_x: int = xconf.field(
default=10,
help="Write current state of grid to disk every X iterations")
every_t_frames: int = xconf.field(
default=1, help="Use every Tth frame as the input cube")
k_modes_vals: list[int] = xconf.field(default_factory=lambda: [15, 100],
help="")
left_extname: FITS_EXT = xconf.field(help="")
left_template: FitsConfig = xconf.field(help="")
right_template: FitsConfig = xconf.field(help="")
right_extname: FITS_EXT = xconf.field(help="")
left_scales: FitsConfig = xconf.field(help="")
right_scales: FitsConfig = xconf.field(help="")
left_mask: FitsConfig = xconf.field(help="")
right_mask: FitsConfig = xconf.field(help="")
angles: Union[FitsConfig, FitsTableColumnConfig] = xconf.field(help="")
sampling: SamplingConfig = xconf.field(
help=
"Configure the sampling of the final derotated field for detection and contrast calibration"
)
min_coverage_frac: float = xconf.field(help="")
ray: Union[LocalRayConfig, RemoteRayConfig] = xconf.field(
default=LocalRayConfig(),
help="Ray distributed framework configuration")
ring_exclude_px: float = xconf.field(
default=12,
help=
"When selecting reference pixel timeseries, determines width of ring centered at radius of interest for which pixel vectors are excluded"
)
resel_px: float = xconf.field(
default=8, help="Resolution element in pixels for these data")
gpu_ram_mb_per_task: Union[float, str, PerTaskConfig, None] = xconf.field(
default=None,
help=
"Maximum amount of GPU RAM used by a stage or grid point, or 'measure'"
)
max_tasks_per_gpu: float = xconf.field(
default=1,
help=
"When GPU is utilized, this is the maximum number of tasks scheduled on the same GPU (RAM permitting)"
)
ram_mb_per_task: Union[float, str, PerTaskConfig, None] = xconf.field(
default=None,
help="Maximum amount of RAM used by a stage or grid point, or 'measure'"
)
use_gpu_decomposition: bool = xconf.field(default=False, help="")
use_gpu_fit: bool = xconf.field(default=False, help="")
use_cgls: bool = xconf.field(default=False, help="")
benchmark: bool = xconf.field(default=False, help="")
benchmark_trials: int = xconf.field(default=2, help="")
efficient_decomp_reuse: bool = xconf.field(
default=True,
help=
"Use a single decomposition for all PAs at given separation by masking a ring"
)
def main(self):
dest_fs = self.get_dest_fs()
coverage_map_fn = f'coverage_t{self.every_t_frames}.fits'
covered_pix_mask_fn = f'coverage_t{self.every_t_frames}_{self.min_coverage_frac}.fits'
output_filepath, coverage_map_path, covered_pix_mask_path = self.get_output_paths(
"grid.fits",
coverage_map_fn,
covered_pix_mask_fn,
)
hdul = iofits.load_fits_from_path(self.input)
log.debug(f"Loaded from {self.input}")
# decimate
left_cube = hdul[self.left_extname].data[::self.every_t_frames]
right_cube = hdul[self.right_extname].data[::self.every_t_frames]
angles = self.angles.load()[::self.every_t_frames]
left_scales = self.left_scales.load()[::self.every_t_frames]
right_scales = self.right_scales.load()[::self.every_t_frames]
log.debug(f"After decimation, {len(angles)} frames remain")
# get masks
left_mask = self.left_mask.load() == 1
right_mask = self.right_mask.load() == 1
mask = left_mask | right_mask
log.debug(f"{np.count_nonzero(mask)=}")
# templates
left_template = self.left_template.load()
right_template = self.right_template.load()
model_inputs = ModelInputs(
data_cube_shape=left_cube.shape,
left_template=left_template,
right_template=right_template,
left_scales=left_scales,
right_scales=right_scales,
angles=angles,
mask=mask,
)
# init ray
ray_init_kwargs = {
'runtime_env': {
'env_vars': {
'RAY_USER_SETUP_FUNCTION':
'xpipeline.commands.vapp_trap.init_worker',
'MKL_NUM_THREADS': '1',
'OMP_NUM_THREADS': '1',
'NUMBA_NUM_THREADS': '1',
}
}
}
if isinstance(self.ray, RemoteRayConfig):
ray.init(self.ray.url, **ray_init_kwargs)
else:
ray.init(num_cpus=self.ray.cpus,
num_gpus=self.ray.gpus,
resources=self.ray.resources,
**ray_init_kwargs)
options = {'resources': {}}
generate_options = options.copy()
decompose_options = options.copy()
evaluate_options = options.copy()
measure_ram = False
if isinstance(self.ram_mb_per_task, PerTaskConfig):
generate_options[
'memory'] = self.ram_mb_per_task.generate * BYTES_PER_MB
decompose_options[
'memory'] = self.ram_mb_per_task.decompose * BYTES_PER_MB
evaluate_options[
'memory'] = self.ram_mb_per_task.evaluate * BYTES_PER_MB
elif self.ram_mb_per_task == "measure":
measure_ram = True
elif self.ram_mb_per_task is not None:
# number or None
generate_options['memory'] = self.ram_mb_per_task * BYTES_PER_MB
decompose_options['memory'] = self.ram_mb_per_task * BYTES_PER_MB
evaluate_options['memory'] = self.ram_mb_per_task * BYTES_PER_MB
if self.use_gpu_decomposition:
gpu_frac = 1 / self.max_tasks_per_gpu
log.debug(
f"Using {self.max_tasks_per_gpu} tasks per GPU, {gpu_frac=}")
# generate_options['num_gpus'] = gpu_frac
decompose_options['num_gpus'] = gpu_frac
# evaluate_options['num_gpus'] = gpu_frac
if self.gpu_ram_mb_per_task is None:
raise RuntimeError(f"Specify GPU RAM per task")
if isinstance(self.gpu_ram_mb_per_task, PerTaskConfig):
generate_options['resources'] = {
'gpu_memory_mb': self.gpu_ram_mb_per_task.generate
}
decompose_options['resources'] = {
'gpu_memory_mb': self.gpu_ram_mb_per_task.decompose
}
evaluate_options['resources'] = {
'gpu_memory_mb': self.gpu_ram_mb_per_task.evaluate
}
elif self.gpu_ram_mb_per_task == "measure":
measure_ram = True
else:
# number or None
generate_options['resources'] = {
'gpu_memory_mb': self.gpu_ram_mb_per_task
}
decompose_options['resources'] = {
'gpu_memory_mb': self.gpu_ram_mb_per_task
}
evaluate_options['resources'] = {
'gpu_memory_mb': self.gpu_ram_mb_per_task
}
if not dest_fs.exists(coverage_map_path):
log.debug(f"Computing coverage map")
final_coverage = pipelines.adi_coverage(mask, angles)
iofits.write_fits(
iofits.DaskHDUList([iofits.DaskHDU(final_coverage)]),
coverage_map_path)
log.debug(f"Wrote coverage map to {coverage_map_path}")
else:
final_coverage = iofits.load_fits_from_path(
coverage_map_path)[0].data
n_frames = len(angles)
covered_pix_mask = final_coverage > int(
n_frames * self.min_coverage_frac)
from skimage.morphology import binary_closing
covered_pix_mask = binary_closing(covered_pix_mask)
log.debug(
f"Coverage map with {self.min_coverage_frac} fraction gives {np.count_nonzero(covered_pix_mask)} possible pixels to analyze"
)
if not dest_fs.exists(covered_pix_mask_path):
iofits.write_fits(
iofits.DaskHDUList(
[iofits.DaskHDU(covered_pix_mask.astype(int))]),
covered_pix_mask_path)
log.debug(f"Wrote covered pix mask to {covered_pix_mask_path}")
log.debug("Generating grid")
grid = grid_generate(
self.k_modes_vals,
covered_pix_mask.shape,
self.sampling,
)
if not self.benchmark:
if self.checkpoint is not None and utils.get_fs(
self.checkpoint).exists(self.checkpoint):
try:
hdul = iofits.load_fits_from_path(self.checkpoint)
grid = np.asarray(hdul['grid'].data)
log.debug(f"Loaded checkpoint successfully")
except Exception as e:
log.exception(
"Checkpoint loading failed, starting with empty grid")
else:
# select most costly points
bench_mask = grid['k_modes'] == max(self.k_modes_vals)
grid = grid[bench_mask][:self.benchmark_trials]
# submit tasks for every grid point
start_time = time.time()
result_refs = launch_grid(
grid,
left_cube,
right_cube,
model_inputs,
self.ring_exclude_px,
self.use_gpu_decomposition,
self.use_gpu_fit,
generate_options=generate_options,
decompose_options=decompose_options,
evaluate_options=evaluate_options,
measure_ram=measure_ram,
efficient_decomp_reuse=self.efficient_decomp_reuse,
resel_px=self.resel_px,
coverage_mask=covered_pix_mask,
)
if self.benchmark:
log.debug(f"Running {self.benchmark_trials} trials...")
for i in range(self.benchmark_trials):
print(ray.get(result_refs[i]))
ray.shutdown()
return 0
# Wait for results, checkpointing as we go
pending = result_refs
total = len(grid)
restored_from_checkpoint = np.count_nonzero(
grid['time_total_sec'] != 0)
def make_hdul(grid):
return iofits.DaskHDUList([
iofits.DaskHDU(None, kind="primary"),
iofits.DaskHDU(grid, kind="bintable", name="grid")
])
with tqdm(total=total) as pbar:
pbar.update(restored_from_checkpoint)
while pending:
complete, pending = ray.wait(pending,
timeout=5,
num_returns=min(
self.checkpoint_every_x,
len(pending)))
for (idx, result) in ray.get(complete):
grid[idx] = result
if len(complete):
pbar.update(len(complete))
if self.checkpoint is not None:
iofits.write_fits(make_hdul(grid),
self.checkpoint,
overwrite=True)
iofits.write_fits(make_hdul(grid), output_filepath, overwrite=True)
ray.shutdown()
class RemoteRayConfig(CommonRayConfig):
url: str = xconf.field(help="URL to existing Ray cluster head node")
