Change sorting traits to use Order enum

Author: YuhanLiin

src/eigh.rs

@@ -3,7 +3,8 @@
 use ndarray::{s, Array1, Array2, ArrayBase, Data, DataMut, Ix2, NdFloat};
 
 use crate::{
-    check_square, givens::GivensRotation, index::*, tridiagonal::SymmetricTridiagonal, Result,
+    check_square, givens::GivensRotation, index::*, tridiagonal::SymmetricTridiagonal, Order,
+    Result,
 };
 
 fn symmetric_eig<A: NdFloat, S: DataMut<Elem = A>>(
@@ -272,44 +273,42 @@ impl<A: NdFloat, S: Data<Elem = A>> EigValsh for ArrayBase<S, Ix2> {
 ///
 /// Will panic if shape or layout of inputs differ from eigen output, or if input contains NaN.
 pub trait EigSort: Sized {
-    fn sort_eig(self, descending: bool) -> Self;
+    fn sort_eig(self, order: Order) -> Self;
 
     /// Sort eigendecomposition by the eigenvalues in ascending order
     fn sort_eig_asc(self) -> Self {
-        self.sort_eig(false)
+        self.sort_eig(Order::Smallest)
     }
 
     /// Sort eigendecomposition by the eigenvalues in descending order
     fn sort_eig_desc(self) -> Self {
-        self.sort_eig(true)
+        self.sort_eig(Order::Largest)
     }
 }
 
 /// Implementation on output of `EigValsh` traits
 impl<A: NdFloat> EigSort for Array1<A> {
-    fn sort_eig(mut self, descending: bool) -> Self {
+    fn sort_eig(mut self, order: Order) -> Self {
         // Panics on non-standard layouts, which is fine because our eigenvals have standard layout
         let slice = self.as_slice_mut().unwrap();
         // Panic only happens with NaN values
-        if descending {
-            slice.sort_by(|a, b| b.partial_cmp(a).unwrap());
-        } else {
-            slice.sort_by(|a, b| a.partial_cmp(b).unwrap());
+        match order {
+            Order::Largest => slice.sort_by(|a, b| b.partial_cmp(a).unwrap()),
+            Order::Smallest => slice.sort_by(|a, b| a.partial_cmp(b).unwrap()),
         }
         self
     }
 }
 
 /// Implementation on output of `Eigh` traits
 impl<A: NdFloat> EigSort for (Array1<A>, Array2<A>) {
-    fn sort_eig(self, descending: bool) -> Self {
+    fn sort_eig(self, order: Order) -> Self {
         let (mut vals, vecs) = self;
         let mut value_idx: Vec<_> = vals.iter().copied().enumerate().collect();
         // Panic only happens with NaN values
-        if descending {
-            value_idx.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
-        } else {
-            value_idx.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
+        match order {
+            Order::Largest => value_idx.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()),
+            Order::Smallest => value_idx.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap()),
         }
 
         let mut out = Array2::zeros(vecs.dim());

src/lib.rs

@@ -70,6 +70,8 @@ pub(crate) fn check_square<S: RawData>(arr: &ArrayBase<S, Ix2>) -> Result<usize>
 }
 
 /// Find largest or smallest eigenvalues
+///
+/// Corresponds to descending and ascending order
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum Order {
     Largest,

src/lobpcg/algorithm.rs

@@ -31,25 +31,16 @@ fn sorted_eig<A: NdFloat>(
     size: usize,
     order: Order,
 ) -> Result<(Array1<A>, Array2<A>)> {
-    let n = a.len_of(Axis(0));
-
     let res = match b {
         Some(b) => generalized_eig(a, b)?,
         _ => a.eigh_into()?,
     };
 
     // sort and ensure that signs are deterministic
-    let (vals, vecs) = res.sort_eig(false);
+    let (vals, vecs) = res.sort_eig(order);
     let s = vecs.row(0).mapv(|x| x.signum());
     let vecs = vecs * s;
-
-    Ok(match order {
-        Order::Largest => (
-            vals.slice_move(s![n-size..; -1]),
-            vecs.slice_move(s![.., n-size..; -1]),
-        ),
-        Order::Smallest => (vals.slice_move(s![..size]), vecs.slice_move(s![.., ..size])),
-    })
+    Ok((vals.slice_move(s![..size]), vecs.slice_move(s![.., ..size])))
 }
 
 /// Masks a matrix with the given `matrix`

src/svd.rs

@@ -8,7 +8,7 @@ use ndarray::{s, Array1, Array2, ArrayBase, Axis, Data, DataMut, Ix2, NdFloat};
 
 use crate::{
     bidiagonal::Bidiagonal, eigh::wilkinson_shift, givens::GivensRotation, index::*, LinalgError,
-    Result,
+    Order, Result,
 };
 
 fn svd<A: NdFloat, S: DataMut<Elem = A>>(
@@ -482,29 +482,28 @@ impl<A: NdFloat, S: Data<Elem = A>> SVD for ArrayBase<S, Ix2> {
 ///
 /// Will panic if shape of inputs differs from shape of SVD output, or if input contains NaN.
 pub trait SvdSort: Sized {
-    fn sort_svd(self, descending: bool) -> Self;
+    fn sort_svd(self, order: Order) -> Self;
 
     /// Sort SVD decomposition by the singular values in ascending order
     fn sort_svd_asc(self) -> Self {
-        self.sort_svd(false)
+        self.sort_svd(Order::Smallest)
     }
 
     /// Sort SVD decomposition by the singular values in descending order
     fn sort_svd_desc(self) -> Self {
-        self.sort_svd(true)
+        self.sort_svd(Order::Largest)
     }
 }
 
 /// Implemented on the output of the `SVD` traits
 impl<A: NdFloat> SvdSort for (Option<Array2<A>>, Array1<A>, Option<Array2<A>>) {
-    fn sort_svd(self, descending: bool) -> Self {
+    fn sort_svd(self, order: Order) -> Self {
         let (u, mut s, vt) = self;
         let mut value_idx: Vec<_> = s.iter().copied().enumerate().collect();
         // Panic only happens with NaN values
-        if descending {
-            value_idx.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
-        } else {
-            value_idx.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
+        match order {
+            Order::Largest => value_idx.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()),
+            Order::Smallest => value_idx.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap()),
         }
 
         let apply_ordering = |arr: &Array2<A>, ax, values_idx: &Vec<_>| {

tests/lobpcg.rs

@@ -67,7 +67,7 @@ fn test_marchenko_pastur() {
 }
 
 fn run_lobpcg_eig_test(arr: Array2<f64>, num: usize, ordering: Order) {
-    let (eigvals, _) = arr.eigh().unwrap().sort_eig(ordering == Order::Largest);
+    let (eigvals, _) = arr.eigh().unwrap().sort_eig(ordering);
     let res = TruncatedEig::new_with_rng(arr.clone(), ordering, Xoshiro256Plus::seed_from_u64(42))
         .precision(1e-3)
         .decompose(num)
@@ -108,10 +108,7 @@ fn problematic_eig_matrix() {
 }
 
 fn run_lobpcg_svd_test(arr: Array2<f64>, ordering: Order) {
-    let (_, s, _) = arr
-        .svd(false, false)
-        .unwrap()
-        .sort_svd(ordering == Order::Largest);
+    let (_, s, _) = arr.svd(false, false).unwrap().sort_svd(ordering);
     let (u, ts, vt) =
         TruncatedSvd::new_with_rng(arr.clone(), ordering, Xoshiro256Plus::seed_from_u64(42))
             .precision(1e-3)

tests/svd.rs

@@ -61,7 +61,7 @@ fn run_svd_test(arr: Array2<f64>) {
 proptest! {
     #![proptest_config(ProptestConfig::with_cases(1000))]
     #[test]
-    fn bidiagonal_test(arr in common::rect_arr()) {
+    fn svd_test(arr in common::rect_arr()) {
         run_svd_test(arr);
     }
 }