Fix pyright lints, update test, fix for 1D

Author: sgreenbury

autoemulate/core/metrics.py

@@ -83,8 +83,12 @@ def __call__(
                 y_pred = y_pred.rsample((n_samples,)).mean(dim=0)
         metric = self.metric()
         metric.to(y_pred.device)
-        # Assume first dim is a batch dim, flatten others for metric calculation
-        metric.update(y_pred.flatten(start_dim=1), y_true.flatten(start_dim=1))
+
+        # Assume first dim is a batch dim if >=2D, flatten others for metric calculation
+        metric.update(
+            y_pred.flatten(start_dim=1) if y_pred.ndim > 1 else y_pred,
+            y_true.flatten(start_dim=1) if y_true.ndim > 1 else y_true,
+        )
         return metric.compute()
 
 

autoemulate/core/model_selection.py

@@ -10,7 +10,7 @@
     get_torch_device,
     move_tensors_to_device,
 )
-from autoemulate.core.metrics import R2, Metric, TorchMetrics, get_metric_configs
+from autoemulate.core.metrics import R2, Metric, get_metric_configs
 from autoemulate.core.types import (
     DeviceLike,
     ModelParams,
@@ -63,7 +63,7 @@ def cross_validate(
     y_transforms: list[Transform] | None = None,
     device: DeviceLike = "cpu",
     random_seed: int | None = None,
-    metrics: list[TorchMetrics] | None = None,
+    metrics: list[Metric] | None = None,
 ):
     """
     Cross validate model performance using the given `cv` strategy.
@@ -158,7 +158,7 @@ def bootstrap(
     n_bootstraps: int | None = 100,
     n_samples: int = 1000,
     device: str | torch.device = "cpu",
-    metrics: list[TorchMetrics] | None = None,
+    metrics: list[Metric] | None = None,
 ) -> dict[str, tuple[float, float]]:
     """
     Get bootstrap estimates of metrics.

autoemulate/core/tuner.py

@@ -6,7 +6,7 @@
 from torch.distributions import Transform
 
 from autoemulate.core.device import TorchDeviceMixin
-from autoemulate.core.metrics import TorchMetrics, get_metric_config
+from autoemulate.core.metrics import Metric, get_metric_config
 from autoemulate.core.model_selection import cross_validate
 from autoemulate.core.types import (
     DeviceLike,
@@ -46,7 +46,7 @@ def __init__(
         n_iter: int = 10,
         device: DeviceLike | None = None,
         random_seed: int | None = None,
-        tuning_metric: str | TorchMetrics = "r2",
+        tuning_metric: str | Metric = "r2",
     ):
         TorchDeviceMixin.__init__(self, device=device)
         self.n_iter = n_iter

tests/core/test_metrics.py

@@ -396,18 +396,13 @@ def test_metric_with_multidimensional_tensors():
 def test_metric_configs_workflow():
     """Test complete workflow of getting and using metric configs."""
     # Get configs from strings
-    configs = get_metric_configs(["r2", "rmse"])
+    metrics = get_metric_configs(["r2", "rmse"])
 
     # Use configs to compute metrics
     y_pred = torch.tensor([1.0, 2.0, 3.0])
     y_true = torch.tensor([1.0, 2.0, 3.0])
 
-    results = {}
-    for config in configs:
-        metric = config.metric()
-        metric.update(y_pred, y_true)
-        results[config.name] = metric.compute()
-
+    results = {metric.name: metric(y_pred, y_true) for metric in metrics}
     assert "r2" in results
     assert "rmse" in results
     assert torch.isclose(results["r2"], torch.tensor(1.0))  # Perfect R2