Merge pull request #62 from alan-turing-institute/metrics-dev

Add metrics (#20)

Author: ContiPaolo

notebooks/00_exploration.ipynb

@@ -72,17 +72,25 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "id": "5",
    "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
    "source": [
     "### Example batch\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -94,7 +102,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7",
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -126,7 +134,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -135,7 +143,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9",
+   "id": "10",
    "metadata": {},
    "source": [
     "### Run trainer\n"
@@ -144,21 +152,21 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "10",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
     "import lightning as L\n",
     "\n",
     "device = \"mps\"  # \"cpu\"\n",
     "# device = \"cpu\"\n",
-    "trainer = L.Trainer(max_epochs=1, accelerator=device, log_every_n_steps=10)\n",
+    "trainer = L.Trainer(max_epochs=10, accelerator=device, log_every_n_steps=10)\n",
     "trainer.fit(model, datamodule.train_dataloader(), datamodule.val_dataloader())"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "11",
+   "id": "12",
    "metadata": {},
    "source": [
     "### Run the evaluation\n"
@@ -167,7 +175,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "12",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -176,7 +184,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "13",
+   "id": "14",
    "metadata": {},
    "source": [
     "### Example rollout\n"
@@ -185,7 +193,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "14",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -196,7 +204,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "15",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -209,7 +217,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "16",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -220,7 +228,18 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "17",
+   "id": "18",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert preds.shape == trues.shape\n",
+    "mse_error = MSE()(preds, trues, trues)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -230,18 +249,49 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "18",
+   "id": "20",
    "metadata": {},
    "outputs": [],
    "source": [
     "assert trues is not None\n",
     "print(trues.shape)\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from the_well.benchmark.metrics import RMSE, MAE, MSE\n",
+    "from the_well.data.datasets import WellMetadata\n",
+    "\n",
+    "rmse_error = RMSE.eval(preds, trues, WellMetadata)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "22",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "WellMetadata.n_spatial_dims"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "auto-cast",
    "language": "python",
    "name": "python3"
   },
@@ -255,7 +305,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.12"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,

src/auto_cast/metrics/__init__.py

@@ -0,0 +1,5 @@
+from .spatiotemporal import MAE, MSE, NMAE, NMSE, NRMSE, RMSE, VMSE, VRMSE, LInfinity
+
+__all__ = ["MAE", "MSE", "NMAE", "NMSE", "NRMSE", "RMSE", "VMSE", "VRMSE", "LInfinity"]
+
+ALL_METRICS = (MSE, MAE, NMAE, NMSE, RMSE, NRMSE, VMSE, VRMSE, LInfinity)

src/auto_cast/metrics/base.py

@@ -0,0 +1,57 @@
+import numpy as np
+import torch
+from torch import nn
+
+from auto_cast.types import TensorBCTSPlus
+
+
+class Metric(nn.Module):
+    """
+    Base class for metrics.
+
+    This class standardizes the input arguments and
+    checks the dimensions of the input tensors.
+
+    Args:
+        f: function
+            Metric function that takes in the following arguments:
+            y_pred: torch.Tensor | np.ndarray
+                Predicted values tensor.
+            y_true: torch.Tensor | np.ndarray
+                Target values tensor.
+            **kwargs : dict
+                Additional arguments for the metric.
+    """
+
+    def forward(self, *args, **kwargs):
+        assert len(args) >= 2, (
+            "At least two arguments required (y_pred, y_true, n_spatial_dims)"
+        )
+        y_pred, y_true, n_spatial_dims = args[:3]
+
+        # Convert y_pred and y_true to torch.Tensor if they are np.ndarray
+        if isinstance(y_pred, np.ndarray):
+            y_pred = torch.from_numpy(y_pred)
+        if isinstance(y_true, np.ndarray):
+            y_true = torch.from_numpy(y_true)
+        assert isinstance(y_pred, torch.Tensor), (
+            "y_pred must be a torch.Tensor or np.ndarray"
+        )
+        assert isinstance(y_true, torch.Tensor), (
+            "y_true must be a torch.Tensor or np.ndarray"
+        )
+
+        # Check dimensions
+        assert y_pred.ndim >= n_spatial_dims + 1, (
+            "y_pred must have at least n_spatial_dims + 1 dimensions"
+        )
+        assert y_true.ndim >= n_spatial_dims + 1, (
+            "y_true must have at least n_spatial_dims + 1 dimensions"
+        )
+        return self.score(y_pred, y_true, n_spatial_dims, **kwargs)
+
+    @staticmethod
+    def score(
+        y_pred: TensorBCTSPlus, y_true: TensorBCTSPlus, n_spatial_dims: int, **kwargs
+    ):
+        raise NotImplementedError