Merge branch 'master' into 0225-lmdb-dataloader

Author: iProzd

.github/workflows/build_wheel.yml

@@ -138,12 +138,12 @@ jobs:
 
       - name: Extract metadata (tags, labels) for Docker
         id: meta
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@v6
         with:
           images: ghcr.io/deepmodeling/deepmd-kit
 
       - name: Build and push Docker image
-        uses: docker/build-push-action@v6
+        uses: docker/build-push-action@v7
         with:
           context: source/install/docker
           push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' && github.actor != 'dependabot[bot]' }}

.github/workflows/test_cuda.yml

@@ -47,7 +47,7 @@ jobs:
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple "paddlepaddle-gpu==3.3.0.dev20251204"
+          pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310"
           source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
           # See https://github.com/jax-ml/jax/issues/29042
           source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'

.github/workflows/test_python.yml

@@ -32,7 +32,7 @@ jobs:
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py --group pin_jax
-          source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple paddlepaddle==3.3.0.dev20251204
+          source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn paddlepaddle==3.4.0.dev20260310
         env:
           # Please note that uv has some issues with finding
           # existing TensorFlow package. Currently, it uses

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.4
+    rev: v0.15.6
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -39,14 +39,6 @@ repos:
       - id: ruff-format
         exclude: ^source/3rdparty
         types_or: [python, pyi, jupyter]
-  - repo: https://github.com/pycqa/flake8
-    # flake8 cannot autofix
-    rev: "7.3.0"
-    hooks:
-      - id: flake8
-        additional_dependencies:
-          - torchfix==0.7.0
-          - flake8-pyproject==1.2.3
   # numpydoc
   - repo: https://github.com/Carreau/velin
     rev: 0.0.12
@@ -70,7 +62,7 @@ repos:
           - mdformat-gfm-alerts==2.0.0
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v22.1.0
+    rev: v22.1.1
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)

