Repository: sunpy/ndcube

Issue:
Robustly ensure the behaviour of `crop` is correct.
SunPy's Map has had unlimited trouble with `submap` we need to make sure our semantics are correct before release: https://github.com/sunpy/sunpy-SEP/pull/51#discussion_r533597518

Discussion:

Patch:
diff --git a/changelog/438.bugfix.rst b/changelog/438.bugfix.rst
new file mode 100644
index 0000000..0be5cff
--- /dev/null
+++ b/changelog/438.bugfix.rst
@@ -0,0 +1,3 @@
+Change the implementation of `.NDCube.crop` so that it takes into account all
+the corners of the world region specified by the upper and lower corners, not
+just those two points.
diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py
index 778a84f..0da23dd 100644
--- a/ndcube/ndcube.py
+++ b/ndcube/ndcube.py
@@ -1,6 +1,7 @@
 import abc
 import textwrap
 import warnings
+import itertools
 from copy import deepcopy
 from collections import namedtuple
 
@@ -22,6 +23,7 @@ from ndcube.extra_coords import ExtraCoords
 from ndcube.global_coords import GlobalCoords
 from ndcube.mixins import NDCubeSlicingMixin
 from ndcube.ndcube_sequence import NDCubeSequence
+from ndcube.utils.wcs_high_level_conversion import high_level_objects_to_values
 from ndcube.visualization import PlotterDescriptor
 from ndcube.wcs.wrappers import CompoundLowLevelWCS
 
@@ -45,8 +47,8 @@ class NDCubeABC(astropy.nddata.NDData, metaclass=abc.ABCMeta):
         coordinates of the corners to pixel coordinates and then cropping the
         smallest pixel region which contains the corners specified.
         This means that the edges of the world coordinate region specified by
-        the coordinates are not guaranteed to be included in the cropped output.
-        This is normally noticeable when cropping a celestial coordinate in a
+        the coordinates are not gaureented to be included in the cropped output.
+        This is normally noiticable when cropping a celestial coordinate in a
         frame which differs from the native frame of the coordinates in the WCS.
 
         Parameters
@@ -93,8 +95,8 @@ class NDCubeABC(astropy.nddata.NDData, metaclass=abc.ABCMeta):
         coordinates of the corners to pixel coordinates and then cropping the
         smallest pixel region which contains the corners specified.
         This means that the edges of the world coordinate region specified by
-        the coordinates are not guaranteed to be included in the cropped output.
-        This is normally noticeable when cropping a celestial coordinate in a
+        the coordinates are not gaureented to be included in the cropped output.
+        This is normally noiticable when cropping a celestial coordinate in a
         frame which differs from the native frame of the coordinates in the WCS.
 
         Parameters
@@ -514,18 +516,44 @@ class NDCubeBase(NDCubeSlicingMixin, NDCubeABC):
     def crop(self, lower_corner, upper_corner, wcs=None):
         # The docstring is defined in NDCubeABC
         lower_corner, upper_corner = utils.misc.sanitize_corners(lower_corner, upper_corner)
-        return self._crop(lower_corner, upper_corner, wcs, False)
+
+        # Quit out early if we are no-op
+        lower_nones = np.array([lower is None for lower in lower_corner])
+        upper_nones = np.array([upper is None for upper in upper_corner])
+        if (lower_nones & upper_nones).all():
+            return self
+
+        lower_corner, upper_corner = self._fill_in_crop_nones(lower_corner, upper_corner, wcs, False)
+
+        if isinstance(wcs, BaseHighLevelWCS):
+            wcs = wcs.low_level_wcs
+
+        lower_corner_values = high_level_objects_to_values(*lower_corner, low_level_wcs=wcs)
+        upper_corner_values = high_level_objects_to_values(*upper_corner, low_level_wcs=wcs)
+        lower_corner_values = [v << u.Unit(unit) for v, unit in zip(lower_corner_values, wcs.world_axis_units)]
+        upper_corner_values = [v << u.Unit(unit) for v, unit in zip(upper_corner_values, wcs.world_axis_units)]
+
+        points = self._bounding_box_to_points(lower_corner_values, upper_corner_values, wcs)
+        return self._crop_from_points(*points, wcs=wcs)
 
     @utils.misc.sanitise_wcs
     def crop_by_values(self, lower_corner, upper_corner, units=None, wcs=None):
         # The docstring is defined in NDCubeABC
         # Sanitize inputs.
         lower_corner, upper_corner = utils.misc.sanitize_corners(lower_corner, upper_corner)
