pd: update loc_mapping for dpa3 in paddle backend (#4797)

adapt preparing input spec in paddle backend to support `loc_mapping`
for dpa3

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

- **New Features**
- Added support for dynamic neighbor and angle selection, enabling
flexible graph-based computations.
- Introduced options for local atom index mapping and alternative edge
feature initialization.
  - Added exponential switch function for neighbor smoothing.
- Introduced JIT-compiled SiLUT activation with higher-order gradient
support, configurable via environment variable.
  - New utility functions for aggregation and graph index computation.

- **Improvements**
  - Enhanced documentation for descriptor parameters.
  - Defaulted local mapping to enabled for relevant descriptors.
- Improved device handling and tensor operations for better
compatibility and clarity.

- **Bug Fixes**
  - Simplified test skipping logic to rely on a unified flag.

- **Chores**
- Refactored and clarified internal logic for smoother weight and switch
function computations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: HydrogenSulfate

deepmd/pd/model/descriptor/dpa3.py

@@ -91,7 +91,7 @@ class DescrptDPA3(BaseDescriptor, paddle.nn.Layer):
         Whether to use bias in the type embedding layer.
     use_loc_mapping : bool, Optional
         Whether to use local atom index mapping in training or non-parallel inference.
-        Not supported yet in Paddle.
+        When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.
     type_map : list[str], Optional
         A list of strings. Give the name to each type of atoms.
 
@@ -117,7 +117,7 @@ def __init__(
         seed: Optional[Union[int, list[int]]] = None,
         use_econf_tebd: bool = False,
         use_tebd_bias: bool = False,
-        use_loc_mapping: bool = False,
+        use_loc_mapping: bool = True,
         type_map: Optional[list[str]] = None,
     ) -> None:
         super().__init__()
@@ -160,6 +160,8 @@ 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,
+            edge_init_use_dist=self.repflow_args.edge_init_use_dist,
+            use_exp_switch=self.repflow_args.use_exp_switch,
             use_dynamic_sel=self.repflow_args.use_dynamic_sel,
             sel_reduce_factor=self.repflow_args.sel_reduce_factor,
             use_loc_mapping=use_loc_mapping,
@@ -170,8 +172,8 @@ def init_subclass_params(sub_data, sub_class):
         )
 
         self.use_econf_tebd = use_econf_tebd
-        self.use_tebd_bias = use_tebd_bias
         self.use_loc_mapping = use_loc_mapping
+        self.use_tebd_bias = use_tebd_bias
         self.type_map = type_map
         self.tebd_dim = self.repflow_args.n_dim
         self.type_embedding = TypeEmbedNet(
@@ -487,12 +489,16 @@ def forward(
             The smooth switch function. shape: nf x nloc x nnei
 
         """
+        parallel_mode = comm_dict is not None
         # cast the input to internal precsion
         extended_coord = extended_coord.to(dtype=self.prec)
         nframes, nloc, nnei = nlist.shape
         nall = extended_coord.reshape([nframes, -1]).shape[1] // 3
 
-        node_ebd_ext = self.type_embedding(extended_atype)
+        if not parallel_mode and self.use_loc_mapping:
+            node_ebd_ext = self.type_embedding(extended_atype[:, :nloc])
+        else:
+            node_ebd_ext = self.type_embedding(extended_atype)
         node_ebd_inp = node_ebd_ext[:, :nloc, :]
         # repflows
         node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(

deepmd/pd/model/descriptor/env_mat.py

@@ -3,6 +3,7 @@
 import paddle
 
 from deepmd.pd.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
@@ -24,15 +26,20 @@ def _make_env_mat(
     nlist = paddle.where(mask, nlist, nall - 1)
     coord_l = coord[:, :natoms].reshape([bsz, -1, 1, 3])
     index = nlist.reshape([bsz, -1]).unsqueeze(-1).expand([-1, -1, 3])
-    coord_r = paddle.take_along_axis(coord, axis=1, indices=index)
+    coord_pad = paddle.concat([coord, coord[:, -1:, :] + rcut], axis=1)
+    coord_r = paddle.take_along_axis(coord_pad, axis=1, indices=index)
     coord_r = coord_r.reshape([bsz, natoms, nnei, 3])
     diff = coord_r - coord_l
     length = paddle.linalg.norm(diff, axis=-1, keepdim=True)
     # for index 0 nloc atom
     length = length + (~mask.unsqueeze(-1)).astype(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 * mask.unsqueeze(-1).astype(weight.dtype)
     if radial_only:
         env_mat = t0 * weight
@@ -51,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.
 
@@ -63,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
     -------
@@ -75,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]