Add FactorizedCholesky

This provides additional operations based on Cholesky decomposition,
including matrix inverses, determinants, and solutions to linear
systems.

Author: jturner314

src/cholesky.rs

@@ -6,33 +6,45 @@
 //!
 //! # Example
 //!
-//! Calculate `L` in the Cholesky decomposition `A = L * L^H`, where `A` is a
-//! Hermitian (or real symmetric) positive definite matrix:
+//! Using the Cholesky decomposition of `A` for various operations, where `A`
+//! is a Hermitian (or real symmetric) positive definite matrix:
 //!
 //! ```
 //! #[macro_use]
 //! extern crate ndarray;
 //! extern crate ndarray_linalg;
 //!
 //! use ndarray::prelude::*;
-//! use ndarray_linalg::{CholeskyInto, UPLO};
+//! use ndarray_linalg::{Cholesky, UPLO};
 //! # fn main() {
 //!
 //! let a: Array2<f64> = array![
 //!     [  4.,  12., -16.],
 //!     [ 12.,  37., -43.],
 //!     [-16., -43.,  98.]
 //! ];
-//! let lower = a.cholesky_into(UPLO::Lower).unwrap();
-//! assert!(lower.all_close(&array![
+//! let chol_lower = a.cholesky(UPLO::Lower).unwrap();
+//!
+//! // Examine `L`
+//! assert!(chol_lower.factor.all_close(&array![
 //!     [ 2., 0., 0.],
 //!     [ 6., 1., 0.],
 //!     [-8., 5., 3.]
 //! ], 1e-9));
+//!
+//! // Find the determinant of `A`
+//! let det = chol_lower.det();
+//! assert!((det - 36.).abs() < 1e-9);
+//!
+//! // Solve `A * x = b`
+//! let b = array![4., 13., -11.];
+//! let x = chol_lower.solve(&b).unwrap();
+//! assert!(x.all_close(&array![-2., 1., 0.], 1e-9));
 //! # }
 //! ```
 
 use ndarray::*;
+use num_traits::Float;
 
 use super::convert::*;
 use super::error::*;
@@ -42,79 +54,212 @@ use super::types::*;
 
 pub use lapack_traits::UPLO;
 
+/// Cholesky decomposition of Hermitian (or real symmetric) positive definite matrix
+pub struct FactorizedCholesky<S>
+where
+    S: Data,
+{
+    /// `L` from the decomposition `A = L * L^H` or `U` from the decomposition
+    /// `A = U^H * U`.
+    pub factor: ArrayBase<S, Ix2>,
+    /// If this is `UPLO::Lower`, then `self.factor` is `L`. If this is
+    /// `UPLO::Upper`, then `self.factor` is `U`.
+    pub uplo: UPLO,
+}
+
+impl<A, S> FactorizedCholesky<S>
+where
+    A: Scalar,
+    S: DataMut<Elem = A>,
+{
+    /// Returns `L` from the Cholesky decomposition `A = L * L^H`.
+    ///
+    /// If `self.uplo == UPLO::Lower`, then no computations need to be
+    /// performed; otherwise, the conjugate transpose of `self.factor` is
+    /// calculated.
+    pub fn into_lower(self) -> ArrayBase<S, Ix2> {
+        match self.uplo {
+            UPLO::Lower => self.factor,
+            UPLO::Upper => self.factor.reversed_axes().mapv_into(|elem| elem.conj()),
+        }
+    }
+
+    /// Returns `U` from the Cholesky decomposition `A = U^H * U`.
+    ///
+    /// If `self.uplo == UPLO::Upper`, then no computations need to be
+    /// performed; otherwise, the conjugate transpose of `self.factor` is
+    /// calculated.
+    pub fn into_upper(self) -> ArrayBase<S, Ix2> {
+        match self.uplo {
+            UPLO::Lower => self.factor.reversed_axes().mapv_into(|elem| elem.conj()),
+            UPLO::Upper => self.factor,
+        }
+    }
+
+    /// Computes the inverse of the Cholesky-factored matrix.
+    ///
+    /// **Warning: The inverse is stored only in the triangular portion of the
+    /// result matrix corresponding to `self.uplo`!** If you want the other
+    /// triangular portion to be correct, you must fill it in yourself.
+    pub fn into_inverse(mut self) -> Result<ArrayBase<S, Ix2>> {
+        unsafe {
+            A::inv_cholesky(
+                self.factor.square_layout()?,
+                self.uplo,
+                self.factor.as_allocated_mut()?,
+            )?
+        };
+        Ok(self.factor)
+    }
+}
+
+impl<A, S> FactorizedCholesky<S>
+where
+    A: Absolute,
+    S: Data<Elem = A>,
+{
+    /// Computes the natural log of the determinant of the Cholesky-factored
+    /// matrix.
+    pub fn ln_det(&self) -> <A as AssociatedReal>::Real {
+        self.factor
+            .diag()
+            .iter()
+            .map(|elem| elem.abs_sqr().ln())
+            .sum()
+    }
+
+    /// Computes the determinant of the Cholesky-factored matrix.
+    pub fn det(&self) -> <A as AssociatedReal>::Real {
+        self.ln_det().exp()
+    }
+}
+
+impl<A, S> FactorizedCholesky<S>
+where
+    A: Scalar,
+    S: Data<Elem = A>,
+{
+    /// Solves a system of linear equations `A * x = b`, where `self` is the
+    /// Cholesky factorization of `A`, `b` is the argument, and `x` is the
+    /// successful result.
+    pub fn solve<Sb>(&self, b: &ArrayBase<Sb, Ix1>) -> Result<Array1<A>>
+    where
+        Sb: Data<Elem = A>,
+    {
+        let mut b = replicate(b);
+        self.solve_mut(&mut b)?;
+        Ok(b)
+    }
+
+    /// Solves a system of linear equations `A * x = b`, where `self` is the
+    /// Cholesky factorization `A`, `b` is the argument, and `x` is the
+    /// successful result.
+    pub fn solve_into<Sb>(&self, mut b: ArrayBase<Sb, Ix1>) -> Result<ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        self.solve_mut(&mut b)?;
+        Ok(b)
+    }
+
+    /// Solves a system of linear equations `A * x = b`, where `self` is the
+    /// Cholesky factorization of `A`, `b` is the argument, and `x` is the
+    /// successful result. The value of `x` is also assigned to the argument.
+    pub fn solve_mut<'a, Sb>(&self, b: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        unsafe {
+            A::solve_cholesky(
+                self.factor.square_layout()?,
+                self.uplo,
+                self.factor.as_allocated()?,
+                b.as_slice_mut().unwrap(),
+            )?
+        };
+        Ok(b)
+    }
+}
+
 /// Cholesky decomposition of Hermitian (or real symmetric) positive definite matrix reference
