First commit (#1)

Signed-off-by: Jinzhe Zeng <[email protected]>

.clang-format

@@ -0,0 +1,4 @@
+---
+BasedOnStyle: Google
+BinPackParameters: false
+InsertBraces: true

.github/workflows/release.yaml

@@ -17,7 +17,7 @@ jobs:
       - name: Set up uv
         run: curl -LsSf https://astral.sh/uv/install.sh | sh
       - name: Build dist
-        run: uv tool run --with build[uv] --from build python -m build --installer uv
+        run: uv tool run --with build[uv] --from build python -m build --installer uv --sdist
       - name: Upload release distributions
         uses: actions/upload-artifact@v4
         with:

.github/workflows/test.yaml

@@ -10,13 +10,12 @@ env:
   UV_SYSTEM_PYTHON: 1
   # Configure a constant location for the uv cache
   UV_CACHE_DIR: /tmp/.uv-cache
-permissions:
-  contents: read
 jobs:
   build:
     runs-on: ubuntu-latest
     permissions:
       id-token: write
+      contents: read
     steps:
       - uses: actions/checkout@v4
       - name: Set up uv

.gitignore

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

.pre-commit-config.yaml

@@ -29,9 +29,19 @@ repos:
     rev: v4.0.0-alpha.8
     hooks:
       - id: prettier
-        types_or: [javascript, css, html, markdown, json, yaml]
+        types_or: [javascript, css, html, markdown, yaml]
   - repo: https://github.com/abravalheri/validate-pyproject
     rev: v0.19
     hooks:
       - id: validate-pyproject
         additional_dependencies: ["validate-pyproject-schema-store[all]"]
+  # C++
+  - repo: https://github.com/pre-commit/mirrors-clang-format
+    rev: v18.1.8
+    hooks:
+      - id: clang-format
+  # CMake
+  - repo: https://github.com/cheshirekow/cmake-format-precommit
+    rev: v0.6.13
+    hooks:
+      - id: cmake-format

CMakeLists.txt

@@ -0,0 +1,67 @@
+cmake_minimum_required(VERSION 3.15)
+project(DeePMD-MACE CXX)
+
+set(CMAKE_CXX_STANDARD 14)
+macro(set_if_higher VARIABLE VALUE)
+  # ${VARIABLE} is a variable name, not a string
+  if(${VARIABLE} LESS "${VALUE}")
+    set(${VARIABLE} ${VALUE})
+  endif()
+endmacro()
+
+# build cpp or python interfaces
+option(BUILD_CPP_IF "Build C++ interfaces" ON)
+option(BUILD_PY_IF "Build Python interfaces" OFF)
+option(USE_PT_PYTHON_LIBS "Use PyTorch Python libraries" OFF)
+
+if((NOT BUILD_PY_IF) AND (NOT BUILD_CPP_IF))
+  # nothing to do
+  message(FATAL_ERROR "Nothing to build.")
+endif()
+
+if(BUILD_CPP_IF
+   AND USE_PT_PYTHON_LIBS
+   AND NOT CMAKE_CROSSCOMPILING
+   AND NOT SKBUILD)
+  find_package(
+    Python
+    COMPONENTS Interpreter
+    REQUIRED)
+  execute_process(
+    COMMAND ${Python_EXECUTABLE} -c
+            "import torch;print(torch.utils.cmake_prefix_path)"
+    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
+    OUTPUT_VARIABLE PYTORCH_CMAKE_PREFIX_PATH
+    RESULT_VARIABLE PYTORCH_CMAKE_PREFIX_PATH_RESULT_VAR
+    ERROR_VARIABLE PYTORCH_CMAKE_PREFIX_PATH_ERROR_VAR
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
+  if(NOT ${PYTORCH_CMAKE_PREFIX_PATH_RESULT_VAR} EQUAL 0)
+    message(
+      FATAL_ERROR
+        "Cannot determine PyTorch CMake prefix path, error code: $PYTORCH_CMAKE_PREFIX_PATH_RESULT_VAR}, error message: ${PYTORCH_CMAKE_PREFIX_PATH_ERROR_VAR}"
+    )
+  endif()
+  list(APPEND CMAKE_PREFIX_PATH ${PYTORCH_CMAKE_PREFIX_PATH})
+endif()
+find_package(Torch REQUIRED)
+if(NOT Torch_VERSION VERSION_LESS "2.1.0")
+  set_if_higher(CMAKE_CXX_STANDARD 17)
+elseif(NOT Torch_VERSION VERSION_LESS "1.5.0")
+  set_if_higher(CMAKE_CXX_STANDARD 14)
+endif()
+string(REGEX MATCH "_GLIBCXX_USE_CXX11_ABI=([0-9]+)" CXXABI_PT_MATCH
+             "${TORCH_CXX_FLAGS}")
+if(CXXABI_PT_MATCH)
+  set(OP_CXX_ABI_PT ${CMAKE_MATCH_1})
+  message(STATUS "PyTorch CXX11 ABI: ${CMAKE_MATCH_1}")
+else()
+  # Maybe in macos/windows
+  set(OP_CXX_ABI_PT 0)
+endif()
+
+# define build type
+if((NOT DEFINED CMAKE_BUILD_TYPE) OR CMAKE_BUILD_TYPE STREQUAL "")
+  set(CMAKE_BUILD_TYPE release)
+endif()
+
+add_subdirectory(op)

README.md

@@ -1,5 +1,108 @@
-# Python template
+# MACE plugin for DeePMD-kit
 
 <!-- [![PyPI - Version](https://img.shields.io/pypi/v/python-template)](https://pypi.org/p/python-template) -->
 
-`python-template` is a template to start a new Python project quickly.
+`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.
+
+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).
+
+## Installation
+
+First, clone this repository:
+
+```sh
+git clone https://github.com/njzjz/deepmd-mace
+cd deepmd-mace
+```
+
+### Python interface plugin
+
+Python 3.9 or above is required. A C++ compiler that supports C++ 14 (for PyTorch 2.0) or C++ 17 (for PyTorch 2.1 or above) is required.
+
+Assume you have installed [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit) (v3.0.0b2 or above) and [PyTorch](https://github.com/pytorch/pytorch) in an environment, then execute
+
+```sh
+# expose PyTorch CMake modules
+export CMAKE_PREFIX_PATH=$(python -c "import torch;print(torch.utils.cmake_prefix_path)")
+
+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).
+
+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
+
+```sh
+# Assume libtorch has been contained in CMAKE_PREFIX_PATH
+mkdir -p build
+cd build
+cmake .. -D CMAKE_INSTALL_PREFIX=/prefix/to/install
+cmake --build . -j8
+cmake --install .
+```
+
+`libdeepmd_mace.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
+```
+
+## Usage
+
+Follow [Parameters section](#parameters) to prepare a DeePMD-kit input file.
+
+```sh
+dp --pt train input.json
+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/).
+
+## Parameters
+
+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:
+
+```json
+"model": {
+  "type": "mace",
+  "type_map": [
+    "O",
+    "H"
+  ],
+  "r_max": 5.0,
+  "sel": "auto",
+  "num_radial_basis": 8,
+  "num_cutoff_basis": 5,
+  "max_ell": 3,
+  "interaction": "RealAgnosticResidualInteractionBlock",
+  "num_interactions": 2,
+  "hidden_irreps": "128x0e + 128x1o",
+  "pair_repulsion": false,
+  "distance_transform": "None",
+  "correlation": 3,
+  "gate": "silu",
+  "MLP_irreps": "16x0e",
+  "radial_type": "bessel",
+  "radial_MLP": [64, 64, 64],
+  "std": null
+}
+```
+
+## 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.
+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.
+
+## Examples
+
+- [examples/water]
+- [examples/dprc]

python_template/__init__.py → deepmd_mace/__init__.py

@@ -1,11 +1,11 @@
-"""A Python template (please revise this docstring)."""
-
-from __future__ import annotations
+"""MACE plugin for DeePMD-kit."""
 
 from ._version import __version__
+from .argcheck import mace_model_args
 
 __email__ = "[email protected]"
 
 __all__ = [
     "__version__",
+    "mace_model_args",
 ]

python_template/__main__.py → deepmd_mace/__main__.py

@@ -1,7 +1,5 @@
 """Main entry point for the command line interface."""
 
-from __future__ import annotations
-
 if __name__ == "__main__":
     msg = "This module is not meant to be executed directly."
     raise NotImplementedError(msg)

deepmd_mace/argcheck.py

@@ -0,0 +1,113 @@
+"""Argument check for the MACE model."""
+
+from dargs import Argument
+from deepmd.utils.argcheck import model_args_plugin
+
+
+@model_args_plugin.register("mace")
+def mace_model_args() -> Argument:
+    """Arguments for the MACE model.
+
+    Returns
+    -------
+    Argument
+        Arguments for the MACE model.
+    """
+    doc_r_max = "distance cutoff (in Ang)"
+    doc_num_radial_basis = "number of radial basis functions"
+    doc_num_cutoff_basis = "number of basis functions for smooth cutoff"
+    doc_max_ell = "highest ell of spherical harmonics"
+    doc_interaction = "name of interaction block"
+    doc_num_interactions = "number of interactions"
+    doc_hidden_irreps = "hidden irreps"
+    doc_pair_repulsion = "use pair repulsion term with ZBL potential"
+    doc_distance_transform = "distance transform"
+    doc_correlation = "correlation order at each layer"
+    doc_gate = "non linearity for last readout"
+    doc_mlp_irreps = "hidden irreps of the MLP in last readout"
+    doc_radial_type = "type of radial basis functions"
+    doc_radial_mlp = "width of the radial MLP"
+    doc_std = "Standard deviation of force components in the training set"
+    return Argument(
+        "mace",
+        dict,
+        [
+            Argument("sel", [int, str], optional=False),
+            Argument("r_max", float, optional=True, default=5.0, doc=doc_r_max),
+            Argument(
+                "num_radial_basis",
+                int,
+                optional=True,
+                default=8,
+                doc=doc_num_radial_basis,
+            ),
+            Argument(
+                "num_cutoff_basis",
+                int,
+                optional=True,
+                default=5,
+                doc=doc_num_cutoff_basis,
+            ),
+            Argument("max_ell", int, optional=True, default=3, doc=doc_max_ell),
+            Argument(
+                "interaction",
+                str,
+                optional=True,
+                default="RealAgnosticResidualInteractionBlock",
+                doc=doc_interaction,
+            ),
+            Argument(
+                "num_interactions",
+                int,
+                optional=True,
+                default=2,
+                doc=doc_num_interactions,
+            ),
+            Argument(
+                "hidden_irreps",
+                str,
+                optional=True,
+                default="128x0e + 128x1o",
+                doc=doc_hidden_irreps,
+            ),
+            Argument(
+                "pair_repulsion",
+                bool,
+                optional=True,
+                default=False,
+                doc=doc_pair_repulsion,
+            ),
+            Argument(
+                "distance_transform",
+                str,
+                optional=True,
+                default="None",
+                doc=doc_distance_transform,
+            ),
+            Argument("correlation", int, optional=True, default=3, doc=doc_correlation),
+            Argument("gate", str, optional=True, default="silu", doc=doc_gate),
+            Argument(
+                "MLP_irreps",
+                str,
+                optional=True,
+                default="16x0e",
+                doc=doc_mlp_irreps,
+            ),
+            Argument(
+                "radial_type",
+                str,
+                optional=True,
+                default="bessel",
+                doc=doc_radial_type,
+            ),
+            Argument(
+                "radial_MLP",
+                list[int],
+                optional=True,
+                default=[64, 64, 64],
+                doc=doc_radial_mlp,
+            ),
+            Argument("std", float, optional=True, doc=doc_std),
+        ],
+        doc="MACE model",
+    )