feat(dpa4): add fixed radial presets and improve export correctness

deepmd/dpmodel/descriptor/dpa4.py

@@ -307,14 +307,18 @@ class DescrptDPA4(NativeOP, BaseDescriptor):
     rcut
         Cutoff radius in Å.
     env_exp
-        C^3 cutoff envelope exponents `[rbf_env_exp, edge_env_exp]`.
-        - `rbf_env_exp`: Controls radial basis function envelope decay.
-        - `edge_env_exp`: Controls message passing edge weight envelope decay.
+        C^3 cutoff envelope exponents. A list `[rbf_env_exp, edge_env_exp]`
+        specifies the radial-basis and message-passing envelopes separately.
+        A zero radial-basis exponent disables that envelope.
+        An integer specifies only the message-passing envelope exponent and
+        disables the radial-basis envelope.
         Larger values give weaker suppression (values stay near 1.0 longer).
     channels
         Total channels per (l,m) coefficient.
     basis_type
-        Radial basis type. Supported values are ``"bessel"`` and ``"gaussian"``.
+        Radial basis type. Supported values are ``"bessel"``, ``"gaussian"``,
+        ``"bessel/fix"`` and ``"gaussian/fix"``; the ``/fix`` forms keep the
+        frequencies or centres fixed during training.
     n_radial
         Number of radial basis functions.
     radial_mlp
