tiny_solver/manifold/
so3.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
use std::{num::NonZero, ops::Mul};

use nalgebra as na;

use super::{AutoDiffManifold, Manifold};

pub struct SO3<T: na::RealField> {
    qx: T,
    qy: T,
    qz: T,
    qw: T,
}

impl<T: na::RealField> SO3<T> {
    /// [x, y, z, w]
    pub fn from_vec(xyzw: na::DVectorView<T>) -> Self {
        SO3 {
            qx: xyzw[0].clone(),
            qy: xyzw[1].clone(),
            qz: xyzw[2].clone(),
            qw: xyzw[3].clone(),
        }
    }
    pub fn from_xyzw(x: T, y: T, z: T, w: T) -> Self {
        SO3 {
            qx: x,
            qy: y,
            qz: z,
            qw: w,
        }
    }
    pub fn identity() -> Self {
        SO3 {
            qx: T::zero(),
            qy: T::zero(),
            qz: T::zero(),
            qw: T::one(),
        }
    }

    pub fn exp(xi: na::DVectorView<T>) -> Self {
        let mut xyzw = na::Vector4::zeros();

        let theta2 = xi.norm_squared();

        if theta2 < T::from_f64(1e-6).unwrap() {
            // cos(theta / 2) \approx 1 - theta^2 / 8
            xyzw.w = T::one() - theta2 / T::from_f64(8.0).unwrap();
            // Complete the square so that norm is one
            let tmp = T::from_f64(0.5).unwrap();
            xyzw.x = xi[0].clone() * tmp.clone();
            xyzw.y = xi[1].clone() * tmp.clone();
            xyzw.z = xi[2].clone() * tmp;
        } else {
            let theta = theta2.sqrt();
            xyzw.w = (theta.clone() * T::from_f64(0.5).unwrap()).cos();

            let omega = xi / theta;
            let sin_theta_half = (T::one() - xyzw.w.clone() * xyzw.w.clone()).sqrt();
            xyzw.x = omega[0].clone() * sin_theta_half.clone();
            xyzw.y = omega[1].clone() * sin_theta_half.clone();
            xyzw.z = omega[2].clone() * sin_theta_half;
        }

        SO3::from_vec(xyzw.as_view())
    }

    pub fn log(&self) -> na::DVector<T> {
        const EPS: f64 = 1e-6;
        let ivec = na::dvector![self.qx.clone(), self.qy.clone(), self.qz.clone()];

        let squared_n = ivec.norm_squared();
        let w = self.qw.clone();

        let near_zero = squared_n.le(&T::from_f64(EPS * EPS).unwrap());

        let w_sq = w.clone() * w.clone();
        let t0 = T::from_f64(2.0).unwrap() / w.clone()
            - T::from_f64(2.0 / 3.0).unwrap() * squared_n.clone() / (w_sq * w.clone());

        let n = squared_n.sqrt();

        let sign = T::from_f64(-1.0)
            .unwrap()
            .select(w.le(&T::zero()), T::one());
        let atan_nbyw = sign.clone() * n.clone().atan2(sign * w);

        let t = T::from_f64(2.0).unwrap() * atan_nbyw / n;

        let two_atan_nbyd_by_n = t0.select(near_zero, t);

        ivec * two_atan_nbyd_by_n
    }

    pub fn hat(xi: na::VectorView3<T>) -> na::Matrix3<T> {
        let mut xi_hat = na::Matrix3::zeros();
        xi_hat[(0, 1)] = -xi[2].clone();
        xi_hat[(0, 2)] = xi[1].clone();
        xi_hat[(1, 0)] = xi[2].clone();
        xi_hat[(1, 2)] = -xi[0].clone();
        xi_hat[(2, 0)] = -xi[1].clone();
        xi_hat[(2, 1)] = xi[0].clone();

        xi_hat
    }

