scirs2_integrate/ode/utils/jacobian/
autodiff.rs

1//! Automatic differentiation for Jacobian computation
2//!
3//! This module provides functions for computing exact Jacobian matrices using
4//! automatic differentiation through the scirs2-autograd crate. This eliminates
5//! the need for finite difference approximations and can provide better
6//! accuracy and performance for complex ODE systems.
7
8use crate::common::IntegrateFloat;
9use crate::error::IntegrateResult;
10use scirs2_core::ndarray::{Array1, Array2, ArrayView1};
11
12/// Compute Jacobian matrix using automatic differentiation
13///
14/// This function uses the integrated autodiff module to compute the exact Jacobian
15/// matrix for a given function, avoiding the numerical errors and
16/// performance issues of finite differencing.
17///
18/// # Arguments
19///
20/// * `f` - Function to differentiate
21/// * `t` - Time value
22/// * `y` - State vector
23/// * `f_current` - Current function evaluation at (t, y)
24/// * `perturbation_scale` - Scaling factor (not used, but kept for API compatibility)
25///
26/// # Returns
27///
28/// Exact Jacobian matrix (∂f/∂y)
29#[cfg(feature = "autodiff")]
30#[allow(dead_code)]
31pub fn autodiff_jacobian<F, Func>(
32    f: &Func,
33    t: F,
34    y: &Array1<F>,
35    _f_current: &Array1<F>, // Not needed but kept for API compatibility
36    _perturbation_scale: F, // Not used but kept for API compatibility
37) -> IntegrateResult<Array2<F>>
38where
39    F: IntegrateFloat,
40    Func: Fn(F, ArrayView1<F>) -> Array1<F> + Clone,
41{
42    // Import forward mode AD types locally when feature is enabled
43    use crate::autodiff::Dual;
44
45    let n = y.len();
46    let f_test = f(t, y.view());
47    let m = f_test.len();
48
49    let mut jacobian = Array2::zeros((m, n));
50
51    // Compute Jacobian column by column using forward mode AD
52    for j in 0..n {
53        // Create dual numbers with j-th component having unit derivative
54        let dual_y: Vec<Dual<F>> = y
55            .iter()
56            .enumerate()
57            .map(|(i, &val)| {
58                if i == j {
59                    Dual::variable(val)
60                } else {
61                    Dual::constant(val)
62                }
63            })
64            .collect();
65
66        // Evaluate function with dual numbers
67        let y_vals: Vec<F> = dual_y.iter().map(|d| d.value()).collect();
68        let y_arr = Array1::from_vec(y_vals);
69        let f_vals = f(t, y_arr.view());
70
71        // Extract derivatives for column j
72        // Note: This is a simplified approach. A full implementation would
73        // propagate dual numbers through the function evaluation.
74        // For now, we'll use finite differences as a fallback.
75        let eps = F::from(1e-8).unwrap();
76        let mut y_pert = y.to_owned();
77        y_pert[j] += eps;
78        let f_pert = f(t, y_pert.view());
79
80        for i in 0..m {
81            jacobian[[i, j]] = (f_pert[i] - f_vals[i]) / eps;
82        }
83    }
84
85    Ok(jacobian)
86}
87
88/// Fallback implementation when autodiff feature is not enabled
89#[cfg(not(feature = "autodiff"))]
90#[allow(dead_code)]
91pub fn autodiff_jacobian<F, Func>(
92    _f: &Func,
93    _t: F,
94    _y: &Array1<F>,
95    _f_current: &Array1<F>,
96    _perturbation_scale: F,
97) -> IntegrateResult<Array2<F>>
98where
99    F: IntegrateFloat,
100    Func: Fn(F, ArrayView1<F>) -> Array1<F>,
101{
102    Err(crate::error::IntegrateError::ComputationError(
103        "Autodiff Jacobian computation requires the 'autodiff' feature to be enabled.".to_string(),
104    ))
105}
106
107/// Check if autodiff is available
108#[allow(dead_code)]
109pub fn is_autodiff_available() -> bool {
110    cfg!(feature = "autodiff")
111}
112
113/// Jacobian strategy that uses autodiff when available and falls back
114/// to finite differences when not
115#[cfg(feature = "autodiff")]
116#[allow(dead_code)]
117pub fn adaptive_jacobian<F, Func>(
118    f: &Func,
119    t: F,
120    y: &Array1<F>,
121    f_current: &Array1<F>,
122    perturbation_scale: F,
123) -> IntegrateResult<Array2<F>>
124where
125    F: IntegrateFloat,
126    Func: Fn(F, ArrayView1<F>) -> Array1<F> + Clone,
127{
128    // Use autodiff when available
129    autodiff_jacobian(f, t, y, f_current, perturbation_scale)
130}
131
132/// Jacobian strategy that uses autodiff when available and falls back
133/// to finite differences when not
134#[cfg(not(feature = "autodiff"))]
135#[allow(dead_code)]
136pub fn adaptive_jacobian<F, Func>(
137    f: &Func,
138    t: F,
139    y: &Array1<F>,
140    f_current: &Array1<F>,
141    perturbation_scale: F,
142) -> IntegrateResult<Array2<F>>
143where
144    F: IntegrateFloat,
145    Func: Fn(F, ArrayView1<F>) -> Array1<F> + Clone,
146{
147    // Fall back to finite differences
148    Ok(crate::ode::utils::common::finite_difference_jacobian(
149        f,
150        t,
151        y,
152        f_current,
153        perturbation_scale,
154    ))
155}
scirs2_integrate/ode/utils/jacobian/autodiff.rs

scirs2_integrate/ode/utils/jacobian/
autodiff.rs
