Test/add test for arithmetic

.gitignore

@@ -3,5 +3,7 @@ env
 __pycache__
 grassmann_tensor/_version.py
 build
+.idea
+.venv
 .coverage
 coverage.xml

grassmann_tensor/tensor.py

@@ -243,7 +243,9 @@ def __iadd__(self, other: typing.Any) -> GrassmannTensor:
             self._tensor += other._tensor
         else:
             self._tensor += other
-        return self
+        if isinstance(self._tensor, torch.Tensor):
+            return self
+        return NotImplemented
 
     def __sub__(self, other: typing.Any) -> GrassmannTensor:
         if isinstance(other, GrassmannTensor):
@@ -281,7 +283,9 @@ def __isub__(self, other: typing.Any) -> GrassmannTensor:
             self._tensor -= other._tensor
         else:
             self._tensor -= other
-        return self
+        if isinstance(self._tensor, torch.Tensor):
+            return self
+        return NotImplemented
 
     def __mul__(self, other: typing.Any) -> GrassmannTensor:
         if isinstance(other, GrassmannTensor):
@@ -319,7 +323,9 @@ def __imul__(self, other: typing.Any) -> GrassmannTensor:
             self._tensor *= other._tensor
         else:
             self._tensor *= other
-        return self
+        if isinstance(self._tensor, torch.Tensor):
+            return self
+        return NotImplemented
 
     def __truediv__(self, other: typing.Any) -> GrassmannTensor:
         if isinstance(other, GrassmannTensor):
@@ -357,4 +363,23 @@ def __itruediv__(self, other: typing.Any) -> GrassmannTensor:
             self._tensor /= other._tensor
         else:
             self._tensor /= other
-        return self
+        if isinstance(self._tensor, torch.Tensor):
+            return self
+        return NotImplemented
+
+    def clone(self) -> GrassmannTensor:
+        """
+        Create a deep copy of the Grassmann tensor.
+        """
+        return dataclasses.replace(
+            self,
+            _tensor=self._tensor.clone(),
+            _parity=tuple(parity.clone() for parity in self._parity) if self._parity is not None else None,
+            _mask=self._mask.clone() if self._mask is not None else None,
+        )
+
+    def __copy__(self) -> GrassmannTensor:
+        return self.clone()
+
+    def __deepcopy__(self, memo: dict) -> GrassmannTensor:
+        return self.clone()

