Source code for suboptimumg.sweep.pack_optimizer_models

from pyexpat import model
from typing import List, Optional, Tuple

import numpy.typing as npt
import pandas as pd
from attr import field
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    field_validator,
    model_validator,
)

from ..compsim.models import CompetitionResults
from .models import SweepParamConfig




class PackOptimizerConfig(BaseModel):
    """Configuration for the pack optimizer."""

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

    # Parameters to optimize (if None, defaults are chosen based on mode)
    parameters: List[SweepParamConfig] = Field(
        description="Explicit list of parameters to optimize",
    )

    max_iterations: int = Field(
        default=100, description="Maximum number of DE generations", gt=0
    )
    population_size: int = Field(
        default=15, description="DE population size multiplier", gt=0
    )
    seed: Optional[int] = Field(
        default=42, description="Random seed for reproducibility"
    )

    capacity_margin: float = Field(
        default=0.1,
        description="Safety margin subtracted from battery capacity (kWh)",
        ge=0,
    )





class FixedPackConfig(PackOptimizerConfig):
    """Configuration for fixed-mass optimization mode."""

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

    capacity: float = Field(
        description="Fixed battery capacity (kWh) when fixed_pack_mode=True",
    )

    @field_validator("parameters")
    @classmethod
    def _validate_parameters(
        cls, parameters: List[SweepParamConfig]
    ) -> List[SweepParamConfig]:
        for param in parameters:
            if "accum" in param.name:
                raise ValueError(
                    f"Invalid parameter '{param.name}' in fixed pack mode. "
                    "Only driving parameters can be optimized when capacity is fixed."
                )
        return parameters





class EvaluationRecord(BaseModel):
    """Record of a single objective function evaluation."""

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

    param_names: List[str] = Field(description="Names of the modified parameters")
    params: npt.NDArray = Field(description="List of modified parameters")

    mass: float = Field(description="Vehicle mass at this evaluation (kg)")
    capacity: float = Field(description="Battery capacity at this evaluation (kWh)")
    score: float = Field(
        description="Penalized score used by optimizer (points - penalty)"
    )
    feasible: bool = Field(
        description="Whether the solution is feasible (i.e., consumed energy < capacity)"
    )
    competition_results: Optional[CompetitionResults] = Field(
        default=None,
        description="Competition Results",
    )



    @model_validator(mode="after")
    def check_consistency(self):
        if len(self.params) != len(self.param_names):
            raise ValueError(
                "Length of params and param_names must be the same. "
                f"Got {len(self.params)} params and {len(self.param_names)} names."
            )
        return self






class PackOptimizationData(BaseModel):
    """Output of a single pack optimization run."""

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

    optimal_evaluation: Optional[EvaluationRecord] = Field(
        default=None,
        description="Best evaluation record by penalized score",
    )
    evaluation_history: List[EvaluationRecord] = Field(
        description="History of all evaluations during the optimization"
    )