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#[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 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 pub fn pack_dirty(matrix: OMatrix<T, D, D>) -> Self {
82 Cholesky { chol: matrix }
83 }
84
85 pub fn unpack(mut self) -> OMatrix<T, D, D> {
88 self.chol.fill_upper_triangle(T::zero(), 1);
89 self.chol
90 }
91
92 pub fn unpack_dirty(self) -> OMatrix<T, D, D> {
98 self.chol
99 }
100
101 #[must_use]
104 pub fn l(&self) -> OMatrix<T, D, D> {
105 self.chol.lower_triangle()
106 }
107
108 #[must_use]
114 pub fn l_dirty(&self) -> &OMatrix<T, D, D> {
115 &self.chol
116 }
117
118 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 #[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 #[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 #[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 pub fn new(matrix: OMatrix<T, D, D>) -> Option<Self> {
175 Self::new_internal(matrix, None)
176 }
177
178 pub fn new_with_substitute(matrix: OMatrix<T, D, D>, substitute: T) -> Option<Self> {
195 Self::new_internal(matrix, Some(substitute))
196 }
197
198 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 return None;
239 }
240
241 Some(Cholesky { chol: matrix })
242 }
243
244 #[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 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 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 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 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 let center_element = T::sqrt(col_j - T::from_real(new_rowj_adjoint.norm_squared()));
309 chol[(j, j)] = center_element.clone();
310
311 let bottom_left_corner = self.chol.slice_range(j.., ..j);
313 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 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 #[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 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 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 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 Dm: Dim,
381 Rx: Dim,
382 Sm: StorageMut<T, Dm, Dm>,
383 Sx: StorageMut<T, Rx, U1>,
384 {
385 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 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 let mut xjplus = x.rows_range_mut(j + 1..);
407 let mut col_j = chol.slice_range_mut(j + 1.., j);
408 xjplus.axpy(-xj.clone() / T::from_real(diag.clone()), &col_j, T::one());
410 if gamma != crate::zero::<T::RealField>() {
411 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}