Now working with the conda lammps

pfebrer · pfebrer · commit b1f3e52611d8 · 2026-03-13T22:25:26.000+01:00
diff --git a/metatomic-torch/include/metatomic/torch/model.hpp b/metatomic-torch/include/metatomic/torch/model.hpp
@@ -49,23 +49,23 @@ class METATOMIC_TORCH_EXPORT ModelOutputHolder: public torch::CustomClassHolder
         std::string quantity,
         std::string unit,
         torch::optional<bool> per_atom_,
-        torch::optional<std::string> sample_kind_,
         std::vector<std::string> explicit_gradients_,
-        std::string description_
+        std::string description_,
+        torch::optional<std::string> sample_kind = torch::nullopt
     ):
         description(std::move(description_)),
         explicit_gradients(std::move(explicit_gradients_))
     {
         this->set_quantity(std::move(quantity));
         this->set_unit(std::move(unit));
 
-        if (per_atom_.has_value() && sample_kind_.has_value()) {
+        if (per_atom_.has_value() && sample_kind.has_value()) {
             C10_THROW_ERROR(ValueError, "Cannot set both `per_atom` and `sample_kind` for a ModelOutput");
-        } else if (per_atom_.has_value()) {
-            this->set_per_atom(std::move(per_atom_.value()));
+        } else if (sample_kind.has_value()) {
+            this->set_sample_kind(sample_kind.value());
         } else {
-            this->set_sample_kind(std::move(sample_kind_.value_or("system")));
-        } 
+            this->set_per_atom(per_atom_.value_or(false));
+        }
 
     }
 
@@ -91,16 +91,43 @@ class METATOMIC_TORCH_EXPORT ModelOutputHolder: public torch::CustomClassHolder
     /// set the unit of the output
     void set_unit(std::string unit);
 
-    /// Which kind of sample this output is. E.g. "system", "atom", "atom_pair"...
-    std::string sample_kind;
-    const std::string& get_sample_kind() const {
-        return sample_kind;
+    /// Although we are moving to using `sample_kind` instead of `per_atom`, 
+    /// for now we need to keep the `per_atom` data member for backward
+    /// compatibility in the C++ API and ABI.
+    bool per_atom;
+
+    /// The setter and getter for `per_atom` that are used in python, which
+    /// allow us to raise an error if `sample_kind` can't be mapped to
+    /// a boolean value for `per_atom`.
+    void set_per_atom(bool per_atom_) {
+        this->per_atom = per_atom_;
+        /// Unset sample_kind
+        this->sample_kind_ = torch::nullopt;
+    }
+    bool get_per_atom() const {
+        if (sample_kind_.has_value()) {
+            C10_THROW_ERROR(
+                ValueError,
+                "Can't infer `per_atom` from `sample_kind` '" + this->sample_kind() + "'. "
+                "`per_atom` only makes sense for `sample_kind` 'atom' and 'system'."
+            );
+        }
+        return per_atom;
+    }
+
+    /// Sample kind of the output. For now it relies on the `per_atom` boolean for
+    /// backward compatibility, but in the future we will simply have a variable
+    /// storing the sample kind.
+    const std::string sample_kind() const {
+        if (sample_kind_.has_value()) {
+            return sample_kind_.value();
+        } else if (per_atom) {
+            return "atom";
+        } else {
+            return "system";
+        }
     }
     void set_sample_kind(std::string sample_kind);
-    
-    // For backward compatibility.
-    void set_per_atom(bool per_atom);
-    bool get_per_atom() const;
 
     /// Which gradients should be computed eagerly and stored inside the output
     /// `TensorMap`
@@ -114,6 +141,12 @@ class METATOMIC_TORCH_EXPORT ModelOutputHolder: public torch::CustomClassHolder
 private:
     std::string quantity_;
     std::string unit_;
+    torch::optional<std::string> sample_kind_;
+    std::vector<std::string> supported_sample_kinds_ = {
+        "system",
+        "atom",
+        "atom_pair",
+    };
 };
 
 
diff --git a/metatomic-torch/src/model.cpp b/metatomic-torch/src/model.cpp
@@ -66,49 +66,48 @@ void ModelOutputHolder::set_unit(std::string unit) {
     this->unit_ = std::move(unit);
 }
 
