WIP: EigWork

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,201 @@ pub trait Eig_: Scalar {
     ) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
 }
 
+/// Working memory for [Eig_]
+#[derive(Debug, Clone)]
+pub struct EigWork<T: Scalar> {
+    pub n: i32,
+    pub jobvr: JobEv,
+    pub jobvl: JobEv,
+
+    /// Eigenvalues used in complex routines
+    pub eigs: Option<Vec<MaybeUninit<T>>>,
+    /// Real part of eigenvalues used in real routines
+    pub eigs_re: Option<Vec<MaybeUninit<T>>>,
+    /// Imaginary part of eigenvalues used in real routines
+    pub eigs_im: Option<Vec<MaybeUninit<T>>>,
+
+    /// Left eigenvectors
+    pub vl: Option<Vec<MaybeUninit<T>>>,
+    /// Right eigenvectors
+    pub vr: Option<Vec<MaybeUninit<T>>>,
+
+    /// Working memory
+    pub work: Vec<MaybeUninit<T>>,
+    /// Working memory with `T::Real`
+    pub rwork: Option<Vec<MaybeUninit<T::Real>>>,
+}
+
+impl EigWork<c64> {
+    pub 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: Vec<MaybeUninit<c64>> = vec_uninit(n as usize);
+        let mut rwork: Vec<MaybeUninit<f64>> = vec_uninit(2 * n as usize);
+
+        let mut vl: Option<Vec<MaybeUninit<c64>>> = jobvl.then(|| vec_uninit((n * n) as usize));
+        let mut vr: Option<Vec<MaybeUninit<c64>>> = 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_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),
+                AsPtr::as_mut_ptr(&mut rwork),
+                &mut info,
+            )
+        };
+        info.as_lapack_result()?;
+
+        let lwork = work_size[0].to_usize().unwrap();
+        let work: Vec<MaybeUninit<c64>> = 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_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,
+                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<f64> {
+    pub 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,
+        })
+    }
+}
+
 macro_rules! impl_eig_complex {
     ($scalar:ty, $ev:path) => {
         impl Eig_ for $scalar {