Add pick_output function

Author: PicoCentauri

metatomic-torch/include/metatomic/torch/misc.hpp

@@ -7,6 +7,7 @@
 #include <vector>
 #include <torch/torch.h>
 
+#include "metatomic/torch/model.hpp"
 #include "metatomic/torch/system.hpp"
 
 #include <torch/types.h>
@@ -26,6 +27,14 @@ METATOMIC_TORCH_EXPORT std::string pick_device(
 	torch::optional<std::string> desired_device = torch::nullopt
 );
 
+/// Pick the output for the given ``requested_output`` from the availabilities of the
+/// model's ``outputs``, according to the optional ``desired_variant``.
+METATOMIC_TORCH_EXPORT std::string pick_output(
+	std::string requested_output,
+	torch::Dict<std::string, ModelOutput> outputs,
+	torch::optional<std::string> desired_variant = torch::nullopt
+);
+
 // ===== File-based =====
 void   save(const std::string& path, const System& system);
 System load_system(const std::string& path);
@@ -46,7 +55,7 @@ inline System        load_system_buffer(const torch::Tensor& data) {
       throw std::runtime_error("System pickle: expected 1D torch.uint8 buffer");
   }
   const uint8_t* ptr = t.data_ptr<uint8_t>();
-  const size_t n = static_cast<size_t>(t.numel());
+  const auto n = static_cast<size_t>(t.numel());
   return load_system_buffer(ptr, n);
 }
 

metatomic-torch/include/metatomic/torch/model.hpp

@@ -39,19 +39,19 @@ bool valid_quantity(const std::string& quantity);
 void validate_unit(const std::string& quantity, const std::string& unit);
 
 
-/// Metadata of one of the quantity a model can compute
+/// Information about one of the quantity a model can compute
 class METATOMIC_TORCH_EXPORT ModelOutputHolder: public torch::CustomClassHolder {
 public:
     ModelOutputHolder() = default;
 
     /// Initialize `ModelOutput` with the given data
     ModelOutputHolder(
-        std::string description_,
         std::string quantity,
         std::string unit,
         bool per_atom_,
-        std::vector<std::string> explicit_gradients_
-    ):  
+        std::vector<std::string> explicit_gradients_,
+        std::string description_
+    ):
         description(std::move(description_)),
         per_atom(per_atom_),
         explicit_gradients(std::move(explicit_gradients_))

metatomic-torch/src/misc.cpp

@@ -1,8 +1,10 @@
 #include <torch/torch.h>
 
+#include "metatomic/torch/model.hpp"
 #include "metatomic/torch/version.h"
 #include "metatomic/torch/misc.hpp"
 
+#include <algorithm>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -75,6 +77,66 @@ std::string pick_device(
     return selected_device;
 }
 
