diff --git a/cmake/Modules/Packages/ML-METATOMIC.cmake b/cmake/Modules/Packages/ML-METATOMIC.cmake index 56de4e4d06c..eb902e77c3b 100644 --- a/cmake/Modules/Packages/ML-METATOMIC.cmake +++ b/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) diff --git a/examples/PACKAGES/metatomic/nickel-lj.pt b/examples/PACKAGES/metatomic/nickel-lj.pt index cf2f830039c..f395ef98943 100644 Binary files a/examples/PACKAGES/metatomic/nickel-lj.pt and b/examples/PACKAGES/metatomic/nickel-lj.pt differ diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 1bb01b7f344..8db02f5379f 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" @@ -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 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() + ); + } // 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, 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}); diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 11b56b40bbf..8376623d082 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,55 @@ static std::array cell_shifts( return {shift_a, shift_b, shift_c}; } +static metatensor_torch::TensorMap make_per_atom_tensormap( + const torch::Tensor& values, + 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 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( + 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.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 +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>& 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(); - for (int i=0; i(); + for (int i=0; i(); + for (int i=0; i(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( - 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} - ); - - 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(); - for (int i=0; i(); + for (int i=0; i(); + for (int i=0; i(m * v[i][0]); + momenta_accessor[i][1] = static_cast(m * v[i][1]); + momenta_accessor[i][2] = static_cast(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( - 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 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 tensor = torch::make_intrusive( - keys, - std::vector{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(); + for (int i=0; i(); + for (int i=0; i(v[i][0]); + velocities_accessor[i][1] = static_cast(v[i][1]); + velocities_accessor[i][2] = static_cast(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); } diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index 98e093955fe..45c594299ec 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" @@ -72,15 +73,19 @@ class MetatomicSystemAdaptor : public Pointers { NeighList* list, bool do_virial, torch::ScalarType dtype, - torch::Device device + torch::Device device, + const std::map>& requested_inputs = {} ); // Add masses as extra data to this system, only for atoms which are not // periodic images of other atoms - virtual void add_masses(metatomic_torch::System& system, std::string name, double unit_conversion); + virtual void add_masses(metatomic_torch::System& system, double unit_conversion); // 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, std::string name, double unit_conversion); + 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..24e53d32c1f --- /dev/null +++ b/src/ML-METATOMIC/metatomic_units.cpp @@ -0,0 +1,78 @@ +/* ---------------------------------------------------------------------- + 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> metatomic_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"} + }}, +}; +// 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..e2f73ab9dcc --- /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> metatomic_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 48051039401..66a8bdd679e 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 = metatomic_unit_map.at("position").at(update->unit_style); + this->energy_unit = metatomic_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 @@ -678,12 +671,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", /*use_new_names=*/ true).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, vflag_global && !do_nc_stress, dtype, - mta_data->device + mta_data->device, + input_holders ); // only run the calculation for atoms actually in the current domain