Merge pull request #36 from termoshtt/hermite

Re-implementation of Eigh and Cholesky

Author: termoshtt

src/bin/main.rs

@@ -7,7 +7,7 @@ use ndarray_linalg::prelude::*;
 
 fn main() {
     let a = arr2(&[[3.0, 1.0, 1.0], [1.0, 3.0, 1.0], [1.0, 1.0, 3.0]]);
-    let (e, vecs) = a.clone().eigh().unwrap();
+    let (e, vecs): (Array1<_>, Array2<_>) = a.clone().eigh(UPLO::Upper).unwrap();
     println!("eigenvalues = \n{:?}", e);
     println!("V = \n{:?}", vecs);
     let av = a.dot(&vecs);

src/cholesky.rs

@@ -0,0 +1,46 @@
+
+use ndarray::*;
+use num_traits::Zero;
+
+use super::error::*;
+use super::layout::*;
+use super::triangular::IntoTriangular;
+
+use impl2::LapackScalar;
+pub use impl2::UPLO;
+
+pub trait Cholesky<K> {
+    fn cholesky(self, UPLO) -> Result<K>;
+}
+
+impl<A, S> Cholesky<ArrayBase<S, Ix2>> for ArrayBase<S, Ix2>
+    where A: LapackScalar + Zero,
+          S: DataMut<Elem = A>
+{
+    fn cholesky(mut self, uplo: UPLO) -> Result<ArrayBase<S, Ix2>> {
+        A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)?;
+        Ok(self.into_triangular(uplo))
+    }
+}
+
+impl<'a, A, S> Cholesky<&'a mut ArrayBase<S, Ix2>> for &'a mut ArrayBase<S, Ix2>
+    where A: LapackScalar + Zero,
+          S: DataMut<Elem = A>
+{
+    fn cholesky(mut self, uplo: UPLO) -> Result<&'a mut ArrayBase<S, Ix2>> {
+        A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)?;
+        Ok(self.into_triangular(uplo))
+    }
+}
+
+impl<'a, A, Si, So> Cholesky<ArrayBase<So, Ix2>> for &'a ArrayBase<Si, Ix2>
+    where A: LapackScalar + Copy + Zero,
+          Si: Data<Elem = A>,
+          So: DataMut<Elem = A> + DataOwned
+{
+    fn cholesky(self, uplo: UPLO) -> Result<ArrayBase<So, Ix2>> {
+        let mut a = replicate(self);
+        A::cholesky(a.square_layout()?, uplo, a.as_allocated_mut()?)?;
+        Ok(a.into_triangular(uplo))
+    }
+}

src/eigh.rs

@@ -0,0 +1,85 @@
+
+use ndarray::*;
+
+use super::error::*;
+use super::layout::*;
+
+use impl2::LapackScalar;
+pub use impl2::UPLO;
+
+pub trait Eigh<EigVal, EigVec> {
+    fn eigh(self, UPLO) -> Result<(EigVal, EigVec)>;
+}
+
+impl<A, S, Se> Eigh<ArrayBase<Se, Ix1>, ArrayBase<S, Ix2>> for ArrayBase<S, Ix2>
+    where A: LapackScalar,
+          S: DataMut<Elem = A>,
+          Se: DataOwned<Elem = A::Real>
+{
+    fn eigh(mut self, uplo: UPLO) -> Result<(ArrayBase<Se, Ix1>, ArrayBase<S, Ix2>)> {
+        let s = A::eigh(true, self.square_layout()?, uplo, self.as_allocated_mut()?)?;
+        Ok((ArrayBase::from_vec(s), self))
+    }
+}
+
+impl<'a, A, S, Se, So> Eigh<ArrayBase<Se, Ix1>, ArrayBase<So, Ix2>> for &'a ArrayBase<S, Ix2>
+    where A: LapackScalar + Copy,
+          S: Data<Elem = A>,
+          Se: DataOwned<Elem = A::Real>,
+          So: DataOwned<Elem = A> + DataMut
+{
+    fn eigh(self, uplo: UPLO) -> Result<(ArrayBase<Se, Ix1>, ArrayBase<So, Ix2>)> {
+        let mut a = replicate(self);
+        let s = A::eigh(true, a.square_layout()?, uplo, a.as_allocated_mut()?)?;
+        Ok((ArrayBase::from_vec(s), a))
+    }
+}
+
+impl<'a, A, S, Se> Eigh<ArrayBase<Se, Ix1>, &'a mut ArrayBase<S, Ix2>> for &'a mut ArrayBase<S, Ix2>
+    where A: LapackScalar,
+          S: DataMut<Elem = A>,
+          Se: DataOwned<Elem = A::Real>
+{
+    fn eigh(mut self, uplo: UPLO) -> Result<(ArrayBase<Se, Ix1>, &'a mut ArrayBase<S, Ix2>)> {
+        let s = A::eigh(true, self.square_layout()?, uplo, self.as_allocated_mut()?)?;
+        Ok((ArrayBase::from_vec(s), self))
+    }
+}
+
+pub trait EigValsh<EigVal> {
+    fn eigvalsh(self, UPLO) -> Result<EigVal>;
+}
+
+impl<A, S, Se> EigValsh<ArrayBase<Se, Ix1>> for ArrayBase<S, Ix2>
+    where A: LapackScalar,
+          S: DataMut<Elem = A>,
+          Se: DataOwned<Elem = A::Real>
+{
+    fn eigvalsh(mut self, uplo: UPLO) -> Result<ArrayBase<Se, Ix1>> {
+        let s = A::eigh(false, self.square_layout()?, uplo, self.as_allocated_mut()?)?;
+        Ok(ArrayBase::from_vec(s))
+    }
+}
+
+impl<'a, A, S, Se> EigValsh<ArrayBase<Se, Ix1>> for &'a ArrayBase<S, Ix2>
+    where A: LapackScalar + Copy,
+          S: Data<Elem = A>,
+          Se: DataOwned<Elem = A::Real>
+{
+    fn eigvalsh(self, uplo: UPLO) -> Result<ArrayBase<Se, Ix1>> {
+        let mut a = self.to_owned();
+        let s = A::eigh(false, a.square_layout()?, uplo, a.as_allocated_mut()?)?;
+        Ok(ArrayBase::from_vec(s))
+    }
+}
+
+impl<'a, A, S, Se> EigValsh<ArrayBase<Se, Ix1>> for &'a mut ArrayBase<S, Ix2>
+    where A: LapackScalar,
+          S: DataMut<Elem = A>,
+          Se: DataOwned<Elem = A::Real>
+{
+    fn eigvalsh(mut self, uplo: UPLO) -> Result<ArrayBase<Se, Ix1>> {
+        let s = A::eigh(true, self.square_layout()?, uplo, self.as_allocated_mut()?)?;
+        Ok(ArrayBase::from_vec(s))
+    }
+}

