From 3d1071c411954819b0cfd46027159a949a2979db Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 21 May 2025 16:01:46 +0800
Subject: [PATCH 1/9] feat(pt/dp): add exponential switch function

---
 deepmd/dpmodel/descriptor/dpa3.py             | 14 +++++++++
 deepmd/dpmodel/descriptor/repflows.py         | 23 ++++++++++++--
 deepmd/dpmodel/utils/env_mat.py               | 30 ++++++++++++++++++-
 deepmd/pt/model/descriptor/dpa3.py            |  1 +
 deepmd/pt/model/descriptor/env_mat.py         | 11 ++++++-
 deepmd/pt/model/descriptor/repflows.py        | 14 +++++++++
 deepmd/pt/utils/preprocess.py                 | 12 ++++++++
 deepmd/utils/argcheck.py                      | 19 ++++++++++++
 .../tests/consistent/descriptor/test_dpa3.py  | 14 ++++++++-
 .../dpmodel/descriptor/test_descriptor.py     |  3 ++
 10 files changed, 136 insertions(+), 5 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py
index 7d3292dff2..7bd83c48bb 100644
--- a/deepmd/dpmodel/descriptor/dpa3.py
+++ b/deepmd/dpmodel/descriptor/dpa3.py
@@ -123,6 +123,16 @@ class RepFlowArgs:
     smooth_edge_update : bool, optional
         Whether to make edge update smooth.
         If True, the edge update from angle message will not use self as padding.
+    use_exp_switch : bool, optional
+        Whether to use an exponential switch function instead of a polynomial one in the neighbor update.
+        The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance
+        `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as:
+        s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
+        Here, `rcut_smth` is an adjustable smoothing factor,
+        while `C` is a tunable parameter controlling the decay rate (default: 20).
+        `rcut_smth` should be chosen carefully according to `rcut`,
+        ensuring s(r) approaches zero smoothly at the cutoff.
+        Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0.
     """
 
     def __init__(
@@ -150,6 +160,7 @@ def __init__(
         skip_stat: bool = False,
         optim_update: bool = True,
         smooth_edge_update: bool = False,
+        use_exp_switch: bool = False,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -176,6 +187,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.use_exp_switch = use_exp_switch
 
     def __getitem__(self, key):
         if hasattr(self, key):
@@ -207,6 +219,7 @@ def serialize(self) -> dict:
             "fix_stat_std": self.fix_stat_std,
             "optim_update": self.optim_update,
             "smooth_edge_update": self.smooth_edge_update,
+            "use_exp_switch": self.use_exp_switch,
         }
 
     @classmethod
@@ -303,6 +316,7 @@ 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,
+            use_exp_switch=self.repflow_args.use_exp_switch,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py
index 7273ba8ebe..e4c0bcc16c 100644
--- a/deepmd/dpmodel/descriptor/repflows.py
+++ b/deepmd/dpmodel/descriptor/repflows.py
@@ -122,6 +122,16 @@ class DescrptBlockRepflows(NativeOP, DescriptorBlock):
     smooth_edge_update : bool, optional
         Whether to make edge update smooth.
         If True, the edge update from angle message will not use self as padding.
+    use_exp_switch : bool, optional
+        Whether to use an exponential switch function instead of a polynomial one in the neighbor update.
+        The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance
+        `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as:
+        s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
+        Here, `rcut_smth` is an adjustable smoothing factor,
+        while `C` is a tunable parameter controlling the decay rate (default: 20).
+        `rcut_smth` should be chosen carefully according to `rcut`,
+        ensuring s(r) approaches zero smoothly at the cutoff.
+        Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0.
     ntypes : int
         Number of element types
     activation_function : str, optional
