Merge branch 'devel' into lmp_units

.github/workflows/build_cc.yml

@@ -20,20 +20,26 @@ jobs:
         - variant: clang
           dp_variant: clang
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
       with:
         submodules: true
     - uses: actions/setup-python@v4
       with:
         python-version: '3.11'
+        cache: 'pip'
+    - uses: lukka/get-cmake@latest
     - run: python -m pip install tensorflow
-    - run: sudo apt-get update && sudo apt-get install -y nvidia-cuda-toolkit
+    - run: |
+         wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
+         && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
+         && sudo apt-get update \
+         && sudo apt-get -y install cuda-cudart-dev-11-8 cuda-nvcc-11-8
       if: matrix.variant == 'cuda'
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
          && sudo apt-get update \
-         && sudo apt-get -y install cuda-12-0
+         && sudo apt-get -y install cuda-cudart-dev-12-0 cuda-nvcc-12-0
       if: matrix.variant == 'cuda120'
       env:
         DEBIAN_FRONTEND: noninteractive
@@ -44,19 +50,19 @@ jobs:
          && sudo apt-get update \
          && sudo apt-get install -y rocm-dev hipcub-dev
       if: matrix.variant == 'rocm'
-    - run: sudo apt-get update && sudo apt-get install -y clang
-      if: matrix.variant == 'clang'
     - run: source/install/build_cc.sh
       env:
         DP_VARIANT: ${{ matrix.dp_variant }}
         DOWNLOAD_TENSORFLOW: "FALSE"
+        CMAKE_GENERATOR: Ninja
       if: matrix.variant != 'clang'
     - run: source/install/build_cc.sh
       env:
         DP_VARIANT: cpu
         DOWNLOAD_TENSORFLOW: "FALSE"
         CC: clang
         CXX: clang++
+        CMAKE_GENERATOR: Ninja
       if: matrix.variant == 'clang'
     - name: Test files exist
       run: |

.github/workflows/build_wheel.yml

@@ -38,7 +38,7 @@ jobs:
             platform_id: manylinux_aarch64
             dp_variant: cpu
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           submodules: true
           # https://github.com/pypa/setuptools_scm/issues/480
@@ -60,13 +60,14 @@ jobs:
     name: Build source distribution
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           submodules: true
       - uses: actions/setup-python@v4
         name: Install Python
         with:
           python-version: '3.11'
+          cache: 'pip'
       - run: python -m pip install build
       - name: Build sdist
         run: python -m build --sdist
@@ -94,7 +95,7 @@ jobs:
     needs: [build_wheels]
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - uses: actions/download-artifact@v3
         with:
           name: artifact
@@ -113,7 +114,7 @@ jobs:
           images: ghcr.io/deepmodeling/deepmd-kit
 
       - name: Build and push Docker image
-        uses: docker/build-push-action@2eb1c1961a95fc15694676618e422e8ba1d63825
+        uses: docker/build-push-action@0a97817b6ade9f46837855d676c4cca3a2471fc9
         with:
           context: source/install/docker
           push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' }}

.github/workflows/package_c.yml

@@ -9,7 +9,7 @@ jobs:
     name: Build C library
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Package C library
         run: ./source/install/docker_package_c.sh
       # for download and debug
@@ -30,7 +30,7 @@ jobs:
     needs: [build_c]
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Download artifact
         uses: actions/download-artifact@v3
         with:

.github/workflows/test_cc.yml

@@ -7,21 +7,24 @@ jobs:
     name: Test C++
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - uses: actions/setup-python@v4
       with:
         python-version: '3.11'
+        cache: 'pip'
     - name: Setup MPI
       uses: mpi4py/setup-mpi@v1
       with:
         mpi: mpich
+    - uses: lukka/get-cmake@latest
     - run: python -m pip install tensorflow
     - run: source/install/test_cc_local.sh
       env:
         OMP_NUM_THREADS: 1
         TF_INTRA_OP_PARALLELISM_THREADS: 1
         TF_INTER_OP_PARALLELISM_THREADS: 1
         LMP_CXX11_ABI_0: 1
+        CMAKE_GENERATOR: Ninja
     # test lammps
     # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843
     # TODO: remove ase version when ase has new release

.github/workflows/test_python.yml

@@ -5,48 +5,43 @@ name: Test Python
 jobs:
   testpython:
     name: Test Python
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         include:
           - python: 3.7
