Split error message for per-atom targets

metatensor · Nov 25, 2024 · d002543 · d002543
1 parent c8c53e1
commit d002543
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/src/metatrain/experimental/alchemical_model/model.py b/src/metatrain/experimental/alchemical_model/model.py
@@ -39,13 +39,17 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         if not (
             target.is_scalar
             and target.quantity == "energy"
-            and "atom" not in target.layout.block(0).samples.names
             and len(target.layout.block(0).properties) == 1
         ):
             raise ValueError(
                 "The Alchemical Model only supports total-energy-like outputs, "
                 f"but a {target.quantity} was provided"
             )
+        if target.per_atom:
+            raise ValueError(
+                "Alchemical Model only supports per-structure outputs, "
+                "but a per-atom output was provided"
+            )
 
         self.outputs = {
             key: ModelOutput(

diff --git a/src/metatrain/experimental/gap/model.py b/src/metatrain/experimental/gap/model.py
@@ -40,7 +40,6 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
             if not (
                 target.is_scalar
                 and target.quantity == "energy"
-                and "atom" not in target.layout.block(0).samples.names
                 and len(target.layout.block(0).properties) == 1
             ):
                 raise ValueError(

diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py
@@ -41,13 +41,17 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         if not (
             target.is_scalar
             and target.quantity == "energy"
-            and "atom" not in target.layout.block(0).samples.names
             and len(target.layout.block(0).properties) == 1
         ):
             raise ValueError(
                 "PET only supports total-energy-like outputs, "
                 f"but a {target.quantity} was provided"
             )
+        if target.per_atom:
+            raise ValueError(
+                "PET only supports per-structure outputs, "
+                "but a per-atom output was provided"
+            )
 
         model_hypers["D_OUTPUT"] = 1
         model_hypers["TARGET_TYPE"] = "atomic"