Allow different metrics than the loss for best model selection

Author: frostedoyster

src/metatrain/experimental/phace/default-hypers.yaml

@@ -39,3 +39,4 @@ architecture:
       type: mse
       weights: {}
       reduction: sum
+    best_model_metric: "rmse_prod"

src/metatrain/experimental/phace/schema-hypers.json

@@ -133,6 +133,10 @@
         "log_mae": {
           "type": "boolean"
         },
+        "best_model_metric": {
+          "type": "string",
+          "enum": ["rmse_prod", "mae_prod", "loss"]
+        },
         "loss": {
           "type": "object",
           "properties": {

src/metatrain/experimental/phace/trainer.py

@@ -17,7 +17,7 @@
 from ...utils.io import check_file_extension
 from ...utils.logging import MetricLogger
 from ...utils.loss import TensorMapDictLoss
-from ...utils.metrics import MAEAccumulator, RMSEAccumulator
+from ...utils.metrics import MAEAccumulator, RMSEAccumulator, get_selected_metric
 from ...utils.neighbor_lists import (
     get_requested_neighbor_lists,
     get_system_with_neighbor_lists,
@@ -430,8 +430,11 @@ def train(
                         patience=self.hypers["scheduler_patience"],
                     )
 
-            if val_loss < self.best_loss:
-                self.best_loss = val_loss
+            metric = get_selected_metric(
+                finalized_val_info, self.hypers["best_model_metric"]
+            )
+            if metric < self.best_loss:
+                self.best_loss = metric
                 self.best_model_state_dict = copy.deepcopy(
                     (
                         scripted_model.module if is_distributed else scripted_model

src/metatrain/utils/metrics.py

@@ -178,3 +178,24 @@ def finalize(
             finalized_info[out_key] = value[0] / value[1]
 
         return finalized_info
+
+
+def get_selected_metric(metric_dict: Dict[str, float], selected_metric: str) -> float:
+    if selected_metric == "loss":
+        metric = metric_dict["loss"]
+    elif selected_metric == "rmse_prod":
+        metric = 1
+        for key in metric_dict:
+            if "RMSE" in key:
+                metric *= metric_dict[key]
+    elif selected_metric == "mae_prod":
+        metric = 1
+        for key in metric_dict:
+            if "MAE" in key:
+                metric *= metric_dict[key]
+    else:
+        raise ValueError(
+            f"Selected metric {selected_metric} not recognized. "
+            "Please select from 'loss', 'rmse_prod', or 'mae_prod'."
+        )
+    return metric

tests/resources/options.yaml

@@ -1,15 +1,10 @@
 seed: 42
 
 architecture:
-  name: experimental.phace
-  # model:
-  #   radial_basis:
-  #     optimizable_lengthscales: true
+  name: experimental.soap_bpnn
   training:
-    batch_size: 8
-    num_epochs: 10
-    gradient_clipping: 10
-    log_interval: 1
+    batch_size: 2
+    num_epochs: 1
 
 training_set:
   systems: