square_transpose

Author: termoshtt

lax/src/cholesky.rs

@@ -1,7 +1,7 @@
 //! Cholesky decomposition
 
 use super::*;
-use crate::{error::*, layout::MatrixLayout};
+use crate::{error::*, layout::*};
 use cauchy::*;
 
 pub trait Cholesky_: Sized {
@@ -24,45 +24,48 @@ macro_rules! impl_cholesky {
         impl Cholesky_ for $scalar {
             fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
                 let (n, _) = l.size();
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
                 let mut info = 0;
-                let uplo = match l {
-                    MatrixLayout::F { .. } => uplo,
-                    MatrixLayout::C { .. } => uplo.t(),
-                };
                 unsafe {
                     $trf(uplo as u8, n, a, n, &mut info);
                 }
                 info.as_lapack_result()?;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
                 Ok(())
             }
 
             fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
                 let (n, _) = l.size();
+                if matches!(l, MatrixLayout::C { .. }) {
+                    square_transpose(l, a);
+                }
                 let mut info = 0;
-                let uplo = match l {
-                    MatrixLayout::F { .. } => uplo,
-                    MatrixLayout::C { .. } => uplo.t(),
-                };
                 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(())
             }
 
             fn solve_cholesky(
                 l: MatrixLayout,
-                uplo: UPLO,
+                mut uplo: UPLO,
                 a: &[Self],
                 b: &mut [Self],
             ) -> Result<()> {
                 let (n, _) = l.size();
                 let nrhs = 1;
-                let uplo = match l {
-                    MatrixLayout::F { .. } => uplo,
-                    MatrixLayout::C { .. } => uplo.t(),
-                };
                 let mut info = 0;
+                if matches!(l, MatrixLayout::C { .. }) {
+                    uplo = uplo.t();
+                }
                 unsafe {
                     $trs(uplo as u8, n, nrhs, a, l.lda(), b, n, &mut info);
                 }

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();
+        }
+    }
+}