Skip to main content

conspire/math/integrate/ode/implicit/rates/
mod.rs

1#[cfg(test)]
2mod test;
3
4use crate::{
5    math::{
6        Derivative, Differentiable, Quantity, Tensor, TensorVec,
7        integrate::{BackwardEuler, FixedStep, IntegrationError, Times, fixed_step_times},
8        optimize::{EqualityConstraints, RootFinding},
9        sparse::SparseSolver,
10    },
11    units::Time,
12};
13use std::{
14    fmt::Debug,
15    ops::{Div, Sub},
16};
17
18/// Implicit integrators for first-order systems.
19/// ```math
20/// C\dot{x} + f(t, x) = 0
21/// ```
22pub trait ImplicitRates<X, F, J, UX, UR, T = Time>
23where
24    Self: Debug,
25    X: Differentiable<T> + Tensor,
26    Derivative<X, T>: Tensor,
27    UX: TensorVec<Item = X>,
28    UR: TensorVec<Item = Derivative<X, T>>,
29{
30    #[doc = include_str!("doc.md")]
31    #[expect(clippy::too_many_arguments)]
32    fn integrate(
33        &self,
34        rate: impl FnMut(Quantity<T>, &X) -> Result<Derivative<X, T>, String>,
35        residual: impl FnMut(Quantity<T>, &X, &Derivative<X, T>) -> Result<F, String>,
36        tangent: impl FnMut(Quantity<T>, &X, &Derivative<X, T>, Quantity<T>) -> Result<J, String>,
37        solver: impl RootFinding<F, J, X>,
38        equality_constraint: impl Fn(Quantity<T>) -> EqualityConstraints,
39        sparse: Option<SparseSolver>,
40        time: &[Quantity<T>],
41        initial_state: X,
42    ) -> Result<(Times<T>, UX, UR), IntegrationError>;
43}
44
45impl<X, F, J, UX, UR, T> ImplicitRates<X, F, J, UX, UR, T> for BackwardEuler
46where
47    X: Differentiable<T> + Tensor,
48    Derivative<X, T>: Tensor,
49    for<'a> &'a X: Sub<&'a X, Output = X>,
50    for<'a> &'a X: Div<Quantity<T>, Output = Derivative<X, T>>,
51    UX: TensorVec<Item = X>,
52    UR: TensorVec<Item = Derivative<X, T>>,
53{
54    fn integrate(
55        &self,
56        mut rate: impl FnMut(Quantity<T>, &X) -> Result<Derivative<X, T>, String>,
57        mut residual: impl FnMut(Quantity<T>, &X, &Derivative<X, T>) -> Result<F, String>,
58        mut tangent: impl FnMut(Quantity<T>, &X, &Derivative<X, T>, Quantity<T>) -> Result<J, String>,
59        solver: impl RootFinding<F, J, X>,
60        equality_constraint: impl Fn(Quantity<T>) -> EqualityConstraints,
61        sparse: Option<SparseSolver>,
62        time: &[Quantity<T>],
63        initial_state: X,
64    ) -> Result<(Times<T>, UX, UR), IntegrationError> {
65        if time.len() < 2 {
66            return Err(IntegrationError::LengthTimeLessThanTwo);
67        }
68        let t_0 = time[0];
69        let t_f = time[time.len() - 1];
70        let dt_fixed = FixedStep::<T>::dt(self);
71        let t_sol: Times<T>;
72        if t_0 >= t_f {
73            return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
74        } else if time.len() == 2 {
75            if dt_fixed <= Quantity::default() || dt_fixed.is_nan() {
76                return Err(IntegrationError::TimeStepNotSet(
77                    time[0].value(),
78                    time[1].value(),
79                    format!("{self:?}"),
80                ));
81            } else {
82                t_sol = fixed_step_times(t_0, t_f, dt_fixed);
83            }
84        } else {
85            t_sol = time.iter().copied().collect();
86        }
87        let mut index = 0;
88        let mut t = t_0;
89        let mut x = initial_state;
90        let mut x_sol = UX::new();
91        x_sol.push(x.clone());
92        let mut r_sol = UR::new();
93        r_sol.push(rate(t, &x)?);
94        while t < t_f {
95            let t_next = t_sol[index + 1];
96            let dt = t_next - t;
97            let x_next = solver
98                .root(
99                    |x_trial: &X| residual(t_next, x_trial, &(&(x_trial - &x) / dt)),
100                    |x_trial: &X| tangent(t_next, x_trial, &(&(x_trial - &x) / dt), dt),
101                    x.clone(),
102                    equality_constraint(t_next),
103                    sparse.clone(),
104                )
105                .map_err(|error| IntegrationError::upstream(error, self))?;
106            r_sol.push(&(&x_next - &x) / dt);
107            x = x_next;
108            t = t_next;
109            x_sol.push(x.clone());
110            index += 1;
111        }
112        Ok((t_sol, x_sol, r_sol))
113    }
114}