Implement Superpixel on NDCube

changelog/450.feature.rst

@@ -0,0 +1 @@
+Implement a new `ndcube.NDCube.superpixel` method to combine integer numbers of pixels along each axis into larger pixels.

ndcube/conftest.py

@@ -56,11 +56,15 @@ def gen_ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l, time_axis, time_base, global_co
     wcs_3d_lt_ln_l.array_shape = shape
     data_cube = data_nd(shape)
     mask = data_cube < 0
+    meta = {"message": "hello world"}
+    unit = u.ph
     extra_coords = time_extra_coords(shape, time_axis, time_base)
     cube = NDCube(data_cube,
                   wcs_3d_lt_ln_l,
                   mask=mask,
-                  uncertainty=data_cube)
+                  uncertainty=data_cube,
+                  meta=meta,
+                  unit=unit)
     cube._extra_coords = extra_coords
 
     if global_coords:
@@ -375,6 +379,7 @@ def ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l):
                                         1,
                                         Time('2000-01-01', format='fits', scale='utc'))
 
+
 @pytest.fixture
 def ndcube_2d_ln_lt(wcs_2d_lt_ln):
     shape = (10, 12)

ndcube/extra_coords/extra_coords.py

@@ -5,6 +5,7 @@
 
 import astropy.units as u
 import numpy as np
+import scipy.interpolate
 from astropy.coordinates import SkyCoord
 from astropy.time import Time
 from astropy.wcs.wcsapi import BaseHighLevelWCS
@@ -215,7 +216,7 @@ def add(self, name, array_dimension, lookup_table, physical_types=None, **kwargs
 
         # Sort the LUTs so that the mapping and the wcs are ordered in pixel dim order
         self._lookup_tables = list(sorted(self._lookup_tables,
-                                          key=lambda x: x[0] if isinstance(x[0], int) else x[0][0]))
+                                          key=lambda x: x[0] if isinstance(x[0], Integral) else x[0][0]))
 
     @property
     def _name_lut_map(self):
@@ -391,6 +392,64 @@ def dropped_world_dimensions(self):
 
         return dict()
 
+    def interpolate(self, new_array_grids, ndcube=None, **kwargs):
+        """Interpolate all extra coords to new array index grids.
+
+        One new array index grid must be supplied for each array axis in array-axis order.
+
+        Parameters
+        ----------
+        new_array_grids: iterable of array-likes or `None`
+            The array index values at which the the new values of the coords are desired.
+            An array for each array axis must be given.
+            If None given for an axis, no interpolation is performed for coords
+            corresponding to that axis.
+
+        ndcube: `~ndcube.NDCube `
+            The NDCube instance with which the output ExtraCoords object is associated.
+
+        Returns
+        -------
+        new_ec: `~ndcube.extra_coords.ExtraCoords`
+            A new ExtraCoords object holding the interpolated coords.
+
+        """
+        cube_dims = self._ndcube.dimensions.value.astype(int)
+        naxes = len(cube_dims)
+        if len(new_array_grids) != naxes:
+            raise ValueError("new_array_grids must have an entry for each array axis."
+                             f"num. array axes: {naxes}; "
+                             f"num. new array grids: {len(new_array_grids)}")
+        array_order_mapping = convert_between_array_and_pixel_axes(np.asarray(self.mapping), naxes)
+        old_array_grids = [None if new_grid is None else np.arange(dim)
+                           for new_grid, dim in zip(new_array_grids, cube_dims)]
+        new_ec = type(self)(ndcube)
+        for key, aom, lut in zip(self.keys(), array_order_mapping, self._lookup_tables):
+            if new_array_grids[aom] is None:
+                new_coord = lut[1]
+            else:
+                table = lut[1].table
+                if isinstance(table, tuple):
+                    table = table[0]
+                if isinstance(table, Time):
+                    table_values = table.mjd
+                elif isinstance(table, SkyCoord):
+                    pass
+                elif isinstance(table, u.Quantity):
+                    table_values = table.value
+                else:
+                    raise TypeError(f"Unrecognized lookup table type: {table}")
+                lut_interp = scipy.interpolate.interp1d(old_array_grids[aom], table_values)
+                new_table_values = lut_interp(new_array_grids[aom])
+                if isinstance(table, Time):
+                    new_coord = Time(new_table_values, scale=table.scale, format="mjd")
+                elif isinstance(table, SkyCoord):
+                    pass
+                else:
+                    new_coord = u.Quantity(new_table_values, unit=table.unit)
+            new_ec.add(key, int(aom), new_coord, physical_types=lut[1].physical_types)
+        return new_ec
+
     def __str__(self):
         elements = [f"{', '.join(table.names)} ({axes}): {table}" for axes, table in self._lookup_tables]
         return f"ExtraCoords({', '.join(elements)})"

