nalgebra/linalg/
cholesky.rs

1#[cfg(feature = "serde-serialize-no-std")]
2use serde::{Deserialize, Serialize};
3
4use num::One;
5use simba::scalar::ComplexField;
6use simba::simd::SimdComplexField;
7
8use crate::allocator::Allocator;
9use crate::base::{Const, DefaultAllocator, Matrix, OMatrix, Vector};
10use crate::constraint::{SameNumberOfRows, ShapeConstraint};
11use crate::dimension::{Dim, DimAdd, DimDiff, DimSub, DimSum, U1};
12use crate::storage::{Storage, StorageMut};
13
14/// The Cholesky decomposition of a symmetric-definite-positive matrix.
15#[cfg_attr(feature = "serde-serialize-no-std", derive(Serialize, Deserialize))]
16#[cfg_attr(
17    feature = "serde-serialize-no-std",
18    serde(bound(serialize = "DefaultAllocator: Allocator<T, D>,
19         OMatrix<T, D, D>: Serialize"))
20)]
21#[cfg_attr(
22    feature = "serde-serialize-no-std",
23    serde(bound(deserialize = "DefaultAllocator: Allocator<T, D>,
24         OMatrix<T, D, D>: Deserialize<'de>"))
25)]
26#[derive(Clone, Debug)]
27pub struct Cholesky<T: SimdComplexField, D: Dim>
28where
29    DefaultAllocator: Allocator<T, D, D>,
30{
31    chol: OMatrix<T, D, D>,
32}
33
34impl<T: SimdComplexField, D: Dim> Copy for Cholesky<T, D>
35where
36    DefaultAllocator: Allocator<T, D, D>,
37    OMatrix<T, D, D>: Copy,
38{
39}
40
41impl<T: SimdComplexField, D: Dim> Cholesky<T, D>
42where
43    DefaultAllocator: Allocator<T, D, D>,
44{
45    /// Computes the Cholesky decomposition of `matrix` without checking that the matrix is definite-positive.
46    ///
47    /// If the input matrix is not definite-positive, the decomposition may contain trash values (Inf, NaN, etc.)
48    pub fn new_unchecked(mut matrix: OMatrix<T, D, D>) -> Self {
49        assert!(matrix.is_square(), "The input matrix must be square.");
50
51        let n = matrix.nrows();
52
53        for j in 0..n {
54            for k in 0..j {
55                let factor = unsafe { -matrix.get_unchecked((j, k)).clone() };
56
57                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
58                let mut col_j = col_j.rows_range_mut(j..);
59                let col_k = col_k.rows_range(j..);
60                col_j.axpy(factor.simd_conjugate(), &col_k, T::one());
61            }
62
63            let diag = unsafe { matrix.get_unchecked((j, j)).clone() };
64            let denom = diag.simd_sqrt();
65
66            unsafe {
67                *matrix.get_unchecked_mut((j, j)) = denom.clone();
68            }
69
70            let mut col = matrix.slice_range_mut(j + 1.., j);
71            col /= denom;
72        }
73
74        Cholesky { chol: matrix }
75    }
76
77    /// Uses the given matrix as-is without any checks or modifications as the
78    /// Cholesky decomposition.
79    ///
80    /// It is up to the user to ensure all invariants hold.
81    pub fn pack_dirty(matrix: OMatrix<T, D, D>) -> Self {
82        Cholesky { chol: matrix }
83    }
84
85    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
86    /// upper-triangular part filled with zeros.
87    pub fn unpack(mut self) -> OMatrix<T, D, D> {
88        self.chol.fill_upper_triangle(T::zero(), 1);
89        self.chol
90    }
91
92    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
93    /// its strict upper-triangular part.
94    ///
95    /// The values of the strict upper-triangular part are garbage and should be ignored by further
96    /// computations.
97    pub fn unpack_dirty(self) -> OMatrix<T, D, D> {
98        self.chol
99    }
100
101    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
102    /// uppen-triangular part filled with zeros.
103    #[must_use]
104    pub fn l(&self) -> OMatrix<T, D, D> {
105        self.chol.lower_triangle()
106    }
107
108    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
109    /// its strict upper-triangular part.
110    ///
111    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
112    /// part are garbage and should be ignored by further computations.
113    #[must_use]
114    pub fn l_dirty(&self) -> &OMatrix<T, D, D> {
115        &self.chol
116    }
117
118    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
119    ///
120    /// The result is stored on `b`.
121    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<T, R2, C2, S2>)
122    where
123        S2: StorageMut<T, R2, C2>,
124        ShapeConstraint: SameNumberOfRows<R2, D>,
125    {
126        self.chol.solve_lower_triangular_unchecked_mut(b);
127        self.chol.ad_solve_lower_triangular_unchecked_mut(b);
128    }
129
130    /// Returns the solution of the system `self * x = b` where `self` is the decomposed matrix and
131    /// `x` the unknown.
132    #[must_use = "Did you mean to use solve_mut()?"]
133    pub fn solve<R2: Dim, C2: Dim, S2>(&self, b: &Matrix<T, R2, C2, S2>) -> OMatrix<T, R2, C2>
134    where
135        S2: Storage<T, R2, C2>,
136        DefaultAllocator: Allocator<T, R2, C2>,
137        ShapeConstraint: SameNumberOfRows<R2, D>,
138    {
139        let mut res = b.clone_owned();
140        self.solve_mut(&mut res);
141        res
142    }
143
144    /// Computes the inverse of the decomposed matrix.
145    #[must_use]
146    pub fn inverse(&self) -> OMatrix<T, D, D> {
147        let shape = self.chol.shape_generic();
148        let mut res = OMatrix::identity_generic(shape.0, shape.1);
149
150        self.solve_mut(&mut res);
151        res
152    }
153
154    /// Computes the determinant of the decomposed matrix.
155    #[must_use]
156    pub fn determinant(&self) -> T::SimdRealField {
157        let dim = self.chol.nrows();
158        let mut prod_diag = T::one();
159        for i in 0..dim {
160            prod_diag *= unsafe { self.chol.get_unchecked((i, i)).clone() };
161        }
162        prod_diag.simd_modulus_squared()
163    }
164}
165
166impl<T: ComplexField, D: Dim> Cholesky<T, D>
167where
168    DefaultAllocator: Allocator<T, D, D>,
169{
170    /// Attempts to compute the Cholesky decomposition of `matrix`.
171    ///
172    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
173    /// to be symmetric and only the lower-triangular part is read.
174    pub fn new(matrix: OMatrix<T, D, D>) -> Option<Self> {
175        Self::new_internal(matrix, None)
176    }
177
178    /// Attempts to approximate the Cholesky decomposition of `matrix` by
179    /// replacing non-positive values on the diagonals during the decomposition
180    /// with the given `substitute`.
181    ///
182    /// [`try_sqrt`](ComplexField::try_sqrt) will be applied to the `substitute`
183    /// when it has to be used.
184    ///
185    /// If your input matrix results only in positive values on the diagonals
186    /// during the decomposition, `substitute` is unused and the result is just
187    /// the same as if you used [`new`](Cholesky::new).
188    ///
189    /// This method allows to compensate for matrices with very small or even
190    /// negative values due to numerical errors but necessarily results in only
191    /// an approximation: it is basically a hack. If you don't specifically need
192    /// Cholesky, it may be better to consider alternatives like the
193    /// [`LU`](crate::linalg::LU) decomposition/factorization.
194    pub fn new_with_substitute(matrix: OMatrix<T, D, D>, substitute: T) -> Option<Self> {
195        Self::new_internal(matrix, Some(substitute))
196    }
197
198    /// Common implementation for `new` and `new_with_substitute`.
199    fn new_internal(mut matrix: OMatrix<T, D, D>, substitute: Option<T>) -> Option<Self> {
200        assert!(matrix.is_square(), "The input matrix must be square.");
201
202        let n = matrix.nrows();
203
204        for j in 0..n {
205            for k in 0..j {
206                let factor = unsafe { -matrix.get_unchecked((j, k)).clone() };
207
208                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
209                let mut col_j = col_j.rows_range_mut(j..);
210                let col_k = col_k.rows_range(j..);
211
212                col_j.axpy(factor.conjugate(), &col_k, T::one());
213            }
214
215            let sqrt_denom = |v: T| {
216                if v.is_zero() {
217                    return None;
218                }
219                v.try_sqrt()
220            };
221
222            let diag = unsafe { matrix.get_unchecked((j, j)).clone() };
223
224            if let Some(denom) =
225                sqrt_denom(diag).or_else(|| substitute.clone().and_then(sqrt_denom))
226            {
227                unsafe {
228                    *matrix.get_unchecked_mut((j, j)) = denom.clone();
229                }
230
231                let mut col = matrix.slice_range_mut(j + 1.., j);
232                col /= denom;
233                continue;
234            }
235
236            // The diagonal element is either zero or its square root could not
237            // be taken (e.g. for negative real numbers).
238            return None;
239        }
240
241        Some(Cholesky { chol: matrix })
242    }
243
244    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
245    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (v * v.adjoint())`.
246    #[inline]
247    pub fn rank_one_update<R2: Dim, S2>(&mut self, x: &Vector<T, R2, S2>, sigma: T::RealField)
248    where
249        S2: Storage<T, R2, U1>,
250        DefaultAllocator: Allocator<T, R2, U1>,
251        ShapeConstraint: SameNumberOfRows<R2, D>,
252    {
253        Self::xx_rank_one_update(&mut self.chol, &mut x.clone_owned(), sigma)
254    }
255
256    /// Updates the decomposition such that we get the decomposition of a matrix with the given column `col` in the `j`th position.
257    /// Since the matrix is square, an identical row will be added in the `j`th row.
258    pub fn insert_column<R2, S2>(
259        &self,
260        j: usize,
261        col: Vector<T, R2, S2>,
262    ) -> Cholesky<T, DimSum<D, U1>>
263    where
264        D: DimAdd<U1>,
265        R2: Dim,
266        S2: Storage<T, R2, U1>,
267        DefaultAllocator: Allocator<T, DimSum<D, U1>, DimSum<D, U1>> + Allocator<T, R2>,
268        ShapeConstraint: SameNumberOfRows<R2, DimSum<D, U1>>,
269    {
270        let mut col = col.into_owned();
271        // for an explanation of the formulas, see https://en.wikipedia.org/wiki/Cholesky_decomposition#Updating_the_decomposition
272        let n = col.nrows();
273        assert_eq!(
274            n,
275            self.chol.nrows() + 1,
276            "The new column must have the size of the factored matrix plus one."
277        );
278        assert!(j < n, "j needs to be within the bound of the new matrix.");
279
280        // loads the data into a new matrix with an additional jth row/column
281        // TODO: would it be worth it to avoid the zero-initialization?
282        let mut chol = Matrix::zeros_generic(
283            self.chol.shape_generic().0.add(Const::<1>),
284            self.chol.shape_generic().1.add(Const::<1>),
285        );
286        chol.slice_range_mut(..j, ..j)
287            .copy_from(&self.chol.slice_range(..j, ..j));
288        chol.slice_range_mut(..j, j + 1..)
289            .copy_from(&self.chol.slice_range(..j, j..));
290        chol.slice_range_mut(j + 1.., ..j)
291            .copy_from(&self.chol.slice_range(j.., ..j));
292        chol.slice_range_mut(j + 1.., j + 1..)
293            .copy_from(&self.chol.slice_range(j.., j..));
294
295        // update the jth row
296        let top_left_corner = self.chol.slice_range(..j, ..j);
297
298        let col_j = col[j].clone();
299        let (mut new_rowj_adjoint, mut new_colj) = col.rows_range_pair_mut(..j, j + 1..);
300        assert!(
301            top_left_corner.solve_lower_triangular_mut(&mut new_rowj_adjoint),
302            "Cholesky::insert_column : Unable to solve lower triangular system!"
303        );
304
305        new_rowj_adjoint.adjoint_to(&mut chol.slice_range_mut(j, ..j));
306
307        // update the center element
308        let center_element = T::sqrt(col_j - T::from_real(new_rowj_adjoint.norm_squared()));
309        chol[(j, j)] = center_element.clone();
310
311        // update the jth column
312        let bottom_left_corner = self.chol.slice_range(j.., ..j);
313        // new_colj = (col_jplus - bottom_left_corner * new_rowj.adjoint()) / center_element;
314        new_colj.gemm(
315            -T::one() / center_element.clone(),
316            &bottom_left_corner,
317            &new_rowj_adjoint,
318            T::one() / center_element,
319        );
320        chol.slice_range_mut(j + 1.., j).copy_from(&new_colj);
321
322        // update the bottom right corner
323        let mut bottom_right_corner = chol.slice_range_mut(j + 1.., j + 1..);
324        Self::xx_rank_one_update(
325            &mut bottom_right_corner,
326            &mut new_colj,
327            -T::RealField::one(),
328        );
329
330        Cholesky { chol }
331    }
332
333    /// Updates the decomposition such that we get the decomposition of the factored matrix with its `j`th column removed.
334    /// Since the matrix is square, the `j`th row will also be removed.
335    #[must_use]
336    pub fn remove_column(&self, j: usize) -> Cholesky<T, DimDiff<D, U1>>
337    where
338        D: DimSub<U1>,
339        DefaultAllocator: Allocator<T, DimDiff<D, U1>, DimDiff<D, U1>> + Allocator<T, D>,
340    {
341        let n = self.chol.nrows();
342        assert!(n > 0, "The matrix needs at least one column.");
343        assert!(j < n, "j needs to be within the bound of the matrix.");
344
345        // loads the data into a new matrix except for the jth row/column
346        // TODO: would it be worth it to avoid this zero initialization?
347        let mut chol = Matrix::zeros_generic(
348            self.chol.shape_generic().0.sub(Const::<1>),
349            self.chol.shape_generic().1.sub(Const::<1>),
350        );
351        chol.slice_range_mut(..j, ..j)
352            .copy_from(&self.chol.slice_range(..j, ..j));
353        chol.slice_range_mut(..j, j..)
354            .copy_from(&self.chol.slice_range(..j, j + 1..));
355        chol.slice_range_mut(j.., ..j)
356            .copy_from(&self.chol.slice_range(j + 1.., ..j));
357        chol.slice_range_mut(j.., j..)
358            .copy_from(&self.chol.slice_range(j + 1.., j + 1..));
359
360        // updates the bottom right corner
361        let mut bottom_right_corner = chol.slice_range_mut(j.., j..);
362        let mut workspace = self.chol.column(j).clone_owned();
363        let mut old_colj = workspace.rows_range_mut(j + 1..);
364        Self::xx_rank_one_update(&mut bottom_right_corner, &mut old_colj, T::RealField::one());
365
366        Cholesky { chol }
367    }
368
369    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `x`,
370    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (x * x.adjoint())`.
371    ///
372    /// This helper method is called by `rank_one_update` but also `insert_column` and `remove_column`
373    /// where it is used on a square slice of the decomposition
374    fn xx_rank_one_update<Dm, Sm, Rx, Sx>(
375        chol: &mut Matrix<T, Dm, Dm, Sm>,
376        x: &mut Vector<T, Rx, Sx>,
377        sigma: T::RealField,
378    ) where
379        //T: ComplexField,
380        Dm: Dim,
381        Rx: Dim,
382        Sm: StorageMut<T, Dm, Dm>,
383        Sx: StorageMut<T, Rx, U1>,
384    {
385        // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html
386        let n = x.nrows();
387        assert_eq!(
388            n,
389            chol.nrows(),
390            "The input vector must be of the same size as the factorized matrix."
391        );
392
393        let mut beta = crate::one::<T::RealField>();
394
395        for j in 0..n {
396            // updates the diagonal
397            let diag = T::real(unsafe { chol.get_unchecked((j, j)).clone() });
398            let diag2 = diag.clone() * diag.clone();
399            let xj = unsafe { x.get_unchecked(j).clone() };
400            let sigma_xj2 = sigma.clone() * T::modulus_squared(xj.clone());
401            let gamma = diag2.clone() * beta.clone() + sigma_xj2.clone();
402            let new_diag = (diag2.clone() + sigma_xj2.clone() / beta.clone()).sqrt();
403            unsafe { *chol.get_unchecked_mut((j, j)) = T::from_real(new_diag.clone()) };
404            beta += sigma_xj2 / diag2;
405            // updates the terms of L
406            let mut xjplus = x.rows_range_mut(j + 1..);
407            let mut col_j = chol.slice_range_mut(j + 1.., j);
408            // temp_jplus -= (wj / T::from_real(diag)) * col_j;
409            xjplus.axpy(-xj.clone() / T::from_real(diag.clone()), &col_j, T::one());
410            if gamma != crate::zero::<T::RealField>() {
411                // col_j = T::from_real(nljj / diag) * col_j  + (T::from_real(nljj * sigma / gamma) * T::conjugate(wj)) * temp_jplus;
412                col_j.axpy(
413                    T::from_real(new_diag.clone() * sigma.clone() / gamma) * T::conjugate(xj),
414                    &xjplus,
415                    T::from_real(new_diag / diag),
416                );
417            }
418        }
419    }
420}