from typing import List, Optional
from ..plotting.overlay_plots import overlay_grid_plot_2D
from ..plotting.plotting_constants import *
from ..sweep.sweep_results_1var import SweepResults1Var
from .constants import STANDARD_GRID_PLOT, TIME_GRID_PLOT, TITLE_MAPPING
from .models import MotorSweepData, SweepData1D
from .types import SweepDatatype
[docs]
class MotorSweepResults:
def __init__(self, sweep_data: MotorSweepData):
self.sweep_data = sweep_data
[docs]
def get_single_motor_sweep_data(self, motor_name: str) -> SweepData1D:
"""
Retrieve sweep results for a specific motor by name.
Parameters
----------
motor_name : str
Name of the motor to retrieve
Returns
-------
SweepData1D
Sweep data for the specified motor
"""
if motor_name in self.sweep_data.data_by_motor:
return self.sweep_data.data_by_motor[motor_name]
else:
raise ValueError(f"Motor {motor_name} not found in sweep results")
[docs]
def plot_single_motor_sweep_result(
self,
motor_name: str,
y_var: SweepDatatype = SweepDatatype.TOTAL_PTS,
*,
title: Optional[str] = None,
**kawrgs,
) -> None:
"""
Plot a single motor's sweep results.
Parameters
----------
motor_name : str
Name of the motor to plot
y_var : SweepDatatype, optional
SweepDatatype to plot on y-axis (default is SweepDatatype.TOTAL_PTS)
title : str, optional
Plot title
kawrgs : dict
Additional keyword arguments for the plot method, including: subtitle, x_label, y_label,
fit_curve, show_points, h_line, v_line, theme, font_config, layout_config, smoothing_config
Returns
-------
go.Figure
Plotly figure object
"""
sweep_results = SweepResults1Var(self.get_single_motor_sweep_data(motor_name))
return sweep_results.plot(
y_var=y_var,
title=title + " - " + motor_name if title else None,
**kawrgs,
)
[docs]
def plot_all(
self,
show_event_times: bool = False,
*,
title: str | None = None,
x_label: str | None = None,
y_label: str | None = None,
**kwargs,
) -> None:
"""
Plot sweep results of all motors on a single overlay grid.
Parameters
----------
show_event_times : bool, optional
Whether to show event times on the plot instead of points (default is False)
title : str, optional
Overall title for the grid
x_label : str, optional
Common x-axis label. Defaults to the first motor's `var_name`.
y_label : str, optional
Common y-axis label (e.g., "Points" or "Time (s)")
kwargs : dict
Additional keyword arguments for the overlay_grid_plot_2D function, including: rows, cols,
show_legend, fit_curve, layout_config, font_config, smoothing_config
Returns
-------
go.Figure
Plotly figure object
"""
motor_names = list(self.sweep_data.data_by_motor.keys())
# Build data dicts
x_data_dict = {}
y_data_dict = {}
variables = STANDARD_GRID_PLOT if not show_event_times else TIME_GRID_PLOT
for var in variables:
x_data_dict[var.value] = {}
y_data_dict[var.value] = {}
for motor_name in motor_names:
motor_data = self.sweep_data.data_by_motor[motor_name]
x_data_dict[var.value][motor_name] = motor_data.sweep_values
y_data_dict[var.value][motor_name] = getattr(motor_data, var.value)
# Build subplot titles using TITLE_MAPPING
subplot_titles = [TITLE_MAPPING.get(var.value) for var in variables]
# Title and labels
final_x_label = (
x_label or self.sweep_data.data_by_motor[motor_names[0]].var_name
)
final_y_label = y_label or ""
final_title = title or f"Motor Comparison (varying {final_x_label})"
# Call the generic overlay plotting function
return overlay_grid_plot_2D(
x_data_dict=x_data_dict,
y_data_dict=y_data_dict,
subplot_titles=subplot_titles,
title=final_title,
x_label=final_x_label,
y_label=final_y_label,
rows=2,
cols=3,
**kwargs,
)