rapier3d/dynamics/solver/joint_constraint/
joint_constraint_builder.rs

1use crate::dynamics::solver::ConstraintsCounts;
2use crate::dynamics::solver::MotorParameters;
3use crate::dynamics::solver::joint_constraint::JointSolverBody;
4use crate::dynamics::solver::joint_constraint::joint_velocity_constraint::{
5    JointConstraint, WritebackId,
6};
7use crate::dynamics::solver::solver_body::SolverBodies;
8use crate::dynamics::{GenericJoint, ImpulseJoint, IntegrationParameters, JointIndex};
9use crate::math::{DIM, Real};
10use crate::prelude::RigidBodySet;
11use crate::utils;
12#[cfg(feature = "dim3")]
13use crate::utils::OrthonormalBasis;
14use crate::utils::{
15    AngularInertiaOps, ComponentMul, CrossProductMatrix, DotProduct, IndexMut2, MatrixColumn,
16    PoseOps, RotationOps, ScalarType, SimdLength,
17};
18
19#[cfg(feature = "dim2")]
20use crate::num::One;
21
22#[cfg(feature = "dim3")]
23use parry::math::Rot3;
24#[cfg(feature = "simd-is-enabled")]
25use {
26    crate::dynamics::SpringCoefficients,
27    crate::math::{SIMD_WIDTH, SimdPose, SimdReal},
28};
29
30pub struct JointConstraintBuilder {
31    body1: u32,
32    body2: u32,
33    joint_id: JointIndex,
34    joint: GenericJoint,
35    constraint_id: usize,
36}
37
38impl JointConstraintBuilder {
39    pub fn generate(
40        joint: &ImpulseJoint,
41        bodies: &RigidBodySet,
42        joint_id: JointIndex,
43        out_builder: &mut Self,
44        out_constraint_id: &mut usize,
45    ) {
46        let rb1 = &bodies[joint.body1];
47        let rb2 = &bodies[joint.body2];
48        let solver_body1 = rb1.effective_active_set_offset();
49        let solver_body2 = rb2.effective_active_set_offset();
50
51        *out_builder = Self {
52            body1: solver_body1,
53            body2: solver_body2,
54            joint_id,
55            joint: joint.data,
56            constraint_id: *out_constraint_id,
57        };
58        // Since solver body poses are given in center-of-mass space,
59        // we need to transform the anchors to that space.
60        out_builder.joint.transform_to_solver_body_space(rb1, rb2);
61
62        let count = ConstraintsCounts::from_joint(joint);
63        *out_constraint_id += count.num_constraints;
64    }
65
66    pub fn update(
67        &self,
68        params: &IntegrationParameters,
69        bodies: &SolverBodies,
70        out: &mut [JointConstraint<Real, 1>],
71    ) {
72        // NOTE: right now, the "update", is basically reconstructing all the
73        //       constraints. Could we make this more incremental?
74
75        let rb1 = bodies.get_pose(self.body1);
76        let rb2 = bodies.get_pose(self.body2);
77        let frame1 = rb1.pose() * self.joint.local_frame1;
78        let frame2 = rb2.pose() * self.joint.local_frame2;
79        let world_com1 = rb1.translation;
80        let world_com2 = rb2.translation;
81
82        let joint_body1 = JointSolverBody {
83            im: rb1.im,
84            ii: rb1.ii,
85            world_com: world_com1,
86            solver_vel: [self.body1],
87        };
88        let joint_body2 = JointSolverBody {
89            im: rb2.im,
90            ii: rb2.ii,
91            world_com: world_com2,
92            solver_vel: [self.body2],
93        };
94
95        JointConstraint::<Real, 1>::update(
96            params,
97            self.joint_id,
98            &joint_body1,
99            &joint_body2,
100            &frame1,
101            &frame2,
102            &self.joint,
103            &mut out[self.constraint_id..],
104        );
105    }
106}
107
108#[cfg(feature = "simd-is-enabled")]
109pub struct JointConstraintBuilderSimd {
110    body1: [u32; SIMD_WIDTH],
111    body2: [u32; SIMD_WIDTH],
112    joint_id: [JointIndex; SIMD_WIDTH],
113    local_frame1: SimdPose<SimdReal>,
114    local_frame2: SimdPose<SimdReal>,
115    locked_axes: u8,
116    softness: SpringCoefficients<SimdReal>,
117    constraint_id: usize,
118}
119
120#[cfg(feature = "simd-is-enabled")]
121impl JointConstraintBuilderSimd {
122    pub fn generate(
123        joint: [&ImpulseJoint; SIMD_WIDTH],
124        bodies: &RigidBodySet,
125        joint_id: [JointIndex; SIMD_WIDTH],
126        out_builder: &mut Self,
127        out_constraint_id: &mut usize,
128    ) {
129        let rb1 = array![|ii| &bodies[joint[ii].body1]];
130        let rb2 = array![|ii| &bodies[joint[ii].body2]];
131
132        let body1 = array![|ii| if rb1[ii].is_dynamic_or_kinematic() {
133            rb1[ii].ids.active_set_id as u32
134        } else {
135            u32::MAX
136        }];
137        let body2 = array![|ii| if rb2[ii].is_dynamic_or_kinematic() {
138            rb2[ii].ids.active_set_id as u32
139        } else {
140            u32::MAX
141        }];
142
143        let local_frame1 = array![|ii| if body1[ii] != u32::MAX {
144            (joint[ii].data.local_frame1).into()
145        } else {
146            (rb1[ii].pos.position * joint[ii].data.local_frame1).into()
147        }]
148        .into();
149        let local_frame2 = array![|ii| if body2[ii] != u32::MAX {
150            (joint[ii].data.local_frame2).into()
151        } else {
152            (rb2[ii].pos.position * joint[ii].data.local_frame2).into()
153        }]
154        .into();
155
156        *out_builder = Self {
157            body1,
158            body2,
159            joint_id,
160            local_frame1,
161            local_frame2,
162            locked_axes: joint[0].data.locked_axes.bits(),
163            softness: SpringCoefficients {
164                natural_frequency: array![|ii| joint[ii].data.softness.natural_frequency].into(),
165                damping_ratio: array![|ii| joint[ii].data.softness.damping_ratio].into(),
166            },
167            constraint_id: *out_constraint_id,
168        };
169
170        let count = ConstraintsCounts::from_joint(joint[0]);
171        *out_constraint_id += count.num_constraints;
172    }
173
174    pub fn update(
175        &mut self,
176        params: &IntegrationParameters,
177        bodies: &SolverBodies,
178        out: &mut [JointConstraint<SimdReal, SIMD_WIDTH>],
179    ) {
180        // NOTE: right now, the "update", is basically reconstructing all the
181        //       constraints. Could we make this more incremental?
182
183        let rb1 = bodies.gather_poses(self.body1);
184        let rb2 = bodies.gather_poses(self.body2);
185        let frame1 = rb1.pose() * self.local_frame1;
186        let frame2 = rb2.pose() * self.local_frame2;
187
188        let joint_body1 = JointSolverBody {
189            im: rb1.im,
190            ii: rb1.ii,
191            world_com: rb1.translation,
192            solver_vel: self.body1,
193        };
194        let joint_body2 = JointSolverBody {
195            im: rb2.im,
196            ii: rb2.ii,
197            world_com: rb2.translation,
198            solver_vel: self.body2,
199        };
200
201        JointConstraint::<SimdReal, SIMD_WIDTH>::update(
202            params,
203            self.joint_id,
204            &joint_body1,
205            &joint_body2,
206            &frame1,
207            &frame2,
208            self.locked_axes,
209            self.softness,
210            &mut out[self.constraint_id..],
211        );
212    }
213}
214
215#[derive(Debug, Copy, Clone)]
216pub struct JointConstraintHelper<N: ScalarType> {
217    pub basis: N::Matrix,
218    #[cfg(feature = "dim3")]
219    pub basis2: N::Matrix, // TODO: used for angular coupling. Can we avoid storing this?
220    #[cfg(feature = "dim3")]
221    pub cmat1_basis: N::Matrix,
222    #[cfg(feature = "dim3")]
223    pub cmat2_basis: N::Matrix,
224    #[cfg(feature = "dim3")]
225    pub ang_basis: N::Matrix,
226    #[cfg(feature = "dim2")]
227    pub cmat1_basis: [N::AngVector; 2],
228    #[cfg(feature = "dim2")]
229    pub cmat2_basis: [N::AngVector; 2],
230    pub lin_err: N::Vector,
231    pub ang_err: N::Rotation,
232}
233
234impl<N: ScalarType> JointConstraintHelper<N> {
235    pub fn new(
236        frame1: &N::Pose,
237        frame2: &N::Pose,
238        world_com1: &N::Vector,
239        world_com2: &N::Vector,
240        locked_lin_axes: u8,
241    ) -> Self {
242        let mut frame1 = *frame1;
243        let basis = frame1.rotation().to_mat();
244        let lin_err = frame2.translation() - frame1.translation();
245
246        // Adjust the point of application of the force for the first body,
247        // by snapping free axes to the second frame's center (to account for
248        // the allowed relative movement).
249        {
250            let mut new_center1 = frame2.translation(); // First, assume all dofs are free.
251
252            // Then snap the locked ones.
253            for i in 0..DIM {
254                if locked_lin_axes & (1 << i) != 0 {
255                    let axis = basis.column(i);
256                    new_center1 -= axis * lin_err.gdot(axis);
257                }
258            }
259            frame1.set_translation(new_center1);
260        }
261
262        let r1 = frame1.translation() - *world_com1;
263        let r2 = frame2.translation() - *world_com2;
264
265        let cmat1 = r1.gcross_matrix();
266        let cmat2 = r2.gcross_matrix();
267
268        #[cfg(feature = "dim3")]
269        let mut ang_basis = frame1.rotation().diff_conj1_2_tr(&frame2.rotation());
270        #[allow(unused_mut)] // The mut is needed for 3D
271        let mut ang_err = frame1.rotation().inverse() * frame2.rotation();
272
273        #[cfg(feature = "dim3")]
274        {
275            let sgn = N::one().simd_copysign(frame1.rotation().dot(&frame2.rotation()));
276            ang_basis *= sgn;
277            ang_err.mul_assign_unchecked(sgn);
278        }
279
280        #[cfg(feature = "dim2")]
281        return Self {
282            basis,
283            cmat1_basis: [
284                cmat1.gdot(basis.column(0)).into(),
285                cmat1.gdot(basis.column(1)).into(),
286            ],
287            cmat2_basis: [
288                cmat2.gdot(basis.column(0)).into(),
289                cmat2.gdot(basis.column(1)).into(),
290            ],
291            lin_err,
292            ang_err,
293        };
294        #[cfg(feature = "dim3")]
295        return Self {
296            basis,
297            basis2: frame2.rotation().to_mat(),
298            cmat1_basis: cmat1 * basis,
299            cmat2_basis: cmat2 * basis,
300            ang_basis,
301            lin_err,
302            ang_err,
303        };
304    }
305
306    pub fn limit_linear<const LANES: usize>(
307        &self,
308        params: &IntegrationParameters,
309        joint_id: [JointIndex; LANES],
310        body1: &JointSolverBody<N, LANES>,
311        body2: &JointSolverBody<N, LANES>,
312        limited_axis: usize,
313        limits: [N; 2],
314        writeback_id: WritebackId,
315        erp_inv_dt: N,
316        cfm_coeff: N,
317    ) -> JointConstraint<N, LANES> {
318        let zero = N::zero();
319        let mut constraint = self.lock_linear(
320            params,
321            joint_id,
322            body1,
323            body2,
324            limited_axis,
325            writeback_id,
326            erp_inv_dt,
327            cfm_coeff,
328        );
329
330        let dist = self.lin_err.gdot(constraint.lin_jac);
331        let min_enabled = dist.simd_le(limits[0]);
332        let max_enabled = limits[1].simd_le(dist);
333
334        let rhs_bias =
335            ((dist - limits[1]).simd_max(zero) - (limits[0] - dist).simd_max(zero)) * erp_inv_dt;
336        constraint.rhs = constraint.rhs_wo_bias + rhs_bias;
337        constraint.cfm_coeff = cfm_coeff;
338        constraint.impulse_bounds = [
339            N::splat(-Real::INFINITY).select(min_enabled, zero),
340            N::splat(Real::INFINITY).select(max_enabled, zero),
341        ];
342
343        constraint
344    }
345
346    pub fn limit_linear_coupled<const LANES: usize>(
347        &self,
348        params: &IntegrationParameters,
349        joint_id: [JointIndex; LANES],
350        body1: &JointSolverBody<N, LANES>,
351        body2: &JointSolverBody<N, LANES>,
352        coupled_axes: u8,
353        limits: [N; 2],
354        writeback_id: WritebackId,
355        erp_inv_dt: N,
356        cfm_coeff: N,
357    ) -> JointConstraint<N, LANES> {
358        let zero = N::zero();
359        let mut lin_jac: N::Vector = Default::default();
360        let mut ang_jac1: N::AngVector = Default::default();
361        let mut ang_jac2: N::AngVector = Default::default();
362
363        for i in 0..DIM {
364            if coupled_axes & (1 << i) != 0 {
365                let coeff = self.basis.column(i).gdot(self.lin_err);
366                lin_jac += self.basis.column(i) * coeff;
367                #[cfg(feature = "dim2")]
368                {
369                    ang_jac1 += self.cmat1_basis[i] * coeff;
370                    ang_jac2 += self.cmat2_basis[i] * coeff;
371                }
372                #[cfg(feature = "dim3")]
373                {
374                    ang_jac1 += self.cmat1_basis.column(i).into() * coeff;
375                    ang_jac2 += self.cmat2_basis.column(i).into() * coeff;
376                }
377            }
378        }
379
380        // FIXME: handle min limit too.
381
382        let dist = lin_jac.simd_length();
383        let inv_dist = crate::utils::simd_inv(dist);
384        lin_jac *= inv_dist;
385        ang_jac1 *= inv_dist;
386        ang_jac2 *= inv_dist;
387
388        let rhs_wo_bias = (dist - limits[1]).simd_min(zero) * N::splat(params.inv_dt());
389
390        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac1);
391        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac2);
392
393        let rhs_bias = (dist - limits[1]).simd_max(zero) * erp_inv_dt;
394        let rhs = rhs_wo_bias + rhs_bias;
395        let impulse_bounds = [N::zero(), N::splat(Real::INFINITY)];
396
397        JointConstraint {
398            joint_id,
399            solver_vel1: body1.solver_vel,
400            solver_vel2: body2.solver_vel,
401            im1: body1.im,
402            im2: body2.im,
403            impulse: N::zero(),
404            impulse_bounds,
405            lin_jac,
406            ang_jac1,
407            ang_jac2,
408            ii_ang_jac1,
409            ii_ang_jac2,
410            inv_lhs: N::zero(), // Will be set during orthogonalization.
411            cfm_coeff,
412            cfm_gain: N::zero(),
413            rhs,
414            rhs_wo_bias,
415            writeback_id,
416        }
417    }
418
419    pub fn motor_linear<const LANES: usize>(
420        &self,
421        params: &IntegrationParameters,
422        joint_id: [JointIndex; LANES],
423        body1: &JointSolverBody<N, LANES>,
424        body2: &JointSolverBody<N, LANES>,
425        motor_axis: usize,
426        motor_params: &MotorParameters<N>,
427        limits: Option<[N; 2]>,
428        writeback_id: WritebackId,
429    ) -> JointConstraint<N, LANES> {
430        let inv_dt = N::splat(params.inv_dt());
431        let mut constraint = self.lock_linear(
432            params,
433            joint_id,
434            body1,
435            body2,
436            motor_axis,
437            writeback_id,
438            // Set regularization factors to zero.
439            // The motor impl. will overwrite them after.
440            N::zero(),
441            N::zero(),
442        );
443
444        let mut rhs_wo_bias = N::zero();
445        if motor_params.erp_inv_dt != N::zero() {
446            let dist = self.lin_err.gdot(constraint.lin_jac);
447            rhs_wo_bias += (dist - motor_params.target_pos) * motor_params.erp_inv_dt;
448        }
449
450        let mut target_vel = motor_params.target_vel;
451        if let Some(limits) = limits {
452            let dist = self.lin_err.gdot(constraint.lin_jac);
453            target_vel =
454                target_vel.simd_clamp((limits[0] - dist) * inv_dt, (limits[1] - dist) * inv_dt);
455        };
456
457        rhs_wo_bias += -target_vel;
458
459        constraint.cfm_coeff = motor_params.cfm_coeff;
460        constraint.cfm_gain = motor_params.cfm_gain;
461        constraint.impulse_bounds = [-motor_params.max_impulse, motor_params.max_impulse];
462        constraint.rhs = rhs_wo_bias;
463        constraint.rhs_wo_bias = rhs_wo_bias;
464        constraint
465    }
466
467    pub fn motor_linear_coupled<const LANES: usize>(
468        &self,
469        params: &IntegrationParameters,
470        joint_id: [JointIndex; LANES],
471        body1: &JointSolverBody<N, LANES>,
472        body2: &JointSolverBody<N, LANES>,
473        coupled_axes: u8,
474        motor_params: &MotorParameters<N>,
475        limits: Option<[N; 2]>,
476        writeback_id: WritebackId,
477    ) -> JointConstraint<N, LANES> {
478        let inv_dt = N::splat(params.inv_dt());
479
480        let mut lin_jac: N::Vector = Default::default();
481        let mut ang_jac1: N::AngVector = Default::default();
482        let mut ang_jac2: N::AngVector = Default::default();
483
484        for i in 0..DIM {
485            if coupled_axes & (1 << i) != 0 {
486                let coeff = self.basis.column(i).gdot(self.lin_err);
487                lin_jac += self.basis.column(i) * coeff;
488                #[cfg(feature = "dim2")]
489                {
490                    ang_jac1 += self.cmat1_basis[i] * coeff;
491                    ang_jac2 += self.cmat2_basis[i] * coeff;
492                }
493                #[cfg(feature = "dim3")]
494                {
495                    ang_jac1 += self.cmat1_basis.column(i).into() * coeff;
496                    ang_jac2 += self.cmat2_basis.column(i).into() * coeff;
497                }
498            }
499        }
500
501        let dist = lin_jac.simd_length();
502        let inv_dist = crate::utils::simd_inv(dist);
503        lin_jac *= inv_dist;
504        ang_jac1 *= inv_dist;
505        ang_jac2 *= inv_dist;
506
507        let mut rhs_wo_bias = N::zero();
508        if motor_params.erp_inv_dt != N::zero() {
509            rhs_wo_bias += (dist - motor_params.target_pos) * motor_params.erp_inv_dt;
510        }
511
512        let mut target_vel = motor_params.target_vel;
513        if let Some(limits) = limits {
514            target_vel =
515                target_vel.simd_clamp((limits[0] - dist) * inv_dt, (limits[1] - dist) * inv_dt);
516        };
517
518        rhs_wo_bias += -target_vel;
519
520        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac1);
521        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac2);
522
523        JointConstraint {
524            joint_id,
525            solver_vel1: body1.solver_vel,
526            solver_vel2: body2.solver_vel,
527            im1: body1.im,
528            im2: body2.im,
529            impulse: N::zero(),
530            impulse_bounds: [-motor_params.max_impulse, motor_params.max_impulse],
531            lin_jac,
532            ang_jac1,
533            ang_jac2,
534            ii_ang_jac1,
535            ii_ang_jac2,
536            inv_lhs: N::zero(), // Will be set during orthogonalization.
537            cfm_coeff: motor_params.cfm_coeff,
538            cfm_gain: motor_params.cfm_gain,
539            rhs: rhs_wo_bias,
540            rhs_wo_bias,
541            writeback_id,
542        }
543    }
544
545    pub fn lock_linear<const LANES: usize>(
546        &self,
547        _params: &IntegrationParameters,
548        joint_id: [JointIndex; LANES],
549        body1: &JointSolverBody<N, LANES>,
550        body2: &JointSolverBody<N, LANES>,
551        locked_axis: usize,
552        writeback_id: WritebackId,
553        erp_inv_dt: N,
554        cfm_coeff: N,
555    ) -> JointConstraint<N, LANES> {
556        let lin_jac = self.basis.column(locked_axis);
557        #[cfg(feature = "dim2")]
558        let ang_jac1 = self.cmat1_basis[locked_axis];
559        #[cfg(feature = "dim2")]
560        let ang_jac2 = self.cmat2_basis[locked_axis];
561        #[cfg(feature = "dim3")]
562        let ang_jac1 = self.cmat1_basis.column(locked_axis).into();
563        #[cfg(feature = "dim3")]
564        let ang_jac2 = self.cmat2_basis.column(locked_axis).into();
565
566        let rhs_wo_bias = N::zero();
567        let rhs_bias = lin_jac.gdot(self.lin_err) * erp_inv_dt;
568
569        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac1);
570        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac2);
571
572        JointConstraint {
573            joint_id,
574            solver_vel1: body1.solver_vel,
575            solver_vel2: body2.solver_vel,
576            im1: body1.im,
577            im2: body2.im,
578            impulse: N::zero(),
579            impulse_bounds: [-N::splat(Real::MAX), N::splat(Real::MAX)],
580            lin_jac,
581            ang_jac1,
582            ang_jac2,
583            ii_ang_jac1,
584            ii_ang_jac2,
585            inv_lhs: N::zero(), // Will be set during orthogonalization.
586            cfm_coeff,
587            cfm_gain: N::zero(),
588            rhs: rhs_wo_bias + rhs_bias,
589            rhs_wo_bias,
590            writeback_id,
591        }
592    }
593
594    pub fn limit_angular<const LANES: usize>(
595        &self,
596        _params: &IntegrationParameters,
597        joint_id: [JointIndex; LANES],
598        body1: &JointSolverBody<N, LANES>,
599        body2: &JointSolverBody<N, LANES>,
600        _limited_axis: usize,
601        limits: [N; 2],
602        writeback_id: WritebackId,
603        erp_inv_dt: N,
604        cfm_coeff: N,
605    ) -> JointConstraint<N, LANES> {
606        let zero = N::zero();
607        let half = N::splat(0.5);
608        let s_limits = [(limits[0] * half).simd_sin(), (limits[1] * half).simd_sin()];
609        #[cfg(feature = "dim2")]
610        let s_ang = (self.ang_err.angle() * half).simd_sin();
611        #[cfg(feature = "dim3")]
612        let s_ang = self.ang_err.imag()[_limited_axis];
613        let min_enabled = s_ang.simd_le(s_limits[0]);
614        let max_enabled = s_limits[1].simd_le(s_ang);
615
616        let impulse_bounds = [
617            N::splat(-Real::INFINITY).select(min_enabled, zero),
618            N::splat(Real::INFINITY).select(max_enabled, zero),
619        ];
620
621        #[cfg(feature = "dim2")]
622        let ang_jac = N::AngVector::one();
623        #[cfg(feature = "dim3")]
624        let ang_jac = self.ang_basis.column(_limited_axis).into();
625        let rhs_wo_bias = N::zero();
626        let rhs_bias = ((s_ang - s_limits[1]).simd_max(zero)
627            - (s_limits[0] - s_ang).simd_max(zero))
628            * erp_inv_dt;
629
630        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
631        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
632
633        JointConstraint {
634            joint_id,
635            solver_vel1: body1.solver_vel,
636            solver_vel2: body2.solver_vel,
637            im1: body1.im,
638            im2: body2.im,
639            impulse: N::zero(),
640            impulse_bounds,
641            lin_jac: Default::default(),
642            ang_jac1: ang_jac,
643            ang_jac2: ang_jac,
644            ii_ang_jac1,
645            ii_ang_jac2,
646            inv_lhs: N::zero(), // Will be set during orthogonalization.
647            cfm_coeff,
648            cfm_gain: N::zero(),
649            rhs: rhs_wo_bias + rhs_bias,
650            rhs_wo_bias,
651            writeback_id,
652        }
653    }
654
655    pub fn motor_angular<const LANES: usize>(
656        &self,
657        joint_id: [JointIndex; LANES],
658        body1: &JointSolverBody<N, LANES>,
659        body2: &JointSolverBody<N, LANES>,
660        _motor_axis: usize,
661        motor_params: &MotorParameters<N>,
662        writeback_id: WritebackId,
663    ) -> JointConstraint<N, LANES> {
664        #[cfg(feature = "dim2")]
665        let ang_jac = N::AngVector::one();
666        #[cfg(feature = "dim3")]
667        let ang_jac = self.basis.column(_motor_axis).into();
668
669        let mut rhs_wo_bias = N::zero();
670        if motor_params.erp_inv_dt != N::zero() {
671            let ang_dist;
672
673            #[cfg(feature = "dim2")]
674            {
675                ang_dist = self.ang_err.angle();
676            }
677
678            #[cfg(feature = "dim3")]
679            {
680                // Clamp the component from -1.0 to 1.0 to account for slight imprecision
681                let clamped_err = self.ang_err.imag()[_motor_axis].simd_clamp(-N::one(), N::one());
682                ang_dist = clamped_err.simd_asin() * N::splat(2.0);
683            }
684
685            let target_ang = motor_params.target_pos;
686            rhs_wo_bias += utils::smallest_abs_diff_between_angles(ang_dist, target_ang)
687                * motor_params.erp_inv_dt;
688        }
689
690        rhs_wo_bias += -motor_params.target_vel;
691
692        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
693        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
694
695        JointConstraint {
696            joint_id,
697            solver_vel1: body1.solver_vel,
698            solver_vel2: body2.solver_vel,
699            im1: body1.im,
700            im2: body2.im,
701            impulse: N::zero(),
702            impulse_bounds: [-motor_params.max_impulse, motor_params.max_impulse],
703            lin_jac: Default::default(),
704            ang_jac1: ang_jac,
705            ang_jac2: ang_jac,
706            ii_ang_jac1,
707            ii_ang_jac2,
708            inv_lhs: N::zero(), // Will be set during orthogonalization.
709            cfm_coeff: motor_params.cfm_coeff,
710            cfm_gain: motor_params.cfm_gain,
711            rhs: rhs_wo_bias,
712            rhs_wo_bias,
713            writeback_id,
714        }
715    }
716
717    pub fn lock_angular<const LANES: usize>(
718        &self,
719        _params: &IntegrationParameters,
720        joint_id: [JointIndex; LANES],
721        body1: &JointSolverBody<N, LANES>,
722        body2: &JointSolverBody<N, LANES>,
723        _locked_axis: usize,
724        writeback_id: WritebackId,
725        erp_inv_dt: N,
726        cfm_coeff: N,
727    ) -> JointConstraint<N, LANES> {
728        #[cfg(feature = "dim2")]
729        let ang_jac = N::AngVector::one();
730        #[cfg(feature = "dim3")]
731        let ang_jac = self.ang_basis.column(_locked_axis).into();
732
733        let rhs_wo_bias = N::zero();
734        #[cfg(feature = "dim2")]
735        let rhs_bias = self.ang_err.imag() * erp_inv_dt;
736        #[cfg(feature = "dim3")]
737        let rhs_bias = self.ang_err.imag()[_locked_axis] * erp_inv_dt;
738
739        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
740        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
741
742        JointConstraint {
743            joint_id,
744            solver_vel1: body1.solver_vel,
745            solver_vel2: body2.solver_vel,
746            im1: body1.im,
747            im2: body2.im,
748            impulse: N::zero(),
749            impulse_bounds: [-N::splat(Real::MAX), N::splat(Real::MAX)],
750            lin_jac: Default::default(),
751            ang_jac1: ang_jac,
752            ang_jac2: ang_jac,
753            ii_ang_jac1,
754            ii_ang_jac2,
755            inv_lhs: N::zero(), // Will be set during orthogonalization.
756            cfm_coeff,
757            cfm_gain: N::zero(),
758            rhs: rhs_wo_bias + rhs_bias,
759            rhs_wo_bias,
760            writeback_id,
761        }
762    }
763
764    /// Orthogonalize the constraints and set their inv_lhs field.
765    pub fn finalize_constraints<const LANES: usize>(constraints: &mut [JointConstraint<N, LANES>]) {
766        let len = constraints.len();
767
768        if len == 0 {
769            return;
770        }
771
772        let imsum = constraints[0].im1 + constraints[0].im2;
773
774        // Use the modified Gram-Schmidt orthogonalization.
775        for j in 0..len {
776            let c_j = &mut constraints[j];
777            let dot_jj = c_j.lin_jac.gdot(imsum.component_mul(&c_j.lin_jac))
778                + c_j.ii_ang_jac1.gdot(c_j.ang_jac1)
779                + c_j.ii_ang_jac2.gdot(c_j.ang_jac2);
780            let cfm_gain = dot_jj * c_j.cfm_coeff + c_j.cfm_gain;
781            let inv_dot_jj = crate::utils::simd_inv(dot_jj);
782            c_j.inv_lhs = crate::utils::simd_inv(dot_jj + cfm_gain); // Don’t forget to update the inv_lhs.
783            c_j.cfm_gain = cfm_gain;
784
785            if c_j.impulse_bounds != [-N::splat(Real::MAX), N::splat(Real::MAX)] {
786                // Don't remove constraints with limited forces from the others
787                // because they may not deliver the necessary forces to fulfill
788                // the removed parts of other constraints.
789                continue;
790            }
791
792            for i in (j + 1)..len {
793                let (c_i, c_j) = constraints.index_mut_const(i, j);
794
795                let dot_ij = c_i.lin_jac.gdot(imsum.component_mul(&c_j.lin_jac))
796                    + c_i.ii_ang_jac1.gdot(c_j.ang_jac1)
797                    + c_i.ii_ang_jac2.gdot(c_j.ang_jac2);
798                let coeff = dot_ij * inv_dot_jj;
799
800                c_i.lin_jac -= c_j.lin_jac * coeff;
801                c_i.ang_jac1 -= c_j.ang_jac1 * coeff;
802                c_i.ang_jac2 -= c_j.ang_jac2 * coeff;
803                c_i.ii_ang_jac1 -= c_j.ii_ang_jac1 * coeff;
804                c_i.ii_ang_jac2 -= c_j.ii_ang_jac2 * coeff;
805                c_i.rhs_wo_bias -= c_j.rhs_wo_bias * coeff;
806                c_i.rhs -= c_j.rhs * coeff;
807            }
808        }
809    }
810}
811
812impl JointConstraintHelper<Real> {
813    #[cfg(feature = "dim3")]
814    pub fn limit_angular_coupled(
815        &self,
816        _params: &IntegrationParameters,
817        joint_id: [JointIndex; 1],
818        body1: &JointSolverBody<Real, 1>,
819        body2: &JointSolverBody<Real, 1>,
820        coupled_axes: u8,
821        limits: [Real; 2],
822        writeback_id: WritebackId,
823        erp_inv_dt: Real,
824        cfm_coeff: Real,
825    ) -> JointConstraint<Real, 1> {
826        // NOTE: right now, this only supports exactly 2 coupled axes.
827        let ang_coupled_axes = coupled_axes >> DIM;
828        assert_eq!(ang_coupled_axes.count_ones(), 2);
829        let not_coupled_index = ang_coupled_axes.trailing_ones() as usize;
830        let axis1 = self.basis.column(not_coupled_index);
831        let axis2 = self.basis2.column(not_coupled_index);
832
833        let rot = Rot3::from_rotation_arc(axis1, axis2);
834        let (mut ang_jac, angle) = rot.to_axis_angle();
835
836        if angle == 0.0 {
837            ang_jac = axis1.orthonormal_basis()[0];
838        }
839
840        let min_enabled = angle <= limits[0];
841        let max_enabled = limits[1] <= angle;
842
843        let impulse_bounds = [
844            if min_enabled { -Real::INFINITY } else { 0.0 },
845            if max_enabled { Real::INFINITY } else { 0.0 },
846        ];
847
848        let rhs_wo_bias = 0.0;
849
850        let rhs_bias = ((angle - limits[1]).max(0.0) - (limits[0] - angle).max(0.0)) * erp_inv_dt;
851
852        let ii_ang_jac1 = body1.ii.transform_vector(ang_jac);
853        let ii_ang_jac2 = body2.ii.transform_vector(ang_jac);
854
855        JointConstraint {
856            joint_id,
857            solver_vel1: body1.solver_vel,
858            solver_vel2: body2.solver_vel,
859            im1: body1.im,
860            im2: body2.im,
861            impulse: 0.0,
862            impulse_bounds,
863            lin_jac: Default::default(),
864            ang_jac1: ang_jac,
865            ang_jac2: ang_jac,
866            ii_ang_jac1,
867            ii_ang_jac2,
868            inv_lhs: 0.0, // Will be set during orthogonalization.
869            cfm_coeff,
870            cfm_gain: 0.0,
871            rhs: rhs_wo_bias + rhs_bias,
872            rhs_wo_bias,
873            writeback_id,
874        }
875    }
876}