conspire/math/integrate/ode/explicit/fixed_step/
mod.rs1#[cfg(test)]
2mod test;
3
4use crate::{
5 math::{
6 Derivative, Differentiable, Quantity, Scalar, Tensor, TensorVec,
7 integrate::{ButcherTableau, Explicit, FixedStep, IntegrationError, Times},
8 },
9 units::Time,
10};
11use std::ops::Mul;
12
13pub(crate) mod bogacki_shampine;
14pub(crate) mod dormand_prince;
15pub(crate) mod euler;
16pub(crate) mod heun;
17pub(crate) mod midpoint;
18pub(crate) mod ralston;
19pub(crate) mod verner_8;
20pub(crate) mod verner_9;
21
22pub trait FixedStepExplicit<Y, U, V, T = Time>
24where
25 Self: Explicit<Y, U, V, T> + FixedStep<T>,
26 Y: Differentiable<T> + Tensor,
27 for<'a> &'a Derivative<Y, T>: Mul<Quantity<T>, Output = Y>,
28 U: TensorVec<Item = Y>,
29 V: TensorVec<Item = Derivative<Y, T>>,
30{
31 type Tableau: ButcherTableau;
33 fn integrate_fixed_step(
34 &self,
35 mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
36 time: &[Quantity<T>],
37 initial_condition: Y,
38 ) -> Result<(Times<T>, U, V), IntegrationError> {
39 let t_0 = time[0];
40 let t_f = time[time.len() - 1];
41 let mut t_sol: Times<T>;
42 if time.len() < 2 {
43 return Err(IntegrationError::LengthTimeLessThanTwo);
44 } else if t_0 >= t_f {
45 return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
46 } else if time.len() == 2 {
47 if self.dt() <= Quantity::default() || self.dt().is_nan() {
48 return Err(IntegrationError::TimeStepNotSet(
49 time[0].value(),
50 time[1].value(),
51 format!("{self:?}"),
52 ));
53 } else {
54 let max_steps = ((t_f - t_0).value() / self.dt().value()).ceil() as usize;
55 t_sol = (0..max_steps)
56 .map(|step| t_0 + self.dt() * (step as Scalar))
57 .collect();
58 t_sol.push(t_f);
59 }
60 } else {
61 t_sol = time.iter().copied().collect();
62 }
63 let mut index = 0;
64 let mut t = t_0;
65 let mut dt;
66 let mut t_trial;
67 let mut k = vec![Derivative::<Y, T>::default(); Self::SLOPES];
68 k[0] = function(t, &initial_condition)?;
69 let mut y = initial_condition.clone();
70 let mut y_sol = U::new();
71 y_sol.push(initial_condition.clone());
72 let mut dydt_sol = V::new();
73 dydt_sol.push(function(t, &y.clone())?);
74 let mut y_trial = Y::default();
75 while t < t_f {
76 t_trial = t_sol[index + 1];
77 dt = t_trial - t;
78 if let Err(error) = self.step(&mut function, &y, t, dt, &mut k, &mut y_trial) {
79 return Err(IntegrationError::upstream(error, self));
80 } else {
81 t += dt;
82 y = y_trial.clone();
83 y_sol.push(y.clone());
84 dydt_sol.push(k[0].clone());
85 index += 1;
86 }
87 }
88 Ok((t_sol, y_sol, dydt_sol))
89 }
90 fn step(
91 &self,
92 mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
93 y: &Y,
94 t: Quantity<T>,
95 dt: Quantity<T>,
96 k: &mut [Derivative<Y, T>],
97 y_trial: &mut Y,
98 ) -> Result<(), String> {
99 k[0] = function(t, y)?;
100 for i in 1..Self::Tableau::STAGES.min(k.len()) {
101 let row = Self::Tableau::A[i];
102 let mut stage = &k[0] * (row[0] * dt);
103 for j in 1..i {
104 stage += &k[j] * (row[j] * dt);
105 }
106 k[i] = function(t + Self::Tableau::C[i] * dt, &(stage + y))?;
107 }
108 let mut sum = &k[0] * (Self::Tableau::B[0] * dt);
109 for (b, slope) in Self::Tableau::B.iter().zip(k.iter()).skip(1) {
110 sum += slope * (*b * dt);
111 }
112 *y_trial = sum + y;
113 Ok(())
114 }
115}