Merge pull request #744 from sunpy/rebin

Rough attempt at ignoring units for rebin

Author: nabobalis

changelog/744.bugfix.rst

@@ -0,0 +1 @@
+:meth:`ndcube.NDCube.rebin` ``bin_shape`` argument now accepts a `astropy.units.Quantity` as input if the units are convertible to pixels.

ndcube/ndcube.py

@@ -1019,11 +1019,12 @@ def rebin(self, bin_shape, operation=np.mean, operation_ignores_mask=False, hand
 
         Parameters
         ----------
-        bin_shape : array-like
+        bin_shape : array-like, `astropy.units.Quantity`
             The number of pixels in a bin in each dimension.
             Must be the same length as number of dimensions in data.
             Each element must be in int. If they are not they will be rounded
-            to the nearest int.
+            to the nearest int. If provided as a `~astropy.units.Quantity` the
+            units have to be convertible to pixels.
         operation : function
             Function applied to the data to derive values of the bins.
             Default is `numpy.mean`
@@ -1122,6 +1123,8 @@ def my_propagate(uncertainty, data, mask, **kwargs):
         """
         # Sanitize input.
         new_unit = new_unit or self.unit
+        if isinstance(bin_shape, u.Quantity):
+            bin_shape = bin_shape.to_value(u.pixel)
         # Make sure the input bin dimensions are integers.
         bin_shape = np.rint(bin_shape).astype(int)
         if np.all(bin_shape == 1):

ndcube/tests/test_ndcube.py

@@ -1,3 +1,4 @@
+import re
 from inspect import signature
 from textwrap import dedent
 
@@ -863,6 +864,17 @@ def test_rebin_dask(ndcube_2d_dask):
     assert isinstance(output.mask, dask_type)
 
 
+def test_rebin_bin_shape_quantity(ndcube_3d_l_ln_lt_ectime):
+    # Confirm rebin's bin_shape argument handles being a astropy unit
+    cube = ndcube_3d_l_ln_lt_ectime[:, 1:]
+    cube._extra_coords = ExtraCoords(cube)
+    bin_shape = (10, 2, 1) * u.pix
+    output = cube.rebin(bin_shape)
+    np.testing.assert_allclose(output.shape, cube.shape / bin_shape.to_value())
+    with pytest.raises(u.UnitConversionError, match=re.escape("'m' (length) and 'pix' are not convertible")):
+        cube.rebin((10, 2, 1) * u.m)
+
+
 def test_rebin_no_ec(ndcube_3d_l_ln_lt_ectime):
     # Confirm rebin does not try to handle extra coords when there aren't any.
     cube = ndcube_3d_l_ln_lt_ectime[:, 1:]