diff --git a/grassmann_tensor/tensor.py b/grassmann_tensor/tensor.py index 7e68fc8..47a75f9 100644 --- a/grassmann_tensor/tensor.py +++ b/grassmann_tensor/tensor.py @@ -10,6 +10,7 @@ import functools import typing import math + import torch @@ -785,13 +786,14 @@ def get_inv_order(dim: int, order: tuple[int, ...]) -> tuple[int, ...]: def _get_inv_order(self, order: tuple[int, ...]) -> tuple[int, ...]: return self.get_inv_order(self.tensor.dim(), order) - @staticmethod def contract( - a: GrassmannTensor, + self, b: GrassmannTensor, a_leg: int | tuple[int, ...], b_leg: int | tuple[int, ...], ) -> GrassmannTensor: + a = self + contract_lengths = [] for leg, tensor in ((a_leg, a), (b_leg, b)): if isinstance(leg, int): diff --git a/tests/conjugate_test.py b/tests/conjugate_test.py index c0263c7..1d505b5 100644 --- a/tests/conjugate_test.py +++ b/tests/conjugate_test.py @@ -84,13 +84,13 @@ def test_conjugate_reverse_order_of_contraction( a = create_random_tensor(arrow_edges, dtype=torch.complex128).update_mask() b = create_random_tensor(arrow_edges, dtype=torch.complex128).update_mask() - contrast_a = GrassmannTensor.contract(a, b, a.tensor.dim() - 1, 0) + contrast_a = a.contract(b, a.tensor.dim() - 1, 0) contrast_a = contrast_a.conj() a_conj = a.conj() b_conj = b.conj() - contrast_b = GrassmannTensor.contract(a_conj, b_conj, a_conj.tensor.dim() - 1, 0) + contrast_b = a_conj.contract(b_conj, a_conj.tensor.dim() - 1, 0) assert contrast_a.arrow == contrast_b.arrow assert torch.allclose(contrast_a.tensor, contrast_b.tensor) diff --git a/tests/contract_test.py b/tests/contract_test.py index 6c494d8..e89f8af 100644 --- a/tests/contract_test.py +++ b/tests/contract_test.py @@ -15,9 +15,9 @@ def test_contract() -> None: ((8, 8), (4, 4), (4, 4), (32, 32)), torch.randn(16, 8, 8, 64, dtype=torch.float64), ) - _ = GrassmannTensor.contract(a, b, 3, 0) - _ = GrassmannTensor.contract(a, b, (0, 2), 3) - _ = GrassmannTensor.contract(a, b, (0, 2), (1, 2)) + _ = a.contract(b, 3, 0) + _ = a.contract(b, (0, 2), 3) + _ = a.contract(b, (0, 2), (1, 2)) def test_contract_assertion() -> None: @@ -28,8 +28,8 @@ def test_contract_assertion() -> None: torch.randn(4, 8, 16, 32, dtype=torch.float64), ) with pytest.raises(AssertionError, match="Contract requires arrow"): - _ = a.contract(a, b, 0, 0) + _ = 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(a, b, 0, (0, 1)) + _ = a.contract(b, 0, (0, 1)) diff --git a/tests/svd_test.py b/tests/svd_test.py index dcff3b9..174d8ef 100644 --- a/tests/svd_test.py +++ b/tests/svd_test.py @@ -106,8 +106,8 @@ def test_svd( gt = GrassmannTensor(arrow, edges, tensor) U, S, Vh = gt.svd(free_names_u, cutoff=cutoff) - US = GrassmannTensor.contract(U, S, U.tensor.dim() - 1, 0) - USV = GrassmannTensor.contract(US, Vh, US.tensor.dim() - 1, 0) + US = U.contract(S, U.tensor.dim() - 1, 0) + USV = US.contract(Vh, US.tensor.dim() - 1, 0) left_legs, right_legs = gt.get_legs_pair(len(edges), free_names_u) order = left_legs + right_legs