Allow single inputs to find_optimal_celestial_wcs and add ability to specify HDU

docs/installation.rst

@@ -9,25 +9,20 @@ Requirements
 
 This package has the following hard run time dependencies:
 
-* `Python <http://www.python.org/>`__ 3.7 or later
+* `Python <http://www.python.org/>`__ 3.8 or later
 
-* `Numpy <http://www.numpy.org/>`__ 1.14 or later
+* `Numpy <http://www.numpy.org/>`__ 1.20 or later
 
-* `Astropy <http://www.astropy.org/>`__ 3.2 or later
+* `Astropy <http://www.astropy.org/>`__ 5.0 or later
 
-* `Scipy <http://www.scipy.org/>`__ 1.1 or later
+* `Scipy <http://www.scipy.org/>`__ 1.5 or later
 
 * `astropy-healpix <https://astropy-healpix.readthedocs.io>`_ 0.6 or later for HEALPIX image reprojection
 
 and the following optional dependencies:
 
 * `shapely <https://toblerity.org/shapely/project.html>`_ 1.6 or later for some of the mosaicking functionality
 
-If you build the package from the source, the following additional packages
-are required:
-
-* `Cython <http://cython.org>`__
-
 and to run the tests, you will also need:
 
 * `Matplotlib <http://matplotlib.org/>`__

reproject/adaptive/high_level.py

