Skip to main content

conspire/math/integrate/ode/explicit/fixed_step/
mod.rs

1#[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
22/// Fixed-step explicit integrators for ordinary differential equations.
23pub 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    /// Butcher tableau of this method.
32    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}