forked from AmbaPant/mantid
-
Notifications
You must be signed in to change notification settings - Fork 1
/
plot_functions.py
110 lines (90 loc) · 4.13 KB
/
plot_functions.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
# Mantid Repository : https://github.com/mantidproject/mantid
#
# Copyright © 2018 ISIS Rutherford Appleton Laboratory UKRI,
# NScD Oak Ridge National Laboratory, European Spallation Source,
# Institut Laue - Langevin & CSNS, Institute of High Energy Physics, CAS
# SPDX - License - Identifier: GPL - 3.0 +
"""
Plotting functions for use in IPython notebooks that are generated by MantidPlot
"""
import matplotlib.pyplot as plt
# Import Mantid
from mantid.simpleapi import *
import mantid.api as mapi
def _plot_with_options(axes_option, workspace, options_list, plot_number):
"""
Enable/disable legend, grid, limits according to
options (ops) for the given axes (ax).
Plot with or without errorbars.
"""
ws_plot = ConvertToPointData(workspace)
if options_list['errorbars']:
axes_option.errorbar(ws_plot.readX(0), ws_plot.readY(0),
yerr=ws_plot.readE(0), label=workspace.name())
else:
axes_option.plot(ws_plot.readX(0),
ws_plot.readY(0),
label=workspace.name())
axes_option.grid(options_list['grid'])
axes_option.set_xscale(options_list['xScale'])
axes_option.set_yscale(options_list['yScale'])
if options_list['xLimits'] != 'auto':
axes_option.set_xlim(options_list['xLimits'])
if options_list['yLimits'] != 'auto':
axes_option.set_ylim(options_list['yLimits'])
# If a list of titles was given, use it to title each subplot
if hasattr(options_list['title'], "__iter__"):
axes_option.set_title(options_list['title'][plot_number])
if options_list['legend'] and hasattr(options_list['legendLocation'], "__iter__"):
axes_option.legend(loc=options_list['legendLocation'][plot_number])
elif options_list['legend']:
axes_option.legend(loc=options_list['legendLocation'])
def plots(list_of_workspaces, *args, **kwargs):
"""
Create a figure with a subplot for each workspace given.
Workspaces within a group workspace are plotted together in the same subplot.
Examples:
plots(rr)
plots(rr, 'TheGraphTitle')
plots(rr, 'TheGraphTitle', grid=True, legend=True,
xScale='linear', yScale='log', xLimits=[0.008, 0.16])
plots(rr, sharedAxes = False, xLimits = [0, 0.1], yLimits = [1e-5, 2],
Title='ASF070_07 I=1A T=3K dq/q=2%',
legend=True, legendLocation=3, errorbars=False)
"""
if not hasattr(list_of_workspaces, "__iter__"):
list_of_workspaces = [list_of_workspaces]
ops = _process_arguments(args, kwargs)
# Create subplots for workspaces in the list
fig, axes_handle = plt.subplots(1,
len(list_of_workspaces),
sharey=ops['sharedAxes'],
figsize=(6 * len(list_of_workspaces), 4))
if not hasattr(axes_handle, "__iter__"):
axes_handle = [axes_handle]
for plot_number, workspace in enumerate(list_of_workspaces):
if isinstance(workspace, mapi.WorkspaceGroup):
# Plot grouped workspaces on the same axes
for sub_ws in workspace:
_plot_with_options(axes_handle[plot_number], sub_ws, ops, plot_number)
else:
_plot_with_options(axes_handle[plot_number], workspace, ops, plot_number)
# If a single title was given, use it to title the whole figure
if not hasattr(ops['title'], "__iter__"):
fig.suptitle(ops['title'])
plt.show()
return plt.gcf()
def _process_arguments(input_args, input_kwargs):
"""
Build a dictionary of plotting options
"""
key_list = ['title', 'grid', 'legend', 'legendLocation',
'xScale', 'yScale', 'xLimits', 'yLimits', 'sharedAxes', 'errorbars']
default_values = ['', True, True, 1, 'log', 'log', 'auto', 'auto', True, 'True']
# Fill ops with the default values
for i in range(len(input_args)): # copy in values provided in args
default_values[i] = input_args[i]
ops = dict(zip(key_list, default_values))
for k in ops.keys(): # copy in any key word given arguments
ops[k] = input_kwargs.get(k, ops[k])
return ops