From 63044486b369ae6abbcab628057342d3eb159ca0 Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Tue, 25 Jan 2022 11:18:34 -0500 Subject: [PATCH 1/5] Fix bug in translating 3D Spectrum1D with only spectral axis defined (rather than full WCS) --- glue_astronomy/translators/spectrum1d.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/glue_astronomy/translators/spectrum1d.py b/glue_astronomy/translators/spectrum1d.py index 65875c1..d096e63 100644 --- a/glue_astronomy/translators/spectrum1d.py +++ b/glue_astronomy/translators/spectrum1d.py @@ -126,7 +126,10 @@ def to_data(self, obj): # Swap the spectral axis to first here. to_object doesn't need this because # Spectrum1D does it automatically on initialization. if len(obj.flux.shape) == 3: - data = Data(coords=obj.wcs.swapaxes(-1, 0)) + # It's possible to have a 3D Spectrum1D with only a spectral axis defined + # rather than a full WCS, in which case we don't need to swap the WCS + if obj.wcs.world_n_dim == 3: + data = Data(coords=obj.wcs.swapaxes(-1, 0)) data['flux'] = np.swapaxes(obj.flux, -1, 0) data.get_component('flux').units = str(obj.unit) else: From 4883daca3331994af4ec5b5e172149ae1f411391 Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Wed, 26 Jan 2022 10:44:50 -0500 Subject: [PATCH 2/5] Generalizing WCS padding with NDCube CompoundLowLevelWCS --- glue_astronomy/translators/spectrum1d.py | 124 ++++++----------------- 1 file changed, 31 insertions(+), 93 deletions(-) diff --git a/glue_astronomy/translators/spectrum1d.py b/glue_astronomy/translators/spectrum1d.py index d096e63..4174ad3 100644 --- a/glue_astronomy/translators/spectrum1d.py +++ b/glue_astronomy/translators/spectrum1d.py @@ -4,6 +4,7 @@ from glue.core import Data, Subset from gwcs import WCS as GWCS +from gwcs.coordinate_frames import CoordinateFrame from astropy.wcs import WCS from astropy import units as u @@ -11,6 +12,9 @@ from astropy.nddata import StdDevUncertainty, InverseVariance, VarianceUncertainty from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper from astropy.wcs.wcsapi import HighLevelWCSMixin, BaseHighLevelWCS +from astropy.modeling import models + +from ndcube.wcs.wrappers import CompoundLowLevelWCS from glue_astronomy.spectral_coordinates import SpectralCoordinates @@ -32,89 +36,23 @@ 'custom:spect.doplerVeloc.beta': 'Beta'} -class PaddedSpectrumWCS(BaseWCSWrapper, HighLevelWCSMixin): - - # Spectrum1D can use a 1D spectral WCS even for n-dimensional - # datasets while glue always needs the dimensionality to match, - # so this class pads the WCS so that it is n-dimensional. - - # NOTE: for now this only handles padding the WCS into 2D WCS. Rather than - # generalize this we can just remove this class and use CompoundLowLevelWCS - # from NDCube once it is in a released version. - - def __init__(self, wcs): - self.spectral_wcs = wcs - - @property - def pixel_n_dim(self): - return 2 - - @property - def world_n_dim(self): - return 2 - - @property - def world_axis_physical_types(self): - return [self.spectral_wcs.world_axis_physical_types[0], None] - - @property - def world_axis_units(self): - return (self.spectral_wcs.world_axis_units[0], None) - - def pixel_to_world_values(self, *pixel_arrays): - # The ravel and reshape are needed because of - # https://github.com/astropy/astropy/issues/12154 - px = np.array(pixel_arrays[0]) - world_arrays = [self.spectral_wcs.pixel_to_world_values(px.ravel()).reshape(px.shape), - pixel_arrays[1]] - return tuple(world_arrays) - - def world_to_pixel_values(self, *world_arrays): - # The ravel and reshape are needed because of - # https://github.com/astropy/astropy/issues/12154 - wx = np.array(world_arrays[0]) - pixel_arrays = [self.spectral_wcs.world_to_pixel_values(wx.ravel()).reshape(wx.shape), - world_arrays[1]] - return tuple(pixel_arrays) - - @property - def world_axis_object_components(self): - return [ - self.spectral_wcs.world_axis_object_components[0], - ('spatial', 'value', 'value') - ] +class PaddedSpectrumWCS(CompoundLowLevelWCS): - @property - def world_axis_object_classes(self): - spectral_key = self.spectral_wcs.world_axis_object_components[0][0] - return { - spectral_key: self.spectral_wcs.world_axis_object_classes[spectral_key], - 'spatial': (u.Quantity, (), {'unit': u.pixel}) - } + def __init__(self, spectral_wcs, n_extra_axes): + self.spectral_wcs = spectral_wcs + frame1 = CoordinateFrame(n_extra_axes, ['SPATIAL']*n_extra_axes, + np.arange(n_extra_axes), unit=[u.pix]*n_extra_axes, + name="Dummy1") + frame2 = CoordinateFrame(n_extra_axes, ['SPATIAL']*n_extra_axes, + np.arange(n_extra_axes), unit=[u.pix]*n_extra_axes, + name="Dummy2") + frame2frame = models.Multiply(1) + if n_extra_axes > 1: + for i in range(n_extra_axes-1): + frame2frame = frame2frame & models.Multiply(1) - @property - def pixel_shape(self): - return None - - @property - def pixel_bounds(self): - return None - - @property - def pixel_axis_names(self): - return tuple([self.spectral_wcs.pixel_axis_names[0], 'spatial']) - - @property - def world_axis_names(self): - return (UCD_TO_SPECTRAL_NAME[self.spectral_wcs.world_axis_physical_types[0]], 'Offset') - - @property - def axis_correlation_matrix(self): - return np.identity(2).astype('bool') - - @property - def serialized_classes(self): - return False + pad_wcs = GWCS([(frame1, frame2frame), (frame2, None)]) + super().__init__(pad_wcs, spectral_wcs) @data_translator(Spectrum1D) @@ -125,18 +63,20 @@ def to_data(self, obj): # Glue expects spectral axis first for cubes (opposite of specutils). # Swap the spectral axis to first here. to_object doesn't need this because # Spectrum1D does it automatically on initialization. - if len(obj.flux.shape) == 3: + if obj.flux.ndim > 1: # It's possible to have a 3D Spectrum1D with only a spectral axis defined - # rather than a full WCS, in which case we don't need to swap the WCS - if obj.wcs.world_n_dim == 3: + # rather than a full WCS, in which case we need to pad the WCS to match + # the dimensionality of the flux array. + if obj.wcs.world_n_dim == obj.flux.ndim: data = Data(coords=obj.wcs.swapaxes(-1, 0)) + else: + n_extra = obj.flux.ndim - obj.wcs.world_n_dim + data = Data(coords=PaddedSpectrumWCS(obj.wcs, n_extra)) data['flux'] = np.swapaxes(obj.flux, -1, 0) data.get_component('flux').units = str(obj.unit) else: if obj.flux.ndim == 1 and obj.wcs.world_n_dim == 1 and isinstance(obj.wcs, GWCS): data = Data(coords=SpectralCoordinates(obj.spectral_axis)) - elif obj.flux.ndim == 2 and obj.wcs.world_n_dim == 1: - data = Data(coords=PaddedSpectrumWCS(obj.wcs)) else: data = Data(coords=obj.wcs) data['flux'] = obj.flux @@ -153,7 +93,7 @@ def to_data(self, obj): # Include mask if it exists if obj.mask is not None: - if len(obj.flux.shape) == 3: + if len(obj.flux.shape) > 1: data['mask'] = np.swapaxes(obj.mask, -1, 0) else: data['mask'] = obj.mask @@ -186,12 +126,10 @@ def to_object(self, data_or_subset, attribute=None, statistic='mean'): if data.ndim < 2 and statistic is not None: statistic = None - if statistic is None and isinstance(data.coords, BaseHighLevelWCS): - - if isinstance(data.coords, PaddedSpectrumWCS): - kwargs = {'wcs': data.coords.spectral_wcs} - else: - kwargs = {'wcs': data.coords} + if statistic is None and isinstance(data.coords, PaddedSpectrumWCS): + kwargs = {'wcs': data.coords.spectral_wcs} + elif statistic is None and isinstance(data.coords, BaseHighLevelWCS): + kwargs = {'wcs': data.coords} elif statistic is not None: From 548171964874bf1765af633343cc33768119247b Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Wed, 26 Jan 2022 10:48:56 -0500 Subject: [PATCH 3/5] Fix codestyle checks --- glue_astronomy/translators/spectrum1d.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/glue_astronomy/translators/spectrum1d.py b/glue_astronomy/translators/spectrum1d.py index 4174ad3..48b3ec5 100644 --- a/glue_astronomy/translators/spectrum1d.py +++ b/glue_astronomy/translators/spectrum1d.py @@ -10,8 +10,7 @@ from astropy import units as u from astropy.wcs import WCSSUB_SPECTRAL from astropy.nddata import StdDevUncertainty, InverseVariance, VarianceUncertainty -from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper -from astropy.wcs.wcsapi import HighLevelWCSMixin, BaseHighLevelWCS +from astropy.wcs.wcsapi import BaseHighLevelWCS from astropy.modeling import models from ndcube.wcs.wrappers import CompoundLowLevelWCS @@ -41,11 +40,11 @@ class PaddedSpectrumWCS(CompoundLowLevelWCS): def __init__(self, spectral_wcs, n_extra_axes): self.spectral_wcs = spectral_wcs frame1 = CoordinateFrame(n_extra_axes, ['SPATIAL']*n_extra_axes, - np.arange(n_extra_axes), unit=[u.pix]*n_extra_axes, - name="Dummy1") + np.arange(n_extra_axes), unit=[u.pix]*n_extra_axes, + name="Dummy1") frame2 = CoordinateFrame(n_extra_axes, ['SPATIAL']*n_extra_axes, - np.arange(n_extra_axes), unit=[u.pix]*n_extra_axes, - name="Dummy2") + np.arange(n_extra_axes), unit=[u.pix]*n_extra_axes, + name="Dummy2") frame2frame = models.Multiply(1) if n_extra_axes > 1: for i in range(n_extra_axes-1): @@ -128,7 +127,7 @@ def to_object(self, data_or_subset, attribute=None, statistic='mean'): if statistic is None and isinstance(data.coords, PaddedSpectrumWCS): kwargs = {'wcs': data.coords.spectral_wcs} - elif statistic is None and isinstance(data.coords, BaseHighLevelWCS): + elif statistic is None and isinstance(data.coords, BaseHighLevelWCS): kwargs = {'wcs': data.coords} elif statistic is not None: From d772700dc69fb80eacbe0899cf7ffee793835a1c Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Wed, 26 Jan 2022 11:47:48 -0500 Subject: [PATCH 4/5] Swap spectral axis back to last manually if using PaddedSpectrumWCS --- glue_astronomy/translators/spectrum1d.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/glue_astronomy/translators/spectrum1d.py b/glue_astronomy/translators/spectrum1d.py index 48b3ec5..2d360aa 100644 --- a/glue_astronomy/translators/spectrum1d.py +++ b/glue_astronomy/translators/spectrum1d.py @@ -125,8 +125,12 @@ def to_object(self, data_or_subset, attribute=None, statistic='mean'): if data.ndim < 2 and statistic is not None: statistic = None + manual_swap = None + if statistic is None and isinstance(data.coords, PaddedSpectrumWCS): kwargs = {'wcs': data.coords.spectral_wcs} + if data.ndim > 1: + manual_swap = True elif statistic is None and isinstance(data.coords, BaseHighLevelWCS): kwargs = {'wcs': data.coords} @@ -202,6 +206,9 @@ def parse_attributes(attributes): mask = np.all(mask, collapse_axes) else: values = data.get_data(attribute) + if manual_swap: + # In this case we need to move the spectral axis back to last + values = np.swapaxes(values, -1, 0) attribute_label = attribute.label From 312a840bf6ef46553f1699f72aecdb19ce143398 Mon Sep 17 00:00:00 2001 From: Derek Homeier Date: Tue, 22 Feb 2022 03:30:56 +0100 Subject: [PATCH 5/5] DNM: attempt to fix spectrum1d_2d; catch WCS warnings --- .../translators/tests/test_spectrum1d.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/glue_astronomy/translators/tests/test_spectrum1d.py b/glue_astronomy/translators/tests/test_spectrum1d.py index 88685b4..a77d0de 100644 --- a/glue_astronomy/translators/tests/test_spectrum1d.py +++ b/glue_astronomy/translators/tests/test_spectrum1d.py @@ -1,4 +1,5 @@ import pytest +import warnings import numpy as np from numpy.testing import assert_allclose, assert_equal @@ -125,7 +126,7 @@ def test_to_spectrum1d_default_attribute(): @pytest.mark.parametrize('mode', ('wcs1d', 'wcs3d', 'lookup')) -def test_from_spectrum1d(mode): +def test_from_spectrum1d(mode, recwarn): if mode == 'wcs3d': # This test is intended to be run with the version of Spectrum1D based @@ -195,10 +196,16 @@ def test_from_spectrum1d(mode): print(uncertainty) assert_quantity_allclose(spec_new.uncertainty.quantity, np.ones((5, 4, 4))*0.01*u.Jy**2) + + assert len(recwarn) == 3 + for w in recwarn: + assert issubclass(w.category, UserWarning) + assert "Input WCS indicates that the spectral axis is not last." in str(w.message) else: assert_quantity_allclose(spec_new.flux, [2, 3, 4, 5] * u.Jy) assert spec_new.uncertainty is not None assert_quantity_allclose(spec_new.uncertainty.quantity, [0.1, 0.1, 0.1, 0.1] * u.Jy**2) + assert len(recwarn) == 0 def test_spectrum1d_2d_data(): @@ -207,6 +214,7 @@ def test_spectrum1d_2d_data(): # Note that Spectrum1D will typically have a 1D spectral WCS even if the # data is N-dimensional, so we need to pad the WCS before passing it to # glue and un-pad it when translating back. + # Also Spectrum1D.flux has the spectral axis along last dimension, not first. # We test both the case where the WCS is 2D and the case where it is 1D @@ -215,7 +223,7 @@ def test_spectrum1d_2d_data(): wcs.wcs.cdelt = [10] wcs.wcs.set() - flux = np.ones((3, 2)) * u.Unit('Jy') + flux = np.arange(1, 7).reshape((3, 2)) * u.Unit('Jy') spec = Spectrum1D(flux, wcs=wcs, meta={'instrument': 'spamcam'}) @@ -231,7 +239,7 @@ def test_spectrum1d_2d_data(): assert isinstance(data, Data) assert len(data.main_components) == 1 assert data.main_components[0].label == 'flux' - assert_allclose(data['flux'], flux.value) + assert_allclose(data['flux'], flux.value.swapaxes(-1, 0)) assert data.coords.pixel_n_dim == 2 assert data.coords.world_n_dim == 2 @@ -240,11 +248,12 @@ def test_spectrum1d_2d_data(): assert data.coordinate_components[0].label == 'Pixel Axis 0 [y]' assert data.coordinate_components[1].label == 'Pixel Axis 1 [x]' - assert data.coordinate_components[2].label == 'Offset' - assert data.coordinate_components[3].label == 'Frequency' + assert data.coordinate_components[2].label == 'World 0' + assert data.coordinate_components[3].label == 'World 1' - assert_equal(data['Offset'], [[0, 0], [1, 1], [2, 2]]) - assert_equal(data['Frequency'], [[10, 20], [10, 20], [10, 20]]) + assert_equal(data['World 0'], [[10, 10, 10], [10, 10, 10]]) + assert data['World 1'].shape == (2, 3) + assert_equal(data['World 1'], [[20, 20, 20], [20, 20, 20]]) s, o = data.coords.pixel_to_world(1, 2) assert isinstance(s, SpectralCoord)