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 Mar 13, 2023
1 parent 8747c37 commit 0cc1125
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"

Check warning on line 1681 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1678-L1681

Added lines #L1678 - L1681 were not covered by tests
" installed.")

reproj_kwargs = kwargs

Check warning on line 1684 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1684

Added line #L1684 was not covered by tests
# 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 "

Check warning on line 1688 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1686-L1688

Added lines #L1686 - L1688 were not covered by tests
"version is: {}".format(version))
elif LooseVersion(version) >= "0.6":
pass # no additional kwargs, no warning either

Check warning on line 1691 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1690-L1691

Added lines #L1690 - L1691 were not covered by tests
else:
reproj_kwargs['independent_celestial_slices'] = True

Check warning on line 1693 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1693

Added line #L1693 was not covered by tests

from reproject import reproject_interp

Check warning on line 1695 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1695

Added line #L1695 was not covered by tests

# 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

Check warning on line 1701 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1700-L1701

Added lines #L1700 - L1701 were not covered by tests
range(header['NAXIS'])][::-1])

def reproject_interp_wrapper(img_slice):

Check warning on line 1704 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1704

Added line #L1704 was not covered by tests
# 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[:]

Check warning on line 1708 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1707-L1708

Added lines #L1707 - L1708 were not covered by tests
else:
data = img_slice._data
return reproject_interp((data, img_slice.header),

Check warning on line 1711 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1710-L1711

Added lines #L1710 - L1711 were not covered by tests
newwcs, shape_out=shape_out, **kwargs)

newcube, newcube_valid = self.apply_function_parallel_spatial(

Check warning on line 1714 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1714

Added line #L1714 was not covered by tests
reproject_interp_wrapper,
accepts_chunks=True,
order=order,
**reproj_kwargs)

return self._new_cube_with(data=newcube,

Check warning on line 1720 in spectral_cube/dask_spectral_cube.py

View check run for this annotation

Codecov / codecov/patch

spectral_cube/dask_spectral_cube.py#L1720

Added line #L1720 was not covered by tests
wcs=newwcs,
mask=BooleanArrayMask(newcube_valid.astype('bool'),
newwcs),
meta=self.meta,
)

0 comments on commit 0cc1125

Please sign in to comment.