sophus_calculus/maps/
matrix_valued_maps.rs

1use crate::dual::dual_matrix::DualM;
2use crate::dual::dual_vector::DualV;
3use crate::types::matrix::IsMatrix;
4use crate::types::MatF64;
5use crate::types::VecF64;
6
7use sophus_tensor::element::SMat;
8use sophus_tensor::mut_tensor::MutTensorDDRC;
9use sophus_tensor::mut_tensor::MutTensorDRC;
10use sophus_tensor::mut_view::IsMutTensorLike;
11
12use std::marker::PhantomData;
13
14/// Matrix-valued map on a vector space.
15///
16/// This is a function which takes a vector and returns a matrix:
17///
18///  f: ℝᵐ -> ℝʳ x ℝᶜ
19///
20pub struct MatrixValuedMapFromVector;
21
22impl MatrixValuedMapFromVector {
23    /// Finite difference quotient of the matrix-valued map.
24    ///
25    /// The derivative is a rank-3 tensor with shape (Rₒ x Cₒ x Rᵢ).
26    ///
27    /// For efficiency reasons, we return Rᵢ x [Rₒ x Cₒ]
28    pub fn sym_diff_quotient<TFn, const OUTROWS: usize, const OUTCOLS: usize, const INROWS: usize>(
29        matrix_valued: TFn,
30        a: VecF64<INROWS>,
31        eps: f64,
32    ) -> MutTensorDRC<f64, OUTROWS, OUTCOLS>
33    where
34        TFn: Fn(VecF64<INROWS>) -> MatF64<OUTROWS, OUTCOLS>,
35    {
36        let mut out = MutTensorDRC::<f64, OUTROWS, OUTCOLS>::from_shape([INROWS]);
37        for i1 in 0..INROWS {
38            let mut a_plus = a;
39
40            a_plus[i1] += eps;
41
42            let mut a_minus = a;
43            a_minus[i1] -= eps;
44
45            let val = (matrix_valued(a_plus) - matrix_valued(a_minus)).scaled(1.0 / (2.0 * eps));
46
47            *out.mut_view().get_mut([i1]) = val;
48        }
49        out
50    }
51
52    /// Auto differentiation of the matrix-valued map.
53    pub fn fw_autodiff<TFn, const OUTROWS: usize, const OUTCOLS: usize, const INROWS: usize>(
54        matrix_valued: TFn,
55        a: VecF64<INROWS>,
56    ) -> MutTensorDRC<f64, OUTROWS, OUTCOLS>
57    where
58        TFn: Fn(DualV<INROWS>) -> DualM<OUTROWS, OUTCOLS>,
59    {
60        MutTensorDRC {
61            mut_array: matrix_valued(DualV::v(a))
62                .dij_val
63                .unwrap()
64                .mut_array
65                .into_shape([INROWS])
66                .unwrap(),
67            phantom: PhantomData,
68        }
69    }
70}
71
72/// Matrix-valued map on a product space (=matrices).
73///
74/// This is a function which takes a matrix and returns a matrix:
75///
76///  f: ℝᵐ x ℝⁿ -> ℝʳ x ℝᶜ
77///
78pub struct MatrixValuedMapFromMatrix;
79
80impl MatrixValuedMapFromMatrix {
81    /// Finite difference quotient of the matrix-valued map.
82    ///
83    /// The derivative is a rank-4 tensor with shape (Rₒ x Cₒ x Rᵢ x Cᵢ).
84    ///
85    /// For efficiency reasons, we return Rᵢ x Cᵢ x [Rₒ x Cₒ]
86    pub fn sym_diff_quotient<
87        TFn,
88        const OUTROWS: usize,
89        const OUTCOLS: usize,
90        const INROWS: usize,
91        const INCOLS: usize,
92    >(
93        vector_field: TFn,
94        a: MatF64<INROWS, INCOLS>,
95        eps: f64,
96    ) -> MutTensorDDRC<f64, OUTROWS, OUTCOLS>
97    where
98        TFn: Fn(MatF64<INROWS, INCOLS>) -> MatF64<OUTROWS, OUTCOLS>,
99    {
100        let mut out = MutTensorDDRC::<f64, OUTROWS, OUTCOLS>::from_shape_and_val(
101            [INROWS, INCOLS],
102            SMat::<f64, OUTROWS, OUTCOLS>::zeros(),
103        );
104        for i1 in 0..INROWS {
105            for i0 in 0..INCOLS {
106                let mut a_plus = a;
107
108                a_plus[(i1, i0)] += eps;
109
110                let mut a_minus = a;
111                a_minus[(i1, i0)] -= eps;
112
113                let val = (vector_field(a_plus) - vector_field(a_minus)) / (2.0 * eps);
114
115                *out.mut_view().get_mut([i1, i0]) = val;
116            }
117        }
118        out
119    }
120
121    /// Auto differentiation of the matrix-valued map.
122    pub fn fw_autodiff<
123        TFn,
124        const OUTROWS: usize,
125        const OUTCOLS: usize,
126        const INROWS: usize,
127        const INCOLS: usize,
128    >(
129        matrix_valued: TFn,
130        a: MatF64<INROWS, INCOLS>,
131    ) -> MutTensorDDRC<f64, OUTROWS, OUTCOLS>
132    where
133        TFn: Fn(DualM<INROWS, INCOLS>) -> DualM<OUTROWS, OUTCOLS>,
134    {
135        matrix_valued(DualM::v(a)).dij_val.unwrap()
136    }
137}
138
139mod test {
140
141    #[test]
142    fn test_batched_matrix_valued_map_from_vector() {
143        use crate::maps::matrix_valued_maps::MatrixValuedMapFromVector;
144        use crate::types::matrix::IsMatrix;
145        use crate::types::scalar::IsScalar;
146        use crate::types::vector::IsVector;
147        use crate::types::VecF64;
148        use sophus_tensor::view::IsTensorLike;
149
150        //      [[ i ]]
151        //      [[   ]]
152        //      [[ j ]]      [[                      ]]
153        //      [[   ]]      [[   0   -k    j    x   ]]
154        //      [[ k ]]      [[                      ]]
155        //  hat [[   ]]   =  [[   k    0   -i    y   ]]
156        //      [[ x ]]      [[                      ]]
157        //      [[   ]]      [[  -j    i    0    z   ]]
158        //      [[ y ]]      [[                      ]]
159        //      [[   ]]
160        //      [[ z ]]
161        fn hat_fn<S: IsScalar<1>, M34: IsMatrix<S, 3, 4, 1>, V6: IsVector<S, 6, 1>>(v: V6) -> M34 {
162            let i = v.get(0);
163            let j = v.get(1);
164            let k = v.get(2);
165            let ni = -i.clone();
166            let nj = -j.clone();
167            let nk = -k.clone();
168            let x = v.get(3);
169            let y = v.get(4);
170            let z = v.get(5);
171
172            let ret: M34 = M34::from_array2([
173                [S::c(0.0), nk, j, x],
174                [k, S::c(0.0), ni, y],
175                [nj, i, S::c(0.0), z],
176            ]);
177            ret
178        }
179
180        let a = VecF64::<6>::new(0.1, 0.2, 0.4, 0.7, 0.8, 0.9);
181
182        let finite_diff = MatrixValuedMapFromVector::sym_diff_quotient(hat_fn, a, 1e-6);
183        let auto_grad = MatrixValuedMapFromVector::fw_autodiff(hat_fn, a);
184        approx::assert_abs_diff_eq!(
185            finite_diff.view().elem_view(),
186            auto_grad.view().elem_view(),
187            epsilon = 0.0001
188        );
189    }
190
191    #[test]
192    fn test_batched_matrix_valued_map_from_matrix() {
193        use crate::maps::matrix_valued_maps::MatrixValuedMapFromMatrix;
194        use crate::types::matrix::IsMatrix;
195        use crate::types::scalar::IsScalar;
196        use crate::types::MatF64;
197        use sophus_tensor::view::IsTensorLike;
198
199        //      [[ a   b ]]       1    [[  d  -b ]]
200        //  inv [[       ]] =  ------- [[        ]]
201        //      [[ c   d ]]    ad - bc [[ -c   a ]]
202
203        fn f<S: IsScalar<1>, M22: IsMatrix<S, 2, 2, 1>>(m: M22) -> M22 {
204            let a = m.get((0, 0));
205            let b = m.get((0, 1));
206
207            let c = m.get((1, 0));
208            let d = m.get((1, 1));
209
210            let det = S::c(1.0) / (a.clone() * d.clone() - (b.clone() * c.clone()));
211
212            M22::from_array2([
213                [det.clone() * d, -det.clone() * b],
214                [-det.clone() * c, det * a],
215            ])
216        }
217        let a = MatF64::<2, 2>::new(0.1, 0.2, 0.4, 0.7);
218
219        let finite_diff = MatrixValuedMapFromMatrix::sym_diff_quotient(f, a, 1e-6);
220        let auto_grad = MatrixValuedMapFromMatrix::fw_autodiff(f, a);
221
222        approx::assert_abs_diff_eq!(
223            finite_diff.view().elem_view(),
224            auto_grad.view().elem_view(),
225            epsilon = 2.0
226        );
227    }
228}