dev(svd): improve code readability

- Improve code readability
- Remove redundant code

Author: gausshj

grassmann_tensor/tensor.py

@@ -615,40 +615,37 @@ def svd(
             Vh_odd = odd_tensor.new_zeros((0, odd_right))
 
         n_even, n_odd = S_even.shape[0], S_odd.shape[0]
-        total = n_even + n_odd
 
         if cutoff is None:
-            cutoff = (n_even, n_odd)
+            k_even, k_odd = n_even, n_odd
         elif isinstance(cutoff, int):
-            assert total >= 0, "Invalid total singular values."
-            if total == 0:
+            if n_even == 0 and n_odd == 0:
                 raise RuntimeError("Both parity block are empty. Can not form SVD.")
-            k = min(int(cutoff), total)
-            assert k > 0, f"Cutoff must be greater than 0, but got {k}"
-
-            cutoff = (cutoff, cutoff)
-            cutoff = (min(cutoff[0], n_even), min(cutoff[1], n_odd))
-
-        if isinstance(cutoff, tuple):
-            k_even = max(0, min(int(cutoff[0]), n_even))
-            k_odd = max(0, min(int(cutoff[1]), n_odd))
+            assert cutoff > 0, f"Cutoff must be greater than 0, but got {cutoff}"
+            k_even = min(cutoff, n_even)
+            k_odd = min(cutoff, n_odd)
+        elif isinstance(cutoff, tuple):
+            assert len(cutoff) == 2, "The length of cutoff must be 2 if cutoff is a tuple."
             if n_even == 0 and n_odd == 0:
                 raise RuntimeError("Both parity block are empty. Can not form SVD.")
-            assert (k_even > 0 or n_even == 0) and (k_odd > 0 or n_odd == 0), (
-                "Per-block cutoff must be compatible with available singulars"
-            )
-
-            keep_even = torch.zeros(n_even, dtype=torch.bool, device=S_even.device)
-            keep_odd = torch.zeros(n_odd, dtype=torch.bool, device=S_odd.device)
-            if k_even > 0:
-                keep_even[:k_even] = True
-            if k_odd > 0:
-                keep_odd[:k_odd] = True
+            k_even = max(0, min(int(cutoff[0]), n_even))
+            k_odd = max(0, min(int(cutoff[1]), n_odd))
         else:
             raise ValueError(
                 f"Cutoff must be an integer or a tuple of two integers, but got {cutoff}"
             )
 
+        assert (k_even > 0 or n_even == 0) and (k_odd > 0 or n_odd == 0), (
+            "Per-block cutoff must be compatible with available singulars"
+        )
+
+        keep_even = torch.zeros(n_even, dtype=torch.bool, device=S_even.device)
+        keep_odd = torch.zeros(n_odd, dtype=torch.bool, device=S_odd.device)
+        if k_even > 0:
+            keep_even[:k_even] = True
+        if k_odd > 0:
+            keep_odd[:k_odd] = True
+
         U_even_trunc = U_even[:, keep_even]
         S_even_trunc = S_even[keep_even]
         Vh_even_trunc = Vh_even[keep_even, :]

tests/svd_test.py

@@ -18,30 +18,17 @@
 
 
 def get_total_singular(edges: Edges, free_names_u: FreeNamesU) -> tuple[int, int]:
-    even, odd = edges[free_names_u[0]]
-    for i in range(1, len(free_names_u)):
-        e, o = edges[free_names_u[i]]
-        even, odd = even * e + odd * o, even * o + odd * e
-    even_singular = min(even, odd)
-
+    even_singular = min(GrassmannTensor.calculate_even_odd(tuple(edges[i] for i in free_names_u)))
     set_all = set(range(len(edges)))
-    right_idx = sorted(set_all - set(free_names_u))
-    even, odd = edges[right_idx[0]]
-    for i in range(1, len(right_idx)):
-        e, o = edges[right_idx[i]]
-        even, odd = even * e + odd * o, even * o + odd * e
-    odd_singular = min(even, odd)
+    remain_idx = sorted(set_all - set(free_names_u))
+    odd_singular = min(GrassmannTensor.calculate_even_odd(tuple(edges[i] for i in remain_idx)))
     return even_singular, odd_singular
 
 
-def tau_for_cutoff(
-    c: int | tuple[int, int], total: int, alpha: float = 0.8, slack: float = 1.05
-) -> float:
+def tau_for_cutoff(c: int, total: int, alpha: float = 0.8, slack: float = 1.05) -> float:
     cut = 0
     if isinstance(c, int):
         cut = c
-    elif isinstance(c, tuple):
-        cut = sum(c)
     lo, hi = 1e-8, 1e-1
     x = (total - cut) / max(1, total - 1)
     return (lo + (hi - lo) * (x**alpha)) * slack
@@ -80,12 +67,6 @@ def svd_cases() -> SVDCases:
                 (even_singular, odd_singular),
             ]
             for cutoff in cutoff_list:
-                if isinstance(cutoff, int):
-                    if cutoff is not None and cutoff < 1:
-                        continue
-                elif isinstance(cutoff, tuple):
-                    if cutoff[0] < 1 or cutoff[1] < 1:
-                        continue
                 if cutoff is None:
                     kept = total
                 elif isinstance(cutoff, int):