ndcube/ndcube.py

@@ -26,7 +26,7 @@
 from ndcube.ndcube_sequence import NDCubeSequence
 from ndcube.utils.wcs_high_level_conversion import high_level_objects_to_values, values_to_high_level_objects
 from ndcube.visualization import PlotterDescriptor
-from ndcube.wcs.wrappers import CompoundLowLevelWCS
+from ndcube.wcs.wrappers import CompoundLowLevelWCS, ResampledLowLevelWCS
 
 __all__ = ['NDCubeABC', 'NDCubeBase', 'NDCube']
 
@@ -782,6 +782,127 @@ def reproject_to(self, target_wcs, shape_out=None, order='bilinear', output_arra
 
         return resampled_cube
 
+    def superpixel(self, superpixel_shape, func=np.sum, new_unit=False):
+        """Downsample array by creating non-overlapping superpixels.
+
+        Values in superpixels are determined applying a function to the pixel
+        values within it.  The number of pixels in each superpixel in each
+        dimension is given is given by the superpixel_shape input.
+        This must be an integer fraction of the cube's array size in each dimension.
+
+        This method currently does not handle uncertainty or extra_coords.
+        They are dropped from the output.
+
+        Parameters
+        ----------
+        data : array-like
+            The data array to be downsampled
+
+        superpixel_shape : array-like
+            The number of pixels in a superpixel 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.
+
+        func : (optional)
+            Function applied to the data to derive values of the superpixels.
+            The function must take an array as its first argument and
+            support the axis kwarg with the same meaning as a numpy axis
+            kward (see the description of `~numpy.sum` for an example.)
+            Default: `~numpy.sum`
+
+        new_unit : `str` or `astropy.units.Unit` (optional)
+            The unit of the new data. This might need to change based on how
+            func alters the data.
+            Default: False, i.e. self.unit is used. (False is used as default
+            flag because None is a valid new_unit option.)
+
+        Returns
+        -------
+        new_cube: `NDCube`
+            The resolution-degraded cube.
+
+        References
+        ----------
+        https://mail.scipy.org/pipermail/numpy-discussion/2010-July/051760.html
+
+        Notes
+        -----
+        Superpixels are created by reshaping the N-D array to a 2N-D array and
+        applying the function over the odd-numbered axes. An example of this
+        applied the data only is the following.  Let's say you have an array::
+
+             x = np.array([[0, 0, 0, 1, 1, 1],
+                           [0, 0, 1, 1, 0, 0],
+                           [1, 1, 0, 0, 1, 1],
+                           [0, 0, 0, 0, 1, 1],
+                           [1, 0, 1, 0, 1, 1],
+                           [0, 0, 1, 0, 0, 0]])
+
+        and you want to sum over 2x2 non-overlapping sub-arrays.  This summing can
+        be done by reshaping the array::
+
+             y = x.reshape(3,2,3,2)
+
+        and then summing over the 1st and third directions::
+
+             y2 = y.sum(axis=3).sum(axis=1)
+
+        which gives the expected array::
+
+             array([[0, 3, 2],
+                   [2, 0, 4],
+                   [1, 2, 2]])
+        """
+        # Sanitize input.
+        # Make sure the input superpixel dimensions are integers.
+        superpixel_shape = np.rint(superpixel_shape).astype(int)
+        offsets = (superpixel_shape - 1) / 2
+        if all(superpixel_shape == 1):
+            return deepcopy(self)
+        # Ensure superpixel_size has right number of entries and each entry is an
+        # integer fraction of the array shape in each dimension.
+        data_shape = self.dimensions.value.astype(int)
+        naxes = len(data_shape)
+        if len(superpixel_shape) != naxes:
+            raise ValueError("superpixel_shape must have an entry for each array axis.")
+        if (np.mod(data_shape, superpixel_shape) != 0).any():
+            raise ValueError(
+                "superpixel shape must be an integer fraction of the data shape in each dimension. "
+                f"data shape: {data_shape};  superpixel shape: {superpixel_shape}")
+        if new_unit is False:
+            new_unit = self.unit
+
+        # Reshape array and apply function over odd axes to generate array of superpixels.
+        if self.mask is None:
+            data = self.data
+            new_mask = None
+        else:
+            data = np.ma.masked_array(self.data, self.mask)
+        reshape = np.empty(data_shape.size + superpixel_shape.size, dtype=int)
+        reshape[0::2] = data_shape / superpixel_shape
+        reshape[1::2] = superpixel_shape
+        new_data = data.reshape(tuple(reshape))
+        for i in range(len(reshape) - 1, 0, -2):
+            new_data = func(new_data, axis=i)
+        if self.mask is not None:
+            new_mask = new_data.mask
+            new_data = new_data.data
+
+        # Resample WCS
+        new_wcs = ResampledLowLevelWCS(self.wcs.low_level_wcs, superpixel_shape[::-1])
+
+        # Reform NDCube.
+        new_cube = type(self)(new_data, new_wcs, mask=new_mask, meta=self.meta, unit=new_unit)
+
+        # Reconstitute extra coords
+        new_array_grids = [None if superpixel_shape[i] == 1 else
+                           np.arange(offsets[i], data_shape[i] + offsets[i], superpixel_shape[i])
+                           for i in range(naxes)]
+        new_cube._extra_coords = self.extra_coords.interpolate(new_array_grids, new_cube)
+
+        return new_cube
+
 
 class NDCube(NDCubeBase, astropy.nddata.NDArithmeticMixin):
     """

ndcube/tests/test_ndcube.py

@@ -630,3 +630,44 @@ def test_wcs_type_after_init(ndcube_3d_ln_lt_l, wcs_3d_l_lt_ln):
     cube = NDCube(ndcube_3d_ln_lt_l.data[slices], low_level_wcs)
     # Check the WCS has been converted to high level but NDCube init.
     assert isinstance(cube.wcs, BaseHighLevelWCS)
+
+
+def test_superpixel(ndcube_3d_l_ln_lt_ectime):
+    # Execute superpixel.
+    cube = ndcube_3d_l_ln_lt_ectime[:, 1:]
+    superpixel_shape = (10, 2, 1)
+    output = cube.superpixel(superpixel_shape, func=np.sum)
+    output_sc, output_spec = output.axis_world_coords(wcs=output.wcs)
+    output_time, = output.axis_world_coords(wcs=output.extra_coords)
+
+    # Build expected cube contents.
+    expected_data = np.array([[3840, 3860, 3880, 3900, 3920, 3940, 3960, 3980],
+                              [4160, 4180, 4200, 4220, 4240, 4260, 4280, 4300]])
+    expected_mask = np.zeros(expected_data.shape, dtype=bool)
+    expected_uncertainty = None
+    expected_unit = cube.unit
+    expected_meta = cube.meta
+    expected_Tx = np.array([[9.99999999, 19.99999994, 29.99999979, 39.9999995,
+                             49.99999902, 59.99999831, 69.99999731, 79.99999599],
+                            [9.99999999, 19.99999994, 29.99999979, 39.9999995,
+                             49.99999902, 59.99999831, 69.99999731, 79.99999599]]) * u.arcsec
+    expected_Ty = np.array([[-14.99999996, -14.9999999, -14.99999981, -14.99999969,
+                             -14.99999953, -14.99999934, -14.99999911, -14.99999885],
+                            [-4.99999999, -4.99999998, -4.99999995, -4.9999999,
+                             -4.99999985, -4.99999979, -4.99999971, -4.99999962]]) * u.arcsec
+    expected_spec = SpectralCoord([1.02e-09], unit=u.m)
+    expected_time = Time([51544.00104167, 51544.00243056], format="mjd", scale="utc")
+
+    # Confirm output is as expected.
+    assert (output.dimensions.value == np.array([1, 2, 8])).all()
+    assert (output.data == expected_data).all()
+    assert (output.mask == expected_mask).all()
+    assert output.uncertainty == expected_uncertainty
+    assert output.unit == expected_unit
+    assert output.meta == expected_meta
+    assert u.allclose(output_sc.Tx, expected_Tx)
+    assert u.allclose(output_sc.Ty, expected_Ty)
+    assert u.allclose(output_spec, expected_spec)
+    assert output_time.scale == expected_time.scale
+    assert output_time.format == expected_time.format
+    assert np.allclose(output_time.value, expected_time.value)