diff --git a/lax/src/eig.rs b/lax/src/eig.rs index f73ed3f6..4e52a7fe 100644 --- a/lax/src/eig.rs +++ b/lax/src/eig.rs @@ -40,11 +40,11 @@ pub struct EigWork { pub jobvl: JobEv, /// Eigenvalues used in complex routines - pub eigs: Option>>, + pub eigs: Option>>, /// Real part of eigenvalues used in real routines - pub eigs_re: Option>>, + pub eigs_re: Option>>, /// Imaginary part of eigenvalues used in real routines - pub eigs_im: Option>>, + pub eigs_im: Option>>, /// Left eigenvectors pub vl: Option>>, @@ -71,6 +71,26 @@ pub struct EigRef<'work, T: Scalar> { pub vl: Option<&'work [T::Complex]>, } +impl EigWork { + /// Create complex eigenvalues from real and imaginary parts. + unsafe fn reconstruct_eigs(&mut self) { + let n = self.n as usize; + if self.eigs.is_none() { + self.eigs = Some(vec_uninit(n)); + } + let eigs = self.eigs.as_mut().unwrap(); + if let (Some(re), Some(im)) = (self.eigs_re.as_ref(), self.eigs_im.as_ref()) { + assert_eq!(re.len(), n); + assert_eq!(im.len(), n); + for i in 0..n { + eigs[i].write(T::complex(re[i].assume_init(), im[i].assume_init())); + } + } else { + panic!("Called without constructing eigs_re and eigs_im"); + } + } +} + pub trait EigWorkImpl: Sized { type Elem: Scalar; /// Create new working memory for eigenvalues compution. @@ -227,6 +247,109 @@ impl EigWorkImpl for EigWork { } } +impl EigWorkImpl for EigWork { + type Elem = f64; + + 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_deref_mut().unwrap_or(&mut [])), + &n, + AsPtr::as_mut_ptr(vr.as_deref_mut().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, + }) + } + + fn calc<'work>(&'work mut self, _a: &mut [f64]) -> Result> { + todo!() + } + + fn eval(mut self, a: &mut [f64]) -> Result> { + let lwork = self.work.len().to_i32().unwrap(); + let mut info = 0; + unsafe { + lapack_sys::dgeev_( + self.jobvl.as_ptr(), + self.jobvr.as_ptr(), + &self.n, + AsPtr::as_mut_ptr(a), + &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.vl.as_deref_mut().unwrap_or(&mut [])), + &self.n, + AsPtr::as_mut_ptr(self.vr.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 { self.eigs_re.unwrap().assume_init() }; + let eig_im = unsafe { self.eigs_im.unwrap().assume_init() }; + let vl = unsafe { self.vl.map(|v| v.assume_init()) }; + let vr = unsafe { self.vr.map(|v| v.assume_init()) }; + + // reconstruct eigenvalues + let eigs: Vec = eig_re + .iter() + .zip(eig_im.iter()) + .map(|(&re, &im)| c64::new(re, im)) + .collect(); + + Ok(Eig { eigs, vl, vr }) + } +} + macro_rules! impl_eig_complex { ($scalar:ty, $ev:path) => { impl Eig_ for $scalar {