-void ModelOutputHolder::set_sample_kind(std::string sample_kind) {
-    /// Sample kind has to be one of "system", "atom", "atom_pair"
-    if (sample_kind != "system" && sample_kind != "atom" && sample_kind != "atom_pair") {
-        C10_THROW_ERROR(ValueError,
-            "invalid sample kind '" + sample_kind + "'. Only 'system', 'atom' and 'atom_pair' are supported"
-        );
+std::string join_strings(const std::vector<std::string>& strings, const std::string& separator) {
+    std::ostringstream oss;
+    for (size_t i = 0; i < strings.size(); ++i) {
+        if (i > 0) {
+            oss << separator;
+        }
+        oss << strings[i];
     }
-
-    this->sample_kind = std::move(sample_kind);
+    return oss.str();
 }
 
-/// For backward compatibility, we keep the `per_atom` property, which
-/// is not really stored, it is simply derived from `sample_kind`.
-bool ModelOutputHolder::get_per_atom() const {
+void ModelOutputHolder::set_sample_kind(std::string sample_kind) {
     if (sample_kind == "atom") {
-        return true;
+        this->set_per_atom(true);
     } else if (sample_kind == "system") {
-        return false;
+        this->set_per_atom(false);
     } else {
-        C10_THROW_ERROR( ValueError,
-            "Cannot determine `per_atom` from sample_kind '" + sample_kind + "'. "
-        );
-    }
-}
+        /// If sample_kind is not a value that can be mapped to per_atom,
+        /// we just store the value in the sample_kind_ private field.
 
-/// If the user sets `per_atom`, we update the `sample_kind` accordingly.
-void ModelOutputHolder::set_per_atom(bool per_atom) {
-    if (per_atom) {
-        this->sample_kind = "atom";
-    } else {
-        this->sample_kind = "system";
+        // Warn if the sample_kind is not one of the supported ones.
+        if (std::find(supported_sample_kinds_.begin(), supported_sample_kinds_.end(), sample_kind) == supported_sample_kinds_.end()) {
+            TORCH_WARN(
+                "Sample_kind '", sample_kind, "' is not officially supported. ",
+                "This means that metatomic doesn't natively understand how to deal ",
+                "with such outputs. If this is a mistake, pass one of the supported ",
+                "sample kinds instead: [", join_strings(supported_sample_kinds_, ", "), "]. "
+            );
+        }
+
+        this->sample_kind_ = std::move(sample_kind);
     }
 }
 
-
-
 static nlohmann::json model_output_to_json(const ModelOutputHolder& self) {
     nlohmann::json result;
 
     result["class"] = "ModelOutput";
     result["quantity"] = self.quantity();
     result["unit"] = self.unit();
-    result["sample_kind"] = self.sample_kind;
+    result["per_atom"] = self.per_atom;
+    result["sample_kind"] = self.sample_kind();
     result["explicit_gradients"] = self.explicit_gradients;
     result["description"] = self.description;
 
diff --git a/metatomic-torch/src/outputs.cpp b/metatomic-torch/src/outputs.cpp
@@ -86,12 +86,18 @@ static void validate_atomic_samples(
 
     // Check if the samples names are as expected based on the sample_kind
     std::vector<std::string> expected_samples_names;
-    if (request->sample_kind == "atom") {
+    if (request->sample_kind() == "atom") {
         expected_samples_names = {"system", "atom"};
-    } else if (request->sample_kind == "atom_pair") {
-        expected_samples_names = {"system", "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"};  
-    } else {
+    } else if (request->sample_kind() == "system") {
         expected_samples_names = {"system"};
+    } else if (request->sample_kind() == "atom_pair") {
+        expected_samples_names = {"system", "first_atom", "second_atom"};
+    } else {
+        C10_THROW_ERROR(ValueError,
+            "Metatomic does not support validating samples for sample_kind"
+            "other than 'system', 'atom' or 'atom_pair' at the moment."
+            " Received sample_kind '" + request->sample_kind()
+        );
     }
 
     if (block->samples()->names() != expected_samples_names) {
@@ -104,7 +110,7 @@ static void validate_atomic_samples(
 
     // Check if the samples match the systems and selected_atoms
     Labels expected_samples;
-    if (request->sample_kind == "atom") {
+    if (request->sample_kind() == "atom") {
         std::vector<int64_t> expected_values_flat;
         for (size_t s; s < systems.size(); s++) {
             for (size_t a; a < systems[s]->size(); a++) {
@@ -123,7 +129,7 @@ static void validate_atomic_samples(
         if (selected_atoms) {
             expected_samples = expected_samples->set_intersection(selected_atoms.value());
         }
-    } else if (request->sample_kind == "system") {
+    } else if (request->sample_kind() == "system") {
         expected_samples = torch::make_intrusive<LabelsHolder>(
             "system",
             torch::arange(static_cast<int64_t>(systems.size()), tensor_options).reshape({-1, 1}),
@@ -139,6 +145,9 @@ static void validate_atomic_samples(
             );
             expected_samples = expected_samples->set_intersection(selected_systems);
         }
+    } else {
+        /// We don't validate values for other cases for now
+        return;
     }
 
     if (expected_samples->set_union(block->samples())->size() != expected_samples->size()) {
diff --git a/metatomic-torch/src/register.cpp b/metatomic-torch/src/register.cpp
@@ -154,24 +154,24 @@ TORCH_LIBRARY(metatomic, m) {
                 std::string,
                 std::string,
                 torch::optional<bool>,
-                torch::optional<std::string>,
                 std::vector<std::string>,
-                std::string
+                std::string,
+                torch::optional<std::string>
             >(),
             DOCSTRING, {
                 torch::arg("quantity") = "",
                 torch::arg("unit") = "",
                 torch::arg("per_atom") = std::nullopt,
-                torch::arg("sample_kind") = std::nullopt,
                 torch::arg("explicit_gradients") = std::vector<std::string>(),
                 torch::arg("description") = "",
+                torch::arg("sample_kind") = std::nullopt,
             }
         )
         .def_readwrite("description", &ModelOutputHolder::description)
         .def_property("quantity", &ModelOutputHolder::quantity, &ModelOutputHolder::set_quantity)
         .def_property("unit", &ModelOutputHolder::unit, &ModelOutputHolder::set_unit)
         .def_property("per_atom", &ModelOutputHolder::get_per_atom, &ModelOutputHolder::set_per_atom)
-        .def_property("sample_kind", &ModelOutputHolder::get_sample_kind, &ModelOutputHolder::set_sample_kind)
+        .def_property("sample_kind", &ModelOutputHolder::sample_kind, &ModelOutputHolder::set_sample_kind)
         .def_readwrite("explicit_gradients", &ModelOutputHolder::explicit_gradients)
         .def_pickle(
             [](const ModelOutput& self) -> std::string {
diff --git a/metatomic-torch/tests/models.cpp b/metatomic-torch/tests/models.cpp
@@ -65,7 +65,7 @@ TEST_CASE("Models metadata") {
         output->description = "my awesome energy";
         output->set_quantity("energy");
         output->set_unit("kJ / mol");
-        output->set_per_atom(false);
+        output->per_atom = false;
         output->explicit_gradients = {"baz", "not.this-one_"};
 
         const auto* expected = R"({
@@ -75,6 +75,7 @@ TEST_CASE("Models metadata") {
         "baz",
         "not.this-one_"
     ],
+    "per_atom": false,
     "quantity": "energy",
     "sample_kind": "system",
     "unit": "kJ / mol"
@@ -90,7 +91,7 @@ TEST_CASE("Models metadata") {
         output = ModelOutputHolder::from_json(json);
         CHECK(output->quantity() == "length");
         CHECK(output->unit().empty());
-        CHECK(output->get_per_atom() == false);
+        CHECK(output->per_atom == false);
         CHECK(output->explicit_gradients.empty());
 
         CHECK_THROWS_WITH(
@@ -132,11 +133,11 @@ TEST_CASE("Models metadata") {
         options->set_length_unit("nanometer");
 
         auto output1 = torch::make_intrusive<ModelOutputHolder>();
-        output1->set_per_atom(false);
+        output1->per_atom = false;
         options->outputs.insert("output_1", output1);
 
         auto output2 = torch::make_intrusive<ModelOutputHolder>();
-        output2->set_per_atom(true);
+        output2->per_atom = true;
         output2->set_quantity("something");
         output2->set_unit("something");
         options->outputs.insert("output_2", output2);
@@ -149,6 +150,7 @@ TEST_CASE("Models metadata") {
             "class": "ModelOutput",
             "description": "",
             "explicit_gradients": [],
+            "per_atom": false,
             "quantity": "",
             "sample_kind": "system",
             "unit": ""
@@ -157,6 +159,7 @@ TEST_CASE("Models metadata") {
             "class": "ModelOutput",
             "description": "",
             "explicit_gradients": [],
+            "per_atom": true,
             "quantity": "something",
             "sample_kind": "atom",
             "unit": "something"
@@ -194,7 +197,7 @@ TEST_CASE("Models metadata") {
         auto output = options->outputs.at("foo");
         CHECK(output->quantity().empty());
         CHECK(output->unit().empty());
-        CHECK(output->get_per_atom() == false);
+        CHECK(output->per_atom == false);
         CHECK(output->explicit_gradients == std::vector<std::string>{"test"});
 
         CHECK_THROWS_WITH(
@@ -221,7 +224,7 @@ TEST_CASE("Models metadata") {
         capabilities->supported_devices = {"cuda", "xla", "cpu"};
 
         auto output = torch::make_intrusive<ModelOutputHolder>();
-        output->set_per_atom(true);
+        output->per_atom = true;
         output->set_quantity("length");
         output->explicit_gradients.emplace_back("µ-λ");
 
@@ -246,6 +249,7 @@ TEST_CASE("Models metadata") {
             "explicit_gradients": [
                 "\u00b5-\u03bb"
             ],
+            "per_atom": true,
             "quantity": "length",
             "sample_kind": "atom",
             "unit": ""
@@ -284,7 +288,7 @@ TEST_CASE("Models metadata") {
         output = capabilities->outputs().at("tests::foo");
         CHECK(output->quantity().empty());
         CHECK(output->unit().empty());
-        CHECK(output->get_per_atom() == false);
+        CHECK(output->per_atom == false);
         CHECK(output->explicit_gradients == std::vector<std::string>{"µ-test"});
 
         CHECK_THROWS_WITH(
@@ -302,7 +306,7 @@ TEST_CASE("Models metadata") {
 
         auto capabilities_variants = torch::make_intrusive<ModelCapabilitiesHolder>();
         auto output_variant = torch::make_intrusive<ModelOutputHolder>();
-        output_variant->set_per_atom(true);
+        output_variant->per_atom = true;
         output_variant->description = "variant output";
 
         auto outputs_variant = torch::Dict<std::string, ModelOutput>();
diff --git a/python/metatomic_torch/metatomic/torch/documentation.py b/python/metatomic_torch/metatomic/torch/documentation.py
@@ -254,9 +254,9 @@ def __init__(
         quantity: str = "",
         unit: str = "",
         per_atom: Optional[bool] = None,
-        sample_kind: Optional[Literal["system", "atom", "atom_pair"]] = None,
         explicit_gradients: List[str] = [],  # noqa B006
         description: str = "",
+        sample_kind: Optional[Literal["system", "atom", "atom_pair"]] = None,
     ):
         pass
 
diff --git a/python/metatomic_torch/tests/outputs.py b/python/metatomic_torch/tests/outputs.py
@@ -53,9 +53,10 @@ def test_sample_kind():
     with pytest.raises(ValueError):
         ModelOutput(per_atom=True, sample_kind="system")
 
-    # Check that setting sample_kind to an invalid value raises an error
-    with pytest.raises(ValueError):
-        ModelOutput(sample_kind="invalid_value")
+    # Arbitrary sample_kind values are allowed, although they will not be
+    # supported by metatomic interfaces to engines. Setting sample_kind
+    # to an arbitrary value will issue a warning.
+    ModelOutput(sample_kind="arbitrary_value")
 
     # Initialize model output with sample_kind="atom_pair"
     # and check that per_atom can not be retrieved
