Add displacements and momenta outputs

Author: frostedoyster

docs/src/outputs/displacements-and-momenta.rst

@@ -0,0 +1,101 @@
+.. _displacements-output:
+
+Displacements
+^^^^^^^^^^^^^
+
+Displacements are differences between atomic positions at two different times.
+They can be used to predict the next configuration in molecular dynamics
+(see, e.g., https://arxiv.org/pdf/2505.19350).
+
+In metatomic models, they are associated with the ``"displacements"``
+key in the model outputs, and must adhere to the following metadata schema:
+
+.. list-table:: Metadata for displacements
+  :widths: 2 3 7
+  :header-rows: 1
+
+  * - Metadata
+    - Names
+    - Description
+
+  * - keys
+    - ``"_"``
+    - the keys must have a single dimension named ``"_"``, with a single
+      entry set to ``0``. Displacements are always a
+      :py:class:`metatensor.torch.TensorMap` with a single block.
+
+  * - samples
+    - ``["system", "atom"]``
+    - the samples must be named ``["system", "atom"]``, since
+      displacements are always per-atom.
+
+      ``"system"`` must range from 0 to the number of systems given as an input
+      to the model. ``"atom"`` must range between 0 and the number of
+      atoms/particles in the corresponding system. If ``selected_atoms`` is
+      provided, then only the selected atoms for each system should be part of
+      the samples.
+
+  * - components
+    - ``"xyz"``
+    - displacements must have a single component dimension named
+      ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``.  The
+      displacements are always 3D vectors, and the order of the
+      components is x, y, z.
+
+  * - properties
+    - ``"displacements"``
+    - displacements must have a single property dimension named
+      ``"displacements"``, with a single entry set to ``0``.
+
+At the moment, displacements are not integrated into any simulation engines.
+
+.. _momenta-output:
+
+Momenta
+^^^^^^^
+
+The momentum of a particle is a vector defined as its mass times its velocity.
+Predictions of momenta can be used, for example, to predict a future step in molecular
+dynamics (see, e.g., https://arxiv.org/pdf/2505.19350).
+
+In metatomic models, they are associated with the ``"momenta"``
+key in the model outputs, and must adhere to the following metadata schema:
+
+.. list-table:: Metadata for momenta
+  :widths: 2 3 7
+  :header-rows: 1
+
+  * - Metadata
+    - Names
+    - Description
+
+  * - keys
+    - ``"_"``
+    - the keys must have a single dimension named ``"_"``, with a single
+      entry set to ``0``. Momenta are always a
+      :py:class:`metatensor.torch.TensorMap` with a single block.
+
+  * - samples
+    - ``["system", "atom"]``
+    - the samples must be named ``["system", "atom"]``, since
+      momenta are always per-atom.
+
+      ``"system"`` must range from 0 to the number of systems given as an input
+      to the model. ``"atom"`` must range between 0 and the number of
+      atoms/particles in the corresponding system. If ``selected_atoms`` is
+      provided, then only the selected atoms for each system should be part of
+      the samples.
+
+  * - components
+    - ``"xyz"``
+    - momenta must have a single component dimension named
+      ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``.  The
+      momenta are always 3D vectors, and the order of the
+      components is x, y, z.
+
+  * - properties
+    - ``"momenta"``
+    - momenta must have a single property dimension named
+      ``"momenta"``, with a single entry set to ``0``.
+
+At the moment, momenta are not integrated into any simulation engines.

docs/src/outputs/index.rst

@@ -20,6 +20,7 @@ section to these pages.
   energy
   features
   non_conservative
+  displacements-and-momenta
 
 
 Physical quantities
@@ -76,6 +77,22 @@ quantities, i.e. quantities with a well-defined physical meaning.
       Stress directly predicted by the model, not derived from the potential
       energy.
 
+    .. grid-item-card:: Displacements
+      :link: displacements-output
+      :link-type: ref
+
+      .. image:: /../static/images/displacements-output.png
+
+      Atomic displacements predicted by the model, to be used in ML-driven simulations.
+
+    .. grid-item-card:: Momenta
+      :link: momenta-output
+      :link-type: ref
+
+      .. image:: /../static/images/momenta-output.png
+
+      Atomic momenta predicted by the model, to be used in ML-driven simulations.
+
 
 Machine learning outputs
 ^^^^^^^^^^^^^^^^^^^^^^^^

docs/src/torch/reference/models/index.rst

@@ -34,7 +34,7 @@ In the mean time, you can create :py:class:`metatomic.torch.ModelOutput` with
 quantities that are not in this table. A warning will be issued and no unit
 conversion will be performed.
 
-When working with one of the quantity in this table, the unit you use must be
+When working with one of the quantities in this table, the unit you use must be
 one of the registered unit.
 
 +----------------+---------------------------------------------------------------------------------------------------+
@@ -48,3 +48,5 @@ one of the registered unit.
 +----------------+---------------------------------------------------------------------------------------------------+
 |   **pressure** | eV/Angstrom^3 (eV/A^3, eV/Angstrom^3)                                                             |
 +----------------+---------------------------------------------------------------------------------------------------+
+|   **momentum** | sqrt(eV*u)                                                                                        |
++----------------+---------------------------------------------------------------------------------------------------+

docs/static/images/displacements-output.png

@@ -141,7 +141,9 @@ std::unordered_set<std::string> KNOWN_OUTPUTS = {
     "energy_uncertainty",
     "features",
     "non_conservative_forces",
-    "non_conservative_stress"
+    "non_conservative_stress",
+    "displacements",
+    "momenta"
 };
 
 void ModelCapabilitiesHolder::set_outputs(torch::Dict<std::string, ModelOutput> outputs) {
@@ -1082,6 +1084,11 @@ static std::map<std::string, Quantity> KNOWN_QUANTITIES = {
         // alternative names
         {"eV/A^3", "eV/Angstrom^3"},
     }}},
+    {"momentum", Quantity{/* name */ "momentum", /* baseline */ "sqrt(eV*u)", {
+        {"sqrt(eV*u)", 1.0},
+    }, {
+        // alternative names
+    }}},
 };
 
 bool metatomic_torch::valid_quantity(const std::string& quantity) {

docs/static/images/momenta-output.png

@@ -109,7 +109,7 @@ TEST_CASE("Models metadata") {
         struct WarningHandler: public torch::WarningHandler {
             virtual ~WarningHandler() override = default;
             void process(const torch::Warning& warning) override {
-                CHECK(warning.msg() == "unknown quantity 'unknown', only [energy force length pressure] are supported");
+                CHECK(warning.msg() == "unknown quantity 'unknown', only [energy force length momentum pressure] are supported");
             }
         };
 

metatomic-torch/src/model.cpp

@@ -51,6 +51,10 @@ def _check_outputs(
             _check_non_conservative_forces(value, systems, request, selected_atoms)
         elif name == "non_conservative_stress":
             _check_non_conservative_stress(value, systems, request)
+        elif name == "displacements":
+            _check_displacements(value, systems, request)
+        elif name == "momenta":
+            _check_momenta(value, systems, request)
         else:
             # this is a non-standard output, there is nothing to check
             continue
@@ -263,6 +267,106 @@ def _check_non_conservative_stress(
         )
 
 
+def _check_displacements(
+    value: TensorMap,
+    systems: List[System],
+    request: ModelOutput,
+):
+    """
+    Check output metadata for displacements.
+    """
+    # Ensure the output contains a single block with the expected key
+    _validate_single_block("displacements", value)
+
+    # Check samples values from systems
+    _validate_atomic_samples(
+        "displacements", value, systems, request, selected_atoms=None
+    )
+
+    displacements_block = value.block_by_id(0)
+
+    # Check that the block has correct "Cartesian-form" components
+    if len(displacements_block.components) != 1:
+        raise ValueError(
+            "invalid components for 'displacements' output: "
+            f"expected one component, got {len(displacements_block.components)}"
+        )
+    expected_component = Labels(
+        "xyz", torch.tensor([[0], [1], [2]], device=value.device)
+    )
+    if displacements_block.components[0] != expected_component:
+        raise ValueError(
+            f"invalid components for 'displacements' output: "
+            f"expected {expected_component}, got {displacements_block.components[0]}"
+        )
+
+    expected_properties = Labels(
+        "displacements", torch.tensor([[0]], device=value.device)
+    )
+    message = "`Labels('displacements', [[0]])`"
+
+    if displacements_block.properties != expected_properties:
+        raise ValueError(
+            f"invalid properties for 'displacements' output: expected {message}, "
+            f"got {displacements_block.properties}"
+        )
+
+    # Should not have any gradients
+    if len(displacements_block.gradients_list()) > 0:
+        raise ValueError(
+            "invalid gradients for 'displacements' output: "
+            f"expected no gradients, found {displacements_block.gradients_list()}"
+        )
+
+
+def _check_momenta(
+    value: TensorMap,
+    systems: List[System],
+    request: ModelOutput,
+):
+    """
+    Check output metadata for momenta.
+    """
+    # Ensure the output contains a single block with the expected key
+    _validate_single_block("momenta", value)
+
+    # Check samples values from systems
+    _validate_atomic_samples("momenta", value, systems, request, selected_atoms=None)
+
+    momenta_block = value.block_by_id(0)
+
+    # Check that the block has correct "Cartesian-form" components
+    if len(momenta_block.components) != 1:
+        raise ValueError(
+            "invalid components for 'momenta' output: "
+            f"expected one component, got {len(momenta_block.components)}"
+        )
+    expected_component = Labels(
+        "xyz", torch.tensor([[0], [1], [2]], device=value.device)
+    )
+    if momenta_block.components[0] != expected_component:
+        raise ValueError(
+            f"invalid components for 'momenta' output: "
+            f"expected {expected_component}, got {momenta_block.components[0]}"
+        )
+
+    expected_properties = Labels("momenta", torch.tensor([[0]], device=value.device))
+    message = "`Labels('momenta', [[0]])`"
+
+    if momenta_block.properties != expected_properties:
+        raise ValueError(
+            f"invalid properties for 'momenta' output: expected {message}, "
+            f"got {momenta_block.properties}"
+        )
+
+    # Should not have any gradients
+    if len(momenta_block.gradients_list()) > 0:
+        raise ValueError(
+            "invalid gradients for 'momenta' output: "
+            f"expected no gradients, found {momenta_block.gradients_list()}"
+        )
+
+
 def _validate_atomic_samples(
     name: str,
     value: TensorMap,