chore(svd): improve svd function and add test cases

- Remove redundant parameters full matrics.
- Remove redundant code.
- Add support for new type of cutoff: tuple[int, int]
- Add support for grassmann tensor with empty parity block

Author: gausshj

grassmann_tensor/tensor.py

@@ -221,6 +221,19 @@ def _reorder_indices(
         sign = (count & 2).to(dtype=torch.bool)
         return len(even), len(odd), reorder, sign.flatten()
 
+    def _calculate_even_odd(self) -> tuple[int, int]:
+        even, odd = 1, 0
+        for e, o in self.edges:
+            even, odd = even * e + odd * o, even * o + odd * e
+        return even, odd
+
+    @staticmethod
+    def calculate_even_odd(edges: tuple[tuple[int, int], ...]) -> tuple[int, int]:
+        even, odd = 1, 0
+        for e, o in edges:
+            even, odd = even * e + odd * o, even * o + odd * e
+        return even, odd
+
     def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTensor:
         """
         Reshape the Grassmann tensor, which may split or merge edges.
@@ -253,12 +266,33 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens
         merging_reorder: list[tuple[int, torch.Tensor]] = []
         merging_sign: list[tuple[int, torch.Tensor]] = []
 
+        original_self_is_scalar = self.tensor.dim() == 0
+        if original_self_is_scalar:
+            normalized: list[tuple[int, int]] = []
+            for item in new_shape:
+                if item == -1:
+                    raise AssertionError("Cannot use -1 when reshaping from a scalar")
+                if isinstance(item, int):
+                    if item != 1:
+                        raise AssertionError(
+                            f"Ambiguous integer dim {item} from scalar. "
+                            "Use explicit (even, odd) pairs, or only use 1 for trivial edges."
+                        )
+                    normalized.append((1, 0))
+                else:
+                    normalized.append(item)
+            new_shape = tuple(normalized)
+            edges_only = typing.cast(tuple[tuple[int, int], ...], new_shape)
+            assert self.calculate_even_odd(edges_only) == (1, 0), (
+                "Cannot split none edges into illegal edges"
+            )
+
         cursor_plan: int = 0
         cursor_self: int = 0
         while cursor_plan != len(new_shape) or cursor_self != self.tensor.dim():
             if len(new_shape) == 0:
-                assert all(edge == (0, 1) or edge == (1, 0) for edge in self.edges), (
-                    f"Edge must be (0, 1) or (1, 0) but got {self.edges}"
+                assert self._calculate_even_odd() == (1, 0), (
+                    "Only pure even edges can be merged into none edges"
                 )
                 cursor_self = self.tensor.dim() - 1
             elif cursor_plan != len(new_shape) and new_shape[cursor_plan] == -1:
@@ -291,9 +325,7 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens
                     else cursor_new_shape[0] + cursor_new_shape[1]
                 )
             # one of total and shape[cursor_self] is not trivial, otherwise it should be handled before
-            if self.tensor.dim() == 0:
-                merging = False
-            elif total == self.tensor.shape[cursor_self]:
+            if total == self.tensor.shape[cursor_self]:
                 # We do not know whether it is merging or splitting, check more
                 if isinstance(cursor_new_shape, int) or cursor_new_shape == self.edges[cursor_self]:
                     # If the new shape is exactly the same as the current edge, we treat it as no change
@@ -387,23 +419,16 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens
                     plan_total *= new_cursor_new_shape[0] + new_cursor_new_shape[1]
                     new_cursor_plan += 1
                     # One dimension included, check if we can stop
-                    if self.tensor.dim() == 0:
+                    if plan_total == self.tensor.shape[cursor_self]:
+                        # new_shape block has been verified to be always tuple[int, int] before
                         even, odd, reorder, sign = self._reorder_indices(
-                            typing.cast(tuple[tuple[int, int], ...], new_shape)
-                        )
-                        new_cursor_plan = len(new_shape)
-                        break
-                    else:
-                        if plan_total == self.tensor.shape[cursor_self]:
-                            # new_shape block has been verified to be always tuple[int, int] before
-                            even, odd, reorder, sign = self._reorder_indices(
-                                typing.cast(
-                                    tuple[tuple[int, int], ...],
-                                    new_shape[cursor_plan:new_cursor_plan],
-                                )
+                            typing.cast(
+                                tuple[tuple[int, int], ...],
+                                new_shape[cursor_plan:new_cursor_plan],
                             )
-                            if (even, odd) == self.edges[cursor_self]:
-                                break
+                        )
+                        if (even, odd) == self.edges[cursor_self]:
+                            break
                     # For some reason we cannot stop here, continue to include more dimension, check something before continue
                     assert plan_total <= self.tensor.shape[cursor_self], (
                         f"Dimension mismatch in splitting with edges {self.edges} and new shape {new_shape}."
@@ -415,10 +440,7 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens
                 for i in range(cursor_plan, new_cursor_plan):
                     # new_shape block has been verified to be always tuple[int, int] in the loop
                     new_cursor_new_shape = typing.cast(tuple[int, int], new_shape[i])
-                    if self.tensor.dim() == 0:
-                        arrow.append(False)
-                    else:
-                        arrow.append(self.arrow[cursor_self])
+                    arrow.append(self.arrow[cursor_self])
                     edges.append(new_cursor_new_shape)
                     shape.append(new_cursor_new_shape[0] + new_cursor_new_shape[1])
                 splitting_reorder.append((cursor_self, reorder))

tests/reshape_test.py

@@ -180,5 +180,20 @@ def test_reshape_equal_edges_nontrivial_merging_with_other_edge() -> None:
 def test_reshape_with_none() -> None:
     a = GrassmannTensor((), (), torch.tensor(2333)).reshape(((1, 0), (1, 0))).reshape(())
     assert len(a.arrow) == 0 and len(a.edges) == 0 and a.tensor.dim() == 0
-    b = GrassmannTensor((), (), torch.tensor(2333)).reshape(((0, 1), (0, 1))).reshape(())
+    b = GrassmannTensor((), (), torch.tensor(2333)).reshape(((1, 0), (1, 0))).reshape(())
     assert len(b.arrow) == 0 and len(b.edges) == 0 and b.tensor.dim() == 0
+    c = GrassmannTensor((), (), torch.tensor(2333)).reshape((1, 1))
+    assert len(c.arrow) == 2 and len(c.edges) == 2 and c.tensor.dim() == 2
+
+
+def test_reshape_with_none_edge_assertion() -> None:
+    with pytest.raises(AssertionError, match="Only pure even edges can be merged into none edges"):
+        _ = GrassmannTensor((True, True), ((0, 1), (1, 0)), torch.tensor([[2333]])).reshape(())
+    with pytest.raises(AssertionError, match="Cannot split none edges into illegal edges"):
+        _ = GrassmannTensor((), (), torch.tensor(2333)).reshape(((0, 1),))
+    with pytest.raises(AssertionError, match="Cannot split none edges into illegal edges"):
+        _ = GrassmannTensor((), (), torch.tensor(2333)).reshape(((0, 1), (1, 0)))
+    with pytest.raises(AssertionError, match="Cannot use -1 when reshaping from a scalar"):
+        _ = GrassmannTensor((), (), torch.tensor(2333)).reshape((1, -1))
+    with pytest.raises(AssertionError, match="Ambiguous integer dim"):
+        _ = GrassmannTensor((), (), torch.tensor(2333)).reshape((2, 2))