Merge pull request #225 from rust-ndarray/lapack-cholesky

Cholesky factorization by LAPACK

Author: termoshtt

lax/src/cholesky.rs

@@ -1,56 +1,90 @@
 //! Cholesky decomposition
 
 use super::*;
-use crate::{error::*, layout::MatrixLayout};
+use crate::{error::*, layout::*};
 use cauchy::*;
 
 pub trait Cholesky_: Sized {
     /// Cholesky: wrapper of `*potrf`
     ///
     /// **Warning: Only the portion of `a` corresponding to `UPLO` is written.**
-    unsafe fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
+    fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
+
     /// Wrapper of `*potri`
     ///
     /// **Warning: Only the portion of `a` corresponding to `UPLO` is written.**
-    unsafe fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
+    fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
+
     /// Wrapper of `*potrs`
-    unsafe fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self])
-        -> Result<()>;
+    fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self]) -> Result<()>;
 }
 
 macro_rules! impl_cholesky {
     ($scalar:ty, $trf:path, $tri:path, $trs:path) => {
         impl Cholesky_ for $scalar {
-            unsafe fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
+            fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
                 let (n, _) = l.size();
-                $trf(l.lapacke_layout(), uplo as u8, n, a, n).as_lapack_result()?;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
+                let mut info = 0;
+                unsafe {
+                    $trf(uplo as u8, n, a, n, &mut info);
+                }
+                info.as_lapack_result()?;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
                 Ok(())
             }
 
-            unsafe fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
+            fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
                 let (n, _) = l.size();
-                $tri(l.lapacke_layout(), uplo as u8, n, a, l.lda()).as_lapack_result()?;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
+                let mut info = 0;
+                unsafe {
+                    $tri(uplo as u8, n, a, l.lda(), &mut info);
+                }
+                info.as_lapack_result()?;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
                 Ok(())
             }
 
-            unsafe fn solve_cholesky(
+            fn solve_cholesky(
                 l: MatrixLayout,
-                uplo: UPLO,
+                mut uplo: UPLO,
                 a: &[Self],
                 b: &mut [Self],
             ) -> Result<()> {
                 let (n, _) = l.size();
                 let nrhs = 1;
-                let ldb = 1;
-                $trs(l.lapacke_layout(), uplo as u8, n, nrhs, a, l.lda(), b, ldb)
-                    .as_lapack_result()?;
+                let mut info = 0;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    uplo = uplo.t();
+                    for val in b.iter_mut() {
+                        *val = val.conj();
+                    }
+                }
+                unsafe {
+                    $trs(uplo as u8, n, nrhs, a, l.lda(), b, n, &mut info);
+                }
+                info.as_lapack_result()?;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    for val in b.iter_mut() {
+                        *val = val.conj();
+                    }
+                }
                 Ok(())
             }
         }
     };
 } // end macro_rules
 
-impl_cholesky!(f64, lapacke::dpotrf, lapacke::dpotri, lapacke::dpotrs);
-impl_cholesky!(f32, lapacke::spotrf, lapacke::spotri, lapacke::spotrs);
-impl_cholesky!(c64, lapacke::zpotrf, lapacke::zpotri, lapacke::zpotrs);
-impl_cholesky!(c32, lapacke::cpotrf, lapacke::cpotri, lapacke::cpotrs);
+impl_cholesky!(f64, lapack::dpotrf, lapack::dpotri, lapack::dpotrs);
+impl_cholesky!(f32, lapack::spotrf, lapack::spotri, lapack::spotrs);
+impl_cholesky!(c64, lapack::zpotrf, lapack::zpotri, lapack::zpotrs);
+impl_cholesky!(c32, lapack::cpotrf, lapack::cpotri, lapack::cpotrs);

lax/src/layout.rs

@@ -37,6 +37,8 @@
 //! This `S` for a matrix `A` is called "leading dimension of the array A" in LAPACK document, and denoted by `lda`.
 //!
 
+use cauchy::Scalar;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum MatrixLayout {
     C { row: i32, lda: i32 },
@@ -96,3 +98,44 @@ impl MatrixLayout {
         }
     }
 }
+
+/// In-place transpose of a square matrix by keeping F/C layout
+///
+/// Transpose for C-continuous array
+///
+/// ```rust
+/// # use lax::layout::*;
+/// let layout = MatrixLayout::C { row: 2, lda: 2 };
+/// let mut a = vec![1., 2., 3., 4.];
+/// square_transpose(layout, &mut a);
+/// assert_eq!(a, &[1., 3., 2., 4.]);
+/// ```
+///
+/// Transpose for F-continuous array
+///
+/// ```rust
+/// # use lax::layout::*;
+/// let layout = MatrixLayout::F { col: 2, lda: 2 };
+/// let mut a = vec![1., 3., 2., 4.];
+/// square_transpose(layout, &mut a);
+/// assert_eq!(a, &[1., 2., 3., 4.]);
+/// ```
+///
+/// Panics
+/// ------
+/// - If size of `a` and `layout` size mismatch
+///
+pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
+    let (m, n) = layout.size();
+    let n = n as usize;
+    let m = m as usize;
+    assert_eq!(a.len(), n * m);
+    for i in 0..m {
+        for j in (i + 1)..n {
+            let a_ij = a[i * n + j];
+            let a_ji = a[j * m + i];
+            a[i * n + j] = a_ji.conj();
+            a[j * m + i] = a_ij.conj();
+        }
+    }
+}

lax/src/lib.rs

@@ -126,6 +126,15 @@ pub enum UPLO {
     Lower = b'L',
 }
 
+impl UPLO {
+    pub fn t(self) -> Self {
+        match self {
+            UPLO::Upper => UPLO::Lower,
+            UPLO::Lower => UPLO::Upper,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy)]
 #[repr(u8)]
 pub enum Transpose {

ndarray-linalg/src/cholesky.rs

@@ -155,7 +155,7 @@ where
 
     fn invc_into(self) -> Result<Self::Output> {
         let mut a = self.factor;
-        unsafe { A::inv_cholesky(a.square_layout()?, self.uplo, a.as_allocated_mut()?)? };
+        A::inv_cholesky(a.square_layout()?, self.uplo, a.as_allocated_mut()?)?;
         triangular_fill_hermitian(&mut a, self.uplo);
         Ok(a)
     }
@@ -173,14 +173,12 @@ where
     where
         Sb: DataMut<Elem = A>,
     {
-        unsafe {
-            A::solve_cholesky(
-                self.factor.square_layout()?,
-                self.uplo,
-                self.factor.as_allocated()?,
-                b.as_slice_mut().unwrap(),
-            )?
-        };
+        A::solve_cholesky(
+            self.factor.square_layout()?,
+            self.uplo,
+            self.factor.as_allocated()?,
+            b.as_slice_mut().unwrap(),
+        )?;
         Ok(b)
     }
 }
@@ -259,7 +257,7 @@ where
     S: DataMut<Elem = A>,
 {
     fn cholesky_inplace(&mut self, uplo: UPLO) -> Result<&mut Self> {
-        unsafe { A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)? };
+        A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)?;
         Ok(self.into_triangular(uplo))
     }
 }