fix(svd): correct cutoff logic and add test cases for svd

- Correct cutoff logic and support None as no-trunction mode
- Add parameterized test cases for svd
- Fix missing coverage on exception and boundary paths

Signed-off-by: Gauss <[email protected]>

Author: Gauss

grassmann_tensor/tensor.py

@@ -503,7 +503,7 @@ def svd(
         2. Merge the tensor with two groups.
         3. Split the block tensor into two parts.
         4. Compute the singular value decomposition.
-        5. Use the cutoff to cut off the lower singular values and corresponding U and V.
+        5. Use cutoff to keep the largest cutoff singular values (globally across even/odd blocks).
         6. Contract U, S and Vh.
         7. Split the legs into original left and right.
         The returned tensors U and V are not unique, nor are they continuous with respect to self.
@@ -534,31 +534,34 @@ def svd(
         U_even, S_even, Vh_even = torch.linalg.svd(even_tensor, full_matrices=full_matrices)
         U_odd, S_odd, Vh_odd = torch.linalg.svd(odd_tensor, full_matrices=full_matrices)
 
-        if cutoff is None:
-            cutoff = 0
         total = S_even.numel() + S_odd.numel()
-        cutoff = max(0, min(cutoff, total))
 
-        if cutoff == 0:
-            keep_even = torch.ones_like(S_even, dtype=torch.bool)
-            keep_odd = torch.ones_like(S_odd, dtype=torch.bool)
+        if cutoff is None:
+            cutoff = total
+        else:
+            cutoff = min(int(cutoff), total)
+
+        assert cutoff > 0, f"Cutoff must be greater than 0, but got {cutoff}"
+
+        if cutoff == total:
+            keep_even = torch.ones_like(S_even, dtype=torch.bool, device=S_even.device)
+            keep_odd = torch.ones_like(S_odd, dtype=torch.bool, device=S_odd.device)
         else:
             S_cat = torch.cat([S_even, S_odd])
-            sorted_vals, sorted_indices = S_cat.sort()
-            cutoff_indices = sorted_indices[:cutoff]
+            top_vals, top_indices = torch.topk(S_cat, k=cutoff, largest=True, sorted=False)
 
-            mask_even_indices = cutoff_indices < S_even.shape[0]
-            mask_odd_indices = ~mask_even_indices
-            even_cutoff_indices = cutoff_indices[mask_even_indices]
-            odd_cutoff_indices = cutoff_indices[mask_odd_indices] - S_even.shape[0]
+            even_mask = top_indices < S_even.shape[0]
+            odd_mask = ~even_mask
+            even_keep_indices = top_indices[even_mask]
+            odd_keep_indices = top_indices[odd_mask] - S_even.shape[0]
 
-            keep_even = torch.ones_like(S_even, dtype=torch.bool)
-            if even_cutoff_indices.numel() > 0:
-                keep_even[even_cutoff_indices] = False
+            keep_even = torch.zeros_like(S_even, dtype=torch.bool).to(S_even.device)
+            if even_keep_indices.numel() > 0:
+                keep_even[even_keep_indices] = True
 
-            keep_odd = torch.ones_like(S_odd, dtype=torch.bool)
-            if odd_cutoff_indices.numel() > 0:
-                keep_odd[odd_cutoff_indices] = False
+            keep_odd = torch.ones_like(S_odd, dtype=torch.bool).to(S_odd.device)
+            if odd_keep_indices.numel() > 0:
+                keep_odd[odd_keep_indices] = True
 
         U_even_trunc = U_even[:, keep_even]
         S_even_trunc = S_even[keep_even]

tests/svd_test.py

@@ -1,34 +1,126 @@
 import torch
+import pytest
 import math
+import itertools
+from typing import TypeAlias, Iterable
+
 from grassmann_tensor import GrassmannTensor
 
+Arrow: TypeAlias = tuple[bool, ...]
+Edges: TypeAlias = tuple[tuple[int, int], ...]
+Tensor: TypeAlias = torch.Tensor
+Cutoff: TypeAlias = int
+Tau: TypeAlias = float
+FreeNamesU: TypeAlias = tuple[int, ...]
+
+SVDCases = Iterable[tuple[Arrow, Edges, Tensor, Cutoff, Tau, FreeNamesU]]
+
+def get_total_singular(edges: Edges, free_names_u: FreeNamesU) -> 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
+    total_singular = min(even, odd)
+
+    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
+    total_singular += min(even, odd)
+    return total_singular
+
+def tau_for_cutoff(c: int, total: int, alpha: float = 0.8) -> float:
+    lo, hi = 1e-8, 1e-1
+    x = (total - c) / max(1, total - 1)
+    return lo + (hi - lo) * (x ** alpha)
+
+def choose_free_names(n_edges: int, limit: int = 8) -> list[FreeNamesU]:
+    combos = [tuple(c) for r in range(1, n_edges) for c in itertools.combinations(range(n_edges), r)]
+    return combos[:limit]
 
-def test_svd() -> None:
-    gt = GrassmannTensor(
-        (True, True, True, True),
-        ((8, 8), (4, 4), (2, 2), (1, 1)),
-        torch.randn([16, 8, 4, 2], dtype=torch.float64),
-    )
-    U, S, Vh = gt.svd((0, 3), cutoff=1)
+BASE_GT_CASES: list[tuple[Arrow, Edges, Tensor]] = [
+    ((True, True), ((2, 2), (4, 4)), torch.randn(4, 8, dtype=torch.float64)),
+    ((True, True, True), ((2, 2), (4, 4), (8, 8)), torch.randn(4, 8, 16, dtype=torch.float64)),
+    ((True, True, True, True), ((2, 2), (4, 4), (8, 8), (16, 16)), torch.randn(4, 8, 16, 32, dtype=torch.float64)),
+]
+
+def svd_cases() -> SVDCases:
+    params = []
+    for arrow, edges, tensor in BASE_GT_CASES:
+        for fnu in choose_free_names(len(edges)):
+            total = get_total_singular(edges, fnu)
+            for cutoff in [None, total, total - 1, total - 2]:
+                if cutoff is not None and cutoff < 1:
+                    continue
+                tau = tau_for_cutoff(cutoff or total, total)
+                params.append(
+                    pytest.param(
+                        arrow, edges, tensor, cutoff, tau, fnu,
+                        id=f"edges={tuple(edges)}|fnu={fnu}|cut={cutoff}|tau={tau:.2e}"
+                    )
+                )
+    return params
+
+@pytest.mark.parametrize(
+    "arrow, edges, tensor, cutoff, tau, free_names_u",
+    svd_cases(),
+)
+@pytest.mark.repeat(20)
+def test_svd(
+        arrow: Arrow,
+        edges: Edges,
+        tensor: Tensor,
+        cutoff: Cutoff,
+        tau: Tau,
+        free_names_u: FreeNamesU,
+) -> None:
+    gt = GrassmannTensor(arrow, edges, tensor)
+    U, S, Vh = gt.svd(free_names_u, cutoff=cutoff)
 
     # reshape U
-    # left_arrow = U.arrow[:-1]
     left_dim = math.prod(U.tensor.shape[:-1])
     left_edge = list(U.edges[:-1])
     U = U.reshape((left_dim, -1))
 
     # reshape Vh
-    # right_arrow = Vh.arrow[1:]
     right_dim = math.prod(Vh.tensor.shape[1:])
     right_edge = list(Vh.edges[1:])
     Vh = Vh.reshape((-1, right_dim))
 
     US = GrassmannTensor.matmul(U, S)
     USV = GrassmannTensor.matmul(US, Vh)
 
+    set_all = set(range(len(edges)))
+    set_u = set(free_names_u)
+    set_v = sorted(set_all - set_u)
+    perm_order = list(free_names_u) + list(set_v)
+    inv_perm = [perm_order.index(i) for i in range(len(edges))]
+
     USV = USV.reshape(tuple(left_edge + right_edge))
-    USV = USV.permute((0, 2, 3, 1))
+    USV = USV.permute(tuple(inv_perm))
+
+    masked = gt.update_mask().tensor
+    den = masked.norm()
+    eps = torch.finfo(masked.dtype).eps
+    rel_err = (masked - USV.tensor).norm() / max(den, eps)
+    assert rel_err <= tau
+
+@pytest.mark.parametrize(
+    "arrow, edges, tensor, cutoff, tau, free_names_u",
+    svd_cases(),
+)
+def test_svd_with_zero_cutoff(
+        arrow: Arrow,
+        edges: Edges,
+        tensor: Tensor,
+        cutoff: Cutoff,
+        tau: Tau,
+        free_names_u: FreeNamesU,
+) -> None:
+    gt = GrassmannTensor(arrow, edges, tensor)
+    with pytest.raises(AssertionError, match="Cutoff must be greater than 0"):
+        _, _, _ = gt.svd(free_names_u, cutoff=0)
+
 
-    # assert torch.allclose(gt.update_mask().tensor, USV.tensor)
-    rel_err = (gt.update_mask().tensor - USV.tensor).norm() / gt.update_mask().tensor.norm()
-    assert rel_err < 1e-2