feat: descriptor provide get_rcut_smth and get_env_protection (#3761)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Introduced methods to enhance the calculation of smooth decay radius
and environment matrix protection across various descriptor modules.
This will improve the precision and reliability of environmental
interactions in models.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Han Wang <[email protected]>
Co-authored-by: Han Wang <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

deepmd/dpmodel/descriptor/dpa1.py

@@ -302,6 +302,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()
@@ -336,6 +340,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.
@@ -635,6 +643,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)
@@ -687,6 +699,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/dpmodel/descriptor/dpa2.py

@@ -369,6 +369,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)
@@ -404,6 +408,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
+
     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,12 @@ def get_rcut(self) -> float:
         """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...
+        # Note: Using the minimum rcut_smth might not be appropriate in all scenarios. Consider using a different approach or provide detailed documentation on why the minimum value is chosen.
+        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 +134,16 @@ 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 same environment matrix protection for all descriptors. Found differing values."
+            )
+        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 @@ def get_rcut(self) -> float:
             """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 @@ def mixed_types(self) -> bool:
             """
             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

@@ -237,6 +237,10 @@ def get_rcut(self):
         """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
@@ -247,6 +251,10 @@ def mixed_types(self):
         """
         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

@@ -193,6 +193,10 @@ def get_rcut(self):
         """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
@@ -203,6 +207,10 @@ def mixed_types(self):
         """
         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 @@ def get_rcut(self) -> float:
         """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 @@ def get_dim_emb(self) -> int:
         """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

@@ -282,6 +282,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()
@@ -316,6 +320,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

@@ -363,6 +363,7 @@ def __init__(
 
         self.tebd_dim = repinit_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
@@ -373,6 +374,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)
@@ -408,6 +413,11 @@ def mixed_types(self) -> bool:
         """
         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,12 @@ 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...
+        # Note: Using the minimum rcut_smth might not be appropriate in all scenarios. Consider using a different approach or provide detailed documentation on why the minimum value is chosen.
+        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 +139,16 @@ 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 same environment matrix protection for all descriptors. Found differing values."
+            )
+        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

@@ -306,6 +306,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)
@@ -358,6 +362,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.