use dualis_units::Time;
pub trait State: Clone {
fn axpy(&mut self, a: f64, other: &Self);
fn scale(&mut self, a: f64);
fn zeros_like(&self) -> Self;
}
pub trait Dynamics {
type S: State;
fn derivative(&self, s: &Self::S, t: Time) -> Self::S;
}
pub trait Newtonian {
type Coords: State;
fn acceleration(&self, x: &Self::Coords, t: Time) -> Self::Coords;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Integrator {
Euler,
Midpoint,
Rk4,
}
impl Integrator {
pub fn step<D: Dynamics>(&self, system: &D, s: &D::S, t: Time, dt: Time) -> D::S {
let h = dt.to_si();
match self {
Integrator::Euler => {
let mut next = s.clone();
next.axpy(h, &system.derivative(s, t));
next
}
Integrator::Midpoint => {
let k1 = system.derivative(s, t);
let mut mid = s.clone();
mid.axpy(h / 2.0, &k1);
let k2 = system.derivative(&mid, t + dt / 2.0);
let mut next = s.clone();
next.axpy(h, &k2);
next
}
Integrator::Rk4 => {
let k1 = system.derivative(s, t);
let mut y = s.clone();
y.axpy(h / 2.0, &k1);
let k2 = system.derivative(&y, t + dt / 2.0);
let mut y = s.clone();
y.axpy(h / 2.0, &k2);
let k3 = system.derivative(&y, t + dt / 2.0);
let mut y = s.clone();
y.axpy(h, &k3);
let k4 = system.derivative(&y, t + dt);
let mut slope = k1;
slope.axpy(2.0, &k2);
slope.axpy(2.0, &k3);
slope.axpy(1.0, &k4);
slope.scale(1.0 / 6.0);
let mut next = s.clone();
next.axpy(h, &slope);
next
}
}
}
pub fn advance<D: Dynamics>(&self, system: &D, s: &D::S, t: Time, dt: Time, n: u32) -> D::S {
let mut state = s.clone();
let mut now = t;
for _ in 0..n {
state = self.step(system, &state, now, dt);
now += dt;
}
state
}
}
pub fn velocity_verlet<N: Newtonian>(
system: &N,
x: &mut N::Coords,
v: &mut N::Coords,
t: Time,
dt: Time,
) {
let h = dt.to_si();
let a0 = system.acceleration(x, t);
v.axpy(h / 2.0, &a0);
x.axpy(h, v);
let a1 = system.acceleration(x, t + dt);
v.axpy(h / 2.0, &a1);
}
pub fn substeps_for(dt: Time, limit: Time) -> u32 {
let (dt, limit) = (dt.to_si(), limit.to_si());
if !limit.is_finite() || limit <= 0.0 || dt <= 0.0 {
return 1;
}
let n = (dt / limit).ceil();
if !n.is_finite() || n < 1.0 {
1
} else {
n.min(u32::MAX as f64) as u32
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Clone, Debug, PartialEq)]
struct Pair(f64, f64);
impl State for Pair {
fn axpy(&mut self, a: f64, other: &Self) {
self.0 += a * other.0;
self.1 += a * other.1;
}
fn scale(&mut self, a: f64) {
self.0 *= a;
self.1 *= a;
}
fn zeros_like(&self) -> Self {
Pair(0.0, 0.0)
}
}
struct Spring;
impl Dynamics for Spring {
type S = Pair;
fn derivative(&self, s: &Pair, _t: Time) -> Pair {
Pair(s.1, -s.0)
}
}
#[derive(Clone, Debug)]
struct Scalar(f64);
impl State for Scalar {
fn axpy(&mut self, a: f64, other: &Self) {
self.0 += a * other.0;
}
fn scale(&mut self, a: f64) {
self.0 *= a;
}
fn zeros_like(&self) -> Self {
Scalar(0.0)
}
}
impl Newtonian for Spring {
type Coords = Scalar;
fn acceleration(&self, x: &Scalar, _t: Time) -> Scalar {
Scalar(-x.0)
}
}
fn energy(x: f64, v: f64) -> f64 {
(x * x + v * v) / 2.0
}
#[test]
fn each_integrator_shows_its_order() {
let exact = |t: f64| t.cos();
let error_at = |method: Integrator, steps: u32| {
let dt = Time::s(1.0 / steps as f64);
let end = method.advance(&Spring, &Pair(1.0, 0.0), Time::ZERO, dt, steps);
(end.0 - exact(1.0)).abs()
};
for (method, expected_ratio) in [
(Integrator::Euler, 2.0),
(Integrator::Midpoint, 4.0),
(Integrator::Rk4, 16.0),
] {
let coarse = error_at(method, 200);
let fine = error_at(method, 400);
let ratio = coarse / fine;
assert!(
(ratio - expected_ratio).abs() / expected_ratio < 0.15,
"{method:?}: halving the step changed the error by {ratio:.2}, \
expected about {expected_ratio}"
);
}
}
#[test]
fn only_the_symplectic_integrator_keeps_its_energy() {
const STEPS: u32 = 200_000;
let dt = Time::s(0.2);
let rk4 = Integrator::Rk4.advance(&Spring, &Pair(1.0, 0.0), Time::ZERO, dt, STEPS);
let rk4_drift = (energy(rk4.0, rk4.1) - 0.5).abs() / 0.5;
let (mut x, mut v) = (Scalar(1.0), Scalar(0.0));
let mut t = Time::ZERO;
let mut worst = 0.0f64;
for _ in 0..STEPS {
velocity_verlet(&Spring, &mut x, &mut v, t, dt);
t += dt;
worst = worst.max((energy(x.0, v.0) - 0.5).abs() / 0.5);
}
assert!(
rk4_drift > 0.05,
"RK4 should have leaked energy over {STEPS} steps, drift {rk4_drift:.3e}"
);
assert!(
worst < 0.05,
"velocity-Verlet's energy error should stay bounded, worst {worst:.3e}"
);
assert!(
worst < rk4_drift / 5.0,
"the symplectic method should hold energy far better: {worst:.3e} vs {rk4_drift:.3e}"
);
}
#[test]
fn integration_is_bit_reproducible() {
let dt = Time::s(0.01);
let once = Integrator::Rk4.advance(&Spring, &Pair(1.0, 0.0), Time::ZERO, dt, 500);
let twice = {
let half = Integrator::Rk4.advance(&Spring, &Pair(1.0, 0.0), Time::ZERO, dt, 250);
Integrator::Rk4.advance(&Spring, &half, Time::s(2.5), dt, 250)
};
assert_eq!(once, twice, "restarting mid-run must change nothing");
let again = Integrator::Rk4.advance(&Spring, &Pair(1.0, 0.0), Time::ZERO, dt, 500);
assert_eq!(once, again);
}
#[test]
fn substep_counts_are_deterministic_integers() {
assert_eq!(substeps_for(Time::s(1.0), Time::s(0.3)), 4);
assert_eq!(substeps_for(Time::s(1.0), Time::s(0.5)), 2);
assert_eq!(substeps_for(Time::s(1.0), Time::s(2.0)), 1);
assert_eq!(substeps_for(Time::s(1.0), Time::from_si(f64::INFINITY)), 1);
assert_eq!(substeps_for(Time::s(1.0), Time::ZERO), 1);
assert_eq!(substeps_for(Time::s(1.0), Time::s(-1.0)), 1);
assert_eq!(substeps_for(Time::ZERO, Time::s(0.1)), 1);
for (dt, limit) in [(1.0, 0.3), (7.3, 0.11), (1e-3, 1e-7)] {
let n = substeps_for(Time::s(dt), Time::s(limit));
assert!(n as f64 * limit >= dt - 1e-12, "{n} x {limit} < {dt}");
}
}
}