diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md index 19bed4f72c..137c0e4069 100644 --- a/doc/third-party/lammps-command.md +++ b/doc/third-party/lammps-command.md @@ -47,7 +47,7 @@ pair_style deepmd models ... keyword value ... - models = frozen model(s) to compute the interaction. If multiple models are provided, then only the first model serves to provide energy and force prediction for each timestep of molecular dynamics, and the model deviation will be computed among all models every `out_freq` timesteps. -- keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _fparam_from_fix_ or _aparam_from_compute_ or _charge_spin_ or _atomic_ or _relative_ or _relative_v_ or _aparam_ or _ttm_ +- keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _fparam_from_fix_ or _aparam_from_compute_ or _charge_spin_ or _atomic_ or _relative_ or _relative_v_ or _aparam_ or _ttm_ or _center_group_ or _environment_cutoff_ or _include_molecule_
     out_file value = filename
@@ -75,6 +75,13 @@ pair_style deepmd models ... keyword value ...
         parameters = one or more atomic parameters of each atom required for model evaluation
     ttm value = id
         id = fix ID of fix ttm
+    center_group value = group-ID
+        group-ID = atoms that are always included and define the centers of compact subsystem selection
+    environment_cutoff value = distance
+        distance = cutoff around center_group atoms, in the current LAMMPS distance units
+    include_molecule value = yes or no
+        yes = include every DeePMD-mapped atom with the same positive molecule ID as an atom inside environment_cutoff (default)
+        no = include only individual atoms inside environment_cutoff
 
