Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
A simple harmonic oscillator simulation for testing out different numerical integration methods. Logs state at each step in CSV format.
// Just position and velocity along one dimension. We're assuming the particle has a mass of 1.
#[derive(Clone)]
struct State {
x: f64,
v: f64,
}
trait Integrator {
fn step(&mut self, state: State, sim: &Simulation, dt: f64) -> State;
}
trait Simulation {
fn acc(&self, state: &State) -> f64;
fn energy(&self, state: &State) -> f64;
fn exact(&self, time: f64) -> f64;
}
struct HarmonicOscillator {
k: f64,
start_state: State
}
impl Simulation for HarmonicOscillator {
fn acc(&self, state: &State) -> f64 {
// https://en.wikipedia.org/wiki/Hooke's_law#For_linear_springs
return -self.k * state.x;
}
fn energy(&self, state: &State) -> f64 {
// Potential + kinetic energy
return 0.5 * self.k * state.x * state.x + 0.5 * state.v * state.v;
}
fn exact(&self, time: f64) -> f64 {
// See https://en.wikipedia.org/wiki/Simple_harmonic_motion#Dynamics for the corresponding equations.
let start_state = &self.start_state;
let omega = self.k.sqrt();
let magnitude = (start_state.x * start_state.x + (start_state.v / omega) * (start_state.v / omega)).sqrt();
let phase_offset =
// Prevent divide by 0
if start_state.v == 0.0 {
3.14159265358979 / 2.0
} else {
(start_state.x * omega / start_state.v).atan()
};
magnitude * (omega * time + phase_offset).sin()
}
}
// https://en.wikipedia.org/wiki/Euler_method
#[allow(dead_code)]
struct ExplicitEuler {}
impl Integrator for ExplicitEuler {
fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
let a = sim.acc(&state);
state.x += state.v * dt;
state.v += a * dt;
state
}
}
// https://en.wikipedia.org/wiki/Semi-implicit_Euler_method
#[allow(dead_code)]
struct SemiImplicitEuler {}
impl Integrator for SemiImplicitEuler {
fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
let a = sim.acc(&state);
state.v += a * dt;
state.x += state.v * dt;
state
}
}
// https://en.wikipedia.org/wiki/Verlet_integration#Velocity_Verlet
#[allow(dead_code)]
struct VelocityVerlet {}
impl Integrator for VelocityVerlet {
fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
let a = sim.acc(&state);
state.x = state.x + state.v * dt + 0.5 * a * dt * dt;
let new_a = sim.acc(&state);
state.v = state.v + ((a + new_a) / 2.0) * dt;
state
}
}
// Integration method given in "Building a Better Jump" - a simplified version of Velocity Verlet that doesn't evaulate acceleration at the updated position.
// http://www.mathforgameprogrammers.com/gdc2016/GDC2016_Pittman_Kyle_BuildingABetterJump.pdf, https://www.youtube.com/watch?v=hG9SzQxaCm8 for a recording of the talk.
#[allow(dead_code)]
struct BABJ {}
impl Integrator for BABJ {
fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
let a = sim.acc(&state);
state.x = state.x + state.v * dt + 0.5 * a * dt * dt;
state.v = state.v + a * dt;
state
}
}
// https://en.wikipedia.org/wiki/Beeman%27s_algorithm
// I don't recommend using my implementation of this method! In my testing it gains energy over time, so I may have misinterpreted the equations. I wasn't able to find the issue with it though..
#[allow(dead_code)]
struct Beeman {
prev_state: State
}
impl Integrator for Beeman {
fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
let a_prev = sim.acc(&self.prev_state);
let a_cur = sim.acc(&state);
state.x += state.v * dt + (2.0 / 3.0) * a_cur * dt * dt - (1.0 / 6.0) * a_prev * dt * dt;
// predicted velocity, in case acceleration calculation depends on velocity
let original_v = state.v;
state.v += (3.0 / 2.0) * a_cur * dt - 0.5 * a_prev * dt;
let a_next = sim.acc(&state);
state.v = original_v + (5.0 / 12.0) * a_next * dt + (2.0 / 3.0) * a_cur * dt - (1.0 / 12.0) * a_prev * dt;
self.prev_state = state.clone();
state
}
}
struct Derivative {
dx: f64,
dv: f64
}
// https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
// This implementation is a almost direct translation of the one given in https://gafferongames.com/post/integration_basics/
struct RK4 {}
impl RK4 {
fn eval(state: &State, sim: &Simulation, dt: f64, derivative: &Derivative) -> Derivative {
let advanced_state = State {
x: state.x + derivative.dx * dt,
v: state.v + derivative.dv * dt,
};
Derivative {
dx: advanced_state.v,
dv: sim.acc(&advanced_state)
}
}
}
impl Integrator for RK4 {
fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
let a = RK4::eval(&state, sim, dt, &Derivative { dx: 0.0, dv: 0.0});
let b = RK4::eval(&state, sim, dt * 0.5, &a);
let c = RK4::eval(&state, sim, dt * 0.5, &b);
let d = RK4::eval(&state, sim, dt, &c);
let dxdt = (1.0 / 6.0) * (a.dx + 2.0 * (b.dx + c.dx) + d.dx);
let dvdt = (1.0 / 6.0) * (a.dv + 2.0 * (b.dv + c.dv) + d.dv);
state.x += dxdt * dt;
state.v += dvdt * dt;
state
}
}
fn main() {
// pick your delta time
let dt = 1.0/64.0;
let total_time = 128.0;
let steps = (total_time / dt) as usize;
let mut state = State {
x: 50.0,
v: 0.0
};
let sim = HarmonicOscillator {
k: 10.0, // spring stiffness
start_state: state.clone(),
};
// Uncomment the line for whatever integration method you want to use
let mut integrator = ExplicitEuler {};
// let mut integrator = SemiImplicitEuler {};
// let mut integrator = VelocityVerlet {};
// let mut integrator = BABJ {};
// let mut integrator = Beeman { prev_state: state.clone() };
// let mut integrator = RK4 {};
print_header();
for i in 0..steps {
state = integrator.step(state, &sim, dt);
log(((i+1) as f64) * dt, &state, &sim);
}
}
fn print_header() {
println!("step,x,exact,energy");
}
fn log(time: f64, state: &State, sim: &Simulation) {
let energy = sim.energy(state);
let exact_x = sim.exact(time);
println!("{},{},{},{}", time, state.x, exact_x, energy);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment