import numpy as np
import plotly.graph_objects as go
from numpy.typing import NDArray
from plotly.subplots import make_subplots
from ..plotting.binned_scatter import plot_binned_scatter
from ..plotting.plot_3d import plot3D_surface
from ..plotting.two_panel_spectrum import plot_two_panel_spectrum
from ..vehicle.irl_vehicle import IrlCar
from .yaw_fit import bicycle_coeffs
from .yaw_fit_models import ChassisParams, Prediction, WindowArrays, YawFit
RESIDUAL_X_AXES = {
"input": ("steering input (deg)", "Residual vs steering input (signed)"),
"amplitude": ("|steering input| (deg)", "Residual vs steering amplitude"),
"speed": ("speed V (m/s)", "Residual vs speed"),
}
[docs]
def eval_lpv_bode(
chassis: ChassisParams,
irl_car: IrlCar,
V_grid: NDArray[np.float64],
f_grid: NDArray[np.float64],
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
"""Evaluate the analytic Bode surface over a (speed, frequency) grid.
Parameters
----------
chassis : ChassisParams
Chassis scalars to evaluate the transfer function at.
irl_car : IrlCar
Source of the mass, wheelbase, and CG split.
V_grid : NDArray[np.float64]
Speed grid (m/s), length ``nV``.
f_grid : NDArray[np.float64]
Frequency grid (Hz), length ``nf``.
Returns
-------
mag : NDArray[np.float64]
Dimensionless magnitude ``|H(j 2 pi f; V)|``, shape ``(nV, nf)``.
phase_deg : NDArray[np.float64]
Phase, unwrapped along the frequency axis (deg), shape ``(nV, nf)``.
"""
c = bicycle_coeffs(np.asarray(V_grid, dtype=float), chassis, irl_car)
K = c.K[:, None]
Tz = c.T_z[:, None]
wn = c.omega_n[:, None]
zt = c.zeta[:, None]
s = 1j * 2.0 * np.pi * np.asarray(f_grid, dtype=float)[None, :]
H = (K * wn * wn) * (1.0 + Tz * s) / (s * s + 2.0 * zt * wn * s + wn * wn)
return np.abs(H), np.unwrap(np.angle(H), axis=1) * 180.0 / np.pi
def _fit_bode(
fit: YawFit, V_grid: NDArray[np.float64], f_grid: NDArray[np.float64]
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
"""Bode surface of a fit's linear-region coefficients."""
return eval_lpv_bode(fit.chassis_params(), fit.ctx.irl_car, V_grid, f_grid)
[docs]
def plot_bode_surface(
fit: YawFit,
V_grid: NDArray[np.float64],
f_grid: NDArray[np.float64],
log_mag: bool = True,
) -> go.Figure:
"""Magnitude Bode surface of a fit over (speed, frequency).
Parameters
----------
fit : YawFit
Supplies the linear-region chassis parameters and geometry.
V_grid : NDArray[np.float64]
Speed grid (m/s).
f_grid : NDArray[np.float64]
Frequency grid (Hz).
log_mag : bool, optional
Plot magnitude in dB rather than linear.
Returns
-------
go.Figure
"""
mag, _ = _fit_bode(fit, V_grid, f_grid)
z = 20.0 * np.log10(np.maximum(mag, 1e-12)) if log_mag else mag
return plot3D_surface(
x_list=np.asarray(V_grid, dtype=float),
y_list=np.asarray(f_grid, dtype=float),
z_list=z,
title=f"LPV yaw-response magnitude ({fit.spec.name})",
x_axis="speed V (m/s)",
y_axis="frequency (Hz)",
z_axis="|H| (dB)" if log_mag else "|H|",
)
[docs]
def plot_bode_slices(
fits: dict[str, YawFit],
V_picks: NDArray[np.float64] | list[float],
f_grid: NDArray[np.float64],
linear_scale: bool = False,
) -> go.Figure:
"""Two-panel magnitude/phase Bode of one or more fits, at fixed speeds.
Every (fit, speed) pair becomes its own series. Fits after the first are
drawn dashed, so a comparison model reads clearly against the headline one.
Parameters
----------
fits : dict[str, YawFit]
Fits keyed by display name.
V_picks : NDArray[np.float64] or list of float
Speeds to slice at (m/s).
f_grid : NDArray[np.float64]
Frequency grid (Hz).
linear_scale : bool, optional
Plot magnitude linearly rather than in dB.
Returns
-------
go.Figure
"""
V_picks = np.asarray(V_picks, dtype=float)
mag_by: dict[str, NDArray[np.float64]] = {}
phase_by: dict[str, NDArray[np.float64]] = {}
dashed: set[str] = set()
for fit_idx, (key, fit) in enumerate(fits.items()):
mag, phase = _fit_bode(fit, V_picks, f_grid)
for v_idx, V in enumerate(V_picks):
label = f"{key}, V={V:.1f}"
row = mag[v_idx]
mag_by[label] = row if linear_scale else 20.0 * np.log10(np.maximum(row, 1e-12))
phase_by[label] = phase[v_idx]
if fit_idx > 0:
dashed.add(label)
return plot_two_panel_spectrum(
np.asarray(f_grid, dtype=float),
mag_by,
phase_by,
title="LPV yaw-response Bode (linear-region, Ca decay = 0)",
magnitude_label="|H|" if linear_scale else "|H| (dB)",
dashed_series=dashed,
)
def _display_idx(n_points: int, duration_s: float, max_display_hz: float) -> NDArray[np.intp]:
"""Indices that cap the plotted sample rate at ``max_display_hz``."""
if not max_display_hz or duration_s <= 0 or n_points <= 1:
return np.arange(n_points)
actual_hz = n_points / duration_s
if actual_hz <= max_display_hz:
return np.arange(n_points)
return np.arange(0, n_points, max(1, int(round(actual_hz / max_display_hz))))
def _resolve_selection(sel: list[int] | range | slice | None, n: int) -> list[int]:
"""Resolve a sub-window selector to an explicit index list."""
if sel is None:
return list(range(n))
if isinstance(sel, slice):
return list(range(*sel.indices(n)))
return list(sel)
[docs]
def plot_fit_overlay(
window_arrays: list[WindowArrays],
prediction: Prediction,
subwindow_indices: list[int] | range | slice | None = None,
max_display_hz: float = 50.0,
height_per_subplot: int = 220,
error_ylim: tuple[float, float] = (-1.0, 1.0),
relative_time: bool = True,
) -> go.Figure:
"""Measured vs simulated yaw rate per sub-window, residual on the second axis.
Works for a training fit (``fit.train``) or a held-out evaluation, since
both are a ``Prediction``.
Parameters
----------
window_arrays : list of WindowArrays
Source sub-windows.
prediction : Prediction
Predictions to overlay; must be aligned with ``window_arrays``.
subwindow_indices : list of int, range, slice, or None, optional
Which sub-windows to draw. ``None`` draws all of them.
max_display_hz : float, optional
Downsample the traces to at most this rate, for browser responsiveness.
height_per_subplot : int, optional
Pixel height per sub-window row.
error_ylim : tuple[float, float], optional
Fixed initial range for the residual axis (rad/s).
relative_time : bool, optional
Start each sub-window's time axis at zero.
Returns
-------
go.Figure
Raises
------
ValueError
If ``subwindow_indices`` selects nothing.
"""
sel = _resolve_selection(subwindow_indices, len(window_arrays))
if not sel:
raise ValueError("subwindow_indices selected zero sub-windows.")
fig = make_subplots(
rows=len(sel),
cols=1,
shared_xaxes=False,
subplot_titles=[f"sub-window {window_arrays[i].sw_id} (pos {i})" for i in sel],
vertical_spacing=min(0.08, 0.6 / max(len(sel), 1)),
specs=[[{"secondary_y": True}] for _ in sel],
)
traces = (
("measured", "#1f77b4", 1.5, "solid", 1.0, False),
("simulated", "#d62728", 1.5, "dot", 1.0, False),
("error", "#7f7f7f", 1.0, "solid", 0.6, True),
)
for row, i in enumerate(sel, start=1):
w = window_arrays[i]
t = w.t - w.t[0] if relative_time and len(w.t) else w.t
idx = _display_idx(len(t), w.duration_s, max_display_hz)
series = (
w.yaw_rate,
prediction.y_hat_per_window[i],
prediction.err_per_window[i],
)
for y, (name, color, width, dash, opacity, on_y2) in zip(series, traces):
fig.add_trace(
go.Scattergl(
x=t[idx],
y=y[idx],
name=name,
mode="lines",
line={"color": color, "width": width, "dash": dash},
opacity=opacity,
legendgroup=name,
showlegend=row == 1,
),
row=row,
col=1,
secondary_y=on_y2,
)
fig.update_yaxes(title_text="yaw rate (rad/s)", row=row, col=1, secondary_y=False)
fig.update_yaxes(
title_text="error (rad/s)",
range=list(error_ylim),
row=row,
col=1,
secondary_y=True,
)
fig.update_xaxes(title_text="time (s)", row=len(sel), col=1)
fig.update_layout(
height=height_per_subplot * len(sel) + 80,
title="Yaw fit overlay (measured / simulated / residual)",
legend={
"orientation": "h",
"y": -0.06,
"x": 0.5,
"xanchor": "center",
"yanchor": "top",
},
margin={"b": 80},
)
return fig
[docs]
def plot_residual_vs(
window_arrays: list[WindowArrays],
prediction: Prediction,
x_of: str = "amplitude",
n_bins: int = 24,
) -> go.Figure:
"""Residual structure against steering, steering amplitude, or speed.
A flat binned mean means the model has extracted everything systematic.
Residual that grows with ``"amplitude"`` is the classic tire-saturation
signature, and the reason the Ca-decay scheduler exists.
Parameters
----------
window_arrays : list of WindowArrays
Source sub-windows.
prediction : Prediction
Supplies the residuals; must be aligned with ``window_arrays``.
x_of : {"input", "amplitude", "speed"}, optional
What to put on the x axis: signed steering (deg), absolute steering
(deg), or speed (m/s).
n_bins : int, optional
Number of quantile bins for the mean/sigma overlay.
Returns
-------
go.Figure
Raises
------
ValueError
If ``x_of`` is unrecognised.
"""
if x_of not in RESIDUAL_X_AXES:
raise ValueError(f"unknown x_of '{x_of}'; expected one of {sorted(RESIDUAL_X_AXES)}")
def x_values(w: WindowArrays) -> NDArray[np.float64]:
if x_of == "input":
return np.rad2deg(w.u_rad)
if x_of == "amplitude":
return np.abs(np.rad2deg(w.u_rad))
return w.speed
x = np.concatenate([x_values(w) for w in window_arrays])
y = np.concatenate([np.asarray(e, dtype=float) for e in prediction.err_per_window])
sw_id = np.concatenate([np.full(len(w.t), w.sw_id, dtype=float) for w in window_arrays])
x_label, title = RESIDUAL_X_AXES[x_of]
return plot_binned_scatter(
x,
y,
color_values=sw_id,
color_label="sub-window",
title=title,
x_label=x_label,
y_label="residual y_meas - y_hat (rad/s)",
n_bins=n_bins,
)