Refactor contract api

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):

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)

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))

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