From 9fda21c93b7dc2adb2e7b3f58bf5800d89011a3a Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Fri, 10 Jul 2026 19:44:59 +0200 Subject: [PATCH 1/5] Extract common logics to helper functions --- src/ML-METATOMIC/compute_metatomic.cpp | 108 +++---------------------- src/ML-METATOMIC/compute_metatomic.h | 1 - src/ML-METATOMIC/fix_metatomic.cpp | 97 ++-------------------- src/ML-METATOMIC/fix_metatomic.h | 2 - src/ML-METATOMIC/metatomic_types.cpp | 101 +++++++++++++++++++++++ src/ML-METATOMIC/metatomic_types.h | 16 ++++ src/ML-METATOMIC/pair_metatomic.cpp | 91 +-------------------- src/ML-METATOMIC/pair_metatomic.h | 4 - 8 files changed, 139 insertions(+), 281 deletions(-) diff --git a/src/ML-METATOMIC/compute_metatomic.cpp b/src/ML-METATOMIC/compute_metatomic.cpp index 8eea2efbbe6..702a50f76bf 100644 --- a/src/ML-METATOMIC/compute_metatomic.cpp +++ b/src/ML-METATOMIC/compute_metatomic.cpp @@ -251,16 +251,8 @@ ComputeMetatomic::ComputeMetatomic(LAMMPS *lmp, int narg, char **arg): Compute(l ); auto capabilities = mta_data->model->run_method("capabilities").toCustomClass(); - c10::ScalarType dtype; - if (capabilities->dtype() == "float64") { - dtype = torch::kFloat64; - } else if (capabilities->dtype() == "float32") { - dtype = torch::kFloat32; - } else { - error->all(FLERR, - "the model requested an unsupported dtype '" + capabilities->dtype() + "'" - ); - } + // validate that the model requests a supported dtype + mta_data->model_dtype(lmp); auto model_outputs = capabilities->outputs(); if (!model_outputs.contains(this->output_name)) { error->all(FLERR, @@ -278,14 +270,8 @@ ComputeMetatomic::ComputeMetatomic(LAMMPS *lmp, int narg, char **arg): Compute(l this->mta_data->evaluation_options->outputs.insert(this->output_name, this->mta_data->requested_output); // add the required additional inputs - auto requested_inputs = mta_data->model->run_method("requested_inputs", /*use_new_names=*/ true).toGenericDict(); - for (const auto& entry : requested_inputs) { - mta_data->requested_inputs.emplace( - entry.key().toStringRef(), - entry.value().toCustomClass() - ); - } - + mta_data->requested_inputs = mta_data->collect_requested_inputs(); + // Initialize the output layout if (strcmp(sample_kind.c_str(), "atom") == 0) { peratom_flag = 1; @@ -316,9 +302,10 @@ ComputeMetatomic::ComputeMetatomic(LAMMPS *lmp, int narg, char **arg): Compute(l // Select the device to use based on the model's preference, the user choice // and what's available. - this->pick_device( - mta_data->device, - this->requested_device ? this->requested_device->c_str() : nullptr + mta_data->pick_device( + lmp, + this->requested_device ? this->requested_device->c_str() : nullptr, + "compute metatomic" ); // move all data to the correct device @@ -356,27 +343,7 @@ void ComputeMetatomic::init() { } // get the model's interaction range - auto range = mta_data->capabilities->engine_interaction_range(mta_data->evaluation_options->length_unit()); - if (range < 0) { - error->all(FLERR, "interaction_range is negative for this model"); - } else if (!std::isfinite(range)) { - if (comm->nprocs > 1) { - error->all(FLERR, - "interaction_range is infinite for this model, " - "using multiple MPI domains is not supported" - ); - } - - // determine the maximal cutoff in the NL - auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); - for (const auto& ivalue: requested_nl.toList()) { - auto options = ivalue.get().toCustomClass(); - auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); - mta_data->max_cutoff = std::max(mta_data->max_cutoff, cutoff); - } - } else { - mta_data->max_cutoff = range; - } + mta_data->resolve_max_cutoff(lmp); // Initialize metatensor system object auto options = MetatomicSystemOptions{ @@ -427,54 +394,6 @@ void ComputeMetatomic::init() { // END HACK } -void ComputeMetatomic::pick_device(c10::Device& device, const char* requested) { - torch::optional requested_string; - torch::DeviceType device_type; - - if (requested != nullptr) { - requested_string = std::string(requested); - } else { - requested_string = torch::nullopt; - } - - try { - device_type = metatomic_torch::pick_device( - this->mta_data->capabilities->supported_devices, - requested_string - ); - } catch (const c10::Error& e) { - error->one(FLERR, "compute metatomic: {}", e.what()); - } - - if (device_type == torch::DeviceType::CUDA) { - // distribute GPUs between multiple MPI processes on the same node - - // (1) get a MPI communicator for all processes on the current node - MPI_Comm local; - MPI_Comm_split_type(world, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local); - // (2) get the rank of this MPI process on the current node - int local_rank; - MPI_Comm_rank(local, &local_rank); - - int size; - MPI_Comm_size(local, &size); - if (size < torch::cuda::device_count()) { - if (comm->me == 0) { - error->warning(FLERR, - "found {} CUDA-capable GPUs, but only {} MPI processes on the current node; the remaining GPUs will not be used", - torch::cuda::device_count(), size - ); - } - } - - // (3) split GPUs between node-local processes using round-robin allocation - auto device_index = local_rank % torch::cuda::device_count(); - device = torch::Device(device_type, static_cast(device_index)); - } else { - device = torch::Device(device_type); - } -} - void ComputeMetatomic::init_list(int id, NeighList *ptr) { mta_list = ptr; } @@ -489,14 +408,7 @@ void ComputeMetatomic::compute() { int *mask = atom->mask; // Determine the dtype of the system based on the model's capabilities - auto dtype = torch::kFloat64; - if (mta_data->capabilities->dtype() == "float64") { - dtype = torch::kFloat64; - } else if (mta_data->capabilities->dtype() == "float32") { - dtype = torch::kFloat32; - } else { - error->all(FLERR, "the model requested an unsupported dtype '{}'", mta_data->capabilities->dtype()); - } + auto dtype = mta_data->model_dtype(lmp); auto system = this->system_adaptor->system_from_lmp( mta_list, diff --git a/src/ML-METATOMIC/compute_metatomic.h b/src/ML-METATOMIC/compute_metatomic.h index 382aa9d049b..3dab2a48b05 100644 --- a/src/ML-METATOMIC/compute_metatomic.h +++ b/src/ML-METATOMIC/compute_metatomic.h @@ -47,7 +47,6 @@ class ComputeMetatomic : public Compute { void compute_peratom() override; protected: - virtual void pick_device(c10::Device& device, const char* requested); std::string model_path; std::optional requested_device; std::optional extensions_directory; diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 0e8df144541..6f9df7f4eb5 100644 --- a/src/ML-METATOMIC/fix_metatomic.cpp +++ b/src/ML-METATOMIC/fix_metatomic.cpp @@ -249,9 +249,10 @@ void FixMetatomic::init() { // Select the device to use based on the model's preference, the user choice // and what's available. - this->pick_device( - mta_data->device, - this->requested_device ? this->requested_device->c_str() : nullptr + mta_data->pick_device( + lmp, + this->requested_device ? this->requested_device->c_str() : nullptr, + "fix metatomic" ); // move all data to the correct device @@ -267,28 +268,7 @@ void FixMetatomic::init() { } // get the model's interaction range - auto range = mta_data->capabilities->engine_interaction_range(mta_data->evaluation_options->length_unit()); - if (range < 0) { - error->all(FLERR, "interaction_range is negative for this model"); - } else if (!std::isfinite(range)) { - if (comm->nprocs > 1) { - error->all(FLERR, - "interaction_range is infinite for this model, " - "using multiple MPI domains is not supported" - ); - } - - // determine the maximal cutoff in the NL - auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); - for (const auto& ivalue: requested_nl.toList()) { - auto options = ivalue.get().toCustomClass(); - auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); - - mta_data->max_cutoff = std::max(mta_data->max_cutoff, cutoff); - } - } else { - mta_data->max_cutoff = range; - } + mta_data->resolve_max_cutoff(lmp); // Initialize metatensor system object auto options = MetatomicSystemOptions{ @@ -339,54 +319,6 @@ void FixMetatomic::init() { // END HACK } -void FixMetatomic::pick_device(c10::Device& device, const char* requested) { - torch::optional requested_string; - torch::DeviceType device_type; - - if (requested != nullptr) { - requested_string = std::string(requested); - } else { - requested_string = torch::nullopt; - } - - try { - device_type = metatomic_torch::pick_device( - this->mta_data->capabilities->supported_devices, - requested_string - ); - } catch (const c10::Error& e) { - error->one(FLERR, "fix metatomic: {}", e.what()); - } - - if (device_type == torch::DeviceType::CUDA) { - // distribute GPUs between multiple MPI processes on the same node - - // (1) get a MPI communicator for all processes on the current node - MPI_Comm local; - MPI_Comm_split_type(world, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local); - // (2) get the rank of this MPI process on the current node - int local_rank; - MPI_Comm_rank(local, &local_rank); - - int size; - MPI_Comm_size(local, &size); - if (size < torch::cuda::device_count()) { - if (comm->me == 0) { - error->warning(FLERR, - "found {} CUDA-capable GPUs, but only {} MPI processes on the current node; the remaining GPUs will not be used", - torch::cuda::device_count(), size - ); - } - } - - // (3) split GPUs between node-local processes using round-robin allocation - auto device_index = local_rank % torch::cuda::device_count(); - device = torch::Device(device_type, static_cast(device_index)); - } else { - device = torch::Device(device_type); - } -} - void FixMetatomic::init_list(int id, NeighList *ptr) { mta_list = ptr; } @@ -426,24 +358,11 @@ void FixMetatomic::initial_integrate(int /*vflag*/) { } } - auto dtype = torch::kFloat64; - if (mta_data->capabilities->dtype() == "float64") { - dtype = torch::kFloat64; - } else if (mta_data->capabilities->dtype() == "float32") { - dtype = torch::kFloat32; - } else { - error->all(FLERR, "the model requested an unsupported dtype '{}'", mta_data->capabilities->dtype()); - } + auto dtype = mta_data->model_dtype(lmp); // 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() - ); - } + auto input_holders = mta_data->collect_requested_inputs(); + // transform from LAMMPS to metatomic System auto system = this->system_adaptor->system_from_lmp( mta_list, diff --git a/src/ML-METATOMIC/fix_metatomic.h b/src/ML-METATOMIC/fix_metatomic.h index 5c7e87048e4..fd70235e936 100644 --- a/src/ML-METATOMIC/fix_metatomic.h +++ b/src/ML-METATOMIC/fix_metatomic.h @@ -54,8 +54,6 @@ class FixMetatomic : public Fix { void init_list(int id, NeighList *ptr) override; protected: - virtual void pick_device(c10::Device& device, const char* requested); - double momentum_conversion_factor; // Conversion factor for momenta double dt; // Timestep std::string model_path; // Path to ML model file diff --git a/src/ML-METATOMIC/metatomic_types.cpp b/src/ML-METATOMIC/metatomic_types.cpp index 6ed5d910db4..cc77a759d65 100644 --- a/src/ML-METATOMIC/metatomic_types.cpp +++ b/src/ML-METATOMIC/metatomic_types.cpp @@ -20,6 +20,11 @@ #include "comm.h" #include "error.h" +#include +#include + +#include + using namespace LAMMPS_NS; CommonMetatomicData::CommonMetatomicData(std::string length_unit): device(torch::kCPU) { @@ -82,3 +87,99 @@ void CommonMetatomicData::load_model( } } } + +void CommonMetatomicData::pick_device(LAMMPS* lmp, const char* requested, const char* cmd_name) { + torch::optional requested_string; + if (requested != nullptr) { + requested_string = std::string(requested); + } else { + requested_string = torch::nullopt; + } + + torch::DeviceType device_type; + try { + device_type = metatomic_torch::pick_device( + this->capabilities->supported_devices, + requested_string + ); + } catch (const c10::Error& e) { + lmp->error->one(FLERR, "{}: {}", cmd_name, e.what()); + } + + if (device_type == torch::DeviceType::CUDA) { + // distribute GPUs between multiple MPI processes on the same node + + // (1) get a MPI communicator for all processes on the current node + MPI_Comm local; + MPI_Comm_split_type(lmp->world, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local); + // (2) get the rank of this MPI process on the current node + int local_rank; + MPI_Comm_rank(local, &local_rank); + + int size; + MPI_Comm_size(local, &size); + if (size < torch::cuda::device_count()) { + if (lmp->comm->me == 0) { + lmp->error->warning(FLERR, + "found {} CUDA-capable GPUs, but only {} MPI processes on the current node; the remaining GPUs will not be used", + torch::cuda::device_count(), size + ); + } + } + + // (3) split GPUs between node-local processes using round-robin allocation + auto device_index = local_rank % torch::cuda::device_count(); + this->device = torch::Device(device_type, static_cast(device_index)); + } else { + this->device = torch::Device(device_type); + } +} + +c10::ScalarType CommonMetatomicData::model_dtype(LAMMPS* lmp) const { + const auto dtype = this->capabilities->dtype(); + if (dtype == "float64") { + return torch::kFloat64; + } else if (dtype == "float32") { + return torch::kFloat32; + } else { + lmp->error->all(FLERR, "the model requested an unsupported dtype '{}'", dtype); + } + return torch::kFloat64; // unreachable, error->all does not return +} + +void CommonMetatomicData::resolve_max_cutoff(LAMMPS* lmp) { + // get the model's interaction range + auto range = this->capabilities->engine_interaction_range(this->evaluation_options->length_unit()); + if (range < 0) { + lmp->error->all(FLERR, "interaction_range is negative for this model"); + } else if (!std::isfinite(range)) { + if (lmp->comm->nprocs > 1) { + lmp->error->all(FLERR, + "interaction_range is infinite for this model, " + "using multiple MPI domains is not supported" + ); + } + + // determine the maximal cutoff in the NL + auto requested_nl = this->model->run_method("requested_neighbor_lists"); + for (const auto& ivalue: requested_nl.toList()) { + auto options = ivalue.get().toCustomClass(); + auto cutoff = options->engine_cutoff(this->evaluation_options->length_unit()); + this->max_cutoff = std::max(this->max_cutoff, cutoff); + } + } else { + this->max_cutoff = range; + } +} + +std::map CommonMetatomicData::collect_requested_inputs() const { + std::map input_holders; + auto requested_inputs = this->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() + ); + } + return input_holders; +} diff --git a/src/ML-METATOMIC/metatomic_types.h b/src/ML-METATOMIC/metatomic_types.h index 7eb98552b89..93b83aa87d9 100644 --- a/src/ML-METATOMIC/metatomic_types.h +++ b/src/ML-METATOMIC/metatomic_types.h @@ -13,6 +13,7 @@ #include "lammps.h" +#include #include #include @@ -29,6 +30,21 @@ struct CommonMetatomicData { CommonMetatomicData(std::string length_unit); void load_model(LAMMPS* lmp, const char* path, const char* extensions_directory); + // pick the compute device from the model's supported devices and the user + // request (or nullptr), storing the result in `this->device`. `cmd_name` is + // used as a prefix in error messages (e.g. "pair_style metatomic"). + void pick_device(LAMMPS* lmp, const char* requested, const char* cmd_name); + + // resolve the torch dtype (float32/float64) requested by the model + c10::ScalarType model_dtype(LAMMPS* lmp) const; + + // compute `this->max_cutoff` from the model's interaction range, falling + // back to the requested neighbor lists when the range is infinite + void resolve_max_cutoff(LAMMPS* lmp); + + // collect the additional inputs requested by the model + std::map collect_requested_inputs() const; + // the metatomic model std::unique_ptr model; // the path used to load the model diff --git a/src/ML-METATOMIC/pair_metatomic.cpp b/src/ML-METATOMIC/pair_metatomic.cpp index 4fd5eab7f0d..09982596b49 100644 --- a/src/ML-METATOMIC/pair_metatomic.cpp +++ b/src/ML-METATOMIC/pair_metatomic.cpp @@ -400,7 +400,7 @@ void PairMetatomic::settings(int argc, char ** argv) { // Select the device to use based on the model's preference, the user choice // and what's available. - this->pick_device(mta_data->device, requested_device); + mta_data->pick_device(lmp, requested_device, "pair_style metatomic"); // move all data to the correct device mta_data->model->to(mta_data->device); @@ -419,54 +419,6 @@ void PairMetatomic::settings(int argc, char ** argv) { } } -void PairMetatomic::pick_device(torch::Device& device, const char* requested) { - torch::optional requested_string; - torch::DeviceType device_type; - - if (requested != nullptr) { - requested_string = std::string(requested); - } else { - requested_string = torch::nullopt; - } - - try { - device_type = metatomic_torch::pick_device( - this->mta_data->capabilities->supported_devices, - requested_string - ); - } catch (const c10::Error& e) { - error->one(FLERR, "pair_style metatomic: {}", e.what()); - } - - if (device_type == torch::DeviceType::CUDA) { - // distribute GPUs between multiple MPI processes on the same node - - // (1) get a MPI communicator for all processes on the current node - MPI_Comm local; - MPI_Comm_split_type(world, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local); - // (2) get the rank of this MPI process on the current node - int local_rank; - MPI_Comm_rank(local, &local_rank); - - int size; - MPI_Comm_size(local, &size); - if (size < torch::cuda::device_count()) { - if (comm->me == 0) { - error->warning(FLERR, - "found {} CUDA-capable GPUs, but only {} MPI processes on the current node; the remaining GPUs will not be used", - torch::cuda::device_count(), size - ); - } - } - - // (3) split GPUs between node-local processes using round-robin allocation - auto device_index = local_rank % torch::cuda::device_count(); - device = torch::Device(device_type, static_cast(device_index)); - } else { - device = torch::Device(device_type); - } -} - void PairMetatomic::allocate() { allocated = 1; @@ -550,28 +502,7 @@ void PairMetatomic::init_style() { } // get the model's interaction range - auto range = mta_data->capabilities->engine_interaction_range(mta_data->evaluation_options->length_unit()); - if (range < 0) { - error->one(FLERR, "interaction_range is negative for this model"); - } else if (!std::isfinite(range)) { - if (comm->nprocs > 1) { - error->one(FLERR, - "interaction_range is infinite for this model, " - "using multiple MPI domains is not supported" - ); - } - - // determine the maximal cutoff in the NL - auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); - for (const auto& ivalue: requested_nl.toList()) { - auto options = ivalue.get().toCustomClass(); - auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); - - mta_data->max_cutoff = std::max(mta_data->max_cutoff, cutoff); - } - } else { - mta_data->max_cutoff = range; - } + mta_data->resolve_max_cutoff(lmp); if (!std::isfinite(mta_data->max_cutoff)) { error->one(FLERR, @@ -662,24 +593,10 @@ void PairMetatomic::compute(int eflag, int vflag) { mta_data->evaluation_options->outputs.insert(mta_data->nc_stress_key, mta_data->nc_stress_output); } - auto dtype = torch::kFloat64; - if (mta_data->capabilities->dtype() == "float64") { - dtype = torch::kFloat64; - } else if (mta_data->capabilities->dtype() == "float32") { - dtype = torch::kFloat32; - } else { - error->one(FLERR, "the model requested an unsupported dtype '{}'", mta_data->capabilities->dtype()); - } + auto dtype = mta_data->model_dtype(lmp); // 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() - ); - } + auto input_holders = mta_data->collect_requested_inputs(); // transform from LAMMPS to metatomic System auto system = this->system_adaptor->system_from_lmp( diff --git a/src/ML-METATOMIC/pair_metatomic.h b/src/ML-METATOMIC/pair_metatomic.h index 8b4d85835ff..dab5d5ef39e 100644 --- a/src/ML-METATOMIC/pair_metatomic.h +++ b/src/ML-METATOMIC/pair_metatomic.h @@ -64,10 +64,6 @@ class PairMetatomic : public Pair { virtual void store_forces(const at::Tensor& forces_tensor); protected: - // pick the correct device to use from the user request (or nullptr) in - // `pair_style metatomic` - virtual void pick_device(c10::Device& device, const char* requested); - PairMetatomicData* mta_data; NeighList *mta_list; From 764221c7553dc782253cc5d305dbe54625785330 Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Sun, 12 Jul 2026 16:51:00 +0200 Subject: [PATCH 2/5] Replace the atom selection logic with `CommonMetatomicData::set_selected_atoms` --- src/ML-METATOMIC/compute_metatomic.cpp | 17 +-------------- src/ML-METATOMIC/fix_metatomic.cpp | 19 +--------------- src/ML-METATOMIC/metatomic_types.cpp | 30 ++++++++++++++++++++++++++ src/ML-METATOMIC/metatomic_types.h | 5 +++++ src/ML-METATOMIC/pair_metatomic.cpp | 30 +------------------------- 5 files changed, 38 insertions(+), 63 deletions(-) diff --git a/src/ML-METATOMIC/compute_metatomic.cpp b/src/ML-METATOMIC/compute_metatomic.cpp index 702a50f76bf..5c64cdb0c30 100644 --- a/src/ML-METATOMIC/compute_metatomic.cpp +++ b/src/ML-METATOMIC/compute_metatomic.cpp @@ -404,8 +404,6 @@ void ComputeMetatomic::compute() { return; } clear_cache(); - int nlocal = atom->nlocal; - int *mask = atom->mask; // Determine the dtype of the system based on the model's capabilities auto dtype = mta_data->model_dtype(lmp); @@ -420,20 +418,7 @@ void ComputeMetatomic::compute() { // 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}); - mta_data->selected_atoms_values.index_put_({torch::indexing::Slice(), 0}, 0); - int64_t idx = 0; - for (int i = 0; i < nlocal; i++) { - if (mask[i] & groupbit) { - mta_data->selected_atoms_values.index_put_({idx, 1}, i); - idx++; - } - } - - auto selected_atoms = torch::make_intrusive( - std::vector{"system", "atom"}, mta_data->selected_atoms_values - ); - mta_data->evaluation_options->set_selected_atoms(selected_atoms); + mta_data->set_selected_atoms(atom, groupbit); // Call the ML model to predict the requested output torch::IValue result_ivalue; diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 6f9df7f4eb5..080e461ea7d 100644 --- a/src/ML-METATOMIC/fix_metatomic.cpp +++ b/src/ML-METATOMIC/fix_metatomic.cpp @@ -374,24 +374,7 @@ void FixMetatomic::initial_integrate(int /*vflag*/) { // Configure selected atoms for evaluation // Only run the calculation for atoms in the current group - auto n_selected = group->count(igroup); - mta_data->selected_atoms_values_cpu.resize_({n_selected, 2}); - auto accessor = mta_data->selected_atoms_values_cpu.accessor(); - int64_t idx = 0; - for (int i = 0; i < nlocal; i++) { - if (mask[i] & groupbit) { - accessor[idx][0] = 0; - accessor[idx][1] = i; - idx++; - } - } - mta_data->selected_atoms_values.resize_({n_selected, 2}); - mta_data->selected_atoms_values.copy_(mta_data->selected_atoms_values_cpu); - - auto selected_atoms = torch::make_intrusive( - std::vector{"system", "atom"}, mta_data->selected_atoms_values - ); - mta_data->evaluation_options->set_selected_atoms(selected_atoms); + mta_data->set_selected_atoms(atom, groupbit); // Call the ML model to predict new positions and momenta torch::IValue result_ivalue; diff --git a/src/ML-METATOMIC/metatomic_types.cpp b/src/ML-METATOMIC/metatomic_types.cpp index cc77a759d65..73f00259f65 100644 --- a/src/ML-METATOMIC/metatomic_types.cpp +++ b/src/ML-METATOMIC/metatomic_types.cpp @@ -16,6 +16,7 @@ ------------------------------------------------------------------------- */ #include "metatomic_types.h" +#include "atom.h" #include "citeme.h" #include "comm.h" #include "error.h" @@ -172,6 +173,35 @@ void CommonMetatomicData::resolve_max_cutoff(LAMMPS* lmp) { } } +void CommonMetatomicData::set_selected_atoms(Atom* atom, int groupbit) { + int nlocal = atom->nlocal; + int *mask = atom->mask; + + int64_t n_selected = 0; + for (int i = 0; i < nlocal; i++) { + if (mask[i] & groupbit) { + n_selected++; + } + } + this->selected_atoms_values_cpu.resize_({n_selected, 2}); + auto accessor = this->selected_atoms_values_cpu.accessor(); + int64_t idx = 0; + for (int i = 0; i < nlocal; i++) { + if (mask[i] & groupbit) { + accessor[idx][0] = 0; + accessor[idx][1] = i; + idx++; + } + } + this->selected_atoms_values.resize_({n_selected, 2}); + this->selected_atoms_values.copy_(this->selected_atoms_values_cpu); + + auto selected_atoms = torch::make_intrusive( + std::vector{"system", "atom"}, this->selected_atoms_values + ); + this->evaluation_options->set_selected_atoms(selected_atoms); +} + std::map CommonMetatomicData::collect_requested_inputs() const { std::map input_holders; auto requested_inputs = this->model->run_method("requested_inputs", /*use_new_names=*/ true).toGenericDict(); diff --git a/src/ML-METATOMIC/metatomic_types.h b/src/ML-METATOMIC/metatomic_types.h index 93b83aa87d9..f2296e23a58 100644 --- a/src/ML-METATOMIC/metatomic_types.h +++ b/src/ML-METATOMIC/metatomic_types.h @@ -26,6 +26,8 @@ namespace LAMMPS_NS { +class Atom; + struct CommonMetatomicData { CommonMetatomicData(std::string length_unit); void load_model(LAMMPS* lmp, const char* path, const char* extensions_directory); @@ -42,6 +44,9 @@ struct CommonMetatomicData { // back to the requested neighbor lists when the range is infinite void resolve_max_cutoff(LAMMPS* lmp); + // set the selected atoms in the evaluation options according to the request of lmp + void set_selected_atoms(Atom* atom, int groupbit); + // collect the additional inputs requested by the model std::map collect_requested_inputs() const; diff --git a/src/ML-METATOMIC/pair_metatomic.cpp b/src/ML-METATOMIC/pair_metatomic.cpp index 09982596b49..b99a96eecd6 100644 --- a/src/ML-METATOMIC/pair_metatomic.cpp +++ b/src/ML-METATOMIC/pair_metatomic.cpp @@ -610,35 +610,7 @@ void PairMetatomic::compute(int eflag, int vflag) { // build selected atoms on CPU, then copy to device // only include atoms in the "all" group; atoms temporarily removed from // the all group (e.g. by fix_gcmc) are excluded. - int nlocal = atom->nlocal; - int *mask = atom->mask; - - int64_t n_selected = 0; - for (int i = 0; i < nlocal; i++) { - if (mask[i] & 1) { - n_selected++; - } - } - - mta_data->selected_atoms_values_cpu.resize_({n_selected, 2}); - auto accessor = mta_data->selected_atoms_values_cpu.accessor(); - int64_t idx = 0; - for (int i = 0; i < nlocal; i++) { - if (mask[i] & 1) { - accessor[idx][0] = 0; - accessor[idx][1] = i; - idx++; - } - } - mta_data->selected_atoms_values.resize_({n_selected, 2}); - mta_data->selected_atoms_values.copy_(mta_data->selected_atoms_values_cpu); - - auto selected_atoms = torch::make_intrusive( - std::vector{"system", "atom"}, - mta_data->selected_atoms_values, - metatensor::assume_unique{} - ); - mta_data->evaluation_options->set_selected_atoms(selected_atoms); + mta_data->set_selected_atoms(atom, 1); torch::IValue results_ivalue; try { From 806bb300b6497a72adb5a290a23bbaecd331e8ca Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Sun, 12 Jul 2026 17:50:36 +0200 Subject: [PATCH 3/5] Replace also the neighbor list requesting --- src/ML-METATOMIC/compute_metatomic.cpp | 24 +-------------------- src/ML-METATOMIC/fix_metatomic.cpp | 24 +-------------------- src/ML-METATOMIC/metatomic_system.cpp | 30 ++++++++++++++++++++++++++ src/ML-METATOMIC/metatomic_system.h | 4 ++++ src/ML-METATOMIC/metatomic_types.h | 3 ++- src/ML-METATOMIC/pair_metatomic.cpp | 14 +----------- 6 files changed, 39 insertions(+), 60 deletions(-) diff --git a/src/ML-METATOMIC/compute_metatomic.cpp b/src/ML-METATOMIC/compute_metatomic.cpp index 5c64cdb0c30..edff5566ade 100644 --- a/src/ML-METATOMIC/compute_metatomic.cpp +++ b/src/ML-METATOMIC/compute_metatomic.cpp @@ -358,29 +358,7 @@ void ComputeMetatomic::init() { // ALL pairs, even if options->full_list() is false. We will then filter // the pairs to only include each pair once where needed. auto request = neighbor->add_request(this, NeighConst::REQ_FULL | NeighConst::REQ_GHOST); - request->set_cutoff(mta_data->max_cutoff); - - auto mincut = mta_data->max_cutoff + neighbor->skin; - if (comm->get_comm_cutoff() < mincut) { - if (comm->me == 0) { - error->warning(FLERR, - "Increasing communication cutoff to {:.8} for compute metatomic", - mincut - ); - } - comm->cutghostuser = mincut; - } - - // Translate from the metatomic neighbor lists requests to LAMMPS neighbor - // lists requests. - auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); - for (const auto& ivalue: requested_nl.toList()) { - auto options = ivalue.get().toCustomClass(); - auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); - assert(cutoff <= mta_data->max_cutoff); - - this->system_adaptor->add_nl_request(cutoff, options); - } + this->system_adaptor->configure_neighbor_lists(request, mta_data, "compute metatomic"); // HACK: Explicitly set the binsize for the neighbor list if there is no // pair_style that would set it instead. diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 080e461ea7d..a3da10a4301 100644 --- a/src/ML-METATOMIC/fix_metatomic.cpp +++ b/src/ML-METATOMIC/fix_metatomic.cpp @@ -283,29 +283,7 @@ void FixMetatomic::init() { // ALL pairs, even if options->full_list() is false. We will then filter // the pairs to only include each pair once where needed. auto request = neighbor->add_request(this, NeighConst::REQ_FULL | NeighConst::REQ_GHOST); - request->set_cutoff(mta_data->max_cutoff); - - auto mincut = mta_data->max_cutoff + neighbor->skin; - if (comm->get_comm_cutoff() < mincut) { - if (comm->me == 0) { - error->warning(FLERR, - "Increasing communication cutoff to {:.8} for fix metatomic", - mincut - ); - } - comm->cutghostuser = mincut; - } - - // Translate from the metatomic neighbor lists requests to LAMMPS neighbor - // lists requests. - auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); - for (const auto& ivalue: requested_nl.toList()) { - auto options = ivalue.get().toCustomClass(); - auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); - assert(cutoff <= mta_data->max_cutoff); - - this->system_adaptor->add_nl_request(cutoff, options); - } + this->system_adaptor->configure_neighbor_lists(request, mta_data, "fix metatomic"); // HACK: Explicitly set the binsize for the neighbor list if there is no // pair_style that would set it instead. diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 2f03802c2a7..cbc8291c027 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -16,6 +16,7 @@ ------------------------------------------------------------------------- */ #include "metatomic_system.h" #include "metatomic_timer.h" +#include "metatomic_types.h" #include "metatomic_quantities.h" #include "atom.h" @@ -24,7 +25,9 @@ #include "error.h" #include "update.h" +#include "neighbor.h" #include "neigh_list.h" +#include "neigh_request.h" #include #include @@ -185,6 +188,33 @@ void MetatomicSystemAdaptor::add_nl_request(double cutoff, metatomic_torch::Neig }); } +// Translate from the metatomic neighbor lists requests to LAMMPS neighbor lists requests +void MetatomicSystemAdaptor::configure_neighbor_lists(NeighRequest* request, CommonMetatomicData* mta_data, const char* requester) { + request->set_cutoff(mta_data->max_cutoff); + + auto mincut = mta_data->max_cutoff + neighbor->skin; + if (comm->get_comm_cutoff() < mincut) { + if (comm->me == 0) { + error->warning(FLERR, + "Increasing communication cutoff to {:.8} for {}", + mincut, requester + ); + } + comm->cutghostuser = mincut; + } + + // Translate from the metatomic neighbor lists requests to LAMMPS neighbor + // lists requests. + auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); + for (const auto& ivalue: requested_nl.toList()) { + auto options = ivalue.get().toCustomClass(); + auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); + assert(cutoff <= mta_data->max_cutoff); + + this->add_nl_request(cutoff, options); + } + +} static std::array cell_shifts( const std::array, 3>& cell_inv, diff --git a/src/ML-METATOMIC/metatomic_system.h b/src/ML-METATOMIC/metatomic_system.h index 5d0865862df..43f30ec24d2 100644 --- a/src/ML-METATOMIC/metatomic_system.h +++ b/src/ML-METATOMIC/metatomic_system.h @@ -27,6 +27,8 @@ namespace LAMMPS_NS { +struct CommonMetatomicData; + struct MetatomicSystemOptions { // Mapping from LAMMPS types to metatomic types. // If used with kokkos, this should be a device pointer @@ -68,6 +70,8 @@ class MetatomicSystemAdaptor : public Pointers { double cutoff, metatomic_torch::NeighborListOptions request ); + virtual void configure_neighbor_lists(NeighRequest* request, CommonMetatomicData* mta_data, const char* requester); + // Create a metatomic system matching the LAMMPS system data virtual metatomic_torch::System system_from_lmp( NeighList* list, diff --git a/src/ML-METATOMIC/metatomic_types.h b/src/ML-METATOMIC/metatomic_types.h index f2296e23a58..2f90992d954 100644 --- a/src/ML-METATOMIC/metatomic_types.h +++ b/src/ML-METATOMIC/metatomic_types.h @@ -44,7 +44,8 @@ struct CommonMetatomicData { // back to the requested neighbor lists when the range is infinite void resolve_max_cutoff(LAMMPS* lmp); - // set the selected atoms in the evaluation options according to the request of lmp + // set the selected atoms in the evaluation options according to the request of lmp, + // build selected atoms on CPU, then copy to device void set_selected_atoms(Atom* atom, int groupbit); // collect the additional inputs requested by the model diff --git a/src/ML-METATOMIC/pair_metatomic.cpp b/src/ML-METATOMIC/pair_metatomic.cpp index b99a96eecd6..95737f26762 100644 --- a/src/ML-METATOMIC/pair_metatomic.cpp +++ b/src/ML-METATOMIC/pair_metatomic.cpp @@ -524,18 +524,7 @@ void PairMetatomic::init_style() { // ALL pairs, even if options->full_list() is false. We will then filter // the pairs to only include each pair once where needed. auto request = neighbor->add_request(this, NeighConst::REQ_FULL | NeighConst::REQ_GHOST); - request->set_cutoff(mta_data->max_cutoff); - - // Translate from the metatomic neighbor lists requests to LAMMPS neighbor - // lists requests. - auto requested_nl = mta_data->model->run_method("requested_neighbor_lists"); - for (const auto& ivalue: requested_nl.toList()) { - auto options = ivalue.get().toCustomClass(); - auto cutoff = options->engine_cutoff(mta_data->evaluation_options->length_unit()); - assert(cutoff <= mta_data->max_cutoff); - - this->system_adaptor->add_nl_request(cutoff, options); - } + this->system_adaptor->configure_neighbor_lists(request, mta_data, "pair metatomic"); } void PairMetatomic::init_list(int id, NeighList *ptr) { @@ -607,7 +596,6 @@ void PairMetatomic::compute(int eflag, int vflag) { input_holders ); - // build selected atoms on CPU, then copy to device // only include atoms in the "all" group; atoms temporarily removed from // the all group (e.g. by fix_gcmc) are excluded. mta_data->set_selected_atoms(atom, 1); From 70061648d9e798f4746752d061969325b989a716 Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Sun, 12 Jul 2026 18:06:52 +0200 Subject: [PATCH 4/5] Fix a bug in creating the tensor of velocities --- src/ML-METATOMIC/metatomic_system.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index cbc8291c027..9a77e1f8432 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -813,7 +813,7 @@ void MetatomicSystemAdaptor::add_velocities(metatomic_torch::System& system, std // gather velocities (per-atom) in a CPU tensor and ship to device auto velocities = torch::from_blob( - v, + *v, {total_n_atoms, 3}, torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU).requires_grad(false) ); From 553c6758aa503ef90f63be78cda8d5bf327d882b Mon Sep 17 00:00:00 2001 From: GardevoirX Date: Mon, 13 Jul 2026 10:02:53 +0200 Subject: [PATCH 5/5] Disable modifying ghost cutoff for `pair metatomic` --- src/ML-METATOMIC/metatomic_system.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ML-METATOMIC/metatomic_system.cpp b/src/ML-METATOMIC/metatomic_system.cpp index 9a77e1f8432..631b229f4cd 100644 --- a/src/ML-METATOMIC/metatomic_system.cpp +++ b/src/ML-METATOMIC/metatomic_system.cpp @@ -193,7 +193,7 @@ void MetatomicSystemAdaptor::configure_neighbor_lists(NeighRequest* request, Com request->set_cutoff(mta_data->max_cutoff); auto mincut = mta_data->max_cutoff + neighbor->skin; - if (comm->get_comm_cutoff() < mincut) { + if (strcmp(requester, "pair metatomic") != 0 && comm->get_comm_cutoff() < mincut) { if (comm->me == 0) { error->warning(FLERR, "Increasing communication cutoff to {:.8} for {}",