pub const G: f64 = 6.674e-11;
#[derive(Debug, Clone)]
pub struct Body {
pub name: String,
pub mass: f64,
pub position: [f64; 3],
pub velocity: [f64; 3],
}
impl Body {
pub fn kinetic_energy(&self) -> f64 {
let v_sq: f64 = self.velocity.iter().map(|v| v * v).sum();
0.5 * self.mass * v_sq
}
}
#[derive(Debug, Clone)]
pub struct OrbitalElements {
pub semi_major_axis: f64,
pub eccentricity: f64,
pub inclination_deg: f64,
pub period_s: f64,
}
pub struct NBodySimulator {
bodies: Vec<Body>,
}
impl NBodySimulator {
pub fn new() -> Self {
Self { bodies: Vec::new() }
}
pub fn add_body(&mut self, body: Body) {
self.bodies.push(body);
}
pub fn body_count(&self) -> usize {
self.bodies.len()
}
pub fn body(&self, name: &str) -> Option<&Body> {
self.bodies.iter().find(|b| b.name == name)
}
pub fn body_mut(&mut self, name: &str) -> Option<&mut Body> {
self.bodies.iter_mut().find(|b| b.name == name)
}
pub fn step(&mut self, dt_s: f64) {
let n = self.bodies.len();
let mut accelerations: Vec<[f64; 3]> = vec![[0.0; 3]; n];
for (i, acc_i) in accelerations.iter_mut().enumerate() {
for (j, body_j) in self.bodies.iter().enumerate() {
if i == j {
continue;
}
let dx = body_j.position[0] - self.bodies[i].position[0];
let dy = body_j.position[1] - self.bodies[i].position[1];
let dz = body_j.position[2] - self.bodies[i].position[2];
let r_sq = dx * dx + dy * dy + dz * dz;
if r_sq < 1.0 {
continue;
}
let r = r_sq.sqrt();
let factor = G * body_j.mass / (r_sq * r);
acc_i[0] += factor * dx;
acc_i[1] += factor * dy;
acc_i[2] += factor * dz;
}
}
for (body, acc) in self.bodies.iter_mut().zip(accelerations.iter()) {
for (vel_k, acc_k) in body.velocity.iter_mut().zip(acc.iter()) {
*vel_k += acc_k * dt_s;
}
for (pos_k, vel_k) in body.position.iter_mut().zip(body.velocity.iter()) {
*pos_k += vel_k * dt_s;
}
}
}
pub fn step_n(&mut self, dt_s: f64, n: usize) {
for _ in 0..n {
self.step(dt_s);
}
}
pub fn kinetic_energy(&self) -> f64 {
self.bodies.iter().map(|b| b.kinetic_energy()).sum()
}
pub fn potential_energy(&self) -> f64 {
let n = self.bodies.len();
let mut u = 0.0;
for i in 0..n {
for j in (i + 1)..n {
let dx = self.bodies[i].position[0] - self.bodies[j].position[0];
let dy = self.bodies[i].position[1] - self.bodies[j].position[1];
let dz = self.bodies[i].position[2] - self.bodies[j].position[2];
let r = (dx * dx + dy * dy + dz * dz).sqrt();
if r > 0.0 {
u -= G * self.bodies[i].mass * self.bodies[j].mass / r;
}
}
}
u
}
pub fn total_energy(&self) -> f64 {
self.kinetic_energy() + self.potential_energy()
}
pub fn bodies(&self) -> &[Body] {
&self.bodies
}
pub fn centre_of_mass(&self) -> [f64; 3] {
let total_mass: f64 = self.bodies.iter().map(|b| b.mass).sum();
if total_mass == 0.0 {
return [0.0; 3];
}
let mut com = [0.0f64; 3];
for body in &self.bodies {
for (com_k, pos_k) in com.iter_mut().zip(body.position.iter()) {
*com_k += body.mass * pos_k;
}
}
for com_k in &mut com {
*com_k /= total_mass;
}
com
}
pub fn total_momentum(&self) -> [f64; 3] {
let mut p = [0.0f64; 3];
for body in &self.bodies {
for (p_k, vel_k) in p.iter_mut().zip(body.velocity.iter()) {
*p_k += body.mass * vel_k;
}
}
p
}
}
impl Default for NBodySimulator {
fn default() -> Self {
Self::new()
}
}
pub struct OrbitalCalculator;
impl OrbitalCalculator {
pub fn circular_orbit_velocity(central_mass: f64, radius_m: f64) -> f64 {
(G * central_mass / radius_m).sqrt()
}
pub fn orbital_period(central_mass: f64, semi_major_axis_m: f64) -> f64 {
let a3 = semi_major_axis_m.powi(3);
2.0 * std::f64::consts::PI * (a3 / (G * central_mass)).sqrt()
}
pub fn escape_velocity(central_mass: f64, radius_m: f64) -> f64 {
(2.0 * G * central_mass / radius_m).sqrt()
}
pub fn vis_viva(central_mass: f64, r: f64, a: f64) -> f64 {
let v_sq = G * central_mass * (2.0 / r - 1.0 / a);
if v_sq >= 0.0 {
v_sq.sqrt()
} else {
f64::NAN
}
}
pub fn specific_orbital_energy(central_mass: f64, semi_major_axis_m: f64) -> f64 {
-G * central_mass / (2.0 * semi_major_axis_m)
}
pub fn apoapsis(periapsis_m: f64, eccentricity: f64) -> Option<f64> {
if eccentricity >= 1.0 {
None } else {
Some(periapsis_m * (1.0 + eccentricity) / (1.0 - eccentricity))
}
}
pub fn periapsis(apoapsis_m: f64, eccentricity: f64) -> f64 {
apoapsis_m * (1.0 - eccentricity) / (1.0 + eccentricity)
}
pub fn semi_major_axis(apoapsis_m: f64, periapsis_m: f64) -> f64 {
(apoapsis_m + periapsis_m) / 2.0
}
pub fn hill_sphere(semi_major_axis_m: f64, body_mass: f64, parent_mass: f64) -> f64 {
semi_major_axis_m * (body_mass / (3.0 * parent_mass)).cbrt()
}
}
#[cfg(test)]
mod tests {
use super::*;
const SOLAR_MASS: f64 = 1.989e30; const EARTH_MASS: f64 = 5.972e24; const AU: f64 = 1.496e11; const EARTH_ORBITAL_V: f64 = 29_780.0; const ONE_YEAR_S: f64 = 3.156e7;
#[test]
fn test_circular_orbit_velocity_earth() {
let v = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU);
assert!(
(v - EARTH_ORBITAL_V).abs() / EARTH_ORBITAL_V < 0.01,
"circular velocity {v} deviates too much from {EARTH_ORBITAL_V}"
);
}
#[test]
fn test_circular_orbit_velocity_increases_with_mass() {
let v1 = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU);
let v2 = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS * 4.0, AU);
assert!((v2 - 2.0 * v1).abs() / v1 < 1e-9);
}
#[test]
fn test_circular_orbit_velocity_decreases_with_radius() {
let v1 = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU);
let v2 = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU * 4.0);
assert!((v2 - v1 / 2.0).abs() / v1 < 1e-9);
}
#[test]
fn test_circular_orbit_velocity_positive() {
assert!(OrbitalCalculator::circular_orbit_velocity(1e20, 1e9) > 0.0);
}
#[test]
fn test_orbital_period_earth_approx_one_year() {
let t = OrbitalCalculator::orbital_period(SOLAR_MASS, AU);
assert!(
(t - ONE_YEAR_S).abs() / ONE_YEAR_S < 0.02,
"period {t} s deviates too much from {ONE_YEAR_S} s"
);
}
#[test]
fn test_orbital_period_kepler_third_law_ratio() {
let t1 = OrbitalCalculator::orbital_period(SOLAR_MASS, AU);
let t2 = OrbitalCalculator::orbital_period(SOLAR_MASS, AU * 4.0);
let ratio_t = t2 / t1;
assert!((ratio_t - 8.0).abs() < 1e-6);
}
#[test]
fn test_orbital_period_positive() {
assert!(OrbitalCalculator::orbital_period(1e20, 1e9) > 0.0);
}
#[test]
fn test_escape_velocity_is_sqrt2_times_circular() {
let v_c = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU);
let v_e = OrbitalCalculator::escape_velocity(SOLAR_MASS, AU);
assert!((v_e - (2.0_f64).sqrt() * v_c).abs() / v_c < 1e-9);
}
#[test]
fn test_escape_velocity_earth_surface() {
let v_esc = OrbitalCalculator::escape_velocity(EARTH_MASS, 6.371e6);
assert!(
(v_esc - 11_186.0).abs() / 11_186.0 < 0.01,
"escape velocity {v_esc} m/s deviates too much from 11,186 m/s"
);
}
#[test]
fn test_escape_velocity_greater_than_circular() {
let v_c = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU);
let v_e = OrbitalCalculator::escape_velocity(SOLAR_MASS, AU);
assert!(v_e > v_c);
}
#[test]
fn test_vis_viva_circular_equals_circular_velocity() {
let v_vv = OrbitalCalculator::vis_viva(SOLAR_MASS, AU, AU);
let v_c = OrbitalCalculator::circular_orbit_velocity(SOLAR_MASS, AU);
assert!((v_vv - v_c).abs() / v_c < 1e-9);
}
#[test]
fn test_vis_viva_at_periapsis() {
let a = 1.5 * AU;
let e = 0.5;
let r_peri = a * (1.0 - e);
let v = OrbitalCalculator::vis_viva(SOLAR_MASS, r_peri, a);
assert!(v > 0.0);
}
#[test]
fn test_vis_viva_at_apoapsis_less_than_periapsis_speed() {
let a = 1.5 * AU;
let e = 0.5;
let r_peri = a * (1.0 - e);
let r_apo = a * (1.0 + e);
let v_peri = OrbitalCalculator::vis_viva(SOLAR_MASS, r_peri, a);
let v_apo = OrbitalCalculator::vis_viva(SOLAR_MASS, r_apo, a);
assert!(
v_peri > v_apo,
"speed at periapsis should be greater than at apoapsis"
);
}
#[test]
fn test_vis_viva_parabolic_orbit_at_peri_equals_escape() {
let large_a = 1e30;
let v_vv = OrbitalCalculator::vis_viva(SOLAR_MASS, AU, large_a);
let v_esc = OrbitalCalculator::escape_velocity(SOLAR_MASS, AU);
assert!((v_vv - v_esc).abs() / v_esc < 1e-3);
}
#[test]
fn test_apoapsis_circular_orbit() {
let apo = OrbitalCalculator::apoapsis(AU, 0.0).expect("circular has apoapsis");
assert!((apo - AU).abs() < 1.0);
}
#[test]
fn test_apoapsis_elliptical() {
let apo = OrbitalCalculator::apoapsis(AU, 0.5).expect("elliptical has apoapsis");
assert!((apo - 3.0 * AU).abs() / AU < 1e-9);
}
#[test]
fn test_apoapsis_hyperbolic_is_none() {
assert!(OrbitalCalculator::apoapsis(AU, 1.0).is_none());
assert!(OrbitalCalculator::apoapsis(AU, 1.5).is_none());
}
#[test]
fn test_periapsis_from_apoapsis() {
let peri = OrbitalCalculator::periapsis(3.0 * AU, 0.5);
assert!((peri - AU).abs() / AU < 1e-9);
}
#[test]
fn test_semi_major_axis_from_apo_peri() {
let a = OrbitalCalculator::semi_major_axis(3.0 * AU, AU);
assert!((a - 2.0 * AU).abs() / AU < 1e-9);
}
#[test]
fn test_new_simulator_is_empty() {
let sim = NBodySimulator::new();
assert_eq!(sim.body_count(), 0);
}
#[test]
fn test_add_body() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "Sun".to_string(),
mass: SOLAR_MASS,
position: [0.0; 3],
velocity: [0.0; 3],
});
assert_eq!(sim.body_count(), 1);
}
#[test]
fn test_body_lookup_found() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "Mars".to_string(),
mass: 6.39e23,
position: [0.0; 3],
velocity: [0.0; 3],
});
assert!(sim.body("Mars").is_some());
assert_eq!(sim.body("Mars").expect("body").name, "Mars");
}
#[test]
fn test_body_lookup_not_found() {
let sim = NBodySimulator::new();
assert!(sim.body("Venus").is_none());
}
#[test]
fn test_single_body_no_force() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "Loner".to_string(),
mass: 1e20,
position: [1e9, 0.0, 0.0],
velocity: [1000.0, 0.0, 0.0],
});
let dt = 10.0;
sim.step(dt);
let b = sim.body("Loner").expect("body");
assert!((b.velocity[0] - 1000.0).abs() < 1e-9);
assert!((b.position[0] - (1e9 + 1000.0 * dt)).abs() < 1e-9);
}
#[test]
fn test_two_body_mutual_attraction() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 1e20,
position: [0.0, 0.0, 0.0],
velocity: [0.0; 3],
});
sim.add_body(Body {
name: "B".to_string(),
mass: 1e20,
position: [1e9, 0.0, 0.0],
velocity: [0.0; 3],
});
sim.step(1.0);
let a = sim.body("A").expect("A");
let b = sim.body("B").expect("B");
assert!(a.velocity[0] > 0.0, "A should accelerate towards B");
assert!(b.velocity[0] < 0.0, "B should accelerate towards A");
}
#[test]
fn test_step_n_equivalent_to_n_steps() {
let make_sim = || {
let mut s = NBodySimulator::new();
s.add_body(Body {
name: "Star".to_string(),
mass: SOLAR_MASS,
position: [0.0; 3],
velocity: [0.0; 3],
});
s.add_body(Body {
name: "Planet".to_string(),
mass: EARTH_MASS,
position: [AU, 0.0, 0.0],
velocity: [0.0, EARTH_ORBITAL_V, 0.0],
});
s
};
let mut sim_step = make_sim();
let mut sim_step_n = make_sim();
let dt = 100.0;
let n = 20;
for _ in 0..n {
sim_step.step(dt);
}
sim_step_n.step_n(dt, n);
let b1 = sim_step.body("Planet").expect("p");
let b2 = sim_step_n.body("Planet").expect("p");
assert!((b1.position[0] - b2.position[0]).abs() < 1e-3);
assert!((b1.velocity[1] - b2.velocity[1]).abs() < 1e-3);
}
#[test]
fn test_kinetic_energy_single_body() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 2.0,
position: [0.0; 3],
velocity: [3.0, 0.0, 0.0],
});
assert!((sim.kinetic_energy() - 9.0).abs() < 1e-9);
}
#[test]
fn test_kinetic_energy_zero_velocity() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 1e20,
position: [0.0; 3],
velocity: [0.0; 3],
});
assert_eq!(sim.kinetic_energy(), 0.0);
}
#[test]
fn test_potential_energy_two_bodies() {
let mut sim = NBodySimulator::new();
let m1 = 1e20;
let m2 = 1e20;
let r = 1e9;
sim.add_body(Body {
name: "A".to_string(),
mass: m1,
position: [0.0, 0.0, 0.0],
velocity: [0.0; 3],
});
sim.add_body(Body {
name: "B".to_string(),
mass: m2,
position: [r, 0.0, 0.0],
velocity: [0.0; 3],
});
let expected = -G * m1 * m2 / r;
assert!((sim.potential_energy() - expected).abs() / expected.abs() < 1e-9);
}
#[test]
fn test_potential_energy_single_body_zero() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 1e20,
position: [0.0; 3],
velocity: [0.0; 3],
});
assert_eq!(sim.potential_energy(), 0.0);
}
#[test]
fn test_total_energy_is_sum() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 2.0,
position: [0.0, 0.0, 0.0],
velocity: [3.0, 0.0, 0.0],
});
sim.add_body(Body {
name: "B".to_string(),
mass: 2.0,
position: [1e9, 0.0, 0.0],
velocity: [0.0; 3],
});
let ke = sim.kinetic_energy();
let pe = sim.potential_energy();
assert!((sim.total_energy() - (ke + pe)).abs() < 1e-9);
}
#[test]
fn test_energy_roughly_conserved_short_run() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "Sun".to_string(),
mass: SOLAR_MASS,
position: [0.0; 3],
velocity: [0.0; 3],
});
sim.add_body(Body {
name: "Earth".to_string(),
mass: EARTH_MASS,
position: [AU, 0.0, 0.0],
velocity: [0.0, EARTH_ORBITAL_V, 0.0],
});
let e0 = sim.total_energy();
sim.step_n(60.0, 10); let e1 = sim.total_energy();
let rel_change = (e1 - e0).abs() / e0.abs();
assert!(rel_change < 0.01, "energy drift {rel_change:.2e} too large");
}
#[test]
fn test_centre_of_mass_two_equal_masses() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 1.0,
position: [-1.0, 0.0, 0.0],
velocity: [0.0; 3],
});
sim.add_body(Body {
name: "B".to_string(),
mass: 1.0,
position: [1.0, 0.0, 0.0],
velocity: [0.0; 3],
});
let com = sim.centre_of_mass();
assert!((com[0]).abs() < 1e-9);
}
#[test]
fn test_total_momentum_two_bodies() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "A".to_string(),
mass: 2.0,
position: [0.0; 3],
velocity: [3.0, 0.0, 0.0],
});
sim.add_body(Body {
name: "B".to_string(),
mass: 2.0,
position: [1.0, 0.0, 0.0],
velocity: [-3.0, 0.0, 0.0],
});
let p = sim.total_momentum();
assert!(p[0].abs() < 1e-9); }
#[test]
fn test_body_kinetic_energy() {
let b = Body {
name: "t".to_string(),
mass: 2.0,
position: [0.0; 3],
velocity: [3.0, 4.0, 0.0],
};
assert!((b.kinetic_energy() - 25.0).abs() < 1e-9);
}
#[test]
fn test_specific_orbital_energy_negative() {
let e = OrbitalCalculator::specific_orbital_energy(SOLAR_MASS, AU);
assert!(e < 0.0, "bound orbit must have negative energy");
}
#[test]
fn test_hill_sphere_earth_in_solar_orbit() {
let r_h = OrbitalCalculator::hill_sphere(AU, EARTH_MASS, SOLAR_MASS);
assert!(
r_h > 1e9 && r_h < 5e9,
"Hill sphere radius {r_h:.2e} m out of expected range"
);
}
#[test]
fn test_default_simulator() {
let sim = NBodySimulator::default();
assert_eq!(sim.body_count(), 0);
}
#[test]
fn test_bodies_accessor() {
let mut sim = NBodySimulator::new();
sim.add_body(Body {
name: "X".to_string(),
mass: 1e10,
position: [0.0; 3],
velocity: [0.0; 3],
});
assert_eq!(sim.bodies().len(), 1);
assert_eq!(sim.bodies()[0].name, "X");
}
#[test]
fn test_g_constant_value() {
assert!((G - 6.674e-11).abs() < 1e-15);
}
}