Source code for suboptimumg.sweep.models

from typing import Dict, List, Optional, Union

import numpy as np
from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator

from ..compsim.models import CompetitionData




class SweepParamConfig(BaseModel):
    """
    Contains the fields needed to specify a sweep variable
    """

    name: str = Field(description="Name of the parameter")
    min: Union[float, int] = Field(description="Lower bound")
    max: Union[float, int] = Field(description="Upper bound")
    steps: int = Field(description="Number of steps to sweep across")



    @model_validator(mode="after")
    def check_valid(self):
        if self.min >= self.max:
            raise ValueError(f"`min` = {self.min} must be less than `max` = {self.max}")
        return self


    model_config = ConfigDict(frozen=True, extra="forbid")





class SweepProcessInput1D(BaseModel):
    """Input model for 1D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    comp_data: CompetitionData = Field(
        description="Competition configuration data model"
    )
    var_1_name: str = Field(description="Name of the variable being swept")
    var_1_value: float = Field(description="Value of the variable for this iteration")
    dep_vals: Dict[str, float] = Field(
        description="Dictionary of dependent parameter values",
        default_factory=dict,
    )
    idx: int = Field(description="Index of this sweep iteration")





class SweepProcessInput2D(BaseModel):
    """Input model for 2D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    comp_data: CompetitionData = Field(
        description="Competition configuration data model"
    )
    var_1_name: str = Field(description="Name of the first variable being swept")
    var_1_value: float = Field(
        description="Value of the first variable for this iteration"
    )
    var_2_name: str = Field(description="Name of the second variable being swept")
    var_2_value: float = Field(
        description="Value of the second variable for this iteration"
    )
    dep_vals: Dict[str, float] = Field(
        description="Dictionary of dependent parameter values",
        default_factory=dict,
    )
    x_idx: int = Field(description="X-axis index of this sweep iteration")
    y_idx: int = Field(description="Y-axis index of this sweep iteration")





class SweepProcessOutput1D(BaseModel):
    """Output model for 1D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    idx: int = Field(description="Index of this sweep iteration")
    accel_pts: float = Field(description="Acceleration event points")
    skidpad_pts: float = Field(description="Skidpad event points")
    autoX_pts: float = Field(description="Autocross event points")
    endurance_pts: float = Field(description="Endurance event points")
    efficiency_pts: float = Field(description="Efficiency points")
    accel_t: float = Field(description="Acceleration event time")
    skidpad_t: float = Field(description="Skidpad event time")
    autoX_t: float = Field(description="Autocross event time")
    endurance_t: float = Field(description="Endurance event time")
    warnings: str = Field(
        description="Warning messages from the simulation", default=""
    )
    error: Optional[str] = Field(
        description="Error message if simulation failed", default=None
    )





class SweepProcessOutput2D(BaseModel):
    """Output model for 2D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    x_idx: int = Field(description="Index of first variable of this sweep iteration")
    y_idx: int = Field(description="Index of second variable of this sweep iteration")
    accel_pts: float = Field(description="Acceleration event points")
    skidpad_pts: float = Field(description="Skidpad event points")
    autoX_pts: float = Field(description="Autocross event points")
    endurance_pts: float = Field(description="Endurance event points")
    efficiency_pts: float = Field(description="Efficiency points")
    accel_t: float = Field(description="Acceleration event time")
    skidpad_t: float = Field(description="Skidpad event time")
    autoX_t: float = Field(description="Autocross event time")
    endurance_t: float = Field(description="Endurance event time")
    warnings: str = Field(
        description="Warning messages from the simulation", default=""
    )
    error: Optional[str] = Field(
        description="Error message if simulation failed", default=None
    )





class ArraySweepProcessInput(BaseModel):
    """Input model for array sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    comp_data: CompetitionData = Field(
        description="Competition configuration data model"
    )
    var_name: str = Field(description="Name of the variable being swept")
    step: float = Field(description="Value of the variable for this iteration")
    var_idx: int = Field(description="Index of the variable in the sweep array")
    step_idx: int = Field(description="Index of the step in the sweep")





class ArraySweepProcessOutput(BaseModel):
    """Output model for array sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    var_idx: int = Field(description="Index of the variable in the sweep array")
    step_idx: int = Field(description="Index of the step in the sweep")
    plausible: bool = Field(
        description="Whether the parameter configuration is plausible"
    )
    accel_pts: float = Field(description="Acceleration event points")
    skidpad_pts: float = Field(description="Skidpad event points")
    autoX_pts: float = Field(description="Autocross event points")
    endurance_pts: float = Field(description="Endurance event points")
    efficiency_pts: float = Field(description="Efficiency points")
    accel_t: float = Field(description="Acceleration event time")
    skidpad_t: float = Field(description="Skidpad event time")
    autoX_t: float = Field(description="Autocross event time")
    endurance_t: float = Field(description="Endurance event time")
    warnings: str = Field(
        description="Warning messages from the simulation", default=""
    )
    error: Optional[str] = Field(
        description="Error message if simulation failed", default=None
    )