+
+        # Quit out early if we are no-op
+        lower_nones = np.array([lower is None for lower in lower_corner])
+        upper_nones = np.array([upper is None for upper in upper_corner])
+        if (lower_nones & upper_nones).all():
+            return self
+
         n_coords = len(lower_corner)
         if units is None:
             units = [None] * n_coords
         elif len(units) != n_coords:
             raise ValueError("units must be None or have same length as corner inputs.")
+
         # Convert float inputs to quantities using units.
         types_with_units = (u.Quantity, type(None))
         for i, (lower, upper, unit) in enumerate(zip(lower_corner, upper_corner, units)):
@@ -544,34 +572,31 @@ class NDCubeBase(NDCubeSlicingMixin, NDCubeABC):
             if lower_corner[i] is not None and upper_corner[i] is not None:
                 upper_corner[i] = upper_corner[i].to(lower_corner[i].unit)
 
-        return self._crop(lower_corner, upper_corner, wcs, True)
+        lower_corner, upper_corner = self._fill_in_crop_nones(lower_corner, upper_corner, wcs, True)
+
+        # Convert coordinates to units used by WCS as WCS.world_to_array_index
+        # does not handle quantities.
+        lower_corner = utils.misc.convert_quantities_to_units(lower_corner,
+                                                              self.wcs.world_axis_units)
+        upper_corner = utils.misc.convert_quantities_to_units(upper_corner,
+                                                              self.wcs.world_axis_units)
 
-    def _crop(self, lower_corner, upper_corner, wcs, crop_by_values):
-        lower_corner = list(lower_corner)
-        upper_corner = list(upper_corner)
-        # If no intervals provided, return NDCube without slicing.
+        points = self._bounding_box_to_points(lower_corner, upper_corner, wcs)
+        return self._crop_from_points(*points, wcs=wcs)
+
+    def _fill_in_crop_nones(self, lower_corner, upper_corner, wcs, crop_by_values):
+        """
+        Replace any instance of None in the inputs with the bounds for that axis.
+        """
         lower_nones = np.array([lower is None for lower in lower_corner])
         upper_nones = np.array([upper is None for upper in upper_corner])
-        if (lower_nones & upper_nones).all():
-            return self
 
-        # Define functions to be used in converting between array indices and world coords
-        # based in input kwarg.
         if crop_by_values:
             if isinstance(wcs, BaseHighLevelWCS):
-                world_to_array_index = wcs.low_level_wcs.world_to_array_index_values
                 array_index_to_world = wcs.low_level_wcs.array_index_to_world_values
             else:
-                world_to_array_index = wcs.world_to_array_index_values
                 array_index_to_world = wcs.array_index_to_world_values
-            # Convert coordinates to units used by WCS as WCS.world_to_array_index
-            # does not handle quantities.
-            lower_corner = utils.misc.convert_quantities_to_units(lower_corner,
-                                                                  self.wcs.world_axis_units)
-            upper_corner = utils.misc.convert_quantities_to_units(upper_corner,
-                                                                  self.wcs.world_axis_units)
         else:
-            world_to_array_index = wcs.world_to_array_index
             array_index_to_world = wcs.array_index_to_world
 
         # If user did not provide all intervals,
