Source code for

from numbers import Integral

import numpy as np
from astropy.wcs import WCS
from astropy.wcs.wcsapi import SlicedLowLevelWCS
from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper

from ndcube.wcs.wrappers import ResampledLowLevelWCS

__all__ = ["unwrap_wcs_to_fitswcs"]

[docs] def unwrap_wcs_to_fitswcs(wcs): """ Create FITS-WCS equivalent to (nested) WCS wrapper object. Underlying WCS must be FITS-WCS. No axes are dropped from original FITS-WCS, even if sliced by an integer. Instead, integer-sliced axes is sliced to length-1 and marked True in the ``dropped_data_axes`` output. Currently supported wrapper classes include `astropy.wcs.wcsapi.SlicedLowLevelWCS` and `ndcube.wcs.wrappers.ResampledLowLevelWCS`. Parameters ---------- wcs: `~astropy.wcs.wcsapi.BaseWCSWrapper` The WCS Wrapper object. Base level WCS implementation must be FITS-WCS. Returns ------- fitswcs: `astropy.wcs.WCS` The equivalent FITS-WCS object. dropped_data_axes: 1-D `numpy.ndarray` Denotes which axes must have been dropped from the data array by slicing wrappers. Axes are in array/numpy order, reversed compared to WCS. """ # If wcs is already a FITS-WCS, return it. low_level_wrapper = wcs.low_level_wcs if hasattr(wcs, "low_level_wcs") else wcs if isinstance(low_level_wrapper, WCS): return low_level_wrapper, np.zeros(low_level_wrapper.naxis, dtype=bool) # Determine chain of wrappers down to the FITS-WCS. wrapper_chain = [] while isinstance(low_level_wrapper, BaseWCSWrapper): wrapper_chain.append(low_level_wrapper) low_level_wrapper = low_level_wrapper._wcs if hasattr(low_level_wrapper, "low_level_wcs"): low_level_wrapper = low_level_wrapper.low_level_wcs if not isinstance(low_level_wrapper, WCS): raise TypeError(f"Base-level WCS must be type {type(WCS)}. Found: {type(low_level_wcs)}") fitswcs = low_level_wrapper dropped_data_axes = np.zeros(fitswcs.naxis, dtype=bool) # Unwrap each wrapper in reverse order and edit fitswcs. for low_level_wrapper in wrapper_chain[::-1]: if isinstance(low_level_wrapper, SlicedLowLevelWCS): slice_items = np.array([slice(None)] * fitswcs.naxis) slice_items[dropped_data_axes == False] = low_level_wrapper._slices_array # numpy order fitswcs, dda = _slice_fitswcs(fitswcs, slice_items, numpy_order=True) dropped_data_axes[dda] = True elif isinstance(low_level_wrapper, ResampledLowLevelWCS): factor = np.ones(fitswcs.naxis) offset = np.zeros(fitswcs.naxis) kept_wcs_axes = dropped_data_axes[::-1] == False # WCS-order factor[kept_wcs_axes] = low_level_wrapper._factor offset[kept_wcs_axes] = low_level_wrapper._offset fitswcs = _resample_fitswcs(fitswcs, factor, offset) else: raise TypeError("Unrecognized/unsupported WCS Wrapper type: {type(low_level_wrapper)}") return fitswcs, dropped_data_axes
def _slice_fitswcs(fitswcs, slice_items, numpy_order=True, shape=None): """ Slice a FITS-WCS. If an `int` is given in ``slice_items``, the corresponding axis is not dropped. But the new 0th pixel will correspond the index given by the `int` in the original WCS. Parameters ---------- fitswcs: `astropy.wcs.WCS` The FITS-WCS object to be sliced. slice_items: iterable of `slice` objects or `int` The slices to by applied to each axis. If an `int` is provided, the axis is sliced to length-1, but not dropped. However, its corresponding entry in the ``dropped_data_axes`` output is marked True. numpy_order: `bool` If True, slices in ``slice_items`` are in array/numpy order, which is reversed compared to the WCS order. shape: sequence of `int`, optional The length of each axis. Only used if negative indices are supplied in ``slice_items``. If not supplied, set to ``fitswcs._naxis``. Order defined by numpy_order kwarg. Returns ------- sliced_wcs: `astropy.wcs.WCS` The sliced FITS-WCS. dropped_data_axes: 1-D `numpy.ndarray` Denotes which axes must have been dropped from the data array by slicing wrappers. Order of axes (numpy or WCS) is dictated by ``numpy_order`` kwarg. """ def negative_index_error_msg(x): return ( f"Negative indexing not supported as {x}th axis length is 0 in " "underlying FITS-WCS. Supply axes lengths via shape kwarg.") naxis = fitswcs.naxis dropped_data_axes = np.zeros(naxis, dtype=bool) # Sanitize inputs if shape is None: shape = fitswcs._naxis if numpy_order: shape = shape[::-1] else: if len(shape) != naxis: raise ValueError("shape kwarg must be same length as number of pixel axes " f"in FITS-WCS, i.e. {naxis}") if not all(isinstance(s, Integral) for s in shape): raise TypeError("All elements of ``shape`` must be integers. " f"shapes types = {[type(s) for s in shape]}") slice_items = list(slice_items) for i, (item, len_axis) in enumerate(zip(slice_items, shape)): if isinstance(item, Integral): # Mark axis corresponding to int item as dropped from data array. dropped_data_axes[i] = True # Convert negative indices to positive equivalent. if item < 0: if len_axis == 0: raise ValueError(negative_index_error_msg(i)) item = len_axis + item # Convert int item to slice so a FITS-WCS is returned after slicing. slice_items[i] = slice(item, item + 1) elif isinstance(item, slice): # Convert negative indices inside slice item to positive equivalent. start_neg = item.start is not None and item.start < 0 stop_neg = item.stop is not None and item.stop < 0 if start_neg or stop_neg: if len_axis == 0: raise ValueError(negative_index_error_msg(i)) start = len_axis + item.start if start_neg else item.start stop = len_axis + item.stop if stop_neg else item.stop slice_items[i] = slice(start, stop, item.step) else: raise TypeError("All slice_items must be a slice or an int. " f"type(slice_items[{i}]) = {type(slice_items[i])}") # Slice WCS sliced_wcs = fitswcs.slice(slice_items, numpy_order=numpy_order) return sliced_wcs, dropped_data_axes def _resample_fitswcs(fitswcs, factor, offset=0): """ Resample the plate scale of a FITS-WCS by a given factor. ``factor`` and ``offset`` inputs are in pixel order. Parameters ---------- fitswcs: `astropy.wcs.WCS` The FITS-WCS object to be resampled. factor: 1-D array-like or scalar The factor by which the FITS-WCS is resampled. Must be same length as number of axes in ``fitswcs``. If scalar, the same factor is applied to all axes. Factors must be given in WCS-order (opposite to data axes order). offset: 1-D array-like or scalar The location on the initial pixel grid which corresponds to zero on the resampled pixel grid. If scalar, the same offset is applied to all axes. Offsets must be given in WCS-order (opposite to data axes order). Returns ------- resampled_wcs: `astropy.wcs.WCS` The resampled FITS-WCS. """ # Sanitize inputs. factor = np.asarray(factor) if len(factor) != fitswcs.naxis: raise ValueError(f"Length of factor must equal number of dimensions {fitswcs.naxis}.") offset = np.asarray(offset) if len(offset) != fitswcs.naxis: raise ValueError(f"Length of offset must equal number of dimensions {fitswcs.naxis}.") # Scale plate scale and shift by offset. fitswcs.wcs.cdelt *= factor fitswcs.wcs.crpix = (fitswcs.wcs.crpix + offset) / factor fitswcs._naxis = list(np.round(np.array(fitswcs._naxis) / factor).astype(int)) return fitswcs