# Licensed under a 3-clause BSD style license - see LICENSE.rst
import itertools
import warnings
import numpy as np
from astropy import units as u
from astropy.coordinates import SkyCoord, frame_transform_graph
from astropy.io.fits import Header
from astropy.wcs import WCS
from astropy.wcs.utils import (
celestial_frame_to_wcs,
pixel_to_pixel,
pixel_to_skycoord,
skycoord_to_pixel,
wcs_to_celestial_frame,
)
from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS, SlicedLowLevelWCS
from astropy.wcs.wcsapi.high_level_wcs_wrapper import HighLevelWCSWrapper
from .._array_utils import sample_array_edges
from .._wcs_utils import pixel_scale
from ..utils import parse_input_shape
__all__ = ["find_optimal_celestial_wcs"]
# Note that if this is modified, the docstring should be updated
NEGATIVE_CDELT_CTYPES = ["RA--", "GLON", "ELON", "HLON", "SLON"]
[docs]
def find_optimal_celestial_wcs(
input_data,
hdu_in=None,
frame=None,
auto_rotate=False,
projection="TAN",
resolution=None,
reference=None,
negative_lon_cdelt=None,
):
"""
Given one or more images, return an optimal WCS projection object and
shape.
This currently only works with 2-d images with celestial WCS.
Parameters
----------
input_data : iterable
One or more input data specifications to include in the calculation of
the final WCS. This should be an iterable containing one entry for each
specification, where a single data specification is one of:
* The name of a FITS file as a `str` or a `pathlib.Path` object
* An `~astropy.io.fits.HDUList` object
* An image HDU object such as a `~astropy.io.fits.PrimaryHDU`,
`~astropy.io.fits.ImageHDU`, or `~astropy.io.fits.CompImageHDU`
instance
* A tuple where the first element is an Numpy array shape tuple and
the second element is either a
`~astropy.wcs.wcsapi.BaseLowLevelWCS`,
`~astropy.wcs.wcsapi.BaseHighLevelWCS`, or a
`~astropy.io.fits.Header` object
* A tuple where the first element is a `~numpy.ndarray` and the
second element is either a
`~astropy.wcs.wcsapi.BaseLowLevelWCS`,
`~astropy.wcs.wcsapi.BaseHighLevelWCS`, or a
`~astropy.io.fits.Header` object
* An `~astropy.nddata.NDData` object from which the ``.data`` and
``.wcs`` attributes will be used as the input data.
* A `~astropy.wcs.wcsapi.BaseLowLevelWCS` object with ``array_shape`` set
or a `~astropy.wcs.wcsapi.BaseHighLevelWCS` object whose
underlying low level WCS object has ``array_shape`` set.
If only one input data needs to be provided, it is also possible to
pass it in without including it in an iterable.
hdu_in : int or str, optional
If ``input_data`` is a FITS file or an `~astropy.io.fits.HDUList`
instance, specifies the HDU to use.
frame : str or `~astropy.coordinates.BaseCoordinateFrame`
The coordinate system for the final image (defaults to the frame of
the first image specified).
auto_rotate : bool
Whether to rotate the header to minimize the final image area (if
`True`, requires shapely>=1.6 to be installed).
projection : str
Three-letter code for the WCS projection.
resolution : `~astropy.units.Quantity`
The resolution of the final image. If not specified, this is the
smallest resolution of the input images.
reference : `~astropy.coordinates.SkyCoord`
The reference coordinate for the final header. If not specified, this
is determined automatically from the input images.
negative_lon_cdelt : bool or str, optional
Whether the CDELT value for the longitude coordinate should be negative
(`True`) or positive (`False`), or determined automatically (``'auto'``).
For astronomical observations of the sky CDELT is usually negative,
while for coordinate systems used in solar physics this is usually
positive. If this is ``'auto'``, the value will be `True` if the
first four characters for CTYPE for the longitude is ``RA--``,
``GLON``, ``ELON``, ``HLON``, or ``SLON``, and `False` otherwise.
The default is currently ``True``, and will become ``'auto'`` in
future.
Returns
-------
wcs : :class:`~astropy.wcs.WCS`
The optimal WCS determined from the input images.
shape : tuple
The optimal shape required to cover all the output.
"""
# TODO: support higher-dimensional datasets in future
# TODO: take into account NaN values when determining the extent of the
# final WCS
if isinstance(frame, str):
frame = frame_transform_graph.lookup_name(frame)()
# Determine whether an iterable of input values was given or a single
# input data.
if isinstance(input_data, str):
# Handle this explicitly as str is iterable too
iterable = False
elif np.iterable(input_data):
if len(input_data) == 2 and isinstance(
input_data[1], BaseLowLevelWCS | BaseHighLevelWCS | Header
):
# Since 2-element tuples are valid single inputs we need to check for this
iterable = False
else:
iterable = True
else:
iterable = False
if iterable:
input_shapes = [parse_input_shape(shape, hdu_in=hdu_in) for shape in input_data]
else:
input_shapes = [parse_input_shape(input_data, hdu_in=hdu_in)]
# We start off by looping over images, checking that they are indeed
# celestial images, and building up a list of all corners and all reference
# coordinates in the frame of reference of the first image.
corners = []
references = []
resolutions = []
for shape, wcs in input_shapes:
if len(shape) > wcs.pixel_n_dim:
shape = shape[-wcs.pixel_n_dim :]
if len(shape) != 2:
raise ValueError(f"Input data is not 2-dimensional (got shape {shape!r})")
if wcs.pixel_n_dim != 2 or wcs.world_n_dim != 2:
raise ValueError("Input WCS is not 2-dimensional")
if isinstance(wcs, WCS):
if not wcs.has_celestial:
raise TypeError("WCS does not have celestial components")
# Determine frame if it wasn't specified
if frame is None:
frame = wcs_to_celestial_frame(wcs)
else:
# Convert a single position to determine type of output and make
# sure there is only a single SkyCoord returned.
coord = wcs.pixel_to_world(0, 0)
if not isinstance(coord, SkyCoord):
raise TypeError("WCS does not have celestial components")
if frame is None:
frame = coord.frame.replicate_without_data()
# Find pixel coordinates of corners. In future if we are worried about
# significant distortions of the edges in the reprojection process we
# could simply add arbitrary numbers of midpoints to this list.
ny, nx = shape
xc = np.array([-0.5, nx - 0.5, nx - 0.5, -0.5])
yc = np.array([-0.5, -0.5, ny - 0.5, ny - 0.5])
# We have to do .frame here to make sure that we get a frame object
# without any 'hidden' attributes, otherwise the stacking below won't
# work.
corners.append(wcs.pixel_to_world(xc, yc).transform_to(frame).frame)
if isinstance(wcs, WCS):
# We now figure out the reference coordinate for the image in the
# frame of the first image. The easiest way to do this is actually
# to use pixel_to_skycoord with the reference position in pixel
# coordinates. We have to set origin=1 because crpix values are
# 1-based.
xp, yp = wcs.wcs.crpix
references.append(pixel_to_skycoord(xp, yp, wcs, origin=1).transform_to(frame).frame)
else:
xp, yp = (nx - 1) / 2, (ny - 1) / 2
references.append(wcs.pixel_to_world(xp, yp).transform_to(frame).frame)
resolutions.append(pixel_scale(wcs, shape))
# We now stack the coordinates - however the frame classes can't do this
# so we have to use the high-level SkyCoord class.
corners = SkyCoord(corners)
references = SkyCoord(references)
# If no reference coordinate has been passed in for the final header, we
# determine the reference coordinate as the mean of all the reference
# positions. This choice is as good as any and if the user really cares,
# they can set it manually.
if reference is None:
reference = SkyCoord(references.data.mean(), frame=references.frame)
# In any case, we need to convert the reference coordinate (either
# specified or automatically determined) to the requested final frame.
reference = reference.transform_to(frame)
# Determine resolution if not specified
if resolution is None:
resolution = np.min(u.Quantity(resolutions))
# Construct WCS object centered on position
wcs_final = celestial_frame_to_wcs(frame, projection=projection)
negative_lon_cdelt_auto = wcs_final.wcs.ctype[0][:4] in NEGATIVE_CDELT_CTYPES
if negative_lon_cdelt == "auto":
negative_lon_cdelt = negative_lon_cdelt_auto
elif negative_lon_cdelt is None:
if not negative_lon_cdelt_auto:
warnings.warn(
"negative_lon_cdelt is not set, and currently defaults to True, "
"but in future will change to 'auto', and for this WCS this will "
"evaluate to False in future. It is recommended that you set "
"negative_lon_cdelt explicitly, either to 'auto', or to True/False.",
DeprecationWarning,
stacklevel=2,
)
negative_lon_cdelt = True
if wcs_final.wcs.cunit[0] == "":
wcs_final.wcs.cunit[0] = "deg"
if wcs_final.wcs.cunit[1] == "":
wcs_final.wcs.cunit[1] = "deg"
rep = reference.represent_as("unitspherical")
wcs_final.wcs.crval = (
rep.lon.to_value(wcs_final.wcs.cunit[0]),
rep.lat.to_value(wcs_final.wcs.cunit[1]),
)
lon_factor = -1 if negative_lon_cdelt else 1
wcs_final.wcs.cdelt = (
lon_factor * resolution.to_value(wcs_final.wcs.cunit[0]),
resolution.to_value(wcs_final.wcs.cunit[1]),
)
# For now, set crpix to (1, 1) and we'll then figure out where all the
# images fall in this projection, then we'll adjust crpix.
wcs_final.wcs.crpix = (1, 1)
# Find pixel coordinates of all corners in the final WCS projection. We use
# origin=1 since we are trying to determine crpix values.
xp, yp = skycoord_to_pixel(corners, wcs_final, origin=1)
if auto_rotate:
# Use shapely to represent the points and find the minimum rotated
# rectangle
from shapely.geometry import MultiPoint
mp = MultiPoint(list(zip(xp, yp, strict=True)))
# The following returns a list of rectangle vertices - in fact there
# are 5 coordinates because shapely represents it as a closed polygon
# with the same first/last vertex.
xr, yr = mp.minimum_rotated_rectangle.exterior.coords.xy
xr, yr = xr[:4], yr[:4]
# The order of the vertices is not guaranteed to be constant so we
# take the vertices with the two smallest y values (which, for a
# rectangle, guarantees that the vertices are neighboring)
order = np.argsort(yr)
x1, y1, x2, y2 = xr[order[0]], yr[order[0]], xr[order[1]], yr[order[1]]
# Determine angle between two of the vertices. It doesn't matter which
# ones they are, we just want to know how far from being straight the
# rectangle is.
angle = np.arctan2(y2 - y1, x2 - x1)
# Determine the smallest angle that would cause the rectangle to be
# lined up with the axes.
angle = angle % (np.pi / 2)
if angle > np.pi / 4:
angle -= np.pi / 2
# Set rotation matrix (use PC instead of CROTA2 since PC is the
# recommended approach)
pc = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
wcs_final.wcs.pc = pc
# Recompute pixel coordinates (more accurate than simply rotating xp, yp)
xp, yp = skycoord_to_pixel(corners, wcs_final, origin=1)
# Find the full range of values
xmin = xp.min()
xmax = xp.max()
ymin = yp.min()
ymax = yp.max()
# Update crpix so that the lower range falls on the bottom and left. We add
# 0.5 because in the final image the bottom left corner should be at (0.5,
# 0.5) not (1, 1).
wcs_final.wcs.crpix = (1 - xmin) + 0.5, (1 - ymin) + 0.5
# Return the final image shape too
naxis1 = int(round(xmax - xmin))
naxis2 = int(round(ymax - ymin))
return wcs_final, (naxis2, naxis1)
def _pixel_to_pixel_list(wcs_in, wcs_out, *inputs):
# pixel_to_pixel returns a bare array rather than a list of per-dimension
# arrays when the output WCS has a single pixel dimension, which would make
# the [::-1] reversals in sample_input_edges_in_output reverse the samples
# instead of the dimensions.
outputs = pixel_to_pixel(wcs_in, wcs_out, *inputs)
if wcs_out.low_level_wcs.pixel_n_dim == 1:
outputs = [outputs]
return outputs
def sample_input_edges_in_output(array_shape, wcs_in, wcs_out, n_samples=11):
"""
Sample the edges of an input array and return their pixel coordinates in the
output WCS, for the reprojected (trailing) dimensions.
This is used to determine the minimal region of the output that an input
image covers. If the input WCS has more pixel dimensions than the output WCS
(for example when using ``non_reprojected_dims`` to reproject a cube into a
celestial-only output), the input WCS is sliced down to its reprojected
(trailing) dimensions before being related to the output, since
``pixel_to_pixel`` requires the two WCS to describe the same number of world
coordinates. Because the reprojected WCS may vary along the non-reprojected
axes (for example a drifting pointing, possibly non-linear), the input WCS is
sliced at ``n_samples`` positions along each of those axes and the resulting
footprints are combined; variation between the sampled positions is not
captured. Axes that do not affect the reprojected world coordinates
(according to the axis correlation matrix) are sliced at a single position.
Parameters
----------
array_shape : tuple
The shape of the input array.
wcs_in : `~astropy.wcs.wcsapi.BaseHighLevelWCS`
The WCS of the input array.
wcs_out : `~astropy.wcs.wcsapi.BaseHighLevelWCS`
The WCS of the output array.
n_samples : int, optional
The number of samples to take along each edge.
Returns
-------
list of `~numpy.ndarray`
The output pixel coordinates of the sampled edges, in array (numpy)
dimension order.
"""
n_extra_in = wcs_in.low_level_wcs.pixel_n_dim - wcs_out.low_level_wcs.pixel_n_dim
if n_extra_in <= 0:
edges = sample_array_edges(
array_shape[-wcs_in.low_level_wcs.pixel_n_dim :], n_samples=n_samples
)[::-1]
return _pixel_to_pixel_list(wcs_in, wcs_out, *edges)[::-1]
n_reproject = wcs_out.low_level_wcs.pixel_n_dim
edges = sample_array_edges(array_shape[-n_reproject:], n_samples=n_samples)[::-1]
# Trim the array shape to the dimensions the input WCS describes before
# taking the non-reprojected leading sizes, since the array may have extra
# leading broadcast dimensions beyond the WCS (as in the branch above).
leading_shape = array_shape[-wcs_in.low_level_wcs.pixel_n_dim : -n_reproject]
# Sample positions along each non-reprojected axis (not just its end points)
# so that non-linear variation of the reprojected WCS along that axis is
# captured, provided the variation is smooth on the scale of the sample
# spacing. Axes that the correlation matrix shows do not affect the
# reprojected world coordinates are sampled at a single position, since
# every slice along them gives the same footprint. Use integer pixel
# indices and de-duplicate for short axes.
matrix = wcs_in.low_level_wcs.axis_correlation_matrix
world_reprojected = matrix[:, :n_reproject].any(axis=1)
leading_samples = []
for iaxis, size in enumerate(leading_shape):
pixel_axis = wcs_in.low_level_wcs.pixel_n_dim - 1 - iaxis
if matrix[world_reprojected, pixel_axis].any():
samples = sorted({int(round(idx)) for idx in np.linspace(0, size - 1, n_samples)})
else:
samples = [0]
leading_samples.append(samples)
edges_out_corners = []
for corner in itertools.product(*leading_samples):
slices = list(corner) + [slice(None)] * n_reproject
wcs_in_reproject = HighLevelWCSWrapper(
SlicedLowLevelWCS(wcs_in.low_level_wcs, slices=slices)
)
edges_out_corners.append(_pixel_to_pixel_list(wcs_in_reproject, wcs_out, *edges)[::-1])
return [
np.concatenate([corner[idim] for corner in edges_out_corners])
for idim in range(n_reproject)
]