@@ -587,26 +612,58 @@ class NDCubeBase(NDCubeSlicingMixin, NDCubeABC):
                     lower_corner[i] = interval[0]
                 if upper_is_none:
                     upper_corner[i] = interval[-1]
-        # Convert lower and upper corner coords to array indices.
-        lower_indices = world_to_array_index(*lower_corner)
-        upper_indices = world_to_array_index(*upper_corner)
-        # Ensure return type is tuple of lists, even if only one axis returned.
-        if not isinstance(lower_indices, tuple):
-            lower_indices = (lower_indices,)
-            upper_indices = (upper_indices,)
-        # Construct item which which to slice NDCube.
-        item = []
-        for i, (lower, upper) in enumerate(zip(lower_indices, upper_indices)):
-            # If upper limit index less than zero,
-            # then interval does not overlap with cube range. Raise error.
-            if upper < 0:
-                raise IndexError("Input real world interval beyond range of NDCube. "
-                                 f"Axis: {i}; Axis length: {self.dimensions[i]}; "
-                                 f"Derived array indices: {(lower, upper)}; ")
-            # Construct slice for this axis and append to item.
-            # Increment upper idex by 1 to ensure the upper world coord
-            # is included in sliced cube.
-            item.append(slice(max(0, lower), upper + 1))
+
+        return lower_corner, upper_corner
+
+    def _bounding_box_to_points(self, lower_corner_values, upper_corner_values, wcs):
+        """
+        Convert two corners of a bounding box to the points of all corners.
+        """
+        return tuple(itertools.product(*zip(lower_corner_values, upper_corner_values)))
+
+    def _crop_from_points(self, *world_points_values, wcs):
+        """
+        Crop to the minimum cube in pixel-space that contains all the specified world points.
+
+        Parameters
+        ----------
+        world_points_values
+            The world coordinates in wcsapi "values" form (i.e. arrays /
+            floats), for however many world points should be contained in the
+            output cube. Each argument should be a tuple with number of
+            coordinates equal to the number of world axes.
+
+        wcs : `~astropy.wcs.wcsapi.BaseHighLevelWCS`, `~astropy.wcs.wcsapi.BaseLowLevelWCS`
+            The WCS to use to convert the world coordinates to array indices.
+
+        Returns
+        -------
+        new_cube
+            A new ``NDCube`` instance.
+        """
+        if isinstance(wcs, BaseHighLevelWCS):
+            wcs = wcs.low_level_wcs
+
+        # Convert all points to array indices.
+        point_indices = []
+        for point in world_points_values:
+            indices = wcs.world_to_array_index_values(*point)
+
+            if not isinstance(indices, tuple):
+                indices = (indices,)
+
+            point_indices.append(indices)
+
+        point_indices = np.array(point_indices)
+        lower = np.min(point_indices, axis=0)
+        upper = np.max(point_indices, axis=0) + 1
+
+        # Wrap the limits to the size of the array
+        lower = [int(np.clip(index, 0, self.data.shape[i])) for i, index in enumerate(lower)]
+        upper = [int(np.clip(index, 0, self.data.shape[i])) for i, index in enumerate(upper)]
+
+        item = tuple(slice(l, u) for l, u in zip(lower, upper))
+
         return self[tuple(item)]
 
     def __str__(self):
