diff --git a/README.md b/README.md index 7abacbc..43326c9 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,34 @@ dp --pt freeze A frozen model file named `frozen_model.pth` will be generated. You can use it in the MD packages or other interfaces. For details, follow [DeePMD-kit documentation](https://docs.deepmodeling.com/projects/deepmd/en/latest/). +### Running LAMMPS + MACE with period boundary conditions + +GNN models use message passing neural networks, +so the neighbor list built with traditional cutoff radius will not work, +since the ghost atoms also need to build neighbor list. +By default, the model requests the neighbor list with a cutoff radius of $r_c \times N_{L}$, +where $r_c$ is set by `r_max` and $N_L$ is set by `num_interactions` (MACE) / `num_layers` (NequIP), +and rebuilds the neighbor list for ghost atoms. +However, this approach is very inefficient. + +The alternative approach for the MACE model (note: NequIP doesn't support such approach) is to use the mapping passed from LAMMPS, which does not support MPI. +One needs to set `DP_GNN_USE_MAPPING` when freezing the models, + +```sh +DP_GNN_USE_MAPPING=1 dp --pt freeze +``` + +and request the mapping when using LAMMPS (also requires DeePMD-kit v3.0.0rc0 or above). +By using the mapping, the ghost atoms will be mapped to the real atoms, +so the regular neighbor list with a cutoff radius of $r_c$ can be used. + +```lammps +atom_modify map array +``` + +In the future, we will explore utilizing the MPI to communicate the neighbor list, +while this approach requires a deep hack for external packages. + ## Parameters ### MACE diff --git a/deepmd_gnn/env.py b/deepmd_gnn/env.py new file mode 100644 index 0000000..efe5641 --- /dev/null +++ b/deepmd_gnn/env.py @@ -0,0 +1,5 @@ +"""Configurations read from environment variables.""" + +import os + +DP_GNN_USE_MAPPING = os.environ.get("DP_GNN_USE_MAPPING", "0") == "1" diff --git a/deepmd_gnn/mace.py b/deepmd_gnn/mace.py index 95d9a67..dd553cf 100644 --- a/deepmd_gnn/mace.py +++ b/deepmd_gnn/mace.py @@ -54,6 +54,7 @@ ) import deepmd_gnn.op # noqa: F401 +from deepmd_gnn import env as deepmd_gnn_env ELEMENTS = [ "H", @@ -368,6 +369,8 @@ def fitting_output_def(self) -> FittingOutputDef: @torch.jit.export def get_rcut(self) -> float: """Get the cut-off radius.""" + if deepmd_gnn_env.DP_GNN_USE_MAPPING: + return self.rcut return self.rcut * self.num_interactions @torch.jit.export @@ -527,6 +530,14 @@ def forward_lower( nf, nall = extended_atype.shape # calculate nlist for ghost atoms, as LAMMPS does not calculate it if mapping is None and self.num_interactions > 1 and nloc < nall: + if deepmd_gnn_env.DP_GNN_USE_MAPPING: + # when setting DP_GNN_USE_MAPPING, ghost atoms are only built + # for one message-passing layer + msg = ( + "When setting DP_GNN_USE_MAPPING, mapping is required. " + "If you are using LAMMPS, set `atom_modify map yes`." + ) + raise ValueError(msg) nlist = build_neighbor_list( extended_coord.view(nf, -1), extended_atype,