diff --git a/lax/src/eig.rs b/lax/src/eig.rs index f8747ee2..d3119919 100644 --- a/lax/src/eig.rs +++ b/lax/src/eig.rs @@ -4,6 +4,17 @@ use num_traits::{ToPrimitive, Zero}; #[cfg_attr(doc, katexit::katexit)] /// Eigenvalue problem for general matrix +/// +/// 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`. pub trait Eig_: Scalar { /// Compute right eigenvalue and eigenvectors $Ax = \lambda x$ /// @@ -21,6 +32,205 @@ pub trait Eig_: Scalar { ) -> Result<(Vec, Vec)>; } +/// Working memory for [Eig_] +#[derive(Debug, Clone)] +pub struct EigWork { + pub n: i32, + pub jobvr: JobEv, + pub jobvl: JobEv, + + /// Eigenvalues used in complex routines + pub eigs: Option>>, + /// 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 vl: Option>>, + /// Right eigenvectors + pub vr: Option>>, + + /// Working memory + pub work: Vec>, + /// Working memory with `T::Real` + pub rwork: Option>>, +} + +impl EigWork { + pub fn new(calc_v: bool, l: MatrixLayout) -> Result { + let (n, _) = l.size(); + let (jobvl, jobvr) = if calc_v { + match l { + MatrixLayout::C { .. } => (JobEv::All, JobEv::None), + MatrixLayout::F { .. } => (JobEv::None, JobEv::All), + } + } 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 vl: Option>> = jobvl.then(|| vec_uninit((n * n) as usize)); + let mut vr: Option>> = jobvr.then(|| vec_uninit((n * n) as usize)); + + // calc work size + let mut info = 0; + let mut work_size = [c64::zero()]; + unsafe { + lapack_sys::zgeev_( + jobvl.as_ptr(), + jobvr.as_ptr(), + &n, + 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 [])), + &n, + AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])), + &n, + 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> = vec_uninit(lwork); + Ok(Self { + n, + jobvl, + jobvr, + eigs: Some(eigs), + eigs_re: None, + eigs_im: None, + rwork: Some(rwork), + vl, + vr, + work, + }) + } + + /// Compute eigenvalues and vectors on this working memory. + pub fn calc<'work>( + &'work mut self, + a: &mut [c64], + ) -> Result<(&'work [c64], Option<&'work [c64]>)> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + lapack_sys::zgeev_( + self.jobvl.as_ptr(), + self.jobvr.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &self.n, + AsPtr::as_mut_ptr(self.eigs.as_mut().unwrap()), + AsPtr::as_mut_ptr( + self.vl + .as_mut() + .map(|v| v.as_mut_slice()) + .unwrap_or(&mut []), + ), + &self.n, + AsPtr::as_mut_ptr( + self.vr + .as_mut() + .map(|v| v.as_mut_slice()) + .unwrap_or(&mut []), + ), + &self.n, + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + &mut info, + ) + }; + info.as_lapack_result()?; + + let eigs = self + .eigs + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }) + .unwrap(); + + // Hermite conjugate + if let Some(vl) = self.vl.as_mut() { + for value in vl { + let value = unsafe { value.assume_init_mut() }; + value.im = -value.im; + } + } + let v = match (self.vl.as_ref(), self.vr.as_ref()) { + (Some(v), None) | (None, Some(v)) => Some(unsafe { v.slice_assume_init_ref() }), + (None, None) => None, + _ => unreachable!(), + }; + Ok((eigs, v)) + } +} + +impl EigWork { + pub fn new(calc_v: bool, l: MatrixLayout) -> Result { + let (n, _) = l.size(); + let (jobvl, jobvr) = if calc_v { + match l { + MatrixLayout::C { .. } => (JobEv::All, JobEv::None), + MatrixLayout::F { .. } => (JobEv::None, JobEv::All), + } + } else { + (JobEv::None, JobEv::None) + }; + let mut eigs_re: Vec> = vec_uninit(n as usize); + let mut eigs_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)); + + // calc work size + let mut info = 0; + let mut work_size: [f64; 1] = [0.0]; + unsafe { + lapack_sys::dgeev_( + jobvl.as_ptr(), + jobvr.as_ptr(), + &n, + std::ptr::null_mut(), + &n, + AsPtr::as_mut_ptr(&mut eigs_re), + AsPtr::as_mut_ptr(&mut eigs_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_size), + &(-1), + &mut info, + ) + }; + info.as_lapack_result()?; + + // actual ev + let lwork = work_size[0].to_usize().unwrap(); + let work: Vec> = vec_uninit(lwork); + + Ok(Self { + n, + jobvr, + jobvl, + eigs: None, + eigs_re: Some(eigs_re), + eigs_im: Some(eigs_im), + rwork: None, + vl, + vr, + work, + }) + } +} + macro_rules! impl_eig_complex { ($scalar:ty, $ev:path) => { impl Eig_ for $scalar {