diff --git a/ndcube/utils/wcs_high_level_conversion.py b/ndcube/utils/wcs_high_level_conversion.py
new file mode 100644
index 0000000..a78df94
--- /dev/null
+++ b/ndcube/utils/wcs_high_level_conversion.py
@@ -0,0 +1,187 @@
+"""
+This module provides helpers introduced in https://github.com/astropy/astropy/pull/11950.
+"""
+from collections import OrderedDict
+
+from astropy.wcs.wcsapi.high_level_api import default_order, rec_getattr
+
+__all__ = ['values_to_high_level_objects', 'high_level_objects_to_values']
+
+try:
+    from astropy.wcs.wcsapi.high_level_api import high_level_objects_to_values, values_to_high_level_objects
+
+except ImportError:
+
+    def high_level_objects_to_values(*world_objects, low_level_wcs):
+        """
+        Convert the input high level object to low level values.
+
+        This function uses the information in ``wcs.world_axis_object_classes`` and
+        ``wcs.world_axis_object_components`` to convert the high level objects
+        (such as `~.SkyCoord`) to low level "values" `~.Quantity` objects.
+
+        This is used in `.HighLevelWCSMixin.world_to_pixel`, but provided as a
+        separate function for use in other places where needed.
+
+        Parameters
+        ----------
+        *world_objects: object
+            High level coordinate objects.
+
+        low_level_wcs: `.BaseLowLevelWCS`
+            The WCS object to use to interpret the coordinates.
+        """
+        # Cache the classes and components since this may be expensive
+        serialized_classes = low_level_wcs.world_axis_object_classes
+        components = low_level_wcs.world_axis_object_components
+
+        # Deserialize world_axis_object_classes using the default order
+        classes = OrderedDict()
+        for key in default_order(components):
+            if low_level_wcs.serialized_classes:
+                classes[key] = deserialize_class(serialized_classes[key],
+                                                 construct=False)
+            else:
+                classes[key] = serialized_classes[key]
+
+        # Check that the number of classes matches the number of inputs
+        if len(world_objects) != len(classes):
+            raise ValueError("Number of world inputs ({}) does not match "
+                             "expected ({})".format(len(world_objects), len(classes)))
+
+        # Determine whether the classes are uniquely matched, that is we check
+        # whether there is only one of each class.
+        world_by_key = {}
+        unique_match = True
+        for w in world_objects:
+            matches = []
+            for key, (klass, *_) in classes.items():
+                if isinstance(w, klass):
+                    matches.append(key)
+            if len(matches) == 1:
+                world_by_key[matches[0]] = w
+            else:
+                unique_match = False
+                break
+
+        # If the match is not unique, the order of the classes needs to match,
+        # whereas if all classes are unique, we can still intelligently match
+        # them even if the order is wrong.
+
+        objects = {}
+
+        if unique_match:
+
+            for key, (klass, args, kwargs, *rest) in classes.items():
+
+                if len(rest) == 0:
+                    klass_gen = klass
+                elif len(rest) == 1:
+                    klass_gen = rest[0]
+                else:
+                    raise ValueError("Tuples in world_axis_object_classes should have length 3 or 4")
+
+                # FIXME: For now SkyCoord won't auto-convert upon initialization
+                # https://github.com/astropy/astropy/issues/7689
+                from astropy.coordinates import SkyCoord
+                if isinstance(world_by_key[key], SkyCoord):
+                    if 'frame' in kwargs:
+                        objects[key] = world_by_key[key].transform_to(kwargs['frame'])
+                    else:
+                        objects[key] = world_by_key[key]
+                else:
+                    objects[key] = klass_gen(world_by_key[key], *args, **kwargs)
+
+        else:
+
+            for ikey, key in enumerate(classes):
+
+                klass, args, kwargs, *rest = classes[key]
+
+                if len(rest) == 0:
+                    klass_gen = klass
+                elif len(rest) == 1:
+                    klass_gen = rest[0]
+                else:
+                    raise ValueError("Tuples in world_axis_object_classes should have length 3 or 4")
+
+                w = world_objects[ikey]
+                if not isinstance(w, klass):
+                    raise ValueError("Expected the following order of world "
+                                     "arguments: {}".format(', '.join([k.__name__ for (k, _, _) in classes.values()])))
+
+                # FIXME: For now SkyCoord won't auto-convert upon initialization
+                # https://github.com/astropy/astropy/issues/7689
+                from astropy.coordinates import SkyCoord
+                if isinstance(w, SkyCoord):
+                    if 'frame' in kwargs:
+                        objects[key] = w.transform_to(kwargs['frame'])
+                    else:
+                        objects[key] = w
+                else:
+                    objects[key] = klass_gen(w, *args, **kwargs)
+
+        # We now extract the attributes needed for the world values
+        world = []
+        for key, _, attr in components:
+            if callable(attr):
+                world.append(attr(objects[key]))
+            else:
+                world.append(rec_getattr(objects[key], attr))
+
+        return world
+
+    def values_to_high_level_objects(*world_values, low_level_wcs):
+        """
+        Convert low level values into high level objects.
+
+        This function uses the information in ``wcs.world_axis_object_classes`` and
+        ``wcs.world_axis_object_components`` to convert the high level objects
+        (such as `~.SkyCoord`) to low level "values" `~.Quantity` objects.
+
+        This is used in `.HighLevelWCSMixin.pixel_to_world`, but provided as a
+        separate function for use in other places where needed.
+
+        Parameters
+        ----------
+        *world_values: object
+            Low level, "values" representations of the world coordinates.
+
+        low_level_wcs: `.BaseLowLevelWCS`
+            The WCS object to use to interpret the coordinates.
+        """
+        # Cache the classes and components since this may be expensive
+        components = low_level_wcs.world_axis_object_components
+        classes = low_level_wcs.world_axis_object_classes
+
+        # Deserialize classes
+        if low_level_wcs.serialized_classes:
+            classes_new = {}
+            for key, value in classes.items():
+                classes_new[key] = deserialize_class(value, construct=False)
+            classes = classes_new
+
+        args = defaultdict(list)
+        kwargs = defaultdict(dict)
+
+        for i, (key, attr, _) in enumerate(components):
+            if isinstance(attr, str):
+                kwargs[key][attr] = world_values[i]
+            else:
+                while attr > len(args[key]) - 1:
+                    args[key].append(None)
+                args[key][attr] = world_values[i]
+
+        result = []
+
+        for key in default_order(components):
+            klass, ar, kw, *rest = classes[key]
+            if len(rest) == 0:
+                klass_gen = klass
+            elif len(rest) == 1:
+                klass_gen = rest[0]
+            else:
+                raise ValueError("Tuples in world_axis_object_classes should have length 3 or 4")
+            result.append(klass_gen(*args[key], *ar, **kwargs[key], **kw))
+
+        return result


