Skip to content

Commit

Permalink
Merge pull request #2335 from dhomeier/python_export_wcs
Browse files Browse the repository at this point in the history 
Fix export to Python script to correctly use `Coordinates` or `WCS` in `reference_data`
  • Loading branch information
@astrofrog
astrofrog authored Oct 28, 2022
2 parents 66550d2 + 0cbfbd1 commit 68c5bca
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 29 deletions.
75 changes: 57 additions & 18 deletions glue/viewers/image/qt/tests/test_python_export.py
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)
19 changes: 8 additions & 11 deletions glue/viewers/image/viewer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit 68c5bca

Please sign in to comment.