From ced212f233b17b5e017bc8db174a14f909fdca71 Mon Sep 17 00:00:00 2001 From: Adam Wick Date: Fri, 27 Dec 2019 08:57:13 -0600 Subject: [PATCH] Add nbody problem bits. --- src/main.rs | 2 + src/nbody.rs | 193 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 195 insertions(+) create mode 100644 src/nbody.rs diff --git a/src/main.rs b/src/main.rs index 61ae8aa..866a407 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,8 @@ mod args; mod fuel; mod image; mod machine; +#[cfg(test)] +mod nbody; mod orbits; #[cfg(test)] mod repair; diff --git a/src/nbody.rs b/src/nbody.rs new file mode 100644 index 0000000..5be49b5 --- /dev/null +++ b/src/nbody.rs @@ -0,0 +1,193 @@ +use std::fmt; +use std::cmp::Ordering; + +#[derive(Clone, Debug, PartialEq)] +struct Body { + x: i64, + y: i64, + z: i64, + x_vel: i64, + y_vel: i64, + z_vel: i64, +} + +impl fmt::Display for Body { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "pos=, vel=", + self.x, self.y, self.z, self.x_vel, self.y_vel, self.z_vel) + } +} + +impl Body { + fn new(x: i64, y: i64, z: i64) -> Body { + Body { + x, y, z, + x_vel: 0, y_vel: 0, z_vel: 0, + } + } + + fn apply_velocity(&mut self) { + self.x += self.x_vel; + self.y += self.y_vel; + self.z += self.z_vel; + } + + fn apply_gravity(&mut self, others: &[Body]) { + for other in others.iter() { + self.x_vel += velocity_adjustment(self.x, other.x); + self.y_vel += velocity_adjustment(self.y, other.y); + self.z_vel += velocity_adjustment(self.z, other.z); + } + } + + fn potential_energy(&self) -> u64 { + self.x.abs() as u64 + + self.y.abs() as u64 + + self.z.abs() as u64 + } + + fn kinetic_energy(&self) -> u64 { + self.x_vel.abs() as u64 + + self.y_vel.abs() as u64 + + self.z_vel.abs() as u64 + } + + fn total_energy(&self) -> u64 { + self.potential_energy() * self.kinetic_energy() + } +} + +fn time_step(bodies: &mut [Body]) { + let mut copy_bodies = Vec::new(); + + copy_bodies.extend(bodies.iter().map(|x| x.clone())); + for body in bodies.iter_mut() { + body.apply_gravity(©_bodies); + } + for body in bodies.iter_mut() { + body.apply_velocity(); + } +} + +fn velocity_adjustment(v1: i64, v2: i64) -> i64 { + match v1.cmp(&v2) { + Ordering::Greater => -1, + Ordering::Equal => 0, + Ordering::Less => 1, + } +} + +#[test] +fn example1() { + let mut example_bodies = [ Body::new(-1, 0, 2), + Body::new(2, -10, -7), + Body::new(4, -8, 8), + Body::new(3, 5, -1), + ]; + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 2, y:-1, z: 1, x_vel: 3, y_vel:-1, z_vel:-1 }, + Body{ x: 3, y:-7, z:-4, x_vel: 1, y_vel: 3, z_vel: 3 }, + Body{ x: 1, y:-7, z: 5, x_vel:-3, y_vel: 1, z_vel:-3 }, + Body{ x: 2, y: 2, z: 0, x_vel:-1, y_vel:-3, z_vel: 1 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 5, y:-3, z:-1, x_vel: 3, y_vel:-2, z_vel:-2 }, + Body{ x: 1, y:-2, z: 2, x_vel:-2, y_vel: 5, z_vel: 6 }, + Body{ x: 1, y:-4, z:-1, x_vel: 0, y_vel: 3, z_vel:-6 }, + Body{ x: 1, y:-4, z: 2, x_vel:-1, y_vel:-6, z_vel: 2 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 5, y:-6, z:-1, x_vel: 0, y_vel:-3, z_vel: 0 }, + Body{ x: 0, y: 0, z: 6, x_vel:-1, y_vel: 2, z_vel: 4 }, + Body{ x: 2, y: 1, z:-5, x_vel: 1, y_vel: 5, z_vel:-4 }, + Body{ x: 1, y:-8, z: 2, x_vel: 0, y_vel:-4, z_vel: 0 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 2, y:-8, z: 0, x_vel:-3, y_vel:-2, z_vel: 1 }, + Body{ x: 2, y: 1, z: 7, x_vel: 2, y_vel: 1, z_vel: 1 }, + Body{ x: 2, y: 3, z:-6, x_vel: 0, y_vel: 2, z_vel:-1 }, + Body{ x: 2, y:-9, z: 1, x_vel: 1, y_vel:-1, z_vel:-1 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x:-1, y:-9, z: 2, x_vel:-3, y_vel:-1, z_vel: 2 }, + Body{ x: 4, y: 1, z: 5, x_vel: 2, y_vel: 0, z_vel:-2 }, + Body{ x: 2, y: 2, z:-4, x_vel: 0, y_vel:-1, z_vel: 2 }, + Body{ x: 3, y:-7, z:-1, x_vel: 1, y_vel: 2, z_vel:-2 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x:-1, y:-7, z: 3, x_vel: 0, y_vel: 2, z_vel: 1 }, + Body{ x: 3, y: 0, z: 0, x_vel:-1, y_vel:-1, z_vel:-5 }, + Body{ x: 3, y:-2, z: 1, x_vel: 1, y_vel:-4, z_vel: 5 }, + Body{ x: 3, y:-4, z:-2, x_vel: 0, y_vel: 3, z_vel:-1 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 2, y:-2, z: 1, x_vel: 3, y_vel: 5, z_vel:-2 }, + Body{ x: 1, y:-4, z:-4, x_vel:-2, y_vel:-4, z_vel:-4 }, + Body{ x: 3, y:-7, z: 5, x_vel: 0, y_vel:-5, z_vel: 4 }, + Body{ x: 2, y: 0, z: 0, x_vel:-1, y_vel: 4, z_vel: 2 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 5, y: 2, z:-2, x_vel: 3, y_vel: 4, z_vel:-3 }, + Body{ x: 2, y:-7, z:-5, x_vel: 1, y_vel:-3, z_vel:-1 }, + Body{ x: 0, y:-9, z: 6, x_vel:-3, y_vel:-2, z_vel: 1 }, + Body{ x: 1, y: 1, z: 3, x_vel:-1, y_vel: 1, z_vel: 3 }, ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 5, y: 3, z:-4, x_vel: 0, y_vel: 1, z_vel:-2 }, + Body{ x: 2, y:-9, z:-3, x_vel: 0, y_vel:-2, z_vel: 2 }, + Body{ x: 0, y:-8, z: 4, x_vel: 0, y_vel: 1, z_vel:-2 }, + Body{ x: 1, y: 1, z: 5, x_vel: 0, y_vel: 0, z_vel: 2 } ]); + + time_step(&mut example_bodies); + assert_eq!(example_bodies, [ + Body{ x: 2, y: 1, z:-3, x_vel:-3, y_vel:-2, z_vel: 1 }, + Body{ x: 1, y:-8, z: 0, x_vel:-1, y_vel: 1, z_vel: 3 }, + Body{ x: 3, y:-6, z: 1, x_vel: 3, y_vel: 2, z_vel:-3 }, + Body{ x: 2, y: 0, z: 4, x_vel: 1, y_vel:-1, z_vel:-1 } ]); + + assert_eq!(36, example_bodies[0].total_energy()); + assert_eq!(45, example_bodies[1].total_energy()); + assert_eq!(80, example_bodies[2].total_energy()); + assert_eq!(18, example_bodies[3].total_energy()); +} + +#[test] +fn example2() { + let mut example_bodies = [ Body::new(-8, -10, 0), + Body::new(5, 5, 10), + Body::new(2, -7, 3), + Body::new(9, -8, -3), ]; + + for _ in 0..100 { + time_step(&mut example_bodies); + } + + assert_eq!(1940, example_bodies.iter() + .map(|x| x.total_energy()) + .fold(0, |acc,x| acc + x)); +} + +#[test] +fn day12() { + let mut example_bodies = [ Body::new(3, -6, 6), + Body::new(10, 7, -9), + Body::new(-3, -7, 9), + Body::new(-8, 0, 4) ]; + + for _ in 0..1000 { + time_step(&mut example_bodies); + } + + assert_eq!(6849, example_bodies.iter() + .map(|x| x.total_energy()) + .fold(0, |acc,x| acc + x)); +} \ No newline at end of file