+std::string pick_output(
+    std::string requested_output,
+    torch::Dict<std::string, ModelOutput> outputs,
+    torch::optional<std::string> desired_variant
+) {
+    std::vector<std::string> matching_keys;
+    bool has_exact = false;
+
+    for (const auto& output: outputs) {
+        const auto& key = output.key();
+
+        // match either exact name or "name/variant"
+        if (key == requested_output
+            || (key.size() > requested_output.size()
+                && key.compare(0, requested_output.size(), requested_output) == 0
+                && key[requested_output.size()] == '/')) {
+            matching_keys.emplace_back(key);
+
+            if (key == requested_output) {
+                has_exact = true;
+            }
+        }
+    }
+
+    if (matching_keys.empty()) {
+        C10_THROW_ERROR(ValueError,
+            "output '" + requested_output + "' not found in outputs"
+        );
+    }
+
+    if (desired_variant != torch::nullopt) {
+        const auto& output = requested_output + "/" + desired_variant.value();
+        auto it = std::find(matching_keys.begin(), matching_keys.end(), output);
+        if (it != matching_keys.end()) {
+            return *it;
+        }
+        C10_THROW_ERROR(ValueError,
+            "variant '" + desired_variant.value() + "' for output '" + requested_output +
+            "' not found in outputs"
+        );
+    } else if (has_exact) {
+        return requested_output;
+    } else {
+        std::ostringstream oss;
+        oss << "output '" << requested_output << "' has no default variant and no `desired_variant` was given. Available variants are:";
+
+        size_t maxlen = 0;
+        for (const auto& key: matching_keys) {
+            maxlen = std::max(key.size(), maxlen);
+        }
+
+        for (const auto& key: matching_keys) {
+            auto description = outputs.at(key)->description;
+            std::string padding(maxlen - key.size(), ' ');
+            oss << "\n - '" << key << "'" << padding << ": " << description;
+        }
+        C10_THROW_ERROR(ValueError, oss.str());
+    }
+}
+
 
 static bool ends_with(const std::string& s, const std::string& suff) {
     return s.size() >= suff.size() && s.compare(s.size() - suff.size(), suff.size(), suff) == 0;

metatomic-torch/src/register.cpp

@@ -124,18 +124,18 @@ TORCH_LIBRARY(metatomic, m) {
     m.class_<ModelOutputHolder>("ModelOutput")
         .def(
             torch::init<
-                std::string,
                 std::string,
                 std::string,
                 bool,
-                std::vector<std::string>
+                std::vector<std::string>,
+                std::string
             >(),
             DOCSTRING, {
-                torch::arg("description") = "",
                 torch::arg("quantity") = "",
                 torch::arg("unit") = "",
                 torch::arg("per_atom") = false,
-                torch::arg("explicit_gradients") = std::vector<std::string>()
+                torch::arg("explicit_gradients") = std::vector<std::string>(),
+                torch::arg("description") = "",
             }
         )
         .def_readwrite("description", &ModelOutputHolder::description)
@@ -220,6 +220,7 @@ TORCH_LIBRARY(metatomic, m) {
     // standalone functions
     m.def("version() -> str", version);
     m.def("pick_device(str[] model_devices, str? requested_device = None) -> str", pick_device);
+    m.def("pick_output(str requested_output, Dict(str, __torch__.torch.classes.metatomic.ModelOutput) outputs, str? desired_variant = None) -> str", pick_output);
 
     m.def("read_model_metadata(str path) -> __torch__.torch.classes.metatomic.ModelMetadata", read_model_metadata);
     m.def("unit_conversion_factor(str quantity, str from_unit, str to_unit) -> float", unit_conversion_factor);

metatomic-torch/tests/misc.cpp

@@ -1,3 +1,4 @@
+#include <c10/util/intrusive_ptr.h>
 #include <torch/torch.h>
 
 #include "metatomic/torch.hpp"
@@ -61,3 +62,35 @@ TEST_CASE("Pick device") {
     std::vector<std::string> supported_devices_cpu = {"cpu"};
     CHECK_THROWS_WITH(metatomic_torch::pick_device(supported_devices_cpu, "cuda"), StartsWith("failed to find requested device"));
 }
+
+
+TEST_CASE("Pick variant") {
+    auto output_base = torch::make_intrusive<metatomic_torch::ModelOutputHolder>();
+    output_base->description = "my awesome energy";
+    output_base->set_quantity("energy");
+
+    auto variantA = torch::make_intrusive<metatomic_torch::ModelOutputHolder>();
+    variantA->set_quantity("energy");
+    variantA->description = "Variant A of the output";
+
+    auto variantfoo = torch::make_intrusive<metatomic_torch::ModelOutputHolder>();
+    variantfoo->set_quantity("energy");
+    variantfoo->description = "Variant foo of the output";
+
+    auto outputs = torch::Dict<std::string, metatomic_torch::ModelOutput>();
+    outputs.insert("energy", output_base);
+    outputs.insert("energy/A", variantA);
+    outputs.insert("energy/foo", variantfoo);
+
+    CHECK(metatomic_torch::pick_output("energy", outputs) == "energy");
+    CHECK(metatomic_torch::pick_output("energy", outputs, "A") == "energy/A");
+    CHECK_THROWS_WITH(metatomic_torch::pick_output("foo", outputs), StartsWith("output 'foo' not found in outputs"));
+    CHECK_THROWS_WITH(metatomic_torch::pick_output("energy", outputs, "C"), StartsWith("variant 'C' for output 'energy' not found in outputs"));
+
+    (void)outputs.erase("energy");
+    const auto *err = "output 'energy' has no default variant and no `desired_variant` was given. "
+        "Available variants are:\n"
+        " - 'energy/A'  : Variant A of the output\n"
+        " - 'energy/foo': Variant foo of the output";
+    CHECK_THROWS_WITH(metatomic_torch::pick_output("energy", outputs), StartsWith(err));
+}

python/metatomic_torch/metatomic/torch/__init__.py

@@ -17,6 +17,7 @@
         check_atomistic_model,
         load_model_extensions,
         pick_device,
+        pick_output,
         read_model_metadata,
         register_autograd_neighbors,
         unit_conversion_factor,
@@ -40,6 +41,7 @@
     register_autograd_neighbors = torch.ops.metatomic.register_autograd_neighbors
     unit_conversion_factor = torch.ops.metatomic.unit_conversion_factor
     pick_device = torch.ops.metatomic.pick_device
+    pick_output = torch.ops.metatomic.pick_output
 
 from .model import (  # noqa: F401
     AtomisticModel,

python/metatomic_torch/metatomic/torch/ase_calculator.py

@@ -19,6 +19,7 @@
     System,
     load_atomistic_model,
     pick_device,
+    pick_output,
     register_autograd_neighbors,
 )
 
@@ -166,50 +167,58 @@ def __init__(
                 f"found unexpected dtype in model capabilities: {capabilities.dtype}"
             )
 
-        self._energy_key = "energy"
-        self._energy_uq_key = "energy_uncertainty"
-        self._nc_forces_key = "non_conservative_forces"
-        self._nc_stress_key = "non_conservative_stress"
-
-        if variants:
-            if "energy" in variants:
-                self._energy_key += f"/{variants['energy']}"
-                self._energy_uq_key += f"/{variants['energy']}"
-                self._nc_forces_key += f"/{variants['energy']}"
-                self._nc_stress_key += f"/{variants['energy']}"
-
-            if "energy_uncertainty" in variants:
-                if variants["energy_uncertainty"] is None:
-                    self._energy_uq_key = "energy_uncertainty"
-                else:
-                    self._energy_uq_key += f"/{variants['energy_uncertainty']}"
-
-            if non_conservative:
-                if (
-                    "non_conservative_stress" in variants
-                    and "non_conservative_forces" in variants
-                    and (
-                        (variants["non_conservative_stress"] is None)
-                        != (variants["non_conservative_forces"] is None)
-                    )
-                ):
-                    raise ValueError(
-                        "if both 'non_conservative_stress' and "
-                        "'non_conservative_forces' are present in `variants`, they "
-                        "must either be both `None` or both not `None`."
-                    )
+        # resolve the output keys to use based on the requested variants
+        variants = variants or {}
+        default_variant = variants.get("energy")
+
+        resolved_variants = {
+            key: variants.get(key, default_variant)
+            for key in [
+                "energy",
+                "energy_uncertainty",
+                "non_conservative_forces",
+                "non_conservative_stress",
+            ]
+        }
+
+        outputs = capabilities.outputs
+        self._energy_key = pick_output("energy", outputs, resolved_variants["energy"])
 
-                if "non_conservative_forces" in variants:
-                    if variants["non_conservative_forces"] is None:
-                        self._nc_forces_key = "non_conservative_forces"
-                    else:
-                        self._nc_forces_key += f"/{variants['non_conservative_forces']}"
-
-                if "non_conservative_stress" in variants:
-                    if variants["non_conservative_stress"] is None:
-                        self._nc_stress_key = "non_conservative_stress"
-                    else:
-                        self._nc_stress_key += f"/{variants['non_conservative_stress']}"
+        if uncertainty_threshold is not None:
+            self._energy_uq_key = pick_output(
+                "energy_uncertainty", outputs, resolved_variants["energy_uncertainty"]
+            )
+        else:
+            self._energy_uq_key = "energy_uncertainty"
+
+        if non_conservative:
+            if (
+                "non_conservative_stress" in variants
+                and "non_conservative_forces" in variants
+                and (
+                    (variants["non_conservative_stress"] is None)
+                    != (variants["non_conservative_forces"] is None)
+                )
+            ):
+                raise ValueError(
+                    "if both 'non_conservative_stress' and "
+                    "'non_conservative_forces' are present in `variants`, they "
+                    "must either be both `None` or both not `None`."
+                )
+
+            self._nc_forces_key = pick_output(
+                "non_conservative_forces",
+                outputs,
+                resolved_variants["non_conservative_forces"],
+            )
+            self._nc_stress_key = pick_output(
+                "non_conservative_stress",
+                outputs,
+                resolved_variants["non_conservative_stress"],
+            )
+        else:
+            self._nc_forces_key = "non_conservative_forces"
+            self._nc_stress_key = "non_conservative_stress"
 
         if additional_outputs is None:
             self._additional_output_requests = {}

python/metatomic_torch/metatomic/torch/documentation.py

@@ -247,24 +247,17 @@ def __ne__(self, other: "NeighborListOptions") -> bool:
 
 
 class ModelOutput:
-    """Description of one of the quantity a model can compute."""
+    """Information about one of the quantity a model can compute."""
 
     def __init__(
         self,
-        description: str = "",
         quantity: str = "",
         unit: str = "",
         per_atom: bool = False,
         explicit_gradients: List[str] = [],  # noqa B006
+        description: str = "",
     ):
         pass
-    
-    @property
-    def description(self) -> str:
-        """
-        A description of this output. Especially recommended for non-standard outputs
-        and variants of the one unit.
-        """
 
     @property
     def quantity(self) -> str:
@@ -294,6 +287,13 @@ def unit(self) -> str:
     :py:class:`TensorMap`.
     """
 
+    @property
+    def description(self) -> str:
+        """
+        A description of this output. Especially recommended for non-standard outputs
+        and variants of the one unit.
+        """
+
 
 class ModelCapabilities:
     """Description of a model capabilities, i.e. everything a model can do."""
@@ -544,3 +544,18 @@ def pick_device(model_devices: List[str], desired_device: Optional[str]) -> str:
     :param desired_device: user-provided desired device. If ``None`` or not available,
         the first available device from ``model_devices`` will be picked.
     """
+
+
+def pick_output(
+    requested_output: str,
+    outputs: Dict[str, ModelOutput],
+    desired_variant: Optional[str] = None,
+) -> str:
+    """
+    Pick the output for the given ``requested_output`` from the availabilities of the
+    model's ``outputs``, according to the optional ``desired_variant``.
+
+    :param requested_output: name of the output to pick a variant for
+    :param outputs: all available outputs from the model
+    :param desired_variant: if provided, try to pick this specific variant
+    """