Merge Solveh_ into Lapack trait

Author: termoshtt

lax/src/lib.rs

@@ -123,9 +123,7 @@ pub type Pivot = Vec<i32>;
 
 #[cfg_attr(doc, katexit::katexit)]
 /// Trait for primitive types which implements LAPACK subroutines
-pub trait Lapack:
-    OperatorNorm_ + Solveh_ + Cholesky_ + Triangular_ + Tridiagonal_ + Rcond_
-{
+pub trait Lapack: OperatorNorm_ + Cholesky_ + Triangular_ + Tridiagonal_ + Rcond_ {
     /// Compute right eigenvalue and eigenvectors for a general matrix
     fn eig(
         calc_v: bool,
@@ -217,6 +215,30 @@ pub trait Lapack:
 
     /// Solve linear equations $Ax = b$ using the output of LU-decomposition
     fn solve(l: MatrixLayout, t: Transpose, a: &[Self], p: &Pivot, b: &mut [Self]) -> Result<()>;
+
+    /// Factorize symmetric/Hermitian matrix using Bunch-Kaufman diagonal pivoting method
+    ///
+    ///
+    /// For a given symmetric matrix $A$,
+    /// this method factorizes $A = U^T D U$ or $A = L D L^T$ where
+    ///
+    /// - $U$ (or $L$) are is a product of permutation and unit upper (lower) triangular matrices
+    /// - $D$ is symmetric and block diagonal with 1-by-1 and 2-by-2 diagonal blocks.
+    ///
+    /// This takes two-step approach based in LAPACK:
+    ///
+    /// 1. Factorize given matrix $A$ into upper ($U$) or lower ($L$) form with diagonal matrix $D$
+    /// 2. Then solve linear equation $Ax = b$, and/or calculate inverse matrix $A^{-1}$
+    ///
+    /// [BK]: https://doi.org/10.2307/2005787
+    ///
+    fn bk(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<Pivot>;
+
+    /// Compute inverse matrix $A^{-1}$ of symmetric/Hermitian matrix using factroized result
+    fn invh(l: MatrixLayout, uplo: UPLO, a: &mut [Self], ipiv: &Pivot) -> Result<()>;
+
+    /// Solve symmetric/Hermitian linear equation $Ax = b$ using factroized result
+    fn solveh(l: MatrixLayout, uplo: UPLO, a: &[Self], ipiv: &Pivot, b: &mut [Self]) -> Result<()>;
 }
 
 macro_rules! impl_lapack {
@@ -335,6 +357,29 @@ macro_rules! impl_lapack {
                 use solve::*;
                 SolveImpl::solve(l, t, a, p, b)
             }
+
+            fn bk(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<Pivot> {
+                use solveh::*;
+                let work = BkWork::<$s>::new(l)?;
+                work.eval(uplo, a)
+            }
+
+            fn invh(l: MatrixLayout, uplo: UPLO, a: &mut [Self], ipiv: &Pivot) -> Result<()> {
+                use solveh::*;
+                let mut work = InvhWork::<$s>::new(l)?;
+                work.calc(uplo, a, ipiv)
+            }
+
+            fn solveh(
+                l: MatrixLayout,
+                uplo: UPLO,
+                a: &[Self],
+                ipiv: &Pivot,
+                b: &mut [Self],
+            ) -> Result<()> {
+                use solveh::*;
+                SolvehImpl::solveh(l, uplo, a, ipiv, b)
+            }
         }
     };
 }

lax/src/solveh.rs

@@ -8,6 +8,15 @@ pub struct BkWork<T: Scalar> {
     pub ipiv: Vec<MaybeUninit<i32>>,
 }
 
+/// Factorize symmetric/Hermitian matrix using Bunch-Kaufman diagonal pivoting method
+///
+/// LAPACK correspondance
+/// ----------------------
+///
+/// | f32    | f64    | c32    | c64    |
+/// |:-------|:-------|:-------|:-------|
+/// | ssytrf | dsytrf | chetrf | zhetrf |
+///
 pub trait BkWorkImpl: Sized {
     type Elem: Scalar;
     fn new(l: MatrixLayout) -> Result<Self>;
@@ -80,6 +89,15 @@ pub struct InvhWork<T: Scalar> {
     pub work: Vec<MaybeUninit<T>>,
 }
 
+/// Compute inverse matrix of symmetric/Hermitian matrix
+///
+/// LAPACK correspondance
+/// ----------------------
+///
+/// | f32    | f64    | c32    | c64    |
+/// |:-------|:-------|:-------|:-------|
+/// | ssytri | dsytri | chetri | zhetri |
+///
 pub trait InvhWorkImpl: Sized {
     type Elem;
     fn new(layout: MatrixLayout) -> Result<Self>;
@@ -122,6 +140,15 @@ impl_invh_work!(c32, lapack_sys::chetri_);
 impl_invh_work!(f64, lapack_sys::dsytri_);
 impl_invh_work!(f32, lapack_sys::ssytri_);
 
+/// Solve symmetric/Hermitian linear equation
+///
+/// LAPACK correspondance
+/// ----------------------
+///
+/// | f32    | f64    | c32    | c64    |
+/// |:-------|:-------|:-------|:-------|
+/// | ssytrs | dsytrs | chetrs | zhetrs |
+///
 pub trait SolvehImpl: Scalar {
     fn solveh(l: MatrixLayout, uplo: UPLO, a: &[Self], ipiv: &Pivot, b: &mut [Self]) -> Result<()>;
 }
@@ -162,174 +189,3 @@ impl_solveh_!(c64, lapack_sys::zhetrs_);
 impl_solveh_!(c32, lapack_sys::chetrs_);
 impl_solveh_!(f64, lapack_sys::dsytrs_);
 impl_solveh_!(f32, lapack_sys::ssytrs_);
-
-#[cfg_attr(doc, katexit::katexit)]
-/// Solve symmetric/hermite indefinite linear problem using the [Bunch-Kaufman diagonal pivoting method][BK].
-///
-/// For a given symmetric matrix $A$,
-/// this method factorizes $A = U^T D U$ or $A = L D L^T$ where
-///
-/// - $U$ (or $L$) are is a product of permutation and unit upper (lower) triangular matrices
-/// - $D$ is symmetric and block diagonal with 1-by-1 and 2-by-2 diagonal blocks.
-///
-/// This takes two-step approach based in LAPACK:
-///
-/// 1. Factorize given matrix $A$ into upper ($U$) or lower ($L$) form with diagonal matrix $D$
-/// 2. Then solve linear equation $Ax = b$, and/or calculate inverse matrix $A^{-1}$
-///
-/// [BK]: https://doi.org/10.2307/2005787
-///
-pub trait Solveh_: Sized {
-    /// Factorize input matrix using Bunch-Kaufman diagonal pivoting method
-    ///
-    /// LAPACK correspondance
-    /// ----------------------
-    ///
-    /// | f32    | f64    | c32    | c64    |
-    /// |:-------|:-------|:-------|:-------|
-    /// | ssytrf | dsytrf | chetrf | zhetrf |
-    ///
-    fn bk(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<Pivot>;
-
-    /// Compute inverse matrix $A^{-1}$ from factroized result
-    ///
-    /// LAPACK correspondance
-    /// ----------------------
-    ///
-    /// | f32    | f64    | c32    | c64    |
-    /// |:-------|:-------|:-------|:-------|
-    /// | ssytri | dsytri | chetri | zhetri |
-    ///
-    fn invh(l: MatrixLayout, uplo: UPLO, a: &mut [Self], ipiv: &Pivot) -> Result<()>;
-
-    /// Solve linear equation $Ax = b$ using factroized result
-    ///
-    /// LAPACK correspondance
-    /// ----------------------
-    ///
-    /// | f32    | f64    | c32    | c64    |
-    /// |:-------|:-------|:-------|:-------|
-    /// | ssytrs | dsytrs | chetrs | zhetrs |
-    ///
-    fn solveh(l: MatrixLayout, uplo: UPLO, a: &[Self], ipiv: &Pivot, b: &mut [Self]) -> Result<()>;
-}
-
-macro_rules! impl_solveh {
-    ($scalar:ty, $trf:path, $tri:path, $trs:path) => {
-        impl Solveh_ for $scalar {
-            fn bk(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<Pivot> {
-                let (n, _) = l.size();
-                let mut ipiv = vec_uninit(n as usize);
-                if n == 0 {
-                    return Ok(Vec::new());
-                }
-
-                // calc work size
-                let mut info = 0;
-                let mut work_size = [Self::zero()];
-                unsafe {
-                    $trf(
-                        uplo.as_ptr(),
-                        &n,
-                        AsPtr::as_mut_ptr(a),
-                        &l.lda(),
-                        AsPtr::as_mut_ptr(&mut ipiv),
-                        AsPtr::as_mut_ptr(&mut work_size),
-                        &(-1),
-                        &mut info,
-                    )
-                };
-                info.as_lapack_result()?;
-
-                // actual
-                let lwork = work_size[0].to_usize().unwrap();
-                let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
-                unsafe {
-                    $trf(
-                        uplo.as_ptr(),
-                        &n,
-                        AsPtr::as_mut_ptr(a),
-                        &l.lda(),
-                        AsPtr::as_mut_ptr(&mut ipiv),
-                        AsPtr::as_mut_ptr(&mut work),
-                        &(lwork as i32),
-                        &mut info,
-                    )
-                };
-                info.as_lapack_result()?;
-                let ipiv = unsafe { ipiv.assume_init() };
-                Ok(ipiv)
-            }
-
-            fn invh(l: MatrixLayout, uplo: UPLO, a: &mut [Self], ipiv: &Pivot) -> Result<()> {
-                let (n, _) = l.size();
-                let mut info = 0;
-                let mut work: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
-                unsafe {
-                    $tri(
-                        uplo.as_ptr(),
-                        &n,
-                        AsPtr::as_mut_ptr(a),
-                        &l.lda(),
-                        ipiv.as_ptr(),
-                        AsPtr::as_mut_ptr(&mut work),
-                        &mut info,
-                    )
-                };
-                info.as_lapack_result()?;
-                Ok(())
-            }
-
-            fn solveh(
-                l: MatrixLayout,
-                uplo: UPLO,
-                a: &[Self],
-                ipiv: &Pivot,
-                b: &mut [Self],
-            ) -> Result<()> {
-                let (n, _) = l.size();
-                let mut info = 0;
-                unsafe {
-                    $trs(
-                        uplo.as_ptr(),
-                        &n,
-                        &1,
-                        AsPtr::as_ptr(a),
-                        &l.lda(),
-                        ipiv.as_ptr(),
-                        AsPtr::as_mut_ptr(b),
-                        &n,
-                        &mut info,
-                    )
-                };
-                info.as_lapack_result()?;
-                Ok(())
-            }
-        }
-    };
-} // impl_solveh!
-
-impl_solveh!(
-    f64,
-    lapack_sys::dsytrf_,
-    lapack_sys::dsytri_,
-    lapack_sys::dsytrs_
-);
-impl_solveh!(
-    f32,
-    lapack_sys::ssytrf_,
-    lapack_sys::ssytri_,
-    lapack_sys::ssytrs_
-);
-impl_solveh!(
-    c64,
-    lapack_sys::zhetrf_,
-    lapack_sys::zhetri_,
-    lapack_sys::zhetrs_
-);
-impl_solveh!(
-    c32,
-    lapack_sys::chetrf_,
-    lapack_sys::chetri_,
-    lapack_sys::chetrs_
-);