HouseholderWork for c64

rust-ndarray · Sep 25, 2022 · dd5cfba · dd5cfba
1 parent 10ae296
commit dd5cfba
Showing 1 changed file with 102 additions and 0 deletions.
diff --git a/lax/src/qr.rs b/lax/src/qr.rs
@@ -18,6 +18,108 @@ pub trait QR_: Sized {
     fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
 }
 
+pub struct HouseholderWork<T: Scalar> {
+    pub m: i32,
+    pub n: i32,
+    pub layout: MatrixLayout,
+    pub tau: Vec<MaybeUninit<T>>,
+    pub work: Vec<MaybeUninit<T>>,
+}
+
+pub trait HouseholderWorkImpl: Sized {
+    type Elem: Scalar;
+    fn new(l: MatrixLayout) -> Result<Self>;
+    fn calc(&mut self, a: &mut [Self::Elem]) -> Result<&[Self::Elem]>;
+    fn eval(self, a: &mut [Self::Elem]) -> Result<Vec<Self::Elem>>;
+}
+
+impl HouseholderWorkImpl for HouseholderWork<c64> {
+    type Elem = c64;
+
+    fn new(layout: MatrixLayout) -> Result<Self> {
+        let m = layout.lda();
+        let n = layout.len();
+        let k = m.min(n);
+        let mut tau = vec_uninit(k as usize);
+        let mut info = 0;
+        let mut work_size = [Self::Elem::zero()];
+        match layout {
+            MatrixLayout::F { .. } => unsafe {
+                lapack_sys::zgeqrf_(
+                    &m,
+                    &n,
+                    std::ptr::null_mut(),
+                    &m,
+                    AsPtr::as_mut_ptr(&mut tau),
+                    AsPtr::as_mut_ptr(&mut work_size),
+                    &(-1),
+                    &mut info,
+                )
+            },
+            MatrixLayout::C { .. } => unsafe {
+                lapack_sys::zgelqf_(
+                    &m,
+                    &n,
+                    std::ptr::null_mut(),
+                    &m,
+                    AsPtr::as_mut_ptr(&mut tau),
+                    AsPtr::as_mut_ptr(&mut work_size),
+                    &(-1),
+                    &mut info,
+                )
+            },
+        }
+        info.as_lapack_result()?;
+        let lwork = work_size[0].to_usize().unwrap();
+        let work = vec_uninit(lwork);
+        Ok(HouseholderWork {
+            n,
+            m,
+            layout,
+            tau,
+            work,
+        })
+    }
+
+    fn calc(&mut self, a: &mut [Self::Elem]) -> Result<&[Self::Elem]> {
+        let lwork = self.work.len().to_i32().unwrap();
+        let mut info = 0;
+        match self.layout {
+            MatrixLayout::F { .. } => unsafe {
+                lapack_sys::zgeqrf_(
+                    &self.m,
+                    &self.n,
+                    AsPtr::as_mut_ptr(a),
+                    &self.m,
+                    AsPtr::as_mut_ptr(&mut self.tau),
+                    AsPtr::as_mut_ptr(&mut self.work),
+                    &lwork,
+                    &mut info,
+                );
+            },
+            MatrixLayout::C { .. } => unsafe {
+                lapack_sys::zgelqf_(
+                    &self.m,
+                    &self.n,
+                    AsPtr::as_mut_ptr(a),
+                    &self.m,
+                    AsPtr::as_mut_ptr(&mut self.tau),
+                    AsPtr::as_mut_ptr(&mut self.work),
+                    &lwork,
+                    &mut info,
+                );
+            },
+        }
+        info.as_lapack_result()?;
+        Ok(unsafe { self.tau.slice_assume_init_ref() })
+    }
+
+    fn eval(mut self, a: &mut [Self::Elem]) -> Result<Vec<Self::Elem>> {
+        let _eig = self.calc(a)?;
+        Ok(unsafe { self.tau.assume_init() })
+    }
+}
+
 macro_rules! impl_qr {
     ($scalar:ty, $qrf:path, $lqf:path, $gqr:path, $glq:path) => {
         impl QR_ for $scalar {