Example notebooks

Introductory examples

If you haven't already, check out the tutorials, which explain how to build and simulate models in Jaxonomy.

Primitive blocks and composability

Shows how to build systems with primitive blocks and how to compose them into larger diagrams.

Block diagram visualization

Render a DiagramBuilder composition as a block diagram to inspect wiring, ports, and hierarchy before simulating.

Custom block authoring: from `LeafSystem` to gradient-correctness CI

Walks through authoring a custom LeafSystem block end-to-end — declaring state, parameters, and ports; wiring the callback signatures — and standing up the gradient-correctness test that guards it.

Bouncing ball

Shows hybrid dynamics modeling of a bouncing ball.

Hybrid thermostat (marimo tutorial)

A marimo notebook implementing a thermostat-controlled room: continuous temperature dynamics with discrete HEAT/OFF modes and hysteresis. Run with marimo edit hybrid_thermostat_tutorial.py or marimo run hybrid_thermostat_tutorial.py (requires marimo and jaxonomy).

Triple inverted pendulum: Jaxonomy LQR + MuJoCo (marimo)

A marimo app that stabilizes a triple pendulum on a cart. Jaxonomy's LinearQuadraticRegulator synthesizes the stabilizing gain (a continuous-time algebraic-Riccati solve on the MuJoCo-linearized plant), and MuJoCo mj_step is the high-fidelity physics it is validated against: a short open-loop kick-up impulse, then closed-loop LQR balance, with a matplotlib + MuJoCo animation. Run with marimo run triple_inverted_pendulum_mujoco_marimo.py (requires marimo, jaxonomy, mujoco).

Bouncing balls on stairs: Jaxonomy dynamics + MuJoCo visualization (marimo)

A marimo app where three balls (bocce, golf, ping-pong) cascade down a five-step staircase. Jaxonomy computes the motion as a hybrid dynamical system — each ball a point mass in free fall with a zero-crossing event at every tread and a Newtonian restitution reset (\(\dot z \mapsto -e\,\dot z\)) — and MuJoCo renders it kinematically (ball positions set from the Jaxonomy trajectory each frame; mj_forward, no contact solver). The restitution coefficient \(e\) alone (0.85 / 0.70 / 0.55) produces the light/medium/heavy bounce characters, and the energy panel shows each bounce keeping a fraction \(e^2\) of the vertical KE. A multi-ball, staircase counterpart to the Jaxonomy-native bouncing ball example above. Run with marimo run bouncing_ball_stairs_marimo.py (requires marimo, jaxonomy, mujoco).

Linear Quadratic Regulator (LQR)

Demonstrates the LQR for a pendulum and a planar quadrotor model.

Energy shaping and LQR stabilization

Demonstrates energy shaping control to swing a pendulum to the vertically 'up' orientation and then stabilize it in the 'up' orientation via LQR.

Linear Model Predictive Control (MPC)

Demonstrates MPC on a linearized model of the Cessna Citation aircraft and a pendulum model.

Differentiable Predictive Control (DPC) of a two-tank system

Trains a neural control policy to make the bottom level of a cascaded two-tank rig track a setpoint, by differentiating straight through the simulated ODE. Shows the policy-as-a-block API (PolicyBlock, PlantBlock, build_closed_loop, simulate_closed_loop): the same policy object trains as a differentiable function and deploys as a composable LeafSystem under jaxonomy.simulate. Deployment note for policies imported from discrete-time training loops (torch/NEUROMANCER-style, via ONNXJax or the predictor blocks): sample-and-hold controllers must be followed by ZeroOrderHold(dt=ts) with SimulatorOptions(max_major_step_length=ts, max_minor_step_size=ts) for step-grid parity — otherwise the policy is silently re-evaluated at every RK4 stage and step-for-step parity with the exporting framework is lost (with the ZOH: parity ~4e-8 over 400 steps on this benchmark).

Multi-layer perceptron (MLP)

Demonstrates training of a multi-layer perceptron (MLP), a class of feedforward artificial neural networks, for a regression task.

Advanced examples

Trajectory optimization and stabilization

Shows trajectory optimization for the problem of swinging an Acrobot to the vertically 'up' orientation and then stabilizing the trajectory via finite-horizon LQR.

Robotic arm control

Implement a controller for a "pick-and-place" task with a robotic arm using MuJoCo as a multibody physics engine. Download the necessary files from here.

Automatic tuning of a PID controller

Demonstrates automatic differentiation and optimization capabilities of Jaxonomy to automatically tune the gains of a discrete-time PID controller.

Interactive and automatic tuning of a PID controller with sensitivity constraints

Showcases fast compiled simulations in Jaxonomy for interactive applications and automatic tuning of a continuous-time PID controller with maximum sensitivity and complementary sensitivity constraints.

Finding limit cycles

Demonstrates how to find limit cycles and assess their stability by leveraging the automatic differentiation capabilities of Jaxonomy.

Kalman Filters: linear and nonlinear extensions

Demonstrates the use of Kalman filters (finite and infinite-horizon) and nonlinear extensions (Extended Kalman Filter and Unscented Kalman Filter) for state estimation in a pendulum model. Where necessary, the nonlinear Pendulum plant is automatically linearized and discretized by Jaxonomy for the construction of the filters.

Engine map fitting to MPC: differentiable lookup tables end-to-end

Fit a noisy 2-D engine torque map with fit_lookup_table_2d, optimise breakpoint placement with fit_table_1d_with_grid, drop the fitted LookupTable2d into a 1-DOF vehicle plant, run a short-horizon shooting MPC tracking a velocity reference, and take a single jax.grad of closed-loop tracking RMSE with respect to the table values, the grid placement, and the MPC weights — all in one call. Showcases the differentiable lookup-table family.

Linearization workflow: from findop to Bode / Nyquist / pzmap and empirical FRE

The full linearization toolchain on one canonical plant: trim with findop, take Jacobians with linearize, evaluate frequency response with frequency_response / bode_data, plot Nyquist contours with nyquist_data, scatter poles and zeros with pole_zero_map, and overlay closed-form step_response / impulse_response against the nonlinear simulator. Finishes with estimate_frequency_response driven by a chirp (matching the analytic Bode to within ~1 dB) and a jax.grad of the Bode-peak magnitude w.r.t. damping — the differentiability a JAX-native stack adds on top of a classical linearization workflow.

Aleatoric vs epistemic uncertainty: Sobol decomposition on a noisy plant

End-to-end tour of jaxonomy.uq: sanity-check Sobol indices on the Ishigami benchmark, compare IID / Latin-hypercube / quasi-Monte-Carlo convergence on a damped pendulum, rank parameters by first- and total-order Sobol index, screen with Morris elementary effects, and finish with the headline decompose_variance_sobol call that splits the QoI variance into aleatoric (irreducible) and epistemic (reducible-with-better-data) shares plus their cross-interaction — a distinction that is often hard to surface cleanly. Closes with a value-at-risk / CVaR computation that connects the variance budget to a concrete design margin.

Bit-exact reproducibility: capturing and replaying a simulation

