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")