From 2119604ad9e013368489caf0542828442607b431 Mon Sep 17 00:00:00 2001 From: selvam Date: Sat, 17 Apr 2021 18:18:01 +0000 Subject: [PATCH] Added eigg function which is a wrapper to ggev Added doc comments for eigg function Added testcase --- lax/Cargo.toml | 2 +- lax/src/eig.rs | 280 +++++++++++++++++++++++++++++++-- ndarray-linalg/examples/eig.rs | 25 ++- ndarray-linalg/src/eig.rs | 33 ++++ ndarray-linalg/tests/eig.rs | 37 +++++ 5 files changed, 365 insertions(+), 12 deletions(-) diff --git a/lax/Cargo.toml b/lax/Cargo.toml index a8e74044..14c3c265 100644 --- a/lax/Cargo.toml +++ b/lax/Cargo.toml @@ -32,7 +32,7 @@ intel-mkl-system = ["intel-mkl-src/mkl-dynamic-lp64-seq"] thiserror = "1.0.24" cauchy = "0.4.0" num-traits = "0.2.14" -lapack = "0.18.0" +lapack = "0.19.0" [dependencies.intel-mkl-src] version = "0.6.0" diff --git a/lax/src/eig.rs b/lax/src/eig.rs index 53245de7..b39c223a 100644 --- a/lax/src/eig.rs +++ b/lax/src/eig.rs @@ -12,10 +12,16 @@ pub trait Eig_: Scalar { l: MatrixLayout, a: &mut [Self], ) -> Result<(Vec, Vec)>; + fn eigg( + calc_v: bool, + l: MatrixLayout, + a: &mut [Self], + b: &mut [Self], + ) -> Result<(Vec, Vec)>; } macro_rules! impl_eig_complex { - ($scalar:ty, $ev:path) => { + ($scalar:ty, $ev1:path, $ev2:path) => { impl Eig_ for $scalar { fn eig( calc_v: bool, @@ -51,7 +57,7 @@ macro_rules! impl_eig_complex { let mut info = 0; let mut work_size = [Self::zero()]; unsafe { - $ev( + $ev1( jobvl, jobvr, n, @@ -74,7 +80,7 @@ macro_rules! impl_eig_complex { let lwork = work_size[0].to_usize().unwrap(); let mut work = unsafe { vec_uninit(lwork) }; unsafe { - $ev( + $ev1( jobvl, jobvr, n, @@ -102,15 +108,115 @@ macro_rules! impl_eig_complex { Ok((eigs, vr.or(vl).unwrap_or(Vec::new()))) } + + fn eigg( + calc_v: bool, + l: MatrixLayout, + mut a: &mut [Self], + mut b: &mut [Self], + ) -> Result<(Vec, Vec)> { + let (n, _) = l.size(); + // Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate. + // However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A + let (jobvl, jobvr) = if calc_v { + match l { + MatrixLayout::C { .. } => (b'V', b'N'), + MatrixLayout::F { .. } => (b'N', b'V'), + } + } else { + (b'N', b'N') + }; + let mut alpha = unsafe { vec_uninit(n as usize) }; + let mut beta = unsafe { vec_uninit(n as usize) }; + let mut rwork = unsafe { vec_uninit(8 * n as usize) }; + + let mut vl = if jobvl == b'V' { + Some(unsafe { vec_uninit((n * n) as usize) }) + } else { + None + }; + let mut vr = if jobvr == b'V' { + Some(unsafe { vec_uninit((n * n) as usize) }) + } else { + None + }; + + // calc work size + let mut info = 0; + let mut work_size = [Self::zero()]; + unsafe { + $ev2( + jobvl, + jobvr, + n, + &mut a, + n, + &mut b, + n, + &mut alpha, + &mut beta, + &mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + &mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + &mut work_size, + -1, + &mut rwork, + &mut info, + ) + }; + info.as_lapack_result()?; + + // actal ev + let lwork = work_size[0].to_usize().unwrap(); + let mut work = unsafe { vec_uninit(lwork) }; + unsafe { + $ev2( + jobvl, + jobvr, + n, + &mut a, + n, + &mut b, + n, + &mut alpha, + &mut beta, + &mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + &mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + &mut work, + lwork as i32, + &mut rwork, + &mut info, + ) + }; + info.as_lapack_result()?; + + // Hermite conjugate + if jobvl == b'V' { + for c in vl.as_mut().unwrap().iter_mut() { + c.im = -c.im + } + } + // reconstruct eigenvalues + let eigs: Vec = alpha + .iter() + .zip(beta.iter()) + .map(|(&a, &b)| a/b) + .collect(); + + Ok((eigs, vr.or(vl).unwrap_or(Vec::new()))) + } } }; } -impl_eig_complex!(c64, lapack::zgeev); -impl_eig_complex!(c32, lapack::cgeev); +impl_eig_complex!(c64, lapack::zgeev, lapack::zggev); +impl_eig_complex!(c32, lapack::cgeev, lapack::cggev); macro_rules! impl_eig_real { - ($scalar:ty, $ev:path) => { + ($scalar:ty, $ev1:path, $ev2:path) => { impl Eig_ for $scalar { fn eig( calc_v: bool, @@ -146,7 +252,7 @@ macro_rules! impl_eig_real { let mut info = 0; let mut work_size = [0.0]; unsafe { - $ev( + $ev1( jobvl, jobvr, n, @@ -169,7 +275,7 @@ macro_rules! impl_eig_real { let lwork = work_size[0].to_usize().unwrap(); let mut work = unsafe { vec_uninit(lwork) }; unsafe { - $ev( + $ev1( jobvl, jobvr, n, @@ -250,9 +356,163 @@ macro_rules! impl_eig_real { Ok((eigs, eigvecs)) } + + fn eigg( + calc_v: bool, + l: MatrixLayout, + mut a: &mut [Self], + mut b: &mut [Self], + ) -> Result<(Vec, Vec)> { + let (n, _) = l.size(); + // Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate. + // However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A + let (jobvl, jobvr) = if calc_v { + match l { + MatrixLayout::C { .. } => (b'V', b'N'), + MatrixLayout::F { .. } => (b'N', b'V'), + } + } else { + (b'N', b'N') + }; + let mut alphar = unsafe { vec_uninit(n as usize) }; + let mut alphai = unsafe { vec_uninit(n as usize) }; + let mut beta = unsafe { vec_uninit(n as usize) }; + + let mut vl = if jobvl == b'V' { + Some(unsafe { vec_uninit((n * n) as usize) }) + } else { + None + }; + let mut vr = if jobvr == b'V' { + Some(unsafe { vec_uninit((n * n) as usize) }) + } else { + None + }; + + // calc work size + let mut info = 0; + let mut work_size = [0.0]; + unsafe { + $ev2( + jobvl, + jobvr, + n, + &mut a, + n, + &mut b, + n, + &mut alphar, + &mut alphai, + &mut beta, + vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + &mut work_size, + -1, + &mut info, + ) + }; + info.as_lapack_result()?; + + // actual ev + let lwork = work_size[0].to_usize().unwrap(); + let mut work = unsafe { vec_uninit(lwork) }; + unsafe { + $ev2( + jobvl, + jobvr, + n, + &mut a, + n, + &mut b, + n, + &mut alphar, + &mut alphai, + &mut beta, + vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []), + n, + &mut work, + lwork as i32, + &mut info, + ) + }; + info.as_lapack_result()?; + + // reconstruct eigenvalues + let alpha: Vec = alphar + .iter() + .zip(alphai.iter()) + .map(|(&re, &im)| Self::complex(re, im)) + .collect(); + + let eigs: Vec = alpha + .iter() + .zip(beta.iter()) + .map(|(&a, &b)| a/b) + .collect(); + + if !calc_v { + return Ok((eigs, Vec::new())); + } + + + // Reconstruct eigenvectors into complex-array + // -------------------------------------------- + // + // From LAPACK API https://software.intel.com/en-us/node/469230 + // + // - If the j-th eigenvalue is real, + // - v(j) = VR(:,j), the j-th column of VR. + // + // - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair, + // - v(j) = VR(:,j) + i*VR(:,j+1) + // - v(j+1) = VR(:,j) - i*VR(:,j+1). + // + // ``` + // j -> <----pair----> <----pair----> + // [ ... (real), (imag), (imag), (imag), (imag), ... ] : eigs + // ^ ^ ^ ^ ^ + // false false true false true : is_conjugate_pair + // ``` + let n = n as usize; + let v = vr.or(vl).unwrap(); + let mut eigvecs = unsafe { vec_uninit(n * n) }; + let mut is_conjugate_pair = false; // flag for check `j` is complex conjugate + for j in 0..n { + if alphai[j] == 0.0 { + // j-th eigenvalue is real + for i in 0..n { + eigvecs[i + j * n] = Self::complex(v[i + j * n], 0.0); + } + } else { + // j-th eigenvalue is complex + // complex conjugated pair can be `j-1` or `j+1` + if is_conjugate_pair { + let j_pair = j - 1; + assert!(j_pair < n); + for i in 0..n { + eigvecs[i + j * n] = Self::complex(v[i + j_pair * n], v[i + j * n]); + } + } else { + let j_pair = j + 1; + assert!(j_pair < n); + for i in 0..n { + eigvecs[i + j * n] = + Self::complex(v[i + j * n], -v[i + j_pair * n]); + } + } + is_conjugate_pair = !is_conjugate_pair; + } + } + + Ok((eigs, eigvecs)) + } } }; } -impl_eig_real!(f64, lapack::dgeev); -impl_eig_real!(f32, lapack::sgeev); +impl_eig_real!(f64, lapack::dgeev, lapack::dggev); +impl_eig_real!(f32, lapack::sgeev, lapack::sggev); diff --git a/ndarray-linalg/examples/eig.rs b/ndarray-linalg/examples/eig.rs index 3e41556a..d4fbd923 100644 --- a/ndarray-linalg/examples/eig.rs +++ b/ndarray-linalg/examples/eig.rs @@ -1,7 +1,7 @@ use ndarray::*; use ndarray_linalg::*; -fn main() { +fn eig() { let a = arr2(&[[2.0, 1.0, 2.0], [-2.0, 2.0, 1.0], [1.0, 2.0, -2.0]]); let (e, vecs) = a.clone().eig().unwrap(); println!("eigenvalues = \n{:?}", e); @@ -10,3 +10,26 @@ fn main() { let av = a_c.dot(&vecs); println!("AV = \n{:?}", av); } + +fn eigg_real() { + let a = arr2(&[[1.0/2.0.sqrt(), 0.0], [0.0, 1.0]]); + let b = arr2(&[[0.0, 1.0], [-1.0/2.0.sqrt(), 0.0]]); + let (e, vecs) = a.clone().eigg(&b).unwrap(); + println!("eigenvalues = \n{:?}", e); + println!("V = \n{:?}", vecs); +} + +fn eigg_complex() { + let a = arr2(&[[c64::complex(-3.84, 2.25), c64::complex(-3.84, 2.25)], [c64::complex(-3.84, 2.25), c64::complex(-3.84, 2.25)]]); + let b = a.clone(); + let (e, vecs) = a.clone().eigg(&b).unwrap(); + println!("eigenvalues = \n{:?}", &e); + println!("V = \n{:?}", vecs); + +} + +fn main() { + eig(); + eigg_real(); + eigg_complex(); +} diff --git a/ndarray-linalg/src/eig.rs b/ndarray-linalg/src/eig.rs index 17f5a1e8..5da13a03 100644 --- a/ndarray-linalg/src/eig.rs +++ b/ndarray-linalg/src/eig.rs @@ -35,6 +35,26 @@ pub trait Eig { /// } /// ``` fn eig(&self) -> Result<(Self::EigVal, Self::EigVec)>; + /// Calculate generalised eigenvalues with the right eigenvector + /// + /// $$ A u_i = \lambda_i B u_i $$ + /// + /// ``` + /// use ndarray::*; + /// use ndarray_linalg::*; + /// + /// let a: Array2 = array![ + /// [ 1.0/2.0.sqrt(), 0.0], + /// [ 0.0, 1.0], + /// ]; + /// let b: Array2 = array![ + /// [ 0.0, 1.0], + /// [-1.0/2.0.sqrt(), 0.0], + /// ]; + /// let (eigs, vecs) = a.eigg(&b).unwrap(); + /// + /// ``` + fn eigg(&self, b: &Self) -> Result<(Self::EigVal, Self::EigVec)>; } impl Eig for ArrayBase @@ -55,6 +75,19 @@ where Array2::from_shape_vec((n, n).f(), t).unwrap(), )) } + + fn eigg(&self, b: &Self) -> Result<(Self::EigVal, Self::EigVec)> { + let mut a = self.to_owned(); + let layout_a = a.square_layout()?; + let mut b = b.to_owned(); + let _ = b.square_layout()?; + let (s, t) = A::eigg(true, layout_a, a.as_allocated_mut()?, b.as_allocated_mut()?)?; + let n = layout_a.len() as usize; + Ok(( + ArrayBase::from(s), + Array2::from_shape_vec((n, n).f(), t).unwrap(), + )) + } } /// Calculate eigenvalues without eigenvectors diff --git a/ndarray-linalg/tests/eig.rs b/ndarray-linalg/tests/eig.rs index 28314b8a..7b706fed 100644 --- a/ndarray-linalg/tests/eig.rs +++ b/ndarray-linalg/tests/eig.rs @@ -19,6 +19,27 @@ where } } +fn test_eigg(a: Array2, b: Array2, eigs: Array1, vecs: Array2) +where + T::Complex: Lapack, +{ + println!("a\n{:+.4}", &a); + println!("a\n{:+.4}", &b); + println!("eigs\n{:+.4}", &eigs); + println!("vec\n{:+.4}", &vecs); + let a: Array2 = a.map(|v| v.as_c()); + let b: Array2 = b.map(|v| v.as_c()); + for (&e, v) in eigs.iter().zip(vecs.axis_iter(Axis(1))) { + let av = a.dot(&v); + let bv = b.dot(&v); + let bve = bv.mapv(|val| val * e); + // let bev = be.mapv(|val| val * v); + println!("av = {:+.4}", &av); + println!("ev = {:+.4}", &bve); + assert_close_l2!(&av, &bve, T::real(1e-3)); + } +} + // Test case for real Eigenvalue problem // // -1.01 0.86 -4.60 3.31 -4.81 @@ -230,6 +251,14 @@ macro_rules! impl_test_real { test_eig(a, e, vecs); } + #[test] + fn [<$real _eigg_t>]() { + let a = test_matrix_real_t::<$real>(); + let b = test_matrix_real_t::<$real>(); + let (e, vecs) = a.eigg(&b).unwrap(); + test_eigg(a, b, e, vecs); + } + } // paste::item! }; } @@ -281,6 +310,14 @@ macro_rules! impl_test_complex { let (e, vecs) = a.eig().unwrap(); test_eig(a, e, vecs); } + + #[test] + fn [<$complex _eigg_t>]() { + let a = test_matrix_complex_t::<$complex>(); + let b = test_matrix_complex_t::<$complex>(); + let (e, vecs) = a.eigg(&b).unwrap(); + test_eigg(a, b, e, vecs); + } } // paste::item! }; }