Merge branch 'devel' into dp-model-format

.github/labeler.yml

@@ -1,15 +1,39 @@
 Python:
-- deepmd/**/*
-- deepmd_utils/**/*
-- source/tests/**/*
-Docs: doc/**/*
-Examples: examples/**/*
-Core: source/lib/**/*
-CUDA: source/lib/src/gpu/**/*
-ROCM: source/lib/src/gpu/**/*
-OP: source/op/**/*
-C++: source/api_cc/**/*
-C: source/api_c/**/*
-LAMMPS: source/lmp/**/*
-Gromacs: source/gmx/**/*
-i-Pi: source/ipi/**/*
+- changed-files:
+  - any-glob-to-any-file:
+    - deepmd/**/*
+    - deepmd_utils/**/*
+    - source/tests/**/*
+Docs:
+- changed-files:
+  - any-glob-to-any-file: doc/**/*
+Examples:
+- changed-files:
+  - any-glob-to-any-file: examples/**/*
+Core:
+- changed-files:
+  - any-glob-to-any-file: source/lib/**/*
+CUDA:
+- changed-files:
+  - any-glob-to-any-file: source/lib/src/gpu/**/*
+ROCM:
+- changed-files:
+  - any-glob-to-any-file: source/lib/src/gpu/**/*
+OP:
+- changed-files:
+  - any-glob-to-any-file: source/op/**/*
+C++:
+- changed-files:
+  - any-glob-to-any-file: source/api_cc/**/*
+C:
+- changed-files:
+  - any-glob-to-any-file: source/api_c/**/*
+LAMMPS:
+- changed-files:
+  - any-glob-to-any-file: source/lmp/**/*
+Gromacs:
+- changed-files:
+  - any-glob-to-any-file: source/gmx/**/*
+i-Pi:
+- changed-files:
+  - any-glob-to-any-file: source/ipi/**/*

.github/workflows/build_cc.yml

@@ -21,7 +21,7 @@ jobs:
           dp_variant: clang
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.11'
         cache: 'pip'

.github/workflows/build_wheel.yml

@@ -87,7 +87,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         name: Install Python
         with:
           python-version: '3.11'
@@ -141,15 +141,15 @@ jobs:
 
       - name: Extract metadata (tags, labels) for Docker
         id: meta
-        uses: docker/metadata-action@96383f45573cb7f253c731d3b3ab81c87ef81934
+        uses: docker/metadata-action@31cebacef4805868f9ce9a0cb03ee36c32df2ac4
         with:
           images: ghcr.io/deepmodeling/deepmd-kit
 
       - name: Build and push Docker image
         uses: docker/build-push-action@4a13e500e55cf31b7a5d59a38ab2040ab0f42f56
         with:
           context: source/install/docker
-          push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' }}
+          push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' && github.actor != 'dependabot[bot]' }}
           tags: ${{ steps.meta.outputs.tags }}${{ matrix.variant }}
           labels: ${{ steps.meta.outputs.labels }}
           build-args: |
@@ -164,7 +164,7 @@ jobs:
         with:
           name: artifact
           path: dist/packages
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         name: Install Python
         with:
           python-version: '3.11'
@@ -189,7 +189,7 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v2
+        uses: actions/deploy-pages@v3
 
   pass:
     name: Pass testing build wheels

.github/workflows/labeler.yml

@@ -9,6 +9,6 @@ jobs:
       pull-requests: write
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/labeler@v4
+    - uses: actions/labeler@v5
       with:
         repo-token: "${{ secrets.GITHUB_TOKEN }}"

.github/workflows/test_cc.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.11'
         cache: 'pip'

.github/workflows/test_cuda.yml

@@ -18,7 +18,7 @@ jobs:
     - name: Make sudo and git work
       run: apt-get update && apt-get install -y sudo git
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.11'
         # cache: 'pip'

.github/workflows/test_python.yml

@@ -18,7 +18,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python }}
         cache: 'pip'

.pre-commit-config.yaml

@@ -23,14 +23,14 @@ repos:
     -   id: check-toml
 # Python
 -   repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
+    rev: 5.13.0
     hooks:
     - id: isort
       files: \.py$
       exclude: ^source/3rdparty
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.1.6
+    rev: v0.1.7
     hooks:
     - id: ruff
       args: ["--fix"]

deepmd/calculator.py

@@ -45,6 +45,8 @@ class DP(Calculator):
     type_dict : Dict[str, int], optional
         mapping of element types and their numbers, best left None and the calculator
         will infer this information from model, by default None
+    neighbor_list : ase.neighborlist.NeighborList, optional
+        The neighbor list object. If None, then build the native neighbor list.
 
     Examples
     --------
@@ -83,10 +85,11 @@ def __init__(
         model: Union[str, "Path"],
         label: str = "DP",
         type_dict: Optional[Dict[str, int]] = None,
+        neighbor_list=None,
         **kwargs,
     ) -> None:
         Calculator.__init__(self, label=label, **kwargs)
-        self.dp = DeepPotential(str(Path(model).resolve()))
+        self.dp = DeepPotential(str(Path(model).resolve()), neighbor_list=neighbor_list)
         if type_dict:
             self.type_dict = type_dict
         else:

deepmd/descriptor/se_a.py

@@ -783,16 +783,16 @@ def _pass_filter(
             type_i = -1
             if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor:
                 inputs_i = descrpt2r4(inputs_i, natoms)
+            self.atype_nloc = tf.reshape(
+                tf.slice(atype, [0, 0], [-1, natoms[0]]), [-1]
+            )  # when nloc != nall, pass nloc to mask
             if len(self.exclude_types):
-                atype_nloc = tf.reshape(
-                    tf.slice(atype, [0, 0], [-1, natoms[0]]), [-1]
-                )  # when nloc != nall, pass nloc to mask
                 mask = self.build_type_exclude_mask(
                     self.exclude_types,
                     self.ntypes,
                     self.sel_a,
                     self.ndescrpt,
-                    atype_nloc,
+                    self.atype_nloc,
                     tf.shape(inputs_i)[0],
                 )
                 inputs_i *= mask
@@ -957,7 +957,7 @@ def _filter_lower(
                     extra_embedding_index = self.nei_type_vec
                 else:
                     padding_ntypes = type_embedding.shape[0]
-                    atype_expand = tf.reshape(self.atype, [-1, 1])
+                    atype_expand = tf.reshape(self.atype_nloc, [-1, 1])
                     idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei])
                     idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei])
                     idx = idx_i + idx_j
@@ -1003,13 +1003,6 @@ def _filter_lower(
                             [-1, two_side_type_embedding.shape[-1]],
                         )
 
-                        atype_expand = tf.reshape(self.atype, [-1, 1])
-                        idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei])
-                        idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei])
-                        idx = idx_i + idx_j
-                        index_of_two_side = tf.reshape(idx, [-1])
-                        self.extra_embedding_index = index_of_two_side
-
                         net_output = embedding_net(
                             two_side_type_embedding,
                             self.filter_neuron,

deepmd/fit/ener.py

@@ -521,6 +521,11 @@ def build(
                     self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
             self.bias_atom_e = self.bias_atom_e[:ntypes_atom]
 
+        if nvnmd_cfg.enable:
+            # fix the bug: CNN and QNN have different t_bias_atom_e.
+            if "t_bias_atom_e" in nvnmd_cfg.weight.keys():
+                self.bias_atom_e = nvnmd_cfg.weight["t_bias_atom_e"]
+
         with tf.variable_scope("fitting_attr" + suffix, reuse=reuse):
             t_dfparam = tf.constant(self.numb_fparam, name="dfparam", dtype=tf.int32)
             t_daparam = tf.constant(self.numb_aparam, name="daparam", dtype=tf.int32)

deepmd/infer/__init__.py

@@ -58,6 +58,7 @@ def DeepPotential(
     load_prefix: str = "load",
     default_tf_graph: bool = False,
     input_map: Optional[dict] = None,
+    neighbor_list=None,
 ) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepDOS, DeepWFC]:
     """Factory function that will inialize appropriate potential read from `model_file`.
 
@@ -71,6 +72,8 @@ def DeepPotential(
         If uses the default tf graph, otherwise build a new tf graph for evaluation
     input_map : dict, optional
         The input map for tf.import_graph_def. Only work with default tf graph
+    neighbor_list : ase.neighborlist.NeighborList, optional
+        The neighbor list object. If None, then build the native neighbor list.
 
     Returns
     -------
@@ -97,6 +100,7 @@ def DeepPotential(
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
             input_map=input_map,
+            neighbor_list=neighbor_list,
         )
     elif model_type == "dos":
         dp = DeepDOS(
@@ -111,20 +115,23 @@ def DeepPotential(
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
             input_map=input_map,
+            neighbor_list=neighbor_list,
         )
     elif model_type == "polar":
         dp = DeepPolar(
             mf,
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
             input_map=input_map,
+            neighbor_list=neighbor_list,
         )
     elif model_type == "global_polar":
         dp = DeepGlobalPolar(
             mf,
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
             input_map=input_map,
+            neighbor_list=neighbor_list,
         )
     elif model_type == "wfc":
         dp = DeepWFC(

deepmd/infer/deep_dipole.py

@@ -27,6 +27,8 @@ class DeepDipole(DeepTensor):
         If uses the default tf graph, otherwise build a new tf graph for evaluation
     input_map : dict, optional
         The input map for tf.import_graph_def. Only work with default tf graph
+    neighbor_list : ase.neighborlist.NeighborList, optional
+        The neighbor list object. If None, then build the native neighbor list.
 
     Warnings
     --------
@@ -41,6 +43,7 @@ def __init__(
         load_prefix: str = "load",
         default_tf_graph: bool = False,
         input_map: Optional[dict] = None,
+        neighbor_list=None,
     ) -> None:
         # use this in favor of dict update to move attribute from class to
         # instance namespace
@@ -58,6 +61,7 @@ def __init__(
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
             input_map=input_map,
+            neighbor_list=neighbor_list,
         )
 
     def get_dim_fparam(self) -> int: