feat(pt_expt): add descriptors dpa1 dpa2 dpa3 and hybrid

Author: Han Wang

deepmd/dpmodel/array_api.py

@@ -29,6 +29,9 @@ def xp_swapaxes(a: Array, axis1: int, axis2: int) -> Array:
 
 def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
     xp = array_api_compat.array_namespace(arr)
+    # torch.take_along_dim requires int64 indices
+    if array_api_compat.is_torch_array(indices):
+        indices = xp.astype(indices, xp.int64)
     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)

deepmd/dpmodel/descriptor/dpa2.py

@@ -573,6 +573,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
         self.smooth = smooth
         self.exclude_types = exclude_types
         self.env_protection = env_protection
+        self.rcut_smth = self.repinit.get_rcut_smth()
         self.trainable = trainable
         self.add_tebd_to_repinit_out = add_tebd_to_repinit_out
 

deepmd/dpmodel/descriptor/repflows.py

@@ -595,8 +595,12 @@ def call(
             # n_angle x 1
             a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask]
         else:
-            edge_index = xp.zeros([2, 1], dtype=nlist.dtype)
-            angle_index = xp.zeros([3, 1], dtype=nlist.dtype)
+            edge_index = xp.zeros(
+                [2, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist)
+            )
+            angle_index = xp.zeros(
+                [3, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist)
+            )
 
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
@@ -1711,6 +1715,7 @@ def call(
                         xp.zeros(
                             (nb, nloc, self.nnei - self.a_sel, self.e_dim),
                             dtype=edge_ebd.dtype,
+                            device=array_api_compat.device(edge_ebd),
                         ),
                     ],
                     axis=2,
@@ -1741,6 +1746,7 @@ def call(
                         xp.zeros(
                             (nb, nloc, self.nnei - self.a_sel),
                             dtype=a_nlist_mask.dtype,
+                            device=array_api_compat.device(a_nlist_mask),
                         ),
                     ],
                     axis=-1,

deepmd/dpmodel/descriptor/repformers.py