deepmd/backend/pretrained.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from collections.abc import (
+    Callable,
+)
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+)
+
+from deepmd.backend.backend import (
+    Backend,
+)
+from deepmd.pretrained.registry import (
+    available_model_names,
+)
+
+if TYPE_CHECKING:
+    from argparse import (
+        Namespace,
+    )
+
+    from deepmd.infer.deep_eval import (
+        DeepEvalBackend,
+    )
+    from deepmd.utils.neighbor_stat import (
+        NeighborStat,
+    )
+
+
+@Backend.register("pretrained")
+class PretrainedBackend(Backend):
+    """Internal virtual backend for pretrained model-name alias dispatch.
+
+    This backend is not intended to be selected explicitly by users as a real
+    compute backend (such as TensorFlow/PyTorch/Paddle/JAX). It only bridges
+    built-in pretrained model names into the regular deep-eval loading path.
+
+    For convenience, all built-in pretrained model names are registered as
+    suffix-like aliases, so users can pass model names directly, e.g.
+    ``DeepPot("DPA-3.2-5M")``.
+    """
+
+    name = "Pretrained"
+    features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL
+    suffixes: ClassVar[list[str]] = [
+        *[model_name.lower() for model_name in available_model_names()],
+    ]
+
+    def is_available(self) -> bool:
+        return True
+
+    @property
+    def entry_point_hook(self) -> Callable[["Namespace"], None]:
+        raise NotImplementedError("Unsupported backend: pretrained")
+
+    @property
+    def deep_eval(self) -> type["DeepEvalBackend"]:
+        from deepmd.pretrained.deep_eval import (
+            PretrainedDeepEvalBackend,
+        )
+
+        return PretrainedDeepEvalBackend
+
+    @property
+    def neighbor_stat(self) -> type["NeighborStat"]:
+        raise NotImplementedError("Unsupported backend: pretrained")
+
+    @property
+    def serialize_hook(self) -> Callable[[str], dict]:
+        raise NotImplementedError("Unsupported backend: pretrained")
+
+    @property
+    def deserialize_hook(self) -> Callable[[str, dict], None]:
+        raise NotImplementedError("Unsupported backend: pretrained")

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -74,6 +74,9 @@ def __init__(
         if hasattr(self.fitting_net, "reinit_exclude"):
             self.fitting_net.reinit_exclude(self.atom_exclude_types)
         self.type_map = type_map
+        self.add_chg_spin_ebd: bool = getattr(
+            self.descriptor, "add_chg_spin_ebd", False
+        )
         super().init_out_stat()
 
     def fitting_output_def(self) -> FittingOutputDef:
@@ -180,11 +183,38 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         atype = xp_take_first_n(extended_atype, 1, nloc)
+
+        # Handle default fparam if fitting net supports it
+        if (
+            hasattr(self.fitting_net, "get_dim_fparam")
+            and self.fitting_net.get_dim_fparam() > 0
+            and fparam is None
+        ):
+            # use default fparam
+            from deepmd.dpmodel.array_api import (
+                array_api_compat,
+            )
+
+            default_fparam = self.fitting_net.get_default_fparam()
+            assert default_fparam is not None
+            xp = array_api_compat.array_namespace(extended_coord)
+            default_fparam_array = xp.asarray(
+                default_fparam,
+                dtype=extended_coord.dtype,
+                device=array_api_compat.device(extended_coord),
+            )
+            fparam_input_for_des = xp.tile(
+                xp.reshape(default_fparam_array, (1, -1)), (nframes, 1)
+            )
+        else:
+            fparam_input_for_des = fparam
+
         descriptor, rot_mat, g2, h2, sw = self.descriptor(
             extended_coord,
             extended_atype,
             nlist,
             mapping=mapping,
+            fparam=fparam_input_for_des if self.add_chg_spin_ebd else None,
         )
         ret = self.fitting_net(
             descriptor,

deepmd/dpmodel/descriptor/dpa1.py

@@ -493,6 +493,7 @@ def call(
         atype_ext: Array,
         nlist: Array,
         mapping: Array | None = None,
+        fparam: Array | None = None,
     ) -> Array:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/dpa2.py

@@ -828,6 +828,7 @@ def call(
         atype_ext: Array,
         nlist: Array,
         mapping: Array | None = None,
+        fparam: Array | None = None,
     ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/dpa3.py

@@ -20,6 +20,7 @@
 )
 from deepmd.dpmodel.utils.network import (
     NativeLayer,
+    get_activation_fn,
 )
 from deepmd.dpmodel.utils.seed import (
     child_seed,
@@ -356,6 +357,7 @@ def __init__(
         use_tebd_bias: bool = False,
         use_loc_mapping: bool = True,
         type_map: list[str] | None = None,
+        add_chg_spin_ebd: bool = False,
     ) -> None:
         super().__init__()
 
@@ -410,6 +412,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
         )
 
         self.use_econf_tebd = use_econf_tebd
+        self.add_chg_spin_ebd = add_chg_spin_ebd
         self.use_tebd_bias = use_tebd_bias
         self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
@@ -428,6 +431,38 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
         )
         self.concat_output_tebd = concat_output_tebd
         self.precision = precision
+
+        if self.add_chg_spin_ebd:
+            self.cs_activation_fn = get_activation_fn(activation_function)
+            # -100 ~ 100 is a conservative bound
+            self.chg_embedding = TypeEmbedNet(
+                ntypes=200,
+                neuron=[self.tebd_dim],
+                padding=True,
+                activation_function="Linear",
+                precision=precision,
+                seed=child_seed(seed, 3),
+            )
+            # 100 is a conservative upper bound
+            self.spin_embedding = TypeEmbedNet(
+                ntypes=100,
+                neuron=[self.tebd_dim],
+                padding=True,
+                activation_function="Linear",
+                precision=precision,
+                seed=child_seed(seed, 4),
+            )
+            self.mix_cs_mlp = NativeLayer(
+                2 * self.tebd_dim,
+                self.tebd_dim,
+                precision=precision,
+                seed=child_seed(seed, 5),
+            )
+        else:
+            self.chg_embedding = None
+            self.spin_embedding = None
+            self.mix_cs_mlp = None
+
         self.exclude_types = exclude_types
         self.env_protection = env_protection
         self.trainable = trainable
@@ -579,6 +614,7 @@ def call(
         atype_ext: Array,
         nlist: Array,
         mapping: Array | None = None,
+        fparam: Array | None = None,
     ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 
@@ -625,6 +661,27 @@ def call(
                 xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
                 (nframes, nall, self.tebd_dim),
             )
+
+        if self.add_chg_spin_ebd:
+            assert fparam is not None
+            assert self.chg_embedding is not None
+            assert self.spin_embedding is not None
+            chg_tebd = self.chg_embedding.call()
+            spin_tebd = self.spin_embedding.call()
+            charge = xp.astype(fparam[:, 0], xp.int64) + 100
+            spin = xp.astype(fparam[:, 1], xp.int64)
+            chg_ebd = xp.reshape(
+                xp.take(chg_tebd, xp.reshape(charge, (-1,)), axis=0),
+                (nframes, self.tebd_dim),
+            )
+            spin_ebd = xp.reshape(
+                xp.take(spin_tebd, xp.reshape(spin, (-1,)), axis=0),
+                (nframes, self.tebd_dim),
+            )
+            cs_cat = xp.concat([chg_ebd, spin_ebd], axis=-1)
+            sys_cs_embd = self.cs_activation_fn(self.mix_cs_mlp.call(cs_cat))
+            node_ebd_ext = node_ebd_ext + xp.expand_dims(sys_cs_embd, axis=1)
+
         node_ebd_inp = node_ebd_ext[:, :nloc, :]
         # repflows
         node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(
@@ -655,9 +712,14 @@ def serialize(self) -> dict:
             "use_econf_tebd": self.use_econf_tebd,
             "use_tebd_bias": self.use_tebd_bias,
             "use_loc_mapping": self.use_loc_mapping,
+            "add_chg_spin_ebd": self.add_chg_spin_ebd,
             "type_map": self.type_map,
             "type_embedding": self.type_embedding.serialize(),
         }
+        if self.add_chg_spin_ebd:
+            data["chg_embedding"] = self.chg_embedding.serialize()
+            data["spin_embedding"] = self.spin_embedding.serialize()
+            data["mix_cs_mlp"] = self.mix_cs_mlp.serialize()
         repflow_variable = {
             "edge_embd": repflows.edge_embd.serialize(),
             "angle_embd": repflows.angle_embd.serialize(),
@@ -684,10 +746,18 @@ def deserialize(cls, data: dict) -> "DescrptDPA3":
         data.pop("type")
         repflow_variable = data.pop("repflow_variable").copy()
         type_embedding = data.pop("type_embedding")
+        chg_embedding = data.pop("chg_embedding", None)
+        spin_embedding = data.pop("spin_embedding", None)
+        mix_cs_mlp = data.pop("mix_cs_mlp", None)
         data["repflow"] = RepFlowArgs(**data.pop("repflow_args"))
         obj = cls(**data)
         obj.type_embedding = TypeEmbedNet.deserialize(type_embedding)
 
+        if obj.add_chg_spin_ebd and chg_embedding is not None:
+            obj.chg_embedding = TypeEmbedNet.deserialize(chg_embedding)
+            obj.spin_embedding = TypeEmbedNet.deserialize(spin_embedding)
+            obj.mix_cs_mlp = NativeLayer.deserialize(mix_cs_mlp)
+
         # deserialize repflow
         statistic_repflows = repflow_variable.pop("@variables")
         env_mat = repflow_variable.pop("env_mat")

deepmd/dpmodel/descriptor/hybrid.py

@@ -245,7 +245,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
@@ -275,6 +275,7 @@ def call(
         atype_ext: Array,
         nlist: Array,
         mapping: Array | None = None,
+        fparam: Array | None = None,
     ) -> tuple[
         Array,
         Array | None,