From fa5f0b06f90325ee1876f5f8e72db5ac2d1c1a68 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 00:41:38 -0500 Subject: [PATCH 1/9] feat: allow using mapping from LAMMPS Signed-off-by: Jinzhe Zeng --- README.md | 23 +++++++++++++++++++++++ deepmd_gnn/env.py | 3 +++ deepmd_gnn/mace.py | 10 ++++++++++ deepmd_gnn/nequip.py | 12 +++++++++++- 4 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 deepmd_gnn/env.py diff --git a/README.md b/README.md index 7abacbc..02d7338 100644 --- a/README.md +++ b/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 diff --git a/deepmd_gnn/env.py b/deepmd_gnn/env.py new file mode 100644 index 0000000..25c000f --- /dev/null +++ b/deepmd_gnn/env.py @@ -0,0 +1,3 @@ +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..a553161 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 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, diff --git a/deepmd_gnn/nequip.py b/deepmd_gnn/nequip.py index 81ff3bf..a227263 100644 --- a/deepmd_gnn/nequip.py +++ b/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, From 8f5f477309815f87014558577be05d66f0c5bf3b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 23 Nov 2024 05:42:05 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- README.md | 3 +++ deepmd_gnn/mace.py | 2 +- deepmd_gnn/nequip.py | 3 ++- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 02d7338..b86a2b5 100644 --- a/README.md +++ b/README.md @@ -89,10 +89,13 @@ 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 ``` diff --git a/deepmd_gnn/mace.py b/deepmd_gnn/mace.py index a553161..8f4565e 100644 --- a/deepmd_gnn/mace.py +++ b/deepmd_gnn/mace.py @@ -535,7 +535,7 @@ def forward_lower( # 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`." + "If you are using LAMMPS, set `atom_modify map yes`.", ) nlist = build_neighbor_list( extended_coord.view(nf, -1), diff --git a/deepmd_gnn/nequip.py b/deepmd_gnn/nequip.py index a227263..c9e6f7d 100644 --- a/deepmd_gnn/nequip.py +++ b/deepmd_gnn/nequip.py @@ -48,6 +48,7 @@ from deepmd_gnn import env + @BaseModel.register("nequip") class NequipModel(BaseModel): """Nequip model. @@ -411,7 +412,7 @@ def forward_lower( # 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`." + "If you are using LAMMPS, set `atom_modify map yes`.", ) nlist = build_neighbor_list( extended_coord.view(nf, -1), From b649e350aa5277345e29b265df26ca13bb7c2dd1 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 00:44:07 -0500 Subject: [PATCH 3/9] improve README --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index b86a2b5..da06598 100644 --- a/README.md +++ b/README.md @@ -87,7 +87,7 @@ where $r_c$ is set by `r_max` and $N_L$ is set by `num_interactions` (MACE) / `n 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. +The alternative 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 @@ -95,6 +95,8 @@ 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 From e638befe210802e4a4d83bdafecb62a6deb8cd89 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 00:46:46 -0500 Subject: [PATCH 4/9] fix Signed-off-by: Jinzhe Zeng --- deepmd_gnn/env.py | 1 + deepmd_gnn/mace.py | 7 ++++--- deepmd_gnn/nequip.py | 7 ++++--- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/deepmd_gnn/env.py b/deepmd_gnn/env.py index 25c000f..b4fd89b 100644 --- a/deepmd_gnn/env.py +++ b/deepmd_gnn/env.py @@ -1,3 +1,4 @@ +"""Configurations read from environement 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 8f4565e..7880734 100644 --- a/deepmd_gnn/mace.py +++ b/deepmd_gnn/mace.py @@ -54,7 +54,7 @@ ) import deepmd_gnn.op # noqa: F401 -from deepmd_gnn import env +from deepmd_gnn import env as deepmd_gnn_env ELEMENTS = [ "H", @@ -369,7 +369,7 @@ 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: + if deepmd_gnn_env.DP_GNN_USE_MAPPING: return self.rcut return self.rcut * self.num_interactions @@ -533,10 +533,11 @@ def forward_lower( if env.DP_GNN_USE_MAPPING: # when setting DP_GNN_USE_MAPPING, ghost atoms are only built # for one message-passing layer - raise ValueError( + 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, diff --git a/deepmd_gnn/nequip.py b/deepmd_gnn/nequip.py index c9e6f7d..d2836b5 100644 --- a/deepmd_gnn/nequip.py +++ b/deepmd_gnn/nequip.py @@ -46,7 +46,7 @@ ) from nequip.model import model_from_config -from deepmd_gnn import env +from deepmd_gnn import env as deepmd_gnn_env @BaseModel.register("nequip") @@ -245,7 +245,7 @@ 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: + if deepmd_gnn_env.DP_GNN_USE_MAPPING: return self.rcut return self.rcut * self.num_layers @@ -410,10 +410,11 @@ def forward_lower( if env.DP_GNN_USE_MAPPING: # when setting DP_GNN_USE_MAPPING, ghost atoms are only built # for one message-passing layer - raise ValueError( + 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, From 1a86778ba82088051031353000eaee9748686c5b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 23 Nov 2024 05:46:58 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd_gnn/env.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd_gnn/env.py b/deepmd_gnn/env.py index b4fd89b..71b8369 100644 --- a/deepmd_gnn/env.py +++ b/deepmd_gnn/env.py @@ -1,4 +1,5 @@ """Configurations read from environement variables.""" + import os DP_GNN_USE_MAPPING = os.environ.get("DP_GNN_USE_MAPPING", "0") == "1" From b8c1c7f534076ee72aadb6436230af998114a281 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 00:51:55 -0500 Subject: [PATCH 6/9] fix Signed-off-by: Jinzhe Zeng --- deepmd_gnn/mace.py | 2 +- deepmd_gnn/nequip.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd_gnn/mace.py b/deepmd_gnn/mace.py index 7880734..6a72486 100644 --- a/deepmd_gnn/mace.py +++ b/deepmd_gnn/mace.py @@ -530,7 +530,7 @@ 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: + 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 = ( diff --git a/deepmd_gnn/nequip.py b/deepmd_gnn/nequip.py index d2836b5..a68a70d 100644 --- a/deepmd_gnn/nequip.py +++ b/deepmd_gnn/nequip.py @@ -407,7 +407,7 @@ def forward_lower( nf, nall = extended_atype.shape # recalculate nlist for ghost atoms if mapping is None and self.num_layers > 1 and nloc < nall: - if env.DP_GNN_USE_MAPPING: + 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 = ( From d220cdba4cfbf99ca48e01863175e6ee47ce2bde Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 01:00:12 -0500 Subject: [PATCH 7/9] revert nequip Signed-off-by: Jinzhe Zeng --- deepmd_gnn/nequip.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/deepmd_gnn/nequip.py b/deepmd_gnn/nequip.py index a68a70d..81ff3bf 100644 --- a/deepmd_gnn/nequip.py +++ b/deepmd_gnn/nequip.py @@ -46,8 +46,6 @@ ) from nequip.model import model_from_config -from deepmd_gnn import env as deepmd_gnn_env - @BaseModel.register("nequip") class NequipModel(BaseModel): @@ -245,8 +243,6 @@ 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_layers @torch.jit.export @@ -406,15 +402,7 @@ def forward_lower( nloc = nlist.shape[1] nf, nall = extended_atype.shape # recalculate nlist for ghost atoms - if mapping is None and self.num_layers > 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) + if self.num_layers > 1 and nloc < nall: nlist = build_neighbor_list( extended_coord.view(nf, -1), extended_atype, From 135f5a505d12dc44bbf743cf907219e81f7100c9 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 01:01:35 -0500 Subject: [PATCH 8/9] update README Signed-off-by: Jinzhe Zeng --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index da06598..7de56a8 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,7 @@ 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 +### Running LAMMPS + MACE with period boundry conditions GNN models use message passing neural networks, so the neighbor list built with traditional cutoff radius will not work, @@ -87,7 +87,7 @@ where $r_c$ is set by `r_max` and $N_L$ is set by `num_interactions` (MACE) / `n 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 does not support MPI. +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 From 48edc0f22b76478364318ee19b2b75e2aca426d1 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 23 Nov 2024 01:08:35 -0500 Subject: [PATCH 9/9] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- README.md | 2 +- deepmd_gnn/env.py | 2 +- deepmd_gnn/mace.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 7de56a8..43326c9 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,7 @@ 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 boundry conditions +### 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, diff --git a/deepmd_gnn/env.py b/deepmd_gnn/env.py index 71b8369..efe5641 100644 --- a/deepmd_gnn/env.py +++ b/deepmd_gnn/env.py @@ -1,4 +1,4 @@ -"""Configurations read from environement variables.""" +"""Configurations read from environment variables.""" import os diff --git a/deepmd_gnn/mace.py b/deepmd_gnn/mace.py index 6a72486..dd553cf 100644 --- a/deepmd_gnn/mace.py +++ b/deepmd_gnn/mace.py @@ -535,7 +535,7 @@ def forward_lower( # 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`.", + "If you are using LAMMPS, set `atom_modify map yes`." ) raise ValueError(msg) nlist = build_neighbor_list(