diff --git a/src/linalg/cholesky.rs b/src/linalg/cholesky.rs index 7431d666..67baefb1 100644 --- a/src/linalg/cholesky.rs +++ b/src/linalg/cholesky.rs @@ -1,33 +1,31 @@ #[cfg(feature = "serde-serialize")] use serde::{Deserialize, Serialize}; +use num::One; use alga::general::ComplexField; use crate::allocator::Allocator; -use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, SquareMatrix}; +use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, SquareMatrix, Vector}; use crate::constraint::{SameNumberOfRows, ShapeConstraint}; -use crate::dimension::{Dim, DimSub, Dynamic}; +use crate::dimension::{Dim, DimAdd, DimSum, DimDiff, DimSub, Dynamic, U1}; use crate::storage::{Storage, StorageMut}; /// The Cholesky decomposition of a symmetric-definite-positive matrix. #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))] #[cfg_attr( feature = "serde-serialize", - serde(bound( - serialize = "DefaultAllocator: Allocator, - MatrixN: Serialize" - )) + serde(bound(serialize = "DefaultAllocator: Allocator, + MatrixN: Serialize")) )] #[cfg_attr( feature = "serde-serialize", - serde(bound( - deserialize = "DefaultAllocator: Allocator, - MatrixN: Deserialize<'de>" - )) + serde(bound(deserialize = "DefaultAllocator: Allocator, + MatrixN: Deserialize<'de>")) )] #[derive(Clone, Debug)] pub struct Cholesky -where DefaultAllocator: Allocator +where + DefaultAllocator: Allocator, { chol: MatrixN, } @@ -36,10 +34,12 @@ impl Copy for Cholesky where DefaultAllocator: Allocator, MatrixN: Copy, -{} +{ +} impl> Cholesky -where DefaultAllocator: Allocator +where + DefaultAllocator: Allocator, { /// Attempts to compute the Cholesky decomposition of `matrix`. /// @@ -146,10 +146,155 @@ where DefaultAllocator: Allocator self.solve_mut(&mut res); res } + + /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`, + /// performs a rank one update such that we end up with the decomposition of `M + sigma * (v * v.adjoint())`. + #[inline] + pub fn rank_one_update(&mut self, x: &Vector, sigma: N::RealField) + where + S2: Storage, + DefaultAllocator: Allocator, + ShapeConstraint: SameNumberOfRows, + { + Self::xx_rank_one_update(&mut self.chol, &mut x.clone_owned(), sigma) + } + + /// Updates the decomposition such that we get the decomposition of a matrix with the given column `col` in the `j`th position. + /// Since the matrix is square, an identical row will be added in the `j`th row. + pub fn insert_column( + &self, + j: usize, + col: Vector, + ) -> Cholesky> + where + D: DimAdd, + R2: Dim, + S2: Storage, + DefaultAllocator: Allocator, DimSum> + Allocator, + ShapeConstraint: SameNumberOfRows>, + { + let mut col = col.into_owned(); + // for an explanation of the formulas, see https://en.wikipedia.org/wiki/Cholesky_decomposition#Updating_the_decomposition + let n = col.nrows(); + assert_eq!(n, self.chol.nrows() + 1, "The new column must have the size of the factored matrix plus one."); + assert!(j < n, "j needs to be within the bound of the new matrix."); + + // loads the data into a new matrix with an additional jth row/column + let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.add(U1), self.chol.data.shape().1.add(U1)) }; + chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j)); + chol.slice_range_mut(..j, j + 1..).copy_from(&self.chol.slice_range(..j, j..)); + chol.slice_range_mut(j + 1.., ..j).copy_from(&self.chol.slice_range(j.., ..j)); + chol.slice_range_mut(j + 1.., j + 1..).copy_from(&self.chol.slice_range(j.., j..)); + + // update the jth row + let top_left_corner = self.chol.slice_range(..j, ..j); + + let col_j = col[j]; + let (mut new_rowj_adjoint, mut new_colj) = col.rows_range_pair_mut(..j, j + 1..); + assert!(top_left_corner.solve_lower_triangular_mut(&mut new_rowj_adjoint), "Cholesky::insert_column : Unable to solve lower triangular system!"); + + new_rowj_adjoint.adjoint_to(&mut chol.slice_range_mut(j, ..j)); + + // update the center element + let center_element = N::sqrt(col_j - N::from_real(new_rowj_adjoint.norm_squared())); + chol[(j, j)] = center_element; + + // update the jth column + let bottom_left_corner = self.chol.slice_range(j.., ..j); + // new_colj = (col_jplus - bottom_left_corner * new_rowj.adjoint()) / center_element; + new_colj.gemm(-N::one() / center_element, &bottom_left_corner, &new_rowj_adjoint, N::one() / center_element); + chol.slice_range_mut(j + 1.., j).copy_from(&new_colj); + + // update the bottom right corner + let mut bottom_right_corner = chol.slice_range_mut(j + 1.., j + 1..); + Self::xx_rank_one_update(&mut bottom_right_corner, &mut new_colj, -N::RealField::one()); + + Cholesky { chol } + } + + /// Updates the decomposition such that we get the decomposition of the factored matrix with its `j`th column removed. + /// Since the matrix is square, the `j`th row will also be removed. + pub fn remove_column( + &self, + j: usize, + ) -> Cholesky> + where + D: DimSub, + DefaultAllocator: Allocator, DimDiff> + Allocator + { + let n = self.chol.nrows(); + assert!(n > 0, "The matrix needs at least one column."); + assert!(j < n, "j needs to be within the bound of the matrix."); + + // loads the data into a new matrix except for the jth row/column + let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.sub(U1), self.chol.data.shape().1.sub(U1)) }; + chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j)); + chol.slice_range_mut(..j, j..).copy_from(&self.chol.slice_range(..j, j + 1..)); + chol.slice_range_mut(j.., ..j).copy_from(&self.chol.slice_range(j + 1.., ..j)); + chol.slice_range_mut(j.., j..).copy_from(&self.chol.slice_range(j + 1.., j + 1..)); + + // updates the bottom right corner + let mut bottom_right_corner = chol.slice_range_mut(j.., j..); + let mut workspace = self.chol.column(j).clone_owned(); + let mut old_colj = workspace.rows_range_mut(j + 1..); + Self::xx_rank_one_update(&mut bottom_right_corner, &mut old_colj, N::RealField::one()); + + Cholesky { chol } + } + + /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `x`, + /// performs a rank one update such that we end up with the decomposition of `M + sigma * (x * x.adjoint())`. + /// + /// This helper method is called by `rank_one_update` but also `insert_column` and `remove_column` + /// where it is used on a square slice of the decomposition + fn xx_rank_one_update(chol : &mut Matrix, x: &mut Vector, sigma: N::RealField) + where + //N: ComplexField, + Dm: Dim, + Rx: Dim, + Sm: StorageMut, + Sx: StorageMut, + { + // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html + let n = x.nrows(); + assert_eq!( + n, + chol.nrows(), + "The input vector must be of the same size as the factorized matrix." + ); + + let mut beta = crate::one::(); + + for j in 0..n { + // updates the diagonal + let diag = N::real(unsafe { *chol.get_unchecked((j, j)) }); + let diag2 = diag * diag; + let xj = unsafe { *x.get_unchecked(j) }; + let sigma_xj2 = sigma * N::modulus_squared(xj); + let gamma = diag2 * beta + sigma_xj2; + let new_diag = (diag2 + sigma_xj2 / beta).sqrt(); + unsafe { *chol.get_unchecked_mut((j, j)) = N::from_real(new_diag) }; + beta += sigma_xj2 / diag2; + // updates the terms of L + let mut xjplus = x.rows_range_mut(j + 1..); + let mut col_j = chol.slice_range_mut(j + 1.., j); + // temp_jplus -= (wj / N::from_real(diag)) * col_j; + xjplus.axpy(-xj / N::from_real(diag), &col_j, N::one()); + if gamma != crate::zero::() { + // col_j = N::from_real(nljj / diag) * col_j + (N::from_real(nljj * sigma / gamma) * N::conjugate(wj)) * temp_jplus; + col_j.axpy( + N::from_real(new_diag * sigma / gamma) * N::conjugate(xj), + &xjplus, + N::from_real(new_diag / diag), + ); + } + } + } } impl, S: Storage> SquareMatrix -where DefaultAllocator: Allocator +where + DefaultAllocator: Allocator, { /// Attempts to compute the Cholesky decomposition of this matrix. /// diff --git a/src/linalg/solve.rs b/src/linalg/solve.rs index f10b1d00..a6b9196f 100644 --- a/src/linalg/solve.rs +++ b/src/linalg/solve.rs @@ -15,7 +15,7 @@ impl> SquareMatrix { b: &Matrix, ) -> Option> where - S2: StorageMut, + S2: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { @@ -35,7 +35,7 @@ impl> SquareMatrix { b: &Matrix, ) -> Option> where - S2: StorageMut, + S2: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { @@ -191,7 +191,7 @@ impl> SquareMatrix { b: &Matrix, ) -> Option> where - S2: StorageMut, + S2: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { @@ -211,7 +211,7 @@ impl> SquareMatrix { b: &Matrix, ) -> Option> where - S2: StorageMut, + S2: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { @@ -273,7 +273,7 @@ impl> SquareMatrix { b: &Matrix, ) -> Option> where - S2: StorageMut, + S2: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { @@ -293,7 +293,7 @@ impl> SquareMatrix { b: &Matrix, ) -> Option> where - S2: StorageMut, + S2: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { diff --git a/tests/linalg/cholesky.rs b/tests/linalg/cholesky.rs index cefc2630..a89802b2 100644 --- a/tests/linalg/cholesky.rs +++ b/tests/linalg/cholesky.rs @@ -1,6 +1,5 @@ #![cfg(all(feature = "arbitrary", feature = "debug"))] - macro_rules! gen_tests( ($module: ident, $scalar: ty) => { mod $module { @@ -78,6 +77,58 @@ macro_rules! gen_tests( id1.is_identity(1.0e-7) && id2.is_identity(1.0e-7) } + + fn cholesky_rank_one_update(_n: usize) -> bool { + let mut m = RandomSDP::new(U4, || random::<$scalar>().0).unwrap(); + let x = Vector4::<$scalar>::new_random().map(|e| e.0); + + // this is dirty but $scalar is not a scalar type (its a Rand) in this file + let zero = random::<$scalar>().0 * 0.; + let one = zero + 1.; + let sigma = random::(); // needs to be a real + let sigma_scalar = zero + sigma; + + // updates cholesky decomposition and reconstructs m updated + let mut chol = m.clone().cholesky().unwrap(); + chol.rank_one_update(&x, sigma); + let m_chol_updated = chol.l() * chol.l().adjoint(); + + // updates m manually + m.gerc(sigma_scalar, &x, &x, one); // m += sigma * x * x.adjoint() + + relative_eq!(m, m_chol_updated, epsilon = 1.0e-7) + } + + fn cholesky_insert_column(n: usize) -> bool { + let n = n.max(1).min(10); + let j = random::() % n; + let m_updated = RandomSDP::new(Dynamic::new(n), || random::<$scalar>().0).unwrap(); + + // build m and col from m_updated + let col = m_updated.column(j); + let m = m_updated.clone().remove_column(j).remove_row(j); + + // remove column from cholesky decomposition and rebuild m + let chol = m.clone().cholesky().unwrap().insert_column(j, col); + let m_chol_updated = chol.l() * chol.l().adjoint(); + + relative_eq!(m_updated, m_chol_updated, epsilon = 1.0e-7) + } + + fn cholesky_remove_column(n: usize) -> bool { + let n = n.max(1).min(10); + let j = random::() % n; + let m = RandomSDP::new(Dynamic::new(n), || random::<$scalar>().0).unwrap(); + + // remove column from cholesky decomposition and rebuild m + let chol = m.clone().cholesky().unwrap().remove_column(j); + let m_chol_updated = chol.l() * chol.l().adjoint(); + + // remove column from m + let m_updated = m.remove_column(j).remove_row(j); + + relative_eq!(m_updated, m_chol_updated, epsilon = 1.0e-7) + } } } }