From c51eee0e00f6fba1cc2106d98b8c8c544669aa62 Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Fri, 27 Mar 2026 23:24:36 +0100 Subject: [PATCH 1/4] Reuse the shared code of `add_masses` and `add_momenta` --- src/ML-METATOMIC/metatomic_system.cpp | 75 ++++++++++++++------------- src/ML-METATOMIC/metatomic_system.h | 9 ++++ 2 files changed, 47 insertions(+), 37 deletions(-) diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 4652449d3b9..265797ac789 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -628,29 +628,7 @@ void MetatomicSystemAdaptor::add_masses(metatomic_torch::System& system, double masses = masses.index_select(0, mta_to_lmp_tensor); masses = masses * unit_conversion; - - 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( - std::vector{"system","atom"}, - samples_values, - metatensor::assume_unique{} - ); - auto properties = metatensor_torch::LabelsHolder::single()->to(device); - auto block = torch::make_intrusive( - masses.to(dtype).to(device).unsqueeze(-1), // add property dimension - samples, - std::vector{}, - properties - ); - auto tensor = torch::make_intrusive( - keys, - std::vector{block} - ); + auto tensor = make_per_atom_tensormap(masses, dtype, device); system->add_data("masses", tensor, /*override=*/true); } @@ -685,37 +663,60 @@ void MetatomicSystemAdaptor::add_momenta(metatomic_torch::System& system, double momenta = momenta.index_select(0, mta_to_lmp_tensor); momenta = momenta * unit_conversion; + auto tensor = make_per_atom_tensormap(momenta, dtype, device, {"xyz"}); - auto keys = metatensor_torch::LabelsHolder::single()->to(device); + system->add_data("momenta", tensor, /*override=*/true); +} +metatensor_torch::TensorMap MetatomicSystemAdaptor::make_per_atom_tensormap( + const torch::Tensor& values, + torch::ScalarType dtype, + const torch::Device& device, + const std::vector& component_names +) { + if (values.dim() != static_cast(component_names.size() + 1)) { + error->all( + FLERR, + "metatomic per-atom TensorMap helper expected rank {} tensor, got rank {}", + component_names.size() + 1, + values.dim() + ); + } + + auto n_atoms = values.size(0); 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(system->size(), label_tensor_options).unsqueeze(1), - torch::arange(system->size(), label_tensor_options).unsqueeze(1) + torch::zeros(n_atoms, label_tensor_options).unsqueeze(1), + torch::arange(n_atoms, label_tensor_options).unsqueeze(1) }); auto samples = torch::make_intrusive( - std::vector{"system","atom"}, + std::vector{"system", "atom"}, samples_values, metatensor::assume_unique{} ); - auto component_values = torch::arange(3, label_tensor_options).unsqueeze(1); - auto component = torch::make_intrusive( - std::vector{"xyz"}, component_values - ); + auto components = std::vector{}; + 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( + std::vector{component_names[axis]}, + component_values + ); + components.push_back(component); + } auto properties = metatensor_torch::LabelsHolder::single()->to(device); - auto block = torch::make_intrusive( - momenta.to(dtype).to(device).unsqueeze(-1), + values.to(dtype).to(device).unsqueeze(-1), samples, - std::vector{component}, + components, properties ); - auto tensor = torch::make_intrusive( + + return torch::make_intrusive( keys, std::vector{block} ); - - system->add_data("momenta", tensor, /*override=*/true); } diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index 38fa9ed62d8..a7dbf7b0888 100644 --- a/src/ML-METATOMIC/metatomic_system.h +++ b/src/ML-METATOMIC/metatomic_system.h @@ -16,6 +16,7 @@ #include #include +#include #include "pointers.h" #include "pair.h" @@ -81,6 +82,14 @@ class MetatomicSystemAdaptor : public Pointers { // Add momenta as extra data to this system, only for atoms which are not // periodic images of other atoms virtual void add_momenta(metatomic_torch::System& system, double unit_conversion); + // Build a per-atom TensorMap from values shaped as [atoms, components...]. + // A property dimension is appended automatically. + virtual metatensor_torch::TensorMap make_per_atom_tensormap( + const torch::Tensor& values, + torch::ScalarType dtype, + const torch::Device& device, + const std::vector& component_names = {} + ); // Explicit strain for virial calculations. This uses the same dtype/device // as LAMMPS data (positions, …) From ea83d148e92b13b12ace9ed3bac6c8448e43bdbf Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Sat, 28 Mar 2026 20:14:54 +0100 Subject: [PATCH 2/4] Give names to the property label of the requested input's tensormap --- src/ML-METATOMIC/metatomic_system.cpp | 133 +++++++++++++++++--------- src/ML-METATOMIC/metatomic_system.h | 21 ++-- 2 files changed, 100 insertions(+), 54 deletions(-) diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 265797ac789..3fe5bb615e4 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -206,6 +206,56 @@ static std::array cell_shifts( return {shift_a, shift_b, shift_c}; } +metatensor_torch::TensorMap LAMMPS_NS::make_per_atom_tensormap( + const torch::Tensor& values, + const torch::ScalarType& dtype, + const torch::Device& device, + const std::string& property_name, + const std::vector& component_names +) { + assert (values.dim() == static_cast(component_names.size() + 1)); + + auto n_atoms = values.size(0); + 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( + std::vector{"system", "atom"}, + samples_values, + metatensor::assume_unique{} + ); + + auto components = std::vector{}; + 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( + std::vector{component_names[axis]}, + component_values + ); + components.push_back(component); + } + + auto properties = torch::make_intrusive( + std::vector{property_name}, + torch::tensor({{0}}, label_tensor_options) + ); + auto block = torch::make_intrusive( + values.to(dtype).to(device).unsqueeze(-1), + samples, + components, + properties + ); + + return torch::make_intrusive( + keys, + std::vector{block} + ); +} + void MetatomicSystemAdaptor::guess_periodic_ghosts() { auto _ = MetatomicTimer("identifying periodic ghosts"); auto total_n_atoms = atom->nlocal + atom->nghost; @@ -589,6 +639,20 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp( this->setup_neighbors(system, list); + for (const auto& input: inputs) { + const auto& quantity = input->quantity().c_str(); + const auto& unit = input->unit().c_str(); + if (strcmp(quantity, "mass") == 0) { + add_masses(system, metatomic_torch::unit_conversion_factor(quantity, "u", unit)); + } else if (strcmp(quantity, "momentum") == 0) { + add_momenta(system, metatomic_torch::unit_conversion_factor(quantity, "u*A/ps", unit)); + } else if (strcmp(quantity, "velocity") == 0) { + add_velocities(system, metatomic_torch::unit_conversion_factor(quantity, "A/ps", unit)); + } else { + error->all(FLERR, "compute metatomic: the model requested an unsupported additional input of quantity '{}'", quantity); + } + } + return system; } @@ -628,7 +692,7 @@ void MetatomicSystemAdaptor::add_masses(metatomic_torch::System& system, double masses = masses.index_select(0, mta_to_lmp_tensor); masses = masses * unit_conversion; - auto tensor = make_per_atom_tensormap(masses, dtype, device); + auto tensor = make_per_atom_tensormap(masses, dtype, device, "mass"); system->add_data("masses", tensor, /*override=*/true); } @@ -663,60 +727,37 @@ void MetatomicSystemAdaptor::add_momenta(metatomic_torch::System& system, double momenta = momenta.index_select(0, mta_to_lmp_tensor); momenta = momenta * unit_conversion; - auto tensor = make_per_atom_tensormap(momenta, dtype, device, {"xyz"}); + auto tensor = make_per_atom_tensormap(momenta, dtype, device, "momentum", {"xyz"}); system->add_data("momenta", tensor, /*override=*/true); } -metatensor_torch::TensorMap MetatomicSystemAdaptor::make_per_atom_tensormap( - const torch::Tensor& values, - torch::ScalarType dtype, - const torch::Device& device, - const std::vector& component_names -) { - if (values.dim() != static_cast(component_names.size() + 1)) { - error->all( - FLERR, - "metatomic per-atom TensorMap helper expected rank {} tensor, got rank {}", - component_names.size() + 1, - values.dim() - ); - } +void MetatomicSystemAdaptor::add_velocities(metatomic_torch::System& system, double unit_conversion) { + double** v = atom->v; - auto n_atoms = values.size(0); - auto label_tensor_options = torch::TensorOptions().dtype(torch::kInt32).device(device); + auto total_n_atoms = atom->nlocal + atom->nghost; - 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( - std::vector{"system", "atom"}, - samples_values, - metatensor::assume_unique{} + auto device = system->device(); + auto dtype = system->scalar_type(); + + auto mta_to_lmp_tensor = torch::from_blob( + mta_to_lmp.data(), + {static_cast(mta_to_lmp.size())}, + torch::TensorOptions().dtype(torch::kInt).device(torch::kCPU) ); - auto components = std::vector{}; - 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( - std::vector{component_names[axis]}, - component_values - ); - components.push_back(component); + // gather momenta (per-atom) in a tensor and ship to device + torch::Tensor velocities = torch::zeros({total_n_atoms, 3}, torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU)); + auto velocities_accessor = velocities.accessor(); + for (int i=0; ito(device); - auto block = torch::make_intrusive( - values.to(dtype).to(device).unsqueeze(-1), - samples, - components, - properties - ); + velocities = velocities.index_select(0, mta_to_lmp_tensor); + velocities = velocities * unit_conversion; + auto tensor = make_per_atom_tensormap(velocities, dtype, device, "velocity", {"xyz"}); - return torch::make_intrusive( - keys, - std::vector{block} - ); + system->add_data("velocities", tensor, /*override=*/true); } diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index a7dbf7b0888..1f1e71deb9f 100644 --- a/src/ML-METATOMIC/metatomic_system.h +++ b/src/ML-METATOMIC/metatomic_system.h @@ -58,6 +58,16 @@ struct MetatomicNeighborsData { std::vector> distances_f32; }; +// Build a per-atom TensorMap from values shaped as [atoms, components...]. +// A property dimension is appended automatically. +metatensor_torch::TensorMap make_per_atom_tensormap( + const torch::Tensor& values, + const torch::ScalarType& dtype, + const torch::Device& device, + const std::string& property_name, + const std::vector& component_names = {} +); + class MetatomicSystemAdaptor : public Pointers { public: MetatomicSystemAdaptor(LAMMPS *lmp, MetatomicSystemOptions options); @@ -82,14 +92,9 @@ class MetatomicSystemAdaptor : public Pointers { // Add momenta as extra data to this system, only for atoms which are not // periodic images of other atoms virtual void add_momenta(metatomic_torch::System& system, double unit_conversion); - // Build a per-atom TensorMap from values shaped as [atoms, components...]. - // A property dimension is appended automatically. - virtual metatensor_torch::TensorMap make_per_atom_tensormap( - const torch::Tensor& values, - torch::ScalarType dtype, - const torch::Device& device, - const std::vector& component_names = {} - ); + // Add velocities as extra data to this system, only for atoms which are not + // periodic images of other atoms + virtual void add_velocities(metatomic_torch::System& system, double unit_conversion); // Explicit strain for virial calculations. This uses the same dtype/device // as LAMMPS data (positions, …) From d8c68524ec5fd3e8c788596d36fbd1019c450f72 Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Fri, 17 Apr 2026 13:47:37 +0200 Subject: [PATCH 3/4] Update the unit conversion logic --- src/ML-METATOMIC/fix_metatomic.cpp | 27 +++---- src/ML-METATOMIC/metatomic_system.cpp | 12 ++- src/ML-METATOMIC/metatomic_system.h | 3 +- src/ML-METATOMIC/metatomic_units.cpp | 110 ++++++++++++++++++++++++++ src/ML-METATOMIC/metatomic_units.h | 27 +++++++ src/ML-METATOMIC/pair_metatomic.cpp | 17 +--- 6 files changed, 160 insertions(+), 36 deletions(-) create mode 100644 src/ML-METATOMIC/metatomic_units.cpp create mode 100644 src/ML-METATOMIC/metatomic_units.h diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 2f302fc153f..1004f4ce68a 100644 --- a/src/ML-METATOMIC/fix_metatomic.cpp +++ b/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" @@ -62,25 +63,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= unit_map.at("energy").at(update->unit_style); + std::string length_unit = unit_map.at("position").at(update->unit_style); + std::string mass_unit = unit_map.at("mass").at(update->unit_style); + std::string velocity_unit = 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", momentum_unit, "(u*eV)^(1/2)"); if (narg < 4) { error->all(FLERR, @@ -454,7 +445,7 @@ void FixMetatomic::initial_integrate(int /*vflag*/) { ); // add the required additional inputs - this->system_adaptor->add_masses(system, 1.0); + this->system_adaptor->add_masses(system, metatomic_torch::unit_conversion_factor("mass", unit_map.at("mass").at(update->unit_style), "u")); this->system_adaptor->add_momenta(system, this->momentum_conversion_factor); // Configure selected atoms for evaluation diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 3fe5bb615e4..74aa1a2a5fe 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -16,11 +16,13 @@ ------------------------------------------------------------------------- */ #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" @@ -570,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::vector>& inputs ) { auto _ = MetatomicTimer("creating System from LAMMPS data"); @@ -643,11 +646,12 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp( const auto& quantity = input->quantity().c_str(); const auto& unit = input->unit().c_str(); if (strcmp(quantity, "mass") == 0) { - add_masses(system, metatomic_torch::unit_conversion_factor(quantity, "u", unit)); + add_masses(system, metatomic_torch::unit_conversion_factor(quantity, unit_map.at("mass").at(update->unit_style), unit)); } else if (strcmp(quantity, "momentum") == 0) { - add_momenta(system, metatomic_torch::unit_conversion_factor(quantity, "u*A/ps", unit)); + const auto& momentum_unit = unit_map.at("mass").at(update->unit_style) + "*" + unit_map.at("velocity").at(update->unit_style); + add_momenta(system, metatomic_torch::unit_conversion_factor(quantity, momentum_unit, unit)); } else if (strcmp(quantity, "velocity") == 0) { - add_velocities(system, metatomic_torch::unit_conversion_factor(quantity, "A/ps", unit)); + add_velocities(system, metatomic_torch::unit_conversion_factor(quantity, unit_map.at("velocity").at(update->unit_style), unit)); } else { error->all(FLERR, "compute metatomic: the model requested an unsupported additional input of quantity '{}'", quantity); } diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index 1f1e71deb9f..d0ec0dff583 100644 --- a/src/ML-METATOMIC/metatomic_system.h +++ b/src/ML-METATOMIC/metatomic_system.h @@ -83,7 +83,8 @@ class MetatomicSystemAdaptor : public Pointers { NeighList* list, bool do_virial, torch::ScalarType dtype, - torch::Device device + torch::Device device, + const std::vector>& inputs = {} ); // Add masses as extra data to this system, only for atoms which are not diff --git a/src/ML-METATOMIC/metatomic_units.cpp b/src/ML-METATOMIC/metatomic_units.cpp new file mode 100644 index 00000000000..bfe1aa30789 --- /dev/null +++ b/src/ML-METATOMIC/metatomic_units.cpp @@ -0,0 +1,110 @@ +/* ---------------------------------------------------------------------- + LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator + https://www.lammps.org/, Sandia National Laboratories + LAMMPS Development team: developers@lammps.org + + Copyright (2003) Sandia Corporation. Under the terms of Contract + DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains + certain rights in this software. This software is distributed under + the GNU General Public License. + + See the README file in the top-level LAMMPS directory. +------------------------------------------------------------------------- */ + +#include "metatomic_units.h" + +#include +#include + +namespace LAMMPS_NS { + +const std::map quantity_map = { + {"energy", "energy"}, + {"energy_ensemble", "energy"}, + {"energy_uncertainty", "energy"}, + {"forces", "force"}, + {"non_conservative_forces", "force"}, + {"stress", "stress"}, + {"non_conservative_stress", "stress"}, + {"masses", "mass"}, + {"position", "position"}, + {"momenta", "momentum"}, + {"velocities", "velocity"}, + {"charges", "charge"}, + {"heat_flux", "heat_flux"}, +}; + +const std::map> unit_map = { + {"mass", { + {"real", "g/mol"}, + {"metal", "g/mol"}, + {"si", "kg"}, + {"cgs", "g"}, + {"electron", "u"}, + {"micro", "pg"}, + {"nano", "ag"} + }}, + {"position", { + {"real", "A"}, + {"metal", "A"}, + {"si", "m"}, + {"cgs", "cm"}, + {"electron", "Bohr"}, + {"micro", "micrometer"}, + {"nano", "nanometer"} + }}, + {"energy", { + {"real", "kcal/mol"}, + {"metal", "eV"}, + {"si", "J"}, + {"cgs", "erg"}, + {"electron", "Hartree"}, + {"micro", "pg*micrometer^2/microsecond^2"}, + {"nano", "ag*nm^2/ns^2"} + }}, + {"velocity", { + {"real", "A/fs"}, + {"metal", "A/ps"}, + {"si", "m/s"}, + {"cgs", "cm/s"}, + {"electron", "Bohr*Hartree/hbar"}, + {"micro", "micrometer/microsecond"}, + {"nano", "nm/ns"} + }}, + {"force", { + {"real", "kcal/mol/A"}, + {"metal", "eV/A"}, + {"si", "kg*m/s^2"}, + {"cgs", "g*cm/s^2"}, + {"electron", "Hartree/Bohr"}, + {"micro", "pg*micrometer/microsecond^2"}, + {"nano", "ag*nm/ns^2"} + }}, + {"charge", { + {"real", "e"}, + {"metal", "e"}, + {"si", "C"}, + {"cgs", "esu"}, + {"electron", "e"}, + {"micro", "pC"}, + {"nano", "e"} + }}, + {"lmp::density", { + {"real", ""}, + {"metal", ""}, + {"si", ""}, + {"cgs", ""}, + {"electron", ""}, + {"micro", ""}, + {"nano", ""} + }}, + {"lmp::dipole_moment", {}}, + {"lmp::dynamic_viscosity", {}}, + {"lmp::electric_field", {}}, + {"lmp::pressure", {}}, + {"lmp::temperature", {}}, + {"lmp::time", {}}, + {"lmp::torque", {}}, +}; +// simple struct to hold unit conversion factors for a given unit style +} // namespace LAMMPS_NS diff --git a/src/ML-METATOMIC/metatomic_units.h b/src/ML-METATOMIC/metatomic_units.h new file mode 100644 index 00000000000..c8133e78ad8 --- /dev/null +++ b/src/ML-METATOMIC/metatomic_units.h @@ -0,0 +1,27 @@ +/* -*- c++ -*- ---------------------------------------------------------- + LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator + https://www.lammps.org/, Sandia National Laboratories + LAMMPS development team: developers@lammps.org + + Copyright (2003) Sandia Corporation. Under the terms of Contract + DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains + certain rights in this software. This software is distributed under + the GNU General Public License. + + See the README file in the top-level LAMMPS directory. +------------------------------------------------------------------------- */ + +#ifndef LMP_METATOMIC_UNITS_H +#define LMP_METATOMIC_UNITS_H + +#include +#include + +namespace LAMMPS_NS { + +extern const std::map quantity_map; + +extern const std::map> unit_map; +} // namespace LAMMPS_NS + +#endif \ No newline at end of file diff --git a/src/ML-METATOMIC/pair_metatomic.cpp b/src/ML-METATOMIC/pair_metatomic.cpp index bd067026e86..f0dc265a350 100644 --- a/src/ML-METATOMIC/pair_metatomic.cpp +++ b/src/ML-METATOMIC/pair_metatomic.cpp @@ -44,6 +44,7 @@ #include "metatomic_system.h" #include "metatomic_timer.h" +#include "metatomic_units.h" using namespace LAMMPS_NS; @@ -72,21 +73,11 @@ PairMetatomic::PairMetatomic(LAMMPS *lmp): system_adaptor(nullptr), scale(1.0) { - if (strcmp(update->unit_style, "real") == 0) { - this->length_unit = "angstrom"; - this->energy_unit = "kcal/mol"; - } else if (strcmp(update->unit_style, "metal") == 0) { - this->length_unit = "angstrom"; - this->energy_unit = "eV"; - } else if (strcmp(update->unit_style, "si") == 0) { - this->length_unit = "meter"; - this->energy_unit = "joule"; - } else if (strcmp(update->unit_style, "electron") == 0) { - this->length_unit = "Bohr"; - this->energy_unit = "Hartree"; - } else { + if (strcmp(update->unit_style, "lj") == 0) { error->one(FLERR, "unsupported units '{}' for pair metatomic ", update->unit_style); } + this->length_unit = unit_map.at("position").at(update->unit_style); + this->energy_unit = unit_map.at("energy").at(update->unit_style); // we might not be running a pure pair potential, // so we can not compute virial as fdotr From 1ed3a33c80c389b966418e7132537d2aa512c7df Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Sat, 18 Apr 2026 11:09:52 +0200 Subject: [PATCH 4/4] Allow for additional inputs, and update to the 2-parameter `unit_conversion_factor` --- src/ML-METATOMIC/fix_metatomic.cpp | 16 +++++++++++++--- src/ML-METATOMIC/metatomic_system.cpp | 21 +++++++++++---------- src/ML-METATOMIC/metatomic_system.h | 2 +- src/ML-METATOMIC/metatomic_units.cpp | 16 ---------------- src/ML-METATOMIC/metatomic_units.h | 3 +-- src/ML-METATOMIC/pair_metatomic.cpp | 17 +++++++++++++++-- 6 files changed, 41 insertions(+), 34 deletions(-) diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 1004f4ce68a..62c0e6687c0 100644 --- a/src/ML-METATOMIC/fix_metatomic.cpp +++ b/src/ML-METATOMIC/fix_metatomic.cpp @@ -71,7 +71,7 @@ FixMetatomic::FixMetatomic(LAMMPS *lmp, int narg, char **arg): Fix(lmp, narg, ar std::string mass_unit = unit_map.at("mass").at(update->unit_style); std::string velocity_unit = 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", momentum_unit, "(u*eV)^(1/2)"); + this->momentum_conversion_factor = metatomic_torch::unit_conversion_factor(momentum_unit, "(u*eV)^(1/2)"); if (narg < 4) { error->all(FLERR, @@ -436,16 +436,26 @@ 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 input_holders; + auto requested_inputs = mta_data->model->run_method("requested_inputs").toGenericDict(); + for (const auto& entry : requested_inputs) { + input_holders.emplace( + entry.key().toStringRef(), + entry.value().toCustomClass() + ); + } // transform from LAMMPS to metatomic System auto system = this->system_adaptor->system_from_lmp( mta_list, static_cast(vflag_global), dtype, - mta_data->device + mta_data->device, + input_holders ); // add the required additional inputs - this->system_adaptor->add_masses(system, metatomic_torch::unit_conversion_factor("mass", unit_map.at("mass").at(update->unit_style), "u")); + this->system_adaptor->add_masses(system, metatomic_torch::unit_conversion_factor(unit_map.at("mass").at(update->unit_style), "u")); this->system_adaptor->add_momenta(system, this->momentum_conversion_factor); // Configure selected atoms for evaluation diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 74aa1a2a5fe..7d0b6b6b7ad 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -26,6 +26,7 @@ #include "neigh_list.h" +#include #include #include @@ -573,7 +574,7 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp( bool do_virial, torch::ScalarType dtype, torch::Device device, - const std::vector>& inputs + const std::map>& inputs ) { auto _ = MetatomicTimer("creating System from LAMMPS data"); @@ -642,18 +643,18 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp( this->setup_neighbors(system, list); - for (const auto& input: inputs) { - const auto& quantity = input->quantity().c_str(); + for (const auto& [property, input]: inputs) { + const auto& property_name = property.c_str(); const auto& unit = input->unit().c_str(); - if (strcmp(quantity, "mass") == 0) { - add_masses(system, metatomic_torch::unit_conversion_factor(quantity, unit_map.at("mass").at(update->unit_style), unit)); - } else if (strcmp(quantity, "momentum") == 0) { + if (strcmp(property_name, "masses") == 0) { + add_masses(system, metatomic_torch::unit_conversion_factor(unit_map.at("mass").at(update->unit_style), unit)); + } else if (strcmp(property_name, "momenta") == 0) { const auto& momentum_unit = unit_map.at("mass").at(update->unit_style) + "*" + unit_map.at("velocity").at(update->unit_style); - add_momenta(system, metatomic_torch::unit_conversion_factor(quantity, momentum_unit, unit)); - } else if (strcmp(quantity, "velocity") == 0) { - add_velocities(system, metatomic_torch::unit_conversion_factor(quantity, unit_map.at("velocity").at(update->unit_style), unit)); + add_momenta(system, metatomic_torch::unit_conversion_factor(momentum_unit, unit)); + } else if (strcmp(property_name, "velocities") == 0) { + add_velocities(system, metatomic_torch::unit_conversion_factor(unit_map.at("velocity").at(update->unit_style), unit)); } else { - error->all(FLERR, "compute metatomic: the model requested an unsupported additional input of quantity '{}'", quantity); + error->all(FLERR, "compute metatomic: the model requested an unsupported additional input of '{}'", property_name); } } diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index d0ec0dff583..41f29ebd712 100644 --- a/src/ML-METATOMIC/metatomic_system.h +++ b/src/ML-METATOMIC/metatomic_system.h @@ -84,7 +84,7 @@ class MetatomicSystemAdaptor : public Pointers { bool do_virial, torch::ScalarType dtype, torch::Device device, - const std::vector>& inputs = {} + const std::map>& inputs = {} ); // Add masses as extra data to this system, only for atoms which are not diff --git a/src/ML-METATOMIC/metatomic_units.cpp b/src/ML-METATOMIC/metatomic_units.cpp index bfe1aa30789..1806889d50d 100644 --- a/src/ML-METATOMIC/metatomic_units.cpp +++ b/src/ML-METATOMIC/metatomic_units.cpp @@ -18,22 +18,6 @@ namespace LAMMPS_NS { -const std::map quantity_map = { - {"energy", "energy"}, - {"energy_ensemble", "energy"}, - {"energy_uncertainty", "energy"}, - {"forces", "force"}, - {"non_conservative_forces", "force"}, - {"stress", "stress"}, - {"non_conservative_stress", "stress"}, - {"masses", "mass"}, - {"position", "position"}, - {"momenta", "momentum"}, - {"velocities", "velocity"}, - {"charges", "charge"}, - {"heat_flux", "heat_flux"}, -}; - const std::map> unit_map = { {"mass", { {"real", "g/mol"}, diff --git a/src/ML-METATOMIC/metatomic_units.h b/src/ML-METATOMIC/metatomic_units.h index c8133e78ad8..f38aec82209 100644 --- a/src/ML-METATOMIC/metatomic_units.h +++ b/src/ML-METATOMIC/metatomic_units.h @@ -19,9 +19,8 @@ namespace LAMMPS_NS { -extern const std::map quantity_map; - extern const std::map> unit_map; + } // namespace LAMMPS_NS #endif \ No newline at end of file diff --git a/src/ML-METATOMIC/pair_metatomic.cpp b/src/ML-METATOMIC/pair_metatomic.cpp index f0dc265a350..d87faddc902 100644 --- a/src/ML-METATOMIC/pair_metatomic.cpp +++ b/src/ML-METATOMIC/pair_metatomic.cpp @@ -37,7 +37,9 @@ #include #endif +#include #include +#include #include #include @@ -85,7 +87,7 @@ PairMetatomic::PairMetatomic(LAMMPS *lmp): this->mta_data = new PairMetatomicData(this->length_unit); // use a default uncertainty threshold of 100 meV/atom - this->mta_data->uncertainty_threshold = 0.1 * metatomic_torch::unit_conversion_factor("energy", "eV", energy_unit); + this->mta_data->uncertainty_threshold = 0.1 * metatomic_torch::unit_conversion_factor("eV", energy_unit); // settings for metatomic pair style this->single_enable = 0; @@ -623,12 +625,23 @@ void PairMetatomic::compute(int eflag, int vflag) { error->one(FLERR, "the model requested an unsupported dtype '{}'", mta_data->capabilities->dtype()); } + // deal with the model requested inputs + std::map input_holders; + auto requested_inputs = mta_data->model->run_method("requested_inputs").toGenericDict(); + for (const auto& entry : requested_inputs) { + input_holders.emplace( + entry.key().toStringRef(), + entry.value().toCustomClass() + ); + } + // transform from LAMMPS to metatomic System auto system = this->system_adaptor->system_from_lmp( mta_list, static_cast(vflag_global), dtype, - mta_data->device + mta_data->device, + input_holders ); // only run the calculation for atoms actually in the current domain