def test_generate_workspace_retrieval_commands_does_not_repeat_retrieval_if_workspace_plotted_twice(
self, mock_plotted_name):
mock_plotted_name.return_value = ['ws1', 'ws2', 'ws1', 'ws-3']
expected_output = sorted([
"from mantid.api import AnalysisDataService as ADS\n",
"ws1 = ADS.retrieve('ws1')", "ws2 = ADS.retrieve('ws2')",
"ws_3 = ADS.retrieve('ws-3')"
])
self.assertEqual(expected_output,
sorted(generate_workspace_retrieval_commands(None)))
def generate_script(fig, exclude_headers=False):
"""
Generate a script to recreate a figure.
This currently only supports recreating artists that were plotted
from a Workspace. The format of the outputted script is as follows:
<Default Workbench script contents (imports)>
<Workspace retrieval from ADS>
fig, axes = plt.subplots()
axes.plot() or axes.errorbar()
axes.set_title()
axes.set_xlabel and axes.set_ylabel()
axes.set_xlim() and axes.set_ylim()
axes.set_xscale() and axes.set_yscale()
axes.legend().set_draggable(True) (if legend present, or just draggable() for earlier matplotlib versions)
plt.show()
:param fig: A matplotlib.pyplot.Figure object you want to create a script from
:param exclude_headers: Boolean. Set to True to ignore imports/headers
:return: A String. A script to recreate the given figure
"""
plot_commands = []
for ax in fig.get_axes():
if not isinstance(ax, MantidAxes) or not curve_in_ax(ax):
continue
ax_object_var = get_axes_object_variable(ax)
plot_commands.extend(get_plot_cmds(ax, ax_object_var)) # ax.plot
plot_commands.extend(get_title_cmds(ax, ax_object_var)) # ax.set_title
plot_commands.extend(get_axis_label_cmds(
ax, ax_object_var)) # ax.set_label
plot_commands.extend(get_axis_limit_cmds(ax,
ax_object_var)) # ax.set_lim
plot_commands.extend(get_axis_scale_cmds(
ax, ax_object_var)) # ax.set_scale
plot_commands.extend(get_legend_cmds(ax, ax_object_var)) # ax.legend
plot_commands.append('')
if not plot_commands:
return
cmds = [] if exclude_headers else [DEFAULT_SCRIPT_CONTENT]
cmds.extend(generate_workspace_retrieval_commands(fig) + [''])
cmds.append("{}, {} = {}".format(FIG_VARIABLE, AXES_VARIABLE,
generate_subplots_command(fig)))
cmds.extend(plot_commands)
cmds.append("fig.show()")
cmds.append(
"# Scripting Plots in Mantid: https://docs.mantidproject.org/nightly/plotting/scripting_plots.html"
)
return '\n'.join(cmds)
def generate_script(fig, exclude_headers=False):
"""
Generate a script to recreate a figure.
This currently only supports recreating artists that were plotted
from a Workspace. The format of the outputted script is as follows:
<Default Workbench script contents (imports)>
<Workspace retrieval from ADS>
fig, axes = plt.subplots()
axes.plot() or axes.errorbar()
axes.set_title()
axes.set_xlabel and axes.set_ylabel()
axes.set_xlim() and axes.set_ylim()
axes.set_xscale() and axes.set_yscale()
axes.legend().set_draggable(True) (if legend present, or just draggable() for earlier matplotlib versions)
plt.show()
:param fig: A matplotlib.pyplot.Figure object you want to create a script from
:param exclude_headers: Boolean. Set to True to ignore imports/headers
:return: A String. A script to recreate the given figure
"""
plot_commands = []
plot_headers = ['import matplotlib.pyplot as plt']
for ax in fig.get_axes():
if not isinstance(ax, MantidAxes):
continue
ax_object_var = get_axes_object_variable(ax)
if axes_type(ax) in [FigureType.Image]:
colormap_lines, colormap_headers = get_plot_2d_cmd(ax, ax_object_var) # ax.imshow or pcolormesh
plot_commands.extend(colormap_lines)
plot_headers.extend(colormap_headers)
else:
if not curve_in_ax(ax):
continue
plot_commands.extend(get_plot_cmds(ax, ax_object_var)) # ax.plot
plot_commands.extend(generate_tick_commands(ax))
plot_commands.extend(generate_grid_commands(ax))
plot_commands.extend(get_title_cmds(ax, ax_object_var)) # ax.set_title
plot_commands.extend(get_axis_label_cmds(ax, ax_object_var)) # ax.set_label
plot_commands.extend(get_axis_limit_cmds(ax, ax_object_var)) # ax.set_lim
plot_commands.extend(get_axis_scale_cmds(ax, ax_object_var)) # ax.set_scale
plot_commands.extend(get_legend_cmds(ax, ax_object_var)) # ax.legend
plot_commands.append('')
if not plot_commands:
return
fit_commands, fit_headers = get_fit_cmds(fig)
cmds = [] if exclude_headers else [DEFAULT_SCRIPT_CONTENT]
if exclude_headers and fit_headers:
cmds.extend(fit_headers)
if plot_headers:
cmds.extend(plot_headers)
if fit_commands:
cmds.append(FIT_DOCUMENTATION_STRING)
cmds.extend(fit_commands + [''])
cmds.extend(generate_workspace_retrieval_commands(fig) + [''])
cmds.append("{}, {} = {}".format(FIG_VARIABLE, AXES_VARIABLE, generate_subplots_command(fig)))
cmds.extend(plot_commands)
cmds.append("plt.show()")
cmds.append("# Scripting Plots in Mantid:")
cmds.append("# https://docs.mantidproject.org/tutorials/python_in_mantid/plotting/02_scripting_plots.html")
return '\n'.join(cmds)