Correctly inherit properties from targets in SOAP-BPNN

metatensor · Nov 20, 2024 · e158719 · e158719
1 parent 0c207fc
commit e158719
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 13 deletions.
diff --git a/src/metatrain/experimental/soap_bpnn/model.py b/src/metatrain/experimental/soap_bpnn/model.py
@@ -432,14 +432,10 @@ def _add_output(self, target_name: str, target: TargetInfo) -> None:
                     values=torch.tensor(self.atomic_types).reshape(-1, 1),
                 ),
                 in_features=self.n_inputs_last_layer,
-                out_features=1,
+                out_features=len(target.layout.block().properties.values),
                 bias=False,
                 out_properties=[
-                    Labels(
-                        names=["energy"],
-                        values=torch.tensor([[0]]),
-                    )
-                    for _ in self.atomic_types
+                    target.layout.block().properties for _ in self.atomic_types
                 ],
             )
         else:
@@ -456,14 +452,10 @@ def _add_output(self, target_name: str, target: TargetInfo) -> None:
                     ),
                 ),
                 in_features=self.n_inputs_last_layer,
-                out_features=1,
+                out_features=len(target.layout.block().properties.values),
                 bias=False,
                 out_properties=[
-                    Labels(
-                        names=["properties"],
-                        values=torch.tensor([[0]]),
-                    )
-                    for _ in self.atomic_types
+                    target.layout.block().properties for _ in self.atomic_types
                 ],
             )
 

diff --git a/src/metatrain/utils/additive/composition.py b/src/metatrain/utils/additive/composition.py
@@ -44,8 +44,9 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo):
                 per_atom=True,
             )
             for key, target_info in dataset_info.targets.items()
-            if target_info.is_scalar
+            if target_info.is_scalar and len(target_info.layout.block().properties) == 1
             # important: only scalars can have composition contributions
+            # for now, we also require that only one property is present
         }
 
         n_types = len(self.atomic_types)