Redesign the contract method to improve robustness and performance

grassmann_tensor/tensor.py

@@ -790,99 +790,165 @@ def get_inv_order(order: tuple[int, ...]) -> tuple[int, ...]:
     def contract(
         self,
         b: GrassmannTensor,
-        a_leg: int | tuple[int, ...],
-        b_leg: int | tuple[int, ...],
+        leg_a: int | tuple[int, ...],
+        leg_b: int | tuple[int, ...],
     ) -> GrassmannTensor:
+        """
+        This function is used to contract the GrassmannTensor with other GrassmannTensors.
+        For two fermion tensors A and B, their edges are divided into two groups: common edges and free edges.
+        Common edges connect each other, while the other edges are free edges.
+        The following operations are performed in sequence:
+        1. `Permutation`
+        Place all free edges in A on the left and common edges on the right, while place all common edges in B on
+        left and free edges on the right.
+        2. `Reverse`
+        Reverse the fermionic arrows of all free edges in A to False. If a sign is generated, the sign is not placed in
+        this tensor. Reverse the fermionic arrows of all common edges in A to True. If a sign is generated, the sign is
+        placed in this tensor. Reverse the fermionic arrows of all free edges in B to False. If a sign is generated, the
+        sign is not placed in this tensor. Reverse the fermionic arrows of all common edges in B to False. If a sign is
+        generated, the sign is not placed in this tensor.
+        3. `Merge`
+        If merging free edges of A generates a sign, it is not included in this tensor; If merging common edges of A
+        generates a sign, it is included in this tensor; If merging free edges of B generates a sign, it is not included
+        in this tensor; If merging common edges of B generates a sign, it is not included in this tensor.
+        4. `Matrix multiplication`
+        5. `Split`
+        Split the remaining two free edges to restore the original shape of the tensor, and do not place any of the
+        generated signs in this tensor.
+        6. `Reverse`
+        Reverse the fermionic arrows back to their original orientation in A and B, and the resulting sign is not placed
+        in this tensor.
+        """
         a = self
 
-        contract_lengths = []
-        for leg, tensor in ((a_leg, a), (b_leg, b)):
-            if isinstance(leg, int):
-                contract_lengths.append(1)
-            else:
-                contract_lengths.append(len(leg))
-                assert all(tensor.arrow[i] == tensor.arrow[leg[0]] for i in leg), (
-                    "All the legs that need to be contracted must have the same arrow"
-                )
+        leg_tuple_a = (leg_a,) if isinstance(leg_a, int) else leg_a
+        leg_tuple_b = (leg_b,) if isinstance(leg_b, int) else leg_b
 
-        contract_length_a, contract_length_b = contract_lengths
+        for tensor, leg in ((a, leg_tuple_a), (b, leg_tuple_b)):
+            assert len(set(leg)) == len(leg), f"Indices must be unique. Got {leg}."
+            assert all(0 <= i < tensor.tensor.dim() for i in leg), (
+                f"Indices must be within tensor dimensions. Got {leg}."
+            )
+
+        contract_length_a, contract_length_b = (
+            1 if isinstance(leg, int) else len(leg) for leg in (leg_a, leg_b)
+        )
 
-        a_leg_tuple = (a_leg,) if isinstance(a_leg, int) else a_leg
-        b_leg_tuple = (b_leg,) if isinstance(b_leg, int) else b_leg
+        dim_a = a.tensor.dim()
+        dim_b = b.tensor.dim()
 
-        a_range_list = tuple(range(a.tensor.dim()))
-        b_range_list = tuple(range(b.tensor.dim()))
+        right_leg_a, left_leg_a = self.get_legs_pair(dim_a, leg_tuple_a)
+        left_leg_b, right_leg_b = self.get_legs_pair(dim_b, leg_tuple_b)
 
-        a_contract_set = set(a_leg_tuple)
-        b_contract_set = set(b_leg_tuple)
+        order_a = left_leg_a + right_leg_a
+        order_b = left_leg_b + right_leg_b
 
