conspire/math/integrate/ode/implicit/rates/
mod.rs1#[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
18pub 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}