-            gcc: 5
-            tf: 1.14
-          - python: 3.7
-            gcc: 6
-            tf: 1.14
-          - python: 3.7
-            gcc: 7
-            tf: 1.14
-          - python: 3.7
-            gcc: 8
             tf: 1.14
           - python: 3.8
-            gcc: 8
             tf:
           - python: "3.11"
-            gcc: 8
             tf:
 
-    container: ghcr.io/deepmodeling/deepmd-kit-test-environment:py${{ matrix.python }}-gcc${{ matrix.gcc }}-tf${{ matrix.tf }}
     steps:
-    - name: work around permission issue
-      run: git config --global --add safe.directory /__w/deepmd-kit/deepmd-kit
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python }}
+        cache: 'pip'
+    - uses: mpi4py/setup-mpi@v1
+      if: ${{ matrix.tf == '' }}
+      with:
+        mpi: openmpi
     # https://github.com/pypa/pip/issues/11770
     - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
     - run: pip install -e .[cpu,test]
       env:
-        CC: gcc-${{ matrix.gcc }}
-        CXX: g++-${{ matrix.gcc }}
         TENSORFLOW_VERSION: ${{ matrix.tf }}
         DP_BUILD_TESTING: 1
+    - run: pip install horovod mpi4py
+      if: ${{ matrix.tf == '' }}
+      env:
+        HOROVOD_WITH_TENSORFLOW: 1
+        HOROVOD_WITHOUT_GLOO: 1
     - run: dp --version
     - run: pytest --cov=deepmd --cov=deepmd_cli source/tests --durations=0
     - uses: codecov/codecov-action@v3
       with:
         gcov: true
-        gcov_executable: gcov-${{ matrix.gcc }}
   pass:
     name: Pass testing Python
     needs: [testpython]

.pre-commit-config.yaml

@@ -22,21 +22,21 @@ repos:
     -   id: check-symlinks
     -   id: check-toml
 # Python
--   repo: https://github.com/psf/black
-    rev: 23.7.0
-    hooks:
-    -   id: black-jupyter
 -   repo: https://github.com/PyCQA/isort
     rev: 5.12.0
     hooks:
     - id: isort
       files: \.py$
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.287
+    rev: v0.0.288
     hooks:
     - id: ruff
       args: ["--fix"]
+-   repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 23.9.1
+    hooks:
+    -   id: black-jupyter
 # numpydoc
 -   repo: https://github.com/Carreau/velin
     rev: 0.0.12

CITATIONS.bib

@@ -105,6 +105,19 @@ @misc{Zhang_2022_DPA1
     doi = {10.48550/arXiv.2208.08236},
 }
 
