Fix a bug with custom model inputs in ASE

The unit conversion was using quantity instead of name when
looking up the data

Author: Luthaf

python/metatomic_torch/metatomic/torch/ase_calculator.py

@@ -60,31 +60,32 @@
         "unit": "nm/fs",
     },
     "ase::initial_magmoms": {
-        "quantity": "",
+        "quantity": "magnetic_moment",
         "getter": ase.Atoms.get_initial_magnetic_moments,
         "unit": "",
     },
     "ase::magnetic_moment": {
-        "quantity": "",
+        "quantity": "magnetic_moment",
         "getter": ase.Atoms.get_magnetic_moment,
         "unit": "",
     },
     "ase::magnetic_moments": {
-        "quantity": "",
+        "quantity": "magnetic_moment",
         "getter": ase.Atoms.get_magnetic_moments,
         "unit": "",
     },
     "ase::initial_charges": {
-        "quantity": "",
+        "quantity": "charge",
         "getter": ase.Atoms.get_initial_charges,
         "unit": "",
     },
     "ase::charges": {
-        "quantity": "",
+        "quantity": "charge",
         "getter": ase.Atoms.get_charges,
         "unit": "",
     },
     "ase::dipole_moment": {
+        "quantity": "dipole_moment",
         "getter": ase.Atoms.get_dipole_moment,
         "unit": "",
     },
@@ -972,18 +973,13 @@ def _get_ase_input(
     dtype: torch.dtype,
     device: torch.device,
 ) -> "TensorMap":
-    if name in ARRAY_QUANTITIES:
-        infos = ARRAY_QUANTITIES[name]
-        if infos["quantity"] != option.quantity:
-            raise ValueError(
-                f"The model requested '{name}' with quantity '{option.quantity}', "
-                f"but the quantity is '{infos['quantity']}' in `ase`."
-            )
-    else:
+    if name not in ARRAY_QUANTITIES:
         raise ValueError(
             f"The model requested '{name}', which is not available in `ase`."
         )
 