Test patch:
diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py
index e8d1430..da1107c 100644
--- a/ndcube/tests/test_ndcube.py
+++ b/ndcube/tests/test_ndcube.py
@@ -1,9 +1,14 @@
+from textwrap import dedent
+
 import astropy.units as u
 import astropy.wcs
 import numpy as np
 import pytest
 from astropy.coordinates import SkyCoord, SpectralCoord
+from astropy.io import fits
 from astropy.time import Time
+from astropy.wcs import WCS
+from astropy.wcs.utils import wcs_to_celestial_frame
 from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS
 from astropy.wcs.wcsapi.wrappers import SlicedLowLevelWCS
 
@@ -471,23 +476,19 @@ def test_crop_by_values_all_nones(ndcube_4d_ln_lt_l_t):
     helpers.assert_cubes_equal(output, ndcube_4d_ln_lt_l_t)
 
 
-def test_crop_by_values_indexerror(ndcube_4d_ln_lt_l_t):
-    intervals = ndcube_4d_ln_lt_l_t.wcs.array_index_to_world_values([1, 2], [0, 1], [0, 1], [0, 2])
-    units = [u.min, u.m, u.deg, u.deg]
-    lower_corner = [coord[0] * unit for coord, unit in zip(intervals, units)]
-    upper_corner = [coord[-1] * unit for coord, unit in zip(intervals, units)]
-    lower_corner[1] *= -1
-    upper_corner[1] *= -1
-    with pytest.raises(IndexError):
-        ndcube_4d_ln_lt_l_t.crop_by_values(lower_corner, upper_corner)
-
-
 def test_crop_by_values_valueerror1(ndcube_4d_ln_lt_l_t):
