From 43ffbc3967edc1d5abdbcd12e7c21de89a141532 Mon Sep 17 00:00:00 2001 From: Vince Buffalo Date: Wed, 8 Jul 2026 10:03:56 -0700 Subject: [PATCH] perf(ode-grad): hoist state-Jacobian eval out of per-parameter loop (gh#400) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In `sensitivity_derivs`, the state-Jacobian term ∂rate_r/∂x_j (`eval_deriv_entry`) was re-evaluated inside the `for p in 0..d` loop over estimated parameters, but its value is independent of the parameter column `p` — only the `S[j,p]` multiplier varies. Evaluate each transition's J_x entries once during the first column and reuse them for the remaining columns, cutting the J_x expression-eval work from P·nnz to nnz per transition per RK4 stage. The FOI-type rates (dense J_x, ~8-9 nonzeros) benefit most. Byte-neutral: column 0 keeps the exact original J_θ-then-J_x evaluation order, so the shared binding cache is populated identically; columns p>0 previously re-called `eval_deriv_entry` against the warm cache and got exactly the values now stashed in `jx`. An earlier fully-hoisted variant that evaluated all J_x before J_θ reordered cache population and perturbed one gradient component by 1 ULP — the column-0 interleave avoids that. A/B gate (det_grad on the seir_observations FD-oracle fixture, f64::to_bits): baseline ll=c0598ed08da0c428 grad=c05a4bbe00ed4375 4070de40d6693ee5 bffbd780d1241668 405e2ba1ec971515 after identical — all four grad words + loglik bit-for-bit. FD oracles green (3/3); sim --lib green (323/0). --- rust/crates/sim/src/ode.rs | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/rust/crates/sim/src/ode.rs b/rust/crates/sim/src/ode.rs index cd292d9c..b79bf668 100644 --- a/rust/crates/sim/src/ode.rs +++ b/rust/crates/sim/src/ode.rs @@ -183,14 +183,36 @@ fn sensitivity_derivs( let _cache = crate::resolved_expr::CacheScope::enter(model.resolved.bindings.len()); for v in d_state.iter_mut() { *v = 0.0; } + // Scratch for a transition's state-Jacobian entries ∂rate_r/∂x_j (gh#400). Each + // entry is independent of the parameter column `p` — only the `S[j,p]` multiplier + // varies — so we evaluate it during the FIRST column (`p == 0`), stash it, and + // reuse it for columns `p > 0` instead of re-walking each derivative expression + // `d` times. Reused across transitions (one amortized allocation). + // + // Byte-identical to the pre-hoist loop: column 0 keeps the exact original + // evaluation order (J_θ then J_x), so the shared binding cache is populated in + // the same order and to the same values; columns p>0 previously re-called + // `eval_deriv_entry` and hit the warm cache, returning exactly the `jx[k]` we + // stashed. Evaluating all J_x *before* J_θ would reorder cache population and + // perturb a binding by ~1 ULP (verified — hence the column-0 interleave). + let mut jx: Vec = Vec::new(); for (tr_idx, stoich) in model.transition_stoich.iter().enumerate() { 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 + jx.clear(); 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); - for (j, entry) in rate_state_grad { - total += eval_deriv_entry(entry, &ctx) * s[j * d + p]; + if p == 0 { + for (j, entry) in rate_state_grad.iter() { + let v = eval_deriv_entry(entry, &ctx); + jx.push(v); + total += v * s[j * d + p]; + } + } else { + for (k, (j, _entry)) in rate_state_grad.iter().enumerate() { + total += jx[k] * s[j * d + p]; + } } // ∂flow_r/∂θ_p = total_dr_dθ[p] — the raw per-transition flow // derivative (dc_r/dt = rate_r, so d(∂c_r/∂θ)/dt = ∂rate_r/∂θ), NO