add distinguish types in base descriptor. fix docstrings

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -36,15 +36,15 @@ def get_rcut(self) -> float:
 
         @abstractmethod
         def get_sel(self) -> List[int]:
-            """Returns the number of selected atoms for each type."""
+            """Returns the number of selected neighboring atoms for each type."""
             pass
 
         def get_nsel(self) -> int:
-            """Returns the total number of selected atoms in the cut-off radius."""
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
             return sum(self.get_sel())
 
         def get_nnei(self) -> int:
-            """Returns the total number of neighboring atoms in the cut-off radius."""
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
             return self.get_nsel()
 
         @abstractmethod
@@ -62,6 +62,13 @@ def get_dim_emb(self) -> int:
             """Returns the embedding dimension of g2."""
             pass
 
+        @abstractmethod
+        def distinguish_types(self) -> bool:
+            """Returns if the descriptor requires a neighbor list that distinguish different
+            atomic types or not.
+            """
+            pass
+
         @abstractmethod
         def compute_input_stats(self, merged):
             """Update mean and stddev for descriptor elements."""

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -215,8 +215,8 @@ def get_sel(self):
         return self.sel
 
     def distinguish_types(self):
-        """Returns if the descriptor uses different nets for
-        different atomic types.
+        """Returns if the descriptor requires a neighbor list that distinguish different
+        atomic types or not.
         """
         return True
 

deepmd/dpmodel/model/dp_atomic_model.py

@@ -63,7 +63,9 @@ def get_sel(self) -> List[int]:
         return self.descriptor.get_sel()
 
     def distinguish_types(self) -> bool:
-        """If distinguish different types by sorting."""
+        """Returns if model requires a neighbor list that distinguish different
+        atomic types or not.
+        """
         return self.descriptor.distinguish_types()
 
     def forward_atomic(

deepmd/dpmodel/model/make_base_atomic_model.py

@@ -21,18 +21,32 @@ class BAM(ABC):
 
         @abstractmethod
         def fitting_output_def(self) -> FittingOutputDef:
+            """Get the fitting output def."""
             pass
 
         @abstractmethod
         def get_rcut(self) -> float:
+            """Get the cut-off radius."""
             pass
 
         @abstractmethod
         def get_sel(self) -> List[int]:
+            """Returns the number of selected atoms for each type."""
             pass
 
+        def get_nsel(self) -> int:
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
+            return sum(self.get_sel())
+
+        def get_nnei(self) -> int:
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
+            return self.get_nsel()
+
         @abstractmethod
         def distinguish_types(self) -> bool:
+            """Returns if the model requires a neighbor list that distinguish different
+            atomic types or not.
+            """
             pass
 
         @abstractmethod

deepmd/pt/model/descriptor/dpa1.py

@@ -106,8 +106,8 @@ def get_dim_emb(self) -> int:
         return self.se_atten.dim_emb
 
     def distinguish_types(self) -> bool:
-        """Returns if the descriptor uses different nets for
-        different atomic types.
+        """Returns if the descriptor requires a neighbor list that distinguish different
+        atomic types or not.
         """
         return False
 

deepmd/pt/model/descriptor/dpa2.py

@@ -272,8 +272,8 @@ def get_dim_emb(self) -> int:
         return self.repformers.dim_emb
 
     def distinguish_types(self) -> bool:
-        """Returns if the descriptor uses different nets for
-        different atomic types.
+        """Returns if the descriptor requires a neighbor list that distinguish different
+        atomic types or not.
         """
         return False
 

deepmd/pt/model/descriptor/se_a.py

@@ -94,8 +94,8 @@ def get_dim_emb(self) -> int:
         return self.sea.get_dim_emb()
 
     def distinguish_types(self):
-        """Returns if the descriptor uses different nets for
-        different atomic types.
+        """Returns if the descriptor requires a neighbor list that distinguish different
+        atomic types or not.
         """
         return True