+    # Test units not being the same length as the inputs
+    lower_corner = [None] * 4
+    lower_corner[0] = 0.5
+    upper_corner = [None] * 4
+    upper_corner[0] = 1.1
+
     with pytest.raises(ValueError):
-        ndcube_4d_ln_lt_l_t.crop_by_values([None, None], [None, None], units=["m"])
+        ndcube_4d_ln_lt_l_t.crop_by_values(lower_corner, upper_corner, units=["m"])
 
 
 def test_crop_by_values_valueerror2(ndcube_4d_ln_lt_l_t):
+    # Test upper and lower coordinates not being the same length
     with pytest.raises(ValueError):
         ndcube_4d_ln_lt_l_t.crop_by_values([None], [None, None])
 
@@ -502,6 +503,49 @@ def test_crop_by_values_1d_dependent(ndcube_4d_ln_lt_l_t):
     helpers.assert_cubes_equal(output, expected)
 
 
+def test_crop_rotated_celestial():
+    # This is a regression test for a highly rotated image where all 4 corners
+    # of the spatial ROI have to be used.
+
+    header = dedent("""\
+        WCSAXES =                    2 / Number of coordinate axes
+        CRPIX1  =          2053.459961 / Pixel coordinate of reference point
+        CRPIX2  =          2047.880005 / Pixel coordinate of reference point
+        PC1_1   =     0.70734471922412 / Coordinate transformation matrix element
+        PC1_2   =     0.70686876305701 / Coordinate transformation matrix element
+        PC2_1   =    -0.70686876305701 / Coordinate transformation matrix element
+        PC2_2   =     0.70734471922412 / Coordinate transformation matrix element
+        CDELT1  =  0.00016652472222222 / [deg] Coordinate increment at reference point
+        CDELT2  =  0.00016652472222222 / [deg] Coordinate increment at reference point
+        CUNIT1  = 'deg'                / Units of coordinate increment and value
+        CUNIT2  = 'deg'                / Units of coordinate increment and value
+        CTYPE1  = 'HPLN-TAN'           / Coordinate type codegnomonic projection
+        CTYPE2  = 'HPLT-TAN'           / Coordinate type codegnomonic projection
+        CRVAL1  =                  0.0 / [deg] Coordinate value at reference point
+        CRVAL2  =                  0.0 / [deg] Coordinate value at reference point
+        LONPOLE =                180.0 / [deg] Native longitude of celestial pole
+        LATPOLE =                  0.0 / [deg] Native latitude of celestial pole
+        MJDREF  =                  0.0 / [d] MJD of fiducial time
+        DATE-OBS= '2014-04-09T06:00:12.970' / ISO-8601 time of observation
+        MJD-OBS =      56756.250150116 / [d] MJD of observation
+        RSUN_REF=          696000000.0 / [m] Solar radius
+        DSUN_OBS=      149860273889.04 / [m] Distance from centre of Sun to observer
+        HGLN_OBS=  -0.0058904803279347 / [deg] Stonyhurst heliographic lng of observer
+        HGLT_OBS=     -6.0489216362492 / [deg] Heliographic latitude of observer
+        """)
+    wcs = WCS(fits.Header.fromstring(header, sep="\n"))
+    data = np.zeros((4096, 4096))
+
+    cube = NDCube(data, wcs=wcs)
+
+    bottom_left = SkyCoord(-100, -100, unit=u.arcsec, frame=wcs_to_celestial_frame(wcs))
+    top_right = SkyCoord(600, 600, unit=u.arcsec, frame=wcs_to_celestial_frame(wcs))
+
+    small = cube.crop(bottom_left, top_right)
+
+    assert small.data.shape == (1652, 1652)
+
+
 def test_initialize_from_ndcube(ndcube_3d_l_ln_lt_ectime):
     cube = ndcube_3d_l_ln_lt_ectime
     cube.global_coords.add('distance', 'pos.distance', 1 * u.m)