Run a damped oscillator with SimulatorOptions(record_provenance=True), persist the resulting ProvenanceManifest and parameter pytree as sibling JSON files, replay the simulation on the same machine and byte-compare the outputs via array.tobytes() equality, then deliberately tamper with the model — first by a single-ULP parameter perturbation, then by a sign flip in the dynamics — and watch the manifest's stable fingerprint and the parameter SHA each catch the right failure mode. Honest about the contract: shows precisely what the manifest fingerprints (versions, options, system type, sorted parameter names), what it deliberately doesn't (parameter values, source code, cross-device floats), and how to layer a parameter hash next to the manifest for full value-drift detection. Few simulation tools emit this artifact by default — it's the trust-building layer for reproducible results.

Hybrid trajectory optimization through events: gradients across the bounce

A bouncing-ball plant with record_event_times=True exposes touchdown instants to autodiff: jax.grad flows through the bounce via the saltation rule, so a scalar objective like "ball energy at \(t=T\)" becomes differentiable w.r.t. the restitution coefficient. Demonstrates closed-form vs simulator bounce-time agreement to \(\sim 10^{-8}\) s, a single-point dt_event/dh0 = 0.2258 s/m matching analytic \(1/\sqrt{2 g h_0}\), a multi-event tilted-floor case with one gradient per bounce, and vmap_event_time_gradient for batched gradients across a 16-point initial-height sweep. Honest about two cracks surfaced while authoring: with_parameter("e", float(e)) triggers a fresh JIT trace per call (filed as a follow-up finding), and the multi-event saltation gradient disagrees with finite differences by 0.16 s/(grade) on the tilted-floor first bounce — a correctness regression also filed. The headline capability: gradients across discrete contact events, which most block-diagram tools don't offer.

Networked control: identifying actuator delay from bench data

Synthesize noisy open-loop bench data on a 1st-order servo with a known transport delay τ_true, define a mean-squared prediction error against the bench data as a function of τ, and watch jax.grad flow cleanly through the VariableTransportDelay block (the linear-interpolation lookup is what makes the gradient finite and nonzero). A handful of scipy iterations recover τ to within ~5% of ground truth; a higher-noise run shows the estimator gets noisier but stays unbiased; and a closing closed-loop demo contrasts the unmodeled-delay baseline (oscillating, saturated) against a Smith-predictor controller built around the identified delay. Showcases the continuous + variable transport delay blocks and the differentiability story behind the saltation rule.

Multirate controller: 1 kHz inner / 100 Hz outer / 10 Hz supervisor

Build the canonical embedded-control cascade on a PM DC motor — 1 kHz current loop, 100 Hz velocity loop, 10 Hz position controller plus state-machine supervisor — and route a four-field measurement bus between them with BusCreator / BusSelector. Place RateTransition blocks explicitly between the layers and also exercise DiagramBuilder(auto_insert_rate_transitions=True); inspect what the scheduler actually built with Diagram.print_schedule() and rate_summary_dot(). Closes with analyze_saturation and analyze_phase_activity on the closed-loop trajectory, calling out the benign FAULT-never-fired warning and pointing the reader at the FAULT-trip exercise. Showcases multirate scheduling + auto-inserted rate transitions, RateTransition / Decimator, and BusCreator / BusSelector.

Real-time fixed-step controller: the embedded deployment story

The same notebook that runs the offline simulator is the reference for the deployed binary — fixed step, deterministic byte-for-byte, microseconds per tick. Builds a 1 kHz discrete PID plus a Luenberger velocity observer wrapped around a lightly-damped 2nd-order spring-mass plant with noisy position measurement, then compares SimulatorOptions(ode_solver_method="rk4") against the adaptive dopri5 solver on three embedded-engineering metrics: per-tick wall-clock jitter (the WCET budget the deadline scheduler must accommodate), bit-identical determinism on re-runs, and post-JIT throughput. Headline live numbers: RK4 mean per-tick wall-clock = 24.1 μs, 99/1 jitter ratio = 1.60×; Dopri5 jitter ratio under the same max_major_step_length=DT constraint = 1.52× (the host-callback synchronisation floor dominates per-tick timing on a developer laptop — the real algorithmic wedge is in the RHS-evaluation count). Under natural adaptive operation, Dopri5 takes 1.20× more RHS evaluations than RK4 for the same controller horizon (6.01 vs 5.00 calls/tick) — and this is the embedded WCET metric that translates faithfully to an MCU because RHS count is fixed floating-point work. Post-JIT throughput in steady state: 45.7 μs / tick (= ~21 kHz on a developer laptop, 4.6% CPU budget at the 1 kHz control rate). Bit-identical re-runs verified on both solvers via tobytes(); RK4 vs Dopri5 trajectories agree to 1.5 μm at final time (= 0.015% of the 10 mm setpoint). Closes with analyze_saturation on the actuator (0.01% saturated, well under the 50% threshold) and five exercises (10 kHz rate, float32 opt-out, fault-detection state machine, XLA HLO inspection, per-tick budget breakdown on a real embedded project). Honest about the JIT-warmup cost, the XLA-CPU vs bare-metal-MCU gap, and the difference between bit-exact reproducibility of the host development sim and of the deployed binary. Showcases fixed-step RK4 via SimulatorOptions(ode_solver_method="rk4"), byte-exact reproducibility, and the precision policy (exercise-only).

Fast restart and batched sweeps: when JIT amortization pays off

Time a 1000-trial damping sweep on a damped harmonic oscillator three ways — a naive simulate() loop, FastRestartSimulator's warm-cached kernel, and simulate_batch(use_vmap=True)'s vectorised XLA launch — and watch the JIT-amortisation wedge open up across two orders of magnitude. On CPU, FastRestartSimulator recovers the closed-form ISE-optimal damping (\(\zeta = 0.5\)) from a brute-force grid in ~0.3 s versus ~130 s for the naive loop — a ~430× speedup with the four code paths agreeing bit-for-bit on the per-trial ISE. The tutorial is candid about where vmap surprises a CPU user (the per-row finalize is linear in \(N\), so vmap loses to the kernel path on CPU at these scales but wins on GPU/TPU where the parallel launch dominates), exercises three concrete failure modes — FastRestartSimulator's structural-change warning on a dtype swap, simulate_batch(use_vmap=True)'s host-callback ValueError, and the simulate(...).buffer_length ring-buffer pitfall (now warned-on as of 7a24e31) — and closes with a decision table that names the regime where each API wins. Showcases FastRestartSimulator + simulate_batch (kernel and vmap paths).

Product-family modeling with Variants: one diagram, three drivetrains

Ship three trim levels — gasoline ICE, parallel hybrid, battery-electric — of the same longitudinal vehicle model from one DiagramBuilder and one Variant at the powertrain attachment point. select_variant resolves the topology at build time (unselected branches never enter the JIT trace); Diagram.with_config(powertrain=...) swaps the realised choice post-build; dump_variant_config_to_json / load_variant_config_from_json round-trip the active binding (bytes-equal on replay); the python -m jaxonomy.cli.run_variants CLI surfaces the same data from release-pipeline scripts. Closes with a 3 × 5 Cartesian-product sweep of 0-60 time over (drivetrain, driver-aggressiveness) via expand_all_variant_configs + simulate_batch, all 15 sims under one JIT-compiled kernel per variant. Showcases the Variants system (phases 1–4) plus the with_config / introspection / runtime-switch follow-ups, and cross-links the substrate submodel_function the variant DSL extends.

