feat!: correctly compute jacobian derivative

 - Make C/J pure functions of a constraint
 - Approximate J' using finite differences (dC'/dq)
 - fn enforce_constraints() no longer depends on dt
 - Use stable-but-slow finite differences method for J by default
 - Provide analytically derived Jacobian implementations for existing
   constraints

3 files changed, 217 insertions, 101 deletions

diff --git a/physics/src/constraint/beam.rs b/physics/src/constraint/beam.rs
index b353e13..f6907c0 100644
--- a/physics/src/constraint/beam.rs
+++ b/physics/src/constraint/beam.rs
@@ -14,8 +14,6 @@ enum BeamPoint {
 pub struct BeamConstraint {
     points: [BeamPoint; 2],
     pub length: Scalar,
-
-    jacobian: RowDVector<Scalar>,
 }
 
 impl ParticleSystem {
@@ -29,7 +27,6 @@ impl ParticleSystem {
                 BeamPoint::FromParticle(particle_ids[1]),
             ],
             length: (a.position - b.position).norm(),
-            jacobian: RowDVector::zeros(self.particles.len() * N),
         }));
     }
 
@@ -42,7 +39,6 @@ impl ParticleSystem {
                 BeamPoint::Static(point),
             ],
             length: (position - point).norm(),
-            jacobian: RowDVector::zeros(self.particles.len() * N),
         }));
     }
 }
@@ -66,11 +62,34 @@ impl Constraint for BeamConstraint {
         (a - b).norm() - self.length
     }
 
-    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>) {
-        self.jacobian = jacobian
-    }
+    fn jacobian(&self, q: &DVector<Scalar>) -> RowDVector<Scalar> {
+        let [a, b] = self.points.map(|p| {
+            match p {
+                BeamPoint::Static(p) => p.coords,
+                BeamPoint::FromParticle(id) => q.fixed_rows::<N>(id * N).into(),
+            }
+        });
+        let mut result = RowDVector::zeros(q.len());
+        match (a - b).try_normalize(0.0) {
+            Some(normal) => {
+                for i in 0..N {
+                    match self.points[0] {
+                        BeamPoint::FromParticle(id) => {
+                            result[id * N + i] = normal[i];
+                        }
+                        _ => {}
+                    }
+                    match self.points[1] {
+                        BeamPoint::FromParticle(id) => {
+                            result[id * N + i] = -normal[i];
+                        }
+                        _ => {}
+                    }
+                }
+            }
+            None => ()
+        };
 
-    fn jacobian_prev(&self) -> RowDVector<Scalar> {
-        self.jacobian.clone()
+        result
     }
 }
diff --git a/physics/src/constraint/mod.rs b/physics/src/constraint/mod.rs
index 15ca986..4df1c8e 100644
--- a/physics/src/constraint/mod.rs
+++ b/physics/src/constraint/mod.rs
@@ -1,135 +1,165 @@
 use nalgebra::{DMatrix, DVector, RowDVector};
 
+use crate::algebra::{Scalar, N};
 use crate::particle_system::ParticleSystem;
-use crate::algebra::{Scalar, Vector, N};
-pub mod slider;
 pub mod beam;
+pub mod slider;
+pub mod subspace_distance;
 
-const SPRING_CONSTANT: Scalar = 0.65;
-const DAMPING_CONSTANT: Scalar = 0.85;
+pub const SPRING_CONSTANT: Scalar = 16.0;
+pub const DAMPING_CONSTANT: Scalar = 4.0;
 
-/// 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>;
 
+    /// Constraint function: C = 0 <=> constrained satisfied
     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.00001;