### Examples @@ -93,6 +100,9 @@ pair_style deepmd ener.pb aparam_from_compute 1 compute 1 all ke/atom pair_style deepmd dpa3.pth charge_spin 1.0 2.0 + +group qm id 1:16 +pair_style deepmd dprc.pb center_group qm environment_cutoff 6.0 include_molecule yes ``` ### Description @@ -121,6 +131,12 @@ If the keyword `aparam` is set, the given atomic parameter(s) will be fed to the If the keyword `charge_spin` is set, the given per-frame charge/spin value(s) will be fed to models that were trained with a charge/spin embedding (e.g. DPA-3 with `add_chg_spin_ebd`). If the keyword is not set, the model's stored `default_chg_spin` (if any) is used. If the keyword `ttm` is set, electronic temperatures from [fix ttm command](https://docs.lammps.org/fix_ttm.html) will be fed to the model as the atomic parameters. +If `center_group` and `environment_cutoff` are set, `pair_style deepmd` evaluates the model on a dynamically selected compact subsystem. Every atom in `center_group` is included. Every other model atom within `environment_cutoff` of a center atom is included using periodic minimum-image distances; orthogonal and restricted triclinic cells are supported. With the default `include_molecule yes`, a cutoff hit promotes every DeePMD-mapped atom with the same positive LAMMPS molecule ID, so a solvent molecule represented by the model is not truncated. With `include_molecule no`, only the individual atom is selected. + +Atoms outside the compact subsystem are removed before the model backend is called. Their DeePMD force, atomic energy, atomic virial, and atomic model-deviation output are zero. Energy and global virial are accumulated from the compact subsystem normally. When multiple models are supplied, every model receives the same selection, and the force and virial deviation summaries are normalized by the number of selected model atoms rather than the full LAMMPS atom count. + +Compact evaluation is intended to be physically equivalent to full-system evaluation when atoms outside `environment_cutoff` have no contribution to the model. The user must choose an environment cutoff that covers the relevant model interaction range and ensure that excluded atom-type energy biases are zero where required by the model construction. For whole-molecule selection, every environment atom that can enter the cutoff must have a positive molecule ID. The center group must contain only atom types mapped to the model, not `NULL` types. + Only a single `pair_coeff` command is used with the deepmd style which specifies atom names. These are mapped to LAMMPS atom types (integers from 1 to Ntypes) by specifying Ntypes additional arguments after `* *` in the `pair_coeff` command. If atom names are not set in the `pair_coeff` command, the training parameter {ref}`type_map ` will be used by default. If a mapping value is specified as `NULL`, the mapping is not performed. This can be used when a deepmd potential is used as part of the hybrid pair style. The `NULL` values are placeholders for atom types that will be used with other potentials. @@ -129,6 +145,7 @@ If the training parameter {ref}`type_map ` is not set, atom name ### Restrictions - The `deepmd` pair style is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information. +- Compact `center_group` evaluation is currently supported by `pair_style deepmd`; `pair_style deepmd/kk` diagnoses this mode as unsupported. ## pair_style `deepspin` diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index 0bfd04781c..7a8742f1b7 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -576,6 +576,16 @@ double PairDeepBaseModel::init_one(int i, int j) { void* PairDeepBaseModel::extract(const char* str, int& dim) { if (strcmp(str, "cut_coul") == 0) { + // A regular Deep Potential cutoff is not a Coulomb cutoff. Advertising + // it as one makes pair_style hybrid/overlay reject a legitimate + // combination with a long-range Coulomb sub-style whenever their cutoffs + // differ (for example, a 6 A DPRc model and 9 A TIP4P electrostatics). + // PPPM/DPLR is the exception: that solver intentionally uses the DeepMD + // model cutoff to split its short- and long-range contributions. + if (force->kspace_style == nullptr || + strcmp(force->kspace_style, "pppm/dplr") != 0) { + return nullptr; + } dim = 0; return (void*)&cutoff; } diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 06d9f9976a..e20e8543c8 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -1,14 +1,17 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include #include +#include #include #include #include #include #include #include +#include #include "atom.h" #include "citeme.h" @@ -18,6 +21,7 @@ #include "error.h" #include "fix.h" #include "force.h" +#include "group.h" #include "memory.h" #include "modify.h" #include "neigh_list.h" @@ -123,6 +127,17 @@ static const char cite_user_deepmd_package[] = PairDeepMD::PairDeepMD(LAMMPS* lmp) : PairDeepBaseModel( lmp, cite_user_deepmd_package, deep_pot, deep_pot_model_devi), + compact_selection_enabled_(false), + compact_include_molecule_(true), + compact_center_group_dynamic_(false), + compact_center_group_bit_(0), + compact_environment_cutoff_(0.0), + compact_natoms_(0), + compact_packing_disabled_( + std::getenv("DP_LAMMPS_DISABLE_COMPACT_PACKING") != nullptr), + compact_packing_valid_(false), + compact_packing_nlocal_(0), + compact_packing_nghost_(0), commdata_(nullptr) { print_summary(" "); } @@ -131,6 +146,451 @@ PairDeepMD::~PairDeepMD() { // Ensure base class destructor is called } +std::vector PairDeepMD::allgather_unique_tagints( + std::vector local_values) const { + std::sort(local_values.begin(), local_values.end()); + local_values.erase(std::unique(local_values.begin(), local_values.end()), + local_values.end()); + if (comm->nprocs == 1) { + return local_values; + } + + // LAMMPS serial MPI stubs declare send buffers as void*, so MPI send + // scalars must remain mutable even though the collective does not alter them. + int local_count = static_cast(local_values.size()); + std::vector counts(comm->nprocs, 0); + std::vector displacements(comm->nprocs, 0); + MPI_Allgather(&local_count, 1, MPI_INT, counts.data(), 1, MPI_INT, world); + int total_count = 0; + for (int rank = 0; rank < comm->nprocs; ++rank) { + displacements[rank] = total_count; + total_count += counts[rank]; + } + if (total_count == 0) { + return {}; + } + + std::vector gathered(total_count); + MPI_Allgatherv(local_values.data(), local_count, MPI_LMP_TAGINT, + gathered.data(), counts.data(), displacements.data(), + MPI_LMP_TAGINT, world); + std::sort(gathered.begin(), gathered.end()); + gathered.erase(std::unique(gathered.begin(), gathered.end()), gathered.end()); + return gathered; +} + +void PairDeepMD::refresh_compact_center_tags() { + std::vector local_center_tags; + local_center_tags.reserve(atom->nlocal); + for (int ii = 0; ii < atom->nlocal; ++ii) { + if (atom->mask[ii] & compact_center_group_bit_) { + local_center_tags.push_back(atom->tag[ii]); + } + } + compact_center_tags_ = allgather_unique_tagints(std::move(local_center_tags)); + if (compact_center_tags_.empty()) { + error->all(FLERR, "center_group for pair_style deepmd is empty"); + } +} + +bool PairDeepMD::apply_compact_selection(std::vector& model_types) { + if (!compact_selection_enabled_) { + return false; + } + + if (compact_center_tags_.empty() || compact_center_group_dynamic_ || + neighbor->ago == 0) { + refresh_compact_center_tags(); + } + + const int nlocal = atom->nlocal; + const int nall = nlocal + atom->nghost; + const int model_ntypes = deep_pot.numb_types(); + double** const x = atom->x; + tagint* const tag = atom->tag; + tagint* const molecule = atom->molecule; + const auto is_model_atom = [&model_types, model_ntypes](int index) { + return model_types[index] >= 0 && model_types[index] < model_ntypes; + }; + + // Atom order and the ghost set remain stable between neighbor rebuilds. + // Dynamic groups are the exception because their membership may change on + // any step even when the neighbor topology does not. + if (compact_is_center_.size() != static_cast(nall) || + compact_center_group_dynamic_ || neighbor->ago == 0) { + compact_is_center_.resize(nall); + for (int ii = 0; ii < nall; ++ii) { + compact_is_center_[ii] = std::binary_search( + compact_center_tags_.begin(), compact_center_tags_.end(), tag[ii]); + } + } + + int invalid_center_local = 0; + for (int ii = 0; ii < nlocal; ++ii) { + if (compact_is_center_[ii] && !is_model_atom(ii)) { + invalid_center_local = 1; + } + } + int invalid_center = 0; + MPI_Allreduce(&invalid_center_local, &invalid_center, 1, MPI_INT, MPI_MAX, + world); + if (invalid_center) { + error->all(FLERR, + "center_group for pair_style deepmd contains an atom whose " + "type is not represented by the DeepMD model"); + } + + if (!list) { + error->all(FLERR, + "compact pair_style deepmd requires an available pair " + "neighbor list"); + } + + // The DeepMD full neighbor list already contains every ordinary atom pair + // within the environment cutoff (plus skin), including the correct ghost + // image in triclinic cells. Walking only center rows avoids an O(Ncenter*N) + // all-atom scan on every step. Pairs removed by special_bonds are handled + // separately below so compact selection remains independent of force-field + // exclusions. + const double environment_cutsq = + compact_environment_cutoff_ * compact_environment_cutoff_; + std::vector local_selection_keys; + int invalid_molecule_local = 0; + const auto select_environment_atom = [&](int center, int environment, + bool apply_minimum_image) { + if (compact_is_center_[environment] || !is_model_atom(environment)) { + return; + } + double dx = x[environment][0] - x[center][0]; + double dy = x[environment][1] - x[center][1]; + double dz = x[environment][2] - x[center][2]; + if (apply_minimum_image) { + domain->minimum_image(FLERR, dx, dy, dz); + } + if (dx * dx + dy * dy + dz * dz >= environment_cutsq) { + return; + } + if (compact_include_molecule_) { + if (molecule[environment] <= 0) { + invalid_molecule_local = 1; + return; + } + local_selection_keys.push_back(molecule[environment]); + } else { + local_selection_keys.push_back(tag[environment]); + } + }; + + for (int ii = 0; ii < nlocal; ++ii) { + if (!compact_is_center_[ii]) { + continue; + } + const int jnum = list->numneigh[ii]; + int* const jlist = list->firstneigh[ii]; + for (int jj = 0; jj < jnum; ++jj) { + select_environment_atom(ii, jlist[jj] & NEIGHMASK, false); + } + } + + // Both zero-valued special_bonds factors remove the corresponding pair + // from the neighbor list. Recover only those few pairs by tag. Building + // the tag lookup is deferred until a non-center excluded partner is found; + // the usual DPRc case has an entire bonded QM molecule in center_group and + // therefore pays no hash-table cost. + if (atom->molecular != Atom::ATOMIC && atom->special && atom->nspecial) { + std::unordered_map tag_to_index; + const auto find_atom_by_tag = [&](tagint atom_tag) { + if (tag_to_index.empty()) { + tag_to_index.reserve(nall); + for (int jj = 0; jj < nall; ++jj) { + tag_to_index.emplace(tag[jj], jj); + } + } + const auto found = tag_to_index.find(atom_tag); + return found == tag_to_index.end() ? -1 : found->second; + }; + + for (int ii = 0; ii < nlocal; ++ii) { + if (!compact_is_center_[ii]) { + continue; + } + for (int level = 1; level <= 3; ++level) { + if (force->special_lj[level] != 0.0 || + force->special_coul[level] != 0.0) { + continue; + } + const int begin = level == 1 ? 0 : atom->nspecial[ii][level - 2]; + const int end = atom->nspecial[ii][level - 1]; + for (int jj = begin; jj < end; ++jj) { + const tagint special_tag = atom->special[ii][jj]; + if (std::binary_search(compact_center_tags_.begin(), + compact_center_tags_.end(), special_tag)) { + continue; + } + const int special_index = find_atom_by_tag(special_tag); + if (special_index >= 0) { + select_environment_atom(ii, special_index, true); + } + } + } + } + } + int invalid_molecule = 0; + MPI_Allreduce(&invalid_molecule_local, &invalid_molecule, 1, MPI_INT, MPI_MAX, + world); + if (invalid_molecule) { + error->all(FLERR, + "include_molecule yes requires positive molecule IDs for all " + "environment atoms selected by pair_style deepmd"); + } + + const std::vector selected_keys = + allgather_unique_tagints(std::move(local_selection_keys)); + std::vector selected(nall, 0); + int selected_nlocal = 0; + for (int ii = 0; ii < nall; ++ii) { + const tagint key = compact_include_molecule_ ? molecule[ii] : tag[ii]; + const bool selected_environment = + std::binary_search(selected_keys.begin(), selected_keys.end(), key); + const bool active = + is_model_atom(ii) && (compact_is_center_[ii] || selected_environment); + selected[ii] = active; + if (!active) { + model_types[ii] = -1; + } else if (ii < nlocal) { + ++selected_nlocal; + } + } + + // Keep MPI send scalars mutable for compatibility with LAMMPS STUBS/mpi.h. + bigint selected_local = selected_nlocal; + MPI_Allreduce(&selected_local, &compact_natoms_, 1, MPI_LMP_BIGINT, MPI_SUM, + world); + if (compact_natoms_ == 0) { + error->all(FLERR, "compact pair_style deepmd selected no model atoms"); + } + + int local_changed = compact_selected_ != selected; + compact_selected_ = std::move(selected); + int global_changed = 0; + MPI_Allreduce(&local_changed, &global_changed, 1, MPI_INT, MPI_MAX, world); + return global_changed != 0; +} + +bool PairDeepMD::can_use_compact_packing() const { + // MPI communication metadata must be remapped together with the atom and + // neighbor indices. Keep the optimized path deliberately single-rank for + // now; all unsupported configurations retain the backend's generic compact + // selection implementation. + return compact_selection_enabled_ && !compact_packing_disabled_ && list && + comm->nprocs == 1 && dim_aparam == 0 && !do_compute_aparam && + aparam.empty() && !do_ttm; +} + +void PairDeepMD::rebuild_compact_packing() { + const int nlocal = atom->nlocal; + const int nall = nlocal + atom->nghost; + if (compact_selected_.size() != static_cast(nall)) { + error->all(FLERR, + "internal error while packing compact pair_style deepmd " + "selection"); + } + + compact_packing_old_to_new_.assign(nall, -1); + compact_packing_new_to_old_.clear(); + compact_packing_new_to_old_.reserve(nall); + + // DeepMD requires owned atoms before ghosts so nghost unambiguously defines + // the local/ghost boundary in the packed coordinate and type arrays. + for (int old_index = 0; old_index < nlocal; ++old_index) { + if (!compact_selected_[old_index]) { + continue; + } + compact_packing_old_to_new_[old_index] = + static_cast(compact_packing_new_to_old_.size()); + compact_packing_new_to_old_.push_back(old_index); + } + compact_packing_nlocal_ = + static_cast(compact_packing_new_to_old_.size()); + for (int old_index = nlocal; old_index < nall; ++old_index) { + if (!compact_selected_[old_index]) { + continue; + } + compact_packing_old_to_new_[old_index] = + static_cast(compact_packing_new_to_old_.size()); + compact_packing_new_to_old_.push_back(old_index); + } + compact_packing_nghost_ = + static_cast(compact_packing_new_to_old_.size()) - + compact_packing_nlocal_; + + compact_packing_ilist_.resize(compact_packing_nlocal_); + compact_packing_numneigh_.resize(compact_packing_nlocal_); + compact_packing_neighbors_.clear(); + compact_packing_neighbors_.resize(compact_packing_nlocal_); + for (int new_index = 0; new_index < compact_packing_nlocal_; ++new_index) { + compact_packing_ilist_[new_index] = new_index; + const int old_index = compact_packing_new_to_old_[new_index]; + const int old_numneigh = list->numneigh[old_index]; + int* const old_neighbors = list->firstneigh[old_index]; + auto& packed_neighbors = compact_packing_neighbors_[new_index]; + packed_neighbors.reserve(old_numneigh); + for (int jj = 0; jj < old_numneigh; ++jj) { + const int old_neighbor = old_neighbors[jj] & NEIGHMASK; + if (old_neighbor < 0 || old_neighbor >= nall) { + error->all(FLERR, + "invalid neighbor index while packing compact pair_style " + "deepmd data"); + } + const int new_neighbor = compact_packing_old_to_new_[old_neighbor]; + if (new_neighbor >= 0) { + packed_neighbors.push_back(new_neighbor); + } + } + compact_packing_numneigh_[new_index] = + static_cast(packed_neighbors.size()); + } + + compact_packing_firstneigh_.resize(compact_packing_nlocal_); + for (int ii = 0; ii < compact_packing_nlocal_; ++ii) { + compact_packing_firstneigh_[ii] = compact_packing_neighbors_[ii].data(); + } + + compact_packing_mapping_.clear(); + if (atom->map_style != Atom::MAP_NONE) { + compact_packing_mapping_.assign(compact_packing_new_to_old_.size(), -1); + for (int new_index = 0; + new_index < static_cast(compact_packing_new_to_old_.size()); + ++new_index) { + const int old_index = compact_packing_new_to_old_[new_index]; + const int old_owner = atom->map(atom->tag[old_index]); + if (old_owner >= 0 && old_owner < nall) { + compact_packing_mapping_[new_index] = + compact_packing_old_to_new_[old_owner]; + } + } + } + compact_packing_valid_ = true; +} + +void PairDeepMD::pack_compact_inputs( + const std::vector& full_types, + std::vector& packed_types, + std::vector& packed_coordinates) const { + const int packed_nall = static_cast(compact_packing_new_to_old_.size()); + packed_types.resize(packed_nall); + packed_coordinates.resize(3 * packed_nall); + for (int new_index = 0; new_index < packed_nall; ++new_index) { + const int old_index = compact_packing_new_to_old_[new_index]; + packed_types[new_index] = full_types[old_index]; + for (int dd = 0; dd < 3; ++dd) { + packed_coordinates[3 * new_index + dd] = + (atom->x[old_index][dd] - domain->boxlo[dd]) / dist_unit_cvt_factor; + } + } +} + +deepmd_compat::InputNlist PairDeepMD::make_compact_packing_nlist() { + deepmd_compat::InputNlist packed_list( + compact_packing_nlocal_, compact_packing_ilist_.data(), + compact_packing_numneigh_.data(), compact_packing_firstneigh_.data()); + if (!compact_packing_mapping_.empty()) { + packed_list.set_mapping(compact_packing_mapping_.data()); + } + return packed_list; +} + +void PairDeepMD::scatter_compact_output(std::vector& values, + int stride, + int full_nall) const { + const size_t packed_size = compact_packing_new_to_old_.size() * stride; + if (values.size() != packed_size) { + error->all(FLERR, + "unexpected DeepMD output size from compact packed " + "evaluation"); + } + std::vector full_values(static_cast(full_nall) * stride, 0.0); + for (int new_index = 0; + new_index < static_cast(compact_packing_new_to_old_.size()); + ++new_index) { + const int old_index = compact_packing_new_to_old_[new_index]; + for (int dd = 0; dd < stride; ++dd) { + full_values[stride * old_index + dd] = values[stride * new_index + dd]; + } + } + values.swap(full_values); +} + +void PairDeepMD::analyze_model_deviation(double& max, + double& min, + double& sum, + const std::vector& deviation, + int nlocal) const { + if (!compact_selection_enabled_) { + ana_st(max, min, sum, deviation, nlocal); + return; + } + + // If this rank owns no selected atom, preserve the caller's reduction-neutral + // seeds (min = max double, max = 0, sum = 0) for the following MPI_Reduce. + bool found = false; + for (int ii = 0; ii < nlocal; ++ii) { + if (!compact_selected_[ii]) { + continue; + } + const double value = deviation[ii]; + if (!found) { + max = min = sum = value; + found = true; + } else { + max = std::max(max, value); + min = std::min(min, value); + sum += value; + } + } +} + +void PairDeepMD::init_style() { + PairDeepBaseModel::init_style(); + if (!compact_selection_enabled_) { + return; + } + // Compact selection reconstructs the environment from the pair neighbor + // list. User exclusions remove otherwise eligible pairs before that list + // reaches this style, and unlike special_bonds there is no bounded topology + // list from which every excluded partner can be recovered. + if (neighbor->nex_type || neighbor->nex_group || neighbor->nex_mol) { + error->all( + FLERR, + "compact pair_style deepmd does not support neigh_modify exclude"); + } + if (!atom->tag_enable) { + error->all(FLERR, + "compact pair_style deepmd requires atom IDs to be enabled"); + } + const int center_group_index = group->find(compact_center_group_id_.c_str()); + if (center_group_index < 0) { + error->all(FLERR, "center_group " + compact_center_group_id_ + + " for pair_style deepmd does not exist"); + } + compact_center_group_bit_ = group->bitmask[center_group_index]; + compact_center_group_dynamic_ = group->dynamic[center_group_index] != 0; + if (compact_include_molecule_ && !atom->molecule_flag) { + error->all(FLERR, + "include_molecule yes requires an atom style with molecule " + "IDs"); + } + refresh_compact_center_tags(); +} + +double PairDeepMD::init_one(int i, int j) { + const double model_neighbor_cutoff = PairDeepBaseModel::init_one(i, j); + if (!compact_selection_enabled_) { + return model_neighbor_cutoff; + } + return std::max(model_neighbor_cutoff, compact_environment_cutoff_); +} + double PairDeepMD::eval_energy_with_fparam( const std::vector& fparam_override) { if (numb_models != 1) { @@ -159,6 +619,13 @@ double PairDeepMD::eval_energy_with_fparam( for (int ii = 0; ii < nall; ++ii) { dtype[ii] = type_idx_map[type[ii] - 1]; } + const bool compact_selection_changed = apply_compact_selection(dtype); + if (compact_selection_changed) { + // This auxiliary energy path deliberately keeps the generic backend + // representation. Invalidate any packed topology so a following force + // evaluation cannot reuse atom indices from the previous selection. + compact_packing_valid_ = false; + } double dener(0); std::vector dforce(nall * 3); @@ -208,6 +675,9 @@ double PairDeepMD::eval_energy_with_fparam( #endif } int ago = neighbor->ago; + if (compact_selection_changed) { + ago = 0; + } if (do_ghost) { if (!list) { @@ -282,11 +752,23 @@ void PairDeepMD::compute(int eflag, int vflag) { for (int ii = 0; ii < nall; ++ii) { dtype[ii] = type_idx_map[type[ii] - 1]; } + const bool compact_selection_changed = apply_compact_selection(dtype); + + const bool compact_packing = can_use_compact_packing(); + if (compact_packing && (!compact_packing_valid_ || + compact_selection_changed || neighbor->ago == 0)) { + rebuild_compact_packing(); + } + const int deepmd_nall = + compact_packing ? static_cast(compact_packing_new_to_old_.size()) + : nall; + const int deepmd_nghost = compact_packing ? compact_packing_nghost_ : nghost; double dener(0); - vector dforce(nall * 3); + vector dforce(deepmd_nall * 3); + bool dforce_is_packed = compact_packing; vector dvirial(9, 0); - vector dcoord(nall * 3, 0.); + vector dcoord; vector dbox(9, 0); vector daparam; @@ -298,18 +780,26 @@ void PairDeepMD::compute(int eflag, int vflag) { dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx - // get coord - for (int ii = 0; ii < nall; ++ii) { - for (int dd = 0; dd < 3; ++dd) { - dcoord[ii * 3 + dd] = - (x[ii][dd] - domain->boxlo[dd]) / dist_unit_cvt_factor; + if (compact_packing) { + vector packed_types; + pack_compact_inputs(dtype, packed_types, dcoord); + dtype.swap(packed_types); + } else { + dcoord.resize(nall * 3, 0.0); + for (int ii = 0; ii < nall; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + dcoord[ii * 3 + dd] = + (x[ii][dd] - domain->boxlo[dd]) / dist_unit_cvt_factor; + } } } // Owner mapping for message-passing .pt2 models that gather ghost features // through the LAMMPS atom map; unused by other models. - std::vector mapping_vec(nall, -1); - if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { + std::vector mapping_vec; + if (!compact_packing && comm->nprocs == 1 && + atom->map_style != Atom::MAP_NONE) { + mapping_vec.assign(nall, -1); for (size_t ii = 0; ii < nall; ++ii) { mapping_vec[ii] = atom->map(atom->tag[ii]); } @@ -337,6 +827,9 @@ void PairDeepMD::compute(int eflag, int vflag) { } int ago = neighbor->ago; + if (compact_selection_changed) { + ago = 0; + } if (numb_models > 1) { if (multi_models_no_mod_devi && (out_freq > 0 && update->ntimestep % out_freq == 0)) { @@ -353,13 +846,19 @@ void PairDeepMD::compute(int eflag, int vflag) { multi_models_mod_devi = (numb_models > 1 && (out_freq > 0 && update->ntimestep % out_freq == 0)); if (do_ghost) { - deepmd_compat::InputNlist lmp_list( - list->inum, list->ilist, list->numneigh, list->firstneigh, - commdata_->nswap, commdata_->sendnum, commdata_->recvnum, - commdata_->firstrecv, commdata_->sendlist, commdata_->sendproc, - commdata_->recvproc, &world, comm->nprocs); + deepmd_compat::InputNlist lmp_list; + if (compact_packing) { + lmp_list = make_compact_packing_nlist(); + } else { + lmp_list = deepmd_compat::InputNlist( + list->inum, list->ilist, list->numneigh, list->firstneigh, + commdata_->nswap, commdata_->sendnum, commdata_->recvnum, + commdata_->firstrecv, commdata_->sendlist, commdata_->sendproc, + commdata_->recvproc, &world, comm->nprocs); + } lmp_list.set_mask(NEIGHMASK); - if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { + if (!compact_packing && comm->nprocs == 1 && + atom->map_style != Atom::MAP_NONE) { lmp_list.set_mapping(mapping_vec.data()); } deepmd_compat::InputNlist extend_lmp_list; @@ -367,23 +866,28 @@ void PairDeepMD::compute(int eflag, int vflag) { // cvflag_atom is the right flag for the cvatom matrix if (!(eflag_atom || cvflag_atom)) { try { - deep_pot.compute(dener, dforce, dvirial, dcoord, dtype, dbox, nghost, - lmp_list, ago, fparam, daparam, charge_spin); + deep_pot.compute(dener, dforce, dvirial, dcoord, dtype, dbox, + deepmd_nghost, lmp_list, ago, fparam, daparam, + charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } // do atomic energy and virial else { - vector deatom(nall * 1, 0); - vector dvatom(nall * 9, 0); + vector deatom(deepmd_nall, 0); + vector dvatom(deepmd_nall * 9, 0); try { deep_pot.compute(dener, dforce, dvirial, deatom, dvatom, dcoord, - dtype, dbox, nghost, lmp_list, ago, fparam, daparam, - charge_spin); + dtype, dbox, deepmd_nghost, lmp_list, ago, fparam, + daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } + if (compact_packing) { + scatter_compact_output(deatom, 1, nall); + scatter_compact_output(dvatom, 9, nall); + } if (eflag_atom) { for (int ii = 0; ii < nlocal; ++ii) { eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; @@ -415,8 +919,8 @@ void PairDeepMD::compute(int eflag, int vflag) { } } } else if (multi_models_mod_devi) { - vector deatom(nall * 1, 0); - vector dvatom(nall * 9, 0); + vector deatom(deepmd_nall, 0); + vector dvatom(deepmd_nall * 9, 0); vector> all_virial; vector all_energy; vector> all_atom_energy; @@ -424,7 +928,7 @@ void PairDeepMD::compute(int eflag, int vflag) { if (!(eflag_atom || cvflag_atom)) { try { deep_pot_model_devi.compute(all_energy, all_force, all_virial, dcoord, - dtype, dbox, nghost, lmp_list, ago, + dtype, dbox, deepmd_nghost, lmp_list, ago, fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); @@ -433,7 +937,7 @@ void PairDeepMD::compute(int eflag, int vflag) { try { deep_pot_model_devi.compute(all_energy, all_force, all_virial, all_atom_energy, all_atom_virial, dcoord, - dtype, dbox, nghost, lmp_list, ago, + dtype, dbox, deepmd_nghost, lmp_list, ago, fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); @@ -444,11 +948,20 @@ void PairDeepMD::compute(int eflag, int vflag) { // deep_pot_model_devi.compute_avg (dvirial, all_virial); // deep_pot_model_devi.compute_avg (deatom, all_atom_energy); // deep_pot_model_devi.compute_avg (dvatom, all_atom_virial); + if (compact_packing) { + for (auto& model_force : all_force) { + scatter_compact_output(model_force, 3, nall); + } + } dener = all_energy[0]; dforce = all_force[0]; + dforce_is_packed = false; dvirial = all_virial[0]; if (eflag_atom) { deatom = all_atom_energy[0]; + if (compact_packing) { + scatter_compact_output(deatom, 1, nall); + } for (int ii = 0; ii < nlocal; ++ii) { eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; } @@ -457,6 +970,9 @@ void PairDeepMD::compute(int eflag, int vflag) { // per-atom virial (xx, yy, zz, xy, xz, yz, yx, zx, zy). if (cvflag_atom) { dvatom = all_atom_virial[0]; + if (compact_packing) { + scatter_compact_output(dvatom, 9, nall); + } for (int ii = 0; ii < nall; ++ii) { cvatom[ii][0] += scale[1][1] * dvatom[9 * ii + 0] * ener_unit_cvt_factor; // xx @@ -496,18 +1012,21 @@ void PairDeepMD::compute(int eflag, int vflag) { deep_pot_model_devi.compute_relative_std_f(std_f, tmp_avg_f, eps); } double min = numeric_limits::max(), max = 0, avg = 0; - ana_st(max, min, avg, std_f, nlocal); + analyze_model_deviation(max, min, avg, std_f, nlocal); double all_f_min = 0, all_f_max = 0, all_f_avg = 0; MPI_Reduce(&min, &all_f_min, 1, MPI_DOUBLE, MPI_MIN, 0, world); MPI_Reduce(&max, &all_f_max, 1, MPI_DOUBLE, MPI_MAX, 0, world); MPI_Reduce(&avg, &all_f_avg, 1, MPI_DOUBLE, MPI_SUM, 0, world); - all_f_avg /= double(atom->natoms); + const double deviation_natoms = + compact_selection_enabled_ ? static_cast(compact_natoms_) + : static_cast(atom->natoms); + all_f_avg /= deviation_natoms; // std v std::vector send_v(9 * numb_models); std::vector recv_v(9 * numb_models); for (int kk = 0; kk < numb_models; ++kk) { for (int ii = 0; ii < 9; ++ii) { - send_v[kk * 9 + ii] = all_virial[kk][ii] / double(atom->natoms); + send_v[kk * 9 + ii] = all_virial[kk][ii] / deviation_natoms; } } MPI_Reduce(&send_v[0], &recv_v[0], 9 * numb_models, MPI_DOUBLE, MPI_SUM, @@ -600,10 +1119,25 @@ void PairDeepMD::compute(int eflag, int vflag) { } } - // get force - for (int ii = 0; ii < nall; ++ii) { - for (int dd = 0; dd < 3; ++dd) { - f[ii][dd] += scale[1][1] * dforce[3 * ii + dd] * force_unit_cvt_factor; + // Add compact forces directly on ordinary steps. Expanding to a zero-filled + // full-system array is necessary only on model-deviation steps, where LAMMPS + // reverse communication and per-atom statistics address the native atom + // indices through all_force. + if (dforce_is_packed) { + for (int new_index = 0; + new_index < static_cast(compact_packing_new_to_old_.size()); + ++new_index) { + const int old_index = compact_packing_new_to_old_[new_index]; + for (int dd = 0; dd < 3; ++dd) { + f[old_index][dd] += + scale[1][1] * dforce[3 * new_index + dd] * force_unit_cvt_factor; + } + } + } else { + for (int ii = 0; ii < nall; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + f[ii][dd] += scale[1][1] * dforce[3 * ii + dd] * force_unit_cvt_factor; + } } } @@ -637,6 +1171,9 @@ static bool is_key(const string& input) { keys.push_back("relative_v"); keys.push_back("virtual_len"); keys.push_back("spin_norm"); + keys.push_back("center_group"); + keys.push_back("environment_cutoff"); + keys.push_back("include_molecule"); for (int ii = 0; ii < keys.size(); ++ii) { if (input == keys[ii]) { @@ -705,6 +1242,28 @@ void PairDeepMD::settings(int narg, char** arg) { fparam.clear(); aparam.clear(); charge_spin.clear(); + compact_selection_enabled_ = false; + compact_include_molecule_ = true; + compact_center_group_dynamic_ = false; + compact_center_group_bit_ = 0; + compact_environment_cutoff_ = 0.0; + compact_natoms_ = 0; + compact_center_group_id_.clear(); + compact_center_tags_.clear(); + compact_selected_.clear(); + compact_packing_valid_ = false; + compact_packing_nlocal_ = 0; + compact_packing_nghost_ = 0; + compact_packing_old_to_new_.clear(); + compact_packing_new_to_old_.clear(); + compact_packing_ilist_.clear(); + compact_packing_numneigh_.clear(); + compact_packing_neighbors_.clear(); + compact_packing_firstneigh_.clear(); + compact_packing_mapping_.clear(); + bool center_group_set = false; + bool environment_cutoff_set = false; + bool include_molecule_set = false; while (iarg < narg) { if (!is_key(arg[iarg])) { error->all(FLERR, @@ -722,6 +1281,45 @@ void PairDeepMD::settings(int narg, char** arg) { } out_file = string(arg[iarg + 1]); iarg += 2; + } else if (string(arg[iarg]) == string("center_group")) { + if (center_group_set) { + error->all(FLERR, "center_group may be specified only once"); + } + if (iarg + 1 >= narg || is_key(arg[iarg + 1])) { + error->all(FLERR, "Illegal center_group, group ID is not provided"); + } + compact_selection_enabled_ = true; + compact_center_group_id_ = arg[iarg + 1]; + center_group_set = true; + iarg += 2; + } else if (string(arg[iarg]) == string("environment_cutoff")) { + if (environment_cutoff_set) { + error->all(FLERR, "environment_cutoff may be specified only once"); + } + if (iarg + 1 >= narg || is_key(arg[iarg + 1])) { + error->all(FLERR, "Illegal environment_cutoff, value is not provided"); + } + compact_environment_cutoff_ = + utils::numeric(FLERR, arg[iarg + 1], false, lmp); + if (!std::isfinite(compact_environment_cutoff_) || + compact_environment_cutoff_ <= 0.0) { + error->all(FLERR, + "environment_cutoff must be a finite value greater than " + "zero"); + } + environment_cutoff_set = true; + iarg += 2; + } else if (string(arg[iarg]) == string("include_molecule")) { + if (include_molecule_set) { + error->all(FLERR, "include_molecule may be specified only once"); + } + if (iarg + 1 >= narg || is_key(arg[iarg + 1])) { + error->all(FLERR, "Illegal include_molecule, yes/no is not provided"); + } + compact_include_molecule_ = + utils::logical(FLERR, arg[iarg + 1], false, lmp) != 0; + include_molecule_set = true; + iarg += 2; } else if (string(arg[iarg]) == string("fparam")) { for (int ii = 0; ii < dim_fparam; ++ii) { if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) { @@ -853,6 +1451,17 @@ void PairDeepMD::settings(int narg, char** arg) { if (out_freq < 0) { error->all(FLERR, "Illegal out_freq, should be >= 0"); } + if (compact_selection_enabled_ && !environment_cutoff_set) { + error->all(FLERR, + "center_group requires environment_cutoff in pair_style " + "deepmd"); + } + if (!compact_selection_enabled_ && + (environment_cutoff_set || include_molecule_set)) { + error->all(FLERR, + "environment_cutoff and include_molecule require center_group " + "in pair_style deepmd"); + } if ((int)do_ttm + (int)do_compute_aparam + (int)(aparam.size() > 0) > 1) { error->all(FLERR, "aparam, aparam_from_compute, and ttm should NOT be set " @@ -905,6 +1514,14 @@ void PairDeepMD::settings(int narg, char** arg) { cout << endl << pre << "rcut in model: " << cutoff << endl << pre << "ntypes in model: " << numb_types << endl; + if (compact_selection_enabled_) { + cout << pre << "compact center group: " << compact_center_group_id_ + << endl + << pre << "environment cutoff: " << compact_environment_cutoff_ + << endl + << pre << "include molecules: " + << (compact_include_molecule_ ? "yes" : "no") << endl; + } if (fparam.size() > 0) { cout << pre << "using fparam(s): "; for (int ii = 0; ii < dim_fparam; ++ii) { diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index c7c40d9b18..83a08c9226 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -47,9 +47,12 @@ class PairDeepMD : public PairDeepBaseModel { void settings(int, char**) override; void coeff(int, char**) override; void compute(int, int) override; + void init_style() override; + double init_one(int, int) override; int pack_reverse_comm(int, int, double*) override; void unpack_reverse_comm(int, int*, double*) override; double eval_energy_with_fparam(const std::vector& fparam_override); + bool compact_selection_enabled() const { return compact_selection_enabled_; } protected: deepmd_compat::DeepPot deep_pot; @@ -60,6 +63,59 @@ class PairDeepMD : public PairDeepBaseModel { deepmd_compat::InputNlist make_comm_nlist(); private: + // Compact evaluation is implemented by assigning type -1 to atoms outside + // the selected subsystem. Every supported DeepPot backend already compacts + // such atoms, remaps its neighbor/communication data, and scatters outputs + // back to the original atom order. + bool compact_selection_enabled_; + bool compact_include_molecule_; + bool compact_center_group_dynamic_; + int compact_center_group_bit_; + double compact_environment_cutoff_; + bigint compact_natoms_; + std::string compact_center_group_id_; + std::vector compact_center_tags_; + // Center membership is stable between neighbor rebuilds for a static + // group. Cache it in LAMMPS atom order so the per-step cutoff search can + // walk only the pair neighbor rows owned by center atoms. + std::vector compact_is_center_; + std::vector compact_selected_; + + // Single-rank compact packing removes atoms excluded by compact selection + // before crossing the LAMMPS/DeepMD API boundary. The atom map and filtered + // neighbor rows are stable until LAMMPS rebuilds its neighbor list; only the + // compact coordinate buffer must be refreshed on ordinary MD steps. + bool compact_packing_disabled_; + bool compact_packing_valid_; + int compact_packing_nlocal_; + int compact_packing_nghost_; + std::vector compact_packing_old_to_new_; + std::vector compact_packing_new_to_old_; + std::vector compact_packing_ilist_; + std::vector compact_packing_numneigh_; + std::vector > compact_packing_neighbors_; + std::vector compact_packing_firstneigh_; + std::vector compact_packing_mapping_; + + std::vector allgather_unique_tagints( + std::vector local_values) const; + void refresh_compact_center_tags(); + bool apply_compact_selection(std::vector& model_types); + bool can_use_compact_packing() const; + void rebuild_compact_packing(); + void pack_compact_inputs(const std::vector& full_types, + std::vector& packed_types, + std::vector& packed_coordinates) const; + deepmd_compat::InputNlist make_compact_packing_nlist(); + void scatter_compact_output(std::vector& values, + int stride, + int full_nall) const; + void analyze_model_deviation(double& max, + double& min, + double& sum, + const std::vector& deviation, + int nlocal) const; + CommBrickDeepMD* commdata_; }; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 9ecd68f055..3b50a665d0 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -176,6 +176,12 @@ void PairDeepMDKokkos::init_style() { // Base setup and the full neighbor-list request. PairDeepMD::init_style(); + if (compact_selection_enabled()) { + error->all(FLERR, + "pair style deepmd/kk does not yet support center_group compact " + "evaluation; use pair style deepmd"); + } + // The device edge path requires a GPU execution space and a single model. if (std::is_same::value) { error->all(FLERR, "pair style deepmd/kk runs on the GPU backend only."); diff --git a/source/lmp/tests/test_dplr.py b/source/lmp/tests/test_dplr.py index 7c912a036d..ef42d52724 100644 --- a/source/lmp/tests/test_dplr.py +++ b/source/lmp/tests/test_dplr.py @@ -375,6 +375,19 @@ def test_pair_deepmd_sr(lammps) -> None: lammps.run(1) +def test_pair_deepmd_hybrid_long_range(lammps) -> None: + """A model cutoff must not masquerade as a hybrid Coulomb cutoff.""" + # The model cutoff is 4 A and the actual Coulomb cutoff is 5 A. Plain + # PPPM must obtain the latter from coul/long without treating the DeepMD + # neighbor cutoff as a conflicting electrostatic cutoff. + lammps.pair_style(f"hybrid/overlay deepmd {pb_file.resolve()} coul/long 5.0") + lammps.pair_coeff("* * deepmd") + lammps.pair_coeff("* * coul/long") + lammps.kspace_style("pppm 1e-5") + lammps.run(0) + assert np.isfinite(lammps.eval("pe")) + + def test_pair_deepmd_sr_virial(lammps) -> None: lammps.group("real_atom type 1 2") lammps.pair_style(f"deepmd {pb_file.resolve()}") diff --git a/source/lmp/tests/test_lammps_compact.py b/source/lmp/tests/test_lammps_compact.py new file mode 100644 index 0000000000..a53e6caba0 --- /dev/null +++ b/source/lmp/tests/test_lammps_compact.py @@ -0,0 +1,526 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for cutoff-based compact ``pair_style deepmd`` evaluation.""" + +import os +import shutil +import subprocess as sp +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from lammps_test_utils import ( + require_backend, +) +from model_convert import ( + ensure_converted_pb, +) + +pbtxt_file = Path(__file__).parents[2] / "tests" / "infer" / "deeppot.pbtxt" +pbtxt_file2 = Path(__file__).parents[2] / "tests" / "infer" / "deeppot-1.pbtxt" + + +def setup_module() -> None: + require_backend("ENABLE_TENSORFLOW", "TensorFlow") + + +@pytest.fixture(scope="module") +def compact_models(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + """Convert the two standard energy models used by the LAMMPS tests.""" + model_dir = tmp_path_factory.mktemp("lammps_compact") + model = model_dir / "graph.pb" + model2 = model_dir / "graph2.pb" + ensure_converted_pb(pbtxt_file, model) + ensure_converted_pb(pbtxt_file2, model2) + return model, model2 + + +def _make_system( + models: tuple[Path, ...], + *, + with_inactive_molecule: bool, + compact: bool, + triclinic: bool, + deviation_file: Path | None = None, +) -> PyLammps: + """Build a core, one selected molecule, and an optional distant molecule. + + Atom 2 is across the periodic x boundary from the core. Atom 3 belongs to + the same molecule but lies outside both the environment and model cutoffs; + retaining its nonzero atomic bias verifies whole-molecule promotion. + """ + # At y=z=5, the triclinic tilt shifts the x origin by 19/30. Apply that + # shift to the boundary-crossing pair so both points remain inside the + # tilted prism while preserving their one-angstrom minimum-image distance. + boundary_shift = 19.0 / 30.0 if triclinic else 0.0 + atom_rows = [ + (1, 0.5 + boundary_shift, 5.0, 5.0, 1), + (2, 29.5 + boundary_shift, 5.0, 5.0, 10), + (1, 20.0, 10.0, 5.0, 10), + ] + if with_inactive_molecule: + atom_rows.extend( + [ + (2, 12.0, 20.0, 5.0, 20), + (1, 13.0, 20.0, 5.0, 20), + ] + ) + + lammps = PyLammps() + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + lammps.lmp.command(f"plugin load {plugin}") + lammps.units("metal") + lammps.boundary("p p p") + lammps.atom_style("molecular") + lammps.atom_modify("map array") + if triclinic: + lammps.region("box prism 0 30 0 30 0 30 3 1 2 units box") + else: + lammps.region("box block 0 30 0 30 0 30 units box") + lammps.create_box("2 box") + for atom_type, x, y, z, _ in atom_rows: + lammps.create_atoms(f"{atom_type} single {x} {y} {z} units box") + for atom_id, (_, _, _, _, molecule_id) in enumerate(atom_rows, start=1): + lammps.set(f"atom {atom_id} mol {molecule_id}") + assert lammps.lmp.get_natoms() == len(atom_rows) + lammps.group("qm id 1") + lammps.mass("1 16") + lammps.mass("2 2") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + + style = "deepmd " + " ".join(str(model.resolve()) for model in models) + if deviation_file is not None: + style += f" out_file {deviation_file.resolve()} out_freq 1 atomic" + if compact: + style += " center_group qm environment_cutoff 1.5 include_molecule yes" + lammps.pair_style(style) + lammps.pair_coeff("* *") + lammps.compute("peatom all pe/atom pair") + lammps.variable("peatom atom c_peatom") + lammps.run(0) + return lammps + + +def _snapshot(lammps: PyLammps, natoms: int) -> tuple[float, np.ndarray, np.ndarray]: + """Return energy, force, and atomic energy ordered by global atom ID.""" + atom_ids = np.array( + lammps.lmp.numpy.extract_atom("id")[:natoms], dtype=np.int64, copy=True + ) + order = np.argsort(atom_ids) + force = np.array( + lammps.lmp.numpy.extract_atom("f")[:natoms], dtype=np.float64, copy=True + )[order] + atom_energy = np.asarray(lammps.variables["peatom"].value, dtype=np.float64)[order] + return float(lammps.eval("pe")), force, atom_energy + + +def _run_compact_mpi_scenario( + model: Path, tmp_path: Path, scenario: str, nprocs: int +) -> tuple[float, np.ndarray]: + """Run a compact selection scenario through the MPI LAMMPS executable.""" + mpirun = shutil.which("mpirun") + lmp = shutil.which("lmp") + if mpirun is None or lmp is None: + pytest.skip("MPI compact tests require mpirun and the lmp executable") + + if scenario == "cross_domain_molecule": + # The cutoff hit crosses the x-domain boundary, while atom 3 verifies + # that the complete molecule is selected on the remote rank. + atom_rows = [ + (1, 14.5, 5.0, 5.0, 1), + (2, 15.5, 5.0, 5.0, 10), + (1, 25.0, 5.0, 5.0, 10), + (2, 2.0, 20.0, 5.0, 20), + (1, 3.0, 20.0, 5.0, 20), + ] + elif scenario == "empty_selected_rank": + # With a 2x1x1 processor grid, all selected atoms are owned by rank 0; + # rank 1 must still participate in the backend call and reductions. + atom_rows = [ + (1, 2.0, 5.0, 5.0, 1), + (2, 3.0, 5.0, 5.0, 10), + (1, 4.0, 5.0, 5.0, 10), + (2, 22.0, 20.0, 5.0, 20), + (1, 23.0, 20.0, 5.0, 20), + ] + else: + raise ValueError(f"unknown compact MPI scenario: {scenario}") + + run_dir = tmp_path / f"{scenario}_{nprocs}" + run_dir.mkdir() + input_file = run_dir / "in.compact" + energy_file = run_dir / "energy.out" + force_file = run_dir / "forces.dump" + commands = [] + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + commands.append(f"plugin load {Path(plugin).resolve()}") + commands.extend( + [ + "units metal", + f"processors {nprocs} 1 1", + "boundary p p p", + "atom_style molecular", + "atom_modify map array", + "region box block 0 30 0 30 0 30 units box", + "create_box 2 box", + ] + ) + for atom_type, x, y, z, _ in atom_rows: + commands.append(f"create_atoms {atom_type} single {x} {y} {z} units box") + for atom_id, (_, _, _, _, molecule_id) in enumerate(atom_rows, start=1): + commands.append(f"set atom {atom_id} mol {molecule_id}") + commands.extend( + [ + "group qm id 1", + "mass 1 16", + "mass 2 2", + "neighbor 2.0 bin", + "neigh_modify every 10 delay 0 check no", + f"pair_style deepmd {model.resolve()} center_group qm " + "environment_cutoff 1.5 include_molecule yes", + "pair_coeff * *", + "run 0", + "variable compact_energy equal pe", + f'print "${{compact_energy}}" file {energy_file} screen no', + f"write_dump all custom {force_file} id fx fy fz modify sort id", + ] + ) + input_file.write_text("\n".join(commands) + "\n") + + result = sp.run( + [mpirun, "-n", str(nprocs), lmp, "-in", str(input_file)], + cwd=run_dir, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + energy = float(energy_file.read_text().strip()) + force_rows = np.loadtxt(force_file, skiprows=9, ndmin=2) + assert force_rows[:, 0].astype(np.int64).tolist() == list( + range(1, len(atom_rows) + 1) + ) + return energy, force_rows[:, 1:] + + +def _make_selection_transition_system( + model: Path, environment_x: float, *, move_environment: bool +) -> PyLammps: + """Create a two-atom system for selection cache-invalidation coverage.""" + lammps = PyLammps() + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + lammps.lmp.command(f"plugin load {plugin}") + lammps.lmp.commands_list( + [ + "units metal", + "boundary p p p", + "atom_style molecular", + "atom_modify map array", + "region box block 0 20 0 20 0 20 units box", + "create_box 2 box", + "create_atoms 1 single 5 5 5 units box", + f"create_atoms 2 single {environment_x} 5 5 units box", + "set atom 1 mol 1", + "set atom 2 mol 10", + "group qm id 1", + "group environment id 2", + "mass 1 16", + "mass 2 2", + "neighbor 2.0 bin", + "neigh_modify every 10 delay 0 check no", + f"pair_style deepmd {model.resolve()} center_group qm " + "environment_cutoff 1.5 include_molecule yes", + "pair_coeff * *", + "compute peatom all pe/atom pair", + "variable peatom atom c_peatom", + ] + ) + if move_environment: + lammps.lmp.commands_list( + [ + "timestep 1.0", + "fix mover environment move linear -1.0 0.0 0.0 units box", + "run 1", + ] + ) + else: + lammps.run(0) + return lammps + + +def _make_special_bond_system(model: Path, *, compact: bool) -> PyLammps: + """Create a bonded center/environment pair excluded from its neighbor list.""" + lammps = PyLammps() + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + lammps.lmp.command(f"plugin load {plugin}") + style = f"deepmd {model.resolve()}" + if compact: + style += " center_group qm environment_cutoff 1.5 include_molecule no" + lammps.lmp.commands_list( + [ + "units metal", + "boundary p p p", + "atom_style molecular", + "atom_modify map array", + "region box block 0 20 0 20 0 20 units box", + "create_box 2 box bond/types 1 extra/bond/per/atom 1 extra/special/per/atom 1", + "create_atoms 1 single 5 5 5 units box", + "create_atoms 2 single 6 5 5 units box", + "set atom 1 mol 1", + "set atom 2 mol 2", + "group qm id 1", + "mass 1 16", + "mass 2 2", + "bond_style zero", + "bond_coeff 1", + "create_bonds single/bond 1 1 2", + "special_bonds lj 0 0 0 coul 0 0 0", + "neighbor 2.0 bin", + style, + "pair_coeff * *", + "compute peatom all pe/atom pair", + "variable peatom atom c_peatom", + "run 0", + ] + ) + return lammps + + +@pytest.mark.parametrize("triclinic", [False, True]) +def test_compact_matches_explicit_selected_subsystem( + compact_models: tuple[Path, Path], triclinic: bool +) -> None: + """Compact evaluation preserves whole molecules and scatters zero outputs.""" + model, _ = compact_models + compact_lmp = _make_system( + (model,), + with_inactive_molecule=True, + compact=True, + triclinic=triclinic, + ) + reference_lmp = _make_system( + (model,), + with_inactive_molecule=False, + compact=False, + triclinic=triclinic, + ) + try: + compact_energy, compact_force, compact_atom_energy = _snapshot(compact_lmp, 5) + reference_energy, reference_force, reference_atom_energy = _snapshot( + reference_lmp, 3 + ) + assert compact_energy == pytest.approx(reference_energy) + assert compact_force[:3] == pytest.approx(reference_force) + assert compact_atom_energy[:3] == pytest.approx(reference_atom_energy) + assert compact_force[3:] == pytest.approx(0.0) + assert compact_atom_energy[3:] == pytest.approx(0.0) + # Atom 3 is outside the cutoff, so its nonzero atomic bias can enter + # only through whole-molecule promotion from atom 2. + assert abs(compact_atom_energy[2]) > 1.0 + finally: + compact_lmp.close() + reference_lmp.close() + + +def test_compact_packing_matches_generic_backend_selection( + compact_models: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + """LAMMPS-side packing matches the generic type-minus-one fallback.""" + model, _ = compact_models + monkeypatch.delenv("DP_LAMMPS_DISABLE_COMPACT_PACKING", raising=False) + packed_lmp = _make_system( + (model,), + with_inactive_molecule=True, + compact=True, + triclinic=True, + ) + monkeypatch.setenv("DP_LAMMPS_DISABLE_COMPACT_PACKING", "1") + generic_lmp = _make_system( + (model,), + with_inactive_molecule=True, + compact=True, + triclinic=True, + ) + try: + packed = _snapshot(packed_lmp, 5) + generic = _snapshot(generic_lmp, 5) + assert packed[0] == pytest.approx(generic[0]) + assert packed[1] == pytest.approx(generic[1]) + assert packed[2] == pytest.approx(generic[2]) + finally: + packed_lmp.close() + generic_lmp.close() + + +def test_compact_model_deviation_uses_selected_atoms_only( + compact_models: tuple[Path, Path], tmp_path: Path +) -> None: + """Deviation summaries use the compact count while excluded atoms stay zero.""" + compact_output = tmp_path / "compact_devi.out" + reference_output = tmp_path / "reference_devi.out" + compact_lmp = _make_system( + compact_models, + with_inactive_molecule=True, + compact=True, + triclinic=False, + deviation_file=compact_output, + ) + reference_lmp = _make_system( + compact_models, + with_inactive_molecule=False, + compact=False, + triclinic=False, + deviation_file=reference_output, + ) + try: + compact_deviation = np.loadtxt(compact_output, ndmin=1) + reference_deviation = np.loadtxt(reference_output, ndmin=1) + assert compact_deviation[:7] == pytest.approx(reference_deviation[:7]) + assert compact_deviation[7:10] == pytest.approx(reference_deviation[7:]) + assert compact_deviation[10:] == pytest.approx(0.0) + finally: + compact_lmp.close() + reference_lmp.close() + + +def test_compact_selection_change_rebuilds_backend_cache( + compact_models: tuple[Path, Path], +) -> None: + """An atom may enter the compact set between LAMMPS neighbor rebuilds.""" + model, _ = compact_models + moving_lmp = _make_selection_transition_system(model, 7.2, move_environment=True) + reference_lmp = _make_selection_transition_system( + model, 6.2, move_environment=False + ) + try: + moving = _snapshot(moving_lmp, 2) + reference = _snapshot(reference_lmp, 2) + assert moving[0] == pytest.approx(reference[0]) + assert moving[1] == pytest.approx(reference[1]) + assert moving[2] == pytest.approx(reference[2]) + finally: + moving_lmp.close() + reference_lmp.close() + + +def test_compact_selection_recovers_special_bonds_excluded_from_neighbor_list( + compact_models: tuple[Path, Path], +) -> None: + """Selection remains independent of zero-valued special_bonds factors.""" + model, _ = compact_models + compact_lmp = _make_special_bond_system(model, compact=True) + reference_lmp = _make_special_bond_system(model, compact=False) + try: + compact = _snapshot(compact_lmp, 2) + reference = _snapshot(reference_lmp, 2) + assert compact[0] == pytest.approx(reference[0]) + assert compact[1] == pytest.approx(reference[1]) + assert compact[2] == pytest.approx(reference[2]) + assert abs(compact[2][1]) > 1.0 + finally: + compact_lmp.close() + reference_lmp.close() + + +def test_compact_rejects_center_type_not_represented_by_model( + compact_models: tuple[Path, Path], +) -> None: + """A nonnegative pair mapping is insufficient if the model lacks the type.""" + model, _ = compact_models + lammps = PyLammps() + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + lammps.lmp.command(f"plugin load {plugin}") + lammps.lmp.commands_list( + [ + "units metal", + "boundary p p p", + "atom_style atomic", + "region box block 0 20 0 20 0 20 units box", + "create_box 3 box", + "create_atoms 3 single 5 5 5 units box", + "group qm id 1", + "mass 1 16", + "mass 2 2", + "mass 3 1", + f"pair_style deepmd {model.resolve()} center_group qm " + "environment_cutoff 1.5 include_molecule no", + "pair_coeff * *", + ] + ) + try: + with pytest.raises( + Exception, + match=r"center_group.*type is not represented by the DeepMD model", + ): + lammps.run(0) + finally: + lammps.close() + + +@pytest.mark.parametrize( + "exclusion", + [ + "exclude type 1 2", + "exclude group qm environment", + "exclude molecule/intra all", + "exclude molecule/inter all", + ], +) +def test_compact_rejects_neighbor_exclusions( + compact_models: tuple[Path, Path], exclusion: str +) -> None: + """Compact selection must not silently omit user-excluded environments.""" + model, _ = compact_models + lammps = PyLammps() + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + lammps.lmp.command(f"plugin load {plugin}") + lammps.lmp.commands_list( + [ + "units metal", + "boundary p p p", + "atom_style molecular", + "atom_modify map array", + "region box block 0 20 0 20 0 20 units box", + "create_box 2 box", + "create_atoms 1 single 5 5 5 units box", + "create_atoms 2 single 6 5 5 units box", + "set atom 1 mol 1", + "set atom 2 mol 2", + "group qm id 1", + "group environment id 2", + "mass 1 16", + "mass 2 2", + "neighbor 2.0 bin", + f"neigh_modify {exclusion}", + f"pair_style deepmd {model.resolve()} center_group qm " + "environment_cutoff 1.5 include_molecule no", + "pair_coeff * *", + ] + ) + try: + with pytest.raises( + Exception, + match=r"compact pair_style deepmd does not support neigh_modify exclude", + ): + lammps.run(0) + finally: + lammps.close() + + +@pytest.mark.parametrize("scenario", ["cross_domain_molecule", "empty_selected_rank"]) +def test_compact_mpi_matches_serial( + compact_models: tuple[Path, Path], tmp_path: Path, scenario: str +) -> None: + """Two-rank selection matches serial, including remote and empty ranks.""" + model, _ = compact_models + serial_energy, serial_force = _run_compact_mpi_scenario( + model, tmp_path, scenario, 1 + ) + mpi_energy, mpi_force = _run_compact_mpi_scenario(model, tmp_path, scenario, 2) + assert mpi_energy == pytest.approx(serial_energy) + assert mpi_force == pytest.approx(serial_force) diff --git a/source/lmp/tests/test_lammps_option_parsers.py b/source/lmp/tests/test_lammps_option_parsers.py index 6a9ffda706..aa94a54550 100644 --- a/source/lmp/tests/test_lammps_option_parsers.py +++ b/source/lmp/tests/test_lammps_option_parsers.py @@ -78,6 +78,21 @@ def _assert_lammps_error(result: sp.CompletedProcess[str], message: str) -> None [ ("deepmd", "relative", "Illegal relative, not provided"), ("deepmd", "relative_v", "Illegal relative_v, not provided"), + ( + "deepmd", + "center_group", + "Illegal center_group, group ID is not provided", + ), + ( + "deepmd", + "environment_cutoff", + "Illegal environment_cutoff, value is not provided", + ), + ( + "deepmd", + "include_molecule", + "Illegal include_molecule, yes/no is not provided", + ), ("deepspin", "relative", "Illegal relative, not provided"), ("deepspin", "relative_v", "Illegal relative_v, not provided"), ( @@ -146,6 +161,38 @@ def test_deepspin_accepts_complete_option_values(spin_model: Path) -> None: assert result.returncode == 0, result.stdout + result.stderr +@pytest.mark.parametrize( + ("options", "message"), + [ + ( + "center_group qm", + "center_group requires environment_cutoff in pair_style deepmd", + ), + ( + "environment_cutoff 2.0", + "environment_cutoff and include_molecule require center_group in " + "pair_style deepmd", + ), + ( + "include_molecule no", + "environment_cutoff and include_molecule require center_group in " + "pair_style deepmd", + ), + ( + "center_group qm environment_cutoff 0", + "environment_cutoff must be a finite value greater than zero", + ), + ], +) +def test_deepmd_rejects_incomplete_compact_options( + spin_model: Path, options: str, message: str +) -> None: + result = _run_lammps( + "units metal", f"pair_style deepmd {spin_model.resolve()} {options}" + ) + _assert_lammps_error(result, message) + + @pytest.mark.parametrize( ("fix_options", "message"), [