Merge pull request #579 from dimforge/joints-improvements

Feat: implement a "small-steps" velocity-based constraints solver + joint improvements

1 files changed, 237 insertions, 169 deletions

diff --git a/src/dynamics/solver/velocity_solver.rs b/src/dynamics/solver/velocity_solver.rs
index d15ea68..8e3ebcd 100644
--- a/src/dynamics/solver/velocity_solver.rs
+++ b/src/dynamics/solver/velocity_solver.rs
@@ -1,48 +1,99 @@
-use super::AnyJointVelocityConstraint;
+use super::{JointConstraintTypes, SolverConstraintsSet};
+use crate::dynamics::solver::solver_body::SolverBody;
 use crate::dynamics::{
-    solver::{AnyVelocityConstraint, DeltaVel},
-    IntegrationParameters, IslandManager, JointGraphEdge, MultibodyJointSet, RigidBodySet,
+    solver::{ContactConstraintTypes, SolverVel},
+    IntegrationParameters, IslandManager, JointGraphEdge, JointIndex, MultibodyJointSet,
+    MultibodyLinkId, RigidBodySet,
 };
-use crate::geometry::ContactManifold;
+use crate::geometry::{ContactManifold, ContactManifoldIndex};
 use crate::math::Real;
-use crate::utils::WAngularInertia;
+use crate::prelude::RigidBodyVelocity;
+use crate::utils::SimdAngularInertia;
 use na::DVector;
 
 pub(crate) struct VelocitySolver {
-    pub mj_lambdas: Vec<DeltaVel<Real>>,
-    pub generic_mj_lambdas: DVector<Real>,
+    pub solver_bodies: Vec<SolverBody>,
+    pub solver_vels: Vec<SolverVel<Real>>,
+    pub solver_vels_increment: Vec<SolverVel<Real>>,
+    pub generic_solver_vels: DVector<Real>,
+    pub generic_solver_vels_increment: DVector<Real>,
+    pub multibody_roots: Vec<MultibodyLinkId>,
 }
 
 impl VelocitySolver {
     pub fn new() -> Self {
         Self {
-            mj_lambdas: Vec::new(),
-            generic_mj_lambdas: DVector::zeros(0),
+            solver_bodies: Vec::new(),
+            solver_vels: Vec::new(),
+            solver_vels_increment: Vec::new(),
+            generic_solver_vels: DVector::zeros(0),
+            generic_solver_vels_increment: DVector::zeros(0),
+            multibody_roots: Vec::new(),
         }
     }
 
-    pub fn solve(
-        &mut self,
+    pub fn init_constraints(
+        &self,
         island_id: usize,
-        params: &IntegrationParameters,
         islands: &IslandManager,
         bodies: &mut RigidBodySet,
         multibodies: &mut MultibodyJointSet,
         manifolds_all: &mut [&mut ContactManifold],
+        manifold_indices: &[ContactManifoldIndex],
         joints_all: &mut [JointGraphEdge],
-        contact_constraints: &mut [AnyVelocityConstraint],
-        generic_contact_jacobians: &DVector<Real>,
-        joint_constraints: &mut [AnyJointVelocityConstraint],
-        generic_joint_jacobians: &DVector<Real>,
+        joint_indices: &[JointIndex],
+        contact_constraints: &mut SolverConstraintsSet<ContactConstraintTypes>,
+        joint_constraints: &mut SolverConstraintsSet<JointConstraintTypes>,
     ) {
-        self.mj_lambdas.clear();
-        self.mj_lambdas
-            .resize(islands.active_island(island_id).len(), DeltaVel::zero());
+        contact_constraints.init(
+            island_id,
+            islands,
+            bodies,
+            multibodies,
+            manifolds_all,
+            manifold_indices,
+        );
 
-        let total_multibodies_ndofs = multibodies.multibodies.iter().map(|m| m.1.ndofs()).sum();
-        self.generic_mj_lambdas = DVector::zeros(total_multibodies_ndofs);
+        joint_constraints.init(
+            island_id,
+            islands,
+            bodies,
+            multibodies,
+            joints_all,
+            joint_indices,
+        );
+    }
 
-        // Initialize delta-velocities (`mj_lambdas`) with external forces (gravity etc):
+    pub fn init_solver_velocities_and_solver_bodies(
+        &mut self,
+        params: &IntegrationParameters,
+        island_id: usize,
+        islands: &IslandManager,
+        bodies: &mut RigidBodySet,
+        multibodies: &mut MultibodyJointSet,
+    ) {
+        self.multibody_roots.clear();
+        self.solver_bodies.clear();
+        self.solver_bodies.resize(
+            islands.active_island(island_id).len(),
+            SolverBody::default(),
+        );
+
+        self.solver_vels_increment.clear();
+        self.solver_vels_increment
+            .resize(islands.active_island(island_id).len(), SolverVel::zero());
+        self.solver_vels.clear();
+        self.solver_vels
+            .resize(islands.active_island(island_id).len(), SolverVel::zero());
+
+        /*
+         * Initialize solver bodies and delta-velocities (`solver_vels_increment`) with external forces (gravity etc):
+         * NOTE: we compute this only once by neglecting changes of mass matrices.
+         */
+
+        // Assign solver ids to multibodies, and collect the relevant roots.
+        // And init solver_vels for rigidb-bodies.
+        let mut multibody_solver_id = 0;
         for handle in islands.active_island(island_id) {
             if let Some(link) = multibodies.rigid_body_link(*handle).copied() {
                 let multibody = multibodies
@@ -50,196 +101,213 @@ impl VelocitySolver {
                     .unwrap();
 
                 if link.id == 0 || link.id == 1 && !multibody.root_is_dynamic {
-                    let mut mj_lambdas = self
-                        .generic_mj_lambdas
-                        .rows_mut(multibody.solver_id, multibody.ndofs());
-                    mj_lambdas.axpy(params.dt, &multibody.accelerations, 0.0);
+                    multibody.solver_id = multibody_solver_id;
+                    multibody_solver_id += multibody.ndofs();
+                    self.multibody_roots.push(link);
                 }
             } else {
                 let rb = &bodies[*handle];
-                let dvel = &mut self.mj_lambdas[rb.ids.active_set_offset];
+                let solver_vel = &mut self.solver_vels[rb.ids.active_set_offset];
+                let solver_vel_incr = &mut self.solver_vels_increment[rb.ids.active_set_offset];
+                let solver_body = &mut self.solver_bodies[rb.ids.active_set_offset];
+                solver_body.copy_from(rb);
 
                 // NOTE: `dvel.angular` is actually storing angular velocity delta multiplied
                 //       by the square root of the inertia tensor:
-                dvel.angular +=
+                solver_vel_incr.angular =
                     rb.mprops.effective_world_inv_inertia_sqrt * rb.forces.torque * params.dt;
-                dvel.linear +=
+                solver_vel_incr.linear =
                     rb.forces.force.component_mul(&rb.mprops.effective_inv_mass) * params.dt;
+
+                solver_vel.linear = rb.vels.linvel;
+                // PERF: can we avoid the call to effective_angular_inertia_sqrt?
+                solver_vel.angular = rb.mprops.effective_angular_inertia_sqrt() * rb.vels.angvel;
             }
         }
 
-        /*
-         * Solve constraints.
-         */
-        for i in 0..params.max_velocity_iterations {
-            let solve_friction = params.interleave_restitution_and_friction_resolution
-                && params.max_velocity_friction_iterations + i >= params.max_velocity_iterations;
-
-            for constraint in &mut *joint_constraints {
-                constraint.solve(
-                    generic_joint_jacobians,
-                    &mut self.mj_lambdas[..],
-                    &mut self.generic_mj_lambdas,
-                );
-            }
+        // PERF: don’t reallocate at each iteration.
+        self.generic_solver_vels_increment = DVector::zeros(multibody_solver_id);
+        self.generic_solver_vels = DVector::zeros(multibody_solver_id);
 
-            for constraint in &mut *contact_constraints {
-                constraint.solve(
-                    generic_contact_jacobians,
-                    &mut self.mj_lambdas[..],
-                    &mut self.generic_mj_lambdas,
-                    true,
-                    false,
-                );
-            }
+        // init solver_vels for multibodies.
+        for link in &self.multibody_roots {
+            let multibody = multibodies
+                .get_multibody_mut_internal(link.multibody)
+                .unwrap();
+            multibody.update_dynamics(params.dt, bodies);
+            multibody.update_acceleration(bodies);
 
-            if solve_friction {
-                for constraint in &mut *contact_constraints {
-                    constraint.solve(
-                        generic_contact_jacobians,
-                        &mut self.mj_lambdas[..],
-                        &mut self.generic_mj_lambdas,
-                        false,
-                        true,
-                    );
-                }
-            }
+            let mut solver_vels_incr = self
+                .generic_solver_vels_increment
+                .rows_mut(multibody.solver_id, multibody.ndofs());
+            let mut solver_vels = self
+                .generic_solver_vels
+                .rows_mut(multibody.solver_id, multibody.ndofs());
+
+            solver_vels_incr.axpy(params.dt, &multibody.accelerations, 0.0);
+            solver_vels.copy_from(&multibody.velocities);
         }
+    }
 
-        let remaining_friction_iterations =
-            if !params.interleave_restitution_and_friction_resolution {
-                params.max_velocity_friction_iterations
-            } else if params.max_velocity_friction_iterations > params.max_velocity_iterations {
-                params.max_velocity_friction_iterations - params.max_velocity_iterations
-            } else {
-                0
-            };
+    pub fn solve_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        num_solver_iterations: usize,
+        bodies: &mut RigidBodySet,
+        multibodies: &mut MultibodyJointSet,
+        contact_constraints: &mut SolverConstraintsSet<ContactConstraintTypes>,
+        joint_constraints: &mut SolverConstraintsSet<JointConstraintTypes>,
+    ) {
+        for small_step_id in 0..num_solver_iterations {
+            let is_last_small_step = small_step_id == num_solver_iterations - 1;
 
-        for _ in 0..remaining_friction_iterations {
-            for constraint in &mut *contact_constraints {
-                constraint.solve(
-                    generic_contact_jacobians,
-                    &mut self.mj_lambdas[..],
-                    &mut self.generic_mj_lambdas,
-                    false,
-                    true,
-                );
+            for (solver_vels, incr) in self
+                .solver_vels
+                .iter_mut()
+                .zip(self.solver_vels_increment.iter())
+            {
+                solver_vels.linear += incr.linear;
+                solver_vels.angular += incr.angular;
             }
-        }
 
-        // Integrate positions.
-        for handle in islands.active_island(island_id) {
-            if let Some(link) = multibodies.rigid_body_link(*handle).copied() {
-                let multibody = multibodies
-                    .get_multibody_mut_internal(link.multibody)
-                    .unwrap();
+            self.generic_solver_vels += &self.generic_solver_vels_increment;
 
-                if link.id == 0 || link.id == 1 && !multibody.root_is_dynamic {
-                    let mj_lambdas = self
-                        .generic_mj_lambdas
-                        .rows(multibody.solver_id, multibody.ndofs());
-                    let prev_vels = multibody.velocities.clone(); // FIXME: avoid allocations.
-                    multibody.velocities += mj_lambdas;
-                    multibody.integrate(params.dt);
-                    multibody.forward_kinematics(bodies, false);
-                    multibody.velocities = prev_vels;
-                }
+            /*
+             * Update & solve constraints.
+             */
+            joint_constraints.update(&params, multibodies, &self.solver_bodies);
+            contact_constraints.update(&params, small_step_id, multibodies, &self.solver_bodies);
+
+            joint_constraints.solve(&mut self.solver_vels, &mut self.generic_solver_vels);
+            contact_constraints
+                .solve_restitution(&mut self.solver_vels, &mut self.generic_solver_vels);
+
+            let num_friction_iterations = if is_last_small_step {
+                params.num_friction_iteration_per_solver_iteration * num_solver_iterations
+                    - (num_solver_iterations - 1)
             } else {
-                let rb = bodies.index_mut_internal(*handle);
+                1
+            };
 
-                let dvel = self.mj_lambdas[rb.ids.active_set_offset];
-                let dangvel = rb
-                    .mprops
-                    .effective_world_inv_inertia_sqrt
-                    .transform_vector(dvel.angular);
-
-                // Update positions.
-                let mut new_pos = rb.pos;
-                let mut new_vels = rb.vels;
-                new_vels.linvel += dvel.linear;
-                new_vels.angvel += dangvel;
-                new_vels = new_vels.apply_damping(params.dt, &rb.damping);
-                new_pos.next_position = new_vels.integrate(
-                    params.dt,
-                    &rb.pos.position,
-                    &rb.mprops.local_mprops.local_com,
-                );
-                rb.integrated_vels = new_vels;
-                rb.pos = new_pos;
+            for _ in 0..num_friction_iterations {
+                contact_constraints
+                    .solve_friction(&mut self.solver_vels, &mut self.generic_solver_vels);
             }
-        }
 
-        for joint in &mut *joint_constraints {
-            joint.remove_bias_from_rhs();
+            /*
+             * Integrate positions.
+             */
+            self.integrate_positions(&params, is_last_small_step, bodies, multibodies);
+
+            /*
+             * Resolution without bias.
+             */
+            joint_constraints.solve_wo_bias(&mut self.solver_vels, &mut self.generic_solver_vels);
+            contact_constraints
+                .solve_restitution_wo_bias(&mut self.solver_vels, &mut self.generic_solver_vels);
         }
-        for constraint in &mut *contact_constraints {
-            constraint.remove_bias_from_rhs();
+    }
+
+    pub fn integrate_positions(
+        &mut self,
+        params: &IntegrationParameters,
+        is_last_small_step: bool,
+        bodies: &mut RigidBodySet,
+        multibodies: &mut MultibodyJointSet,
+    ) {
+        // Integrate positions.
+        for (solver_vels, solver_body) in self.solver_vels.iter().zip(self.solver_bodies.iter_mut())
+        {
+            let linvel = solver_vels.linear;
+            let angvel = solver_body.sqrt_ii.transform_vector(solver_vels.angular);
+
+            let mut new_vels = RigidBodyVelocity { linvel, angvel };
+            new_vels = new_vels.apply_damping(params.dt, &solver_body.damping);
+            let new_pos =
+                new_vels.integrate(params.dt, &solver_body.position, &solver_body.local_com);
+            solver_body.integrated_vels += new_vels;
+            solver_body.position = new_pos;
+            solver_body.world_com = new_pos * solver_body.local_com;
         }
 
-        for _ in 0..params.max_stabilization_iterations {
-            for constraint in &mut *joint_constraints {
-                constraint.solve(
-                    generic_joint_jacobians,
-                    &mut self.mj_lambdas[..],
-                    &mut self.generic_mj_lambdas,
-                );
-            }
+        // Integrate multibody positions.
+        for link in &self.multibody_roots {
+            let multibody = multibodies
+                .get_multibody_mut_internal(link.multibody)
+                .unwrap();
+            let solver_vels = self
+                .generic_solver_vels
+                .rows(multibody.solver_id, multibody.ndofs());
+            multibody.velocities.copy_from(&solver_vels);
+            multibody.integrate(params.dt);
+            // PERF: we could have a mode where it doesn’t write back to the `bodies` yet.
+            multibody.forward_kinematics(bodies, !is_last_small_step);
 
-            for constraint in &mut *contact_constraints {
-                constraint.solve(
-                    generic_contact_jacobians,
-                    &mut self.mj_lambdas[..],
-                    &mut self.generic_mj_lambdas,
-                    true,
-                    false,
-                );
-            }
+            if !is_last_small_step {
+                // These are very expensive and not needed if we don’t
+                // have to run another step.
+                multibody.update_dynamics(params.dt, bodies);
+                multibody.update_acceleration(bodies);
 
-            for constraint in &mut *contact_constraints {
-                constraint.solve(
-                    generic_contact_jacobians,
-                    &mut self.mj_lambdas[..],
-                    &mut self.generic_mj_lambdas,
-                    false,
-                    true,
-                );
+                let mut solver_vels_incr = self
+                    .generic_solver_vels_increment
+                    .rows_mut(multibody.solver_id, multibody.ndofs());
+                solver_vels_incr.axpy(params.dt, &multibody.accelerations, 0.0);
             }
         }
+    }
 
-        // Update velocities.
+    pub fn writeback_bodies(
+        &mut self,
+        params: &IntegrationParameters,
+        num_solver_iterations: usize,
+        islands: &IslandManager,
+        island_id: usize,
+        bodies: &mut RigidBodySet,
+        multibodies: &mut MultibodyJointSet,
+    ) {
         for handle in islands.active_island(island_id) {
-            if let Some(link) = multibodies.rigid_body_link(*handle).copied() {
+            let link = if self.multibody_roots.is_empty() {
+                None
+            } else {
+                multibodies.rigid_body_link(*handle).copied()
+            };
+
+            if let Some(link) = link {
                 let multibody = multibodies
                     .get_multibody_mut_internal(link.multibody)
                     .unwrap();
 
                 if link.id == 0 || link.id == 1 && !multibody.root_is_dynamic {
-                    let mj_lambdas = self
-                        .generic_mj_lambdas
+                    let solver_vels = self
+                        .generic_solver_vels
                         .rows(multibody.solver_id, multibody.ndofs());
-                    multibody.velocities += mj_lambdas;
+                    multibody.velocities.copy_from(&solver_vels);
                 }
             } else {
                 let rb = bodies.index_mut_internal(*handle);
-                let dvel = self.mj_lambdas[rb.ids.active_set_offset];
-                let dangvel = rb
-                    .mprops
-                    .effective_world_inv_inertia_sqrt
-                    .transform_vector(dvel.angular);
-
-                rb.vels.linvel += dvel.linear;
-                rb.vels.angvel += dangvel;
-                rb.vels = rb.vels.apply_damping(params.dt, &rb.damping);
-            }
-        }
+                let solver_body = &self.solver_bodies[rb.ids.active_set_offset];
+                let solver_vels = &self.solver_vels[rb.ids.active_set_offset];
 
-        // Write impulses back into the manifold structures.
-        for constraint in &*joint_constraints {
-            constraint.writeback_impulses(joints_all);
-        }
+                let dangvel = solver_body.sqrt_ii.transform_vector(solver_vels.angular);
 
-        for constraint in &*contact_constraints {
-            constraint.writeback_impulses(manifolds_all);
+                let mut new_vels = RigidBodyVelocity {
+                    linvel: solver_vels.linear,
+                    angvel: dangvel,
+                };
+                new_vels = new_vels.apply_damping(params.dt, &solver_body.damping);
+
+                // NOTE: using integrated_vels instead of interpolation is interesting for
+                //       high angular velocities. However, it is a bit inexact due to the
+                //       solver integrating at intermediate sub-steps. Should we just switch
+                //       to interpolation?
+                rb.integrated_vels.linvel =
+                    solver_body.integrated_vels.linvel / num_solver_iterations as Real;
+                rb.integrated_vels.angvel =
+                    solver_body.integrated_vels.angvel / num_solver_iterations as Real;
+                rb.vels = new_vels;
+                rb.pos.next_position = solver_body.position;
+            }
         }
     }
 }