Merge branch 'devel' into enerhess

.github/workflows/test_cc.yml

@@ -70,7 +70,7 @@ jobs:
         PATH: ${{ github.workspace }}/dp_test/bin:$PATH
         LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib
       if: ${{ !matrix.check_memleak }}
-    - uses: codecov/codecov-action@v4
+    - uses: codecov/codecov-action@v5
       env:
         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
   pass:

.github/workflows/test_python.yml

@@ -71,7 +71,7 @@ jobs:
         name: split-${{ matrix.python }}-${{ matrix.group }}
         path: .test_durations_${{ matrix.group }}
         include-hidden-files: true
-    - uses: codecov/codecov-action@v4
+    - uses: codecov/codecov-action@v5
       env:
         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
   update_durations:

.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.7.3
+    rev: v0.7.4
     hooks:
       - id: ruff
         args: ["--fix"]

backend/find_pytorch.py

@@ -133,7 +133,8 @@ def get_pt_requirement(pt_version: str = "") -> dict:
             # https://peps.python.org/pep-0440/#version-matching
             f"torch=={Version(pt_version).base_version}.*"
             if pt_version != ""
-            else "torch>=2a",
+            # https://github.com/pytorch/pytorch/commit/7e0c26d4d80d6602aed95cb680dfc09c9ce533bc
+            else "torch>=2.1.0"
         ],
     }
 

deepmd/dpmodel/array_api.py

@@ -2,6 +2,9 @@
 """Utilities for the array API."""
 
 import array_api_compat
+from packaging.version import (
+    Version,
+)
 
 
 def support_array_api(version: str) -> callable:
@@ -45,6 +48,9 @@ def xp_swapaxes(a, axis1, axis2):
 
 def xp_take_along_axis(arr, indices, axis):
     xp = array_api_compat.array_namespace(arr)
+    if Version(xp.__array_api_version__) >= Version("2024.12"):
+        # see: https://github.com/data-apis/array-api-strict/blob/d086c619a58f35c38240592ef994aa19ca7beebc/array_api_strict/_indexing_functions.py#L30-L39
+        return xp.take_along_axis(arr, indices, axis=axis)
     arr = xp_swapaxes(arr, axis, -1)
     indices = xp_swapaxes(indices, axis, -1)
 

deepmd/dpmodel/atomic_model/__init__.py

@@ -17,9 +17,18 @@
 from .base_atomic_model import (
     BaseAtomicModel,
 )
+from .dipole_atomic_model import (
+    DPDipoleAtomicModel,
+)
+from .dos_atomic_model import (
+    DPDOSAtomicModel,
+)
 from .dp_atomic_model import (
     DPAtomicModel,
 )
