From ede3cd9da761b957a10ebe3cd7c485d46b01f7f4 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Fri, 23 Sep 2022 18:08:56 +0900 Subject: [PATCH] WIP: EigWorkImpl for f64 --- lax/src/eig.rs | 183 +++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 153 insertions(+), 30 deletions(-) diff --git a/lax/src/eig.rs b/lax/src/eig.rs index f73ed3f6..93465cd9 100644 --- a/lax/src/eig.rs +++ b/lax/src/eig.rs @@ -40,16 +40,18 @@ pub struct EigWork { pub jobvl: JobEv, /// Eigenvalues used in complex routines - pub eigs: Option>>, + pub eigs: Vec>, /// 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>>, + pub vc_l: Option>>, + pub vr_l: Option>>, /// Right eigenvectors - pub vr: Option>>, + pub vc_r: Option>>, + pub vr_r: Option>>, /// Working memory pub work: Vec>, @@ -97,8 +99,8 @@ impl EigWorkImpl for EigWork { 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)); + let mut vc_l: Option>> = jobvl.then(|| vec_uninit((n * n) as usize)); + let mut vc_r: Option>> = jobvr.then(|| vec_uninit((n * n) as usize)); // calc work size let mut info = 0; @@ -111,9 +113,9 @@ impl EigWorkImpl for EigWork { std::ptr::null_mut(), &n, AsPtr::as_mut_ptr(&mut eigs), - AsPtr::as_mut_ptr(vl.as_deref_mut().unwrap_or(&mut [])), + AsPtr::as_mut_ptr(vc_l.as_deref_mut().unwrap_or(&mut [])), &n, - AsPtr::as_mut_ptr(vr.as_deref_mut().unwrap_or(&mut [])), + AsPtr::as_mut_ptr(vc_r.as_deref_mut().unwrap_or(&mut [])), &n, AsPtr::as_mut_ptr(&mut work_size), &(-1), @@ -129,12 +131,14 @@ impl EigWorkImpl for EigWork { n, jobvl, jobvr, - eigs: Some(eigs), + eigs, eigs_re: None, eigs_im: None, rwork: Some(rwork), - vl, - vr, + vc_l, + vc_r, + vr_l: None, + vr_r: None, work, }) } @@ -149,10 +153,10 @@ impl EigWorkImpl for EigWork { &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_deref_mut().unwrap_or(&mut [])), + 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.vr.as_deref_mut().unwrap_or(&mut [])), + AsPtr::as_mut_ptr(self.vc_r.as_deref_mut().unwrap_or(&mut [])), &self.n, AsPtr::as_mut_ptr(&mut self.work), &lwork, @@ -162,14 +166,10 @@ impl EigWorkImpl for EigWork { }; info.as_lapack_result()?; - let eigs = self - .eigs - .as_ref() - .map(|v| unsafe { v.slice_assume_init_ref() }) - .unwrap(); + let eigs = unsafe { self.eigs.slice_assume_init_ref() }; // Hermite conjugate - if let Some(vl) = self.vl.as_mut() { + if let Some(vl) = self.vc_l.as_mut() { for value in vl { let value = unsafe { value.assume_init_mut() }; value.im = -value.im; @@ -178,11 +178,11 @@ impl EigWorkImpl for EigWork { Ok(EigRef { eigs, vl: self - .vl + .vc_l .as_ref() .map(|v| unsafe { v.slice_assume_init_ref() }), vr: self - .vr + .vc_r .as_ref() .map(|v| unsafe { v.slice_assume_init_ref() }), }) @@ -198,10 +198,10 @@ impl EigWorkImpl for EigWork { &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_deref_mut().unwrap_or(&mut [])), + 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.vr.as_deref_mut().unwrap_or(&mut [])), + AsPtr::as_mut_ptr(self.vc_r.as_deref_mut().unwrap_or(&mut [])), &self.n, AsPtr::as_mut_ptr(&mut self.work), &lwork, @@ -210,10 +210,10 @@ impl EigWorkImpl for EigWork { ) }; info.as_lapack_result()?; - let eigs = self.eigs.map(|v| unsafe { v.assume_init() }).unwrap(); + let eigs = unsafe { self.eigs.assume_init() }; // Hermite conjugate - if let Some(vl) = self.vl.as_mut() { + if let Some(vl) = self.vc_l.as_mut() { for value in vl { let value = unsafe { value.assume_init_mut() }; value.im = -value.im; @@ -221,10 +221,123 @@ impl EigWorkImpl for EigWork { } Ok(Eig { eigs, - vl: self.vl.map(|v| unsafe { v.assume_init() }), - vr: self.vr.map(|v| unsafe { v.assume_init() }), + vl: self.vc_l.map(|v| unsafe { v.assume_init() }), + vr: self.vc_r.map(|v| unsafe { v.assume_init() }), + }) + } +} + +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 vr_l: Option>> = jobvl.then(|| vec_uninit((n * n) as usize)); + let mut vr_r: 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(vr_l.as_deref_mut().unwrap_or(&mut [])), + &n, + AsPtr::as_mut_ptr(vr_r.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: vec_uninit(n as usize), + eigs_re: Some(eigs_re), + eigs_im: Some(eigs_im), + rwork: None, + vr_l, + vr_r, + vc_l: None, + vc_r: None, + 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.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 eigs_re = unsafe { self.eigs_re.unwrap().assume_init() }; + let eigs_im = unsafe { self.eigs_im.unwrap().assume_init() }; + + let n = self.n as usize; + let vl = self.vr_l.map(|v| { + let v = unsafe { v.assume_init() }; + let mut vc = vec_uninit(n * n); + reconstruct_eigenvectors(false, &eigs_im, &v, &mut vc); + unsafe { vc.assume_init() } + }); + let vr = self.vr_r.map(|v| { + let v = unsafe { v.assume_init() }; + let mut vc = vec_uninit(n * n); + reconstruct_eigenvectors(true, &eigs_im, &v, &mut vc); + unsafe { vc.assume_init() } + }); + + reconstruct_eigs(&eigs_re, &eigs_im, &mut self.eigs); + let eigs = unsafe { self.eigs.assume_init() }; + + Ok(Eig { eigs, vl, vr }) + } } macro_rules! impl_eig_complex { @@ -497,3 +610,13 @@ fn reconstruct_eigenvectors( } } } + +/// 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])); + } +}