+@article{Zhang_PhysPlasmas_2020_v27_p122704,
+    annote = {frame-specific parameters (e.g. electronic temperature)},
+    author = {Zhang, Yuzhi and Gao, Chang and Liu, Qianrui and Zhang, Linfeng and Wang, Han and Chen, Mohan},
+    title = {{Warm dense matter simulation via electron temperature dependent deep potential molecular dynamics}},
+    journal = {Phys. Plasmas},
+    volume = {27},
+    number = {12},
+    pages = {122704},
+    year = {2020},
+    month = {12},
+    doi = {10.1063/5.0023265},
+}
+
 @article{Zhang_PhysRevB_2020_v102_p41121,
     annote = {fit dipole},
     title={{Deep neural network for the dielectric response of insulators}},

backend/find_tensorflow.py

@@ -112,39 +112,60 @@ def get_tf_requirement(tf_version: str = "") -> dict:
     if tf_version == "":
         tf_version = os.environ.get("TENSORFLOW_VERSION", "")
 
+    extra_requires = []
+    extra_select = {}
+    if not (tf_version == "" or tf_version in SpecifierSet(">=2.12")):
+        extra_requires.append("protobuf<3.20")
+    if tf_version == "" or tf_version in SpecifierSet(">=1.15"):
+        extra_select["mpi"] = [
+            "horovod",
+            "mpi4py",
+        ]
+    else:
+        extra_select["mpi"] = []
+
     if tf_version == "":
         return {
             "cpu": [
                 "tensorflow-cpu; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
                 "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')",
+                *extra_requires,
             ],
             "gpu": [
                 "tensorflow",
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
+                *extra_requires,
             ],
+            **extra_select,
         }
     elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(
         ">=2.0,<2.1"
     ):
         return {
             "cpu": [
                 f"tensorflow=={tf_version}",
+                *extra_requires,
             ],
             "gpu": [
                 f"tensorflow-gpu=={tf_version}; platform_machine!='aarch64'",
                 f"tensorflow=={tf_version}; platform_machine=='aarch64'",
+                *extra_requires,
             ],
+            **extra_select,
         }
     else:
         return {
             "cpu": [
                 f"tensorflow-cpu=={tf_version}; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
                 f"tensorflow=={tf_version}; platform_machine=='aarch64'  or (platform_machine=='arm64' and platform_system == 'Darwin')",
+                *extra_requires,
             ],
             "gpu": [
                 f"tensorflow=={tf_version}",
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
+                *extra_requires,
             ],
+            **extra_select,
         }
 
 

data/raw/shuffle_raw.py

@@ -69,7 +69,8 @@ def _main():
     tmp = np.reshape(tmp, [nframe, -1])
     nframe = tmp.shape[0]
     idx = np.arange(nframe)
-    np.random.shuffle(idx)
+    rng = np.random.default_rng()
+    rng.shuffle(idx)
 
     for ii in raws:
         data = np.loadtxt(inpath + "/" + ii)

deepmd/calculator.py

@@ -6,6 +6,7 @@
 )
 from typing import (
     TYPE_CHECKING,
+    ClassVar,
     Dict,
     List,
     Optional,
@@ -69,13 +70,19 @@ class DP(Calculator):
     """
 
     name = "DP"
-    implemented_properties = ["energy", "free_energy", "forces", "virial", "stress"]
+    implemented_properties: ClassVar[List[str]] = [
+        "energy",
+        "free_energy",
+        "forces",
+        "virial",
+        "stress",
+    ]
 
     def __init__(
         self,
         model: Union[str, "Path"],
         label: str = "DP",
-        type_dict: Dict[str, int] = None,
+        type_dict: Optional[Dict[str, int]] = None,
         **kwargs,
     ) -> None:
         Calculator.__init__(self, label=label, **kwargs)

deepmd/descriptor/se_a.py

@@ -144,7 +144,7 @@ class DescrptSeA(DescrptSe):
     .. [1] Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018.
        End-to-end symmetry preserving inter-atomic potential energy model for finite and extended
        systems. In Proceedings of the 32nd International Conference on Neural Information Processing
-       Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441–4451.
+       Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
     """
 
     def __init__(
@@ -890,7 +890,7 @@ def _filter_lower(
         suffix="",
     ):
         """Input env matrix, returns R.G."""
-        outputs_size = [1] + self.filter_neuron
+        outputs_size = [1, *self.filter_neuron]
         # cut-out inputs
         # with natom x (nei_type_i x 4)
         inputs_i = tf.slice(inputs, [0, start_index * 4], [-1, incrs_index * 4])
@@ -1006,7 +1006,7 @@ def _filter(
         nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0]
         # natom x (nei x 4)
         shape = inputs.get_shape().as_list()
-        outputs_size = [1] + self.filter_neuron
+        outputs_size = [1, *self.filter_neuron]
         outputs_size_2 = self.n_axis_neuron
         all_excluded = all(
             (type_input, type_i) in self.exclude_types for type_i in range(self.ntypes)
@@ -1054,8 +1054,8 @@ def _filter(
                         # add zero is meaningless; skip
                         rets.append(ret)
                     start_index += self.sel_a[type_i]
-                # faster to use accumulate_n than multiple add
-                xyz_scatter_1 = tf.accumulate_n(rets)
+                # faster to use add_n than multiple add
+                xyz_scatter_1 = tf.add_n(rets)
             else:
                 xyz_scatter_1 = self._filter_lower(
                     type_i,

deepmd/descriptor/se_a_ebd.py

@@ -230,7 +230,7 @@ def _embedding_net(
         # natom x (nei x 4)
         inputs = tf.reshape(inputs, [-1, self.ndescrpt])
         shape = inputs.get_shape().as_list()
-        outputs_size = [1] + filter_neuron
+        outputs_size = [1, *filter_neuron]
         with tf.variable_scope(name, reuse=reuse):
             xyz_scatter_total = []
             # with natom x (nei x 4)