diff options
Diffstat (limited to 'src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs')
| -rw-r--r-- | src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs | 608 |
1 files changed, 608 insertions, 0 deletions
diff --git a/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs new file mode 100644 index 0000000..3e31256 --- /dev/null +++ b/src/dynamics/solver/joint_constraint/joint_velocity_constraint.rs @@ -0,0 +1,608 @@ +use crate::dynamics::solver::joint_constraint::JointVelocityConstraintBuilder; +use crate::dynamics::solver::DeltaVel; +use crate::dynamics::{IntegrationParameters, JointData, JointGraphEdge, JointIndex}; +use crate::math::{AngVector, AngularInertia, Isometry, Point, Real, Vector, DIM, SPATIAL_DIM}; +use crate::utils::{WDot, WReal}; +use simba::simd::SimdRealField; + +#[cfg(feature = "simd-is-enabled")] +use { + crate::math::{SimdReal, SIMD_WIDTH}, + na::SimdValue, +}; + +#[derive(Copy, Clone, PartialEq, Debug)] +pub struct MotorParameters<N: SimdRealField> { + pub stiffness: N, + pub damping: N, + pub gamma: N, + // pub keep_lhs: bool, + pub target_pos: N, + pub target_vel: N, + pub max_impulse: N, +} + +impl<N: SimdRealField> Default for MotorParameters<N> { + fn default() -> Self { + Self { + stiffness: N::zero(), + damping: N::zero(), + gamma: N::zero(), + // keep_lhs: true, + target_pos: N::zero(), + target_vel: N::zero(), + max_impulse: N::zero(), + } + } +} + +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub enum WritebackId { + Dof(usize), + Limit(usize), + Motor(usize), +} + +// TODO: right now we only use this for impulse_joints. +// However, it may actually be a good idea to use this everywhere in +// the solver, to avoid fetching data from the rigid-body set +// every time. +#[derive(Copy, Clone)] +pub struct SolverBody<N: SimdRealField, const LANES: usize> { + pub linvel: Vector<N>, + pub angvel: AngVector<N>, + pub im: N, + pub sqrt_ii: AngularInertia<N>, + pub world_com: Point<N>, + pub mj_lambda: [usize; LANES], +} + +#[derive(Debug, Copy, Clone)] +pub struct JointVelocityConstraint<N: SimdRealField, const LANES: usize> { + pub mj_lambda1: [usize; LANES], + pub mj_lambda2: [usize; LANES], + + pub joint_id: [JointIndex; LANES], + + pub impulse: N, + pub impulse_bounds: [N; 2], + pub lin_jac: Vector<N>, + pub ang_jac1: AngVector<N>, + pub ang_jac2: AngVector<N>, + + pub inv_lhs: N, + pub rhs: N, + pub rhs_wo_bias: N, + + pub im1: N, + pub im2: N, + + pub writeback_id: WritebackId, +} + +impl<N: WReal, const LANES: usize> JointVelocityConstraint<N, LANES> { + pub fn invalid() -> Self { + Self { + mj_lambda1: [crate::INVALID_USIZE; LANES], + mj_lambda2: [crate::INVALID_USIZE; LANES], + joint_id: [crate::INVALID_USIZE; LANES], + impulse: N::zero(), + impulse_bounds: [N::zero(), N::zero()], + lin_jac: Vector::zeros(), + ang_jac1: na::zero(), + ang_jac2: na::zero(), + inv_lhs: N::zero(), + rhs: N::zero(), + rhs_wo_bias: N::zero(), + im1: N::zero(), + im2: N::zero(), + writeback_id: WritebackId::Dof(0), + } + } + + pub fn solve_generic(&mut self, mj_lambda1: &mut DeltaVel<N>, mj_lambda2: &mut DeltaVel<N>) { + let dlinvel = self.lin_jac.dot(&(mj_lambda2.linear - mj_lambda1.linear)); + let dangvel = + self.ang_jac2.gdot(mj_lambda2.angular) - self.ang_jac1.gdot(mj_lambda1.angular); + + let rhs = dlinvel + dangvel + self.rhs; + let total_impulse = (self.impulse + self.inv_lhs * rhs) + .simd_clamp(self.impulse_bounds[0], self.impulse_bounds[1]); + let delta_impulse = total_impulse - self.impulse; + self.impulse = total_impulse; + + let lin_impulse = self.lin_jac * delta_impulse; + 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.angular += ang_impulse1; + mj_lambda2.linear -= lin_impulse * self.im2; + mj_lambda2.angular -= ang_impulse2; + } + + pub fn remove_bias_from_rhs(&mut self) { + self.rhs = self.rhs_wo_bias; + } +} + +impl JointVelocityConstraint<Real, 1> { + pub fn lock_axes( + params: &IntegrationParameters, + joint_id: JointIndex, + body1: &SolverBody<Real, 1>, + body2: &SolverBody<Real, 1>, + frame1: &Isometry<Real>, + frame2: &Isometry<Real>, + joint: &JointData, + out: &mut [Self], + ) -> usize { + let mut len = 0; + let locked_axes = joint.locked_axes.bits(); + let motor_axes = joint.motor_axes.bits(); + let limit_axes = joint.limit_axes.bits(); + + let builder = JointVelocityConstraintBuilder::new( + frame1, + frame2, + &body1.world_com, + &body2.world_com, + locked_axes, + ); + + for i in 0..DIM { + if locked_axes & (1 << i) != 0 { + out[len] = + builder.lock_linear(params, [joint_id], body1, body2, i, WritebackId::Dof(i)); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if locked_axes & (1 << i) != 0 { + out[len] = builder.lock_angular( + params, + [joint_id], + body1, + body2, + i - DIM, + WritebackId::Dof(i), + ); + len += 1; + } + } + + for i in 0..DIM { + if motor_axes & (1 << i) != 0 { + out[len] = builder.motor_linear( + params, + [joint_id], + body1, + body2, + locked_axes >> DIM, + i, + &joint.motors[i].motor_params(params.dt), + WritebackId::Motor(i), + ); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if motor_axes & (1 << i) != 0 { + out[len] = builder.motor_angular( + [joint_id], + body1, + body2, + i - DIM, + &joint.motors[i].motor_params(params.dt), + WritebackId::Motor(i), + ); + len += 1; + } + } + + for i in 0..DIM { + if limit_axes & (1 << i) != 0 { + out[len] = builder.limit_linear( + params, + [joint_id], + body1, + body2, + i, + [joint.limits[i].min, joint.limits[i].max], + WritebackId::Limit(i), + ); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if limit_axes & (1 << i) != 0 { + out[len] = builder.limit_angular( + params, + [joint_id], + body1, + body2, + i - DIM, + [joint.limits[i].min, joint.limits[i].max], + WritebackId::Limit(i), + ); + len += 1; + } + } + + JointVelocityConstraintBuilder::finalize_constraints(&mut out[..len]); + len + } + + pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) { + let mut mj_lambda1 = mj_lambdas[self.mj_lambda1[0] as usize]; + let mut mj_lambda2 = mj_lambdas[self.mj_lambda2[0] as usize]; + + self.solve_generic(&mut mj_lambda1, &mut mj_lambda2); + + mj_lambdas[self.mj_lambda1[0] as usize] = mj_lambda1; + mj_lambdas[self.mj_lambda2[0] as usize] = mj_lambda2; + } + + pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) { + let joint = &mut joints_all[self.joint_id[0]].weight; + match self.writeback_id { + WritebackId::Dof(i) => joint.impulses[i] = self.impulse, + WritebackId::Limit(i) => joint.data.limits[i].impulse = self.impulse, + WritebackId::Motor(i) => joint.data.motors[i].impulse = self.impulse, + } + } +} +#[cfg(feature = "simd-is-enabled")] +impl JointVelocityConstraint<SimdReal, SIMD_WIDTH> { + pub fn lock_axes( + params: &IntegrationParameters, + joint_id: [JointIndex; SIMD_WIDTH], + body1: &SolverBody<SimdReal, SIMD_WIDTH>, + body2: &SolverBody<SimdReal, SIMD_WIDTH>, + frame1: &Isometry<SimdReal>, + frame2: &Isometry<SimdReal>, + locked_axes: u8, + out: &mut [Self], + ) -> usize { + let builder = JointVelocityConstraintBuilder::new( + frame1, + frame2, + &body1.world_com, + &body2.world_com, + locked_axes, + ); + + let mut len = 0; + for i in 0..DIM { + if locked_axes & (1 << i) != 0 { + out[len] = + builder.lock_linear(params, joint_id, body1, body2, i, WritebackId::Dof(i)); + len += 1; + } + } + + for i in DIM..SPATIAL_DIM { + if locked_axes & (1 << i) != 0 { + out[len] = builder.lock_angular( + params, + joint_id, + body1, + body2, + i - DIM, + WritebackId::Dof(i), + ); + len += 1; + } + } + + JointVelocityConstraintBuilder::finalize_constraints(&mut out[..len]); + len + } + + pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) { + let mut mj_lambda1 = DeltaVel { + linear: Vector::from(gather![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].linear]), + angular: AngVector::from(gather![ + |ii| mj_lambdas[self.mj_lambda1[ii] as usize].angular + ]), + }; + let mut mj_lambda2 = DeltaVel { + linear: Vector::from(gather![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear]), + angular: AngVector::from(gather![ + |ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular + ]), + }; + + self.solve_generic(&mut mj_lambda1, &mut mj_lambda2); + + for ii in 0..SIMD_WIDTH { + mj_lambdas[self.mj_lambda1[ii] as usize].linear = mj_lambda1.linear.extract(ii); + mj_lambdas[self.mj_lambda1[ii] as usize].angular = mj_lambda1.angular.extract(ii); + mj_lambdas[self.mj_lambda2[ii] as usize].linear = mj_lambda2.linear.extract(ii); + mj_lambdas[self.mj_lambda2[ii] as usize].angular = mj_lambda2.angular.extract(ii); + } + } + + pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) { + let impulses: [_; SIMD_WIDTH] = self.impulse.into(); + + // TODO: should we move the iteration on ii deeper in the mested match? + for ii in 0..SIMD_WIDTH { + let joint = &mut joints_all[self.joint_id[ii]].weight; + match self.writeback_id { + WritebackId::Dof(i) => joint.impulses[i] = impulses[ii], + WritebackId::Limit(i) => joint.data.limits[i].impulse = impulses[ii], + WritebackId::Motor(i) => joint.data.motors[i].impulse = impulses[ii], + } + } + } +} + +#[derive(Debug, Copy, Clone)] +pub struct JointVelocityGroundConstraint<N: SimdRealField, const LANES: usize> { + pub mj_lambda2: [usize; LANES], + + pub joint_id: [JointIndex; LANES], + + pub impulse: N, + pub impulse_bounds: [N; 2], + pub lin_jac: Vector<N>, + pub ang_jac2: AngVector<N>, + + pub inv_lhs: N, + pub rhs: N, + pub rhs_wo_bias: N, + + pub im2: N, + + pub writeback_id: WritebackId, +} + +impl<N: WReal, const LANES: usize> JointVelocityGroundConstraint<N, LANES> { + pub fn invalid() -> Self { + Self { + mj_lambda2: [crate::INVALID_USIZE; LANES], + joint_id: [crate::INVALID_USIZE; LANES], + impulse: N::zero(), + impulse_bounds: [N::zero(), N::zero()], + lin_jac: Vector::zeros(), + ang_jac2: na::zero(), + inv_lhs: N::zero(), + rhs: N::zero(), + rhs_wo_bias: N::zero(), + im2: N::zero(), + writeback_id: WritebackId::Dof(0), + } + } + + pub fn solve_generic(&mut self, mj_lambda2: &mut DeltaVel<N>) { + let dlinvel = mj_lambda2.linear; + let dangvel = mj_lambda2.angular; + + let dvel = self.lin_jac.dot(&dlinvel) + self.ang_jac2.gdot(dangvel) + self.rhs; + let total_impulse = (self.impulse + self.inv_lhs * dvel) + .simd_clamp(self.impulse_bounds[0], self.impulse_bounds[1]); + let delta_impulse = total_impulse - self.impulse; + self.impulse = total_impulse; + + 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.angular -= ang_impulse; + } + + pub fn remove_bias_from_rhs(&mut self) { + self.rhs = self.rhs_wo_bias; + } +} + +impl JointVelocityGroundConstraint<Real, 1> { + pub fn lock_axes( + params: &IntegrationParameters, + joint_id: JointIndex, + body1: &SolverBody<Real, 1>, + body2: &SolverBody<Real, 1>, + frame1: &Isometry<Real>, + frame2: &Isometry<Real>, + joint: &JointData, + out: &mut [Self], + ) -> usize { + let mut len = 0; + let locked_axes = joint.locked_axes.bits() as u8; + let motor_axes = joint.motor_axes.bits() as u8; + let limit_axes = joint.limit_axes.bits() as u8; + + let builder = JointVelocityConstraintBuilder::new( + frame1, + frame2, + &body1.world_com, + &body2.world_com, + locked_axes, + ); + + for i in 0..DIM { + if locked_axes & (1 << i) != 0 { + out[len] = builder.lock_linear_ground( + params, + [joint_id], + body1, + body2, + i, + WritebackId::Dof(i), + ); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if locked_axes & (1 << i) != 0 { + out[len] = builder.lock_angular_ground( + params, + [joint_id], + body1, + body2, + i - DIM, + WritebackId::Dof(i), + ); + len += 1; + } + } + + for i in 0..DIM { + if motor_axes & (1 << i) != 0 { + out[len] = builder.motor_linear_ground( + [joint_id], + body1, + body2, + i, + &joint.motors[i].motor_params(params.dt), + WritebackId::Motor(i), + ); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if motor_axes & (1 << i) != 0 { + out[len] = builder.motor_angular_ground( + [joint_id], + body1, + body2, + i - DIM, + &joint.motors[i].motor_params(params.dt), + WritebackId::Motor(i), + ); + len += 1; + } + } + + for i in 0..DIM { + if limit_axes & (1 << i) != 0 { + out[len] = builder.limit_linear_ground( + params, + [joint_id], + body1, + body2, + i, + [joint.limits[i].min, joint.limits[i].max], + WritebackId::Limit(i), + ); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if limit_axes & (1 << i) != 0 { + out[len] = builder.limit_angular_ground( + params, + [joint_id], + body1, + body2, + i - DIM, + [joint.limits[i].min, joint.limits[i].max], + WritebackId::Limit(i), + ); + len += 1; + } + } + + JointVelocityConstraintBuilder::finalize_ground_constraints(&mut out[..len]); + len + } + + pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) { + let mut mj_lambda2 = mj_lambdas[self.mj_lambda2[0] as usize]; + self.solve_generic(&mut mj_lambda2); + mj_lambdas[self.mj_lambda2[0] as usize] = mj_lambda2; + } + + pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) { + let joint = &mut joints_all[self.joint_id[0]].weight; + match self.writeback_id { + WritebackId::Dof(i) => joint.impulses[i] = self.impulse, + WritebackId::Limit(i) => joint.data.limits[i].impulse = self.impulse, + WritebackId::Motor(i) => joint.data.motors[i].impulse = self.impulse, + } + } +} + +#[cfg(feature = "simd-is-enabled")] +impl JointVelocityGroundConstraint<SimdReal, SIMD_WIDTH> { + pub fn lock_axes( + params: &IntegrationParameters, + joint_id: [JointIndex; SIMD_WIDTH], + body1: &SolverBody<SimdReal, SIMD_WIDTH>, + body2: &SolverBody<SimdReal, SIMD_WIDTH>, + frame1: &Isometry<SimdReal>, + frame2: &Isometry<SimdReal>, + locked_axes: u8, + out: &mut [Self], + ) -> usize { + let mut len = 0; + let builder = JointVelocityConstraintBuilder::new( + frame1, + frame2, + &body1.world_com, + &body2.world_com, + locked_axes, + ); + + for i in 0..DIM { + if locked_axes & (1 << i) != 0 { + out[len] = builder.lock_linear_ground( + params, + joint_id, + body1, + body2, + i, + WritebackId::Dof(i), + ); + len += 1; + } + } + for i in DIM..SPATIAL_DIM { + if locked_axes & (1 << i) != 0 { + out[len] = builder.lock_angular_ground( + params, + joint_id, + body1, + body2, + i - DIM, + WritebackId::Dof(i), + ); + len += 1; + } + } + + JointVelocityConstraintBuilder::finalize_ground_constraints(&mut out[..len]); + len + } + + pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) { + let mut mj_lambda2 = DeltaVel { + linear: Vector::from(gather![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear]), + angular: AngVector::from(gather![ + |ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular + ]), + }; + + self.solve_generic(&mut mj_lambda2); + + for ii in 0..SIMD_WIDTH { + mj_lambdas[self.mj_lambda2[ii] as usize].linear = mj_lambda2.linear.extract(ii); + mj_lambdas[self.mj_lambda2[ii] as usize].angular = mj_lambda2.angular.extract(ii); + } + } + + pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) { + let impulses: [_; SIMD_WIDTH] = self.impulse.into(); + + // TODO: should we move the iteration on ii deeper in the mested match? + for ii in 0..SIMD_WIDTH { + let joint = &mut joints_all[self.joint_id[ii]].weight; + match self.writeback_id { + WritebackId::Dof(i) => joint.impulses[i] = impulses[ii], + WritebackId::Limit(i) => joint.data.limits[i].impulse = impulses[ii], + WritebackId::Motor(i) => joint.data.motors[i].impulse = impulses[ii], + } + } + } +} |