Truth tables and state machines: vehicle gear-selection logic

A four-bit eligibility truth table (brake, stopped, req_b1, req_b0) drives a four-state Park / Reverse / Neutral / Drive FSM with ten prioritised transitions and — deliberately — no Drive ↔ Reverse direct edge. Builds the table with the fluent TruthTable.builder(...).row(...).build() API (wildcards on the driver-request bits collapse 16 input combinations to 3 rows), runs validate(strict_completeness=True) and watches the strict check catch a deliberately-removed row, byte-equal round-trips the table through to_csv / from_csv, and runs a 14-second scripted driving trace where the FSM silently refuses an illegal Drive → Reverse request (caught post-hoc by analyze_phase_activity on a truncated trace). Closes with jax.grad flowing through a standalone gain-schedule table whose row outputs are callables of the raw inputs (the numeric-output extension). Honest about the four DX papercuts the tutorial surfaced — TruthTable being branchless not vectorised (a follow-up finding), TruthTableBuilder(input_names=...) labels not propagating to port names, no time_mode= knob on StateMachineBuilder.build(), and overlapping_pairs=[] being counter-intuitive on wildcard tables. Showcases the TruthTable API + the existing StateMachineBuilder API.

Two-degree-of-freedom PID with classical auto-tuning rules

Hand the practicing engineer a starting point from two minutes of bench time. Identify the ultimate-cycle parameters \((K_u, T_u)\) via a relay-feedback experiment, fit the FOPDT triple \((K, \tau, \theta)\) off the open-loop step's inflection tangent, then apply PIDController2DOF.ziegler_nichols(Ku, Tu, dt, mode="PID"), .cohen_coon(K, tau, theta, dt, mode="PID"), and .tyreus_luyben(Ku, Tu, dt) to the same FOPDT plant (\(K=2\), \(\tau=10\ \mathrm{s}\), \(\theta=2\ \mathrm{s}\)). Compare on a setpoint step + load disturbance with quantitative metrics (rise time, overshoot, 2% settling time, windowed ISE) — Z-N moderate, Cohen-Coon most aggressive (best disturbance ISE, worst overshoot), Tyreus-Luyben most conservative (no overshoot, slowest recovery). The 2-DOF wedge then shows the same Z-N gains under three \((b, c)\) settings: setpoint overshoot drops from 48% → 16% → 12% while the disturbance maximum stays pinned at 1.194 to four significant figures — the empirical demonstration that \((b, c)\) change only the setpoint-tracking numerator, not the closed-loop poles. Closes with jax.grad refinement of the Cohen-Coon starting point (25 steps, 22% ISE reduction in ~2.5 s), analyze_saturation / analyze_control_oscillation diagnostics on every controller (all silent), and a concrete failure-mode demo (Cohen-Coon on a misidentified second-order plant overshoots 68% with visible ringing). Showcases PIDController2DOF + the three classmethod tuning rules, and bridges to the gradient-based and sensitivity-constrained stories in pid_tuning.ipynb and pid_autotuning_interactive.ipynb.

Unit-safe wiring: dimensional consistency at build time

Annotate ports with Unit (meter, newton, joule, …) and watch DiagramBuilder.connect() refuse a force-into-displacement wire before any kernel launches; auto-convert mm→m at the wire (with unit_conversion="auto" / "warn" / "error" modes); run propagate_diagram_units(diagram) and see the math-block algebra catch Constant(units=newton) → Integrator → output declared as joule because N·s ≠ J — the headline beat that closes Simulink's "units dropped after computational blocks" gap. Tags BusCreator(field_units=...) for compound bus signals, demonstrates offset-aware °C↔K via convert_offset_aware, sets a USD→EUR FX rate to exercise the currency axis, and shows the acausal connector library's canonical flow_units / pot_units per domain plus the physical_quantity tag that distinguishes N·m@torque from N·m@energy. Closes with JSON round-trip of unit annotations. Showcases the units system (phases 1–3), the BusUnit cross-link, and the currency / temperature / pint / source-block follow-ups.

Conservation laws as CI: patterns for ensuring physical validity

Every real-world simulator drifts. Energy bleeds from an oscillator, angular momentum tilts on a tumbling body, probability mass leaks out of a diffusion — the drift is usually invisible in the trajectory but real and shipping. Pedagogical guide to jaxonomy's property-test framework for asserting that physical invariants hold within tolerance over long simulation horizons. Walks four canonical patterns (energy on an undamped SHO, angular momentum on a torque-free rigid body, probability mass on a 5000-path Brownian ensemble, electrical energy on an acausal LC oscillator under BDF) and shows the test catching deliberately-broken plants in each one. Headline numbers: SHO energy conserved to relative drift \(1.45 \times 10^{-9}\) over 50 oscillation periods under Dopri5 at rtol=1e-10; a stray \(-cv\) damping bug (\(c = 10^{-3}\)) pushes the drift to \(1.87 \times 10^{-2}\) at the endpoint — seven decades above the floor — and the test catches it. Rigid-body \(\lVert \mathbf{H} \rVert^2\) drift under the same solver = \(5.75 \times 10^{-10}\) over 20 s of intermediate-axis tumble; a \(5\%\) asymmetric scaling on one cross-product component pushes drift to \(3.0 \times 10^{-2}\) (caught). Brownian variance: empirical \(\langle x^2 \rangle(T) = 4.975\) vs analytic \(2DT = 5.000\) for \(N = 5000\) paths (z-score \(-0.25\), within \(\pm 3\sigma\)); a \(10\%\) noise-amplitude scaling bug produces \(z = +10.2\) (caught at the \(3\sigma\) envelope). LC oscillator under BDF: energy drift \(1.94 \times 10^{-8}\) over 10 oscillation periods (the BDF \(10^{-4}\) envelope is loose because BDF dissipates). The pedagogical climax: in section 4 a small visible-bug, invisible-trajectory comparison forces the reader to internalise why visual inspection is insufficient and the property test is strictly stronger. Closes with the "where to put these tests in your project" pattern (test/<your_block>/test_conservation.py), five failure modes (tolerance-too-tight, statistical false positives, damped-system mis-test, chaotic-system non-invariants, endpoint-only vs max-over-trajectory), and five exercises (anti-conservation test for a damped SHO, charge conservation on the LC, jax.grad of drift w.r.t. solver tolerance, Lorenz divergence as a flow invariant, planar-pendulum constraint residual under Baumgarte/SSP/none). The marketing wedge: the property-test framework is to physical validity what gradient-correctness CI is to autodiff — defense-in-depth against silent failure modes the simulator cannot detect on its own. Runtime: ~6 s end-to-end. Cells: 21 code + 29 markdown = 50 total, 4 matplotlib figures. Showcases the conservation-test framework (test/conservation/_framework.py::assert_conserved), the acausal LC oscillator under BDF, and the stochastic-source / probability-mass conservation pattern.

Differentiable acausal DAEs: learning unmodeled drag inside a constrained pendulum

