Skip to content

Extract common logics to helper functions #55

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
GardevoirX wants to merge 5 commits into
base: metatomic
Choose a base branch
from
Open

Extract common logics to helper functions #55

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9fda21c
Extract common logics to helper functions
GardevoirX Jul 10, 2026
764221c
Replace the atom selection logic with `CommonMetatomicData::set_selec…
GardevoirX Jul 12, 2026
806bb30
Replace also the neighbor list requesting
GardevoirX Jul 12, 2026
7006164
Fix a bug in creating the tensor of velocities
GardevoirX Jul 12, 2026
553c675
Disable modifying ghost cutoff for `pair metatomic`
GardevoirX Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 12 additions & 137 deletions src/ML-METATOMIC/compute_metatomic.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,16 +251,8 @@ ComputeMetatomic::ComputeMetatomic(LAMMPS *lmp, int narg, char **arg): Compute(l
);

auto capabilities = mta_data->model->run_method("capabilities").toCustomClass<metatomic_torch::ModelCapabilitiesHolder>();
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,
Expand All @@ -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<metatomic_torch::ModelOutputHolder>()
);
}

mta_data->requested_inputs = mta_data->collect_requested_inputs();

// Initialize the output layout
if (strcmp(sample_kind.c_str(), "atom") == 0) {
peratom_flag = 1;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<metatomic_torch::NeighborListOptionsHolder>();
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{
Expand All @@ -391,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<metatomic_torch::NeighborListOptionsHolder>();
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.
Expand All @@ -427,54 +372,6 @@ void ComputeMetatomic::init() {
// END HACK
}

void ComputeMetatomic::pick_device(c10::Device& device, const char* requested) {
torch::optional<std::string> 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<torch::DeviceIndex>(device_index));
} else {
device = torch::Device(device_type);
}
}

void ComputeMetatomic::init_list(int id, NeighList *ptr) {
mta_list = ptr;
}
Expand All @@ -485,18 +382,9 @@ 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 = 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,
Expand All @@ -508,20 +396,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<metatensor_torch::LabelsHolder>(
std::vector<std::string>{"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;
Expand Down
1 change: 0 additions & 1 deletion src/ML-METATOMIC/compute_metatomic.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::string> requested_device;
std::optional<std::string> extensions_directory;
Expand Down
140 changes: 10 additions & 130 deletions src/ML-METATOMIC/fix_metatomic.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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<metatomic_torch::NeighborListOptionsHolder>();
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{
Expand All @@ -303,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<metatomic_torch::NeighborListOptionsHolder>();
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.
Expand All @@ -339,54 +297,6 @@ void FixMetatomic::init() {
// END HACK
}

void FixMetatomic::pick_device(c10::Device& device, const char* requested) {
torch::optional<std::string> 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<torch::DeviceIndex>(device_index));
} else {
device = torch::Device(device_type);
}
}

void FixMetatomic::init_list(int id, NeighList *ptr) {
mta_list = ptr;
}
Expand Down Expand Up @@ -426,24 +336,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<std::string, metatomic_torch::ModelOutput> input_holders;
auto requested_inputs = mta_data->model->run_method("requested_inputs", /*use_new_names=*/ true).toGenericDict();
for (const auto& entry : requested_inputs) {
input_holders.emplace(
entry.key().toStringRef(),
entry.value().toCustomClass<metatomic_torch::ModelOutputHolder>()
);
}
auto input_holders = mta_data->collect_requested_inputs();

// transform from LAMMPS to metatomic System
auto system = this->system_adaptor->system_from_lmp(
mta_list,
Expand All @@ -455,24 +352,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<int32_t, 2>();
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<metatensor_torch::LabelsHolder>(
std::vector<std::string>{"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;
Expand Down
Loading