1use phyz_math::{DMat, SpatialMat, SpatialTransform};
4use phyz_model::{Model, State};
5
6pub 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 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 let mut i_c: Vec<SpatialMat> = model.bodies.iter().map(|b| b.inertia.to_matrix()).collect();
32
33 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 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 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; 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 let s_i = joint_i.motion_subspace_matrix(); let f_i_mat = &i_c[i].data * &s_i; 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 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}