use simba::simd::SimdValue;
use crate::dynamics::solver::DeltaVel;
use crate::dynamics::{
IntegrationParameters, JointGraphEdge, JointIndex, JointParams, RevoluteJoint, RigidBody,
};
use crate::math::{
AngVector, AngularInertia, Isometry, Point, Real, Rotation, SimdReal, Vector, SIMD_WIDTH,
};
use crate::utils::{WAngularInertia, WCross, WCrossMatrix};
use na::{Cholesky, Matrix3x2, Matrix5, Unit, Vector5, U2, U3};
#[derive(Debug)]
pub(crate) struct WRevoluteVelocityConstraint {
mj_lambda1: [usize; SIMD_WIDTH],
mj_lambda2: [usize; SIMD_WIDTH],
joint_id: [JointIndex; SIMD_WIDTH],
r1: Vector<SimdReal>,
r2: Vector<SimdReal>,
inv_lhs: Matrix5<SimdReal>,
rhs: Vector5<SimdReal>,
impulse: Vector5<SimdReal>,
axis1: [Vector<Real>; SIMD_WIDTH],
basis1: Matrix3x2<SimdReal>,
basis2: Matrix3x2<SimdReal>,
im1: SimdReal,
im2: SimdReal,
ii1_sqrt: AngularInertia<SimdReal>,
ii2_sqrt: AngularInertia<SimdReal>,
}
impl WRevoluteVelocityConstraint {
pub fn from_params(
params: &IntegrationParameters,
joint_id: [JointIndex; SIMD_WIDTH],
rbs1: [&RigidBody; SIMD_WIDTH],
rbs2: [&RigidBody; SIMD_WIDTH],
joints: [&RevoluteJoint; SIMD_WIDTH],
) -> Self {
let position1 = Isometry::from(array![|ii| rbs1[ii].position; SIMD_WIDTH]);
let linvel1 = Vector::from(array![|ii| rbs1[ii].linvel; SIMD_WIDTH]);
let angvel1 = AngVector::<SimdReal>::from(array![|ii| rbs1[ii].angvel; SIMD_WIDTH]);
let world_com1 = Point::from(array![|ii| rbs1[ii].world_com; SIMD_WIDTH]);
let im1 = SimdReal::from(array![|ii| rbs1[ii].effective_inv_mass; SIMD_WIDTH]);
let ii1_sqrt = AngularInertia::<SimdReal>::from(
array![|ii| rbs1[ii].effective_world_inv_inertia_sqrt; SIMD_WIDTH],
);
let mj_lambda1 = array![|ii| rbs1[ii].active_set_offset; SIMD_WIDTH];
let position2 = Isometry::from(array![|ii| rbs2[ii].position; SIMD_WIDTH]);
let linvel2 = Vector::from(array![|ii| rbs2[ii].linvel; SIMD_WIDTH]);
let angvel2 = AngVector::<SimdReal>::from(array![|ii| rbs2[ii].angvel; SIMD_WIDTH]);
let world_com2 = Point::from(array![|ii| rbs2[ii].world_com; SIMD_WIDTH]);
let im2 = SimdReal::from(array![|ii| rbs2[ii].effective_inv_mass; SIMD_WIDTH]);
let ii2_sqrt = AngularInertia::<SimdReal>::from(
array![|ii| rbs2[ii].effective_world_inv_inertia_sqrt; SIMD_WIDTH],
);
let mj_lambda2 = array![|ii| rbs2[ii].active_set_offset; SIMD_WIDTH];
let local_anchor1 = Point::from(array![|ii| joints[ii].local_anchor1; SIMD_WIDTH]);
let local_anchor2 = Point::from(array![|ii| joints[ii].local_anchor2; SIMD_WIDTH]);
let local_basis1 = [
Vector::from(array![|ii| joints[ii].basis1[0]; SIMD_WIDTH]),
Vector::from(array![|ii| joints[ii].basis1[1]; SIMD_WIDTH]),
];
let local_basis2 = [
Vector::from(array![|ii| joints[ii].basis2[0]; SIMD_WIDTH]),
Vector::from(array![|ii| joints[ii].basis2[1]; SIMD_WIDTH]),
];
let impulse = Vector5::from(array![|ii| joints[ii].impulse; SIMD_WIDTH]);
let anchor1 = position1 * local_anchor1;
let anchor2 = position2 * local_anchor2;
let basis1 =
Matrix3x2::from_columns(&[position1 * local_basis1[0], position1 * local_basis1[1]]);
let basis2 =
Matrix3x2::from_columns(&[position2 * local_basis2[0], position2 * local_basis2[1]]);
let basis_projection2 = basis2 * basis2.transpose();
let basis2 = basis_projection2 * basis1;
let ii1 = ii1_sqrt.squared();
let r1 = anchor1 - world_com1;
let r1_mat = r1.gcross_matrix();
let ii2 = ii2_sqrt.squared();
let r2 = anchor2 - world_com2;
let r2_mat = r2.gcross_matrix();
let mut lhs = Matrix5::zeros();
let lhs00 =
ii2.quadform(&r2_mat).add_diagonal(im2) + ii1.quadform(&r1_mat).add_diagonal(im1);
let lhs10 = basis1.tr_mul(&(ii2 * r2_mat)) + basis2.tr_mul(&(ii1 * r1_mat));
let lhs11 = (ii1.quadform3x2(&basis1) + ii2.quadform3x2(&basis2)).into_matrix();
lhs.fixed_slice_mut::<U3, U3>(0, 0)
.copy_from(&lhs00.into_matrix());
lhs.fixed_slice_mut::<U2, U3>(3, 0).copy_from(&lhs10);
lhs.fixed_slice_mut::<U2, U2>(3, 3).copy_from(&lhs11);
let inv_lhs = Cholesky::new_unchecked(lhs).inverse();
let linvel_err = linvel2 + angvel2.gcross(r2) - linvel1 - angvel1.gcross(r1);
let angvel_err = basis2.tr_mul(&angvel2) - basis1.tr_mul(&angvel1);
let mut rhs = Vector5::new(
linvel_err.x,
linvel_err.y,
linvel_err.z,
angvel_err.x,
angvel_err.y,
) * SimdReal::splat(params.velocity_solve_fraction);
let velocity_based_erp_inv_dt = params.velocity_based_erp_inv_dt();
if velocity_based_erp_inv_dt != 0.0 {
let velocity_based_erp_inv_dt = SimdReal::splat(velocity_based_erp_inv_dt);
let lin_err = anchor2 - anchor1;
let local_axis1 =
Unit::<Vector<_>>::from(array![|ii| joints[ii].local_axis1; SIMD_WIDTH]);
let local_axis2 =
Unit::<Vector<_>>::from(array![|ii| joints[ii].local_axis2; SIMD_WIDTH]);
let axis1 = position1 * local_axis1;
let axis2 = position2 * local_axis2;
let axis_error = axis1.cross(&axis2);
let ang_err =
(basis2.tr_mul(&axis_error) + basis1.tr_mul(&axis_error)) * SimdReal::splat(0.5);
rhs += Vector5::new(lin_err.x, lin_err.y, lin_err.z, ang_err.x, ang_err.y)
* velocity_based_erp_inv_dt;
}
let warmstart_coeff = SimdReal::splat(params.warmstart_coeff);
let mut impulse = impulse * warmstart_coeff;
let axis1 = array![|ii| rbs1[ii].position * *joints[ii].local_axis1; SIMD_WIDTH];
let rotated_impulse = Vector::from(array![|ii| {
let axis_rot = Rotation::rotation_between(&joints[ii].prev_axis1, &axis1[ii])
.unwrap_or_else(Rotation::identity);
axis_rot * joints[ii].world_ang_impulse
}; SIMD_WIDTH]);
let rotated_basis_impulse = basis1.tr_mul(&rotated_impulse);
impulse[3] = rotated_basis_impulse.x * warmstart_coeff;
impulse[4] = rotated_basis_impulse.y * warmstart_coeff;
WRevoluteVelocityConstraint {
joint_id,
mj_lambda1,
mj_lambda2,
im1,
ii1_sqrt,
axis1,
basis1,
basis2,
im2,
ii2_sqrt,
impulse,
inv_lhs,
rhs,
r1,
r2,
}
}
pub fn warmstart(&self, mj_lambdas: &mut [DeltaVel<Real>]) {
let mut mj_lambda1 = DeltaVel {
linear: Vector::from(
array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].linear; SIMD_WIDTH],
),
angular: AngVector::from(
array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].angular; SIMD_WIDTH],
),
};
let mut mj_lambda2 = DeltaVel {
linear: Vector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
),
angular: AngVector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
),
};
let lin_impulse1 = self.impulse.fixed_rows::<U3>(0).into_owned();
let lin_impulse2 = self.impulse.fixed_rows::<U3>(0).into_owned();
let ang_impulse1 = self.basis1 * self.impulse.fixed_rows::<U2>(3).into_owned();
let ang_impulse2 = self.basis2 * self.impulse.fixed_rows::<U2>(3).into_owned();
mj_lambda1.linear += lin_impulse1 * self.im1;
mj_lambda1.angular += self
.ii1_sqrt
.transform_vector(ang_impulse1 + self.r1.gcross(lin_impulse1));
mj_lambda2.linear -= lin_impulse2 * self.im2;
mj_lambda2.angular -= self
.ii2_sqrt
.transform_vector(ang_impulse2 + self.r2.gcross(lin_impulse2));
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);
}
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 solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
let mut mj_lambda1 = DeltaVel {
linear: Vector::from(
array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].linear; SIMD_WIDTH],
),
angular: AngVector::from(
array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].angular; SIMD_WIDTH],
),
};
let mut mj_lambda2 = DeltaVel {
linear: Vector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
),
angular: AngVector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
),
};
let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
let lin_dvel = (mj_lambda2.linear + ang_vel2.gcross(self.r2))
- (mj_lambda1.linear + ang_vel1.gcross(self.r1));
let ang_dvel = self.basis2.tr_mul(&ang_vel2) - self.basis1.tr_mul(&ang_vel1);
let rhs =
Vector5::new(lin_dvel.x, lin_dvel.y, lin_dvel.z, ang_dvel.x, ang_dvel.y) + self.rhs;
let impulse = self.inv_lhs * rhs;
self.impulse += impulse;
let lin_impulse1 = impulse.fixed_rows::<U3>(0).into_owned();
let lin_impulse2 = impulse.fixed_rows::<U3>(0).into_owned();
let ang_impulse1 = self.basis1 * impulse.fixed_rows::<U2>(3).into_owned();
let ang_impulse2 = self.basis2 * impulse.fixed_rows::<U2>(3).into_owned();
mj_lambda1.linear += lin_impulse1 * self.im1;
mj_lambda1.angular += self
.ii1_sqrt
.transform_vector(ang_impulse1 + self.r1.gcross(lin_impulse1));
mj_lambda2.linear -= lin_impulse2 * self.im2;
mj_lambda2.angular -= self
.ii2_sqrt
.transform_vector(ang_impulse2 + self.r2.gcross(lin_impulse2));
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);
}
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 rot_part = self.impulse.fixed_rows::<U2>(3).into_owned();
let world_ang_impulse = self.basis1 * rot_part;
for ii in 0..SIMD_WIDTH {
let joint = &mut joints_all[self.joint_id[ii]].weight;
if let JointParams::RevoluteJoint(rev) = &mut joint.params {
rev.impulse = self.impulse.extract(ii);
rev.world_ang_impulse = world_ang_impulse.extract(ii);
rev.prev_axis1 = self.axis1[ii];
}
}
}
}
#[derive(Debug)]
pub(crate) struct WRevoluteVelocityGroundConstraint {
mj_lambda2: [usize; SIMD_WIDTH],
joint_id: [JointIndex; SIMD_WIDTH],
r2: Vector<SimdReal>,
inv_lhs: Matrix5<SimdReal>,
rhs: Vector5<SimdReal>,
impulse: Vector5<SimdReal>,
basis2: Matrix3x2<SimdReal>,
im2: SimdReal,
ii2_sqrt: AngularInertia<SimdReal>,
}
impl WRevoluteVelocityGroundConstraint {
pub fn from_params(
params: &IntegrationParameters,
joint_id: [JointIndex; SIMD_WIDTH],
rbs1: [&RigidBody; SIMD_WIDTH],
rbs2: [&RigidBody; SIMD_WIDTH],
joints: [&RevoluteJoint; SIMD_WIDTH],
flipped: [bool; SIMD_WIDTH],
) -> Self {
let position1 = Isometry::from(array![|ii| rbs1[ii].position; SIMD_WIDTH]);
let linvel1 = Vector::from(array![|ii| rbs1[ii].linvel; SIMD_WIDTH]);
let angvel1 = AngVector::<SimdReal>::from(array![|ii| rbs1[ii].angvel; SIMD_WIDTH]);
let world_com1 = Point::from(array![|ii| rbs1[ii].world_com; SIMD_WIDTH]);
let position2 = Isometry::from(array![|ii| rbs2[ii].position; SIMD_WIDTH]);
let linvel2 = Vector::from(array![|ii| rbs2[ii].linvel; SIMD_WIDTH]);
let angvel2 = AngVector::<SimdReal>::from(array![|ii| rbs2[ii].angvel; SIMD_WIDTH]);
let world_com2 = Point::from(array![|ii| rbs2[ii].world_com; SIMD_WIDTH]);
let im2 = SimdReal::from(array![|ii| rbs2[ii].effective_inv_mass; SIMD_WIDTH]);
let ii2_sqrt = AngularInertia::<SimdReal>::from(
array![|ii| rbs2[ii].effective_world_inv_inertia_sqrt; SIMD_WIDTH],
);
let mj_lambda2 = array![|ii| rbs2[ii].active_set_offset; SIMD_WIDTH];
let impulse = Vector5::from(array![|ii| joints[ii].impulse; SIMD_WIDTH]);
let local_anchor1 = Point::from(
array![|ii| if flipped[ii] { joints[ii].local_anchor2 } else { joints[ii].local_anchor1 }; SIMD_WIDTH],
);
let local_anchor2 = Point::from(
array![|ii| if flipped[ii] { joints[ii].local_anchor1 } else { joints[ii].local_anchor2 }; SIMD_WIDTH],
);
let basis1 = Matrix3x2::from_columns(&[
position1
* Vector::from(
array![|ii| if flipped[ii] { joints[ii].basis2[0] } else { joints[ii].basis1[0] }; SIMD_WIDTH],
),
position1
* Vector::from(
array![|ii| if flipped[ii] { joints[ii].basis2[1] } else { joints[ii].basis1[1] }; SIMD_WIDTH],
),
]);
let basis2 = Matrix3x2::from_columns(&[
position2
* Vector::from(
array![|ii| if flipped[ii] { joints[ii].basis1[0] } else { joints[ii].basis2[0] }; SIMD_WIDTH],
),
position2
* Vector::from(
array![|ii| if flipped[ii] { joints[ii].basis1[1] } else { joints[ii].basis2[1] }; SIMD_WIDTH],
),
]);
let basis_projection2 = basis2 * basis2.transpose();
let basis2 = basis_projection2 * basis1;
let anchor1 = position1 * local_anchor1;
let anchor2 = position2 * local_anchor2;
let ii2 = ii2_sqrt.squared();
let r1 = anchor1 - world_com1;
let r2 = anchor2 - world_com2;
let r2_mat = r2.gcross_matrix();
let mut lhs = Matrix5::zeros();
let lhs00 = ii2.quadform(&r2_mat).add_diagonal(im2);
let lhs10 = basis2.tr_mul(&(ii2 * r2_mat));
let lhs11 = ii2.quadform3x2(&basis2).into_matrix();
lhs.fixed_slice_mut::<U3, U3>(0, 0)
.copy_from(&lhs00.into_matrix());
lhs.fixed_slice_mut::<U2, U3>(3, 0).copy_from(&lhs10);
lhs.fixed_slice_mut::<U2, U2>(3, 3).copy_from(&lhs11);
let inv_lhs = Cholesky::new_unchecked(lhs).inverse();
let linvel_err = (linvel2 + angvel2.gcross(r2)) - (linvel1 + angvel1.gcross(r1));
let angvel_err = basis2.tr_mul(&angvel2) - basis1.tr_mul(&angvel1);
let mut rhs = Vector5::new(
linvel_err.x,
linvel_err.y,
linvel_err.z,
angvel_err.x,
angvel_err.y,
) * SimdReal::splat(params.velocity_solve_fraction);
let velocity_based_erp_inv_dt = params.velocity_based_erp_inv_dt();
if velocity_based_erp_inv_dt != 0.0 {
let velocity_based_erp_inv_dt = SimdReal::splat(velocity_based_erp_inv_dt);
let lin_err = anchor2 - anchor1;
let local_axis1 = Unit::<Vector<_>>::from(
array![|ii| if flipped[ii] { joints[ii].local_axis2 } else { joints[ii].local_axis1 }; SIMD_WIDTH],
);
let local_axis2 = Unit::<Vector<_>>::from(
array![|ii| if flipped[ii] { joints[ii].local_axis1 } else { joints[ii].local_axis2 }; SIMD_WIDTH],
);
let axis1 = position1 * local_axis1;
let axis2 = position2 * local_axis2;
let axis_error = axis1.cross(&axis2);
let ang_err = basis2.tr_mul(&axis_error) - basis1.tr_mul(&axis_error);
rhs += Vector5::new(lin_err.x, lin_err.y, lin_err.z, ang_err.x, ang_err.y)
* velocity_based_erp_inv_dt;
}
WRevoluteVelocityGroundConstraint {
joint_id,
mj_lambda2,
im2,
ii2_sqrt,
impulse: impulse * SimdReal::splat(params.warmstart_coeff),
basis2,
inv_lhs,
rhs,
r2,
}
}
pub fn warmstart(&self, mj_lambdas: &mut [DeltaVel<Real>]) {
let mut mj_lambda2 = DeltaVel {
linear: Vector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
),
angular: AngVector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
),
};
let lin_impulse = self.impulse.fixed_rows::<U3>(0).into_owned();
let ang_impulse = self.basis2 * self.impulse.fixed_rows::<U2>(3).into_owned();
mj_lambda2.linear -= lin_impulse * self.im2;
mj_lambda2.angular -= self
.ii2_sqrt
.transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
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 solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
let mut mj_lambda2 = DeltaVel {
linear: Vector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
),
angular: AngVector::from(
array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
),
};
let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
let lin_dvel = mj_lambda2.linear + ang_vel2.gcross(self.r2);
let ang_dvel = self.basis2.tr_mul(&ang_vel2);
let rhs =
Vector5::new(lin_dvel.x, lin_dvel.y, lin_dvel.z, ang_dvel.x, ang_dvel.y) + self.rhs;
let impulse = self.inv_lhs * rhs;
self.impulse += impulse;
let lin_impulse = impulse.fixed_rows::<U3>(0).into_owned();
let ang_impulse = self.basis2 * impulse.fixed_rows::<U2>(3).into_owned();
mj_lambda2.linear -= lin_impulse * self.im2;
mj_lambda2.angular -= self
.ii2_sqrt
.transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
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]) {
for ii in 0..SIMD_WIDTH {
let joint = &mut joints_all[self.joint_id[ii]].weight;
if let JointParams::RevoluteJoint(rev) = &mut joint.params {
rev.impulse = self.impulse.extract(ii)
}
}
}
}