Merge branch 'devel' into D0520_dpa3_dist

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>

Author: iProzd

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.9
+    rev: v0.11.10
     hooks:
       - id: ruff
         args: ["--fix"]

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

@@ -126,6 +126,18 @@ class RepFlowArgs:
     edge_init_use_dist : bool, optional
         Whether to use direct distance r to initialize the edge features instead of 1/r.
         Note that when using this option, the activation function will not be used when initializing edge features.
+    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__(
@@ -154,6 +166,8 @@ def __init__(
         optim_update: bool = True,
         smooth_edge_update: bool = False,
         edge_init_use_dist: bool = False,
+        use_dynamic_sel: bool = False,
+        sel_reduce_factor: float = 10.0,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -181,6 +195,8 @@ def __init__(
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
         self.edge_init_use_dist = edge_init_use_dist
+        self.use_dynamic_sel = use_dynamic_sel
+        self.sel_reduce_factor = sel_reduce_factor
 
     def __getitem__(self, key):
         if hasattr(self, key):
@@ -213,6 +229,8 @@ def serialize(self) -> dict:
             "optim_update": self.optim_update,
             "smooth_edge_update": self.smooth_edge_update,
             "edge_init_use_dist": self.edge_init_use_dist,
+            "use_dynamic_sel": self.use_dynamic_sel,
+            "sel_reduce_factor": self.sel_reduce_factor,
         }
 
     @classmethod
@@ -310,6 +328,8 @@ def init_subclass_params(sub_data, sub_class):
             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_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,