feat(pt): add use_loc_mapping

Author: iProzd

deepmd/pt/model/descriptor/dpa3.py

@@ -150,13 +150,48 @@ def init_subclass_params(sub_data, sub_class):
             fix_stat_std=self.repflow_args.fix_stat_std,
             optim_update=self.repflow_args.optim_update,
             smooth_edge_update=self.repflow_args.smooth_edge_update,
+            angle_multi_freq=self.repflow_args.angle_multi_freq,
+            use_dynamic_sel=self.repflow_args.use_dynamic_sel,
+            sel_reduce_factor=self.repflow_args.sel_reduce_factor,
+            use_env_envelope=self.repflow_args.use_env_envelope,
+            use_new_sw=self.repflow_args.use_new_sw,
+            update_dihedral=self.repflow_args.update_dihedral,
+            d_dim=self.repflow_args.d_dim,
+            d_sel=self.repflow_args.d_sel,
+            d_rcut=self.repflow_args.d_rcut,
+            d_rcut_smth=self.repflow_args.d_rcut_smth,
+            use_ffn_node_edge_message=self.repflow_args.use_ffn_node_edge_message,
+            use_ffn_edge_edge_message=self.repflow_args.use_ffn_edge_edge_message,
+            use_ffn_edge_angle_message=self.repflow_args.use_ffn_edge_angle_message,
+            use_ffn_angle_angle_message=self.repflow_args.use_ffn_angle_angle_message,
+            ffn_hidden_dim=self.repflow_args.ffn_hidden_dim,
+            edge_use_concat_rbf=self.repflow_args.edge_use_concat_rbf,
+            edge_use_rbf=self.repflow_args.edge_use_rbf,
+            edge_use_dist=self.repflow_args.edge_use_dist,
+            embed_use_bias=self.repflow_args.embed_use_bias,
+            edge_use_attn=self.repflow_args.edge_use_attn,
+            edge_attn_hidden=self.repflow_args.edge_attn_hidden,
+            edge_attn_head=self.repflow_args.edge_attn_head,
+            edge_attn_use_ln=self.repflow_args.edge_attn_use_ln,
+            edge_rbf_dot_self=self.repflow_args.edge_rbf_dot_self,
+            edge_rbf_dot_message=self.repflow_args.edge_rbf_dot_message,
+            edge_use_esen_rbf=self.repflow_args.edge_use_esen_rbf,
+            edge_use_esen_atom_ebd=self.repflow_args.edge_use_esen_atom_ebd,
+            edge_use_esen_env=self.repflow_args.edge_use_esen_env,
+            residual_pref=self.repflow_args.residual_pref,
+            tebd_use_act=self.repflow_args.tebd_use_act,
+            message_use_self_concat=self.repflow_args.message_use_self_concat,
+            use_slim_message=self.repflow_args.use_slim_message,
+            use_combined_output=self.repflow_args.use_combined_output,
+            use_loc_mapping=use_loc_mapping,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
             seed=child_seed(seed, 1),
         )
 
         self.use_econf_tebd = use_econf_tebd
+        self.use_loc_mapping = use_loc_mapping
         self.use_tebd_bias = use_tebd_bias
         self.type_map = type_map
         self.tebd_dim = self.repflow_args.n_dim
@@ -466,12 +501,16 @@ def forward(
             The smooth switch function. shape: nf x nloc x nnei
 
         """
+        parrallel_mode = comm_dict is not None
         # cast the input to internal precsion
         extended_coord = extended_coord.to(dtype=self.prec)
         nframes, nloc, nnei = nlist.shape
         nall = extended_coord.view(nframes, -1).shape[1] // 3
 
-        node_ebd_ext = self.type_embedding(extended_atype)
+        if not parrallel_mode and self.use_loc_mapping:
+            node_ebd_ext = self.type_embedding(extended_atype[:, :nloc])
+        else:
+            node_ebd_ext = self.type_embedding(extended_atype)
         node_ebd_inp = node_ebd_ext[:, :nloc, :]
         # repflows
         node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(

deepmd/pt/model/descriptor/repflow_layer.py

@@ -469,7 +469,7 @@ def optim_edge_update(
 
     def forward(
         self,
-        node_ebd_ext: torch.Tensor,  # nf x nall x n_dim
+        node_ebd_ext: torch.Tensor,  # nf x nall x n_dim [OR] nf x nloc x n_dim when not parrallel_mode
         edge_ebd: torch.Tensor,  # nf x nloc x nnei x e_dim
         h2: torch.Tensor,  # nf x nloc x nnei x 3
         angle_ebd: torch.Tensor,  # nf x nloc x a_nnei x a_nnei x a_dim

deepmd/pt/model/descriptor/repflows.py

@@ -377,9 +377,9 @@ def forward(
         mapping: Optional[torch.Tensor] = None,
         comm_dict: Optional[dict[str, torch.Tensor]] = None,
     ):
-        if comm_dict is None:
+        parrallel_mode = comm_dict is not None
+        if not parrallel_mode:
             assert mapping is not None
-            assert extended_atype_embd is not None
         nframes, nloc, nnei = nlist.shape
         nall = extended_coord.view(nframes, -1).shape[1] // 3
         atype = extended_atype[:, :nloc]
@@ -459,18 +459,23 @@ def forward(
         # if the a neighbor is real or not is indicated by nlist_mask
         nlist[nlist == -1] = 0
         # nb x nall x n_dim
-        if comm_dict is None:
+        if not parrallel_mode:
             assert mapping is not None
             mapping = (
                 mapping.view(nframes, nall).unsqueeze(-1).expand(-1, -1, self.n_dim)
             )
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
-            # node_ebd_ext: nb x nall x n_dim
-            if comm_dict is None:
+            # node_ebd_ext: nb x nall x n_dim [OR] nb x nloc x n_dim when not parrallel_mode
+            if not parrallel_mode:
                 assert mapping is not None
-                node_ebd_ext = torch.gather(node_ebd, 1, mapping)
+                node_ebd_ext = (
+                    torch.gather(node_ebd, 1, mapping)
+                    if not self.use_loc_mapping
+                    else node_ebd
+                )
             else:
+                assert comm_dict is not None
                 has_spin = "has_spin" in comm_dict
                 if not has_spin:
                     n_padding = nall - nloc