From c953001c864a4bb752f64b0abc65c9258f9e16fc Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Sat, 24 Sep 2022 22:57:59 +0900 Subject: [PATCH 1/6] Split generalized eigenvalue routine --- lax/src/eigh.rs | 99 +++------------------------------ lax/src/eigh_generalized.rs | 106 ++++++++++++++++++++++++++++++++++++ lax/src/lib.rs | 3 + 3 files changed, 118 insertions(+), 90 deletions(-) create mode 100644 lax/src/eigh_generalized.rs diff --git a/lax/src/eigh.rs b/lax/src/eigh.rs index 54af63f7..cd83f2b8 100644 --- a/lax/src/eigh.rs +++ b/lax/src/eigh.rs @@ -21,33 +21,16 @@ pub trait Eigh_: Scalar { uplo: UPLO, a: &mut [Self], ) -> Result>; - - /// Compute generalized right eigenvalue and eigenvectors $Ax = \lambda B x$ - /// - /// LAPACK correspondance - /// ---------------------- - /// - /// | f32 | f64 | c32 | c64 | - /// |:------|:------|:------|:------| - /// | ssygv | dsygv | chegv | zhegv | - /// - fn eigh_generalized( - calc_eigenvec: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - b: &mut [Self], - ) -> Result>; } macro_rules! impl_eigh { - (@real, $scalar:ty, $ev:path, $evg:path) => { - impl_eigh!(@body, $scalar, $ev, $evg, ); + (@real, $scalar:ty, $ev:path) => { + impl_eigh!(@body, $scalar, $ev, ); }; - (@complex, $scalar:ty, $ev:path, $evg:path) => { - impl_eigh!(@body, $scalar, $ev, $evg, rwork); + (@complex, $scalar:ty, $ev:path) => { + impl_eigh!(@body, $scalar, $ev, rwork); }; - (@body, $scalar:ty, $ev:path, $evg:path, $($rwork_ident:ident),*) => { + (@body, $scalar:ty, $ev:path, $($rwork_ident:ident),*) => { impl Eigh_ for $scalar { fn eigh( calc_v: bool, @@ -106,75 +89,11 @@ macro_rules! impl_eigh { let eigs = unsafe { eigs.assume_init() }; Ok(eigs) } - - fn eigh_generalized( - calc_v: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - b: &mut [Self], - ) -> Result> { - assert_eq!(layout.len(), layout.lda()); - let n = layout.len(); - let jobz = if calc_v { JobEv::All } else { JobEv::None }; - let mut eigs: Vec> = vec_uninit(n as usize); - - $( - let mut $rwork_ident: Vec> = vec_uninit(3 * n as usize - 2); - )* - - // calc work size - let mut info = 0; - let mut work_size = [Self::zero()]; - unsafe { - $evg( - &1, // ITYPE A*x = (lambda)*B*x - jobz.as_ptr(), - uplo.as_ptr(), - &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(b), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work_size), - &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - - // actual evg - let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); - let lwork = lwork as i32; - unsafe { - $evg( - &1, // ITYPE A*x = (lambda)*B*x - jobz.as_ptr(), - uplo.as_ptr(), - &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(b), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work), - &lwork, - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - let eigs = unsafe { eigs.assume_init() }; - Ok(eigs) - } } }; } // impl_eigh! -impl_eigh!(@real, f64, lapack_sys::dsyev_, lapack_sys::dsygv_); -impl_eigh!(@real, f32, lapack_sys::ssyev_, lapack_sys::ssygv_); -impl_eigh!(@complex, c64, lapack_sys::zheev_, lapack_sys::zhegv_); -impl_eigh!(@complex, c32, lapack_sys::cheev_, lapack_sys::chegv_); +impl_eigh!(@real, f64, lapack_sys::dsyev_); +impl_eigh!(@real, f32, lapack_sys::ssyev_); +impl_eigh!(@complex, c64, lapack_sys::zheev_); +impl_eigh!(@complex, c32, lapack_sys::cheev_); diff --git a/lax/src/eigh_generalized.rs b/lax/src/eigh_generalized.rs new file mode 100644 index 00000000..2a0c5302 --- /dev/null +++ b/lax/src/eigh_generalized.rs @@ -0,0 +1,106 @@ +use super::*; +use crate::{error::*, layout::MatrixLayout}; +use cauchy::*; +use num_traits::{ToPrimitive, Zero}; + +#[cfg_attr(doc, katexit::katexit)] +/// Eigenvalue problem for symmetric/hermite matrix +pub trait EighGeneralized_: Scalar { + /// Compute generalized right eigenvalue and eigenvectors $Ax = \lambda B x$ + /// + /// LAPACK correspondance + /// ---------------------- + /// + /// | f32 | f64 | c32 | c64 | + /// |:------|:------|:------|:------| + /// | ssygv | dsygv | chegv | zhegv | + /// + fn eigh_generalized( + calc_eigenvec: bool, + layout: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + b: &mut [Self], + ) -> Result>; +} + +macro_rules! impl_eigh { + (@real, $scalar:ty, $evg:path) => { + impl_eigh!(@body, $scalar, $evg, ); + }; + (@complex, $scalar:ty, $evg:path) => { + impl_eigh!(@body, $scalar, $evg, rwork); + }; + (@body, $scalar:ty, $evg:path, $($rwork_ident:ident),*) => { + impl EighGeneralized_ for $scalar { + fn eigh_generalized( + calc_v: bool, + layout: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + b: &mut [Self], + ) -> Result> { + assert_eq!(layout.len(), layout.lda()); + let n = layout.len(); + let jobz = if calc_v { JobEv::All } else { JobEv::None }; + let mut eigs: Vec> = vec_uninit(n as usize); + + $( + let mut $rwork_ident: Vec> = vec_uninit(3 * n as usize - 2); + )* + + // calc work size + let mut info = 0; + let mut work_size = [Self::zero()]; + unsafe { + $evg( + &1, // ITYPE A*x = (lambda)*B*x + jobz.as_ptr(), + uplo.as_ptr(), + &n, + AsPtr::as_mut_ptr(a), + &n, + AsPtr::as_mut_ptr(b), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* + &mut info, + ); + } + info.as_lapack_result()?; + + // actual evg + let lwork = work_size[0].to_usize().unwrap(); + let mut work: Vec> = vec_uninit(lwork); + let lwork = lwork as i32; + unsafe { + $evg( + &1, // ITYPE A*x = (lambda)*B*x + jobz.as_ptr(), + uplo.as_ptr(), + &n, + AsPtr::as_mut_ptr(a), + &n, + AsPtr::as_mut_ptr(b), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work), + &lwork, + $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* + &mut info, + ); + } + info.as_lapack_result()?; + let eigs = unsafe { eigs.assume_init() }; + Ok(eigs) + } + } + }; +} // impl_eigh! + +impl_eigh!(@real, f64, lapack_sys::dsygv_); +impl_eigh!(@real, f32, lapack_sys::ssygv_); +impl_eigh!(@complex, c64, lapack_sys::zhegv_); +impl_eigh!(@complex, c32, lapack_sys::chegv_); diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 00393916..b246b665 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -89,6 +89,7 @@ pub mod eig; mod alloc; mod cholesky; mod eigh; +mod eigh_generalized; mod least_squares; mod opnorm; mod qr; @@ -102,6 +103,7 @@ mod tridiagonal; pub use self::cholesky::*; pub use self::eigh::*; +pub use self::eigh_generalized::*; pub use self::flags::*; pub use self::least_squares::*; pub use self::opnorm::*; @@ -130,6 +132,7 @@ pub trait Lapack: + Solveh_ + Cholesky_ + Eigh_ + + EighGeneralized_ + Triangular_ + Tridiagonal_ + Rcond_ From b864638d1a4382eac3336df3f0d3197079a157ac Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Sat, 24 Sep 2022 23:56:02 +0900 Subject: [PATCH 2/6] EighWork --- lax/src/eigh.rs | 92 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/lax/src/eigh.rs b/lax/src/eigh.rs index cd83f2b8..b8e9e9f7 100644 --- a/lax/src/eigh.rs +++ b/lax/src/eigh.rs @@ -23,6 +23,98 @@ pub trait Eigh_: Scalar { ) -> Result>; } +pub struct EighWork { + pub n: i32, + pub jobz: JobEv, + pub eigs: Vec>, + pub work: Vec>, + pub rwork: Option>>, +} + +pub trait EighWorkImpl: Sized { + type Elem: Scalar; + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result; + fn calc(&mut self, uplo: UPLO, a: &mut [Self::Elem]) + -> Result<&[::Real]>; + fn eval(self, uplo: UPLO, a: &mut [Self::Elem]) -> Result::Real>>; +} + +impl EighWorkImpl for EighWork { + type Elem = c64; + + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result { + assert_eq!(layout.len(), layout.lda()); + let n = layout.len(); + let jobz = if calc_eigenvectors { + JobEv::All + } else { + JobEv::None + }; + let mut eigs = vec_uninit(n as usize); + let mut rwork = vec_uninit(3 * n as usize - 2 as usize); + let mut info = 0; + let mut work_size = [c64::zero()]; + unsafe { + lapack_sys::zheev_( + jobz.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO + &n, + std::ptr::null_mut(), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + AsPtr::as_mut_ptr(&mut rwork), + &mut info, + ); + } + info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(EighWork { + n, + eigs, + jobz, + work, + rwork: Some(rwork), + }) + } + + fn calc( + &mut self, + uplo: UPLO, + a: &mut [Self::Elem], + ) -> Result<&[::Real]> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + lapack_sys::zheev_( + self.jobz.as_ptr(), + uplo.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + &mut info, + ); + } + info.as_lapack_result()?; + Ok(unsafe { self.eigs.slice_assume_init_ref() }) + } + + fn eval( + mut self, + uplo: UPLO, + a: &mut [Self::Elem], + ) -> Result::Real>> { + let _eig = self.calc(uplo, a)?; + Ok(unsafe { self.eigs.assume_init() }) + } +} + macro_rules! impl_eigh { (@real, $scalar:ty, $ev:path) => { impl_eigh!(@body, $scalar, $ev, ); From 5d0753879cd8817d259748d087f8f1f240387557 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Sun, 25 Sep 2022 18:06:24 +0900 Subject: [PATCH 3/6] impl EighWorkImpl for EighWork in c32, f32, f64 --- lax/src/eigh.rs | 229 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 157 insertions(+), 72 deletions(-) diff --git a/lax/src/eigh.rs b/lax/src/eigh.rs index b8e9e9f7..a8bbe14c 100644 --- a/lax/src/eigh.rs +++ b/lax/src/eigh.rs @@ -39,81 +39,166 @@ pub trait EighWorkImpl: Sized { fn eval(self, uplo: UPLO, a: &mut [Self::Elem]) -> Result::Real>>; } -impl EighWorkImpl for EighWork { - type Elem = c64; - - fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result { - assert_eq!(layout.len(), layout.lda()); - let n = layout.len(); - let jobz = if calc_eigenvectors { - JobEv::All - } else { - JobEv::None - }; - let mut eigs = vec_uninit(n as usize); - let mut rwork = vec_uninit(3 * n as usize - 2 as usize); - let mut info = 0; - let mut work_size = [c64::zero()]; - unsafe { - lapack_sys::zheev_( - jobz.as_ptr(), - UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO - &n, - std::ptr::null_mut(), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work_size), - &(-1), - AsPtr::as_mut_ptr(&mut rwork), - &mut info, - ); - } - info.as_lapack_result()?; - let lwork = work_size[0].to_usize().unwrap(); - let work = vec_uninit(lwork); - Ok(EighWork { - n, - eigs, - jobz, - work, - rwork: Some(rwork), - }) - } - - fn calc( - &mut self, - uplo: UPLO, - a: &mut [Self::Elem], - ) -> Result<&[::Real]> { - let lwork = self.work.len().to_i32().unwrap(); - let mut info = 0; - unsafe { - lapack_sys::zheev_( - self.jobz.as_ptr(), - uplo.as_ptr(), - &self.n, - AsPtr::as_mut_ptr(a), - &self.n, - AsPtr::as_mut_ptr(&mut self.eigs), - AsPtr::as_mut_ptr(&mut self.work), - &lwork, - AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), - &mut info, - ); +macro_rules! impl_eigh_work_c { + ($c:ty, $ev:path) => { + impl EighWorkImpl for EighWork<$c> { + type Elem = $c; + + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result { + assert_eq!(layout.len(), layout.lda()); + let n = layout.len(); + let jobz = if calc_eigenvectors { + JobEv::All + } else { + JobEv::None + }; + let mut eigs = vec_uninit(n as usize); + let mut rwork = vec_uninit(3 * n as usize - 2 as usize); + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + unsafe { + $ev( + jobz.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO + &n, + std::ptr::null_mut(), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + AsPtr::as_mut_ptr(&mut rwork), + &mut info, + ); + } + info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(EighWork { + n, + eigs, + jobz, + work, + rwork: Some(rwork), + }) + } + + fn calc( + &mut self, + uplo: UPLO, + a: &mut [Self::Elem], + ) -> Result<&[::Real]> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + $ev( + self.jobz.as_ptr(), + uplo.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + &mut info, + ); + } + info.as_lapack_result()?; + Ok(unsafe { self.eigs.slice_assume_init_ref() }) + } + + fn eval( + mut self, + uplo: UPLO, + a: &mut [Self::Elem], + ) -> Result::Real>> { + let _eig = self.calc(uplo, a)?; + Ok(unsafe { self.eigs.assume_init() }) + } } - info.as_lapack_result()?; - Ok(unsafe { self.eigs.slice_assume_init_ref() }) - } + }; +} +impl_eigh_work_c!(c64, lapack_sys::zheev_); +impl_eigh_work_c!(c32, lapack_sys::cheev_); - fn eval( - mut self, - uplo: UPLO, - a: &mut [Self::Elem], - ) -> Result::Real>> { - let _eig = self.calc(uplo, a)?; - Ok(unsafe { self.eigs.assume_init() }) - } +macro_rules! impl_eigh_work_r { + ($f:ty, $ev:path) => { + impl EighWorkImpl for EighWork<$f> { + type Elem = $f; + + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result { + assert_eq!(layout.len(), layout.lda()); + let n = layout.len(); + let jobz = if calc_eigenvectors { + JobEv::All + } else { + JobEv::None + }; + let mut eigs = vec_uninit(n as usize); + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + unsafe { + $ev( + jobz.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO + &n, + std::ptr::null_mut(), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + &mut info, + ); + } + info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(EighWork { + n, + eigs, + jobz, + work, + rwork: None, + }) + } + + fn calc( + &mut self, + uplo: UPLO, + a: &mut [Self::Elem], + ) -> Result<&[::Real]> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + $ev( + self.jobz.as_ptr(), + uplo.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + &mut info, + ); + } + info.as_lapack_result()?; + Ok(unsafe { self.eigs.slice_assume_init_ref() }) + } + + fn eval( + mut self, + uplo: UPLO, + a: &mut [Self::Elem], + ) -> Result::Real>> { + let _eig = self.calc(uplo, a)?; + Ok(unsafe { self.eigs.assume_init() }) + } + } + }; } +impl_eigh_work_r!(f64, lapack_sys::dsyev_); +impl_eigh_work_r!(f32, lapack_sys::ssyev_); macro_rules! impl_eigh { (@real, $scalar:ty, $ev:path) => { From 370c7d20742ff225cf8dc80e95ac6508c5b0e40b Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Sun, 25 Sep 2022 18:16:28 +0900 Subject: [PATCH 4/6] Merge Eigh_ into Lapack trait --- lax/src/eigh.rs | 104 +++++------------------------------------------- lax/src/lib.rs | 23 +++++++++-- 2 files changed, 29 insertions(+), 98 deletions(-) diff --git a/lax/src/eigh.rs b/lax/src/eigh.rs index a8bbe14c..429aaa75 100644 --- a/lax/src/eigh.rs +++ b/lax/src/eigh.rs @@ -1,28 +1,17 @@ +//! Eigenvalue problem for symmetric/Hermite matricies +//! +//! LAPACK correspondance +//! ---------------------- +//! +//! | f32 | f64 | c32 | c64 | +//! |:------|:------|:------|:------| +//! | ssyev | dsyev | cheev | zheev | + use super::*; use crate::{error::*, layout::MatrixLayout}; use cauchy::*; use num_traits::{ToPrimitive, Zero}; -#[cfg_attr(doc, katexit::katexit)] -/// Eigenvalue problem for symmetric/hermite matrix -pub trait Eigh_: Scalar { - /// Compute right eigenvalue and eigenvectors $Ax = \lambda x$ - /// - /// LAPACK correspondance - /// ---------------------- - /// - /// | f32 | f64 | c32 | c64 | - /// |:------|:------|:------|:------| - /// | ssyev | dsyev | cheev | zheev | - /// - fn eigh( - calc_eigenvec: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - ) -> Result>; -} - pub struct EighWork { pub n: i32, pub jobz: JobEv, @@ -199,78 +188,3 @@ macro_rules! impl_eigh_work_r { } impl_eigh_work_r!(f64, lapack_sys::dsyev_); impl_eigh_work_r!(f32, lapack_sys::ssyev_); - -macro_rules! impl_eigh { - (@real, $scalar:ty, $ev:path) => { - impl_eigh!(@body, $scalar, $ev, ); - }; - (@complex, $scalar:ty, $ev:path) => { - impl_eigh!(@body, $scalar, $ev, rwork); - }; - (@body, $scalar:ty, $ev:path, $($rwork_ident:ident),*) => { - impl Eigh_ for $scalar { - fn eigh( - calc_v: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - ) -> Result> { - assert_eq!(layout.len(), layout.lda()); - let n = layout.len(); - let jobz = if calc_v { JobEv::All } else { JobEv::None }; - let mut eigs: Vec> = vec_uninit(n as usize); - - $( - let mut $rwork_ident: Vec> = vec_uninit(3 * n as usize - 2 as usize); - )* - - // calc work size - let mut info = 0; - let mut work_size = [Self::zero()]; - unsafe { - $ev( - jobz.as_ptr() , - uplo.as_ptr(), - &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work_size), - &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - - // actual ev - let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); - let lwork = lwork as i32; - unsafe { - $ev( - jobz.as_ptr(), - uplo.as_ptr(), - &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work), - &lwork, - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - - let eigs = unsafe { eigs.assume_init() }; - Ok(eigs) - } - } - }; -} // impl_eigh! - -impl_eigh!(@real, f64, lapack_sys::dsyev_); -impl_eigh!(@real, f32, lapack_sys::ssyev_); -impl_eigh!(@complex, c64, lapack_sys::zheev_); -impl_eigh!(@complex, c32, lapack_sys::cheev_); diff --git a/lax/src/lib.rs b/lax/src/lib.rs index b246b665..f02532e5 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -131,25 +131,31 @@ pub trait Lapack: + Solve_ + Solveh_ + Cholesky_ - + Eigh_ + EighGeneralized_ + Triangular_ + Tridiagonal_ + Rcond_ + LeastSquaresSvdDivideConquer_ { - /// Compute right eigenvalue and eigenvectors + /// Compute right eigenvalue and eigenvectors for a general matrix fn eig( calc_v: bool, l: MatrixLayout, a: &mut [Self], ) -> Result<(Vec, Vec)>; + + /// Compute right eigenvalue and eigenvectors for a symmetric or hermite matrix + fn eigh( + calc_eigenvec: bool, + layout: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + ) -> Result>; } macro_rules! impl_lapack { ($s:ty) => { impl Lapack for $s { - /// Compute right eigenvalue and eigenvectors fn eig( calc_v: bool, l: MatrixLayout, @@ -160,6 +166,17 @@ macro_rules! impl_lapack { let EigOwned { eigs, vr, vl } = work.eval(a)?; Ok((eigs, vr.or(vl).unwrap_or_default())) } + + fn eigh( + calc_eigenvec: bool, + layout: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + ) -> Result> { + use eigh::*; + let work = EighWork::<$s>::new(calc_eigenvec, layout)?; + work.eval(uplo, a) + } } }; } From c7a45862d184687da5ed73f1e2b0ef1fd5d60425 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Sun, 25 Sep 2022 18:40:30 +0900 Subject: [PATCH 5/6] Add EighGeneralizedWork --- lax/src/eigh_generalized.rs | 212 ++++++++++++++++++++++++++++++++++++ 1 file changed, 212 insertions(+) diff --git a/lax/src/eigh_generalized.rs b/lax/src/eigh_generalized.rs index 2a0c5302..16ded9c6 100644 --- a/lax/src/eigh_generalized.rs +++ b/lax/src/eigh_generalized.rs @@ -1,8 +1,220 @@ +//! Compute generalized right eigenvalue and eigenvectors +//! +//! LAPACK correspondance +//! ---------------------- +//! +//! | f32 | f64 | c32 | c64 | +//! |:------|:------|:------|:------| +//! | ssygv | dsygv | chegv | zhegv | +//! + use super::*; use crate::{error::*, layout::MatrixLayout}; use cauchy::*; use num_traits::{ToPrimitive, Zero}; +pub struct EighGeneralizedWork { + pub n: i32, + pub jobz: JobEv, + pub eigs: Vec>, + pub work: Vec>, + pub rwork: Option>>, +} + +pub trait EighGeneralizedWorkImpl: Sized { + type Elem: Scalar; + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result; + fn calc( + &mut self, + uplo: UPLO, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result<&[::Real]>; + fn eval( + self, + uplo: UPLO, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result::Real>>; +} + +macro_rules! impl_eigh_generalized_work_c { + ($c:ty, $gv:path) => { + impl EighGeneralizedWorkImpl for EighGeneralizedWork<$c> { + type Elem = $c; + + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result { + assert_eq!(layout.len(), layout.lda()); + let n = layout.len(); + let jobz = if calc_eigenvectors { + JobEv::All + } else { + JobEv::None + }; + let mut eigs = vec_uninit(n as usize); + let mut rwork = vec_uninit(3 * n as usize - 2 as usize); + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + unsafe { + $gv( + &1, // ITYPE A*x = (lambda)*B*x + jobz.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO + &n, + std::ptr::null_mut(), + &n, + std::ptr::null_mut(), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + AsPtr::as_mut_ptr(&mut rwork), + &mut info, + ); + } + info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(EighGeneralizedWork { + n, + eigs, + jobz, + work, + rwork: Some(rwork), + }) + } + + fn calc( + &mut self, + uplo: UPLO, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result<&[::Real]> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + $gv( + &1, // ITYPE A*x = (lambda)*B*x + self.jobz.as_ptr(), + uplo.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &self.n, + AsPtr::as_mut_ptr(b), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + &mut info, + ); + } + info.as_lapack_result()?; + Ok(unsafe { self.eigs.slice_assume_init_ref() }) + } + + fn eval( + mut self, + uplo: UPLO, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result::Real>> { + let _eig = self.calc(uplo, a, b)?; + Ok(unsafe { self.eigs.assume_init() }) + } + } + }; +} +impl_eigh_generalized_work_c!(c64, lapack_sys::zhegv_); +impl_eigh_generalized_work_c!(c32, lapack_sys::chegv_); + +macro_rules! impl_eigh_generalized_work_r { + ($f:ty, $gv:path) => { + impl EighGeneralizedWorkImpl for EighGeneralizedWork<$f> { + type Elem = $f; + + fn new(calc_eigenvectors: bool, layout: MatrixLayout) -> Result { + assert_eq!(layout.len(), layout.lda()); + let n = layout.len(); + let jobz = if calc_eigenvectors { + JobEv::All + } else { + JobEv::None + }; + let mut eigs = vec_uninit(n as usize); + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + unsafe { + $gv( + &1, // ITYPE A*x = (lambda)*B*x + jobz.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO + &n, + std::ptr::null_mut(), + &n, + std::ptr::null_mut(), + &n, + AsPtr::as_mut_ptr(&mut eigs), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + &mut info, + ); + } + info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(EighGeneralizedWork { + n, + eigs, + jobz, + work, + rwork: None, + }) + } + + fn calc( + &mut self, + uplo: UPLO, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result<&[::Real]> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + $gv( + &1, // ITYPE A*x = (lambda)*B*x + self.jobz.as_ptr(), + uplo.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &self.n, + AsPtr::as_mut_ptr(b), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + &mut info, + ); + } + info.as_lapack_result()?; + Ok(unsafe { self.eigs.slice_assume_init_ref() }) + } + + fn eval( + mut self, + uplo: UPLO, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result::Real>> { + let _eig = self.calc(uplo, a, b)?; + Ok(unsafe { self.eigs.assume_init() }) + } + } + }; +} +impl_eigh_generalized_work_r!(f64, lapack_sys::dsygv_); +impl_eigh_generalized_work_r!(f32, lapack_sys::ssygv_); + #[cfg_attr(doc, katexit::katexit)] /// Eigenvalue problem for symmetric/hermite matrix pub trait EighGeneralized_: Scalar { From bef10838ecd5266dfbbda3c04017b2f700cd546f Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Sun, 25 Sep 2022 18:51:01 +0900 Subject: [PATCH 6/6] Merge EighGeneralized_ to Lapack trait --- lax/src/eigh_generalized.rs | 102 ------------------------------------ lax/src/lib.rs | 22 +++++++- 2 files changed, 21 insertions(+), 103 deletions(-) diff --git a/lax/src/eigh_generalized.rs b/lax/src/eigh_generalized.rs index 16ded9c6..1771cee8 100644 --- a/lax/src/eigh_generalized.rs +++ b/lax/src/eigh_generalized.rs @@ -214,105 +214,3 @@ macro_rules! impl_eigh_generalized_work_r { } impl_eigh_generalized_work_r!(f64, lapack_sys::dsygv_); impl_eigh_generalized_work_r!(f32, lapack_sys::ssygv_); - -#[cfg_attr(doc, katexit::katexit)] -/// Eigenvalue problem for symmetric/hermite matrix -pub trait EighGeneralized_: Scalar { - /// Compute generalized right eigenvalue and eigenvectors $Ax = \lambda B x$ - /// - /// LAPACK correspondance - /// ---------------------- - /// - /// | f32 | f64 | c32 | c64 | - /// |:------|:------|:------|:------| - /// | ssygv | dsygv | chegv | zhegv | - /// - fn eigh_generalized( - calc_eigenvec: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - b: &mut [Self], - ) -> Result>; -} - -macro_rules! impl_eigh { - (@real, $scalar:ty, $evg:path) => { - impl_eigh!(@body, $scalar, $evg, ); - }; - (@complex, $scalar:ty, $evg:path) => { - impl_eigh!(@body, $scalar, $evg, rwork); - }; - (@body, $scalar:ty, $evg:path, $($rwork_ident:ident),*) => { - impl EighGeneralized_ for $scalar { - fn eigh_generalized( - calc_v: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - b: &mut [Self], - ) -> Result> { - assert_eq!(layout.len(), layout.lda()); - let n = layout.len(); - let jobz = if calc_v { JobEv::All } else { JobEv::None }; - let mut eigs: Vec> = vec_uninit(n as usize); - - $( - let mut $rwork_ident: Vec> = vec_uninit(3 * n as usize - 2); - )* - - // calc work size - let mut info = 0; - let mut work_size = [Self::zero()]; - unsafe { - $evg( - &1, // ITYPE A*x = (lambda)*B*x - jobz.as_ptr(), - uplo.as_ptr(), - &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(b), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work_size), - &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - - // actual evg - let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); - let lwork = lwork as i32; - unsafe { - $evg( - &1, // ITYPE A*x = (lambda)*B*x - jobz.as_ptr(), - uplo.as_ptr(), - &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(b), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work), - &lwork, - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - let eigs = unsafe { eigs.assume_init() }; - Ok(eigs) - } - } - }; -} // impl_eigh! - -impl_eigh!(@real, f64, lapack_sys::dsygv_); -impl_eigh!(@real, f32, lapack_sys::ssygv_); -impl_eigh!(@complex, c64, lapack_sys::zhegv_); -impl_eigh!(@complex, c32, lapack_sys::chegv_); diff --git a/lax/src/lib.rs b/lax/src/lib.rs index f02532e5..b5f9f16c 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -131,7 +131,6 @@ pub trait Lapack: + Solve_ + Solveh_ + Cholesky_ - + EighGeneralized_ + Triangular_ + Tridiagonal_ + Rcond_ @@ -151,6 +150,15 @@ pub trait Lapack: uplo: UPLO, a: &mut [Self], ) -> Result>; + + /// Compute right eigenvalue and eigenvectors for a symmetric or hermite matrix + fn eigh_generalized( + calc_eigenvec: bool, + layout: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + b: &mut [Self], + ) -> Result>; } macro_rules! impl_lapack { @@ -177,6 +185,18 @@ macro_rules! impl_lapack { let work = EighWork::<$s>::new(calc_eigenvec, layout)?; work.eval(uplo, a) } + + fn eigh_generalized( + calc_eigenvec: bool, + layout: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + b: &mut [Self], + ) -> Result> { + use eigh_generalized::*; + let work = EighGeneralizedWork::<$s>::new(calc_eigenvec, layout)?; + work.eval(uplo, a, b) + } } }; }