feat: descriptor provide get_rcut_smth and get_env_protection

deepmd/dpmodel/descriptor/dpa1.py

@@ -386,6 +386,10 @@
         """Returns cutoff radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_sel(self):
         """Returns cutoff radius."""
         return self.sel
@@ -402,6 +406,10 @@
         """
         return True
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/dpmodel/descriptor/hybrid.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import math
 from typing import (
     Any,
     Dict,
@@ -96,6 +97,11 @@
         """Returns the cut-off radius."""
         return np.max([descrpt.get_rcut() for descrpt in self.descrpt_list]).item()
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        # may not be a good idea...
+        return np.min([descrpt.get_rcut_smth() for descrpt in self.descrpt_list]).item()
+
     def get_sel(self) -> List[int]:
         """Returns the number of selected atoms for each type."""
         if self.mixed_types():
@@ -127,6 +133,17 @@
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix. All descriptors should be the same."""
+        all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]
+        same_as_0 = [math.isclose(ii, all_protection[0]) for ii in all_protection]
+        if not all(same_as_0):
+            raise ValueError(
+                "Hybrid descriptor requires the environment matrix protection being the same ",
+                "for all the descriptors.",
+            )
+        return all_protection[0]
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -51,6 +51,11 @@
             """Returns the cut-off radius."""
             pass
 
+        @abstractmethod
+        def get_rcut_smth(self) -> float:
+            """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+            pass
+
         @abstractmethod
         def get_sel(self) -> List[int]:
             """Returns the number of selected neighboring atoms for each type."""
@@ -86,6 +91,11 @@
             """
             pass
 
+        @abstractmethod
+        def get_env_protection(self) -> float:
+            """Returns the protection of building environment matrix."""
+            pass
+
         @abstractmethod
         def share_params(self, base_class, shared_level, resume=False):
             """

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -239,6 +239,10 @@
         """Returns cutoff radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_sel(self):
         """Returns cutoff radius."""
         return self.sel
@@ -249,6 +253,10 @@
         """
         return False
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/dpmodel/descriptor/se_r.py

@@ -195,6 +195,10 @@
         """Returns cutoff radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_sel(self):
         """Returns cutoff radius."""
         return self.sel
@@ -205,6 +209,10 @@
         """
         return False
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/pt/model/descriptor/descriptor.py

@@ -58,6 +58,11 @@
         """Returns the cut-off radius."""
         pass
 
+    @abstractmethod
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        pass
+
     @abstractmethod
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
@@ -88,6 +93,11 @@
         """Returns the embedding dimension."""
         pass
 
+    @abstractmethod
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        pass
+
     def compute_input_stats(
         self,
         merged: Union[Callable[[], List[dict]], List[dict]],

deepmd/pt/model/descriptor/dpa1.py

@@ -281,6 +281,10 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.se_atten.get_rcut()
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.se_atten.get_rcut_smth()
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return self.se_atten.get_nsel()
@@ -315,6 +319,10 @@ def mixed_types(self) -> bool:
         """
         return self.se_atten.mixed_types()
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.se_atten.get_env_protection()
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/pt/model/descriptor/dpa2.py

@@ -263,6 +263,7 @@
         self.concat_output_tebd = concat_output_tebd
         self.tebd_dim = tebd_dim
         self.rcut = self.repinit.get_rcut()
+        self.rcut_smth = self.repinit.get_rcut_smth()
         self.ntypes = ntypes
         self.sel = self.repinit.sel
         # set trainable
@@ -273,6 +274,10 @@
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)
@@ -308,6 +313,11 @@
         """
         return True
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        # the env_protection of repinit is the same as that of the repformer
+        return self.repinit.get_env_protection()
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/pt/model/descriptor/hybrid.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import math
 from typing import (
     Any,
     Dict,
@@ -101,6 +102,11 @@ def get_rcut(self) -> float:
         # do not use numpy here - jit is not happy
         return max([descrpt.get_rcut() for descrpt in self.descrpt_list])
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        # may not be a good idea...
+        return min([descrpt.get_rcut_smth() for descrpt in self.descrpt_list])
+
     def get_sel(self) -> List[int]:
         """Returns the number of selected atoms for each type."""
         if self.mixed_types():
@@ -132,6 +138,17 @@ def mixed_types(self):
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix. All descriptors should be the same."""
+        all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]
+        same_as_0 = [math.isclose(ii, all_protection[0]) for ii in all_protection]
+        if not all(same_as_0):
+            raise ValueError(
+                "Hybrid descriptor requires the environment matrix protection being the same ",
+                "for all the descriptors.",
+            )
+        return all_protection[0]
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/pt/model/descriptor/repformers.py

@@ -190,6 +190,10 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)
@@ -226,6 +230,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""

deepmd/pt/model/descriptor/se_a.py

@@ -106,6 +106,10 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.sea.get_rcut()
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.sea.get_rcut_smth()
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return self.sea.get_nsel()
@@ -132,6 +136,10 @@ def mixed_types(self):
         """
         return self.sea.mixed_types()
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.sea.get_env_protection()
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
@@ -400,6 +408,10 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)
@@ -436,6 +448,10 @@ def mixed_types(self) -> bool:
         """
         return False
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""

deepmd/pt/model/descriptor/se_atten.py

@@ -286,6 +286,10 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)
@@ -338,6 +342,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""

deepmd/pt/model/descriptor/se_r.py

@@ -124,6 +124,10 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)
@@ -160,6 +164,10 @@ def mixed_types(self) -> bool:
         """
         return False
 
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.

deepmd/pt/utils/env_mat_stat.py

@@ -142,9 +142,9 @@ def iter(
                 one_stddev,
                 self.descriptor.get_rcut(),
                 # TODO: export rcut_smth from DescriptorBlock
-                self.descriptor.rcut_smth,
+                self.descriptor.get_rcut_smth(),
                 radial_only,
-                protection=self.descriptor.env_protection,
+                protection=self.descriptor.get_env_protection(),
             )
             # apply excluded_types
             exclude_mask = self.descriptor.emask(nlist, extended_atype)