Skip to main content

conspire/math/integrate/ode/dynamics/explicit/velocity_verlet/
mod.rs

1#[cfg(test)]
2mod test;
3
4use crate::math::{
5    Derivative, Differentiable, Quantity, Scalar, Tensor, TensorVec,
6    integrate::{ExplicitDynamics, FixedStep, IntegrationError, Times},
7};
8use std::ops::Mul;
9
10/// The extent of the stability region of the method along the imaginary axis.
11const EXTENT: Scalar = 2.0;
12
13#[doc = include_str!("doc.md")]
14#[derive(Debug, Default)]
15pub struct VelocityVerlet {
16    /// Fixed value for the time step.
17    dt: Scalar,
18}
19
20impl<T> FixedStep<T> for VelocityVerlet {
21    fn dt(&self) -> Quantity<T> {
22        Quantity::new(self.dt)
23    }
24}
25
26impl VelocityVerlet {
27    /// Constructs the integrator with a fixed time step.
28    pub fn new<T>(dt: Quantity<T>) -> Self {
29        Self { dt: dt.value() }
30    }
31    fn integrate_checked<X, UX, UV, UA, T>(
32        &self,
33        mut function: impl FnMut(
34            Quantity<T>,
35            &X,
36            &Derivative<X, T>,
37        ) -> Result<Derivative<Derivative<X, T>, T>, String>,
38        mut check: impl FnMut(
39            Quantity<T>,
40            &X,
41            &Derivative<X, T>,
42            Quantity<T>,
43        ) -> Result<(), IntegrationError>,
44        time: &[Quantity<T>],
45        initial_position: X,
46        initial_velocity: Derivative<X, T>,
47    ) -> Result<(Times<T>, UX, UV, UA), IntegrationError>
48    where
49        X: Differentiable<T> + Tensor,
50        Derivative<X, T>: Differentiable<T> + Tensor,
51        for<'a> &'a Derivative<X, T>: Mul<Quantity<T>, Output = X>,
52        for<'a> &'a Derivative<Derivative<X, T>, T>: Mul<Quantity<T>, Output = Derivative<X, T>>,
53        UX: TensorVec<Item = X>,
54        UV: TensorVec<Item = Derivative<X, T>>,
55        UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
56    {
57        if time.len() < 2 {
58            return Err(IntegrationError::LengthTimeLessThanTwo);
59        }
60        let t_0 = time[0];
61        let t_f = time[time.len() - 1];
62        let dt_fixed = FixedStep::<T>::dt(self);
63        let mut t_sol: Times<T>;
64        if t_0 >= t_f {
65            return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
66        } else if time.len() == 2 {
67            if dt_fixed <= Quantity::default() || dt_fixed.is_nan() {
68                return Err(IntegrationError::TimeStepNotSet(
69                    time[0].value(),
70                    time[1].value(),
71                    format!("{self:?}"),
72                ));
73            } else {
74                let max_steps = ((t_f - t_0).value() / dt_fixed.value()).ceil() as usize;
75                t_sol = (0..max_steps)
76                    .map(|step| t_0 + dt_fixed * (step as Scalar))
77                    .collect();
78                t_sol.push(t_f);
79            }
80        } else {
81            t_sol = time.iter().copied().collect();
82        }
83        let mut index = 0;
84        let mut t = t_0;
85        let mut x = initial_position;
86        let mut v = initial_velocity;
87        let mut a = function(t, &x, &v)?;
88        let mut x_sol = UX::new();
89        x_sol.push(x.clone());
90        let mut v_sol = UV::new();
91        v_sol.push(v.clone());
92        let mut a_sol = UA::new();
93        a_sol.push(a.clone());
94        while t < t_f {
95            let t_next = t_sol[index + 1];
96            let dt = t_next - t;
97            check(t, &x, &v, dt)?;
98            let v_half = &a * (dt * 0.5) + &v;
99            let x_next = &v_half * dt + &x;
100            a = function(t_next, &x_next, &v_half)
101                .map_err(|error| IntegrationError::upstream(error, self))?;
102            v = &a * (dt * 0.5) + &v_half;
103            x = x_next;
104            t = t_next;
105            x_sol.push(x.clone());
106            v_sol.push(v.clone());
107            a_sol.push(a.clone());
108            index += 1;
109        }
110        Ok((t_sol, x_sol, v_sol, a_sol))
111    }
112}
113
114impl<X, UX, UV, UA, T> ExplicitDynamics<X, UX, UV, UA, T> for VelocityVerlet
115where
116    X: Differentiable<T> + Tensor,
117    Derivative<X, T>: Differentiable<T> + Tensor,
118    for<'a> &'a Derivative<X, T>: Mul<Quantity<T>, Output = X>,
119    for<'a> &'a Derivative<Derivative<X, T>, T>: Mul<Quantity<T>, Output = Derivative<X, T>>,
120    UX: TensorVec<Item = X>,
121    UV: TensorVec<Item = Derivative<X, T>>,
122    UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
123{
124    fn integrate(
125        &self,
126        function: impl FnMut(
127            Quantity<T>,
128            &X,
129            &Derivative<X, T>,
130        ) -> Result<Derivative<Derivative<X, T>, T>, String>,
131        time: &[Quantity<T>],
132        initial_position: X,
133        initial_velocity: Derivative<X, T>,
134    ) -> Result<(Times<T>, UX, UV, UA), IntegrationError> {
135        self.integrate_checked(
136            function,
137            |_, _, _, _| Ok(()),
138            time,
139            initial_position,
140            initial_velocity,
141        )
142    }
143    fn integrate_bounded(
144        &self,
145        function: impl FnMut(
146            Quantity<T>,
147            &X,
148            &Derivative<X, T>,
149        ) -> Result<Derivative<Derivative<X, T>, T>, String>,
150        mut bound: impl FnMut(Quantity<T>, &X, &Derivative<X, T>) -> Result<Quantity<T>, String>,
151        safety: Scalar,
152        time: &[Quantity<T>],
153        initial_position: X,
154        initial_velocity: Derivative<X, T>,
155    ) -> Result<(Times<T>, UX, UV, UA), IntegrationError> {
156        if !(safety > 0.0 && safety <= 1.0) {
157            return Err(IntegrationError::InvalidSafetyFactor(safety));
158        }
159        self.integrate_checked(
160            function,
161            |t, x, v, dt| {
162                let limit = bound(t, x, v)? * (EXTENT * safety);
163                if dt > limit {
164                    Err(IntegrationError::UnstableTimeStep(
165                        dt.value(),
166                        limit.value(),
167                        format!("{self:?}"),
168                    ))
169                } else {
170                    Ok(())
171                }
172            },
173            time,
174            initial_position,
175            initial_velocity,
176        )
177    }
178}