Add heat flux as a standard output, also add momentum and velocity as outputs of i-pi

docs/src/outputs/heat_flux.rst

@@ -0,0 +1,61 @@
+.. _heat-flux-output:
+
+Heat Flux
+^^^^^^^^^
+
+Heat flux is associated with the ``"heat_flux"`` or ``"heat_flux/<variant>"``
+name (see :ref:`output-variants`), and must have the following metadata:
+
+.. list-table:: Metadata for heat fluxes
+  :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``. Heat fluxes are always a
+      :py:class:`metatensor.torch.TensorMap` with a single block.
+
+  * - samples
+    - ``["system"]``
+    - the samples must be named ``["system"]``, since heat fluxes are always
+      per-system.
+
+      ``"system"`` must range from 0 to the number of systems given as an input
+      to the model.
+
+  * - components
+    - ``"xyz"``
+    - heat fluxes must have a single component dimension named
+      ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``.  The
+      heat fluxes are always 3D vectors, and the order of the
+      components is x, y, z.
+
+  * - properties
+    - ``"heat_flux"``
+    - heat fluxes must have a single property dimension named
+      ``"heat_flux"``, with a single entry set to ``0``.
+
+The following simulation engine can use the ``"heat_flux"`` output.
+
+.. grid:: 1 3 3 3
+
+  .. grid-item-card::
+    :text-align: center
+    :padding: 1
+    :link: engine-ase
+    :link-type: ref
+
+    |ase-logo|
+
+  .. grid-item-card::
+    :text-align: center
+    :padding: 1
+    :link: engine-ipi
+    :link-type: ref
+
+    |ipi-logo|

docs/src/outputs/index.rst

@@ -25,6 +25,7 @@ schema they need and add a new section to these pages.
   momenta
   velocities
   charges
+  heat_flux
   features
   variants
 
@@ -131,6 +132,15 @@ quantities, i.e. quantities with a well-defined physical meaning.
 
       Atomic charges, e.g. formal or partial charges on atoms
 
+    .. grid-item-card:: Heat flux
+      :link: heat-flux-output
+      :link-type: ref
+
+      .. image:: /../static/images/heat-flux-output.png
+
+      Heat flux, i.e. the amount of energy transferred per unit time, i.e.
+      :math:`\sum_i E_i \times \vec v_i`
+
 Machine learning quantities
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

docs/src/outputs/momenta.rst

@@ -59,3 +59,11 @@ The following simulation engine can provide ``"momenta"`` as inputs to the model
     :link-type: ref
 
     |ase-logo|
+
+  .. grid-item-card::
+    :text-align: center
+    :padding: 1
+    :link: engine-ipi
+    :link-type: ref
+
+    |ipi-logo|

docs/src/outputs/velocities.rst

@@ -55,3 +55,11 @@ The following simulation engine can provide ``"velocities"`` as inputs to the mo
     :link-type: ref
 
     |ase-logo|
+
+  .. grid-item-card::
+    :text-align: center
+    :padding: 1
+    :link: engine-ipi
+    :link-type: ref
+
+    |ipi-logo|

metatomic-torch/src/misc.cpp

@@ -427,6 +427,7 @@ inline std::unordered_set<std::string> KNOWN_INPUTS_OUTPUTS = {
     "velocities",
     "masses",
     "charges",
+    "heat_flux",
 };
 
 std::tuple<bool, std::string, std::string> details::validate_name_and_check_variant(

metatomic-torch/src/outputs.cpp

@@ -584,6 +584,44 @@ static void check_charges(
     validate_no_gradients("charges", charges_block);
 }
 
+/// Check output metadata for heat flux.
+static void check_heat_flux(
+    const TensorMap& value,
+    const std::vector<System>& systems,
+    const ModelOutput& request
+) {
+    // Ensure the output contains a single block with the expected key
+    validate_single_block("heat_flux", value);
+
+    // Check samples values from systems
+    if (request->per_atom) {
+        C10_THROW_ERROR(ValueError,
+            "invalid 'heat_flux' output: heat flux cannot be per-atom, but the request "
+            "indicates `per_atom=True`"
+        );
+    }
+    validate_atomic_samples("heat_flux", value, systems, request, torch::nullopt);
+
+    auto tensor_options = torch::TensorOptions().device(value->device());
+    auto heat_flux_block = TensorMapHolder::block_by_id(value, 0);
+    std::vector<Labels> expected_component {
+        torch::make_intrusive<LabelsHolder>(
+            "xyz",
+            torch::tensor({{0}, {1}, {2}}, tensor_options)
+        )
+    };
+    validate_components("heat_flux", heat_flux_block->components(), expected_component);
+
+    auto expected_properties = torch::make_intrusive<LabelsHolder>(
+        "heat_flux",
+        torch::tensor({{0}}, tensor_options)
+    );
+    validate_properties("heat_flux", heat_flux_block, expected_properties);
+
+    // Should not have any gradients
+    validate_no_gradients("heat_flux", heat_flux_block);
+}
+
 void metatomic_torch::check_outputs(
     const std::vector<System>& systems,
     const c10::Dict<std::string, ModelOutput>& requested,
@@ -654,6 +692,8 @@ void metatomic_torch::check_outputs(
             check_velocities(value, systems, request);
         } else if (base == "charges") {
             check_charges(value, systems, request);
+        } else if (base == "heat_flux") {
+            check_heat_flux(value, systems, request);
         } else if (name.find("::") != std::string::npos) {
             // this is a non-standard output, there is nothing to check
         } else {

metatomic-torch/src/units.cpp

@@ -586,14 +586,15 @@ static UnitValue parse_unit_expression(const std::string& unit) {
 
 
 static const auto QUANTITY_DIMS = std::unordered_map<std::string, Dimension>{
-    {"length",   DIM_LENGTH},
-    {"energy",   DIM_ENERGY},
-    {"force",    {{1, -2, 1, 0, 0}}},   // energy/length
-    {"pressure", {{-1, -2, 1, 0, 0}}},  // energy/length^3
-    {"momentum", {{1, -1, 1, 0, 0}}},   // mass*length/time
-    {"mass",     DIM_MASS},
-    {"velocity", {{1, -1, 0, 0, 0}}},   // length/time
-    {"charge",   DIM_CHARGE},
+    {"length",    DIM_LENGTH},
+    {"energy",    DIM_ENERGY},
+    {"force",     {{1, -2, 1, 0, 0}}},   // energy/length
+    {"pressure",  {{-1, -2, 1, 0, 0}}},  // energy/length^3
+    {"momentum",  {{1, -1, 1, 0, 0}}},   // mass*length/time
+    {"mass",      DIM_MASS},
+    {"velocity",  {{1, -1, 0, 0, 0}}},   // length/time
+    {"charge",    DIM_CHARGE},
+    {"heat_flux", {{3, -3, 1, 0, 0}}}, // energy*velocity
 };
 
 

metatomic-torch/tests/models.cpp

@@ -110,7 +110,7 @@ TEST_CASE("Models metadata") {
             virtual ~WarningHandler() override = default;
             void process(const torch::Warning& warning) override {
                 auto expected = std::string(
-                    "unknown quantity 'unknown', only [charge energy force "
+                    "unknown quantity 'unknown', only [charge energy force heat_flux "
                     "length mass momentum pressure velocity] are supported"
                 );
                 CHECK(warning.msg() == expected);