1use glam_traits_ext::{FloatAffine, FloatVec, TAffine3, TVec3};
2
3use crate::SRobotQ;
4
5use super::{AAffine3, AVec3, FKChain, FKScalar};
6
7#[derive(Debug, Clone)]
21pub struct PrismaticFK<const M: usize, const N: usize, F: FKScalar, FK: FKChain<N, F>> {
22 inner: FK,
23 axis: AVec3<F>,
24 q_index_first: bool,
25}
26
27impl<const M: usize, const N: usize, F: FKScalar, FK: FKChain<N, F>> PrismaticFK<M, N, F, FK> {
28 pub fn new(inner: FK, axis: AVec3<F>, q_index_first: bool) -> Self {
29 const { assert!(M == N + 1, "M must equal N + 1") };
30 Self {
31 inner,
32 axis,
33 q_index_first,
34 }
35 }
36
37 pub fn inner(&self) -> &FK {
38 &self.inner
39 }
40
41 pub fn axis(&self) -> AVec3<F> {
42 self.axis
43 }
44
45 pub fn q_index_first(&self) -> bool {
46 self.q_index_first
47 }
48
49 fn split_q(&self, q: &SRobotQ<M, F>) -> (F, SRobotQ<N, F>) {
50 let mut inner = [F::zero(); N];
51 if self.q_index_first {
52 inner.copy_from_slice(&q.0[1..M]);
53 (q.0[0], SRobotQ(inner))
54 } else {
55 inner.copy_from_slice(&q.0[..N]);
56 (q.0[M - 1], SRobotQ(inner))
57 }
58 }
59
60 fn prismatic_col(&self) -> usize {
61 if self.q_index_first { 0 } else { N }
62 }
63
64 fn revolute_offset(&self) -> usize {
65 if self.q_index_first { 1 } else { 0 }
66 }
67}
68
69impl<const M: usize, const N: usize, F: FKScalar, FK: FKChain<N, F>> FKChain<M, F>
70 for PrismaticFK<M, N, F, FK>
71{
72 type Error = FK::Error;
73
74 fn base_tf(&self) -> AAffine3<F> {
75 self.inner.base_tf()
76 }
77
78 fn fk(&self, q: &SRobotQ<M, F>) -> Result<[AAffine3<F>; M], Self::Error> {
79 let (q_p, inner_q) = self.split_q(q);
80 let offset = self.axis * q_p;
81 let inner_frames = self.inner.fk(&inner_q)?;
82 let mut out = [AAffine3::<F>::IDENTITY; M];
83
84 out[0] = AAffine3::<F>::from_translation(offset);
85 for i in 0..N {
86 let mut f = inner_frames[i];
87 f = AAffine3::<F>::from_mat3_translation(f.matrix3(), f.translation() + offset);
90 out[i + 1] = f;
91 }
92
93 Ok(out)
94 }
95
96 fn fk_end(&self, q: &SRobotQ<M, F>) -> Result<AAffine3<F>, Self::Error> {
97 let (q_p, inner_q) = self.split_q(q);
98 let end = self.inner.fk_end(&inner_q)?;
99 let offset = self.axis * q_p;
100 Ok(AAffine3::<F>::from_mat3_translation(
101 end.matrix3(),
102 end.translation() + offset,
103 ))
104 }
105
106 fn joint_axes_positions(
107 &self,
108 q: &SRobotQ<M, F>,
109 ) -> Result<([AVec3<F>; M], [AVec3<F>; M], AVec3<F>), Self::Error> {
110 let (q_p, inner_q) = self.split_q(q);
111 let offset = self.axis * q_p;
112 let (inner_axes, inner_pos, inner_p_ee) = self.inner.joint_axes_positions(&inner_q)?;
113
114 let mut axes = [AVec3::<F>::ZERO; M];
115 let mut positions = [AVec3::<F>::ZERO; M];
116
117 axes[0] = self.axis;
118 for i in 0..N {
119 axes[i + 1] = inner_axes[i];
120 positions[i + 1] = inner_pos[i] + offset;
121 }
122
123 Ok((axes, positions, inner_p_ee + offset))
124 }
125
126 fn jacobian(&self, q: &SRobotQ<M, F>) -> Result<[[F; M]; 6], Self::Error> {
127 let (_q_p, inner_q) = self.split_q(q);
128 let inner_j = self.inner.jacobian(&inner_q)?;
129 let p_col = self.prismatic_col();
130 let r_off = self.revolute_offset();
131
132 let mut j = [[F::zero(); M]; 6];
133 j[0][p_col] = self.axis.x();
134 j[1][p_col] = self.axis.y();
135 j[2][p_col] = self.axis.z();
136
137 for row in 0..6 {
138 for col in 0..N {
139 j[row][col + r_off] = inner_j[row][col];
140 }
141 }
142
143 Ok(j)
144 }
145
146 fn jacobian_dot(
147 &self,
148 q: &SRobotQ<M, F>,
149 qdot: &SRobotQ<M, F>,
150 ) -> Result<[[F; M]; 6], Self::Error> {
151 let (_q_p, inner_q) = self.split_q(q);
152 let (_qdot_p, inner_qdot) = self.split_q(qdot);
153 let inner_jd = self.inner.jacobian_dot(&inner_q, &inner_qdot)?;
154 let r_off = self.revolute_offset();
155
156 let mut jd = [[F::zero(); M]; 6];
157 for row in 0..6 {
158 for col in 0..N {
159 jd[row][col + r_off] = inner_jd[row][col];
160 }
161 }
162
163 Ok(jd)
164 }
165
166 fn jacobian_ddot(
167 &self,
168 q: &SRobotQ<M, F>,
169 qdot: &SRobotQ<M, F>,
170 qddot: &SRobotQ<M, F>,
171 ) -> Result<[[F; M]; 6], Self::Error> {
172 let (_q_p, inner_q) = self.split_q(q);
173 let (_qdot_p, inner_qdot) = self.split_q(qdot);
174 let (_qddot_p, inner_qddot) = self.split_q(qddot);
175 let inner_jdd = self
176 .inner
177 .jacobian_ddot(&inner_q, &inner_qdot, &inner_qddot)?;
178 let r_off = self.revolute_offset();
179
180 let mut jdd = [[F::zero(); M]; 6];
181 for row in 0..6 {
182 for col in 0..N {
183 jdd[row][col + r_off] = inner_jdd[row][col];
184 }
185 }
186
187 Ok(jdd)
188 }
189}
190
191impl<const M: usize, const N: usize, FK32, FK64> From<PrismaticFK<M, N, f32, FK32>>
192 for PrismaticFK<M, N, f64, FK64>
193where
194 FK32: FKChain<N, f32>,
195 FK64: FKChain<N, f64> + From<FK32>,
196{
197 #[inline]
198 fn from(p: PrismaticFK<M, N, f32, FK32>) -> Self {
199 PrismaticFK {
200 inner: FK64::from(p.inner),
201 axis: p.axis.as_dvec3(),
202 q_index_first: p.q_index_first,
203 }
204 }
205}
206
207impl<const M: usize, const N: usize, FK64, FK32> From<PrismaticFK<M, N, f64, FK64>>
208 for PrismaticFK<M, N, f32, FK32>
209where
210 FK64: FKChain<N, f64>,
211 FK32: FKChain<N, f32> + From<FK64>,
212{
213 #[inline]
214 fn from(p: PrismaticFK<M, N, f64, FK64>) -> Self {
215 PrismaticFK {
216 inner: FK32::from(p.inner),
217 axis: p.axis.as_vec3a(),
218 q_index_first: p.q_index_first,
219 }
220 }
221}