From 0cc11250b5aa34d0de84c6c06ddcdeae0a077c9a Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Wed, 5 Oct 2022 13:08:51 -0400
Subject: [PATCH] WIP: try making dask do reprojection

---
 spectral_cube/dask_spectral_cube.py | 92 +++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/spectral_cube/dask_spectral_cube.py b/spectral_cube/dask_spectral_cube.py
index 08125222a..fc57e673c 100644
--- a/spectral_cube/dask_spectral_cube.py
+++ b/spectral_cube/dask_spectral_cube.py
@@ -1632,3 +1632,95 @@ def _mask_include(self):
                                               chunks=self._data.chunksize),
                                 wcs=self.wcs,
                                 shape=self.shape)
+
+    def reproject(self, header, order='bilinear', use_memmap=False,
+                  filled=True, **kwargs):
+        """
+        Spatially reproject the cube into a new header.  Fills the data with
+        the cube's ``fill_value`` to replace bad values before reprojection.
+
+        If you want to reproject a cube both spatially and spectrally, you need
+        to use `spectral_interpolate` as well.
+
+        .. warning::
+            The current implementation of ``reproject`` requires that the whole
+            cube be loaded into memory.  Issue #506 notes that this is a
+            problem, and it is on our to-do list to fix.
+
+        Parameters
+        ----------
+        header : `astropy.io.fits.Header`
+            A header specifying a cube in valid WCS
+        order : int or str, optional
+            The order of the interpolation (if ``mode`` is set to
+            ``'interpolation'``). This can be either one of the following
+            strings:
+
+                * 'nearest-neighbor'
+                * 'bilinear'
+                * 'biquadratic'
+                * 'bicubic'
+
+            or an integer. A value of ``0`` indicates nearest neighbor
+            interpolation.
+        use_memmap : bool
+            If specified, a memory mapped temporary file on disk will be
+            written to rather than storing the intermediate spectra in memory.
+        filled : bool
+            Fill the masked values with the cube's fill value before
+            reprojection?  Note that setting ``filled=False`` will use the raw
+            data array, which can be a workaround that prevents loading large
+            data into memory.
+        kwargs : dict
+            Passed to `reproject.reproject_interp`.
+        """
+
+        try:
+            from reproject.version import version
+        except ImportError:
+            raise ImportError("Requires the reproject package to be"
+                              " installed.")
+
+        reproj_kwargs = kwargs
+        # Need version > 0.2 to work with cubes, >= 0.5 for memmap
+        from distutils.version import LooseVersion
+        if LooseVersion(version) < "0.5":
+            raise Warning("Requires version >=0.5 of reproject. The current "
+                          "version is: {}".format(version))
+        elif LooseVersion(version) >= "0.6":
+            pass # no additional kwargs, no warning either
+        else:
+            reproj_kwargs['independent_celestial_slices'] = True
+
+        from reproject import reproject_interp
+
+        # TODO: Find the minimal subcube that contains the header and only reproject that
+        # (see FITS_tools.regrid_cube for a guide on how to do this)
+
+        newwcs = wcs.WCS(header)
+        shape_out = tuple([header['NAXIS{0}'.format(i + 1)] for i in
+                           range(header['NAXIS'])][::-1])
+
+        def reproject_interp_wrapper(img_slice):
+            # What exactly is the wrapper getting here?
+            # I think it is given a _cube_ that is a cutout?
+            if filled:
+                data = img_slice.filled_data[:]
+            else:
+                data = img_slice._data
+            return reproject_interp((data, img_slice.header),
+                                    newwcs, shape_out=shape_out, **kwargs)
+
+        newcube, newcube_valid = self.apply_function_parallel_spatial(
+            reproject_interp_wrapper,
+            accepts_chunks=True,
+            order=order,
+            **reproj_kwargs)
+
+        return self._new_cube_with(data=newcube,
+                                   wcs=newwcs,
+                                   mask=BooleanArrayMask(newcube_valid.astype('bool'),
+                                                         newwcs),
+                                   meta=self.meta,
+                                  )
+