feat: allow using mapping from LAMMPS

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

README.md

@@ -77,6 +77,29 @@ 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 with period boundry 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 is to use the mapping passed from LAMMPS, which only 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).
+```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,3 @@
+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
 
 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 env.DP_GNN_USE_MAPPING:
+            return self.rcut
         return self.rcut * self.num_interactions
 
     @torch.jit.export
@@ -527,6 +530,13 @@ 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 env.DP_GNN_USE_MAPPING:
+                # when setting DP_GNN_USE_MAPPING, ghost atoms are only built
+                # for one message-passing layer
+                raise ValueError(
+                    "When setting DP_GNN_USE_MAPPING, mapping is required. "
+                    "If you are using LAMMPS, set `atom_modify map yes`."
+                )
             nlist = build_neighbor_list(
                 extended_coord.view(nf, -1),
                 extended_atype,

deepmd_gnn/nequip.py

@@ -46,6 +46,7 @@
 )
 from nequip.model import model_from_config
 
+from deepmd_gnn import env
 
 @BaseModel.register("nequip")
 class NequipModel(BaseModel):
@@ -243,6 +244,8 @@ def fitting_output_def(self) -> FittingOutputDef:
     @torch.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
+        if env.DP_GNN_USE_MAPPING:
+            return self.rcut
         return self.rcut * self.num_layers
 
     @torch.jit.export
@@ -402,7 +405,14 @@ def forward_lower(
         nloc = nlist.shape[1]
         nf, nall = extended_atype.shape
         # recalculate nlist for ghost atoms
-        if self.num_layers > 1 and nloc < nall:
+        if mapping is None and self.num_layers > 1 and nloc < nall:
+            if env.DP_GNN_USE_MAPPING:
+                # when setting DP_GNN_USE_MAPPING, ghost atoms are only built
+                # for one message-passing layer
+                raise ValueError(
+                    "When setting DP_GNN_USE_MAPPING, mapping is required. "
+                    "If you are using LAMMPS, set `atom_modify map yes`."
+                )
             nlist = build_neighbor_list(
                 extended_coord.view(nf, -1),
                 extended_atype,