+from .energy_atomic_model import (
+    DPEnergyAtomicModel,
+)
 from .linear_atomic_model import (
     DPZBLLinearEnergyAtomicModel,
     LinearEnergyAtomicModel,
@@ -30,12 +39,19 @@
 from .pairtab_atomic_model import (
     PairTabAtomicModel,
 )
+from .polar_atomic_model import (
+    DPPolarAtomicModel,
+)
 
 __all__ = [
     "make_base_atomic_model",
     "BaseAtomicModel",
     "DPAtomicModel",
+    "DPEnergyAtomicModel",
     "PairTabAtomicModel",
     "LinearEnergyAtomicModel",
     "DPZBLLinearEnergyAtomicModel",
+    "DPDOSAtomicModel",
+    "DPPolarAtomicModel",
+    "DPDipoleAtomicModel",
 ]

deepmd/dpmodel/atomic_model/dipole_atomic_model.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from deepmd.dpmodel.fitting.dipole_fitting import (
+    DipoleFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDipoleAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, DipoleFitting):
+            raise TypeError(
+                "fitting must be an instance of DipoleFitting for DPDipoleAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        # dipole not applying bias
+        return ret

deepmd/dpmodel/atomic_model/dos_atomic_model.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.fitting.dos_fitting import (
+    DOSFittingNet,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDOSAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, DOSFittingNet):
+            raise TypeError(
+                "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/energy_atomic_model.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.fitting.ener_fitting import (
+    EnergyFittingNet,
+    InvarFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPEnergyAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not (
+            isinstance(fitting, EnergyFittingNet) or isinstance(fitting, InvarFitting)
+        ):
+            raise TypeError(
+                "fitting must be an instance of EnergyFittingNet or InvarFitting for DPEnergyAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/polar_atomic_model.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import array_api_compat
+import numpy as np
+
+from deepmd.dpmodel.fitting.polarizability_fitting import (
+    PolarFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPPolarAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, PolarFitting):
+            raise TypeError(
+                "fitting must be an instance of PolarFitting for DPPolarAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        """Apply the stat to each atomic output.
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        xp = array_api_compat.array_namespace(atype)
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+
+        if self.fitting.shift_diag:
+            nframes, nloc = atype.shape
+            dtype = out_bias[self.bias_keys[0]].dtype
+            for kk in self.bias_keys:
+                ntypes = out_bias[kk].shape[0]
+                temp = xp.mean(
+                    xp.diagonal(out_bias[kk].reshape(ntypes, 3, 3), axis1=1, axis2=2),
+                    axis=1,
+                )
+                modified_bias = temp[atype]
+
+                # (nframes, nloc, 1)
+                modified_bias = (
+                    modified_bias[..., xp.newaxis] * (self.fitting.scale[atype])
+                )
+
+                eye = xp.eye(3, dtype=dtype)
+                eye = xp.tile(eye, (nframes, nloc, 1, 1))
+                # (nframes, nloc, 3, 3)
+                modified_bias = modified_bias[..., xp.newaxis] * eye
+
+                # nf x nloc x odims, out_bias: ntypes x odims
+                ret[kk] = ret[kk] + modified_bias
+        return ret

deepmd/dpmodel/atomic_model/property_atomic_model.py

@@ -9,6 +9,9 @@
 
 
 class DPPropertyAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs) -> None:
-        assert isinstance(fitting, PropertyFittingNet)
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, PropertyFittingNet):
+            raise TypeError(
+                "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel"
+            )
         super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/descriptor/descriptor.py

@@ -123,6 +123,7 @@ def call(
         extended_atype: np.ndarray,
         extended_atype_embd: Optional[np.ndarray] = None,
         mapping: Optional[np.ndarray] = None,
+        type_embedding: Optional[np.ndarray] = None,
     ):
         """Calculate DescriptorBlock."""
         pass

deepmd/dpmodel/descriptor/dpa1.py

@@ -494,9 +494,10 @@ def call(
         xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
         nf, nloc, nnei = nlist.shape
         nall = xp.reshape(coord_ext, (nf, -1)).shape[1] // 3
+        type_embedding = self.type_embedding.call()
         # nf x nall x tebd_dim
         atype_embd_ext = xp.reshape(
-            xp.take(self.type_embedding.call(), 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
@@ -507,6 +508,7 @@ def call(
             atype_ext,
             atype_embd_ext,
             mapping=None,
+            type_embedding=type_embedding,
         )
         # nf x nloc x (ng x ng1 + tebd_dim)
         if self.concat_output_tebd:
@@ -874,10 +876,6 @@ def cal_g_strip(
         embedding_idx,
     ):
         assert self.embeddings_strip is not None
-        xp = array_api_compat.array_namespace(ss)
-        nfnl, nnei = ss.shape[0:2]
-        shape2 = math.prod(ss.shape[2:])
-        ss = xp.reshape(ss, (nfnl, nnei, shape2))
         # nfnl x nnei x ng
         gg = self.embeddings_strip[embedding_idx].call(ss)
         return gg
@@ -889,13 +887,15 @@ def call(
         atype_ext: np.ndarray,
         atype_embd_ext: Optional[np.ndarray] = None,
         mapping: Optional[np.ndarray] = None,
+        type_embedding: Optional[np.ndarray] = None,
     ):
         xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext)
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(
             coord_ext, atype_ext, nlist, self.mean, self.stddev
         )
         nf, nloc, nnei, _ = dmatrix.shape
+        atype = atype_ext[:, :nloc]
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # nfnl x nnei
         exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
@@ -906,28 +906,33 @@ def call(
         dmatrix = xp.reshape(dmatrix, (nf * nloc, nnei, 4))
         # nfnl x nnei x 1
         sw = xp.reshape(sw, (nf * nloc, nnei, 1))
-        # nfnl x tebd_dim
-        atype_embd = xp.reshape(atype_embd_ext[:, :nloc, :], (nf * nloc, self.tebd_dim))
-        # nfnl x nnei x tebd_dim
-        atype_embd_nnei = xp.tile(atype_embd[:, xp.newaxis, :], (1, nnei, 1))
         # nfnl x nnei
         nlist_mask = nlist != -1
         # nfnl x nnei x 1
         sw = xp.where(nlist_mask[:, :, None], sw, xp.full_like(sw, 0.0))
         nlist_masked = xp.where(nlist_mask, nlist, xp.zeros_like(nlist))
-        index = xp.tile(xp.reshape(nlist_masked, (nf, -1, 1)), (1, 1, self.tebd_dim))
-        # nfnl x nnei x tebd_dim
-        atype_embd_nlist = xp_take_along_axis(atype_embd_ext, index, axis=1)
-        atype_embd_nlist = xp.reshape(
-            atype_embd_nlist, (nf * nloc, nnei, self.tebd_dim)
-        )
         ng = self.neuron[-1]
+        nt = self.tebd_dim
         # nfnl x nnei x 4
         rr = xp.reshape(dmatrix, (nf * nloc, nnei, 4))
         rr = rr * xp.astype(exclude_mask[:, :, None], rr.dtype)
         # nfnl x nnei x 1
         ss = rr[..., 0:1]
         if self.tebd_input_mode in ["concat"]:
+            # nfnl x tebd_dim
+            atype_embd = xp.reshape(
+                atype_embd_ext[:, :nloc, :], (nf * nloc, self.tebd_dim)
+            )
+            # nfnl x nnei x tebd_dim
+            atype_embd_nnei = xp.tile(atype_embd[:, xp.newaxis, :], (1, nnei, 1))
+            index = xp.tile(
+                xp.reshape(nlist_masked, (nf, -1, 1)), (1, 1, self.tebd_dim)
+            )
+            # nfnl x nnei x tebd_dim
+            atype_embd_nlist = xp_take_along_axis(atype_embd_ext, index, axis=1)
+            atype_embd_nlist = xp.reshape(
+                atype_embd_nlist, (nf * nloc, nnei, self.tebd_dim)
+            )
             if not self.type_one_side:
                 # nfnl x nnei x (1 + 2 * tebd_dim)
                 ss = xp.concat([ss, atype_embd_nlist, atype_embd_nnei], axis=-1)
@@ -941,14 +946,48 @@ def call(
             # nfnl x nnei x ng
             gg_s = self.cal_g(ss, 0)
             assert self.embeddings_strip is not None
-            if not self.type_one_side:
-                # nfnl x nnei x (tebd_dim * 2)
-                tt = xp.concat([atype_embd_nlist, atype_embd_nnei], axis=-1)
+            assert type_embedding is not None
+            ntypes_with_padding = type_embedding.shape[0]
+            # nf x (nl x nnei)
+            nlist_index = xp.reshape(nlist_masked, (nf, nloc * nnei))
+            # nf x (nl x nnei)
+            nei_type = xp_take_along_axis(atype_ext, nlist_index, axis=1)
+            # (nf x nl x nnei) x ng
+            nei_type_index = xp.tile(xp.reshape(nei_type, (-1, 1)), (1, ng))
+            if self.type_one_side:
+                tt_full = self.cal_g_strip(type_embedding, 0)
+                # (nf x nl x nnei) x ng
+                gg_t = xp_take_along_axis(tt_full, nei_type_index, axis=0)
             else:
-                # nfnl x nnei x tebd_dim
-                tt = atype_embd_nlist
-            # nfnl x nnei x ng
-            gg_t = self.cal_g_strip(tt, 0)
+                idx_i = xp.reshape(
+                    xp.tile(
+                        (xp.reshape(atype, (-1, 1)) * ntypes_with_padding), (1, nnei)
+                    ),
+                    (-1),
+                )
+                idx_j = xp.reshape(nei_type, (-1,))
+                # (nf x nl x nnei) x ng
+                idx = xp.tile(xp.reshape((idx_i + idx_j), (-1, 1)), (1, ng))
+                # (ntypes) * ntypes * nt
+                type_embedding_nei = xp.tile(
+                    xp.reshape(type_embedding, (1, ntypes_with_padding, nt)),
+                    (ntypes_with_padding, 1, 1),
+                )
+                # ntypes * (ntypes) * nt
+                type_embedding_center = xp.tile(
+                    xp.reshape(type_embedding, (ntypes_with_padding, 1, nt)),
+                    (1, ntypes_with_padding, 1),
+                )
+                # (ntypes * ntypes) * (nt+nt)
+                two_side_type_embedding = xp.reshape(
+                    xp.concat([type_embedding_nei, type_embedding_center], axis=-1),
+                    (-1, nt * 2),
+                )
+                tt_full = self.cal_g_strip(two_side_type_embedding, 0)
+                # (nf x nl x nnei) x ng
+                gg_t = xp_take_along_axis(tt_full, idx, axis=0)
+            # (nf x nl) x nnei x ng
+            gg_t = xp.reshape(gg_t, (nf * nloc, nnei, ng))
             if self.smooth:
                 gg_t = gg_t * xp.reshape(sw, (-1, self.nnei, 1))
             # nfnl x nnei x ng