WIP: EigWorkImpl for f64

Author: termoshtt

lax/src/eig.rs

@@ -40,11 +40,11 @@ pub struct EigWork<T: Scalar> {
     pub jobvl: JobEv,
 
     /// Eigenvalues used in complex routines
-    pub eigs: Option<Vec<MaybeUninit<T>>>,
+    pub eigs: Option<Vec<MaybeUninit<T::Complex>>>,
     /// Real part of eigenvalues used in real routines
-    pub eigs_re: Option<Vec<MaybeUninit<T>>>,
+    pub eigs_re: Option<Vec<MaybeUninit<T::Real>>>,
     /// Imaginary part of eigenvalues used in real routines
-    pub eigs_im: Option<Vec<MaybeUninit<T>>>,
+    pub eigs_im: Option<Vec<MaybeUninit<T::Real>>>,
 
     /// Left eigenvectors
     pub vl: Option<Vec<MaybeUninit<T>>>,
@@ -71,6 +71,26 @@ pub struct EigRef<'work, T: Scalar> {
     pub vl: Option<&'work [T::Complex]>,
 }
 
+impl<T: Scalar> EigWork<T> {
+    /// 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<c64> {
     }
 }
 
+impl EigWorkImpl for EigWork<f64> {
+    type Elem = f64;
+
+    fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
+        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<MaybeUninit<f64>> = vec_uninit(n as usize);
+        let mut eigs_im: Vec<MaybeUninit<f64>> = vec_uninit(n as usize);
+
+        let mut vl: Option<Vec<MaybeUninit<f64>>> = jobvl.then(|| vec_uninit((n * n) as usize));
+        let mut vr: Option<Vec<MaybeUninit<f64>>> = 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<MaybeUninit<f64>> = 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<EigRef<'work, f64>> {
+        todo!()
+    }
+
+    fn eval(mut self, a: &mut [f64]) -> Result<Eig<f64>> {
+        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<c64> = 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 {