Check the requested inputs with _check_outputs

Author: GardevoirX

python/metatomic_torch/metatomic/torch/ase_calculator.py

@@ -44,20 +44,17 @@
     "float64": torch.float64,
 }
 
-ARRAY_PROPERTIES = {
-    "momenta": {
+ARRAY_QUANTITIES = {
+    "momentum": {
         "getter": lambda atoms: atoms.get_momenta(),
-        "quantity": "momentum",
         "unit": "(eV*u)^(1/2)",
     },
-    "masses": {
+    "mass": {
         "getter": lambda atoms: atoms.get_masses(),
-        "quantity": "mass",
         "unit": "u",
     },
-    "velocities": {
+    "velocity": {
         "getter": lambda atoms: atoms.get_velocities(),
-        "quantity": "velocity",
         "unit": "nm/fs",
     },
     "initial_magmoms": {},
@@ -368,7 +365,7 @@ def run_model(
             # Get the additional inputs requested by the model
             for quantity, option in self._model.requested_inputs().items():
                 input_tensormap = _get_ase_input(
-                    atoms, quantity, option, dtype=self._dtype, device=self._device
+                    atoms, option, dtype=self._dtype, device=self._device
                 )
                 system.add_data(quantity, input_tensormap)
             systems.append(system)
@@ -519,7 +516,7 @@ def calculate(
                 system.add_neighbor_list(options, neighbors)
             for quantity, option in self._model.requested_inputs().items():
                 input_tensormap = _get_ase_input(
-                    atoms, quantity, option, dtype=self._dtype, device=self._device
+                    atoms, option, dtype=self._dtype, device=self._device
                 )
                 system.add_data(quantity, input_tensormap)
 
@@ -944,21 +941,20 @@ def _compute_ase_neighbors(atoms, options, dtype, device):
 
 def _get_ase_input(
     atoms: ase.Atoms,
-    quantity: str,
     option: ModelOutput,
     dtype: torch.dtype,
     device: torch.device,
 ) -> "TensorMap":
-    if quantity in ARRAY_PROPERTIES:
-        if len(ARRAY_PROPERTIES[quantity]) == 0:
+    if option.quantity in ARRAY_QUANTITIES:
+        if len(ARRAY_QUANTITIES[option.quantity]) == 0:
             raise NotImplementedError(
-                f"Though the property {quantity} is available in `ase`, it is "
+                f"Though the quantity {option.quantity} is available in `ase`, it is "
                 "currently not supported by metatomic."
             )
-        infos = ARRAY_PROPERTIES[quantity]
+        infos = ARRAY_QUANTITIES[option.quantity]
     else:
         raise ValueError(
-            f"The model requested '{quantity}', which is not available in `ase`."
+            f"The model requested '{option.quantity}', which is not available in `ase`."
         )
 
     values = infos["getter"](atoms)

python/metatomic_torch/metatomic/torch/model.py

@@ -438,6 +438,18 @@ def forward(
                     options=options,
                     expected_dtype=self._model_dtype,
                 )
+                # check the requested inputs stored in the `systems`
+                for system in systems:
+                    system_inputs: Dict[str, TensorMap] = {}
+                    for name in system.known_data():
+                        system_inputs[name] = system.get_data(name)
+                    _check_outputs(
+                        systems=[system],
+                        requested=self._requested_inputs,
+                        selected_atoms=options.selected_atoms,
+                        outputs=system_inputs,
+                        model_dtype=self._capabilities.dtype,
+                    )
 
         with record_function("AtomisticModel::check_atomic_types"):
             # always (i.e. even if check_consistency=False) check that the atomic types
@@ -898,6 +910,8 @@ def _check_inputs(
                 )
 
         # Check additional inputs
+        # Might be problematic, this requires that only requested inputs are stored as
+        # the data pf the system
         known_additional_inputs = system.known_data()
         for request in requested_inputs:
             found = False

python/metatomic_torch/tests/ase_calculator.py

@@ -26,7 +26,7 @@
     System,
 )
 from metatomic.torch.ase_calculator import (
-    ARRAY_PROPERTIES,
+    ARRAY_QUANTITIES,
     MetatomicCalculator,
     _compute_ase_neighbors,
     _full_3x3_to_voigt_6_stress,
@@ -835,8 +835,8 @@ def forward(
 
 def test_additional_input(atoms):
     inputs = {
-        "masses": ModelOutput(quantity="mass", unit="u", per_atom=True),
-        "velocities": ModelOutput(quantity="velocity", unit="A/fs", per_atom=True),
+        "mass": ModelOutput(quantity="mass", unit="u", per_atom=True),
+        "velocity": ModelOutput(quantity="velocity", unit="A/fs", per_atom=True),
     }
     outputs = {("extra::" + prop): inputs[prop] for prop in inputs}
     capabilities = ModelCapabilities(
@@ -862,6 +862,6 @@ def test_additional_input(atoms):
         shape = v[0].values.numpy().shape
         assert np.allclose(
             v[0].values.numpy(),
-            ARRAY_PROPERTIES[prop]["getter"](atoms).reshape(shape)
-            * (10 if prop == "velocities" else 1),
+            ARRAY_QUANTITIES[prop]["getter"](atoms).reshape(shape)
+            * (10 if prop == "velocity" else 1),
         )