diff --git a/glue/plugins/tools/__init__.py b/glue/plugins/tools/__init__.py index 8b17deb67..13dd6c3e4 100644 --- a/glue/plugins/tools/__init__.py +++ b/glue/plugins/tools/__init__.py @@ -16,3 +16,7 @@ def setup(): from glue.viewers.scatter.qt.data_viewer import ScatterViewer ScatterViewer.subtools = deepcopy(ScatterViewer.subtools) ScatterViewer.subtools['save'].append('save:python') + + from glue.viewers.profile.qt.data_viewer import ProfileViewer + ProfileViewer.subtools = deepcopy(ProfileViewer.subtools) + ProfileViewer.subtools['save'].append('save:python') diff --git a/glue/viewers/profile/layer_artist.py b/glue/viewers/profile/layer_artist.py index b10352858..4ea73a3d1 100644 --- a/glue/viewers/profile/layer_artist.py +++ b/glue/viewers/profile/layer_artist.py @@ -8,11 +8,13 @@ from glue.viewers.profile.state import ProfileLayerState from glue.viewers.matplotlib.layer_artist import MatplotlibLayerArtist from glue.core.exceptions import IncompatibleAttribute, IncompatibleDataException +from glue.viewers.profile.python_export import python_export_profile_layer class ProfileLayerArtist(MatplotlibLayerArtist): _layer_state_cls = ProfileLayerState + _python_exporter = python_export_profile_layer def __init__(self, axes, viewer_state, layer_state=None, layer=None): diff --git a/glue/viewers/profile/python_export.py b/glue/viewers/profile/python_export.py new file mode 100644 index 000000000..dce73519f --- /dev/null +++ b/glue/viewers/profile/python_export.py @@ -0,0 +1,45 @@ +from glue.viewers.common.python_export import serialize_options +from glue.core import Subset + + +def python_export_profile_layer(layer, *args): + + if len(layer.mpl_artists) == 0 or not layer.enabled or not layer.visible: + return [], None + + script = "" + imports = ["import numpy as np"] + + script += "# Calculate the profile of the data\n" + script += "profile_axis = {0}\n".format(layer._viewer_state.x_att_pixel.axis) + script += "collapsed_axes = tuple(i for i in range(layer_data.ndim) if i != profile_axis)\n" + if isinstance(layer.state.layer, Subset): + script += "base_data = layer_data.data\n" + script += "cid = base_data.find_component_id('{0}')\n".format(layer.state.attribute.label) + script += "profile_values = base_data.compute_statistic('{0}', cid, axis=collapsed_axes, subset_state=layer_data.subset_state)\n\n".format(layer._viewer_state.function) + else: + script += "cid = layer_data.find_component_id('{0}')\n".format(layer.state.attribute.label) + script += "profile_values = layer_data.compute_statistic('{0}', cid, axis=collapsed_axes)\n\n".format(layer._viewer_state.function) + + script += "# Extract the values for the x-axis\n" + script += "axis_view = [0] * layer_data.ndim\n" + script += "axis_view[profile_axis] = slice(None)\n" + script += "profile_x_values = layer_data['{0}', tuple(axis_view)]\n".format(layer._viewer_state.x_att) + script += "keep = ~np.isnan(profile_values) & ~np.isnan(profile_x_values)\n\n" + + if layer._viewer_state.normalize: + script += "# Normalize the profile data\n" + script += "vmax = np.nanmax(profile_values)\n" + script += "vmin = np.nanmin(profile_values)\n" + script += "profile_values = (profile_values - vmin)/(vmax - vmin)\n\n" + + script += "# Plot the profile\n" + plot_options = dict(color=layer.state.color, + linewidth=layer.state.linewidth, + alpha=layer.state.alpha, + zorder=layer.state.zorder, + drawstyle='steps-mid') + + script += "ax.plot(profile_x_values[keep], profile_values[keep], '-', {0})\n\n".format(serialize_options(plot_options)) + + return imports, script.strip() diff --git a/glue/viewers/profile/qt/tests/test_python_export.py b/glue/viewers/profile/qt/tests/test_python_export.py new file mode 100644 index 000000000..d7540561e --- /dev/null +++ b/glue/viewers/profile/qt/tests/test_python_export.py @@ -0,0 +1,76 @@ +from astropy.utils import NumpyRNGContext + +from glue.core import Data, DataCollection +from glue.app.qt.application import GlueApplication +from glue.viewers.matplotlib.qt.tests.test_python_export import BaseTestExportPython, random_with_nan +from glue.viewers.profile.qt import ProfileViewer +from glue.viewers.profile.tests.test_state import SimpleCoordinates + + +class TestExportPython(BaseTestExportPython): + + def setup_method(self, method): + + self.data = Data(label='d1') + self.data.coords = SimpleCoordinates() + with NumpyRNGContext(12345): + self.data['x'] = random_with_nan(48, 5).reshape((6, 4, 2)) + self.data['y'] = random_with_nan(48, 12).reshape((6, 4, 2)) + self.data_collection = DataCollection([self.data]) + self.app = GlueApplication(self.data_collection) + self.viewer = self.app.new_data_viewer(ProfileViewer) + self.viewer.add_data(self.data) + + def teardown_method(self, method): + self.viewer.close() + self.viewer = None + self.app.close() + self.app = None + + def test_simple(self, tmpdir): + self.assert_same(tmpdir) + + def test_color(self, tmpdir): + self.viewer.state.layers[0].color = '#ac0567' + self.assert_same(tmpdir) + + def test_linewidth(self, tmpdir): + self.viewer.state.layers[0].linewidth = 7.25 + self.assert_same(tmpdir) + + def test_max(self, tmpdir): + self.viewer.state.function = 'maximum' + self.assert_same(tmpdir) + + def test_min(self, tmpdir): + self.viewer.state.function = 'minimum' + self.assert_same(tmpdir) + + def test_mean(self, tmpdir): + self.viewer.state.function = 'mean' + self.assert_same(tmpdir) + + def test_median(self, tmpdir): + self.viewer.state.function = 'median' + self.assert_same(tmpdir) + + def test_sum(self, tmpdir): + self.viewer.state.function = 'sum' + self.assert_same(tmpdir) + + def test_normalization(self, tmpdir): + self.viewer.state.normalize = True + self.assert_same(tmpdir) + + def test_subset(self, tmpdir): + self.viewer.state.function = 'mean' + self.data_collection.new_subset_group('mysubset', self.data.id['x'] > 0.25) + self.assert_same(tmpdir) + + def test_xatt(self, tmpdir): + self.viewer.x_att = self.data.pixel_component_ids[1] + self.assert_same(tmpdir) + + def test_profile_att(self, tmpdir): + self.viewer.layers[0].state.attribute = self.data.id['y'] + self.assert_same(tmpdir)