Add a heat flux wrapper

docs/src/outputs/heat_fluxes.rst

@@ -0,0 +1,65 @@
+.. _heat-fluxes-output:
+
+Heat Fluxes
+^^^^^^^^^^^
+
+Heat fluxes are 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 not 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"``
+    - 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

@@ -131,6 +131,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 fluxes
+      :link: heat-fluxes-output
+      :link-type: ref
+
+      .. image:: /../static/images/charges-output.png
+
+       Heat fluxes, 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,12 @@ 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/model.cpp

@@ -1157,6 +1157,12 @@ static std::map<std::string, Quantity> KNOWN_QUANTITIES = {
         // alternative names
         {"C", "Coulomb"},
     }}},
+    {"heat_flux", Quantity{/* name */ "heat_flux", /* baseline */ "eV*Angstrom/fs", {
+        {"eV*Angstrom/fs", 1.0},
+    }, {
+        // alternative names
+        {"eV*A/fs", "eV*Angstrom/fs"},
+    }}}
 };
 
 bool metatomic_torch::valid_quantity(const std::string& quantity) {

metatomic-torch/src/outputs.cpp

@@ -584,6 +584,38 @@ 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
+    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 +686,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/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);

python/metatomic_torch/metatomic/torch/ase_calculator.py

@@ -935,8 +935,7 @@ def _get_ase_input(
     tensor.set_info("quantity", infos["quantity"])
     tensor.set_info("unit", infos["unit"])
 
-    tensor = tensor.to(dtype=dtype, device=device)
-    return tensor
+    return tensor.to(dtype=dtype, device=device)
 
 
 def _ase_to_torch_data(atoms, dtype, device):