Skip to main content

phyz_rigid/
crba.rs

1//! Composite Rigid Body Algorithm (CRBA) — mass matrix computation.
2
3use phyz_math::{DMat, SpatialMat, SpatialTransform};
4use phyz_model::{Model, State};
5
6/// Compute the joint-space mass matrix M(q) using CRBA.
7///
8/// Returns an nv × nv symmetric positive-definite matrix.
9pub fn crba(model: &Model, state: &State) -> DMat {
10    let nb = model.nbodies();
11    let mut mass_matrix = DMat::zeros(model.nv, model.nv);
12
13    // Compute tree transforms
14    let mut x_tree = vec![SpatialTransform::identity(); nb];
15    #[allow(clippy::needless_range_loop)]
16    for i in 0..nb {
17        let body = &model.bodies[i];
18        let joint = &model.joints[body.joint_idx];
19        let q_idx = model.q_offsets[body.joint_idx];
20        let ndof = joint.ndof();
21        let x_joint = if ndof == 0 {
22            SpatialTransform::identity()
23        } else {
24            let q_slice = &state.q.as_slice()[q_idx..q_idx + ndof];
25            joint.joint_transform_slice(q_slice)
26        };
27        x_tree[i] = x_joint.compose(&joint.parent_to_joint);
28    }
29
30    // Composite inertias (initialized from body inertias)
31    let mut i_c: Vec<SpatialMat> = model.bodies.iter().map(|b| b.inertia.to_matrix()).collect();
32
33    // Backward pass: accumulate composite inertias
34    for i in (0..nb).rev() {
35        let body = &model.bodies[i];
36        if body.parent >= 0 {
37            let pi = body.parent as usize;
38            let x_mot = x_tree[i].to_motion_matrix();
39            let ic_in_parent = SpatialMat::from_mat6(x_mot.transpose() * i_c[i].data * x_mot);
40            i_c[pi] = i_c[pi] + ic_in_parent;
41        }
42    }
43
44    // Compute mass matrix entries
45    for i in 0..nb {
46        let joint_i = &model.joints[model.bodies[i].joint_idx];
47        let v_i = model.v_offsets[model.bodies[i].joint_idx];
48        let ndof_i = joint_i.ndof();
49
50        if ndof_i == 0 {
51            continue;
52        }
53
54        if ndof_i == 1 {
55            let s_i = joint_i.motion_subspace();
56            let f_i = i_c[i].mul_vec(&s_i);
57            mass_matrix[(v_i, v_i)] = s_i.dot(&f_i);
58
59            // Off-diagonal: walk up the tree
60            let mut f = x_tree[i].inv_apply_force(&f_i);
61            let mut j = model.bodies[i].parent;
62            while j >= 0 {
63                let ju = j as usize;
64                let joint_j = &model.joints[model.bodies[ju].joint_idx];
65                let v_j = model.v_offsets[model.bodies[ju].joint_idx];
66                let ndof_j = joint_j.ndof();
67
68                if ndof_j == 1 {
69                    let s_j = joint_j.motion_subspace();
70                    mass_matrix[(v_i, v_j)] = s_j.dot(&f);
71                    mass_matrix[(v_j, v_i)] = mass_matrix[(v_i, v_j)];
72                } else if ndof_j > 1 {
73                    let s_j = joint_j.motion_subspace_matrix();
74                    let f_vec = nalgebra::DVector::from_column_slice(f.data.as_slice());
75                    let block = s_j.transpose() * &f_vec; // ndof_j x 1
76                    for kj in 0..ndof_j {
77                        mass_matrix[(v_i, v_j + kj)] = block[kj];
78                        mass_matrix[(v_j + kj, v_i)] = block[kj];
79                    }
80                }
81
82                f = x_tree[ju].inv_apply_force(&f);
83                j = model.bodies[ju].parent;
84            }
85        } else {
86            // Multi-DOF joint
87            let s_i = joint_i.motion_subspace_matrix(); // 6 x ndof_i
88            // F_i = I_c * S_i  (6 x ndof_i)
89            let f_i_mat = &i_c[i].data * &s_i; // 6 x ndof_i
90
91            // Diagonal block: S_i^T * I_c * S_i  (ndof_i x ndof_i)
92            let diag = s_i.transpose() * &f_i_mat;
93            for ki in 0..ndof_i {
94                for kj in 0..ndof_i {
95                    mass_matrix[(v_i + ki, v_i + kj)] = diag[(ki, kj)];
96                }
97            }
98
99            // Off-diagonal: walk up tree, one column at a time
100            for col in 0..ndof_i {
101                let f_col = phyz_math::SpatialVec {
102                    data: phyz_math::Vec6::new(
103                        f_i_mat[(0, col)],
104                        f_i_mat[(1, col)],
105                        f_i_mat[(2, col)],
106                        f_i_mat[(3, col)],
107                        f_i_mat[(4, col)],
108                        f_i_mat[(5, col)],
109                    ),
110                };
111                let mut f = x_tree[i].inv_apply_force(&f_col);
112                let mut j = model.bodies[i].parent;
113                while j >= 0 {
114                    let ju = j as usize;
115                    let joint_j = &model.joints[model.bodies[ju].joint_idx];
116                    let v_j = model.v_offsets[model.bodies[ju].joint_idx];
117                    let ndof_j = joint_j.ndof();
118
119                    if ndof_j == 1 {
120                        let s_j = joint_j.motion_subspace();
121                        let val = s_j.dot(&f);
122                        mass_matrix[(v_i + col, v_j)] = val;
123                        mass_matrix[(v_j, v_i + col)] = val;
124                    } else if ndof_j > 1 {
125                        let s_j = joint_j.motion_subspace_matrix();
126                        let f_vec = nalgebra::DVector::from_column_slice(f.data.as_slice());
127                        let block = s_j.transpose() * &f_vec;
128                        for kj in 0..ndof_j {
129                            mass_matrix[(v_i + col, v_j + kj)] = block[kj];
130                            mass_matrix[(v_j + kj, v_i + col)] = block[kj];
131                        }
132                    }
133
134                    f = x_tree[ju].inv_apply_force(&f);
135                    j = model.bodies[ju].parent;
136                }
137            }
138        }
139    }
140
141    mass_matrix
142}