repose_kinematics/
chain.rs

1//! Serial kinematic chain definitions, forward kinematics, and geometric Jacobians.
2
3use bincode::Decode;
4use bincode::Encode;
5use cu_spatial_payloads::Transform3D;
6use cu29::prelude::*;
7use cu29::units::si::angle::radian;
8use cu29::units::si::f64::Angle as Angle64;
9use cu29::units::si::f64::Length as Length64;
10use cu29::units::si::length::meter;
11use serde::Deserialize;
12use serde::Serialize;
13
14/// Principal or arbitrary unit axis of actuation in the joint's local frame.
15#[derive(
16    Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize, Encode, Decode, Reflect,
17)]
18#[serde(rename_all = "snake_case")]
19pub enum CuJointAxis {
20    PosX,
21    NegX,
22    PosY,
23    NegY,
24    #[default]
25    PosZ,
26    NegZ,
27    Arbitrary([f64; 3]),
28}
29
30impl CuJointAxis {
31    /// Returns the unit vector `[x, y, z]` in the local joint frame.
32    pub fn unit_vector(self) -> [f64; 3] {
33        match self {
34            Self::PosX => [1.0, 0.0, 0.0],
35            Self::NegX => [-1.0, 0.0, 0.0],
36            Self::PosY => [0.0, 1.0, 0.0],
37            Self::NegY => [0.0, -1.0, 0.0],
38            Self::PosZ => [0.0, 0.0, 1.0],
39            Self::NegZ => [0.0, 0.0, -1.0],
40            Self::Arbitrary(v) => {
41                let norm = libm::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
42                v.map(|c| c / norm)
43            }
44        }
45    }
46}
47
48/// Actuation mechanism of a single joint in a serial chain.
49#[derive(
50    Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, Encode, Decode, Reflect,
51)]
52#[serde(rename_all = "snake_case")]
53pub enum CuJointType {
54    #[default]
55    Revolute,
56    Prismatic,
57    Fixed,
58}
59
60use crate::joint::CuJointLimits;
61use crate::joint::CuJointPosition;
62
63/// Static kinematic specification of a single joint following URDF conventions.
64#[derive(
65    Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize, Encode, Decode, Reflect,
66)]
67#[serde(deny_unknown_fields)]
68pub struct CuJointSpec {
69    /// Actuation mechanism.
70    pub joint_type: CuJointType,
71    /// Parent-to-joint translation (URDF `origin xyz`).
72    pub origin_translation: [Length64; 3],
73    /// Parent-to-joint roll-pitch-yaw (URDF `origin rpy`).
74    pub origin_rpy: [Angle64; 3],
75    /// Actuation axis in the joint frame.
76    pub axis: CuJointAxis,
77    /// Position and velocity bounds.
78    pub limits: CuJointLimits,
79}
80
81impl CuJointSpec {
82    /// Creates a revolute joint from its parent-frame origin (URDF `xyz`, `rpy`).
83    pub const fn revolute(
84        origin_translation: [Length64; 3],
85        origin_rpy: [Angle64; 3],
86        axis: CuJointAxis,
87        limits: CuJointLimits,
88    ) -> Self {
89        Self::actuated(
90            CuJointType::Revolute,
91            origin_translation,
92            origin_rpy,
93            axis,
94            limits,
95        )
96    }
97
98    /// Creates a prismatic joint from its parent-frame origin (URDF `xyz`, `rpy`).
99    pub const fn prismatic(
100        origin_translation: [Length64; 3],
101        origin_rpy: [Angle64; 3],
102        axis: CuJointAxis,
103        limits: CuJointLimits,
104    ) -> Self {
105        Self::actuated(
106            CuJointType::Prismatic,
107            origin_translation,
108            origin_rpy,
109            axis,
110            limits,
111        )
112    }
113
114    const fn actuated(
115        joint_type: CuJointType,
116        origin_translation: [Length64; 3],
117        origin_rpy: [Angle64; 3],
118        axis: CuJointAxis,
119        limits: CuJointLimits,
120    ) -> Self {
121        Self {
122            joint_type,
123            origin_translation,
124            origin_rpy,
125            axis,
126            limits,
127        }
128    }
129
130    /// Checks coordinate units and validity of joint specifications before numerical work.
131    pub fn validate_coordinate(&self, q: CuJointPosition) -> CuResult<()> {
132        self.limits.validate(q)?;
133        if self
134            .origin_translation
135            .iter()
136            .any(|v| !v.get::<meter>().is_finite())
137            || self
138                .origin_rpy
139                .iter()
140                .any(|v| !v.get::<radian>().is_finite())
141            || (self.joint_type != CuJointType::Fixed
142                && self.axis.unit_vector().iter().any(|v| !v.is_finite()))
143        {
144            return Err(CuError::from(
145                "Non-finite joint origin or invalid joint axis",
146            ));
147        }
148        if !matches!(
149            (self.joint_type, q),
150            (CuJointType::Revolute, CuJointPosition::Revolute(_))
151                | (CuJointType::Prismatic, CuJointPosition::Prismatic(_))
152                | (CuJointType::Fixed, CuJointPosition::Fixed)
153        ) {
154            return Err(CuError::from("Joint type and coordinate units disagree"));
155        }
156        Ok(())
157    }
158
159    /// Computes the static parent-to-joint origin transform.
160    pub fn origin_transform(&self) -> Transform3D<f64> {
161        Transform3D::<f64>::from_translation_euler_xyz(self.origin_translation, self.origin_rpy)
162    }
163
164    /// Computes the local motion transform for joint displacement `q`.
165    pub fn motion_transform(&self, q: CuJointPosition) -> CuResult<Transform3D<f64>> {
166        self.validate_coordinate(q)?;
167        let val = q.si_value();
168        let u = self.axis.unit_vector();
169        Ok(match self.joint_type {
170            CuJointType::Fixed => Transform3D::<f64>::identity(),
171            CuJointType::Revolute => axis_angle_transform(u, val),
172            CuJointType::Prismatic => Transform3D::<f64>::from_matrix([
173                [1.0, 0.0, 0.0, u[0] * val],
174                [0.0, 1.0, 0.0, u[1] * val],
175                [0.0, 0.0, 1.0, u[2] * val],
176                [0.0, 0.0, 0.0, 1.0],
177            ]),
178        })
179    }
180}
181
182/// Compile-time sized $N$-DOF serial robot kinematic chain.
183///
184/// # Example
185/// ```rust
186/// use repose_kinematics::{CuJointAxis, CuJointLimits, CuJointSpec, CuKinematicChain};
187/// use cu_spatial_payloads::Transform3D;
188/// use cu29::units::si::angle::radian;
189/// use cu29::units::si::angular_velocity::radian_per_second;
190/// use cu29::units::si::f64::{Angle, AngularVelocity, Length};
191/// use cu29::units::si::length::meter;
192///
193/// let m = Length::new::<meter>;
194/// let rad = Angle::new::<radian>;
195/// let chain = CuKinematicChain::new([CuJointSpec::revolute(
196///     [m(0.0), m(0.0), m(0.1)],
197///     [rad(0.0); 3],
198///     CuJointAxis::PosZ,
199///     CuJointLimits::revolute(
200///         rad(-core::f64::consts::PI),
201///         rad(core::f64::consts::PI),
202///         AngularVelocity::new::<radian_per_second>(2.0),
203///     ),
204/// )])
205/// .with_tool_offset(Transform3D::<f64>::from_translation_euler_xyz([m(0.0), m(0.0), m(0.05)], [rad(0.0); 3]));
206/// chain.validate().unwrap();
207/// let mut ee = Transform3D::<f64>::identity();
208/// chain.forward_kinematics(&[Angle::new::<radian>(0.0).into()], &mut ee).unwrap();
209/// assert!((ee.to_matrix()[2][3] - 0.15).abs() < 1e-12);
210/// ```
211#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)]
212#[reflect(no_field_bounds, from_reflect = false)]
213#[serde(deny_unknown_fields)]
214pub struct CuKinematicChain<const N: usize> {
215    #[serde(with = "crate::serde_arrays::joint_spec_array")]
216    joints: [CuJointSpec; N],
217    #[serde(default = "Transform3D::<f64>::identity")]
218    tool_offset: Transform3D<f64>,
219}
220
221impl<const N: usize> PartialEq for CuKinematicChain<N> {
222    fn eq(&self, other: &Self) -> bool {
223        self.joints == other.joints && self.tool_offset.to_matrix() == other.tool_offset.to_matrix()
224    }
225}
226
227impl<const N: usize> CuKinematicChain<N> {
228    /// Creates a serial kinematic chain with an identity tool offset.
229    pub const fn new(joints: [CuJointSpec; N]) -> Self {
230        Self {
231            joints,
232            tool_offset: Transform3D::<f64>::identity(),
233        }
234    }
235
236    /// Replaces the flange-to-tool transform (builder form of [`Self::set_tool_offset`]).
237    pub const fn with_tool_offset(mut self, tool_offset: Transform3D<f64>) -> Self {
238        self.tool_offset = tool_offset;
239        self
240    }
241
242    /// Checks that every joint has finite origins, a valid axis, and well-formed limits.
243    pub fn validate(&self) -> CuResult<()> {
244        for joint in &self.joints {
245            joint.validate_coordinate(joint.limits.lower())?;
246        }
247        if self
248            .tool_offset
249            .to_matrix()
250            .iter()
251            .flatten()
252            .any(|v| !v.is_finite())
253        {
254            return Err(CuError::from("Non-finite chain tool offset"));
255        }
256        Ok(())
257    }
258
259    /// Borrows the static joint specifications.
260    pub const fn joints(&self) -> &[CuJointSpec; N] {
261        &self.joints
262    }
263
264    /// Borrows the flange-to-tool transform.
265    pub const fn tool_offset(&self) -> &Transform3D<f64> {
266        &self.tool_offset
267    }
268
269    /// Sets the flange-to-tool transform.
270    pub fn set_tool_offset(&mut self, tool_offset: Transform3D<f64>) {
271        self.tool_offset = tool_offset;
272    }
273
274    /// Returns `true` if every joint angle in `q` satisfies its joint limits.
275    pub fn is_within_limits(&self, q: &[CuJointPosition; N]) -> bool {
276        let mut i = 0;
277        while i < N {
278            if self.joints[i].validate_coordinate(q[i]).is_err()
279                || !self.joints[i].limits.contains(q[i])
280            {
281                return false;
282            }
283            i += 1;
284        }
285        true
286    }
287
288    /// Clamps each joint angle in `q` in place to its configured limits.
289    pub fn clamp_to_limits(&self, q: &mut [CuJointPosition; N]) -> CuResult<()> {
290        let mut i = 0;
291        while i < N {
292            self.joints[i].validate_coordinate(q[i])?;
293            q[i] = self.joints[i].limits.clamp(q[i])?;
294            i += 1;
295        }
296        Ok(())
297    }
298
299    /// Computes the base-to-tool forward kinematics transform in place.
300    pub fn forward_kinematics(
301        &self,
302        q: &[CuJointPosition; N],
303        out: &mut Transform3D<f64>,
304    ) -> CuResult<()> {
305        let mut acc = Transform3D::<f64>::identity();
306        let mut i = 0;
307        while i < N {
308            if !q[i].si_value().is_finite() {
309                return Err(CuError::from("repose_kinematics: non-finite joint angle"));
310            }
311            let step = self.joints[i]
312                .origin_transform()
313                .compose(self.joints[i].motion_transform(q[i])?);
314            acc = acc.compose(step);
315            i += 1;
316        }
317        *out = acc.compose(self.tool_offset);
318        Ok(())
319    }
320
321    /// Computes every intermediate link frame (`0..N`) and the final end-effector frame in place.
322    pub fn forward_frames(
323        &self,
324        q: &[CuJointPosition; N],
325        out_frames: &mut [Transform3D<f64>; N],
326        out_ee: &mut Transform3D<f64>,
327    ) -> CuResult<()> {
328        let mut acc = Transform3D::<f64>::identity();
329        let mut i = 0;
330        while i < N {
331            if !q[i].si_value().is_finite() {
332                return Err(CuError::from("repose_kinematics: non-finite joint angle"));
333            }
334            let step = self.joints[i]
335                .origin_transform()
336                .compose(self.joints[i].motion_transform(q[i])?);
337            acc = acc.compose(step);
338            out_frames[i] = acc;
339            i += 1;
340        }
341        *out_ee = acc.compose(self.tool_offset);
342        Ok(())
343    }
344
345    /// Computes the $6 \times N$ geometric space Jacobian in the base frame in place.
346    ///
347    /// Rows `0..3` correspond to linear velocity $[v_x, v_y, v_z]^T$ in $\text{m/s}$ and
348    /// rows `3..6` correspond to angular velocity $[\omega_x, \omega_y, \omega_z]^T$ in $\text{rad/s}$,
349    /// both measured at the tool origin and expressed in the base frame.
350    pub fn geometric_jacobian(
351        &self,
352        q: &[CuJointPosition; N],
353        out_jacobian: &mut [[f64; N]; 6],
354    ) -> CuResult<()> {
355        let mut ee = Transform3D::<f64>::identity();
356        self.geometric_jacobian_and_pose(q, out_jacobian, &mut ee)
357    }
358
359    /// Computes the $6 \times N$ geometric body Jacobian in the end-effector (tool) frame in place.
360    ///
361    /// Rows `0..3` correspond to linear velocity $[v_x, v_y, v_z]^T$ in $\text{m/s}$ and
362    /// rows `3..6` correspond to angular velocity $[\omega_x, \omega_y, \omega_z]^T$ in $\text{rad/s}$,
363    /// both measured at the tool origin and expressed in the end-effector frame.
364    pub fn body_jacobian(
365        &self,
366        q: &[CuJointPosition; N],
367        out_jacobian: &mut [[f64; N]; 6],
368    ) -> CuResult<()> {
369        let mut space_jac = [[0.0_f64; N]; 6];
370        let mut ee = Transform3D::<f64>::identity();
371        self.geometric_jacobian_and_pose(q, &mut space_jac, &mut ee)?;
372
373        let ee_mat = ee.to_matrix();
374        let mut col = 0;
375        while col < N {
376            let vx = space_jac[0][col];
377            let vy = space_jac[1][col];
378            let vz = space_jac[2][col];
379            let wx = space_jac[3][col];
380            let wy = space_jac[4][col];
381            let wz = space_jac[5][col];
382
383            // Multiply by R_ee^T (columns of R_ee become rows of R_ee^T)
384            out_jacobian[0][col] = ee_mat[0][0] * vx + ee_mat[1][0] * vy + ee_mat[2][0] * vz;
385            out_jacobian[1][col] = ee_mat[0][1] * vx + ee_mat[1][1] * vy + ee_mat[2][1] * vz;
386            out_jacobian[2][col] = ee_mat[0][2] * vx + ee_mat[1][2] * vy + ee_mat[2][2] * vz;
387
388            out_jacobian[3][col] = ee_mat[0][0] * wx + ee_mat[1][0] * wy + ee_mat[2][0] * wz;
389            out_jacobian[4][col] = ee_mat[0][1] * wx + ee_mat[1][1] * wy + ee_mat[2][1] * wz;
390            out_jacobian[5][col] = ee_mat[0][2] * wx + ee_mat[1][2] * wy + ee_mat[2][2] * wz;
391            col += 1;
392        }
393        Ok(())
394    }
395
396    pub(crate) fn geometric_jacobian_and_pose(
397        &self,
398        q: &[CuJointPosition; N],
399        out_jacobian: &mut [[f64; N]; 6],
400        out_ee: &mut Transform3D<f64>,
401    ) -> CuResult<()> {
402        let mut joint_origins = [[0.0_f64; 3]; N];
403        let mut world_axes = [[0.0_f64; 3]; N];
404        let mut acc = Transform3D::<f64>::identity();
405
406        let mut i = 0;
407        while i < N {
408            if !q[i].si_value().is_finite() {
409                return Err(CuError::from("repose_kinematics: non-finite joint angle"));
410            }
411            let joint_frame = acc.compose(self.joints[i].origin_transform());
412            let mat = joint_frame.to_matrix();
413            joint_origins[i] = [mat[0][3], mat[1][3], mat[2][3]];
414
415            let local_axis = self.joints[i].axis.unit_vector();
416            world_axes[i] = [
417                mat[0][0] * local_axis[0] + mat[0][1] * local_axis[1] + mat[0][2] * local_axis[2],
418                mat[1][0] * local_axis[0] + mat[1][1] * local_axis[1] + mat[1][2] * local_axis[2],
419                mat[2][0] * local_axis[0] + mat[2][1] * local_axis[1] + mat[2][2] * local_axis[2],
420            ];
421
422            acc = joint_frame.compose(self.joints[i].motion_transform(q[i])?);
423            i += 1;
424        }
425
426        *out_ee = acc.compose(self.tool_offset);
427        let ee_mat = out_ee.to_matrix();
428        let p_ee = [ee_mat[0][3], ee_mat[1][3], ee_mat[2][3]];
429
430        let mut col = 0;
431        while col < N {
432            let w = world_axes[col];
433            match self.joints[col].joint_type {
434                CuJointType::Revolute => {
435                    let r = [
436                        p_ee[0] - joint_origins[col][0],
437                        p_ee[1] - joint_origins[col][1],
438                        p_ee[2] - joint_origins[col][2],
439                    ];
440                    let v = cross3(w, r);
441                    out_jacobian[0][col] = v[0];
442                    out_jacobian[1][col] = v[1];
443                    out_jacobian[2][col] = v[2];
444                    out_jacobian[3][col] = w[0];
445                    out_jacobian[4][col] = w[1];
446                    out_jacobian[5][col] = w[2];
447                }
448                CuJointType::Prismatic => {
449                    out_jacobian[0][col] = w[0];
450                    out_jacobian[1][col] = w[1];
451                    out_jacobian[2][col] = w[2];
452                    out_jacobian[3][col] = 0.0;
453                    out_jacobian[4][col] = 0.0;
454                    out_jacobian[5][col] = 0.0;
455                }
456                CuJointType::Fixed => {
457                    out_jacobian[0][col] = 0.0;
458                    out_jacobian[1][col] = 0.0;
459                    out_jacobian[2][col] = 0.0;
460                    out_jacobian[3][col] = 0.0;
461                    out_jacobian[4][col] = 0.0;
462                    out_jacobian[5][col] = 0.0;
463                }
464            }
465            col += 1;
466        }
467        Ok(())
468    }
469}
470
471const fn cross3(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
472    [
473        a[1] * b[2] - a[2] * b[1],
474        a[2] * b[0] - a[0] * b[2],
475        a[0] * b[1] - a[1] * b[0],
476    ]
477}
478
479fn axis_angle_transform(axis: [f64; 3], angle_rad: f64) -> Transform3D<f64> {
480    let norm_sq = axis[0] * axis[0] + axis[1] * axis[1] + axis[2] * axis[2];
481    if norm_sq <= 1e-15 {
482        return Transform3D::<f64>::identity();
483    }
484    let inv_norm = 1.0 / libm::sqrt(norm_sq);
485    let ux = axis[0] * inv_norm;
486    let uy = axis[1] * inv_norm;
487    let uz = axis[2] * inv_norm;
488
489    let s = libm::sin(angle_rad);
490    let c = libm::cos(angle_rad);
491    let one_minus_c = 1.0 - c;
492
493    Transform3D::<f64>::from_matrix([
494        [
495            c + ux * ux * one_minus_c,
496            ux * uy * one_minus_c - uz * s,
497            ux * uz * one_minus_c + uy * s,
498            0.0,
499        ],
500        [
501            uy * ux * one_minus_c + uz * s,
502            c + uy * uy * one_minus_c,
503            uy * uz * one_minus_c - ux * s,
504            0.0,
505        ],
506        [
507            uz * ux * one_minus_c - uy * s,
508            uz * uy * one_minus_c + ux * s,
509            c + uz * uz * one_minus_c,
510            0.0,
511        ],
512        [0.0, 0.0, 0.0, 1.0],
513    ])
514}
repose_kinematics/chain.rs

repose_kinematics/
chain.rs
