From 9b87f06a856c4d673642e210f8b0986cfdbac3af Mon Sep 17 00:00:00 2001 From: Sébastien Crozet Date: Sun, 21 Jan 2024 21:02:23 +0100 Subject: feat: implement new "small-steps" solver + joint improvements --- .../joint_constraint/any_joint_constraint.rs | 97 ++ .../solver/joint_constraint/joint_constraint.rs | 540 -------- .../joint_constraint/joint_constraint_builder.rs | 1401 ++++++++++++++++++++ .../joint_constraint/joint_constraints_set.rs | 446 +++++++ .../joint_constraint/joint_generic_constraint.rs | 552 ++++++++ .../joint_generic_constraint_builder.rs | 1286 ++++++++++++++++++ .../joint_generic_velocity_constraint.rs | 562 -------- .../joint_generic_velocity_constraint_builder.rs | 886 ------------- .../joint_constraint/joint_velocity_constraint.rs | 387 +++--- .../joint_velocity_constraint_builder.rs | 1127 ---------------- src/dynamics/solver/joint_constraint/mod.rs | 23 +- 11 files changed, 4014 insertions(+), 3293 deletions(-) create mode 100644 src/dynamics/solver/joint_constraint/any_joint_constraint.rs delete mode 100644 src/dynamics/solver/joint_constraint/joint_constraint.rs create mode 100644 src/dynamics/solver/joint_constraint/joint_constraint_builder.rs create mode 100644 src/dynamics/solver/joint_constraint/joint_constraints_set.rs create mode 100644 src/dynamics/solver/joint_constraint/joint_generic_constraint.rs create mode 100644 src/dynamics/solver/joint_constraint/joint_generic_constraint_builder.rs delete mode 100644 src/dynamics/solver/joint_constraint/joint_generic_velocity_constraint.rs delete mode 100644 src/dynamics/solver/joint_constraint/joint_generic_velocity_constraint_builder.rs delete mode 100644 src/dynamics/solver/joint_constraint/joint_velocity_constraint_builder.rs (limited to 'src/dynamics/solver/joint_constraint') diff --git a/src/dynamics/solver/joint_constraint/any_joint_constraint.rs b/src/dynamics/solver/joint_constraint/any_joint_constraint.rs new file mode 100644 index 0000000..8cbd438 --- /dev/null +++ b/src/dynamics/solver/joint_constraint/any_joint_constraint.rs @@ -0,0 +1,97 @@ +use crate::dynamics::solver::joint_constraint::joint_generic_constraint::{ + JointGenericOneBodyConstraint, JointGenericTwoBodyConstraint, +}; +use crate::dynamics::solver::joint_constraint::joint_velocity_constraint::{ + JointOneBodyConstraint, JointTwoBodyConstraint, +}; +use crate::dynamics::solver::{AnyConstraintMut, ConstraintTypes}; +use crate::dynamics::JointGraphEdge; +use crate::math::Real; +use na::DVector; + +use crate::dynamics::solver::joint_constraint::joint_generic_constraint_builder::{ + JointGenericOneBodyConstraintBuilder, JointGenericTwoBodyConstraintBuilder, +}; +use crate::dynamics::solver::solver_vel::SolverVel; +#[cfg(feature = "simd-is-enabled")] +use crate::{ + dynamics::solver::joint_constraint::joint_constraint_builder::{ + JointOneBodyConstraintBuilderSimd, JointTwoBodyConstraintBuilderSimd, + }, + math::{SimdReal, SIMD_WIDTH}, +}; + +use crate::dynamics::solver::joint_constraint::joint_constraint_builder::{ + JointOneBodyConstraintBuilder, JointTwoBodyConstraintBuilder, +}; + +pub struct JointConstraintTypes; + +impl<'a> AnyConstraintMut<'a, JointConstraintTypes> { + pub fn remove_bias(&mut self) { + match self { + Self::OneBody(c) => c.remove_bias_from_rhs(), + Self::TwoBodies(c) => c.remove_bias_from_rhs(), + Self::GenericOneBody(c) => c.remove_bias_from_rhs(), + Self::GenericTwoBodies(c) => c.remove_bias_from_rhs(), + #[cfg(feature = "simd-is-enabled")] + Self::SimdOneBody(c) => c.remove_bias_from_rhs(), + #[cfg(feature = "simd-is-enabled")] + Self::SimdTwoBodies(c) => c.remove_bias_from_rhs(), + } + } + + pub fn solve( + &mut self, + generic_jacobians: &DVector, + solver_vels: &mut [SolverVel], + generic_solver_vels: &mut DVector, + ) { + match self { + Self::OneBody(c) => c.solve(solver_vels), + Self::TwoBodies(c) => c.solve(solver_vels), + Self::GenericOneBody(c) => c.solve(generic_jacobians, solver_vels, generic_solver_vels), + Self::GenericTwoBodies(c) => { + c.solve(generic_jacobians, solver_vels, generic_solver_vels) + } + #[cfg(feature = "simd-is-enabled")] + Self::SimdOneBody(c) => c.solve(solver_vels), + #[cfg(feature = "simd-is-enabled")] + Self::SimdTwoBodies(c) => c.solve(solver_vels), + } + } + + pub fn writeback_impulses(&mut self, joints_all: &mut [JointGraphEdge]) { + match self { + Self::OneBody(c) => c.writeback_impulses(joints_all), + Self::TwoBodies(c) => c.writeback_impulses(joints_all), + Self::GenericOneBody(c) => c.writeback_impulses(joints_all), + Self::GenericTwoBodies(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "simd-is-enabled")] + Self::SimdOneBody(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "simd-is-enabled")] + Self::SimdTwoBodies(c) => c.writeback_impulses(joints_all), + } + } +} + +impl ConstraintTypes for JointConstraintTypes { + type OneBody = JointOneBodyConstraint; + type TwoBodies = JointTwoBodyConstraint; + type GenericOneBody = JointGenericOneBodyConstraint; + type GenericTwoBodies = JointGenericTwoBodyConstraint; + + #[cfg(feature = "simd-is-enabled")] + type SimdOneBody = JointOneBodyConstraint; + #[cfg(feature = "simd-is-enabled")] + type SimdTwoBodies = JointTwoBodyConstraint; + + type BuilderOneBody = JointOneBodyConstraintBuilder; + type BuilderTwoBodies = JointTwoBodyConstraintBuilder; + type GenericBuilderOneBody = JointGenericOneBodyConstraintBuilder; + type GenericBuilderTwoBodies = JointGenericTwoBodyConstraintBuilder; + #[cfg(feature = "simd-is-enabled")] + type SimdBuilderOneBody = JointOneBodyConstraintBuilderSimd; + #[cfg(feature = "simd-is-enabled")] + type SimdBuilderTwoBodies = JointTwoBodyConstraintBuilderSimd; +} diff --git a/src/dynamics/solver/joint_constraint/joint_constraint.rs b/src/dynamics/solver/joint_constraint/joint_constraint.rs deleted file mode 100644 index 962602f..0000000 --- a/src/dynamics/solver/joint_constraint/joint_constraint.rs +++ /dev/null @@ -1,540 +0,0 @@ -use crate::dynamics::solver::joint_constraint::joint_generic_velocity_constraint::{ - JointGenericVelocityConstraint, JointGenericVelocityGroundConstraint, -}; -use crate::dynamics::solver::joint_constraint::joint_velocity_constraint::{ - JointVelocityConstraint, JointVelocityGroundConstraint, SolverBody, -}; -use crate::dynamics::solver::DeltaVel; -use crate::dynamics::{ - ImpulseJoint, IntegrationParameters, JointGraphEdge, JointIndex, RigidBodySet, -}; -use crate::math::{Real, SPATIAL_DIM}; -use crate::prelude::MultibodyJointSet; -use na::DVector; - -#[cfg(feature = "simd-is-enabled")] -use crate::math::{Isometry, SimdReal, SIMD_WIDTH}; - -#[cfg(feature = "parallel")] -use crate::dynamics::JointAxesMask; - -pub enum AnyJointVelocityConstraint { - JointConstraint(JointVelocityConstraint), - JointGroundConstraint(JointVelocityGroundConstraint), - JointGenericConstraint(JointGenericVelocityConstraint), - JointGenericGroundConstraint(JointGenericVelocityGroundConstraint), - #[cfg(feature = "simd-is-enabled")] - JointConstraintSimd(JointVelocityConstraint), - #[cfg(feature = "simd-is-enabled")] - JointGroundConstraintSimd(JointVelocityGroundConstraint), - Empty, -} - -impl AnyJointVelocityConstraint { - #[cfg(feature = "parallel")] - pub fn num_active_constraints_and_jacobian_lines(joint: &ImpulseJoint) -> (usize, usize) { - let joint = &joint.data; - let locked_axes = joint.locked_axes.bits(); - let motor_axes = joint.motor_axes.bits() & !locked_axes; - let limit_axes = joint.limit_axes.bits() & !locked_axes; - let coupled_axes = joint.coupled_axes.bits(); - - let num_constraints = (motor_axes & !coupled_axes).count_ones() as usize - + ((motor_axes & coupled_axes) & JointAxesMask::ANG_AXES.bits() != 0) as usize - + ((motor_axes & coupled_axes) & JointAxesMask::LIN_AXES.bits() != 0) as usize - + locked_axes.count_ones() as usize - + (limit_axes & !coupled_axes).count_ones() as usize - + ((limit_axes & coupled_axes) & JointAxesMask::ANG_AXES.bits() != 0) as usize - + ((limit_axes & coupled_axes) & JointAxesMask::LIN_AXES.bits() != 0) as usize; - (num_constraints, num_constraints) - } - - pub fn from_joint( - params: &IntegrationParameters, - joint_id: JointIndex, - joint: &ImpulseJoint, - bodies: &RigidBodySet, - multibodies: &MultibodyJointSet, - j_id: &mut usize, - jacobians: &mut DVector, - out: &mut Vec, - insert_at: Option, - ) { - let local_frame1 = joint.data.local_frame1; - let local_frame2 = joint.data.local_frame2; - let rb1 = &bodies[joint.body1]; - let rb2 = &bodies[joint.body2]; - let frame1 = rb1.pos.position * local_frame1; - let frame2 = rb2.pos.position * local_frame2; - - let body1 = SolverBody { - linvel: rb1.vels.linvel, - angvel: rb1.vels.angvel, - im: rb1.mprops.effective_inv_mass, - sqrt_ii: rb1.mprops.effective_world_inv_inertia_sqrt, - world_com: rb1.mprops.world_com, - mj_lambda: [rb1.ids.active_set_offset], - }; - let body2 = SolverBody { - linvel: rb2.vels.linvel, - angvel: rb2.vels.angvel, - im: rb2.mprops.effective_inv_mass, - sqrt_ii: rb2.mprops.effective_world_inv_inertia_sqrt, - world_com: rb2.mprops.world_com, - mj_lambda: [rb2.ids.active_set_offset], - }; - - let mb1 = multibodies - .rigid_body_link(joint.body1) - .map(|link| (&multibodies[link.multibody], link.id)); - let mb2 = multibodies - .rigid_body_link(joint.body2) - .map(|link| (&multibodies[link.multibody], link.id)); - - if mb1.is_some() || mb2.is_some() { - let multibodies_ndof = mb1.map(|m| m.0.ndofs()).unwrap_or(SPATIAL_DIM) - + mb2.map(|m| m.0.ndofs()).unwrap_or(SPATIAL_DIM); - - if multibodies_ndof == 0 { - // Both multibodies are fixed, don’t generate any constraint. - return; - } - - // For each solver contact we generate up to SPATIAL_DIM constraints, and each - // constraints appends the multibodies jacobian and weighted jacobians. - // Also note that for impulse_joints, the rigid-bodies will also add their jacobians - // to the generic DVector. - // TODO: is this count correct when we take both motors and limits into account? - let required_jacobian_len = *j_id + multibodies_ndof * 2 * SPATIAL_DIM; - - if jacobians.nrows() < required_jacobian_len && !cfg!(feature = "parallel") { - jacobians.resize_vertically_mut(required_jacobian_len, 0.0); - } - - // TODO: find a way to avoid the temporary buffer. - let mut out_tmp = [JointGenericVelocityConstraint::invalid(); 12]; - let out_tmp_len = JointGenericVelocityConstraint::lock_axes( - params, - joint_id, - &body1, - &body2, - mb1, - mb2, - &frame1, - &frame2, - &joint.data, - jacobians, - j_id, - &mut out_tmp, - ); - - if let Some(at) = insert_at { - for (i, c) in out_tmp.into_iter().take(out_tmp_len).enumerate() { - out[at + i] = AnyJointVelocityConstraint::JointGenericConstraint(c); - } - } else { - for c in out_tmp.into_iter().take(out_tmp_len) { - out.push(AnyJointVelocityConstraint::JointGenericConstraint(c)); - } - } - } else { - // TODO: find a way to avoid the temporary buffer. - let mut out_tmp = [JointVelocityConstraint::invalid(); 12]; - let out_tmp_len = JointVelocityConstraint::::lock_axes( - params, - joint_id, - &body1, - &body2, - &frame1, - &frame2, - &joint.data, - &mut out_tmp, - ); - - if let Some(at) = insert_at { - for (i, c) in out_tmp.into_iter().take(out_tmp_len).enumerate() { - out[at + i] = AnyJointVelocityConstraint::JointConstraint(c); - } - } else { - for c in out_tmp.into_iter().take(out_tmp_len) { - out.push(AnyJointVelocityConstraint::JointConstraint(c)); - } - } - } - } - - #[cfg(feature = "simd-is-enabled")] - pub fn from_wide_joint( - params: &IntegrationParameters, - joint_id: [JointIndex; SIMD_WIDTH], - impulse_joints: [&ImpulseJoint; SIMD_WIDTH], - bodies: &RigidBodySet, - out: &mut Vec, - insert_at: Option, - ) { - use crate::dynamics::{ - RigidBodyIds, RigidBodyMassProps, RigidBodyPosition, RigidBodyVelocity, - }; - - let rbs1: ( - [&RigidBodyPosition; SIMD_WIDTH], - [&RigidBodyVelocity; SIMD_WIDTH], - [&RigidBodyMassProps; SIMD_WIDTH], - [&RigidBodyIds; SIMD_WIDTH], - ) = ( - gather![|ii| &bodies[impulse_joints[ii].body1].pos], - gather![|ii| &bodies[impulse_joints[ii].body1].vels], - gather![|ii| &bodies[impulse_joints[ii].body1].mprops], - gather![|ii| &bodies[impulse_joints[ii].body1].ids], - ); - let rbs2: ( - [&RigidBodyPosition; SIMD_WIDTH], - [&RigidBodyVelocity; SIMD_WIDTH], - [&RigidBodyMassProps; SIMD_WIDTH], - [&RigidBodyIds; SIMD_WIDTH], - ) = ( - gather![|ii| &bodies[impulse_joints[ii].body2].pos], - gather![|ii| &bodies[impulse_joints[ii].body2].vels], - gather![|ii| &bodies[impulse_joints[ii].body2].mprops], - gather![|ii| &bodies[impulse_joints[ii].body2].ids], - ); - - let (rb_pos1, rb_vel1, rb_mprops1, rb_ids1) = rbs1; - let (rb_pos2, rb_vel2, rb_mprops2, rb_ids2) = rbs2; - let pos1: Isometry = gather![|ii| rb_pos1[ii].position].into(); - let pos2: Isometry = gather![|ii| rb_pos2[ii].position].into(); - - let local_frame1: Isometry = - gather![|ii| impulse_joints[ii].data.local_frame1].into(); - let local_frame2: Isometry = - gather![|ii| impulse_joints[ii].data.local_frame2].into(); - - let frame1 = pos1 * local_frame1; - let frame2 = pos2 * local_frame2; - - let body1: SolverBody = SolverBody { - linvel: gather![|ii| rb_vel1[ii].linvel].into(), - angvel: gather![|ii| rb_vel1[ii].angvel].into(), - im: gather![|ii| rb_mprops1[ii].effective_inv_mass].into(), - sqrt_ii: gather![|ii| rb_mprops1[ii].effective_world_inv_inertia_sqrt].into(), - world_com: gather![|ii| rb_mprops1[ii].world_com].into(), - mj_lambda: gather![|ii| rb_ids1[ii].active_set_offset], - }; - let body2: SolverBody = SolverBody { - linvel: gather![|ii| rb_vel2[ii].linvel].into(), - angvel: gather![|ii| rb_vel2[ii].angvel].into(), - im: gather![|ii| rb_mprops2[ii].effective_inv_mass].into(), - sqrt_ii: gather![|ii| rb_mprops2[ii].effective_world_inv_inertia_sqrt].into(), - world_com: gather![|ii| rb_mprops2[ii].world_com].into(), - mj_lambda: gather![|ii| rb_ids2[ii].active_set_offset], - }; - - // TODO: find a way to avoid the temporary buffer. - let mut out_tmp = [JointVelocityConstraint::invalid(); 12]; - let out_tmp_len = JointVelocityConstraint::::lock_axes( - params, - joint_id, - &body1, - &body2, - &frame1, - &frame2, - impulse_joints[0].data.locked_axes.bits(), - &mut out_tmp, - ); - - if let Some(at) = insert_at { - for (i, c) in out_tmp.into_iter().take(out_tmp_len).enumerate() { - out[at + i] = AnyJointVelocityConstraint::JointConstraintSimd(c); - } - } else { - for c in out_tmp.into_iter().take(out_tmp_len) { - out.push(AnyJointVelocityConstraint::JointConstraintSimd(c)); - } - } - } - - pub fn from_joint_ground( - params: &IntegrationParameters, - joint_id: JointIndex, - joint: &ImpulseJoint, - bodies: &RigidBodySet, - multibodies: &MultibodyJointSet, - j_id: &mut usize, - jacobians: &mut DVector, - out: &mut Vec, - insert_at: Option, - ) { - let mut handle1 = joint.body1; - let mut handle2 = joint.body2; - let flipped = !bodies[handle2].is_dynamic(); - - let (local_frame1, local_frame2) = if flipped { - std::mem::swap(&mut handle1, &mut handle2); - (joint.data.local_frame2, joint.data.local_frame1) - } else { - (joint.data.local_frame1, joint.data.local_frame2) - }; - - let rb1 = &bodies[handle1]; - let rb2 = &bodies[handle2]; - - let frame1 = rb1.pos.position * local_frame1; - let frame2 = rb2.pos.position * local_frame2; - - let body1 = SolverBody { - linvel: rb1.vels.linvel, - angvel: rb1.vels.angvel, - im: rb1.mprops.effective_inv_mass, - sqrt_ii: rb1.mprops.effective_world_inv_inertia_sqrt, - world_com: rb1.mprops.world_com, - mj_lambda: [crate::INVALID_USIZE], - }; - let body2 = SolverBody { - linvel: rb2.vels.linvel, - angvel: rb2.vels.angvel, - im: rb2.mprops.effective_inv_mass, - sqrt_ii: rb2.mprops.effective_world_inv_inertia_sqrt, - world_com: rb2.mprops.world_com, - mj_lambda: [rb2.ids.active_set_offset], - }; - - if let Some(mb2) = multibodies - .rigid_body_link(handle2) - .map(|link| (&multibodies[link.multibody], link.id)) - { - let multibodies_ndof = mb2.0.ndofs(); - - if multibodies_ndof == 0 { - // The multibody is fixed, don’t generate any constraint. - return; - } - - // For each solver contact we generate up to SPATIAL_DIM constraints, and each - // constraints appends the multibodies jacobian and weighted jacobians. - // Also note that for impulse_joints, the rigid-bodies will also add their jacobians - // to the generic DVector. - // TODO: is this count correct when we take both motors and limits into account? - let required_jacobian_len = *j_id + multibodies_ndof * 2 * SPATIAL_DIM; - - if jacobians.nrows() < required_jacobian_len && !cfg!(feature = "parallel") { - jacobians.resize_vertically_mut(required_jacobian_len, 0.0); - } - - // TODO: find a way to avoid the temporary buffer. - let mut out_tmp = [JointGenericVelocityGroundConstraint::invalid(); 12]; - let out_tmp_len = JointGenericVelocityGroundConstraint::lock_axes( - params, - joint_id, - &body1, - &body2, - mb2, - &frame1, - &frame2, - &joint.data, - jacobians, - j_id, - &mut out_tmp, - ); - - if let Some(at) = insert_at { - for (i, c) in out_tmp.into_iter().take(out_tmp_len).enumerate() { - out[at + i] = AnyJointVelocityConstraint::JointGenericGroundConstraint(c); - } - } else { - for c in out_tmp.into_iter().take(out_tmp_len) { - out.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(c)); - } - } - } else { - // TODO: find a way to avoid the temporary buffer. - let mut out_tmp = [JointVelocityGroundConstraint::invalid(); 12]; - let out_tmp_len = JointVelocityGroundConstraint::::lock_axes( - params, - joint_id, - &body1, - &body2, - &frame1, - &frame2, - &joint.data, - &mut out_tmp, - ); - - if let Some(at) = insert_at { - for (i, c) in out_tmp.into_iter().take(out_tmp_len).enumerate() { - out[at + i] = AnyJointVelocityConstraint::JointGroundConstraint(c); - } - } else { - for c in out_tmp.into_iter().take(out_tmp_len) { - out.push(AnyJointVelocityConstraint::JointGroundConstraint(c)); - } - } - } - } - - #[cfg(feature = "simd-is-enabled")] - pub fn from_wide_joint_ground( - params: &IntegrationParameters, - joint_id: [JointIndex; SIMD_WIDTH], - impulse_joints: [&ImpulseJoint; SIMD_WIDTH], - bodies: &RigidBodySet, - out: &mut Vec, - insert_at: Option, - ) { - use crate::dynamics::{ - RigidBodyIds, RigidBodyMassProps, RigidBodyPosition, RigidBodyType, RigidBodyVelocity, - }; - - let mut handles1 = gather![|ii| impulse_joints[ii].body1]; - let mut handles2 = gather![|ii| impulse_joints[ii].body2]; - let status2: [&RigidBodyType; SIMD_WIDTH] = gather![|ii| &bodies[handles2[ii]].body_type]; - let mut flipped = [false; SIMD_WIDTH]; - - for ii in 0..SIMD_WIDTH { - if !status2[ii].is_dynamic() { - std::mem::swap(&mut handles1[ii], &mut handles2[ii]); - flipped[ii] = true; - } - } - - let local_frame1: Isometry = gather![|ii| if flipped[ii] { - impulse_joints[ii].data.local_frame2 - } else { - impulse_joints[ii].data.local_frame1 - }] - .into(); - let local_frame2: Isometry = gather![|ii| if flipped[ii] { - impulse_joints[ii].data.local_frame1 - } else { - impulse_joints[ii].data.local_frame2 - }] - .into(); - - let rbs1: ( - [&RigidBodyPosition; SIMD_WIDTH], - [&RigidBodyVelocity; SIMD_WIDTH], - [&RigidBodyMassProps; SIMD_WIDTH], - ) = ( - gather![|ii| &bodies[handles1[ii]].pos], - gather![|ii| &bodies[handles1[ii]].vels], - gather![|ii| &bodies[handles1[ii]].mprops], - ); - let rbs2: ( - [&RigidBodyPosition; SIMD_WIDTH], - [&RigidBodyVelocity; SIMD_WIDTH], - [&RigidBodyMassProps; SIMD_WIDTH], - [&RigidBodyIds; SIMD_WIDTH], - ) = ( - gather![|ii| &bodies[handles2[ii]].pos], - gather![|ii| &bodies[handles2[ii]].vels], - gather![|ii| &bodies[handles2[ii]].mprops], - gather![|ii| &bodies[handles2[ii]].ids], - ); - - let (rb_pos1, rb_vel1, rb_mprops1) = rbs1; - let (rb_pos2, rb_vel2, rb_mprops2, rb_ids2) = rbs2; - let pos1: Isometry = gather![|ii| rb_pos1[ii].position].into(); - let pos2: Isometry = gather![|ii| rb_pos2[ii].position].into(); - - let frame1 = pos1 * local_frame1; - let frame2 = pos2 * local_frame2; - - let body1: SolverBody = SolverBody { - linvel: gather![|ii| rb_vel1[ii].linvel].into(), - angvel: gather![|ii| rb_vel1[ii].angvel].into(), - im: gather![|ii| rb_mprops1[ii].effective_inv_mass].into(), - sqrt_ii: gather![|ii| rb_mprops1[ii].effective_world_inv_inertia_sqrt].into(), - world_com: gather![|ii| rb_mprops1[ii].world_com].into(), - mj_lambda: [crate::INVALID_USIZE; SIMD_WIDTH], - }; - let body2: SolverBody = SolverBody { - linvel: gather![|ii| rb_vel2[ii].linvel].into(), - angvel: gather![|ii| rb_vel2[ii].angvel].into(), - im: gather![|ii| rb_mprops2[ii].effective_inv_mass].into(), - sqrt_ii: gather![|ii| rb_mprops2[ii].effective_world_inv_inertia_sqrt].into(), - world_com: gather![|ii| rb_mprops2[ii].world_com].into(), - mj_lambda: gather![|ii| rb_ids2[ii].active_set_offset], - }; - - // TODO: find a way to avoid the temporary buffer. - let mut out_tmp = [JointVelocityGroundConstraint::invalid(); 12]; - let out_tmp_len = JointVelocityGroundConstraint::::lock_axes( - params, - joint_id, - &body1, - &body2, - &frame1, - &frame2, - impulse_joints[0].data.locked_axes.bits(), - &mut out_tmp, - ); - - if let Some(at) = insert_at { - for (i, c) in out_tmp.into_iter().take(out_tmp_len).enumerate() { - out[at + i] = AnyJointVelocityConstraint::JointGroundConstraintSimd(c); - } - } else { - for c in out_tmp.into_iter().take(out_tmp_len) { - out.push(AnyJointVelocityConstraint::JointGroundConstraintSimd(c)); - } - } - } - - pub fn remove_bias_from_rhs(&mut self) { - match self { - AnyJointVelocityConstraint::JointConstraint(c) => c.remove_bias_from_rhs(), - AnyJointVelocityConstraint::JointGroundConstraint(c) => c.remove_bias_from_rhs(), - #[cfg(feature = "simd-is-enabled")] - AnyJointVelocityConstraint::JointConstraintSimd(c) => c.remove_bias_from_rhs(), - #[cfg(feature = "simd-is-enabled")] - AnyJointVelocityConstraint::JointGroundConstraintSimd(c) => c.remove_bias_from_rhs(), - AnyJointVelocityConstraint::JointGenericConstraint(c) => c.remove_bias_from_rhs(), - AnyJointVelocityConstraint::JointGenericGroundConstraint(c) => c.remove_bias_from_rhs(), - AnyJointVelocityConstraint::Empty => unreachable!(), - } - } - - pub fn solve( - &mut self, - jacobians: &DVector, - mj_lambdas: &mut [DeltaVel], - generic_mj_lambdas: &mut DVector, - ) { - match self { - AnyJointVelocityConstraint::JointConstraint(c) => c.solve(mj_lambdas), - AnyJointVelocityConstraint::JointGroundConstraint(c) => c.solve(mj_lambdas), - #[cfg(feature = "simd-is-enabled")] - AnyJointVelocityConstraint::JointConstraintSimd(c) => c.solve(mj_lambdas), - #[cfg(feature = "simd-is-enabled")] - AnyJointVelocityConstraint::JointGroundConstraintSimd(c) => c.solve(mj_lambdas), - AnyJointVelocityConstraint::JointGenericConstraint(c) => { - c.solve(jacobians, mj_lambdas, generic_mj_lambdas) - } - AnyJointVelocityConstraint::JointGenericGroundConstraint(c) => { - c.solve(jacobians, mj_lambdas, generic_mj_lambdas) - } - AnyJointVelocityConstraint::Empty => unreachable!(), - } - } - - pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) { - match self { - AnyJointVelocityConstraint::JointConstraint(c) => c.writeback_impulses(joints_all), - AnyJointVelocityConstraint::JointGroundConstraint(c) => { - c.writeback_impulses(joints_all) - } - #[cfg(feature = "simd-is-enabled")] - AnyJointVelocityConstraint::JointConstraintSimd(c) => c.writeback_impulses(joints_all), - #[cfg(feature = "simd-is-enabled")] - AnyJointVelocityConstraint::JointGroundConstraintSimd(c) => { - c.writeback_impulses(joints_all) - } - AnyJointVelocityConstraint::JointGenericConstraint(c) => { - c.writeback_impulses(joints_all) - } - AnyJointVelocityConstraint::JointGenericGroundConstraint(c) => { - c.writeback_impulses(joints_all) - } - AnyJointVelocityConstraint::Empty => unreachable!(), - } - } -} diff --git a/src/dynamics/solver/joint_constraint/joint_constraint_builder.rs b/src/dynamics/solver/joint_constraint/joint_constraint_builder.rs new file mode 100644 index 0000000..00bead1 --- /dev/null +++ b/src/dynamics/solver/joint_constraint/joint_constraint_builder.rs @@ -0,0 +1,1401 @@ +use crate::dynamics::solver::joint_constraint::joint_velocity_constraint::{ + JointFixedSolverBody, JointOneBodyConstraint, JointTwoBodyConstraint, WritebackId, +}; +use crate::dynamics::solver::joint_constraint::JointSolverBody; +use crate::dynamics::solver::solver_body::SolverBody; +use crate::dynamics::solver::ConstraintsCounts; +use crate::dynamics::solver::MotorParameters; +use crate::dynamics::{GenericJoint, ImpulseJoint, IntegrationParameters, JointIndex}; +use crate::math::{AngVector, Isometry, Matrix, Point, Real, Rotation, Vector, ANG_DIM, DIM}; +use crate::prelude::RigidBodySet; +use crate::utils; +use crate::utils::{IndexMut2, SimdCrossMatrix, SimdDot, SimdQuat, SimdRealCopy}; +use na::SMatrix; + +#[cfg(feature = "dim3")] +use crate::utils::SimdBasis; + +#[cfg(feature = "simd-is-enabled")] +use crate::math::{SimdReal, SIMD_WIDTH}; + +pub struct JointTwoBodyConstraintBuilder { + body1: usize, + body2: usize, + joint_id: JointIndex, + joint: GenericJoint, + constraint_id: usize, +} + +impl JointTwoBodyConstraintBuilder { + pub fn generate( + joint: &ImpulseJoint, + bodies: &RigidBodySet, + joint_id: JointIndex, + out_builder: &mut Self, + out_constraint_id: &mut usize, + ) { + let rb1 = &bodies[joint.body1]; + let rb2 = &bodies[joint.body2]; + *out_builder = Self { + body1: rb1.ids.active_set_offset, + body2: rb2.ids.active_set_offset, + joint_id, + joint: joint.data, + constraint_id: *out_constraint_id, + }; + + let count = ConstraintsCounts::from_joint(joint); + *out_constraint_id += count.num_constraints; + } + + pub fn update( + &self, + params: &IntegrationParameters, + bodies: &[SolverBody], + out: &mut [JointTwoBodyConstraint], + ) { + // NOTE: right now, the "update", is basically reconstructing all the + // constraints. Could we make this more incremental? + + let rb1 = &bodies[self.body1]; + let rb2 = &bodies[self.body2]; + let frame1 = rb1.position * self.joint.local_frame1; + let frame2 = rb2.position * self.joint.local_frame2; + + let joint_body1 = JointSolverBody { + im: rb1.im, + sqrt_ii: rb1.sqrt_ii, + world_com: rb1.world_com, + solver_vel: [self.body1], + }; + let joint_body2 = JointSolverBody { + im: rb2.im, + sqrt_ii: rb2.sqrt_ii, + world_com: rb2.world_com, + solver_vel: [self.body2], + }; + + JointTwoBodyConstraint::::lock_axes( + params, + self.joint_id, + &joint_body1, + &joint_body2, + &frame1, + &frame2, + &self.joint, + &mut out[self.constraint_id..], + ); + } +} + +#[cfg(feature = "simd-is-enabled")] +pub struct JointTwoBodyConstraintBuilderSimd { + body1: [usize; SIMD_WIDTH], + body2: [usize; SIMD_WIDTH], + joint_body1: JointSolverBody, + joint_body2: JointSolverBody, + joint_id: [JointIndex; SIMD_WIDTH], + local_frame1: Isometry, + local_frame2: Isometry, + locked_axes: u8, + constraint_id: usize, +} + +#[cfg(feature = "simd-is-enabled")] +impl JointTwoBodyConstraintBuilderSimd { + pub fn generate( + joint: [&ImpulseJoint; SIMD_WIDTH], + bodies: &RigidBodySet, + joint_id: [JointIndex; SIMD_WIDTH], + out_builder: &mut Self, + out_constraint_id: &mut usize, + ) { + let rb1 = gather![|ii| &bodies[joint[ii].body1]]; + let rb2 = gather![|ii| &bodies[joint[ii].body2]]; + + let body1 = gather![|ii| rb1[ii].ids.active_set_offset]; + let body2 = gather![|ii| rb2[ii].ids.active_set_offset]; + + let joint_body1 = JointSolverBody { + im: gather![|ii| rb1[ii].mprops.effective_inv_mass].into(), + sqrt_ii: gather![|ii| rb1[ii].mprops.effective_world_inv_inertia_sqrt].into(), + world_com: Point::origin(), + solver_vel: body1, + }; + let joint_body2 = JointSolverBody { + im: gather![|ii| rb2[ii].mprops.effective_inv_mass].into(), + sqrt_ii: gather![|ii| rb2[ii].mprops.effective_world_inv_inertia_sqrt].into(), + world_com: Point::origin(), + solver_vel: body2, + }; + + *out_builder = Self { + body1, + body2, + joint_body1, + joint_body2, + joint_id, + local_frame1: gather![|ii| joint[ii].data.local_frame1].into(), + local_frame2: gather![|ii| joint[ii].data.local_frame2].into(), + locked_axes: joint[0].data.locked_axes.bits(), + constraint_id: *out_constraint_id, + }; + + let count = ConstraintsCounts::from_joint(joint[0]); + *out_constraint_id += count.num_constraints; + } + + pub fn update( + &mut self, + params: &IntegrationParameters, + bodies: &[SolverBody], + out: &mut [JointTwoBodyConstraint], + ) { + // NOTE: right now, the "update", is basically reconstructing all the + // constraints. Could we make this more incremental? + + let rb1 = gather![|ii| &bodies[self.body1[ii]]]; + let rb2 = gather![|ii| &bodies[self.body2[ii]]]; + let frame1 = Isometry::from(gather![|ii| rb1[ii].position]) * self.local_frame1; + let frame2 = Isometry::from(gather![|ii| rb2[ii].position]) * self.local_frame2; + self.joint_body1.world_com = gather![|ii| rb1[ii].world_com].into(); + self.joint_body2.world_com = gather![|ii| rb2[ii].world_com].into(); + + JointTwoBodyConstraint::::lock_axes( + params, + self.joint_id, + &self.joint_body1, + &self.joint_body2, + &frame1, + &frame2, + self.locked_axes, + &mut out[self.constraint_id..], + ); + } +} + +pub struct JointOneBodyConstraintBuilder { + body1: JointFixedSolverBody, + frame1: Isometry, + body2: usize, + joint_id: JointIndex, + joint: GenericJoint, + constraint_id: usize, +} + +impl JointOneBodyConstraintBuilder { + pub fn generate( + joint: &ImpulseJoint, + bodies: &RigidBodySet, + joint_id: JointIndex, + out_builder: &mut Self, + out_constraint_id: &mut usize, + ) { + let mut joint_data = joint.data; + let mut handle1 = joint.body1; + let mut handle2 = joint.body2; + let flipped = !bodies[handle2].is_dynamic(); + + if flipped { + std::mem::swap(&mut handle1, &mut handle2); + std::mem::swap(&mut joint_data.local_frame1, &mut joint_data.local_frame2); + }; + + let rb1 = &bodies[handle1]; + let rb2 = &bodies[handle2]; + + let frame1 = rb1.pos.position * joint_data.local_frame1; + let joint_body1 = JointFixedSolverBody { + linvel: rb1.vels.linvel, + angvel: rb1.vels.angvel, + world_com: rb1.mprops.world_com, + }; + + *out_builder = Self { + body1: joint_body1, + frame1, + body2: rb2.ids.active_set_offset, + joint_id, + joint: joint_data, + constraint_id: *out_constraint_id, + }; + + let count = ConstraintsCounts::from_joint(joint); + *out_constraint_id += count.num_constraints; + } + + pub fn update( + &self, + params: &IntegrationParameters, + bodies: &[SolverBody], + out: &mut [JointOneBodyConstraint], + ) { + // NOTE: right now, the "update", is basically reconstructing all the + // constraints. Could we make this more incremental? + + let rb2 = &bodies[self.body2]; + let frame2 = rb2.position * self.joint.local_frame2; + + let joint_body2 = JointSolverBody { + im: rb2.im, + sqrt_ii: rb2.sqrt_ii, + world_com: rb2.world_com, + solver_vel: [self.body2], + }; + + JointOneBodyConstraint::::lock_axes( + params, + self.joint_id, + &self.body1, + &joint_body2, + &self.frame1, + &frame2, + &self.joint, + &mut out[self.constraint_id..], + ); + } +} + +#[cfg(feature = "simd-is-enabled")] +pub struct JointOneBodyConstraintBuilderSimd { + body1: JointFixedSolverBody, + frame1: Isometry, + body2: [usize; SIMD_WIDTH], + joint_id: [JointIndex; SIMD_WIDTH], + local_frame2: Isometry, + locked_axes: u8, + constraint_id: usize, +} + +#[cfg(feature = "simd-is-enabled")] +impl JointOneBodyConstraintBuilderSimd { + pub fn generate( + joint: [&ImpulseJoint; SIMD_WIDTH], + bodies: &RigidBodySet, + joint_id: [JointIndex; SIMD_WIDTH], + out_builder: &mut Self, + out_constraint_id: &mut usize, + ) { + let mut rb1 = gather![|ii| &bodies[joint[ii].body1]]; + let mut rb2 = gather![|ii| &bodies[joint[ii].body2]]; + let mut local_frame1 = gather![|ii| joint[ii].data.local_frame1]; + let mut local_frame2 = gather![|ii| joint[ii].data.local_frame2]; + + for ii in 0..SIMD_WIDTH { + if !rb2[ii].is_dynamic() { + std::mem::swap(&mut rb1[ii], &mut rb2[ii]); + std::mem::swap(&mut local_frame1[ii], &mut local_frame2[ii]); + } + } + + let poss1 = Isometry::from(gather![|ii| rb1[ii].pos.position]); + + let joint_body1 = JointFixedSolverBody { + linvel: gather![|ii| rb1[ii].vels.linvel].into(), + angvel: gather![|ii| rb1[ii].vels.angvel].into(), + world_com: gather![|ii| rb1[ii].mprops.world_com].into(), + }; + + *out_builder = Self { + body1: joint_body1, + body2: gather![|ii| rb2[ii].ids.active_set_offset], + joint_id, + frame1: poss1 * Isometry::from(local_frame1), + local_frame2: local_frame2.into(), + locked_axes: joint[0].data.locked_axes.bits(), + constraint_id: *out_constraint_id, + }; + + let count = ConstraintsCounts::from_joint(joint[0]); + *out_constraint_id += count.num_constraints; + } + + pub fn update( + &self, + params: &IntegrationParameters, + bodies: &[SolverBody], + out: &mut [JointOneBodyConstraint], + ) { + // NOTE: right now, the "update", is basically reconstructing all the + // constraints. Could we make this more incremental? + + let rb2 = gather![|ii| &bodies[self.body2[ii]]]; + let frame2 = Isometry::from(gather![|ii| rb2[ii].position]) * self.local_frame2; + + let joint_body2 = JointSolverBody { + im: gather![|ii| rb2[ii].im].into(), + sqrt_ii: gather![|ii| rb2[ii].sqrt_ii].into(), + world_com: gather![|ii| rb2[ii].world_com].into(), + solver_vel: self.body2, + }; + + JointOneBodyConstraint::::lock_axes( + params, + self.joint_id, + &self.body1, + &joint_body2, + &self.frame1, + &frame2, + self.locked_axes, + &mut out[self.constraint_id..], + ); + } +} + +#[derive(Debug, Copy, Clone)] +pub struct JointTwoBodyConstraintHelper { + pub basis: Matrix, + pub basis2: Matrix, // TODO: used for angular coupling. Can we avoid storing this? + pub cmat1_basis: SMatrix, + pub cmat2_basis: SMatrix, + pub ang_basis: SMatrix, + pub lin_err: Vector, + pub ang_err: Rotation, +} + +impl JointTwoBodyConstraintHelper { + pub fn new( + frame1: &Isometry, + frame2: &Isometry, + world_com1: &Point, + world_com2: &Point, + locked_lin_axes: u8, + ) -> Self { + let mut frame1 = *frame1; + let basis = frame1.rotation.to_rotation_matrix().into_inner(); + let lin_err = frame2.translation.vector - frame1.translation.vector; + + // Adjust the point of application of the force for the first body, + // by snapping free axes to the second frame’s center (to account for + // the allowed relative movement). + { + let mut new_center1 = frame2.translation.vector; // First, assume all dofs are free. + + // Then snap the locked ones. + for i in 0..DIM { + if locked_lin_axes & (1 << i) != 0 { + let axis = basis.column(i); + new_center1 -= axis * lin_err.dot(&axis); + } + } + frame1.translation.vector = new_center1; + } + + let r1 = frame1.translation.vector - world_com1.coords; + let r2 = frame2.translation.vector - world_com2.coords; + + let cmat1 = r1.gcross_matrix(); + let cmat2 = r2.gcross_matrix(); + + #[allow(unused_mut)] // The mut is needed for 3D + let mut ang_basis = frame1.rotation.diff_conj1_2(&frame2.rotation).transpose(); + #[allow(unused_mut)] // The mut is needed for 3D + let mut ang_err = frame1.rotation.inverse() * frame2.rotation; + + #[cfg(feature = "dim3")] + { + let sgn = N::one().simd_copysign(frame1.rotation.dot(&frame2.rotation)); + ang_basis *= sgn; + *ang_err.as_mut_unchecked() *= sgn; + } + + Self { + basis, + basis2: frame2.rotation.to_rotation_matrix().into_inner(), + cmat1_basis: cmat1 * basis, + cmat2_basis: cmat2 * basis, + ang_basis, + lin_err, + ang_err, + } + } + + pub fn limit_linear( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + limited_axis: usize, + limits: [N; 2], + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + let zero = N::zero(); + let mut constraint = + self.lock_linear(params, joint_id, body1, body2, limited_axis, writeback_id); + + let dist = self.lin_err.dot(&constraint.lin_jac); + let min_enabled = dist.simd_le(limits[0]); + let max_enabled = limits[1].simd_le(dist); + + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + let rhs_bias = + ((dist - limits[1]).simd_max(zero) - (limits[0] - dist).simd_max(zero)) * erp_inv_dt; + constraint.rhs = constraint.rhs_wo_bias + rhs_bias; + constraint.cfm_coeff = cfm_coeff; + constraint.impulse_bounds = [ + N::splat(-Real::INFINITY).select(min_enabled, zero), + N::splat(Real::INFINITY).select(max_enabled, zero), + ]; + + constraint + } + + pub fn limit_linear_coupled( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + coupled_axes: u8, + limits: [N; 2], + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + let zero = N::zero(); + let mut lin_jac = Vector::zeros(); + let mut ang_jac1: AngVector = na::zero(); + let mut ang_jac2: AngVector = na::zero(); + + for i in 0..DIM { + if coupled_axes & (1 << i) != 0 { + let coeff = self.basis.column(i).dot(&self.lin_err); + lin_jac += self.basis.column(i) * coeff; + #[cfg(feature = "dim2")] + { + ang_jac1 += self.cmat1_basis[i] * coeff; + ang_jac2 += self.cmat2_basis[i] * coeff; + } + #[cfg(feature = "dim3")] + { + ang_jac1 += self.cmat1_basis.column(i) * coeff; + ang_jac2 += self.cmat2_basis.column(i) * coeff; + } + } + } + + // FIXME: handle min limit too. + + let dist = lin_jac.norm(); + let inv_dist = crate::utils::simd_inv(dist); + lin_jac *= inv_dist; + ang_jac1 *= inv_dist; + ang_jac2 *= inv_dist; + + let rhs_wo_bias = (dist - limits[1]).simd_min(zero) * N::splat(params.inv_dt()); + + ang_jac1 = body1.sqrt_ii * ang_jac1; + ang_jac2 = body2.sqrt_ii * ang_jac2; + + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + let rhs_bias = (dist - limits[1]).simd_max(zero) * erp_inv_dt; + let rhs = rhs_wo_bias + rhs_bias; + let impulse_bounds = [N::zero(), N::splat(Real::INFINITY)]; + + JointTwoBodyConstraint { + joint_id, + solver_vel1: body1.solver_vel, + solver_vel2: body2.solver_vel, + im1: body1.im, + im2: body2.im, + impulse: N::zero(), + impulse_bounds, + lin_jac, + ang_jac1, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff, + cfm_gain: N::zero(), + rhs, + rhs_wo_bias, + writeback_id, + } + } + + pub fn motor_linear( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + motor_axis: usize, + motor_params: &MotorParameters, + limits: Option<[N; 2]>, + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + let inv_dt = N::splat(params.inv_dt()); + let mut constraint = + self.lock_linear(params, joint_id, body1, body2, motor_axis, writeback_id); + + let mut rhs_wo_bias = N::zero(); + if motor_params.erp_inv_dt != N::zero() { + let dist = self.lin_err.dot(&constraint.lin_jac); + rhs_wo_bias += (dist - motor_params.target_pos) * motor_params.erp_inv_dt; + } + + let mut target_vel = motor_params.target_vel; + if let Some(limits) = limits { + let dist = self.lin_err.dot(&constraint.lin_jac); + target_vel = + target_vel.simd_clamp((limits[0] - dist) * inv_dt, (limits[1] - dist) * inv_dt); + }; + + rhs_wo_bias += -target_vel; + + constraint.cfm_coeff = motor_params.cfm_coeff; + constraint.cfm_gain = motor_params.cfm_gain; + constraint.impulse_bounds = [-motor_params.max_impulse, motor_params.max_impulse]; + constraint.rhs = rhs_wo_bias; + constraint.rhs_wo_bias = rhs_wo_bias; + constraint + } + + pub fn lock_linear( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + locked_axis: usize, + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + let lin_jac = self.basis.column(locked_axis).into_owned(); + #[cfg(feature = "dim2")] + let mut ang_jac1 = self.cmat1_basis[locked_axis]; + #[cfg(feature = "dim2")] + let mut ang_jac2 = self.cmat2_basis[locked_axis]; + #[cfg(feature = "dim3")] + let mut ang_jac1 = self.cmat1_basis.column(locked_axis).into_owned(); + #[cfg(feature = "dim3")] + let mut ang_jac2 = self.cmat2_basis.column(locked_axis).into_owned(); + + let rhs_wo_bias = N::zero(); + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + let rhs_bias = lin_jac.dot(&self.lin_err) * erp_inv_dt; + + ang_jac1 = body1.sqrt_ii * ang_jac1; + ang_jac2 = body2.sqrt_ii * ang_jac2; + + JointTwoBodyConstraint { + joint_id, + solver_vel1: body1.solver_vel, + solver_vel2: body2.solver_vel, + im1: body1.im, + im2: body2.im, + impulse: N::zero(), + impulse_bounds: [-N::splat(Real::MAX), N::splat(Real::MAX)], + lin_jac, + ang_jac1, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff, + cfm_gain: N::zero(), + rhs: rhs_wo_bias + rhs_bias, + rhs_wo_bias, + writeback_id, + } + } + + pub fn limit_angular( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + _limited_axis: usize, + limits: [N; 2], + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + let zero = N::zero(); + let half = N::splat(0.5); + let s_limits = [(limits[0] * half).simd_sin(), (limits[1] * half).simd_sin()]; + #[cfg(feature = "dim2")] + let s_ang = (self.ang_err.angle() * half).simd_sin(); + #[cfg(feature = "dim3")] + let s_ang = self.ang_err.imag()[_limited_axis]; + let min_enabled = s_ang.simd_le(s_limits[0]); + let max_enabled = s_limits[1].simd_le(s_ang); + + let impulse_bounds = [ + N::splat(-Real::INFINITY).select(min_enabled, zero), + N::splat(Real::INFINITY).select(max_enabled, zero), + ]; + + #[cfg(feature = "dim2")] + let ang_jac = N::one(); + #[cfg(feature = "dim3")] + let ang_jac = self.ang_basis.column(_limited_axis).into_owned(); + let rhs_wo_bias = N::zero(); + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + let rhs_bias = ((s_ang - s_limits[1]).simd_max(zero) + - (s_limits[0] - s_ang).simd_max(zero)) + * erp_inv_dt; + + let ang_jac1 = body1.sqrt_ii * ang_jac; + let ang_jac2 = body2.sqrt_ii * ang_jac; + + JointTwoBodyConstraint { + joint_id, + solver_vel1: body1.solver_vel, + solver_vel2: body2.solver_vel, + im1: body1.im, + im2: body2.im, + impulse: N::zero(), + impulse_bounds, + lin_jac: na::zero(), + ang_jac1, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff, + cfm_gain: N::zero(), + rhs: rhs_wo_bias + rhs_bias, + rhs_wo_bias, + writeback_id, + } + } + + pub fn motor_angular( + &self, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + _motor_axis: usize, + motor_params: &MotorParameters, + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + #[cfg(feature = "dim2")] + let ang_jac = N::one(); + #[cfg(feature = "dim3")] + let ang_jac = self.basis.column(_motor_axis).into_owned(); + + let mut rhs_wo_bias = N::zero(); + if motor_params.erp_inv_dt != N::zero() { + #[cfg(feature = "dim2")] + let ang_dist = self.ang_err.angle(); + #[cfg(feature = "dim3")] + let ang_dist = self.ang_err.imag()[_motor_axis].simd_asin() * N::splat(2.0); + let target_ang = motor_params.target_pos; + rhs_wo_bias += utils::smallest_abs_diff_between_angles(ang_dist, target_ang) + * motor_params.erp_inv_dt; + } + + rhs_wo_bias += -motor_params.target_vel; + + let ang_jac1 = body1.sqrt_ii * ang_jac; + let ang_jac2 = body2.sqrt_ii * ang_jac; + + JointTwoBodyConstraint { + joint_id, + solver_vel1: body1.solver_vel, + solver_vel2: body2.solver_vel, + im1: body1.im, + im2: body2.im, + impulse: N::zero(), + impulse_bounds: [-motor_params.max_impulse, motor_params.max_impulse], + lin_jac: na::zero(), + ang_jac1, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff: motor_params.cfm_coeff, + cfm_gain: motor_params.cfm_gain, + rhs: rhs_wo_bias, + rhs_wo_bias, + writeback_id, + } + } + + pub fn lock_angular( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointSolverBody, + body2: &JointSolverBody, + _locked_axis: usize, + writeback_id: WritebackId, + ) -> JointTwoBodyConstraint { + #[cfg(feature = "dim2")] + let ang_jac = N::one(); + #[cfg(feature = "dim3")] + let ang_jac = self.ang_basis.column(_locked_axis).into_owned(); + + let rhs_wo_bias = N::zero(); + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + #[cfg(feature = "dim2")] + let rhs_bias = self.ang_err.im * erp_inv_dt; + #[cfg(feature = "dim3")] + let rhs_bias = self.ang_err.imag()[_locked_axis] * erp_inv_dt; + + let ang_jac1 = body1.sqrt_ii * ang_jac; + let ang_jac2 = body2.sqrt_ii * ang_jac; + + JointTwoBodyConstraint { + joint_id, + solver_vel1: body1.solver_vel, + solver_vel2: body2.solver_vel, + im1: body1.im, + im2: body2.im, + impulse: N::zero(), + impulse_bounds: [-N::splat(Real::MAX), N::splat(Real::MAX)], + lin_jac: na::zero(), + ang_jac1, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff, + cfm_gain: N::zero(), + rhs: rhs_wo_bias + rhs_bias, + rhs_wo_bias, + writeback_id, + } + } + + /// Orthogonalize the constraints and set their inv_lhs field. + pub fn finalize_constraints( + constraints: &mut [JointTwoBodyConstraint], + ) { + let len = constraints.len(); + + if len == 0 { + return; + } + + let imsum = constraints[0].im1 + constraints[0].im2; + + // Use the modified Gram-Schmidt orthogonalization. + for j in 0..len { + let c_j = &mut constraints[j]; + 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 cfm_gain = dot_jj * c_j.cfm_coeff + c_j.cfm_gain; + let inv_dot_jj = crate::utils::simd_inv(dot_jj); + c_j.inv_lhs = crate::utils::simd_inv(dot_jj + cfm_gain); // Don’t forget to update the inv_lhs. + c_j.cfm_gain = cfm_gain; + + if c_j.impulse_bounds != [-N::splat(Real::MAX), N::splat(Real::MAX)] { + // Don't remove constraints with limited forces from the others + // because they may not deliver the necessary forces to fulfill + // the removed parts of other constraints. + continue; + } + + 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(&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; + + c_i.lin_jac -= c_j.lin_jac * coeff; + c_i.ang_jac1 -= c_j.ang_jac1 * coeff; + c_i.ang_jac2 -= c_j.ang_jac2 * coeff; + c_i.rhs_wo_bias -= c_j.rhs_wo_bias * coeff; + c_i.rhs -= c_j.rhs * coeff; + } + } + } + + pub fn limit_linear_one_body( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointFixedSolverBody, + body2: &JointSolverBody, + limited_axis: usize, + limits: [N; 2], + writeback_id: WritebackId, + ) -> JointOneBodyConstraint { + let zero = N::zero(); + let lin_jac = self.basis.column(limited_axis).into_owned(); + let dist = self.lin_err.dot(&lin_jac); + + let min_enabled = dist.simd_le(limits[0]); + let max_enabled = limits[1].simd_le(dist); + + let impulse_bounds = [ + N::splat(-Real::INFINITY).select(min_enabled, zero), + N::splat(Real::INFINITY).select(max_enabled, zero), + ]; + + let ang_jac1 = self.cmat1_basis.column(limited_axis).into_owned(); + #[cfg(feature = "dim2")] + let mut ang_jac2 = self.cmat2_basis[limited_axis]; + #[cfg(feature = "dim3")] + let mut ang_jac2 = self.cmat2_basis.column(limited_axis).into_owned(); + + let rhs_wo_bias = -lin_jac.dot(&body1.linvel) - ang_jac1.gdot(body1.angvel); + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + let rhs_bias = + ((dist - limits[1]).simd_max(zero) - (limits[0] - dist).simd_max(zero)) * erp_inv_dt; + + ang_jac2 = body2.sqrt_ii * ang_jac2; + + JointOneBodyConstraint { + joint_id, + solver_vel2: body2.solver_vel, + im2: body2.im, + impulse: zero, + impulse_bounds, + lin_jac, + ang_jac2, + inv_lhs: zero, // Will be set during ortogonalization. + cfm_coeff, + cfm_gain: N::zero(), + rhs: rhs_wo_bias + rhs_bias, + rhs_wo_bias, + writeback_id, + } + } + + pub fn limit_linear_coupled_one_body( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointFixedSolverBody, + body2: &JointSolverBody, + coupled_axes: u8, + limits: [N; 2], + writeback_id: WritebackId, + ) -> JointOneBodyConstraint { + let zero = N::zero(); + let mut lin_jac = Vector::zeros(); + let mut ang_jac1: AngVector = na::zero(); + let mut ang_jac2: AngVector = na::zero(); + + for i in 0..DIM { + if coupled_axes & (1 << i) != 0 { + let coeff = self.basis.column(i).dot(&self.lin_err); + lin_jac += self.basis.column(i) * coeff; + #[cfg(feature = "dim2")] + { + ang_jac1 += self.cmat1_basis[i] * coeff; + ang_jac2 += self.cmat2_basis[i] * coeff; + } + #[cfg(feature = "dim3")] + { + ang_jac1 += self.cmat1_basis.column(i) * coeff; + ang_jac2 += self.cmat2_basis.column(i) * coeff; + } + } + } + + let dist = lin_jac.norm(); + let inv_dist = crate::utils::simd_inv(dist); + lin_jac *= inv_dist; + ang_jac1 *= inv_dist; + ang_jac2 *= inv_dist; + + // FIXME: handle min limit too. + let proj_vel1 = -lin_jac.dot(&body1.linvel) - ang_jac1.gdot(body1.angvel); + let rhs_wo_bias = proj_vel1 + (dist - limits[1]).simd_min(zero) * N::splat(params.inv_dt()); + + ang_jac2 = body2.sqrt_ii * ang_jac2; + + let erp_inv_dt = N::splat(params.joint_erp_inv_dt()); + let cfm_coeff = N::splat(params.joint_cfm_coeff()); + let rhs_bias = (dist - limits[1]).simd_max(zero) * erp_inv_dt; + let rhs = rhs_wo_bias + rhs_bias; + let impulse_bounds = [N::zero(), N::splat(Real::INFINITY)]; + + JointOneBodyConstraint { + joint_id, + solver_vel2: body2.solver_vel, + im2: body2.im, + impulse: N::zero(), + impulse_bounds, + lin_jac, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff, + cfm_gain: N::zero(), + rhs, + rhs_wo_bias, + writeback_id, + } + } + + pub fn motor_linear_one_body( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointFixedSolverBody, + body2: &JointSolverBody, + motor_axis: usize, + motor_params: &MotorParameters, + limits: Option<[N; 2]>, + writeback_id: WritebackId, + ) -> JointOneBodyConstraint { + let inv_dt = N::splat(params.inv_dt()); + + let lin_jac = self.basis.column(motor_axis).into_owned(); + let ang_jac1 = self.cmat1_basis.column(motor_axis).into_owned(); + #[cfg(feature = "dim2")] + let mut ang_jac2 = self.cmat2_basis[motor_axis]; + #[cfg(feature = "dim3")] + let mut ang_jac2 = self.cmat2_basis.column(motor_axis).into_owned(); + + let mut rhs_wo_bias = N::zero(); + if motor_params.erp_inv_dt != N::zero() { + let dist = self.lin_err.dot(&lin_jac); + rhs_wo_bias += (dist - motor_params.target_pos) * motor_params.erp_inv_dt; + } + + let mut target_vel = motor_params.target_vel; + if let Some(limits) = limits { + let dist = self.lin_err.dot(&lin_jac); + target_vel = + target_vel.simd_clamp((limits[0] - dist) * inv_dt, (limits[1] - dist) * inv_dt); + }; + + let proj_vel1 = -lin_jac.dot(&body1.linvel) - ang_jac1.gdot(body1.angvel); + rhs_wo_bias += proj_vel1 - target_vel; + + ang_jac2 = body2.sqrt_ii * ang_jac2; + + JointOneBodyConstraint { + joint_id, + solver_vel2: body2.solver_vel, + im2: body2.im, + impulse: N::zero(), + impulse_bounds: [-motor_params.max_impulse, motor_params.max_impulse], + lin_jac, + ang_jac2, + inv_lhs: N::zero(), // Will be set during ortogonalization. + cfm_coeff: motor_params.cfm_coeff, + cfm_gain: motor_params.cfm_gain, + rhs: rhs_wo_bias, + rhs_wo_bias, + writeback_id, + } + } + + pub fn motor_linear_coupled_one_body( + &self, + params: &IntegrationParameters, + joint_id: [JointIndex; LANES], + body1: &JointFixedSolverB