diff --git a/glue/viewers/image/qt/tests/test_python_export.py b/glue/viewers/image/qt/tests/test_python_export.py index 3ee339619..cfee1c89a 100644 --- a/glue/viewers/image/qt/tests/test_python_export.py +++ b/glue/viewers/image/qt/tests/test_python_export.py @@ -1,8 +1,11 @@ +import pytest import numpy as np import matplotlib.pyplot as plt from astropy.utils import NumpyRNGContext +from astropy.wcs import WCS from glue.core import Data, DataCollection +from glue.core.coordinates import AffineCoordinates from glue.app.qt.application import GlueApplication from glue.viewers.image.qt import ImageViewer from glue.viewers.matplotlib.qt.tests.test_python_export import BaseTestExportPython @@ -14,7 +17,13 @@ def setup_method(self, method): with NumpyRNGContext(12345): self.data = Data(cube=np.random.random((30, 50, 20))) - self.data_collection = DataCollection([self.data]) + # Create data versions with WCS and affine coordinates + matrix = np.array([[2, 3, 4, -1], [1, 2, 2, 2], [1, 1, 1, -2], [0, 0, 0, 1]]) + affine = AffineCoordinates(matrix, units=['Mm', 'Mm', 'km'], labels=['xw', 'yw', 'zw']) + + self.data_wcs = Data(label='cube', cube=self.data['cube'], coords=WCS(naxis=3)) + self.data_affine = Data(label='cube', cube=self.data['cube'], coords=affine) + self.data_collection = DataCollection([self.data, self.data_wcs, self.data_affine]) self.app = GlueApplication(self.data_collection) self.viewer = self.app.new_data_viewer(ImageViewer) self.viewer.add_data(self.data) @@ -34,19 +43,34 @@ def teardown_method(self, method): def assert_same(self, tmpdir, tol=0.1): BaseTestExportPython.assert_same(self, tmpdir, tol=tol) - def test_simple(self, tmpdir): + def viewer_load(self, coords): + if coords is not None: + self.viewer.add_data(getattr(self, f'data_{coords}')) + self.viewer.remove_data(self.data) + + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_simple(self, tmpdir, coords): + self.viewer_load(coords) self.assert_same(tmpdir) - def test_simple_legend(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_simple_legend(self, tmpdir, coords): + self.viewer_load(coords) self.viewer.state.show_legend = True self.assert_same(tmpdir) - def test_simple_att(self, tmpdir): - self.viewer.state.x_att = self.data.pixel_component_ids[1] - self.viewer.state.y_att = self.data.pixel_component_ids[0] + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_simple_att(self, tmpdir, coords): + self.viewer_load(coords) + self.viewer.state.x_att = self.viewer.state.reference_data.pixel_component_ids[1] + self.viewer.state.y_att = self.viewer.state.reference_data.pixel_component_ids[0] + if coords == 'affine': + pytest.xfail('Known issue with axis label rendering') self.assert_same(tmpdir) - def test_simple_visual(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_simple_visual(self, tmpdir, coords): + self.viewer_load(coords) self.viewer.state.legend.visible = True self.viewer.state.layers[0].cmap = plt.cm.RdBu self.viewer.state.layers[0].v_min = 0.2 @@ -57,31 +81,46 @@ def test_simple_visual(self, tmpdir): self.viewer.state.layers[0].bias = 0.6 self.assert_same(tmpdir) - def test_slice(self, tmpdir): - self.viewer.state.x_att = self.data.pixel_component_ids[1] - self.viewer.state.y_att = self.data.pixel_component_ids[0] + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_slice(self, tmpdir, coords): + self.viewer_load(coords) + self.viewer.state.x_att = self.viewer.state.reference_data.pixel_component_ids[1] + self.viewer.state.y_att = self.viewer.state.reference_data.pixel_component_ids[0] self.viewer.state.slices = (2, 3, 4) + if coords == 'affine': + pytest.xfail('Known issue with axis label rendering') self.assert_same(tmpdir) - def test_aspect(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_aspect(self, tmpdir, coords): + self.viewer_load(coords) self.viewer.state.aspect = 'auto' self.assert_same(tmpdir) - def test_subset(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_subset(self, tmpdir, coords): + self.viewer_load(coords) self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5) self.assert_same(tmpdir) - def test_subset_legend(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_subset_legend(self, tmpdir, coords): + self.viewer_load(coords) + self.data_collection.new_subset_group('mysubset', + self.viewer.state.reference_data.id['cube'] > 0.5) self.viewer.state.legend.visible = True - self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5) self.assert_same(tmpdir, tol=0.15) # transparency and such - def test_subset_slice(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_subset_slice(self, tmpdir, coords): + self.viewer_load(coords) self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5) - self.test_slice(tmpdir) + self.test_slice(tmpdir, coords) - def test_subset_transposed(self, tmpdir): + @pytest.mark.parametrize('coords', [None, 'wcs', 'affine']) + def test_subset_transposed(self, tmpdir, coords): + self.viewer_load(coords) + self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5) self.viewer.state.x_att = self.data.pixel_component_ids[0] self.viewer.state.y_att = self.data.pixel_component_ids[1] - self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5) self.assert_same(tmpdir) diff --git a/glue/viewers/image/viewer.py b/glue/viewers/image/viewer.py index 7fceee9c8..c4ffb3b0f 100644 --- a/glue/viewers/image/viewer.py +++ b/glue/viewers/image/viewer.py @@ -214,24 +214,21 @@ def _script_header(self): script = "" script += "fig, ax = init_mpl(wcs=True)\n" - script += "ax.set_aspect('{0}')\n".format(self.state.aspect) + script += f"ax.set_aspect('{self.state.aspect}')\n" script += '\ncomposite = CompositeArray()\n' - script += "image = imshow(ax, composite, origin='lower', interpolation='nearest', aspect='{0}')\n\n".format(self.state.aspect) + script += f"image = imshow(ax, composite, origin='lower', interpolation='nearest', aspect='{self.state.aspect}')\n\n" dindex = self.session.data_collection.index(self.state.reference_data) - script += "ref_data = data_collection[{0}]\n".format(dindex) + script += f"ref_data = data_collection[{dindex}]\n" - ref_coords = self.state.reference_data.coords - - if hasattr(ref_coords, 'wcs'): - script += "ax.reset_wcs(slices={0}, wcs=ref_data.coords.wcs)\n".format(self.state.wcsaxes_slice) - elif hasattr(ref_coords, 'wcsaxes_dict'): - raise NotImplementedError() - else: + if isinstance(self.state.reference_data.coords, (LegacyCoordinates, type(None))): imports.append('from glue.viewers.image.viewer import get_identity_wcs') - script += "ax.reset_wcs(slices={0}, wcs=get_identity_wcs(ref_data.ndim))\n".format(self.state.wcsaxes_slice) + ref_wcs = "get_identity_wcs(ref_data.ndim)" + else: + ref_wcs = "ref_data.coords" + script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs={ref_wcs})\n" script += "# for the legend\n" script += "legend_handles = []\n"