diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 2f302fc153f..62c0e6687c0 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_unit, "(u*eV)^(1/2)"); if (narg < 4) { error->all(FLERR, @@ -445,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, 1.0); + 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 4652449d3b9..7d0b6b6b7ad 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/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 #include #include @@ -206,6 +209,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; @@ -520,7 +573,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>& inputs ) { auto _ = MetatomicTimer("creating System from LAMMPS data"); @@ -589,6 +643,21 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp( this->setup_neighbors(system, list); + for (const auto& [property, input]: inputs) { + const auto& property_name = property.c_str(); + const auto& unit = input->unit().c_str(); + 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(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 '{}'", property_name); + } + } + return system; } @@ -628,29 +697,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, "mass"); system->add_data("masses", tensor, /*override=*/true); } @@ -685,37 +732,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, "momentum", {"xyz"}); - auto keys = metatensor_torch::LabelsHolder::single()->to(device); + system->add_data("momenta", tensor, /*override=*/true); +} - 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{} - ); +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( - std::vector{"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( - momenta.to(dtype).to(device).unsqueeze(-1), - samples, - std::vector{component}, - properties - ); - auto tensor = torch::make_intrusive( - keys, - std::vector{block} + 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) ); - system->add_data("momenta", tensor, /*override=*/true); + // 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; iadd_data("velocities", tensor, /*override=*/true); } diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index 38fa9ed62d8..41f29ebd712 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" @@ -57,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); @@ -72,7 +83,8 @@ class MetatomicSystemAdaptor : public Pointers { NeighList* list, bool do_virial, torch::ScalarType dtype, - torch::Device device + torch::Device device, + const std::map>& inputs = {} ); // Add masses as extra data to this system, only for atoms which are not @@ -81,6 +93,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); + // 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, …) diff --git a/src/ML-METATOMIC/metatomic_units.cpp b/src/ML-METATOMIC/metatomic_units.cpp new file mode 100644 index 00000000000..1806889d50d --- /dev/null +++ b/src/ML-METATOMIC/metatomic_units.cpp @@ -0,0 +1,94 @@ +/* ---------------------------------------------------------------------- + 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> 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..f38aec82209 --- /dev/null +++ b/src/ML-METATOMIC/metatomic_units.h @@ -0,0 +1,26 @@ +/* -*- 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> 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..d87faddc902 100644 --- a/src/ML-METATOMIC/pair_metatomic.cpp +++ b/src/ML-METATOMIC/pair_metatomic.cpp @@ -37,13 +37,16 @@ #include #endif +#include #include +#include #include #include #include "metatomic_system.h" #include "metatomic_timer.h" +#include "metatomic_units.h" using namespace LAMMPS_NS; @@ -72,21 +75,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 @@ -94,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; @@ -632,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