EigWorkImpl for c64

Author: termoshtt

lax/src/eig.rs

@@ -57,6 +57,170 @@ pub struct EigWork<T: Scalar> {
     pub rwork: Option<Vec<MaybeUninit<T::Real>>>,
 }
 
+pub trait EigWorkImpl: Sized {
+    type Elem: Scalar;
+    /// Create new working memory for eigenvalues compution.
+    fn new(calc_v: bool, l: MatrixLayout) -> Result<Self>;
+    /// Compute eigenvalues and vectors on this working memory.
+    fn calc<'work>(
+        &'work mut self,
+        a: &mut [Self::Elem],
+    ) -> Result<(
+        &'work [<Self::Elem as Scalar>::Complex],
+        Option<&'work [<Self::Elem as Scalar>::Complex]>,
+    )>;
+    /// Compute eigenvalues and vectors by consuming this working memory.
+    fn eval(
+        self,
+        a: &mut [Self::Elem],
+    ) -> Result<(
+        Vec<<Self::Elem as Scalar>::Complex>,
+        Option<Vec<<Self::Elem as Scalar>::Complex>>,
+    )>;
+}
+
+impl EigWorkImpl for EigWork<c64> {
+    type Elem = c64;
+
+    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,
+        })
+    }
+
+    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))
+    }
+
+    fn eval(mut self, a: &mut [c64]) -> Result<(Vec<c64>, Option<Vec<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.map(|v| unsafe { v.assume_init() }).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, self.vr) {
+            (Some(v), None) | (None, Some(v)) => Some(unsafe { v.assume_init() }),
+            (None, None) => None,
+            _ => unreachable!(),
+        };
+        Ok((eigs, v))
+    }
+}
+
 macro_rules! impl_eig_complex {
     ($scalar:ty, $ev:path) => {
         impl Eig_ for $scalar {