Merge branch 'devel' into D0604_fitting_hook

Author: iProzd

.github/workflows/build_wheel.yml

@@ -44,7 +44,7 @@ jobs:
             platform_id: macosx_arm64
             dp_variant: cpu
           # win-64
-          - os: windows-2019
+          - os: windows-2022
             python: 311
             platform_id: win_amd64
             dp_variant: cpu
@@ -70,7 +70,7 @@ jobs:
           rm -rf .git
         if: matrix.dp_pkg_name == 'deepmd-kit-cu11'
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.23
+        uses: pypa/cibuildwheel@v3.0
         env:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_ARCHS: all

.github/workflows/test_cc.yml

@@ -21,17 +21,13 @@ jobs:
       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 uv
     - name: Install Python dependencies
       run: |
         source/install/uv_with_retry.sh pip install --system tensorflow-cpu~=2.18.0 jax==0.5.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 -e .[cpu,test,lmp,jax] mpi4py
+        source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax] mpi4py mpich
         source/install/uv_with_retry.sh pip install --system 'torch==2.7' --index-url https://download.pytorch.org/whl/cpu
     - name: Convert models
       run: source/tests/infer/convert-models.sh
@@ -54,13 +50,13 @@ jobs:
         cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/paddle/lib/*.so ${{ github.workspace }}/dp_test/lib/
         cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/third_party/install/onednn/lib/* ${{ github.workspace }}/dp_test/lib/
         cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/third_party/install/mklml/lib/* ${{ github.workspace }}/dp_test/lib/
+        export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$LD_LIBRARY_PATH
         pytest --cov=deepmd source/lmp/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
-        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib
       if: ${{ !matrix.check_memleak }}
     # test ipi
     - run: |

.github/workflows/test_python.yml

@@ -40,7 +40,6 @@ jobs:
         # changes, setting `TENSORFLOW_ROOT`.
         DP_ENABLE_PYTORCH: 1
         DP_BUILD_TESTING: 1
-        UV_EXTRA_INDEX_URL: "https://pypi.anaconda.org/mpi4py/simple"
         HOROVOD_WITH_TENSORFLOW: 1
         HOROVOD_WITHOUT_PYTORCH: 1
         HOROVOD_WITH_MPI: 1

.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.11.12
+    rev: v0.12.2
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -40,7 +40,7 @@ repos:
         types_or: [python, pyi, jupyter]
   - repo: https://github.com/pycqa/flake8
     # flake8 cannot autofix
-    rev: "7.2.0"
+    rev: "7.3.0"
     hooks:
       - id: flake8
         additional_dependencies:
@@ -60,7 +60,7 @@ repos:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v20.1.5
+    rev: v20.1.7
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$)

README.md

@@ -76,7 +76,7 @@ See [our v2 paper](https://doi.org/10.1063/5.0155600) for details of all feature
 #### v3
 
 - Multiple backends supported. Add PyTorch and JAX backends.
-- The DPA-2 model.
+- The DPA2 and DPA3 models.
 - Plugin mechanisms for external models.
 
 See [our v3 paper](https://doi.org/10.1021/acs.jctc.5c00340) for details of all features until v3.0.

deepmd/calculator.py

@@ -130,7 +130,12 @@ def calculate(
             cell = None
         symbols = self.atoms.get_chemical_symbols()
         atype = [self.type_dict[k] for k in symbols]
-        e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype)[:3]
+
+        fparam = self.atoms.info.get("fparam", None)
+        aparam = self.atoms.info.get("aparam", None)
+        e, f, v = self.dp.eval(
+            coords=coord, cells=cell, atom_types=atype, fparam=fparam, aparam=aparam
+        )[:3]
         self.results["energy"] = e[0][0]
         # see https://gitlab.com/ase/ase/-/merge_requests/2485
         self.results["free_energy"] = e[0][0]

deepmd/dpmodel/array_api.py

@@ -60,7 +60,7 @@ def xp_take_along_axis(arr, indices, axis):
 
     shape = list(arr.shape)
     shape.pop(-1)
-    shape = [*shape, n]
+    shape = (*shape, n)
 
     arr = xp.reshape(arr, (-1,))
     if n != 0:

deepmd/dpmodel/descriptor/dpa1.py

@@ -319,6 +319,7 @@ def __init__(
             trainable_ln=trainable_ln,
             ln_eps=ln_eps,
             seed=child_seed(seed, 0),
+            trainable=trainable,
         )
         self.use_econf_tebd = use_econf_tebd
         self.use_tebd_bias = use_tebd_bias
@@ -333,6 +334,7 @@ def __init__(
             use_tebd_bias=use_tebd_bias,
             type_map=type_map,
             seed=child_seed(seed, 1),
+            trainable=trainable,
         )
         self.tebd_dim = tebd_dim
         self.concat_output_tebd = concat_output_tebd
@@ -520,7 +522,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # nf x nall x tebd_dim
         atype_embd_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nf, nall, self.tebd_dim),
         )
         # nfnl x tebd_dim
@@ -691,6 +693,7 @@ def __init__(
         ln_eps: Optional[float] = 1e-5,
         smooth: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
+        trainable: bool = True,
     ) -> None:
         self.rcut = rcut
         self.rcut_smth = rcut_smth
@@ -741,6 +744,7 @@ def __init__(
             self.resnet_dt,
             self.precision,
             seed=child_seed(seed, 0),
+            trainable=trainable,
         )
         self.embeddings = embeddings
         if self.tebd_input_mode in ["strip"]:
@@ -756,6 +760,7 @@ def __init__(
                 self.resnet_dt,
                 self.precision,
                 seed=child_seed(seed, 1),
+                trainable=trainable,
             )
             self.embeddings_strip = embeddings_strip
         else:
@@ -774,6 +779,7 @@ def __init__(
             smooth=self.smooth,
             precision=self.precision,
             seed=child_seed(seed, 2),
+            trainable=trainable,
         )
 
         wanted_shape = (self.ntypes, self.nnei, 4)
@@ -1027,7 +1033,7 @@ def call(
                     xp.tile(
                         (xp.reshape(atype, (-1, 1)) * ntypes_with_padding), (1, nnei)
                     ),
-                    (-1),
+                    (-1,),
                 )
                 idx_j = xp.reshape(nei_type, (-1,))
                 # (nf x nl x nnei) x ng
@@ -1186,6 +1192,7 @@ def __init__(
         smooth: bool = True,
         precision: str = DEFAULT_PRECISION,
         seed: Optional[Union[int, list[int]]] = None,
+        trainable: bool = True,
     ) -> None:
         """Construct a neighbor-wise attention net."""
         super().__init__()
@@ -1219,6 +1226,7 @@ def __init__(
                 smooth=smooth,
                 precision=precision,
                 seed=child_seed(seed, ii),
+                trainable=trainable,
             )
             for ii in range(layer_num)
         ]
@@ -1314,6 +1322,7 @@ def __init__(
         smooth: bool = True,
         precision: str = DEFAULT_PRECISION,
         seed: Optional[Union[int, list[int]]] = None,
+        trainable: bool = True,
     ) -> None:
         """Construct a neighbor-wise attention layer."""
         super().__init__()
@@ -1340,6 +1349,7 @@ def __init__(
             smooth=smooth,
             precision=precision,
             seed=child_seed(seed, 0),
+            trainable=trainable,
         )
         self.attn_layer_norm = LayerNorm(
             self.embed_dim,
@@ -1420,6 +1430,7 @@ def __init__(
         smooth: bool = True,
         precision: str = DEFAULT_PRECISION,
         seed: Optional[Union[int, list[int]]] = None,
+        trainable: bool = True,
     ) -> None:
         """Construct a multi-head neighbor-wise attention net."""
         super().__init__()
@@ -1449,6 +1460,7 @@ def __init__(
             use_timestep=False,
             precision=precision,
             seed=child_seed(seed, 0),
+            trainable=trainable,
         )
         self.out_proj = NativeLayer(
             hidden_dim,
@@ -1457,6 +1469,7 @@ def __init__(
             use_timestep=False,
             precision=precision,
             seed=child_seed(seed, 1),
+            trainable=trainable,
         )
 
     def call(self, query, nei_mask, input_r=None, sw=None, attnw_shift=20.0):

deepmd/dpmodel/descriptor/dpa2.py

@@ -474,6 +474,7 @@ def init_subclass_params(sub_data, sub_class):
             smooth=smooth,
             type_one_side=self.repinit_args.type_one_side,
             seed=child_seed(seed, 0),
+            trainable=trainable,
         )
         self.use_three_body = self.repinit_args.use_three_body
         if self.use_three_body:
@@ -493,6 +494,7 @@ def init_subclass_params(sub_data, sub_class):
                 resnet_dt=self.repinit_args.resnet_dt,
                 smooth=smooth,
                 seed=child_seed(seed, 5),
+                trainable=trainable,
             )
         else:
             self.repinit_three_body = None
@@ -533,6 +535,7 @@ def init_subclass_params(sub_data, sub_class):
             g1_out_mlp=self.repformer_args.g1_out_mlp,
             ln_eps=self.repformer_args.ln_eps,
             seed=child_seed(seed, 1),
+            trainable=trainable,
         )
         self.rcsl_list = [
             (self.repformers.get_rcut(), self.repformers.get_nsel()),
@@ -562,6 +565,7 @@ def init_subclass_params(sub_data, sub_class):
             use_tebd_bias=use_tebd_bias,
             type_map=type_map,
             seed=child_seed(seed, 2),
+            trainable=trainable,
         )
         self.concat_output_tebd = concat_output_tebd
         self.precision = precision
@@ -585,6 +589,7 @@ def init_subclass_params(sub_data, sub_class):
                 bias=False,
                 precision=precision,
                 seed=child_seed(seed, 3),
+                trainable=trainable,
             )
         self.tebd_transform = None
         if self.add_tebd_to_repinit_out:
@@ -594,6 +599,7 @@ def init_subclass_params(sub_data, sub_class):
                 bias=False,
                 precision=precision,
                 seed=child_seed(seed, 4),
+                trainable=trainable,
             )
         assert self.repinit.rcut > self.repformers.rcut
         assert self.repinit.sel[0] > self.repformers.sel[0]
@@ -841,7 +847,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # repinit
         g1_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nframes, nall, self.tebd_dim),
         )
         g1_inp = g1_ext[:, :nloc, :]

deepmd/dpmodel/descriptor/dpa3.py

@@ -251,7 +251,7 @@ def deserialize(cls, data: dict) -> "RepFlowArgs":
 
 @BaseDescriptor.register("dpa3")
 class DescrptDPA3(NativeOP, BaseDescriptor):
-    r"""The DPA-3 descriptor.
+    r"""The DPA3 descriptor[1]_.
 
     Parameters
     ----------
@@ -282,6 +282,12 @@ class DescrptDPA3(NativeOP, BaseDescriptor):
         When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.
     type_map : list[str], Optional
         A list of strings. Give the name to each type of atoms.
+
+    References
+    ----------
+    .. [1] Zhang, D., Peng, A., Cai, C. et al. Graph neural
+       network model for the era of large atomistic models.
+       arXiv preprint arXiv:2506.01686 (2025).
     """
 
     def __init__(
@@ -351,6 +357,7 @@ def init_subclass_params(sub_data, sub_class):
             env_protection=env_protection,
             precision=precision,
             seed=child_seed(seed, 1),
+            trainable=trainable,
         )
 
         self.use_econf_tebd = use_econf_tebd
@@ -368,6 +375,7 @@ def init_subclass_params(sub_data, sub_class):
             use_tebd_bias=use_tebd_bias,
             type_map=type_map,
             seed=child_seed(seed, 2),
+            trainable=trainable,
         )
         self.concat_output_tebd = concat_output_tebd
         self.precision = precision
@@ -556,12 +564,12 @@ def call(
         type_embedding = self.type_embedding.call()
         if self.use_loc_mapping:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], [-1]), axis=0),
+                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], (-1,)), axis=0),
                 (nframes, nloc, self.tebd_dim),
             )
         else:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+                xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
                 (nframes, nall, self.tebd_dim),
             )
         node_ebd_inp = node_ebd_ext[:, :nloc, :]