add eig

Author: doraneko94

examples/eig.rs

@@ -0,0 +1,14 @@
+extern crate ndarray;
+extern crate ndarray_linalg;
+
+use ndarray::*;
+use ndarray_linalg::*;
+
+fn main() {
+    let a = arr2(&[[0.0, -1.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 3.0, 2.0]]);
+    let (e, vecs) = a.clone().eig().unwrap();
+    println!("eigenvalues = \n{:?}", e);
+    println!("V = \n{:?}", vecs.t());
+    let av = a.dot(&vecs.t());
+    println!("AV = \n{:?}", av);
+}

src/eig.rs

@@ -0,0 +1,62 @@
+//! Eigenvalue decomposition for Hermite matrices
+
+use ndarray::*;
+use crate::error::*;
+use crate::layout::*;
+use crate::types::*;
+
+/// Eigenvalue decomposition of general matrix reference
+pub trait Eig {
+    type EigVal;
+    type EigVec;
+    fn eig(&self) -> Result<(Self::EigVal, Self::EigVec)>;
+}
+
+impl<A, S> Eig for ArrayBase<S, Ix2>
+where
+    A: Scalar + Lapack,
+    S: Data<Elem = A>,
+{
+    type EigVal = Array1<A::Complex>;
+    type EigVec = Array2<A>;
+
+    fn eig(&self) -> Result<(Self::EigVal, Self::EigVec)> {
+        let mut a = self.to_owned();
+        let layout = a.square_layout()?;
+        // XXX Force layout to be Fortran (see #146)
+        match layout {
+            MatrixLayout::C(_) => a.swap_axes(0, 1),
+            MatrixLayout::F(_) => {}
+        }
+        let (s, t) = unsafe { A::eig(true, a.square_layout()?, a.as_allocated_mut()?)? };
+        let (n, _) = layout.size();
+        println!("{:?}", t);
+        Ok((ArrayBase::from(s), ArrayBase::from(t).into_shape((n as usize, n as usize)).unwrap()))
+    }
+}
+
+/// Calculate eigenvalues without eigenvectors
+pub trait EigVals {
+    type EigVal;
+    fn eigvals(&self) -> Result<Self::EigVal>;
+}
+
+impl<A, S> EigVals for ArrayBase<S, Ix2>
+where
+    A: Scalar + Lapack,
+    S: DataMut<Elem = A>,
+{
+    type EigVal = Array1<A::Complex>;
+
+    fn eigvals(&self) -> Result<Self::EigVal> {
+        let mut a = self.to_owned();
+        let layout = a.square_layout()?;
+        // XXX Force layout to be Fortran (see #146)
+        match layout {
+            MatrixLayout::C(_) => a.swap_axes(0, 1),
+            MatrixLayout::F(_) => {}
+        }
+        let (s, _) = unsafe { A::eig(true, a.square_layout()?, a.as_allocated_mut()?)? };
+        Ok(ArrayBase::from(s))
+    }
+}

src/lapack/eig.rs

@@ -0,0 +1,54 @@
+//! Eigenvalue decomposition for general matrices
+
+use lapacke;
+use num_traits::Zero;
+
+use crate::error::*;
+use crate::layout::MatrixLayout;
+use crate::types::*;
+
+use super::into_result;
+
+/// Wraps `*geev` for real/complex
+pub trait Eig_: Scalar {
+    unsafe fn eig(calc_v: bool, l: MatrixLayout, a: &mut [Self]) -> Result<(Vec<Self::Complex>, Vec<Self>)>;
+}
+
+macro_rules! impl_eig_complex {
+    ($scalar:ty, $ev:path) => {
+        impl Eig_ for $scalar {
+            unsafe fn eig(calc_v: bool, l: MatrixLayout, mut a: &mut [Self]) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
+                let (n, _) = l.size();
+                let jobvl = if calc_v { b'V' } else { b'N' };
+                let mut w = vec![Self::Complex::zero(); n as usize];
+                let mut vl = vec![Self::Complex::zero(); (n * n) as usize];
+                let mut vr = Vec::new();
+                let info = $ev(l.lapacke_layout(), jobvl, b'N', n, &mut a, n, &mut w, &mut vl, n, &mut vr, n);
+                into_result(info, (w, vl))
+            }
+        }
+    };
+}
+
+macro_rules! impl_eig_real {
+    ($scalar:ty, $ev:path, $cn:ty) => {
+        impl Eig_ for $scalar {
+            unsafe fn eig(calc_v: bool, l: MatrixLayout, mut a: &mut [Self]) -> Result<(Vec<Self::Complex>, Vec<Self::Real>)> {
+                let (n, _) = l.size();
+                let jobvl = if calc_v { b'V' } else { b'N' };
+                let mut wr = vec![Self::Real::zero(); n as usize];
+                let mut wi = vec![Self::Real::zero(); n as usize];
+                let mut vl = vec![Self::Real::zero(); (n * n) as usize];
+                let mut vr = Vec::new();
+                let info = $ev(l.lapacke_layout(), jobvl, b'N', n, &mut a, n, &mut wr, &mut wi, &mut vl, n, &mut vr, n);
+                let w: Vec<$cn> = wr.iter().zip(wi.iter()).map(|(&r, &i)| <$cn>::new(r, i)).collect();
+                into_result(info, (w, vl))
+            }
+        }
+    };
+}
+
+impl_eig_real!(f64, lapacke::dgeev, c64);
+impl_eig_real!(f32, lapacke::sgeev, c32);
+impl_eig_complex!(c64, lapacke::zgeev);
+impl_eig_complex!(c32, lapacke::cgeev);

src/lapack/mod.rs

@@ -1,6 +1,7 @@
 //! Define traits wrapping LAPACK routines
 
 pub mod cholesky;
+pub mod eig;
 pub mod eigh;
 pub mod opnorm;
 pub mod qr;
@@ -11,6 +12,7 @@ pub mod svddc;
 pub mod triangular;
 
 pub use self::cholesky::*;
+pub use self::eig::*;
 pub use self::eigh::*;
 pub use self::opnorm::*;
 pub use self::qr::*;
@@ -26,7 +28,7 @@ use super::types::*;
 pub type Pivot = Vec<i32>;
 
 /// Trait for primitive types which implements LAPACK subroutines
-pub trait Lapack: OperatorNorm_ + QR_ + SVD_ + SVDDC_ + Solve_ + Solveh_ + Cholesky_ + Eigh_ + Triangular_ {}
+pub trait Lapack: OperatorNorm_ + QR_ + SVD_ + SVDDC_ + Solve_ + Solveh_ + Cholesky_ + Eig_ + Eigh_ + Triangular_ {}
 
 impl Lapack for f32 {}
 impl Lapack for f64 {}

src/lib.rs

@@ -43,6 +43,7 @@ pub mod assert;
 pub mod cholesky;
 pub mod convert;
 pub mod diagonal;
+pub mod eig;
 pub mod eigh;
 pub mod error;
 pub mod generate;
@@ -66,6 +67,7 @@ pub use assert::*;
 pub use cholesky::*;
 pub use convert::*;
 pub use diagonal::*;
+pub use eig::*;
 pub use eigh::*;
 pub use generate::*;
 pub use inner::*;