Add a heat flux wrapper

python/metatomic_ase/tests/calculator.py

@@ -1,4 +1,5 @@
 import os
+import re
 import subprocess
 import sys
 from typing import Dict, List, Optional
@@ -774,7 +775,7 @@ def forward(
 
 
 class AdditionalInputModel(torch.nn.Module):
-    def __init__(self, inputs):
+    def __init__(self, inputs: Dict[str, ModelOutput]):
         super().__init__()
         self._requested_inputs = inputs
 
@@ -793,6 +794,32 @@ def forward(
         }
 
 
+class SimpleWrapperModel(torch.nn.Module):
+    def __init__(self, model: AtomisticModel, inputs: Dict[str, ModelOutput]):
+        super().__init__()
+        self._model = model.module
+        self._requested_inputs = inputs
+        self._capabilities = model.capabilities()
+
+    def requested_inputs(self) -> Dict[str, ModelOutput]:
+        return self._requested_inputs
+
+    def forward(
+        self,
+        systems: List[System],
+        outputs: Dict[str, ModelOutput],
+        selected_atoms: Optional[Labels] = None,
+    ) -> Dict[str, TensorMap]:
+        results = self._model(systems, outputs, selected_atoms)
+        results.update(
+            {
+                ("extra::" + input): systems[0].get_data(input)
+                for input in self._requested_inputs
+            }
+        )
+        return results
+
+
 def test_additional_input(atoms):
     inputs = {
         "masses": ModelOutput(quantity="mass", unit="u", per_atom=True),
@@ -833,6 +860,41 @@ def test_additional_input(atoms):
         assert np.allclose(values, expected)
 
 
+def test_inputs_different_units():
+    inputs = {
+        "masses": ModelOutput(quantity="mass", unit="u", per_atom=True),
+        "velocities": ModelOutput(quantity="velocity", unit="A/fs", per_atom=True),
+        "charges": ModelOutput(quantity="charge", unit="e", per_atom=True),
+        "ase::initial_charges": ModelOutput(quantity="charge", unit="e", per_atom=True),
+    }
+    outputs = {("extra::" + n): inputs[n] for n in inputs}
+    capabilities = ModelCapabilities(
+        outputs=outputs,
+        atomic_types=[28],
+        interaction_range=0.0,
+        supported_devices=["cpu"],
+        dtype="float64",
+    )
+
+    model = AtomisticModel(
+        AdditionalInputModel(inputs).eval(), ModelMetadata(), capabilities
+    )
+
+    inputs_wrapper = {
+        "masses": ModelOutput(quantity="mass", unit="kg", per_atom=True),
+    }
+    wrapper = SimpleWrapperModel(model, inputs_wrapper)
+    with pytest.raises(
+        NotImplementedError,
+        match=re.escape(
+            "Different units for the same quantity `mass` is not supported. "
+            "Requested by 'SimpleWrapperModel._model' (unit='u') and "
+            "'SimpleWrapperModel' (unit='kg')."
+        ),
+    ):
+        AtomisticModel(wrapper.eval(), ModelMetadata(), capabilities)
+
+
 @pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE)
 def test_mixed_pbc(model, device, dtype):
     """Test that the calculator works on a mixed-PBC system"""

python/metatomic_ase/tests/heat_flux.py

@@ -0,0 +1,87 @@
+import numpy as np
+import pytest
+import torch
+from ase import Atoms
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+
+import metatomic_lj_test
+from metatomic.torch import ModelOutput
+from metatomic.torch.heat_flux import (
+    HeatFlux,
+)
+from metatomic_ase import MetatomicCalculator
+
+
+@pytest.fixture
+def model():
+    return metatomic_lj_test.lennard_jones_model(
+        atomic_type=18,
+        cutoff=7.0,
+        sigma=3.405,
+        epsilon=0.01032,
+        length_unit="Angstrom",
+        energy_unit="eV",
+        with_extension=False,
+    )
+
+
+@pytest.fixture
+def model_in_kcal_per_mol():
+    return metatomic_lj_test.lennard_jones_model(
+        atomic_type=18,
+        cutoff=7.0,
+        sigma=3.405,
+        epsilon=0.2380,
+        length_unit="Angstrom",
+        energy_unit="kcal/mol",
+        with_extension=False,
+    )
+
+
+@pytest.fixture
+def atoms(request):
+    if hasattr(request, "param") and request.param == "atoms_triclinic":
+        cell = np.array([[6.0, 3.0, 1.0], [2.0, 6.0, 0.0], [0.0, 0.0, 6.0]])
+        positions = np.array([[0.0, 0.0, 0.0]])
+    else:
+        cell = np.array([[6.0, 0.0, 0.0], [0.0, 6.0, 0.0], [0.0, 0.0, 6.0]])
+        positions = np.array([[3.0, 3.0, 3.0]])
+    atoms = Atoms("Ar", scaled_positions=positions, cell=cell, pbc=True).repeat(
+        (2, 2, 2)
+    )
+    MaxwellBoltzmannDistribution(
+        atoms, temperature_K=300, rng=np.random.default_rng(42)
+    )
+    return atoms
+
+
+@pytest.mark.parametrize("use_script", [True, False])
+@pytest.mark.parametrize(
+    "atoms, expected",
+    [
+        ("atoms", [[8.8238e-05], [-2.5559e-04], [-2.0570e-04]]),
+    ],
+    indirect=["atoms"],
+)
+def test_wrap(model, atoms, expected, use_script):
+    wrapped_model = HeatFlux.wrap(model, scripting=use_script)
+    calc = MetatomicCalculator(
+        wrapped_model,
+        device="cpu",
+        additional_outputs={
+            "heat_flux": ModelOutput(
+                quantity="heat_flux",
+                unit="eV*A/fs",
+                explicit_gradients=[],
+                per_atom=False,
+            )
+        },
+        check_consistency=True,
+    )
+    atoms.calc = calc
+    atoms.get_potential_energy()
+    results = atoms.calc.additional_outputs["heat_flux"].block().values
+    assert torch.allclose(
+        results,
+        torch.tensor(expected, dtype=results.dtype),
+    )