+    infos = ARRAY_QUANTITIES[name]
+
     values = infos["getter"](atoms)
     if values.shape[0] != len(atoms):
         raise NotImplementedError(
@@ -994,7 +990,7 @@ def _get_ase_input(
     # for metatensor
     values = torch.tensor(values[..., None])
 
-    tblock = TensorBlock(
+    block = TensorBlock(
         values,
         samples=Labels(
             ["system", "atom"],
@@ -1005,21 +1001,16 @@ def _get_ase_input(
         components=[Labels(["xyz"], torch.arange(values.shape[1]).reshape(-1, 1))]
         if values.shape[1] != 1
         else [],
-        properties=Labels(
-            [
-                name if "::" not in name else name.split("::")[1],
-            ],
-            torch.tensor([[0]]),
-        ),
-    )
-    tmap = TensorMap(
-        Labels(["_"], torch.tensor([[0]])),
-        [tblock],
+        properties=Labels([option.quantity], torch.tensor([[0]])),
     )
-    tmap.set_info("quantity", option.quantity)
-    tmap.set_info("unit", option.unit)
-    tmap.to(dtype=dtype, device=device)
-    return tmap
+
+    tensor = TensorMap(Labels(["_"], torch.tensor([[0]])), [block])
+
+    tensor.set_info("quantity", infos["quantity"])
+    tensor.set_info("unit", infos["unit"])
+
+    tensor.to(dtype=dtype, device=device)
+    return tensor
 
 
 def _ase_to_torch_data(atoms, dtype, device):

python/metatomic_torch/metatomic/torch/model.py

@@ -477,34 +477,12 @@ def forward(
 
         # convert systems from engine to model units
         with record_function("AtomisticModel::convert_units_input"):
-            if self._capabilities.length_unit != options.length_unit:
-                conversion = unit_conversion_factor(
-                    quantity="length",
-                    from_unit=options.length_unit,
-                    to_unit=self._capabilities.length_unit,
-                )
-
-                systems = _convert_systems_units(
-                    systems,
-                    conversion,
-                    model_length_unit=self._capabilities.length_unit,
-                    system_length_unit=options.length_unit,
-                )
-
-            for name, option in self._requested_inputs.items():
-                system_unit = str(
-                    systems[0].get_data(name).get_info("unit")
-                )  # For torchscript
-                to_unit = option.unit
-                conversion = unit_conversion_factor(
-                    quantity=option.quantity,
-                    from_unit=system_unit,
-                    to_unit=to_unit,
-                )
-
-                _convert_systems_input_units(
-                    systems, option.quantity, conversion, to_unit
-                )
+            systems = _convert_systems_units(
+                systems,
+                model_length_unit=self._capabilities.length_unit,
+                system_length_unit=options.length_unit,
+                requested_inputs=self._requested_inputs,
+            )
 
         # run the actual calculations
         with record_function("Model::forward"):
@@ -948,12 +926,19 @@ def _check_inputs(
 
 def _convert_systems_units(
     systems: List[System],
-    conversion: float,
     model_length_unit: str,
     system_length_unit: str,
+    requested_inputs: Dict[str, ModelOutput],
 ) -> List[System]:
-    if conversion == 1.0:
-        return systems
+    if model_length_unit == "" or system_length_unit == "":
+        # no conversion for positions/cell/NL
+        conversion = 1.0
+    else:
+        conversion = unit_conversion_factor(
+            quantity="length",
+            from_unit=system_length_unit,
+            to_unit=model_length_unit,
+        )
 
     new_systems: List[System] = []
     for system in systems:
@@ -978,41 +963,58 @@ def _convert_systems_units(
             )
 
         known_data = system.known_data()
-        if len(known_data) != 0:
-            warnings.warn(
-                "the model requires a different length unit "
-                f"({model_length_unit}) than the system ({system_length_unit}), "
-                f"but we don't know how to convert custom data ({known_data}) "
-                "accordingly",
-                stacklevel=2,
-            )
+        for name in known_data:
+            if name not in requested_inputs:
+                # not a requested input, just copy as is
+                new_system.add_data(name, system.get_data(name))
 
-        for data in known_data:
-            new_system.add_data(data, system.get_data(data))
+            else:
+                requested = requested_inputs[name]
+                tensor = system.get_data(name)
+                unit = tensor.get_info("unit")
+
+                if requested.quantity != "" and unit is not None:
+                    conversion = unit_conversion_factor(
+                        quantity=requested.quantity,
+                        from_unit=unit,
+                        to_unit=requested.unit,
+                    )
+                else:
+                    conversion = 1.0
+
+                new_blocks: List[TensorBlock] = []
+                for block in tensor.blocks():
+                    new_values = conversion * block.values
+                    new_block = TensorBlock(
+                        values=new_values,
+                        samples=block.samples,
+                        components=block.components,
+                        properties=block.properties,
+                    )
 
-        new_systems.append(new_system)
+                    for parameter, gradient in block.gradients():
+                        if len(gradient.gradients_list()) != 0:
+                            raise NotImplementedError(
+                                "nested gradients are not supported"
+                            )
+
+                        new_gradient = TensorBlock(
+                            values=conversion * gradient.values,
+                            samples=gradient.samples,
+                            components=gradient.components,
+                            properties=gradient.properties,
+                        )
+                        new_block.add_gradient(parameter, new_gradient)
+                    new_blocks.append(new_block)
 
-    return new_systems
+                new_tensor = TensorMap(
+                    keys=tensor.keys,
+                    blocks=new_blocks,
+                )
+                new_tensor.set_info("unit", requested.unit)
+                new_tensor.set_info("quantity", requested.quantity)
+                new_system.add_data(name, new_tensor)
 
+        new_systems.append(new_system)
 
-def _convert_systems_input_units(
-    systems: List[System], quantity: str, conversion: float, to_unit: str
-) -> None:
-    if conversion != 1.0:
-        for system in systems:
-            tensor = system.get_data(quantity)
-            tblock = tensor.block()
-            new_tensor = TensorMap(
-                Labels("_", torch.tensor([[0]])),
-                [
-                    TensorBlock(
-                        values=conversion * tblock.values,
-                        samples=tblock.samples,
-                        components=tblock.components,
-                        properties=tblock.properties,
-                    )
-                ],
-            )
-            new_tensor.set_info("unit", to_unit)
-            new_tensor.set_info("quantity", quantity)
-            system.add_data(quantity, new_tensor, override=True)
+    return new_systems

python/metatomic_torch/tests/ase_calculator.py

@@ -837,7 +837,7 @@ 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),
-        "ase::initial_charges": ModelOutput(quantity="", unit="", per_atom=True),
+        "ase::initial_charges": ModelOutput(quantity="charge", unit="", per_atom=True),
     }
     outputs = {("extra::" + prop): inputs[prop] for prop in inputs}
     capabilities = ModelCapabilities(
@@ -855,17 +855,16 @@ def test_additional_input(atoms):
     atoms.set_initial_charges([0.0] * len(atoms))
     calculator = MetatomicCalculator(model, check_consistency=True)
     results = calculator.run_model(atoms, outputs)
-    for k, v in results.items():
-        head, prop = k.split("::", maxsplit=1)
+    for name, tensor in results.items():
+        head, name = name.split("::", maxsplit=1)
         assert head == "extra"
-        assert prop in inputs
-        assert len(v.keys.names) == 1
-        assert v.get_info("quantity") == inputs[prop].quantity
-        values = v[0].values.numpy()
-        shape = values.shape
-        assert shape[0] == len(atoms), f"Expected {len(atoms)} values, got {shape[0]}"
-        assert np.allclose(
-            values,
-            ARRAY_QUANTITIES[prop]["getter"](atoms).reshape(shape)
-            * (10 if prop == "velocity" else 1),  # ase velocity is in nm/fs, not A/fs
-        )
+        assert name in inputs
+
+        assert tensor.get_info("quantity") == inputs[name].quantity
+        values = tensor[0].values.numpy()
+
+        expected = ARRAY_QUANTITIES[name]["getter"](atoms).reshape(values.shape)
+        if name == "velocities":
+            expected *= 10.0  # ase velocity is in nm/fs
+
+        assert np.allclose(values, expected)