Skip to main content

levenberg_marquardt_sparse/
problem.rs

1use alloc::vec::Vec;
2use nalgebra::{
3    ComplexField, DefaultAllocator, Dim, Matrix, OMatrix, Vector,
4    allocator::Allocator,
5    storage::{IsContiguous, RawStorage, RawStorageMut, Storage},
6};
7
8/// Sparse Jacobian in triplet (COO) format.
9#[derive(Clone, Debug, PartialEq)]
10pub struct SparseJacobian<F> {
11    pub rows: usize,
12    pub cols: usize,
13    pub entries: Vec<(usize, usize, F)>,
14}
15
16impl<F: ComplexField + Copy> SparseJacobian<F> {
17    #[must_use]
18    pub fn new(rows: usize, cols: usize) -> Self {
19        Self {
20            rows,
21            cols,
22            entries: Vec::new(),
23        }
24    }
25
26    #[must_use]
27    pub fn from_triplets(rows: usize, cols: usize, entries: Vec<(usize, usize, F)>) -> Self {
28        Self {
29            rows,
30            cols,
31            entries,
32        }
33    }
34
35    #[must_use]
36    pub fn from_dense<M, N, S>(matrix: Matrix<F, M, N, S>) -> Self
37    where
38        M: Dim,
39        N: Dim,
40        S: RawStorage<F, M, N>,
41    {
42        let mut entries = Vec::new();
43        let rows = matrix.nrows();
44        let cols = matrix.ncols();
45        for i in 0..rows {
46            for j in 0..cols {
47                let value = matrix[(i, j)];
48                if !value.is_zero() {
49                    entries.push((i, j, value));
50                }
51            }
52        }
53        Self {
54            rows,
55            cols,
56            entries,
57        }
58    }
59
60    #[must_use]
61    pub fn to_dense<M, N>(&self) -> OMatrix<F, M, N>
62    where
63        M: Dim,
64        N: Dim,
65        DefaultAllocator: Allocator<M, N>,
66    {
67        let mut out = OMatrix::<F, M, N>::zeros_generic(
68            Dim::from_usize(self.rows),
69            Dim::from_usize(self.cols),
70        );
71        for &(i, j, value) in &self.entries {
72            out[(i, j)] += value;
73        }
74        out
75    }
76}
77
78/// A least squares minimization problem.
79///
80/// This is what [`LevenbergMarquardt`](struct.LevenbergMarquardt.html) needs
81/// to compute the residuals and the Jacobian. See the [module documentation](index.html)
82/// for a usage example.
83pub trait LeastSquaresProblem<F, M, N>
84where
85    F: ComplexField + Copy,
86    N: Dim,
87    M: Dim,
88{
89    /// Storage type used for the residuals. Use `nalgebra::storage::Owned<F, M>`
90    /// if you want to use `VectorN` or `MatrixMN`.
91    type ResidualStorage: RawStorageMut<F, M> + Storage<F, M> + IsContiguous;
92    type ParameterStorage: RawStorageMut<F, N> + Storage<F, N> + IsContiguous + Clone;
93
94    /// Set the stored parameters `$\vec{x}$`.
95    fn set_params(&mut self, x: &Vector<F, N, Self::ParameterStorage>);
96
97    /// Get the current parameter vector `$\vec{x}$`.
98    fn params(&self) -> Vector<F, N, Self::ParameterStorage>;
99
100    /// Compute the residual vector.
101    fn residuals(&self) -> Option<Vector<F, M, Self::ResidualStorage>>;
102
103    /// Compute the Jacobian of the residual vector.
104    fn jacobian(&self) -> Option<SparseJacobian<F>>;
105}