feat: allow using mapping from LAMMPS (#38)

Signed-off-by: Jinzhe Zeng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <[email protected]>

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

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"

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,