The differentiable-acausal capability in one notebook. A PlanarPendulum on a rigid link is a holonomic constraint \(x^2+y^2=L^2\) — an index-3 DAE compiled through Pantelides index reduction — and the real rig has an unmodeled quadratic aerodynamic drag \(-c\,v\,|v|\) the modeler never wrote. We bolt a small MLP \(f_{NN}(v;\boldsymbol\theta)\) onto the compiled DAE's differential rows as a learned correction via the phase-2 NeuralDAEBlock (targets=[(pend, "v")], gather-in/scatter-out, injected after index reduction so the non-symbolic neural term never enters the symbolic Pantelides path), then fit \(\boldsymbol\theta\) by gradient descent through the BDF-DAE adjoint — jax.grad flows from a terminal-state loss, through three implicit constrained-DAE integrations, into the network weights. This is the wedge no single tool spans: Modelica expresses the constraint but can't autodiff through it; causal UDE tools (Neuromancer, DiffEqFlux) autodiff but can't express the algebraic constraint coupling. Headline live numbers (no checkpoint — the fit runs live to demonstrate the adjoint): the multi-horizon terminal-state loss drops from \(4.97\times10^{-2}\) to \(1.85\times10^{-4}\), a ~269× reduction in 25 Adam steps (~7.3 s/grad through 3 BDF-DAE solves, ~3 min fit on laptop CPU); the fitted model's angle RMS vs the damped truth is 0.0118 rad vs 0.1088 rad for the no-drag baseline (9.2× better) over a 6 s window it was never trained on. The conservation validation confirms the compile is correct: total mechanical energy is conserved without drag and decays monotonically with it, and the holonomic constraint residual \(|x^2+y^2-L^2|\) holds at \(\sim10^{-14}\) m² (the DAE-projection invariant). Closes with an honest identifiability beat — the learned \(f_{NN}(v)\) tracks the true drag inside the velocity band the trajectory explored (RMS 0.07 m/s²) but diverges outside it (RMS 1.35 m/s²), because you can only learn the dynamics where the data took you. Surfaces the genuine sharp edge it had to work around: index reduction's choice of which symmetric velocity (\(v_x\) vs \(v_y\)) survives as the differential state is PYTHONHASHSEED-dependent, so the demo resolves it at runtime by following alias_map into sed.x (filed as a follow-up finding); plus the recorded_signals ⊥ autodiff constraint that forces the terminal-state loss, the buffer_length overflow on the visualization path, and the Adam-overshoot early-stop. Five exercises (linear-drag swap, harder excitation to widen the identifiable band, hand-derive \(\lambda\) from the acceleration-level constraint, compose a second block on the position row, symbolic-regress \(f_{NN}\) back to closed form). Runtime: ~4–5 min end-to-end (live fit). Showcases the phase-2 NeuralDAEBlock + AcausalDiagram.add_neural_correction_block, the acausal Pantelides → BDF pipeline, and the autodiff-through-the-DAE-adjoint story; sibling to the causal-ODE UDE notebook and the conservation-laws-as-CI notebook.

Stiff chemistry: the Robertson problem under BDF

The canonical stiff-ODE test case in numerical analysis (Robertson 1966; Hairer & Wanner 1996 Vol II §IV.1; Shampine 1994 Ch. 6): three coupled species with rate constants \(k_1=0.04\), \(k_2=3 \times 10^7\), \(k_3=10^4\) spanning 9 orders of magnitude. Walks through what "stiff" actually means (the stability step is much smaller than the accuracy step — Curtiss & Hirschfelder 1952), why explicit Dopri5 catastrophically over-resolves the fast B-mode, and why implicit BDF cruises to steady state in a few hundred adaptive steps. Headline numbers: BDF integrates the full \(t \in [0, 10^{11}]\) Robertson horizon in ~0.4 s wall-time (854 recorded samples, mass conservation \(|y_1+y_2+y_3-1| < 10^{-15}\)); Dopri5 at \(T=1\) s already needs 713 minor steps (~0.13 s wall), giving a linear extrapolation of ~600 years of wall-time to clear the full horizon. The stiffness crossover where Dopri5 step count first exceeds BDF sits at \(k_2 \approx 10^7\) (ratio \(k_2/k_1 \sim 2.5 \times 10^8\)). Closes with a Jacobian-spectrum visualization (jax.jacrev on the RHS, six probe points, spectral ratio from 1 to \(10^{19}\)), the canonical Robertson log-log trajectory plot, BDF-vs-Dopri5 horizon scaling (BDF flat across 11 decades, Dopri5 linear in \(T\)), \(k_2\)-stiffness sweep, a one-paragraph "why BDF works" explanation with the Newton-iteration math, three practical rules of thumb, a central-difference gradient cross-check (\(\partial y_3(10^4) / \partial k_1 = 3.438\)), and five exercises (less-stiff Dopri5 sweep, source-flux extension, multi-parameter gradient, HIRES/OREGO port, Radau IIA implementation). 3 follow-up findings: (i) int_time_scale not auto-promoted for \(t > 9 \times 10^6\) s; (ii) max_major_step_length is JIT-static so horizon sweeps re-compile per iteration; (iii) BDF's recorded n_samples measures major-step recording cadence, not solver step count — solver-agnostic step counting requires instrumentation. Cells: 12 code + 22 markdown = 34 total, 4 matplotlib figures. Runtime: ~10 s on warm JIT cache, ~60-90 s cold. Showcases BDF as a first-class jaxonomy solver + the canonical stiff-ODE diagnosis workflow + the gradient-through-implicit-solver story.

Container blocks tour: `ForEach`, `EnabledSubsystem`, `TriggeredSubsystem`

Three canonical container-block patterns end-to-end on three real-engineering motivations: (a) ForEach-vectorised 100-cell battery sweep with per-cell capacity/resistance tolerance, (b) EnabledSubsystem-wrapped 4-wheel ADCS where one reaction wheel fails mid-mission and the remaining three redistribute disturbance-rejection torque, (c) TriggeredSubsystem gating a Kalman measurement update on the rising edge of a 1 Hz star-tracker pulse. Each pattern closes with a jax.grad through the container boundary — the autodiff wedge over Simulink. Headline numbers: (A) per-cell ∂(mean SOC)/∂C_i = 1.528e-3 matches the closed-form Coulomb counter to 0.01% in one backward pass; (B) post-failure per-wheel torque grows from 10.5 mN·m to 13.3 mN·m (ratio 1.27, theoretical 4/3 = 1.333), failed wheel held at exactly 0 mN·m by EnabledSubsystem(mode="reset"); (C) d(ISE)/dK = −7.39 autodiff matches central-difference within 0.10%. Honest about the phase-1 TriggeredSubsystem running its child on every sample step (gates the latch, not the kernel) and the requires_inputs=False declaration needed to break the cycle checker on no-feedthrough output ports — both written up as follow-up findings. Runtime: ~7 s end-to-end. Showcases EnabledSubsystem + TriggeredSubsystem + ForEach.

Multi-domain HVAC: a heat pump heating a Stockholm apartment

