Skip to main content

deke_types/fk/
prismatic.rs

1use glam_traits_ext::{FloatAffine, FloatVec, TAffine3, TVec3};
2
3use crate::SRobotQ;
4
5use super::{AAffine3, AVec3, FKChain, FKScalar};
6
7/// Wraps an `FKChain<N, F>` and prepends a prismatic (linear) joint, producing
8/// an `FKChain<M, F>` where `M = N + 1`.
9///
10/// The prismatic joint always acts first in the kinematic chain — it
11/// translates the entire arm along `axis` (world frame).  The
12/// `q_index_first` flag only controls where the prismatic value is read
13/// from in `SRobotQ<M, F>`: when `true` it is `q[0]`, when `false` it is
14/// `q[M-1]`.
15///
16/// Jacobian columns for the prismatic joint are `[axis; 0]` (pure linear,
17/// no angular contribution).  Because the prismatic uniformly shifts all
18/// positions, the revolute Jacobian columns are identical to the inner
19/// chain's.
20#[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            // Apply translation by `offset` to inner frame: post-multiply with translate(offset).
88            // Equivalent to setting f.translation += offset.
89            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}