Spin lmp

.devcontainer/Dockerfile

@@ -0,0 +1,3 @@
+FROM mcr.microsoft.com/devcontainers/cpp:1-ubuntu-24.04
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

.devcontainer/READMD.md

@@ -0,0 +1,35 @@
+# DeePMD-kit devcontainer environment
+
+This [devcontainer](https://vscode.js.cn/docs/devcontainers/devcontainer-cli) environment setups Python and C++ environment to develop DeePMD-kit.
+One can setup locally or use [GitHub Codespaces](https://docs.github.com/en/codespaces) by clicking the Code button on the DeePMD-kit repository page.
+The whole setup process requires about 10 minutes, so one needs to be patient.
+
+## Python environment
+
+The following packages are installed into the Python environment `.venv`:
+
+- DeePMD-kit (in edit mode)
+- Backends including TensorFlow, PyTorch, JAX
+- LAMMPS
+- MPICH
+- CMake
+- pre-commit (including hooks)
+- Test packages including pytest
+- Doc packages including sphinx
+
+## C++ interface
+
+The C++ interface with TensorFlow and PyTorch support is installed into `dp` directory.
+
+When calling and debuging LAMMPS with DeePMD-kit, use the following scripts instead of the regular `lmp`:
+
+- `.devcontainer/lmp`
+- `.devcontainer/gdb_lmp`
+
+## Rebuild
+
+Usually the Python package does not need to reinstall.
+But when one wants to recompile the C++ code, the following scripts can be executed.
+
+- `.devcontainer/build_cxx.sh`
+- `.devcontainer/build_py.sh`

.devcontainer/build_cxx.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+set -ev
+
+NPROC=$(nproc --all)
+SCRIPT_PATH=$(dirname $(realpath -s $0))
+
+export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
+TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+
+mkdir -p ${SCRIPT_PATH}/../buildcxx/
+cd ${SCRIPT_PATH}/../buildcxx/
+cmake -D ENABLE_TENSORFLOW=ON \
+	-D ENABLE_PYTORCH=ON \
+	-D CMAKE_INSTALL_PREFIX=${SCRIPT_PATH}/../dp/ \
+	-D LAMMPS_VERSION=stable_29Aug2024_update1 \
+	-D CMAKE_BUILD_TYPE=Debug \
+	-D BUILD_TESTING:BOOL=TRUE \
+	-D TENSORFLOW_ROOT=${TENSORFLOW_ROOT} \
+	${SCRIPT_PATH}/../source
+cmake --build . -j${NPROC}
+cmake --install .

.devcontainer/build_py.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+set -ev
+
+SCRIPT_PATH=$(dirname $(realpath -s $0))
+cd ${SCRIPT_PATH}/..
+
+uv sync --dev --python 3.12 --extra cpu --extra torch --extra jax --extra lmp --extra test --extra docs
+pre-commit install

.devcontainer/devcontainer.json

@@ -0,0 +1,17 @@
+{
+  "name": "DeePMD-kit",
+  "build": {
+    "dockerfile": "Dockerfile"
+  },
+  "features": {
+    "ghcr.io/devcontainers/features/github-cli:1": {}
+  },
+  "postCreateCommand": ".devcontainer/build_py.sh && .devcontainer/download_libtorch.sh && .devcontainer/build_cxx.sh && pre-commit install-hooks",
+  "remoteEnv": {
+    "PATH": "${containerEnv:PATH}:${containerWorkspaceFolder}/.venv/bin",
+    "DP_ENABLE_PYTORCH": "1",
+    "DP_VARIANT": "cpu",
+    "LMP_CXX11_ABI_0": "1",
+    "UV_EXTRA_INDEX_URL": "https://download.pytorch.org/whl/cpu"
+  }
+}

.devcontainer/download_libtorch.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+set -ev
+
+SCRIPT_PATH=$(dirname $(realpath -s $0))
+cd ${SCRIPT_PATH}/..
+
+wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcpu.zip -O ~/libtorch.zip
+unzip ~/libtorch.zip

.devcontainer/gdb_lmp

@@ -0,0 +1,9 @@
+#!/bin/bash
+SCRIPT_PATH=$(dirname $(realpath -s $0))
+
+export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
+TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+
+env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
+	LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
+	gdb ${SCRIPT_PATH}/../.venv/lib/python3.12/site-packages/lammps/lmp "$@"

.devcontainer/lmp

@@ -0,0 +1,9 @@
+#!/bin/bash
+SCRIPT_PATH=$(dirname $(realpath -s $0))
+
+export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
+TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+
+env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
+	LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
+	${SCRIPT_PATH}/../.venv/bin/lmp "$@"

.github/workflows/build_cc.yml

@@ -11,7 +11,7 @@ name: Build C++
 jobs:
   buildcc:
     name: Build C++
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         include:

.github/workflows/build_wheel.yml

@@ -148,6 +148,7 @@ jobs:
       - uses: actions/download-artifact@v4
         with:
           path: source/install/docker/dist
+          pattern: cibw-*-manylinux_x86_64-cu${{ matrix.cuda_version }}*
           merge-multiple: true
       - name: Log in to the Container registry
         uses: docker/login-action@v3
@@ -180,6 +181,7 @@ jobs:
       - uses: actions/download-artifact@v4
         with:
           path: dist/packages
+          pattern: cibw-*
           merge-multiple: true
       - uses: actions/setup-python@v5
         name: Install Python

.github/workflows/codeql.yml

@@ -13,7 +13,7 @@ concurrency:
 jobs:
   analyze:
     name: Analyze
-    runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }}
+    runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-22.04' }}
     timeout-minutes: ${{ (matrix.language == 'swift' && 120) || 360 }}
     permissions:
       actions: read