@@ -475,7 +475,7 @@ def call(
         g1 = self.act(atype_embd)
         # nf x nloc x nnei x 1,  nf x nloc x nnei x 3
         if not self.direct_dist:
-            g2, h2 = xp.split(dmatrix, [1], axis=-1)
+            g2, h2 = dmatrix[..., :1], dmatrix[..., 1:]
         else:
             g2, h2 = safe_for_vector_norm(diff, axis=-1, keepdims=True), diff
             g2 = g2 / self.rcut
@@ -756,10 +756,12 @@ def _cal_hg(
     else:
         g = _apply_switch(g, sw)
         if not use_sqrt_nnei:
-            invnnei = (1.0 / float(nnei)) * xp.ones((nf, nloc, 1, 1), dtype=g.dtype)
+            invnnei = (1.0 / float(nnei)) * xp.ones(
+                (nf, nloc, 1, 1), dtype=g.dtype, device=array_api_compat.device(g)
+            )
         else:
             invnnei = (1.0 / (float(nnei) ** 0.5)) * xp.ones(
-                (nf, nloc, 1, 1), dtype=g.dtype
+                (nf, nloc, 1, 1), dtype=g.dtype, device=array_api_compat.device(g)
             )
     # nf x nloc x 3 x ng
     hg = xp.matmul(xp.matrix_transpose(h), g) * invnnei
@@ -1655,7 +1657,9 @@ def _update_g1_conv(
             invnnei = invnnei[:, :, xp.newaxis]
         else:
             gg1 = _apply_switch(gg1, sw)
-            invnnei = (1.0 / float(nnei)) * xp.ones((nf, nloc, 1), dtype=gg1.dtype)
+            invnnei = (1.0 / float(nnei)) * xp.ones(
+                (nf, nloc, 1), dtype=gg1.dtype, device=array_api_compat.device(gg1)
+            )
         if not self.g1_out_conv:
             # nf x nloc x ng2
             g1_11 = xp.sum(g2 * gg1, axis=2) * invnnei

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -427,45 +427,64 @@ def call(
             The smooth switch function.
         """
         del mapping
+        xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
+        input_dtype = coord_ext.dtype
         # nf x nloc x nnei x 4
         rr, diff, ww = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.davg, self.dstd
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.davg[...],
+            self.dstd[...],
         )
         nf, nloc, nnei, _ = rr.shape
-        sec = np.append([0], np.cumsum(self.sel))
+        sec = self.sel_cumsum
 
         ng = self.neuron[-1]
-        gr = np.zeros([nf * nloc, ng, 4], dtype=PRECISION_DICT[self.precision])
+        gr = xp.zeros(
+            [nf * nloc, ng, 4],
+            dtype=input_dtype,
+            device=array_api_compat.device(coord_ext),
+        )
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # merge nf and nloc axis, so for type_one_side == False,
         # we don't require atype is the same in all frames
-        exclude_mask = exclude_mask.reshape(nf * nloc, nnei)
-        rr = rr.reshape(nf * nloc, nnei, 4)
+        exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
+        rr = xp.reshape(rr, (nf * nloc, nnei, 4))
+        rr = xp.astype(rr, self.dstd.dtype)
+
+        if not self.type_one_side:
+            # nf x nloc -> (nf * nloc)
+            atype_loc = xp.reshape(atype_ext[:, :nloc], (nf * nloc,))
 
         for embedding_idx in itertools.product(
             range(self.ntypes), repeat=self.embeddings.ndim
         ):
             if self.type_one_side:
                 (tt,) = embedding_idx
-                ti_mask = np.s_[:]
             else:
                 ti, tt = embedding_idx
-                ti_mask = atype_ext[:, :nloc].ravel() == ti
-            mm = exclude_mask[ti_mask, sec[tt] : sec[tt + 1]]
-            tr = rr[ti_mask, sec[tt] : sec[tt + 1], :]
-            tr = tr * mm[:, :, None]
+            mm = exclude_mask[:, sec[tt] : sec[tt + 1]]
+            tr = rr[:, sec[tt] : sec[tt + 1], :]
+            tr = tr * xp.astype(mm[:, :, None], tr.dtype)
             ss = tr[..., 0:1]
             gg = self.cal_g(ss, embedding_idx)
-            gr_tmp = np.einsum("lni,lnj->lij", gg, tr)
-            gr[ti_mask] += gr_tmp
-        gr = gr.reshape(nf, nloc, ng, 4)
+            gr_tmp = xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1)
+            if not self.type_one_side:
+                # (nf * nloc) x 1 x 1
+                ti_mask = xp.astype(
+                    xp.reshape(atype_loc == ti, (nf * nloc, 1, 1)), gr_tmp.dtype
+                )
+                gr_tmp = gr_tmp * ti_mask
+            gr += gr_tmp
+        gr = xp.reshape(gr, (nf, nloc, ng, 4))
         # nf x nloc x ng x 4
         gr /= self.nnei
         gr1 = gr[:, :, : self.axis_neuron, :]
         # nf x nloc x ng x ng1
-        grrg = np.einsum("flid,fljd->flij", gr, gr1)
+        grrg = xp.sum(gr[:, :, :, None, :] * gr1[:, :, None, :, :], axis=4)
         # nf x nloc x (ng x ng1)
-        grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron)
+        grrg = xp.reshape(grrg, (nf, nloc, ng * self.axis_neuron))
         return grrg, gr[..., 1:], None, None, ww
 
     def serialize(self) -> dict:
@@ -553,94 +572,4 @@ def update_sel(
         return local_jdata_cpy, min_nbor_dist
 
 
-class DescrptSeAArrayAPI(DescrptSeA):
-    @cast_precision
-    def call(
-        self,
-        coord_ext: Array,
-        atype_ext: Array,
-        nlist: Array,
-        mapping: Array | None = None,
-    ) -> Array:
-        """Compute the descriptor.
-
-        Parameters
-        ----------
-        coord_ext
-            The extended coordinates of atoms. shape: nf x (nallx3)
-        atype_ext
-            The extended aotm types. shape: nf x nall
-        nlist
-            The neighbor list. shape: nf x nloc x nnei
-        mapping
-            The index mapping from extended to local region. not used by this descriptor.
-
-        Returns
-        -------
-        descriptor
-            The descriptor. shape: nf x nloc x (ng x axis_neuron)
-        gr
-            The rotationally equivariant and permutationally invariant single particle
-            representation. shape: nf x nloc x ng x 3
-        g2
-            The rotationally invariant pair-partical representation.
-            this descriptor returns None
-        h2
-            The rotationally equivariant pair-partical representation.
-            this descriptor returns None
-        sw
-            The smooth switch function.
-        """
-        if not self.type_one_side:
-            raise NotImplementedError(
-                "type_one_side == False is not supported in DescrptSeAArrayAPI"
-            )
-        del mapping
-        xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
-        input_dtype = coord_ext.dtype
-        # nf x nloc x nnei x 4
-        rr, diff, ww = self.env_mat.call(
-            coord_ext,
-            atype_ext,
-            nlist,
-            self.davg[...],
-            self.dstd[...],
-        )
-        nf, nloc, nnei, _ = rr.shape
-        sec = self.sel_cumsum
-
-        ng = self.neuron[-1]
-        gr = xp.zeros(
-            [nf * nloc, ng, 4],
-            dtype=input_dtype,
-            device=array_api_compat.device(coord_ext),
-        )
-        exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
-        # merge nf and nloc axis, so for type_one_side == False,
-        # we don't require atype is the same in all frames
-        exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
-        rr = xp.reshape(rr, (nf * nloc, nnei, 4))
-        rr = xp.astype(rr, self.dstd.dtype)
-
-        for embedding_idx in itertools.product(
-            range(self.ntypes), repeat=self.embeddings.ndim
-        ):
-            (tt,) = embedding_idx
-            mm = exclude_mask[:, sec[tt] : sec[tt + 1]]
-            tr = rr[:, sec[tt] : sec[tt + 1], :]
-            tr = tr * xp.astype(mm[:, :, None], tr.dtype)
-            ss = tr[..., 0:1]
-            gg = self.cal_g(ss, embedding_idx)
-            # gr_tmp = xp.einsum("lni,lnj->lij", gg, tr)
-            gr_tmp = xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1)
-            gr += gr_tmp
-        gr = xp.reshape(gr, (nf, nloc, ng, 4))
-        # nf x nloc x ng x 4
-        gr /= self.nnei
-        gr1 = gr[:, :, : self.axis_neuron, :]
-        # nf x nloc x ng x ng1
-        # grrg = xp.einsum("flid,fljd->flij", gr, gr1)
-        grrg = xp.sum(gr[:, :, :, None, :] * gr1[:, :, None, :, :], axis=4)
-        # nf x nloc x (ng x ng1)
-        grrg = xp.reshape(grrg, (nf, nloc, ng * self.axis_neuron))
-        return grrg, gr[..., 1:], None, None, ww
+DescrptSeAArrayAPI = DescrptSeA