diff --git a/lax/src/eigh.rs b/lax/src/eigh.rs index 54af63f7..429aaa75 100644 --- a/lax/src/eigh.rs +++ b/lax/src/eigh.rs @@ -1,180 +1,190 @@ +//! 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, + pub eigs: Vec>, + pub work: Vec>, + pub rwork: Option>>, +} - /// 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>; +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>>; } -macro_rules! impl_eigh { - (@real, $scalar:ty, $ev:path, $evg:path) => { - impl_eigh!(@body, $scalar, $ev, $evg, ); - }; - (@complex, $scalar:ty, $ev:path, $evg:path) => { - impl_eigh!(@body, $scalar, $ev, $evg, rwork); - }; - (@body, $scalar:ty, $ev:path, $evg:path, $($rwork_ident:ident),*) => { - impl Eigh_ for $scalar { - fn eigh( - calc_v: bool, - layout: MatrixLayout, - uplo: UPLO, - a: &mut [Self], - ) -> Result> { +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_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 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::zero()]; + let mut work_size = [Self::Elem::zero()]; unsafe { $ev( - jobz.as_ptr() , - uplo.as_ptr(), + jobz.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO &n, - AsPtr::as_mut_ptr(a), + 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_ident),)* + AsPtr::as_mut_ptr(&mut rwork), &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; + 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( - jobz.as_ptr(), + self.jobz.as_ptr(), uplo.as_ptr(), - &n, + &self.n, AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(&mut work), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), &lwork, - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), &mut info, ); } info.as_lapack_result()?; - - let eigs = unsafe { eigs.assume_init() }; - Ok(eigs) + Ok(unsafe { self.eigs.slice_assume_init_ref() }) } - fn eigh_generalized( - calc_v: bool, - layout: MatrixLayout, + fn eval( + mut self, 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); + a: &mut [Self::Elem], + ) -> Result::Real>> { + let _eig = self.calc(uplo, a)?; + Ok(unsafe { self.eigs.assume_init() }) + } + } + }; +} +impl_eigh_work_c!(c64, lapack_sys::zheev_); +impl_eigh_work_c!(c32, lapack_sys::cheev_); - $( - let mut $rwork_ident: Vec> = vec_uninit(3 * n as usize - 2); - )* +macro_rules! impl_eigh_work_r { + ($f:ty, $ev:path) => { + impl EighWorkImpl for EighWork<$f> { + type Elem = $f; - // calc work size + 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::zero()]; + let mut work_size = [Self::Elem::zero()]; unsafe { - $evg( - &1, // ITYPE A*x = (lambda)*B*x + $ev( jobz.as_ptr(), - uplo.as_ptr(), + UPLO::Upper.as_ptr(), // dummy, working memory is not affected by UPLO &n, - AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(b), + 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_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; + 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 { - $evg( - &1, // ITYPE A*x = (lambda)*B*x - jobz.as_ptr(), + $ev( + self.jobz.as_ptr(), uplo.as_ptr(), - &n, + &self.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), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(&mut self.work), &lwork, - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* &mut info, ); } info.as_lapack_result()?; - let eigs = unsafe { eigs.assume_init() }; - Ok(eigs) + 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! - -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_work_r!(f64, lapack_sys::dsyev_); +impl_eigh_work_r!(f32, lapack_sys::ssyev_); diff --git a/lax/src/eigh_generalized.rs b/lax/src/eigh_generalized.rs new file mode 100644 index 00000000..1771cee8 --- /dev/null +++ b/lax/src/eigh_generalized.rs @@ -0,0 +1,216 @@ +//! 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_); diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 00393916..b5f9f16c 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::*; @@ -129,24 +131,39 @@ pub trait Lapack: + Solve_ + Solveh_ + Cholesky_ - + Eigh_ + 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>; + + /// 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 { ($s:ty) => { impl Lapack for $s { - /// Compute right eigenvalue and eigenvectors fn eig( calc_v: bool, l: MatrixLayout, @@ -157,6 +174,29 @@ 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) + } + + 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) + } } }; }