src/hermite.rs

@@ -0,0 +1,29 @@
+//! implement Cholesky decomposition
+
+use lapack::c;
+
+use types::*;
+use error::*;
+use layout::Layout;
+
+use super::{into_result, UPLO};
+
+pub trait Cholesky_: Sized {
+    fn cholesky(Layout, UPLO, a: &mut [Self]) -> Result<()>;
+}
+
+macro_rules! impl_cholesky {
+    ($scalar:ty, $potrf:path) => {
+impl Cholesky_ for $scalar {
+    fn cholesky(l: Layout, uplo: UPLO, mut a: &mut [Self]) -> Result<()> {
+        let (n, _) = l.size();
+        let info = $potrf(l.lapacke_layout(), uplo as u8, n, &mut a, n);
+        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);

src/impl2/cholesky.rs

@@ -0,0 +1,31 @@
+
+use lapack::c;
+use num_traits::Zero;
+
+use types::*;
+use error::*;
+use layout::Layout;
+
+use super::{into_result, UPLO};
+
+pub trait Eigh_: AssociatedReal {
+    fn eigh(calc_eigenvec: bool, Layout, UPLO, a: &mut [Self]) -> Result<Vec<Self::Real>>;
+}
+
+macro_rules! impl_eigh {
+    ($scalar:ty, $ev:path) => {
+impl Eigh_ for $scalar {
+    fn eigh(calc_v: bool, l: Layout, uplo: UPLO, mut a: &mut [Self]) -> Result<Vec<Self::Real>> {
+        let (n, _) = l.size();
+        let jobz = if calc_v { b'V' } else { b'N' };
+        let mut w = vec![Self::Real::zero(); n as usize];
+        let info = $ev(l.lapacke_layout(), jobz, uplo as u8, n, &mut a, n, &mut w);
+        into_result(info, w)
+    }
+}
+}} // impl_eigh!
+
+impl_eigh!(f64, c::dsyev);
+impl_eigh!(f32, c::ssyev);
+impl_eigh!(c64, c::zheev);
+impl_eigh!(c32, c::cheev);

src/impl2/eigh.rs

@@ -3,16 +3,20 @@ pub mod opnorm;
 pub mod qr;
 pub mod svd;
 pub mod solve;
+pub mod cholesky;
+pub mod eigh;
 
 pub use self::opnorm::*;
 pub use self::qr::*;
 pub use self::svd::*;
 pub use self::solve::*;
+pub use self::cholesky::*;
+pub use self::eigh::*;
 
 use super::error::*;
 
-pub trait LapackScalar: OperatorNorm_ + QR_ + SVD_ + Solve_ {}
-impl<A> LapackScalar for A where A: OperatorNorm_ + QR_ + SVD_ + Solve_ {}
+pub trait LapackScalar: OperatorNorm_ + QR_ + SVD_ + Solve_ + Cholesky_ + Eigh_ {}
+impl<A> LapackScalar for A where A: OperatorNorm_ + QR_ + SVD_ + Solve_ + Cholesky_ + Eigh_ {}
 
 pub fn into_result<T>(info: i32, val: T) -> Result<T> {
     if info == 0 {
@@ -21,3 +25,10 @@ pub fn into_result<T>(info: i32, val: T) -> Result<T> {
         Err(LapackError::new(info).into())
     }
 }
+
+#[derive(Debug, Clone, Copy)]
+#[repr(u8)]
+pub enum UPLO {
+    Upper = b'U',
+    Lower = b'L',
+}

src/impl2/mod.rs

@@ -1,5 +1,3 @@
 //! Implement trait bindings of LAPACK
 pub mod outer;
-pub mod eigh;
 pub mod solve;
-pub mod cholesky;