Skip to content

Commit

Permalink
Merge pull request #940 from mantidproject/912_fix_mslice_scripting_i…
Browse files Browse the repository at this point in the history
…ssue

Fix mslice scripting issue when overplotting bragg peaks
  • Loading branch information
SilkeSchomann authored Sep 27, 2023
2 parents e85a876 + 2737647 commit 98719f9
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 55 deletions.
1 change: 1 addition & 0 deletions src/mslice/plotting/plot_window/cut_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()):
Expand Down
78 changes: 40 additions & 38 deletions src/mslice/scripting/helperfunctions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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")

Expand All @@ -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, "
Expand All @@ -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
Expand All @@ -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"):
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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")
Expand All @@ -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):
Expand Down
Loading

0 comments on commit 98719f9

Please sign in to comment.