.github/workflows/package_c.yml

@@ -18,7 +18,7 @@ jobs:
     strategy:
       matrix:
         include:
-          - tensorflow_build_version: "2.15"
+          - tensorflow_build_version: "2.18"
             tensorflow_version: ""
             filename: libdeepmd_c.tar.gz
           - tensorflow_build_version: "2.14"

.github/workflows/test_cc.yml

@@ -11,7 +11,7 @@ name: Test C++
 jobs:
   testcc:
     name: Test C++
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         check_memleak: [true, false]

.github/workflows/test_cuda.yml

@@ -47,11 +47,11 @@ jobs:
          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
     - run: python -m pip install -U uv
-    - run: source/install/uv_with_retry.sh pip install --system "tensorflow>=2.15.0rc0" "torch==2.3.1.*"
+    - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.5.0" "jax[cuda12]"
     - run: |
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-        source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch] mpi4py
+        source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1
@@ -61,9 +61,11 @@ jobs:
       env:
         NUM_WORKERS: 0
         CUDA_VISIBLE_DEVICES: 0
+        # See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
+        XLA_PYTHON_CLIENT_PREALLOCATE: false
     - name: Download libtorch
       run: |
-         wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.1%2Bcu121.zip -O libtorch.zip
+         wget https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcu124.zip -O libtorch.zip
          unzip libtorch.zip
     - run: |
         export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch

.github/workflows/test_python.yml

@@ -16,7 +16,7 @@ jobs:
       fail-fast: false
       matrix:
         group: [1, 2, 3, 4, 5, 6]
-        python: ["3.8", "3.12"]
+        python: ["3.9", "3.12"]
 
     steps:
     - uses: actions/checkout@v4
@@ -25,19 +25,23 @@ jobs:
         python-version: ${{ matrix.python }}
     - run: python -m pip install -U uv
     - run: |
-        source/install/uv_with_retry.sh pip install --system mpich
-        source/install/uv_with_retry.sh pip install --system "torch==2.3.0+cpu.cxx11.abi" -i https://download.pytorch.org/whl/
+        source/install/uv_with_retry.sh pip install --system openmpi tensorflow-cpu
+        source/install/uv_with_retry.sh pip install --system torch -i https://download.pytorch.org/whl/cpu
+        export TENSORFLOW_ROOT=$(python -c 'import tensorflow;print(tensorflow.__path__[0])')
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-        source/install/uv_with_retry.sh pip install --system --only-binary=horovod -e .[cpu,test] horovod[tensorflow-cpu] mpi4py
+        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py
+        source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation
       env:
         # Please note that uv has some issues with finding
         # existing TensorFlow package. Currently, it uses
         # TensorFlow in the build dependency, but if it
         # changes, setting `TENSORFLOW_ROOT`.
-        TENSORFLOW_VERSION: ${{ matrix.python == '3.8' && '2.13.1' || '2.16.1' }}
         DP_ENABLE_PYTORCH: 1
         DP_BUILD_TESTING: 1
-        UV_EXTRA_INDEX_URL: "https://pypi.anaconda.org/njzjz/simple https://pypi.anaconda.org/mpi4py/simple"
+        UV_EXTRA_INDEX_URL: "https://pypi.anaconda.org/mpi4py/simple"
+        HOROVOD_WITH_TENSORFLOW: 1
+        HOROVOD_WITHOUT_PYTORCH: 1
+        HOROVOD_WITH_MPI: 1
     - run: dp --version
     - name: Get durations from cache
       uses: actions/cache@v4
@@ -53,6 +57,12 @@ jobs:
     - run: pytest --cov=deepmd source/tests --durations=0 --splits 6 --group ${{ matrix.group }} --store-durations --durations-path=.test_durations --splitting-algorithm least_duration
       env:
         NUM_WORKERS: 0
