fixing docs for metrics

Author: ContiPaolo

src/auto_cast/metrics/__init__.py

@@ -2,6 +2,4 @@
 
 __all__ = ["MAE", "MSE", "NMAE", "NMSE", "NRMSE", "RMSE", "VMSE", "VRMSE", "LInfinity"]
 
-ALL_METRICS = (
-    MSE, MAE, NMAE, NMSE, RMSE, NRMSE, VMSE, VRMSE, LInfinity
-)
+ALL_METRICS = (MSE, MAE, NMAE, NMSE, RMSE, NRMSE, VMSE, VRMSE, LInfinity)

src/auto_cast/metrics/base.py

@@ -2,10 +2,15 @@
 import torch
 from torch import nn
 
+from auto_cast.types import TensorBCTSPlus
+
 
 class Metric(nn.Module):
     """
-    Decorator for metrics that standardizes the input arguments and checks the dimensions of the input tensors.
+    Base class for metrics.
+
+    This class standardizes the input arguments and
+    checks the dimensions of the input tensors.
 
     Args:
         f: function
@@ -46,5 +51,7 @@ def forward(self, *args, **kwargs):
         return self.score(y_pred, y_true, n_spatial_dims, **kwargs)
 
     @staticmethod
-    def score(y_pred, y_true, n_spatial_dims, **kwargs):
+    def score(
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
+    ):
         raise NotImplementedError

src/auto_cast/metrics/spatiotemporal.py

@@ -1,19 +1,18 @@
-import numpy as np
 import torch
 
 from auto_cast.metrics.base import Metric
-from auto_cast.types import TensorBTC, TensorBTSPlusC
+from auto_cast.types import TensorBCTSPlus, TensorBTC
 
 
 class MSE(Metric):
+    """Mean Squared Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
-        n_spatial_dims: int,
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
     ) -> TensorBTC:
         """
-        Mean Squared Error
+        Compute Mean Squared Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -30,14 +29,14 @@ def score(
 
 
 class MAE(Metric):
+    """Mean Absolute Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
-        n_spatial_dims: int,
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
     ) -> TensorBTC:
         """
-        Mean Absolute Error
+        Compute Mean Absolute Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -54,15 +53,18 @@ def score(
 
 
 class NMAE(Metric):
+    """Normalized Mean Absolute Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
+        y_pred: TensorBCTSPlus,
+        y_true: TensorBCTSPlus,
         n_spatial_dims: int,
         eps: float = 1e-7,
+        **kwargs,
     ) -> TensorBTC:
         """
-        Normalized Mean Absolute Error
+        Compute Normalized Mean Absolute Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -80,16 +82,19 @@ def score(
 
 
 class NMSE(Metric):
+    """Normalized Mean Squared Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
+        y_pred: TensorBCTSPlus,
+        y_true: TensorBCTSPlus,
         n_spatial_dims: int,
         eps: float = 1e-7,
         norm_mode: str = "norm",
+        **kwargs,
     ) -> TensorBTC:
         """
-        Normalized Mean Squared Error
+        Compute Normalized Mean Squared Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -98,7 +103,8 @@ def score(
                 Number of spatial dimensions.
             eps: Small value to avoid division by zero. Default is 1e-7.
             norm_mode:
-                Mode for computing the normalization factor. Can be 'norm' or 'std'. Default is 'norm'.
+                Mode for computing the normalization factor. Can be 'norm' or 'std'.
+                Default is 'norm'.
 
         Returns
         -------
@@ -115,14 +121,14 @@ def score(
 
 
 class RMSE(Metric):
+    """Root Mean Squared Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
-        n_spatial_dims: int,
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
     ) -> TensorBTC:
         """
-        Root Mean Squared Error
+        Compute Root Mean Squared Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -138,24 +144,28 @@ def score(
 
 
 class NRMSE(Metric):
+    """Normalized Root Mean Squared Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
+        y_pred: TensorBCTSPlus,
+        y_true: TensorBCTSPlus,
         n_spatial_dims: int,
         eps: float = 1e-7,
         norm_mode: str = "norm",
+        **kwargs,
     ) -> TensorBTC:
         """
-        Normalized Root Mean Squared Error
+        Compute Normalized Root Mean Squared Error.
 
         Args:
             y_pred: Predicted values tensor.
             y_true: Target values tensor.
             n_spatial_dims: int
                 Number of spatial dimensions.
             eps: Small value to avoid division by zero. Default is 1e-7.
-            norm_mode : Mode for computing the normalization factor. Can be 'norm' or 'std'. Default is 'norm'.
+            norm_mode : Mode for computing the normalization factor.
+                Can be 'norm' or 'std'. Default is 'norm'.
 
         Returns
         -------
@@ -168,14 +178,14 @@ def score(
 
 
 class VMSE(Metric):
+    """Variance Scaled Mean Squared Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
-        n_spatial_dims: int,
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
     ) -> TensorBTC:
         """
-        Variance Scaled Mean Squared Error
+        Compute Variance Scaled Mean Squared Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -191,14 +201,14 @@ def score(
 
 
 class VRMSE(Metric):
+    """Variance Scaled Root Mean Squared Error."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
-        n_spatial_dims: int,
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
     ) -> TensorBTC:
         """
-        Root Variance Scaled Mean Squared Error
+        Compute Root Variance Scaled Mean Squared Error.
 
         Args:
             y_pred: Predicted values tensor.
@@ -214,14 +224,14 @@ def score(
 
 
 class LInfinity(Metric):
+    """L-Infinity Norm."""
+
     @staticmethod
     def score(
-        y_pred: TensorBTSPlusC,
-        y_true: TensorBTSPlusC,
-        n_spatial_dims: int,
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
     ) -> TensorBTC:
         """
-        L-Infinity Norm
+        Compute L-Infinity Norm.
 
         Args:
             x: Input tensor.

tests/metrics/test_metrics.py

@@ -6,11 +6,10 @@
 
 
 @pytest.mark.parametrize("metric", ALL_METRICS)
-
 def test_spatiotemporal_metrics(metric):
     y_pred: TensorBTSPlusC = torch.ones((2, 3, 4, 5))
     y_true: TensorBTSPlusC = torch.ones((2, 3, 4, 5))
     n_spatial_dims = 1
 
-    error =  metric()(y_pred, y_true, n_spatial_dims)
+    error = metric()(y_pred, y_true, n_spatial_dims)
     assert torch.allclose(error.nansum(), torch.tensor(0.0))