Module sensitivity 

Forward sensitivity analysis for parameterised ODE systems.

Given an ODE system parameterised by a vector $p \in \mathbb{R}^{N_s}$,

dy/dt = f(t, y, p),    y(t_0) = y_0(p),

the forward sensitivity matrix S(t) = ∂y(t) / ∂p ∈ ℝ^{N × N_s} satisfies the variational equations

dS/dt = J_y(t, y, p) · S + J_p(t, y, p),    S(t_0) = ∂y_0/∂p

where J_y = ∂f/∂y ∈ ℝ^{N×N} and J_p = ∂f/∂p ∈ ℝ^{N×N_s}.

Numra integrates the augmented state z = [y; vec(S)] ∈ ℝ^{N(1+N_s)} with the user’s chosen crate::Solver (DoPri5, Tsit5, Vern*, Radau5, BDF, Esdirk*, Auto). The augmented Jacobian is block_diag(J_y, …, J_y) — the standard CVODES simultaneous-corrector formulation.

Layout conventions

Numra mixes row-major and column-major flattening so that each Jacobian’s “natural” axis is contiguous. Implementors of ParametricOdeSystem must follow these rules exactly:

Quantity	Layout	Index map
state y	flat, length N	y[i]
state Jacobian J_y	row-major, N × N	jy[i*N + j] = ∂f_i/∂y_j
parameter Jacobian J_p	column-major, N × N_s	jp[k*N + i] = ∂f_i/∂p_k
sensitivity S	column-major, N × N_s	s[k*N + i] = ∂y_i/∂p_k
augmented state z	[y; vec(S)]	z[N + k*N + i] = S_{i,k}
initial sensitivity ∂y_0/∂p	column-major, N × N_s	s0[k*N + i] = ∂y_{0,i}/∂p_k

Choosing column-major for J_p and S matches CVODES indexing, makes the per-parameter sub-vector S_{:,k} contiguous (and free to slice via SensitivityResult::sensitivity_for_param), and improves cache locality in the inner loop Ṡ_{:,k} = J_y · S_{:,k} + J_p_{:,k}. J_y stays row-major because state-row access is the standard OdeSystem::jacobian convention used by every implicit solver in this crate.

Worked example: 2-state, 2-parameter linear ODE

Consider the system

ẏ_0 = -p_0 · y_0 + p_1 · y_1
ẏ_1 =          - p_1 · y_1

Then

J_y = [ -p_0    p_1 ]      J_p = [ -y_0    y_1 ]
      [   0    -p_1 ]            [   0    -y_1 ]

With N = N_s = 2, J_y is row-major and J_p is column-major. Implementing the trait by hand:

use numra_ode::sensitivity::ParametricOdeSystem;

struct Lin2 { p: [f64; 2] }

impl ParametricOdeSystem<f64> for Lin2 {
    fn n_states(&self) -> usize { 2 }
    fn n_params(&self) -> usize { 2 }
    fn params(&self) -> &[f64] { &self.p }

    fn rhs_with_params(&self, _t: f64, y: &[f64], p: &[f64], dy: &mut [f64]) {
        dy[0] = -p[0] * y[0] + p[1] * y[1];
        dy[1] =                -p[1] * y[1];
    }

    fn jacobian_y(&self, _t: f64, _y: &[f64], jy: &mut [f64]) {
        // Row-major (N × N): jy[i*N + j] = ∂f_i/∂y_j
        jy[0*2 + 0] = -self.p[0];   // ∂f_0/∂y_0
        jy[0*2 + 1] =  self.p[1];   // ∂f_0/∂y_1
        jy[1*2 + 0] =  0.0;         // ∂f_1/∂y_0
        jy[1*2 + 1] = -self.p[1];   // ∂f_1/∂y_1
    }

    fn jacobian_p(&self, _t: f64, y: &[f64], jp: &mut [f64]) {
        // Column-major (N × N_s): jp[k*N + i] = ∂f_i/∂p_k
        jp[0*2 + 0] = -y[0];        // ∂f_0/∂p_0
        jp[0*2 + 1] =  0.0;         // ∂f_1/∂p_0
        jp[1*2 + 0] =  y[1];        // ∂f_0/∂p_1
        jp[1*2 + 1] = -y[1];        // ∂f_1/∂p_1
    }

    // Flag analytical so AugmentedSystem honors the overrides above
    // instead of falling through to its inline-FD fast path.
    fn has_analytical_jacobian_y(&self) -> bool { true }
    fn has_analytical_jacobian_p(&self) -> bool { true }
}

let sys = Lin2 { p: [1.0, 0.5] };
assert_eq!(sys.n_states(), 2);
assert_eq!(sys.n_params(), 2);

Note the asymmetry: J_y[0*2 + 1] is row 0, column 1 of J_y (state Jacobian, row-major), while J_p[1*2 + 0] is row 0, column 1 of J_p (parameter Jacobian, column-major). They look indexed the same way in Rust syntax but address different mathematical entries by design.

Send / Sync

ParametricOdeSystem does not require Send + Sync, matching the crate::OdeSystem precedent. Add the bounds at the call site (where Sys: ParametricOdeSystem<S> + Send + Sync) when sharing the system across threads — for instance in parallel Monte Carlo. Tightening a trait bound is a breaking change; loosening it is not, so the default is permissive.

Performance note (v1)

Numra’s v1 ships correct forward sensitivity through every solver, with analytical J_y and J_p overrides replacing FD on the dominant cost path. The block-diagonal structure of the augmented Jacobian — which would let implicit solvers reuse a single LU of M = (1/γh)·I_N - J_y across all N_s + 1 sub-systems — is not yet exploited: Radau5 and BDF currently perform dense LU on the full N(N_s+1) × N(N_s+1) matrix. Block-aware factorisation is tracked as a named follow-up (see docs/internal-followups.md).

Author: Moussa Leblouba Date: 4 February 2026 Modified: 6 May 2026

Structs

AugmentedSystem

Wraps a ParametricOdeSystem as an OdeSystem over the augmented state z = [y; vec(S)] (column-major sensitivity flattening).

ClosureSystem

Closure-shaped adapter for solve_forward_sensitivity_with.

SensitivityResult

Result of an ODE forward-sensitivity solve.

Traits

ParametricOdeSystem

An ODE system parameterised by a parameter vector p.

Functions

solve_forward_sensitivity

Compute the forward sensitivity matrix S(t) = ∂y(t) / ∂p of an ODE solution using any Solver.

solve_forward_sensitivity_with

Closure-shaped convenience wrapper around solve_forward_sensitivity.