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

1#[cfg(test)]
2mod test;
3
4use crate::{
5    math::{
6        Derivative, Differentiable, Quantity, Scalar, Tensor, TensorVec,
7        integrate::{FixedStep, IntegrationError, OdeIntegrator, Times},
8        optimize::{EqualityConstraint, FirstOrderRootFinding, ZerothOrderRootFinding},
9    },
10    units::Time,
11};
12
13pub(crate) mod backward_euler;
14pub(crate) mod midpoint;
15pub(crate) mod trapezoidal;
16
17/// Implicit integrators for ordinary differential equations using zeroth-order root-finding.
18pub trait ImplicitZerothOrder<Y, U, V, T = Time>
19where
20    Self: FixedStep<T> + OdeIntegrator<Y, U>,
21    Y: Differentiable<T> + Tensor,
22    U: TensorVec<Item = Y>,
23    V: TensorVec<Item = Derivative<Y, T>>,
24{
25    #[doc = include_str!("doc.md")]
26    fn integrate(
27        &self,
28        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, IntegrationError>,
29        time: &[Quantity<T>],
30        initial_condition: Y,
31        solver: impl ZerothOrderRootFinding<Y, Y>,
32    ) -> Result<(Times<T>, U, V), IntegrationError> {
33        let t_0 = time[0];
34        let t_f = time[time.len() - 1];
35        let mut t_sol: Times<T>;
36        if time.len() < 2 {
37            return Err(IntegrationError::LengthTimeLessThanTwo);
38        } else if t_0 >= t_f {
39            return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
40        } else if time.len() == 2 {
41            if self.dt() <= Quantity::default() || self.dt().is_nan() {
42                return Err(IntegrationError::TimeStepNotSet(
43                    time[0].value(),
44                    time[1].value(),
45                    format!("{self:?}"),
46                ));
47            } else {
48                let max_steps = ((t_f - t_0).value() / self.dt().value()).ceil() as usize;
49                t_sol = (0..max_steps)
50                    .map(|step| t_0 + self.dt() * (step as Scalar))
51                    .collect();
52                t_sol.push(t_f);
53            }
54        } else {
55            t_sol = time.iter().copied().collect();
56        }
57        let mut index = 0;
58        let mut t = t_0;
59        let mut dt;
60        let mut t_trial;
61        let mut y = initial_condition.clone();
62        let mut y_sol = U::new();
63        y_sol.push(initial_condition.clone());
64        let mut dydt_sol = V::new();
65        dydt_sol.push(function(t, &y.clone())?);
66        let mut y_trial;
67        while t < t_f {
68            t_trial = t_sol[index + 1];
69            dt = t_trial - t;
70            y_trial = match solver.root(
71                |y_trial: &Y| self.residual(&mut function, t, &y, t_trial, y_trial, dt),
72                y.clone(),
73                EqualityConstraint::None,
74            ) {
75                Ok(solution) => solution,
76                Err(error) => {
77                    return Err(IntegrationError::upstream(error, self));
78                }
79            };
80            t = t_trial;
81            y = y_trial;
82            y_sol.push(y.clone());
83            dydt_sol.push(function(t, &y)?);
84            index += 1;
85        }
86        Ok((t_sol, y_sol, dydt_sol))
87    }
88    fn residual(
89        &self,
90        function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, IntegrationError>,
91        t: Quantity<T>,
92        y: &Y,
93        t_trial: Quantity<T>,
94        y_trial: &Y,
95        dt: Quantity<T>,
96    ) -> Result<Y, String>;
97}
98
99/// Implicit integrators for ordinary differential equations using first-order root-finding.
100pub trait ImplicitFirstOrder<Y, J, U, V, T = Time>
101where
102    Self: ImplicitZerothOrder<Y, U, V, T>,
103    Y: Differentiable<T> + Tensor,
104    J: Differentiable<T> + Tensor,
105    U: TensorVec<Item = Y>,
106    V: TensorVec<Item = Derivative<Y, T>>,
107{
108    #[doc = include_str!("doc.md")]
109    fn integrate(
110        &self,
111        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, IntegrationError>,
112        mut jacobian: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<J, T>, IntegrationError>,
113        time: &[Quantity<T>],
114        initial_condition: Y,
115        solver: impl FirstOrderRootFinding<Y, J, Y>,
116    ) -> Result<(Times<T>, U, V), IntegrationError> {
117        let t_0 = time[0];
118        let t_f = time[time.len() - 1];
119        let mut t_sol: Times<T>;
120        if time.len() < 2 {
121            return Err(IntegrationError::LengthTimeLessThanTwo);
122        } else if t_0 >= t_f {
123            return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
124        } else if time.len() == 2 {
125            if self.dt() <= Quantity::default() || self.dt().is_nan() {
126                return Err(IntegrationError::TimeStepNotSet(
127                    time[0].value(),
128                    time[1].value(),
129                    format!("{self:?}"),
130                ));
131            } else {
132                let max_steps = ((t_f - t_0).value() / self.dt().value()).ceil() as usize;
133                t_sol = (0..max_steps)
134                    .map(|step| t_0 + self.dt() * (step as Scalar))
135                    .collect();
136                t_sol.push(t_f);
137            }
138        } else {
139            t_sol = time.iter().copied().collect();
140        }
141        let mut index = 0;
142        let mut t = t_0;
143        let mut dt;
144        let mut t_trial;
145        let mut y = initial_condition.clone();
146        let mut y_sol = U::new();
147        y_sol.push(initial_condition.clone());
148        let mut dydt_sol = V::new();
149        dydt_sol.push(function(t, &y.clone())?);
150        let mut y_trial;
151        while t < t_f {
152            t_trial = t_sol[index + 1];
153            dt = t_trial - t;
154            y_trial = match solver.root(
155                |y_trial: &Y| self.residual(&mut function, t, &y, t_trial, y_trial, dt),
156                |y_trial: &Y| self.hessian(&mut jacobian, t, &y, t_trial, y_trial, dt),
157                y.clone(),
158                EqualityConstraint::None,
159                None,
160            ) {
161                Ok(solution) => solution,
162                Err(error) => {
163                    return Err(IntegrationError::upstream(error, self));
164                }
165            };
166            t = t_trial;
167            y = y_trial;
168            y_sol.push(y.clone());
169            dydt_sol.push(function(t, &y)?);
170            index += 1;
171        }
172        Ok((t_sol, y_sol, dydt_sol))
173    }
174    fn hessian(
175        &self,
176        jacobian: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<J, T>, IntegrationError>,
177        t: Quantity<T>,
178        y: &Y,
179        t_trial: Quantity<T>,
180        y_trial: &Y,
181        dt: Quantity<T>,
182    ) -> Result<J, String>;
183}
conspire/math/integrate/ode/implicit/mod.rs

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