Merge branch 'devel' into dataset

.github/workflows/build_cc.yml

@@ -1,6 +1,12 @@
 on:
   push:
+    branches-ignore:
+      - "gh-readonly-queue/*"
   pull_request:
+  merge_group:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 name: Build C++
 jobs:
   buildcc:

.github/workflows/build_wheel.yml

@@ -2,7 +2,14 @@ name: Build and upload to PyPI
 
 on:
   push:
+    branches-ignore:
+      - "gh-readonly-queue/*"
   pull_request:
+  merge_group:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 
 jobs:
   determine-arm64-runner:

.github/workflows/codeql.yml

@@ -2,10 +2,14 @@ name: "CodeQL"
 
 on:
   push:
+    branches-ignore:
+      - "gh-readonly-queue/*"
   pull_request:
   schedule:
     - cron: '45 2 * * 2'
-
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 jobs:
   analyze:
     name: Analyze

.github/workflows/labeler.yml

@@ -11,4 +11,4 @@ jobs:
     steps:
     - uses: actions/labeler@v5
       with:
-        repo-token: "${{ secrets.GITHUB_TOKEN }}"
+        repo-token: "${{ secrets.GITHUB_TOKEN }}"

.github/workflows/package_c.yml

@@ -2,8 +2,13 @@ name: Build C library
 
 on:
   push:
+    branches-ignore:
+      - "gh-readonly-queue/*"
   pull_request:
-
+  merge_group:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 jobs:
   build_c:
     name: Build C library

.github/workflows/test_cc.yml

@@ -1,6 +1,12 @@
 on:
   push:
+    branches-ignore:
+      - "gh-readonly-queue/*"
   pull_request:
+  merge_group:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 name: Test C++
 jobs:
   testcc:

.github/workflows/test_cuda.yml

@@ -4,16 +4,23 @@ on:
   pull_request:
     types:
       - "labeled"
+      # to let the PR pass the test
+      - "created"
+      - "synchronize"
+  merge_group:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 name: Test CUDA
 jobs:
   test_cuda:
     name: Test Python and C++ on CUDA
     runs-on: nvidia
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
     container:
-      image: nvidia/cuda:12.2.0-devel-ubuntu22.04
+      image: nvidia/cuda:12.3.1-devel-ubuntu22.04
       options: --gpus all
-    if: github.repository_owner == 'deepmodeling' && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch'
+    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
     - name: Make sudo and git work
       run: apt-get update && apt-get install -y sudo git
@@ -27,24 +34,24 @@ jobs:
       with:
         mpi: mpich
     - uses: lukka/get-cmake@latest
+      with:
+        useLocalCache: true
+        useCloudCache: false
     - 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-2 libcudnn8=8.9.5.*-1+cuda12.2
+         && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
-    - name: Set PyPI mirror for Aliyun cloud machine
-      run: python -m pip config --user set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple/
     - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
     - run: python -m pip install "tensorflow>=2.15.0rc0" "torch>=2.2.0"
     - run: python -m pip install -v -e .[gpu,test,lmp,cu12,torch] "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz"
       env:
-        DP_BUILD_TESTING: 1
         DP_VARIANT: cuda
-        CUDA_PATH: /usr/local/cuda-12.2
         NUM_WORKERS: 0
+        DP_ENABLE_NATIVE_OPTIMIZATION: 1
     - run: dp --version
-    - run: python -m pytest --cov=deepmd source/tests --durations=0
+    - run: python -m pytest source/tests --durations=0
     - run: source/install/test_cc_local.sh
       env:
         OMP_NUM_THREADS: 1
@@ -54,18 +61,24 @@ jobs:
         CMAKE_GENERATOR: Ninja
         DP_VARIANT: cuda
         DP_USE_MPICH2: 1
-        CUDA_PATH: /usr/local/cuda-12.2
     - run: |
         export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
         export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
-        python -m pytest --cov=deepmd source/lmp/tests
-        python -m pytest --cov=deepmd source/ipi/tests
+        python -m pytest source/lmp/tests
+        python -m pytest source/ipi/tests
       env:
         OMP_NUM_THREADS: 1
         TF_INTRA_OP_PARALLELISM_THREADS: 1
         TF_INTER_OP_PARALLELISM_THREADS: 1
         LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
-        CUDA_PATH: /usr/local/cuda-12.2
-    - uses: codecov/codecov-action@v4
-      env:
-        CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+  pass:
+    name: Pass testing on CUDA
+    needs: [test_cuda]
+    runs-on: ubuntu-latest
+    if: always()
+    steps:
+    - name: Decide whether the needed jobs succeeded or failed
+      uses: re-actors/alls-green@release/v1
+      with:
+        jobs: ${{ toJSON(needs) }}
+        allowed-skips: test_cuda

.github/workflows/test_python.yml

@@ -1,6 +1,12 @@
 on:
   push:
+    branches-ignore:
+      - "gh-readonly-queue/*"
   pull_request:
+  merge_group:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: true
 name: Test Python
 jobs:
   testpython:
@@ -27,8 +33,6 @@ jobs:
         mpi: openmpi
     # https://github.com/pypa/pip/issues/11770
     - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
-    - run: python -m pip install -U "torch==${{ matrix.torch }}" "numpy<1.20"
-      if: matrix.torch != ''
     - run: pip install -e .[cpu,test,torch]
       env:
         TENSORFLOW_VERSION: ${{ matrix.tf }}

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -22,14 +22,12 @@
     EmbeddingNet,
     EnvMat,
     NetworkCollection,
