from __future__ import annotations
from typing import Literal
import numpy as np
from numpy.typing import NDArray
from pydantic import BaseModel, ConfigDict, Field
from ..vehicle.irl_vehicle import IrlCar
from .kinematics import BODY_SLIP_ANGLE
from .steering import FRONT_BICYCLE_MODEL_STEER_ANGLE
CaScheduler = Literal["none", "abs_alpha_per_axle"]
Relaxation = Literal["none", "front"]
# Parameters carried in radians, shown in degrees.
DEGREE_PARAMS = ("u_off",)
[docs]
class WindowArrays(BaseModel):
"""Aligned numpy arrays for one sub-window, as the optimizer sees it.
Built once per sub-window by ``FitContext.arrays``; every array shares the
yaw-rate timestamp grid. ``u_rad`` is the steering input already converted
to radians. ``body_slip_rad`` and ``bicycle_steer_rad`` are ``None`` when
the corresponding channel is absent from the sub-window.
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
sw_id: int = Field(description="Sub-window identifier (position in the parent sub-window list)")
t: NDArray[np.float64] = Field(description="Uniformly-sampled time vector (s)")
speed: NDArray[np.float64] = Field(description="Vehicle speed (m/s)")
yaw_rate: NDArray[np.float64] = Field(description="Measured yaw rate (rad/s)")
u_rad: NDArray[np.float64] = Field(description="Steering input (rad)")
body_slip_rad: NDArray[np.float64] | None = Field(
description="Body slip angle (rad), or None if the channel is absent"
)
bicycle_steer_rad: NDArray[np.float64] | None = Field(
description="Front bicycle-model steer angle (rad), or None if the channel is absent"
)
@property
def duration_s(self) -> float:
"""Time spanned by the sub-window (s)."""
return float(self.t[-1] - self.t[0]) if len(self.t) > 1 else 0.0
[docs]
class FitContext(BaseModel):
"""Channel names, units, and the car geometry shared by every fit.
Built once per notebook run. Holds everything the predictor and optimizer
need that depends on neither the fit spec nor the parameter vector.
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
irl_car: IrlCar = Field(description="IrlCar supplying mass/geometry for the bicycle model")
input_col: str = Field(
default=FRONT_BICYCLE_MODEL_STEER_ANGLE,
description="Steering-input channel name",
)
input_units: Literal["rad", "deg"] = Field(default="deg", description="Units of input_col")
speed_col: str = Field(default="pcm.vnav.velocityBody.x", description="Speed channel name")
yaw_rate_col: str = Field(
default="pcm.vnav.compensatedAngularRate.z", description="Yaw-rate channel name"
)
body_slip_col: str = Field(default=BODY_SLIP_ANGLE, description="Body-slip-angle channel name")
body_slip_units: Literal["rad", "deg"] = Field(
default="deg", description="Units of body_slip_col"
)
bicycle_steer_col: str = Field(
default=FRONT_BICYCLE_MODEL_STEER_ANGLE,
description="Front bicycle-model steer-angle channel name",
)
bicycle_steer_units: Literal["rad", "deg"] = Field(
default="deg", description="Units of bicycle_steer_col"
)
[docs]
class FitSpec(BaseModel):
"""Selects the Ca-decay scheduler and the tire-relaxation flavor.
Together with a ``ParamLayout`` (which decides which scalars are free)
this fully determines the prediction pipeline.
"""
model_config = ConfigDict(frozen=True)
name: str = Field(description="Identifying label for this fit variant")
ca_scheduler: CaScheduler = Field(
default="none", description="Per-axle cornering-stiffness decay law"
)
relaxation: Relaxation = Field(default="none", description="Front-tire relaxation-lag flavor")
[docs]
class Param(BaseModel):
"""One entry of the optimizer's parameter vector.
A ``pinned`` param is held at ``init`` and never handed to the optimizer,
so it carries no bounds.
"""
model_config = ConfigDict(frozen=True)
name: str = Field(description="Parameter name")
init: float = Field(description="Initial guess (also the held value when pinned)")
lo: float = Field(description="Lower bound; ignored when pinned")
hi: float = Field(description="Upper bound; ignored when pinned")
scale: float = Field(description="Characteristic magnitude, for the optimizer's x_scale")
pinned: bool = Field(default=False, description="Hold at init instead of fitting")
[docs]
class ParamLayout(BaseModel):
"""The ordered parameter vector for a fit: what is free, and within what box.
Free params appear in the optimizer's ``x`` in ``free`` order; pinned
params are held at their initial value and never reach the optimizer.
"""
params: list[Param] = Field(description="Every parameter of this fit, free and pinned")
@property
def free(self) -> list[Param]:
"""The parameters handed to the optimizer, in ``x`` order."""
return [p for p in self.params if not p.pinned]
@property
def n_free(self) -> int:
"""Length of the optimizer's ``x`` vector."""
return len(self.free)
@property
def x0(self) -> NDArray[np.float64]:
"""Initial guess, aligned with ``free``."""
return np.array([p.init for p in self.free], dtype=float)
@property
def bounds(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
"""``(lo, hi)`` arrays, aligned with ``free``."""
return (
np.array([p.lo for p in self.free], dtype=float),
np.array([p.hi for p in self.free], dtype=float),
)
@property
def x_scale(self) -> NDArray[np.float64]:
"""Characteristic magnitudes, aligned with ``free``."""
return np.array([p.scale for p in self.free], dtype=float)
[docs]
def has(self, name: str) -> bool:
"""Whether ``name`` is a parameter of this fit, free or pinned."""
return any(p.name == name for p in self.params)
[docs]
def get(self, x: NDArray[np.float64], name: str) -> float:
"""Value of ``name`` given the optimizer's current ``x``.
Pinned parameters ignore ``x`` and return their held value.
"""
for i, p in enumerate(self.free):
if p.name == name:
return float(x[i])
for p in self.params:
if p.name == name:
return float(p.init)
raise KeyError(f"'{name}' is not a parameter of this fit")
[docs]
def named(self, x: NDArray[np.float64]) -> dict[str, float]:
"""Every parameter keyed by name, free values taken from ``x``."""
return {p.name: self.get(x, p.name) for p in self.params}
def __str__(self) -> str:
lines = ["ParamLayout:"]
for p in self.params:
if p.pinned:
lines.append(f" {p.name:>8s} pinned={p.format(p.init):>16s}")
else:
lines.append(
f" {p.name:>8s} init={p.format(p.init):>16s} "
f"bounds=[{p.format(p.lo)}, {p.format(p.hi)}]"
)
return "\n".join(lines)
[docs]
class ChassisParams(BaseModel):
"""The three free chassis scalars of the linear bicycle model.
Transfer-function form: ``H(s; V) = K wn^2 (1 + T_z s) /
(s^2 + 2 zeta wn s + wn^2)``, evaluated by ``bicycle_coeffs``.
"""
model_config = ConfigDict(frozen=True)
Ca_f: float = Field(description="Front per-axle cornering stiffness (N/rad)")
Ca_r: float = Field(description="Rear per-axle cornering stiffness (N/rad)")
Izz: float = Field(description="Yaw inertia about the CG (kg.m^2)")
[docs]
class BicycleCoeffs(BaseModel):
"""LPV coefficients of the yaw-rate-to-front-tire-steer transfer function.
Evaluated at one or more speeds by ``bicycle_coeffs``; every field is
shaped like the input speed(s).
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
K: NDArray[np.float64] = Field(description="DC gain (rad/s per rad steer)")
T_z: NDArray[np.float64] = Field(description="Numerator (zero) time constant (s)")
omega_n: NDArray[np.float64] = Field(description="Undamped natural frequency (rad/s)")
zeta: NDArray[np.float64] = Field(description="Damping ratio (unitless)")
[docs]
class BicyclePoint(BaseModel):
"""Linear-region bicycle coefficients evaluated at a single speed."""
V: float = Field(description="Speed the coefficients were evaluated at (m/s)")
K: float = Field(description="DC gain (rad/s per rad steer)")
T_z_ms: float = Field(description="Numerator (zero) time constant (ms)")
wn_hz: float = Field(description="Undamped natural frequency (Hz)")
zeta: float = Field(description="Damping ratio (unitless)")
def __str__(self) -> str:
return (
f"V={self.V:>4.1f} K={self.K:+.3f} Tz={self.T_z_ms:+8.1f} ms "
f"wn={self.wn_hz:.2f} Hz zeta={self.zeta:.3f}"
)
[docs]
class WindowMetrics(BaseModel):
"""Residual diagnostics for one sub-window.
``ac`` metrics have the per-window DC offset removed, isolating how well
the *dynamics* are captured from any steady-state bias.
"""
subwindow_id: int = Field(description="sw_id of the sub-window these metrics describe")
duration_s: float = Field(description="Sub-window duration (s)")
mean_speed_mps: float = Field(description="Mean speed over the sub-window (m/s)")
dc_offset_rad_s: float = Field(description="Mean residual, i.e. steady-state bias (rad/s)")
rmse_raw_rad_s: float = Field(description="RMSE of the raw residual (rad/s)")
rmse_ac_rad_s: float = Field(description="RMSE of the DC-removed residual (rad/s)")
yaw_rate_std_meas_rad_s: float = Field(
description="Standard deviation of the measured yaw rate (rad/s)"
)
vaf_ac_pct: float = Field(description="Variance accounted for by the DC-removed fit (%)")
@property
def rmse_ac_deg_s(self) -> float:
"""``rmse_ac_rad_s`` in deg/s."""
return float(np.rad2deg(self.rmse_ac_rad_s))
def __str__(self) -> str:
return (
f"sub {self.subwindow_id:>2d}: {self.duration_s:5.1f}s "
f"V={self.mean_speed_mps:5.1f} "
f"RMSE(ac)={self.rmse_ac_deg_s:5.2f} deg/s "
f"VAF={self.vaf_ac_pct:5.1f}%"
)
[docs]
class AggregateMetrics(BaseModel):
"""RMSE / VAF pooled across a list of ``WindowMetrics``."""
rmse_raw_rad_s: float = Field(
description="Root-mean-square of the per-window raw RMSEs (rad/s)"
)
rmse_ac_rad_s: float = Field(description="Root-mean-square of the per-window AC RMSEs (rad/s)")
vaf_ac_pct: float = Field(description="Mean per-window AC variance-accounted-for (%)")
@property
def rmse_ac_deg_s(self) -> float:
"""``rmse_ac_rad_s`` in deg/s."""
return float(np.rad2deg(self.rmse_ac_rad_s))
[docs]
@classmethod
def from_rows(cls, rows: list[WindowMetrics]) -> AggregateMetrics:
"""Pool per-window rows; all-NaN when ``rows`` is empty."""
if not rows:
nan = float("nan")
return cls(rmse_raw_rad_s=nan, rmse_ac_rad_s=nan, vaf_ac_pct=nan)
return cls(
rmse_raw_rad_s=float(np.sqrt(np.mean([r.rmse_raw_rad_s**2 for r in rows]))),
rmse_ac_rad_s=float(np.sqrt(np.mean([r.rmse_ac_rad_s**2 for r in rows]))),
vaf_ac_pct=float(np.mean([r.vaf_ac_pct for r in rows])),
)
def __str__(self) -> str:
return (
f"Aggregate raw RMSE = {self.rmse_raw_rad_s:.4f} rad/s "
f"({np.rad2deg(self.rmse_raw_rad_s):.2f} deg/s)\n"
f"Aggregate AC RMSE = {self.rmse_ac_rad_s:.4f} rad/s "
f"({self.rmse_ac_deg_s:.2f} deg/s)\n"
f"Mean VAF (AC) = {self.vaf_ac_pct:.2f}%"
)
[docs]
class Prediction(BaseModel):
"""Per-sub-window predictions and residual diagnostics.
Shared by a training fit and a held-out evaluation, which differ only in
which sub-windows they were computed over.
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
y_hat_per_window: list[NDArray[np.float64]] = Field(
description="Predicted yaw rate per sub-window (rad/s)"
)
err_per_window: list[NDArray[np.float64]] = Field(
description="Residual (measured - predicted) per sub-window (rad/s)"
)
window_metrics: list[WindowMetrics] = Field(
description="Residual diagnostics, one row per sub-window"
)
n_samples: int = Field(description="Total samples across every sub-window")
aic: float = Field(description="Akaike information criterion")
bic: float = Field(description="Bayesian information criterion")
[docs]
def aggregate(self) -> AggregateMetrics:
"""RMSE / VAF pooled across this prediction's sub-windows."""
return AggregateMetrics.from_rows(self.window_metrics)
[docs]
class YawFit(BaseModel):
"""A fitted yaw-response model: the spec, the parameters, and how well it did.
Carries everything needed to evaluate the model on new data
(``spec`` / ``layout`` / ``ctx`` / ``x_hat``) and everything needed to
report on the fit that produced it (``train``).
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
spec: FitSpec = Field(description="The variant that was fit")
layout: ParamLayout = Field(description="Parameter vector layout ``x_hat`` is aligned with")
ctx: FitContext = Field(description="Channel/geometry metadata the fit was run against")
x_hat: NDArray[np.float64] = Field(description="Fitted free-parameter vector")
train: Prediction = Field(description="Predictions and metrics on the training sub-windows")
final_cost: float = Field(description="least_squares final cost, 0.5 * sum(residual^2)")
n_eval: int = Field(description="Number of residual-function evaluations")
status: int = Field(description="least_squares termination status code")
message: str = Field(description="least_squares termination message")
balance_mode: str = Field(description="Sign-balancing mode used to build the fit views")
fit_view_count: int = Field(description="Number of views the residual was summed over")
source_window_count: int = Field(description="Number of source sub-windows behind those views")
@property
def n_params(self) -> int:
"""Number of free parameters."""
return self.layout.n_free
[docs]
def param(self, name: str) -> float:
"""Fitted (or pinned) value of ``name``."""
return self.layout.get(self.x_hat, name)
[docs]
def has_param(self, name: str) -> bool:
"""Whether ``name`` is a parameter of this fit."""
return self.layout.has(name)
[docs]
def params_named(self) -> dict[str, float]:
"""Every fitted and pinned parameter, keyed by name."""
return self.layout.named(self.x_hat)
[docs]
def chassis_params(self) -> ChassisParams:
"""This fit's ``(Ca_f, Ca_r, Izz)`` triple."""
return ChassisParams(
Ca_f=self.param("Ca_f"), Ca_r=self.param("Ca_r"), Izz=self.param("Izz")
)
[docs]
def aggregate(self) -> AggregateMetrics:
"""RMSE / VAF pooled across the training sub-windows."""
return self.train.aggregate()
def __str__(self) -> str:
agg = self.aggregate()
lines = [
f"[{self.spec.name}] final cost = {self.final_cost:.4f}, "
f"n_eval = {self.n_eval}, status = {self.status} ({self.message})",
f" n_samples = {self.train.n_samples}, n_params = {self.n_params}, "
f"AIC = {self.train.aic:.1f}, BIC = {self.train.bic:.1f}",
f" ca_scheduler: {self.spec.ca_scheduler}, relaxation: {self.spec.relaxation}",
f" {'param':>8s} {'init':>16s} {'fit':>16s} {'delta':>10s}",
]
for p in self.layout.params:
fitted = self.param(p.name)
if p.pinned:
delta = "pinned"
elif abs(p.init) < 1e-9:
delta = "n/a"
else:
delta = f"{100.0 * (fitted - p.init) / abs(p.init):+7.1f}%"
lines.append(
f" {p.name:>8s} {p.format(p.init):>16s} "
f"{p.format(fitted):>16s} {delta:>10s}"
)
lines.append(f" RMSE(ac) = {agg.rmse_ac_deg_s:.2f} deg/s, VAF = {agg.vaf_ac_pct:.1f}%")
return "\n".join(lines)
[docs]
class FitComparison(BaseModel):
"""Side-by-side table of several ``YawFit``, one row each."""
rows: list["FitComparisonRow"] = Field(description="One row per fit")
[docs]
@classmethod
def from_fits(cls, fits: dict[str, YawFit]) -> FitComparison:
"""Build a comparison table from fits keyed by display name."""
return cls(rows=[FitComparisonRow.from_fit(k, f) for k, f in fits.items()])
def __str__(self) -> str:
header = (
f"{'model':>10s} {'n_p':>4s} {'cost':>8s} {'RMSE(ac)':>9s} "
f"{'VAF%':>6s} {'AIC':>9s} {'BIC':>9s}"
)
lines = [header, "-" * len(header)]
for r in self.rows:
lines.append(
f"{r.name:>10s} {r.n_params:>4d} {r.final_cost:>8.3f} "
f"{r.rmse_ac_deg_s:>9.2f} {r.vaf_ac_pct:>6.1f} "
f"{r.aic:>9.1f} {r.bic:>9.1f}"
)
return "\n".join(lines)
[docs]
class FitComparisonRow(BaseModel):
"""One fit's row in a ``FitComparison``."""
name: str = Field(description="Display name of the fit")
ca_scheduler: CaScheduler = Field(description="Per-axle cornering-stiffness decay law")
relaxation: Relaxation = Field(description="Front-tire relaxation-lag flavor")
n_params: int = Field(description="Number of free parameters")
final_cost: float = Field(description="least_squares final cost")
rmse_ac_deg_s: float = Field(description="Pooled AC-coupled RMSE (deg/s)")
vaf_ac_pct: float = Field(description="Pooled AC-coupled variance-accounted-for (%)")
aic: float = Field(description="Akaike information criterion")
bic: float = Field(description="Bayesian information criterion")
Ca_f: float = Field(description="Fitted front cornering stiffness (N/rad)")
Ca_r: float = Field(description="Fitted rear cornering stiffness (N/rad)")
Izz: float = Field(description="Fitted yaw inertia (kg.m^2)")
L_f: float = Field(description="Fitted front relaxation length (m); NaN if not modelled")
u_off_deg: float = Field(description="Fitted steering offset (deg)")
[docs]
@classmethod
def from_fit(cls, name: str, fit: YawFit) -> FitComparisonRow:
"""Summarize one fit into a comparison row."""
agg = fit.aggregate()
return cls(
name=name,
ca_scheduler=fit.spec.ca_scheduler,
relaxation=fit.spec.relaxation,
n_params=fit.n_params,
final_cost=fit.final_cost,
rmse_ac_deg_s=agg.rmse_ac_deg_s,
vaf_ac_pct=agg.vaf_ac_pct,
aic=fit.train.aic,
bic=fit.train.bic,
Ca_f=fit.param("Ca_f"),
Ca_r=fit.param("Ca_r"),
Izz=fit.param("Izz"),
L_f=fit.param("L_f") if fit.has_param("L_f") else float("nan"),
u_off_deg=float(np.rad2deg(fit.param("u_off"))),
)