feat(ode): first-class sensitivity analysis API with unified builders

examples/ode/15_forward_sensitivity_analysis/main.rs

@@ -1,65 +1,65 @@
-//! Example 15: Forward Sensitivity Analysis (FSA) using Dual Numbers
+//! Example 15: Forward Sensitivity Analysis (FSA)
 //!
 //! This example demonstrates how to compute sensitivities of an ODE solution
-//! with respect to its parameters using the `num-dual` crate.
-//!
-//! Forward sensitivity analysis calculates the gradient of the solver states y
-//! with respect to system parameters p. By making our ODE implementation generic
-//! over the `Real` trait, we can solve the system using dual numbers, which compute
-//! both the state and its derivatives simultaneously without requiring an
-//! explicitly augmented sensitivity matrix system.
+//! with respect to its parameters using the new structured API for
+//! forward sensitivities.
 //!
 //! We use a simple decay model: dy/dt = -k * y
 //! The parameter is `k`. We want to know how the final state `y(t_f)` changes
 //! with respect to `k` (i.e., dy/dk).
 
 use differential_equations::prelude::*;
-use differential_equations::traits::Real;
 use nalgebra::SVector;
-use num_dual::Dual64;
 
 /// A simple decay model: dy/dt = -k * y
-/// We make the struct generic over `T` to support both `f64` and `Dual64`.
-struct Decay<T> {
-    k: T,
+struct Decay {
+    k: f64,
 }
 
-impl<T: Real> ODE<T, SVector<T, 1>> for Decay<T> {
-    fn diff(&self, _t: T, y: &SVector<T, 1>, dydt: &mut SVector<T, 1>) {
+impl ODE<f64, SVector<f64, 1>> for Decay {
+    fn diff(&self, _t: f64, y: &SVector<f64, 1>, dydt: &mut SVector<f64, 1>) {
         dydt[0] = -self.k * y[0];
     }
 }
 
+impl ParametrizedODE<f64, SVector<f64, 1>, SVector<f64, 1>> for Decay {
+    fn parameters(&self) -> SVector<f64, 1> {
+        SVector::from([self.k])
+    }
+
+    fn jacobian_p(&self, _t: f64, y: &SVector<f64, 1>, j: &mut Matrix<f64>) {
+        // df/dk = -y
+        j[(0, 0)] = -y[0];
+    }
+}
+
 fn main() {
-    // We want to evaluate the sensitivity with respect to `k`. The dual part of
-    // `k` is 1.0 because dk/dk = 1.
-    let k = Dual64::new(1.0, 1.0);
-    let decay = Decay { k };
+    let decay = Decay { k: 1.0 };
 
-    // Initial conditions: y(0) = 1.0.
-    // The dual part is 0.0 because the initial state does not depend on `k`.
-    let y0 = SVector::from([Dual64::from(1.0)]);
+    // Initial condition for the forward state
+    let y0 = SVector::<f64, 1>::from([1.0]);
+
+    // The initial sensitivity S(0) = dy/dk(0) = 0.0 because the initial state does not depend on `k`.
+    // The augmented state is [y, S].
+    let y0_aug = SVector::<f64, 2>::from([1.0, 0.0]);
 
     // Time span
-    let t0 = Dual64::from(0.0);
-    let tf = Dual64::from(2.0);
+    let t0 = 0.0;
+    let tf = 2.0;
 
-    // Specify the DOP853 scalar and state types so the solver uses Dual64.
-    let method = ExplicitRungeKutta::<_, _, Dual64, SVector<Dual64, 1>, 8, 12, 16>::dop853()
-        .rtol(Dual64::from(1e-8))
-        .atol(Dual64::from(1e-8));
+    let method = ExplicitRungeKutta::dop853().rtol(1e-8).atol(1e-8);
 
-    // Solve the ODE. The solver handles Dual64 through the same Real trait used
-    // by f64.
-    let solution = IVP::ode(&decay, t0, tf, y0).method(method).solve().unwrap();
+    // Solve the augmented system using Forward Sensitivity Analysis.
+    let solution = IVP::ode(&decay, t0, tf, y0)
+        .forward_sensitivity(y0_aug)
+        .method(method)
+        .solve()
+        .unwrap();
 
     // Extract the final state
-    let y_final_dual = solution.y.last().unwrap()[0];
-
-    // The real part is the actual solution value: y(tf)
-    let y_final = y_final_dual.re;
-    // The dual part is the derivative with respect to k: dy/dk(tf)
-    let sens_final = y_final_dual.eps;
+    let y_final_aug = solution.y.last().unwrap();
+    let y_final = y_final_aug[0];
+    let sens_final = y_final_aug[1];
 
     // Analytical solution for verification: y(t) = y0 * e^(-k*t)
     // dy/dk = -t * y0 * e^(-k*t)
@@ -68,7 +68,7 @@ fn main() {
     let expected_val = (-k_f64 * tf_f64).exp();
     let expected_sens = -tf_f64 * expected_val;
 
-    println!("Forward Sensitivity Analysis (FSA) using Dual Numbers");
+    println!("Forward Sensitivity Analysis (FSA)");
     println!("-----------------------------------------------------");
     println!("Numerical final state:       {:.8}", y_final);
     println!("Analytical final state:      {:.8}", expected_val);

examples/ode/16_adjoint_sensitivities/main.rs

@@ -1,19 +1,13 @@
 //! Example 16: Adjoint Sensitivity Analysis
 //!
 //! This example demonstrates how to perform Adjoint Sensitivity Analysis (ASA)
-//! using the backward-integration capability of the solvers.
-//!
-//! ASA calculates the gradient of a cost function with respect to parameters
-//! by solving an adjoint ODE backward in time.
+//! using the backward-integration capability of the solvers via the structured
+//! `AdjointOde` API.
 
 use differential_equations::prelude::*;
-use nalgebra::{Matrix2, SVector, vector};
+use nalgebra::{SVector, vector};
 
 // 1. The Forward Problem
-// Consider the problem:
-// dy_0/dt = -p_0 * y_0 + p_1 * y_1
-// dy_1/dt = p_0 * y_0 - p_1 * y_1
-// With y(0) = [1.0, 0.0], p = [0.1, 0.2]
 struct ForwardOde {
     p: SVector<f64, 2>,
 }
@@ -25,81 +19,21 @@ impl ODE<f64, SVector<f64, 2>> for ForwardOde {
     }
 }
 
-// 2. The Adjoint Problem
-// We want to minimize the cost function: G = 0.5 * (y_1(T) - 0.5)^2
-// (Note: This is a discrete cost function at the final time T)
-//
-// Let g(y, p) = 0.5 * (y_1(T) - 0.5)^2
-// dg/dy(T) = [0, y_1(T) - 0.5]
-//
-// The forward Jacobian is:
-// J_y = [-p_0, p_1]
-//       [ p_0,-p_1]
-//
-// J_p = [-y_0, y_1]
-//       [ y_0,-y_1]
-//
-// The adjoint variables lambda satisfy:
-// d lambda / dt = -J_y^T lambda
-// with final condition: lambda(T) = dg/dy(T)^T
-//
-// The parameter-gradient accumulator mu satisfies:
-// d mu / dt = -J_p^T lambda
-// with final condition: mu(T) = 0
-
-struct AdjointOde {
-    p: SVector<f64, 2>,
-    forward_solution: Solution<f64, SVector<f64, 2>>,
-}
-
-// The state for the adjoint problem is composed of lambda and mu
-// Let's use a 4D vector: [lambda_0, lambda_1, mu_0, mu_1]
-impl ODE<f64, SVector<f64, 4>> for AdjointOde {
-    fn diff(&self, t: f64, adjoint_state: &SVector<f64, 4>, dydt: &mut SVector<f64, 4>) {
-        let y = self.interpolate_forward(t);
-
-        let lambda = SVector::<f64, 2>::new(adjoint_state[0], adjoint_state[1]);
-
-        // J_y^T
-        let j_y_t = Matrix2::new(-self.p[0], self.p[0], self.p[1], -self.p[1]);
-
-        // J_p^T
-        let j_p_t = Matrix2::new(-y[0], y[0], y[1], -y[1]);
-
-        // d lambda / dt = -J_y^T lambda
-        let d_lambda_dt = -j_y_t * lambda;
-
-        // d mu / dt = -J_p^T lambda
-        let d_mu_dt = -j_p_t * lambda;
-
-        dydt[0] = d_lambda_dt[0];
-        dydt[1] = d_lambda_dt[1];
-        dydt[2] = d_mu_dt[0];
-        dydt[3] = d_mu_dt[1];
+impl ParametrizedODE<f64, SVector<f64, 2>, SVector<f64, 2>> for ForwardOde {
+    fn parameters(&self) -> SVector<f64, 2> {
+        self.p
     }
-}
-
-impl AdjointOde {
-    fn interpolate_forward(&self, t: f64) -> SVector<f64, 2> {
-        let times = &self.forward_solution.t;
-        let states = &self.forward_solution.y;
-
-        if times.is_empty() {
-            return vector![0.0, 0.0];
-        }
 
-        if t <= times[0] {
-            return states[0];
-        }
-
-        if t >= *times.last().unwrap() {
-            return *states.last().unwrap();
-        }
-
-        let upper = times.partition_point(|ti| *ti < t);
-        let lower = upper - 1;
-        let s = (t - times[lower]) / (times[upper] - times[lower]);
-        states[lower] * (1.0 - s) + states[upper] * s
+    fn jacobian_p(&self, _t: f64, y: &SVector<f64, 2>, j: &mut Matrix<f64>) {
+        // df/dp
+        // df_0/dp_0 = -y_0
+        // df_0/dp_1 = y_1
+        // df_1/dp_0 = y_0
+        // df_1/dp_1 = -y_1
+        j[(0, 0)] = -y[0];
+        j[(0, 1)] = y[1];
+        j[(1, 0)] = y[0];
+        j[(1, 1)] = -y[1];
     }
 }
 
@@ -112,24 +46,20 @@ fn main() {
     // 1. Solve the forward problem with dense output
     let forward_ode = ForwardOde { p };
     let forward_solution = IVP::ode(&forward_ode, t0, tf, y0)
-        .dense(10) // high density for accurate linear interpolation
+        .dense(10) // high density for accurate interpolation
         .method(ExplicitRungeKutta::dop853().rtol(1e-8).atol(1e-8))
         .solve()
         .unwrap();
 
-    // 2. Solve the adjoint problem backwards
-    let adjoint_ode = AdjointOde {
-        p,
-        forward_solution,
-    };
-
+    // 2. Solve the adjoint problem backwards using the new builder method
     // Initial condition for backward pass: lambda(T) = dg/dy(T)^T, mu(T) = 0
-    let y_final = adjoint_ode.forward_solution.y.last().unwrap();
+    let y_final = forward_solution.y.last().unwrap();
     let dg_dy_final = vector![0.0, y_final[1] - 0.5];
     let adjoint_y0 = vector![dg_dy_final[0], dg_dy_final[1], 0.0, 0.0];
 
-    // Integrate backwards from tf to t0
-    let adjoint_solution = IVP::ode(&adjoint_ode, tf, t0, adjoint_y0)
+    // Integrate backwards
+    let adjoint_solution = forward_solution
+        .adjoint_sensitivity(&forward_ode, adjoint_y0)
         .method(ExplicitRungeKutta::dop853().rtol(1e-8).atol(1e-8))
         .solve()
         .unwrap();
@@ -139,9 +69,6 @@ fn main() {
 
     println!("Adjoint Sensitivity Analysis");
     println!("============================");
-    // The backward integration returns the accumulator at t0. With the sign
-    // convention above, mu(t0) is the gradient of the terminal cost with
-    // respect to the parameters.
     println!("Computed Gradient w.r.t parameters: {:?}", gradient);
 
     // Double-check with central finite differences.