Merge pull request #2082 from jsub1/ProfileViewerExport

Profile Viewer Script Export

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')

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):
 

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()

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)