Reconstruct eigenvectors

Author: termoshtt

lax/src/eig.rs

@@ -429,59 +429,71 @@ 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 mut eigvecs = vec_uninit((n * n) as usize);
+                    reconstruct_eigenvectors(
+                        jobvl.is_calc(),
+                        &eig_im,
+                        &vr.or(vl).unwrap(),
+                        &mut eigvecs,
+                    );
+                    Ok((eigs, unsafe { eigvecs.assume_init() }))
+                } 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>(
+    take_hermite_conjugate: bool,
+    eig_im: &[T],
+    vr: &[T],
+    vc: &mut [MaybeUninit<T::Complex>],
+) {
+    let n = eig_im.len();
+    assert_eq!(vr.len(), n * n);
+    assert_eq!(vc.len(), 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 = vr[row + col * n];
+                vc[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 = vr[row + col * n];
+                let mut im = vr[row + (col + 1) * n];
+                if take_hermite_conjugate {
+                    im = -im;
+                }
+                vc[row + col * n].write(T::complex(re, im));
+                vc[row + (col + 1) * n].write(T::complex(re, -im));
+            }
+            col += 2;
+        }
+    }
+}