class SweepData1D(BaseModel):
    """Strongly typed 1D sweep data container"""

    var_name: str = Field(description="Variable name being swept")
    sweep_values: np.ndarray = Field(
        description="Actual parameter values used in the sweep (x-values)"
    )

    accel_pts: np.ndarray = Field(description="Acceleration event points")
    skidpad_pts: np.ndarray = Field(description="Skidpad event points")
    autoX_pts: np.ndarray = Field(description="Autocross event points")
    endurance_pts: np.ndarray = Field(description="Endurance event points")
    efficiency_pts: np.ndarray = Field(description="Efficiency points")
    accel_t: np.ndarray = Field(description="Acceleration event time")
    skidpad_t: np.ndarray = Field(description="Skidpad event time")
    autoX_t: np.ndarray = Field(description="Autocross event time")
    endurance_t: np.ndarray = Field(description="Endurance event time")

    plausible: np.ndarray = Field(
        description="Plausibility check result for each sweep step"
    )

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    @computed_field
    @property
    def total_pts(self) -> np.ndarray:
        """Compute total points as sum of all point categories."""
        return (
            self.accel_pts
            + self.skidpad_pts
            + self.autoX_pts
            + self.endurance_pts
            + self.efficiency_pts
        )



    @classmethod
    def create(cls, var_name: str, sweep_values: np.ndarray) -> "SweepData1D":
        """Create a SweepData1D with pre-allocated numpy arrays of specified length."""
        length = len(sweep_values)
        return cls(
            var_name=var_name,
            sweep_values=sweep_values,
            accel_pts=np.zeros(length, dtype=float),
            skidpad_pts=np.zeros(length, dtype=float),
            autoX_pts=np.zeros(length, dtype=float),
            endurance_pts=np.zeros(length, dtype=float),
            efficiency_pts=np.zeros(length, dtype=float),
            accel_t=np.zeros(length, dtype=float),
            skidpad_t=np.zeros(length, dtype=float),
            autoX_t=np.zeros(length, dtype=float),
            endurance_t=np.zeros(length, dtype=float),
            plausible=np.zeros(length, dtype=bool),
        )






class SweepData2D(BaseModel):
    """Strongly typed 2D sweep data container"""

    var_name_1: str = Field(description="First variable name being swept")
    var_list_1: np.ndarray = Field(
        description="Actual values for first variable used in the sweep (x-values)"
    )
    var_name_2: str = Field(description="Second variable name being swept")
    var_list_2: np.ndarray = Field(
        description="Actual values for second variable used in the sweep (y-values)"
    )

    accel_pts: np.ndarray = Field(description="Acceleration event points")
    skidpad_pts: np.ndarray = Field(description="Skidpad event points")
    autoX_pts: np.ndarray = Field(description="Autocross event points")
    endurance_pts: np.ndarray = Field(description="Endurance event points")
    efficiency_pts: np.ndarray = Field(description="Efficiency points")
    accel_t: np.ndarray = Field(description="Acceleration event time")
    skidpad_t: np.ndarray = Field(description="Skidpad event time")
    autoX_t: np.ndarray = Field(description="Autocross event time")
    endurance_t: np.ndarray = Field(description="Endurance event time")

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    @computed_field
    @property
    def total_pts(self) -> np.ndarray:
        """Compute total points as sum of all point categories."""
        return (
            self.accel_pts
            + self.skidpad_pts
            + self.autoX_pts
            + self.endurance_pts
            + self.efficiency_pts
        )



    @classmethod
    def create(
        cls,
        var_name_1: str,
        var_list_1: np.ndarray,
        var_name_2: str,
        var_list_2: np.ndarray,
    ) -> "SweepData2D":
        """Create a SweepData2D with pre-allocated numpy arrays of specified dimensions."""
        dim1 = len(var_list_1)
        dim2 = len(var_list_2)
        return cls(
            var_name_1=var_name_1,
            var_list_1=var_list_1,
            var_name_2=var_name_2,
            var_list_2=var_list_2,
            accel_pts=np.zeros((dim1, dim2), dtype=float),
            skidpad_pts=np.zeros((dim1, dim2), dtype=float),
            autoX_pts=np.zeros((dim1, dim2), dtype=float),
            endurance_pts=np.zeros((dim1, dim2), dtype=float),
            efficiency_pts=np.zeros((dim1, dim2), dtype=float),
            accel_t=np.zeros((dim1, dim2), dtype=float),
            skidpad_t=np.zeros((dim1, dim2), dtype=float),
            autoX_t=np.zeros((dim1, dim2), dtype=float),
            endurance_t=np.zeros((dim1, dim2), dtype=float),
        )






class ArraySweepData(BaseModel):
    """Strongly typed container for array sweep data"""

    percent_steps: np.ndarray = Field(
        description="Percent step values used in the sweep (shared across all variables)"
    )
    sweep_outputs: List[SweepData1D] = Field(
        description="List of sweep output data for each variable"
    )

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)





class MotorSweepData(BaseModel):
    """
    Strongly typed container for motor sweep data
    """

    data_by_motor: Dict[str, SweepData1D] = Field(
        description="Motor name to sweep data map"
    )

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)