Last active
March 28, 2020 23:59
-
-
Save KeyMaster-/8aa87082083e048ae1a30e7fd754184f to your computer and use it in GitHub Desktop.
A simple harmonic oscillator simulation for testing out different numerical integration methods. Logs state at each step in CSV format.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// 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