use crate::dynamics::solver::MotorParameters;
use crate::dynamics::solver::joint_constraint::JointSolverBody;
use crate::dynamics::solver::joint_constraint::joint_velocity_constraint::{
JointConstraint, WritebackId,
};
use crate::dynamics::{IntegrationParameters, JointIndex};
use crate::math::{DIM, Real};
use crate::utils;
#[cfg(feature = "dim3")]
use crate::utils::OrthonormalBasis;
use crate::utils::{
AngularInertiaOps, ComponentMul, CrossProductMatrix, DotProduct, IndexMut2, MatrixColumn,
PoseOps, RotationOps, ScalarType, SimdLength,
};
#[cfg(feature = "dim2")]
use crate::num::One;
#[cfg(feature = "dim3")]
use parry::math::Rot3;
#[derive(Debug, Copy, Clone)]
pub struct JointConstraintHelper<N: ScalarType> {
pub basis: N::Matrix,
#[cfg(feature = "dim3")]
pub basis2: N::Matrix, #[cfg(feature = "dim3")]
pub cmat1_basis: N::Matrix,
#[cfg(feature = "dim3")]
pub cmat2_basis: N::Matrix,
#[cfg(feature = "dim3")]
pub ang_basis: N::Matrix,
#[cfg(feature = "dim2")]
pub cmat1_basis: [N::AngVector; 2],
#[cfg(feature = "dim2")]
pub cmat2_basis: [N::AngVector; 2],
pub lin_err: N::Vector,
pub ang_err: N::Rotation,
}
impl<N: ScalarType> JointConstraintHelper<N> {
pub fn new(
frame1: &N::Pose,
frame2: &N::Pose,
world_com1: &N::Vector,
world_com2: &N::Vector,
locked_lin_axes: u8,
) -> Self {
let mut frame1 = *frame1;
let basis = frame1.rotation().to_mat();
let lin_err = frame2.translation() - frame1.translation();
{
let mut new_center1 = frame2.translation();
for i in 0..DIM {
if locked_lin_axes & (1 << i) != 0 {
let axis = basis.column(i);
new_center1 -= axis * lin_err.gdot(axis);
}
}
frame1.set_translation(new_center1);
}
let r1 = frame1.translation() - *world_com1;
let r2 = frame2.translation() - *world_com2;
let cmat1 = r1.gcross_matrix();
let cmat2 = r2.gcross_matrix();
#[cfg(feature = "dim3")]
let mut ang_basis = frame1.rotation().diff_conj1_2_tr(&frame2.rotation());
#[allow(unused_mut)] 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.mul_assign_unchecked(sgn);
}
#[cfg(feature = "dim2")]
return Self {
basis,
cmat1_basis: [
cmat1.gdot(basis.column(0)).into(),
cmat1.gdot(basis.column(1)).into(),
],
cmat2_basis: [
cmat2.gdot(basis.column(0)).into(),
cmat2.gdot(basis.column(1)).into(),
],
lin_err,
ang_err,
};
#[cfg(feature = "dim3")]
return Self {
basis,
basis2: frame2.rotation().to_mat(),
cmat1_basis: cmat1 * basis,
cmat2_basis: cmat2 * basis,
ang_basis,
lin_err,
ang_err,
};
}
pub fn limit_linear<const LANES: usize>(
&self,
params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
limited_axis: usize,
limits: [N; 2],
writeback_id: WritebackId,
erp_inv_dt: N,
cfm_coeff: N,
) -> JointConstraint<N, LANES> {
let zero = N::zero();
let mut constraint = self.lock_linear(
params,
joint_id,
body1,
body2,
limited_axis,
writeback_id,
erp_inv_dt,
cfm_coeff,
);
let dist = self.lin_err.gdot(constraint.lin_jac);
let min_enabled = dist.simd_le(limits[0]);
let max_enabled = limits[1].simd_le(dist);
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<const LANES: usize>(
&self,
params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
coupled_axes: u8,
limits: [N; 2],
writeback_id: WritebackId,
erp_inv_dt: N,
cfm_coeff: N,
) -> JointConstraint<N, LANES> {
let zero = N::zero();
let mut lin_jac: N::Vector = Default::default();
let mut ang_jac1: N::AngVector = Default::default();
let mut ang_jac2: N::AngVector = Default::default();
for i in 0..DIM {
if coupled_axes & (1 << i) != 0 {
let coeff = self.basis.column(i).gdot(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).into() * coeff;
ang_jac2 += self.cmat2_basis.column(i).into() * coeff;
}
}
}
let dist = lin_jac.simd_length();
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());
let ii_ang_jac1 = body1.ii.transform_vector(ang_jac1);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac2);
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)];
JointConstraint {
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,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: N::zero(), cfm_coeff,
cfm_gain: N::zero(),
rhs,
rhs_wo_bias,
writeback_id,
}
}
pub fn motor_linear<const LANES: usize>(
&self,
params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
motor_axis: usize,
motor_params: &MotorParameters<N>,
limits: Option<[N; 2]>,
writeback_id: WritebackId,
) -> JointConstraint<N, LANES> {
let inv_dt = N::splat(params.inv_dt());
let mut constraint = self.lock_linear(
params,
joint_id,
body1,
body2,
motor_axis,
writeback_id,
N::zero(),
N::zero(),
);
let mut rhs_wo_bias = N::zero();
if motor_params.erp_inv_dt != N::zero() {
let dist = self.lin_err.gdot(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.gdot(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 motor_linear_coupled<const LANES: usize>(
&self,
params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
coupled_axes: u8,
motor_params: &MotorParameters<N>,
limits: Option<[N; 2]>,
writeback_id: WritebackId,
) -> JointConstraint<N, LANES> {
let inv_dt = N::splat(params.inv_dt());
let mut lin_jac: N::Vector = Default::default();
let mut ang_jac1: N::AngVector = Default::default();
let mut ang_jac2: N::AngVector = Default::default();
for i in 0..DIM {
if coupled_axes & (1 << i) != 0 {
let coeff = self.basis.column(i).gdot(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).into() * coeff;
ang_jac2 += self.cmat2_basis.column(i).into() * coeff;
}
}
}
let dist = lin_jac.simd_length();
let inv_dist = crate::utils::simd_inv(dist);
lin_jac *= inv_dist;
ang_jac1 *= inv_dist;
ang_jac2 *= inv_dist;
let mut rhs_wo_bias = N::zero();
if motor_params.erp_inv_dt != N::zero() {
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 {
target_vel =
target_vel.simd_clamp((limits[0] - dist) * inv_dt, (limits[1] - dist) * inv_dt);
};
rhs_wo_bias += -target_vel;
let ii_ang_jac1 = body1.ii.transform_vector(ang_jac1);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac2);
JointConstraint {
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,
ang_jac1,
ang_jac2,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: N::zero(), cfm_coeff: motor_params.cfm_coeff,
cfm_gain: motor_params.cfm_gain,
rhs: rhs_wo_bias,
rhs_wo_bias,
writeback_id,
}
}
pub fn lock_linear<const LANES: usize>(
&self,
_params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
locked_axis: usize,
writeback_id: WritebackId,
erp_inv_dt: N,
cfm_coeff: N,
) -> JointConstraint<N, LANES> {
let lin_jac = self.basis.column(locked_axis);
#[cfg(feature = "dim2")]
let ang_jac1 = self.cmat1_basis[locked_axis];
#[cfg(feature = "dim2")]
let ang_jac2 = self.cmat2_basis[locked_axis];
#[cfg(feature = "dim3")]
let ang_jac1 = self.cmat1_basis.column(locked_axis).into();
#[cfg(feature = "dim3")]
let ang_jac2 = self.cmat2_basis.column(locked_axis).into();
let rhs_wo_bias = N::zero();
let rhs_bias = lin_jac.gdot(self.lin_err) * erp_inv_dt;
let ii_ang_jac1 = body1.ii.transform_vector(ang_jac1);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac2);
JointConstraint {
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,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: N::zero(), cfm_coeff,
cfm_gain: N::zero(),
rhs: rhs_wo_bias + rhs_bias,
rhs_wo_bias,
writeback_id,
}
}
pub fn limit_angular<const LANES: usize>(
&self,
_params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
_limited_axis: usize,
s_limits: [N; 2],
writeback_id: WritebackId,
erp_inv_dt: N,
cfm_coeff: N,
) -> JointConstraint<N, LANES> {
let zero = N::zero();
#[cfg(feature = "dim2")]
let half = N::splat(0.5);
#[cfg(feature = "dim2")]
let s_ang = ((N::one() - self.ang_err.real()).simd_max(zero) * half)
.simd_sqrt()
.simd_copysign(self.ang_err.imag());
#[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::AngVector::one();
#[cfg(feature = "dim3")]
let ang_jac = self.ang_basis.column(_limited_axis).into();
let rhs_wo_bias = N::zero();
let rhs_bias = ((s_ang - s_limits[1]).simd_max(zero)
- (s_limits[0] - s_ang).simd_max(zero))
* erp_inv_dt;
let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
JointConstraint {
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: Default::default(),
ang_jac1: ang_jac,
ang_jac2: ang_jac,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: N::zero(), cfm_coeff,
cfm_gain: N::zero(),
rhs: rhs_wo_bias + rhs_bias,
rhs_wo_bias,
writeback_id,
}
}
pub fn motor_angular<const LANES: usize>(
&self,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
_motor_axis: usize,
motor_params: &MotorParameters<N>,
writeback_id: WritebackId,
) -> JointConstraint<N, LANES> {
#[cfg(feature = "dim2")]
let ang_jac = N::AngVector::one();
#[cfg(feature = "dim3")]
let ang_jac = self.basis.column(_motor_axis).into();
let mut rhs_wo_bias = N::zero();
if motor_params.erp_inv_dt != N::zero() {
let ang_dist;
#[cfg(feature = "dim2")]
{
ang_dist = self.ang_err.angle();
}
#[cfg(feature = "dim3")]
{
let clamped_err = self.ang_err.imag()[_motor_axis].simd_clamp(-N::one(), N::one());
ang_dist = clamped_err.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 ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
JointConstraint {
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: Default::default(),
ang_jac1: ang_jac,
ang_jac2: ang_jac,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: N::zero(), 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<const LANES: usize>(
&self,
_params: &IntegrationParameters,
joint_id: [JointIndex; LANES],
body1: &JointSolverBody<N, LANES>,
body2: &JointSolverBody<N, LANES>,
_locked_axis: usize,
writeback_id: WritebackId,
erp_inv_dt: N,
cfm_coeff: N,
) -> JointConstraint<N, LANES> {
#[cfg(feature = "dim2")]
let ang_jac = N::AngVector::one();
#[cfg(feature = "dim3")]
let ang_jac = self.ang_basis.column(_locked_axis).into();
let rhs_wo_bias = N::zero();
#[cfg(feature = "dim2")]
let rhs_bias = self.ang_err.imag() * erp_inv_dt;
#[cfg(feature = "dim3")]
let rhs_bias = self.ang_err.imag()[_locked_axis] * erp_inv_dt;
let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
JointConstraint {
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: Default::default(),
ang_jac1: ang_jac,
ang_jac2: ang_jac,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: N::zero(), cfm_coeff,
cfm_gain: N::zero(),
rhs: rhs_wo_bias + rhs_bias,
rhs_wo_bias,
writeback_id,
}
}
pub fn finalize_constraints<const LANES: usize>(constraints: &mut [JointConstraint<N, LANES>]) {
let len = constraints.len();
if len == 0 {
return;
}
let imsum = constraints[0].im1 + constraints[0].im2;
for j in 0..len {
let c_j = &mut constraints[j];
let dot_jj = c_j.lin_jac.gdot(imsum.component_mul(&c_j.lin_jac))
+ c_j.ii_ang_jac1.gdot(c_j.ang_jac1)
+ c_j.ii_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); c_j.cfm_gain = cfm_gain;
if c_j.impulse_bounds != [-N::splat(Real::MAX), N::splat(Real::MAX)] {
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.gdot(imsum.component_mul(&c_j.lin_jac))
+ c_i.ii_ang_jac1.gdot(c_j.ang_jac1)
+ c_i.ii_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.ii_ang_jac1 -= c_j.ii_ang_jac1 * coeff;
c_i.ii_ang_jac2 -= c_j.ii_ang_jac2 * coeff;
c_i.rhs_wo_bias -= c_j.rhs_wo_bias * coeff;
c_i.rhs -= c_j.rhs * coeff;
}
}
}
}
impl JointConstraintHelper<Real> {
#[cfg(feature = "dim3")]
pub fn limit_angular_coupled(
&self,
_params: &IntegrationParameters,
joint_id: [JointIndex; 1],
body1: &JointSolverBody<Real, 1>,
body2: &JointSolverBody<Real, 1>,
coupled_axes: u8,
limits: [Real; 2],
writeback_id: WritebackId,
erp_inv_dt: Real,
cfm_coeff: Real,
) -> JointConstraint<Real, 1> {
let ang_coupled_axes = coupled_axes >> DIM;
assert_eq!(ang_coupled_axes.count_ones(), 2);
let not_coupled_index = ang_coupled_axes.trailing_ones() as usize;
let axis1 = self.basis.column(not_coupled_index);
let axis2 = self.basis2.column(not_coupled_index);
let rot = Rot3::from_rotation_arc(axis1, axis2);
let (mut ang_jac, angle) = rot.to_axis_angle();
if angle == 0.0 {
ang_jac = axis1.orthonormal_basis()[0];
}
let min_enabled = angle <= limits[0];
let max_enabled = limits[1] <= angle;
let impulse_bounds = [
if min_enabled { -Real::INFINITY } else { 0.0 },
if max_enabled { Real::INFINITY } else { 0.0 },
];
let rhs_wo_bias = 0.0;
let rhs_bias = ((angle - limits[1]).max(0.0) - (limits[0] - angle).max(0.0)) * erp_inv_dt;
let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
JointConstraint {
joint_id,
solver_vel1: body1.solver_vel,
solver_vel2: body2.solver_vel,
im1: body1.im,
im2: body2.im,
impulse: 0.0,
impulse_bounds,
lin_jac: Default::default(),
ang_jac1: ang_jac,
ang_jac2: ang_jac,
ii_ang_jac1,
ii_ang_jac2,
inv_lhs: 0.0, cfm_coeff,
cfm_gain: 0.0,
rhs: rhs_wo_bias + rhs_bias,
rhs_wo_bias,
writeback_id,
}
}
}