Merge pull request #375 from WilliamJamieson/feature/compound_boundin…

…g_box

Enable Compound Bounding Box support

CHANGES.rst

@@ -6,6 +6,8 @@ Bug Fixes
 - Updated code in ``region.py`` with latest improvements and bug fixes
   from ``stsci.skypac.regions.py`` [#382]
 
+- Enabled ``CompoundBoundingBox`` support for wcs. [#375]
+
 
 New Features
 ^^^^^^^^^^^^
@@ -52,7 +54,7 @@ API Changes
 - Modified interface to ``wcs_from_points`` function to better match analogous function
   in astropy. [#349]
 
-- ``Model._BoundingBox`` was renamed to ``Model.ModelBoundingBox`. [#376, #377]
+- ``Model._BoundingBox`` was renamed to ``Model.ModelBoundingBox``. [#376, #377]
 
 0.16.1 (2020-12-20)
 -------------------

gwcs/tests/test_wcs.py

@@ -283,6 +283,62 @@ def test_bounding_box():
     assert_allclose(w(-1*u.pix, -1*u.pix), (np.nan, np.nan))
 
 
+def test_compound_bounding_box():
+    trans3 = models.Shift(10) & models.Scale(2) & models.Shift(-1)
+    pipeline = [('detector', trans3), ('sky', None)]
+    w = wcs.WCS(pipeline)
+    cbb = {
+        1: ((-1, 10), (6, 15)),
+        2: ((-1, 5), (3, 17)),
+        3: ((-3, 7), (1, 27)),
+    }
+    if new_bbox:
+        # Test attaching a valid bounding box (ignoring input 'x')
+        w.attach_compound_bounding_box(cbb, [('x',)])
+        from astropy.modeling.bounding_box import CompoundBoundingBox
+        cbb = CompoundBoundingBox.validate(trans3, cbb, selector_args=[('x',)], order='F')
+        assert w.bounding_box == cbb
+        assert w.bounding_box is trans3.bounding_box
+
+        # Test evaluating
+        assert_allclose(w(13, 2, 1), (np.nan, np.nan, np.nan))
+        assert_allclose(w(13, 2, 2), (np.nan, np.nan, np.nan))
+        assert_allclose(w(13, 0, 3), (np.nan, np.nan, np.nan))
+        # No bounding box for selector
+        with pytest.raises(RuntimeError):
+            w(13, 13, 4)
+
+        # Test attaching a invalid bounding box (not ignoring input 'x')
+        with pytest.raises(ValueError):
+            w.attach_compound_bounding_box(cbb, [('x', False)])
+    else:
+        with pytest.raises(NotImplementedError) as err:
+            w.attach_compound_bounding_box(cbb, [('x',)])
+        assert str(err.value) == 'Compound bounding box is not supported for your version of astropy'
+
+    # Test that bounding_box with quantities can be assigned and evaluates
+    trans = models.Shift(10 * u .pix) & models.Shift(2 * u.pix)
+    pipeline = [('detector', trans), ('sky', None)]
+    w = wcs.WCS(pipeline)
+    cbb = {
+        1 * u.pix: (1 * u.pix, 5 * u.pix),
+        2 * u.pix: (2 * u.pix, 6 * u.pix)
+    }
+    if new_bbox:
+        w.attach_compound_bounding_box(cbb, [('x1',)])
+
+        from astropy.modeling.bounding_box import CompoundBoundingBox
+        cbb = CompoundBoundingBox.validate(trans, cbb, selector_args=[('x1',)], order='F')
+        assert w.bounding_box == cbb
+        assert w.bounding_box is trans.bounding_box
+
+        assert_allclose(w(-1*u.pix, 1*u.pix), (np.nan, np.nan))
+        assert_allclose(w(7*u.pix, 2*u.pix), (np.nan, np.nan))
+    else:
+        with pytest.raises(NotImplementedError) as err:
+            w.attach_compound_bounding_box(cbb, [('x1',)])
+
+
 def test_grid_from_bounding_box():
     bb = ((-1, 9.9), (6.5, 15))
     x, y = grid_from_bounding_box(bb, step=[.1, .5], center=False)

gwcs/wcs.py

@@ -24,6 +24,7 @@
 
 try:
     from astropy.modeling.bounding_box import ModelBoundingBox as Bbox
+    from astropy.modeling.bounding_box import CompoundBoundingBox
     new_bbox = True
 except ImportError:
     from astropy.modeling.utils import _BoundingBox as Bbox
@@ -1313,6 +1314,7 @@ def bounding_box(self):
         Return the range of acceptable values for each input axis.
         The order of the axes is `~gwcs.coordinate_frames.CoordinateFrame.axes_order`.
         """
+
         frames = self.available_frames
         transform_0 = self.get_transform(frames[0], frames[1])
         try:
@@ -1353,7 +1355,10 @@ def bounding_box(self, value):
             try:
                 # Make sure the dimensions of the new bbox are correct.
                 if new_bbox:
-                    bbox = Bbox.validate(transform_0, value, order='F')
+                    if isinstance(value, CompoundBoundingBox):
+                        bbox = CompoundBoundingBox.validate(transform_0, value, order='F')
+                    else:
+                        bbox = Bbox.validate(transform_0, value, order='F')
                 else:
                     Bbox.validate(transform_0, value)
             except Exception:
@@ -1373,6 +1378,16 @@ def bounding_box(self, value):
 
         self.set_transform(frames[0], frames[1], transform_0)
 
+    def attach_compound_bounding_box(self, cbbox, selector_args):
+        if new_bbox:
+            frames = self.available_frames
+            transform_0 = self.get_transform(frames[0], frames[1])
+
+            self.bounding_box = CompoundBoundingBox.validate(transform_0, cbbox, selector_args=selector_args,
+                                                             order='F')
+        else:
+            raise NotImplementedError('Compound bounding box is not supported for your version of astropy')
+
     def _get_axes_indices(self):
         try:
             axes_ind = np.argsort(self.input_frame.axes_order)