From a6ffe26ea039bf88149963d3bd60e0c683343dc5 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Wed, 2 Oct 2024 00:47:41 -0400 Subject: [PATCH] refactor --- src/stcal/resample/__init__.py | 3 +- src/stcal/resample/resample.py | 31 +++++++++++------- src/stcal/resample/utils.py | 60 +++++++++++++++++++++++++++------- 3 files changed, 69 insertions(+), 25 deletions(-) diff --git a/src/stcal/resample/__init__.py b/src/stcal/resample/__init__.py index 2baa9834a..991b8abbc 100644 --- a/src/stcal/resample/__init__.py +++ b/src/stcal/resample/__init__.py @@ -1,8 +1,9 @@ from .resample import * __all__ = [ - "LibModelAccess", + "LibModelAccessBase", "OutputTooLargeError", "Resample", "resampled_wcs_from_models", + "UnsupportedWCSError", ] diff --git a/src/stcal/resample/resample.py b/src/stcal/resample/resample.py index bffb13020..88644dbcd 100644 --- a/src/stcal/resample/resample.py +++ b/src/stcal/resample/resample.py @@ -18,8 +18,8 @@ interpret_bit_flags, ) -from .utils import bytes2human, get_tmeasure -from ..alignment.util import ( +from stcal.resample.utils import bytes2human, get_tmeasure +from stcal.alignment.util import ( compute_scale, wcs_bbox_from_shape, wcs_from_footprints, @@ -30,10 +30,11 @@ log.setLevel(logging.DEBUG) __all__ = [ - "LibModelAccess", + "LibModelAccessBase", "OutputTooLargeError", "Resample", "resampled_wcs_from_models", + "UnsupportedWCSError", ] _SUPPORTED_CUSTOM_WCS_PARS = [ @@ -62,7 +63,7 @@ def _resample_range(data_shape, bbox=None): return xmin, xmax, ymin, ymax -class LibModelAccess(abc.ABC): +class LibModelAccessBase(abc.ABC): # list of model attributes needed by this module. While this is not # required, it is helpful for subclasses to check they know how to # access these attributes. @@ -131,8 +132,8 @@ def resampled_wcs_from_models( Parameters ---------- - input_models : LibModelAccess - An object of `LibModelAccess`-derived type. + input_models : LibModelAccessBase + An object of `LibModelAccessBase`-derived type. pixel_scale_ratio : float, optional Desired pixel scale ratio defined as the ratio of the first model's @@ -241,6 +242,12 @@ class OutputTooLargeError(RuntimeError): """Raised when the output is too large for in-memory instantiation""" +class UnsupportedWCSError(RuntimeError): + """ Raised when provided output WCS has an unexpected number of axes + or has an unsupported structure. + """ + + class Resample: """ This is the controlling routine for the resampling process. @@ -282,9 +289,9 @@ def __init__(self, input_models, pixfrac=1.0, kernel="square", """ Parameters ---------- - input_models : LibModelAccess - A `LibModelAccess` object allowing iterating over all contained - models of interest. + input_models : LibModelAccessBase + A `LibModelAccessBase`-based object allowing iterating over + all contained models of interest. kwargs : dict Other parameters. @@ -414,8 +421,8 @@ def check_output_wcs(self, output_wcs, wcs_pars, """ naxes = output_wcs.output_frame.naxes if naxes != 2: - raise RuntimeError( - "Output WCS needs 2 spatial axes but the " + raise UnsupportedWCSError( + "Output WCS needs 2 coordinate axes but the " f"supplied WCS has {naxes} axes." ) @@ -841,7 +848,7 @@ def init_resample_variance(self): 'fillval': np.nan, 'out_shape': self._output_array_shape, # 'exptime': 1.0, - 'no_ctx': True, + 'disable_ctx': True, } self.driz_rnoise = Drizzle(**driz_init_kwargs) self.driz_poisson = Drizzle(**driz_init_kwargs) diff --git a/src/stcal/resample/utils.py b/src/stcal/resample/utils.py index 71d32b123..976010b06 100644 --- a/src/stcal/resample/utils.py +++ b/src/stcal/resample/utils.py @@ -1,26 +1,62 @@ -import os -from pathlib import Path, PurePath +from copy import deepcopy +import asdf import numpy as np from astropy.nddata.bitmask import interpret_bit_flags __all__ = [ - "build_mask", "build_output_model_name", "get_tmeasure", "bytes2human" + "build_mask", "get_tmeasure", "bytes2human", "load_custom_wcs" ] -def build_output_model_name(input_filename_list): - fnames = {f for f in input_filename_list if f is not None} +def load_custom_wcs(asdf_wcs_file, output_shape=None): + """ Load a custom output WCS from an ASDF file. - if not fnames: - return "resampled_data_{resample_suffix}{resample_file_ext}" + Parameters + ---------- + asdf_wcs_file : str + Path to an ASDF file containing a GWCS structure. + + output_shape : tuple of int, optional + Array shape (in ``[x, y]`` order) for the output data. If not provided, + the custom WCS must specify one of: pixel_shape, + array_shape, or bounding_box. + + Returns + ------- + wcs : WCS + The output WCS to resample into. - # TODO: maybe remove ending suffix for single file names? - prefix = os.path.commonprefix( - [PurePath(f).stem.strip('_- ') for f in fnames] - ) + """ + if not asdf_wcs_file: + return None + + with asdf.open(asdf_wcs_file) as af: + wcs = deepcopy(af.tree["wcs"]) + wcs.pixel_area = af.tree.get("pixel_area", None) + wcs.pixel_shape = af.tree.get("pixel_shape", None) + wcs.array_shape = af.tree.get("array_shape", None) + + if output_shape is not None: + wcs.array_shape = output_shape[::-1] + wcs.pixel_shape = output_shape + elif wcs.pixel_shape is not None: + wcs.array_shape = wcs.pixel_shape[::-1] + elif wcs.array_shape is not None: + wcs.pixel_shape = wcs.array_shape[::-1] + elif wcs.bounding_box is not None: + wcs.array_shape = tuple( + int(axs[1] + 0.5) + for axs in wcs.bounding_box.bounding_box(order="C") + ) + else: + raise ValueError( + "Step argument 'output_shape' is required when custom WCS " + "does not have neither of 'array_shape', 'pixel_shape', or " + "'bounding_box' attributes set." + ) - return prefix + "{resample_suffix}{resample_file_ext}" + return wcs def build_mask(dqarr, bitvalue, flag_name_map=None):