Refactor unit conversion and allow additional inputs

cmake/Modules/Packages/ML-METATOMIC.cmake

@@ -41,7 +41,7 @@ set(METATENSOR_CORE_SHA256 "809a799b1c8d58b7ede3868d6ebe4123924ba31e4637481f9ca5
 set(METATENSOR_TORCH_VERSION "0.9.1")
 set(METATENSOR_TORCH_SHA256 "fa21ae9f5111f3b40479e51ed55152154fc2c6eb30f38d9de6adad53938d0444")
 
-set(METATOMIC_TORCH_VERSION "0.1.13")
+set(METATOMIC_TORCH_VERSION "0.1.14")
 set(METATOMIC_TORCH_SHA256 "ca521613d81eb380d8ce17a48ee640867e11c00b5041dd28eaf31dd9d8b3f9c8")
 
 set(DOWNLOAD_METATENSOR_DEFAULT ON)

src/ML-METATOMIC/fix_metatomic.cpp

@@ -25,6 +25,7 @@
 ------------------------------------------------------------------------- */
 #include "metatomic_types.h"
 #include "metatomic_system.h"
+#include "metatomic_units.h"
 
 #include "fix_metatomic.h"
 
@@ -61,26 +62,15 @@ FixMetatomic::FixMetatomic(LAMMPS *lmp, int narg, char **arg): Fix(lmp, narg, ar
     }
 
     // Determine unit system for the ML model
-    // Currently only 'metal' units are fully supported for momenta
-    std::string energy_unit;
-    std::string length_unit;
-    if (strcmp(update->unit_style, "metal") == 0) {
-        length_unit = "angstrom";
-        this->momentum_conversion_factor = 10.1805057179 / 1000.0;
-    } else if (strcmp(update->unit_style, "real") == 0) {
-        length_unit = "angstrom";
-        this->momentum_conversion_factor = 10.1805057179;
-    } else if (strcmp(update->unit_style, "si") == 0) {
-        length_unit = "m";
-        this->momentum_conversion_factor = 10.1805057179 / 1.6605390666e-22;
-    } else {
+    if (strcmp(update->unit_style, "lj") == 0) {
         error->all(FLERR, "unsupported units '{}' for fix metatomic", update->unit_style);
     }
-
-    // For now, only metal units are fully tested and supported
-    if (strcmp(update->unit_style, "metal") != 0) {
-        error->all(FLERR, "fix metatomic currently only supports 'metal' units");
-    }
+    std::string energy_unit= metatomic_unit_map.at("energy").at(update->unit_style);
+    std::string length_unit = metatomic_unit_map.at("position").at(update->unit_style);
+    std::string mass_unit = metatomic_unit_map.at("mass").at(update->unit_style);
+    std::string velocity_unit = metatomic_unit_map.at("velocity").at(update->unit_style);
+    std::string momentum_unit = mass_unit + "*" + velocity_unit;
+    this->momentum_conversion_factor = metatomic_torch::unit_conversion_factor(momentum_unit, "(u*eV)^(1/2)");
 
     if (narg < 4) {
         error->all(FLERR,
@@ -445,19 +435,24 @@ void FixMetatomic::initial_integrate(int /*vflag*/) {
         error->all(FLERR, "the model requested an unsupported dtype '{}'", mta_data->capabilities->dtype());
     }
 
+    // deal with the model requested inputs
+    std::map<std::string, metatomic_torch::ModelOutput> input_holders;
+    auto requested_inputs = mta_data->model->run_method("requested_inputs", /*use_new_names=*/ true).toGenericDict();
+    for (const auto& entry : requested_inputs) {
+        input_holders.emplace(
+            entry.key().toStringRef(),
+            entry.value().toCustomClass<metatomic_torch::ModelOutputHolder>()
+        );
+    }
     // transform from LAMMPS to metatomic System
     auto system = this->system_adaptor->system_from_lmp(
         mta_list,
         static_cast<bool>(vflag_global),
         dtype,
-        mta_data->device
+        mta_data->device,
+        input_holders
     );
 
-    // add the required additional inputs, for now FlashMD uses the old names
-    // and does not go through the requested_inputs mechanism.
-    this->system_adaptor->add_masses(system, "masses", 1.0);
-    this->system_adaptor->add_momenta(system, "momenta", this->momentum_conversion_factor);
-
     // Configure selected atoms for evaluation
     // Only run the calculation for atoms in the current group
     mta_data->selected_atoms_values.resize_({group->count(igroup), 2});

src/ML-METATOMIC/metatomic_system.cpp

@@ -16,14 +16,17 @@
 ------------------------------------------------------------------------- */
 #include "metatomic_system.h"
 #include "metatomic_timer.h"
+#include "metatomic_units.h"
 
 #include "atom.h"
 #include "comm.h"
 #include "domain.h"
 #include "error.h"
+#include "update.h"
 
 #include "neigh_list.h"
 
+#include <map>
 #include <string>
 
 #include <metatensor/torch.hpp>
@@ -206,6 +209,55 @@ static std::array<int32_t, 3> cell_shifts(
     return {shift_a, shift_b, shift_c};
 }
 
+metatensor_torch::TensorMap LAMMPS_NS::make_per_atom_tensormap(
+    const torch::Tensor& values,
+    const std::string& property_name,
+    const std::vector<std::string>& component_names
+) {
+    assert (values.dim() == static_cast<int64_t>(component_names.size() + 1));
+
+    auto n_atoms = values.size(0);
+    auto device = values.device();
+    auto label_tensor_options = torch::TensorOptions().dtype(torch::kInt32).device(device);
+
+    auto keys = metatensor_torch::LabelsHolder::single()->to(device);
+    auto samples_values = torch::column_stack({
+        torch::zeros(n_atoms, label_tensor_options).unsqueeze(1),
+        torch::arange(n_atoms, label_tensor_options).unsqueeze(1)
+    });
+    auto samples = torch::make_intrusive<metatensor_torch::LabelsHolder>(
+        std::vector<std::string>{"system", "atom"},
+        samples_values,
+        metatensor::assume_unique{}
+    );
+
+    auto components = std::vector<metatensor_torch::Labels>{};
+    for (size_t axis = 0; axis < component_names.size(); axis++) {
+        auto component_values = torch::arange(values.size(axis + 1), label_tensor_options).unsqueeze(1);
+        auto component = torch::make_intrusive<metatensor_torch::LabelsHolder>(
+            std::vector<std::string>{component_names[axis]},
+            component_values
+        );
+        components.push_back(component);
+    }
+
+    auto properties = torch::make_intrusive<metatensor_torch::LabelsHolder>(
+        std::vector<std::string>{property_name},
+        torch::tensor({{0}}, label_tensor_options)
+    );
+    auto block = torch::make_intrusive<metatensor_torch::TensorBlockHolder>(
+        values.unsqueeze(-1),
+        samples,
+        components,
+        properties
+    );
+
+    return torch::make_intrusive<metatensor_torch::TensorMapHolder>(
+        keys,
+        std::vector<metatensor_torch::TensorBlock>{block}
+    );
+}
+
 void MetatomicSystemAdaptor::guess_periodic_ghosts() {
     auto _ = MetatomicTimer("identifying periodic ghosts");
     auto total_n_atoms = atom->nlocal + atom->nghost;
@@ -520,7 +572,8 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp(
     NeighList* list,
     bool do_virial,
     torch::ScalarType dtype,
-    torch::Device device
+    torch::Device device,
+    const std::map<std::string, torch::intrusive_ptr<metatomic_torch::ModelOutputHolder>>& requested_inputs
 ) {
     auto _ = MetatomicTimer("creating System from LAMMPS data");
 
@@ -589,10 +642,25 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp(
 
     this->setup_neighbors(system, list);
 
+    for (const auto& [property, input]: requested_inputs) {
+        const auto& property_name = property.c_str();
+        const auto& unit = input->unit().c_str();
+        if (strcmp(property_name, "mass") == 0 || strcmp(property_name, "masses") == 0) {
+            add_masses(system, metatomic_torch::unit_conversion_factor(metatomic_unit_map.at("mass").at(update->unit_style), unit));
+        } else if (strcmp(property_name, "momentum") == 0 || strcmp(property_name, "momenta") == 0) {
+            const auto& momentum_unit = metatomic_unit_map.at("mass").at(update->unit_style) + "*" + metatomic_unit_map.at("velocity").at(update->unit_style);
+            add_momenta(system, metatomic_torch::unit_conversion_factor(momentum_unit, unit));
+        } else if (strcmp(property_name, "velocity") == 0 || strcmp(property_name, "velocities") == 0) {
+            add_velocities(system, metatomic_torch::unit_conversion_factor(metatomic_unit_map.at("velocity").at(update->unit_style), unit));
+        } else {
+            error->all(FLERR, "compute metatomic: the model requested an unsupported additional input of '{}'", property_name);
+        }
+    }
+
     return system;
 }
 
-void MetatomicSystemAdaptor::add_masses(metatomic_torch::System& system, std::string name, double unit_conversion) {
+void MetatomicSystemAdaptor::add_masses(metatomic_torch::System& system, double unit_conversion) {
     double* rmass = atom->rmass;
     double* mass = atom->mass;
     int* type = atom->type;
@@ -608,56 +676,43 @@ void MetatomicSystemAdaptor::add_masses(metatomic_torch::System& system, std::st
         torch::TensorOptions().dtype(torch::kInt).device(torch::kCPU)
     );
 
-    // gather masses (per-atom) in a tensor and ship to device
+    // gather masses (per-atom) in a CPU tensor and ship to device
     torch::Tensor masses;
-    auto float_tensor_options = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU);
+    auto cpu_float_tensor_options = torch::TensorOptions().dtype(dtype).device(torch::kCPU);
     if (rmass) {
         masses = torch::from_blob(
             rmass,
             {total_n_atoms},
-            float_tensor_options.requires_grad(false)
-        );
+            torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU).requires_grad(false)
+        ).to(dtype);
     } else {
         // need to map from atom type to mass
-        masses = torch::empty({total_n_atoms}, float_tensor_options);
-        auto masses_accessor = masses.accessor<double, 1>();
-        for (int i=0; i<total_n_atoms; i++) {
-            masses_accessor[i] = mass[type[i]];
+        masses = torch::empty({total_n_atoms}, cpu_float_tensor_options);
+        if (dtype == torch::kFloat64) {
+            auto masses_accessor = masses.accessor<double, 1>();
+            for (int i=0; i<total_n_atoms; i++) {
+                masses_accessor[i] = mass[type[i]];
+            }
+        } else if (dtype == torch::kFloat32) {
+            auto masses_accessor = masses.accessor<float, 1>();
+            for (int i=0; i<total_n_atoms; i++) {
+                masses_accessor[i] = static_cast<float>(mass[type[i]]);
+            }
+        } else {
+            error->one(FLERR, "invalid dtype, this is a bug");
         }
     }
 
     masses = masses.index_select(0, mta_to_lmp_tensor);
     masses = masses * unit_conversion;
+    masses = masses.to(device);
+    auto tensor = make_per_atom_tensormap(masses, "mass");
 
-    auto keys = metatensor_torch::LabelsHolder::single()->to(device);
-    auto label_tensor_options = torch::TensorOptions().dtype(torch::kInt32).device(device);
-    auto samples_values = torch::column_stack({
-        torch::zeros(system->size(), label_tensor_options).unsqueeze(1),
-        torch::arange(system->size(), label_tensor_options).unsqueeze(1)
-    });
-    auto samples = torch::make_intrusive<metatensor_torch::LabelsHolder>(
-        std::vector<std::string>{"system","atom"},
-        samples_values,
-        metatensor::assume_unique{}
-    );
-    auto properties = metatensor_torch::LabelsHolder::single()->to(device);
-    auto block = torch::make_intrusive<metatensor_torch::TensorBlockHolder>(
-        masses.to(dtype).to(device).unsqueeze(-1),  // add property dimension
-        samples,
-        std::vector<metatensor_torch::Labels>{},
-        properties
-    );
-    auto tensor = torch::make_intrusive<metatensor_torch::TensorMapHolder>(
-        keys,
-        std::vector<metatensor_torch::TensorBlock>{block}
-    );
-
-    assert(name == "mass" || name == "masses");
-    system->add_data(name, tensor);
+    system->add_data("mass", tensor);
 }
 
 
-void MetatomicSystemAdaptor::add_momenta(metatomic_torch::System& system, std::string name, double unit_conversion) {
+void MetatomicSystemAdaptor::add_momenta(metatomic_torch::System& system, double unit_conversion) {
     double* rmass = atom->rmass;
     double* mass = atom->mass;
     double** v = atom->v;
@@ -674,50 +729,80 @@ void MetatomicSystemAdaptor::add_momenta(metatomic_torch::System& system, std::s
         torch::TensorOptions().dtype(torch::kInt).device(torch::kCPU)
     );
 
-    // gather momenta (per-atom) in a tensor and ship to device
-    torch::Tensor momenta = torch::zeros({total_n_atoms, 3}, torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU));
-    auto momenta_accessor = momenta.accessor<double, 2>();
-    for (int i=0; i<total_n_atoms; i++) {
-        double m = rmass ? rmass[i] : mass[type[i]];
-        momenta_accessor[i][0] = m * v[i][0];
-        momenta_accessor[i][1] = m * v[i][1];
-        momenta_accessor[i][2] = m * v[i][2];
+    // gather momenta (per-atom) in a CPU tensor and ship to device
+    torch::Tensor momenta = torch::empty(
+        {total_n_atoms, 3},
+        torch::TensorOptions().dtype(dtype).device(torch::kCPU)
+    );
+    if (dtype == torch::kFloat64) {
+        auto momenta_accessor = momenta.accessor<double, 2>();
+        for (int i=0; i<total_n_atoms; i++) {
+            double m = rmass ? rmass[i] : mass[type[i]];
+            momenta_accessor[i][0] = m * v[i][0];
+            momenta_accessor[i][1] = m * v[i][1];
+            momenta_accessor[i][2] = m * v[i][2];
+        }
+    } else if (dtype == torch::kFloat32) {
+        auto momenta_accessor = momenta.accessor<float, 2>();
+        for (int i=0; i<total_n_atoms; i++) {
+            double m = rmass ? rmass[i] : mass[type[i]];
+            momenta_accessor[i][0] = static_cast<float>(m * v[i][0]);
+            momenta_accessor[i][1] = static_cast<float>(m * v[i][1]);
+            momenta_accessor[i][2] = static_cast<float>(m * v[i][2]);
+        }
+    } else {
+        error->one(FLERR, "invalid dtype, this is a bug");
     }
 
     momenta = momenta.index_select(0, mta_to_lmp_tensor);
     momenta = momenta * unit_conversion;
+    momenta = momenta.to(device);
+    auto tensor = make_per_atom_tensormap(momenta, "momentum", {"xyz"});
 
-    auto keys = metatensor_torch::LabelsHolder::single()->to(device);
+    system->add_data("momentum", tensor);
+}
 
-    auto label_tensor_options = torch::TensorOptions().dtype(torch::kInt32).device(device);
-    auto samples_values = torch::column_stack({
-        torch::zeros(system->size(), label_tensor_options).unsqueeze(1),
-        torch::arange(system->size(), label_tensor_options).unsqueeze(1)
-    });
-    auto samples = torch::make_intrusive<metatensor_torch::LabelsHolder>(
-        std::vector<std::string>{"system","atom"},
-        samples_values,
-        metatensor::assume_unique{}
-    );
+void MetatomicSystemAdaptor::add_velocities(metatomic_torch::System& system, double unit_conversion) {
+    double** v = atom->v;
 
-    auto component_values = torch::arange(3, label_tensor_options).unsqueeze(1);
-    auto component = torch::make_intrusive<metatensor_torch::LabelsHolder>(
-        std::vector<std::string>{"xyz"}, component_values
-    );
+    auto total_n_atoms = atom->nlocal + atom->nghost;
 
-    auto properties = metatensor_torch::LabelsHolder::single()->to(device);
+    auto device = system->device();
+    auto dtype = system->scalar_type();
 
-    auto block = torch::make_intrusive<metatensor_torch::TensorBlockHolder>(
-        momenta.to(dtype).to(device).unsqueeze(-1),
-        samples,
-        std::vector<metatensor_torch::Labels>{component},
-        properties
+    auto mta_to_lmp_tensor = torch::from_blob(
+        mta_to_lmp.data(),
+        {static_cast<int64_t>(mta_to_lmp.size())},
+        torch::TensorOptions().dtype(torch::kInt).device(torch::kCPU)
     );
-    auto tensor = torch::make_intrusive<metatensor_torch::TensorMapHolder>(
-        keys,
-        std::vector<metatensor_torch::TensorBlock>{block}
+
+    // gather velocities (per-atom) in a CPU tensor and ship to device
+    torch::Tensor velocities = torch::empty(
+        {total_n_atoms, 3},
+        torch::TensorOptions().dtype(dtype).device(torch::kCPU)
     );
+    if (dtype == torch::kFloat64) {
+        auto velocities_accessor = velocities.accessor<double, 2>();
+        for (int i=0; i<total_n_atoms; i++) {
+            velocities_accessor[i][0] = v[i][0];
+            velocities_accessor[i][1] = v[i][1];
+            velocities_accessor[i][2] = v[i][2];
+        }
+    } else if (dtype == torch::kFloat32) {
+        auto velocities_accessor = velocities.accessor<float, 2>();
+        for (int i=0; i<total_n_atoms; i++) {
+            velocities_accessor[i][0] = static_cast<float>(v[i][0]);
+            velocities_accessor[i][1] = static_cast<float>(v[i][1]);
+            velocities_accessor[i][2] = static_cast<float>(v[i][2]);
+        }
+    } else {
+        error->one(FLERR, "invalid dtype, this is a bug");
+    }
+
+    velocities = velocities.index_select(0, mta_to_lmp_tensor);
+    velocities = velocities * unit_conversion;
+    velocities = velocities.to(device);
+    auto tensor = make_per_atom_tensormap(velocities, "velocity", {"xyz"});
 
-    assert(name == "momentum" || name == "momenta");
-    system->add_data(name, tensor);
+    system->add_data(std::string{"velocity"}, tensor, /*override=*/true);
 }