Add exclude_types for data stat

deepmd/pt/model/descriptor/dpa1.py

@@ -3,6 +3,7 @@
     Callable,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -62,7 +63,7 @@ def __init__(
         type_one_side: bool = True,
         precision: str = "default",
         trainable: bool = True,
-        exclude_types: Optional[List[List[int]]] = None,
+        exclude_types: List[Tuple[int, int]] = [],
         stripped_type_embedding: bool = False,
         smooth_type_embdding: bool = False,
     ):
@@ -73,8 +74,6 @@ def __init__(
             raise NotImplementedError("type_one_side is not supported.")
         if precision != "default" and precision != "float64":
             raise NotImplementedError("precison is not supported.")
-        if exclude_types is not None and exclude_types != []:
-            raise NotImplementedError("exclude_types is not supported.")
         if stripped_type_embedding:
             raise NotImplementedError("stripped_type_embedding is not supported.")
         if smooth_type_embdding:
@@ -103,6 +102,7 @@ def __init__(
             normalize=normalize,
             temperature=temperature,
             return_rot=return_rot,
+            exclude_types=exclude_types,
             env_protection=env_protection,
         )
         self.type_embedding = TypeEmbedNet(ntypes, tebd_dim)

deepmd/pt/model/descriptor/dpa2.py

@@ -3,6 +3,7 @@
     Callable,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -79,6 +80,7 @@ def __init__(
         repformer_add_type_ebd_to_seq: bool = False,
         env_protection: float = 0.0,
         trainable: bool = True,
+        exclude_types: List[Tuple[int, int]] = [],
         type: Optional[
             str
         ] = None,  # work around the bad design in get_trainer and DpLoaderSet!
@@ -176,6 +178,9 @@ def __init__(
             repformers block: concatenate the type embedding at the output.
         trainable : bool
             If the parameters in the descriptor are trainable.
+        exclude_types : List[Tuple[int, int]] = [],
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
 
         Returns
         -------
@@ -206,6 +211,7 @@ def __init__(
             tebd_input_mode="concat",
             # tebd_input_mode='dot_residual_s',
             set_davg_zero=repinit_set_davg_zero,
+            exclude_types=exclude_types,
             env_protection=env_protection,
             activation_function=repinit_activation,
         )
@@ -238,6 +244,7 @@ def __init__(
             set_davg_zero=repformer_set_davg_zero,
             smooth=True,
             add_type_ebd_to_seq=repformer_add_type_ebd_to_seq,
+            exclude_types=exclude_types,
             env_protection=env_protection,
         )
         self.type_embedding = TypeEmbedNet(ntypes, tebd_dim)

deepmd/pt/model/descriptor/repformers.py

@@ -4,6 +4,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -24,6 +25,9 @@
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSe,
 )
+from deepmd.pt.utils.exclude_mask import (
+    PairExcludeMask,
+)
 from deepmd.pt.utils.utils import (
     get_activation_fn,
 )
@@ -83,6 +87,7 @@ def __init__(
         set_davg_zero: bool = True,  # TODO
         smooth: bool = True,
         add_type_ebd_to_seq: bool = False,
+        exclude_types: List[Tuple[int, int]] = [],
         env_protection: float = 0.0,
         type: Optional[str] = None,
     ):
@@ -115,6 +120,8 @@ def __init__(
         self.act = get_activation_fn(activation_function)
         self.direct_dist = direct_dist
         self.add_type_ebd_to_seq = add_type_ebd_to_seq
+        # order matters, placed after the assignment of self.ntypes
+        self.reinit_exclude(exclude_types)
         self.env_protection = env_protection
 
         self.g2_embd = mylinear(1, self.g2_dim)
@@ -213,6 +220,13 @@ def dim_emb(self):
         """Returns the embedding dimension g2."""
         return self.get_dim_emb()
 
+    def reinit_exclude(
+        self,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
+        self.exclude_types = exclude_types
+        self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)
+
     def forward(
         self,
         nlist: torch.Tensor,

deepmd/pt/model/descriptor/se_atten.py

@@ -4,6 +4,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -26,6 +27,9 @@
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSe,
 )
+from deepmd.pt.utils.exclude_mask import (
+    PairExcludeMask,
+)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
@@ -61,6 +65,7 @@ def __init__(
         normalize=True,
         temperature=None,
         return_rot=False,
+        exclude_types: List[Tuple[int, int]] = [],
         env_protection: float = 0.0,
         type: Optional[str] = None,
     ):
@@ -108,6 +113,8 @@ def __init__(
         self.split_sel = self.sel
         self.nnei = sum(sel)
         self.ndescrpt = self.nnei * 4
+        # order matters, placed after the assignment of self.ntypes
+        self.reinit_exclude(exclude_types)
         self.dpa1_attention = NeighborWiseAttention(
             self.attn_layer,
             self.nnei,
@@ -251,6 +258,13 @@ def get_stats(self) -> Dict[str, StatItem]:
             )
         return self.stats
 
+    def reinit_exclude(
+        self,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
+        self.exclude_types = exclude_types
+        self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)
+
     def forward(
         self,
         nlist: torch.Tensor,

deepmd/pt/model/descriptor/se_r.py

@@ -82,7 +82,8 @@ def __init__(
         self.old_impl = False  # this does not support old implementation.
         self.exclude_types = exclude_types
         self.ntypes = len(sel)
-        self.emask = PairExcludeMask(len(sel), exclude_types=exclude_types)
+        # order matters, placed after the assignment of self.ntypes
+        self.reinit_exclude(exclude_types)
         self.env_protection = env_protection
 
         self.sel = sel
@@ -255,6 +256,13 @@ def __getitem__(self, key):
         else:
             raise KeyError(key)
 
+    def reinit_exclude(
+        self,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
+        self.exclude_types = exclude_types
+        self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)
+
     def forward(
         self,
         coord_ext: torch.Tensor,