@@ -41,6 +41,9 @@ def reproject_adaptive(
             * 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
+              the second element is either a `~astropy.wcs.WCS` 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.WCS` or a
               `~astropy.io.fits.Header` object

reproject/conftest.py

@@ -5,6 +5,12 @@
 
 import os
 
+import numpy as np
+import pytest
+from astropy.io import fits
+from astropy.nddata import NDData
+from astropy.wcs import WCS
+
 try:
     from pytest_astropy_header.display import PYTEST_HEADER_MODULES, TESTED_VERSIONS
 
@@ -29,3 +35,93 @@ def pytest_configure(config):
         from reproject import __version__
 
         TESTED_VERSIONS["reproject"] = __version__
+
+
+def valid_celestial_input(tmp_path, request):
+    array = np.ones((30, 40))
+
+    wcs = WCS(naxis=2)
+    wcs.wcs.ctype = "RA---TAN", "DEC--TAN"
+    wcs.wcs.crpix = (1, 2)
+    wcs.wcs.crval = (30, 40)
+    wcs.wcs.cdelt = (-0.05, 0.04)
+    wcs.wcs.equinox = 2000.0
+
+    hdulist = fits.HDUList(
+        [
+            fits.PrimaryHDU(array, wcs.to_header()),
+            fits.ImageHDU(array, wcs.to_header()),
+            fits.CompImageHDU(array, wcs.to_header()),
+        ]
+    )
+
+    kwargs = {}
+
+    if request.param in ["filename", "path"]:
+        input_value = tmp_path / "test.fits"
+        if request.param == "filename":
+            input_value = str(input_value)
+        hdulist.writeto(input_value)
+        kwargs["hdu_in"] = 0
+    elif request.param == "hdulist":
+        input_value = hdulist
+        kwargs["hdu_in"] = 1
+    elif request.param == "primary_hdu":
+        input_value = hdulist[0]
+    elif request.param == "image_hdu":
+        input_value = hdulist[1]
+    elif request.param == "comp_image_hdu":
+        input_value = hdulist[2]
+    elif request.param == "ape14_wcs":
+        input_value = wcs
+        input_value._naxis = list(array.shape[::-1])
+    elif request.param == "shape_wcs_tuple":
+        input_value = (array.shape, wcs)
+    elif request.param == "data_wcs_tuple":
+        input_value = (array, wcs)
+    elif request.param == "nddata":
+        input_value = NDData(data=array, wcs=wcs)
+    elif request.param == "ape14_wcs":
+        input_value = wcs
+        input_value._naxis = list(array.shape[::-1])
+    elif request.param == "shape_wcs_tuple":
+        input_value = (array.shape, wcs)
+
+    else:
+        raise ValueError(f"Unknown mode: {request.param}")
+
+    return array, wcs, input_value, kwargs
+
+
+@pytest.fixture(
+    params=[
+        "filename",
+        "path",
+        "hdulist",
+        "primary_hdu",
+        "image_hdu",
+        "comp_image_hdu",
+        "data_wcs_tuple",
+        "nddata",
+    ]
+)
+def valid_celestial_input_data(tmp_path, request):
+    return valid_celestial_input(tmp_path, request)
+
+
+@pytest.fixture(
+    params=[
+        "filename",
+        "path",
+        "hdulist",
+        "primary_hdu",
+        "image_hdu",
+        "comp_image_hdu",
+        "data_wcs_tuple",
+        "nddata",
+        "ape14_wcs",
+        "shape_wcs_tuple",
+    ]
+)
+def valid_celestial_input_shapes(tmp_path, request):
+    return valid_celestial_input(tmp_path, request)

reproject/interpolation/high_level.py

@@ -44,6 +44,9 @@ def reproject_interp(
             * 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
+              the second element is either a `~astropy.wcs.WCS` 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.WCS` or a
               `~astropy.io.fits.Header` object

reproject/mosaicking/tests/test_wcs_helpers.py

@@ -4,10 +4,14 @@
 import pytest
 from astropy import units as u
 from astropy.coordinates import FK5, Galactic, SkyCoord
+from astropy.io import fits
+from astropy.nddata import NDData
 from astropy.wcs import WCS
 from astropy.wcs.wcsapi import HighLevelWCSWrapper
 from numpy.testing import assert_allclose, assert_equal
 
+from reproject.tests.helpers import assert_header_allclose
+
 from ..wcs_helpers import find_optimal_celestial_wcs
 
 try:
@@ -223,3 +227,26 @@ def test_args_tuple_header(self):
     frame_projection_expected_shape = 46, 50
     auto_rotate_expected_crpix = 20.520875, 15.503349
     multiple_size_expected_crpix = 27.279739, 17.29016
+
+
+@pytest.mark.parametrize("iterable", [False, True])
+def test_input_types(valid_celestial_input_shapes, iterable):
+    # Test different kinds of inputs and check the result is always the same
+
+    array, wcs, input_value, kwargs = valid_celestial_input_shapes
+
+    wcs_ref, shape_ref = find_optimal_celestial_wcs([(array, wcs)], frame=FK5())
+
+    if iterable:
+        input_value = [input_value]
+
+    wcs_test, shape_test = find_optimal_celestial_wcs(input_value, frame=FK5(), **kwargs)
+    assert_header_allclose(wcs_test.to_header(), wcs_ref.to_header())
+    assert shape_test == shape_ref
+
+    if isinstance(input_value, fits.HDUList) and not iterable:
+        # Also check case of not passing hdu_in and having all HDUs being included
+
+        wcs_test, shape_test = find_optimal_celestial_wcs(input_value, frame=FK5())
+        assert_header_allclose(wcs_test.to_header(), wcs_ref.to_header())
+        assert shape_test == shape_ref

reproject/mosaicking/wcs_helpers.py

@@ -3,6 +3,8 @@
 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.utils import isiterable
 from astropy.wcs import WCS
 from astropy.wcs.utils import (
     celestial_frame_to_wcs,
@@ -18,7 +20,13 @@
 
 
 def find_optimal_celestial_wcs(
-    input_data, frame=None, auto_rotate=False, projection="TAN", resolution=None, reference=None
+    input_data,
+    hdu_in=None,
+    frame=None,
+    auto_rotate=False,
+    projection="TAN",
+    resolution=None,
+    reference=None,
 ):
     """
     Given one or more images, return an optimal WCS projection object and
@@ -44,7 +52,15 @@ def find_optimal_celestial_wcs(
             * A tuple where the first element is a `~numpy.ndarray` and the
               second element is either a `~astropy.wcs.WCS` 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.
 
+        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)
@@ -75,7 +91,25 @@ def find_optimal_celestial_wcs(
     if isinstance(frame, str):
         frame = frame_transform_graph.lookup_name(frame)()
 
-    input_shapes = [parse_input_shape(shape) for shape in input_data]
+    # Determine whether an iterable of input values was given or a single
+    # input data.
+
+    if isinstance(input_data, str):
+        # Handle this explicitly as isiterable(str) is True
+        iterable = False
+    elif isiterable(input_data):
+        if len(input_data) == 2 and isinstance(input_data[1], (WCS, 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

reproject/spherical_intersect/high_level.py

@@ -24,6 +24,9 @@ def reproject_exact(
             * 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
+              the second element is either a `~astropy.wcs.WCS` 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.WCS` or a
               `~astropy.io.fits.Header` object

reproject/tests/helpers.py

@@ -1,8 +1,19 @@
 from astropy.io import fits
+from numpy.testing import assert_allclose
 
 
 def array_footprint_to_hdulist(array, footprint, header):
     hdulist = fits.HDUList()
     hdulist.append(fits.PrimaryHDU(array, header))
     hdulist.append(fits.ImageHDU(footprint, header, name="footprint"))
     return hdulist
+
+
+def assert_header_allclose(header1, header2, **kwargs):
+    assert sorted(header1) == sorted(header2)
+
+    for key1, value1 in header1.items():
+        if isinstance(value1, str):
+            assert value1 == header2[key1]
+        else:
+            assert_allclose(value1, header2[key1], **kwargs)