@@ -607,7 +611,7 @@ def __init__(
         ntypes: int,
         sel: list[int] | int,
         rcut: float = 6.0,
-        env_exp: list[int] | None = None,
+        env_exp: int | list[int] | None = None,
         channels: int = 64,
         basis_type: str = "bessel",
         n_radial: int = 16,
@@ -675,11 +679,17 @@ def __init__(
         self.rcut = float(rcut)
         if env_exp is None:
             env_exp = [7, 5]
-        if len(env_exp) != 2:
-            raise ValueError(
-                "`env_exp` must be a list of two integers: [rbf_env_exp, edge_env_exp]"
-            )
-        self.env_exp = [int(x) for x in env_exp]
+        if isinstance(env_exp, int):
+            self.env_exp = env_exp
+            edge_env_exp = env_exp
+        else:
+            if len(env_exp) != 2:
+                raise ValueError(
+                    "`env_exp` must be an integer or a list of two integers: "
+                    "[rbf_env_exp, edge_env_exp]"
+                )
+            self.env_exp = [int(x) for x in env_exp]
+            edge_env_exp = self.env_exp[1]
         self.eps = float(eps)
         # Floor for the envelope-squared degree normalization (GIE / env_seed).
         # version < 1.1 keeps the tiny ``eps`` floor (legacy path, untouched);
@@ -1068,7 +1078,7 @@ def __init__(
             basis_type=self.basis_type,
             n_radial=self.n_radial,
             precision=self.compute_precision,  # force fp32+
-            exponent=self.env_exp[0],
+            exponent=0 if isinstance(self.env_exp, int) else self.env_exp[0],
         )
 
         # === Shared radial embedding: RBF -> per-l radial features ===
@@ -1088,7 +1098,7 @@ def __init__(
         )
 
         # === C^3 cutoff envelope for edge weight ===
-        self.edge_envelope = C3CutoffEnvelope(rcut=self.rcut, exponent=self.env_exp[1])
+        self.edge_envelope = C3CutoffEnvelope(rcut=self.rcut, exponent=edge_env_exp)
 
         # === Edge-aligned Wigner-D calculator ===
         # Cartesian blocks (degree 1 or 2) skip the SO(2) rotations, so the full

deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py

@@ -148,9 +148,15 @@ def _project_frames(coeff: Any, proj: ChannelLinear, n_frames: int) -> Any:
     than the ``G``-point grid.
     """
     xp = array_api_compat.array_namespace(coeff)
-    n_batch, coeff_dim, n_focus, _ = coeff.shape
-    projected = proj(xp.reshape(coeff, (n_batch, coeff_dim, n_focus, n_frames, -1)))
-    return xp.reshape(projected, (n_batch, coeff_dim, n_focus, -1))
+    n_batch, coeff_dim, n_focus, n_channels = coeff.shape
+    projected = proj(
+        xp.reshape(
+            coeff, (n_batch, coeff_dim, n_focus, n_frames, n_channels // n_frames)
+        )
+    )
+    return xp.reshape(
+        projected, (n_batch, coeff_dim, n_focus, n_frames * projected.shape[-1])
+    )
 
 
 def _project_pair_in_one_transform(
@@ -162,14 +168,14 @@ def _project_pair_in_one_transform(
 ) -> tuple[Any, Any]:
     """Project two equally shaped coefficient operands in one linear transform."""
     xp = array_api_compat.array_namespace(left, right)
-    n_batch, coeff_dim, n_focus, _ = left.shape
-    frame_shape = (n_batch, coeff_dim, n_focus, n_frames, -1)
+    n_batch, coeff_dim, n_focus, n_channels = left.shape
+    frame_shape = (n_batch, coeff_dim, n_focus, n_frames, n_channels // n_frames)
     pair = xp.reshape(
         xp.concat(
             [xp.reshape(left, frame_shape), xp.reshape(right, frame_shape)],
             axis=-1,
         ),
-        (n_batch, coeff_dim, n_focus, -1),
+        (n_batch, coeff_dim, n_focus, 2 * n_channels),
     )
     pair_grid = to_grid(pair)
     split = pair_grid.shape[-1] // 2
@@ -354,10 +360,11 @@ def _project_operands(
         """Apply the two coefficient-space channel projections."""
         xp = array_api_compat.array_namespace(left)
         if self.mode == "self":
-            shape = (*left.shape[:-1], self.n_frames, -1)
+            n_channels = left.shape[-1]
+            shape = (*left.shape[:-1], self.n_frames, n_channels // self.n_frames)
             fused = xp.reshape(
                 xp.concat([xp.reshape(left, shape), xp.reshape(right, shape)], axis=-1),
-                (*left.shape[:-1], -1),
+                (*left.shape[:-1], 2 * n_channels),
             )  # per-frame concat -> (N, D, F, K*2C)
             left = _project_frames(fused, self.left_proj, self.n_frames)
             right = _project_frames(fused, self.right_proj, self.n_frames)
@@ -1254,11 +1261,14 @@ def _project_pair_in_one_transform(
         )
 
     def _to_grid(self, coeff: Any) -> Any:
-        # The per-frame channel width is inferred so the projector also serves
-        # widened operands (e.g. a branch hidden width ``n_branches * C``).
+        # Derive the per-frame width from the channel axis so empty batches
+        # and widened operands (e.g. ``n_branches * C``) are both valid.
         xp = array_api_compat.array_namespace(coeff)
-        n_batch, coeff_dim, n_focus, _ = coeff.shape
-        coeff_view = xp.reshape(coeff, (n_batch, coeff_dim, n_focus, self.n_frames, -1))
+        n_batch, coeff_dim, n_focus, n_channels = coeff.shape
+        n_channels //= self.n_frames
+        coeff_view = xp.reshape(
+            coeff, (n_batch, coeff_dim, n_focus, self.n_frames, n_channels)
+        )
         to_grid = xp_asarray_nodetach(
             xp, self.projector.to_grid_mat[...], device=array_api_compat.device(coeff)
         )
@@ -1274,7 +1284,6 @@ def _to_grid(self, coeff: Any) -> Any:
         # (`xp_einsum("gdk,ndfkc->ngfc")`) costs 3.6 ms. The same contraction
         # is faster there and slower here, so the choice belongs to the graph
         # around it rather than to the contraction itself.
-        n_channels = coeff_view.shape[-1]
         coeff_dk = xp.permute_dims(coeff_view, (0, 1, 3, 2, 4))  # (N, D, K, F, C)
         coeff_flat = xp.reshape(
             coeff_dk, (n_batch, coeff_dim * self.n_frames, n_focus * n_channels)
@@ -1285,7 +1294,7 @@ def _to_grid(self, coeff: Any) -> Any:
     def _from_grid(self, grid: Any) -> Any:
         # Channel width is inferred to match the (possibly widened) grid field.
         xp = array_api_compat.array_namespace(grid)
-        n_batch, _, n_focus, _ = grid.shape
+        n_batch, _, n_focus, n_channels = grid.shape
         coeff_dim = self.projector.coeff_dim // self.n_frames
         from_grid = xp_asarray_nodetach(
             xp, self.projector.from_grid_mat[...], device=array_api_compat.device(grid)
@@ -1294,7 +1303,6 @@ def _from_grid(self, grid: Any) -> Any:
         # einsum "dkg,ngfc->ndfkc" (with from_grid reshaped (D, K, G)) as a
         # broadcast batched matmul, then a reshape to (N, D, F, K*C). from_grid
         # is already stored as (D*K, G); the matmul output is reshaped/permuted.
-        n_channels = grid.shape[-1]
         grid_flat = xp.reshape(
             grid, (n_batch, self.projector.grid_size, n_focus * n_channels)
         )
@@ -1310,12 +1318,11 @@ def _from_grid(self, grid: Any) -> Any:
     def _from_grid_scalar(self, grid: Any) -> Any:
         """Project a grid field to the ``l=0`` coefficient only."""
         xp = array_api_compat.array_namespace(grid)
-        n_batch, _, n_focus, _ = grid.shape
+        n_batch, _, n_focus, n_channels = grid.shape
         from_grid = xp_asarray_nodetach(
             xp, self.projector.from_grid_mat[...], device=array_api_compat.device(grid)
         )
         from_grid = xp.astype(from_grid[: self.n_frames], grid.dtype)
-        n_channels = grid.shape[-1]
         grid_flat = xp.reshape(
             grid, (n_batch, self.projector.grid_size, n_focus * n_channels)
         )
@@ -1334,9 +1341,16 @@ def _scalar_so3_product(self, left: Any, right: Any) -> Any:
             xp, weight[...], device=array_api_compat.device(left)
         )
         weight = xp.astype(weight, left.dtype)
-        n_batch, coeff_dim, n_focus, _ = left.shape
-        left_view = xp.reshape(left, (n_batch, coeff_dim, n_focus, self.n_frames, -1))
-        right_view = xp.reshape(right, (n_batch, coeff_dim, n_focus, self.n_frames, -1))
+        n_batch, coeff_dim, n_focus, n_channels = left.shape
+        frame_shape = (
+            n_batch,
+            coeff_dim,
+            n_focus,
+            self.n_frames,
+            n_channels // self.n_frames,
+        )
+        left_view = xp.reshape(left, frame_shape)
+        right_view = xp.reshape(right, frame_shape)
         scalar = xp.sum(
             left_view * weight[None, :, None, :, None] * right_view,
             axis=(1, 3),
@@ -1355,9 +1369,11 @@ def _to_ndfc(self, value: Any) -> tuple[Any, tuple[int, ...]]:
             return xp.permute_dims(value, (0, 2, 1, 3)), tuple(value.shape)
         if self.layout == "fndc":
             return xp.permute_dims(value, (1, 2, 0, 3)), tuple(value.shape)
-        n_batch, coeff_dim, _ = value.shape
+        n_batch, coeff_dim, n_channels = value.shape
         return (
-            xp.reshape(value, (n_batch, coeff_dim, self.n_focus, -1)),
+            xp.reshape(
+                value, (n_batch, coeff_dim, self.n_focus, n_channels // self.n_focus)
+            ),
             tuple(value.shape),
         )
 
@@ -1377,7 +1393,7 @@ def _restore_layout(
             return xp.permute_dims(value, (2, 0, 1, 3))
         n_batch, input_coeff_dim, _ = shape_info
         coeff_dim = 1 if scalar_only else input_coeff_dim
-        return xp.reshape(value, (n_batch, coeff_dim, -1))
+        return xp.reshape(value, (n_batch, coeff_dim, value.shape[2] * value.shape[3]))
 
     def _slice_scalar_layout(self, value: Any) -> Any:
         """Select the degree axis from a restored full-layout tensor."""