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

1#[cfg(test)]
2mod test;
3
4use crate::math::{
5    Derivative, Differentiable, Quantity, Scalar, Tensor, TensorVec,
6    integrate::{
7        ButcherTableau, EmbeddedTableau, Explicit, Flat, HermiteSegment, IntegrationError, Times,
8        VariableStep, interpolate_hermite,
9    },
10    interpolate::InterpolateSolution,
11};
12use crate::units::Time;
13use std::ops::{Div, Mul, Sub};
14
15pub(crate) mod bogacki_shampine;
16pub(crate) mod dormand_prince;
17pub(crate) mod verner_8;
18pub(crate) mod verner_9;
19
20/// Variable-step explicit integrators for ordinary differential equations.
21pub trait VariableStepExplicit<Y, U, V, T = Time>
22where
23    Self: Explicit<Y, U, V, T> + VariableStep<T>,
24    Y: Differentiable<T> + Tensor,
25    Derivative<Y, T>: Mul<Quantity<T>, Output = Y>,
26    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
27    for<'a> &'a Derivative<Y, T>:
28        Mul<Scalar, Output = Derivative<Y, T>> + Mul<Quantity<T>, Output = Y>,
29    U: TensorVec<Item = Y>,
30    V: TensorVec<Item = Derivative<Y, T>>,
31{
32    fn integrate_variable_step(
33        &self,
34        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
35        time: &[Quantity<T>],
36        initial_condition: Y,
37    ) -> Result<(Times<T>, U, V), IntegrationError>
38    where
39        Self: InterpolateSolution<Y, U, V, T>,
40    {
41        let t_0 = time[0];
42        let t_f = time[time.len() - 1];
43        if time.len() < 2 {
44            return Err(IntegrationError::LengthTimeLessThanTwo);
45        } else if t_0 >= t_f {
46            return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
47        }
48        let mut t = t_0;
49        let mut dt = t_f - t_0;
50        let mut k = vec![Derivative::<Y, T>::default(); Self::SLOPES];
51        k[0] = function(t, &initial_condition)?;
52        let mut t_sol = Times::new();
53        t_sol.push(t_0);
54        let mut y = initial_condition.clone();
55        let mut y_sol = U::new();
56        y_sol.push(initial_condition.clone());
57        let mut dydt_sol = V::new();
58        dydt_sol.push(k[0].clone());
59        let mut k_sol: Vec<V> = Vec::new();
60        let mut y_trial = Y::default();
61        while t < t_f {
62            match self.slopes_and_error(&mut function, &y, t, dt, &mut k, &mut y_trial) {
63                Ok(e) => {
64                    if let Err(error) = self.step(
65                        &mut function,
66                        &mut y,
67                        &mut t,
68                        &mut y_sol,
69                        &mut t_sol,
70                        &mut dydt_sol,
71                        &mut k_sol,
72                        &mut dt,
73                        &mut k,
74                        &y_trial,
75                        e,
76                    ) {
77                        dt *= self.dt_cut();
78                        if dt < self.dt_min() {
79                            return Err(IntegrationError::MinimumStepSizeUpstream(
80                                self.dt_min().value(),
81                                error,
82                                format!("{self:?}"),
83                            ));
84                        }
85                    } else {
86                        dt = dt.min(t_f - t);
87                        if dt < self.dt_min() && t < t_f {
88                            return Err(IntegrationError::MinimumStepSizeReached(
89                                self.dt_min().value(),
90                                format!("{self:?}"),
91                            ));
92                        }
93                    }
94                }
95                Err(error) => {
96                    dt *= self.dt_cut();
97                    if dt < self.dt_min() {
98                        return Err(IntegrationError::MinimumStepSizeUpstream(
99                            self.dt_min().value(),
100                            error,
101                            format!("{self:?}"),
102                        ));
103                    }
104                }
105            }
106        }
107        if time.len() > 2 {
108            let t_int = Times::from(time);
109            let (y_int, dydt_int) =
110                self.interpolate(&t_int, &t_sol, &y_sol, &dydt_sol, &k_sol, function)?;
111            Ok((t_int, y_int, dydt_int))
112        } else {
113            Ok((t_sol, y_sol, dydt_sol))
114        }
115    }
116    fn interpolate_variable_step(
117        time: &Times<T>,
118        tp: &Times<T>,
119        yp: &U,
120        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
121    ) -> Result<(U, V), IntegrationError> {
122        let mut dt;
123        let mut i;
124        let mut k = vec![Derivative::<Y, T>::default(); Self::SLOPES];
125        let mut t;
126        let mut y;
127        let mut y_int = U::new();
128        let mut dydt_int = V::new();
129        let mut y_trial = Y::default();
130        for time_k in time.iter() {
131            i = tp.iter().position(|tp_i| tp_i >= time_k).unwrap();
132            if time_k == &tp[i] {
133                t = tp[i];
134                y_trial = yp[i].clone();
135                dt = Quantity::default();
136            } else {
137                t = tp[i - 1];
138                y = &yp[i - 1];
139                dt = *time_k - t;
140                Self::slopes(&mut function, y, t, dt, &mut k, &mut y_trial)?;
141            }
142            dydt_int.push(function(t + dt, &y_trial)?);
143            y_int.push(y_trial.clone());
144        }
145        Ok((y_int, dydt_int))
146    }
147    /// Butcher tableau of this method's embedded pair.
148    type Tableau: EmbeddedTableau;
149    /// Runge–Kutta stages and the propagating solution.
150    ///
151    /// ```math
152    /// \mathbf{k}_i = \mathbf{f}\!\left(t + c_i h,\ \mathbf{y} + h \sum_{j<i} a_{ij}\, \mathbf{k}_j\right)
153    /// ,\qquad
154    /// \mathbf{y}_{n+1} = \mathbf{y} + h \textstyle\sum_i b_i\,\mathbf{k}_i
155    /// ```
156    fn slopes(
157        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
158        y: &Y,
159        t: Quantity<T>,
160        dt: Quantity<T>,
161        k: &mut [Derivative<Y, T>],
162        y_trial: &mut Y,
163    ) -> Result<(), String> {
164        let last = if Self::Tableau::FSAL {
165            Self::Tableau::STAGES - 1
166        } else {
167            k[0] = function(t, y)?;
168            Self::Tableau::STAGES
169        };
170        for i in 1..last.min(k.len()) {
171            let row = Self::Tableau::A[i];
172            let mut sigma = &k[0] * row[0];
173            for j in 1..i {
174                sigma += &k[j] * row[j];
175            }
176            let stage = &sigma * dt + y;
177            k[i] = function(t + Self::Tableau::C[i] * dt, &stage)?;
178        }
179        let mut sum = &k[0] * Self::Tableau::B[0];
180        for (b, slope) in Self::Tableau::B.iter().zip(k.iter()).skip(1) {
181            sum += slope * *b;
182        }
183        *y_trial = &sum * dt + y;
184        Ok(())
185    }
186    /// Embedded local-error estimate reduced through the error norm.
187    ///
188    /// ```math
189    /// e_{n+1} = \Big\Vert h \textstyle\sum_i d_i\,\mathbf{k}_i \Big\Vert
190    /// ```
191    fn error(&self, dt: Quantity<T>, k: &[Derivative<Y, T>]) -> Result<Scalar, String> {
192        let mut sum = &k[0] * Self::Tableau::D[0];
193        for (d, slope) in Self::Tableau::D.iter().zip(k.iter()).skip(1) {
194            sum += slope * *d;
195        }
196        Ok(self.error_norm().measure(&(&sum * dt)))
197    }
198    fn slopes_and_error(
199        &self,
200        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
201        y: &Y,
202        t: Quantity<T>,
203        dt: Quantity<T>,
204        k: &mut [Derivative<Y, T>],
205        y_trial: &mut Y,
206    ) -> Result<Scalar, String> {
207        Self::slopes(&mut function, y, t, dt, k, y_trial)?;
208        self.error(dt, k)
209    }
210    #[expect(clippy::too_many_arguments)]
211    fn step(
212        &self,
213        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
214        y: &mut Y,
215        t: &mut Quantity<T>,
216        y_sol: &mut U,
217        t_sol: &mut Times<T>,
218        dydt_sol: &mut V,
219        k_sol: &mut Vec<V>,
220        dt: &mut Quantity<T>,
221        k: &mut [Derivative<Y, T>],
222        y_trial: &Y,
223        e: Scalar,
224    ) -> Result<(), String> {
225        let tolerance = self
226            .abs_tol()
227            .max(self.rel_tol() * self.error_norm().measure(y_trial));
228        if e < self.abs_tol() || e < self.rel_tol() * self.error_norm().measure(y_trial) {
229            k_sol.push(k.iter().cloned().collect());
230            *t += *dt;
231            *y = y_trial.clone();
232            t_sol.push(*t);
233            y_sol.push(y.clone());
234            dydt_sol.push(function(*t, y)?);
235        }
236        self.time_step(e, tolerance, dt);
237        Ok(())
238    }
239    /// Provides the adaptive time step as a function of the error.
240    ///
241    /// ```math
242    /// h_{n+1} = \beta h \left(\frac{e_\mathrm{tol}}{e_{n+1}}\right)^{1/p}
243    /// ```
244    fn time_step(&self, error: Scalar, tolerance: Scalar, dt: &mut Quantity<T>) {
245        if error > 0.0 {
246            *dt *= (self.dt_beta() * (tolerance / error).powf(1.0 / self.dt_expn()))
247                .clamp(self.dt_cut(), self.dt_grow())
248        } else {
249            *dt *= self.dt_grow();
250        }
251    }
252}
253
254/// Free (dense-output) interpolant for explicit ordinary differential equation integrators.
255///
256/// Uses cubic Hermite interpolation over the accepted-step values and derivatives already
257/// computed during integration, so it requires no additional evaluations of the right-hand side
258/// function.
259pub trait FreeInterpolant<Y, U, V, T = Time>
260where
261    Self: VariableStepExplicit<Y, U, V, T>,
262    Y: Differentiable<T> + Div<Quantity<T>, Output = Derivative<Y, T>> + Tensor,
263    Derivative<Y, T>: Mul<Quantity<T>, Output = Y>,
264    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
265    for<'a> &'a Derivative<Y, T>:
266        Mul<Scalar, Output = Derivative<Y, T>> + Mul<Quantity<T>, Output = Y>,
267    U: TensorVec<Item = Y>,
268    V: TensorVec<Item = Derivative<Y, T>>,
269{
270    /// The state, via [`HermiteSegment`] built pointwise over `Flat<Y>` — the
271    /// state is a flat vector space, so this is the same cubic Hermite as
272    /// before, just built once instead of duplicated here.
273    fn interpolate_free(time: &Times<T>, tp: &Times<T>, yp: &U, dydtp: &V) -> (U, V) {
274        let segments: Vec<HermiteSegment<Flat<Y>, T>> = (1..tp.len())
275            .map(|i| {
276                let h = tp[i] - tp[i - 1];
277                HermiteSegment::new(
278                    tp[i - 1],
279                    h,
280                    yp[i - 1].clone(),
281                    &yp[i] - &yp[i - 1],
282                    &dydtp[i - 1] * h,
283                    &dydtp[i] * h,
284                )
285            })
286            .collect();
287        let y_int = interpolate_hermite::<Flat<Y>, U, T>(&segments, time.as_slice())
288            .expect("Flat::reconstruct is infallible");
289        let mut dydt_int = V::new();
290        for time_k in time.iter() {
291            let i = tp.iter().position(|tp_i| tp_i >= time_k).unwrap();
292            if time_k == &tp[i] {
293                dydt_int.push(dydtp[i].clone());
294            } else {
295                let t_0 = tp[i - 1];
296                let h = tp[i] - t_0;
297                let theta = (*time_k - t_0).value() / h.value();
298                let theta2 = theta * theta;
299                let dh00 = 6.0 * theta2 - 6.0 * theta;
300                let dh10 = 3.0 * theta2 - 4.0 * theta + 1.0;
301                let dh01 = -6.0 * theta2 + 6.0 * theta;
302                let dh11 = 3.0 * theta2 - 2.0 * theta;
303                dydt_int.push(
304                    (&yp[i - 1] * dh00 + &yp[i] * dh01) / h
305                        + &dydtp[i - 1] * dh10
306                        + &dydtp[i] * dh11,
307                );
308            }
309        }
310        (y_int, dydt_int)
311    }
312}
313
314/// First-same-as-last property for explicit ordinary differential equation integrators.
315pub trait VariableStepExplicitFirstSameAsLast<Y, U, V, T = Time>
316where
317    Self: VariableStepExplicit<Y, U, V, T>,
318    Y: Differentiable<T> + Tensor,
319    Derivative<Y, T>: Mul<Quantity<T>, Output = Y>,
320    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
321    for<'a> &'a Derivative<Y, T>:
322        Mul<Scalar, Output = Derivative<Y, T>> + Mul<Quantity<T>, Output = Y>,
323    U: TensorVec<Item = Y>,
324    V: TensorVec<Item = Derivative<Y, T>>,
325{
326    fn slopes_and_error_fsal(
327        &self,
328        mut function: impl FnMut(Quantity<T>, &Y) -> Result<Derivative<Y, T>, String>,
329        y: &Y,
330        t: Quantity<T>,
331        dt: Quantity<T>,
332        k: &mut [Derivative<Y, T>],
333        y_trial: &mut Y,
334    ) -> Result<Scalar, String> {
335        Self::slopes(&mut function, y, t, dt, k, y_trial)?;
336        k[Self::SLOPES - 1] = function(t + dt, y_trial)?;
337        self.error(dt, k)
338    }
339    #[expect(clippy::too_many_arguments)]
340    fn step_fsal(
341        &self,
342        y: &mut Y,
343        t: &mut Quantity<T>,
344        y_sol: &mut U,
345        t_sol: &mut Times<T>,
346        dydt_sol: &mut V,
347        k_sol: &mut Vec<V>,
348        dt: &mut Quantity<T>,
349        k: &mut [Derivative<Y, T>],
350        y_trial: &Y,
351        e: Scalar,
352    ) -> Result<(), String> {
353        let tolerance = self
354            .abs_tol()
355            .max(self.rel_tol() * self.error_norm().measure(y_trial));
356        if e < self.abs_tol() || e < self.rel_tol() * self.error_norm().measure(y_trial) {
357            k_sol.push(k.iter().cloned().collect());
358            k[0] = k[Self::SLOPES - 1].clone();
359            *t += *dt;
360            *y = y_trial.clone();
361            t_sol.push(*t);
362            y_sol.push(y.clone());
363            dydt_sol.push(k[0].clone());
364        }
365        self.time_step(e, tolerance, dt);
366        Ok(())
367    }
368}
conspire/math/integrate/ode/explicit/variable_step/mod.rs

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