Merge pull request #75 from USTC-KnowledgeComputingLab/dev/add-support-for-exponential

dev(exponential): add support for exponential
Close #75

Author: gausshj

grassmann_tensor/tensor.py

@@ -569,6 +569,27 @@ def matmul(self, other: GrassmannTensor) -> GrassmannTensor:
             _tensor=tensor,
         )
 
+    def _group_edges(
+        self,
+        left_legs: typing.Iterable[int],
+    ) -> tuple[GrassmannTensor, tuple[int, ...], tuple[int, ...]]:
+        left_legs = tuple(int(i) for i in left_legs)
+        right_legs = tuple(i for i in range(self.tensor.dim()) if i not in left_legs)
+        assert set(left_legs) | set(right_legs) == set(range(self.tensor.dim())), (
+            "Left/right must cover all tensor legs."
+        )
+
+        order = left_legs + right_legs
+
+        tensor = self.permute(order)
+
+        left_dim = math.prod(tensor.tensor.shape[: len(left_legs)])
+        right_dim = math.prod(tensor.tensor.shape[len(left_legs) :])
+
+        tensor = tensor.reshape((left_dim, right_dim))
+
+        return tensor, left_legs, right_legs
+
     def svd(
         self,
         free_names_u: tuple[int, ...],
@@ -591,22 +612,10 @@ def svd(
         Furthermore, if the distance between any two singular values is close to zero, the gradient
         will be numerically unstable, as it depends on the singular values
         """
-        left_legs = tuple(int(i) for i in free_names_u)
-        right_legs = tuple(i for i in range(self.tensor.dim()) if i not in left_legs)
-        assert set(left_legs) | set(right_legs) == set(range(self.tensor.dim())), (
-            "Left/right must cover all tensor legs."
-        )
-
         if isinstance(cutoff, tuple):
             assert len(cutoff) == 2, "The length of cutoff must be 2 if cutoff is a tuple."
 
-        order = left_legs + right_legs
-        tensor = self.permute(order)
-
-        left_dim = math.prod(tensor.tensor.shape[: len(left_legs)])
-        right_dim = math.prod(tensor.tensor.shape[len(left_legs) :])
-
-        tensor = tensor.reshape((left_dim, right_dim))
+        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]
@@ -709,6 +718,51 @@ def svd(
 
         return U, S, Vh
 
+    def _get_inv_order(self, order: tuple[int, ...]) -> tuple[int, ...]:
+        inv = [0] * self.tensor.dim()
+        for new_position, origin_idx in enumerate(order):
+            inv[origin_idx] = new_position
+        return tuple(inv)
+
+    def exponential(self, pairs: tuple[int, ...]) -> GrassmannTensor:
+        tensor, left_legs, right_legs = self._group_edges(pairs)
+
+        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
+
+        assert left_dim == right_dim, (
+            f"Exponential requires a square operator, but got {left_dim} x {right_dim}."
+        )
+
+        (even_left, odd_left) = tensor.edges[0]
+        (even_right, odd_right) = tensor.edges[1]
+
+        even_tensor = tensor.tensor[:even_left, :even_right]
+        odd_tensor = tensor.tensor[even_left:, even_right:]
+
+        even_tensor_exp = torch.linalg.matrix_exp(even_tensor)
+        odd_tensor_exp = torch.linalg.matrix_exp(odd_tensor)
+
+        tensor_exp = torch.block_diag(even_tensor_exp, odd_tensor_exp)  # type: ignore[no-untyped-call]
+
+        tensor_exp = dataclasses.replace(tensor, _tensor=tensor_exp)
+
+        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_order = self._get_inv_order(order)
+
+        tensor_exp = tensor_exp.permute(inv_order)
+
+        return tensor_exp
+
     def __post_init__(self) -> None:
         assert len(self._arrow) == self._tensor.dim(), (
             f"Arrow length ({len(self._arrow)}) must match tensor dimensions ({self._tensor.dim()})."

tests/exponential_test.py

@@ -0,0 +1,30 @@
+import torch
+import pytest
+
+from grassmann_tensor import GrassmannTensor
+
+
+def test_exponential() -> None:
+    a = GrassmannTensor(
+        (True, True, True, True),
+        ((4, 4), (8, 8), (4, 4), (8, 8)),
+        torch.randn(8, 16, 8, 16, dtype=torch.float64),
+    )
+    a.exponential((0, 3))
+
+
+def test_exponential_with_empty_parity_block() -> None:
+    a = GrassmannTensor((False, True), ((1, 0), (1, 0)), torch.randn(1, 1))
+    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))