Merge branch 'devel' into hlo

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

.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 "$@"

.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

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.7.0
+    rev: v0.7.1
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -55,7 +55,7 @@ repos:
         exclude: ^source/3rdparty
   # Python inside docs
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.19.0
+    rev: 1.19.1
     hooks:
       - id: blacken-docs
   # C++

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/read_env.py

@@ -43,7 +43,7 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]:
     """
     cmake_args = []
     extra_scripts = {}
-    # get variant option from the environment varibles, available: cpu, cuda, rocm
+    # get variant option from the environment variables, available: cpu, cuda, rocm
     dp_variant = os.environ.get("DP_VARIANT", "cpu").lower()
     if dp_variant == "cpu" or dp_variant == "":
         cmake_minimum_required_version = "3.16"

deepmd/__init__.py

@@ -17,7 +17,7 @@
 
 
 def DeepPotential(*args, **kwargs):
-    """Factory function that forwards to DeepEval (for compatbility
+    """Factory function that forwards to DeepEval (for compatibility
     and performance).
 
     Parameters

deepmd/backend/jax.py

@@ -33,9 +33,9 @@ class JAXBackend(Backend):
     """The formal name of the backend."""
     features: ClassVar[Backend.Feature] = (
         Backend.Feature.IO
-        # Backend.Feature.ENTRY_POINT
+        | Backend.Feature.ENTRY_POINT
         | Backend.Feature.DEEP_EVAL
-        # | Backend.Feature.NEIGHBOR_STAT
+        | Backend.Feature.NEIGHBOR_STAT
     )
     """The features of the backend."""
     suffixes: ClassVar[list[str]] = [".hlo", ".jax"]
@@ -86,7 +86,11 @@ def neighbor_stat(self) -> type["NeighborStat"]:
         type[NeighborStat]
             The neighbor statistics of the backend.
         """
-        raise NotImplementedError
+        from deepmd.jax.utils.neighbor_stat import (
+            NeighborStat,
+        )
+
+        return NeighborStat
 
     @property
     def serialize_hook(self) -> Callable[[str], dict]:

deepmd/backend/suffix.py

@@ -23,7 +23,7 @@ def format_model_suffix(
     """Check and format the suffixes of a filename.
 
     When preferred_backend is not given, this method checks the suffix of the filename
-    is within the suffixes of the any backends (with the given feature) and doesn't do formating.
+    is within the suffixes of the any backends (with the given feature) and doesn't do formatting.
     When preferred_backend is given, strict_prefer must be given.
     If strict_prefer is True and the suffix is not within the suffixes of the preferred backend,
     or strict_prefer is False and the suffix is not within the suffixes of the any backend with the given feature,

deepmd/calculator.py

@@ -32,7 +32,7 @@
 class DP(Calculator):
     """Implementation of ASE deepmd calculator.
 
-    Implemented propertie are `energy`, `forces` and `stress`
+    Implemented properties are `energy`, `forces` and `stress`
 
     Parameters
     ----------

deepmd/common.py

@@ -77,7 +77,7 @@ def select_idx_map(atom_types: np.ndarray, select_types: np.ndarray) -> np.ndarr
     Parameters
     ----------
     atom_types : np.ndarray
-        array specifing type for each atoms as integer
+        array specifying type for each atoms as integer
     select_types : np.ndarray
         types of atoms you want to find indices for
 
@@ -126,7 +126,7 @@ def make_default_mesh(pbc: bool, mixed_type: bool) -> np.ndarray:
 def j_deprecated(
     jdata: dict[str, "_DICT_VAL"], key: str, deprecated_key: list[str] = []
 ) -> "_DICT_VAL":
-    """Assert that supplied dictionary conaines specified key.
+    """Assert that supplied dictionary contains specified key.
 
     Parameters
     ----------
@@ -218,7 +218,7 @@ def get_np_precision(precision: "_PRECISION") -> np.dtype:
     Returns
     -------
     np.dtype
-        numpy presicion constant
+        numpy precision constant
 
     Raises
     ------

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -158,7 +158,7 @@ def forward_common_atomic(
         Parameters
         ----------
         extended_coord
-            extended coodinates, shape: nf x (nall x 3)
+            extended coordinates, shape: nf x (nall x 3)
         extended_atype
             extended atom typs, shape: nf x nall
             for a type < 0 indicating the atomic is virtual.

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -100,7 +100,7 @@ def forward_atomic(
         Parameters
         ----------
         extended_coord
-            coodinates in extended region
+            coordinates in extended region
         extended_atype
             atomic type in extended region
         nlist
@@ -169,7 +169,7 @@ def serialize(self) -> dict:
         )
         return dd
 
-    # for subclass overriden
+    # for subclass overridden
     base_descriptor_cls = BaseDescriptor
     """The base descriptor class."""
     base_fitting_cls = BaseFitting

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -162,7 +162,7 @@ def forward_atomic(
         Parameters
         ----------
         extended_coord
-            coodinates in extended region, (nframes, nall * 3)
+            coordinates in extended region, (nframes, nall * 3)
         extended_atype
             atomic type in extended region, (nframes, nall)
         nlist
@@ -341,7 +341,7 @@ class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel):
         Mapping atom type to the name (str) of the type.
         For example `type_map[1]` gives the name of the type 1.
     smin_alpha
-        The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor.
+        The short-range tabulated interaction will be switched according to the distance of the nearest neighbor.
         This distance is calculated by softmin.
     """
 

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -152,7 +152,7 @@ def make_atom_mask(
             self,
             atype: t_tensor,
         ) -> t_tensor:
-            """The atoms with type < 0 are treated as virutal atoms,
+            """The atoms with type < 0 are treated as virtual atoms,
             which serves as place-holders for multi-frame calculations
             with different number of atoms in different frames.
 
@@ -164,7 +164,7 @@ def make_atom_mask(
             Returns
             -------
             mask
-                True for real atoms and False for virutal atoms.
+                True for real atoms and False for virtual atoms.
 
             """
             # supposed to be supported by all backends

deepmd/dpmodel/common.py

@@ -30,7 +30,7 @@
     "int64": np.int64,
     "bool": bool,
     "default": GLOBAL_NP_FLOAT_PRECISION,
-    # NumPy doesn't have bfloat16 (and does't plan to add)
+    # NumPy doesn't have bfloat16 (and doesn't plan to add)
     # ml_dtypes is a solution, but it seems not supporting np.save/np.load
     # hdf5 hasn't supported bfloat16 as well (see https://forum.hdfgroup.org/t/11975)
     "bfloat16": ml_dtypes.bfloat16,

deepmd/dpmodel/descriptor/descriptor.py

@@ -110,7 +110,7 @@ def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
         If not start from checkpoint (resume is False),
-        some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
+        some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.
         """
         raise NotImplementedError
 

deepmd/dpmodel/descriptor/dpa1.py

@@ -358,11 +358,11 @@ def get_dim_emb(self) -> int:
         return self.se_atten.dim_emb
 
     def mixed_types(self) -> bool:
-        """If true, the discriptor
+        """If true, the descriptor
         1. assumes total number of atoms aligned across frames;
         2. requires a neighbor list that does not distinguish different atomic types.
 
-        If false, the discriptor
+        If false, the descriptor
         1. assumes total number of atoms of each atom type aligned across frames;
         2. requires a neighbor list that distinguishes different atomic types.
 
@@ -385,7 +385,7 @@ def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
         If not start from checkpoint (resume is False),
-        some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
+        some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.
         """
         raise NotImplementedError
 
@@ -459,7 +459,7 @@ def call(
         nlist
             The neighbor list. shape: nf x nloc x nnei
         mapping
-            The index mapping from extended to lcoal region. not used by this descriptor.
+            The index mapping from extended to local region. not used by this descriptor.
 
         Returns
         -------
@@ -602,7 +602,7 @@ def update_sel(
         Parameters
         ----------
         train_data : DeepmdDataSystem
-            data used to do neighbor statictics
+            data used to do neighbor statistics
         type_map : list[str], optional
             The name of each type of atoms
         local_jdata : dict
@@ -793,11 +793,11 @@ def __getitem__(self, key):
             raise KeyError(key)
 
     def mixed_types(self) -> bool:
-        """If true, the discriptor
+        """If true, the descriptor
         1. assumes total number of atoms aligned across frames;
         2. requires a neighbor list that does not distinguish different atomic types.
 
-        If false, the discriptor
+        If false, the descriptor
         1. assumes total number of atoms of each atom type aligned across frames;
         2. requires a neighbor list that distinguishes different atomic types.