from enum import Enum
from typing import Annotated, Literal, Optional
import numpy as np
import numpy.typing as npt
from pydantic import (
BaseModel,
ConfigDict,
Field,
TypeAdapter,
computed_field,
field_validator,
)
from .subsystem_models import *
[docs]
class FittedCurve:
"""Polynomial fit of tabulated [x, y] data with Horner evaluation.
Parameters
----------
data : list of [x, y] pairs
Raw tabulated input data.
poly_order : int
Polynomial degree for the least-squares fit.
Attributes
----------
x_raw, y_raw : np.ndarray
Original input arrays.
poly_order : int
coeffs : np.ndarray
Polynomial coefficients in **descending** degree order
``[c_n, c_{n-1}, ..., c_1, c_0]``, matching ``np.polyfit``
output and the order Horner's method consumes directly.
"""
def __init__(self, data: list[list[float]], poly_order: int):
self.x_raw = np.array([p[0] for p in data], dtype=np.float64)
self.y_raw = np.array([p[1] for p in data], dtype=np.float64)
self.poly_order = poly_order
self.coeffs = np.polyfit(self.x_raw, self.y_raw, poly_order)
def __call__(
self, x: float | npt.NDArray[np.float64]
) -> np.float64 | npt.NDArray[np.float64]:
return self.evaluate(x)
[docs]
def evaluate(
self, x: float | npt.NDArray[np.float64]
) -> np.float64 | npt.NDArray[np.float64]:
"""Evaluate the fitted polynomial via Horner's method.
Parameters
----------
x : float | npt.NDArray[np.float64]
Input value(s) at which to evaluate.
Returns
-------
np.float64 | npt.NDArray[np.float64]
Evaluated polynomial value(s).
"""
x_arr = np.asarray(x, dtype=np.float64)
result = self.coeffs[0]
for c in self.coeffs[1:]:
result = result * x_arr + c
return result
[docs]
class VehicleModel(BaseModel):
model_config = ConfigDict(validate_assignment=True)
type: Literal["base"] = "base"
# Vehicle mass and geometry
mass: float = Field(gt=0, description="Total vehicle mass (kg)")
w_distr_b: float = Field(
gt=0, lt=1, description="Rear weight distribution percentage (0-1)"
)
cg_h: float = Field(gt=0, description="Center of gravity height (m)")
wb: float = Field(gt=0, description="Wheelbase (m)")
front_track: float = Field(gt=0, description="Front track width (m)")
rear_track: float = Field(gt=0, description="Rear track width (m)")
rolling_coeff: float = Field(
ge=0, description="Rolling resistance coefficient (unitless)"
)
@computed_field
@property
def w_distr_front(self) -> float:
return 1 - self.w_distr_b
@computed_field
@property
def track(self) -> float:
return (self.front_track + self.rear_track) / 2
# Subsystems
aero: AeroModel = Field(description="Aerodynamics configuration")
pwrtn: PowertrainModel = Field(description="Powertrain configuration")
accum: AccumulatorModel = Field(description="Accumulator configuration")
sus: SimpleSuspensionModel | ComplexSuspensionModel = Field(
discriminator="type", description="Suspension configuration"
)
tires: TireModel = Field(description="Tire parameters")
dri: DriverInterfaceModel = Field(description="Driver interface configuration")
chass: Optional[ChassisModel] = Field(None, description="Chassis configuration")
[docs]
@field_validator("rolling_coeff")
def validate_rolling_coeff(cls, v):
if v != 0.02:
raise ValueError(
"Currently only rolling coefficient of 0.02 is considered valid"
)
return v
[docs]
class AlignmentModel(BaseModel):
model_config = ConfigDict(validate_assignment=True)
camber_front_deg: float = Field(
description="Front camber angle (degrees, negative = cambered in)"
)
camber_rear_deg: float = Field(description="Rear camber angle (degrees)")
toe_front_deg: float = Field(
description="Front toe angle (degrees, positive = toe-in)"
)
toe_rear_deg: float = Field(description="Rear toe angle (degrees)")
[docs]
class SteeringSetupModel(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
ackermann: FittedCurve = Field(
description="Ackermann steering curve: steering wheel deg -> tire angle deg"
)
[docs]
@field_validator("ackermann", mode="before")
@classmethod
def coerce_ackermann(cls, v):
if isinstance(v, dict):
return FittedCurve(data=v["data"], poly_order=v["poly_order"])
return v
[docs]
class SuspensionSetupModel(BaseModel):
model_config = ConfigDict(validate_assignment=True)
roll_stiffness_front_Nm_per_rad: float = Field(
gt=0, description="Front roll stiffness (Nm/rad)"
)
roll_stiffness_rear_Nm_per_rad: float = Field(
gt=0, description="Rear roll stiffness (Nm/rad)"
)
heave_stiffness_front_N_per_m: float = Field(
gt=0, description="Front heave stiffness (N/m)"
)
heave_stiffness_rear_N_per_m: float = Field(
gt=0, description="Rear heave stiffness (N/m)"
)
anti_dive_pct: float = Field(ge=0, le=1, description="Anti-dive percentage (0-1)")
anti_squat_pct: float = Field(ge=0, le=1, description="Anti-squat percentage (0-1)")
motion_ratio_front: float = Field(
gt=0, description="Front motion ratio (wheel travel / shock travel)"
)
motion_ratio_rear: float = Field(
gt=0, description="Rear motion ratio (wheel travel / shock travel)"
)
[docs]
class ExtendedVehicleModel(VehicleModel):
"""VehicleModel extended with real-world session setup parameters.
Set ``type: irl_setup`` in the YAML to trigger this model.
Includes alignment, steering ackermann, and suspension setup fields
on top of the base car parameters.
"""
model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
type: Literal["irl_setup"] = "irl_setup"
alignment: AlignmentModel = Field(description="Wheel alignment configuration")
steering_setup: SteeringSetupModel = Field(
description="Steering geometry and ackermann curve"
)
suspension_setup: SuspensionSetupModel = Field(
description="Real-world suspension setup parameters"
)
AnyVehicleModel = TypeAdapter(
Annotated[VehicleModel | ExtendedVehicleModel, Field(discriminator="type")]
)