Source code for suboptimumg.loganalysis.track_builder

"""Build QSS-compatible Track objects from processed lap DataFrames."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly.colors import qualitative
from plotly.subplots import make_subplots
from scipy.fft import rfft, rfftfreq
from scipy.interpolate import splprep
from scipy.spatial import cKDTree

from suboptimumg.track import Track

if TYPE_CHECKING:
    from suboptimumg.compsim.models import LapsimResults

# ---------------------------------------------------------------------------
# Resampling
# ---------------------------------------------------------------------------


[docs] def resample_lap_to_distance( lap_df: pd.DataFrame, dx: float = 0.1, curvature_col: str = "body.curvature", velocity_col: str = "groundSpeed", pos_n_col: str = "posN", pos_e_col: str = "posE", dist_col: str = "LapDist", ) -> dict: """Resample a trimmed lap DataFrame to a uniform distance grid. Parameters ---------- lap_df : DataFrame A single-lap DataFrame produced by :func:`split_laps`, which must contain *dist_col*, *curvature_col*, *velocity_col*, *pos_n_col*, and *pos_e_col*. dx : float Distance step in meters. Returns ------- dict Keys: ``distance``, ``velocity``, ``curvature``, ``x_m``, ``y_m``. """ dist = lap_df[dist_col].values distance_grid = np.arange(0, dist[-1], dx) return { "distance": distance_grid, "velocity": np.interp(distance_grid, dist, lap_df[velocity_col].values), "curvature": np.interp(distance_grid, dist, lap_df[curvature_col].values), "x_m": np.interp(distance_grid, dist, lap_df[pos_e_col].values), "y_m": np.interp(distance_grid, dist, lap_df[pos_n_col].values), }
# --------------------------------------------------------------------------- # Spatial curvature filter # ---------------------------------------------------------------------------
[docs] def filter_curvature_spatial( distance: np.ndarray, curvature: np.ndarray, cutoff_freq: float = 0.18, ) -> np.ndarray: """FFT low-pass filter on curvature in the spatial-frequency domain. Parameters ---------- distance : ndarray Uniform distance grid (m). curvature : ndarray Curvature (1/m) on the same grid. cutoff_freq : float Cutoff spatial frequency in 1/m. Returns ------- ndarray Filtered curvature array (same length). """ dx = distance[1] - distance[0] mean_curv = np.mean(curvature) yf = rfft(curvature - mean_curv) xf = rfftfreq(len(curvature), dx) yf_filt = np.where(xf <= cutoff_freq, yf, 0.0) filtered = np.fft.irfft(yf_filt, n=len(distance)) + mean_curv return filtered
# --------------------------------------------------------------------------- # Track builder # ---------------------------------------------------------------------------
[docs] def build_qss_track( lap_df: pd.DataFrame, dx: float = 0.1, curvature_col: str = "body.curvature", velocity_col: str = "groundSpeed", pos_n_col: str = "posN", pos_e_col: str = "posE", dist_col: str = "LapDist", cutoff_freq: float = 0.18, max_radius: float = 300.0, show_plots: bool = False, ) -> Track: """Build a QSS-compatible Track from a trimmed lap DataFrame. Pipeline: resample to uniform *dx* -> FFT spatial low-pass on curvature -> convert to absolute radius -> fit B-spline on XY -> construct :class:`Track`. Parameters ---------- lap_df : DataFrame Single-lap DataFrame from :func:`split_laps` (must have *dist_col*, *curvature_col*, *velocity_col*, position cols). dx : float Distance step in meters (default 0.1). cutoff_freq : float Spatial frequency cutoff for curvature filter (1/m). max_radius : float Radius clamp for near-straight segments. show_plots : bool If True, display verification plots. """ data = resample_lap_to_distance( lap_df, dx, curvature_col, velocity_col, pos_n_col, pos_e_col, dist_col ) distance_grid = data["distance"] curv_raw = data["curvature"] curv_filt = filter_curvature_spatial(distance_grid, curv_raw, cutoff_freq) with np.errstate(divide="ignore", invalid="ignore"): radius = np.where(curv_filt != 0, np.abs(1.0 / curv_filt), max_radius) radius = np.clip(radius, 0.0, max_radius) x_m = data["x_m"] y_m = data["y_m"] tck, _ = splprep([x_m, y_m], u=distance_grid, s=0, k=3) n = len(distance_grid) track = Track( dx=np.full(n, dx), radius=radius, cumulative_dist=distance_grid.copy(), x_m=x_m, y_m=y_m, distance_step=dx, continuous=True, tck=tck, ) if show_plots: _plot_build_verification(distance_grid, curv_raw, curv_filt, radius, x_m, y_m) print( f"[track_builder] Track built: {n} points, " f"total distance = {distance_grid[-1]:.1f} m, dx = {dx} m" ) return track
def _plot_build_verification( distance: np.ndarray, curv_raw: np.ndarray, curv_filt: np.ndarray, radius: np.ndarray, x_m: np.ndarray, y_m: np.ndarray, ) -> None: """Show XY track and curvature/radius before-after filter.""" fig = make_subplots( rows=1, cols=2, subplot_titles=("Track XY", "Curvature (before/after filter)"), specs=[[{}, {"secondary_y": True}]], ) fig.add_trace( go.Scattergl(x=y_m, y=x_m, mode="lines", name="Track"), row=1, col=1, ) fig.update_xaxes(title_text="North (m)", row=1, col=1) fig.update_yaxes(title_text="East (m)", scaleanchor="x", scaleratio=1, row=1, col=1) fig.add_trace( go.Scattergl(x=distance, y=curv_raw, mode="lines", name="Raw κ", opacity=0.4), row=1, col=2, secondary_y=False, ) fig.add_trace( go.Scattergl(x=distance, y=curv_filt, mode="lines", name="Filtered κ"), row=1, col=2, secondary_y=False, ) fig.add_trace( go.Scattergl(x=distance, y=radius, mode="lines", name="Radius"), row=1, col=2, secondary_y=True, ) fig.update_xaxes(title_text="Distance (m)", row=1, col=2) fig.update_yaxes(title_text="Curvature (1/m)", row=1, col=2, secondary_y=False) fig.update_yaxes(title_text="Radius (m)", row=1, col=2, secondary_y=True) fig.update_layout(height=450, showlegend=True) fig.show() # --------------------------------------------------------------------------- # QSS-to-log variable mapping # --------------------------------------------------------------------------- QSS_TO_LOG: dict[str, str | None] = { "lap_vels": "groundSpeed", "lap_accs": "pcm.vnav.linearAccelBody.x", "lap_powers": "bms.pack.power", "lap_eff_motor_torques": "pcm.moc.motor.torqueFeedback", "lap_dxs": None, "lap_t": "time_s", "v_max_profile": None, "cumulative_dist": "LapDist", "radius": "body.radius", "x_m": "posE", "y_m": "posN", } LOG_TO_QSS: dict[str, str] = { "groundSpeed": "qss.velocity", "pcm.vnav.linearAccelBody.x": "qss.accLong", "bms.pack.power": "qss.power", "pcm.moc.motor.torqueFeedback": "qss.motorTorque", "body.radius": "qss.radius", "posN": "qss.posN", "posE": "qss.posE", } def _resolve_qss_col(variable: str, qss_df: pd.DataFrame) -> str: """Translate a log column name to its qss_df equivalent. Lookup order: 1. ``LOG_TO_QSS[variable]`` if present in *qss_df*. 2. ``"qss." + variable`` if present in *qss_df*. 3. *variable* itself if present in *qss_df*. 4. Raise ``KeyError`` with available ``qss.*`` columns. """ mapped = LOG_TO_QSS.get(variable) if mapped and mapped in qss_df.columns: return mapped prefixed = f"qss.{variable}" if prefixed in qss_df.columns: return prefixed if variable in qss_df.columns: return variable qss_cols = [c for c in qss_df.columns if c.startswith("qss.")] raise KeyError( f"Cannot resolve '{variable}' in qss_df. " f"Available qss columns: {qss_cols}" )
[docs] def qss_results_to_dataframe( lapsim_results: LapsimResults, track: Track, ) -> pd.DataFrame: """Convert QSS lapsim output arrays into a distance-indexed DataFrame. Columns use a ``qss.`` prefix so they sit cleanly alongside log columns in a merged comparison DataFrame. Includes position columns (``qss.posN``, ``qss.posE``) from the track. """ dist = track.cumulative_dist n = min(len(dist), len(lapsim_results.lap_vels)) return pd.DataFrame( { "qss.velocity": np.asarray(lapsim_results.lap_vels)[:n], "qss.accLong": np.asarray(lapsim_results.lap_accs)[:n], "qss.power": np.asarray(lapsim_results.lap_powers)[:n] / 1000.0, "qss.motorTorque": np.asarray(lapsim_results.lap_eff_motor_torques)[:n], "qss.time": np.asarray(lapsim_results.lap_t)[:n], "qss.radius": np.asarray(track.radius)[:n], "qss.posN": np.asarray(track.y_m)[:n], "qss.posE": np.asarray(track.x_m)[:n], }, index=pd.Index(dist[:n], name="distance"), )
[docs] def resample_log_to_distance( lap_df: pd.DataFrame, distance_grid: np.ndarray, columns: list[str] | None = None, dist_col: str = "LapDist", ) -> pd.DataFrame: """Resample selected log columns onto a QSS distance grid. Uses linear interpolation via :func:`numpy.interp`. Parameters ---------- lap_df : DataFrame Trimmed single-lap DataFrame. distance_grid : ndarray Target distance array (from a built Track). columns : list[str] or None Columns to resample. If ``None``, uses all non-``None`` values from :data:`QSS_TO_LOG`. dist_col : str Distance column in *lap_df*. """ if columns is None: columns = [ v for v in QSS_TO_LOG.values() if v is not None and v in lap_df.columns ] dist = lap_df[dist_col].values result: dict[str, np.ndarray] = {} for col in columns: if col not in lap_df.columns: continue result[col] = np.interp(distance_grid, dist, lap_df[col].values) return pd.DataFrame(result, index=pd.Index(distance_grid, name="distance"))
[docs] def build_comparison_df( lapsim_results: LapsimResults, track: Track, lap_df: pd.DataFrame, log_columns: list[str] | None = None, dist_col: str = "LapDist", ) -> pd.DataFrame: """Merge QSS results and resampled log data into one DataFrame. The result is indexed by distance and contains both ``qss.*`` columns and the original log column names side by side. """ qss_df = qss_results_to_dataframe(lapsim_results, track) log_df = resample_log_to_distance( lap_df, qss_df.index.values, log_columns, dist_col ) return pd.concat([qss_df, log_df], axis=1)
# --------------------------------------------------------------------------- # Column resolution helpers # --------------------------------------------------------------------------- def _get_col(df: pd.DataFrame, col: str) -> str: """Return the column name to use for *col* in *df*. If *col* is present directly, return it. Otherwise try :func:`_resolve_qss_col` for automatic log->QSS translation. """ if col in df.columns: return col return _resolve_qss_col(col, df) def _get_values(df: pd.DataFrame, col: str) -> np.ndarray: """Extract values for *col* from *df*, auto-resolving QSS names.""" return df[_get_col(df, col)].values # --------------------------------------------------------------------------- # Comparison plots — generalized # ---------------------------------------------------------------------------
[docs] def plot_heatmap( baseline_df: pd.DataFrame, comparison_df: pd.DataFrame, variable: str = "groundSpeed", baseline_label: str = "Baseline", comparison_label: str = "Comparison", pos_cols: tuple[str, str] = ("posN", "posE"), dist_col: str = "LapDist", units: str = "m/s", dx: float = 0.1, ) -> go.Figure: """3-panel 2D track heatmap comparing any two DataFrames. Works for lap-vs-lap or lap-vs-QSS. Column names are auto-resolved via :data:`LOG_TO_QSS` so the caller always uses log-world names (e.g. ``"groundSpeed"``). Parameters ---------- baseline_df, comparison_df : DataFrame Any two lap DataFrames, or a lap and a ``qss_df``. variable : str Column to compare (log-world name). pos_cols : (north_col, east_col) Position columns (log-world names). units : str Label for colorbars and hover. dx : float Resample step when the two DataFrames are on different grids. """ pn, pe = pos_cols bl_n = _get_values(baseline_df, pn) bl_e = _get_values(baseline_df, pe) bl_var = _get_values(baseline_df, variable) bl_xy = np.column_stack([bl_e, bl_n]) if dist_col in baseline_df.columns: bl_dist = baseline_df[dist_col].values else: bl_dist = baseline_df.index.values cmp_n = _get_values(comparison_df, pn) cmp_e = _get_values(comparison_df, pe) cmp_var = _get_values(comparison_df, variable) cmp_xy = np.column_stack([cmp_e, cmp_n]) # Project comparison onto baseline spatial grid via nearest-neighbor tree = cKDTree(bl_xy) _, nearest_idx = tree.query(cmp_xy) n_grid = len(bl_xy) vel_sums = np.zeros(n_grid) vel_counts = np.zeros(n_grid, dtype=int) for i, tidx in enumerate(nearest_idx): vel_sums[tidx] += cmp_var[i] vel_counts[tidx] += 1 has_data = vel_counts > 0 cmp_projected = np.full(n_grid, np.nan) cmp_projected[has_data] = vel_sums[has_data] / vel_counts[has_data] n_empty = int(np.sum(~has_data)) if n_empty > 0: valid = np.where(~np.isnan(cmp_projected))[0] if len(valid) > 0 and len(valid) < n_grid: cmp_projected = np.interp(np.arange(n_grid), valid, cmp_projected[valid]) print( f" [heatmap] {n_empty}/{n_grid} grid points had no comparison " "data — linearly interpolated." ) delta = bl_var - cmp_projected v_min = min(np.nanmin(bl_var), np.nanmin(cmp_projected)) v_max = max(np.nanmax(bl_var), np.nanmax(cmp_projected)) d_abs = max(abs(np.nanmin(delta)), abs(np.nanmax(delta))) fig = make_subplots( rows=1, cols=3, subplot_titles=( baseline_label, comparison_label, f"Δ ({baseline_label}{comparison_label})", ), horizontal_spacing=0.10, ) common = dict(mode="markers", marker=dict(size=3)) cb_kwargs = dict(len=0.75, thickness=12) fig.add_trace( go.Scattergl( x=bl_e, y=bl_n, **common, marker_color=bl_var, marker_cmin=v_min, marker_cmax=v_max, marker_colorscale="Viridis", marker_colorbar=dict(title=units, x=0.28, **cb_kwargs), name=baseline_label, customdata=np.column_stack([bl_var, bl_dist]), hovertemplate=( f"{baseline_label}: %{{customdata[0]:.2f}} {units}<br>" f"Distance: %{{customdata[1]:.1f}} m<extra></extra>" ), ), row=1, col=1, ) fig.add_trace( go.Scattergl( x=bl_e, y=bl_n, **common, marker_color=cmp_projected, marker_cmin=v_min, marker_cmax=v_max, marker_colorscale="Viridis", marker_colorbar=dict(title=units, x=0.65, **cb_kwargs), name=comparison_label, customdata=np.column_stack([cmp_projected, bl_dist]), hovertemplate=( f"{comparison_label}: %{{customdata[0]:.2f}} {units}<br>" f"Distance: %{{customdata[1]:.1f}} m<extra></extra>" ), ), row=1, col=2, ) fig.add_trace( go.Scattergl( x=bl_e, y=bl_n, **common, marker_color=delta, marker_cmin=-d_abs, marker_cmax=d_abs, marker_colorscale="RdBu_r", marker_colorbar=dict(title=f{units}", x=1.02, **cb_kwargs), name="Delta", customdata=np.column_stack([delta, bl_dist]), hovertemplate=( f"Δ: %{{customdata[0]:+.2f}} {units}<br>" f"Distance: %{{customdata[1]:.1f}} m<extra></extra>" ), ), row=1, col=3, ) for col_idx in (1, 2, 3): fig.update_xaxes(title_text="East (m)", row=1, col=col_idx) fig.update_yaxes( title_text="North (m)", scaleanchor=f"x{col_idx if col_idx > 1 else ''}", scaleratio=1, row=1, col=col_idx, ) fig.update_layout( height=500, width=1600, title=f"{variable} Comparison — Track Heatmap", showlegend=False, ) mean_d = np.nanmean(delta) rms_d = np.sqrt(np.nanmean(delta**2)) max_d = np.nanmax(np.abs(delta)) print( f" [heatmap] Mean Δ: {mean_d:.3f} {units}, " f"RMS Δ: {rms_d:.3f} {units}, " f"Max |Δ|: {max_d:.3f} {units}" ) return fig
[docs] def plot_comparison_trace( laps: dict[str | int, pd.DataFrame], variable: str = "groundSpeed", qss_df: pd.DataFrame | None = None, baseline: str | int = "fastest", dist_col: str = "LapDist", vel_col: str = "groundSpeed", units: str = "m/s", dx: float = 0.1, ) -> go.Figure: """Multi-lap comparison trace with cumulative time-delta subplot. Parameters ---------- laps : dict[label, DataFrame] Trimmed lap DataFrames from :func:`split_laps`. variable : str Column to plot on the top subplot (log-world name). qss_df : DataFrame or None If provided, the QSS series is included. Column names are auto-resolved via :data:`LOG_TO_QSS`. baseline : str or int Reference for time delta. ``"fastest"`` picks the lap with the shortest ``LapTime``. ``"first"`` / ``"last"`` use positional order. ``"qss"`` uses the QSS as baseline. An explicit key selects that lap. vel_col : str Velocity column used for time-delta computation (log-world name, auto-resolved for QSS). """ palette = qualitative.Plotly lap_keys = list(laps.keys()) def _label(key): return f"Lap {key}" if isinstance(key, int) else str(key) # --- Build resampled data on a common grid --- # Find the shortest lap distance to define the grid all_dists = [] for k in lap_keys: d = laps[k][dist_col].values all_dists.append(d[-1]) if qss_df is not None: all_dists.append(qss_df.index.values[-1]) grid_max = min(all_dists) grid = np.arange(0, grid_max, dx) # Resample all laps sources_var: dict[str, np.ndarray] = {} sources_vel: dict[str, np.ndarray] = {} for k in lap_keys: label = _label(k) d = laps[k][dist_col].values sources_var[label] = np.interp( grid, d, laps[k][variable].values, left=np.nan, right=np.nan ) sources_vel[label] = np.interp( grid, d, laps[k][vel_col].values, left=np.nan, right=np.nan ) # Add QSS if provided if qss_df is not None: qss_var_col = _resolve_qss_col(variable, qss_df) qss_vel_resolved = _resolve_qss_col(vel_col, qss_df) qd = qss_df.index.values sources_var["QSS"] = np.interp( grid, qd, qss_df[qss_var_col].values, left=np.nan, right=np.nan ) sources_vel["QSS"] = np.interp( grid, qd, qss_df[qss_vel_resolved].values, left=np.nan, right=np.nan ) all_labels = list(sources_var.keys()) # --- Resolve baseline --- if baseline == "fastest": bl_label = _label(min(lap_keys, key=lambda k: laps[k]["LapTime"].iloc[-1])) elif baseline == "first": bl_label = _label(lap_keys[0]) elif baseline == "last": bl_label = _label(lap_keys[-1]) elif baseline == "qss": if qss_df is None: raise ValueError("baseline='qss' requires qss_df") bl_label = "QSS" else: bl_label = _label(baseline) if isinstance(baseline, int) else str(baseline) if bl_label not in sources_var: raise ValueError( f"baseline={baseline!r} not found. Available: {all_labels}" ) bl_var = sources_var[bl_label] bl_vel = sources_vel[bl_label] # --- Assign colors: baseline last, distinct --- non_bl = [lb for lb in all_labels if lb != bl_label] color_map: dict[str, str] = {} for i, lb in enumerate(non_bl): color_map[lb] = palette[i % len(palette)] # Baseline gets a color that is not in the non-baseline set bl_color_idx = len(non_bl) % len(palette) color_map[bl_label] = palette[bl_color_idx] # --- Build figure --- fig = make_subplots( rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.08, row_heights=[0.75, 0.25], subplot_titles=( f"{variable} vs Distance", f"Cumulative Time Delta (ref: {bl_label} — positive = slower)", ), ) # Top: variable traces — non-baseline first, then baseline on top for lb in non_bl: fig.add_trace( go.Scattergl( x=grid, y=sources_var[lb], mode="lines", name=lb, legendgroup=lb, line=dict(color=color_map[lb]), ), row=1, col=1, ) fig.add_trace( go.Scattergl( x=grid, y=bl_var, mode="lines", name=f"{bl_label} (baseline)", legendgroup=bl_label, line=dict(color=color_map[bl_label], width=2.5), ), row=1, col=1, ) # Bottom: time deltas for each non-baseline series with np.errstate(divide="ignore", invalid="ignore"): dt_bl = np.where(bl_vel > 0.5, dx / bl_vel, 0.0) for lb in non_bl: vel = sources_vel[lb] with np.errstate(divide="ignore", invalid="ignore"): dt_cmp = np.where(vel > 0.5, dx / vel, 0.0) dt_cmp = np.where(np.isnan(vel), dt_bl, dt_cmp) cum_delta = np.cumsum(dt_cmp - dt_bl) hover = [] for dt in cum_delta: if dt >= 0: hover.append(f"{lb} slower by {abs(dt):.3f}s") else: hover.append(f"{lb} faster by {abs(dt):.3f}s") fig.add_trace( go.Scattergl( x=grid, y=cum_delta, mode="lines", name=f"Δt {lb}", legendgroup=lb, showlegend=False, line=dict(color=color_map[lb]), hovertext=hover, hoverinfo="text+x", ), row=2, col=1, ) final = cum_delta[-1] if len(cum_delta) > 0 else 0.0 tag = "slower" if final >= 0 else "faster" print(f" [trace] {lb}: {tag} by {abs(final):.3f}s vs {bl_label}") fig.add_hline(y=0, line_dash="dot", line_color="gray", row=2, col=1) fig.update_yaxes(title_text=f"{variable} ({units})", row=1, col=1) fig.update_yaxes(title_text="Δt (s) [+ = slower]", row=2, col=1) fig.update_xaxes(title_text="Distance (m)", row=2, col=1) fig.update_layout( height=650, title=f"{variable} Comparison — Baseline: {bl_label}", showlegend=True, legend=dict(orientation="h", y=-0.08), ) return fig
# --------------------------------------------------------------------------- # Backward-compatible wrappers # ---------------------------------------------------------------------------
[docs] def plot_velocity_heatmap( comp_df: pd.DataFrame, track: Track, log_vel_col: str = "groundSpeed", qss_vel_col: str = "qss.velocity", ) -> go.Figure: """3-panel velocity heatmap (IRL vs QSS). Thin wrapper around :func:`plot_heatmap` for backward compatibility. """ log_cols = [c for c in comp_df.columns if not c.startswith("qss.")] qss_cols = [c for c in comp_df.columns if c.startswith("qss.")] log_part = comp_df[log_cols].copy() log_part["posN"] = track.y_m[: len(comp_df)] log_part["posE"] = track.x_m[: len(comp_df)] qss_part = comp_df[qss_cols].copy() qss_part["qss.posN"] = track.y_m[: len(comp_df)] qss_part["qss.posE"] = track.x_m[: len(comp_df)] return plot_heatmap( log_part, qss_part, variable=log_vel_col, baseline_label="IRL", comparison_label="QSS", units="m/s", )
[docs] def plot_velocity_trace( comp_df: pd.DataFrame, log_vel_col: str = "groundSpeed", qss_vel_col: str = "qss.velocity", ) -> go.Figure: """F1-style velocity trace (IRL vs QSS). Thin wrapper around :func:`plot_comparison_trace` for backward compatibility. """ dist = comp_df.index.values log_cols = [c for c in comp_df.columns if not c.startswith("qss.")] qss_cols = [c for c in comp_df.columns if c.startswith("qss.")] log_part = pd.DataFrame( {log_vel_col: comp_df[log_vel_col].values}, index=comp_df.index, ) log_part["LapDist"] = dist log_part["LapTime"] = np.cumsum( np.where( comp_df[log_vel_col].values > 0.5, np.diff(dist, prepend=dist[0]) / comp_df[log_vel_col].values, 0.0, ) ) qss_part = comp_df[qss_cols].copy() return plot_comparison_trace( {"IRL": log_part}, variable=log_vel_col, qss_df=qss_part, baseline="qss", vel_col=log_vel_col, units="m/s", )