    pub fn to_vec(&self) -> na::Vector4<T> {
        na::Vector4::new(
            self.qx.clone(),
            self.qy.clone(),
            self.qz.clone(),
            self.qw.clone(),
        )
    }
    pub fn to_dvec(&self) -> na::DVector<T> {
        na::dvector![
            self.qx.clone(),
            self.qy.clone(),
            self.qz.clone(),
            self.qw.clone(),
        ]
    }
    pub fn cast<U: na::RealField + simba::scalar::SupersetOf<T>>(&self) -> SO3<U> {
        SO3::from_vec(self.to_vec().cast().as_view())
    }
    pub fn inverse(&self) -> Self {
        SO3 {
            qx: -self.qx.clone(),
            qy: -self.qy.clone(),
            qz: -self.qz.clone(),
            qw: self.qw.clone(),
        }
    }
    pub fn compose(&self, rhs: &Self) -> Self {
        let x0 = self.qx.clone();
        let y0 = self.qy.clone();
        let z0 = self.qz.clone();
        let w0 = self.qw.clone();

        let x1 = rhs.qx.clone();
        let y1 = rhs.qy.clone();
        let z1 = rhs.qz.clone();
        let w1 = rhs.qw.clone();

        // Compute the product of the two quaternions, term by term
        let qx = w0.clone() * x1.clone() + x0.clone() * w1.clone() + y0.clone() * z1.clone()
            - z0.clone() * y1.clone();
        let qy = w0.clone() * y1.clone() - x0.clone() * z1.clone()
            + y0.clone() * w1.clone()
            + z0.clone() * x1.clone();
        let qz = w0.clone() * z1.clone() + x0.clone() * y1.clone() - y0.clone() * x1.clone()
            + z0.clone() * w1.clone();
        let qw = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1;
        // println!("q {}", self.to_vec());
        SO3 { qx, qy, qz, qw }
    }
}

impl<T: na::RealField> Mul for SO3<T> {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        self.compose(&rhs)
    }
}
impl<T: na::RealField> Mul for &SO3<T> {
    type Output = SO3<T>;

    fn mul(self, rhs: Self) -> Self::Output {
        self.compose(rhs)
    }
}

impl<T: na::RealField> Mul<na::VectorView3<'_, T>> for &SO3<T> {
    type Output = na::Vector3<T>;

    fn mul(self, rhs: na::VectorView3<'_, T>) -> Self::Output {
        let qv = SO3::from_xyzw(rhs[0].clone(), rhs[1].clone(), rhs[2].clone(), T::zero());
        let inv = self.inverse();
        let v_rot: SO3<T> = (self * &qv) * inv;
        na::Vector3::new(v_rot.qx, v_rot.qy, v_rot.qz)
    }
}

#[derive(Debug, Clone)]
pub struct QuaternionManifold;
impl<T: na::RealField> AutoDiffManifold<T> for QuaternionManifold {
    fn plus(
        &self,
        x: nalgebra::DVectorView<T>,
        delta: nalgebra::DVectorView<T>,
    ) -> nalgebra::DVector<T> {
        let d: SO3<T> = SO3::exp(delta);
        let x_s03: SO3<T> = SO3::from_xyzw(x[0].clone(), x[1].clone(), x[2].clone(), x[3].clone());
        let x_plus = x_s03 * d;
        na::dvector![
            x_plus.qx.clone(),
            x_plus.qy.clone(),
            x_plus.qz.clone(),
            x_plus.qw.clone(),
        ]
    }

    fn minus(
        &self,
        y: nalgebra::DVectorView<T>,
        x: nalgebra::DVectorView<T>,
    ) -> nalgebra::DVector<T> {
        let y_so3 = SO3::from_vec(y);
        let x_so3_inv = SO3::from_vec(x).inverse();
        let x_inv_y_log = (x_so3_inv * y_so3).log();
        na::dvector![
            x_inv_y_log[0].clone(),
            x_inv_y_log[1].clone(),
            x_inv_y_log[2].clone()
        ]
    }
}
impl Manifold for QuaternionManifold {
    fn tangent_size(&self) -> NonZero<usize> {
        NonZero::new(3).unwrap()
    }
}