+    /// J = dC / dq
+    /// This implementation computes partial derivative via
+    /// finite differences method: stepping +-dQ along each axis
+    ///
+    /// NOTE: this is rather computationally intensive so
+    /// providing constraint-specific analytical implementation is preferred
+    fn jacobian(&self, q: &DVector<Scalar>) -> RowDVector<Scalar> {
+        let dq = 1e-10;
         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 a = q.clone();
-                let mut b = q.clone();
-                a[index] += dq;
-                b[index] -= dq;
+                let mut q = q.clone();
+
+                q[index] += dq;
+                let c_a = self.c(&q);
+
+                q[index] -= 2.0 * dq;
+                let c_b = self.c(&q);
 
-                result[index] = (self.c(&a) - self.c(&b)) / (2.0 * dq)
+                result[index] = (c_a - c_b) / (2.0 * 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);
+    /// C' = J * q'
+    fn c_dot(&self, jacobian: &RowDVector<Scalar>, q_dot: &DVector<Scalar>) -> Scalar {
+        (jacobian * q_dot)[0]
+    }
 
-        let c_dot = (c - c_prev) / dt;
+    /// J' = dC' / dq
+    fn jacobian_dot(&self, q: &DVector<Scalar>, q_dot: &DVector<Scalar>) -> RowDVector<Scalar> {
+        let dq = 1e-4;
+        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 = q.clone();
 
-        let jacobian = self.partial_derivative(&q);
+                q[index] += dq;
+                let j_a = self.jacobian(&q);
 
-        let jacobian_dot = (jacobian.clone() - self.jacobian_prev()) / dt;
-        self.set_jacobian(jacobian.clone());
+                q[index] -= 2.0 * dq;
+                let j_b = self.jacobian(&q);
 
-        (c, c_dot, jacobian, jacobian_dot)
+                result[index] = (self.c_dot(&j_a, q_dot) - self.c_dot(&j_b, q_dot)) / (2.0 * dq)
+            }
+        }
+        result
     }
 }
 
 impl ParticleSystem {
-    pub fn enforce_constraints(&mut self, dt: Scalar) {
-        if self.constraints.len() == 0 {
-            return
-        }
-
-        let size = self.particles.len() * N;
+    /// q is a state vector - concatenated positions of particles
+    /// q' - its derivative - concatenated velocities
+    pub fn collect_q(&self) -> (DVector<Scalar>, DVector<Scalar>) {
+        let size = self.q_dim();
         let mut q = DVector::zeros(size);
-        let mut q_prev = DVector::zeros(size);
+        let mut q_dot = 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;
+                q_dot[p * N + i] = particle.velocity[i];
             }
         }
 
-        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());
+        (q, q_dot)
+    }
+    /// dim(q) = #particles * N
+    fn q_dim(&self) -> usize {
+        self.particles.len() * N
+    }
 