@@ -170,6 +180,7 @@ def __init__(
         fix_stat_std: float = 0.3,
         optim_update: bool = True,
         smooth_edge_update: bool = False,
+        use_exp_switch: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         super().__init__()
@@ -201,6 +212,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.use_exp_switch = use_exp_switch
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -261,10 +273,16 @@ def __init__(
 
         wanted_shape = (self.ntypes, self.nnei, 4)
         self.env_mat_edge = EnvMat(
-            self.e_rcut, self.e_rcut_smth, protection=self.env_protection
+            self.e_rcut,
+            self.e_rcut_smth,
+            protection=self.env_protection,
+            use_exp_switch=self.use_exp_switch,
         )
         self.env_mat_angle = EnvMat(
-            self.a_rcut, self.a_rcut_smth, protection=self.env_protection
+            self.a_rcut,
+            self.a_rcut_smth,
+            protection=self.env_protection,
+            use_exp_switch=self.use_exp_switch,
         )
         self.mean = np.zeros(wanted_shape, dtype=PRECISION_DICT[self.precision])
         self.stddev = np.ones(wanted_shape, dtype=PRECISION_DICT[self.precision])
@@ -577,6 +595,7 @@ def serialize(self):
             "precision": self.precision,
             "fix_stat_std": self.fix_stat_std,
             "optim_update": self.optim_update,
+            "use_exp_switch": self.use_exp_switch,
             # variables
             "edge_embd": self.edge_embd.serialize(),
             "angle_embd": self.angle_embd.serialize(),
diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index 6c17ad7702..0c7f0ff9fb 100644
--- a/deepmd/dpmodel/utils/env_mat.py
+++ b/deepmd/dpmodel/utils/env_mat.py
@@ -35,6 +35,24 @@ def compute_smooth_weight(
     return vv
 
 
+@support_array_api(version="2023.12")
+def compute_exp_sw(
+    distance: np.ndarray,
+    rmin: float,
+    rmax: float,
+):
+    """Compute the exponential switch function for neighbor update."""
+    if rmin >= rmax:
+        raise ValueError("rmin should be less than rmax.")
+    xp = array_api_compat.array_namespace(distance)
+    distance = xp.clip(distance, min=rmin, max=rmax)
+    C = 20
+    a = C / rmin
+    b = rmin
+    exp_sw = xp.exp(-xp.exp(a * (distance - b)))
+    return exp_sw
+
+
 def _make_env_mat(
     nlist,
     coord,
@@ -42,6 +60,7 @@ def _make_env_mat(
     ruct_smth: float,
     radial_only: bool = False,
     protection: float = 0.0,
+    use_exp_switch: bool = False,
 ):
     """Make smooth environment matrix."""
     xp = array_api_compat.array_namespace(nlist)
@@ -66,7 +85,11 @@ def _make_env_mat(
     length = length + xp.astype(~xp.expand_dims(mask, axis=-1), length.dtype)
     t0 = 1 / (length + protection)
     t1 = diff / (length + protection) ** 2
-    weight = compute_smooth_weight(length, ruct_smth, rcut)
+    weight = (
+        compute_smooth_weight(length, ruct_smth, rcut)
+        if not use_exp_switch
+        else compute_exp_sw(length, ruct_smth, rcut)
+    )
     weight = weight * xp.astype(xp.expand_dims(mask, axis=-1), weight.dtype)
     if radial_only:
         env_mat = t0 * weight
@@ -81,10 +104,12 @@ def __init__(
         rcut,
         rcut_smth,
         protection: float = 0.0,
+        use_exp_switch: bool = False,
     ) -> None:
         self.rcut = rcut
         self.rcut_smth = rcut_smth
         self.protection = protection
+        self.use_exp_switch = use_exp_switch
 
     def call(
         self,
@@ -142,6 +167,7 @@ def _call(self, nlist, coord_ext, radial_only):
             self.rcut_smth,
             radial_only=radial_only,
             protection=self.protection,
+            use_exp_switch=self.use_exp_switch,
         )
         return em, diff, ww
 
@@ -151,6 +177,8 @@ def serialize(
         return {
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
+            "protection": self.protection,
+            "use_exp_switch": self.use_exp_switch,
         }
 
     @classmethod
diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py
index 545da962e7..b6bb47a1f3 100644
--- a/deepmd/pt/model/descriptor/dpa3.py
+++ b/deepmd/pt/model/descriptor/dpa3.py
@@ -150,6 +150,7 @@ 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,
+            use_exp_switch=self.repflow_args.use_exp_switch,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index dc7142249a..c57ae209fd 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -3,6 +3,7 @@
 import torch
 
 from deepmd.pt.utils.preprocess import (
+    compute_exp_sw,
     compute_smooth_weight,
 )
 
@@ -14,6 +15,7 @@ def _make_env_mat(
     ruct_smth: float,
     radial_only: bool = False,
     protection: float = 0.0,
+    use_exp_switch: bool = False,
 ):
     """Make smooth environment matrix."""
     bsz, natoms, nnei = nlist.shape
@@ -33,7 +35,11 @@ def _make_env_mat(
     length = length + ~mask.unsqueeze(-1)
     t0 = 1 / (length + protection)
     t1 = diff / (length + protection) ** 2
-    weight = compute_smooth_weight(length, ruct_smth, rcut)
+    weight = (
+        compute_smooth_weight(length, ruct_smth, rcut)
+        if not use_exp_switch
+        else compute_exp_sw(length, ruct_smth, rcut)
+    )
     weight = weight * mask.unsqueeze(-1)
     if radial_only:
         env_mat = t0 * weight
@@ -52,6 +58,7 @@ def prod_env_mat(
     rcut_smth: float,
     radial_only: bool = False,
     protection: float = 0.0,
+    use_exp_switch: bool = False,
 ):
     """Generate smooth environment matrix from atom coordinates and other context.
 
@@ -64,6 +71,7 @@ def prod_env_mat(
     - rcut_smth: Smooth hyper-parameter for pair force & energy.
     - radial_only: Whether to return a full description or a radial-only descriptor.
     - protection: Protection parameter to prevent division by zero errors during calculations.
+    - use_exp_switch: Whether to use the exponential switch function.
 
     Returns
     -------
@@ -76,6 +84,7 @@ def prod_env_mat(
         rcut_smth,
         radial_only,
         protection=protection,
+        use_exp_switch=use_exp_switch,
     )  # shape [n_atom, dim, 4 or 1]
     t_avg = mean[atype]  # [n_atom, dim, 4 or 1]
     t_std = stddev[atype]  # [n_atom, dim, 4 or 1]
diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py
index 330336b1de..32971865f5 100644
--- a/deepmd/pt/model/descriptor/repflows.py
+++ b/deepmd/pt/model/descriptor/repflows.py
@@ -133,6 +133,16 @@ class DescrptBlockRepflows(DescriptorBlock):
     smooth_edge_update : bool, optional
         Whether to make edge update smooth.
         If True, the edge update from angle message will not use self as padding.
+    use_exp_switch : bool, optional
+        Whether to use an exponential switch function instead of a polynomial one in the neighbor update.
+        The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance
+        `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as:
+        s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
+        Here, `rcut_smth` is an adjustable smoothing factor,
+        while `C` is a tunable parameter controlling the decay rate (default: 20).
+        `rcut_smth` should be chosen carefully according to `rcut`,
+        ensuring s(r) approaches zero smoothly at the cutoff.
+        Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0.
     optim_update : bool, optional
         Whether to enable the optimized update method.
         Uses a more efficient process when enabled. Defaults to True
@@ -183,6 +193,7 @@ def __init__(
         precision: str = "float64",
         fix_stat_std: float = 0.3,
         smooth_edge_update: bool = False,
+        use_exp_switch: bool = False,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -215,6 +226,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.use_exp_switch = use_exp_switch
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -396,6 +408,7 @@ def forward(
             self.e_rcut,
             self.e_rcut_smth,
             protection=self.env_protection,
+            use_exp_switch=self.use_exp_switch,
         )
         nlist_mask = nlist != -1
         sw = torch.squeeze(sw, -1)
@@ -432,6 +445,7 @@ def forward(
             self.a_rcut,
             self.a_rcut_smth,
             protection=self.env_protection,
+            use_exp_switch=self.use_exp_switch,
         )
         a_nlist_mask = a_nlist != -1
         a_sw = torch.squeeze(a_sw, -1)
diff --git a/deepmd/pt/utils/preprocess.py b/deepmd/pt/utils/preprocess.py
index 8ab489dede..7161bac692 100644
--- a/deepmd/pt/utils/preprocess.py
+++ b/deepmd/pt/utils/preprocess.py
@@ -15,3 +15,15 @@ def compute_smooth_weight(distance, rmin: float, rmax: float):
     uu2 = uu * uu
     vv = uu2 * uu * (-6 * uu2 + 15 * uu - 10) + 1
     return vv
+
+
+def compute_exp_sw(distance, rmin: float, rmax: float):
+    """Compute the exponential switch function for neighbor update."""
+    if rmin >= rmax:
+        raise ValueError("rmin should be less than rmax.")
+    distance = torch.clamp(distance, min=0.0, max=rmax)
+    C = 20
+    a = C / rmin
+    b = rmin
+    exp_sw = torch.exp(-torch.exp(a * (distance - b)))
+    return exp_sw
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 0260700165..86fbe9fa99 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -1497,6 +1497,17 @@ def dpa3_repflow_args():
         "Whether to make edge update smooth. "
         "If True, the edge update from angle message will not use self as padding."
     )
+    doc_use_exp_switch = (
+        "Whether to use an exponential switch function instead of a polynomial one in the neighbor update. "
+        "The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance "
+        "`r` approaches the cutoff radius `rcut`. Specifically, the function is defined as: "
+        "s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut. "
+        "Here, `rcut_smth` is an adjustable smoothing factor, "
+        "while `C` is a tunable parameter controlling the decay rate (default: 20). "
+        "`rcut_smth` should be chosen carefully according to `rcut`, "
+        "ensuring s(r) approaches zero smoothly at the cutoff. "
+        "Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0."
+    )
 
     return [
         # repflow args
@@ -1597,6 +1608,14 @@ def dpa3_repflow_args():
             default=False,  # For compatability. This will be True in the future
             doc=doc_smooth_edge_update,
         ),
+        Argument(
+            "use_exp_switch",
+            bool,
+            optional=True,
+            default=False,
+            alias=["use_env_envelope"],
+            doc=doc_use_exp_switch,
+        ),
     ]
 
 
diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py
index 56df87f931..9649cb304a 100644
--- a/source/tests/consistent/descriptor/test_dpa3.py
+++ b/source/tests/consistent/descriptor/test_dpa3.py
@@ -65,6 +65,7 @@
     (1, 2),  # a_compress_e_rate
     (True, False),  # a_compress_use_split
     (True, False),  # optim_update
+    (True, False),  # use_exp_switch
     (0.3, 0.0),  # fix_stat_std
     (1, 2),  # n_multi_edge_message
     ("float64",),  # precision
@@ -80,6 +81,7 @@ def data(self) -> dict:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -103,6 +105,7 @@ def data(self) -> dict:
                     "a_compress_e_rate": a_compress_e_rate,
                     "a_compress_use_split": a_compress_use_split,
                     "optim_update": optim_update,
+                    "use_exp_switch": use_exp_switch,
                     "fix_stat_std": fix_stat_std,
                     "n_multi_edge_message": n_multi_edge_message,
                     "axis_neuron": 4,
@@ -130,6 +133,7 @@ def skip_pt(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -146,11 +150,14 @@ def skip_pd(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
         ) = self.param
-        return not INSTALLED_PD or precision == "bfloat16"
+        return (
+            not INSTALLED_PD or precision == "bfloat16" or use_exp_switch
+        )  # not supported yet
 
     @property
     def skip_dp(self) -> bool:
@@ -162,6 +169,7 @@ def skip_dp(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -178,6 +186,7 @@ def skip_tf(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -236,6 +245,7 @@ def setUp(self) -> None:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -315,6 +325,7 @@ def rtol(self) -> float:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -337,6 +348,7 @@ def atol(self) -> float:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            use_exp_switch,
             fix_stat_std,
             n_multi_edge_message,
             precision,
diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
index 4fa0593419..b958017bc8 100644
--- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py
+++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
@@ -482,6 +482,7 @@ def DescriptorParamDPA3(
     a_compress_use_split=False,
     optim_update=True,
     smooth_edge_update=False,
+    use_exp_switch=False,
     fix_stat_std=0.3,
     precision="float64",
 ):
@@ -503,6 +504,7 @@ def DescriptorParamDPA3(
                 "a_compress_e_rate": a_compress_e_rate,
                 "a_compress_use_split": a_compress_use_split,
                 "optim_update": optim_update,
+                "use_exp_switch": use_exp_switch,
                 "smooth_edge_update": smooth_edge_update,
                 "fix_stat_std": fix_stat_std,
                 "n_multi_edge_message": n_multi_edge_message,
@@ -540,6 +542,7 @@ def DescriptorParamDPA3(
             "a_compress_use_split": (True, False),
             "optim_update": (True, False),
             "smooth_edge_update": (True,),
+            "use_exp_switch": (True, False),
             "fix_stat_std": (0.3,),
             "n_multi_edge_message": (1, 2),
             "env_protection": (0.0, 1e-8),

From 44a169e65145ee741b9966808dbb4bcf9b66dfda Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 21 May 2025 21:13:41 +0800
Subject: [PATCH 2/9] fix docstr

---
 deepmd/dpmodel/descriptor/dpa3.py      | 8 +++-----
 deepmd/dpmodel/descriptor/repflows.py  | 8 +++-----
 deepmd/pt/model/descriptor/repflows.py | 8 +++-----
 deepmd/utils/argcheck.py               | 6 ++----
 4 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py
index 7bd83c48bb..df434eb88b 100644
--- a/deepmd/dpmodel/descriptor/dpa3.py
+++ b/deepmd/dpmodel/descriptor/dpa3.py
@@ -127,11 +127,9 @@ class RepFlowArgs:
         Whether to use an exponential switch function instead of a polynomial one in the neighbor update.
         The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance
         `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as:
-        s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
-        Here, `rcut_smth` is an adjustable smoothing factor,
-        while `C` is a tunable parameter controlling the decay rate (default: 20).
-        `rcut_smth` should be chosen carefully according to `rcut`,
-        ensuring s(r) approaches zero smoothly at the cutoff.
+        s(r) = \\exp(-\\exp(20 * (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
+        Here, `rcut_smth` is an adjustable smoothing factor and `rcut_smth` should be chosen carefully
+        according to `rcut`, ensuring s(r) approaches zero smoothly at the cutoff.
         Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0.
     """
 
diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py
index e4c0bcc16c..3a5a8a6ebb 100644
--- a/deepmd/dpmodel/descriptor/repflows.py
+++ b/deepmd/dpmodel/descriptor/repflows.py
@@ -126,11 +126,9 @@ class DescrptBlockRepflows(NativeOP, DescriptorBlock):
         Whether to use an exponential switch function instead of a polynomial one in the neighbor update.
         The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance
         `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as:
-        s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
-        Here, `rcut_smth` is an adjustable smoothing factor,
-        while `C` is a tunable parameter controlling the decay rate (default: 20).
-        `rcut_smth` should be chosen carefully according to `rcut`,
-        ensuring s(r) approaches zero smoothly at the cutoff.
+        s(r) = \\exp(-\\exp(20 * (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
+        Here, `rcut_smth` is an adjustable smoothing factor and `rcut_smth` should be chosen carefully
+        according to `rcut`, ensuring s(r) approaches zero smoothly at the cutoff.
         Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0.
     ntypes : int
         Number of element types
diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py
index 32971865f5..324f1f4324 100644
--- a/deepmd/pt/model/descriptor/repflows.py
+++ b/deepmd/pt/model/descriptor/repflows.py
@@ -137,11 +137,9 @@ class DescrptBlockRepflows(DescriptorBlock):
         Whether to use an exponential switch function instead of a polynomial one in the neighbor update.
         The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance
         `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as:
-        s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
-        Here, `rcut_smth` is an adjustable smoothing factor,
-        while `C` is a tunable parameter controlling the decay rate (default: 20).
-        `rcut_smth` should be chosen carefully according to `rcut`,
-        ensuring s(r) approaches zero smoothly at the cutoff.
+        s(r) = \\exp(-\\exp(20 * (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut.
+        Here, `rcut_smth` is an adjustable smoothing factor and `rcut_smth` should be chosen carefully
+        according to `rcut`, ensuring s(r) approaches zero smoothly at the cutoff.
         Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0.
     optim_update : bool, optional
         Whether to enable the optimized update method.
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 86fbe9fa99..033bb3d0f9 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -1501,10 +1501,8 @@ def dpa3_repflow_args():
         "Whether to use an exponential switch function instead of a polynomial one in the neighbor update. "
         "The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance "
         "`r` approaches the cutoff radius `rcut`. Specifically, the function is defined as: "
-        "s(r) = \\exp(-\\exp(C (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut. "
-        "Here, `rcut_smth` is an adjustable smoothing factor, "
-        "while `C` is a tunable parameter controlling the decay rate (default: 20). "
-        "`rcut_smth` should be chosen carefully according to `rcut`, "
+        "s(r) = \\exp(-\\exp(20 * (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut. "
+        "Here, `rcut_smth` is an adjustable smoothing factor and should be chosen carefully according to `rcut`, "
         "ensuring s(r) approaches zero smoothly at the cutoff. "
         "Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0."
     )

From 364599d81bc26e1641ef77b204cbec792175b4d7 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 21 May 2025 21:40:03 +0800
Subject: [PATCH 3/9] Update env_mat.py

---
 deepmd/dpmodel/utils/env_mat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index 0c7f0ff9fb..ee11678d3a 100644
--- a/deepmd/dpmodel/utils/env_mat.py
+++ b/deepmd/dpmodel/utils/env_mat.py
@@ -45,7 +45,7 @@ def compute_exp_sw(
     if rmin >= rmax:
         raise ValueError("rmin should be less than rmax.")
     xp = array_api_compat.array_namespace(distance)
-    distance = xp.clip(distance, min=rmin, max=rmax)
+    distance = xp.clip(distance, min=0.0, max=rmax)
     C = 20
     a = C / rmin
     b = rmin

From 31d3815d0f62279e2a860566b60adcc42474ef0a Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 21 May 2025 21:42:34 +0800
Subject: [PATCH 4/9] Update env_mat.py

---
 deepmd/dpmodel/utils/env_mat.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index ee11678d3a..e77039ef5c 100644
--- a/deepmd/dpmodel/utils/env_mat.py
+++ b/deepmd/dpmodel/utils/env_mat.py
@@ -177,8 +177,6 @@ def serialize(
         return {
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
-            "protection": self.protection,
-            "use_exp_switch": self.use_exp_switch,
         }
 
     @classmethod

From 3a520b5ad8fd571b7d36240308bd5abf532967d3 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 21 May 2025 21:58:32 +0800
Subject: [PATCH 5/9] Update test_descriptor.py

---
 .../tests/universal/dpmodel/descriptor/test_descriptor.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
index b958017bc8..70f01bcc78 100644
--- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py
+++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
@@ -495,10 +495,14 @@ def DescriptorParamDPA3(
                 "a_dim": 8,
                 "nlayers": 2,
                 "e_rcut": rcut,
-                "e_rcut_smth": rcut_smth,
+                "e_rcut_smth": rcut_smth
+                if not use_exp_switch
+                else (rcut - 1.0),  # suitable for ut
                 "e_sel": sum(sel),
                 "a_rcut": rcut / 2,
-                "a_rcut_smth": rcut_smth / 2,
+                "a_rcut_smth": rcut_smth / 2
+                if not use_exp_switch
+                else (rcut - 1.0) / 2,  # suitable for ut
                 "a_sel": sum(sel) // 4,
                 "a_compress_rate": a_compress_rate,
                 "a_compress_e_rate": a_compress_e_rate,

From fcdd130b4ea3a8b02702b2db761860c2db2e1285 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Date: Wed, 21 May 2025 23:10:16 +0800
Subject: [PATCH 6/9] serialize protection and use_exp_switch for EnvMat

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
---
 deepmd/dpmodel/utils/env_mat.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index e77039ef5c..ee11678d3a 100644
--- a/deepmd/dpmodel/utils/env_mat.py
+++ b/deepmd/dpmodel/utils/env_mat.py
@@ -177,6 +177,8 @@ def serialize(
         return {
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
+            "protection": self.protection,
+            "use_exp_switch": self.use_exp_switch,
         }
 
     @classmethod

From bb717e00cdccefbde32023decde48c6ba3afc724 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Thu, 22 May 2025 15:26:32 +0800
Subject: [PATCH 7/9] fix env_mat serialize for ut

---
 deepmd/pd/model/descriptor/dpa1.py | 4 +++-
 deepmd/pd/model/descriptor/se_a.py | 4 +++-
 deepmd/pt/model/descriptor/dpa1.py | 4 +++-
 deepmd/pt/model/descriptor/se_a.py | 4 +++-
 deepmd/pt/model/descriptor/se_t.py | 4 +++-
 5 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py
index 6af494d5b4..6942b096c9 100644
--- a/deepmd/pd/model/descriptor/dpa1.py
+++ b/deepmd/pd/model/descriptor/dpa1.py
@@ -506,7 +506,9 @@ def serialize(self) -> dict:
             "precision": RESERVED_PRECISION_DICT[obj.prec],
             "embeddings": obj.filter_layers.serialize(),
             "attention_layers": obj.dpa1_attention.serialize(),
-            "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(),
+            "env_mat": DPEnvMat(
+                obj.rcut, obj.rcut_smth, obj.env_protection
+            ).serialize(),
             "type_embedding": self.type_embedding.embedding.serialize(),
             "exclude_types": obj.exclude_types,
             "env_protection": obj.env_protection,
diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py
index 8d84620caa..7b70a742ce 100644
--- a/deepmd/pd/model/descriptor/se_a.py
+++ b/deepmd/pd/model/descriptor/se_a.py
@@ -337,7 +337,9 @@ def serialize(self) -> dict:
             # make deterministic
             "precision": RESERVED_PRECISION_DICT[obj.prec],
             "embeddings": obj.filter_layers.serialize(),
-            "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(),
+            "env_mat": DPEnvMat(
+                obj.rcut, obj.rcut_smth, obj.env_protection
+            ).serialize(),
             "exclude_types": obj.exclude_types,
             "env_protection": obj.env_protection,
             "@variables": {
diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py
index fae50a138c..9c1e144f48 100644
--- a/deepmd/pt/model/descriptor/dpa1.py
+++ b/deepmd/pt/model/descriptor/dpa1.py
@@ -508,7 +508,9 @@ def serialize(self) -> dict:
             "precision": RESERVED_PRECISION_DICT[obj.prec],
             "embeddings": obj.filter_layers.serialize(),
             "attention_layers": obj.dpa1_attention.serialize(),
-            "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(),
+            "env_mat": DPEnvMat(
+                obj.rcut, obj.rcut_smth, obj.env_protection
+            ).serialize(),
             "type_embedding": self.type_embedding.embedding.serialize(),
             "exclude_types": obj.exclude_types,
             "env_protection": obj.env_protection,
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index 9d3b850d0f..fc3e14bd25 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -381,7 +381,9 @@ def serialize(self) -> dict:
             # make deterministic
             "precision": RESERVED_PRECISION_DICT[obj.prec],
             "embeddings": obj.filter_layers.serialize(),
-            "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(),
+            "env_mat": DPEnvMat(
+                obj.rcut, obj.rcut_smth, obj.env_protection
+            ).serialize(),
             "exclude_types": obj.exclude_types,
             "env_protection": obj.env_protection,
             "@variables": {
diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py
index fb444c3d13..6e075a04e4 100644
--- a/deepmd/pt/model/descriptor/se_t.py
+++ b/deepmd/pt/model/descriptor/se_t.py
@@ -415,7 +415,9 @@ def serialize(self) -> dict:
             "activation_function": obj.activation_function,
             "precision": RESERVED_PRECISION_DICT[obj.prec],
             "embeddings": obj.filter_layers.serialize(),
-            "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(),
+            "env_mat": DPEnvMat(
+                obj.rcut, obj.rcut_smth, obj.env_protection
+            ).serialize(),
             "exclude_types": obj.exclude_types,
             "env_protection": obj.env_protection,
             "type_map": self.type_map,

From 11a8e48708288e8ba6d6f45fd958ed17c7346193 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 May 2025 05:35:56 +0000
Subject: [PATCH 8/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/consistent/descriptor/test_dpa3.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py
index 94e81a3d75..2647da52b3 100644
--- a/source/tests/consistent/descriptor/test_dpa3.py
+++ b/source/tests/consistent/descriptor/test_dpa3.py
@@ -162,7 +162,10 @@ def skip_pd(self) -> bool:
             precision,
         ) = self.param
         return (
-            not INSTALLED_PD or precision == "bfloat16" or use_exp_switch or use_dynamic_sel
+            not INSTALLED_PD
+            or precision == "bfloat16"
+            or use_exp_switch
+            or use_dynamic_sel
         )  # not supported yet
 
     @property

From a574ac337bd3f2ac7a2355fea0d504138f9ea2c8 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Mon, 26 May 2025 13:40:09 +0800
Subject: [PATCH 9/9] Update argcheck.py

---
 deepmd/utils/argcheck.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index ce03e69621..6e9663592f 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -1505,6 +1505,7 @@ def dpa3_repflow_args():
         "Here, `rcut_smth` is an adjustable smoothing factor and should be chosen carefully according to `rcut`, "
         "ensuring s(r) approaches zero smoothly at the cutoff. "
         "Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0."
+    )
     doc_use_dynamic_sel = (
         "Whether to dynamically select neighbors within the cutoff radius. "
         "If True, the exact number of neighbors within the cutoff radius is used "