diff --git a/lax/src/alloc.rs b/lax/src/alloc.rs index 5872bab4..63458818 100644 --- a/lax/src/alloc.rs +++ b/lax/src/alloc.rs @@ -33,18 +33,34 @@ impl_as_ptr!(MaybeUninit, lapack_sys::__BindgenComplex); impl_as_ptr!(MaybeUninit, lapack_sys::__BindgenComplex); pub(crate) trait VecAssumeInit { - type Target; - unsafe fn assume_init(self) -> Self::Target; + type Elem; + unsafe fn assume_init(self) -> Vec; + + /// An replacement of unstable API + /// https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#method.slice_assume_init_ref + unsafe fn slice_assume_init_ref(&self) -> &[Self::Elem]; + + /// An replacement of unstable API + /// https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#method.slice_assume_init_mut + unsafe fn slice_assume_init_mut(&mut self) -> &mut [Self::Elem]; } impl VecAssumeInit for Vec> { - type Target = Vec; - unsafe fn assume_init(self) -> Self::Target { + type Elem = T; + unsafe fn assume_init(self) -> Vec { // FIXME use Vec::into_raw_parts instead after stablized // https://doc.rust-lang.org/std/vec/struct.Vec.html#method.into_raw_parts let mut me = std::mem::ManuallyDrop::new(self); Vec::from_raw_parts(me.as_mut_ptr() as *mut T, me.len(), me.capacity()) } + + unsafe fn slice_assume_init_ref(&self) -> &[T] { + std::slice::from_raw_parts(self.as_ptr() as *const T, self.len()) + } + + unsafe fn slice_assume_init_mut(&mut self) -> &mut [T] { + std::slice::from_raw_parts_mut(self.as_mut_ptr() as *mut T, self.len()) + } } /// Create a vector without initialization diff --git a/lax/src/eig.rs b/lax/src/eig.rs index 6fcbc26a..2c24a842 100644 --- a/lax/src/eig.rs +++ b/lax/src/eig.rs @@ -1,9 +1,45 @@ +//! Eigenvalue problem for general matricies + use crate::{error::*, layout::MatrixLayout, *}; use cauchy::*; use num_traits::{ToPrimitive, Zero}; #[cfg_attr(doc, katexit::katexit)] /// Eigenvalue problem for general matrix +/// +/// To manage memory more strictly, use [EigWork]. +/// +/// Right and Left eigenvalue problem +/// ---------------------------------- +/// LAPACK can solve both right eigenvalue problem +/// $$ +/// AV_R = V_R \Lambda +/// $$ +/// where $V_R = \left( v_R^1, \cdots, v_R^n \right)$ are right eigenvectors +/// and left eigenvalue problem +/// $$ +/// V_L^\dagger A = V_L^\dagger \Lambda +/// $$ +/// where $V_L = \left( v_L^1, \cdots, v_L^n \right)$ are left eigenvectors +/// and eigenvalues +/// $$ +/// \Lambda = \begin{pmatrix} +/// \lambda_1 & & 0 \\\\ +/// & \ddots & \\\\ +/// 0 & & \lambda_n +/// \end{pmatrix} +/// $$ +/// which satisfies $A v_R^i = \lambda_i v_R^i$ and +/// $\left(v_L^i\right)^\dagger A = \lambda_i \left(v_L^i\right)^\dagger$ +/// for column-major matrices, although row-major matrices are not supported. +/// Since a row-major matrix can be interpreted +/// as a transpose of a column-major matrix, +/// this transforms right eigenvalue problem to left one: +/// +/// $$ +/// A^\dagger V = V Λ ⟺ V^\dagger A = Λ V^\dagger +/// $$ +/// pub trait Eig_: Scalar { /// Compute right eigenvalue and eigenvectors $Ax = \lambda x$ /// @@ -21,25 +57,116 @@ pub trait Eig_: Scalar { ) -> Result<(Vec, Vec)>; } -macro_rules! impl_eig_complex { - ($scalar:ty, $ev:path) => { - impl Eig_ for $scalar { +macro_rules! impl_eig { + ($s:ty) => { + impl Eig_ for $s { fn eig( calc_v: bool, l: MatrixLayout, a: &mut [Self], ) -> Result<(Vec, Vec)> { + let work = EigWork::<$s>::new(calc_v, l)?; + let EigOwned { eigs, vr, vl } = work.eval(a)?; + Ok((eigs, vr.or(vl).unwrap_or_default())) + } + } + }; +} +impl_eig!(c64); +impl_eig!(c32); +impl_eig!(f64); +impl_eig!(f32); + +/// Working memory for [Eig_] +#[non_exhaustive] +pub struct EigWork { + /// Problem size + pub n: i32, + /// Compute right eigenvectors or not + pub jobvr: JobEv, + /// Compute left eigenvectors or not + pub jobvl: JobEv, + + /// Eigenvalues + pub eigs: Vec>, + /// Real part of eigenvalues used in real routines + pub eigs_re: Option>>, + /// Imaginary part of eigenvalues used in real routines + pub eigs_im: Option>>, + + /// Left eigenvectors + pub vc_l: Option>>, + /// Left eigenvectors used in real routines + pub vr_l: Option>>, + /// Right eigenvectors + pub vc_r: Option>>, + /// Right eigenvectors used in real routines + pub vr_r: Option>>, + + /// Working memory + pub work: Vec>, + /// Working memory with `T::Real` + pub rwork: Option>>, +} + +impl EigWork +where + T: Scalar, + EigWork: EigWorkImpl, +{ + /// Create new working memory for eigenvalues compution. + pub fn new(calc_v: bool, l: MatrixLayout) -> Result { + EigWorkImpl::new(calc_v, l) + } + + /// Compute eigenvalues and vectors on this working memory. + pub fn calc(&mut self, a: &mut [T]) -> Result> { + EigWorkImpl::calc(self, a) + } + + /// Compute eigenvalues and vectors by consuming this working memory. + pub fn eval(self, a: &mut [T]) -> Result> { + EigWorkImpl::eval(self, a) + } +} + +/// Owned result of eigenvalue problem by [EigWork::eval] +#[derive(Debug, Clone, PartialEq)] +pub struct EigOwned { + /// Eigenvalues + pub eigs: Vec, + /// Right eigenvectors + pub vr: Option>, + /// Left eigenvectors + pub vl: Option>, +} + +/// Reference result of eigenvalue problem by [EigWork::calc] +#[derive(Debug, Clone, PartialEq)] +pub struct EigRef<'work, T: Scalar> { + /// Eigenvalues + pub eigs: &'work [T::Complex], + /// Right eigenvectors + pub vr: Option<&'work [T::Complex]>, + /// Left eigenvectors + pub vl: Option<&'work [T::Complex]>, +} + +/// Helper trait for implementing [EigWork] methods +pub trait EigWorkImpl: Sized { + type Elem: Scalar; + fn new(calc_v: bool, l: MatrixLayout) -> Result; + fn calc<'work>(&'work mut self, a: &mut [Self::Elem]) -> Result>; + fn eval(self, a: &mut [Self::Elem]) -> Result>; +} + +macro_rules! impl_eig_work_c { + ($c:ty, $ev:path) => { + impl EigWorkImpl for EigWork<$c> { + type Elem = $c; + + fn new(calc_v: bool, l: MatrixLayout) -> Result { let (n, _) = l.size(); - // LAPACK assumes a column-major input. A row-major input can - // be interpreted as the transpose of a column-major input. So, - // for row-major inputs, we we want to solve the following, - // given the column-major input `A`: - // - // A^T V = V Λ ⟺ V^T A = Λ V^T ⟺ conj(V)^H A = Λ conj(V)^H - // - // So, in this case, the right eigenvectors are the conjugates - // of the left eigenvectors computed with `A`, and the - // eigenvalues are the eigenvalues computed with `A`. let (jobvl, jobvr) = if calc_v { match l { MatrixLayout::C { .. } => (JobEv::All, JobEv::None), @@ -48,28 +175,26 @@ macro_rules! impl_eig_complex { } else { (JobEv::None, JobEv::None) }; - let mut eigs: Vec> = vec_uninit(n as usize); - let mut rwork: Vec> = vec_uninit(2 * n as usize); + let mut eigs = vec_uninit(n as usize); + let mut rwork = vec_uninit(2 * n as usize); - let mut vl: Option>> = - jobvl.then(|| vec_uninit((n * n) as usize)); - let mut vr: Option>> = - jobvr.then(|| vec_uninit((n * n) as usize)); + let mut vc_l = jobvl.then(|| vec_uninit((n * n) as usize)); + let mut vc_r = jobvr.then(|| vec_uninit((n * n) as usize)); // calc work size let mut info = 0; - let mut work_size = [Self::zero()]; + let mut work_size = [<$c>::zero()]; unsafe { $ev( jobvl.as_ptr(), jobvr.as_ptr(), &n, - AsPtr::as_mut_ptr(a), + std::ptr::null_mut(), &n, AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), + AsPtr::as_mut_ptr(vc_l.as_deref_mut().unwrap_or(&mut [])), &n, - AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), + AsPtr::as_mut_ptr(vc_r.as_deref_mut().unwrap_or(&mut [])), &n, AsPtr::as_mut_ptr(&mut work_size), &(-1), @@ -79,75 +204,91 @@ macro_rules! impl_eig_complex { }; info.as_lapack_result()?; - // actal ev let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); - let lwork = lwork as i32; + let work: Vec> = vec_uninit(lwork); + Ok(Self { + n, + jobvl, + jobvr, + eigs, + eigs_re: None, + eigs_im: None, + rwork: Some(rwork), + vc_l, + vc_r, + vr_l: None, + vr_r: None, + work, + }) + } + + fn calc<'work>( + &'work mut self, + a: &mut [Self::Elem], + ) -> Result> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; unsafe { $ev( - jobvl.as_ptr(), - jobvr.as_ptr(), - &n, + self.jobvl.as_ptr(), + self.jobvr.as_ptr(), + &self.n, AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), - &n, - AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), - &n, - AsPtr::as_mut_ptr(&mut work), + &self.n, + AsPtr::as_mut_ptr(&mut self.eigs), + AsPtr::as_mut_ptr(self.vc_l.as_deref_mut().unwrap_or(&mut [])), + &self.n, + AsPtr::as_mut_ptr(self.vc_r.as_deref_mut().unwrap_or(&mut [])), + &self.n, + AsPtr::as_mut_ptr(&mut self.work), &lwork, - AsPtr::as_mut_ptr(&mut rwork), + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), &mut info, ) }; info.as_lapack_result()?; - - let eigs = unsafe { eigs.assume_init() }; - let vr = unsafe { vr.map(|v| v.assume_init()) }; - let mut vl = unsafe { vl.map(|v| v.assume_init()) }; - // Hermite conjugate - if jobvl.is_calc() { - for c in vl.as_mut().unwrap().iter_mut() { - c.im = -c.im; + if let Some(vl) = self.vc_l.as_mut() { + for value in vl { + let value = unsafe { value.assume_init_mut() }; + value.im = -value.im; } } + Ok(EigRef { + eigs: unsafe { self.eigs.slice_assume_init_ref() }, + vl: self + .vc_l + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }), + vr: self + .vc_r + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }), + }) + } - Ok((eigs, vr.or(vl).unwrap_or(Vec::new()))) + fn eval(mut self, a: &mut [Self::Elem]) -> Result> { + let _eig_ref = self.calc(a)?; + Ok(EigOwned { + eigs: unsafe { self.eigs.assume_init() }, + vl: self.vc_l.map(|v| unsafe { v.assume_init() }), + vr: self.vc_r.map(|v| unsafe { v.assume_init() }), + }) } } }; } -impl_eig_complex!(c64, lapack_sys::zgeev_); -impl_eig_complex!(c32, lapack_sys::cgeev_); +impl_eig_work_c!(c32, lapack_sys::cgeev_); +impl_eig_work_c!(c64, lapack_sys::zgeev_); -macro_rules! impl_eig_real { - ($scalar:ty, $ev:path) => { - impl Eig_ for $scalar { - fn eig( - calc_v: bool, - l: MatrixLayout, - a: &mut [Self], - ) -> Result<(Vec, Vec)> { +macro_rules! impl_eig_work_r { + ($f:ty, $ev:path) => { + impl EigWorkImpl for EigWork<$f> { + type Elem = $f; + + fn new(calc_v: bool, l: MatrixLayout) -> Result { let (n, _) = l.size(); - // LAPACK assumes a column-major input. A row-major input can - // be interpreted as the transpose of a column-major input. So, - // for row-major inputs, we we want to solve the following, - // given the column-major input `A`: - // - // A^T V = V Λ ⟺ V^T A = Λ V^T ⟺ conj(V)^H A = Λ conj(V)^H - // - // So, in this case, the right eigenvectors are the conjugates - // of the left eigenvectors computed with `A`, and the - // eigenvalues are the eigenvalues computed with `A`. - // - // We could conjugate the eigenvalues instead of the - // eigenvectors, but we have to reconstruct the eigenvectors - // into new matrices anyway, and by not modifying the - // eigenvalues, we preserve the nice ordering specified by - // `sgeev`/`dgeev`. let (jobvl, jobvr) = if calc_v { match l { MatrixLayout::C { .. } => (JobEv::All, JobEv::None), @@ -156,29 +297,28 @@ macro_rules! impl_eig_real { } else { (JobEv::None, JobEv::None) }; - let mut eig_re: Vec> = vec_uninit(n as usize); - let mut eig_im: Vec> = vec_uninit(n as usize); - - let mut vl: Option>> = - jobvl.then(|| vec_uninit((n * n) as usize)); - let mut vr: Option>> = - jobvr.then(|| vec_uninit((n * n) as usize)); + let mut eigs_re = vec_uninit(n as usize); + let mut eigs_im = vec_uninit(n as usize); + let mut vr_l = jobvl.then(|| vec_uninit((n * n) as usize)); + let mut vr_r = jobvr.then(|| vec_uninit((n * n) as usize)); + let vc_l = jobvl.then(|| vec_uninit((n * n) as usize)); + let vc_r = jobvr.then(|| vec_uninit((n * n) as usize)); // calc work size let mut info = 0; - let mut work_size: [Self; 1] = [0.0]; + let mut work_size: [$f; 1] = [0.0]; unsafe { $ev( jobvl.as_ptr(), jobvr.as_ptr(), &n, - AsPtr::as_mut_ptr(a), + std::ptr::null_mut(), &n, - AsPtr::as_mut_ptr(&mut eig_re), - AsPtr::as_mut_ptr(&mut eig_im), - AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), + AsPtr::as_mut_ptr(&mut eigs_re), + AsPtr::as_mut_ptr(&mut eigs_im), + AsPtr::as_mut_ptr(vr_l.as_deref_mut().unwrap_or(&mut [])), &n, - AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), + AsPtr::as_mut_ptr(vr_r.as_deref_mut().unwrap_or(&mut [])), &n, AsPtr::as_mut_ptr(&mut work_size), &(-1), @@ -189,93 +329,153 @@ macro_rules! impl_eig_real { // 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(Self { + n, + jobvr, + jobvl, + eigs: vec_uninit(n as usize), + eigs_re: Some(eigs_re), + eigs_im: Some(eigs_im), + rwork: None, + vr_l, + vr_r, + vc_l, + vc_r, + work, + }) + } + + fn calc<'work>( + &'work mut self, + a: &mut [Self::Elem], + ) -> Result> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; unsafe { $ev( - jobvl.as_ptr(), - jobvr.as_ptr(), - &n, + self.jobvl.as_ptr(), + self.jobvr.as_ptr(), + &self.n, AsPtr::as_mut_ptr(a), - &n, - AsPtr::as_mut_ptr(&mut eig_re), - AsPtr::as_mut_ptr(&mut eig_im), - AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), - &n, - AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), - &n, - AsPtr::as_mut_ptr(&mut work), + &self.n, + AsPtr::as_mut_ptr(self.eigs_re.as_mut().unwrap()), + AsPtr::as_mut_ptr(self.eigs_im.as_mut().unwrap()), + AsPtr::as_mut_ptr(self.vr_l.as_deref_mut().unwrap_or(&mut [])), + &self.n, + AsPtr::as_mut_ptr(self.vr_r.as_deref_mut().unwrap_or(&mut [])), + &self.n, + AsPtr::as_mut_ptr(&mut self.work), &lwork, &mut info, ) }; info.as_lapack_result()?; - let eig_re = unsafe { eig_re.assume_init() }; - let eig_im = unsafe { eig_im.assume_init() }; - let vl = unsafe { vl.map(|v| v.assume_init()) }; - let vr = unsafe { vr.map(|v| v.assume_init()) }; - - // reconstruct eigenvalues - let eigs: Vec = eig_re - .iter() - .zip(eig_im.iter()) - .map(|(&re, &im)| Self::complex(re, im)) - .collect(); + let eigs_re = self + .eigs_re + .as_ref() + .map(|e| unsafe { e.slice_assume_init_ref() }) + .unwrap(); + let eigs_im = self + .eigs_im + .as_ref() + .map(|e| unsafe { e.slice_assume_init_ref() }) + .unwrap(); + reconstruct_eigs(eigs_re, eigs_im, &mut self.eigs); - if !calc_v { - return Ok((eigs, Vec::new())); + if let Some(v) = self.vr_l.as_ref() { + let v = unsafe { v.slice_assume_init_ref() }; + reconstruct_eigenvectors(true, eigs_im, v, self.vc_l.as_mut().unwrap()); } - - // 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). - // - // In the C-layout case, we need the conjugates of the left - // eigenvectors, so the signs should be reversed. - - let n = n as usize; - let v = vr.or(vl).unwrap(); - let mut eigvecs: Vec> = vec_uninit(n * n); - let mut col = 0; - while col < n { - if eig_im[col] == 0. { - // The corresponding eigenvalue is real. - for row in 0..n { - let re = v[row + col * n]; - eigvecs[row + col * n].write(Self::complex(re, 0.)); - } - col += 1; - } else { - // This is a complex conjugate pair. - assert!(col + 1 < n); - for row in 0..n { - let re = v[row + col * n]; - let mut im = v[row + (col + 1) * n]; - if jobvl.is_calc() { - im = -im; - } - eigvecs[row + col * n].write(Self::complex(re, im)); - eigvecs[row + (col + 1) * n].write(Self::complex(re, -im)); - } - col += 2; - } + if let Some(v) = self.vr_r.as_ref() { + let v = unsafe { v.slice_assume_init_ref() }; + reconstruct_eigenvectors(false, eigs_im, v, self.vc_r.as_mut().unwrap()); } - let eigvecs = unsafe { eigvecs.assume_init() }; - Ok((eigs, eigvecs)) + Ok(EigRef { + eigs: unsafe { self.eigs.slice_assume_init_ref() }, + vl: self + .vc_l + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }), + vr: self + .vc_r + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }), + }) + } + + fn eval(mut self, a: &mut [Self::Elem]) -> Result> { + let _eig_ref = self.calc(a)?; + Ok(EigOwned { + eigs: unsafe { self.eigs.assume_init() }, + vl: self.vc_l.map(|v| unsafe { v.assume_init() }), + vr: self.vc_r.map(|v| unsafe { v.assume_init() }), + }) } } }; } +impl_eig_work_r!(f32, lapack_sys::sgeev_); +impl_eig_work_r!(f64, lapack_sys::dgeev_); -impl_eig_real!(f64, lapack_sys::dgeev_); -impl_eig_real!(f32, lapack_sys::sgeev_); +/// 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). +/// +/// In the C-layout case, we need the conjugates of the left +/// eigenvectors, so the signs should be reversed. +fn reconstruct_eigenvectors( + take_hermite_conjugate: bool, + eig_im: &[T], + vr: &[T], + vc: &mut [MaybeUninit], +) { + let n = eig_im.len(); + assert_eq!(vr.len(), n * n); + assert_eq!(vc.len(), n * n); + + let mut col = 0; + while col < n { + if eig_im[col].is_zero() { + // The corresponding eigenvalue is real. + for row in 0..n { + let re = vr[row + col * n]; + vc[row + col * n].write(T::complex(re, T::zero())); + } + col += 1; + } else { + // This is a complex conjugate pair. + assert!(col + 1 < n); + for row in 0..n { + let re = vr[row + col * n]; + let mut im = vr[row + (col + 1) * n]; + if take_hermite_conjugate { + im = -im; + } + vc[row + col * n].write(T::complex(re, im)); + vc[row + (col + 1) * n].write(T::complex(re, -im)); + } + col += 2; + } + } +} + +/// Create complex eigenvalues from real and imaginary parts. +fn reconstruct_eigs(re: &[T], im: &[T], eigs: &mut [MaybeUninit]) { + let n = eigs.len(); + assert_eq!(re.len(), n); + assert_eq!(im.len(), n); + for i in 0..n { + eigs[i].write(T::complex(re[i], im[i])); + } +} diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 8f697fbc..397970f7 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -84,9 +84,10 @@ pub mod error; pub mod flags; pub mod layout; +pub mod eig; + mod alloc; mod cholesky; -mod eig; mod eigh; mod least_squares; mod opnorm; @@ -100,7 +101,7 @@ mod triangular; mod tridiagonal; pub use self::cholesky::*; -pub use self::eig::*; +pub use self::eig::Eig_; pub use self::eigh::*; pub use self::flags::*; pub use self::least_squares::*;