-        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;
-            }
-        }
+    /// Diagonal matrix of particle masses (each repeated N times to match dim(q))
+    fn mass_matrix(&self) -> DMatrix<Scalar> {
+        DMatrix::from_diagonal(&DVector::from_iterator(
+            self.q_dim(),
+            self.particles
+                .iter()
+                .flat_map(|particle| (0..N).map(|_i| 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];
-            }
+    pub fn enforce_constraints(&mut self) {
+        if self.constraints.len() == 0 {
+            return;
         }
 
-        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);
+        let (q, q_dot) = self.collect_q();
+        let dim_q = self.q_dim();
+        let dim_c = self.constraints.len();
+
+        let c = DVector::from_iterator(
+            dim_c,
+            self.constraints.iter().map(|constraint| constraint.c(&q)),
+        );
+
+        let jacobian = DMatrix::from_rows(
+            &self
+                .constraints
+                .iter()
+                .map(|constraint| constraint.jacobian(&q))
+                .collect::<Vec<_>>(),
+        );
+
+        let jacobian_dot = DMatrix::from_rows(
+            &self
+                .constraints
+                .iter()
+                .map(|constraint| constraint.jacobian_dot(&q, &q_dot))
+                .collect::<Vec<_>>(),
+        );
+
+        let c_dot = &jacobian * &q_dot;
+
+        let forces = DVector::from_iterator(
+            dim_q,
+            self.particles
+                .iter()
+                .flat_map(|particle| particle.force.iter().map(|&v| v)),
+        );
+
+        let jacobian_times_w = &jacobian * self.mass_matrix();
+        let jacobian_transpose = &jacobian.transpose();
+
+        let lhs = &jacobian_times_w * jacobian_transpose;
+
+        let rhs = -(&jacobian_dot * &q_dot
+            + jacobian_times_w * &forces
+            + SPRING_CONSTANT * &c
+            + DAMPING_CONSTANT * &c_dot);
 
         match lhs.lu().solve(&rhs) {
             Some(lambda) => {
-                // println!("Lambda = {}", lambda);
-                let constraint_force = jacobian.transpose() * 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);
+                    let force = constraint_force.fixed_rows::<N>(i * N);
+                    self.particles[i].apply_force(force.into());
                 }
             }
             None => println!("Lambda not found"),
diff --git a/physics/src/constraint/slider.rs b/physics/src/constraint/slider.rs
index fa7e840..be8e261 100644
--- a/physics/src/constraint/slider.rs
+++ b/physics/src/constraint/slider.rs
@@ -9,10 +9,7 @@ use super::Constraint;
 
 pub struct SliderConstraint {
     pub particle_id: usize,
-
     pub line: Line,
-
-    jacobian: RowDVector<Scalar>,
 }
 
 impl ParticleSystem {
@@ -21,7 +18,6 @@ impl ParticleSystem {
         self.constraints.push(Box::new(SliderConstraint {
             particle_id,
             line: Line::new(point, [direction]),
-            jacobian: RowDVector::zeros(self.particles.len() * N),
         }));
     }
 }
@@ -37,11 +33,82 @@ impl Constraint for SliderConstraint {
         self.line.distance(point)
     }
 
-    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>) {
-        self.jacobian = jacobian
+    fn jacobian(&self, q: &DVector<Scalar>) -> RowDVector<Scalar> {
+        let position = q.fixed_rows::<N>(self.particle_id * N);
+        let point = Point::from_coordinates(position.into());
+        let mut result = RowDVector::zeros(q.len());
+
+        match (point - self.line.project_point(point)).try_normalize(0.0) {
+            Some(normal) => {
+                for i in 0..N {
+                    result[self.particle_id * N + i] = normal[i]
+                }
+            }
+            None => ()
+        };
+
+
+        result
     }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{fs::File, io::Write};
+
+    use crate::{algebra::{Point, Vector}, constraint::{DAMPING_CONSTANT, SPRING_CONSTANT}, force::gravity::Gravity, particle_system::{Particle, ParticleSystem}, solver::Solver};
+
+    #[test]
+    fn test_drift_resiliance() {
+        let origin = Point::origin();
+        let mut system = ParticleSystem {
+            particles: vec![Particle::new(origin, 1.0)],
+            constraints: vec![],
+            forces: vec![Box::new(Gravity {
+                vector: Vector::new(-5.0, -1.0, 20.0),
+
+            })],
+            t: 0.0,
+        };
+        system.add_slider_constraint(0, Vector::x());
+
+        let mut errors = vec![];
+        let mut positions = vec![];
+
+        let sim_time = 10.0;
+
+        let dt = 1e-5;
+        for i in 0..(sim_time / dt) as usize {
+            system.apply_forces();
+            system.enforce_constraints();
+            system.step(dt);
+
+            let (q, _) = system.collect_q();
+            let error = system.constraints[0].c(&q);
+            errors.push(error);
+            positions.push(system.particles[0].position);
+            // println!("Iteration {}, error = {}", i, error);
+        }
+
+        let error_max = errors.clone().into_iter().reduce(f64::max).unwrap();
+
+        let mut error_sorted = errors.clone();
+        error_sorted.sort_by(f64::total_cmp);
+        let mean = error_sorted[errors.len() / 2];
 
-    fn jacobian_prev(&self) -> RowDVector<Scalar> {
-        self.jacobian.clone()
+        {
+            let filename = format!("./plots/constraint/errors-s{}-d{}.txt", SPRING_CONSTANT, DAMPING_CONSTANT);
+            let mut f = File::create(filename).unwrap();
+            let s = errors.iter().fold(String::new(), |acc, e| acc + &e.to_string() + "\n");
+            f.write(s.as_bytes()).unwrap();
+        }
+        {
+            let filename = format!("./plots/slider/pos-s{}-d{}.txt", SPRING_CONSTANT, DAMPING_CONSTANT);
+            let mut f = File::create(filename).unwrap();
+            let s = positions.iter().fold(String::new(), |acc, e| acc + &e.y.to_string() + " " + &e.z.to_string() + "\n");
+            f.write(s.as_bytes()).unwrap();
+        }
+        // assert!(error < error_max, "Error must decrease over time");
+        assert!(error_max < 1.0e-5, "Error too high: MAX = {}, MEAN = {}", error_max, mean);
     }
 }