From b2118dbcc3e7a64c28f8d40c5ea9fbe9812300fa Mon Sep 17 00:00:00 2001 From: SpacialTree Date: Wed, 20 Jul 2022 13:10:21 -0400 Subject: [PATCH] parent de76f8a68d87fbff09cbad6a3f727694a88bc3d3 author SpacialTree 1658337021 -0400 committer Adam Ginsburg (keflavich) 1663972401 -0400 Add cube mosaicking Created functions to mosaic cubes together. Made Requested Changes to Pull Request Delete .gitignore Revert "Delete .gitignore" This reverts commit 04ccb30b1df348719d679a5417fa9d0bf3820a81. 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 cube_utils.py Update spectral_cube/cube_utils.py Co-authored-by: Adam Ginsburg Update spectral_cube/cube_utils.py Co-authored-by: Adam Ginsburg Update spectral_cube/cube_utils.py Co-authored-by: Adam Ginsburg Update cube_utils.py += was causing an error Update spectral_cube/tests/test_regrid.py Co-authored-by: Adam Ginsburg Update spectral_cube/tests/test_regrid.py Co-authored-by: Adam Ginsburg Update spectral_cube/tests/test_regrid.py Co-authored-by: Adam Ginsburg 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 Update spectral_cube/cube_utils.py block size argument error Co-authored-by: Adam Ginsburg Update spectral_cube/tests/test_regrid.py Co-authored-by: Adam Ginsburg 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 * Update docs/creating_reading.rst Co-authored-by: Eric Koch * Update docs/creating_reading.rst Co-authored-by: Eric Koch * Update docs/creating_reading.rst Co-authored-by: Eric Koch * Update spectral_cube/lower_dimensional_structures.py Co-authored-by: Eric Koch * Update spectral_cube/lower_dimensional_structures.py * Update spectral_cube/lower_dimensional_structures.py Co-authored-by: Eric Koch * Update spectral_cube/lower_dimensional_structures.py Co-authored-by: Eric Koch 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 --- .gitignore | 1 + spectral_cube/cube_utils.py | 98 ++++++++++++++++++++++++++++++ spectral_cube/tests/test_regrid.py | 28 +++++++++ 3 files changed, 127 insertions(+) diff --git a/.gitignore b/.gitignore index 2eaa0dbc7..81d5fc7ed 100644 --- a/.gitignore +++ b/.gitignore @@ -60,3 +60,4 @@ docs/api */cython_version.py .tmp +.ipynb_checkpoints diff --git a/spectral_cube/cube_utils.py b/spectral_cube/cube_utils.py index c949b46f0..78099348f 100644 --- a/spectral_cube/cube_utils.py +++ b/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,100 @@ 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: + cube_repr = cube.reproject(header, block_size=[spectral_block_size, cube.shape[1], cube.shape[2]], **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 + 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[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 + # TODO: this should use the same cube type as cube1 + cube = cube1.__class__(data=final_array*cube1.unit, wcs=WCS(header)) + return cube diff --git a/spectral_cube/tests/test_regrid.py b/spectral_cube/tests/test_regrid.py index e1f70d642..a990134ef 100644 --- a/spectral_cube/tests/test_regrid.py +++ b/spectral_cube/tests/test_regrid.py @@ -26,6 +26,7 @@ from .. import SpectralCube 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 @@ -585,3 +586,30 @@ def test_reproject_3D_memory(): assert result.wcs.wcs.crval[0] == 0.001 assert result.wcs.wcs.crpix[0] == 2. + +@pytest.mark.parametrize('use_memmap', (True, False)) +def test_mosaic_cubes(use_memmap, data_adv, use_dask): + + pytest.importorskip('reproject') + + # Read in data to use + cube, data = cube_and_raw(data_adv, use_dask=use_dask) + + from reproject.mosaicking import find_optimal_celestial_wcs + + 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.):,:] + + result = mosaic_cubes([part1, part2], order='nearest-neighbor') + + # 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 comaprable + np.testing.assert_almost_equal(result.filled_data[:], cube.filled_data[:]) \ No newline at end of file