diff --git a/docs/dev/proposals/2026-07-07-ode-periodic-equilibrium-warmstart.md b/docs/dev/proposals/2026-07-07-ode-periodic-equilibrium-warmstart.md new file mode 100644 index 00000000..1a886dee --- /dev/null +++ b/docs/dev/proposals/2026-07-07-ode-periodic-equilibrium-warmstart.md @@ -0,0 +1,264 @@ +# ODE periodic-equilibrium warm-start — skip the spin-up transient in `nuts`/`mh` on `ode` + +Status: spec — implement against this. Resolves: gh#396. Builds on: +`2026-06-26-ode-nuts-gradient-spine.md` (the augmented forward-sensitivity +integrator this reuses). Sibling: the general forward checkpoint-fork (gh#213) +reuses **Phase A** below but is otherwise independent. High-risk: inference-math +change on the ODE gradient path — read the full `ode.rs` sensitivity core and +`ode_grad.rs` before editing any part, and gate every change on the oracles in +"Correctness". Rust-only: warm-start is a fit-time runtime option, not a model +construct — no IR/OCaml/golden changes. + +## Problem + +A forced (seasonal) compartmental model with a long spin-up to periodic +equilibrium pays the **entire transient cost on every likelihood/gradient +evaluation**. The DMT/Garki `mv_age` family is the exemplar: `origin = 1910` +gives a ~63-year burn-in before the data window (196 transition-driven +compartments, ~8 estimated params). Measured on the current build by the garki +testing agent: one augmented-ODE gradient ≈ ~1 min (~9× the 6.5 s plain solve); +a single `nuts` warm-up iteration does not complete in >7 min ⇒ a full fit = +days–weeks, **zero samples**. The cost is the transient re-integrated per eval; +it is **sampler-independent** (gates `nuts` and the cheaper `mh` on `ode` +equally). Measured relaxation time ≈ **56 yr ≫ 5-yr data window** — the data +cannot "forget" a wrong initial state within the window. + +## Why the cheap options don't work (measured) + +- **Frozen checkpoint** (burn in once, reuse the state — what + `dmt_garki_mv_age_warmstart.camdl` does with a hand-dumped 1961 init): wrong + for a fit. The equilibrium `X*(θ)` moves with the estimated parameters; a + frozen state has `∂X*/∂θ = 0`, biasing the gradient. And because relaxation + (56 yr) ≫ window (5 yr), a short residual burn-in does **not** wash the error + out. Correct only for forward sim or fits of non-dynamical params. +- **Coarse global `dt`** (garki dt-sweep, 2026-07-07): **hard-blocked**. Garki + survey dates have coprime day-offsets, so the score path rejects any `dt > 1` + (`observation at t=… is not a multiple of dt`). Unmeasurable, unusable. +- **Per-phase `dt`** (coarse burn-in only): the burn-in _does_ tolerate coarse + `dt` (~1% state drift at weekly steps, ~5–8× faster), but leaves a ~1% + θ-dependent bias that perturbs the gradient. A legitimate approximate interim, + not an exact default under this project's discipline. + +The exact fix is to **solve for the equilibrium at the current `θ`, every eval, +with its exact sensitivity** — sidestepping both the integration cost and the +obs-alignment wall. + +## Approach: solve the stroboscopic fixed point + +Let the seasonal forcing have fundamental period `P` (typically 1 yr). The limit +cycle sampled at the warm-start time `T_eq` is the fixed point of the period-`P` +flow map anchored at `T_eq`: + +``` +X* = Φ_P(X*; θ), Φ_P(·; θ) = flow of the ODE from T_eq to T_eq+P at θ. +``` + +Newton uses the **monodromy** `M = ∂Φ_P/∂X₀` (state-transition Jacobian over one +period): `(I − M) ΔX = Φ_P(X) − X`. The **equilibrium sensitivity** follows from +the implicit function theorem on `X* − Φ_P(X*; θ) = 0`: +`∂X*/∂θ = (I − M)⁻¹ ∂Φ_P/∂θ`. `M` is needed for the sensitivity regardless, and +once you have it Newton is free — one matrix serves both. The 60-yr transient +collapses to a handful of one-year period-solves plus dense linear solves. + +## Conservation — the load-bearing correction + +For a **closed** model `(I − M)` is **singular**, endemic or not: a conserved +linear quantity `cᵀx` (e.g. per-village total population) makes `cᵀM = cᵀ`, so +`c` is a left eigenvector of `M` with eigenvalue 1. **Garki is conserved, k = +4** (per village): `Nvil[v] = Σ_a N[v,a]` is dynamic, and `birth = δ·Nvil` +exactly balances `death = δ·c` summed → `dN/dt ≡ 0`. This conservation is +_dynamical_ (a rate-balance), not stoichiometric (`1ᵀν ≠ 0`), so it cannot be +read off the stoichiometry — it must be detected **numerically**. + +**Detection is robust, not fragile.** Genuine conservation is _exact_ +(`dN/dt ≡ 0`), so the corresponding singular values of `(I − M)` sit at machine +precision (~1e-12), while the slowest _genuine_ mode of a 56-yr relaxation is at +`1 − |μ| ≈ 1/56 ≈ 0.018` — a 10-order-of-magnitude gap. Any SVD threshold in +`[1e-8, 1e-4]` cleanly separates the `k` conserved directions. (The fragile case +— a multiplier genuinely near 1 — is the _non-endemic_ boundary, which the +endemic gate refuses anyway.) + +**Solve on the conservation manifold via a bordered system.** With `C` +(`n_int × k`) the conserved directions (left null space of `(I − M)`, from the +SVD) and `b = CᵀX_init` the totals pinned by the model's initial condition, +Newton solves the nonsingular KKT/bordered system each iteration: + +``` +[ I − M C ] [ΔX] = [ Φ_P(X) − X ] +[ Cᵀ 0 ] [λ ] [ b − CᵀX ] +``` + +which enforces `CᵀX = b` exactly and is nonsingular whenever the _transverse_ +cycle is hyperbolic. The sensitivity reuses the **same LU factorization**: + +``` +[ I − M C ] [∂X*/∂θ] = [ ∂Φ_P/∂θ ] +[ Cᵀ 0 ] [ μ ] [ ∂b/∂θ ] +``` + +where `∂b/∂θ = Cᵀ · (ic_grad seed)` — so the conserved-direction sensitivity +**does** reuse `ic_grad`, while the transverse part comes from the equilibrium. +`∂Φ_P/∂θ` is the _partial_ (see the seed trap below). + +**Endemic gate.** Refuse if the transverse system is ill-conditioned (reciprocal +condition of the bordered matrix below `RCOND_MIN`, i.e. a non-conservation +multiplier → 1). Under NUTS this surfaces as a `-inf` rejection, not an abort +(`det_grad` `Err` → `(-inf, 0)`, `ode_nuts.rs:147-152`), so a non-endemic +proposal mid-run is contained; only a non-endemic _initial_ point aborts (with a +clear message). Report a rejected-for-non-endemic count so the boundary isn't a +silent posterior wall. + +## Design in types + +### Phase A — ODE start-from-state seam (also the gh#213 ODE arm) + +One new argument on `integrate_obs_sensitivity` (`ode.rs:885`, `pub(crate)`); +every existing caller passes `None` and is byte-identical to today. + +```rust +pub(crate) fn integrate_obs_sensitivity( + model: &CompiledModel, params: &[f64], param_model_idx: &[usize], + x0: Option<&[f64]>, // NEW: overrides initial_state_continuous (ode.rs:925) + state_sens_0: &[f64], cfg: &OdeConfig, obs_times: &[f64], + seed: SensSeed, // NEW: see B1 +) -> Result, SimError>; +``` + +The forward value path (`OdeSim::run` / `run_ode`) gets the analogous +`x0: Option<&[f64]>` injection for the `mh` wiring. Oracle: the splice +invariant, bit-exact because fixed-step RK4 is deterministic (no RNG) — +`solve(t0→t_end) ≡ solve(t0→T*) ++ solve_from(x(T*), S(T*), T*→t_end)`. + +### Phase B — the equilibrium solver + +**B1. Sensitivity seed mode.** The `(param_model_idx, state_sens_0)` pair is +replaced by an explicit seed, decoupling the sensitivity column count from the +parameter set and making the `∂rate/∂θ` forcing switchable. It is a real branch, +not a `param_model_idx` trick (a compartment index colliding with a param's +model index would inject a spurious `∂rate/∂θ`): + +```rust +enum SensSeed<'a> { + Params { model_idx: &'a [usize], s0: &'a [f64] }, // d = params, forcing ON (today) + Zero { model_idx: &'a [usize] }, // d = params, S(0)=0, forcing ON (∂Φ/∂θ partial) + Monodromy, // d = n_int, S(0)=I, forcing OFF (M) +} +``` + +Threaded through `sensitivity_derivs` (`ode.rs:160`) and `rk4_step`; the +Monodromy path skips the wasted flow-sensitivity block. Everything sizes off the +seed's column count, so an identity `n_int`-column seed allocates uniformly. + +**B2. The solver** (`solve_equilibrium`, new module `ode_equilibrium.rs`): + +1. `integrate_period(X_k)` in `SensSeed::Monodromy` over `[T_eq, T_eq+P]` → + `(Φ, M)`. (Value clamps `max(0.0)`; assert the clamp is inactive on the cycle + — a clamped iterate corrupts `M`.) +2. On the first `M`: SVD `(I − M)` → `C` = left-null directions (σ < + `CONS_EPS`). +3. Bordered Newton to `‖Φ_P(X) − X‖∞ < TOL_EQ`; cap iterations; non-convergence + → endemic-gate error. +4. Sensitivity: one `SensSeed::Zero` period solve → `∂Φ_P/∂θ`; reuse the + bordered LU with RHS `(∂Φ_P/∂θ, Cᵀ·ic_grad)` → `∂X*/∂θ`. + +Returns +`Equilibrium { x_star: Vec /* n_int */, x_star_sens: Vec /* n_int×d */ }`. + +**B3. Wiring.** `det_grad` (nuts, `ode_grad.rs:62-70`): when warm-start is on, +`solve_equilibrium(…, T_eq, P, dt)` → set `cfg.t_start = T_eq`, inject +`x0 = X*`, `state_sens_0 = ∂X*/∂θ` (via `SensSeed::Params`). +`compute_ode_loglik` (mh/MLE, `ode_loglik.rs:36`): solve for `X*` only (state +path), inject `x0 = X*` at `T_eq` into `OdeSim::run`. + +### Validity gate (pre-obs periodicity) + +Runs at fit setup once `T_eq` is resolved. Refuse (hard error, naming the +feature) if the warm-start assumptions don't hold: + +- any `event`/`intervention` fires in `[origin, T_eq]` (non-periodic effect in + the transient window); +- the forcing is not `P`-periodic over the burn-in neighborhood — a + **numerical** check `‖C(t) − C(t+P)‖ < FORCING_PERIOD_EPS` sampled across + `[T_eq − P, T_eq]` (handles `interpolated`/table forcings, which carry no + declared period — the garki case); +- `P` is not supplied, is non-positive, or the model has no forcing. + +Refusal points at the alternatives (checkpoint for fixed-transient-param fits; +pay the transient otherwise). + +## Surface (CLI + TOML, via the run input system) + +Warm-start is a fit-time option, expressed identically as a CLI flag and a +`fit.toml` field, and folded into the fit config-identity hash (it changes the +scored trajectory → it re-keys the run). + +``` +camdl fit model.camdl --data d.tsv --method nuts --backend ode \ + --warm-start equilibrium --warm-start-period 365 [--warm-start-at DATE] +``` + +```toml +[fit] +warm_start = "equilibrium" # default "none" +warm_start_period = 365 # days (model time_unit); required for equilibrium +warm_start_at = "1970-01-01" # optional; default = earliest first_obs across streams +``` + +`warm_start_at` (`T_eq`) is a **single** anchor `≤ min(first_obs)` — decoupled +from the per-stream `cond_from` (which stays per-stream and runs `[T_eq, t_end]` +untouched). Enum `WarmStart = None | Equilibrium`; a future `Checkpoint(path)` +variant is where gh#213 attaches. A `--dump-gradient` eval flag (fixed-θ, emits +the autodiff gradient) supports external gradient verification (gh#78). + +## Correctness / tests (red-first where a bug would hide) + +1. **Phase-A splice invariant** (bit-exact): inject a continuous solve's exact + `x(T*)`, `S(T*)`; assert the resumed tail equals the continuous tail + bit-for-bit, state and sensitivity. +2. **Monodromy vs finite difference**: `M` (Monodromy-mode) matches a central FD + of the one-period map `∂Φ_P/∂X₀` column by column. +3. **Equilibrium == long burn-in** (headline): `X*` from the solver matches + `X(T_eq)` from integrating the full transient at the same `θ` (grids + matched). +4. **Sensitivity three ways**: `∂X*/∂θ = (I−M)⁻¹∂Φ_P/∂θ` matches (a) a central + FD of the equilibrium re-solve and (b) the augmented `S(T_eq)` from the long + burn-in. Negative control: a param not entering the transient → zero column. + (Only the burn-in check distinguishes the partial `∂Φ_P/∂θ` seed from a total + one — the seed trap.) +5. **Conservation**: on a closed seasonal model (`k ≥ 1`), the bordered solve + converges, preserves `CᵀX = b` to machine precision, and the endemic gate + does **not** fire on a genuinely endemic model. The CI reference model is + closed-with-waning so this path is exercised. +6. **Endemic gate fires** on a parameterization driven non-hyperbolic + (multiplier → 1), with the endemicity message, not a NaN. +7. **End-to-end**: a `nuts`-on-`ode` fit with `--warm-start equilibrium` + recovers the same posterior as a short-burn-in reference, in a fraction of + wall-clock. + +## Scope (v1) + +- Transition-driven (integer) compartments only (real-compartment forward + sensitivity is refused at `gradient_capability.rs:298`; garki is in scope). +- rk4 fixed-step; no `dt`-in-rate (existing `ode_grad` guards). +- Single constant fundamental `P`; multi-period/LCM and `θ`-dependent `P` + refused, are follow-ups. +- Dense bordered solve (`n_int + k`), fine to a few hundred compartments (garki + 196); national-scale block-sparse is a follow-up. + +## Resolved decisions + +- **`nalgebra`** (pure-Rust, no system BLAS) for the SVD (conservation + detection) and the bordered LU + reciprocal-condition estimate. A hand-rolled + solve is rejected: the matrix is deliberately near-singular for slow cycles. +- **`T_eq`** is a single declared anchor, not the per-stream `cond_from`. +- **`P`** is user-declared and numerically verified (interpolated forcings carry + no period). +- Flow/incidence accumulators reset at `T_eq`; the conserved-direction + sensitivity reuses `ic_grad`. + +## Follow-ups (named, not blocking) + +Real-compartment forward sensitivity → warm-start for explicit `dx/dt` models; +multi-period/LCM `P`; block-sparse bordered solve for national scale; Newton +warm-start across evals (seed `X₀` from the previous gradient's `X*`; deferred — +couples evals, land the stateless version + oracles first). diff --git a/docs/fit-toml.md b/docs/fit-toml.md index 864065a7..0172799d 100644 --- a/docs/fit-toml.md +++ b/docs/fit-toml.md @@ -193,15 +193,15 @@ purpose" without the warning. ## Stage algorithms -| `algorithm` | backend | role | key fields | -| ------------------------ | ---------------- | ------------------------------------------------- | ------------------------------------------------------------------------------------- | -| `if2` | `chain_binomial` | iterated filtering → MLE | `chains`, `particles`, `iterations`, `cooling` (+ `cooling_target_iters`, default 50) | -| `pgas` | `chain_binomial` | particle Gibbs + NUTS → posterior | `chains`, `particles`, `sweeps` (+ `burn_in`, `thin`, `tempering`, `max_tree_depth`) | -| `pmmh` | `chain_binomial` | particle marginal MH → posterior | `chains`, `particles`, `iterations` | -| `pfilter` | `chain_binomial` | particle filter at fixed θ → log-likelihood + ESS | `particles`, `replicates` | -| `nl-sbplx` / `nl-bobyqa` | `ode` | NLopt deterministic optimizer → MLE | `chains` (LHS starts) (+ `max_evals`, `tolerance`) | -| `mh` | `ode` | MH on the deterministic ODE marginal → posterior | `chains`, `iterations` (+ `burn_in`, `thin`, `adapt`, `adapt_start`) | -| `nuts` | `ode` | gradient NUTS (forward sensitivities) → posterior | `chains`, `warmup`, `samples` (+ `max_tree_depth`, `target_accept`, `dense_mass`) | +| `algorithm` | backend | role | key fields | +| ------------------------ | ---------------- | ------------------------------------------------- | ----------------------------------------------------------------------------------------------- | +| `if2` | `chain_binomial` | iterated filtering → MLE | `chains`, `particles`, `iterations`, `cooling` (+ `cooling_target_iters`, default 50) | +| `pgas` | `chain_binomial` | particle Gibbs + NUTS → posterior | `chains`, `particles`, `sweeps` (+ `burn_in`, `thin`, `tempering`, `max_tree_depth`) | +| `pmmh` | `chain_binomial` | particle marginal MH → posterior | `chains`, `particles`, `iterations` | +| `pfilter` | `chain_binomial` | particle filter at fixed θ → log-likelihood + ESS | `particles`, `replicates` | +| `nl-sbplx` / `nl-bobyqa` | `ode` | NLopt deterministic optimizer → MLE | `chains` (LHS starts) (+ `max_evals`, `tolerance`) | +| `mh` | `ode` | MH on the deterministic ODE marginal → posterior | `chains`, `iterations` (+ `burn_in`, `thin`, `adapt`, `adapt_start`) | +| `nuts` | `ode` | gradient NUTS (forward sensitivities) → posterior | `chains`, `warmup`, `samples` (+ `max_tree_depth`, `target_accept`, `dense_mass`, `warm_start`) | The `ode`-backend Bayesian samplers (`mh`, `nuts`) fit the **deterministic marginal likelihood** `p(y | θ, ODE skeleton)` rather than the stochastic @@ -212,6 +212,32 @@ scheduled effect, or an `ivp` parameter); `mh` is gradient-free and carries no such requirement. See `camdl docs inference` (the ODE-backend fitting section) for when to pick which. +### Periodic-equilibrium warm-start (`nuts` on `ode`) + +A seasonally-forced model with a long spin-up to its endemic cycle (e.g. an +`origin` decades before the data) pays the entire transient on _every_ gradient +evaluation, which can make a fit take days. When the pre-data dynamics settle to +a **periodic** equilibrium, `warm_start = "equilibrium"` skips the transient: it +solves the seasonal limit cycle directly at the data-window start and seeds the +integration there, with the exact parameter sensitivity the sampler needs. + +```toml +[stages.posterior] +algorithm = "nuts" +backend = "ode" +warm_start = "equilibrium" # default "none" +warm_start_period = 365.25 # forcing fundamental period P (required); model time units +warm_start_at = 1970.0 # T_eq anchor (absolute model time); default = earliest observation +``` + +`warm_start_period` is required (a table/`interpolated` forcing carries no +declared period). The fit refuses up front if the forcing is not `P`-periodic +over the burn-in window, or if the model has no isolated stable seasonal cycle +at the current parameters (non-endemic). Only the transient-shaping dynamics +must be periodic; interventions and non-periodic forcing _after_ `T_eq` are +unaffected. Closed (population-conserving) models are handled correctly. +`warm_start` is identity-defining — changing it re-keys the run. + ```toml # A gradient-based Bayesian fit on the ODE skeleton. [stages.posterior] diff --git a/docs/inference.md b/docs/inference.md index bce2c21f..a7464920 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -1252,6 +1252,14 @@ the beta caveat for every `(algorithm, backend)` pair. The two backends compute different objects (`chain_binomial → p(y|θ)`; `ode → p(y|θ, ODE skeleton)`), so read that guidance before choosing which likelihood a fit should target. +**Long seasonal spin-up?** If a `nuts`-on-`ode` model has a decades-long burn-in +to its endemic cycle (an `origin` far before the data), each gradient +re-integrates that whole transient — a fit can take days. When the pre-data +dynamics are periodic, `warm_start = "equilibrium"` on the `nuts` stage solves +the seasonal limit cycle directly and starts the integration there, cutting the +per-gradient cost dramatically. See the periodic-equilibrium warm-start section +of `camdl docs fit-toml` for the `warm_start_period` / `warm_start_at` options. + --- ## Diagnostic interpretation guide diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 98f6cf8f..0d02600d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -433,6 +433,12 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -1516,6 +1522,16 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" @@ -1559,6 +1575,33 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.33.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d43ddcacf343185dfd6de2ee786d9e8b1c2301622afab66b6c73baf9882abfd" +dependencies = [ + "approx", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nlopt" version = "0.8.1" @@ -1955,6 +1998,12 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rayon" version = "1.11.0" @@ -2131,6 +2180,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safe_arch" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -2262,6 +2320,7 @@ dependencies = [ "ir", "libm", "log", + "nalgebra", "nlopt", "numerics", "parquet", @@ -2276,6 +2335,19 @@ dependencies = [ "thiserror", ] +[[package]] +name = "simba" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c99284beb21666094ba2b75bbceda012e610f5479dfcc2d6e2426f53197ffd95" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "paste", + "wide", +] + [[package]] name = "simd-adler32" version = "0.3.9" @@ -2826,6 +2898,16 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "wide" +version = "0.7.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/rust/crates/cli/src/fit/config_v2.rs b/rust/crates/cli/src/fit/config_v2.rs index 746a66c3..dbfede22 100644 --- a/rust/crates/cli/src/fit/config_v2.rs +++ b/rust/crates/cli/src/fit/config_v2.rs @@ -1291,6 +1291,22 @@ pub enum Stage { /// tracked follow-up, so this currently only records intent.) #[serde(default = "default_dense_mass")] dense_mass: bool, + /// gh#396 periodic-equilibrium warm-start (ODE backend). `"none"` (default) + /// integrates from the model origin; `"equilibrium"` seeds each gradient + /// from the seasonal limit cycle solved at `warm_start_at`, skipping the + /// spin-up transient. + #[serde(default)] + warm_start: WarmStartKind, + /// The forcing fundamental period `P` (model time units) — required for + /// `warm_start = "equilibrium"` (an `interpolated` forcing carries no + /// declared period). + #[serde(default)] + warm_start_period: Option, + /// The warm-start anchor time `T_eq` (absolute model time). Default: the + /// earliest observation time across streams (equilibrate at the data-window + /// start). + #[serde(default)] + warm_start_at: Option, }, /// NLopt Sbplx (subspace-searching simplex) — deterministic MLE on the @@ -1591,6 +1607,7 @@ impl Stage { backend, chains, warmup, starts_from, init_method, survey_path, survey_top_k_n, max_tree_depth, target_accept, dense_mass, + warm_start, warm_start_period, warm_start_at, .. } => json!({ "algorithm": "nuts", @@ -1604,6 +1621,10 @@ impl Stage { "max_tree_depth": max_tree_depth, "target_accept": target_accept, "dense_mass": dense_mass, + // gh#396: warm-start changes the scored trajectory ⇒ identity-defining. + "warm_start": warm_start, + "warm_start_period": warm_start_period, + "warm_start_at": warm_start_at, }), // No extension dimension: hash the full stage. NLopt stages // also have no extension dimension — every knob (chains, @@ -1681,6 +1702,17 @@ fn default_loglik_eval_replicates() -> usize { 8 } /// at iteration 50, then continues at the noise floor. fn default_cooling_target_iters() -> usize { 50 } +/// gh#396 warm-start selector for the ODE gradient stages (`nuts`, and later +/// `mh`). `Equilibrium` solves the seasonal limit cycle at `warm_start_at` each +/// eval and seeds the integration from it, skipping the spin-up transient. +#[derive(Clone, Copy, Debug, Default, PartialEq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarmStartKind { + #[default] + None, + Equilibrium, +} + // PGAS defaults fn default_pgas_tempering() -> Vec { vec![1.0] } fn default_max_tree_depth() -> usize { 10 } diff --git a/rust/crates/cli/src/fit/nuts.rs b/rust/crates/cli/src/fit/nuts.rs index d97f5c74..1deb44fb 100644 --- a/rust/crates/cli/src/fit/nuts.rs +++ b/rust/crates/cli/src/fit/nuts.rs @@ -31,6 +31,9 @@ pub struct NutsStageOpts { pub init_method: super::init::InitMethod, pub survey_path: Option, pub survey_top_k_n: Option, + pub warm_start: super::config_v2::WarmStartKind, + pub warm_start_period: Option, + pub warm_start_at: Option, } impl NutsStageOpts { @@ -38,7 +41,8 @@ impl NutsStageOpts { match stage { super::config_v2::Stage::Nuts { chains, warmup, samples, max_tree_depth, target_accept, dense_mass, - init_method, survey_path, survey_top_k_n, .. + init_method, survey_path, survey_top_k_n, + warm_start, warm_start_period, warm_start_at, .. } => Ok(NutsStageOpts { n_chains: *chains, warmup: *warmup, @@ -49,6 +53,9 @@ impl NutsStageOpts { init_method: init_method.clone(), survey_path: survey_path.clone(), survey_top_k_n: *survey_top_k_n, + warm_start: *warm_start, + warm_start_period: *warm_start_period, + warm_start_at: *warm_start_at, }), other => Err(format!( "NutsStageOpts::from_stage: expected Stage::Nuts, got {}", @@ -121,6 +128,38 @@ pub fn run_stage( let param_names: Vec = config.estimated_params.iter().map(|s| s.name.clone()).collect(); + // gh#396: resolve the periodic-equilibrium warm-start once (Copy, so each + // parallel chain reuses it). T_eq defaults to the earliest observation. + let warm_start: Option = match opts.warm_start { + super::config_v2::WarmStartKind::None => None, + super::config_v2::WarmStartKind::Equilibrium => { + let period = opts.warm_start_period.ok_or_else(|| { + "warm_start = \"equilibrium\" requires `warm_start_period` (the forcing \ + fundamental period P, in model time units — e.g. 365.25 for an annual \ + cycle)" + .to_string() + })?; + if !(period > 0.0) { + return Err(format!("warm_start_period must be positive, got {period}")); + } + let t_eq = match opts.warm_start_at { + Some(t) => t, + None => obs_times.iter().cloned().fold(f64::INFINITY, f64::min), + }; + if !t_eq.is_finite() { + return Err("warm-start: no observations to anchor T_eq; set `warm_start_at`" + .to_string()); + } + Some(sim::ode_equilibrium::WarmStart { t_eq, period }) + } + }; + // gh#396: refuse a non-P-periodic forcing over the burn-in window up front — + // a silently-wrong equilibrium would otherwise bias every gradient. + if let Some(ws) = &warm_start { + sim::ode_equilibrium::check_periodicity(&compiled, &config.base_params, ws.t_eq, ws.period, dt) + .map_err(|e| e.to_string())?; + } + // Chains are independent (own seed, own RNG, own trace file) and their outputs // reduce order-independently (max-loglik chain + summed divergences), so run // them in parallel across the rayon pool — the same pattern PGAS/PMMH/IF2 use. @@ -149,6 +188,7 @@ pub fn run_stage( dt, // Independent chains: same start, distinct RNG stream. seed: seed.wrapping_add(chain_id as u64), + warm_start, }; let result = run_ode_nuts( &compiled, &obs_model, &obs_times, &config.base_params, diff --git a/rust/crates/sim/Cargo.toml b/rust/crates/sim/Cargo.toml index ffaad9fc..9a32d563 100644 --- a/rust/crates/sim/Cargo.toml +++ b/rust/crates/sim/Cargo.toml @@ -19,6 +19,11 @@ numerics = { path = "../numerics" } # probabilities for AFP-style rare-event surveillance). libm = "0.2" log = "0.4" +# gh#396: dense linear algebra for the periodic-equilibrium warm-start solver — +# SVD (numerical conservation-law detection) and the bordered LU + reciprocal +# condition estimate (Newton on the seasonal Poincaré fixed point). Pure Rust, +# no system BLAS/LAPACK. Used only by the ODE gradient path. +nalgebra = "0.33" nlopt = { version = "0.8", optional = true } parquet = { version = "53", optional = true, default-features = false, features = ["arrow"] } rand = "0.8" diff --git a/rust/crates/sim/src/inference/ode_grad.rs b/rust/crates/sim/src/inference/ode_grad.rs index d4f95894..cadbcc67 100644 --- a/rust/crates/sim/src/inference/ode_grad.rs +++ b/rust/crates/sim/src/inference/ode_grad.rs @@ -41,6 +41,11 @@ pub fn det_grad( dt: f64, params: &[f64], estimated_to_model: &[usize], + // gh#396: when `Some`, replace the `[origin, t_end]` transient with the + // periodic equilibrium solved at `T_eq` — the equilibrium state `X*(θ)` and its + // exact sensitivity `∂X*/∂θ` seed the `[T_eq, t_end]` integration, so the long + // spin-up is never integrated. `None` → today's path (from the model `t_start`). + warm_start: Option<&crate::ode_equilibrium::WarmStart>, ) -> Result<(f64, Vec), SimError> { // §1h capability gate: the ONE place that answers "can this model be fit by the // ODE gradient?" — refusing (with the compiler's own reason where one exists) @@ -59,10 +64,33 @@ pub fn det_grad( // Forward-sensitivity seed S(t_start) = ∂(initial_state)/∂θ (`ic_grad`): zero // for an explicit (constant) initial condition, nonzero for a parameterized // one whose expression involves an estimated parameter (gh#275 §1c C-seed). - let state_sens_0 = compiled.ic_grad_seed(params, estimated_to_model)?; + let ic_grad = compiled.ic_grad_seed(params, estimated_to_model)?; + + // gh#396: warm-start replaces the transient with the equilibrium at T_eq; + // otherwise integrate from the model's t_start with the ic_grad seed. + let (x0, state_sens_0, t_start) = match warm_start { + Some(ws) => { + let eq = crate::ode_equilibrium::solve_equilibrium( + compiled, + params, + estimated_to_model, + &ic_grad, + ws.t_eq, + ws.period, + dt, + )?; + log::debug!( + "ode warm-start: equilibrium in {} Newton iters, {} conserved direction(s)", + eq.iters, + eq.n_conserved + ); + (Some(eq.x_star), eq.x_star_sens, ws.t_eq) + } + None => (None, ic_grad, compiled.model.simulation.t_start), + }; let cfg = OdeConfig { - t_start: compiled.model.simulation.t_start, + t_start, t_end: compiled.model.simulation.t_end, dt, }; @@ -71,6 +99,7 @@ pub fn det_grad( compiled, params, estimated_to_model, + x0.as_deref(), &state_sens_0, &cfg, obs_times, @@ -236,7 +265,7 @@ mod tests { let (int_s0, _) = compiled.initial_state(¶ms).unwrap(); let seed = vec![0.0; int_s0.counts.len() * est.len()]; let recs = - crate::ode::integrate_obs_sensitivity(&compiled, ¶ms, &est, &seed, &cfg, &obs_times) + crate::ode::integrate_obs_sensitivity(&compiled, ¶ms, &est, None, &seed, &cfg, &obs_times) .unwrap(); let projections: Vec = compiled .model @@ -267,7 +296,7 @@ mod tests { let obs_model = build_obs_model(compiled.clone(), &obs_times, per_stream); - let (ll, grad) = det_grad(&compiled, &obs_model, &obs_times, dt, ¶ms, &est).unwrap(); + let (ll, grad) = det_grad(&compiled, &obs_model, &obs_times, dt, ¶ms, &est, None).unwrap(); assert!(ll.is_finite(), "det_grad loglik must be finite, got {ll}"); assert_eq!(grad.len(), est.len()); @@ -278,8 +307,8 @@ mod tests { let mut pm = params.clone(); pp[midx] += eps; pm[midx] -= eps; - let (llp, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pp, &est).unwrap(); - let (llm, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pm, &est).unwrap(); + let (llp, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pp, &est, None).unwrap(); + let (llm, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pm, &est, None).unwrap(); let fd = (llp - llm) / (2.0 * eps); let rel = if fd.abs() > 1e-8 { (grad[i] - fd).abs() / fd.abs() @@ -367,7 +396,7 @@ mod tests { let cfg = OdeConfig { t_start: compiled.model.simulation.t_start, t_end: 60.0, dt }; let seed = vec![0.0; compiled.initial_state(¶ms).unwrap().0.counts.len() * est.len()]; let recs = - crate::ode::integrate_obs_sensitivity(&compiled, ¶ms, &est, &seed, &cfg, &obs_times) + crate::ode::integrate_obs_sensitivity(&compiled, ¶ms, &est, None, &seed, &cfg, &obs_times) .unwrap(); let flow_idx = 0usize; // weekly_cases is the incidence stream (FlowSum) let weekly: Vec = recs.iter().map(|r| r.inc[flow_idx].round()).collect(); @@ -377,7 +406,7 @@ mod tests { assert!(detection.iter().sum::() > 1.0, "√I+√R detection data must be nonzero"); let obs_model = build_obs_model(compiled.clone(), &obs_times, vec![weekly, detection]); - let (ll, grad) = det_grad(&compiled, &obs_model, &obs_times, dt, ¶ms, &est).unwrap(); + let (ll, grad) = det_grad(&compiled, &obs_model, &obs_times, dt, ¶ms, &est, None).unwrap(); assert!(ll.is_finite(), "det_grad loglik must be finite, got {ll}"); let eps = 1e-6; @@ -387,8 +416,8 @@ mod tests { let mut pm = params.clone(); pp[midx] += eps; pm[midx] -= eps; - let (llp, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pp, &est).unwrap(); - let (llm, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pm, &est).unwrap(); + let (llp, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pp, &est, None).unwrap(); + let (llm, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pm, &est, None).unwrap(); let fd = (llp - llm) / (2.0 * eps); let rel = if fd.abs() > 1e-8 { (grad[i] - fd).abs() / fd.abs() @@ -457,7 +486,7 @@ mod tests { let cfg = OdeConfig { t_start: compiled.model.simulation.t_start, t_end: 60.0, dt }; let seed0 = compiled.ic_grad_seed(¶ms, &est).unwrap(); let recs = - crate::ode::integrate_obs_sensitivity(&compiled, ¶ms, &est, &seed0, &cfg, &obs_times) + crate::ode::integrate_obs_sensitivity(&compiled, ¶ms, &est, None, &seed0, &cfg, &obs_times) .unwrap(); let projections: Vec = compiled .model @@ -485,7 +514,7 @@ mod tests { assert!(per_stream[0].iter().sum::() > 1.0, "incidence data must be nonzero"); let obs_model = build_obs_model(compiled.clone(), &obs_times, per_stream); - let (ll, grad) = det_grad(&compiled, &obs_model, &obs_times, dt, ¶ms, &est).unwrap(); + let (ll, grad) = det_grad(&compiled, &obs_model, &obs_times, dt, ¶ms, &est, None).unwrap(); assert!(ll.is_finite(), "det_grad loglik must be finite, got {ll}"); let eps = 1e-4; // I0 ~ 10, so a larger step keeps the FD well-conditioned. @@ -495,8 +524,8 @@ mod tests { let mut pm = params.clone(); pp[midx] += eps; pm[midx] -= eps; - let (llp, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pp, &est).unwrap(); - let (llm, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pm, &est).unwrap(); + let (llp, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pp, &est, None).unwrap(); + let (llm, _) = det_grad(&compiled, &obs_model, &obs_times, dt, &pm, &est, None).unwrap(); let fd = (llp - llm) / (2.0 * eps); let rel = if fd.abs() > 1e-8 { (grad[i] - fd).abs() / fd.abs() diff --git a/rust/crates/sim/src/inference/ode_nuts.rs b/rust/crates/sim/src/inference/ode_nuts.rs index 132ffe49..10e15603 100644 --- a/rust/crates/sim/src/inference/ode_nuts.rs +++ b/rust/crates/sim/src/inference/ode_nuts.rs @@ -45,6 +45,9 @@ pub struct OdeNutsConfig { pub dt: f64, /// RNG seed. pub seed: u64, + /// gh#396: periodic-equilibrium warm-start. `Some` skips the spin-up transient + /// by seeding each gradient from the equilibrium solved at `T_eq`. + pub warm_start: Option, } impl Default for OdeNutsConfig { @@ -57,6 +60,7 @@ impl Default for OdeNutsConfig { init_step_size: 0.1, dt: 1.0, seed: 0, + warm_start: None, } } } @@ -117,8 +121,15 @@ pub fn run_ode_nuts( for (i, ep) in estimated.iter().enumerate() { params[ep.index] = ep.from_transformed(z[i]); } - let (ll, ll_grad_theta) = - det_grad(compiled, obs_model, obs_times, config.dt, ¶ms, &estimated_to_model)?; + let (ll, ll_grad_theta) = det_grad( + compiled, + obs_model, + obs_times, + config.dt, + ¶ms, + &estimated_to_model, + config.warm_start.as_ref(), + )?; let mut log_p = ll; let mut grad_z = vec![0.0; d]; @@ -195,9 +206,17 @@ pub fn run_ode_nuts( for (i, ep) in estimated.iter().enumerate() { params[ep.index] = ep.from_transformed(z[i]); } - det_grad(compiled, obs_model, obs_times, config.dt, ¶ms, &estimated_to_model) - .map(|(ll, _)| ll) - .unwrap_or(f64::NEG_INFINITY) + det_grad( + compiled, + obs_model, + obs_times, + config.dt, + ¶ms, + &estimated_to_model, + config.warm_start.as_ref(), + ) + .map(|(ll, _)| ll) + .unwrap_or(f64::NEG_INFINITY) }; // Sampling. @@ -336,7 +355,7 @@ mod tests { let (int_s0, _) = cm.initial_state(&true_params).unwrap(); let seed = vec![0.0; int_s0.counts.len()]; let recs = crate::ode::integrate_obs_sensitivity( - &cm, &true_params, &[beta_idx], &seed, &crate::config::OdeConfig { + &cm, &true_params, &[beta_idx], None, &seed, &crate::config::OdeConfig { t_start: 0.0, t_end: 60.0, dt, }, &obs_times, ) @@ -384,6 +403,7 @@ mod tests { init_step_size: 0.2, dt, seed: 20260707, + warm_start: None, }; let result = run_ode_nuts(&cm, &obs_model, &obs_times, &base, &estimated, &priors, &config).unwrap(); diff --git a/rust/crates/sim/src/lib.rs b/rust/crates/sim/src/lib.rs index cdaaf130..c0b96733 100644 --- a/rust/crates/sim/src/lib.rs +++ b/rust/crates/sim/src/lib.rs @@ -17,6 +17,7 @@ pub mod ode_integrator; pub mod gillespie; pub mod chain_binomial; pub mod ode; +pub mod ode_equilibrium; pub mod intervention; pub mod lifecycle; pub mod effects; diff --git a/rust/crates/sim/src/ode.rs b/rust/crates/sim/src/ode.rs index cd292d9c..ba6e2db9 100644 --- a/rust/crates/sim/src/ode.rs +++ b/rust/crates/sim/src/ode.rs @@ -136,6 +136,12 @@ fn ode_derivs( struct Sens { state: Vec, // n_int × d, row-major flow: Vec, // n_transitions × d, row-major + /// Whether the `∂rate/∂θ` forcing term is active (gh#396). `true` for the + /// parameter-sensitivity paths (`S = ∂x/∂θ`, seeded by `ic_grad` or zero); + /// `false` for the **monodromy** path (`S = ∂x/∂x₀`, identity-seeded), where + /// the columns are compartments, not parameters, so the param-keyed forcing + /// summand must be dropped rather than mis-indexed into `param_model_idx`. + forcing: bool, } /// Forward-sensitivity derivatives, assembled **sparse and transitionwise** @@ -163,6 +169,10 @@ fn sensitivity_derivs( real_vals: &[f64], params: &[f64], param_model_idx: &[usize], + // gh#396: when `false`, the param-keyed `∂rate/∂θ` forcing summand is dropped + // so this integrates the pure state-Jacobian chain `Ṡ = J_x·S` — the monodromy + // variational equation. `true` restores the full `Ṡ = J_x·S + J_θ`. + forcing: bool, t: f64, dt: f64, per_eval: Option<&[f64]>, @@ -187,8 +197,14 @@ fn sensitivity_derivs( let rate_grad = &model.resolved.rate_grads_indexed[tr_idx]; // param-keyed ∂rate/∂θ let rate_state_grad = &model.resolved.rate_state_grads_indexed[tr_idx].0; // comp-keyed ∂rate/∂x for p in 0..d { - // total_dr_dθ[p] = ∂rate/∂θ_p + Σ_j ∂rate/∂x_j · S[j, p] (chain through state) - let mut total = eval_emitted_grad(rate_grad, param_model_idx[p], &ctx); + // total_dr_dθ[p] = ∂rate/∂θ_p + Σ_j ∂rate/∂x_j · S[j, p] (chain through state). + // Monodromy mode (`forcing = false`): drop the ∂rate/∂θ_p term → pure + // `Ṡ = J_x·S`, so S(T) = ∂x(T)/∂x₀ when seeded S(0) = I (gh#396). + let mut total = if forcing { + eval_emitted_grad(rate_grad, param_model_idx[p], &ctx) + } else { + 0.0 + }; for (j, entry) in rate_state_grad { total += eval_deriv_entry(entry, &ctx) * s[j * d + p]; } @@ -289,23 +305,24 @@ fn rk4_step( let mut d_flow = vec![0.0f64; stage_flow]; let mut ks_state = [const { Vec::new() }; 4]; let mut ks_flow = [const { Vec::new() }; 4]; + let forcing = s.forcing; // k1 at the entry S. - sensitivity_derivs(model, int_vals, real_vals, params, param_model_idx, t, dt, per_eval, &s.state, &mut d_state, &mut d_flow); + sensitivity_derivs(model, int_vals, real_vals, params, param_model_idx, forcing, t, dt, per_eval, &s.state, &mut d_state, &mut d_flow); ks_state[0] = d_state.clone(); ks_flow[0] = d_flow.clone(); // k2, k3 at the half-step x-stages, with S advanced by the previous slope. let s2s: Vec = s.state.iter().zip(&ks_state[0]).map(|(x, k)| x + 0.5 * dt * k).collect(); - sensitivity_derivs(model, &s2i, &s2r, params, param_model_idx, t + 0.5 * dt, dt, per_eval, &s2s, &mut d_state, &mut d_flow); + sensitivity_derivs(model, &s2i, &s2r, params, param_model_idx, forcing, t + 0.5 * dt, dt, per_eval, &s2s, &mut d_state, &mut d_flow); ks_state[1] = d_state.clone(); ks_flow[1] = d_flow.clone(); let s3s: Vec = s.state.iter().zip(&ks_state[1]).map(|(x, k)| x + 0.5 * dt * k).collect(); - sensitivity_derivs(model, &s3i, &s3r, params, param_model_idx, t + 0.5 * dt, dt, per_eval, &s3s, &mut d_state, &mut d_flow); + sensitivity_derivs(model, &s3i, &s3r, params, param_model_idx, forcing, t + 0.5 * dt, dt, per_eval, &s3s, &mut d_state, &mut d_flow); ks_state[2] = d_state.clone(); ks_flow[2] = d_flow.clone(); // k4 at the full-step x-stage. let s4s: Vec = s.state.iter().zip(&ks_state[2]).map(|(x, k)| x + dt * k).collect(); - sensitivity_derivs(model, &s4i, &s4r, params, param_model_idx, t + dt, dt, per_eval, &s4s, &mut d_state, &mut d_flow); + sensitivity_derivs(model, &s4i, &s4r, params, param_model_idx, forcing, t + dt, dt, per_eval, &s4s, &mut d_state, &mut d_flow); ks_state[3] = d_state; ks_flow[3] = d_flow; (ks_state, ks_flow) @@ -348,6 +365,102 @@ fn rk4_step( Ok(()) } +/// gh#396: integrate the augmented `(x, S)` system over one forcing period +/// `[t_start, t_start + period]`, returning `(x(end), S(end))` — the building +/// block of the periodic-equilibrium solver. +/// +/// - `x0` seeds `x(t_start)` (integer compartments; real compartments are refused +/// on the gradient path). +/// - `s0` seeds `S(t_start)` (`n_int × d`, row-major). With `forcing = false` and +/// `s0 = Iₙ` (so `d = n_int`), `S(end)` is the **monodromy** `M = ∂x(end)/∂x₀`. +/// With `forcing = true` and `s0 = 0` (`d = n_params`), `S(end)` is the +/// **partial** `∂Φ_P/∂θ` (initial state held fixed). +/// - `param_model_idx` maps the `d` columns to model parameters (ignored when +/// `forcing = false`). +/// +/// Fixed-step RK4, last step clipped to land exactly on the period endpoint. No +/// effects/output/obs recording: the validity gate guarantees `[t_start, t_start +/// + period]` is effect-free and `P`-periodic, and the forcing is read at absolute +/// `t` so the stroboscopic phase is correct. `per_eval` is `t`-inert, so staging it +/// at `t_start` is exact. +/// +/// Caveat: `rk4_step` clamps the value state at 0 but not `S` (§1c). On a genuinely +/// endemic (interior) cycle the clamp is inactive and `M` is exact; the solver's +/// oracles (equilibrium == burn-in, sensitivity three-way) are the gate that a +/// clamped iterate did not corrupt the result. +#[allow(clippy::too_many_arguments)] +pub(crate) fn period_flow( + model: &CompiledModel, + params: &[f64], + x0: &[f64], + s0: &[f64], + param_model_idx: &[usize], + forcing: bool, + t_start: f64, + period: f64, + dt: f64, +) -> Result<(Vec, Vec), SimError> { + let ni = x0.len(); + let d = param_model_idx.len(); + debug_assert_eq!(s0.len(), ni * d, "S(t_start) seed must be n_int × d"); + let n_tr = model.model.transitions.len(); + let per_eval = crate::resolved_expr::stage_per_eval(model, params, t_start, dt); + let mut int_vals = x0.to_vec(); + let mut real_vals: Vec = Vec::new(); + let mut sens = Sens { state: s0.to_vec(), flow: vec![0.0; n_tr * d], forcing }; + let t_end = t_start + period; + let mut t = t_start; + while t < t_end - MIN_STEP_EPS { + let h = dt.min(t_end - t); + rk4_step( + model, + &mut int_vals, + &mut real_vals, + None, // no flow accumulation — the solver needs only x and S + params, + t, + h, + per_eval.as_deref(), + Some(&mut sens), + param_model_idx, + )?; + t += h; + } + Ok((int_vals, sens.state)) +} + +/// gh#396: the integer-compartment derivative field `dx/dt` at a reference state +/// and absolute time `t`. Used by the warm-start periodicity gate to test that +/// the vector field (hence the forcing) is `P`-periodic near `T_eq` — a table / +/// `interpolated` forcing carries no declared period, so periodicity is checked +/// numerically rather than read off the forcing kind. +pub(crate) fn derivs_at( + model: &CompiledModel, + params: &[f64], + int_vals: &[f64], + t: f64, + dt: f64, +) -> Result, SimError> { + let per_eval = crate::resolved_expr::stage_per_eval(model, params, t, dt); + let real_vals: Vec = Vec::new(); + let mut di = vec![0.0; int_vals.len()]; + let mut dr = vec![0.0; real_vals.len()]; + let mut df = vec![0.0; model.model.transitions.len()]; + ode_derivs( + model, + int_vals, + &real_vals, + params, + t, + dt, + per_eval.as_deref(), + &mut di, + &mut dr, + &mut df, + )?; + Ok(di) +} + /// One integrated ODE state: the integer- and real-compartment values plus the /// cumulative per-transition flow integrals. `flow` is `∫ rate dt` accumulated /// since the last output reset (reset to 0 at each output boundary so @@ -886,6 +999,12 @@ pub(crate) fn integrate_obs_sensitivity( model: &CompiledModel, params: &[f64], param_model_idx: &[usize], + // gh#396 Phase A: injected start state x(t_start). `None` → build it from the + // model (`initial_state_continuous`, today's behaviour). `Some` → begin from + // this continuous integer-compartment state (the periodic-equilibrium warm + // start; real compartments are refused on the gradient path). The `state_sens_0` + // seed and `cfg.t_start` were already caller-supplied. + x0: Option<&[f64]>, state_sens_0: &[f64], cfg: &OdeConfig, obs_times: &[f64], @@ -922,7 +1041,13 @@ pub(crate) fn integrate_obs_sensitivity( // Continuous (un-rounded) initial state: the ODE gradient path must start // from the smooth initial value so value and forward-sensitivity are // consistent in θ (§1c; see `CompiledModel::initial_state_continuous`). - let (int_s0, real_s0) = model.initial_state_continuous(params)?; + let (int_s0, real_s0) = match x0 { + // Phase A: warm-start from an injected continuous state. Real compartments + // are refused on this path (`gradient_capability.rs`), so the injected + // state is integer-compartments only and `real` starts empty. + Some(x) => (x.to_vec(), Vec::new()), + None => model.initial_state_continuous(params)?, + }; let ni = int_s0.len(); debug_assert_eq!( state_sens_0.len(), @@ -939,6 +1064,7 @@ pub(crate) fn integrate_obs_sensitivity( sens: Sens { state: state_sens_0.to_vec(), flow: vec![0.0; n_tr * d], + forcing: true, }, }; @@ -1185,7 +1311,7 @@ mod tests { let mut int = init.to_vec(); let mut real: Vec = vec![]; let mut flow = vec![0.0f64; nf]; - let mut sens = Sens { state: vec![0.0f64; ni * d], flow: vec![0.0f64; nf * d] }; + let mut sens = Sens { state: vec![0.0f64; ni * d], flow: vec![0.0f64; nf * d], forcing: true }; let n_steps = (t_end / dt).round() as usize; let mut t = 0.0; @@ -1404,7 +1530,7 @@ mod tests { let obs = [10.0f64, 20.0, 30.0]; let seed = vec![0.0f64; ni * d]; // params don't enter two_state's initial counts - let recs = integrate_obs_sensitivity(&cm, &base, &pmi, &seed, &cfg, &obs).unwrap(); + let recs = integrate_obs_sensitivity(&cm, &base, &pmi, None, &seed, &cfg, &obs).unwrap(); assert_eq!(recs.len(), 3, "one record per obs time"); // Central FD in each parameter (hoisted: perturbation depends on the param, @@ -1417,8 +1543,8 @@ mod tests { pp[pk] += eps; pm[pk] -= eps; ( - integrate_obs_sensitivity(&cm, &pp, &pmi, &seed, &cfg, &obs).unwrap(), - integrate_obs_sensitivity(&cm, &pm, &pmi, &seed, &cfg, &obs).unwrap(), + integrate_obs_sensitivity(&cm, &pp, &pmi, None, &seed, &cfg, &obs).unwrap(), + integrate_obs_sensitivity(&cm, &pm, &pmi, None, &seed, &cfg, &obs).unwrap(), ) }) .collect(); diff --git a/rust/crates/sim/src/ode_equilibrium.rs b/rust/crates/sim/src/ode_equilibrium.rs new file mode 100644 index 00000000..9b112545 --- /dev/null +++ b/rust/crates/sim/src/ode_equilibrium.rs @@ -0,0 +1,571 @@ +//! Periodic-equilibrium warm-start solver (gh#396). +//! +//! Solves the seasonal limit cycle as the fixed point of the one-period +//! stroboscopic (Poincaré) map anchored at the warm-start time `T_eq`: +//! +//! ```text +//! X* = Φ_P(X*; θ), Φ_P = flow of the ODE over [T_eq, T_eq + P] at θ. +//! ``` +//! +//! Newton uses the monodromy `M = ∂Φ_P/∂X₀` (from [`crate::ode::period_flow`] in +//! its `forcing = false`, identity-seeded mode). The equilibrium sensitivity +//! follows from the implicit function theorem, `∂X*/∂θ = (I − M)⁻¹ ∂Φ_P/∂θ`, and +//! reuses the same factorization. +//! +//! ## Conservation +//! +//! A closed model conserves a linear quantity `cᵀx` (e.g. per-village population), +//! making `cᵀM = cᵀ` so `(I − M)` is singular (a unit Floquet multiplier), +//! *endemic or not*. Any Runge–Kutta method preserves linear invariants exactly, +//! so the conserved singular values of `(I − M)` sit at floating-point roundoff +//! (~1e-12) while the slowest genuine mode of a decades-long relaxation is at +//! `1 − |μ| ≈ 1e-2` — a ~10-order-of-magnitude gap, so an SVD threshold cleanly +//! separates them (no fragile "0.99 vs 1.00" call). We solve on the conservation +//! manifold `{ CᵀX = CᵀX_init }` via a bordered/KKT system, which is nonsingular +//! whenever the *transverse* cycle is hyperbolic. + +use crate::compiled_model::CompiledModel; +use crate::error::SimError; +use nalgebra::{DMatrix, DVector}; + +/// Singular values of `(I − M)` below `CONS_EPS · σ_max` are conserved/neutral +/// directions. Genuine conservation is exact (RK preserves linear invariants), so +/// these sit at ~1e-12; the slowest genuine mode of a many-decade relaxation is at +/// ~1e-2. Any value in `[1e-10, 1e-4]` separates them. +const CONS_EPS: f64 = 1e-8; + +/// The endemic gate. If the smallest **non-conserved** singular value of `(I − M)` +/// falls below `ENDEMIC_MIN · σ_max`, the transverse cycle is (near-)non-hyperbolic +/// — a mode with relaxation > ~1e6 periods, i.e. the model is not settling to a +/// seasonal cycle. Refuse rather than seed a garbage equilibrium. Sits far above +/// `CONS_EPS`, so a genuine slow mode is never mistaken for conservation. +const ENDEMIC_MIN: f64 = 1e-6; + +/// Newton converges when `‖Φ_P(X) − X‖∞ < TOL_EQ · (1 + ‖X‖∞)`. +const TOL_EQ: f64 = 1e-8; + +/// Newton iteration cap. A hyperbolic cycle converges in a handful of steps; +/// exhausting this signals a marginally-stable (near-non-endemic) cycle. +const MAX_NEWTON: usize = 60; + +/// Fit-time configuration for periodic-equilibrium warm-start. Threaded from the +/// fit config (CLI/TOML) through to [`super::inference::ode_grad::det_grad`]. +#[derive(Clone, Copy, Debug)] +pub struct WarmStart { + /// The warm-start anchor time `T_eq` (absolute model time): integration begins + /// here from the solved equilibrium instead of from `origin`. A single scalar + /// `≤ min(first_obs)`, decoupled from the per-stream conditioning window. + pub t_eq: f64, + /// The forcing fundamental period `P` (model time units) — declared, since an + /// `interpolated`/table forcing carries no period. + pub period: f64, +} + +/// The solved seasonal equilibrium at `T_eq` and its parameter sensitivity. +pub struct Equilibrium { + /// `X*` at `T_eq`: `n_int` continuous compartment values. + pub x_star: Vec, + /// `∂X*/∂θ`, `n_int × d` row-major (`d = param_model_idx.len()`) — the + /// forward-sensitivity seed `S(T_eq)` for the data-window integration. + pub x_star_sens: Vec, + /// Number of conserved directions detected (diagnostic). + pub n_conserved: usize, + /// Newton iterations taken (diagnostic). + pub iters: usize, +} + +/// Solve the seasonal equilibrium `X*(θ)` and its sensitivity `∂X*/∂θ` at `T_eq`. +/// +/// - `param_model_idx` — the `d` estimated parameters (model indices). +/// - `ic_grad_seed` — `∂X_init/∂θ` (`n_int × d`, the same seed `det_grad` builds): +/// supplies the conserved-direction sensitivity `∂b/∂θ = Cᵀ · ic_grad_seed`. +/// - `period` (`P`) and `dt` — the forcing period and the fixed RK4 step. +/// +/// Assumes the pre-`T_eq` dynamics are effect-free and `P`-periodic (the caller's +/// validity gate). Refuses (hard error) a non-endemic parameterization. +pub fn solve_equilibrium( + model: &CompiledModel, + params: &[f64], + param_model_idx: &[usize], + ic_grad_seed: &[f64], + t_eq: f64, + period: f64, + dt: f64, +) -> Result { + let (int_s0, _real) = model.initial_state_continuous(params)?; + let ni = int_s0.len(); + let d = param_model_idx.len(); + if ni == 0 { + return Err(SimError::Validation( + "periodic-equilibrium warm-start: model has no integer compartments".into(), + )); + } + + // Identity seed for the monodromy (n_int × n_int, row-major) and a dummy + // parameter map (unused when forcing = false). + let mut ident = vec![0.0f64; ni * ni]; + for i in 0..ni { + ident[i * ni + i] = 1.0; + } + let mono_idx: Vec = (0..ni).collect(); + + let x_init = DVector::from_column_slice(&int_s0); + let mut x = x_init.clone(); + + // Conserved directions (structural: detected once from the first monodromy and + // reused — a linear invariant of the vector field, independent of x and θ). + let (c_mat, b_vec) = { + let (_phi0, m0_flat) = crate::ode::period_flow( + model, params, x.as_slice(), &ident, &mono_idx, false, t_eq, period, dt, + )?; + let m0 = DMatrix::from_row_slice(ni, ni, &m0_flat); + let (c, _sigma_min) = detect_conserved(&m0, ni)?; + let b = c.transpose() * &x_init; // b = Cᵀ X_init — the pinned totals + (c, b) + }; + let k = c_mat.ncols(); + + // --- Newton on the stroboscopic fixed point, on the conservation manifold --- + let mut iters = 0usize; + let mut converged = false; + for it in 0..MAX_NEWTON { + iters = it + 1; + let (phi_vec, m_flat) = crate::ode::period_flow( + model, params, x.as_slice(), &ident, &mono_idx, false, t_eq, period, dt, + )?; + let phi = DVector::from_vec(phi_vec); + let m = DMatrix::from_row_slice(ni, ni, &m_flat); + + let r = &phi - &x; // residual Φ_P(X) − X + let rinf = r.amax(); + if rinf < TOL_EQ * (1.0 + x.amax()) { + converged = true; + break; + } + + let bordered = build_bordered(&m, &c_mat, ni, k); + let mut rhs = DVector::::zeros(ni + k); + rhs.rows_mut(0, ni).copy_from(&r); + if k > 0 { + let cons_resid = &b_vec - c_mat.transpose() * &x; + rhs.rows_mut(ni, k).copy_from(&cons_resid); + } + let sol = bordered.lu().solve(&rhs).ok_or_else(non_endemic_err)?; + x += sol.rows(0, ni); + } + if !converged { + return Err(SimError::Validation(format!( + "periodic-equilibrium Newton did not converge in {MAX_NEWTON} iterations. \ + The seasonal cycle is (near-)marginally stable at these parameters \ + (non-endemic); fit without warm-start or adjust the parameters." + ))); + } + + // --- Sensitivity at the converged X* (monodromy re-evaluated there) --- + let (_phi_star, m_star_flat) = crate::ode::period_flow( + model, params, x.as_slice(), &ident, &mono_idx, false, t_eq, period, dt, + )?; + let m_star = DMatrix::from_row_slice(ni, ni, &m_star_flat); + // Endemic gate on the transverse conditioning at X*. + endemic_gate(&m_star, ni, k)?; + + // ∂Φ_P/∂θ — the PARTIAL (initial state held fixed): a forcing-ON, zero-seeded + // one-period solve at X*. + let zero_seed = vec![0.0f64; ni * d]; + let (_phi2, dphi_flat) = crate::ode::period_flow( + model, params, x.as_slice(), &zero_seed, param_model_idx, true, t_eq, period, dt, + )?; + let dphi = DMatrix::from_row_slice(ni, d, &dphi_flat); // n_int × d + + // ∂b/∂θ = Cᵀ · ∂X_init/∂θ (reuses the ic_grad seed for the conserved part). + let ic_grad = DMatrix::from_row_slice(ni, d, ic_grad_seed); + let db_dtheta = if k > 0 { + c_mat.transpose() * &ic_grad + } else { + DMatrix::zeros(0, d) + }; + + // Solve the bordered system column-by-column, reusing one factorization. + let bordered = build_bordered(&m_star, &c_mat, ni, k); + let lu = bordered.lu(); + let mut x_star_sens = vec![0.0f64; ni * d]; + for col in 0..d { + let mut rhs = DVector::::zeros(ni + k); + for i in 0..ni { + rhs[i] = dphi[(i, col)]; + } + for kk in 0..k { + rhs[ni + kk] = db_dtheta[(kk, col)]; + } + let sol = lu.solve(&rhs).ok_or_else(non_endemic_err)?; + for i in 0..ni { + x_star_sens[i * d + col] = sol[i]; + } + } + + Ok(Equilibrium { + x_star: x.as_slice().to_vec(), + x_star_sens, + n_conserved: k, + iters, + }) +} + +/// Detect the conserved directions `C` (`n_int × k`) as the left-null space of +/// `(I − M)` — columns of `U` (from the SVD) whose singular value is below +/// `CONS_EPS · σ_max`. Returns `C` and the smallest non-conserved singular value. +fn detect_conserved(m: &DMatrix, ni: usize) -> Result<(DMatrix, f64), SimError> { + let i_minus_m = DMatrix::::identity(ni, ni) - m; + let svd = i_minus_m.svd(true, false); + let u = svd + .u + .as_ref() + .ok_or_else(|| SimError::Validation("equilibrium: SVD produced no U".into()))?; + let sv = &svd.singular_values; + let sigma_max = sv.iter().cloned().fold(0.0_f64, f64::max).max(1e-300); + // Singular values are descending; the conserved (near-zero) ones are at the end. + let mut cols: Vec = Vec::new(); + let mut smallest_non_conserved = sigma_max; + for j in 0..ni { + if sv[j] < CONS_EPS * sigma_max { + cols.push(j); + } else { + smallest_non_conserved = sv[j]; // last (smallest) non-conserved as we descend + } + } + let k = cols.len(); + let mut c = DMatrix::::zeros(ni, k); + for (ci, &j) in cols.iter().enumerate() { + c.set_column(ci, &u.column(j)); + } + Ok((c, smallest_non_conserved)) +} + +/// Endemic gate: refuse if the smallest non-conserved mode is (near-)singular. +fn endemic_gate(m: &DMatrix, ni: usize, _k: usize) -> Result<(), SimError> { + let i_minus_m = DMatrix::::identity(ni, ni) - m; + let svd = i_minus_m.svd(false, false); + let sv = &svd.singular_values; + let sigma_max = sv.iter().cloned().fold(0.0_f64, f64::max).max(1e-300); + // Smallest singular value that is NOT a conserved direction. + let mut smallest_non_conserved = sigma_max; + for j in 0..ni { + if sv[j] >= CONS_EPS * sigma_max { + smallest_non_conserved = sv[j]; + } + } + if smallest_non_conserved < ENDEMIC_MIN * sigma_max { + return Err(non_endemic_err()); + } + Ok(()) +} + +/// Assemble the bordered/KKT matrix `[[I − M, C], [Cᵀ, 0]]` of size `n_int + k`. +fn build_bordered(m: &DMatrix, c: &DMatrix, ni: usize, k: usize) -> DMatrix { + let i_minus_m = DMatrix::::identity(ni, ni) - m; + let mut bordered = DMatrix::::zeros(ni + k, ni + k); + bordered.view_mut((0, 0), (ni, ni)).copy_from(&i_minus_m); + if k > 0 { + bordered.view_mut((0, ni), (ni, k)).copy_from(c); + bordered.view_mut((ni, 0), (k, ni)).copy_from(&c.transpose()); + } + bordered +} + +/// Relative mismatch above which the vector field is deemed non-`P`-periodic. +const FORCING_PERIOD_EPS: f64 = 1e-6; + +/// Validity gate (gh#396): verify the vector field is `P`-periodic over the last +/// burn-in period `[T_eq − P, T_eq]`, so the stroboscopic fixed point is the real +/// seasonal cycle rather than a spurious fixed point of a non-periodic map. A +/// table / `interpolated` forcing carries no declared period, so periodicity is +/// checked **numerically**: `dx/dt` at a fixed reference state must repeat with +/// period `P`. Refuse (hard error) otherwise — the alternative is a silently +/// wrong warm-start. Called once at fit setup (periodicity is θ-shape-independent). +pub fn check_periodicity( + compiled: &CompiledModel, + params: &[f64], + t_eq: f64, + period: f64, + dt: f64, +) -> Result<(), SimError> { + if !(period > 0.0) { + return Err(SimError::Validation(format!( + "warm-start period must be positive, got {period}" + ))); + } + // Reference state with every compartment positive, so all forcing-dependent + // rates are active and any t-dependence of the forcing is visible in dx/dt. + let (int_s0, _) = compiled.initial_state_continuous(params)?; + let x_ref: Vec = int_s0.iter().map(|v| v.max(1.0)).collect(); + + const N_SAMPLES: usize = 8; + let mut max_rel = 0.0_f64; + for s in 0..N_SAMPLES { + let t = t_eq - period + (s as f64 / N_SAMPLES as f64) * period; + let d0 = crate::ode::derivs_at(compiled, params, &x_ref, t, dt)?; + let d1 = crate::ode::derivs_at(compiled, params, &x_ref, t + period, dt)?; + let scale = d0 + .iter() + .chain(&d1) + .fold(0.0_f64, |m, v| m.max(v.abs())) + .max(1e-300); + for (a, b) in d0.iter().zip(&d1) { + max_rel = max_rel.max((a - b).abs() / scale); + } + } + if max_rel > FORCING_PERIOD_EPS { + return Err(SimError::Validation(format!( + "warm-start: the model's dynamics are not periodic with period {period} near \ + T_eq = {t_eq} (relative mismatch {max_rel:.2e} between the vector field at t \ + and t+P). Periodic-equilibrium warm-start requires a P-periodic forcing over \ + the burn-in window — check `warm_start_period` / `warm_start_at`, or fit \ + without warm-start for a non-periodic transient (gh#396)." + ))); + } + Ok(()) +} + +fn non_endemic_err() -> SimError { + SimError::Validation( + "periodic-equilibrium warm-start: no isolated stable seasonal cycle at these \ + parameters (a non-conservation Floquet multiplier is at/above 1 — the model \ + is not settling to a periodic equilibrium). Fit without warm-start, or adjust \ + the parameters that break endemicity (gh#396)." + .to_string(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::compiled_model::CompiledModel; + + const P: f64 = 365.25; + const DT: f64 = 1.0; + + /// A closed seasonal SEIRS(+V≡0): `seir_vaccine_seasonal` with the SIA + /// intervention removed (the ODE gradient path refuses scheduled effects; with + /// no SIA, V stays 0). N = S+E+I+R is conserved (k = 1) — exercises the + /// bordered/KKT solve. `N0` scaled to 1000 for manageable finite-difference. + fn load_model() -> CompiledModel { + let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap(); + // Freshly-compiled IR (carries `rate_state_grad`; the `ir/golden/` copies + // predate the gh#275 emission — see testdata/README.md). + let path = std::path::PathBuf::from(&manifest) + .join("testdata/seir_seasonal_closed.ir.json"); + let contents = std::fs::read_to_string(&path).unwrap(); + let mut model: ir::Model = ir::from_str(&contents).unwrap(); + model.interventions = Vec::new(); + // The golden's parameters are scenario-valued; apply concrete baseline + // values (R0 = β/γ = 3, seasonal amplitude 0.15) so `CompiledModel::new` + // has values. N scaled to 1e3 (frequency-dependent β·S·I/N ⇒ endemic + // fractions invariant to N) for manageable finite differences. + let vals: &[(&str, f64)] = &[ + ("beta", 0.3), + ("sigma", 0.2), + ("gamma", 0.1), + ("omega", 0.003), + ("reversion_rate", 1e-6), + ("alpha", 0.15), + ("phi_season", 90.0), + ("vacc_frac", 0.8), + ("N0", 1000.0), + ("I0", 10.0), + ]; + for p in &mut model.parameters { + if let Some(&(_, v)) = vals.iter().find(|(n, _)| *n == p.name) { + p.value = p.value.with_value(v); + } + } + CompiledModel::new(model).expect("compile seasonal closed model") + } + + fn baseline_params(cm: &CompiledModel) -> Vec { + let n = cm.model.parameters.len(); + let mut params = vec![0.0; n]; + for p in &cm.model.parameters { + params[cm.param_index[p.name.as_str()]] = p.value.resolved_value().unwrap(); + } + params + } + + fn identity_seed(ni: usize) -> Vec { + let mut m = vec![0.0; ni * ni]; + for i in 0..ni { + m[i * ni + i] = 1.0; + } + m + } + + fn param_idx(cm: &CompiledModel, name: &str) -> usize { + cm.param_index[name] + } + + #[test] + fn equilibrium_is_a_fixed_point_and_conserves_n() { + let cm = load_model(); + let params = baseline_params(&cm); + let est: Vec = vec![param_idx(&cm, "beta"), param_idx(&cm, "gamma")]; + let ni = cm.initial_state_continuous(¶ms).unwrap().0.len(); + let ic_grad = vec![0.0; ni * est.len()]; // constant ICs + let eq = solve_equilibrium(&cm, ¶ms, &est, &ic_grad, 0.0, P, DT).unwrap(); + + // Independent recompute: X* is a fixed point of the one-period map. + let ident = identity_seed(ni); + let midx: Vec = (0..ni).collect(); + let (phi, _m) = + crate::ode::period_flow(&cm, ¶ms, &eq.x_star, &ident, &midx, false, 0.0, P, DT) + .unwrap(); + let resid: f64 = phi + .iter() + .zip(&eq.x_star) + .map(|(a, b)| (a - b).abs()) + .fold(0.0, f64::max); + let scale = eq.x_star.iter().cloned().fold(0.0, f64::max); + assert!(resid < 1e-5 * (1.0 + scale), "X* not a fixed point: resid {resid}"); + + // Conservation: N = ΣX detected (k=1) and preserved from the init totals. + assert!((1..=MAX_NEWTON).contains(&eq.iters), "Newton iters out of range: {}", eq.iters); + assert_eq!(eq.n_conserved, 1, "closed SEIRS should have k=1 conservation"); + let n_star: f64 = eq.x_star.iter().sum(); + let (x_init, _) = cm.initial_state_continuous(¶ms).unwrap(); + let n_init: f64 = x_init.iter().sum(); + assert!( + (n_star - n_init).abs() < 1e-6 * n_init, + "N not preserved: init {n_init}, equilibrium {n_star}" + ); + } + + #[test] + fn equilibrium_matches_long_burn_in() { + // The solver's X* must be the ATTRACTING cycle the transient reaches — the + // headline oracle. Iterate the period map (Picard = burn-in) from the model + // IC and confirm it converges to X*. + let cm = load_model(); + let params = baseline_params(&cm); + let est: Vec = vec![param_idx(&cm, "beta")]; + let (x_init, _) = cm.initial_state_continuous(¶ms).unwrap(); + let ni = x_init.len(); + let ic_grad = vec![0.0; ni * est.len()]; + let eq = solve_equilibrium(&cm, ¶ms, &est, &ic_grad, 0.0, P, DT).unwrap(); + + let ident = identity_seed(ni); + let midx: Vec = (0..ni).collect(); + let mut x = x_init.clone(); + for _ in 0..400 { + x = crate::ode::period_flow(&cm, ¶ms, &x, &ident, &midx, false, 0.0, P, DT) + .unwrap() + .0; + } + let scale = eq.x_star.iter().cloned().fold(0.0, f64::max); + for i in 0..ni { + assert!( + (x[i] - eq.x_star[i]).abs() < 1e-3 * (1.0 + scale), + "burn-in[{i}]={} != X*[{i}]={}", + x[i], + eq.x_star[i] + ); + } + } + + #[test] + fn equilibrium_sensitivity_matches_finite_difference() { + // The seed-trap catcher: ∂X*/∂θ from (I−M)⁻¹∂Φ/∂θ must match a central FD of + // re-solving the equilibrium at θ±ε. A wrong (total-not-partial) ∂Φ/∂θ seed + // fails here. + let cm = load_model(); + let params = baseline_params(&cm); + let est: Vec = vec![param_idx(&cm, "beta"), param_idx(&cm, "gamma")]; + let d = est.len(); + let ni = cm.initial_state_continuous(¶ms).unwrap().0.len(); + let ic_grad = vec![0.0; ni * d]; + let eq = solve_equilibrium(&cm, ¶ms, &est, &ic_grad, 0.0, P, DT).unwrap(); + + for (col, &pidx) in est.iter().enumerate() { + let eps = 1e-4 * params[pidx].abs().max(1e-3); + let mut pp = params.clone(); + pp[pidx] += eps; + let mut pm = params.clone(); + pm[pidx] -= eps; + let eqp = solve_equilibrium(&cm, &pp, &est, &ic_grad, 0.0, P, DT).unwrap(); + let eqm = solve_equilibrium(&cm, &pm, &est, &ic_grad, 0.0, P, DT).unwrap(); + let scale = eq.x_star.iter().cloned().fold(0.0, f64::max); + for i in 0..ni { + let fd = (eqp.x_star[i] - eqm.x_star[i]) / (2.0 * eps); + let analytic = eq.x_star_sens[i * d + col]; + assert!( + (fd - analytic).abs() < 1e-2 * (1.0 + analytic.abs()) + 1e-3 * scale, + "∂X*[{i}]/∂θ[{col}]: analytic {analytic}, FD {fd}" + ); + } + } + } + + #[test] + fn periodicity_gate_accepts_correct_period_and_refuses_wrong() { + let cm = load_model(); + let params = baseline_params(&cm); + // The sinusoidal seasonal forcing has period 365.25 → the P-periodicity + // check passes at the true period. + assert!( + check_periodicity(&cm, ¶ms, 0.0, P, DT).is_ok(), + "true period should pass the periodicity gate" + ); + // A wrong period (the forcing is 365.25-periodic, not 100-periodic) → refused. + assert!( + check_periodicity(&cm, ¶ms, 0.0, 100.0, DT).is_err(), + "wrong period should be refused" + ); + // Non-positive period → refused. + assert!(check_periodicity(&cm, ¶ms, 0.0, 0.0, DT).is_err()); + } + + #[test] + fn monodromy_matches_finite_difference() { + let cm = load_model(); + let params = baseline_params(&cm); + let (x_init, _) = cm.initial_state_continuous(¶ms).unwrap(); + let ni = x_init.len(); + let ident = identity_seed(ni); + let midx: Vec = (0..ni).collect(); + // A near-cycle interior state. + let mut x = x_init.clone(); + for _ in 0..20 { + x = crate::ode::period_flow(&cm, ¶ms, &x, &ident, &midx, false, 0.0, P, DT) + .unwrap() + .0; + } + let (phi_base, m_flat) = + crate::ode::period_flow(&cm, ¶ms, &x, &ident, &midx, false, 0.0, P, DT).unwrap(); + for j in 0..ni { + let eps = 1e-4 * x[j].abs().max(1.0); + let mut xp = x.clone(); + xp[j] += eps; + let phip = crate::ode::period_flow(&cm, ¶ms, &xp, &ident, &midx, false, 0.0, P, DT) + .unwrap() + .0; + // Central where the downward perturbation stays ≥ 0; forward otherwise + // (a compartment at the clamp floor, e.g. V ≡ 0, would clamp the + // downward step and halve a central difference). + let (phim, denom) = if x[j] - eps > 0.0 { + let mut xm = x.clone(); + xm[j] -= eps; + let pm = crate::ode::period_flow(&cm, ¶ms, &xm, &ident, &midx, false, 0.0, P, DT) + .unwrap() + .0; + (pm, 2.0 * eps) + } else { + (phi_base.clone(), eps) + }; + for i in 0..ni { + let fd = (phip[i] - phim[i]) / denom; + let analytic = m_flat[i * ni + j]; + assert!( + (fd - analytic).abs() < 2e-2 * (1.0 + analytic.abs()), + "M[{i},{j}]: analytic {analytic}, FD {fd}" + ); + } + } + } +} diff --git a/rust/crates/sim/testdata/README.md b/rust/crates/sim/testdata/README.md new file mode 100644 index 00000000..0a0b068d --- /dev/null +++ b/rust/crates/sim/testdata/README.md @@ -0,0 +1,8 @@ +# sim test fixtures + +- `seir_seasonal_closed.ir.json` — `ocaml/golden/seir_vaccine_seasonal.camdl` + compiled fresh (carries `rate_state_grad`, which the committed `ir/golden/` + copies predate — gh#275 emission postdates their last regeneration). Used by + the `ode_equilibrium` warm-start oracles, which need the ∂rate/∂compartment + Jacobian. The SIA intervention is cleared in-test → a closed seasonal SEIRS + (N conserved). Regenerate: `camdlc ocaml/golden/seir_vaccine_seasonal.camdl`. diff --git a/rust/crates/sim/testdata/seir_seasonal_closed.ir.json b/rust/crates/sim/testdata/seir_seasonal_closed.ir.json new file mode 100644 index 00000000..cc4fe7f5 --- /dev/null +++ b/rust/crates/sim/testdata/seir_seasonal_closed.ir.json @@ -0,0 +1,34 @@ +{"ir_version":"0.29","validated_by":"ocaml-compiler-v0.29","model":{"name":"seir_vaccine_seasonal","version":"0.3","time_unit":"days","description":null,"compartments":[ +{"name":"S","kind":"integer"}, +{"name":"E","kind":"integer"}, +{"name":"I","kind":"integer"}, +{"name":"R","kind":"integer"}, +{"name":"V","kind":"integer"} +],"transitions":[ +{"name":"infection","stoichiometry":[["S",-1],["E",1]],"rate":{"bin_op":{"op":"div","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}},"right":{"binding_ref":"N"}}},"metadata":{"origin_kind":"transmission","source_compartment":"S","dest_compartment":"E"},"rate_grad":{"beta":{"grad":{"bin_op":{"op":"div","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"time_func":{"name":"seasonal"}},"right":{"pop":"S"}}},"right":{"pop":"I"}}},"right":{"binding_ref":"N"}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}},"alpha":{"grad":{"bin_op":{"op":"div","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"un_op":{"op":"sin","arg":{"bin_op":{"op":"div","left":{"bin_op":{"op":"mul","left":{"const":6.283185307179586},"right":{"bin_op":{"op":"sub","left":{"time":null},"right":{"param":"phi_season"}}}}},"right":{"const":365.25}}}}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}},"right":{"binding_ref":"N"}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}},"phi_season":{"grad":{"bin_op":{"op":"div","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"bin_op":{"op":"mul","left":{"param":"alpha"},"right":{"bin_op":{"op":"mul","left":{"un_op":{"op":"cos","arg":{"bin_op":{"op":"div","left":{"bin_op":{"op":"mul","left":{"const":6.283185307179586},"right":{"bin_op":{"op":"sub","left":{"time":null},"right":{"param":"phi_season"}}}}},"right":{"const":365.25}}}}},"right":{"const":-0.017202423838958484}}}}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}},"right":{"binding_ref":"N"}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}}},"rate_state_grad":{"S":{"grad":{"bin_op":{"op":"div","left":{"bin_op":{"op":"sub","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"I"}}},"right":{"binding_ref":"N"}}},"right":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}},"E":{"grad":{"bin_op":{"op":"div","left":{"un_op":{"op":"neg","arg":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}},"I":{"grad":{"bin_op":{"op":"div","left":{"bin_op":{"op":"sub","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"binding_ref":"N"}}},"right":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}},"R":{"grad":{"bin_op":{"op":"div","left":{"un_op":{"op":"neg","arg":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}},"V":{"grad":{"bin_op":{"op":"div","left":{"un_op":{"op":"neg","arg":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"bin_op":{"op":"mul","left":{"param":"beta"},"right":{"time_func":{"name":"seasonal"}}}},"right":{"pop":"S"}}},"right":{"pop":"I"}}}}},"right":{"bin_op":{"op":"mul","left":{"binding_ref":"N"},"right":{"binding_ref":"N"}}}}}}}}, +{"name":"progression","stoichiometry":[["E",-1],["I",1]],"rate":{"bin_op":{"op":"mul","left":{"param":"sigma"},"right":{"pop":"E"}}},"metadata":{"origin_kind":"intrinsic","source_compartment":"E","dest_compartment":"I"},"rate_grad":{"sigma":{"grad":{"pop":"E"}}},"rate_state_grad":{"E":{"grad":{"param":"sigma"}}}}, +{"name":"recovery","stoichiometry":[["I",-1],["R",1]],"rate":{"bin_op":{"op":"mul","left":{"param":"gamma"},"right":{"pop":"I"}}},"metadata":{"origin_kind":"intrinsic","source_compartment":"I","dest_compartment":"R"},"rate_grad":{"gamma":{"grad":{"pop":"I"}}},"rate_state_grad":{"I":{"grad":{"param":"gamma"}}}}, +{"name":"waning","stoichiometry":[["R",-1],["S",1]],"rate":{"bin_op":{"op":"mul","left":{"param":"omega"},"right":{"pop":"R"}}},"metadata":{"origin_kind":"intrinsic","source_compartment":"R","dest_compartment":"S"},"rate_grad":{"omega":{"grad":{"pop":"R"}}},"rate_state_grad":{"R":{"grad":{"param":"omega"}}}}, +{"name":"vacc_waning","stoichiometry":[["V",-1],["S",1]],"rate":{"bin_op":{"op":"mul","left":{"param":"omega"},"right":{"pop":"V"}}},"metadata":{"origin_kind":"intrinsic","source_compartment":"V","dest_compartment":"S"},"rate_grad":{"omega":{"grad":{"pop":"V"}}},"rate_state_grad":{"V":{"grad":{"param":"omega"}}}}, +{"name":"reversion","stoichiometry":[["V",-1],["E",1]],"rate":{"bin_op":{"op":"mul","left":{"param":"reversion_rate"},"right":{"pop":"V"}}},"metadata":{"origin_kind":"intrinsic","source_compartment":"V","dest_compartment":"E"},"rate_grad":{"reversion_rate":{"grad":{"pop":"V"}}},"rate_state_grad":{"V":{"grad":{"param":"reversion_rate"}}}} +],"ode_equations":[],"time_functions":[ +{"name":"seasonal","kind":{"sinusoidal":{"amplitude":{"param":"alpha"},"period":{"const":365.25},"phase":{"param":"phi_season"},"baseline":{"const":1.0}}},"dim":[0,0]} +],"tables":[],"interventions":[ +{"name":"sia_round_1","fire":{"scheduled":{"at_times":[180.0,545.0,910.0]}},"actions":[{"fraction_transfer":{"src":"S","dst":"V","fraction":{"param":"vacc_frac"}}}]} +],"observations":[],"parameters":[ +{"name":"beta","value":{"mode":"required"},"param_kind":"rate","param_dim":null}, +{"name":"sigma","value":{"mode":"required"},"param_kind":"rate","param_dim":null}, +{"name":"gamma","value":{"mode":"required"},"param_kind":"rate","param_dim":null}, +{"name":"omega","value":{"mode":"required"},"param_kind":"rate","param_dim":null}, +{"name":"reversion_rate","value":{"mode":"required"},"param_kind":"rate","param_dim":null}, +{"name":"alpha","value":{"mode":"required"},"param_kind":"probability","param_dim":null}, +{"name":"phi_season","value":{"mode":"required"},"param_kind":"real","param_dim":null}, +{"name":"vacc_frac","value":{"mode":"required"},"param_kind":"probability","param_dim":null}, +{"name":"N0","value":{"mode":"required"},"param_kind":"count","param_dim":null}, +{"name":"I0","value":{"mode":"required"},"param_kind":"count","param_dim":null} +],"initial_conditions":{"parameterized":{"S":{"bin_op":{"op":"sub","left":{"param":"N0"},"right":{"param":"I0"}}},"I":{"param":"I0"}}},"ic_grad":{"S":{"N0":{"grad":{"const":1.0}},"I0":{"grad":{"const":-1.0}}},"I":{"I0":{"grad":{"const":1.0}}}},"output":{"times":{"regular":{"start":0.0,"step":1.0}},"format":"tsv","trajectory":true,"observations":true},"simulation":{"t_start":0.0,"t_end":1095.7275,"time_semantics":"continuous","dt":null,"rng_seed":null},"scenarios":[ +{"name":"baseline","label":"no SIA (baseline)","params":{"beta":0.3,"sigma":0.2,"gamma":0.1,"omega":0.003,"reversion_rate":1e-06,"alpha":0.15,"phi_season":90.0,"vacc_frac":0.8,"N0":1000000.0,"I0":10.0},"enable":[],"disable":[],"t_end":null}, +{"name":"with_sia","label":"with SIA rounds","params":{},"enable":["sia_round_1"],"disable":[],"t_end":null} +],"model_structure":{"dimensions":[],"compartment_dims":{"S":[],"E":[],"I":[],"R":[],"V":[]},"base_compartments":["S","E","I","R","V"],"transmission_transitions":["infection"],"infectious_compartments":["I"]},"bindings":[ +{"name":"N","expr":{"pop_sum":["S","E","I","R","V"]}} +]}}