diff --git a/.gitignore b/.gitignore index 630e96a..b6723c8 100644 --- a/.gitignore +++ b/.gitignore @@ -3,5 +3,7 @@ env __pycache__ grassmann_tensor/_version.py build +.idea +.venv .coverage coverage.xml diff --git a/grassmann_tensor/tensor.py b/grassmann_tensor/tensor.py index 30ec73e..bc849b9 100644 --- a/grassmann_tensor/tensor.py +++ b/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() diff --git a/tests/arithmetic_test.py b/tests/arithmetic_test.py new file mode 100644 index 0000000..3f41ec8 --- /dev/null +++ b/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 diff --git a/tests/clone_test.py b/tests/clone_test.py new file mode 100644 index 0000000..6a67bcc --- /dev/null +++ b/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)