feat(pt/dp): add dynamic sel for DPA3 (#4754)

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

- **New Features**
- Introduced dynamic neighbor selection for the DPA3 descriptor,
enabling handling of variable neighbor counts without fixed padding.
- Added configuration options to enable dynamic selection
(`use_dynamic_sel`) and adjust normalization scaling
(`sel_reduce_factor`).
- Added new utility functions for aggregation and graph index
computation to support dynamic neighbor processing.

- **Bug Fixes**
- Corrected tensor split order in angle update logic for improved
accuracy.

- **Documentation**
  - Enhanced argument documentation for new dynamic selection options.

- **Tests**
- Added comprehensive tests for dynamic selection mode, verifying
consistency across configurations.
- Updated existing test parameterizations to include dynamic selection
options.
  - Integrated new descriptor and parameters into model test suites.

- **Chores**
- Extended public APIs to include new utility functions supporting
dynamic neighbor selection.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: iProzd

deepmd/dpmodel/array_api.py

@@ -92,3 +92,37 @@ def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray
         )
     else:
         raise NotImplementedError("Only JAX arrays are supported.")
+
+
+def xp_add_at(x, indices, values):
+    """Adds values to the specified indices of x in place or returns new x (for JAX)."""
+    xp = array_api_compat.array_namespace(x, indices, values)
+    if array_api_compat.is_numpy_array(x):
+        # NumPy: supports np.add.at (in-place)
+        xp.add.at(x, indices, values)
+        return x
+
+    elif array_api_compat.is_jax_array(x):
+        # JAX: functional update, not in-place
+        return x.at[indices].add(values)
+    else:
+        # Fallback for array_api_strict: use basic indexing only
+        # may need a more efficient way to do this
+        n = indices.shape[0]
+        for i in range(n):
+            idx = int(indices[i])
+            x[idx, ...] = x[idx, ...] + values[i, ...]
+        return x
+
+
+def xp_bincount(x, weights=None, minlength=0):
+    """Counts the number of occurrences of each value in x."""
+    xp = array_api_compat.array_namespace(x)
+    if array_api_compat.is_numpy_array(x) or array_api_compat.is_jax_array(x):
+        result = xp.bincount(x, weights=weights, minlength=minlength)
+    else:
+        if weights is None:
+            weights = xp.ones_like(x)
+        result = xp.zeros((max(minlength, int(xp.max(x)) + 1),), dtype=weights.dtype)
+        result = xp_add_at(result, x, weights)
+    return result

deepmd/dpmodel/descriptor/dpa3.py

@@ -123,6 +123,18 @@ 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_dynamic_sel : bool, optional
+        Whether to dynamically select neighbors within the cutoff radius.
+        If True, the exact number of neighbors within the cutoff radius is used
+        without padding to a fixed selection numbers.
+        When enabled, users can safely set larger values for `e_sel` or `a_sel` (e.g., 1200 or 300, respectively)
+        to guarantee capturing all neighbors within the cutoff radius.
+        Note that when using dynamic selection, the `smooth_edge_update` must be True.
+    sel_reduce_factor : float, optional
+        Reduction factor applied to neighbor-scale normalization when `use_dynamic_sel` is True.
+        In the dynamic selection case, neighbor-scale normalization will use `e_sel / sel_reduce_factor`
+        or `a_sel / sel_reduce_factor` instead of the raw `e_sel` or `a_sel` values,
+        accommodating larger selection numbers.
     """
 
     def __init__(
@@ -150,6 +162,8 @@ def __init__(
         skip_stat: bool = False,
         optim_update: bool = True,
         smooth_edge_update: bool = False,
+        use_dynamic_sel: bool = False,
+        sel_reduce_factor: float = 10.0,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -176,6 +190,8 @@ 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_dynamic_sel = use_dynamic_sel
+        self.sel_reduce_factor = sel_reduce_factor
 
     def __getitem__(self, key):
         if hasattr(self, key):
@@ -207,6 +223,8 @@ def serialize(self) -> dict:
             "fix_stat_std": self.fix_stat_std,
             "optim_update": self.optim_update,
             "smooth_edge_update": self.smooth_edge_update,
+            "use_dynamic_sel": self.use_dynamic_sel,
+            "sel_reduce_factor": self.sel_reduce_factor,
         }
 
     @classmethod
@@ -303,6 +321,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,
+            use_dynamic_sel=self.repflow_args.use_dynamic_sel,
+            sel_reduce_factor=self.repflow_args.sel_reduce_factor,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,