-pub trait Cholesky {
-    type Output;
+pub trait Cholesky<S: Data> {
     /// Computes the Cholesky decomposition of the Hermitian (or real
     /// symmetric) positive definite matrix.
     ///
     /// If the argument is `UPLO::Upper`, then computes the decomposition `A =
-    /// U^H * U` using the upper triangular portion of `A` and returns `U`.
-    /// Otherwise, if the argument is `UPLO::Lower`, computes the decomposition
-    /// `A = L * L^H` using the lower triangular portion of `A` and returns
-    /// `L`.
-    fn cholesky(&self, UPLO) -> Result<Self::Output>;
+    /// U^H * U` using the upper triangular portion of `A` and returns the
+    /// factorization containing `U`. Otherwise, if the argument is
+    /// `UPLO::Lower`, computes the decomposition `A = L * L^H` using the lower
+    /// triangular portion of `A` and returns the factorization containing `L`.
+    fn cholesky(&self, UPLO) -> Result<FactorizedCholesky<S>>;
 }
 
 /// Cholesky decomposition of Hermitian (or real symmetric) positive definite matrix
-pub trait CholeskyInto: Sized {
+pub trait CholeskyInto<S: Data> {
     /// Computes the Cholesky decomposition of the Hermitian (or real
     /// symmetric) positive definite matrix.
     ///
     /// If the argument is `UPLO::Upper`, then computes the decomposition `A =
-    /// U^H * U` using the upper triangular portion of `A` and returns `U`.
-    /// Otherwise, if the argument is `UPLO::Lower`, computes the decomposition
-    /// `A = L * L^H` using the lower triangular portion of `A` and returns
-    /// `L`.
-    fn cholesky_into(self, UPLO) -> Result<Self>;
+    /// U^H * U` using the upper triangular portion of `A` and returns the
+    /// factorization containing `U`. Otherwise, if the argument is
+    /// `UPLO::Lower`, computes the decomposition `A = L * L^H` using the lower
+    /// triangular portion of `A` and returns the factorization containing `L`.
+    fn cholesky_into(self, UPLO) -> Result<FactorizedCholesky<S>>;
 }
 
 /// Cholesky decomposition of Hermitian (or real symmetric) positive definite mutable reference of matrix
-pub trait CholeskyMut {
+pub trait CholeskyMut<'a, S: Data> {
     /// Computes the Cholesky decomposition of the Hermitian (or real
-    /// symmetric) positive definite matrix, storing the result in `self` and
-    /// returning it.
+    /// symmetric) positive definite matrix, storing the result (`L` or `U`
+    /// according to the argument) in `self` and returning the factorization.
     ///
     /// If the argument is `UPLO::Upper`, then computes the decomposition `A =
-    /// U^H * U` using the upper triangular portion of `A` and returns `U`.
-    /// Otherwise, if the argument is `UPLO::Lower`, computes the decomposition
-    /// `A = L * L^H` using the lower triangular portion of `A` and returns
-    /// `L`.
-    fn cholesky_mut(&mut self, UPLO) -> Result<&mut Self>;
+    /// U^H * U` using the upper triangular portion of `A` and returns the
+    /// factorization containing `U`. Otherwise, if the argument is
+    /// `UPLO::Lower`, computes the decomposition `A = L * L^H` using the lower
+    /// triangular portion of `A` and returns the factorization containing `L`.
+    fn cholesky_mut(&'a mut self, UPLO) -> Result<FactorizedCholesky<S>>;
 }
 
-impl<A, S> CholeskyInto for ArrayBase<S, Ix2>
+impl<A, S> CholeskyInto<S> for ArrayBase<S, Ix2>
 where
     A: Scalar,
     S: DataMut<Elem = A>,
 {
-    fn cholesky_into(mut self, uplo: UPLO) -> Result<Self> {
+    fn cholesky_into(mut self, uplo: UPLO) -> Result<FactorizedCholesky<S>> {
         unsafe { A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)? };
-        Ok(self.into_triangular(uplo))
+        Ok(FactorizedCholesky {
+            factor: self.into_triangular(uplo),
+            uplo: uplo,
+        })
     }
 }
 
-impl<A, S> CholeskyMut for ArrayBase<S, Ix2>
+impl<'a, A, Si> CholeskyMut<'a, ViewRepr<&'a mut A>> for ArrayBase<Si, Ix2>
 where
     A: Scalar,
-    S: DataMut<Elem = A>,
+    Si: DataMut<Elem = A>,
 {
-    fn cholesky_mut(&mut self, uplo: UPLO) -> Result<&mut Self> {
+    fn cholesky_mut(&'a mut self, uplo: UPLO) -> Result<FactorizedCholesky<ViewRepr<&'a mut A>>> {
         unsafe { A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)? };
-        Ok(self.into_triangular(uplo))
+        Ok(FactorizedCholesky {
+            factor: self.into_triangular(uplo).view_mut(),
+            uplo: uplo,
+        })
     }
 }
 
-impl<A, S> Cholesky for ArrayBase<S, Ix2>
+impl<A, Si> Cholesky<OwnedRepr<A>> for ArrayBase<Si, Ix2>
 where
     A: Scalar,
-    S: Data<Elem = A>,
+    Si: Data<Elem = A>,
 {
-    type Output = Array2<A>;
-
-    fn cholesky(&self, uplo: UPLO) -> Result<Self::Output> {
+    fn cholesky(&self, uplo: UPLO) -> Result<FactorizedCholesky<OwnedRepr<A>>> {
         let mut a = replicate(self);
         unsafe { A::cholesky(a.square_layout()?, uplo, a.as_allocated_mut()?)? };
-        Ok(a.into_triangular(uplo))
+        Ok(FactorizedCholesky {
+            factor: a.into_triangular(uplo),
+            uplo: uplo,
+        })
     }
 }

src/lapack_traits/cholesky.rs

@@ -9,21 +9,40 @@ use types::*;
 use super::{UPLO, into_result};
 
 pub trait Cholesky_: Sized {
+    /// Cholesky: wrapper of `*potrf`
     unsafe fn cholesky(MatrixLayout, UPLO, a: &mut [Self]) -> Result<()>;
+    /// Wrapper of `*potri`
+    unsafe fn inv_cholesky(MatrixLayout, UPLO, a: &mut [Self]) -> Result<()>;
+    /// Wrapper of `*potrs`
+    unsafe fn solve_cholesky(MatrixLayout, UPLO, a: &[Self], b: &mut [Self]) -> Result<()>;
 }
 
 macro_rules! impl_cholesky {
-    ($scalar:ty, $potrf:path) => {
+    ($scalar:ty, $trf:path, $tri:path, $trs:path) => {
 impl Cholesky_ for $scalar {
     unsafe fn cholesky(l: MatrixLayout, uplo: UPLO, mut a: &mut [Self]) -> Result<()> {
         let (n, _) = l.size();
-        let info = $potrf(l.lapacke_layout(), uplo as u8, n, &mut a, n);
+        let info = $trf(l.lapacke_layout(), uplo as u8, n, a, n);
+        into_result(info, ())
+    }
+
+    unsafe fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
+        let (n, _) = l.size();
+        let info = $tri(l.lapacke_layout(), uplo as u8, n, a, l.lda());
+        into_result(info, ())
+    }
+
+    unsafe fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self]) -> Result<()> {
+        let (n, _) = l.size();
+        let nrhs = 1;
+        let ldb = 1;
+        let info = $trs(l.lapacke_layout(), uplo as u8, n, nrhs, a, l.lda(), b, ldb);
         into_result(info, ())
     }
 }
 }} // end macro_rules
 
-impl_cholesky!(f64, c::dpotrf);
-impl_cholesky!(f32, c::spotrf);
-impl_cholesky!(c64, c::zpotrf);
-impl_cholesky!(c32, c::cpotrf);
+impl_cholesky!(f64, c::dpotrf, c::dpotri, c::dpotrs);
+impl_cholesky!(f32, c::spotrf, c::spotri, c::spotrs);
+impl_cholesky!(c64, c::zpotrf, c::zpotri, c::zpotrs);
+impl_cholesky!(c32, c::cpotrf, c::cpotri, c::cpotrs);