tests/arithmetic_test.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+import typing
+import pytest
+import torch
+from grassmann_tensor import GrassmannTensor
+
+
+@pytest.fixture(params=[
+    (
+        GrassmannTensor((False, False), ((2, 2), (1, 3)), torch.randn([4, 4])),
+        GrassmannTensor((False, False), ((2, 2), (1, 3)), torch.randn([4, 4])),
+    ),
+    (
+        GrassmannTensor((True, False, True), ((1, 1), (2, 2), (3, 1)), torch.randn([2, 4, 4])),
+        GrassmannTensor((True, False, True), ((1, 1), (2, 2), (3, 1)), torch.randn([2, 4, 4])),
+    ),
+    (
+        GrassmannTensor((True, True, False, False), ((1, 2), (2, 2), (1, 1), (3, 1)), torch.randn([3, 4, 2, 4])),
+        GrassmannTensor((True, True, False, False), ((1, 2), (2, 2), (1, 1), (3, 1)), torch.randn([3, 4, 2, 4])),
+    ),
+])
+def tensors(request: pytest.FixtureRequest) -> tuple[GrassmannTensor, GrassmannTensor]:
+    return request.param
+
+
+@pytest.fixture(params=[
+    torch.randn([]),
+    torch.randn([]).item(),
+])
+def scalar(request: pytest.FixtureRequest) -> torch.Tensor | float:
+    return request.param
+
+
+class FakeTensor:
+
+    def __init__(self) -> None:
+        pass
+
+    def __add__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __radd__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __sub__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __rsub__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __mul__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __rmul__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __truediv__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+    def __rtruediv__(self, other: typing.Any) -> FakeTensor:
+        if isinstance(other, torch.Tensor):
+            return self
+        else:
+            return NotImplemented
+
+
+@pytest.mark.parametrize(
+    "unsupported_type",
+    [
+        "string",  #string
+        None,  #NoneType
+        {"key", "value"},  #dict
+        [1, 2, 3],  #list
+        {1, 2},  #set
+        object(),  #arbitrary object
+        FakeTensor(),  #a ill defined tensor-like object
+    ])
+def test_arithmetic(unsupported_type: typing.Any, tensors: tuple[GrassmannTensor, GrassmannTensor], scalar: torch.Tensor | float) -> None:
+    tensor_a, tensor_b = tensors
+
+    # Test __pos__ method.
+    assert torch.equal((+tensor_a).tensor, +tensor_a.tensor)
+
+    # Test __neg__ method.
+    assert torch.equal((-tensor_a).tensor, -tensor_a.tensor)
+
+    # Test __add__ method.
+    assert torch.equal((tensor_a + scalar).tensor, tensor_a.tensor + scalar)
+    assert torch.equal((tensor_a + tensor_b).tensor, tensor_a.tensor + tensor_b.tensor)
+    assert torch.equal((scalar + tensor_a).tensor, scalar + tensor_a.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_a + unsupported_type
+
+    with pytest.raises(TypeError):
+        unsupported_type + tensor_a
+
+    # Test __iadd__ method.
+    tensor_c = tensor_a.clone()
+    tensor_c += scalar
+    assert torch.equal(tensor_c.tensor, tensor_a.tensor + scalar)
+    tensor_c = tensor_a.clone()
+    tensor_c += tensor_b
+    assert torch.equal(tensor_c.tensor, tensor_a.tensor + tensor_b.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_c = tensor_a.clone()
+        tensor_c += unsupported_type
+
+    # Test __sub__ method.
+    assert torch.equal((tensor_a - scalar).tensor, tensor_a.tensor - scalar)
+    assert torch.equal((tensor_a - tensor_b).tensor, tensor_a.tensor - tensor_b.tensor)
+    assert torch.equal((scalar - tensor_a).tensor, scalar - tensor_a.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_a - unsupported_type
+
+    with pytest.raises(TypeError):
+        unsupported_type - tensor_a
+
+    # Test __isub__ method.
+    tensor_c = tensor_a.clone()
+    tensor_c -= scalar
+    assert torch.equal(tensor_c.tensor, tensor_a.tensor - scalar)
+    tensor_c = tensor_a.clone()
+    tensor_c -= tensor_b
+    assert torch.equal(tensor_c.tensor, tensor_a.tensor - tensor_b.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_c = tensor_a.clone()
+        tensor_c -= unsupported_type
+
+    # Test __mul__ method.
+    assert torch.allclose((tensor_a * scalar).tensor, tensor_a.tensor * scalar)
+    assert torch.allclose((tensor_a * tensor_b).tensor, tensor_a.tensor * tensor_b.tensor)
+    assert torch.allclose((scalar * tensor_a).tensor, scalar * tensor_a.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_a * unsupported_type
+
+    with pytest.raises(TypeError):
+        unsupported_type * tensor_a
+
+    # Test __imul__ method.
+    tensor_c = tensor_a.clone()
+    tensor_c *= scalar
+    assert torch.allclose(tensor_c.tensor, tensor_a.tensor * scalar)
+    tensor_c = tensor_a.clone()
+    tensor_c *= tensor_b
+    assert torch.allclose(tensor_c.tensor, tensor_a.tensor * tensor_b.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_c = tensor_a.clone()
+        tensor_c *= unsupported_type
+
+    # Test __truediv__ method.
+    assert torch.allclose((tensor_a / scalar).tensor, tensor_a.tensor / scalar)
+    assert torch.allclose((tensor_a / tensor_b).tensor, tensor_a.tensor / tensor_b.tensor)
+    assert torch.allclose((scalar / tensor_a).tensor, scalar / tensor_a.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_a / unsupported_type
+
+    with pytest.raises(TypeError):
+        unsupported_type / tensor_a
+
+    # Test __itruediv__ method.
+    tensor_c = tensor_a.clone()
+    tensor_c /= scalar
+    assert torch.allclose(tensor_c.tensor, tensor_a.tensor / scalar)
+    tensor_c = tensor_a.clone()
+    tensor_c /= tensor_b
+    assert torch.allclose(tensor_c.tensor, tensor_a.tensor / tensor_b.tensor)
+
+    with pytest.raises(TypeError):
+        tensor_c = tensor_a.clone()
+        tensor_c /= unsupported_type

tests/clone_test.py

@@ -0,0 +1,50 @@
+import typing
+import copy
+import pytest
+import torch
+from grassmann_tensor import GrassmannTensor
+
+
+@pytest.mark.parametrize("parity,mask", [(True, True), (True, False), (False, False)])
+@pytest.mark.parametrize("which", ["clone", "copy", "deepcopy"])
+def test_clone(
+    parity: bool,
+    mask: bool,
+    which: typing.Literal["clone", "copy", "deepcopy"],
+) -> None:
+    original_tensor = GrassmannTensor(
+        _arrow=(False, True),
+        _edges=((2, 2), (1, 3)),
+        _tensor=torch.randn([4, 4]),
+    )
+
+    if parity:
+        _ = original_tensor.parity
+    if mask:
+        _ = original_tensor.mask
+
+    match which:
+        case "clone":
+            cloned_tensor = original_tensor.clone()
+        case "copy":
+            cloned_tensor = copy.copy(original_tensor)
+        case "deepcopy":
+            cloned_tensor = copy.deepcopy(original_tensor)
+
+    assert cloned_tensor._arrow == original_tensor._arrow
+    assert cloned_tensor._edges == original_tensor._edges
+    assert torch.equal(cloned_tensor._tensor, original_tensor._tensor)
+    if parity:
+        assert cloned_tensor._parity is not None
+        assert original_tensor._parity is not None
+        assert all(torch.equal(c, o) for c, o in zip(cloned_tensor._parity, original_tensor._parity))
+    else:
+        assert cloned_tensor._parity is original_tensor._parity
+    if mask:
+        assert cloned_tensor._mask is not None
+        assert original_tensor._mask is not None
+        assert torch.equal(cloned_tensor._mask, original_tensor._mask)
+    else:
+        assert cloned_tensor._mask is original_tensor._mask
+
+    assert id(original_tensor.tensor) != id(cloned_tensor.tensor)