A 100 m² apartment in Stockholm in mid-January under an air-source heat pump whose COP swings 1.5x across the diurnal ambient cycle. Builds the two-state RC envelope as an acausal thermal network (HeatCapacitor + Insulator + TemperatureSource + HeatflowSource(enable_heat_port=True)), wraps the heat pump as a causal LeafSystem whose temperature-dependent COP comes from an npa.interp lookup table, and closes the loop with a discrete-time PI controller updating at 1/60 Hz. Headline unit-safety beat: a deliberate m^3/s (volumetric refrigerant flow) -> kg/s (mass flow) wiring mistake at the causal/acausal seam raises UnitMismatchError at DiagramBuilder.connect() time, before any kernel launches; the fix is an explicit density-conversion Gain(rho) with the algebra verified on both ends. Honest about the units follow-up (canonical acausal units exposed but the Pantelides pass does not yet consume them — filed as a follow-up finding); the causal-seam check is the production-grade boundary that catches the bug today. 24 hours integrate in ~1 s under BDF (~16k samples — see ring-buffer note), tracking RMSE drops to mK during occupied hours, total electricity bill comes in at 12.97 kWh / day, and a four-parameter central-difference sensitivity analysis ranks wall R-value \(\gg\) COP \(>\) set-back depth \(\gg\) thermal mass for retrofit ROI (the autodiff workaround caveat from the battery tutorial applies). Closes with analyze_saturation + analyze_control_oscillation (both silent), failure-modes section (deep-cold backup heater unmodelled, manufacturer-specific COP curve, two-state envelope vs ASHRAE 90.1 multi-zone, simplified refrigerant cycle, mean-ambient wall coupling), and five exercises (COP-curve swap, dawn pre-heat, multi-zone, optimal schedule, hybrid heater). Showcases unit-safety at the cross-domain seam + acausal thermal library + LeafSystem composition + autodiff sensitivities.

FMI export round-trip: build a controller in jaxonomy, co-simulate with an external tool

Build a discrete-time PI controller as a DiagramBuilder composition, export it as a binary FMI 2.0 Co-Simulation .fmu via jaxonomy.library.fmu_export.build_fmu, then re-import the binary three ways: as a jaxonomy.library.ModelicaFMU block inside a fresh diagram (Architecture B), and via raw fmpy.fmi2.FMU2Slave calls orchestrated by a hand-rolled Python loop (Architecture C). Headlines: the manual fmpy orchestration agrees with a pure-Python reference loop on the same difference equations to max abs error 7.8e-4 on the plant state and 4.0e-3 on the control output over 400 steps — the residual is FMI 2.0 §4.2.4's first-step initialization protocol, not a numerical artifact (the per-step error drops to ~1e-4 on x and ~1.6e-5 on u after k=1). Demonstrates JaxonomyDiagramSlave wiring, the auto-exposed Constant-block input convention, and the structural-XML round-trip showing the modelDescription.xml jaxonomy generated is what any commercial FMI host (OpenModelica, AVL CRUISE M, IPG CarMaker, dSPACE SCALEXIO, MATLAB Simulink) reads at import time. The publication/fast-mode pattern caches the heavy ModelicaFMU JIT compile (~22 s) under media/fmi_export_publication.npz so the reader's notebook runs in ~3 seconds. Honest about three sharp edges: the ModelicaFMU block's periodic-update offset=dt introduces a 1-sample phase lag relative to in-process discrete blocks (Architecture A vs B mismatch in the transient), pythonfmu's embedded-Python single-instance dylib limit forces a subprocess for the manual loop, and the FMU boundary is not differentiable (FMI export is for delivery, not for training). One more honest number: each co-simulation doStep re-enters jaxonomy.simulate per segment, costing ~100+ ms per step versus ~tens of µs per step in-process on a small system — the FMI boundary is for validation and interop, not for training loops. Marketing wedge: jaxonomy speaks FMI end-to-end, so it slots into every existing embedded-systems toolchain. Showcases binary FMU export, the Darwin pythonfmu wrapper, auto-exposed Constant block inputs, and FMI 2.0/3.0 import via ModelicaFMU.

Physics-informed learning across stacks, part 1: policy export via ONNX

