diff --git a/src/mslice/plotting/plot_window/cut_plot.py b/src/mslice/plotting/plot_window/cut_plot.py index bc0dd5d0..ca8964b1 100644 --- a/src/mslice/plotting/plot_window/cut_plot.py +++ b/src/mslice/plotting/plot_window/cut_plot.py @@ -515,6 +515,7 @@ def on_newplot(self, plot_over, ws_name): self.plot_window.toggle_waterfall_edit() if not plot_over: self._reset_plot_window_options() + self.ws_name = ws_name all_lines = [line for container in line_containers for line in container.get_children()] for cached_lines in list(self._waterfall_cache.keys()): diff --git a/src/mslice/scripting/helperfunctions.py b/src/mslice/scripting/helperfunctions.py index a2ac5d25..22035e66 100644 --- a/src/mslice/scripting/helperfunctions.py +++ b/src/mslice/scripting/helperfunctions.py @@ -13,7 +13,7 @@ def header(plot_handler): """Creates a list of import statements to be used in the generated script header""" from mslice.plotting.plot_window.cut_plot import CutPlot from mslice.plotting.plot_window.slice_plot import SlicePlot - statements = ["# Python Script Generated by Mslice on {}\n".format(datetime.now().replace(microsecond=0))] + statements = [f"# Python Script Generated by Mslice on {datetime.now().replace(microsecond=0)}\n"] statements.append("\n".join(COMMON_PACKAGES)) if isinstance(plot_handler, SlicePlot) and plot_handler.colorbar_log is True: @@ -47,8 +47,8 @@ def add_plot_statements(script_lines, plot_handler, ax): add_slice_plot_statements(script_lines, plot_handler) add_overplot_statements(script_lines, plot_handler) elif isinstance(plot_handler, CutPlot): - add_cut_plot_statements(script_lines, plot_handler, ax) - add_overplot_statements(script_lines, plot_handler) + return_ws_vars = add_cut_plot_statements(script_lines, plot_handler, ax) + add_overplot_statements(script_lines, plot_handler, return_ws_vars) script_lines.append("mc.Show()\n") @@ -63,25 +63,25 @@ def add_slice_plot_statements(script_lines, plot_handler): energy_axis = str(slice.energy_axis) norm = slice.norm_to_one - script_lines.append('slice_ws = mc.Slice(ws_{}, Axis1="{}", Axis2="{}", NormToOne={})\n\n'.format( - plot_handler.ws_name.replace(".", "_"), momentum_axis, energy_axis, norm)) + script_lines.append(f'slice_ws = mc.Slice(ws_{plot_handler.ws_name.replace(".", "_")}, Axis1="{momentum_axis}", Axis2="{energy_axis}", ' + f'NormToOne={norm})\n\n') if plot_handler.intensity is True: intensity = IntensityCache.get_desc_from_type(plot_handler.intensity_type) if plot_handler.temp_dependent: - script_lines.append('mesh = ax.pcolormesh(slice_ws, cmap="{}", intensity="{}", temperature={})\n'.format( - cache[plot_handler.ws_name].colourmap, intensity, plot_handler.temp)) + script_lines.append(f'mesh = ax.pcolormesh(slice_ws, cmap="{cache[plot_handler.ws_name].colourmap}", intensity="{intensity}", ' + f'temperature={plot_handler.temp})\n') else: - script_lines.append('mesh = ax.pcolormesh(slice_ws, cmap="{}", intensity="{}")\n'.format( - cache[plot_handler.ws_name].colourmap, intensity)) + script_lines.append(f'mesh = ax.pcolormesh(slice_ws, cmap="{cache[plot_handler.ws_name].colourmap}", ' + f'intensity="{intensity}")\n') else: - script_lines.append('mesh = ax.pcolormesh(slice_ws, cmap="{}")\n'.format(cache[plot_handler.ws_name].colourmap)) + script_lines.append(f'mesh = ax.pcolormesh(slice_ws, cmap="{cache[plot_handler.ws_name].colourmap}")\n') - script_lines.append("mesh.set_clim({}, {})\n".format(*plot_handler.colorbar_range)) + script_lines.append(f"mesh.set_clim({plot_handler.colorbar_range[0]}, {plot_handler.colorbar_range[1]})\n") if plot_handler.colorbar_log: min, maximum = plot_handler.colorbar_range[0], plot_handler.colorbar_range[1] min = max(min, LOG_SCALE_MIN) - script_lines.append("mesh.set_norm(colors.LogNorm({}, {}))\n".format(min, maximum)) + script_lines.append(f"mesh.set_norm(colors.LogNorm({min}, {maximum}))\n") script_lines.append("cb = plt.colorbar(mesh, ax=ax)\n") script_lines.append(f"cb.set_label('{plot_handler.colorbar_label}', labelpad=20, rotation=270, picker=5, " @@ -91,7 +91,7 @@ def add_slice_plot_statements(script_lines, plot_handler): add_plot_options(script_lines, plot_handler) -def add_overplot_statements(script_lines, plot_handler): +def add_overplot_statements(script_lines, plot_handler, ws_vars=None): """Adds overplot line statements to the script if they were plotted""" ax = plot_handler._canvas.figure.gca() line_artists = ax.lines @@ -107,23 +107,23 @@ def add_overplot_statements(script_lines, plot_handler): recoil = True if rmm is not None or key in [1, 2, 4] else False cif = None # Does not yet account for CIF files + ws_var = ws_vars.pop(0) if ws_vars else f"'{plot_handler.ws_name}'" if recoil: if element is None: - script_lines.append(f"ax.recoil(workspace='{plot_handler.ws_name}', rmm={rmm}, color='{color}')\n") + script_lines.append(f"ax.recoil(workspace={ws_var}, rmm={rmm}, color='{color}')\n") else: - script_lines.append(f"ax.recoil(workspace='{plot_handler.ws_name}', element='{element}', color='{color}')\n") + script_lines.append(f"ax.recoil(workspace={ws_var}, element='{element}', color='{color}')\n") else: if cif is None: - script_lines.append(f"ax.bragg(workspace='{plot_handler.ws_name}', element='{element}', color='{color}')\n") - + script_lines.append(f"ax.bragg(workspace={ws_var}, element='{element}', color='{color}')\n") else: - script_lines.append(f"ax.bragg(workspace='{plot_handler.ws_name}', cif='{cif}', color='{color}')\n") + script_lines.append(f"ax.bragg(workspace={ws_var}, cif='{cif}', color='{color}')\n") def add_cut_plot_statements(script_lines, plot_handler, ax): """Adds cut specific statements to the script""" - add_cut_lines(script_lines, plot_handler, ax) + return_ws_vars = add_cut_lines(script_lines, plot_handler, ax) add_plot_options(script_lines, plot_handler) if plot_handler.is_changed("x_log"): @@ -133,14 +133,16 @@ def add_cut_plot_statements(script_lines, plot_handler, ax): if plot_handler.is_changed("y_log"): script_lines.append(f"ax.set_yscale('symlog', " f"linthresh=pow(10, np.floor(np.log10({plot_handler.y_axis_min}))))\n") + return return_ws_vars def add_cut_lines(script_lines, plot_handler, ax): cuts = plot_handler._cut_plotter_presenter._cut_cache_dict[ax] errorbars = plot_handler._canvas.figure.gca().containers intensity_correction = plot_handler.intensity_type - add_cut_lines_with_width(errorbars, script_lines, cuts, intensity_correction) + return_ws_vars = add_cut_lines_with_width(errorbars, script_lines, cuts, intensity_correction) hide_lines(script_lines, plot_handler, ax) + return return_ws_vars def hide_lines(script_lines, plot_handler, ax): @@ -173,13 +175,14 @@ def hide_lines(script_lines, plot_handler, ax): def add_cut_lines_with_width(errorbars, script_lines, cuts, intensity_correction): """Adds the cut statements for each interval of the cuts that were plotted""" index = 0 # Required as we run through the loop multiple times for each cut + return_ws_vars = [] for cut in cuts: integration_start = cut.integration_axis.start integration_end = cut.integration_axis.end cut_start, cut_end = integration_start, min(integration_start + cut.width, integration_end) intensity_range = (cut.intensity_start, cut.intensity_end) norm_to_one = cut.norm_to_one - algo_str = '' if 'Rebin' in cut.algorithm else ', Algorithm="{}"'.format(cut.algorithm) + algo_str = '' if 'Rebin' in cut.algorithm else f', Algorithm="{cut.algorithm}"' while cut_start != cut_end and index < len(errorbars): cut.integration_axis.start = cut_start @@ -196,24 +199,25 @@ def add_cut_lines_with_width(errorbars, script_lines, cuts, intensity_correction intensity_correction_arg = f"'{IntensityCache.get_desc_from_type(intensity_correction)}'" \ if not intensity_correction == IntensityType.SCATTERING_FUNCTION else False - script_lines.append('cut_ws_{} = mc.Cut(ws_{}, CutAxis="{}", IntegrationAxis="{}", ' - 'NormToOne={}{}, IntensityCorrection={}, SampleTemperature={})' - '\n'.format(index, replace_ws_special_chars(cut.parent_ws_name), cut_axis, integration_axis, - norm_to_one, algo_str, intensity_correction_arg, cut.raw_sample_temp)) - + cut_ws = f'cut_ws_{index}' + script_lines.append(f'{cut_ws} = mc.Cut(ws_{replace_ws_special_chars(cut.parent_ws_name)}, CutAxis="{cut_axis}", ' + f'IntegrationAxis="{integration_axis}", NormToOne={norm_to_one}{algo_str}, ' + f'IntensityCorrection={intensity_correction_arg}, SampleTemperature={cut.raw_sample_temp})\n') + return_ws_vars.append(cut_ws) + plot_over = False if index == 0 else True if intensity_range != (None, None): script_lines.append( - 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", ' - 'lw={}, intensity_range={})\n\n'.format(index, label, colour, marker, style, width, - intensity_range)) + f'ax.errorbar(cut_ws_{index}, label="{label}", color="{colour}", marker="{marker}", ls="{style}", ' + f'lw={width}, intensity_range={intensity_range}, plot_over={plot_over})\n\n') else: script_lines.append( - 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", ' - 'lw={})\n\n'.format(index, label, colour, marker, style, width)) + f'ax.errorbar(cut_ws_{index}, label="{label}", color="{colour}", marker="{marker}", ls="{style}", ' + f'lw={width}, plot_over={plot_over})\n\n') cut_start, cut_end = cut_end, min(cut_end + cut.width, integration_end) index += 1 cut.reset_integration_axis(cut.start, cut.end) + return return_ws_vars def add_plot_options(script_lines, plot_handler): @@ -227,16 +231,16 @@ def add_plot_options(script_lines, plot_handler): script_lines.append(f"ax.set_xlabel(r'{plot_handler.x_label}', fontsize={plot_handler.x_label_size})\n") if plot_handler.is_changed("y_grid"): - script_lines.append("ax.grid({}, axis='y')\n".format(plot_handler.y_grid)) + script_lines.append(f"ax.grid({plot_handler.y_grid}, axis='y')\n") if plot_handler.is_changed("x_grid"): - script_lines.append("ax.grid({}, axis='x')\n".format(plot_handler.x_grid)) + script_lines.append(f"ax.grid({plot_handler.x_grid}, axis='x')\n") if plot_handler.is_changed("y_range"): - script_lines.append("ax.set_ylim(bottom={}, top={})\n".format(*plot_handler.y_range)) + script_lines.append(f"ax.set_ylim(bottom={plot_handler.y_range[0]}, top={plot_handler.y_range[1]})\n") if plot_handler.is_changed("x_range"): - script_lines.append("ax.set_xlim(left={}, right={})\n".format(*plot_handler.x_range)) + script_lines.append(f"ax.set_xlim(left={plot_handler.x_range[0]}, right={plot_handler.x_range[1]})\n") if plot_handler.is_changed("y_range_font_size"): script_lines.append(f"ax.yaxis.set_tick_params(labelsize={plot_handler.y_range_font_size})\n") @@ -246,9 +250,7 @@ def add_plot_options(script_lines, plot_handler): from mslice.plotting.plot_window.cut_plot import CutPlot if isinstance(plot_handler, CutPlot) and plot_handler.is_changed("waterfall"): - script_lines.append("ax.set_waterfall({}, {}, {})\n".format(plot_handler.waterfall, - plot_handler.waterfall_x, - plot_handler.waterfall_y)) + script_lines.append(f"ax.set_waterfall({plot_handler.waterfall}, {plot_handler.waterfall_x},{plot_handler.waterfall_y})\n") def replace_ws_special_chars(workspace_name): diff --git a/tests/scripting_helperfunctions_test.py b/tests/scripting_helperfunctions_test.py index 117e096b..b4c140f1 100644 --- a/tests/scripting_helperfunctions_test.py +++ b/tests/scripting_helperfunctions_test.py @@ -15,6 +15,9 @@ class ScriptingHelperFunctionsTest(unittest.TestCase): + def tearDown(self) -> None: + plt.gcf().clf() + def assign_slice_parameters(self, plot_handler, intensity=True, temp_dependent=True): plot_handler.colorbar_label = 'colorbar_label' plot_handler.colorbar_label_size = 10 @@ -79,11 +82,12 @@ def test_that_add_plot_statements_works_as_expected_for_cuts(self, add_overplot, fig = mock.MagicMock() ax = fig.add_subplot(111, projection='mslice') + add_cut.return_value = mock.MagicMock() add_plot_statements(script_lines, plot_handler, ax) add_header.assert_called_once_with(script_lines, plot_handler) add_cut.assert_called_once_with(script_lines, plot_handler, ax) - add_overplot.assert_called_once_with(script_lines, plot_handler) + add_overplot.assert_called_once_with(script_lines, plot_handler, add_cut.return_value) self.assertIn("mc.Show()\n", script_lines) self.assertIn('fig = plt.gcf()\n', script_lines) @@ -172,39 +176,51 @@ def test_that_add_slice_plot_statements_works_with_intensity_and_no_temp_depende def test_that_add_overplot_statements_works_as_expected_with_recoil_element(self, gfm): plot_handler = gfm.get_active_figure().plot_handler plot_handler.add_mock_spec(SlicePlot) + plot_handler.ws_name = 'test_ws_name' self.assign_slice_parameters(plot_handler) plot_handler._canvas.figure.gca().lines = [Line2D([1, 2], [1, 2], label="Hydrogen")] - workspace_name = plot_handler.ws_name script_lines = [] add_overplot_statements(script_lines, plot_handler) - self.assertIn(f"ax.recoil(workspace='{workspace_name}', element='Hydrogen', color='C0')\n", script_lines) + self.assertIn(f"ax.recoil(workspace='{plot_handler.ws_name}', element='Hydrogen', color='C0')\n", script_lines) + + @mock.patch('mslice.cli._mslice_commands.GlobalFigureManager') + def test_that_add_overplot_statements_works_as_expected_with_recoil_element_with_cut_ws_var(self, gfm): + plot_handler = gfm.get_active_figure().plot_handler + plot_handler.add_mock_spec(CutPlot) + self.assign_cut_parameters(plot_handler) + plot_handler._canvas.figure.gca().lines = [Line2D([1, 2], [1, 2], label="Hydrogen")] + script_lines = [] + + add_overplot_statements(script_lines, plot_handler, ['test_ws_name']) + + self.assertIn("ax.recoil(workspace=test_ws_name, element='Hydrogen', color='C0')\n", script_lines) @mock.patch('mslice.cli._mslice_commands.GlobalFigureManager') def test_that_add_overplot_statements_works_as_expected_with_arbitrary_nuclei(self, gfm): plot_handler = gfm.get_active_figure().plot_handler plot_handler.add_mock_spec(SlicePlot) + plot_handler.ws_name = 'test_ws_name' self.assign_slice_parameters(plot_handler) plot_handler._canvas.figure.gca().lines = [Line2D([1, 2], [1, 2], label="Relative Mass 55", color="red")] - workspace_name = plot_handler.ws_name script_lines = [] add_overplot_statements(script_lines, plot_handler) - self.assertIn(f"ax.recoil(workspace='{workspace_name}', rmm=55, color='red')\n", script_lines) + self.assertIn(f"ax.recoil(workspace='{plot_handler.ws_name}', rmm=55, color='red')\n", script_lines) @mock.patch('mslice.cli._mslice_commands.GlobalFigureManager') def test_that_add_overplot_statements_works_as_expected_with_bragg_peaks_elements(self, gfm): plot_handler = gfm.get_active_figure().plot_handler plot_handler.add_mock_spec(SlicePlot) plot_handler._canvas.figure.gca().lines = [Line2D([1, 2], [1, 2], label="Tantalum", color="green")] - workspace_name = plot_handler.ws_name + plot_handler.ws_name = 'test_ws_name' script_lines = [] add_overplot_statements(script_lines, plot_handler) - self.assertIn(f"ax.bragg(workspace='{workspace_name}', element='Tantalum', color='green')\n", script_lines) + self.assertIn(f"ax.bragg(workspace='{plot_handler.ws_name}', element='Tantalum', color='green')\n", script_lines) @mock.patch('mslice.scripting.helperfunctions.add_plot_options') @mock.patch('mslice.scripting.helperfunctions.add_cut_lines') @@ -217,7 +233,9 @@ def test_that_add_cut_plot_statements_works_as_expected(self, gfm, add_cut_lines fig = mock.MagicMock() ax = fig.add_subplot(111, projection='mslice') - add_cut_plot_statements(script_lines, plot_handler, ax) + add_cut_lines.return_value = ['test_ws_var'] + + ret_val = add_cut_plot_statements(script_lines, plot_handler, ax) add_cut_lines.assert_called_once_with(script_lines, plot_handler, ax) add_plot_options.assert_called_once_with(script_lines, plot_handler) @@ -226,6 +244,7 @@ def test_that_add_cut_plot_statements_works_as_expected(self, gfm, add_cut_lines plot_handler.x_axis_min), script_lines) self.assertIn("ax.set_yscale('symlog', linthresh=pow(10, np.floor(np.log10({}))))\n".format( plot_handler.y_axis_min), script_lines) + self.assertEqual(['test_ws_var'], ret_val) @mock.patch('mslice.scripting.helperfunctions.add_cut_lines_with_width') @mock.patch('mslice.cli._mslice_commands.GlobalFigureManager') @@ -237,14 +256,17 @@ def test_that_add_cut_lines_works_as_expected(self, gfm, add_cut_lines_with_widt plot_handler.add_mock_spec(CutPlot) plot_handler.intensity_type = IntensityType.SCATTERING_FUNCTION + add_cut_lines_with_width.return_value = ['test_ws_var'] + script_lines = [] - add_cut_lines(script_lines, plot_handler, ax) + ret_val = add_cut_lines(script_lines, plot_handler, ax) cuts = plot_handler._cut_plotter_presenter._cut_cache_dict[ax] errorbars = plot_handler._canvas.figure.gca().containers add_cut_lines_with_width.assert_called_once_with(errorbars, script_lines, cuts, plot_handler.intensity_type) + self.assertEqual(['test_ws_var'], ret_val) @mock.patch('mslice.cli._mslice_commands.GlobalFigureManager') def test_that_add_plot_options_works_as_expected(self, gfm): @@ -291,7 +313,8 @@ def test_that_add_cut_lines_with_width_works_as_expected_without_intensity_range cuts = [cut] script_lines = [] - add_cut_lines_with_width(errorbars, script_lines, cuts, IntensityType.SCATTERING_FUNCTION) + ret_val = add_cut_lines_with_width(errorbars, script_lines, cuts, IntensityType.SCATTERING_FUNCTION) + pass self.assertIn( 'cut_ws_{} = mc.Cut(ws_{}, CutAxis="{}", IntegrationAxis="{}", NormToOne={}, IntensityCorrection={}, ' @@ -299,8 +322,10 @@ def test_that_add_cut_lines_with_width_works_as_expected_without_intensity_range False, None), script_lines) self.assertIn( - 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", lw={})\n\n'.format( - 0, 'errorbar_label', 'blue', None, '-', 1.5), script_lines) + 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", lw={}, plot_over={})\n\n'.format( + 0, 'errorbar_label', 'blue', None, '-', 1.5, False), script_lines) + + self.assertEqual(['cut_ws_0'], ret_val) def test_that_add_cut_lines_with_width_works_as_expected_with_intensity_range(self): x_data, y_data = np.arange(0, 10), np.arange(0, 10) @@ -312,12 +337,13 @@ def test_that_add_cut_lines_with_width_works_as_expected_with_intensity_range(se cuts = [cut] script_lines = [] - add_cut_lines_with_width(errorbars, script_lines, cuts, IntensityType.SCATTERING_FUNCTION) + ret_val = add_cut_lines_with_width(errorbars, script_lines, cuts, IntensityType.SCATTERING_FUNCTION) self.assertIn( 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", lw={}, ' - 'intensity_range={})\n\n'.format(0, 'errorbar_label', 'blue', None, '-', 1.5, (1.0, 2.0)), + 'intensity_range={}, plot_over={})\n\n'.format(0, 'errorbar_label', 'blue', None, '-', 1.5, (1.0, 2.0), False), script_lines) + self.assertEqual(['cut_ws_0'], ret_val) def test_that_add_cut_lines_with_width_works_as_expected_with_multiple_cuts(self): x_data, y_data = np.arange(0, 10), np.arange(0, 10) @@ -333,7 +359,7 @@ def test_that_add_cut_lines_with_width_works_as_expected_with_multiple_cuts(self cuts = [cut_0, cut_2] script_lines = [] - add_cut_lines_with_width(errorbars, script_lines, cuts, IntensityType.SCATTERING_FUNCTION) + ret_val = add_cut_lines_with_width(errorbars, script_lines, cuts, IntensityType.SCATTERING_FUNCTION) self.assertIn( 'cut_ws_{} = mc.Cut(ws_{}, CutAxis="{}", IntegrationAxis="{}", NormToOne={}, IntensityCorrection={}, ' @@ -350,8 +376,22 @@ def test_that_add_cut_lines_with_width_works_as_expected_with_multiple_cuts(self 'SampleTemperature={})\n'.format(2, 'ws_1', cuts[1].cut_axis, cuts[1].integration_axis, cuts[1].norm_to_one, False, None), script_lines) - # Each mc.Cut statement has a corresponding errorbar statement - self.assertEqual(len(script_lines), 6) + self.assertIn( + 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", lw={}, ' + 'intensity_range={}, plot_over={})\n\n'.format(0, 'error_label_0', 'blue', None, '-', 1.5, (1.0, 2.0), False), + script_lines) + + self.assertIn( + 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", lw={}, ' + 'intensity_range={}, plot_over={})\n\n'.format(1, 'error_label_1', 'blue', None, '-', 1.5, (1.0, 2.0), True), + script_lines) + + self.assertIn( + 'ax.errorbar(cut_ws_{}, label="{}", color="{}", marker="{}", ls="{}", lw={}, ' + 'intensity_range={}, plot_over={})\n\n'.format(2, 'error_label_2', 'blue', None, '-', 1.5, (1.0, 2.0), True), + script_lines) + + self.assertEqual(['cut_ws_0', 'cut_ws_1', 'cut_ws_2'], ret_val) def test_show_or_hide_containers_in_script(self): fig, ax = plt.subplots()