Align dpa3 numpy backend with PT loc mapping

Author: caic99

deepmd/dpmodel/descriptor/dpa3.py

@@ -271,6 +271,8 @@ class DescrptDPA3(NativeOP, BaseDescriptor):
         Whether to use electronic configuration type embedding.
     use_tebd_bias : bool, Optional
         Whether to use bias in the type embedding layer.
+    use_loc_mapping : bool, Optional
+        Whether to use local atom index mapping in non-parallel inference.
     type_map : list[str], Optional
         A list of strings. Give the name to each type of atoms.
     """
@@ -290,6 +292,7 @@ def __init__(
         seed: Optional[Union[int, list[int]]] = None,
         use_econf_tebd: bool = False,
         use_tebd_bias: bool = False,
+        use_loc_mapping: bool = True,
         type_map: Optional[list[str]] = None,
     ) -> None:
         super().__init__()
@@ -335,6 +338,7 @@ def init_subclass_params(sub_data, sub_class):
             use_exp_switch=self.repflow_args.use_exp_switch,
             use_dynamic_sel=self.repflow_args.use_dynamic_sel,
             sel_reduce_factor=self.repflow_args.sel_reduce_factor,
+            use_loc_mapping=use_loc_mapping,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
@@ -343,6 +347,7 @@ def init_subclass_params(sub_data, sub_class):
 
         self.use_econf_tebd = use_econf_tebd
         self.use_tebd_bias = use_tebd_bias
+        self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
         self.tebd_dim = self.repflow_args.n_dim
         self.type_embedding = TypeEmbedNet(
@@ -541,10 +546,16 @@ def call(
         nall = xp.reshape(coord_ext, (nframes, -1)).shape[1] // 3
 
         type_embedding = self.type_embedding.call()
-        node_ebd_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
-            (nframes, nall, self.tebd_dim),
-        )
+        if self.use_loc_mapping:
+            node_ebd_ext = xp.reshape(
+                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),
+                (nframes, nall, self.tebd_dim),
+            )
         node_ebd_inp = node_ebd_ext[:, :nloc, :]
         # repflows
         node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(

deepmd/dpmodel/descriptor/repflows.py

@@ -145,6 +145,8 @@ class DescrptBlockRepflows(NativeOP, DescriptorBlock):
         In the dynamic selection case, neighbor-scale normalization will use `e_sel / sel_reduce_factor`
         or `a_sel / sel_reduce_factor` instead of the raw `e_sel` or `a_sel` values,
         accommodating larger selection numbers.
+    use_loc_mapping : bool, optional
+        Whether to use local atom index mapping in non-parallel inference.
     ntypes : int
         Number of element types
     activation_function : str, optional
@@ -196,6 +198,7 @@ def __init__(
         use_exp_switch: bool = False,
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
+        use_loc_mapping: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         super().__init__()
@@ -229,6 +232,7 @@ def __init__(
         self.smooth_edge_update = smooth_edge_update
         self.use_exp_switch = use_exp_switch
         self.use_dynamic_sel = use_dynamic_sel
+        self.use_loc_mapping = use_loc_mapping
         self.sel_reduce_factor = sel_reduce_factor
         if self.use_dynamic_sel and not self.smooth_edge_update:
             raise NotImplementedError(
@@ -527,10 +531,18 @@ def call(
             cosine_ij, (nframes, nloc, self.a_sel, self.a_sel, 1)
         ) / (xp.pi**0.5)
 
+        if self.use_loc_mapping:
+            assert mapping is not None
+            flat_map = xp.reshape(mapping, (nframes, -1))
+            nlist = xp.reshape(
+                xp_take_along_axis(flat_map, xp.reshape(nlist, (nframes, -1)), axis=1),
+                nlist.shape,
+            )
+        
         if self.use_dynamic_sel:
             # get graph index
             edge_index, angle_index = get_graph_index(
-                nlist, nlist_mask, a_nlist_mask, nall
+                nlist, nlist_mask, a_nlist_mask, nall, use_loc_mapping=self.use_loc_mapping
             )
             # flat all the tensors
             # n_edge x 1
@@ -561,7 +573,11 @@ def call(
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
             # node_ebd_ext: nb x nall x n_dim
-            node_ebd_ext = xp_take_along_axis(node_ebd, mapping, axis=1)
+            node_ebd_ext = (
+                node_ebd
+                if self.use_loc_mapping
+                else xp_take_along_axis(node_ebd, mapping, axis=1)
+            )
             node_ebd, edge_ebd, angle_ebd = ll.call(
                 node_ebd_ext,
                 edge_ebd,

deepmd/dpmodel/utils/network.py

@@ -1006,6 +1006,7 @@ def get_graph_index(
     nlist_mask: np.ndarray,
     a_nlist_mask: np.ndarray,
     nall: int,
+    use_loc_mapping: bool = True,
 ):
     """
     Get the index mapping for edge graph and angle graph, ready in `aggregate` or `index_select`.
@@ -1020,6 +1021,8 @@ def get_graph_index(
         Masks of the neighbor list for angle. real nei 1 otherwise 0
     nall
         The number of extended atoms.
+    use_loc_mapping
+        Whether to use local atom index mapping in non-parallel inference.
 
     Returns
     -------
@@ -1060,7 +1063,7 @@ def get_graph_index(
     n2e_index = n2e_index[xp.astype(nlist_mask, xp.bool)]
 
     # node_ext(j) to edge(ij) index_select
-    frame_shift = xp.arange(nf, dtype=nlist.dtype) * nall
+    frame_shift = xp.arange(nf, dtype=nlist.dtype) * (nall if not use_loc_mapping else nloc)
     shifted_nlist = nlist + frame_shift[:, xp.newaxis, xp.newaxis]
     # n_edge
     n_ext2e_index = shifted_nlist[xp.astype(nlist_mask, xp.bool)]