+    PairExcludeMask,
 )
 
 from .base_descriptor import (
     BaseDescriptor,
 )
-from .exclude_mask import (
-    ExcludeMask,
-)
 
 
 class DescrptSeA(NativeOP, BaseDescriptor):
@@ -160,7 +158,7 @@ def __init__(
         self.activation_function = activation_function
         self.precision = precision
         self.spin = spin
-        self.emask = ExcludeMask(self.ntypes, self.exclude_types)
+        self.emask = PairExcludeMask(self.ntypes, self.exclude_types)
 
         in_dim = 1  # not considiering type embedding
         self.embeddings = NetworkCollection(

deepmd/dpmodel/fitting/invar_fitting.py

@@ -19,6 +19,7 @@
     fitting_check_output,
 )
 from deepmd.dpmodel.utils import (
+    AtomExcludeMask,
     FittingNet,
     NetworkCollection,
 )
@@ -126,6 +127,7 @@ def __init__(
         use_aparam_as_mask: bool = False,
         spin: Any = None,
         distinguish_types: bool = False,
+        exclude_types: List[int] = [],
     ):
         # seed, uniform_seed are not included
         if tot_ener_zero:
@@ -159,8 +161,10 @@ def __init__(
         self.use_aparam_as_mask = use_aparam_as_mask
         self.spin = spin
         self.distinguish_types = distinguish_types
+        self.exclude_types = exclude_types
         if self.spin is not None:
             raise NotImplementedError("spin is not supported")
+        self.emask = AtomExcludeMask(self.ntypes, exclude_types=self.exclude_types)
 
         # init constants
         self.bias_atom_e = np.zeros([self.ntypes, self.dim_out])
@@ -260,6 +264,7 @@ def serialize(self) -> dict:
             "precision": self.precision,
             "distinguish_types": self.distinguish_types,
             "nets": self.nets.serialize(),
+            "exclude_types": self.exclude_types,
             "@variables": {
                 "bias_atom_e": self.bias_atom_e,
                 "fparam_avg": self.fparam_avg,
@@ -370,4 +375,10 @@ def call(
                 outs = outs + atom_energy  # Shape is [nframes, natoms[0], 1]
         else:
             outs = self.nets[()](xx) + self.bias_atom_e[atype]
+
+        # nf x nloc
+        exclude_mask = self.emask.build_type_exclude_mask(atype)
+        # nf x nloc x nod
+        outs = outs * exclude_mask[:, :, None]
+
         return {self.var_name: outs}

deepmd/dpmodel/utils/__init__.py

@@ -2,6 +2,10 @@
 from .env_mat import (
     EnvMat,
 )
+from .exclude_mask import (
+    AtomExcludeMask,
+    PairExcludeMask,
+)
 from .network import (
     EmbeddingNet,
     FittingNet,
@@ -53,4 +57,6 @@
     "inter2phys",
     "phys2inter",
     "to_face_distance",
+    "AtomExcludeMask",
+    "PairExcludeMask",
 ]

deepmd/dpmodel/descriptor/exclude_mask.py → deepmd/dpmodel/utils/exclude_mask.py

@@ -7,15 +7,54 @@
 import numpy as np
 
 
-class ExcludeMask:
-    """Computes the atom type exclusion mask."""
+class AtomExcludeMask:
+    """Computes the type exclusion mask for atoms."""
+
+    def __init__(
+        self,
+        ntypes: int,
+        exclude_types: List[int] = [],
+    ):
+        self.ntypes = ntypes
+        self.exclude_types = exclude_types
+        self.type_mask = np.array(
+            [1 if tt_i not in self.exclude_types else 0 for tt_i in range(ntypes)],
+            dtype=np.int32,
+        )
+        # (ntypes)
+        self.type_mask = self.type_mask.reshape([-1])
+
+    def build_type_exclude_mask(
+        self,
+        atype: np.ndarray,
+    ):
+        """Compute type exclusion mask for atoms.
+
+        Parameters
+        ----------
+        atype
+            The extended aotm types. shape: nf x natom
+
+        Returns
+        -------
+        mask
+            The type exclusion mask for atoms. shape: nf x natom
+            Element [ff,ii] being 0 if type(ii) is excluded,
+            otherwise being 1.
+
+        """
+        nf, natom = atype.shape
+        return self.type_mask[atype].reshape(nf, natom)
+
+
+class PairExcludeMask:
+    """Computes the type exclusion mask for atom pairs."""
 
     def __init__(
         self,
         ntypes: int,
         exclude_types: List[Tuple[int, int]] = [],
     ):
-        super().__init__()
         self.ntypes = ntypes
         self.exclude_types = set()
         for tt in exclude_types:
@@ -41,7 +80,7 @@ def build_type_exclude_mask(
         nlist: np.ndarray,
         atype_ext: np.ndarray,
     ):
-        """Compute type exclusion mask.
+        """Compute type exclusion mask for atom pairs.
 
         Parameters
         ----------
@@ -53,7 +92,7 @@ def build_type_exclude_mask(
         Returns
         -------
         mask
-            The type exclusion mask of shape: nf x nloc x nnei.
+            The type exclusion mask for pair atoms of shape: nf x nloc x nnei.
             Element [ff,ii,jj] being 0 if type(ii), type(nlist[ff,ii,jj]) is excluded,
             otherwise being 1.