Allow locking individual translational axes

3 files changed, 17 insertions, 16 deletions

diff --git a/src/dynamics/solver/joint_constraint/joint_generic_velocity_constraint_builder.rs b/src/dynamics/solver/joint_constraint/joint_generic_velocity_constraint_builder.rs
index 92bcc9f..5edc815 100644
--- a/src/dynamics/solver/joint_constraint/joint_generic_velocity_constraint_builder.rs
+++ b/src/dynamics/solver/joint_constraint/joint_generic_velocity_constraint_builder.rs
@@ -33,7 +33,7 @@ impl SolverBody<Real, 1> {
             let mut out_invm_j = jacobians.fixed_rows_mut::<SPATIAL_DIM>(wj_id);
             out_invm_j
                 .fixed_rows_mut::<DIM>(0)
-                .axpy(self.im, &unit_force, 0.0);
+                .copy_from(&self.im.component_mul(&unit_force));
 
             #[cfg(feature = "dim2")]
             {
diff --git a/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs
index 3e31256..a4ec313 100644
--- a/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs
@@ -51,7 +51,7 @@ pub enum WritebackId {
 pub struct SolverBody<N: SimdRealField, const LANES: usize> {
     pub linvel: Vector<N>,
     pub angvel: AngVector<N>,
-    pub im: N,
+    pub im: Vector<N>,
     pub sqrt_ii: AngularInertia<N>,
     pub world_com: Point<N>,
     pub mj_lambda: [usize; LANES],
@@ -74,8 +74,8 @@ pub struct JointVelocityConstraint<N: SimdRealField, const LANES: usize> {
     pub rhs: N,
     pub rhs_wo_bias: N,
 
-    pub im1: N,
-    pub im2: N,
+    pub im1: Vector<N>,
+    pub im2: Vector<N>,
 
     pub writeback_id: WritebackId,
 }
@@ -94,8 +94,8 @@ impl<N: WReal, const LANES: usize> JointVelocityConstraint<N, LANES> {
             inv_lhs: N::zero(),
             rhs: N::zero(),
             rhs_wo_bias: N::zero(),
-            im1: N::zero(),
-            im2: N::zero(),
+            im1: na::zero(),
+            im2: na::zero(),
             writeback_id: WritebackId::Dof(0),
         }
     }
@@ -115,9 +115,9 @@ impl<N: WReal, const LANES: usize> JointVelocityConstraint<N, LANES> {
         let ang_impulse1 = self.ang_jac1 * delta_impulse;
         let ang_impulse2 = self.ang_jac2 * delta_impulse;
 
-        mj_lambda1.linear += lin_impulse * self.im1;
+        mj_lambda1.linear += lin_impulse.component_mul(&self.im1);
         mj_lambda1.angular += ang_impulse1;
-        mj_lambda2.linear -= lin_impulse * self.im2;
+        mj_lambda2.linear -= lin_impulse.component_mul(&self.im2);
         mj_lambda2.angular -= ang_impulse2;
     }
 
@@ -353,7 +353,7 @@ pub struct JointVelocityGroundConstraint<N: SimdRealField, const LANES: usize> {
     pub rhs: N,
     pub rhs_wo_bias: N,
 
-    pub im2: N,
+    pub im2: Vector<N>,
 
     pub writeback_id: WritebackId,
 }
@@ -370,7 +370,7 @@ impl<N: WReal, const LANES: usize> JointVelocityGroundConstraint<N, LANES> {
             inv_lhs: N::zero(),
             rhs: N::zero(),
             rhs_wo_bias: N::zero(),
-            im2: N::zero(),
+            im2: na::zero(),
             writeback_id: WritebackId::Dof(0),
         }
     }
@@ -388,7 +388,7 @@ impl<N: WReal, const LANES: usize> JointVelocityGroundConstraint<N, LANES> {
         let lin_impulse = self.lin_jac * delta_impulse;
         let ang_impulse = self.ang_jac2 * delta_impulse;
 
-        mj_lambda2.linear -= lin_impulse * self.im2;
+        mj_lambda2.linear -= lin_impulse.component_mul(&self.im2);
         mj_lambda2.angular -= ang_impulse;
     }
 
diff --git a/src/dynamics/solver/joint_constraint/joint_velocity_constraint_builder.rs b/src/dynamics/solver/joint_constraint/joint_velocity_constraint_builder.rs
index dd17ecf..55a112a 100644
--- a/src/dynamics/solver/joint_constraint/joint_velocity_constraint_builder.rs
+++ b/src/dynamics/solver/joint_constraint/joint_velocity_constraint_builder.rs
@@ -355,7 +355,7 @@ impl<N: WReal> JointVelocityConstraintBuilder<N> {
         // Use the modified Gram-Schmidt orthogonalization.
         for j in 0..len {
             let c_j = &mut constraints[j];
-            let dot_jj = c_j.lin_jac.norm_squared() * imsum
+            let dot_jj = c_j.lin_jac.dot(&imsum.component_mul(&c_j.lin_jac))
                 + c_j.ang_jac1.gdot(c_j.ang_jac1)
                 + c_j.ang_jac2.gdot(c_j.ang_jac2);
             let inv_dot_jj = crate::utils::simd_inv(dot_jj);
@@ -370,7 +370,7 @@ impl<N: WReal> JointVelocityConstraintBuilder<N> {
 
             for i in (j + 1)..len {
                 let (c_i, c_j) = constraints.index_mut_const(i, j);
-                let dot_ij = c_i.lin_jac.dot(&c_j.lin_jac) * imsum
+                let dot_ij = c_i.lin_jac.dot(&imsum.component_mul(&c_j.lin_jac))
                     + c_i.ang_jac1.gdot(c_j.ang_jac1)
                     + c_i.ang_jac2.gdot(c_j.ang_jac2);
                 let coeff = dot_ij * inv_dot_jj;
@@ -672,7 +672,8 @@ impl<N: WReal> JointVelocityConstraintBuilder<N> {
         // Use the modified Gram-Schmidt orthogonalization.
         for j in 0..len {
             let c_j = &mut constraints[j];
-            let dot_jj = c_j.lin_jac.norm_squared() * imsum + c_j.ang_jac2.gdot(c_j.ang_jac2);
+            let dot_jj = c_j.lin_jac.dot(&imsum.component_mul(&c_j.lin_jac))
+                + c_j.ang_jac2.gdot(c_j.ang_jac2);
             let inv_dot_jj = crate::utils::simd_inv(dot_jj);
             c_j.inv_lhs = inv_dot_jj; // Don’t forget to update the inv_lhs.
 
@@ -685,8 +686,8 @@ impl<N: WReal> JointVelocityConstraintBuilder<N> {
 
             for i in (j + 1)..len {
                 let (c_i, c_j) = constraints.index_mut_const(i, j);
-                let dot_ij =
-                    c_i.lin_jac.dot(&c_j.lin_jac) * imsum + c_i.ang_jac2.gdot(c_j.ang_jac2);
+                let dot_ij = c_i.lin_jac.dot(&imsum.component_mul(&c_j.lin_jac))
+                    + c_i.ang_jac2.gdot(c_j.ang_jac2);
                 let coeff = dot_ij * inv_dot_jj;
 
                 c_i.lin_jac -= c_j.lin_jac * coeff;