From ff3e02ad34b09aed7ebfabdbc09d68325527cfd1 Mon Sep 17 00:00:00 2001
From: Martyn Gigg <martyn.gigg@stfc.ac.uk>
Date: Thu, 2 Aug 2018 14:13:16 +0100
Subject: [PATCH] Add plotting functions from workspace names
---
.../workbench/plotting/figuremanager.py | 2 +-
.../workbench/plotting/figuretype.py | 1 +
.../workbench/plotting/figurewindow.py | 2 +-
.../workbench/workbench/plotting/functions.py | 139 ++++++++++++++----
.../workbench/plugins/workspacewidget.py | 38 +----
.../mantidqt/dialogs/spectraselectordialog.py | 4 +-
6 files changed, 123 insertions(+), 63 deletions(-)
diff --git a/qt/applications/workbench/workbench/plotting/figuremanager.py b/qt/applications/workbench/workbench/plotting/figuremanager.py
index 0fc22823072..d85fd35bd92 100644
--- a/qt/applications/workbench/workbench/plotting/figuremanager.py
+++ b/qt/applications/workbench/workbench/plotting/figuremanager.py
@@ -21,7 +21,7 @@ import matplotlib
from matplotlib.backend_bases import FigureManagerBase
from matplotlib.backends.backend_qt5agg import (FigureCanvasQTAgg, backend_version, draw_if_interactive, show) # noqa
from matplotlib._pylab_helpers import Gcf
-from qtpy.QtCore import Qt, QEvent, QObject, Signal
+from qtpy.QtCore import Qt, QObject
from qtpy.QtWidgets import QApplication, QLabel
from six import text_type
diff --git a/qt/applications/workbench/workbench/plotting/figuretype.py b/qt/applications/workbench/workbench/plotting/figuretype.py
index c034b628bf0..ce918175d54 100644
--- a/qt/applications/workbench/workbench/plotting/figuretype.py
+++ b/qt/applications/workbench/workbench/plotting/figuretype.py
@@ -22,6 +22,7 @@ from __future__ import absolute_import
from mantidqt.py3compat import Enum
+
class FigureType(Enum):
"""Enumerate possible types of Figure"""
# The values are irrelevant
diff --git a/qt/applications/workbench/workbench/plotting/figurewindow.py b/qt/applications/workbench/workbench/plotting/figurewindow.py
index 8433eaee1d9..1afe52c0d14 100644
--- a/qt/applications/workbench/workbench/plotting/figurewindow.py
+++ b/qt/applications/workbench/workbench/plotting/figurewindow.py
@@ -25,7 +25,7 @@ from qtpy.QtCore import QEvent, Signal
from qtpy.QtWidgets import QMainWindow
# local imports
-from .figuretype import axes_type, figure_type, FigureType
+from .figuretype import figure_type, FigureType
class FigureWindow(QMainWindow):
diff --git a/qt/applications/workbench/workbench/plotting/functions.py b/qt/applications/workbench/workbench/plotting/functions.py
index 7873b98f7c3..991948d5a03 100644
--- a/qt/applications/workbench/workbench/plotting/functions.py
+++ b/qt/applications/workbench/workbench/plotting/functions.py
@@ -23,11 +23,15 @@ import collections
import math
# 3rd party imports
-from mantid.api import MatrixWorkspace
+from mantid.api import AnalysisDataService, MatrixWorkspace
import matplotlib.pyplot as plt
from mantidqt.py3compat import is_text_string
+from mantidqt.dialogs.spectraselectordialog import get_spectra_selection
+from matplotlib.gridspec import GridSpec
+import numpy as np
# local imports
+from .figuretype import figure_type, FigureType
# -----------------------------------------------------------------------------
# Constants
@@ -43,6 +47,27 @@ SUBPLOT_HSPACE = 0.5
# 'Public' Functions
# -----------------------------------------------------------------------------
+def can_overplot():
+ """
+ Checks if overplotting on the current figure can proceed
+ with the given options
+
+ :return: A 2-tuple of boolean indicating compatability and
+ a string containing an error message if the current figure is not
+ compatible.
+ """
+ compatible = False
+ msg = "Unable to overplot on currently active plot type.\n" \
+ "Please select another plot."
+ fig = current_figure_or_none()
+ if fig is not None:
+ figtype = figure_type(fig)
+ if figtype is FigureType.Line or figtype is FigureType.Errorbar:
+ compatible, msg = True, None
+
+ return compatible, msg
+
+
def current_figure_or_none():
"""If an active figure exists then return it otherwise return None
@@ -76,8 +101,30 @@ def figure_title(workspaces, fig_num):
return wsname(first) + '-' + str(fig_num)
+def plot_from_names(names, errors, overplot, fig=None):
+ """
+ Given a list of names of workspaces, raise a dialog asking for the
+ a selection of what to plot and then plot it
+
+ :param names: A list of workspace names
+ :param errors: If true then error bars will be plotted on the points
+ :param overplot: If true then the add to the current figure if one
+ exists and it is a compatible figure
+ :param fig: If not None then use this figure object to plot
+ :return: The figure containing the plot or None if selection was cancelled
+ """
+ workspaces = AnalysisDataService.Instance().retrieveWorkspaces(names, unrollGroups=True)
+ selection = get_spectra_selection(workspaces)
+ if selection is None:
+ return None
+
+ return plot(selection.workspaces, spectrum_nums=selection.spectra,
+ wksp_indices=selection.wksp_indices,
+ errors=errors, overplot=overplot, fig=fig)
+
+
def plot(workspaces, spectrum_nums=None, wksp_indices=None, errors=False,
- overplot=False):
+ overplot=False, fig=None):
"""
Create a figure with a single subplot and for each workspace/index add a
line plot to the new axes. show() is called before returning the figure instance. A legend
@@ -88,7 +135,8 @@ def plot(workspaces, spectrum_nums=None, wksp_indices=None, errors=False,
:param wksp_indices: A list of workspace indexes (starts from 0)
:param errors: If true then error bars are added for each plot
:param overplot: If true then overplot over the current figure if one exists
- :returns: The figure containing the plots
+ :param fig: If not None then use this Figure object to plot
+ :return: The figure containing the plots
"""
# check inputs
_validate_plot_inputs(workspaces, spectrum_nums, wksp_indices)
@@ -97,13 +145,16 @@ def plot(workspaces, spectrum_nums=None, wksp_indices=None, errors=False,
else:
kw, nums = 'wkspIndex', wksp_indices
- # get/create the axes to hold the plot
- if overplot:
- ax = plt.gca(projection=PROJECTION)
- fig = ax.figure
+ if fig is None:
+ # get/create the axes to hold the plot
+ if overplot:
+ ax = plt.gca(projection=PROJECTION)
+ fig = ax.figure
+ else:
+ fig = plt.figure()
+ ax = fig.add_subplot(111, projection=PROJECTION)
else:
- fig = plt.figure()
- ax = fig.add_subplot(111, projection=PROJECTION)
+ ax = fig.gca()
# do the plotting
plot_fn = ax.errorbar if errors else ax.plot
@@ -121,14 +172,24 @@ def plot(workspaces, spectrum_nums=None, wksp_indices=None, errors=False,
return fig
-def pcolormesh(workspaces):
+def pcolormesh_from_names(names, fig=None):
"""
- Create a figure containing subplots
+ Create a figure containing pcolor subplots
+
+ :param names: A list of workspace names
+ :param fig: An optional figure to contain the new plots. Its current contents will be cleared
+ :returns: The figure containing the plots
+ """
+ return pcolormesh(AnalysisDataService.retrieveWorkspaces(names, unrollGroups=True),
+ fig=fig)
+
+
+def pcolormesh(workspaces, fig=None):
+ """
+ Create a figure containing pcolor subplots
:param workspaces: A list of workspace handles
- :param spectrum_nums: A list of spectrum number identifiers (general start from 1)
- :param wksp_indices: A list of workspace indexes (starts from 0)
- :param errors: If true then error bars are added for each plot
+ :param fig: An optional figure to contain the new plots. Its current contents will be cleared
:returns: The figure containing the plots
"""
# check inputs
@@ -137,17 +198,8 @@ def pcolormesh(workspaces):
# create a subplot of the appropriate number of dimensions
# extend in number of columns if the number of plottables is not a square number
workspaces_len = len(workspaces)
- square_side_len = int(math.ceil(math.sqrt(workspaces_len)))
- nrows, ncols = square_side_len, square_side_len
- if square_side_len*square_side_len != workspaces_len:
- # not a square number - square_side_len x square_side_len
- # will be large enough but we could end up with an empty
- # row so chop that off
- if workspaces_len <= (nrows-1)*ncols:
- nrows -= 1
+ fig, axes, nrows, ncols = _create_subplots(workspaces_len)
- fig, axes = plt.subplots(nrows, ncols, squeeze=False,
- subplot_kw=dict(projection=PROJECTION))
row_idx, col_idx = 0, 0
for subplot_idx in range(nrows*ncols):
ax = axes[row_idx][col_idx]
@@ -174,8 +226,9 @@ def pcolormesh(workspaces):
fig.show()
return fig
+# ----------------- Compatability functions ---------------------
+
-# Compatibility function for existing MantidPlot functionality
def plotSpectrum(workspaces, indices, distribution=None, error_bars=False,
type=None, window=None, clearWindow=None,
waterfall=False):
@@ -233,3 +286,39 @@ def _validate_pcolormesh_inputs(workspaces):
"""Raises a ValueError if any arguments have the incorrect types"""
if not isinstance(workspaces, MatrixWorkspace):
_raise_if_not_sequence(workspaces, 'Workspaces')
+
+
+def _create_subplots(nplots, fig=None):
+ """
+ Create a set of subplots suitable for a given number of plots. A stripped down
+ version of plt.subplots that can accept an existing figure instance.
+
+ :param nplots: The number of plots required
+ :param fig: An optional figure. It is cleared before plotting the new contents
+ :return: A 2-tuple of (fig, axes)
+ """
+ square_side_len = int(math.ceil(math.sqrt(nplots)))
+ nrows, ncols = square_side_len, square_side_len
+ if square_side_len*square_side_len != nplots:
+ # not a square number - square_side_len x square_side_len
+ # will be large enough but we could end up with an empty
+ # row so chop that off
+ if nplots <= (nrows-1)*ncols:
+ nrows -= 1
+
+ if fig is None:
+ fig = plt.figure()
+ else:
+ fig.clf()
+ # annoyling this repl
+ nplots = nrows*ncols
+ gs = GridSpec(nrows, ncols)
+ axes = np.empty(nplots, dtype=object)
+ ax0 = fig.add_subplot(gs[0, 0], projection=PROJECTION)
+ axes[0] = ax0
+ for i in range(1, nplots):
+ axes[i] = fig.add_subplot(gs[i // ncols, i % ncols],
+ projection=PROJECTION)
+ axes = axes.reshape(nrows, ncols)
+
+ return fig, axes, nrows, ncols
diff --git a/qt/applications/workbench/workbench/plugins/workspacewidget.py b/qt/applications/workbench/workbench/plugins/workspacewidget.py
index 2bdfc8e60e2..c8c7698117e 100644
--- a/qt/applications/workbench/workbench/plugins/workspacewidget.py
+++ b/qt/applications/workbench/workbench/plugins/workspacewidget.py
@@ -20,15 +20,13 @@ from __future__ import (absolute_import, unicode_literals)
import functools
# third-party library imports
-from mantid.api import AnalysisDataService, MatrixWorkspace, WorkspaceGroup
-from mantidqt.dialogs.spectraselectordialog import get_spectra_selection
+from mantid.api import AnalysisDataService
from mantidqt.widgets.workspacewidget.workspacetreewidget import WorkspaceTreeWidget
from qtpy.QtWidgets import QMessageBox, QVBoxLayout
# local package imports
from workbench.plugins.base import PluginWidget
-from workbench.plotting.figuretype import figure_type, FigureType
-from workbench.plotting.functions import current_figure_or_none, pcolormesh, plot
+from workbench.plotting.functions import can_overplot, pcolormesh, plot_from_names
class WorkspaceWidget(PluginWidget):
@@ -79,21 +77,12 @@ class WorkspaceWidget(PluginWidget):
exists and it is a compatible figure
"""
if overplot:
- compatible, error_msg = self._can_overplot()
+ compatible, error_msg = can_overplot()
if not compatible:
QMessageBox.warning(self, "", error_msg)
return
- try:
- workspaces = self._ads.retrieveWorkspaces(names, unrollGroups=True)
- selection = get_spectra_selection(workspaces, self)
- if selection is not None:
- plot(selection.workspaces, spectrum_nums=selection.spectra,
- wksp_indices=selection.wksp_indices,
- errors=errors, overplot=overplot)
- except BaseException:
- import traceback
- traceback.print_exc()
+ plot_from_names(names, errors, overplot)
def _do_plot_colorfill(self, names):
"""
@@ -106,22 +95,3 @@ class WorkspaceWidget(PluginWidget):
except BaseException:
import traceback
traceback.print_exc()
-
- def _can_overplot(self):
- """
- Checks if overplotting can proceed with the given options
-
- :return: A 2-tuple of boolean indicating compatability and
- a string containing an error message if the current figure is not
- compatible.
- """
- compatible = False
- msg = "Unable to overplot on currently active plot type.\n" \
- "Please select another plot."
- fig = current_figure_or_none()
- if fig is not None:
- figtype = figure_type(fig)
- if figtype is FigureType.Line or figtype is FigureType.Errorbar:
- compatible, msg = True, None
-
- return compatible, msg
diff --git a/qt/python/mantidqt/dialogs/spectraselectordialog.py b/qt/python/mantidqt/dialogs/spectraselectordialog.py
index 0bde6df7468..98544861c41 100644
--- a/qt/python/mantidqt/dialogs/spectraselectordialog.py
+++ b/qt/python/mantidqt/dialogs/spectraselectordialog.py
@@ -161,14 +161,14 @@ class SpectraSelectionDialog(SpectraSelectionDialogUIBase):
return self.selection is not None
-def get_spectra_selection(workspaces, parent_widget):
+def get_spectra_selection(workspaces, parent_widget=None):
"""Decides whether it is necessary to request user input
when asked to plot a list of workspaces. The input
dialog will only be shown in the case where all workspaces
have more than 1 spectrum
:param workspaces: A list of MatrixWorkspaces that will be plotted
- :param parent_widget: A parent_widget to use for the input selection dialog
+ :param parent_widget: An optional parent_widget to use for the input selection dialog
:returns: Either a SpectraSelection object containing the details of workspaces to plot or None indicating
the request was cancelled
"""
--
GitLab