ode_integrate/solvers/
fixed_step.rs

1// Copyright: Jonas Pleyer
2// This Source Code Form is subject to the terms of the Mozilla Public
3// License, v. 2.0. If a copy of the MPL was not distributed with this
4// file, You can obtain one at https://mozilla.org/MPL/2.0/.
5
6use crate::concepts::steppers::*;
7use crate::concepts::ode_def::*;
8
9use std::ops::{Mul};
10
11pub enum FixedStepSolvers {
12    Euler,
13    Rk4,
14}
15
16/// # Euler stepper
17/// This stepper is meant as an example and is one that should generally not be used.
18/// Solving of the ODE is done via
19/// \begin{equation}
20///     y_1 = y_0 + dt f(y, t, p)
21/// \end{equation}
22pub struct Euler<'a, I, F, P, Err> {
23    ode_def: OdeDefinition<'a, I, F, P, Err>,
24    dy: I,
25}
26
27/// Create an Euler stepper from a OdeDefinition
28impl<'a, I, F, P, Err> From<OdeDefinition<'a, I, F, P, Err>> for Euler<'a, I, F, P, Err>
29where
30    I: Clone,
31    F: Copy,
32    P: Clone,
33{
34    fn from(input: OdeDefinition<'a, I, F, P, Err>) -> Euler<'a, I, F, P, Err> {
35        let dy = input.y0.clone();
36        Euler {
37            ode_def: input,
38            dy,
39        }
40    }
41}
42
43impl<'a, I, F, P, Err> Stepper<I, F, P, Err> for Euler<'a, I, F, P, Err> {
44    fn do_step_iter
45    (
46        &mut self,
47        y:  &mut I,
48        t:  &F,
49        dt: &F,
50        p:  &P
51    ) -> Result<(), Err>
52    where
53        for<'m>&'m mut I: IntoIterator<Item=&'m mut F>,
54        for<'m>&'m I: IntoIterator<Item=&'m F>,
55        F: FloatLikeType,
56    {
57        (self.ode_def.func)(y, &mut self.dy, t, p)?;
58        for (yi, dyi) in y.into_iter().zip(self.dy.into_iter()) {
59            *yi += *dt * *dyi;
60        }
61        Ok(())
62    }
63
64    fn do_step_add
65    (
66        &mut self,
67        y:  &mut I,
68        t:  &F,
69        dt: &F,
70        p:  &P
71    ) -> Result<(), Err>
72    where
73        I: MathVecLikeType<F>,
74        F: FloatLikeType + Mul<I,Output=I>,
75    {
76        (self.ode_def.func)(y, &mut self.dy, t, p)?;
77        *y += *dt * self.dy.clone();
78        Ok(())
79    }
80}
81
82/// # Runge-Kutta 4th order stepper
83/// The Runge-Kutta 4th order solving scheme works with the following equations
84/// First we compute the assisting variables
85/// \begin{equation}
86///     \begin{alignedat}{7}
87///         k_1 &= &&hf(t_0 &,& y_0&&)\\\\
88///         k_2 &= &&hf(t_0 + \tfrac{1}{2} &h ,& y_0 + \tfrac{1}{2} &k_1&)\\\\
89///         k_3 &= &&hf(t_0 + \tfrac{1}{2} &h ,& y_0 + \tfrac{1}{2} &k_2&)\\\\
90///         k_4 &= &&hf(t_0 + &h,& y_0 + &k_3&)
91///     \end{alignedat}
92/// \end{equation}
93/// and finally combine them with
94/// \begin{equation}
95///     y_1 = y_0 + \tfrac{1}{6} (k_1 + 2 k_2 + 2 k_3 + k_4).
96/// \end{equation}
97pub struct Rk4<'a, I, F, P, Err> {
98    ode_def: OdeDefinition<'a, I, F, P, Err>,
99    // Helper variables
100    k1: I,
101    k2: I,
102    k3: I,
103    k4: I,
104    dy: I,
105    ym: I,
106}
107
108/// Create a Rk4 stepper from a 
109impl<'a, I, F, P, Err> From<OdeDefinition<'a, I, F, P, Err>> for Rk4<'a, I, F, P, Err>
110where
111    I: Clone,
112    F: Copy,
113    P: Clone,
114{
115    fn from(input: OdeDefinition<'a, I, F, P, Err>) -> Rk4<'a, I, F, P, Err> {
116        let dy = input.y0.clone();
117        Rk4 {
118            ode_def: input,
119            k1: dy.clone(),
120            k2: dy.clone(),
121            k3: dy.clone(),
122            k4: dy.clone(),
123            dy: dy.clone(),
124            ym: dy,
125        }
126    }
127}
128
129// Implement the Rk4 stepper
130impl<'a, I, F, P, Err> Stepper<I, F, P, Err> for Rk4<'a, I, F, P, Err> {
131    fn do_step_iter
132    (
133        &mut self,
134        y:  &mut I,
135        t:  &F,
136        dt: &F,
137        p:  &P
138    ) -> Result<(), Err>
139    where
140        for<'m>&'m mut I: IntoIterator<Item=&'m mut F>,
141        for<'m>&'m I: IntoIterator<Item=&'m F>,
142        F: FloatLikeType,
143    {
144        (self.ode_def.func)(y, &mut self.dy, t, p)?;
145        for (yi, dyi) in y.into_iter().zip(self.dy.into_iter()) {
146            // TODO
147            // This is not a Runge-Kutta solver yet!
148            *yi += *dt * *dyi;
149        }
150        Ok(())
151    }
152
153    fn do_step_add
154    (
155        &mut self,
156        y:  &mut I,
157        t:  &F,
158        dt: &F,
159        p:  &P
160    ) -> Result<(), Err>
161    where
162        I: MathVecLikeType<F>,
163        F: FloatLikeType + Mul<I,Output=I>,
164    {
165        let half = F::from(1)/F::from(2);
166
167        (self.ode_def.func)(y, &mut self.dy, t, p)?;
168        // TODO
169        // Find more optimal version of this code
170        self.k1 = *dt * self.dy.clone();
171        self.ym = y.clone() + half * self.k1.clone();
172        (self.ode_def.func)(&self.ym, &mut self.dy, &(*t + half * *dt), p)?;
173        self.k2 = *dt * self.dy.clone();
174        self.ym = y.clone() + half * self.k2.clone();
175        (self.ode_def.func)(&self.ym, &mut self.dy, &(*t + half * *dt), p)?;
176        self.k3 = *dt * self.dy.clone();
177        self.ym = y.clone() + self.k3.clone();
178        (self.ode_def.func)(&self.ym, &mut self.dy, &(*t + *dt), p)?;
179        self.k4 = *dt * self.dy.clone();
180        *y += half / F::from(3) * (self.k1.clone() + F::from(2) * self.k2.clone() + F::from(2) * self.k3.clone() + self.k4.clone());
181        Ok(())
182    }
183}