Add cube mosaicking (#829)

* Add cube mosaicking

Created functions to mosaic cubes together.

* Made Requested Changes to Pull Request

* Delete .gitignore

* Revert "Delete .gitignore"

This reverts commit 04ccb30.

* Update .gitignore

* Fix for Colorbar already exists error

In the case where the FITSFigure has already has a colorbar, will not attempt to add the colorbar again.

* Catch all Exceptions for Quicklook

Should print out all exceptions when encountering them in quicklook, and display an image still instead of crashing.

* Change to Combine Headers

Rename `reproject_together` and change function to use headers instead of cubes.

* mosaic_cubes now takes a list

Rework of mosaic_cubes to take a list of cubes and mosaic them all together into the same field.

* Update cube_utils.py

Change function name to combine_headers where it was called

* Update combine_headers

Fixed problems with find_optimal_celestial_wcs line being passed 3D WCS and wrong order of WCS & shape

* Update cube_utils.py

Fixed masking out additional cubes

* Update test_regrid.py

Added test for mosaic_cubes

* Update spectral_cube/cube_utils.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/cube_utils.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/cube_utils.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/cube_utils.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update cube_utils.py

* Update cube_utils.py

+= was causing an error

* Update spectral_cube/tests/test_regrid.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/tests/test_regrid.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/tests/test_regrid.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update lower_dimensional_structures.py

* Update test_regrid.py

add import statement

* Update test_regrid.py

more specific import statement

* Update spectral_cube/tests/test_regrid.py

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/cube_utils.py

block size argument error

Co-authored-by: Adam Ginsburg <[email protected]>

* Update spectral_cube/tests/test_regrid.py

Co-authored-by: Adam Ginsburg <[email protected]>

* pass kwargs to reproject

* add documentation for reading lower-dimensional objects (#834)

* add documentation for reading lower-dimensional objects

* Update docs/creating_reading.rst

Co-authored-by: Eric Koch <[email protected]>

* Update docs/creating_reading.rst

Co-authored-by: Eric Koch <[email protected]>

* Update docs/creating_reading.rst

Co-authored-by: Eric Koch <[email protected]>

* Update docs/creating_reading.rst

Co-authored-by: Eric Koch <[email protected]>

* Update spectral_cube/lower_dimensional_structures.py

Co-authored-by: Eric Koch <[email protected]>

* Update spectral_cube/lower_dimensional_structures.py

* Update spectral_cube/lower_dimensional_structures.py

Co-authored-by: Eric Koch <[email protected]>

* Update spectral_cube/lower_dimensional_structures.py

Co-authored-by: Eric Koch <[email protected]>

* fix some simple issues

* fix up first part of test (testing WCS)

* remove comments

* minor whitespace

* Fix whitespace

* Fix for missing edges

* Add fix to other reproject

* pass kwargs to mosaic_cubes

* fix tests and add some helpful error messages

* fix bug uncovered with new test

* Update spectral_cube/tests/test_regrid.py

Co-authored-by: Eric Koch <[email protected]>

* Update spectral_cube/cube_utils.py

Co-authored-by: Eric Koch <[email protected]>

* resolve comment from e-koch

* minor flake8 fixes

Co-authored-by: Adam Ginsburg <[email protected]>
Co-authored-by: Eric Koch <[email protected]>

Author: 3 people

.gitignore

@@ -60,3 +60,4 @@ docs/api
 */cython_version.py
 
 .tmp
+.ipynb_checkpoints

spectral_cube/cube_utils.py

@@ -14,6 +14,7 @@
 import numpy as np
 from astropy.wcs.utils import proj_plane_pixel_area
 from astropy.wcs import (WCSSUB_SPECTRAL, WCSSUB_LONGITUDE, WCSSUB_LATITUDE)
+from astropy.wcs import WCS
 from . import wcs_utils
 from .utils import FITSWarning, AstropyUserWarning, WCSCelestialError
 from astropy import log
@@ -750,3 +751,102 @@ def bunit_converters(obj, unit, equivalencies=(), freq=None):
     else:
         # Slice along first axis to return a 1D array.
         return factors[0]
+
+def combine_headers(header1, header2):
+    '''
+    Given two Header objects, this function returns a fits Header of the optimal wcs. 
+
+    Parameters
+    ----------
+    header1 : astropy.io.fits.Header
+        A Header.
+    header2 : astropy.io.fits.Header
+        A Header.
+
+    Returns
+    -------
+    header : astropy.io.fits.Header
+        A header object of a field containing both initial headers. 
+
+    '''
+
+    from reproject.mosaicking import find_optimal_celestial_wcs
+
+    # Get wcs and shape of both headers
+    w1 = WCS(header1).celestial
+    s1 = w1.array_shape
+    w2 = WCS(header2).celestial
+    s2 = w2.array_shape
+
+    # Get the optimal wcs and shape for both fields together
+    wcs_opt, shape_opt = find_optimal_celestial_wcs([(s1, w1), (s2, w2)], auto_rotate=False)
+
+    # Make a new header using the optimal wcs and information from cubes
+    header = header1.copy()
+    header['NAXIS'] = 3
+    header['NAXIS1'] = shape_opt[1]
+    header['NAXIS2'] = shape_opt[0]
+    header['NAXIS3'] = header1['NAXIS3']
+    header.update(wcs_opt.to_header())
+    header['WCSAXES'] = 3
+    return header
+
+def mosaic_cubes(cubes, spectral_block_size=100, **kwargs):
+    '''
+    This function reprojects cubes onto a common grid and combines them to a single field.  
+
+    Parameters
+    ----------
+    cubes : iterable
+        Iterable list of SpectralCube objects to reproject and add together. 
+    spectral_block_size : int
+        Block size so that reproject does not run out of memory. 
+
+    Outputs
+    -------
+    cube : SpectralCube
+        A spectral cube with the list of cubes mosaicked together.
+    '''
+
+    cube1 = cubes[0]
+    header = cube1.header
+
+    # Create a header for a field containing all cubes
+    for cu in cubes[1:]:
+        header = combine_headers(header, cu.header)
+
+    # Prepare an array and mask for the final cube
+    shape_opt = (header['NAXIS3'], header['NAXIS2'], header['NAXIS1'])
+    final_array = np.zeros(shape_opt)
+    mask_opt = np.zeros(shape_opt[1:])
+
+    for cube in cubes:
+        # Reproject cubes to the header
+        try:
+            if spectral_block_size is not None:
+                cube_repr = cube.reproject(header, block_size=[spectral_block_size, cube.shape[1], cube.shape[2]], **kwargs)
+            else:
+                cube_repr = cube.reproject(header, **kwargs)
+        except TypeError:
+            warnings.warn("The block_size argument is not accepted by `reproject`.  A more recent version may be needed.")
+            cube_repr = cube.reproject(header, **kwargs)
+
+        # Create weighting mask (2D)
+        mask = (cube_repr[0:1].get_mask_array()[0])
+        mask_opt += mask.astype(float)
+
+        # Go through each slice of the cube, add it to the final array
+        for ii in range(final_array.shape[0]):
+            slice1 = np.nan_to_num(cube_repr.unitless_filled_data[ii])
+            final_array[ii] = final_array[ii] + slice1
+
+    # Dividing by the mask throws errors where it is zero
+    with np.errstate(divide='ignore'):
+
+        # Use weighting mask to average where cubes overlap
+        for ss in range(final_array.shape[0]):
+            final_array[ss] /= mask_opt
+
+    # Create Cube
+    cube = cube1.__class__(data=final_array * cube1.unit, wcs=WCS(header))
+    return cube

spectral_cube/spectral_cube.py

@@ -2701,6 +2701,13 @@ def reproject(self, header, order='bilinear', use_memmap=False,
                                                   shape_out=shape_out,
                                                   order=order,
                                                   **reproj_kwargs)
+        if np.all(np.isnan(newcube)):
+            raise ValueError("All values in reprojected cube are nan.  This can be caused"
+                             " by an error in which coordinates do not 'round-trip'.  Try "
+                             "setting ``roundtrip_coords=False``.  You might also check "
+                             "whether the WCS transformation produces valid pixel->world "
+                             "and world->pixel coordinates in each axis."
+                             )
 
         return self._new_cube_with(data=newcube,
                                    wcs=newwcs,

spectral_cube/tests/test_regrid.py

@@ -4,7 +4,7 @@
 import numpy as np
 import os
 
-from astropy import units as u
+from astropy import constants, units as u
 from astropy import convolution
 from astropy.wcs import WCS
 from astropy import wcs
@@ -25,7 +25,9 @@
 from radio_beam import beam, Beam
 
 from .. import SpectralCube
+from ..masks import BooleanArrayMask
 from ..utils import WCSCelestialError
+from ..cube_utils import mosaic_cubes, combine_headers
 from .test_spectral_cube import cube_and_raw
 from .test_projection import load_projection
 from . import path, utilities
@@ -113,6 +115,11 @@ def test_reproject(use_memmap, data_adv, use_dask):
     wcs_out.wcs.crval = [134.37608, -31.939241, wcs_in.wcs.crval[2]]
     wcs_out.wcs.crpix = [2., 2., wcs_in.wcs.crpix[2]]
 
+    # cube is doppler-optical by default, which uses the rest wavelength,
+    # which isn't auto-computed, resulting in nan pixels in the WCS transform
+    wcs_out.wcs.restwav = 0.21106114549833
+    cube._wcs.wcs.restwav = 0.21106114549833
+
     header_out = cube.header
     header_out['NAXIS1'] = 4
     header_out['NAXIS2'] = 5
@@ -585,3 +592,44 @@ def test_reproject_3D_memory():
 
     assert result.wcs.wcs.crval[0] == 0.001
     assert result.wcs.wcs.crpix[0] == 2.
+
+@pytest.mark.parametrize('spectral_block_size,use_memmap', ((None, False),
+                                                            (100, False),
+                                                            (None, True),
+                                                            (100, False),
+                                                            (1, True),
+                                                            (1, False),
+                                                            ))
+def test_mosaic_cubes(use_memmap, data_adv, use_dask, spectral_block_size):
+
+    pytest.importorskip('reproject')
+
+    # Read in data to use
+    cube, data = cube_and_raw(data_adv, use_dask=use_dask)
+
+    # cube is doppler-optical by default, which uses the rest wavelength,
+    # which isn't auto-computed, resulting in nan pixels in the WCS transform
+    cube._wcs.wcs.restwav = constants.c.to(u.m/u.s).value / cube.wcs.wcs.restfrq
+
+    expected_wcs = WCS(combine_headers(cube.header, cube.header)).celestial
+
+    # Make two overlapping cubes of the data
+    part1 = cube[:, :round(cube.shape[1]*2./3.), :]
+    part2 = cube[:, round(cube.shape[1]/3.):, :]
+
+    assert part1.wcs.wcs.restwav != 0
+    assert part2.wcs.wcs.restwav != 0
+
+    result = mosaic_cubes([part1, part2], order='nearest-neighbor',
+                          roundtrip_coords=False,
+                          spectral_block_size=spectral_block_size)
+
+    # Check that the shapes are the same
+    assert result.shape == cube.shape
+
+    # Check WCS in reprojected matches wcs_out
+    # (comparing WCS failed for no reason we could discern)
+    assert repr(expected_wcs) == repr(result.wcs.celestial)
+    # Check that values of original and result are comparable
+    np.testing.assert_almost_equal(result.filled_data[:].value, cube.filled_data[:].value, decimal=3)
+    # only good to 3 decimal places is not amazing...