-        order_a = tuple(i for i in a_range_list if i not in a_contract_set) + a_leg_tuple
-        order_b = b_leg_tuple + tuple(i for i in b_range_list if i not in b_contract_set)
+        # 1. Permutation
+        a = a.permute(order_a)
+        b = b.permute(order_b)
 
-        tensor_a = a.permute(order_a)
-        tensor_b = b.permute(order_b)
+        arrow = a.arrow[:-contract_length_a] + b.arrow[contract_length_b:]
+        edges = a.edges[:-contract_length_a] + b.edges[contract_length_b:]
+        shape = a.tensor.shape[:-contract_length_a] + b.tensor.shape[contract_length_b:]
 
-        assert (tensor_a.arrow[-1], tensor_b.arrow[0]) in ((False, True), (True, False)), (
-            f"Contract requires arrow (False, True) or (True, False), but got {tensor_a.arrow[-1], tensor_b.arrow[0]}"
+        # 2. Reverse
+        # Reverse tensor `a`
+        arrow_after_reverse_a = tuple(
+            [False] * (dim_a - contract_length_a) + [True] * contract_length_a
         )
 
-        arrow_after_permute_a = tensor_a.arrow
-        arrow_after_permute_b = tensor_b.arrow
+        tensor_after_reverse = a.tensor
+        # Only calculate the sign of common edges of `a`
+        parity_index = tuple(range(dim_a - contract_length_a, dim_a))
+        reversing_parity = functools.reduce(
+            torch.logical_xor,
+            (
+                self._unsqueeze(parity, index, self.tensor.dim())
+                for index, parity in enumerate(a.parity)
+                if index in parity_index and not a.arrow[index]
+            ),
+            torch.zeros([], dtype=torch.bool, device=self.tensor.device),
+        )
+        tensor_after_reverse = torch.where(
+            reversing_parity, -tensor_after_reverse, +tensor_after_reverse
+        )
 
-        edge_after_permute_a = tensor_a.edges
-        edge_after_permute_b = tensor_b.edges
+        a = dataclasses.replace(
+            a,
+            _arrow=arrow_after_reverse_a,
+            _tensor=tensor_after_reverse,
+        )
 
-        arrow_expected_a = [i >= a.tensor.dim() - contract_length_a for i in range(a.tensor.dim())]
-        arrow_expected_b = [i >= contract_length_b for i in range(b.tensor.dim())]
+        # 3. Merge
+        # Merge tensor `a` free edges
+        arrow_merge_free_edges = tuple([False] + [True] * contract_length_a)
+        edges_merge_free_edges = typing.cast(
+            tuple[tuple[int, int], ...],
+            (a.calculate_even_odd(a.edges[:-contract_length_a]), *a.edges[-contract_length_a:]),
+        )
+        shape_merge_free_edges = (
+            math.prod(a.tensor.shape[:-contract_length_a]),
+            *a.tensor.shape[-contract_length_a:],
+        )
+        tensor_merge_free_edges = a.tensor.reshape(shape_merge_free_edges)
+        a = dataclasses.replace(
+            a,
+            _arrow=arrow_merge_free_edges,
+            _edges=edges_merge_free_edges,
+            _tensor=tensor_merge_free_edges,
+        )
 
-        arrow_reverse_a = tuple(
-            i
-            for i, (cur, exp) in enumerate(zip(arrow_after_permute_a, arrow_expected_a))
-            if cur != exp
+        # Merge tensor `a` common edges
+        arrow_merge_common_edges = (False, True)
+        shape_merge_common_edges = (a.tensor.shape[0], math.prod(a.tensor.shape[1:]))
+        even, odd, reorder, sign = self._reorder_indices(a.edges[1:])
+        edges_merge_common_edges = typing.cast(
+            tuple[tuple[int, int], ...],
+            (a.edges[0], (even, odd)),
+        )
+        tensor_merge_common_edges = a.tensor.reshape(shape_merge_common_edges)
+        merging_parity = functools.reduce(
+            torch.logical_xor,
+            (self._unsqueeze(sign, 1, tensor_merge_common_edges.dim())),
+            torch.zeros([], dtype=torch.bool, device=self.tensor.device),
         )
-        arrow_reverse_b = tuple(
-            i
-            for i, (cur, exp) in enumerate(zip(arrow_after_permute_b, arrow_expected_b))
-            if cur != exp
+        tensor_merge_common_edges = torch.where(
+            merging_parity, -tensor_merge_common_edges, +tensor_merge_common_edges
         )
+        tensor_merge_common_edges = tensor_merge_common_edges.index_select(1, reorder)
 
-        if arrow_reverse_a:
-            tensor_a = (
-                tensor_a.reverse(arrow_reverse_a).reverse(arrow_reverse_a).reverse(arrow_reverse_a)
-            )
-        if arrow_reverse_b:
-            tensor_b = (
-                tensor_b.reverse(arrow_reverse_b).reverse(arrow_reverse_b).reverse(arrow_reverse_b)
-            )
+        a = dataclasses.replace(
+            a,
+            _arrow=arrow_merge_common_edges,
+            _edges=edges_merge_common_edges,
+            _tensor=tensor_merge_common_edges,
+        )
 
-        tensor_a = tensor_a.reshape(
-            (
-                math.prod(tensor_a.tensor.shape[:-contract_length_a]),
-                math.prod(tensor_a.tensor.shape[-contract_length_a:]),
-            )
+        # Reverse and merge tensor `b`
+        arrow_after_reverse_merge_b = (False, True)
+        edges_after_reverse_merge_b = (
+            self.calculate_even_odd(b.edges[:contract_length_b]),
+            self.calculate_even_odd(b.edges[contract_length_b:]),
         )
-        tensor_b = tensor_b.reshape(
+        tensor_after_reverse_merge_b = b.tensor.reshape(
             (
-                math.prod(tensor_b.tensor.shape[:contract_length_b]),
-                math.prod(tensor_b.tensor.shape[contract_length_b:]),
+                math.prod(b.tensor.shape[:contract_length_b]),
+                math.prod(b.tensor.shape[contract_length_b:]),
             )
         )
 
-        c = tensor_a @ tensor_b
-
-        c = c.reshape(
-            (edge_after_permute_a[:-contract_length_a] + edge_after_permute_b[contract_length_b:])
+        b = dataclasses.replace(
+            b,
+            _arrow=arrow_after_reverse_merge_b,
+            _edges=edges_after_reverse_merge_b,
+            _tensor=tensor_after_reverse_merge_b,
         )
 
-        arrow_reverse_c = tuple(
-            [i for i in arrow_reverse_a if i < a.tensor.dim() - contract_length_a]
-            + [
-                (a.tensor.dim() - contract_length_a) + (i - contract_length_b)
-                for i in arrow_reverse_b
-                if i >= contract_length_b
-            ]
-        )
-        c = c.reverse(arrow_reverse_c)
+        # 4. Matrix multiplication
+        c = a @ b
+
+        # Split and reverse back
+        c = dataclasses.replace(c, _arrow=arrow, _edges=edges, _tensor=c.tensor.reshape(shape))
         return c
 
     def exponential(self, pairs: tuple[tuple[int, ...], tuple[int, ...]]) -> GrassmannTensor:

tests/contract_test.py

@@ -1,35 +1,101 @@
 import torch
 import pytest
+from _pytest.mark.structures import ParameterSet
+import random
+from typing import Iterable
 
 from grassmann_tensor import GrassmannTensor
 
+ContractCases = Iterable[ParameterSet]
 
-def test_contract() -> None:
+
+def contract_cases() -> ContractCases:
+    edge_unit = (4, 4)
+    max_dim = 4
+    num_cases = 10
+
+    rng = random.Random(0)
+    gen = torch.Generator().manual_seed(0)
+
+    cases = []
+
+    for case_idx in range(num_cases):
+        dim = rng.randint(1, max_dim)
+
+        edges = tuple(edge_unit for _ in range(dim))
+
+        shape = (sum(edge_unit),) * dim
+
+        arrow = tuple(bool(rng.getrandbits(1)) for _ in range(dim))
+
+        tensor = torch.randn(*shape, dtype=torch.float64, generator=gen)
+
+        a = GrassmannTensor(
+            arrow,
+            edges,
+            tensor,
+        )
+
+        contract_length = rng.randint(1, dim)
+        if contract_length == 1:
+            leg_a: int | tuple[int, ...] = (
+                (sorted(rng.sample(range(dim), contract_length)))[0]
+                if rng.random() < 0.5
+                else tuple(sorted(rng.sample(range(dim), contract_length)))
+            )
+            leg_b: int | tuple[int, ...] = (
+                (sorted(rng.sample(range(dim), contract_length)))[0]
+                if rng.random() < 0.5
+                else tuple(sorted(rng.sample(range(dim), contract_length)))
+            )
+        else:
+            leg_a = tuple(sorted(rng.sample(range(dim), contract_length)))
+            leg_b = tuple(sorted(rng.sample(range(dim), contract_length)))
+
+        cases.append(
+            pytest.param(
+                a,
+                leg_a,
+                leg_b,
+                id=f"arrow={arrow}-dim={dim}-leg_a={leg_a}-leg_b={leg_b}",
+            )
+        )
+
+    return cases
+
+
+@pytest.mark.parametrize("a, leg_a, leg_b", contract_cases())
+def test_contract(
+    a: GrassmannTensor,
+    leg_a: int | tuple[int, ...],
+    leg_b: int | tuple[int, ...],
+) -> None:
+    _ = a.contract(a, leg_a, leg_b)
+
+
+def test_contract_assertion() -> None:
+    a = GrassmannTensor((False, True), ((4, 4), (4, 4)), torch.randn(8, 8, dtype=torch.float64))
+    b = GrassmannTensor(
+        (False, True, False, True),
+        ((4, 4), (4, 4), (4, 4), (4, 4)),
+        torch.randn(8, 8, 8, 8, dtype=torch.float64),
+    )
+    with pytest.raises(AssertionError, match="Indices must be unique"):
+        _ = a.contract(b, (0, 0), 0)
+    with pytest.raises(AssertionError, match="Indices must be within tensor dimensions"):
+        _ = a.contract(b, 0, (0, 4))
+
+
+def test_contract_full_legs() -> None:
     a = GrassmannTensor(
         (False, False, False, True),
         ((2, 2), (4, 4), (8, 8), (8, 8)),
         torch.randn(4, 8, 16, 16, dtype=torch.float64),
     )
     b = GrassmannTensor(
         (False, True, True, True),
-        ((8, 8), (4, 4), (4, 4), (32, 32)),
-        torch.randn(16, 8, 8, 64, dtype=torch.float64),
-    )
-    _ = a.contract(b, 3, 0)
-    _ = a.contract(b, (0, 2), 3)
-    _ = a.contract(b, (0, 2), (1, 2))
-
-
-def test_contract_assertion() -> None:
-    a = GrassmannTensor((False, True), ((1, 0), (1, 0)), torch.randn(1, 1, dtype=torch.float64))
-    b = GrassmannTensor(
-        (False, True, False, True),
-        ((2, 2), (4, 4), (8, 8), (16, 16)),
-        torch.randn(4, 8, 16, 32, dtype=torch.float64),
+        ((8, 8), (8, 8), (4, 4), (2, 2)),
+        torch.randn(16, 16, 8, 4, dtype=torch.float64),
     )
-    with pytest.raises(AssertionError, match="Contract requires arrow"):
-        _ = a.contract(b, 0, 0)
-    with pytest.raises(
-        AssertionError, match="All the legs that need to be contracted must have the same arrow"
-    ):
-        _ = a.contract(b, 0, (0, 1))
+    c = a.contract(b, (0, 1, 2, 3), (0, 1, 2, 3))
+    assert c.tensor.dim() == 0