chore: optimize compute_smooth_weight

deepmd/dpmodel/utils/env_mat.py

@@ -15,7 +15,7 @@
 )
 
 
-@support_array_api(version="2022.12")
+@support_array_api(version="2023.12")
 def compute_smooth_weight(
     distance: np.ndarray,
     rmin: float,
@@ -25,14 +25,11 @@ def compute_smooth_weight(
     if rmin >= rmax:
         raise ValueError("rmin should be less than rmax.")
     xp = array_api_compat.array_namespace(distance)
-    min_mask = distance <= rmin
-    max_mask = distance >= rmax
-    mid_mask = xp.logical_not(xp.logical_or(min_mask, max_mask))
+    distance = xp.clip(distance, min=rmin, max=rmax)
     uu = (distance - rmin) / (rmax - rmin)
-    vv = uu * uu * uu * (-6.0 * uu * uu + 15.0 * uu - 10.0) + 1.0
-    return vv * xp.astype(mid_mask, distance.dtype) + xp.astype(
-        min_mask, distance.dtype
-    )
+    uu2 = uu * uu
+    vv = uu2 * uu * (-6.0 * uu2 + 15.0 * uu - 10.0) + 1.0
+    return vv
 
 
 def _make_env_mat(

deepmd/pt/utils/preprocess.py

@@ -10,9 +10,8 @@ def compute_smooth_weight(distance, rmin: float, rmax: float):
     """Compute smooth weight for descriptor elements."""
     if rmin >= rmax:
         raise ValueError("rmin should be less than rmax.")
-    min_mask = distance <= rmin
-    max_mask = distance >= rmax
-    mid_mask = torch.logical_not(torch.logical_or(min_mask, max_mask))
+    distance = torch.clamp(distance, min=rmin, max=rmax)
     uu = (distance - rmin) / (rmax - rmin)
-    vv = uu * uu * uu * (-6 * uu * uu + 15 * uu - 10) + 1
-    return vv * mid_mask + min_mask
+    uu2 = uu * uu
+    vv = uu2 * uu * (-6 * uu2 + 15 * uu - 10) + 1
+    return vv

source/tests/array_api_strict/__init__.py

@@ -1,2 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Synchronize with deepmd.jax for test purpose only."""
+
+import array_api_strict
+
+# this is the default version in the latest array_api_strict,
+# but in old versions it may be 2022.12
+array_api_strict.set_array_api_strict_flags(api_version="2023.12")