feat: implement constraint system and PPM rendering

8 files changed, 335 insertions, 17 deletions

diff --git a/src/constraint/anchor.rs b/src/constraint/anchor.rs
new file mode 100644
index 0000000..3ef94ff
--- /dev/null
+++ b/src/constraint/anchor.rs
@@ -0,0 +1,42 @@
+use nalgebra::{DVector, RowDVector};
+
+use crate::particle_system::{ParticleSystem, Point, Scalar, N};
+
+use super::Constraint;
+
+pub struct AnchorConstraint {
+    pub particle_id: usize,
+    pub anchor: Point,
+
+    jacobian: RowDVector<Scalar>,
+}
+
+impl ParticleSystem {
+    pub fn add_anchor_constraint(&mut self, particle_id: usize) {
+        let anchor = self.particles[particle_id].position;
+        self.constraints.push(Box::new(AnchorConstraint {
+            particle_id,
+            anchor,
+            jacobian: RowDVector::zeros(self.particles.len() * N),
+        }));
+    }
+}
+
+impl Constraint for AnchorConstraint {
+    fn get_particles(&self) -> Vec<usize> {
+        vec![self.particle_id]
+    }
+
+    fn c(&self, q: &DVector<Scalar>) -> Scalar {
+        let position = q.fixed_rows(self.particle_id * N);
+        (position - self.anchor.coords).norm()
+    }
+
+    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>) {
+        self.jacobian = jacobian
+    }
+
+    fn jacobian_prev(&self) -> RowDVector<Scalar> {
+        self.jacobian.clone()
+    }
+}
diff --git a/src/constraint/beam.rs b/src/constraint/beam.rs
new file mode 100644
index 0000000..c953920
--- /dev/null
+++ b/src/constraint/beam.rs
@@ -0,0 +1,46 @@
+use nalgebra::{DVector, RowDVector};
+
+use crate::particle_system::{ParticleSystem, Scalar, N};
+
+use super::Constraint;
+
+pub struct BeamConstraint {
+    pub particle_ids: [usize; 2],
+    pub length: Scalar,
+
+    jacobian: RowDVector<Scalar>,
+}
+
+impl ParticleSystem {
+    pub fn add_beam_constraint(&mut self, particle_ids: [usize; 2]) {
+        let a = &self.particles[particle_ids[0]];
+        let b = &self.particles[particle_ids[1]];
+
+        self.constraints.push(Box::new(BeamConstraint {
+            particle_ids,
+            length: (a.position - b.position).norm(),
+            jacobian: RowDVector::zeros(self.particles.len() * N),
+        }));
+    }
+}
+
+impl Constraint for BeamConstraint {
+    fn get_particles(&self) -> Vec<usize> {
+        Vec::from(self.particle_ids)
+    }
+
+    fn c(&self, q: &DVector<Scalar>) -> Scalar {
+        let a = q.fixed_rows::<N>(self.particle_ids[0] * N);
+        let b = q.fixed_rows::<N>(self.particle_ids[1] * N);
+
+        (a - b).norm() - self.length
+    }
+
+    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>) {
+        self.jacobian = jacobian
+    }
+
+    fn jacobian_prev(&self) -> RowDVector<Scalar> {
+        self.jacobian.clone()
+    }
+}
diff --git a/src/constraint/mod.rs b/src/constraint/mod.rs
new file mode 100644
index 0000000..99905ba
--- /dev/null
+++ b/src/constraint/mod.rs
@@ -0,0 +1,138 @@
+use nalgebra::{DMatrix, DVector, RowDVector};
+
+use crate::particle_system::{ParticleSystem, Scalar, Vector, N};
+pub mod anchor;
+pub mod beam;
+
+const SPRING_CONSTANT: Scalar = 0.75;
+const DAMPING_CONSTANT: Scalar = 0.45;
+
+/// SIZE is always 3xN, but this operation can't be done at compile time yet:
+/// "generic parameters may not be used in const operations"
+pub trait Constraint {
+    fn get_particles(&self) -> Vec<usize>;
+
+    fn c(&self, q: &DVector<Scalar>) -> Scalar;
+
+    fn jacobian_prev(&self) -> RowDVector<Scalar>;
+    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>);
+
+    fn partial_derivative(&self, q: &DVector<Scalar>) -> RowDVector<Scalar> {
+        let dq = 0.001;
+        let c_original = self.c(&q);
+        let mut result = RowDVector::zeros(q.len());
+        // The only non-zero components of derivative vector are
+        // the particles that this constraint affects
+        for particle_id in self.get_particles() {
+            for i in 0..N {
+                let index = particle_id * N + i;
+                let mut q_partial = q.clone();
+                q_partial[index] += dq;
+
+                let c = self.c(&q_partial);
+
+                result[index] = (c - c_original) / dq
+            }
+        }
+
+        result
+    }
+
+    fn compute(
+        &mut self,
+        q: &DVector<Scalar>,
+        q_prev: &DVector<Scalar>,
+        dt: Scalar,
+    ) -> (Scalar, Scalar, RowDVector<Scalar>, RowDVector<Scalar>) {
+        let c = self.c(q);
+        let c_prev = self.c(q_prev);
+
+        let c_dot = (c - c_prev) / dt;
+
+        let jacobian = self.partial_derivative(&q);
+
+        let jacobian_dot = (jacobian.clone() - self.jacobian_prev()) / dt;
+        self.set_jacobian(jacobian.clone());
+
+        (c, c_dot, jacobian, jacobian_dot)
+    }
+}
+
+impl ParticleSystem {
+    pub fn enforce_constraints(&mut self, dt: Scalar) {
+        if self.constraints.len() == 0 {
+            return
+        }
+
+        let size = self.particles.len() * N;
+        let mut q = DVector::zeros(size);
+        let mut q_prev = DVector::zeros(size);
+
+        for (p, particle) in self.particles.iter().enumerate() {
+            for i in 0..N {
+                q[p * N + i] = particle.position[i];
+                q_prev[p * N + i] = particle.position[i] - particle.velocity[i] * dt;
+            }
+        }
+
+        let q_dot = (q.clone() - q_prev.clone()) / dt;
+
+        let mut c = DVector::zeros(self.constraints.len());
+        let mut c_dot = DVector::zeros(self.constraints.len());
+        let mut jacobian = DMatrix::<Scalar>::zeros(self.constraints.len(), size);
+        let mut jacobian_dot = DMatrix::<Scalar>::zeros(self.constraints.len(), size);
+
+        for (constraint_id, constraint) in self.constraints.iter_mut().enumerate() {
+            let (value, derivative, j, jdot) = constraint.compute(&q, &q_prev, dt);
+            jacobian.set_row(constraint_id, &j);
+            jacobian_dot.set_row(constraint_id, &jdot);
+            c[constraint_id] = value;
+            c_dot[constraint_id] = derivative;
+        }
+
+        // println!("C {}\nC' {}", c, c_dot);
+        // println!("J  = {}", jacobian.clone());
+        // println!("J' = {}", jacobian_dot.clone());
+
+        let mut w = DMatrix::zeros(size, size);
+        for (i, particle) in self.particles.iter().enumerate() {
+            for j in 0..N {
+                *w.index_mut((i * N + j, i * N + j)) = 1.0 / particle.mass;
+            }
+        }
+
+        let mut forces = DVector::zeros(size);
+        for (p, particle) in self.particles.iter().enumerate() {
+            for i in 0..N {
+                forces[p * N + i] = particle.force[i];
+            }
+        }
+
+        let lhs = jacobian.clone() * w.clone() * jacobian.transpose();
+        // println!("LHS {}", lhs);
+
+        let rhs = -(jacobian_dot * q_dot
+            + jacobian.clone() * w * forces
+            + SPRING_CONSTANT * c
+            + DAMPING_CONSTANT * c_dot);
+
+        // println!("RHS {}", rhs);
+
+        match lhs.lu().solve(&rhs) {
+            Some(lambda) => {
+                // println!("Lambda = {}", lambda);
+                let constraint_force = jacobian.transpose() * lambda;
+
+                for i in 0..self.particles.len() {
+                    let mut force = Vector::zeros();
+                    for j in 0..N {
+                        force[j] = constraint_force[i * N + j];
+                    }
+
+                    self.particles[i].apply_force(force);
+                }
+            }
+            None => println!("Lambda not found"),
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 0a51789..718abf8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,30 +1,57 @@
+use std::path::PathBuf;
+
 use particle_system::{Particle, ParticleSystem, Point, Vector};
+use ppm::PPM;
 use solver::Solver;
 
+mod constraint;
 mod particle_system;
+mod ppm;
 mod solver;
 
 fn main() {
+    let ppm = PPM {
+        width: 100,
+        height: 200,
+        prefix: PathBuf::from("./out"),
+    };
+
     let dt = 0.01;
     let mut system = ParticleSystem {
-        particles: vec![Particle::new(Point::origin(), 1.0)],
+        particles: vec![
+            Particle::new(Point::origin(), 4.0),
+            Particle::new(Point::new(-30.0, 0.0), 15.0),
+            Particle::new(Point::new(20.0, 0.0), 30.0),
+            Particle::new(Point::new(5.0, 20.0), 50.0),
+        ],
+        constraints: vec![],
         t: 0.0,
     };
 
+    system.add_anchor_constraint(0);
+    system.add_beam_constraint([0, 2]);
+    system.add_beam_constraint([1, 2]);
+    system.add_beam_constraint([1, 3]);
+    system.add_beam_constraint([2, 3]);
+
     let gravity = -9.8 * Vector::y();
 
-    for i in 0..3 {
-        println!("Iteration #{i}");
+    for i in 0..150_00 {
         for particle in &mut system.particles {
             particle.reset_force();
+            // Apply custom forces
+            particle.apply_force(gravity * particle.mass);
         }
+        system.particles[0].apply_force(gravity); // Extra force on particle 0
 
-        for particle in &mut system.particles {
-            particle.apply_force(gravity);
+        if i % 10 == 0 {
+            println!("Iteration #{i}");
+            println!("{:#?}", system.particles);
+            ppm.save_frame(&system.particles, system.t);
         }
 
-        system.step(dt);
+        system.enforce_constraints(dt);
 
-        println!("{:?}", system);
+        system.step(dt);
     }
 }
diff --git a/src/particle_system.rs b/src/particle_system.rs
index 6f43c35..a3d7a4b 100644
--- a/src/particle_system.rs
+++ b/src/particle_system.rs
@@ -1,7 +1,9 @@
 use nalgebra::{Point as PointBase, SVector};
 
+use crate::{constraint::Constraint, force::Force};
+
 pub const N: usize = 2;
-pub type Scalar = f32;
+pub type Scalar = f64;
 
 pub type Vector = SVector<Scalar, N>;
 pub type Point = PointBase<Scalar, N>;
@@ -34,9 +36,11 @@ impl Particle {
     }
 }
 
-#[derive(Debug)]
+// #[derive(Debug)]
 pub struct ParticleSystem {
     pub particles: Vec<Particle>,
+    pub constraints: Vec<Box<dyn Constraint>>,
+    pub forces: Vec<Box<dyn Force>>,
 
     /// Simulation clock
     pub t: Scalar,
diff --git a/src/ppm.rs b/src/ppm.rs
new file mode 100644
index 0000000..01cdebb
--- /dev/null
+++ b/src/ppm.rs
@@ -0,0 +1,41 @@
+use std::{fs::File, io::Write, path::PathBuf};
+
+use crate::particle_system::{Particle, Vector, N, Scalar};
+
+pub struct PPM {
+    pub prefix: PathBuf,
+    pub width: usize,
+    pub height: usize,
+    // pixels_per_unit: usize,
+}
+
+impl PPM {
+    fn render_particles(&self, particles: &Vec<Particle>) -> String {
+        let mut s = format!("P3\n{} {}\n255\n", self.width, self.height);
+        let white = "255 255 255 ";
+        let black = "0 0 0 ";
+
+        for pixel_row in 0..self.height {
+            for pixel_col in 0..self.width {
+                let point = Vector::new(pixel_col as Scalar, (pixel_row as Scalar) * -1.0)
+                    + Vector::new(self.width as Scalar / -2.0, self.height as Scalar / 2.0);
+                let color = match particles.iter().any(|p| {
+                    (p.position - point).coords.norm() <= (p.mass).powf(1.0 / (N as f64))
+                }) {
+                    true => black,
+                    false => white,
+                };
+                s += color;
+            }
+        }
+        s
+    }
+
+    pub fn save_frame(&self, particles: &Vec<Particle>, time: Scalar) {
+        let file_name = format!("frame-{:08.3}", time);
+        let path = self.prefix.join(file_name);
+        let mut file = File::create(path).unwrap();
+        let data = self.render_particles(particles);
+        file.write(data.as_bytes()).unwrap();
+    }
+}
diff --git a/src/solver/midpoint.rs b/src/solver/midpoint.rs
index c5a01c4..08a3e3c 100644
--- a/src/solver/midpoint.rs
+++ b/src/solver/midpoint.rs
@@ -6,7 +6,7 @@ impl Solver for ParticleSystem {
         let mut state = self.collect_phase_space();
 
         // Shift to the midpoint
-        self.scatter_phase_space(&&PhaseSpace {
+        self.scatter_phase_space(&PhaseSpace {
             0: state.0.clone() + self.compute_derivative().0 * dt / 2.0,
         });
 
diff --git a/src/solver/mod.rs b/src/solver/mod.rs
index 4a5fec5..4bb8747 100644
--- a/src/solver/mod.rs
+++ b/src/solver/mod.rs
@@ -7,10 +7,10 @@ mod midpoint;
 #[derive(Debug, Clone)]
 pub struct PhaseSpace(DVector<Scalar>);
 type ParticleView<'a> = Matrix<
-    f32,
+    Scalar,
     Const<{ PhaseSpace::PARTICLE_DIM }>,
     Const<1>,
-    ViewStorage<'a, f32, Const<{ PhaseSpace::PARTICLE_DIM }>, Const<1>, Const<1>, Dyn>,
+    ViewStorage<'a, Scalar, Const<{ PhaseSpace::PARTICLE_DIM }>, Const<1>, Const<1>, Dyn>,
 >;
 
 impl PhaseSpace {
@@ -63,7 +63,6 @@ impl ParticleSystem {
     fn scatter_phase_space(&mut self, phase_space: &PhaseSpace) {
         for (particle_index, particle) in &mut self.particles.iter_mut().enumerate() {
             let view = phase_space.particle_view(particle_index);
-
             for i in 0..N {
                 particle.position[i] = view[i];
                 particle.velocity[i] = view[i + N];
@@ -85,6 +84,7 @@ mod tests {
     fn test_collect_phase_space() {
         let system = ParticleSystem {
             particles: vec![Particle::new(Point::new(2.0, 3.0), 1.0)],
+            constraints: vec![],
             t: 0.0,
         };
         let phase_space = system.collect_phase_space();
@@ -105,6 +105,7 @@ mod tests {
                 Particle::new(Point::new(0.0, 0.0), 1.0),
                 Particle::new(Point::new(0.0, 0.0), 1.0),
             ],
+            constraints: vec![],
             t: 0.0,
         };
 
@@ -120,11 +121,16 @@ mod tests {
         );
     }
 
-    fn simulate_falling_ball(fall_time: Scalar, dt: Scalar) -> (Vector, Vector) {
+    fn simulate_falling_ball(
+        fall_time: Scalar,
+        dt: Scalar,
+        mass: Scalar,
+    ) -> (Vector, Vector, ParticleSystem) {
         let gravity = -9.8 * Vector::y();
 
         let mut system = ParticleSystem {
-            particles: vec![Particle::new(Point::origin(), 1.0)],
+            particles: vec![Particle::new(Point::origin(), mass)],
+            constraints: vec![],
             t: 0.0,
         };
 
@@ -133,7 +139,7 @@ mod tests {
         for _ in 0..iterations {
             for particle in &mut system.particles {
                 particle.reset_force();
-                particle.apply_force(gravity);
+                particle.apply_force(gravity * particle.mass);
             }
             system.step(dt);
         }
@@ -144,12 +150,13 @@ mod tests {
         (
             system.particles[0].position.coords - expected_position,
             system.particles[0].velocity - expected_velocity,
+            system,
         )
     }
 
     #[test]
     fn ball_should_fall() {
-        let (position_error, velocity_error) = simulate_falling_ball(10.0, 0.01);
+        let (position_error, velocity_error, _) = simulate_falling_ball(10.0, 0.01, 3.0);
         assert!(
             position_error.norm() < 0.01,
             "Position error is too high: {}",
@@ -161,4 +168,17 @@ mod tests {
             velocity_error,
         );
     }
+
+    #[test]
+    fn freefall_different_masses() {
+        let (_, _, system1) = simulate_falling_ball(10.0, 0.01, 2.0);
+        let (_, _, system2) = simulate_falling_ball(10.0, 0.01, 10.0);
+
+        let diff = system1.particles[0].position - system2.particles[0].position;
+        assert!(
+            diff.norm() < 0.01,
+            "Points with different masses should fall with the same speed, diff: {}",
+            diff
+        );
+    }
 }