diff options
Diffstat (limited to 'src/dynamics/solver/joint_constraint/joint_constraint.rs')
| -rw-r--r-- | src/dynamics/solver/joint_constraint/joint_constraint.rs | 340 |
1 files changed, 340 insertions, 0 deletions
diff --git a/src/dynamics/solver/joint_constraint/joint_constraint.rs b/src/dynamics/solver/joint_constraint/joint_constraint.rs new file mode 100644 index 0000000..8b1f8bf --- /dev/null +++ b/src/dynamics/solver/joint_constraint/joint_constraint.rs @@ -0,0 +1,340 @@ +use super::{ + BallVelocityConstraint, BallVelocityGroundConstraint, FixedVelocityConstraint, + FixedVelocityGroundConstraint, PrismaticVelocityConstraint, PrismaticVelocityGroundConstraint, +}; +#[cfg(feature = "dim3")] +use super::{RevoluteVelocityConstraint, RevoluteVelocityGroundConstraint}; +#[cfg(feature = "simd-is-enabled")] +use super::{ + WBallVelocityConstraint, WBallVelocityGroundConstraint, WFixedVelocityConstraint, + WFixedVelocityGroundConstraint, WPrismaticVelocityConstraint, + WPrismaticVelocityGroundConstraint, +}; +#[cfg(feature = "dim3")] +#[cfg(feature = "simd-is-enabled")] +use super::{WRevoluteVelocityConstraint, WRevoluteVelocityGroundConstraint}; +use crate::dynamics::solver::DeltaVel; +use crate::dynamics::{ + IntegrationParameters, Joint, JointGraphEdge, JointIndex, JointParams, RigidBodySet, +}; +#[cfg(feature = "simd-is-enabled")] +use crate::math::SIMD_WIDTH; + +pub(crate) enum AnyJointVelocityConstraint { + BallConstraint(BallVelocityConstraint), + BallGroundConstraint(BallVelocityGroundConstraint), + #[cfg(feature = "simd-is-enabled")] + WBallConstraint(WBallVelocityConstraint), + #[cfg(feature = "simd-is-enabled")] + WBallGroundConstraint(WBallVelocityGroundConstraint), + FixedConstraint(FixedVelocityConstraint), + FixedGroundConstraint(FixedVelocityGroundConstraint), + #[cfg(feature = "simd-is-enabled")] + WFixedConstraint(WFixedVelocityConstraint), + #[cfg(feature = "simd-is-enabled")] + WFixedGroundConstraint(WFixedVelocityGroundConstraint), + PrismaticConstraint(PrismaticVelocityConstraint), + PrismaticGroundConstraint(PrismaticVelocityGroundConstraint), + #[cfg(feature = "simd-is-enabled")] + WPrismaticConstraint(WPrismaticVelocityConstraint), + #[cfg(feature = "simd-is-enabled")] + WPrismaticGroundConstraint(WPrismaticVelocityGroundConstraint), + #[cfg(feature = "dim3")] + RevoluteConstraint(RevoluteVelocityConstraint), + #[cfg(feature = "dim3")] + RevoluteGroundConstraint(RevoluteVelocityGroundConstraint), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + WRevoluteConstraint(WRevoluteVelocityConstraint), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + WRevoluteGroundConstraint(WRevoluteVelocityGroundConstraint), + #[allow(dead_code)] // The Empty variant is only used with parallel code. + Empty, +} + +impl AnyJointVelocityConstraint { + #[cfg(feature = "parallel")] + pub fn num_active_constraints(_: &Joint) -> usize { + 1 + } + + pub fn from_joint( + params: &IntegrationParameters, + joint_id: JointIndex, + joint: &Joint, + bodies: &RigidBodySet, + ) -> Self { + let rb1 = &bodies[joint.body1]; + let rb2 = &bodies[joint.body2]; + + match &joint.params { + JointParams::BallJoint(p) => AnyJointVelocityConstraint::BallConstraint( + BallVelocityConstraint::from_params(params, joint_id, rb1, rb2, p), + ), + JointParams::FixedJoint(p) => AnyJointVelocityConstraint::FixedConstraint( + FixedVelocityConstraint::from_params(params, joint_id, rb1, rb2, p), + ), + JointParams::PrismaticJoint(p) => AnyJointVelocityConstraint::PrismaticConstraint( + PrismaticVelocityConstraint::from_params(params, joint_id, rb1, rb2, p), + ), + #[cfg(feature = "dim3")] + JointParams::RevoluteJoint(p) => AnyJointVelocityConstraint::RevoluteConstraint( + RevoluteVelocityConstraint::from_params(params, joint_id, rb1, rb2, p), + ), + } + } + + #[cfg(feature = "simd-is-enabled")] + pub fn from_wide_joint( + params: &IntegrationParameters, + joint_id: [JointIndex; SIMD_WIDTH], + joints: [&Joint; SIMD_WIDTH], + bodies: &RigidBodySet, + ) -> Self { + let rbs1 = array![|ii| &bodies[joints[ii].body1]; SIMD_WIDTH]; + let rbs2 = array![|ii| &bodies[joints[ii].body2]; SIMD_WIDTH]; + + match &joints[0].params { + JointParams::BallJoint(_) => { + let joints = array![|ii| joints[ii].params.as_ball_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WBallConstraint(WBallVelocityConstraint::from_params( + params, joint_id, rbs1, rbs2, joints, + )) + } + JointParams::FixedJoint(_) => { + let joints = array![|ii| joints[ii].params.as_fixed_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WFixedConstraint(WFixedVelocityConstraint::from_params( + params, joint_id, rbs1, rbs2, joints, + )) + } + JointParams::PrismaticJoint(_) => { + let joints = + array![|ii| joints[ii].params.as_prismatic_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WPrismaticConstraint( + WPrismaticVelocityConstraint::from_params(params, joint_id, rbs1, rbs2, joints), + ) + } + #[cfg(feature = "dim3")] + JointParams::RevoluteJoint(_) => { + let joints = + array![|ii| joints[ii].params.as_revolute_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WRevoluteConstraint( + WRevoluteVelocityConstraint::from_params(params, joint_id, rbs1, rbs2, joints), + ) + } + } + } + + pub fn from_joint_ground( + params: &IntegrationParameters, + joint_id: JointIndex, + joint: &Joint, + bodies: &RigidBodySet, + ) -> Self { + let mut rb1 = &bodies[joint.body1]; + let mut rb2 = &bodies[joint.body2]; + let flipped = !rb2.is_dynamic(); + + if flipped { + std::mem::swap(&mut rb1, &mut rb2); + } + + match &joint.params { + JointParams::BallJoint(p) => AnyJointVelocityConstraint::BallGroundConstraint( + BallVelocityGroundConstraint::from_params(params, joint_id, rb1, rb2, p, flipped), + ), + JointParams::FixedJoint(p) => AnyJointVelocityConstraint::FixedGroundConstraint( + FixedVelocityGroundConstraint::from_params(params, joint_id, rb1, rb2, p, flipped), + ), + JointParams::PrismaticJoint(p) => { + AnyJointVelocityConstraint::PrismaticGroundConstraint( + PrismaticVelocityGroundConstraint::from_params( + params, joint_id, rb1, rb2, p, flipped, + ), + ) + } + #[cfg(feature = "dim3")] + JointParams::RevoluteJoint(p) => AnyJointVelocityConstraint::RevoluteGroundConstraint( + RevoluteVelocityGroundConstraint::from_params( + params, joint_id, rb1, rb2, p, flipped, + ), + ), + } + } + + #[cfg(feature = "simd-is-enabled")] + pub fn from_wide_joint_ground( + params: &IntegrationParameters, + joint_id: [JointIndex; SIMD_WIDTH], + joints: [&Joint; SIMD_WIDTH], + bodies: &RigidBodySet, + ) -> Self { + let mut rbs1 = array![|ii| &bodies[joints[ii].body1]; SIMD_WIDTH]; + let mut rbs2 = array![|ii| &bodies[joints[ii].body2]; SIMD_WIDTH]; + let mut flipped = [false; SIMD_WIDTH]; + + for ii in 0..SIMD_WIDTH { + if !rbs2[ii].is_dynamic() { + std::mem::swap(&mut rbs1[ii], &mut rbs2[ii]); + flipped[ii] = true; + } + } + + match &joints[0].params { + JointParams::BallJoint(_) => { + let joints = array![|ii| joints[ii].params.as_ball_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WBallGroundConstraint( + WBallVelocityGroundConstraint::from_params( + params, joint_id, rbs1, rbs2, joints, flipped, + ), + ) + } + JointParams::FixedJoint(_) => { + let joints = array![|ii| joints[ii].params.as_fixed_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WFixedGroundConstraint( + WFixedVelocityGroundConstraint::from_params( + params, joint_id, rbs1, rbs2, joints, flipped, + ), + ) + } + JointParams::PrismaticJoint(_) => { + let joints = + array![|ii| joints[ii].params.as_prismatic_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WPrismaticGroundConstraint( + WPrismaticVelocityGroundConstraint::from_params( + params, joint_id, rbs1, rbs2, joints, flipped, + ), + ) + } + #[cfg(feature = "dim3")] + JointParams::RevoluteJoint(_) => { + let joints = + array![|ii| joints[ii].params.as_revolute_joint().unwrap(); SIMD_WIDTH]; + AnyJointVelocityConstraint::WRevoluteGroundConstraint( + WRevoluteVelocityGroundConstraint::from_params( + params, joint_id, rbs1, rbs2, joints, flipped, + ), + ) + } + } + } + + pub fn warmstart(&self, mj_lambdas: &mut [DeltaVel<f32>]) { + match self { + AnyJointVelocityConstraint::BallConstraint(c) => c.warmstart(mj_lambdas), + AnyJointVelocityConstraint::BallGroundConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WBallConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WBallGroundConstraint(c) => c.warmstart(mj_lambdas), + AnyJointVelocityConstraint::FixedConstraint(c) => c.warmstart(mj_lambdas), + AnyJointVelocityConstraint::FixedGroundConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WFixedConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WFixedGroundConstraint(c) => c.warmstart(mj_lambdas), + AnyJointVelocityConstraint::PrismaticConstraint(c) => c.warmstart(mj_lambdas), + AnyJointVelocityConstraint::PrismaticGroundConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WPrismaticConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WPrismaticGroundConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "dim3")] + AnyJointVelocityConstraint::RevoluteConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "dim3")] + AnyJointVelocityConstraint::RevoluteGroundConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WRevoluteConstraint(c) => c.warmstart(mj_lambdas), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WRevoluteGroundConstraint(c) => c.warmstart(mj_lambdas), + AnyJointVelocityConstraint::Empty => unreachable!(), + } + } + + pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<f32>]) { + match self { + AnyJointVelocityConstraint::BallConstraint(c) => c.solve(mj_lambdas), + AnyJointVelocityConstraint::BallGroundConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WBallConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WBallGroundConstraint(c) => c.solve(mj_lambdas), + AnyJointVelocityConstraint::FixedConstraint(c) => c.solve(mj_lambdas), + AnyJointVelocityConstraint::FixedGroundConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WFixedConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WFixedGroundConstraint(c) => c.solve(mj_lambdas), + AnyJointVelocityConstraint::PrismaticConstraint(c) => c.solve(mj_lambdas), + AnyJointVelocityConstraint::PrismaticGroundConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WPrismaticConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WPrismaticGroundConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "dim3")] + AnyJointVelocityConstraint::RevoluteConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "dim3")] + AnyJointVelocityConstraint::RevoluteGroundConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WRevoluteConstraint(c) => c.solve(mj_lambdas), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WRevoluteGroundConstraint(c) => c.solve(mj_lambdas), + AnyJointVelocityConstraint::Empty => unreachable!(), + } + } + + pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) { + match self { + AnyJointVelocityConstraint::BallConstraint(c) => c.writeback_impulses(joints_all), + + AnyJointVelocityConstraint::BallGroundConstraint(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WBallConstraint(c) => c.writeback_impulses(joints_all), + + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WBallGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + AnyJointVelocityConstraint::FixedConstraint(c) => c.writeback_impulses(joints_all), + AnyJointVelocityConstraint::FixedGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WFixedConstraint(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WFixedGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + AnyJointVelocityConstraint::PrismaticConstraint(c) => c.writeback_impulses(joints_all), + AnyJointVelocityConstraint::PrismaticGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WPrismaticConstraint(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WPrismaticGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + #[cfg(feature = "dim3")] + AnyJointVelocityConstraint::RevoluteConstraint(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "dim3")] + AnyJointVelocityConstraint::RevoluteGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WRevoluteConstraint(c) => c.writeback_impulses(joints_all), + #[cfg(feature = "dim3")] + #[cfg(feature = "simd-is-enabled")] + AnyJointVelocityConstraint::WRevoluteGroundConstraint(c) => { + c.writeback_impulses(joints_all) + } + AnyJointVelocityConstraint::Empty => unreachable!(), + } + } +} |