Skip to content

Commit

Permalink
WIP: try making dask do reprojection
Browse files Browse the repository at this point in the history
  • Loading branch information
keflavich committed Oct 5, 2022
1 parent c27f0de commit 6def462
Showing 1 changed file with 92 additions and 0 deletions.
92 changes: 92 additions & 0 deletions spectral_cube/dask_spectral_cube.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

0 comments on commit 6def462

Please sign in to comment.