Hi folks, I'm trying to make an "idiomatic, fast and lightweight" implementation of the Nbabel nbody integrator based on this python version. I would appreciate very much any feedback on it:
/// The Computer Language Benchmarks Game
/// https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
///
/// Contributed by Ilia Schelokov and modified by Darley Barreto
use std::fs;
use std::path::Path;
use std::ops::{Add, Sub, Mul, AddAssign, SubAssign};
use std::default::Default;
#[derive(Clone, Copy)]
struct Vec3D(f64, f64, f64);
impl Vec3D {
fn sum_squares(&self) -> f64 {
self.0 * self.0
+ self.1 * self.1
+ self.2 * self.2
}
}
impl Default for Vec3D {
fn default() -> Vec3D {
Vec3D(0.0, 0.0, 0.0)
}
}
impl Add for Vec3D {
type Output = Vec3D;
fn add(self, rhs: Self) -> Self::Output {
Vec3D(
self.0 + rhs.0,
self.1 + rhs.1,
self.2 + rhs.2
)
}
}
impl Sub for &Vec3D {
type Output = Vec3D;
fn sub(self, rhs: Self) -> Self::Output {
Vec3D(
self.0 - rhs.0,
self.1 - rhs.1,
self.2 - rhs.2
)
}
}
impl Mul<&Vec3D> for f64 {
type Output = Vec3D;
fn mul(self, rhs: &Vec3D) -> Self::Output {
Vec3D(
self * rhs.0,
self * rhs.1,
self * rhs.2
)
}
}
impl Mul<Vec3D> for f64 {
type Output = Vec3D;
fn mul(self, rhs: Vec3D) -> Self::Output {
Vec3D(
self * rhs.0,
self * rhs.1,
self * rhs.2
)
}
}
impl AddAssign for Vec3D {
fn add_assign(&mut self, rhs: Self) {
self.0 += rhs.0;
self.1 += rhs.1;
self.2 += rhs.2;
}
}
impl SubAssign for Vec3D {
fn sub_assign(&mut self, rhs: Self) {
self.0 -= rhs.0;
self.1 -= rhs.1;
self.2 -= rhs.2;
}
}
#[derive(Clone)]
struct Particle {
position: Vec3D,
velocity: Vec3D,
acceleration: [Vec3D;2],
mass: f64,
}
fn compute_energy(particles: &mut Vec<Particle>, pe: f64) -> f64{
pe + particles.iter().map(|p| 0.5 * p.mass * p.velocity.sum_squares() ).sum::<f64>()
}
fn advance_velocities(particles: &mut Vec<Particle>, dt: f64) {
for p in particles {
p.velocity += 0.5 * dt * (p.acceleration[0] +
p.acceleration[1]);
}
}
fn advance_positions(particles: &mut Vec<Particle>, dt: f64) {
for p in particles {
p.position += dt * p.velocity + 0.5 *
dt.powi(2) * p.acceleration[0];
}
}
fn accelerate(particles: &mut Vec<Particle>) -> f64 {
for p in particles.iter_mut() {
p.acceleration[1] = p.acceleration[0];
p.acceleration[0] = Default::default();
}
let mut pe = 0.0;
for (i, p1) in unsafe { &mut *(particles as *mut Vec<Particle>) }.iter_mut().enumerate() {
for p2 in &mut particles[i+1 ..] {
let vector = &p1.position - &p2.position;
let distance = vector.sum_squares().sqrt();
let distance_cube = distance.powi(3);
p1.acceleration[0] -= (p2.mass / distance_cube) * &vector;
p2.acceleration[0] += (p1.mass / distance_cube) * &vector;
pe -= (p1.mass * p2.mass) / distance;
}
}
pe
}
fn parse_row(line: &str) -> Particle {
let row_vec: Vec<f64> = line
.split(" ")
.filter(|s| s.len() > 2)
.map(|s| s.parse().unwrap())
.collect();
Particle {
position: Vec3D(row_vec[1], row_vec[2], row_vec[3]),
velocity: Vec3D(row_vec[4], row_vec[5], row_vec[6]),
acceleration: Default::default(),
mass: row_vec[0],
}
}
fn main() {
let file_path= Path::new("data/input1k");
let mut particles: Vec<Particle> = fs::read_to_string(file_path)
.expect("File not found!")
.lines()
.filter(|x| x.len() > 0)
.map(|x| parse_row(x))
.collect();
let mut pe;
let mut energy = 0.;
let (tend, dt) = (10.0, 0.001); // end time, timestep
let (mut old_energy, energy0) = (-0.25, -0.25);
accelerate(&mut particles);
for step in 1 .. (tend / dt + 1.0) as i64 {
advance_positions(&mut particles, dt);
pe = accelerate(&mut particles);
advance_velocities(&mut particles,dt);
if step % 100 == 0 {
energy = compute_energy(&mut particles, pe);
println!("t = {}, E = {}, dE/E = {}", dt * step as f64, energy, (energy - old_energy) / old_energy);
old_energy = energy;
}
}
println!("Final dE/E = {}", (energy - energy0) / energy0);
}
PS: If you want to run, I am testing on the file data/input1k inside of the same repo, I can't put the link here because of the limitation for new users.
PSS: This code is a modified version of the nbody Rust #8 of the benchmarks game.
.