breaking: rename package to deepmd_gnn

.gitignore

@@ -162,4 +162,4 @@ cython_debug/
 #.idea/
 .ruff_cache/
 node_modules/
-deepmd_mace/_version.py
+deepmd_gnn/_version.py

CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.15)
-project(DeePMD-MACE CXX)
+project(deepmd-gnn CXX)
 
 set(CMAKE_CXX_STANDARD 14)
 macro(set_if_higher VARIABLE VALUE)

README.md

@@ -1,19 +1,24 @@
-# MACE plugin for DeePMD-kit
+# DeePMD-kit plugin for various graph neural network models
 
 <!-- [![PyPI - Version](https://img.shields.io/pypi/v/python-template)](https://pypi.org/p/python-template) -->
 
-`deepmd-mace` is a [MACE](https://github.com/ACEsuit/mace) plugin for [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit), which connects MACE (PyTorch version) and DeePMD-kit by enabling MACE models in DeePMD-kit PyTorch backend.
+`deepmd-gnn` is a [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit) plugin for various graph neural network (GNN) models, which connects DeePMD-kit and atomistic GNN packages by enabling GNN models in DeePMD-kit.
 
-After [installing the plugin](#installation), you can train the MACE model using DeePMD-kit, run active learning cycles for the MACE model using [DP-GEN](https://github.com/deepmodeling/dpgen), perform simulations with the MACE model using molecular dynamic packages supported by DeePMD-kit, such as [LAMMPS](https://github.com/lammps/lammps) and [AMBER](https://ambermd.org/).
-You can follow [DeePMD-kit documentation](https://docs.deepmodeling.com/projects/deepmd/en/latest/) to train the MACE models using its PyTorch backend, after using the specific [MACE parameters](#parameters).
+Supported packages and models include:
+
+- [MACE](https://github.com/ACEsuit/mace) (PyTorch version)
+- [NequIP](https://github.com/mir-group/nequip) (PyTorch version)
+
+After [installing the plugin](#installation), you can train the GNN models using DeePMD-kit, run active learning cycles for the GNN models using [DP-GEN](https://github.com/deepmodeling/dpgen), perform simulations with the MACE model using molecular dynamic packages supported by DeePMD-kit, such as [LAMMPS](https://github.com/lammps/lammps) and [AMBER](https://ambermd.org/).
+You can follow [DeePMD-kit documentation](https://docs.deepmodeling.com/projects/deepmd/en/latest/) to train the GNN models using its PyTorch backend, after using the specific [model parameters](#parameters).
 
 ## Installation
 
 First, clone this repository:
 
 ```sh
-git clone https://github.com/njzjz/deepmd-mace
-cd deepmd-mace
+git clone https://github.com/njzjz/deepmd-gnn
+cd deepmd-gnn
 ```
 
 ### Python interface plugin
@@ -31,7 +36,7 @@ pip install .
 
 ### C++ interface plugin
 
-DeePMD-kit needs to support [customized OP library in C++ interface](https://github.com/deepmodeling/deepmd-kit/pull/4073) (available after Aug 23, 2024).
+DeePMD-kit version should be v3.0.0b4 or later.
 
 Follow [DeePMD-kit documentation](https://docs.deepmodeling.com/projects/deepmd/en/latest/install/install-from-source.html#install-the-c-interface) to install DeePMD-kit C++ interface with PyTorch backend support and other related MD packages.
 After that, you can build the plugin
@@ -45,11 +50,11 @@ cmake --build . -j8
 cmake --install .
 ```
 
-`libdeepmd_mace.so` will be installed into the directory you assign.
+`libdeepmd_gnn.so` will be installed into the directory you assign.
 When using any DeePMD-kit C++ interface, set the following environment variable in advance:
 
 ```sh
-export DP_PLUGIN_PATH=/prefix/to/install/lib/libdeepmd_mace.so
+export DP_PLUGIN_PATH=/prefix/to/install/lib/libdeepmd_gnn.so
 ```
 
 ## Usage
@@ -66,6 +71,8 @@ For details, follow [DeePMD-kit documentation](https://docs.deepmodeling.com/pro
 
 ## Parameters
 
+### MACE
+
 To use the MACE model, set `"type": "mace"` in the `model` section of the training script.
 Below is default values for the MACE model, most of which follows default values in the MACE package:
 
@@ -96,12 +103,41 @@ Below is default values for the MACE model, most of which follows default values
 }
 ```
 
+### NequIP
+
+```json
+"model": {
+  "type": "nequip",
+  "type_map": [
+    "O",
+    "H"
+  ],
+  "r_max": 5.0,
+  "sel": "auto",
+  "num_layers": 4,
+  "l_max": 2,
+  "num_features": 32,
+  "nonlinearity_type": "gate",
+  "parity": true,
+  "num_basis": 8,
+  "BesselBasis_trainable": true,
+  "PolynomialCutoff_p": 6,
+  "invariant_layers": 2,
+  "invariant_neurons": 64,
+  "use_sc": true,
+  "irreps_edge_sh": "0e + 1e",
+  "feature_irreps_hidden": "32x0o + 32x0e + 32x1o + 32x1e",
+  "chemical_embedding_irreps_out": "32x0e",
+  "conv_to_output_hidden_irreps_out": "16x0e"
+}
+```
+
 ## DPRc support
 
-In `deepmd-mace`, the MACE model can be used in a [DPRc](https://docs.deepmodeling.com/projects/deepmd/en/latest/model/dprc.html) way.
+In `deepmd-gnn`, the GNN model can be used in a [DPRc](https://docs.deepmodeling.com/projects/deepmd/en/latest/model/dprc.html) way.
 Type maps that starts with `m` (such as `mH`) or `OW` or `HW` will be recognized as MM types.
 Two MM atoms will not build edges with each other.
-Such MACE+DPRc model can be directly used in AmberTools24.
+Such GNN+DPRc model can be directly used in AmberTools24.
 
 ## Examples
 

deepmd_mace/mace.py → deepmd_gnn/mace.py

@@ -52,7 +52,7 @@
     interaction_classes,
 )
 
-import deepmd_mace.op  # noqa: F401
+import deepmd_gnn.op  # noqa: F401
 
 ELEMENTS = [
     "H",
@@ -614,7 +614,7 @@ def forward_lower_common(
             extended_coord_ff = extended_coord[ff]
             extended_atype_ff = extended_atype[ff]
             nlist_ff = nlist[ff]
-            edge_index = torch.ops.deepmd_mace.mace_edge_index(
+            edge_index = torch.ops.deepmd_gnn.mace_edge_index(
                 nlist_ff,
                 extended_atype_ff,
                 torch.tensor(self.mm_types, dtype=torch.int64, device="cpu"),

deepmd_mace/nequip.py → deepmd_gnn/nequip.py

@@ -495,7 +495,7 @@ def forward_lower_common(
             extended_coord_ff = extended_coord[ff]
             extended_atype_ff = extended_atype[ff]
             nlist_ff = nlist[ff]
-            edge_index = torch.ops.deepmd_mace.mace_edge_index(
+            edge_index = torch.ops.deepmd_gnn.mace_edge_index(
                 nlist_ff,
                 extended_atype_ff,
                 torch.tensor(self.mm_types, dtype=torch.int64, device="cpu"),

deepmd_mace/op.py → deepmd_gnn/op.py

@@ -8,9 +8,9 @@
 
 import torch
 
-import deepmd_mace.lib
+import deepmd_gnn.lib
 
-SHARED_LIB_DIR = Path(deepmd_mace.lib.__path__[0])
+SHARED_LIB_DIR = Path(deepmd_gnn.lib.__path__[0])
 
 
 def load_library(module_name: str) -> None:
@@ -38,4 +38,4 @@ def load_library(module_name: str) -> None:
     torch.ops.load_library(module_file)
 
 
-load_library("deepmd_mace")
+load_library("deepmd_gnn")

op/CMakeLists.txt

@@ -1,22 +1,22 @@
 file(GLOB OP_SRC mace.cc)
 
-add_library(deepmd_mace MODULE ${OP_SRC})
+add_library(deepmd_gnn MODULE ${OP_SRC})
 # link: libdeepmd libtorch
-target_link_libraries(deepmd_mace PRIVATE ${TORCH_LIBRARIES})
+target_link_libraries(deepmd_gnn PRIVATE ${TORCH_LIBRARIES})
 target_compile_definitions(
-  deepmd_mace
+  deepmd_gnn
   PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_PT}>")
 if(APPLE)
-  set_target_properties(deepmd_mace PROPERTIES INSTALL_RPATH "@loader_path")
+  set_target_properties(deepmd_gnn PROPERTIES INSTALL_RPATH "@loader_path")
 else()
-  set_target_properties(deepmd_mace PROPERTIES INSTALL_RPATH "$ORIGIN")
+  set_target_properties(deepmd_gnn PROPERTIES INSTALL_RPATH "$ORIGIN")
 endif()
 
 if(BUILD_PY_IF)
-  install(TARGETS deepmd_mace DESTINATION deepmd_mace/lib/)
+  install(TARGETS deepmd_gnn DESTINATION deepmd_gnn/lib/)
   file(TOUCH ${CMAKE_CURRENT_BINARY_DIR}/__init__.py)
   install(FILES ${CMAKE_CURRENT_BINARY_DIR}/__init__.py
-          DESTINATION deepmd_mace/lib)
+          DESTINATION deepmd_gnn/lib)
 else(BUILD_PY_IF)
-  install(TARGETS deepmd_mace DESTINATION lib/)
+  install(TARGETS deepmd_gnn DESTINATION lib/)
 endif(BUILD_PY_IF)

op/mace.cc

@@ -58,4 +58,6 @@ torch::Tensor mace_edge_index(const torch::Tensor &nlist_tensor,
   return edge_index_tensor.to(nlist_tensor.device());
 }
 
+TORCH_LIBRARY(deepmd_gnn, m) { m.def("mace_edge_index", mace_edge_index); }
+// compatbility with old models freezed by deepmd_mace package
 TORCH_LIBRARY(deepmd_mace, m) { m.def("mace_edge_index", mace_edge_index); }

pyproject.toml

@@ -5,9 +5,9 @@ requires = [
 build-backend = "scikit_build_core.build"
 
 [project]
-name = "deepmd-mace"
+name = "deepmd-gnn"
 dynamic = ["version"]
-description = "MACE plugin for DeePMD-kit."
+description = "DeePMD-kit plugin for graph neural network models."
 authors = [
     { name = "Jinzhe Zeng", email = "[email protected]"},
 ]
@@ -38,11 +38,11 @@ keywords = [
 [project.scripts]
 
 [project.entry-points."deepmd.pt"]
-mace = "deepmd_mace.mace:MaceModel"
-nequip = "deepmd_mace.nequip:NequipModel"
+mace = "deepmd_gnn.mace:MaceModel"
+nequip = "deepmd_gnn.nequip:NequipModel"
 
 [project.urls]
-repository = "https://github.com/njzjz/deepmd_mace"
+repository = "https://github.com/njzjz/deepmd-gnn"
 
 [project.optional-dependencies]
 test = [
@@ -55,15 +55,15 @@ test = [
 wheel.py-api = "py2.py3"
 metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
 sdist.include = [
-    "/deepmd_mace/_version.py",
+    "/deepmd_gnn/_version.py",
 ]
 
 [tool.scikit-build.cmake.define]
 BUILD_PY_IF = true
 BUILD_CPP_IF = false
 
 [tool.setuptools_scm]
-version_file = "deepmd_mace/_version.py"
+version_file = "deepmd_gnn/_version.py"
 
 [tool.ruff.lint]
 select = [
@@ -97,4 +97,4 @@ convention = "numpy"
 ]
 
 [tool.coverage.report]
-include = ["deepmd_mace/*"]
+include = ["deepmd_gnn/*"]

tests/test_examples.py

@@ -7,7 +7,7 @@
 from dargs.check import check
 from deepmd.utils.argcheck import gen_args
 
-from deepmd_mace.argcheck import mace_model_args  # noqa: F401
+from deepmd_gnn.argcheck import mace_model_args  # noqa: F401
 
 example_path = Path(__file__).parent.parent / "examples"
 

tests/test_model.py

@@ -24,8 +24,8 @@
     to_torch_tensor,
 )
 
-from deepmd_mace.mace import MaceModel
-from deepmd_mace.nequip import NequipModel
+from deepmd_gnn.mace import MaceModel
+from deepmd_gnn.nequip import NequipModel
 
 GLOBAL_SEED = 20240822
 

tests/test_version.py

@@ -4,9 +4,9 @@
 
 from importlib.metadata import version
 
-from deepmd_mace import __version__
+from deepmd_gnn import __version__
 
 
 def test_version() -> None:
     """Test version."""
-    assert version("deepmd-mace") == __version__
+    assert version("deepmd-gnn") == __version__