Merge pull request #167 from nden/footprint

generalize the footprint calculation to more than 2 axes

CHANGES.rst

@@ -26,8 +26,10 @@ New Features
 
 - Add a ``StokesFrame`` which converts from 'I', 'Q', 'U', 'V' to 0-3. [#133]
 
-- Support serialising the base ``CoordinateFrame`` class to asdf, by making
-  a specific tag and schema for ``Frame2D`` [#150]
+- Support serializing the base ``CoordinateFrame`` class to asdf, by making
+  a specific tag and schema for ``Frame2D``. [#150]
+
+- Generalized the footrpint calculation to all output axes. [#167]
 
 
 API Changes
@@ -36,6 +38,8 @@ API Changes
 - The argument ``output="numerical_plus"`` was replaced by a bool
   argument ``with_units``. [#156]
 
+- Added a new flag ``axis_type`` to the footprint method. It controls what
+  type of footprint to calculate. [#167]
 
 Bug Fixes
 ^^^^^^^^^

gwcs/tests/test_wcs.py

@@ -1,6 +1,6 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 import numpy as np
-from numpy.testing import assert_allclose
+from numpy.testing import assert_allclose, assert_equal
 from astropy.modeling import models
 from astropy import coordinates as coord
 from astropy.io import fits
@@ -295,6 +295,36 @@ def test_available_frames():
     assert w.available_frames == ['detector', 'focal', 'icrs']
 
 
+def test_footprint():
+    icrs = cf.CelestialFrame(name='icrs', reference_frame=coord.ICRS(),
+                             axes_order=(0, 1))
+    spec = cf.SpectralFrame(name='freq', unit=[u.Hz, ], axes_order=(2, ))
+    world = cf.CompositeFrame([icrs, spec])
+    transform = (models.Shift(10) & models.Shift(-1)) & models.Scale(2)
+    pipe = [('det', transform), (world, None)]
+    w = wcs.WCS(pipe)
+
+    with pytest.raises(TypeError):
+        w.footprint()
+
+    w.bounding_box = ((1,5), (1,3), (1, 6))
+
+    assert_equal(w.footprint(), np.array([[11, 0, 2],
+                                          [11, 0, 12],
+                                          [11, 2, 2],
+                                          [11, 2, 12],
+                                          [15, 0, 2],
+                                          [15, 0, 12],
+                                          [15, 2, 2],
+                                          [15, 2, 12]]))
+    assert_equal(w.footprint(axis_type='spatial'), np.array([[ 11.,   0.],
+                                                             [ 11.,   2.],
+                                                             [ 15.,   2.],
+                                                             [ 15.,   0.]]))
+
+    assert_equal(w.footprint(axis_type='spectral'), np.array([2, 12]))
+
+
 class TestImaging(object):
     def setup_class(self):
         hdr = fits.Header.fromtextfile(get_pkg_data_filename("data/acs.hdr"), endcard=False)
@@ -368,7 +398,7 @@ def test_backward(self):
 
     def test_footprint(self):
         bb = ((1, 4096), (1, 2048))
-        footprint = (self.wcs.footprint(bb)).T
+        footprint = (self.wcs.footprint(bb))
         fits_footprint = self.fitsw.calc_footprint(axes=(4096, 2048))
         assert_allclose(footprint, fits_footprint)
 

gwcs/wcs.py

@@ -1,6 +1,6 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 import functools
-
+import itertools
 import numpy as np
 from astropy.modeling.core import Model
 from astropy.modeling import utils as mutils
@@ -519,31 +519,65 @@ def __repr__(self):
             self.output_frame, self.input_frame, self.forward_transform)
         return fmt
 
-    def footprint(self, bounding_box=None, center=False):
+    def footprint(self, bounding_box=None, center=False, axis_type="all"):
         """
-        Return the footprint of the observation in world coordinates.
+        Return the footprint in world coordinates.
 
         Parameters
         ----------
         bounding_box : tuple of floats: (start, stop)
             `prop: bounding_box`
         center : bool
             If `True` use the center of the pixel, otherwise use the corner.
+        axis_type : str
+            A supported ``output_frame.axes_type`` or "all" (default).
+            One of ['spatial', 'spectral', 'temporal'] or a custom type.
 
         Returns
         -------
-        coord : array of coordinates in the output_frame.
-            The order is counter-clockwise starting with the bottom left corner.
+        coord : ndarray
+            Array of coordinates in the output_frame mapping
+            corners to the output frame. For spatial coordinates the order
+            is clockwise, starting from the bottom left corner.
+
         """
+        def _order_clockwise(v):
+            return np.asarray([[v[0][0], v[1][0]], [v[0][0], v[1][1]],
+                               [v[0][1], v[1][1]], [v[0][1], v[1][0]]]).T
+
         if bounding_box is None:
             if self.bounding_box is None:
                 raise TypeError("Need a valid bounding_box to compute the footprint.")
             bb = self.bounding_box
         else:
             bb = bounding_box
-        vertices = np.asarray([[bb[0][0], bb[1][0]], [bb[0][0], bb[1][1]],
-                               [bb[0][1], bb[1][1]], [bb[0][1], bb[1][0]]]).T
+
+        all_spatial = all([t.lower() == "spatial" for t in self.output_frame.axes_type])
+
+        if all_spatial:
+            vertices = _order_clockwise(bb)
+        else:
+            vertices = np.array(list(itertools.product(*bb))).T
+
         if center:
             vertices = _toindex(vertices)
-        result = self.__call__(*vertices, **{'with_bounding_box': False})
-        return np.asarray(result)
+
+        result = np.asarray(self.__call__(*vertices, **{'with_bounding_box': False}))
+
+        axis_type = axis_type.lower()
+        if axis_type == 'spatial' and all_spatial:
+            return result.T
+
+        if axis_type != "all":
+            axtyp_ind = np.array([t.lower() for t in self.output_frame.axes_type]) == axis_type
+            if not axtyp_ind.any():
+                raise ValueError('This WCS does not have axis of type "{}".'.format(axis_type))
+            result = np.asarray([(r.min(), r.max()) for r in result[axtyp_ind]])
+
+            if axis_type == "spatial":
+                result = _order_clockwise(result)
+            else:
+                result.sort()
+                result = np.squeeze(result)
+
+        return result.T