First of a three-part series pairing jaxonomy with NEUROMANCER (PNNL's PyTorch differentiable-programming/DPC library) around one physical system — a two-tank pump-and-valve network from NEUROMANCER's own psl benchmark set. This part trains a bounded-MLP DPC policy in NEUROMANCER, exports it to ONNX, and embeds it in a jaxonomy diagram via the differentiable ONNXJax block plus a ZeroOrderHold (the load-bearing detail: without it the policy is re-evaluated at every RK4 stage instead of once per sample). Headline numbers (executed outputs): torch↔onnxruntime export parity 2.68e-7; closed-loop parity against the exporting framework 3.84e-8 over 400 steps; and the payoff neither stack does alone — a jit+vmap Monte-Carlo robustness sweep of 512 plant-parameter samples × 200-step rollouts in ~7 ms steady-state, with the one fragile sample explained exactly by an analytic pump-authority boundary overlaid on the scatter. Ships the trained policy artifacts in media/ so the notebook runs in ~5 s without torch installed; with NEUROMANCER present the training reproduces bit-for-bit. Showcases ONNXJax, ZeroOrderHold, adapter LeafSystems, and jaxonomy.diagnostics on the closed loop.

Physics-informed learning across stacks, part 2: a neural DAE and gradients across the framework boundary

The series centerpiece. Builds the tank network as it really is — reservoir, pump, three-way manifold, two gravity tanks — by writing two custom acausal hydraulic components (~30 lines each) and letting AcausalCompiler produce a genuine semi-explicit DAE: 2 differential states plus 18 algebraic unknowns (the manifold pressure among them), integrated by BDF with a singular mass matrix. Then gradients cross the stack boundary in both directions. Direction 1: the tank1→tank2 orifice is secretly clogged (55% of nominal); a physics-structured neural correction — one scalar flow, scattered with the fixed −1/ρA₁ : +1/ρA₂ ratio — is trained by jax.value_and_grad through the implicit BDF solve and recovers the true flow-deficit law to 1.43% relative error, correlation 0.996 (AD matches finite differences to ~5 significant digits; the unstructured 2-in/2-out variant drops the loss ~90× yet fails to identify the residual — the rank-1 physics prior is the lesson). Direction 2: the DAE plant becomes a NEUROMANCER Node via a torch.autograd.Function whose backward is the JAX VJP over zero-copy dlpack — steady batched DAE steps of ~3 ms — and NEUROMANCER's unmodified Trainer optimizes a policy through it, with honest hold-state/NaN-sanitization guard counters (24 events, all one knife-edge point, root-caused to projection non-convergence). Cites Koch, Shapiro, Sharma, Vrabie & Drgoňa (CDC 2025) for the operator-splitting approach to the same problem class. Showcases custom acausal components, add_neural_correction, project_constraints, dae_initial_projection, and torch↔JAX interop.

Physics-informed learning across stacks, part 3: the plant as an FMI co-simulation FMU

Closes the series at the tool-neutral boundary: the two-tank plant is exported as a binary FMI 2.0 Co-Simulation FMU via build_fmu (passes fmpy.validate_fmu with zero findings, tri-platform binaries) and driven in lockstep by the part-1 policy — rebuilt master-side in plain torch from the weight file, parity-checked against the shipped ONNX at 4.2e-7. Headline numbers: 200-step lockstep loop matches the fully in-process closed loop to max |Δh| = 1.783e-8 m, with the residual explained by solver independence (adaptive integration inside the FMU vs a single RK4 step outside) — which is precisely what co-simulation buys. Honest cost accounting: ~140 ms per doStep through the FMI boundary vs ~69 µs in-process (~2,000×) — the boundary is for validation, certification workflows, and HIL prep, not training loops (part 2 exists because gradients don't cross FMI). Documents the slave sharp edges (output priming, Constant-block input convention, one-instance-per-process) and runs jaxonomy.diagnostics on both actuators. Showcases JaxonomyDiagramSlave, build_fmu, fmpy mastering, and the gained/lost trade-offs of the FMI boundary.

Hybrid ML + physics: dropping a pre-trained predictor into a control loop

The drop-in story for teams who own a PyTorch / TensorFlow training pipeline. Trains a small residual MLP on a damped pendulum with hidden Coulomb friction, loads it as a jaxonomy.library.PyTorch predictor block when torch is installed (falls back to TensorFlow, then to a JAX/Equinox MLP — the public CI configuration), and closes a PD loop around three predictors: pure-physics (no Coulomb), pure-ML (no inductive bias), and hybrid (physics + learned residual). Headline three-way RMSE comparison: in-distribution pure-physics 0.0146 rad, pure-ML 0.0008 rad, hybrid 0.0004 rad (35.5x better than pure-physics); out-of-distribution (release from \(\theta_0 = 2.5\) rad, 5x training envelope) pure-physics 0.0292 rad, pure-ML 0.0330 rad (catastrophic extrapolation), hybrid 0.0019 rad (~15x better than either alone — the structural argument for the hybrid). The autodiff bonus beat takes jax.grad of closed-loop tracking ISE w.r.t. PD gain Kp against the (differentiable) hybrid plant, cross-checks against central-difference to 0.0000% relative error. Honest about the limit: jax.grad does NOT flow through the PyTorch / TensorFlow predictor blocks (same jax.pure_callback boundary that breaks gradient flow through ModelicaFMU); the workaround is to fine-tune against an in-process JAX surrogate of the predictor and validate on the real predictor in a second pass. The publication/fast-mode pattern caches the 4000-epoch training + 4 closed-loop simulations under media/hybrid_ml_physics_publication.npz so the notebook runs in ~5 s; in fast mode the inline retrain runs at 1500 epochs in ~30 s. Marketing wedge: "trained your model in PyTorch on your team's data? Drop it in." Closes Wave 4. Showcases the PyTorch / TensorFlow predictor blocks + the jaxonomy.library.MLP Equinox block as the JAX-native equivalent + the hybrid-ML-physics pattern from Rackauckas et al. 2020 and Karniadakis et al. 2021.

Co-simulating a Modelica plant FMU under a jaxonomy controller

The inverse direction of fmi_export_roundtrip.ipynb: instead of exporting a jaxonomy controller as a binary FMU, import a Modelica-style plant FMU as the closed-loop plant under a jaxonomy PI controller. The combined story (controller-out + plant-in) makes jaxonomy a credible drop-in partner for OpenModelica, Dymola, AVL CRUISE M, IPG CarMaker, and dSPACE workflows. Builds the plant — a damped 2nd-order mass-spring (\(m\ddot{x} + c\dot{x} + kx = F\), \(\omega_n = 1\) rad/s, \(\zeta = 0.25\)) — as a JaxonomyDiagramSlave-wrapped LeafSystem and build_fmu's it to a binary FMI 2.0 Co-Simulation .fmu (~21 ms / 899 KiB on darwin); then closes the loop with the same jaxonomy PI controller around it via ModelicaFMU. Headline numbers (all live or NPZ-cached): closed-loop position trajectory of the FMU-plant Architecture B matches the in-process Architecture A to max abs error 2.46e-2 m over the full 12 s horizon, dominated by the documented ModelicaFMU offset=dt first-step phase lag. The autodiff bonus beat — the marketing wedge over Simulink's FMI Import block — uses the cost-as-Integrator pattern to take jax.grad of the closed-loop tracking ISE w.r.t. the controller gains: dJ/dKp = -0.2437 and dJ/dKi = -1.4586, cross-checked against central-difference to 0.179% and 0.018% relative error respectively. Honest about the limit: jax.grad does NOT flow through the ModelicaFMU block (the FMU is a host C call); the autodiff path therefore runs against an in-process surrogate of the plant, and the tuned gains are then validated on the FMU plant in a second simulation step. The publication/fast-mode pattern caches the heavy Architecture B JIT-compile (~65 s offline) under media/openmodelica_plant_fmu_publication.npz so the reader's notebook runs in ~5 seconds. Three production routes documented: Reference-FMU corpus (Dahlquist / VanDerPol / BouncingBall — exercise 1), OpenModelica .mo → translateModelFMU export (exercise 2), and the always-runnable pythonfmu-built synthetic plant (the path the tutorial takes). Showcases ModelicaFMU import + JaxonomyDiagramSlave / build_fmu + the autodiff-across-the-FMU-boundary wedge.

Universal Differential Equations (UDEs) and symbolic regression (SR)

Demonstrates training a Universal Differential Equation (UDE) to fit the observations produced by the Lotka-Volterra predator-prey system. Subsequently, the UDE is symbolically regressed to learn a closed-form model.

Nonlinear MPC

See thematic series on modeling and control of 3D quadcopter below, which showcases trajectory tracking by nonlinear MPC.

Wind turbine control: MPPT below rated, blade-pitch regulation above rated

A variable-speed ~1.75 MW turbine built from primitive LeafSystem blocks: a standard \(C_p(\lambda,\beta)\) aerodynamic map, a two-mass torsional drivetrain, and a controller that does the two jobs wind-turbine control is made of. Below rated wind the maximum-power-point-tracking torque law \(T_e=K_{opt}\omega_g^2\) sits the rotor at the aerodynamic optimum — the validation gate confirms the steady tip-speed ratio reaches \(\lambda_{opt}=8.10\) and the captured \(C_p\) reaches \(C_{p,max}=0.480\) exactly. Above rated wind the torque saturates and an anti-windup PI pitch loop feathers the blades to hold the generator within 2% of rated speed (1.79 MW). The MPPT gain is derived from the aerodynamics and gearbox, not guessed; a failure-mode cell shows a 2× mistuned gain costs 25% of capture. Notes the slow-rotor buffer_length trap (the adaptive recorder silently truncates to the tail if undersized).

HL-20 lifting body: an angle-of-attack-hold glide autopilot

The NASA HL-20's longitudinal flight dynamics (a 3-DOF model: range, altitude, pitch) with a representative analytical stability-derivative aero set, flown by an AoA-hold autopilot (\(\delta_e=\delta_{e,trim}+K_\alpha(\alpha_{cmd}-\alpha)+K_q q\)). The lesson is the lifting body's low lift-to-drag ratio (\(L/D_{max}\approx5.8\)): the shallowest steady glide it can hold is \(\gamma=-\arctan(1/(L/D_{max}))\approx-10^\circ\). Commanding best-L/D angle of attack, the body settles at the analytic trim angle (\(\gamma\approx-9.1^\circ\), glide ratio 6.25) — validated against \(-\arctan(C_D/C_L)\) averaged over the lightly-damped phugoid. The failure-mode cell makes the point concrete: a naive controller commanding a shallow \(-3^\circ\) glide slope rails the elevator at \(-25^\circ\) and runs the angle of attack to \(66^\circ\), because \(-3^\circ\) is aerodynamically infeasible.

Differentiable vehicle handling: yaw-plane dynamics and design sensitivity

A yaw-plane (single-track) handling model with a Pacejka magic-formula tyre, validated against classical understeer theory — at low lateral acceleration the nonlinear step-steer yaw rate matches \(r/\delta=U/(L+K_{us}U^2)\) to 1.3%. The headline is what a JAX-native simulator gives a chassis engineer: the steady yaw rate's sensitivity to a design knob (rear tyre stiffness) is differentiated through the whole simulation — simulate_jacfwd (forward) and jax.grad (reverse) both return \(\partial r/\partial k_r=-0.0824\), matching central differences to five significant figures, and the gradient is drawn as a tangent on the design sweep. The same gradient is a stability early-warning: softening the rear drives \(K_{us}\) negative into oversteer, where the step-steer yaw rate diverges.

Satellite ADCS: attitude control with an EKF and reaction wheels

A spacecraft attitude-determination-and-control loop: reaction-wheel actuation, an Extended Kalman Filter for attitude/rate estimation from noisy sensors, and momentum management under disturbance torques.

An artificial pancreas: closed-loop glucose control under uncertainty

Model-predictive insulin dosing around a glucose-insulin plant, closing the loop under meal disturbances and parameter uncertainty — a safety-critical biomedical control example.

A grid-forming microgrid: droop control and fault ride-through

Grid-forming inverter droop control on an islanded microgrid, including fault ride-through behaviour — an example from the power-electronics / energy-systems domain.

Reinforcement learning environment from a jaxonomy diagram

Wrap a DiagramBuilder model as a vectorised RL environment, using JAX vmap/jit to batch rollouts for policy training.

Differentiable digital audio: parametric EQ + dynamic-range compressor

A parametric equaliser and a dynamic-range compressor built as jaxonomy blocks, with jax.grad flowing through the DSP chain for gradient-based audio parameter fitting.

DAE constraint projection on a 1-hour pendulum simulation

Long-horizon constrained-DAE integration with constraint projection, showing how the holonomic constraint residual is held near machine precision over an hour of simulated time.

Thematic examples

Battery modeling

Equivalent circuit model (ECM) for a battery
ECM parameter estimation: synthetic data
ECM parameter estimation: experimental data
Data-driven modeling: Dynamic Mode Decomposition (DMD)
Data-driven modeling: Extended DMD
Data-driven modeling: SINDy with control
Data-driven modeling: Neural Networks
Pack-level modeling: cell, module, and thermally-coupled pack — compose BatteryCellECM cells into series modules and a parallel pack via the acausal electrical layer, wrap each cell in a HeatCapacitor with Insulator-based lateral conduction and TemperatureSource ambient boundaries to couple electrical and thermal domains, and run projected gradient descent on a fixed cooling budget (gradient via parameter-rebinding finite differences, since the BDF-DAE adjoint returns the wrong sign on this DAE — filed as a follow-up finding) to redistribute cooling capacity onto the bottlenecked module. Showcases BatteryCellECM + BatteryCellTabular and the cross-domain acausal libraries; the marketing wedge is production-grade pack modelling in open source, no toolbox license required.
Scaling a battery pack from 8 cells to 100,000 — takes the same acausal-ECM pack construction to a 100k-cell pack, documenting what scales cleanly under JAX (vmaped cell dynamics, JIT-compiled kernels) and what breaks (compile time, memory footprint, DAE solve cost) as the cell count grows five orders of magnitude.

Electric drives (PMSM)

A six-part series taking an interior permanent-magnet synchronous motor from a differentiable electromagnetic model to an Arm Cortex-M binary — modelling, control, thermal derating, calibration-from-data, robustness, and embedded deployment.

3D quadcopter modeling and control

Quanser Qube Servo hardware control

Returning rocket booster (Falcon-9-class)

A six-part series modelling the propulsive landing of a returning rocket booster from atmospheric apogee to a soft touchdown on a drone-ship pad. Each part progressively adds fidelity — atmospheric aerodynamics and multi-phase guidance, a multi-engine cluster with bandwidth-limited actuators and variable inertia, noisy sensors with Kalman filtering — and culminates in cinematic MuJoCo rendering. Part 6 fills the GNC validation tooling that production aerospace work expects: linearised Bode / Nyquist / eigenvalue analysis at hover, a 1000-trial Monte Carlo dispersion sweep, the autodiff-vs-finite-difference timing comparison that quantifies what jaxonomy uniquely buys you over MATLAB / Modelica / hand-rolled scipy, plus prose on real-time scheduling, lossless convex powered-descent guidance, and the ITAR realities of production flight software.

Bonus: the cinematic Falcon-9-class landing demo — a 14-second 1280×720 video produced by render_booster.py, which solves Part 1's trajectory optimisation and renders it through a richly-instrumented MuJoCo scene (drone-ship pad, ocean, sun + fill lighting, four landing legs, four grid fins, four-layer plume with visible Mach diamonds, and a hinged engine assembly that visibly tilts on the gimbal command — the visceral cue that there is active control of the descent).

F1 race car (6 parts)

A six-part series taking a 2022+-era ground-effect Formula 1 car from a credible whole-lap simulator (Part 1) to a lap-time-aware aero shape-optimisation architecture that co-simulates an external CFD adjoint solver (Parts 5–6, demonstrated live with an in-process placeholder solver — the full SU2 run needs an external toolchain; see the note under Series C). The property that runs end-to-end: every layer of the stack — chassis, tire, powertrain, driver, aero map, CFD — is differentiable, so jax.grad(lap_time) w.r.t. any setup or aero parameter is one backward pass. That end-to-end differentiability is what sets this apart from conventional lap-time simulators, which fall back on finite differences that scale poorly as the parameter count grows.

The series is structured as three sub-arcs:

Series A — Lap-time-aware setup optimisation (pure jaxonomy). Part 1 builds the bicycle + Pacejka + powertrain + QSS hot-lap driver, validates against analytic cornering equilibrium via findop, and renders the lap in MuJoCo. Part 2 wraps it as lap_time(setup) → float and takes jax.grad through it: 20 Adam steps drop the lap by ~0.4 s from a deliberately-bad baseline.
Series B — CFD budget allocation under the FIA ATR (synthetic aero map). Part 3 fits a noisy 5-D aero map from sparse CFD samples. Part 4 runs Sobol-decomposition over the map to allocate the FIA-ATR-limited CFD budget (the leader gets the least testing-hours; the optimisation question is where to spend them).
Series C — Co-sim with SU2 for lap-time-aware shape optimisation (3D adjoint CFD external). Part 5 closes the loop on a NACA airfoil for the tractable proof. Part 6 swaps it for a parametrised rear-wing assembly on the Perrinn 424 / DrivAer baseline, with PyVista + Blender Cycles for the hero shape-optimisation MP4. Note: Parts 5–6 demonstrate the co-sim architecture live with an in-process solver; the real SU2 RANS + discrete-adjoint results and the Blender hero render require external executables (pysu2, SU2_AD, OpenVSP, ParaView, Blender) not bundled with the repo, so those specific CFD figures are not yet included (each notebook flags this to the reader).
Lap-time simulator: vehicle dynamics, Pacejka tire, powertrain — a 4-state longitudinal–lateral–yaw bicycle with a Pacejka 5.2 magic-formula tire and a friction-ellipse closure; a LookupTable1d engine map + 7-speed gearbox with finite-time shifts; a synthetic 4-corner GP-style track expressed as \(\kappa(s)\); a quasi-steady-state hot-lap driver that respects the friction-ellipse forward and backward; integration-to-steady-state validation against the analytic cornering equilibrium \(V^2 = \mu_{\text{eff}}(V)\,g\,R\) (agreement within 1.8% on a \(R = 100\) m corner; findop filed as failing on this stiff Jacobian + passive-integrators case as a follow-up finding); and a MuJoCo lap render with overlaid HUD. The whole stack is jit- and grad-able, sized for Part 2's setup-optimisation gradient.
Setup optimisation via jax.grad through lap time — wraps Part 1's LTS as lap_time(setup) → float via the cost-as-Integrator pattern under SimulatorOptions(enable_autodiff=True) and with_parameters on the car block. Headline live beat: jax.grad(forward)(SETUP_BASELINE) prints all 8 setup sensitivities in one backward pass at a tractable T_END = 15 s horizon; FD validation on 2 components (k_f, h_f) cross-checks within 5%. The compute-heavy beats (29-iter L-BFGS-B from a deliberately-bad starting setup, 3^6 = 729-point FD-vs-AD grid head-to-head, 16-start LHS multi-start) follow the publication / fast-mode pattern: default MODE = "publication" loads from media/f1_part_2_publication.npz (currently placeholder, awaiting media/f1_part_2_publication_offline.py's offline run at full T_END = 60 fidelity — a lengthy offline run that did not complete on the reference hardware); set MODE = "fast" for a 3-step projected gradient descent verifier in ~30 s. The full-fidelity optimisation figures await a completed offline run (media/f1_part_2_publication_offline.py at T_END = 60, which did not finish on the reference hardware), so the notebook ships the fast-mode verifier above rather than unverified headline numbers. MuJoCo before/after render: two cars (baseline red + optimised yellow) through corner 1 with HUD overlay.
Fitting a 5-D aero map from sparse CFD samples — kicks off Series B. Synthesises a 5-D ground-truth aero map \((h_F, h_R, \phi, \beta, \delta) \to (C_L A, C_D A, x_{\text{CoP}})\) from a closed-form drag-bucket + downforce-vs-rake + asymmetric yaw response, adds heteroscedastic CFD noise (~2% on \(C_L\), ~3% on \(C_D\)), draws 64 Latin-hypercube probes (the realistic CFD budget for one F1 design iteration under the FIA ATR), and fits a multilinear LookupTableND surrogate. Then drops the surrogate into a slim copy of the Part-1 LTS and takes jax.grad(-arc_length) w.r.t. the 5-D aero state at the nominal trim point. Headline live numbers: \(C_L A\) RMS-vs-truth ~ 0.23 m² (bias-dominated by the multilinear-on-quadratic-truth fit, not the 0.06 m² noise floor — honest framing in the prose); dominant-component gradient through the fit matches truth to within ~28% (the multilinear-fit bias dominates; the sign and direction of every component are preserved, which is what matters for setup-search). Pareto sweep over \(N \in \{8, 16, 32, 64, 128, 256\}\) follows the publication / fast-mode pattern: default loads media/f1_part_3_publication.npz, whose fit-error curve is real (measured against a 5000-point validation set) but whose gradient-error curve is a disclosed structural estimate (flagged in-notebook as an estimate) pending a full-fidelity LTS gradient sweep — offline script at media/f1_part_3_publication_offline.py. The fit-error curve flattens to a bias floor past \(N \sim 64\), while the gradient-error curve continues falling — the structural Pareto trade-off Part 4 will exploit. Surfaces 3 follow-up findings: (a) no fit_lookup_table_nd shipped (every author hand-rolls the multilinear LS design matrix), (b) LapTimeAccumulator's integral-of-smoothed-indicator readout gives zero gradient when the lap doesn't finish within T_END (Part 2's live cell shows all-zeros for this reason; Part 3 swaps to a -arc_length proxy), (c) analyze_phase_activity lacks the name= kwarg that analyze_saturation / analyze_control_oscillation accept.
Sobol-driven CFD budget allocation under the FIA ATR — closes Series B. Inherits Part 3's fitted 5-D LookupTableND aero surrogate and Part 1's LTS substrate, then promotes the surrogate from a gradient oracle (Part 2) to a spend-allocation oracle via sobol_indices, decompose_variance_sobol, vmap_qoi, and morris_screening. A live Sobol pass ranks the 5 aero axes (\(h_F\) dominant at \(S_T \approx 0.77\), then \(h_R \approx 0.32\), \(\delta \approx 0.18\), \(\phi \approx 0.11\), with \(\beta \approx 0.06\) smallest), the decompose_variance_sobol call splits Var(lap_time) into aleatoric (~1%, CFD noise floor — irreducible) and epistemic (~99%, fit residual — reducible-with-more-CFD), and a per-cell variance-reduction-per-CFD-hour heatmap on the dominant \((h_F, \beta)\) slice points the next batch at the informative corners. Strategy comparison (the counter-intuitive headline): greedy Sobol-weighted allocation (Strategy B) loses to uniform-LHS (Strategy A) at the 50-probe budget — roughly ~12% vs ~34% variance reduction (~3× worse) — because chasing the single high-\(S_T\) axis starves the fit everywhere else. The Sobol indices stay the right diagnostic of what matters (\(h_F\) dominates); they are the wrong rule for where to sample. Tied back to the FIA ATR sliding scale (positions 1 through 10). Uses the publication / fast-mode pattern (default loads the real media/f1_part_4_publication.npz, generated by media/f1_part_4_publication_offline.py). Closes with Morris screening at \(n_\text{traj} = 16\) (96 evaluations — 10× cheaper than Sobol) that recovers the same headline ranking (the Series-C tractability lever for the SU2-coupled pipeline), a proxy-vs-Sobol sanity check on the §11 axis-sweep variance, five exercises (bootstrap-CI on Sobol, sixth-axis rake, multi-objective \(\alpha = w_V \Delta V + w_G \|\nabla T\|\) acquisition), and the explicit Series B is complete hand-off to Series C. Surfaces 1 follow-up finding on sobol_indices ranking inconsistency under tiny-\(N\) MC noise (small enough at \(N = 4096\) to be benign, but worth a documented n_samples floor).
NACA airfoil SU2 co-sim: jax.custom_vjp wrapping an external CFD solver — demonstrates the co-simulation architecture end-to-end: a jax.custom_vjp + jax.pure_callback wrapper lets jax.grad(lap_time) flow through an external CFD solver on a NACA airfoil. Runs live with an in-process panel-method stand-in solver; the SU2 v8.5.0 RANS + discrete-adjoint numbers need an external toolchain (pysu2, SU2_AD) not bundled here, so those specific figures are not yet included and the notebook flags this to the reader.
Full lap-time-aware aero shape optimisation: DrivAerML + OpenVSP + SU2 + hero MP4 — the same architecture on a parametrised rear-wing assembly. The full pipeline (OpenVSP geometry → SU2_DEF/CFD/adjoint → ParaView/Blender hero render) needs external executables not available here; the headline shape-optimisation numbers are therefore not yet included, pending an offline run of media/f1_part_6_publication_offline.sh (the notebook flags this to the reader).