Composable numerical methods for Rust.

use numra::ode::{Bdf, SolverOptions};
use numra::ode::sensitivity::{solve_forward_sensitivity, ParametricOdeSystem};
use numra::optim::{lm_minimize, LmOptions};

// Observed: a sample of the decaying state at t = 4.
let observed = 0.14;

// Parametric model: dy/dt = -k·y, with k unknown.
struct Decay { k: f64 }
impl ParametricOdeSystem<f64> for Decay {
    fn n_states(&self) -> usize  { 1 }
    fn n_params(&self) -> usize  { 1 }
    fn params(&self)   -> &[f64] { std::slice::from_ref(&self.k) }
    fn rhs_with_params(&self, _t: f64, y: &[f64], p: &[f64], dy: &mut [f64]) {
        dy[0] = -p[0] * y[0];
    }
}

// One Bdf solve of the augmented system gives state AND analytical sensitivity dy/dk.
let solve_with_sens = |k: &[f64]| {
    let sys = Decay { k: k[0] };
    solve_forward_sensitivity::<Bdf, f64, _>(
        &sys, 0.0, 4.0, &[1.0],
        &SolverOptions::default().rtol(1e-8).atol(1e-12),
    ).expect("augmented solve failed")
};

let residual = |k: &[f64], r: &mut [f64]| {
    let sol  = solve_with_sens(k);
    let last = sol.len() - 1;
    r[0] = sol.y_at(last)[0] - observed;
};
let jacobian = |k: &[f64], j: &mut [f64]| {
    let sol  = solve_with_sens(k);
    let last = sol.len() - 1;
    j[0] = sol.dyi_dpj(last, 0, 0);
};

let fit = lm_minimize(residual, jacobian, &[0.3], 1, &LmOptions::default())?;
println!("k = {:.4}  (final cost {:.2e})", fit.x[0], fit.f);