+    - name: Test TF2 eager mode
+      run: pytest --cov=deepmd source/tests/consistent/io/test_io.py --durations=0
+      env:
+        NUM_WORKERS: 0
+        DP_TEST_TF2_ONLY: 1
+      if: matrix.group == 1
     - run: mv .test_durations .test_durations_${{ matrix.group }}
     - name: Upload partial durations
       uses: actions/upload-artifact@v4
@@ -69,7 +79,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: ["3.8", "3.12"]
+        python: ["3.9", "3.12"]
     needs: testpython
     steps:
     - name: Get durations from cache

.gitignore

@@ -45,3 +45,8 @@ build_c_tests
 build_c/
 libdeepmd_c/
 .uv/
+libtorch/
+uv.lock
+buildcxx/
+node_modules/
+*.bib.original

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
         exclude: "^.+\\.pbtxt$"
@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.6.5
+    rev: v0.7.2
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -38,6 +38,14 @@ repos:
       - id: ruff-format
         exclude: ^source/3rdparty
         types_or: [python, pyi, jupyter]
+  - repo: https://github.com/pycqa/flake8
+    # flake8 cannot autofix
+    rev: "7.1.1"
+    hooks:
+      - id: flake8
+        additional_dependencies:
+          - torchfix==0.6.0
+          - flake8-pyproject==1.2.3
   # numpydoc
   - repo: https://github.com/Carreau/velin
     rev: 0.0.12
@@ -47,15 +55,15 @@ repos:
         exclude: ^source/3rdparty
   # Python inside docs
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.18.0
+    rev: 1.19.1
     hooks:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.8
+    rev: v19.1.3
     hooks:
       - id: clang-format
-        exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc
+        exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$)
   # markdown, yaml, CSS, javascript
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v4.0.0-alpha.8
@@ -66,7 +74,7 @@ repos:
         exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
   # Shell
   - repo: https://github.com/scop/pre-commit-shfmt
-    rev: v3.9.0-1
+    rev: v3.10.0-1
     hooks:
       - id: shfmt
   # CMake
@@ -146,7 +154,7 @@ repos:
         exclude: .pre-commit-config.yaml|source/lmp
   # customized pylint rules
   - repo: https://github.com/pylint-dev/pylint/
-    rev: v3.2.7
+    rev: v3.3.1
     hooks:
       - id: pylint
         entry: env PYTHONPATH=source/checker pylint

README.md

@@ -19,7 +19,7 @@ For more information, check the [documentation](https://deepmd.readthedocs.io/).
 
 ### Highlighted features
 
-- **interfaced with multiple backends**, including TensorFlow and PyTorch, the most popular deep learning frameworks, making the training process highly automatic and efficient.
+- **interfaced with multiple backends**, including TensorFlow, PyTorch, and JAX, the most popular deep learning frameworks, making the training process highly automatic and efficient.
 - **interfaced with high-performance classical MD and quantum (path-integral) MD packages**, including LAMMPS, i-PI, AMBER, CP2K, GROMACS, OpenMM, and ABUCUS.
 - **implements the Deep Potential series models**, which have been successfully applied to finite and extended systems, including organic molecules, metals, semiconductors, insulators, etc.
 - **implements MPI and GPU supports**, making it highly efficient for high-performance parallel and distributed computing.
@@ -72,7 +72,7 @@ See [our latest paper](https://doi.org/10.1063/5.0155600) for details of all fea
 
 #### v3
 
-- Multiple backends supported. Add a PyTorch backend.
+- Multiple backends supported. Add PyTorch and JAX backends.
 - The DPA-2 model.
 
 ## Install and use DeePMD-kit

backend/dp_backend.py

@@ -1,10 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """A PEP-517 backend to find TensorFlow."""
 
-from typing import (
-    List,
-)
-
 from scikit_build_core import build as _orig
 
 from .find_pytorch import (
@@ -26,7 +22,7 @@
 ]
 
 
-def __dir__() -> List[str]:
+def __dir__() -> list[str]:
     return __all__
 
 
@@ -42,7 +38,7 @@ def __dir__() -> List[str]:
 
 def get_requires_for_build_wheel(
     config_settings: dict,
-) -> List[str]:
+) -> list[str]:
     return (
         _orig.get_requires_for_build_wheel(config_settings)
         + find_tensorflow()[1]
@@ -52,7 +48,7 @@ def get_requires_for_build_wheel(
 
 def get_requires_for_build_editable(
     config_settings: dict,
-) -> List[str]:
+) -> list[str]:
     return (
         _orig.get_requires_for_build_editable(config_settings)
         + find_tensorflow()[1]