dev(exponential): unify terminology to edges and add test cases

- Unify terminology to `edges` to maintain consistency
- Add test cases for assertation

Author: gausshj

grassmann_tensor/tensor.py

@@ -569,7 +569,7 @@ def matmul(self, other: GrassmannTensor) -> GrassmannTensor:
             _tensor=tensor,
         )
 
-    def _group_legs(
+    def _group_edges(
         self,
         left_legs: typing.Iterable[int],
     ) -> tuple[GrassmannTensor, tuple[int, ...], tuple[int, ...]]:
@@ -615,7 +615,7 @@ def svd(
         if isinstance(cutoff, tuple):
             assert len(cutoff) == 2, "The length of cutoff must be 2 if cutoff is a tuple."
 
-        tensor, left_legs, right_legs = self._group_legs(free_names_u)
+        tensor, left_legs, right_legs = self._group_edges(free_names_u)
 
         (even_left, odd_left) = tensor.edges[0]
         (even_right, odd_right) = tensor.edges[1]
@@ -718,12 +718,18 @@ def svd(
 
         return U, S, Vh
 
+    def _get_inv_order(self, order: tuple[int, ...]) -> tuple[int, ...]:
+        inv = [0] * self.tensor.dim()
+        for new_pos, orig_idx in enumerate(order):
+            inv[orig_idx] = new_pos
+        return tuple(inv)
+
     def exponential(self, pairs: tuple[int, ...]) -> GrassmannTensor:
-        tensor, left_legs, right_legs = self._group_legs(pairs)
+        tensor, left_legs, right_legs = self._group_edges(pairs)
 
-        axes_to_reverse = [i for i in range(2) if tensor.arrow[i]]
-        if axes_to_reverse:
-            tensor = tensor.reverse(tuple(axes_to_reverse))
+        edges_to_reverse = [i for i in range(2) if tensor.arrow[i]]
+        if edges_to_reverse:
+            tensor = tensor.reverse(tuple(edges_to_reverse))
 
         left_dim, right_dim = tensor.tensor.shape
 
@@ -744,17 +750,14 @@ def exponential(self, pairs: tuple[int, ...]) -> GrassmannTensor:
 
         tensor_exp = dataclasses.replace(tensor, _tensor=tensor_exp)
 
-        if axes_to_reverse:
-            tensor_exp = tensor_exp.reverse(tuple(axes_to_reverse))
+        if edges_to_reverse:
+            tensor_exp = tensor_exp.reverse(tuple(edges_to_reverse))
 
         order = left_legs + right_legs
         edges_after_permute = tuple(self.edges[i] for i in order)
         tensor_exp = tensor_exp.reshape(edges_after_permute)
 
-        inv = [0] * self.tensor.dim()
-        for new_pos, orig_idx in enumerate(order):
-            inv[orig_idx] = new_pos
-        inv_order = tuple(inv)
+        inv_order = self._get_inv_order(order)
 
         tensor_exp = tensor_exp.permute(inv_order)
 

tests/exponential_test.py

@@ -1,4 +1,5 @@
 import torch
+import pytest
 
 from grassmann_tensor import GrassmannTensor
 
@@ -17,3 +18,13 @@ def test_exponential_with_empty_parity_block() -> None:
     a.exponential((0,))
     b = GrassmannTensor((False, True), ((0, 1), (0, 1)), torch.randn(1, 1))
     b.exponential((0,))
+
+
+def test_exponential_assertation() -> None:
+    a = GrassmannTensor(
+        (True, True, True, True),
+        ((2, 2), (4, 4), (8, 8), (16, 16)),
+        torch.randn(4, 8, 16, 32, dtype=torch.float64),
+    )
+    with pytest.raises(AssertionError, match="Exponential requires a square operator"):
+        a.exponential((0, 2))