scirs2_integrate/symbolic/
jacobian.rs

1//! Automatic Jacobian generation using symbolic differentiation
2//!
3//! This module provides functionality for automatically generating
4//! Jacobian matrices from symbolic expressions, eliminating the need
5//! for finite difference approximations.
6
7use super::expression::{simplify, SymbolicExpression, Variable};
8use crate::common::IntegrateFloat;
9use crate::error::{IntegrateError, IntegrateResult};
10use scirs2_core::ndarray::{Array2, ArrayView1};
11use std::collections::HashMap;
12
13// Helper functions for creating symbolic expressions
14#[allow(dead_code)]
15fn var<F: IntegrateFloat>(name: &str) -> SymbolicExpression<F> {
16    SymbolicExpression::var(name)
17}
18
19#[allow(dead_code)]
20fn indexed_var<F: IntegrateFloat>(name: &str, index: usize) -> SymbolicExpression<F> {
21    SymbolicExpression::indexedvar(name, index)
22}
23
24#[allow(dead_code)]
25fn constant<F: IntegrateFloat>(value: F) -> SymbolicExpression<F> {
26    SymbolicExpression::constant(value)
27}
28
29/// Represents a symbolic Jacobian matrix
30pub struct SymbolicJacobian<F: IntegrateFloat> {
31    /// The symbolic expressions for each element of the Jacobian
32    pub elements: Array2<SymbolicExpression<F>>,
33    /// The state variables with respect to which we differentiate
34    pub state_vars: Vec<Variable>,
35    /// The time variable (if time-dependent)
36    pub time_var: Option<Variable>,
37}
38
39impl<F: IntegrateFloat> SymbolicJacobian<F> {
40    /// Create a new symbolic Jacobian
41    pub fn new(
42        elements: Array2<SymbolicExpression<F>>,
43        state_vars: Vec<Variable>,
44        time_var: Option<Variable>,
45    ) -> Self {
46        SymbolicJacobian {
47            elements,
48            state_vars,
49            time_var,
50        }
51    }
52
53    /// Evaluate the Jacobian at given state values
54    pub fn evaluate(&self, t: F, y: ArrayView1<F>) -> IntegrateResult<Array2<F>> {
55        let n = self.state_vars.len();
56        if y.len() != n {
57            return Err(IntegrateError::DimensionMismatch(format!(
58                "Expected {} states, got {}",
59                n,
60                y.len()
61            )));
62        }
63
64        // Build value map
65        let mut values = HashMap::new();
66        for (i, var) in self.state_vars.iter().enumerate() {
67            values.insert(var.clone(), y[i]);
68        }
69        if let Some(ref t_var) = self.time_var {
70            values.insert(t_var.clone(), t);
71        }
72
73        // Evaluate each element
74        let (rows, cols) = self.elements.dim();
75        let mut result = Array2::zeros((rows, cols));
76
77        for i in 0..rows {
78            for j in 0..cols {
79                result[[i, j]] = self.elements[[i, j]].evaluate(&values)?;
80            }
81        }
82
83        Ok(result)
84    }
85
86    /// Simplify all expressions in the Jacobian
87    pub fn simplify(&mut self) {
88        let (rows, cols) = self.elements.dim();
89        for i in 0..rows {
90            for j in 0..cols {
91                self.elements[[i, j]] = simplify(&self.elements[[i, j]]);
92            }
93        }
94    }
95}
96
97/// Generate a symbolic Jacobian from a vector of symbolic expressions
98///
99/// # Arguments
100/// * `expressions` - Vector of symbolic expressions representing the ODE system
101/// * `state_vars` - Variables with respect to which to differentiate
102/// * `time_var` - Optional time variable
103///
104/// # Returns
105/// A symbolic Jacobian matrix where J[i,j] = ∂f[i]/∂y[j]
106#[allow(dead_code)]
107pub fn generate_jacobian<F: IntegrateFloat>(
108    expressions: &[SymbolicExpression<F>],
109    state_vars: &[Variable],
110    time_var: Option<Variable>,
111) -> IntegrateResult<SymbolicJacobian<F>> {
112    let n = expressions.len();
113    let m = state_vars.len();
114
115    if n == 0 || m == 0 {
116        return Err(IntegrateError::ValueError(
117            "Empty expressions or state variables".to_string(),
118        ));
119    }
120
121    let mut jacobian = Array2::from_elem((n, m), SymbolicExpression::Constant(F::zero()));
122
123    // Compute partial derivatives
124    for (i, expr) in expressions.iter().enumerate() {
125        for (j, var) in state_vars.iter().enumerate() {
126            jacobian[[i, j]] = expr.differentiate(var);
127        }
128    }
129
130    Ok(SymbolicJacobian::new(
131        jacobian,
132        state_vars.to_vec(),
133        time_var,
134    ))
135}
136
137/// Builder for creating symbolic ODE systems
138pub struct SymbolicODEBuilder<F: IntegrateFloat> {
139    expressions: Vec<SymbolicExpression<F>>,
140    state_vars: Vec<Variable>,
141    time_var: Option<Variable>,
142}
143
144impl<F: IntegrateFloat> SymbolicODEBuilder<F> {
145    /// Create a new builder
146    pub fn new() -> Self {
147        SymbolicODEBuilder {
148            expressions: Vec::new(),
149            state_vars: Vec::new(),
150            time_var: None,
151        }
152    }
153
154    /// Set the number of state variables
155    pub fn with_state_vars(mut self, n: usize) -> Self {
156        self.state_vars = (0..n).map(|i| Variable::indexed("y", i)).collect();
157        self
158    }
159
160    /// Set custom state variable names
161    pub fn with_named_vars(mut self, names: Vec<String>) -> Self {
162        self.state_vars = names.into_iter().map(Variable::new).collect();
163        self
164    }
165
166    /// Enable time dependence
167    pub fn with_time(mut self) -> Self {
168        self.time_var = Some(Variable::new("t"));
169        self
170    }
171
172    /// Add an ODE expression
173    pub fn add_equation(&mut self, expr: SymbolicExpression<F>) -> &mut Self {
174        self.expressions.push(expr);
175        self
176    }
177
178    /// Build the symbolic Jacobian
179    pub fn build_jacobian(&self) -> IntegrateResult<SymbolicJacobian<F>> {
180        generate_jacobian(&self.expressions, &self.state_vars, self.time_var.clone())
181    }
182}
183
184impl<F: IntegrateFloat> Default for SymbolicODEBuilder<F> {
185    fn default() -> Self {
186        Self::new()
187    }
188}
189
190/// Example: Create a symbolic Jacobian for the Van der Pol oscillator
191#[allow(dead_code)]
192pub fn example_van_der_pol<F: IntegrateFloat>(mu: F) -> IntegrateResult<SymbolicJacobian<F>> {
193    use SymbolicExpression::*;
194
195    // Variables: y[0] = x, y[1] = x'
196    let y0 = Var(Variable::indexed("y", 0));
197    let y1 = Var(Variable::indexed("y", 1));
198
199    // Van der Pol equations:
200    // dy[0]/dt = y[1]
201    // dy[1]/dt = _mu * (1 - y[0]^2) * y[1] - y[0]
202
203    let expr1 = y1.clone();
204    let expr2 = Sub(
205        Box::new(Mul(
206            Box::new(Mul(
207                Box::new(Constant(mu)),
208                Box::new(Sub(
209                    Box::new(Constant(F::one())),
210                    Box::new(Pow(
211                        Box::new(y0.clone()),
212                        Box::new(Constant(F::from(2.0).unwrap())),
213                    )),
214                )),
215            )),
216            Box::new(y1),
217        )),
218        Box::new(y0),
219    );
220
221    SymbolicODEBuilder::new()
222        .with_state_vars(2)
223        .add_equation(expr1)
224        .add_equation(expr2)
225        .build_jacobian()
226}
227
228/// Example: Create a symbolic Jacobian for a stiff chemical reaction system
229#[allow(dead_code)]
230pub fn example_stiff_chemical<F: IntegrateFloat>() -> IntegrateResult<SymbolicJacobian<F>> {
231    // Robertson's chemical reaction problem (stiff ODE)
232    // dy1/dt = -0.04*y1 + 1e4*y2*y3
233    // dy2/dt = 0.04*y1 - 1e4*y2*y3 - 3e7*y2^2
234    // dy3/dt = 3e7*y2^2
235
236    let y1 = SymbolicExpression::indexedvar("y", 0);
237    let y2 = SymbolicExpression::indexedvar("y", 1);
238    let y3 = SymbolicExpression::indexedvar("y", 2);
239
240    let k1 = SymbolicExpression::constant(F::from(0.04).unwrap());
241    let k2 = SymbolicExpression::constant(F::from(1e4).unwrap());
242    let k3 = SymbolicExpression::constant(F::from(3e7).unwrap());
243
244    // Using operator overloading for cleaner syntax
245    let expr1 = -k1.clone() * y1.clone() + k2.clone() * y2.clone() * y3.clone();
246    let expr2 = k1 * y1 - k2 * y2.clone() * y3 - k3.clone() * y2.clone() * y2.clone();
247    let expr3 = k3 * y2.clone() * y2;
248
249    SymbolicODEBuilder::new()
250        .with_state_vars(3)
251        .add_equation(expr1)
252        .add_equation(expr2)
253        .add_equation(expr3)
254        .build_jacobian()
255}
256
257/// Example: Create a symbolic Jacobian for a predator-prey system with seasonal effects
258#[allow(dead_code)]
259pub fn example_seasonal_predator_prey<F: IntegrateFloat>() -> IntegrateResult<SymbolicJacobian<F>> {
260    // Lotka-Volterra with seasonal variation
261    // dx/dt = a*x*(1 + b*sin(2π*t)) - c*x*y
262    // dy/dt = -d*y + e*x*y
263
264    let x = indexed_var("y", 0);
265    let y = indexed_var("y", 1);
266    let t = var("t");
267
268    let a = constant(F::from(1.5).unwrap());
269    let b = constant(F::from(0.1).unwrap());
270    let c = constant(F::from(0.5).unwrap());
271    let d = constant(F::from(0.75).unwrap());
272    let e = constant(F::from(0.25).unwrap());
273    let two_pi = constant(F::from(std::f64::consts::TAU).unwrap());
274
275    // Seasonal growth term: 1 + b*sin(2π*t)
276    let seasonal = constant(F::one()) + b * SymbolicExpression::Sin(Box::new(two_pi * t));
277
278    let expr1 = a * x.clone() * seasonal - c * x.clone() * y.clone();
279    let expr2 = -d * y.clone() + e * x * y;
280
281    let mut builder = SymbolicODEBuilder::new().with_state_vars(2).with_time();
282    builder.add_equation(expr1);
283    builder.add_equation(expr2);
284
285    builder.build_jacobian()
286}
287
288/// Integration with the ODE solver autodiff module
289#[cfg(feature = "autodiff")]
290#[allow(dead_code)]
291pub fn create_autodiff_jacobian<F, Func>(
292    symbolic_jacobian: &SymbolicJacobian<F>,
293) -> impl Fn(F, ArrayView1<F>) -> IntegrateResult<Array2<F>>
294where
295    F: IntegrateFloat,
296    Func: Fn(F, ArrayView1<F>) -> IntegrateResult<ArrayView1<F>>,
297{
298    let jac = symbolic_jacobian.clone();
299    move |t: F, y: ArrayView1<F>| jac.evaluate(t, y)
300}
301
302impl<F: IntegrateFloat> Clone for SymbolicJacobian<F> {
303    fn clone(&self) -> Self {
304        SymbolicJacobian {
305            elements: self.elements.clone(),
306            state_vars: self.state_vars.clone(),
307            time_var: self.time_var.clone(),
308        }
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use crate::{
316        generate_jacobian,
317        SymbolicExpression::{Neg, Var},
318        Variable,
319    };
320    use scirs2_core::ndarray::ArrayView1;
321    use std::collections::HashMap;
322
323    #[test]
324    fn test_simple_jacobian() {
325        // System: dy/dt = -y
326        let y = Var(Variable::new("y"));
327        let expr = Neg(Box::new(y));
328
329        let jacobian = generate_jacobian(&[expr], &[Variable::new("y")], None).unwrap();
330
331        // Jacobian should be [[-1]]
332        let mut values = HashMap::new();
333        values.insert(Variable::new("y"), 1.0);
334
335        let j = jacobian.evaluate(0.0, ArrayView1::from(&[1.0])).unwrap();
336        assert_eq!(j.dim(), (1, 1));
337        assert!((j[[0, 0]] + 1.0_f64).abs() < 1e-10);
338    }
339}
scirs2_integrate/symbolic/jacobian.rs

scirs2_integrate/symbolic/
jacobian.rs
