tiny_solver/manifold/
se3.rs

1use std::{num::NonZero, ops::Mul};
2
3use nalgebra as na;
4
5use super::{AutoDiffManifold, Manifold, so3::SO3};
6
7pub struct SE3<T: na::RealField> {
8    pub xyz: na::Vector3<T>,
9    pub rot: SO3<T>,
10}
11
12impl<T: na::RealField> SE3<T> {
13    /// [qx, qy, qz, qw, tx, ty, tz]
14    pub fn from_vec(qxyzw_txyz: na::DVectorView<T>) -> Self {
15        let rot = SO3::from_xyzw(
16            qxyzw_txyz[0].clone(),
17            qxyzw_txyz[1].clone(),
18            qxyzw_txyz[2].clone(),
19            qxyzw_txyz[3].clone(),
20        );
21        let xyz = na::Vector3::new(
22            qxyzw_txyz[4].clone(),
23            qxyzw_txyz[5].clone(),
24            qxyzw_txyz[6].clone(),
25        );
26        SE3 { xyz, rot }
27    }
28    pub fn from_qvec_tvec(qxyzw: na::DVectorView<T>, tvec: na::DVectorView<T>) -> Self {
29        let rot = SO3::from_vec(qxyzw);
30        let xyz = na::Vector3::from_row_slice(tvec.as_slice());
31        SE3 { xyz, rot }
32    }
33    pub fn identity() -> Self {
34        let rot = SO3::identity();
35        let xyz = na::Vector3::zeros();
36        SE3 { xyz, rot }
37    }
38
39    pub fn exp(xi: na::DVectorView<T>) -> Self {
40        // fake exp
41        let rot = SO3::<T>::exp(xi.rows(0, 3).as_view());
42        let xyz = na::Vector3::new(xi[3].clone(), xi[4].clone(), xi[5].clone());
43        SE3 { xyz, rot }
44    }
45
46    pub fn log(&self) -> na::DVector<T> {
47        let mut xi = na::DVector::zeros(6);
48        let xi_theta = self.rot.log();
49        // let xyz = self.xyz;
50        xi.as_mut_slice()[0..3].clone_from_slice(xi_theta.as_slice());
51        xi.as_mut_slice()[3..6].clone_from_slice(self.xyz.as_slice());
52        xi
53    }
54
55    pub fn to_dvec(&self) -> na::DVector<T> {
56        let quat = self.rot.to_vec();
57        na::dvector![
58            quat[0].clone(),
59            quat[1].clone(),
60            quat[2].clone(),
61            quat[3].clone(),
62            self.xyz[0].clone(),
63            self.xyz[1].clone(),
64            self.xyz[2].clone(),
65        ]
66    }
67
68    // pub fn hat(xi: na::VectorView3<T>) -> na::Matrix3<T> {
69    //     let mut xi_hat = na::Matrix3::zeros();
70    //     xi_hat[(0, 1)] = -xi[2].clone();
71    //     xi_hat[(0, 2)] = xi[1].clone();
72    //     xi_hat[(1, 0)] = xi[2].clone();
73    //     xi_hat[(1, 2)] = -xi[0].clone();
74    //     xi_hat[(2, 0)] = -xi[1].clone();
75    //     xi_hat[(2, 1)] = xi[0].clone();
76
77    //     xi_hat
78    // }
79
80    pub fn cast<U: na::RealField + simba::scalar::SupersetOf<T>>(&self) -> SE3<U> {
81        SE3 {
82            rot: self.rot.cast(),
83            xyz: self.xyz.clone().cast(),
84        }
85    }
86    pub fn inverse(&self) -> Self {
87        let inv = self.rot.inverse();
88        let xyz = -(&inv * self.xyz.as_view());
89        SE3 { xyz, rot: inv }
90    }
91    pub fn compose(&self, rhs: &Self) -> Self {
92        SE3 {
93            rot: &self.rot * &rhs.rot,
94            xyz: (&self.rot * rhs.xyz.as_view()) + self.xyz.clone(),
95        }
96    }
97}
98
99impl<T: na::RealField> Mul for SE3<T> {
100    type Output = Self;
101
102    fn mul(self, rhs: Self) -> Self::Output {
103        self.compose(&rhs)
104    }
105}
106impl<T: na::RealField> Mul for &SE3<T> {
107    type Output = SE3<T>;
108
109    fn mul(self, rhs: Self) -> Self::Output {
110        self.compose(rhs)
111    }
112}
113
114impl<T: na::RealField> Mul<na::VectorView3<'_, T>> for SE3<T> {
115    type Output = na::Vector3<T>;
116
117    fn mul(self, rhs: na::VectorView3<'_, T>) -> Self::Output {
118        let qv = SO3::from_xyzw(rhs[0].clone(), rhs[1].clone(), rhs[2].clone(), T::zero());
119        let rinv = self.rot.inverse();
120        let v_rot = ((self.rot * qv) * rinv).to_vec();
121        let v = na::Vector3::new(v_rot[0].clone(), v_rot[1].clone(), v_rot[2].clone());
122        v + self.xyz
123    }
124}
125
126impl<T: na::RealField> Mul<na::VectorView3<'_, T>> for &SE3<T> {
127    type Output = na::Vector3<T>;
128
129    fn mul(self, rhs: na::VectorView3<'_, T>) -> Self::Output {
130        let qv = SO3::from_xyzw(rhs[0].clone(), rhs[1].clone(), rhs[2].clone(), T::zero());
131        let rinv = self.rot.inverse();
132        let v_rot = ((&self.rot * &qv) * rinv).to_vec();
133        let v = na::Vector3::new(v_rot[0].clone(), v_rot[1].clone(), v_rot[2].clone());
134        v + self.xyz.clone()
135    }
136}
137
138#[derive(Debug, Clone)]
139pub struct SE3Manifold;
140impl<T: na::RealField> AutoDiffManifold<T> for SE3Manifold {
141    fn plus(
142        &self,
143        x: nalgebra::DVectorView<T>,
144        delta: nalgebra::DVectorView<T>,
145    ) -> nalgebra::DVector<T> {
146        let d = SE3::exp(delta);
147        let x_se3 = SE3::from_vec(x);
148        let x_plus = x_se3 * d;
149        x_plus.to_dvec()
150    }
151
152    fn minus(
153        &self,
154        y: nalgebra::DVectorView<T>,
155        x: nalgebra::DVectorView<T>,
156    ) -> nalgebra::DVector<T> {
157        let y_se3 = SE3::from_vec(y);
158        let x_se3_inv = SE3::from_vec(x).inverse();
159        let x_inv_y_log = (x_se3_inv * y_se3).log();
160        na::dvector![
161            x_inv_y_log[0].clone(),
162            x_inv_y_log[1].clone(),
163            x_inv_y_log[2].clone(),
164            x_inv_y_log[3].clone(),
165            x_inv_y_log[4].clone(),
166            x_inv_y_log[5].clone()
167        ]
168    }
169}
170impl Manifold for SE3Manifold {
171    fn tangent_size(&self) -> NonZero<usize> {
172        NonZero::new(6).unwrap()
173    }
174}