Source code for suboptimumg.sweep.array_sweep_results

from typing import Optional

from ..plotting.bar_chart import plot_bar_chart
from .models import ArraySweepData, SweepData1D
from .types import SweepDatatype




class ArraySweepResults:
    def __init__(
        self,
        sweep_data: ArraySweepData,
    ):
        self.sweep_data = sweep_data



    def get_one_var_sweep_result(self, var_name: str) -> SweepData1D:
        for output in self.sweep_data.sweep_outputs:
            if output.var_name == var_name:
                return output
        raise ValueError(f"Variable {var_name} not found in sweep results")




    def plot(
        self,
        y_var: SweepDatatype = SweepDatatype.TOTAL_PTS,
        *,
        title: Optional[str] = None,
        y_label: Optional[str] = None,
        **kwargs,
    ):
        """
        Uses Plotly to create an interactive categorical dot plot showing points vs variable names.

        Parameters
        ----------
        y_var : SweepDatatype, optional
            Sweep datatype to plot on y-axis (default is SweepDatatype.TOTAL_PTS)
        title : str, optional
            Plot title (default is "Independent Variable Sweep Results")
        y_label : str, optional
            Y-axis label (default is y_var)
        **kwargs : dict
            Additional arguments passed to plot_bar_chart

        Returns
        -------
        go.Figure
            Plotly figure object
        """
        data_by_variable = []

        baseline_value = getattr(self.sweep_data.sweep_outputs[0], y_var.value)[0]
        for idx, p in enumerate(self.sweep_data.percent_steps):
            if p == 0:
                baseline_value = getattr(self.sweep_data.sweep_outputs[0], y_var.value)[
                    idx
                ]
                break

        # Build data for plotting from strongly-typed sweep data
        for sweep_output in self.sweep_data.sweep_outputs:
            y_data = getattr(sweep_output, y_var.value)

            data_by_variable.append(
                {
                    "var_name": sweep_output.var_name,
                    "points": y_data,
                    "plausible": sweep_output.plausible,
                }
            )

        # Generate final labels and title
        final_title = title if title else "Independent Variable Sweep Results"
        final_y_label = y_label if y_label is not None else y_var

        # Use generic bar chart function
        return plot_bar_chart(
            data_by_variable=data_by_variable,
            percent_steps=self.sweep_data.percent_steps,
            baseline_value=baseline_value,
            title=final_title,
            y_label=final_y_label,
            **kwargs,
        )