Reconstruct eigenvectors

rust-ndarray · Sep 8, 2022 · 802149b · 802149b
1 parent 4737b9c
commit 802149b
Showing 1 changed file with 56 additions and 48 deletions.
diff --git a/lax/src/eig.rs b/lax/src/eig.rs
@@ -223,59 +223,67 @@ macro_rules! impl_eig_real {
                     .map(|(&re, &im)| Self::complex(re, im))
                     .collect();
 
-                if !calc_v {
-                    return Ok((eigs, Vec::new()));
-                }
-
-                // Reconstruct eigenvectors into complex-array
-                // --------------------------------------------
-                //
-                // From LAPACK API https://software.intel.com/en-us/node/469230
-                //
-                // - If the j-th eigenvalue is real,
-                //   - v(j) = VR(:,j), the j-th column of VR.
-                //
-                // - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
-                //   - v(j)   = VR(:,j) + i*VR(:,j+1)
-                //   - v(j+1) = VR(:,j) - i*VR(:,j+1).
-                //
-                // In the C-layout case, we need the conjugates of the left
-                // eigenvectors, so the signs should be reversed.
-
-                let n = n as usize;
-                let v = vr.or(vl).unwrap();
-                let mut eigvecs: Vec<MaybeUninit<Self::Complex>> = vec_uninit(n * n);
-                let mut col = 0;
-                while col < n {
-                    if eig_im[col] == 0. {
-                        // The corresponding eigenvalue is real.
-                        for row in 0..n {
-                            let re = v[row + col * n];
-                            eigvecs[row + col * n].write(Self::complex(re, 0.));
-                        }
-                        col += 1;
-                    } else {
-                        // This is a complex conjugate pair.
-                        assert!(col + 1 < n);
-                        for row in 0..n {
-                            let re = v[row + col * n];
-                            let mut im = v[row + (col + 1) * n];
-                            if jobvl.is_calc() {
-                                im = -im;
-                            }
-                            eigvecs[row + col * n].write(Self::complex(re, im));
-                            eigvecs[row + (col + 1) * n].write(Self::complex(re, -im));
-                        }
-                        col += 2;
-                    }
+                if calc_v {
+                    let eigvecs = reconstruct_eigenvectors(jobvl, n, &eig_im, vr, vl);
+                    Ok((eigs, eigvecs))
+                } else {
+                    Ok((eigs, Vec::new()))
                 }
-                let eigvecs = unsafe { eigvecs.assume_init() };
-
-                Ok((eigs, eigvecs))
             }
         }
     };
 }
 
 impl_eig_real!(f64, lapack_sys::dgeev_);
 impl_eig_real!(f32, lapack_sys::sgeev_);
+
+/// Reconstruct eigenvectors into complex-array
+/// --------------------------------------------
+///
+/// From LAPACK API https://software.intel.com/en-us/node/469230
+///
+/// - If the j-th eigenvalue is real,
+///   - v(j) = VR(:,j), the j-th column of VR.
+///
+/// - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
+///   - v(j)   = VR(:,j) + i*VR(:,j+1)
+///   - v(j+1) = VR(:,j) - i*VR(:,j+1).
+///
+/// In the C-layout case, we need the conjugates of the left
+/// eigenvectors, so the signs should be reversed.
+fn reconstruct_eigenvectors<T: Scalar>(
+    jobvl: JobEv,
+    n: i32,
+    eig_im: &[T],
+    vr: Option<Vec<T>>,
+    vl: Option<Vec<T>>,
+) -> Vec<T::Complex> {
+    let n = n as usize;
+    let v = vr.or(vl).unwrap();
+    let mut eigvecs: Vec<MaybeUninit<T::Complex>> = vec_uninit(n * n);
+    let mut col = 0;
+    while col < n {
+        if eig_im[col].is_zero() {
+            // The corresponding eigenvalue is real.
+            for row in 0..n {
+                let re = v[row + col * n];
+                eigvecs[row + col * n].write(T::complex(re, T::zero()));
+            }
+            col += 1;
+        } else {
+            // This is a complex conjugate pair.
+            assert!(col + 1 < n);
+            for row in 0..n {
+                let re = v[row + col * n];
+                let mut im = v[row + (col + 1) * n];
+                if jobvl.is_calc() {
+                    im = -im;
+                }
+                eigvecs[row + col * n].write(T::complex(re, im));
+                eigvecs[row + (col + 1) * n].write(T::complex(re, -im));
+            }
+            col += 2;
+        }
+    }
+    unsafe { eigvecs.assume_init() }
+}