Support complex case

Author: termoshtt

lax/src/least_squares.rs

@@ -28,8 +28,15 @@ pub trait LeastSquaresSvdDivideConquer_: Scalar {
     ) -> Result<LeastSquaresOutput<Self>>;
 }
 
-macro_rules! impl_least_squares_real {
-    ($scalar:ty, $gelsd:path) => {
+macro_rules! impl_least_squares {
+    (@real, $scalar:ty, $gelsd:path) => {
+        impl_least_squares!(@body, $scalar, $gelsd, );
+    };
+    (@complex, $scalar:ty, $gelsd:path) => {
+        impl_least_squares!(@body, $scalar, $gelsd, rwork);
+    };
+
+    (@body, $scalar:ty, $gelsd:path, $($rwork:ident),*) => {
         impl LeastSquaresSvdDivideConquer_ for $scalar {
             unsafe fn least_squares(
                 l: MatrixLayout,
@@ -85,6 +92,9 @@ macro_rules! impl_least_squares_real {
                 let mut info = 0;
                 let mut work_size = [Self::zero()];
                 let mut iwork_size = [0];
+                $(
+                let mut $rwork = [Self::Real::zero()];
+                )*
                 $gelsd(
                     m,
                     n,
@@ -98,6 +108,7 @@ macro_rules! impl_least_squares_real {
                     &mut rank,
                     &mut work_size,
                     -1,
+                    $(&mut $rwork,)*
                     &mut iwork_size,
                     &mut info,
                 );
@@ -108,6 +119,10 @@ macro_rules! impl_least_squares_real {
                 let mut work = vec![Self::zero(); lwork];
                 let liwork = iwork_size[0].to_usize().unwrap();
                 let mut iwork = vec![0; liwork];
+                $(
+                let lrwork = $rwork[0].to_usize().unwrap();
+                let mut $rwork = vec![Self::Real::zero(); lrwork];
+                )*
                 $gelsd(
                     m,
                     n,
@@ -121,6 +136,7 @@ macro_rules! impl_least_squares_real {
                     &mut rank,
                     &mut work,
                     lwork as i32,
+                    $(&mut $rwork,)*
                     &mut iwork,
                     &mut info,
                 );
@@ -141,93 +157,7 @@ macro_rules! impl_least_squares_real {
     };
 }
 
-impl_least_squares_real!(f64, lapack::dgelsd);
-impl_least_squares_real!(f32, lapack::sgelsd);
-
-macro_rules! impl_least_squares {
-    ($scalar:ty, $gelsd:path) => {
-        impl LeastSquaresSvdDivideConquer_ for $scalar {
-            unsafe fn least_squares(
-                a_layout: MatrixLayout,
-                a: &mut [Self],
-                b: &mut [Self],
-            ) -> Result<LeastSquaresOutput<Self>> {
-                let (m, n) = a_layout.size();
-                if (m as usize) > b.len() || (n as usize) > b.len() {
-                    return Err(Error::InvalidShape);
-                }
-                let k = ::std::cmp::min(m, n);
-                let nrhs = 1;
-                let ldb = match a_layout {
-                    MatrixLayout::F { .. } => m.max(n),
-                    MatrixLayout::C { .. } => 1,
-                };
-                let rcond: Self::Real = -1.;
-                let mut singular_values: Vec<Self::Real> = vec![Self::Real::zero(); k as usize];
-                let mut rank: i32 = 0;
-
-                $gelsd(
-                    a_layout.lapacke_layout(),
-                    m,
-                    n,
-                    nrhs,
-                    a,
-                    a_layout.lda(),
-                    b,
-                    ldb,
-                    &mut singular_values,
-                    rcond,
-                    &mut rank,
-                )
-                .as_lapack_result()?;
-
-                Ok(LeastSquaresOutput {
-                    singular_values,
-                    rank,
-                })
-            }
-
-            unsafe fn least_squares_nrhs(
-                a_layout: MatrixLayout,
-                a: &mut [Self],
-                b_layout: MatrixLayout,
-                b: &mut [Self],
-            ) -> Result<LeastSquaresOutput<Self>> {
-                let (m, n) = a_layout.size();
-                if (m as usize) > b.len()
-                    || (n as usize) > b.len()
-                    || a_layout.lapacke_layout() != b_layout.lapacke_layout()
-                {
-                    return Err(Error::InvalidShape);
-                }
-                let k = ::std::cmp::min(m, n);
-                let nrhs = b_layout.size().1;
-                let rcond: Self::Real = -1.;
-                let mut singular_values: Vec<Self::Real> = vec![Self::Real::zero(); k as usize];
-                let mut rank: i32 = 0;
-
-                $gelsd(
-                    a_layout.lapacke_layout(),
-                    m,
-                    n,
-                    nrhs,
-                    a,
-                    a_layout.lda(),
-                    b,
-                    b_layout.lda(),
-                    &mut singular_values,
-                    rcond,
-                    &mut rank,
-                )
-                .as_lapack_result()?;
-                Ok(LeastSquaresOutput {
-                    singular_values,
-                    rank,
-                })
-            }
-        }
-    };
-}
-
-impl_least_squares!(c64, lapacke::zgelsd);
-impl_least_squares!(c32, lapacke::cgelsd);
+impl_least_squares!(@real, f64, lapack::dgelsd);
+impl_least_squares!(@real, f32, lapack::sgelsd);
+impl_least_squares!(@complex, c64, lapack::zgelsd);
+impl_least_squares!(@complex, c32, lapack::cgelsd);