clean up

src/graph_pes/training/loss.py

@@ -82,36 +82,6 @@ def __repr__(self) -> str:
             metric=self.metric,
         )
 
-    ## Methods for creating total losses ##
-    def __mul__(self, weight: float | int) -> TotalLoss:
-        if not isinstance(weight, (int, float)):
-            raise TypeError(f"Cannot multiply Loss and {type(weight)}")
-
-        return TotalLoss([self], [weight])
-
-    def __rmul__(self, weight: float) -> TotalLoss:
-        if not isinstance(weight, (int, float)):
-            raise TypeError(f"Cannot multiply Loss and {type(weight)}")
-
-        return TotalLoss([self], [weight])
-
-    def __truediv__(self, weight: float | int) -> TotalLoss:
-        if not isinstance(weight, (int, float)):
-            raise TypeError(f"Cannot divide Loss and {type(weight)}")
-
-        return TotalLoss([self], [1 / weight])
-
-    def __add__(self, loss: Loss | TotalLoss) -> TotalLoss:
-        if isinstance(loss, Loss):
-            return TotalLoss([self, loss], [1, 1])
-        elif isinstance(loss, TotalLoss):
-            return TotalLoss([self] + list(loss.losses), [1] + loss.weights)
-        else:
-            raise TypeError(f"Cannot add Loss and {type(loss)}")
-
-    def __radd__(self, other: Loss | TotalLoss) -> TotalLoss:
-        return self.__add__(other)
-
 
 class SubLossPair(NamedTuple):
     loss_value: torch.Tensor
@@ -157,34 +127,6 @@ def __init__(
         self.losses = UniformModuleList(losses)
         self.weights = weights or [1.0] * len(losses)
 
-    def __add__(self, other: TotalLoss | Loss) -> TotalLoss:
-        if isinstance(other, Loss):
-            return TotalLoss(self.losses + [other], self.weights + [1.0])
-        elif isinstance(other, TotalLoss):
-            return TotalLoss(
-                self.losses + other.losses, self.weights + other.weights
-            )
-        else:
-            raise TypeError(f"Cannot add TotalLoss and {type(other)}")
-
-    def __mul__(self, other: float | int) -> TotalLoss:
-        if not isinstance(other, (int, float)):
-            raise TypeError(f"Cannot multiply TotalLoss and {type(other)}")
-
-        return TotalLoss(self.losses, [w * other for w in self.weights])
-
-    def __rmul__(self, other: float | int) -> TotalLoss:
-        if not isinstance(other, (int, float)):
-            raise TypeError(f"Cannot multiply TotalLoss and {type(other)}")
-
-        return TotalLoss(self.losses, [w * other for w in self.weights])
-
-    def __true_div__(self, other: float | int) -> TotalLoss:
-        if not isinstance(other, (int, float)):
-            raise TypeError(f"Cannot divide TotalLoss and {type(other)}")
-
-        return TotalLoss(self.losses, [w / other for w in self.weights])
-
     def forward(
         self,
         predictions: dict[keys.LabelKey, torch.Tensor],
@@ -310,27 +252,6 @@ def __init__(self):
         super().__init__()
 
 
-class MeanVectorPercentageError(torch.nn.Module):
-    r"""
-    Mean vector percentage error metric:
-
-    .. math::
-        \frac{1}{N} \sum_i^N \frac{\left{||} \hat{v}_i - v_i \right{||}}
-        {||v_i|| + \varepsilon}
-    """
-
-    def __init__(self, epsilon: float = 1e-6):
-        super().__init__()
-        self.epsilon = epsilon
-
-    def forward(
-        self, input: torch.Tensor, target: torch.Tensor
-    ) -> torch.Tensor:
-        return (
-            (input - target).norm(dim=-1) / (target.norm(dim=-1) + self.epsilon)
-        ).mean()
-
-
 def _get_metric_name(metric: Callable[[Tensor, Tensor], Tensor]) -> str:
     # if metric is a function, we want the function's name, otherwise
     # we want the metric's class name, all lowercased

tests/test_loss.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import pytest
 import torch
-from graph_pes.training.loss import MAE, RMSE, Loss, TotalLoss
+from graph_pes.training.loss import MAE, RMSE
 
 
 def test_metrics():
@@ -15,42 +14,3 @@ def test_metrics():
     c = torch.tensor([0, 0, 0]).float()
     assert torch.allclose(MAE()(a, c), torch.tensor(2.0))
     assert torch.allclose(RMSE()(a, c), torch.tensor((1 + 4 + 9) / 3).sqrt())
-
-
-def test_loss_ops():
-    l1 = Loss("energy")
-    l2 = Loss("forces")
-
-    weighted = 10 * l1 + l2 * 1
-    assert isinstance(weighted, TotalLoss)
-    assert set(weighted.losses) == {l1, l2}
-    assert set(weighted.weights) == {10, 1}
-
-    weighted = l1 / 2 + l2 / 3
-    assert isinstance(weighted, TotalLoss)
-    assert set(weighted.losses) == {l1, l2}
-    assert set(weighted.weights) == {0.5, 1 / 3}
-
-    weighted = l1 + l1 + l1
-    assert isinstance(weighted, TotalLoss)
-    assert len(weighted.losses) == 3
-
-
-def test_warnings():
-    l1 = Loss("energy")
-    l2 = Loss("forces")
-
-    with pytest.raises(TypeError, match="Cannot multiply Loss and"):
-        l1 * l2  # type: ignore
-    with pytest.raises(TypeError, match="Cannot divide Loss and"):
-        l1 / l2  # type: ignore
-    with pytest.raises(TypeError, match="Cannot multiply Loss and"):
-        l1 * "hello"  # type: ignore
-    with pytest.raises(TypeError, match="Cannot add Loss and"):
-        l1 + "hello"  # type: ignore
-
-    total = TotalLoss([l1, l2], [1, 1])
-    with pytest.raises(TypeError, match="Cannot multiply TotalLoss and"):
-        total * l1  # type: ignore
-    with pytest.raises(TypeError, match="Cannot add TotalLoss and"):
-        total + "hello"  # type: ignore