diff --git a/deepmd/dpmodel/descriptor/__init__.py b/deepmd/dpmodel/descriptor/__init__.py index 765e3069fd..0b3570b4fe 100644 --- a/deepmd/dpmodel/descriptor/__init__.py +++ b/deepmd/dpmodel/descriptor/__init__.py @@ -8,6 +8,9 @@ from .dpa3 import ( DescrptDPA3, ) +from .dpa4 import ( + DescrptDPA4, +) from .hybrid import ( DescrptHybrid, ) @@ -34,6 +37,7 @@ "DescrptDPA1", "DescrptDPA2", "DescrptDPA3", + "DescrptDPA4", "DescrptHybrid", "DescrptSeA", "DescrptSeAttenV2", diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py new file mode 100644 index 0000000000..ad8c30a1d4 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -0,0 +1,1263 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +DPA4 (SeZM) descriptor: dpmodel (array-API) backend. + +This is the dpmodel port of ``deepmd.pt.model.descriptor.sezm.DescrptSeZM``. +It orchestrates the dpa4_nn building blocks on the padded, frame-explicit +edge layout (``E = nf * nloc * nnei``; no ``torch.nonzero``-style sparse +edge extraction anywhere; see ``dpa4_nn.edge_cache``). + +Scope notes (vs pt): + +- Only the standard DeePMD ``call(coord_ext, atype_ext, nlist, mapping)`` + path is ported. The pt-only paths (sparse ``edge_index`` inputs, + ``forward_with_edges``, zone bridging / InnerClamp, charge/spin condition + embedding, AMP autocast) are out of core scope; out-of-core construction + flags raise ``NotImplementedError`` at ``__init__`` (either here or in the + owning submodule). +- ``random_gamma`` is a training-only augmentation in pt + (``random_gamma and self.training``); dpmodel evaluates in inference mode, + so the roll is never applied (the config value is still serialized). +- ``use_amp`` is accepted and ignored: it is a pt-runtime (CUDA autocast) + switch with no dpmodel counterpart. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, + Any, + NoReturn, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + NativeOP, +) +from deepmd.dpmodel.common import ( + PRECISION_DICT, + get_xp_precision, + to_numpy_array, +) +from deepmd.dpmodel.utils import ( + EnvMat, +) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_descriptor import ( + BaseDescriptor, +) +from .dpa4_nn.block import ( + SeZMInteractionBlock, +) +from .dpa4_nn.edge_cache import ( + EdgeCache, + build_edge_cache, +) +from .dpa4_nn.embedding import ( + EnvironmentInitialEmbedding, + GeometricInitialEmbedding, + SeZMTypeEmbedding, +) +from .dpa4_nn.ffn import ( + EquivariantFFN, +) +from .dpa4_nn.indexing import ( + get_so3_dim_of_lmax, +) +from .dpa4_nn.norm import ( + ScalarRMSNorm, +) +from .dpa4_nn.radial import ( + C3CutoffEnvelope, + RadialBasis, + RadialMLP, +) +from .dpa4_nn.utils import ( + get_promoted_dtype, +) +from .dpa4_nn.wignerd import ( + WignerDCalculator, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.utils.data_system import ( + DeepmdDataSystem, + ) + from deepmd.utils.path import ( + DPPath, + ) + +ATTN_RES_MODES = ("none", "independent", "dependent") + + +@BaseDescriptor.register("SeZM") +@BaseDescriptor.register("sezm") +@BaseDescriptor.register("DPA4") +@BaseDescriptor.register("dpa4") +class DescrptDPA4(NativeOP, BaseDescriptor): + """ + DPA4 (SeZM) descriptor, dpmodel backend. + + See the pt ``DescrptSeZM`` docstring + (``deepmd/pt/model/descriptor/sezm.py``) for the full per-parameter + description; the constructor mirrors the pt signature and defaults + exactly. Parameters whose machinery is not ported to dpmodel raise + ``NotImplementedError`` at construction (some directly here, the rest + delegated to the owning submodule, e.g. ``layer_scale`` and the + ``*_attn_res`` / SO(2) attention projection flags). + + Execution outline (pt ``forward`` standard path): + + 1. Type embedding and pair-exclusion keep mask. + 2. ``build_edge_cache`` once (geometry, envelope, RBF, Wigner-D) on the + padded edge layout. + 3. Radial features once; optional environment FiLM seeding and geometric + initial embedding. + 4. ``SeZMInteractionBlock`` stack with the per-block l/m schedules. + 5. Final scalar (l=0) FFN readout to ``(nf, nloc, channels)``. + """ + + LATEST_VERSION: float = 1.1 + + def __init__( + self, + ntypes: int, + sel: list[int] | int, + rcut: float = 6.0, + env_exp: list[int] | None = None, + channels: int = 64, + basis_type: str = "bessel", + n_radial: int = 16, + radial_mlp: list[int] | None = None, + use_env_seed: bool = True, + random_gamma: bool = True, + lmax: int = 3, + l_schedule: list[int] | None = None, + mmax: int | None = 1, + kmax: int = 1, + m_schedule: list[int] | None = None, + extra_node_l: int = 0, + n_blocks: int = 3, + so2_norm: bool = False, + so2_layers: int = 4, + so2_attn_res: str = "none", + radial_so2_mode: str = "degree_channel", + radial_so2_rank: int = 1, + n_focus: int = 1, + focus_dim: int = 0, + n_atten_head: int = 1, + atten_f_mix: bool = False, + atten_v_proj: bool = False, + atten_o_proj: bool = False, + ffn_neurons: int = 0, + grid_mlp: bool | list[bool] = False, + grid_branch: int | list[int] = 0, + ffn_blocks: int = 1, + sandwich_norm: list[bool] | None = None, + mlp_bias: bool = False, + layer_scale: bool = False, + full_attn_res: str = "none", + block_attn_res: str = "none", + s2_activation: list[bool] | None = None, + ffn_so3_grid: bool = False, + node_wise_s2: bool = False, + node_wise_so3: bool = False, + message_node_s2: bool = False, + message_node_so3: bool = False, + lebedev_quadrature: bool | list[bool] | None = True, + activation_function: str = "silu", + glu_activation: bool = True, + use_amp: bool = True, + exclude_types: list[tuple[int, int]] | None = None, + precision: str = "float32", + eps: float = 1e-7, + trainable: bool = True, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + inner_clamp_r_inner: float | None = None, + inner_clamp_r_outer: float | None = None, + add_chg_spin_ebd: bool = False, + default_chg_spin: list[float] | None = None, + **kwargs: Any, + ) -> None: + self.version = float(self.LATEST_VERSION) + 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] + self.eps = float(eps) + # version >= 1.1 O(1) floor for the envelope-squared degree + # normalization (see pt sezm.py). + self.deg_norm_floor = 0.25 + + if isinstance(sel, int): + sel = [sel] + self.ntypes = int(ntypes) + self.sel = [int(x) for x in sel] + self.type_map = type_map + self.nnei = int(sum(self.sel)) + + self.channels = int(channels) + self.n_focus = int(n_focus) + if self.n_focus < 1: + raise ValueError("`n_focus` must be >= 1") + self.focus_dim = int(focus_dim) + if self.focus_dim < 0: + raise ValueError("`focus_dim` must be >= 0") + self.basis_type = str(basis_type).lower() + self.n_radial = int(n_radial) + if radial_mlp is None: + radial_mlp = [0] + self.radial_mlp = [self.channels if x == 0 else int(x) for x in radial_mlp] + if sandwich_norm is None: + sandwich_norm = [False, True, True, False] + if not isinstance(sandwich_norm, (list, tuple)) or len(sandwich_norm) != 4: + raise ValueError( + "sandwich_norm must be a list[bool] of length 4: " + "[so2_pre, so2_post, ffn_pre, ffn_post]" + ) + self.sandwich_norm = [bool(x) for x in sandwich_norm] + ( + self.so2_pre_norm, + self.so2_post_norm, + self.ffn_pre_norm, + self.ffn_post_norm, + ) = self.sandwich_norm + if s2_activation is None: + s2_activation = [False, True] + if not isinstance(s2_activation, list) or len(s2_activation) != 2: + raise ValueError( + "`s2_activation` must be a list[bool] of length 2: " + "[so2_activation, ffn_activation]" + ) + if any(not isinstance(flag, bool) for flag in s2_activation): + raise ValueError( + "`s2_activation` must be a list[bool] of length 2: " + "[so2_activation, ffn_activation]" + ) + self.s2_activation = list(s2_activation) + self.ffn_so3_grid = bool(ffn_so3_grid) + self.node_wise_s2 = bool(node_wise_s2) + self.node_wise_so3 = bool(node_wise_so3) + self.message_node_s2 = bool(message_node_s2) + self.message_node_so3 = bool(message_node_so3) + if lebedev_quadrature is None: + lebedev_quadrature = [True, True] + elif isinstance(lebedev_quadrature, bool): + lebedev_quadrature = [lebedev_quadrature, lebedev_quadrature] + if not isinstance(lebedev_quadrature, list) or len(lebedev_quadrature) != 2: + raise ValueError( + "`lebedev_quadrature` must be a bool or a list[bool] of length 2: " + "[so2_quadrature, ffn_quadrature]" + ) + if any(not isinstance(flag, bool) for flag in lebedev_quadrature): + raise ValueError( + "`lebedev_quadrature` must be a bool or a list[bool] of length 2: " + "[so2_quadrature, ffn_quadrature]" + ) + self.lebedev_quadrature = list(lebedev_quadrature) + # The tensor-product (e3nn-style) sphere grid is not ported to + # dpmodel; only the packaged Lebedev quadrature path exists + # (see dpa4_nn.projection). + if not all(self.lebedev_quadrature): + raise NotImplementedError( + "lebedev_quadrature entries with False (tensor-product S2 " + "grid) are not ported to dpmodel" + ) + self.activation_function = str(activation_function) + self.glu_activation = bool(glu_activation) + + # === Split effective activation config by branch (pt sezm.py) === + self.so2_s2_activation = self.s2_activation[0] + self.ffn_s2_activation = False if self.ffn_so3_grid else self.s2_activation[1] + self.so2_lebedev_quadrature = self.lebedev_quadrature[0] + self.ffn_lebedev_quadrature = self.lebedev_quadrature[1] + self.so2_activation_function = ( + "silu" if self.so2_s2_activation else self.activation_function + ) + self.ffn_activation_function = ( + "silu" if self.ffn_s2_activation else self.activation_function + ) + self.ffn_glu_activation = ( + True + if (self.ffn_s2_activation or self.ffn_so3_grid) + else self.glu_activation + ) + self.out_activation_function = self.activation_function + self.out_glu_activation = self.glu_activation + self.precision = str(precision) + # Geometry / seeding paths run in promoted ("fp32+") precision (pt + # uses compute_dtype = get_promoted_dtype(dtype) there). + self.compute_precision = str( + np.dtype(get_promoted_dtype(PRECISION_DICT[self.precision])).name + ) + self.mlp_bias = bool(mlp_bias) + self.layer_scale = bool(layer_scale) + # pt-runtime-only switch (CUDA bfloat16 autocast during training); + # accepted for config compatibility and ignored by dpmodel. + self.use_amp = bool(use_amp) + self.trainable = bool(trainable) + self.seed = seed + self.random_gamma = bool(random_gamma) + self.add_chg_spin_ebd = bool(add_chg_spin_ebd) + if self.add_chg_spin_ebd: + raise NotImplementedError( + "add_chg_spin_ebd=True (ChargeSpinEmbedding) is not ported to dpmodel" + ) + if default_chg_spin is not None and len(default_chg_spin) != 2: + raise ValueError("`default_chg_spin` must contain [charge, spin].") + self.default_chg_spin = ( + None if default_chg_spin is None else [float(x) for x in default_chg_spin] + ) + + # === Zone bridging (InnerClamp + BridgingSwitch): not ported === + self.inner_clamp_r_inner = ( + float(inner_clamp_r_inner) if inner_clamp_r_inner is not None else None + ) + self.inner_clamp_r_outer = ( + float(inner_clamp_r_outer) if inner_clamp_r_outer is not None else None + ) + if self.inner_clamp_r_inner is not None or self.inner_clamp_r_outer is not None: + raise NotImplementedError( + "inner_clamp_r_inner/inner_clamp_r_outer (zone bridging) are " + "not ported to dpmodel" + ) + + # === Env seed derived dimensions (pt sezm.py) === + self.use_env_seed = bool(use_env_seed) + self.env_seed_embed_dim = min(self.channels, 128) + self.env_seed_type_dim = min(32, max(8, self.channels // 4)) + axis_dim = 4 if self.env_seed_embed_dim < 64 else 8 + self.env_seed_axis_dim = min(axis_dim, max(1, self.env_seed_embed_dim - 1)) + rbf_out_dim = max(32, self.env_seed_embed_dim - 2 * self.env_seed_type_dim) + g_in_dim = rbf_out_dim + 2 * self.env_seed_type_dim + self.env_seed_hidden_dim = min(256, max(2 * self.env_seed_embed_dim, g_in_dim)) + + # === Deterministic seed split (same indices as pt) === + seed_type_embedding = child_seed(self.seed, 0) + seed_blocks = child_seed(self.seed, 1) + seed_out = child_seed(self.seed, 2) + seed_radial_embedding = child_seed(self.seed, 3) + seed_env_seed = child_seed(self.seed, 4) + + # === L/M schedules === + self._init_lm_schedules(lmax, n_blocks, l_schedule, mmax, m_schedule) + self.kmax = int(kmax) + if self.kmax < 0: + raise ValueError("`kmax` must be non-negative") + if self.kmax > self.lmax: + raise ValueError("`kmax` must be <= `lmax`") + self.ebed_dims = [get_so3_dim_of_lmax(l) for l in self.l_schedule] + self._init_node_l_schedules(extra_node_l) + self.rad_sizes_per_block = [l + 1 for l in self.l_schedule] + + self.so2_norm = bool(so2_norm) + self.so2_layers = int(so2_layers) + self.so2_attn_res_mode = str(so2_attn_res).lower() + if self.so2_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`so2_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + self.radial_so2_mode = str(radial_so2_mode).lower() + if self.radial_so2_mode not in {"none", "degree", "degree_channel"}: + raise ValueError( + "`radial_so2_mode` must be one of 'none', 'degree', or 'degree_channel'" + ) + self.radial_so2_rank = int(radial_so2_rank) + if self.radial_so2_rank < 0: + raise ValueError("`radial_so2_rank` must be non-negative") + self.ffn_neurons = int(ffn_neurons) + self.block_ffn_neurons = self._resolve_ffn_neurons( + self.ffn_neurons, glu_activation=self.ffn_glu_activation + ) + self.out_ffn_neurons = self._resolve_ffn_neurons( + self.ffn_neurons, glu_activation=self.out_glu_activation + ) + self.grid_mlp = self._broadcast_grid_setting( + grid_mlp, name="grid_mlp", cast=bool + ) + self.grid_branch = self._broadcast_grid_setting( + grid_branch, name="grid_branch", cast=int, non_negative=True + ) + ( + self.node_wise_grid_mlp, + self.message_node_grid_mlp, + self.ffn_grid_mlp, + ) = self.grid_mlp + ( + self.node_wise_grid_branch, + self.message_node_grid_branch, + self.ffn_grid_branch, + ) = self.grid_branch + self.ffn_blocks = int(ffn_blocks) + if self.ffn_blocks < 1: + raise ValueError("`ffn_blocks` must be >= 1") + self.full_attn_res_mode = str(full_attn_res).lower() + if self.full_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`full_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + self.block_attn_res_mode = str(block_attn_res).lower() + if self.block_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`block_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + self.use_full_attn_res = self.full_attn_res_mode != "none" + self.use_block_attn_res = self.block_attn_res_mode != "none" + if self.use_full_attn_res and self.use_block_attn_res: + raise ValueError( + "`full_attn_res` and `block_attn_res` cannot both be enabled" + ) + self.n_atten_head = int(n_atten_head) + self.atten_f_mix = bool(atten_f_mix) + self.use_atten_v_proj = bool(atten_v_proj) + self.use_atten_o_proj = bool(atten_o_proj) + so2_focus_dim = self.channels if self.focus_dim == 0 else self.focus_dim + attn_focus_dim = ( + self.n_focus * so2_focus_dim if self.atten_f_mix else so2_focus_dim + ) + if self.n_atten_head > 0 and attn_focus_dim % self.n_atten_head != 0: + raise ValueError( + "`n_atten_head` must divide the attention width " + "(`focus_dim` or `n_focus * focus_dim` when `atten_f_mix=True`)" + ) + + # === Excluded type pairs === + self.reinit_exclude(exclude_types) + + # === Type embedding (fp32+) === + self.type_embedding = SeZMTypeEmbedding( + ntypes=self.ntypes, + embed_dim=self.channels, + precision=self.compute_precision, + seed=seed_type_embedding, + trainable=self.trainable, + ) + + # === Env FiLM embedding (optional, fp32+) === + compute_np_prec = PRECISION_DICT[self.compute_precision] + if self.use_env_seed: + self.env_seed_embedding: EnvironmentInitialEmbedding | None = ( + EnvironmentInitialEmbedding( + ntypes=self.ntypes, + n_radial=self.n_radial, + channels=self.channels, + embed_dim=self.env_seed_embed_dim, + axis_dim=self.env_seed_axis_dim, + type_dim=self.env_seed_type_dim, + hidden_dim=self.env_seed_hidden_dim, + mlp_bias=self.mlp_bias, + activation_function=self.activation_function, + eps=self.eps, + precision=self.compute_precision, + trainable=self.trainable, + seed=seed_env_seed, + ) + ) + self.film_scale_norm: ScalarRMSNorm | None = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + precision=self.compute_precision, + trainable=self.trainable, + ) + self.film_shift_norm: ScalarRMSNorm | None = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + precision=self.compute_precision, + trainable=self.trainable, + ) + film_strength_init = 0.01 + self.film_scale_strength_log: np.ndarray | None = np.full( + (1,), math.log(film_strength_init), dtype=compute_np_prec + ) + self.film_shift_strength_log: np.ndarray | None = np.full( + (1,), math.log(film_strength_init), dtype=compute_np_prec + ) + else: + self.env_seed_embedding = None + self.film_scale_norm = None + self.film_shift_norm = None + self.film_scale_strength_log = None + self.film_shift_strength_log = None + + self.radial_basis = RadialBasis( + rcut=self.rcut, + basis_type=self.basis_type, + n_radial=self.n_radial, + precision=self.compute_precision, + exponent=self.env_exp[0], + ) + + # === Shared radial embedding: RBF -> per-l radial features (fp32+) === + radial_out_dim = (self.node_l_schedule[0] + 1) * self.channels + radial_mlp_layers = [self.n_radial, *self.radial_mlp, radial_out_dim] + self.radial_embedding = RadialMLP( + radial_mlp_layers, + activation_function=self.activation_function, + precision=self.compute_precision, + trainable=self.trainable, + seed=seed_radial_embedding, + ) + + # === C^3 cutoff envelope for edge weight === + self.edge_envelope = C3CutoffEnvelope( + self.rcut, self.env_exp[1], precision=self.compute_precision + ) + + wigner_lmax = self.l_schedule[0] + self.wigner_calc = WignerDCalculator( + wigner_lmax, eps=self.eps, precision=self.compute_precision + ) + + # === Geometric initial embedding (optional, fp32+) === + self.use_gie = self.use_env_seed and self.node_l_schedule[0] > 0 + if self.use_gie: + self.gie: GeometricInitialEmbedding | None = GeometricInitialEmbedding( + lmax=self.node_l_schedule[0], + channels=self.channels, + precision=self.compute_precision, + ) + if self.extra_node_l > 0: + self.gie_zonal_wigner_calc: WignerDCalculator | None = ( + WignerDCalculator( + self.node_l_schedule[0], + eps=self.eps, + precision=self.compute_precision, + ) + ) + else: + self.gie_zonal_wigner_calc = None + else: + self.gie = None + self.gie_zonal_wigner_calc = None + + # === Interaction blocks === + blocks: list[SeZMInteractionBlock] = [] + for block_idx, (l_b, node_l_b, m_b) in enumerate( + zip( + self.l_schedule, + self.node_l_schedule, + self.m_schedule, + strict=True, + ) + ): + k_b = min(self.kmax, l_b) + blocks.append( + SeZMInteractionBlock( + lmax=l_b, + node_lmax=node_l_b, + mmax=m_b, + channels=self.channels, + n_focus=self.n_focus, + focus_dim=self.focus_dim, + so2_norm=self.so2_norm, + so2_layers=self.so2_layers, + so2_attn_res=self.so2_attn_res_mode, + radial_so2_mode=self.radial_so2_mode, + radial_so2_rank=self.radial_so2_rank, + ffn_neurons=self.block_ffn_neurons, + node_wise_grid_mlp=self.node_wise_grid_mlp, + node_wise_grid_branch=self.node_wise_grid_branch, + message_node_grid_mlp=self.message_node_grid_mlp, + message_node_grid_branch=self.message_node_grid_branch, + ffn_grid_mlp=self.ffn_grid_mlp, + ffn_grid_branch=self.ffn_grid_branch, + ffn_blocks=self.ffn_blocks, + layer_scale=self.layer_scale, + full_attn_res=self.full_attn_res_mode, + block_attn_res=self.block_attn_res_mode, + so2_s2_activation=self.so2_s2_activation, + node_wise_s2=self.node_wise_s2, + node_wise_so3=self.node_wise_so3, + message_node_s2=self.message_node_s2, + message_node_so3=self.message_node_so3, + ffn_s2_activation=self.ffn_s2_activation, + ffn_so3_grid=self.ffn_so3_grid, + kmax=k_b, + so2_lebedev_quadrature=self.so2_lebedev_quadrature, + ffn_lebedev_quadrature=self.ffn_lebedev_quadrature, + n_atten_head=self.n_atten_head, + atten_f_mix=self.atten_f_mix, + atten_v_proj=self.use_atten_v_proj, + atten_o_proj=self.use_atten_o_proj, + so2_pre_norm=self.so2_pre_norm, + so2_post_norm=self.so2_post_norm, + so2_activation_function=self.so2_activation_function, + ffn_pre_norm=self.ffn_pre_norm, + ffn_post_norm=self.ffn_post_norm, + ffn_activation_function=self.ffn_activation_function, + ffn_glu_activation=self.ffn_glu_activation, + mlp_bias=self.mlp_bias, + eps=self.eps, + precision=self.precision, + seed=child_seed(seed_blocks, block_idx), + trainable=self.trainable, + ) + ) + self.blocks = blocks + + # === Final FFN for l=0 output mixing (fp32+) === + self.output_ffn = EquivariantFFN( + lmax=0, + channels=self.channels, + hidden_channels=self.out_ffn_neurons, + grid_mlp=False, + s2_activation=False, + activation_function=self.out_activation_function, + glu_activation=self.out_glu_activation, + mlp_bias=self.mlp_bias, + precision=self.compute_precision, + trainable=self.trainable, + seed=seed_out, + ) + + # === Statistics buffers (interface compatibility, unused in call) === + model_np_prec = PRECISION_DICT[self.precision] + self.mean = np.zeros((0,), dtype=model_np_prec) + self.stddev = np.ones((0,), dtype=model_np_prec) + + # ========================================================================= + # Construction helpers (mirroring pt) + # ========================================================================= + + @staticmethod + def _broadcast_grid_setting( + value: bool | int | list[bool] | list[int], + *, + name: str, + cast: type, + non_negative: bool = False, + ) -> list: + """Normalize a grid-path setting to ``[node_wise, message_node, ffn]``.""" + entries = list(value) if isinstance(value, list) else [value, value, value] + if len(entries) != 3: + raise ValueError( + f"`{name}` must be a {cast.__name__} or a list[{cast.__name__}] " + "of length 3: [node_wise, message_node, ffn]" + ) + normalized = [cast(entry) for entry in entries] + if non_negative and any(entry < 0 for entry in normalized): + raise ValueError(f"`{name}` entries must be non-negative") + return normalized + + def _resolve_ffn_neurons(self, ffn_neurons: int, *, glu_activation: bool) -> int: + """Resolve one FFN hidden width from the descriptor config.""" + resolved = int(ffn_neurons) + if resolved < 0: + raise ValueError("`ffn_neurons` must be >= 0") + if resolved > 0: + return resolved + base_width = ( + (8.0 * float(self.channels) / 3.0) + if glu_activation + else (4.0 * float(self.channels)) + ) + return int(32 * math.ceil(base_width / 32.0)) + + def _init_lm_schedules( + self, + lmax: int, + n_blocks: int, + l_schedule: list[int] | None, + mmax: int | None, + m_schedule: list[int] | None, + ) -> None: + """Parse and validate L/M schedules (pt ``_init_lm_schedules``).""" + if l_schedule is None: + self.l_schedule = [int(lmax)] * int(n_blocks) + else: + self.l_schedule = [int(x) for x in l_schedule] + if len(self.l_schedule) == 0: + raise ValueError("`l_schedule` must be non-empty") + if any(x < 0 for x in self.l_schedule): + raise ValueError("`l_schedule` entries must be non-negative") + if any( + self.l_schedule[i] < self.l_schedule[i + 1] + for i in range(len(self.l_schedule) - 1) + ): + raise ValueError("`l_schedule` must be non-increasing (pyramid schedule)") + + self.lmax = int(self.l_schedule[0]) + self.n_blocks = len(self.l_schedule) + + if m_schedule is None: + if mmax is None: + self.m_schedule = [int(l) for l in self.l_schedule] + else: + mmax_i = int(mmax) + if mmax_i < 0: + raise ValueError("`mmax` must be non-negative") + self.m_schedule = [min(mmax_i, int(l)) for l in self.l_schedule] + else: + self.m_schedule = [int(x) for x in m_schedule] + if len(self.m_schedule) == 0: + raise ValueError("`m_schedule` must be non-empty") + if len(self.m_schedule) != len(self.l_schedule): + raise ValueError("`m_schedule` must have the same length as `l_schedule`") + if any(x < 0 for x in self.m_schedule): + raise ValueError("`m_schedule` entries must be non-negative") + if any(m > l for m, l in zip(self.m_schedule, self.l_schedule, strict=True)): + raise ValueError( + "`m_schedule` entries must satisfy `m_schedule[i] <= l_schedule[i]`" + ) + self.mmax = int(self.m_schedule[0]) + + def _init_node_l_schedules(self, extra_node_l: int) -> None: + """Parse node degree schedules derived from message-passing schedules.""" + self.extra_node_l = int(extra_node_l) + if self.extra_node_l < 0: + raise ValueError("`extra_node_l` must be non-negative") + self.node_l_schedule = [ + int(l_value) + self.extra_node_l for l_value in self.l_schedule + ] + self.node_ebed_dims = [ + get_so3_dim_of_lmax(l_value) for l_value in self.node_l_schedule + ] + self.node_lmax = int(self.node_l_schedule[0]) + self.node_ebed_dim = int(self.node_ebed_dims[0]) + + def reinit_exclude( + self, exclude_types: list[tuple[int, int]] | None = None + ) -> None: + if exclude_types is None: + exclude_types = [] + self.exclude_types = exclude_types + self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) + + # ========================================================================= + # Forward + # ========================================================================= + + def call( + self, + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + comm_dict: dict | None = None, + ) -> tuple[Array, Any, Any, Any, Any]: + """Compute the DPA4 descriptor. + + Parameters + ---------- + coord_ext + Extended coordinates with shape (nf, nall*3) or (nf, nall, 3). + atype_ext + Extended atom types with shape (nf, nall). + nlist + Neighbor list with shape (nf, nloc, nnei); -1 marks padding. + mapping + Extended-to-local mapping with shape (nf, nall), or None when the + neighbor indices are already local. + fparam + Frame parameters; not used by DPA4 (interface compatibility). + comm_dict + MPI communication metadata; not used (interface compatibility). + + Returns + ------- + descriptor + Scalar descriptor with shape (nf, nloc, channels). + rot_mat, g2, h2, sw + ``None`` placeholders (pt returns empty tensors for these). + """ + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + nf, nloc, nnei = nlist.shape + nall = xp.reshape(coord_ext, (nf, -1)).shape[1] // 3 + extended_coord = xp.reshape(coord_ext, (nf, nall, 3)) + extended_coord = xp.astype( + extended_coord, get_xp_precision(xp, self.compute_precision) + ) + n_nodes = nf * nloc + + # === Step 1. Excluded type pairs (keep mask, True means keep) === + # The dpmodel PairExcludeMask returns an int mask; build_edge_cache + # expects a boolean keep mask. + pair_keep_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) != 0 + + # === Step 2. Type embedding (l=0) === + atype_loc = atype_ext[:, :nloc] + type_ebed = xp.reshape( + self.type_embedding(atype_loc), (n_nodes, self.channels) + ) # (N, C) + + # === Step 3. Build edge cache once (geometry + RBF + Wigner-D) === + # Random local-Z roll is a training-only augmentation in pt; the + # dpmodel descriptor evaluates in inference mode, so gamma is fixed. + edge_cache = build_edge_cache( + type_ebed=type_ebed, + extended_coord=extended_coord, + nlist=nlist, + mapping=mapping, + pair_keep_mask=pair_keep_mask, + eps=self.eps, + deg_norm_floor=(self.deg_norm_floor if self.version >= 1.1 else self.eps), + edge_envelope=self.edge_envelope, + radial_basis=self.radial_basis, + n_radial=self.n_radial, + random_gamma=False, + wigner_calc=self.wigner_calc, + ) + + # === Step 4. Compute radial features once (fp32+) === + # Padded layout: E = nf * nloc * nnei is shape-determined, so there is + # no pt-style empty-edge special case. + radial_feat = xp.reshape( + self.radial_embedding(edge_cache.edge_rbf), + (-1, self.node_l_schedule[0] + 1, self.channels), + ) # (E, node_lmax+1, C) + if self.version >= 1.1: + radial_feat = radial_feat * xp.reshape(edge_cache.edge_env, (-1, 1, 1)) + + # === Step 5. Env FiLM conditioning (optional, fp32+) === + x0_out = type_ebed # (N, C) + if self.use_env_seed: + atype_flat = xp.reshape(atype_loc, (-1,)) + film = self.env_seed_embedding( + edge_cache=edge_cache, + atype_flat=atype_flat, + n_nodes=n_nodes, + ) # (N, 2*C) + scale_logits = film[:, : self.channels] + shift_logits = film[:, self.channels :] + scale_hat = self.film_scale_norm(scale_logits) + shift_hat = self.film_shift_norm(shift_logits) + device = array_api_compat.device(scale_hat) + scale_strength = xp.exp( + xp.asarray(self.film_scale_strength_log, device=device) + ) + shift_strength = xp.exp( + xp.asarray(self.film_shift_strength_log, device=device) + ) + scale = 1.0 + scale_strength * xp.tanh(scale_hat) + shift = shift_strength * xp.tanh(shift_hat) + x0_out = type_ebed * scale + shift + + # === Step 6. Build backbone l=0 features === + # pt scatters x0_out into x[:, 0, 0, :] of a zeros tensor; here this + # is a concat with zero rows for l >= 1 (no fancy __setitem__). + ebed_dim_0 = self.node_ebed_dims[0] + x = xp.concat( + [ + x0_out[:, None, :], + xp.zeros( + (n_nodes, ebed_dim_0 - 1, self.channels), + dtype=x0_out.dtype, + device=array_api_compat.device(x0_out), + ), + ] + if ebed_dim_0 > 1 + else [x0_out[:, None, :]], + axis=1, + ) # (N, D, C) + + # === Step 7. Geometric initial embedding (fp32+) === + if self.use_gie: + zonal_coupling = self._build_gie_zonal_coupling(edge_cache) + x = x + self.gie( + n_nodes=n_nodes, + edge_cache=edge_cache, + radial_feat=radial_feat[:, 1:, :], + zonal_coupling=zonal_coupling, + ) + x = x[:, :, None, :] # (N, D, 1, C) + + # === Step 8. Fuse edge type features into radial features === + radial_feat = radial_feat + edge_cache.edge_type_feat[:, None, :] + rad_feat_per_block = [ + radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block + ] + + # === Step 9. Run interaction blocks (residual baseline path) === + for i, block in enumerate(self.blocks): + x = x[:, : self.node_ebed_dims[i], :, :] + x = block(x, edge_cache, rad_feat_per_block[i])[0] + + # === Step 10. Final l=0 output mixing === + x_scalar = xp.reshape( + x[:, 0:1, :, :], (n_nodes, 1, 1, self.channels) + ) # (N, 1, 1, C) + x_scalar = x_scalar + self.output_ffn(x_scalar) + + # === Step 11. Reshape and return === + descriptor = xp.reshape(x_scalar, (nf, nloc, self.channels)) + descriptor = xp.astype(descriptor, get_xp_precision(xp, "global")) + return descriptor, None, None, None, None + + def _build_gie_zonal_coupling(self, edge_cache: EdgeCache) -> Any: + """ + Build node-level zonal coupling for GIE when node degrees exceed MP + degrees (pt ``_build_gie_zonal_coupling``). + + Returns ``None`` when ``extra_node_l == 0``, letting GIE gather from + the MP Wigner-D cache. + """ + if self.gie_zonal_wigner_calc is None: + return None + xp = array_api_compat.array_namespace(edge_cache.edge_quat) + device = array_api_compat.device(edge_cache.edge_quat) + n_edge = edge_cache.dst.shape[0] + mp_row_count = self.ebed_dims[0] - 1 + mp_rows = self.gie.non_scalar_row_index[:mp_row_count] + mp_cols = self.gie.zonal_m0_col_index_for_row[:mp_row_count] + Dt_full = edge_cache.Dt_full + dim_full = Dt_full.shape[-1] + flat_index = xp.asarray(mp_rows * dim_full + mp_cols, device=device) + mp_coupling = xp.take( + xp.reshape(Dt_full, (n_edge, dim_full * dim_full)), + flat_index, + axis=1, + ) # (E, D_mp - 1) + extra_coupling = self.gie_zonal_wigner_calc.forward_zonal( + edge_cache.edge_quat, + lmin=self.lmax + 1, + ) + return xp.concat([mp_coupling, extra_coupling], axis=1) + + # ========================================================================= + # DeePMD descriptor interface + # ========================================================================= + + def get_rcut(self) -> float: + return self.rcut + + def get_rcut_smth(self) -> float: + return self.rcut + + def get_sel(self) -> list[int]: + return self.sel + + def get_nsel(self) -> int: + return sum(self.sel) + + def get_ntypes(self) -> int: + return self.ntypes + + def get_type_map(self) -> list[str]: + return self.type_map if self.type_map is not None else [] + + def get_dim_out(self) -> int: + return self.channels + + def get_dim_emb(self) -> int: + return self.get_dim_out() + + def mixed_types(self) -> bool: + """DPA4 uses SeZMTypeEmbedding, no type-distinguished nlist needed.""" + return True + + def has_message_passing(self) -> bool: + return bool(len(self.blocks) > 0 and self.lmax > 0) + + def has_message_passing_across_ranks(self) -> bool: + return self.has_message_passing() + + def need_sorted_nlist_for_lower(self) -> bool: + return False + + def get_env_protection(self) -> float: + return self.eps + + @property + def dim_out(self) -> int: + return self.get_dim_out() + + @property + def dim_emb(self) -> int: + return self.get_dim_emb() + + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> NoReturn: + """Parameter sharing is a pt-backend training feature.""" + raise NotImplementedError + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat: Any = None + ) -> NoReturn: + raise NotImplementedError("change_type_map is not supported for SeZM") + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> NoReturn: + raise NotImplementedError("Compression is unsupported for SeZM.") + + # === Statistics interface (interface compatibility only) === + # SeZM normalizes with learnable RMS norms; mean/stddev are kept only for + # interface and checkpoint-format compatibility (see pt sezm.py). + + def compute_input_stats( + self, merged: list[dict], path: DPPath | None = None + ) -> None: + """No-op: statistics are not used by the DPA4 forward pass.""" + + def set_stat_mean_and_stddev(self, mean: Array, stddev: Array) -> None: + """Set mean and stddev (interface compatibility, unused in call).""" + self.mean = mean + self.stddev = stddev + + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: + """Get mean and stddev (interface compatibility, unused in call).""" + return self.mean, self.stddev + + # ========================================================================= + # Serialization (pt state_dict-key compatible) + # ========================================================================= + + def _variables(self) -> dict[str, np.ndarray]: + """Variables keyed exactly by the pt ``state_dict()`` key names.""" + model_np_prec = PRECISION_DICT[self.precision] + variables: dict[str, np.ndarray] = { + # pt interface-compatibility buffers + "version_tensor": np.asarray(self.version, dtype=np.float64), + "_empty_tensor": np.zeros((0,), dtype=np.float64), + "mean": np.asarray(self.mean, dtype=model_np_prec), + "stddev": np.asarray(self.stddev, dtype=model_np_prec), + } + + def add(prefix: str, sub_vars: dict[str, Any]) -> None: + for key, value in sub_vars.items(): + variables[f"{prefix}{key}"] = to_numpy_array(value) + + add("type_embedding.", self.type_embedding.serialize()["@variables"]) + if self.use_env_seed: + add( + "env_seed_embedding.", + self.env_seed_embedding.serialize()["@variables"], + ) + add("film_scale_norm.", self.film_scale_norm.serialize()["@variables"]) + add("film_shift_norm.", self.film_shift_norm.serialize()["@variables"]) + variables["film_scale_strength_log"] = to_numpy_array( + self.film_scale_strength_log + ) + variables["film_shift_strength_log"] = to_numpy_array( + self.film_shift_strength_log + ) + add("radial_basis.", self.radial_basis.serialize()["@variables"]) + add("radial_embedding.net.", self.radial_embedding.serialize()["@variables"]) + + # Static pt WignerDCalculator buffers (rebuilt at construction here; + # emitted so pt's strict load_state_dict finds every key). + def wigner_buffers(calc: WignerDCalculator) -> dict[str, np.ndarray]: + return { + "l1_perm": np.asarray([1, 2, 0], dtype=np.int64), + "l1_sign_outer": np.asarray(calc.l1_sign_outer, dtype=np.float64), + } + + add("wigner_calc.", wigner_buffers(self.wigner_calc)) + if self.use_gie: + add( + "gie.", + { + "non_scalar_row_index": self.gie.non_scalar_row_index, + "zonal_m0_col_index_for_row": self.gie.zonal_m0_col_index_for_row, + "radial_slot_index_for_row": self.gie.radial_slot_index_for_row, + }, + ) + if self.gie_zonal_wigner_calc is not None: + add( + "gie_zonal_wigner_calc.", + wigner_buffers(self.gie_zonal_wigner_calc), + ) + for i, block in enumerate(self.blocks): + add(f"blocks.{i}.", block._variables()) + add("output_ffn.", self.output_ffn._variables()) + return variables + + def _load_variables(self, variables: dict[str, Any]) -> None: + """Load variables keyed by the pt ``state_dict()`` key names.""" + variables = dict(variables) + + def take_prefix(prefix: str) -> dict[str, Any]: + sub = { + key[len(prefix) :]: value + for key, value in variables.items() + if key.startswith(prefix) + } + for key in list(variables): + if key.startswith(prefix): + del variables[key] + return sub + + # Transient / static pt buffers rebuilt at construction. + for key in ("version_tensor", "_empty_tensor"): + variables.pop(key, None) + take_prefix("wigner_calc.") + take_prefix("gie.") + take_prefix("gie_zonal_wigner_calc.") + + model_np_prec = PRECISION_DICT[self.precision] + compute_np_prec = PRECISION_DICT[self.compute_precision] + if "mean" in variables: + self.mean = np.asarray(variables.pop("mean"), dtype=model_np_prec) + if "stddev" in variables: + self.stddev = np.asarray(variables.pop("stddev"), dtype=model_np_prec) + + def load_via_serialize(attr: str, prefix: str) -> None: + sub = getattr(self, attr) + sv = take_prefix(prefix) + if sub is None: + if sv: + raise KeyError(f"Unexpected variables with prefix: {prefix}") + return + if not sv: + raise KeyError(f"Missing variables with prefix: {prefix}") + data = sub.serialize() + data["@variables"] = sv + setattr(self, attr, type(sub).deserialize(data)) + + load_via_serialize("type_embedding", "type_embedding.") + load_via_serialize("env_seed_embedding", "env_seed_embedding.") + load_via_serialize("film_scale_norm", "film_scale_norm.") + load_via_serialize("film_shift_norm", "film_shift_norm.") + load_via_serialize("radial_basis", "radial_basis.") + load_via_serialize("radial_embedding", "radial_embedding.net.") + if self.use_env_seed: + for name in ("film_scale_strength_log", "film_shift_strength_log"): + value = np.asarray(variables.pop(name), dtype=compute_np_prec) + setattr(self, name, value.reshape((1,))) + for i, block in enumerate(self.blocks): + block._load_variables(take_prefix(f"blocks.{i}.")) + self.output_ffn._load_variables(take_prefix("output_ffn.")) + if variables: + raise KeyError(f"Unknown variables: {sorted(variables)}") + + def serialize(self) -> dict[str, Any]: + """Serialize the descriptor (pt ``DescrptSeZM.serialize`` format).""" + return { + "@class": "Descriptor", + "type": "SeZM", + "@version": self.version, + "config": { + "ntypes": self.ntypes, + "sel": self.sel, + "rcut": self.rcut, + "env_exp": self.env_exp, + "type_map": self.type_map, + "lmax": self.lmax, + "n_blocks": self.n_blocks, + "l_schedule": self.l_schedule, + "mmax": self.mmax, + "kmax": self.kmax, + "m_schedule": self.m_schedule, + "extra_node_l": self.extra_node_l, + "channels": self.channels, + "basis_type": self.basis_type, + "n_radial": self.n_radial, + "radial_mlp": self.radial_mlp, + "use_env_seed": self.use_env_seed, + "random_gamma": self.random_gamma, + "so2_norm": self.so2_norm, + "so2_layers": self.so2_layers, + "so2_attn_res": self.so2_attn_res_mode, + "radial_so2_mode": self.radial_so2_mode, + "radial_so2_rank": self.radial_so2_rank, + "n_focus": self.n_focus, + "focus_dim": self.focus_dim, + "ffn_neurons": self.ffn_neurons, + "grid_mlp": self.grid_mlp, + "grid_branch": self.grid_branch, + "ffn_blocks": self.ffn_blocks, + "layer_scale": self.layer_scale, + "n_atten_head": self.n_atten_head, + "atten_f_mix": self.atten_f_mix, + "atten_v_proj": self.use_atten_v_proj, + "atten_o_proj": self.use_atten_o_proj, + "sandwich_norm": self.sandwich_norm, + "full_attn_res": self.full_attn_res_mode, + "block_attn_res": self.block_attn_res_mode, + "s2_activation": self.s2_activation, + "ffn_so3_grid": self.ffn_so3_grid, + "node_wise_s2": self.node_wise_s2, + "node_wise_so3": self.node_wise_so3, + "message_node_s2": self.message_node_s2, + "message_node_so3": self.message_node_so3, + "lebedev_quadrature": self.lebedev_quadrature, + "activation_function": self.activation_function, + "glu_activation": self.glu_activation, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "mlp_bias": self.mlp_bias, + "exclude_types": self.exclude_types, + "eps": self.eps, + "trainable": self.trainable, + "seed": self.seed, + "inner_clamp_r_inner": self.inner_clamp_r_inner, + "inner_clamp_r_outer": self.inner_clamp_r_outer, + "add_chg_spin_ebd": self.add_chg_spin_ebd, + "default_chg_spin": self.default_chg_spin, + }, + "@variables": self._variables(), + "env_mat": EnvMat(self.rcut, self.rcut, self.eps).serialize(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> DescrptDPA4: + """Deserialize from a dict (accepts the pt ``serialize()`` output).""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "Descriptor": + raise ValueError(f"Invalid class for DescrptDPA4: {data_cls}") + type_val = data.pop("type") + if type_val not in ("SeZM", "sezm", "dpa4"): + raise ValueError(f"Invalid type for DescrptDPA4: {type_val}") + version = float(data.pop("@version")) + check_version_compatibility(version, cls.LATEST_VERSION, 1) + config = dict(data.pop("config")) + variables = data.pop("@variables") + data.pop("env_mat", None) + config.pop("s2_grid_resolution", None) + obj = cls(**config) + obj.version = version + obj._load_variables(variables) + return obj + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Update the selection and perform neighbor statistics.""" + local_jdata_cpy = local_jdata.copy() + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, + type_map, + local_jdata_cpy["rcut"], + local_jdata_cpy["sel"], + True, # mixed_type=True for unified sel + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py new file mode 100644 index 0000000000..2a0f5a8616 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-agnostic (array-API) building blocks for the DPA4/SeZM descriptor. + +This package is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn``. +""" diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/activation.py b/deepmd/dpmodel/descriptor/dpa4_nn/activation.py new file mode 100644 index 0000000000..23fbeada4a --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/activation.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Activation helper modules for DPA4/SeZM. + +This module is the dpmodel port of +``deepmd.pt.model.descriptor.sezm_nn.activation``. It contains the +coefficient-space nonlinear operators. Both pt classes are ported: +``GatedActivation`` (used by ``so2``, ``ffn``) and ``SwiGLU`` (used by +``grid_net``, which is consumed by ``ffn``). + +Serialization contract: ``GatedActivation`` mirrors the pt ``serialize()`` +format exactly (same config and ``@variables`` keys, including the nested +``gate_linear.weight``/``gate_linear.bias`` state-dict names), so pt +``serialize()`` output deserializes directly. ``SwiGLU`` is parameter-free in +pt (no ``serialize()``, no state-dict entries), so no serialization is +implemented for it. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + xp_sigmoid, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.network import ( + get_activation_fn, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .indexing import ( + build_m_major_l_index, + map_degree_idx, +) +from .so3 import ( + FocusLinear, +) + + +class GatedActivation(NativeOP): + """ + Gated activation for SO(3) equivariant features with per-l independent gates. + + Standard mode (gate=None in call): + - l=0: Uses the specified activation function + - l>0: Each degree l has an independent gate derived from the l=0 scalar + features. The gate for each l is expanded to all m components within + that l-block. + + GLU mode (gate provided in call, e.g., from split linear output): + - l=0: x0 * act(g0) (SwiGLU-style when act=silu, GeGLU when act=gelu, etc.) + - l>0: Uses gate's scalar (g0) to generate sigmoid gates for x's vector + components. This preserves SO(3) equivariance (scalar gates vector, + not vector gates vector). + + This module also supports the m-major reduced layout used inside SO(2) + blocks. If `mmax` is provided, the coefficient axis is assumed to follow + the truncated m-major order built by `build_m_major_index(lmax, mmax)`; + otherwise, it is assumed to be the full packed (l, m) layout with + D=(lmax+1)^2. + + Parameters + ---------- + lmax : int + Maximum spherical harmonic degree. + mmax : int | None + Maximum order (|m|) for the m-major reduced layout. If None, use the + full packed layout with D=(lmax+1)^2. + channels : int + Number of channels per focus stream. + n_focus : int + Number of focus streams. + precision : str + Internal compute precision used by the gate projection and sigmoid path. + activation_function : str + Activation function for l=0 components (e.g., "silu", "tanh", "gelu"). + mlp_bias : bool + Whether to use bias in the gate linear layer. + layout : str + Tensor layout convention. ``"nfdc"`` means input shape (N, F, D, C); + ``"ndfc"`` means input shape (N, D, F, C). + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for weight initialization. + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None = None, + channels: int, + n_focus: int = 1, + precision: str = DEFAULT_PRECISION, + activation_function: str = "silu", + mlp_bias: bool = False, + layout: str = "nfdc", + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + self.lmax = int(lmax) + self.mmax = None if mmax is None else int(mmax) + if self.mmax is not None: + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.channels = int(channels) + self.n_focus = int(n_focus) + self.precision = precision + self.mlp_bias = bool(mlp_bias) + self.layout = str(layout).lower() + if self.layout not in {"nfdc", "ndfc"}: + raise ValueError("`layout` must be either 'nfdc' or 'ndfc'") + self.trainable = bool(trainable) + self.activation_function = str(activation_function) + prec = PRECISION_DICT[self.precision.lower()] + + # === Build expand_index for mapping per-l gates to all m components === + if self.lmax > 0: + if self.mmax is None: + expand_index = map_degree_idx(self.lmax)[1:] - 1 + else: + degree_index = build_m_major_l_index(self.lmax, self.mmax) + expand_index = degree_index[1:] - 1 + self.gate_linear: FocusLinear | None = FocusLinear( + in_channels=self.channels, + out_channels=self.lmax * self.channels, + n_focus=self.n_focus, + precision=self.precision, + bias=self.mlp_bias, + seed=seed, + trainable=self.trainable, + ) + # pt re-initializes the gate weight with normal(0, 0.01) seeded + # by child_seed(seed, 1) and zeroes the bias (bias is already + # zero-initialized here). + rng = np.random.default_rng(child_seed(seed, 1)) + self.gate_linear.weight = rng.normal( + 0.0, 0.01, size=self.gate_linear.weight.shape + ).astype(prec) + else: + # pt uses nn.Identity() here (parameter-free, no state-dict keys); + # the dpmodel equivalent is no gate module at all. + expand_index = np.zeros((0,), dtype=np.int64) + self.gate_linear = None + self.expand_index = expand_index + + def call(self, x: Any, gate: Any = None) -> Any: + """ + Apply the gated activation. + + Parameters + ---------- + x : Array + Value features. Shape is (N, F, D, C) when ``layout='nfdc'``, + or (N, D, F, C) when ``layout='ndfc'``. + gate : Array | None + Optional gate features with the same layout as ``x``. + When provided, enables GLU mode: + - l=0: x0 * act(g0) (e.g., SwiGLU when act=silu) + - l>0: sigmoid(Linear(g0)) gates x's vector components + When None (default), uses standard mode where gates are derived + from x itself. + + Returns + ------- + Array + Gated features with the same layout as ``x``. + """ + xp = array_api_compat.array_namespace(x) + degree_axis = 1 if self.layout == "ndfc" else 2 + + gate_source = x if gate is None else gate + if degree_axis == 1: + gate_scalar_source = gate_source[:, 0, :, :] # (N, F, C) + g0 = gate_source[:, :1, :, :] + x0_in = x[:, :1, :, :] + else: + gate_scalar_source = gate_source[:, :, 0, :] # (N, F, C) + g0 = gate_source[:, :, :1, :] + x0_in = x[:, :, :1, :] + + scalar_act = get_activation_fn(self.activation_function) + if gate is not None: + x0 = x0_in * scalar_act(g0) + else: + x0 = scalar_act(x0_in) + + if self.lmax == 0: + return x0 + + gate_weight = xp.asarray( + self.gate_linear.weight[...], device=array_api_compat.device(x) + ) + input_dtype = gate_scalar_source.dtype + if input_dtype != gate_weight.dtype: + gate_scalar_source = xp.astype(gate_scalar_source, gate_weight.dtype) + gating_scalars = xp_sigmoid(self.gate_linear.call(gate_scalar_source)) + if gating_scalars.dtype != input_dtype: + gating_scalars = xp.astype(gating_scalars, input_dtype) + gating_scalars = xp.reshape( + gating_scalars, + (x.shape[0], gate_scalar_source.shape[1], self.lmax, self.channels), + ) + expand_index = xp.asarray(self.expand_index, device=array_api_compat.device(x)) + gates = xp.take(gating_scalars, expand_index, axis=2) # (N, F, D-1, C) + if self.layout == "ndfc": + gates = xp.permute_dims(gates, (0, 2, 1, 3)) # (N, D-1, F, C) + xt = x[:, 1:, :, :] * gates + else: + xt = x[:, :, 1:, :] * gates + return xp.concat([x0, xt], axis=degree_axis) + + def serialize(self) -> dict[str, Any]: + """Serialize the GatedActivation to a dict (pt-compatible format).""" + variables = {"expand_index": to_numpy_array(self.expand_index)} + if self.gate_linear is not None: + variables["gate_linear.weight"] = to_numpy_array(self.gate_linear.weight) + if self.mlp_bias: + variables["gate_linear.bias"] = to_numpy_array(self.gate_linear.bias) + return { + "@class": "GatedActivation", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.mmax, + "channels": self.channels, + "n_focus": self.n_focus, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "activation_function": self.activation_function, + "mlp_bias": self.mlp_bias, + "layout": self.layout, + "trainable": self.trainable, + "seed": None, + }, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> GatedActivation: + """Deserialize a GatedActivation from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "GatedActivation": + raise ValueError(f"Invalid class for GatedActivation: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + mmax = config["mmax"] + obj = cls( + lmax=int(config["lmax"]), + mmax=None if mmax is None else int(mmax), + channels=int(config["channels"]), + n_focus=int(config["n_focus"]), + precision=str(config["precision"]), + activation_function=str(config["activation_function"]), + mlp_bias=bool(config["mlp_bias"]), + layout=str(config["layout"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + prec = PRECISION_DICT[obj.precision.lower()] + expand_index = np.asarray(variables["expand_index"], dtype=np.int64) + if not np.array_equal(expand_index, obj.expand_index): + raise ValueError("expand_index does not match the lmax/mmax tables") + if obj.gate_linear is not None: + weight = np.asarray(variables["gate_linear.weight"], dtype=prec) + if weight.shape != obj.gate_linear.weight.shape: + raise ValueError( + f"gate_linear.weight shape {weight.shape} does not match " + f"the expected shape {obj.gate_linear.weight.shape}" + ) + obj.gate_linear.weight = weight + if obj.mlp_bias: + obj.gate_linear.bias = np.asarray( + variables["gate_linear.bias"], dtype=prec + ).reshape(obj.gate_linear.bias.shape) + return obj + + +class SwiGLU(NativeOP): + """Point-wise SwiGLU on the last feature axis. + + Parameter-free, matching the pt version (which defines no ``serialize()`` + and contributes no state-dict entries). + """ + + def call(self, inputs: Any) -> Any: + """ + Apply point-wise SwiGLU. + + Parameters + ---------- + inputs : Array + Input array with shape ``(..., 2*C)``; the first half of the last + axis is the gate, the second half the value. + + Returns + ------- + Array + Gated array with shape ``(..., C)``. + """ + # torch.chunk(inputs, 2, dim=-1): first chunk gets ceil(C/2) entries + nc = (inputs.shape[-1] + 1) // 2 + gate = inputs[..., :nc] + value = inputs[..., nc:] + return gate * xp_sigmoid(gate) * value diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py new file mode 100644 index 0000000000..c27b6174b8 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Attention utilities for DPA4/SeZM message passing. + +This module is the dpmodel port of +``deepmd.pt.model.descriptor.sezm_nn.attention``. It implements the +destination-wise envelope-gated softmax used by the SO(2) attention path. + +Padded-edge adaptation +---------------------- +The pt version consumes a sparse edge list and reduces per destination node +with ``scatter_reduce(amax)`` / ``scatter_add`` keyed by ``dst``. In the +dpmodel padded layout (see ``edge_cache.EdgeCache``) the edge axis is +``E = n_nodes * nnei`` with slot ``(i, j)`` belonging to node ``i``, so every +destination-wise reduction becomes a plain reduction over the ``nnei`` axis +after a ``(n_nodes, nnei, ...)`` reshape, and invalid (padded) slots are +removed by folding ``edge_mask`` into the non-negative per-edge weight. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.utils.network import ( + softplus_t, +) + + +def segment_envelope_gated_softmax( + logits: Any, + edge_env: Any, + n_nodes: int, + z_bias_raw: Any, + eps: float, + src_weight: Any = None, + edge_mask: Any = None, +) -> Any: + """ + Compute destination-wise envelope-gated softmax attention. + + All array arguments must live in the same array namespace. + + Parameters + ---------- + logits + Attention logits with shape (E, F, H), padded-edge layout with + ``E = n_nodes * nnei``. + edge_env + Cutoff envelope weights with shape (E, 1) or (E,). + n_nodes + Number of nodes. The pt ``dst`` argument is dropped: in the padded + layout the destination of edge slot ``(i, j)`` is implicitly node + ``i``. + z_bias_raw + Unconstrained denominator bias with shape (F, H). + Softplus is applied to keep the bias strictly positive. + eps + Small epsilon for denominator stability. + src_weight + Optional per-edge source-side multiplier with shape (E, 1) or + (E,). When provided the per-edge weight becomes + ``edge_env**2 * src_weight`` and the attention reduces to + ``edge_env**2 * src_weight * exp(logits) / + (zeta + sum(edge_env**2 * src_weight * exp(logits)))``. + ``src_weight = 0`` therefore removes the source from both the + numerator and the denominator, which is what SFPG needs so that + a muted source does not even leak through the softmax + normalization. + edge_mask + Optional padded-edge validity mask with shape (E,) or (E, 1); + zero marks invalid slots. Folded into the non-negative per-edge + weight, so invalid slots drop out of the group max, the numerator, + and the denominator exactly like absent edges in the pt sparse + layout. + + Returns + ------- + Array + Normalized edge weights with shape (E, F, H). Zero on invalid slots. + """ + xp = array_api_compat.array_namespace(logits) + n_edge, n_focus, n_head = logits.shape + n_channel = n_focus * n_head + eps_f = float(eps) + if n_nodes <= 0 or n_edge % int(n_nodes) != 0: + raise ValueError( + "padded-edge layout requires E to be a multiple of n_nodes; " + f"got E={n_edge}, n_nodes={n_nodes}" + ) + nnei = n_edge // int(n_nodes) + device = array_api_compat.device(logits) + + # === Step 1. Flatten (F, H) and build the effective per-edge weight === + logits_2d = xp.reshape(logits, (n_edge, n_channel)) + zeros_e = xp.zeros((n_edge,), dtype=logits.dtype, device=device) + edge_env_1d = xp.astype(xp.reshape(edge_env, (n_edge,)), logits.dtype) + edge_env_1d = xp.where(edge_env_1d > 0.0, edge_env_1d, zeros_e) + # edge_weight_sq acts as the non-negative multiplier applied to every + # ``exp(logit)`` term. Folding ``src_weight`` (and, in the padded + # layout, ``edge_mask``) here guarantees that any edge with zero weight + # is excluded from the group max, the numerator, and the denominator in + # a single pass. + edge_weight_sq = edge_env_1d * edge_env_1d + if src_weight is not None: + src_weight_1d = xp.astype(xp.reshape(src_weight, (n_edge,)), logits.dtype) + src_weight_1d = xp.where(src_weight_1d > 0.0, src_weight_1d, zeros_e) + edge_weight_sq = edge_weight_sq * src_weight_1d + if edge_mask is not None: + mask_1d = xp.astype(xp.reshape(edge_mask, (n_edge,)), logits.dtype) + edge_weight_sq = edge_weight_sq * mask_1d + zeta = xp.astype(xp.reshape(softplus_t(z_bias_raw), (1, n_channel)), logits.dtype) + has_weight = edge_weight_sq > 0.0 + minus_inf = xp.full( + (n_edge, n_channel), + float("-inf"), + dtype=logits.dtype, + device=device, + ) + logits_for_max = xp.where( + has_weight[:, None], + logits_2d, + minus_inf, + ) + + # === Step 2. Destination-wise max for stable exponentials === + # pt: scatter_reduce(amax) over dst — padded-edge max over the nnei axis. + group_max = xp.max( + xp.reshape(logits_for_max, (n_nodes, nnei, n_channel)), axis=1 + ) # (N, n_channel) + edge_max = xp.reshape( + xp.broadcast_to(group_max[:, None, :], (n_nodes, nnei, n_channel)), + (n_edge, n_channel), + ) + zeros_en = xp.zeros((n_edge, n_channel), dtype=logits.dtype, device=device) + zeros_nn = xp.zeros((n_nodes, n_channel), dtype=logits.dtype, device=device) + edge_max = xp.where(xp.isfinite(edge_max), edge_max, zeros_en) + group_max_safe = xp.where(xp.isfinite(group_max), group_max, zeros_nn) + + # === Step 3. Envelope/SFPG-gated exponential terms === + exp_shifted = xp.exp(logits_2d - edge_max) + edge_weighted_exp = edge_weight_sq[:, None] * exp_shifted + + # === Step 4. Destination-wise normalization with positive denominator bias === + # pt: scatter_add over dst — padded-edge masked sum over the nnei axis + # (invalid slots already carry zero weight). + denom_sum = xp.sum( + xp.reshape(edge_weighted_exp, (n_nodes, nnei, n_channel)), axis=1 + ) # (N, n_channel) + denom = denom_sum + zeta * xp.exp(-group_max_safe) + + denom_edge = xp.reshape( + xp.broadcast_to(denom[:, None, :], (n_nodes, nnei, n_channel)), + (n_edge, n_channel), + ) + alpha = edge_weighted_exp / (denom_edge + eps_f) + return xp.reshape(alpha, (n_edge, n_focus, n_head)) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/block.py b/deepmd/dpmodel/descriptor/dpa4_nn/block.py new file mode 100644 index 0000000000..ce529c3cd2 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/block.py @@ -0,0 +1,606 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Interaction blocks for DPA4/SeZM. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.block``. +It defines the SeZM interaction block that combines SO(2) message passing and +equivariant feed-forward subblocks with residual shortcuts. + +Branches guarded with ``NotImplementedError`` at this level (flags consumed by +block.py itself, all unused by the core DPA4 config): + +- ``full_attn_res != "none"`` / ``block_attn_res != "none"`` — the pt block + builds ``DepthAttnRes`` aggregators and switches the forward implementation + (pt block.py:514/541); only the baseline residual-shortcut path + (pt block.py:756) is ported. +- ``layer_scale=True`` — the pt block builds per-channel + ``adam_ffn_layer_scales`` on the FFN residual branches (pt block.py:500) + in addition to the SO(2)-internal scales; not ported. + +Flags merely forwarded to sub-components keep their guards there (delegated, +not duplicated here): ``so2_attn_res``, ``so2_s2_activation``, +``node_wise_s2/so3``, ``message_node_s2/so3``, ``atten_f_mix``, +``atten_v_proj``, ``atten_o_proj`` (raised by ``SO2Convolution``) and +``ffn_so3_grid``, ``ffn_grid_mlp`` with the grid path active (raised by +``EquivariantFFN`` / ``S2GridNet``). + +The pt eval-time activation-checkpoint / nvtx instrumentation +(``DP_ACT_INFER``, ``DP_COMPILE_INFER``, ``nvtx_range``) is pt-runtime-only +and intentionally not ported. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .ffn import ( + EquivariantFFN, +) +from .norm import ( + EquivariantRMSNorm, +) +from .so2 import ( + SO2Convolution, + _compute_precision, +) +from .utils import ( + ATTN_RES_MODES, +) + +if TYPE_CHECKING: + from .edge_cache import ( + EdgeCache, + ) + + +class SeZMInteractionBlock(NativeOP): + """ + SeZM interaction block with SO(2) message passing and equivariant FFN stack. + + Branch order: + 1. SO(2) branch: optional pre-norm -> `SO2Convolution` -> optional post-norm. + 2. FFN branch: repeated subblocks of + optional pre-norm -> `EquivariantFFN` -> optional post-norm. + + Outer residual shortcuts are applied around the SO(2) unit and each FFN + subblock (the pt AttnRes paths are not ported; see the module docstring). + + `SO2Convolution` internally handles the real multi-focus expansion, so this + block keeps a singleton-focus backbone layout `(N, D, 1, C)` at boundaries. + + Parameters mirror the pt ``SeZMInteractionBlock`` (pt block.py:227) with + ``precision`` replacing ``dtype``; see the pt docstring for the full + per-parameter description. + """ + + def __init__( + self, + *, + lmax: int, + node_lmax: int | None = None, + mmax: int | None = None, + kmax: int = 1, + channels: int, + n_focus: int = 1, + focus_dim: int = 0, + focus_compete: bool = True, + so2_norm: bool = False, + so2_layers: int = 4, + so2_attn_res: str = "none", + radial_so2_mode: str = "none", + radial_so2_rank: int = 0, + n_atten_head: int = 1, + atten_f_mix: bool = False, + atten_v_proj: bool = False, + atten_o_proj: bool = False, + so2_pre_norm: bool = True, + so2_post_norm: bool = False, + ffn_pre_norm: bool = True, + ffn_post_norm: bool = False, + ffn_neurons: int = 96, + node_wise_grid_mlp: bool = False, + node_wise_grid_branch: int = 0, + message_node_grid_mlp: bool = False, + message_node_grid_branch: int = 0, + ffn_grid_mlp: bool = False, + ffn_grid_branch: int = 0, + ffn_blocks: int = 1, + layer_scale: bool = False, + full_attn_res: str = "none", + block_attn_res: str = "none", + so2_s2_activation: bool = False, + node_wise_s2: bool = False, + node_wise_so3: bool = False, + message_node_s2: bool = False, + message_node_so3: bool = False, + ffn_s2_activation: bool = False, + ffn_so3_grid: bool = False, + so2_lebedev_quadrature: bool = False, + ffn_lebedev_quadrature: bool = False, + so2_activation_function: str = "silu", + ffn_activation_function: str, + ffn_glu_activation: bool = True, + mlp_bias: bool = False, + eps: float = 1e-7, + precision: str = DEFAULT_PRECISION, + seed: int | list[int] | None = None, + trainable: bool = True, + ) -> None: + self.lmax = int(lmax) + self.node_lmax = self.lmax if node_lmax is None else int(node_lmax) + if self.node_lmax < self.lmax: + raise ValueError("`node_lmax` must be >= `lmax`") + self.mp_ebed_dim = (self.lmax + 1) ** 2 + self.node_ebed_dim = (self.node_lmax + 1) ** 2 + self.mmax = int(self.lmax if mmax is None else mmax) + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.kmax = int(kmax) + if self.kmax < 0: + raise ValueError("`kmax` must be non-negative") + self.channels = int(channels) + self.n_focus = int(n_focus) + if self.n_focus < 1: + raise ValueError("`n_focus` must be >= 1") + self.focus_dim = int(focus_dim) + if self.focus_dim < 0: + raise ValueError("`focus_dim` must be >= 0") + self.focus_compete = bool(focus_compete) + self.so2_norm = bool(so2_norm) + self.so2_layers = int(so2_layers) + self.so2_attn_res_mode = str(so2_attn_res).lower() + if self.so2_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`so2_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + self.radial_so2_mode = str(radial_so2_mode).lower() + self.radial_so2_rank = int(radial_so2_rank) + self.n_atten_head = int(n_atten_head) + self.atten_f_mix = bool(atten_f_mix) + self.use_atten_v_proj = bool(atten_v_proj) + self.use_atten_o_proj = bool(atten_o_proj) + self.so2_pre_norm = bool(so2_pre_norm) + self.so2_post_norm = bool(so2_post_norm) + self.ffn_pre_norm = bool(ffn_pre_norm) + self.ffn_post_norm = bool(ffn_post_norm) + self.ffn_neurons = int(ffn_neurons) + self.node_wise_grid_mlp = bool(node_wise_grid_mlp) + self.node_wise_grid_branch = int(node_wise_grid_branch) + self.message_node_grid_mlp = bool(message_node_grid_mlp) + self.message_node_grid_branch = int(message_node_grid_branch) + self.ffn_grid_mlp = bool(ffn_grid_mlp) + self.ffn_grid_branch = int(ffn_grid_branch) + if ( + min( + self.node_wise_grid_branch, + self.message_node_grid_branch, + self.ffn_grid_branch, + ) + < 0 + ): + raise ValueError("grid branch counts must be non-negative") + self.ffn_blocks = int(ffn_blocks) + if self.ffn_blocks < 1: + raise ValueError("`ffn_blocks` must be >= 1") + self.layer_scale = bool(layer_scale) + if self.layer_scale: + # consumed by block.py itself (FFN-branch adam_ffn_layer_scales) + raise NotImplementedError("layer_scale=True is not ported to dpmodel") + self.full_attn_res_mode = str(full_attn_res).lower() + if self.full_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`full_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + self.block_attn_res_mode = str(block_attn_res).lower() + if self.block_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`block_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + if self.full_attn_res_mode != "none": + raise NotImplementedError( + "full_attn_res != 'none' (DepthAttnRes) is not ported to dpmodel" + ) + if self.block_attn_res_mode != "none": + raise NotImplementedError( + "block_attn_res != 'none' (DepthAttnRes) is not ported to dpmodel" + ) + self.so2_s2_activation = bool(so2_s2_activation) + self.node_wise_s2 = bool(node_wise_s2) + self.node_wise_so3 = bool(node_wise_so3) + self.message_node_s2 = bool(message_node_s2) + self.message_node_so3 = bool(message_node_so3) + self.ffn_s2_activation = bool(ffn_s2_activation) + self.ffn_so3_grid = bool(ffn_so3_grid) + self.so2_lebedev_quadrature = bool(so2_lebedev_quadrature) + self.ffn_lebedev_quadrature = bool(ffn_lebedev_quadrature) + self.so2_activation_function = str(so2_activation_function) + self.ffn_activation_function = str(ffn_activation_function) + self.ffn_glu_activation = bool(ffn_glu_activation) + self.mlp_bias = bool(mlp_bias) + self.eps = float(eps) + self.precision = precision + self.compute_precision = _compute_precision(precision) + self.trainable = bool(trainable) + + # === Step 0. Split deterministic seeds at the block top-level === + # pt also splits seed_full_attn / seed_block_attn (block.py:378-379); + # those consumers are guarded above, so the splits are unused here. + seed_so2_conv = child_seed(seed, 0) + seed_ffn = child_seed(seed, 1) + + # === Step 1. SO(2) convolution branch norms === + # pt uses nn.Identity() for disabled norms (parameter-free); the + # dpmodel equivalent is None. + self.pre_so2_norm: EquivariantRMSNorm | None = ( + EquivariantRMSNorm( + self.lmax, + self.channels, + n_focus=1, + precision=self.compute_precision, + trainable=self.trainable, + ) + if self.so2_pre_norm + else None + ) + self.post_so2_norm: EquivariantRMSNorm | None = ( + EquivariantRMSNorm( + self.lmax, + self.channels, + n_focus=1, + precision=self.compute_precision, + trainable=self.trainable, + ) + if self.so2_post_norm + else None + ) + + self.so2_conv = SO2Convolution( + lmax=self.lmax, + mmax=self.mmax, + kmax=self.kmax, + channels=self.channels, + n_focus=self.n_focus, + focus_dim=self.focus_dim, + focus_compete=self.focus_compete, + so2_norm=self.so2_norm, + so2_layers=self.so2_layers, + so2_attn_res=self.so2_attn_res_mode, + radial_so2_mode=self.radial_so2_mode, + radial_so2_rank=self.radial_so2_rank, + layer_scale=self.layer_scale, + n_atten_head=self.n_atten_head, + atten_f_mix=self.atten_f_mix, + atten_v_proj=self.use_atten_v_proj, + atten_o_proj=self.use_atten_o_proj, + s2_activation=self.so2_s2_activation, + node_wise_grid_mlp=self.node_wise_grid_mlp, + node_wise_grid_branch=self.node_wise_grid_branch, + message_node_grid_mlp=self.message_node_grid_mlp, + message_node_grid_branch=self.message_node_grid_branch, + node_wise_s2=self.node_wise_s2, + node_wise_so3=self.node_wise_so3, + message_node_s2=self.message_node_s2, + message_node_so3=self.message_node_so3, + lebedev_quadrature=self.so2_lebedev_quadrature, + activation_function=self.so2_activation_function, + mlp_bias=self.mlp_bias, + eps=self.eps, + precision=self.precision, + seed=seed_so2_conv, + trainable=self.trainable, + ) + + # === Step 2. FFN subblock sequence === + pre_ffn_norms: list[EquivariantRMSNorm | None] = [] + post_ffn_norms: list[EquivariantRMSNorm | None] = [] + ffns: list[EquivariantFFN] = [] + + for i in range(self.ffn_blocks): + seed_ffn_i = child_seed(seed_ffn, i) + pre_ffn_norms.append( + EquivariantRMSNorm( + self.node_lmax, + self.channels, + n_focus=1, + precision=self.compute_precision, + trainable=self.trainable, + ) + if self.ffn_pre_norm + else None + ) + post_ffn_norms.append( + EquivariantRMSNorm( + self.node_lmax, + self.channels, + n_focus=1, + precision=self.compute_precision, + trainable=self.trainable, + ) + if self.ffn_post_norm + else None + ) + ffns.append( + EquivariantFFN( + lmax=self.node_lmax, + channels=self.channels, + hidden_channels=self.ffn_neurons, + kmax=self.kmax, + grid_mlp=self.ffn_grid_mlp, + grid_branch=self.ffn_grid_branch, + s2_activation=self.ffn_s2_activation, + ffn_so3_grid=self.ffn_so3_grid, + lebedev_quadrature=self.ffn_lebedev_quadrature, + activation_function=self.ffn_activation_function, + glu_activation=self.ffn_glu_activation, + mlp_bias=self.mlp_bias, + precision=self.precision, + trainable=self.trainable, + seed=seed_ffn_i, + ) + ) + self.pre_ffn_norms = pre_ffn_norms + self.post_ffn_norms = post_ffn_norms + self.ffns = ffns + + def _run_so2_unit( + self, + x: Any, + edge_cache: EdgeCache, + radial_feat: Any, + ) -> Any: + """ + Run the SO(2) unit without an outer block-level residual shortcut. + + Parameters + ---------- + x + Canonical node features with shape `(N, D, 1, C)`. + edge_cache + Edge cache (padded layout; see ``edge_cache.EdgeCache``). + radial_feat + Per-edge radial features with shape (E, lmax+1, C). + + Returns + ------- + Array + SO(2) unit output with shape `(N, D, 1, C)`. + """ + xp = array_api_compat.array_namespace(x) + n_node = x.shape[0] + channels = self.channels + use_full_node = self.node_lmax == self.lmax + x_so2 = x if use_full_node else x[:, : self.mp_ebed_dim, :, :] + x_pre = x_so2 if self.pre_so2_norm is None else self.pre_so2_norm(x_so2) + so2_unit_output = self.so2_conv( + xp.reshape(x_pre, (n_node, x_so2.shape[1], channels)), + edge_cache, + radial_feat, + ) + so2_unit_output = so2_unit_output[:, :, None, :] + if self.post_so2_norm is not None: + so2_unit_output = self.post_so2_norm(so2_unit_output) + if use_full_node: + return so2_unit_output + # zero-pad the degrees above lmax (pt writes into x.new_zeros) + pad = xp.zeros( + (n_node, self.node_ebed_dim - self.mp_ebed_dim, 1, channels), + dtype=x.dtype, + device=array_api_compat.device(x), + ) + return xp.concat([so2_unit_output, pad], axis=1) + + def _run_ffn_unit(self, x: Any, unit_idx: int) -> Any: + """ + Run one FFN subblock without the outer unit-level residual shortcut. + + Parameters + ---------- + x + Canonical node features with shape `(N, D, 1, C)`. + unit_idx + FFN subblock index. + + Returns + ------- + Array + FFN unit output with shape `(N, D, 1, C)`. + """ + pre_norm = self.pre_ffn_norms[unit_idx] + post_norm = self.post_ffn_norms[unit_idx] + x_pre = x if pre_norm is None else pre_norm(x) + y = self.ffns[unit_idx](x_pre) + if post_norm is not None: + y = post_norm(y) + return y + + def call( + self, + x: Any, + edge_cache: EdgeCache, + radial_feat: Any, + unit_history: list[Any] | None = None, + ) -> tuple[Any, None, None, None]: + """ + Run the residual-connected block path (pt baseline path). + + Parameters + ---------- + x + Features with shape `(N, D, 1, C)`. + edge_cache + Edge cache (padded layout). + radial_feat + Per-edge radial features with shape (E, lmax+1, C). + unit_history + Unused in the residual-connected path (the pt AttnRes paths that + consume it are not ported). + + Returns + ------- + tuple[Array, None, None, None] + Tuple `(block_output, None, None, None)` matching the pt + baseline-path return convention. + """ + so2_unit_output = self._run_so2_unit(x, edge_cache, radial_feat) + ffn_state = x + so2_unit_output + for i in range(self.ffn_blocks): + ffn_state = ffn_state + self._run_ffn_unit(ffn_state, i) + return ffn_state, None, None, None + + def _sub_modules(self) -> list[tuple[str, NativeOP | None]]: + """Sub-modules with their pt module names (None = pt nn.Identity).""" + subs: list[tuple[str, NativeOP | None]] = [ + ("pre_so2_norm", self.pre_so2_norm), + ("post_so2_norm", self.post_so2_norm), + ("so2_conv", self.so2_conv), + ] + for i in range(self.ffn_blocks): + subs.append((f"pre_ffn_norms.{i}", self.pre_ffn_norms[i])) + subs.append((f"post_ffn_norms.{i}", self.post_ffn_norms[i])) + subs.append((f"ffns.{i}", self.ffns[i])) + return subs + + def _variables(self) -> dict[str, Any]: + """Variables keyed by the pt ``state_dict`` key names.""" + variables: dict[str, Any] = {} + for prefix, sub in self._sub_modules(): + if sub is None: + continue + if isinstance(sub, SO2Convolution): + sub_vars = sub._variables() + else: + sub_vars = sub.serialize()["@variables"] + for key, value in sub_vars.items(): + variables[f"{prefix}.{key}"] = value + return variables + + def _load_variables(self, variables: dict[str, Any]) -> None: + """Load variables keyed by the pt ``state_dict`` key names.""" + variables = dict(variables) + for name, sub in self._sub_modules(): + if sub is None: + continue + full = f"{name}." + sv = { + key[len(full) :]: value + for key, value in variables.items() + if key.startswith(full) + } + for key in list(variables): + if key.startswith(full): + del variables[key] + if not sv: + raise KeyError(f"Missing variables with prefix: {full}") + if isinstance(sub, SO2Convolution): + sub._load_variables(sv) + elif isinstance(sub, EquivariantFFN): + sub._load_variables(sv) + else: + # norms: rebuild through the shape-checking deserialize + data = sub.serialize() + data["@variables"] = sv + new_sub = type(sub).deserialize(data) + attr, _, idx = name.partition(".") + if idx: + getattr(self, attr)[int(idx)] = new_sub + else: + setattr(self, attr, new_sub) + if variables: + raise KeyError(f"Unknown variables: {sorted(variables)}") + + def serialize(self) -> dict[str, Any]: + """Serialize the SeZMInteractionBlock to a dict (pt-compatible format).""" + return { + "@class": "SeZMInteractionBlock", + "@version": 1, + "config": { + "lmax": self.lmax, + "node_lmax": self.node_lmax, + "mmax": self.mmax, + "kmax": self.kmax, + "channels": self.channels, + "n_focus": self.n_focus, + "focus_dim": self.focus_dim, + "focus_compete": self.focus_compete, + "so2_norm": self.so2_norm, + "so2_layers": self.so2_layers, + "so2_attn_res": self.so2_attn_res_mode, + "radial_so2_mode": self.radial_so2_mode, + "radial_so2_rank": self.radial_so2_rank, + "n_atten_head": self.n_atten_head, + "atten_f_mix": self.atten_f_mix, + "atten_v_proj": self.use_atten_v_proj, + "atten_o_proj": self.use_atten_o_proj, + "so2_pre_norm": self.so2_pre_norm, + "so2_post_norm": self.so2_post_norm, + "ffn_pre_norm": self.ffn_pre_norm, + "ffn_post_norm": self.ffn_post_norm, + "ffn_neurons": self.ffn_neurons, + "node_wise_grid_mlp": self.node_wise_grid_mlp, + "node_wise_grid_branch": self.node_wise_grid_branch, + "message_node_grid_mlp": self.message_node_grid_mlp, + "message_node_grid_branch": self.message_node_grid_branch, + "ffn_grid_mlp": self.ffn_grid_mlp, + "ffn_grid_branch": self.ffn_grid_branch, + "ffn_blocks": self.ffn_blocks, + "full_attn_res": self.full_attn_res_mode, + "block_attn_res": self.block_attn_res_mode, + "so2_s2_activation": self.so2_s2_activation, + "node_wise_s2": self.node_wise_s2, + "node_wise_so3": self.node_wise_so3, + "message_node_s2": self.message_node_s2, + "message_node_so3": self.message_node_so3, + "ffn_s2_activation": self.ffn_s2_activation, + "ffn_so3_grid": self.ffn_so3_grid, + "so2_lebedev_quadrature": self.so2_lebedev_quadrature, + "ffn_lebedev_quadrature": self.ffn_lebedev_quadrature, + "so2_activation_function": self.so2_activation_function, + "ffn_activation_function": self.ffn_activation_function, + "ffn_glu_activation": self.ffn_glu_activation, + "mlp_bias": self.mlp_bias, + "layer_scale": self.layer_scale, + "eps": self.eps, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "seed": None, + }, + "@variables": self._variables(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SeZMInteractionBlock: + """Deserialize a SeZMInteractionBlock from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "SeZMInteractionBlock": + raise ValueError(f"Invalid class for SeZMInteractionBlock: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = dict(data.pop("config")) + variables = data.pop("@variables") + config["precision"] = str(config.pop("precision")) + obj = cls(**config) + obj._load_variables(variables) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py new file mode 100644 index 0000000000..e2afba1670 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Edge cache construction for the dpmodel DPA4/SeZM descriptor. + +This module defines the :class:`EdgeCache` dataclass (the dpmodel +counterpart of the pt ``EdgeFeatureCache`` NamedTuple from +``deepmd.pt.model.descriptor.sezm_nn.edge_cache``) and +:func:`build_edge_cache`, the padded-layout counterpart of pt's +sparse ``build_edge_cache``. + +Padded-edge layout +------------------ +The pt implementation extracts a *sparse* edge list with ``torch.nonzero``: +only valid neighbor slots become edges, and per-edge tensors have a +data-dependent length ``E``. The dpmodel implementation instead uses a +*padded* and frame-explicit edge layout: every neighbor slot of the DeePMD +neighbor list contributes one edge, so + + ``E = nf * nloc * nnei`` + +with per-edge tensors flattened from ``(nf, nloc, nnei, ...)`` in row-major +order. Invalid slots (``nlist == -1`` padding, excluded type pairs) stay in +the arrays and are marked by ``edge_mask == 0``. Edge slot ``(f, i, j)`` +always belongs to destination node ``f * nloc + i``, so destination +aggregation is a masked sum over the ``nnei`` axis instead of a scatter. +""" + +from __future__ import ( + annotations, +) + +import math +from dataclasses import ( + dataclass, + field, +) +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from .utils import ( + safe_norm, +) +from .wignerd import ( + build_edge_quaternion, + quaternion_multiply, + quaternion_z_rotation, +) + + +@dataclass +class EdgeCache: + """ + Global edge feature cache created once per forward(). + + All per-edge arrays are aligned on the same padded edge axis + (``E = nf * nloc * nnei``); see the module docstring for the layout + contract. Node-level arrays use the local node axis ``N = nf * nloc``. + + An ``EdgeCache`` must not be reused across forward passes: + ``D_to_m_cache``/``Dt_from_m_cache`` are keyed only by ``"lmax:mmax"``, + not by the contents of ``D_full``, so reuse with different Wigner blocks + would silently return stale projections. + + Parameters + ---------- + src + Source (neighbor) node indices with shape (E,), pointing into the + local node axis ``N = nf * nloc``. Invalid slots must hold a safe + in-range index (their contribution is masked out by ``edge_mask``). + dst + Destination (center) node indices with shape (E,). In the padded + layout this is slot-implicit and MUST equal + ``arange(nf * nloc)`` with each index repeated ``nnei`` consecutive + times (i.e. ``np.repeat(np.arange(nf * nloc), nnei)``; + node-contiguous order); aggregation code relies on this ordering. + edge_type_feat + Per-edge type embeddings with shape (E, C), computed as src+dst. + edge_vec + Edge vectors with shape (E, 3) in Å. + edge_rbf + Radial basis with shape (E, n_radial). + The C^3 cutoff envelope is already baked in. + edge_env + C^3 cutoff envelope weights with shape (E, 1). Zero on invalid slots. + deg + Envelope-squared smooth degree with shape (N,), computed as the + masked ``sum(edge_env**2)`` over each node's ``nnei`` slots. + inv_sqrt_deg + Inverse square root smooth degree normalization with shape (N, 1, 1). + D_full + Block-diagonal Wigner-D matrix with shape (E, D, D) where D=(lmax+1)^2. + Used for efficient batched rotation. None if not available. + Dt_full + Transpose of D_full with shape (E, D, D). None if not available. + D_to_m_cache + Lazy cache for projected D matrices keyed by a normalized + ``"lmax:mmax"`` identifier. The key does not capture the contents + of ``D_full``, so the cache is only valid for the forward pass + that created this ``EdgeCache`` (see the class docstring). + Dt_from_m_cache + Lazy cache for projected Dt matrices keyed by a normalized + ``"lmax:mmax"`` identifier. Same single-forward-pass validity + caveat as ``D_to_m_cache``. + edge_src_gate + Optional per-edge Source Freeze Propagation Gate (SFPG) weight with + shape (E, 1). Present only in bridging mode; ``None`` otherwise. + edge_quat + Per-edge global-to-local quaternion used to build ``D_full`` and + ``Dt_full`` with shape (E, 4). None if not available. + edge_mask + Validity mask for the padded-edge layout with shape (E,) or (E, 1); + nonzero (1) marks a real edge, zero marks a padded/invalid slot. + ``None`` means all slots are valid. This field has no pt counterpart: + pt's sparse edge list contains valid edges only, while dpmodel keeps + the padded ``nf * nloc * nnei`` slots and masks the invalid ones. + """ + + src: Any + dst: Any + edge_type_feat: Any + edge_vec: Any + edge_rbf: Any + edge_env: Any + deg: Any + inv_sqrt_deg: Any + D_full: Any = None + Dt_full: Any = None + D_to_m_cache: dict[str, Any] = field(default_factory=dict) + Dt_from_m_cache: dict[str, Any] = field(default_factory=dict) + edge_src_gate: Any = None + edge_quat: Any = None + edge_mask: Any = None + + +def _build_edge_mask_and_src( + xp: Any, + nlist: Any, + mapping: Any, + pair_keep_mask: Any, + nall: int, +) -> tuple[Any, Any, Any]: + """ + Build the padded edge validity mask and safe source-local indices. + + Mirrors the pt edge-keep semantics of + ``sezm_nn.edge_cache._build_standard_edge_index`` exactly: + + - padding slots (``nlist == -1``) are invalid; + - excluded type pairs (``pair_keep_mask == False``) are invalid; + - after mapping the neighbor's extended index to a local index, slots + whose source falls outside ``[0, nloc)`` are invalid (pt's ``src_ok`` + filter; e.g. broken mapping or ghost-only neighbors); + - no distance-based filtering: edges beyond ``rcut`` stay valid and are + zeroed naturally by the smooth envelope. + + Instead of dropping invalid slots (pt's ``torch.nonzero``), they are kept + with ``mask == False`` and safe (index 0) placeholder indices. + + Parameters + ---------- + xp + Array namespace. + nlist + Neighbor list with shape (nf, nloc, nnei); -1 marks padding. + mapping + Extended-to-local mapping with shape (nf, nall), or None if the + neighbor indices are already local. + pair_keep_mask + Pair exclusion keep mask with shape (nf, nloc, nnei). True means keep. + nall + Number of atoms on the extended axis per frame. + + Returns + ------- + tuple[Any, Any, Any] + ``(mask, nlist_safe, src_local_safe)``, all with shape + (nf, nloc, nnei). ``mask`` is boolean; the two index arrays are int64 + with 0 substituted on invalid slots. + """ + nf, nloc, nnei = nlist.shape + nlist = xp.astype(nlist, xp.int64) + mask = (nlist >= 0) & pair_keep_mask + nlist_safe = xp.where(mask, nlist, xp.zeros_like(nlist)) + + if mapping is None: + # Neighbor indices are already local indices in [0, nloc). + src_local = nlist_safe + else: + # Map extended index -> local index for each frame. + mapping_flat = xp.astype(xp.reshape(mapping, (-1,)), xp.int64) + frame_idx = xp.reshape( + xp.arange(nf, dtype=xp.int64, device=array_api_compat.device(nlist)), + (nf, 1, 1), + ) + flat_idx = xp.reshape(frame_idx * nall + nlist_safe, (-1,)) + src_local = xp.reshape(xp.take(mapping_flat, flat_idx, axis=0), nlist.shape) + + # pt's src_ok filter: drop (here: mask) edges mapping outside [0, nloc). + mask = mask & (src_local >= 0) & (src_local < nloc) + src_local_safe = xp.where(mask, src_local, xp.zeros_like(src_local)) + # Re-zero nlist_safe after the src_ok update so coordinate gathers stay + # in-bounds when callers pass local nlists with out-of-range entries. + nlist_safe = xp.where(mask, nlist_safe, xp.zeros_like(nlist_safe)) + return mask, nlist_safe, src_local_safe + + +def build_edge_cache( + *, + type_ebed: Any, + extended_coord: Any, + nlist: Any, + mapping: Any, + pair_keep_mask: Any, + eps: float, + deg_norm_floor: float, + edge_envelope: Any, + radial_basis: Any, + n_radial: int, # unused: kept for pt signature parity (pt sizes its empty cache) + random_gamma: bool, + wigner_calc: Any, + gamma: Any = None, +) -> EdgeCache: + """ + Build the global padded edge cache from a DeePMD padded neighbor list. + + Padded counterpart of pt ``sezm_nn.edge_cache.build_edge_cache``. Instead + of extracting a sparse edge list with ``torch.nonzero`` (data-dependent + length), every neighbor slot becomes one edge slot: + ``E = nf * nloc * nnei`` flattened row-major, with invalid slots marked by + ``edge_mask == 0`` (see the :class:`EdgeCache` layout contract). In + particular ``dst == np.repeat(arange(nf * nloc), nnei)`` always, and there + is no empty-cache special case (E is shape-determined). + + Masked-slot safety: gathered edge vectors on invalid slots are garbage + (placeholder index 0), and could even be exactly zero (self-difference), + which would produce a 0/0 in the normalization inside the quaternion + construction. Although the *forward* contribution of such slots is masked + out downstream, a NaN there would still poison the *backward* pass + (``where`` propagates NaN gradients from the unselected branch). Invalid + slots are therefore rewritten to the safe dummy unit vector ``+z`` BEFORE + any norm/quaternion/Wigner evaluation, and their envelope, radial basis, + and type features are multiplied by the mask so they are exactly zero. + + Parameters + ---------- + type_ebed + Per-node type embedding with shape (N, C), where N = nf * nloc. + extended_coord + Extended coordinates with shape (nf, nall, 3). + nlist + Neighbor list with shape (nf, nloc, nnei); -1 marks padding. + mapping + Mapping from extended to local indices with shape (nf, nall), or None + when the neighbor indices are already local. + pair_keep_mask + Pair keep mask from ``PairExcludeMask`` with shape (nf, nloc, nnei). + True means keep. + eps + Small positive epsilon for safe norm / quaternion construction. + deg_norm_floor + Floor added to the envelope-squared degree before the inverse-sqrt + normalization. + edge_envelope + C^3 edge envelope callable ``(E, 1) -> (E, 1)``. + radial_basis + Radial basis callable ``(E, 1) -> (E, n_radial)`` (envelope baked in). + n_radial + Number of radial basis channels. Unused in the padded layout (kept + for signature parity with pt, where it sizes the empty cache). + random_gamma + Whether to apply a random roll around the local +Z axis before + constructing Wigner-D blocks. + wigner_calc + Callable converting edge quaternions (E, 4) into packed Wigner-D + blocks ``(D_full, Dt_full)``. + gamma + Optional per-edge roll angles with shape (E,), used only when + ``random_gamma`` is True. pt draws gamma internally with + ``torch.rand`` and the draw cannot be reproduced here, so callers + needing determinism (e.g. tests) inject the angles explicitly. When + None, angles are drawn from ``numpy.random.default_rng()`` uniformly + in ``[0, 2*pi)``, matching pt's distribution. + + Returns + ------- + EdgeCache + Padded per-edge cache with ``edge_mask`` set. + """ + xp = array_api_compat.array_namespace(type_ebed, extended_coord, nlist) + device = array_api_compat.device(extended_coord) + nf, nloc, nnei = nlist.shape + nall = extended_coord.shape[1] + n_nodes = nf * nloc + n_edge = n_nodes * nnei + + # === Step 1. Validity mask and safe indices (pt edge_keep semantics) === + mask, nlist_safe, src_local_safe = _build_edge_mask_and_src( + xp, nlist, mapping, pair_keep_mask, nall + ) + mask_flat = xp.reshape(mask, (-1,)) + + # === Step 2. Node indices === + # dst is slot-implicit: arange(nf * nloc) repeated nnei times (contract). + frame_idx = xp.reshape(xp.arange(nf, dtype=xp.int64, device=device), (nf, 1, 1)) + src = xp.reshape(frame_idx * nloc + src_local_safe, (-1,)) + node_idx = xp.arange(n_nodes, dtype=xp.int64, device=device) + dst = xp.reshape(xp.broadcast_to(node_idx[:, None], (n_nodes, nnei)), (-1,)) + + # === Step 3. Gather per-edge geometry from extended coordinates === + coord_flat = xp.reshape(extended_coord, (nf * nall, 3)) + neighbor_coord_index = xp.reshape(frame_idx * nall + nlist_safe, (-1,)) + loc_idx = xp.reshape(xp.arange(nloc, dtype=xp.int64, device=device), (1, nloc, 1)) + center_ext = xp.broadcast_to(frame_idx * nall + loc_idx, (nf, nloc, nnei)) + center_coord_index = xp.reshape(center_ext, (-1,)) + neighbor_pos = xp.take(coord_flat, neighbor_coord_index, axis=0) + center_pos = xp.take(coord_flat, center_coord_index, axis=0) + vec = neighbor_pos - center_pos # (E, 3) + + # === Step 4. Rewrite invalid slots to the safe +z dummy vector === + # Gradient safety: see the function docstring. + maskf = xp.astype(mask_flat, vec.dtype)[:, None] # (E, 1) + z_unit = xp.asarray(np.array([[0.0, 0.0, 1.0]]), dtype=vec.dtype, device=device) + edge_vec = vec * maskf + (1.0 - maskf) * z_unit + edge_len = safe_norm(edge_vec, eps) # (E, 1) + + # === Step 5. Envelope and radial basis, masked to zero on invalid slots === + edge_env = edge_envelope(edge_len) * maskf # (E, 1) + edge_rbf = radial_basis(edge_len) * maskf # (E, n_radial) + + # === Step 6. Edge quaternion -> Wigner-D blocks === + edge_quat = build_edge_quaternion(edge_vec, edge_len=edge_len, eps=eps) + if random_gamma: + if gamma is None: + gamma = np.random.default_rng().uniform(0.0, 2.0 * math.pi, n_edge) + gamma = xp.astype(xp.asarray(gamma, device=device), edge_quat.dtype) + edge_quat = quaternion_multiply(quaternion_z_rotation(gamma), edge_quat) + D_full, Dt_full = wigner_calc(edge_quat) + + # === Step 7. Edge type features (src + dst), masked === + edge_type_feat = ( + xp.take(type_ebed, src, axis=0) + xp.take(type_ebed, dst, axis=0) + ) * xp.astype(maskf, type_ebed.dtype) + + # === Step 8. Smooth destination degrees === + # pt accumulates env^2 with index_add_ over dst (edge_cache.py:622); in the + # padded node-contiguous layout this is a plain sum over the nnei axis. + # edge_env is already exactly zero on invalid slots. + env_sq = xp.reshape(edge_env[:, 0] * edge_env[:, 0], (n_nodes, nnei)) + deg = xp.sum(env_sq, axis=1) # (N,) + inv_sqrt_deg = xp.reshape(1.0 / xp.sqrt(deg + deg_norm_floor), (n_nodes, 1, 1)) + + return EdgeCache( + src=src, + dst=dst, + edge_type_feat=edge_type_feat, + edge_vec=edge_vec, + edge_rbf=edge_rbf, + edge_env=edge_env, + deg=deg, + inv_sqrt_deg=inv_sqrt_deg, + D_full=D_full, + Dt_full=Dt_full, + D_to_m_cache={}, + Dt_from_m_cache={}, + edge_src_gate=None, + edge_quat=edge_quat, + edge_mask=mask_flat, + ) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py new file mode 100644 index 0000000000..e20e69fc87 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -0,0 +1,728 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Embedding layers for the dpmodel DPA4/SeZM descriptor. + +This module is the dpmodel port of +``deepmd.pt.model.descriptor.sezm_nn.embedding``. It defines the type +embedding, geometric initial embedding, and environment-seed embedding used +to initialize SeZM node features. + +Padded-edge layout +------------------ +The pt implementation aggregates sparse per-edge messages into nodes with +``index_add_``. The dpmodel port uses the padded, frame-explicit edge layout +of :class:`~deepmd.dpmodel.descriptor.dpa4_nn.edge_cache.EdgeCache` +(``E = nf * nloc * nnei`` with invalid slots marked by ``edge_mask == 0``), +so every destination aggregation becomes a masked sum over the ``nnei`` axis +of the ``(N, nnei, ...)`` reshape. Each rewrite is commented with the pt +line it replaces. + +Ported / skipped classes +------------------------ +- ``SeZMTypeEmbedding``, ``GeometricInitialEmbedding`` and + ``EnvironmentInitialEmbedding`` are ported (core consumers: + ``sezm.py:710``, ``sezm.py:826`` and ``sezm.py:733`` respectively). +- ``ChargeSpinEmbedding`` (pt ``embedding.py:591``) is NOT ported: it is + constructed only when ``add_chg_spin_ebd=True`` (``sezm.py:717``), and the + flag defaults to ``False`` (``sezm.py:440``), so it is outside the core + DPA4 configuration targeted by this port. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.network import ( + NativeLayer, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .indexing import ( + build_gie_zonal_index, + get_so3_dim_of_lmax, +) + +if TYPE_CHECKING: + from .edge_cache import ( + EdgeCache, + ) + + +def _edge_layout(n_edge: int, n_nodes: int) -> int: + """Validate the padded-edge layout and return ``nnei = E // N``.""" + if n_nodes <= 0 or n_edge % n_nodes != 0: + raise ValueError( + "padded-edge layout requires E to be a multiple of N; " + f"got E={n_edge}, N={n_nodes}" + ) + return n_edge // n_nodes + + +class SeZMTypeEmbedding(NativeOP): + """ + Minimal SeZM type embedding with Adam-routed parameter naming. + + Parameters + ---------- + ntypes + Number of atom types. + embed_dim + Embedding dimension. + precision + Floating-point precision of the embedding table. + seed + Random seed for initialization. + trainable + Whether parameters are trainable. + padding + Whether to append one all-zero padding row. + + Notes + ----- + The parameter is named with ``adam_`` prefix so HybridMuon routes it to + Adam (the name matches the pt ``state_dict`` key ``adam_type_embedding``). + """ + + def __init__( + self, + *, + ntypes: int, + embed_dim: int, + precision: str = DEFAULT_PRECISION, + seed: int | list[int] | None = None, + trainable: bool = True, + padding: bool = True, + ) -> None: + self.ntypes = int(ntypes) + self.embed_dim = int(embed_dim) + self.precision = precision + self.seed = seed + self.trainable = bool(trainable) + self.padding = bool(padding) + if self.ntypes <= 0: + raise ValueError("`ntypes` must be positive") + if self.embed_dim <= 0: + raise ValueError("`embed_dim` must be positive") + prec = PRECISION_DICT[self.precision.lower()] + + # === Step 1+2. Build the table; active rows N(0, init_std), padding + # row zero (pt embedding.py:103-124). The numpy RNG stream differs + # from pt's torch generator; weight values are not bit-compatible. + init_std = 1.0 / math.sqrt(float(self.ntypes + self.embed_dim)) + rng = np.random.default_rng(child_seed(seed, 0)) + table = rng.normal(scale=init_std, size=(self.ntypes, self.embed_dim)) + if self.padding: + table = np.concatenate( + [table, np.zeros((1, self.embed_dim), dtype=table.dtype)], axis=0 + ) + self.adam_type_embedding = table.astype(prec) + + def call(self, atype: Any) -> Any: + """ + Gather type embeddings. + + Parameters + ---------- + atype + Atom types with shape (...,). Valid type range is [0, ntypes-1] + (plus the padding row index ``ntypes`` when ``padding=True``). + Negative type ids are invalid input and are NOT validated here + (caller contract). + + Returns + ------- + Array + Type embeddings with shape (..., embed_dim). + """ + xp = array_api_compat.array_namespace(atype) + weight = xp.asarray( + self.adam_type_embedding[...], device=array_api_compat.device(atype) + ) + # pt embedding.py:143 torch.embedding -> flat int64 take + reshape. + index = xp.astype(xp.reshape(atype, (-1,)), xp.int64) + out = xp.take(weight, index, axis=0) + return xp.reshape(out, (*atype.shape, self.embed_dim)) + + def serialize(self) -> dict[str, Any]: + """Serialize to a dict. + + The pt class has no ``serialize()``; the ``@variables`` key here + matches the pt ``state_dict()`` key (``adam_type_embedding``). + """ + return { + "@class": "SeZMTypeEmbedding", + "@version": 1, + "config": { + "ntypes": self.ntypes, + "embed_dim": self.embed_dim, + "padding": self.padding, + "precision": self.precision.lower(), + "trainable": self.trainable, + "seed": None, + }, + "@variables": { + "adam_type_embedding": to_numpy_array(self.adam_type_embedding) + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SeZMTypeEmbedding: + """Deserialize from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "SeZMTypeEmbedding": + raise ValueError(f"Invalid class for SeZMTypeEmbedding: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls(**config) + prec = PRECISION_DICT[obj.precision.lower()] + table = np.asarray(variables["adam_type_embedding"], dtype=prec) + if table.shape != obj.adam_type_embedding.shape: + raise ValueError( + f"adam_type_embedding shape {table.shape} does not match " + f"the expected shape {obj.adam_type_embedding.shape}" + ) + obj.adam_type_embedding = table + return obj + + +class GeometricInitialEmbedding(NativeOP): + """ + Geometric initial embedding that adds zonal (m=0) rotated features. + + This module rotates pre-computed radial features for each degree l >= 1 + using the zonal (m=0) column of the cached inverse Wigner-D blocks + (local->global). The l=0 component is not computed here since it comes + from type embedding. + + Parameters + ---------- + lmax + Maximum node degree for the initial embedding. + channels + Number of channels per (l, m) coefficient. + precision + Floating-point precision label (kept for config parity with pt; the + computation follows the input dtype). + """ + + def __init__( + self, + *, + lmax: int, + channels: int, + precision: str = DEFAULT_PRECISION, + ) -> None: + self.lmax = int(lmax) + self.channels = int(channels) + self.ebed_dim = get_so3_dim_of_lmax(self.lmax) + self.precision = precision + # One aligned entry per non-scalar node row: output row, local m=0 + # column, and the matching radial degree slot (static int64 tables; + # pt registers them as persistent buffers, embedding.py:185-195). + ( + self.non_scalar_row_index, + self.zonal_m0_col_index_for_row, + self.radial_slot_index_for_row, + ) = build_gie_zonal_index(self.lmax) + + def call( + self, + *, + n_nodes: int, + edge_cache: EdgeCache, + radial_feat: Any, + zonal_coupling: Any = None, + ) -> Any: + """ + Parameters + ---------- + n_nodes + Number of nodes (nf*nloc). + edge_cache + Per-edge cache containing geometry, weights, and Wigner-D blocks + in the padded layout (``E = n_nodes * nnei``). + radial_feat + Per-edge radial features with shape (E, lmax, C) for l=1..lmax. + zonal_coupling + Optional precomputed zonal coupling with shape (E, D-1). If None, + it is gathered from ``edge_cache.Dt_full``. + + Returns + ------- + Array + Initial features to add with shape (N, D, C). l=0 is guaranteed + zero. + """ + # === Step 1. Initialize output === + xp = array_api_compat.array_namespace(edge_cache.edge_vec) + device = array_api_compat.device(edge_cache.edge_vec) + dtype = edge_cache.edge_vec.dtype + if self.lmax == 0: + # pt embedding.py:226-230: zeros short-circuit. + return xp.zeros( + (n_nodes, self.ebed_dim, self.channels), dtype=dtype, device=device + ) + n_edge = int(edge_cache.dst.shape[0]) + nnei = _edge_layout(n_edge, int(n_nodes)) + + # === Step 2. Gather all m=0 columns (l >= 1) in one shot === + # pt embedding.py:235-241 pairs one packed non-scalar row with the + # zonal m=0 column from the same degree block via advanced indexing + # Dt_full[:, rows, cols]; here this becomes a flat row-major take. + if zonal_coupling is None: + Dt_full = edge_cache.Dt_full # (E, D, D) + dim_full = Dt_full.shape[-1] + flat_index = xp.asarray( + self.non_scalar_row_index * dim_full + self.zonal_m0_col_index_for_row, + device=device, + ) + zonal_coupling = xp.take( + xp.reshape(Dt_full, (n_edge, dim_full * dim_full)), + flat_index, + axis=1, + ) # (E, D-1) + + # === Step 3. Broadcast radial features per row === + # Each non-scalar packed row reuses the radial feature of its degree l + # (pt embedding.py:245-250, index_select on axis 1). + radial_slot_index = xp.asarray(self.radial_slot_index_for_row, device=device) + radial_value_for_row = xp.take( + radial_feat, radial_slot_index, axis=1 + ) # (E, D-1, C) + non_scalar_message = ( + zonal_coupling[:, :, None] * radial_value_for_row + ) # (E, D-1, C) + + # === Step 4. Source Freeze Propagation Gate (optional) === + # pt embedding.py:256-260: mute messages emitted by nodes whose local + # neighborhood enters the frozen zone; ``edge_src_gate`` is ``None`` + # outside bridging mode so this is a no-op in normal training. + src_gate = edge_cache.edge_src_gate + if src_gate is not None: + non_scalar_message = non_scalar_message * xp.astype( + xp.reshape(src_gate, (n_edge, 1, 1)), non_scalar_message.dtype + ) + + # === Step 5. Aggregate to nodes and normalize === + # pt embedding.py:264-267: non_scalar_out.index_add_(0, dst, msg) — + # padded-edge masked sum over the nnei axis (dst is slot-implicit). + edge_mask = edge_cache.edge_mask + if edge_mask is not None: + non_scalar_message = non_scalar_message * xp.astype( + xp.reshape(edge_mask, (n_edge, 1, 1)), non_scalar_message.dtype + ) + non_scalar_out = xp.sum( + xp.reshape( + non_scalar_message, + (n_nodes, nnei, self.ebed_dim - 1, self.channels), + ), + axis=1, + ) # (N, D-1, C) + # pt embedding.py:268: out[:, non_scalar_row_index, :] = non_scalar_out + # with row 0 (l=0) left at its zeros init (pt embedding.py:226). + # ``non_scalar_row_index`` is the contiguous arange(1, D), so the + # writeback is a concat with a zero l=0 row. + out = xp.concat( + [ + xp.zeros( + (n_nodes, 1, self.channels), + dtype=non_scalar_out.dtype, + device=device, + ), + non_scalar_out, + ], + axis=1, + ) # (N, D, C) + # pt embedding.py:269: out.mul_(inv_sqrt_deg). + out = out * xp.astype(edge_cache.inv_sqrt_deg, out.dtype) + return xp.astype(out, dtype) + + def serialize(self) -> dict[str, Any]: + """Serialize to a dict (config only; same flat layout as pt).""" + return { + "@class": "GeometricInitialEmbedding", + "@version": 1, + "lmax": self.lmax, + "channels": self.channels, + "precision": self.precision.lower(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> GeometricInitialEmbedding: + """Deserialize from a dict (accepts the pt ``serialize()`` output).""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "GeometricInitialEmbedding": + raise ValueError(f"Invalid class for GeometricInitialEmbedding: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + return cls( + lmax=int(data.pop("lmax")), + channels=int(data.pop("channels")), + precision=str(data.pop("precision")), + ) + + +class EnvironmentInitialEmbedding(NativeOP): + """ + Environment matrix initial embedding for l=0 features. + + Computes an initial embedding based on the 4D environment matrix:: + + [s, s * rx, s * ry, s * rz] + + Combined with independent type embeddings (individual type embedding), + providing physical inductive bias for l=0 features. + + The computation follows the environment matrix approach where:: + + 1. Build `r_tilde = [s, s*r_hat]` where `s = edge_env / r` and + `r_hat = edge_vec / r` + 2. G network: `g = G(rbf_proj(edge_rbf), type_src, type_dst)` produces + per-edge features + - Uses independent `env_type_embed` instead of projecting from the + main type embedding + - Uses `rbf_proj` to project edge_rbf to `rbf_out_dim` + 3. env_agg: aggregate outer product `r_tilde ⊗ g` by destination node + 4. D matrix: `D = env_agg^T @ env_agg[:, :, :axis_dim]` + 5. Output: projection of flattened D matrix into FiLM logits + + Parameters + ---------- + ntypes : int + Number of atom types. + n_radial : int + Number of radial basis functions. + channels : int + Output channel dimension per FiLM branch (final output is 2*channels). + embed_dim : int + G network output dimension (filter width). + axis_dim : int + D matrix axis dimension (must be < embed_dim). + type_dim : int + Dimension for independent type embeddings in env_seed. + hidden_dim : int + Hidden layer size for G network. + mlp_bias : bool + Whether to enable bias terms in env-seed MLP layers + (`rbf_proj_layer1/2` and `g_layer1/2`). + activation_function : str + Activation function for G network hidden layer. + eps : float + Small epsilon for numerical stability. + precision : str + Floating-point precision of the parameters. + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for reproducibility. + """ + + def __init__( + self, + *, + ntypes: int, + n_radial: int, + channels: int, + embed_dim: int = 64, + axis_dim: int = 8, + type_dim: int = 16, + hidden_dim: int = 64, + mlp_bias: bool = False, + activation_function: str = "silu", + eps: float = 1e-7, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + # === Validate parameters === + if axis_dim >= embed_dim: + raise ValueError( + f"`axis_dim` ({axis_dim}) must be < `embed_dim` ({embed_dim})" + ) + + self.ntypes = int(ntypes) + self.n_radial = int(n_radial) + self.channels = int(channels) + self.embed_dim = int(embed_dim) + self.axis_dim = int(axis_dim) + self.type_dim = int(type_dim) + self.hidden_dim = int(hidden_dim) + self.mlp_bias = bool(mlp_bias) + self.activation_function = str(activation_function) + self.eps = float(eps) + self.precision = precision + self.trainable = bool(trainable) + + # === RBF projection: n_radial -> rbf_out_dim (two-layer MLP) === + # rbf_out_dim = max(32, embed_dim - 2*type_dim) to align G-network + # width to embed_dim. First layer: n_radial -> rbf_out_dim with + # activation. Second layer: rbf_out_dim -> rbf_out_dim linear. + self.rbf_out_dim = max(32, self.embed_dim - 2 * self.type_dim) + seed_rbf_proj = child_seed(seed, 0) + self.rbf_proj_layer1 = NativeLayer( + self.n_radial, + self.rbf_out_dim, + bias=self.mlp_bias, + activation_function=self.activation_function, + precision=self.precision, + seed=child_seed(seed_rbf_proj, 0), + trainable=self.trainable, + ) + self.rbf_proj_layer2 = NativeLayer( + self.rbf_out_dim, + self.rbf_out_dim, + bias=self.mlp_bias, + activation_function=None, + precision=self.precision, + seed=child_seed(seed_rbf_proj, 1), + trainable=self.trainable, + ) + + # === Independent type embedding: ntypes -> type_dim === + # Individual type embedding + self.env_type_embed = SeZMTypeEmbedding( + ntypes=self.ntypes, + embed_dim=self.type_dim, + precision=self.precision, + seed=child_seed(seed, 1), + trainable=self.trainable, + ) + + # === G network: (rbf_out_dim + 2*type_dim) -> hidden_dim -> embed_dim === + seed_g_net = child_seed(seed, 2) + g_in_dim = self.rbf_out_dim + 2 * self.type_dim + self.g_layer1 = NativeLayer( + g_in_dim, + self.hidden_dim, + bias=self.mlp_bias, + activation_function=self.activation_function, + precision=self.precision, + seed=child_seed(seed_g_net, 0), + trainable=self.trainable, + ) + self.g_layer2 = NativeLayer( + self.hidden_dim, + self.embed_dim, + bias=self.mlp_bias, + activation_function=None, + precision=self.precision, + seed=child_seed(seed_g_net, 1), + trainable=self.trainable, + ) + + # === Output projection: embed_dim * axis_dim -> 2*channels === + # Zero init so FiLM logits start at zero (pt init="final", + # embedding.py:447-455); strengths control magnitude. + self.output_proj = NativeLayer( + self.embed_dim * self.axis_dim, + 2 * self.channels, + bias=False, + activation_function=None, + precision=self.precision, + seed=child_seed(seed, 3), + trainable=self.trainable, + ) + self.output_proj.w = np.zeros_like(self.output_proj.w) + + def call( + self, + *, + edge_cache: EdgeCache, + atype_flat: Any, + n_nodes: int, + ) -> Any: + """ + Compute environment FiLM logits for l=0 conditioning. + + Parameters + ---------- + edge_cache : EdgeCache + Edge cache containing src, dst, edge_vec, edge_rbf, edge_env in + the padded layout (``E = n_nodes * nnei``). + atype_flat : Array + Flattened atom types with shape (N,), where N = nf * nloc. + n_nodes : int + Number of nodes (N = nf * nloc). + + Returns + ------- + Array + FiLM logits with shape (N, 2*channels). + """ + xp = array_api_compat.array_namespace(edge_cache.edge_vec) + src, dst = edge_cache.src, edge_cache.dst + edge_vec = edge_cache.edge_vec # (E, 3) + edge_rbf = edge_cache.edge_rbf # (E, n_radial) + edge_env = edge_cache.edge_env # (E, 1) + n_edge = int(dst.shape[0]) + nnei = _edge_layout(n_edge, int(n_nodes)) + + # === Step 1. Construct r_tilde = [s, s*r_hat] === + # s = edge_env * (1/r), r_hat = edge_vec / r (pt embedding.py:489-495) + r_sq = xp.sum(edge_vec * edge_vec, axis=-1, keepdims=True) # (E, 1) + inv_r = 1.0 / xp.sqrt(r_sq + self.eps * self.eps) # (E, 1) + s = edge_env * inv_r # (E, 1) + r_hat = edge_vec * inv_r # (E, 3) + r_tilde = xp.concat([s, s * r_hat], axis=-1) # (E, 4) + + # === Step 2. Compute G network input and output === + # Use independent type embeddings (decoupled from main type embedding) + src_index = xp.astype(xp.reshape(src, (n_edge,)), xp.int64) + dst_index = xp.astype(xp.reshape(dst, (n_edge,)), xp.int64) + atype_src = xp.take(atype_flat, src_index, axis=0) # (E,) + atype_dst = xp.take(atype_flat, dst_index, axis=0) # (E,) + type_src = self.env_type_embed(atype_src) # (E, type_dim) + type_dst = self.env_type_embed(atype_dst) # (E, type_dim) + + # Project edge_rbf to rbf_out_dim (two-layer MLP) + rbf_proj = self.rbf_proj_layer2( + self.rbf_proj_layer1(edge_rbf) + ) # (E, rbf_out_dim) + + # G network input: concat projected RBF and type embeddings + g_input = xp.concat([rbf_proj, type_src, type_dst], axis=-1) # (E, g_in_dim) + g = self.g_layer2(self.g_layer1(g_input)) # (E, embed_dim) + + # === Step 3. Aggregate outer product by destination node === + # pt embedding.py:515 einsum("ei,ej->eij") -> broadcast product. + outer = r_tilde[:, :, None] * g[:, None, :] # (E, 4, embed_dim) + outer_flat = xp.reshape(outer, (n_edge, 4 * self.embed_dim)) + # Source Freeze Propagation Gate (pt embedding.py:519-521): mute the + # outer-product contribution of any edge whose source node has a + # neighbor in the frozen zone. + src_gate = edge_cache.edge_src_gate + if src_gate is not None: + outer_flat = outer_flat * xp.astype( + xp.reshape(src_gate, (n_edge, 1)), outer_flat.dtype + ) + # pt embedding.py:522-523: env_agg.index_add_(0, dst, outer_flat) — + # padded-edge masked sum over the nnei axis (dst is slot-implicit). + edge_mask = edge_cache.edge_mask + if edge_mask is not None: + outer_flat = outer_flat * xp.astype( + xp.reshape(edge_mask, (n_edge, 1)), outer_flat.dtype + ) + env_agg = xp.sum( + xp.reshape(outer_flat, (n_nodes, nnei, 4 * self.embed_dim)), + axis=1, + ) # (N, 4*embed_dim) + env_agg = xp.reshape(env_agg, (n_nodes, 4, self.embed_dim)) + + # === Step 4. Smooth normalization by envelope-squared degree === + # Reuse the cache's inverse-sqrt degree so the version-aware + # ``deg_norm_floor`` is applied consistently with GIE. + env_agg = env_agg * xp.astype(edge_cache.inv_sqrt_deg, env_agg.dtype) + + # === Step 5. D matrix: D = env_agg^T @ env_agg[:, :, :axis_dim] === + env_agg_t = xp.permute_dims(env_agg, (0, 2, 1)) # (N, embed_dim, 4) + env_agg_axis = env_agg[:, :, : self.axis_dim] # (N, 4, axis_dim) + mat_d = xp.matmul(env_agg_t, env_agg_axis) # (N, embed_dim, axis_dim) + + # === Step 6. Output projection for FiLM logits === + d_flat = xp.reshape( + mat_d, (n_nodes, self.embed_dim * self.axis_dim) + ) # (N, embed_dim*axis_dim) + return self.output_proj(d_flat) + + def _variable_slots(self) -> dict[str, tuple[Any, str]]: + """Map pt ``state_dict`` keys to (owner object, attribute name).""" + slots: dict[str, tuple[Any, str]] = {} + for name in ("rbf_proj_layer1", "rbf_proj_layer2", "g_layer1", "g_layer2"): + layer = getattr(self, name) + slots[f"{name}.matrix"] = (layer, "w") + if self.mlp_bias: + slots[f"{name}.bias"] = (layer, "b") + slots["env_type_embed.adam_type_embedding"] = ( + self.env_type_embed, + "adam_type_embedding", + ) + slots["output_proj.matrix"] = (self.output_proj, "w") + return slots + + def serialize(self) -> dict[str, Any]: + """Serialize to a dict. + + The ``@variables`` keys match the pt ``state_dict()`` key names, so + the pt ``serialize()`` output deserializes directly into this class + (and vice versa). + """ + variables = { + key: to_numpy_array(getattr(owner, attr)) + for key, (owner, attr) in self._variable_slots().items() + } + return { + "@class": "EnvironmentInitialEmbedding", + "@version": 1, + "config": { + "ntypes": self.ntypes, + "n_radial": self.n_radial, + "channels": self.channels, + "embed_dim": self.embed_dim, + "axis_dim": self.axis_dim, + "type_dim": self.type_dim, + "hidden_dim": self.hidden_dim, + "mlp_bias": self.mlp_bias, + "activation_function": self.activation_function, + "eps": self.eps, + "precision": self.precision.lower(), + "trainable": self.trainable, + "seed": None, + }, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> EnvironmentInitialEmbedding: + """Deserialize from a dict (accepts the pt ``serialize()`` output).""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "EnvironmentInitialEmbedding": + raise ValueError(f"Invalid class: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls(**config) + prec = PRECISION_DICT[obj.precision.lower()] + slots = obj._variable_slots() + if set(variables) != set(slots): + raise ValueError( + f"variable keys {sorted(variables)} do not match the expected " + f"keys {sorted(slots)}" + ) + for key, (owner, attr) in slots.items(): + value = np.asarray(variables[key], dtype=prec) + expected_shape = getattr(owner, attr).shape + if value.shape != expected_shape: + raise ValueError( + f"shape of {key} {value.shape} does not match " + f"the expected shape {expected_shape}" + ) + setattr(owner, attr, value) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/ffn.py b/deepmd/dpmodel/descriptor/dpa4_nn/ffn.py new file mode 100644 index 0000000000..a54cf7c4eb --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/ffn.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Equivariant feed-forward layers for DPA4/SeZM. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.ffn``. +It defines the full SO(3)-equivariant feed-forward network used inside SeZM +interaction blocks. + +Branches guarded with ``NotImplementedError`` (flags unused by the core DPA4 +config): + +- ``ffn_so3_grid=True`` — the pt path instantiates ``SO3GridNet`` + (pt ffn.py:209), which is not ported to dpmodel. +- ``grid_mlp=True`` together with the grid path active selects + ``op_type='mlp'`` for ``S2GridNet`` (pt ffn.py:206); the delegate + ``S2GridNet`` constructor raises for that op type (``GridMLP`` is not + ported), so no duplicate guard is added here. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .activation import ( + GatedActivation, +) +from .grid_net import ( + S2GridNet, +) +from .projection import ( + resolve_s2_grid_resolution, +) +from .so2 import ( + _compute_precision, +) +from .so3 import ( + SO3Linear, +) + + +class EquivariantFFN(NativeOP): + """ + Full equivariant FFN operating on all spherical harmonic degrees. + + Default structure (glu_activation=False): + SO3 linear (in -> hidden) -> GatedActivation -> SO3 linear (hidden -> out) + + Default structure (glu_activation=True): + SO3 linear (in -> 2*hidden) -> split -> GatedActivation(val, gate) -> SO3 linear (hidden -> out) + + Optional grid-FFN structure (s2_activation=True): + SO3 linear (in -> 2*hidden) + -> project packed SO(3) coefficients to the S2 grid + -> grid GLU or scalar-routed polynomial branch on hidden features + -> project grid features back to packed SO(3) coefficients + -> add scalar LinearSwiGLU branch to l=0 + -> SO3 linear (hidden -> out) + + Parameters + ---------- + lmax + Maximum degree. + channels + Number of channels per (l, m) coefficient. + hidden_channels + Hidden dimension for the FFN. + kmax + Maximum Wigner-D frame order (|k|) used by the SO3 Wigner-D FFN grid. + grid_mlp + If True, select the polynomial grid MLP operation when the + block-internal FFN grid path is enabled. Not ported: the delegate + ``S2GridNet`` raises ``NotImplementedError`` for ``op_type='mlp'``. + grid_branch + Number of scalar-routed polynomial product branches used when the + block-internal FFN grid path is enabled. ``0`` disables this branch + mixer. Positive values take precedence over ``grid_mlp``. + s2_activation + If True, enable the S2 FFN grid path. + ffn_so3_grid + If True, enable the SO3 Wigner-D FFN grid path (not ported). + lebedev_quadrature + If True, use Lebedev quadrature for the S2 projector in this FFN. + activation_function + Activation function for l=0 components (e.g., "silu", "tanh", "gelu"). + glu_activation + If True, use GLU-style gating (e.g., silu -> swiglu, gelu -> geglu). + mlp_bias + Whether to use bias in SO3Linear (l=0 bias), GatedActivation + (gate linear bias). + precision + Parameter precision. + trainable + Whether parameters are trainable. + seed + Random seed for weight initialization. + """ + + def __init__( + self, + *, + lmax: int, + channels: int, + hidden_channels: int, + kmax: int = 1, + grid_mlp: bool = False, + grid_branch: int = 0, + s2_activation: bool = False, + ffn_so3_grid: bool = False, + lebedev_quadrature: bool = False, + activation_function: str = "silu", + glu_activation: bool = True, + mlp_bias: bool = False, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + self.lmax = int(lmax) + self.channels = int(channels) + self.hidden_channels = int(hidden_channels) + self.kmax = int(kmax) + if self.kmax < 0: + raise ValueError("`kmax` must be non-negative") + self.use_grid_mlp = bool(grid_mlp) + self.grid_branch = int(grid_branch) + if self.grid_branch < 0: + raise ValueError("`grid_branch` must be non-negative") + self.use_grid_branch = self.grid_branch > 0 + self.s2_activation = bool(s2_activation) + self.ffn_so3_grid = bool(ffn_so3_grid) + if self.ffn_so3_grid: + raise NotImplementedError( + "ffn_so3_grid=True (SO3GridNet) is not ported to dpmodel" + ) + self.lebedev_quadrature = bool(lebedev_quadrature) + self.s2_grid_method = "lebedev" if self.lebedev_quadrature else "e3nn" + base_grid = resolve_s2_grid_resolution( + self.lmax, + self.lmax, + method=self.s2_grid_method, + ) + self.s2_grid_resolution = ( + [max(base_grid), max(base_grid)] + if self.s2_grid_method == "e3nn" + else base_grid + ) + self.activation_function = str(activation_function) + self.glu_activation = bool(glu_activation) + self.mlp_bias = bool(mlp_bias) + self.precision = precision + self.compute_precision = _compute_precision(precision) + self.trainable = bool(trainable) + # pt: grid_n_frames = 2 * kmax + 1 only when ffn_so3_grid (guarded above) + self.grid_n_frames = 1 + + # === Step 0. Split deterministic seeds at the module top-level === + seed_so3_in = child_seed(seed, 0) + seed_act = child_seed(seed, 1) + seed_so3_out = child_seed(seed, 2) + + # === First SO3Linear for channel mixing === + self.use_grid_net = self.s2_activation or self.ffn_so3_grid + if self.use_grid_net: + linear1_out_channels = 2 * self.grid_n_frames * self.hidden_channels + else: + linear1_out_channels = ( + 2 * self.hidden_channels + if self.glu_activation + else self.hidden_channels + ) + self.so3_linear_1 = SO3Linear( + lmax=self.lmax, + in_channels=self.channels, + out_channels=linear1_out_channels, + n_focus=1, + precision=self.precision, + mlp_bias=self.mlp_bias, + trainable=self.trainable, + seed=seed_so3_in, + ) + + # === Equivariant nonlinearity path === + if self.use_grid_net: + grid_op = ( + "branch" + if self.use_grid_branch + else ("mlp" if self.use_grid_mlp else "glu") + ) + # op_type='mlp' raises NotImplementedError inside S2GridNet + self.act: NativeOP = S2GridNet( + lmax=self.lmax, + channels=self.hidden_channels, + n_focus=1, + mode="self", + op_type=grid_op, + precision=self.compute_precision, + layout="ndfc", + grid_resolution_list=self.s2_grid_resolution, + coefficient_layout="packed", + grid_method=self.s2_grid_method, + grid_branches=max(1, self.grid_branch), + mlp_bias=self.mlp_bias, + trainable=self.trainable, + seed=seed_act, + ) + else: + self.act = GatedActivation( + lmax=self.lmax, + channels=self.hidden_channels, + precision=self.compute_precision, + activation_function=self.activation_function, + mlp_bias=self.mlp_bias, + layout="ndfc", + trainable=self.trainable, + seed=seed_act, + ) + + # === Second SO3Linear for channel mixing === + # Zero-initialized so residual path starts near-identity. + self.so3_linear_2 = SO3Linear( + lmax=self.lmax, + in_channels=self.grid_n_frames * self.hidden_channels, + out_channels=self.channels, + n_focus=1, + precision=self.precision, + mlp_bias=self.mlp_bias, + trainable=self.trainable, + seed=seed_so3_out, + init_std=0.0, + ) + + def call(self, x: Any) -> Any: + """ + Parameters + ---------- + x + Input with shape (N, D, F, C) where D=(lmax+1)^2. + + Returns + ------- + Array + Output with shape (N, D, F, C). + """ + # === Step 1. Input up projection === + x = self.so3_linear_1(x) + + # === Step 2. Equivariant nonlinearity === + if self.use_grid_net: + x = self.act(x) + elif self.glu_activation: + # Split into value and gate branches along channel dimension + # (pt uses x.chunk(2, dim=-1); slicing is array-API portable) + x_val = x[..., : self.hidden_channels] + x_gate = x[..., self.hidden_channels :] + # Pass gate to GatedActivation for GLU-style gating + x = self.act(x_val, gate=x_gate) + else: + x = self.act(x) + + # === Step 3. Per-degree output projection === + x = self.so3_linear_2(x) + + return x + + def _sub_modules(self) -> list[tuple[str, NativeOP]]: + """Sub-modules with their pt module names.""" + return [ + ("so3_linear_1", self.so3_linear_1), + ("act", self.act), + ("so3_linear_2", self.so3_linear_2), + ] + + def _variables(self) -> dict[str, Any]: + """Variables keyed by the pt ``state_dict`` key names.""" + variables: dict[str, Any] = {} + for prefix, sub in self._sub_modules(): + for key, value in sub.serialize()["@variables"].items(): + variables[f"{prefix}.{key}"] = value + return variables + + def _load_variables(self, variables: dict[str, Any]) -> None: + """Load variables keyed by the pt ``state_dict`` key names.""" + variables = dict(variables) + for attr, sub in self._sub_modules(): + full = f"{attr}." + sv = { + key[len(full) :]: value + for key, value in variables.items() + if key.startswith(full) + } + for key in list(variables): + if key.startswith(full): + del variables[key] + if not sv: + raise KeyError(f"Missing variables with prefix: {full}") + # rebuild the sub-module through its own (shape-checking) + # deserialize, reusing its serialized config + data = sub.serialize() + data["@variables"] = sv + setattr(self, attr, type(sub).deserialize(data)) + if variables: + raise KeyError(f"Unknown variables: {sorted(variables)}") + + def serialize(self) -> dict[str, Any]: + """Serialize the EquivariantFFN to a dict (pt-compatible format).""" + return { + "@class": "EquivariantFFN", + "@version": 1, + "config": { + "lmax": self.lmax, + "channels": self.channels, + "hidden_channels": self.hidden_channels, + "kmax": self.kmax, + "grid_mlp": self.use_grid_mlp, + "grid_branch": self.grid_branch, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "s2_activation": self.s2_activation, + "ffn_so3_grid": self.ffn_so3_grid, + "lebedev_quadrature": self.lebedev_quadrature, + "activation_function": self.activation_function, + "glu_activation": self.glu_activation, + "mlp_bias": self.mlp_bias, + "trainable": self.trainable, + "seed": None, + }, + "@variables": self._variables(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> EquivariantFFN: + """Deserialize an EquivariantFFN from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "EquivariantFFN": + raise ValueError(f"Invalid class for EquivariantFFN: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = dict(data.pop("config")) + variables = data.pop("@variables") + config["precision"] = str(config.pop("precision")) + obj = cls(**config) + obj._load_variables(variables) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py new file mode 100644 index 0000000000..647d112722 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py @@ -0,0 +1,646 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Grid-space nonlinearities for DPA4/SeZM coefficient tensors. + +This module is the dpmodel port of +``deepmd.pt.model.descriptor.sezm_nn.grid_net``, restricted to the S2/Lebedev +path used by the core DPA4 configuration. A grid net receives coefficient +tensors, converts them to quadrature values, applies one point-wise grid +operation, and projects the result back to coefficients. The public shapes +are: + +* ``mode='self'``: one input ``(N, D, F, 2*C)`` or ``(N, F, D, 2*C)``. +* grid values: ``(N, G, F, C)`` after S2 projection. + +Ported names: ``BaseGridNet`` (``mode='self'``; ``op_type`` 'glu'/'branch'), +``S2GridNet``, ``GridBranch``. + +Skipped names, with consumer evidence from the pt sources: + +- ``SO3GridNet``: only constructed by ``so2.py`` (``node_wise_so3``, + ``message_node_so3``) and ``ffn.py`` (``ffn_so3_grid``) — all disabled in + the core DPA4 config. +- ``FrameContract``, ``FrameExpand``, ``_build_frame_degree_index``: only + constructed by ``SO3GridNet`` (``mode='cross'``); the S2 projector always + has ``n_frames == 1``, so the frame machinery is unreachable here. +- ``GridMLP``: only selected via ``op_type='mlp'`` (``grid_mlp=True`` paths); + the core config has ``grid_mlp=[False, False, False]``. ``BaseGridNet`` + raises ``NotImplementedError`` for ``op_type='mlp'``. + +Guarded (routable from the shared ``S2GridNet`` entry point but only used by +the disabled ``node_wise_s2``/``message_node_s2`` grid products in +``so2.py``): ``mode='cross'`` (and with it ``layout='flat'``) and +``residual_scale_init is not None`` raise ``NotImplementedError``. + +Serialization contract: the pt ``S2GridNet`` and ``GridBranch`` define no +``serialize()`` (they only appear nested inside larger modules' +state-dicts); the dpmodel ``serialize()``/``deserialize()`` use +``@variables`` keys equal to the pt ``state_dict`` key names +(``scalar_gate.weight``, ``grid_op.left_proj.weight``, ...) so that pt +state-dict fragments load directly. The fixed projector matrices are +non-persistent buffers in pt (not in the state dict) and are rebuilt from +the config on deserialization. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + xp_sigmoid, +) +from deepmd.dpmodel.common import ( + get_xp_precision, + to_numpy_array, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .activation import ( + SwiGLU, +) +from .projection import ( + BaseGridProjector, + S2GridProjector, +) +from .so3 import ( + ChannelLinear, + FocusLinear, +) + + +def _softmax_last_axis(x: Any) -> Any: + """Numerically stable softmax on the last axis (matches torch.softmax).""" + xp = array_api_compat.array_namespace(x) + e_x = xp.exp(x - xp.max(x, axis=-1, keepdims=True)) + return e_x / xp.sum(e_x, axis=-1, keepdims=True) + + +class GridBranch(NativeOP): + """ + Scalar-routed polynomial mixer over grid product branches. + + The softmax sees only invariant scalar inputs. Each branch is a + quadratic product of grid fields, so rotations only act through the grid + argument and the operation remains as band-limited as the product path. + + Parameters + ---------- + channels : int + Number of channels per grid point. + n_branches : int + Number of scalar-routed product branches. + precision : str + Parameter precision. + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for weight initialization. + """ + + def __init__( + self, + *, + channels: int, + n_branches: int, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + self.channels = int(channels) + self.n_branches = int(n_branches) + if self.n_branches < 1: + raise ValueError("`n_branches` must be positive") + self.precision = precision + self.trainable = bool(trainable) + self.left_proj = ChannelLinear( + in_channels=self.channels, + out_channels=self.n_branches * self.channels, + precision=precision, + bias=False, + trainable=trainable, + seed=child_seed(seed, 0), + ) + self.right_proj = ChannelLinear( + in_channels=self.channels, + out_channels=self.n_branches * self.channels, + precision=precision, + bias=False, + trainable=trainable, + seed=child_seed(seed, 1), + ) + self.router = ChannelLinear( + in_channels=2 * self.channels, + out_channels=self.n_branches, + precision=precision, + bias=False, + trainable=trainable, + seed=child_seed(seed, 2), + ) + self.out_proj = ChannelLinear( + in_channels=self.channels, + out_channels=self.channels, + precision=precision, + bias=False, + trainable=trainable, + seed=child_seed(seed, 3), + ) + + def call( + self, + query_grid: Any, + context_grid: Any, + scalar_pair: Any, + ) -> Any: + """ + Apply scalar-routed grid branch mixing. + + Parameters + ---------- + query_grid + First grid source with shape ``(N, G, F, C)``. + context_grid + Second grid source with shape ``(N, G, F, C)``. + scalar_pair + Invariant router source with shape ``(N, F, 2*C)``. + """ + xp = array_api_compat.array_namespace(query_grid) + n_batch, n_grid, n_focus, _ = query_grid.shape + left = self.left_proj(query_grid) + right = self.right_proj(context_grid) + value = xp.reshape( + left * right, + (n_batch, n_grid, n_focus, self.n_branches, self.channels), + ) # (N, G, F, N_branches, C) + router = _softmax_last_axis(self.router(scalar_pair)) # (N, F, N_branches) + # einsum "ngfhc,nfh->ngfc" as a broadcast sum over the branch axis + out = xp.sum(value * router[:, None, :, :, None], axis=3) # (N, G, F, C) + return self.out_proj(out) + + def serialize(self) -> dict[str, Any]: + """Serialize the GridBranch to a dict. + + The pt ``GridBranch`` has no ``serialize()``; the ``@variables`` keys + here match the pt ``state_dict`` key names. + """ + return { + "@class": "GridBranch", + "@version": 1, + "config": { + "channels": self.channels, + "n_branches": self.n_branches, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "seed": None, + }, + "@variables": { + "left_proj.weight": to_numpy_array(self.left_proj.weight), + "right_proj.weight": to_numpy_array(self.right_proj.weight), + "router.weight": to_numpy_array(self.router.weight), + "out_proj.weight": to_numpy_array(self.out_proj.weight), + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> GridBranch: + """Deserialize a GridBranch from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "GridBranch": + raise ValueError(f"Invalid class for GridBranch: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + channels=int(config["channels"]), + n_branches=int(config["n_branches"]), + precision=str(config["precision"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + obj._load_variables(variables) + return obj + + def _load_variables(self, variables: dict[str, Any]) -> None: + prec = PRECISION_DICT[self.precision.lower()] + for name, proj in ( + ("left_proj", self.left_proj), + ("right_proj", self.right_proj), + ("router", self.router), + ("out_proj", self.out_proj), + ): + weight = np.asarray(variables[f"{name}.weight"], dtype=prec) + if weight.shape != proj.weight.shape: + raise ValueError( + f"{name}.weight shape {weight.shape} does not match " + f"the expected shape {proj.weight.shape}" + ) + proj.weight = weight + + +class BaseGridNet(NativeOP): + """ + Shared implementation for S2 grid nets (``mode='self'`` only). + + ``mode='self'`` expects one input whose last channel axis contains two + branches; the first half supplies the SwiGLU gates of the scalar path. + + The pt ``mode='cross'`` path (with ``layout='flat'``, + ``residual_scale_init``, and the SO(3) frame machinery) backs the + ``node_wise_s2``/``message_node_s2`` grid products only, which are + disabled in the core DPA4 config; it is not ported. + """ + + def __init__( + self, + *, + projector: BaseGridProjector, + channels: int, + n_focus: int, + mode: str, + op_type: str, + precision: str = DEFAULT_PRECISION, + layout: str, + mlp_bias: bool, + trainable: bool = True, + grid_branches: int = 1, + residual_scale_init: float | None = None, + seed: int | list[int] | None = None, + ) -> None: + self.projector = projector + self.lmax = int(projector.lmax) + self.channels = int(channels) + self.n_focus = int(n_focus) + self.n_frames = int(projector.n_frames) + if self.n_frames != 1: + raise ValueError( + "dpmodel BaseGridNet only supports S2 projectors (n_frames == 1)" + ) + self.mode = str(mode).lower() + if self.mode not in {"self", "cross"}: + raise ValueError("`mode` must be either 'self' or 'cross'") + if self.mode == "cross": + raise NotImplementedError( + "mode='cross' (node_wise_s2/message_node_s2 grid products) " + "is not ported to dpmodel" + ) + self.op_type = str(op_type).lower() + if self.op_type not in {"glu", "mlp", "branch"}: + raise ValueError("`op_type` must be one of 'glu', 'mlp', or 'branch'") + if self.op_type == "mlp": + raise NotImplementedError( + "op_type='mlp' (grid_mlp=True paths) is not ported to dpmodel" + ) + self.precision = precision + self.layout = str(layout).lower() + if self.layout not in {"ndfc", "nfdc", "flat"}: + raise ValueError("`layout` must be one of 'ndfc', 'nfdc', or 'flat'") + if self.mode == "self" and self.layout == "flat": + raise ValueError("`layout='flat'` is only supported for cross grid nets") + self.mlp_bias = bool(mlp_bias) + self.trainable = bool(trainable) + self.expanded_channels = self.n_frames * self.channels + self.query_channels = 2 * self.expanded_channels + self.output_channels = self.expanded_channels + self.frame_zero_index = 0 + if residual_scale_init is not None: + raise NotImplementedError( + "`residual_scale_init` is only used by the cross-mode " + "node_wise_s2/message_node_s2 grid products, which are not " + "ported to dpmodel" + ) + self.residual_scale = None + + self.scalar_act = SwiGLU() + self.scalar_gate = FocusLinear( + in_channels=2 * self.channels, + out_channels=self.channels, + n_focus=self.n_focus, + precision=self.precision, + bias=self.mlp_bias, + trainable=trainable, + seed=child_seed(seed, 0), + init_std=0.01, + ) + if self.op_type == "branch": + self.grid_op: GridBranch | None = GridBranch( + channels=self.channels, + n_branches=grid_branches, + precision=self.precision, + trainable=trainable, + seed=child_seed(seed, 1), + ) + else: + # pt uses nn.Identity() here (parameter-free, no state-dict keys) + self.grid_op = None + + def call(self, query: Any, context: Any = None) -> Any: + """Apply the configured grid net and restore the input layout.""" + xp = array_api_compat.array_namespace(query) + input_dtype = query.dtype + compute_dtype = get_xp_precision(xp, self.precision) + query_ndfc = self._to_ndfc(query) + left, right = self._split_self_query(query_ndfc) + scalar_pair = self._make_scalar_pair(left, right, compute_dtype) + grid_out = self._apply_grid_op(left, right, scalar_pair, compute_dtype) + coeff_out = self._from_grid(grid_out) + coeff_out = self._apply_scalar_path(coeff_out, scalar_pair) + if coeff_out.dtype != input_dtype: + coeff_out = xp.astype(coeff_out, input_dtype) + return self._restore_layout(coeff_out) + + def _apply_grid_op( + self, + left: Any, + right: Any, + scalar_pair: Any, + compute_dtype: Any, + ) -> Any: + xp = array_api_compat.array_namespace(left) + if left.dtype != compute_dtype: + left = xp.astype(left, compute_dtype) + if right.dtype != compute_dtype: + right = xp.astype(right, compute_dtype) + left_grid = self._to_grid(left) + right_grid = self._to_grid(right) + if self.op_type == "glu": + return left_grid * right_grid + return self.grid_op(left_grid, right_grid, scalar_pair) + + def _apply_scalar_path(self, coeff: Any, scalar_pair: Any) -> Any: + xp = array_api_compat.array_namespace(coeff) + scalar_out = self.scalar_act(scalar_pair) # (N, F, C) + scalar_gate = xp_sigmoid(self.scalar_gate(scalar_pair)) # (N, F, C) + coeff = coeff * scalar_gate[:, None, :, :] + # gradient-safe equivalent of the pt in-place + # ``coeff_view[:, 0, :, 0, :].add_(scalar_out)`` (n_frames == 1) + head = coeff[:, :1, :, :] + scalar_out[:, None, :, :] + return xp.concat([head, coeff[:, 1:, :, :]], axis=1) + + def _split_self_query(self, query: Any) -> tuple[Any, Any]: + self._check_last_dim(query, self.query_channels, "query") + # torch.chunk(query, 2, dim=-1) with an even channel count + return ( + query[..., : self.expanded_channels], + query[..., self.expanded_channels :], + ) + + def _make_scalar_pair(self, left: Any, right: Any, compute_dtype: Any) -> Any: + xp = array_api_compat.array_namespace(left) + scalar_pair = xp.concat( + [ + self._extract_scalar(left), + self._extract_scalar(right), + ], + axis=-1, + ) + if scalar_pair.dtype != compute_dtype: + scalar_pair = xp.astype(scalar_pair, compute_dtype) + return scalar_pair + + def _extract_scalar(self, coeff: Any) -> Any: + # (N, D, F, C) -> the (l=0, m=0) scalar slice (N, F, C); n_frames == 1 + return coeff[:, 0, :, :] + + def _to_grid(self, coeff: Any) -> Any: + # einsum "gd,ndfc->ngfc" (n_frames == 1) as a broadcast batched matmul + xp = array_api_compat.array_namespace(coeff) + n_batch, coeff_dim, n_focus, _ = coeff.shape + to_grid_mat = xp.asarray( + self.projector.to_grid_mat[...], device=array_api_compat.device(coeff) + ) + if to_grid_mat.dtype != coeff.dtype: + to_grid_mat = xp.astype(to_grid_mat, coeff.dtype) + flat = xp.reshape(coeff, (n_batch, coeff_dim, n_focus * self.channels)) + out = xp.matmul(to_grid_mat[None, ...], flat) # (N, G, F*C) + return xp.reshape( + out, (n_batch, self.projector.grid_size, n_focus, self.channels) + ) + + def _from_grid(self, grid: Any) -> Any: + # einsum "dg,ngfc->ndfc" (n_frames == 1) as a broadcast batched matmul + xp = array_api_compat.array_namespace(grid) + n_batch, n_grid, n_focus, _ = grid.shape + coeff_dim = self.projector.coeff_dim + from_grid_mat = xp.asarray( + self.projector.from_grid_mat[...], device=array_api_compat.device(grid) + ) + if from_grid_mat.dtype != grid.dtype: + from_grid_mat = xp.astype(from_grid_mat, grid.dtype) + flat = xp.reshape(grid, (n_batch, n_grid, n_focus * self.channels)) + out = xp.matmul(from_grid_mat[None, ...], flat) # (N, D, F*C) + return xp.reshape(out, (n_batch, coeff_dim, n_focus, self.expanded_channels)) + + def _to_ndfc(self, value: Any) -> Any: + if self.layout == "ndfc": + return value + # "nfdc": (N, F, D, C) -> (N, D, F, C); "flat" is cross-only (blocked) + xp = array_api_compat.array_namespace(value) + return xp.permute_dims(value, (0, 2, 1, 3)) + + def _restore_layout(self, value: Any) -> Any: + if self.layout == "ndfc": + return value + xp = array_api_compat.array_namespace(value) + return xp.permute_dims(value, (0, 2, 1, 3)) + + def _check_last_dim(self, value: Any, expected: int, name: str) -> None: + if value.shape[-1] != expected: + raise ValueError( + f"`{name}` last dimension must be {expected}, got {value.shape[-1]}" + ) + + +class S2GridNet(BaseGridNet): + """Grid net using an S2 spherical-harmonic projector (Lebedev only). + + Parameters + ---------- + lmax : int + Maximum spherical harmonic degree. + mmax : int | None + Maximum order kept in the coefficient layout. If None, use ``lmax``. + channels : int + Number of channels per (l, m) coefficient. + n_focus : int + Number of focus streams. + mode : str + Pairing mode; only ``"self"`` is ported. + op_type : str + Point-wise grid operation; ``"glu"`` or ``"branch"`` (``"mlp"`` is + not ported). + precision : str + Parameter precision. + layout : str + Tensor layout convention: ``"ndfc"`` or ``"nfdc"``. + grid_resolution_list : list[int] | None + Lebedev ``[precision, n_points]`` pair; resolved automatically if None. + coefficient_layout : str + ``"packed"`` or ``"m_major"`` coefficient ordering. + grid_method : str + S2 quadrature backend; only ``"lebedev"`` is ported. + grid_branches : int + Number of scalar-routed branches when ``op_type='branch'``. + residual_scale_init : float | None + Not ported (cross-mode only); must be None. + mlp_bias : bool + Whether to use bias in the scalar gate projection. + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for weight initialization. + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None = None, + channels: int, + n_focus: int = 1, + mode: str, + op_type: str, + precision: str = DEFAULT_PRECISION, + layout: str, + grid_resolution_list: list[int] | None = None, + coefficient_layout: str = "packed", + # Deliberate divergence from pt's default ("e3nn"): the e3nn + # product-grid branch is not ported to dpmodel and always raises, so + # the only usable default here is "lebedev". Checkpoint compatibility + # is unaffected because serialize always records the explicit value. + grid_method: str = "lebedev", + grid_branches: int = 1, + residual_scale_init: float | None = None, + mlp_bias: bool = False, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + projector = S2GridProjector( + lmax=lmax, + mmax=mmax, + precision=precision, + grid_resolution_list=grid_resolution_list, + coefficient_layout=coefficient_layout, + grid_method=grid_method, + ) + self.grid_resolution_list = projector.grid_resolution_list + self.grid_method = projector.grid_method + self.grid_branches = int(grid_branches) + super().__init__( + projector=projector, + channels=channels, + n_focus=n_focus, + mode=mode, + op_type=op_type, + precision=precision, + layout=layout, + mlp_bias=mlp_bias, + trainable=trainable, + grid_branches=grid_branches, + residual_scale_init=residual_scale_init, + seed=seed, + ) + + def serialize(self) -> dict[str, Any]: + """Serialize the S2GridNet to a dict. + + The pt ``S2GridNet`` has no ``serialize()``; the ``@variables`` keys + here match the pt ``state_dict`` key names (the projector matrices + are non-persistent buffers in pt and are rebuilt from the config). + """ + variables = {"scalar_gate.weight": to_numpy_array(self.scalar_gate.weight)} + if self.mlp_bias: + variables["scalar_gate.bias"] = to_numpy_array(self.scalar_gate.bias) + if self.op_type == "branch": + grid_op_data = self.grid_op.serialize()["@variables"] + for key, value in grid_op_data.items(): + variables[f"grid_op.{key}"] = value + return { + "@class": "S2GridNet", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.projector.mmax, + "channels": self.channels, + "n_focus": self.n_focus, + "mode": self.mode, + "op_type": self.op_type, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "layout": self.layout, + "grid_resolution_list": self.grid_resolution_list, + "coefficient_layout": self.projector.coefficient_layout, + "grid_method": self.grid_method, + "grid_branches": self.grid_branches, + "mlp_bias": self.mlp_bias, + "trainable": self.trainable, + "seed": None, + }, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> S2GridNet: + """Deserialize an S2GridNet from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "S2GridNet": + raise ValueError(f"Invalid class for S2GridNet: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + lmax=int(config["lmax"]), + mmax=int(config["mmax"]), + channels=int(config["channels"]), + n_focus=int(config["n_focus"]), + mode=str(config["mode"]), + op_type=str(config["op_type"]), + precision=str(config["precision"]), + layout=str(config["layout"]), + grid_resolution_list=config["grid_resolution_list"], + coefficient_layout=str(config["coefficient_layout"]), + grid_method=str(config["grid_method"]), + grid_branches=int(config["grid_branches"]), + mlp_bias=bool(config["mlp_bias"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + prec = PRECISION_DICT[obj.precision.lower()] + weight = np.asarray(variables["scalar_gate.weight"], dtype=prec) + if weight.shape != obj.scalar_gate.weight.shape: + raise ValueError( + f"scalar_gate.weight shape {weight.shape} does not match " + f"the expected shape {obj.scalar_gate.weight.shape}" + ) + obj.scalar_gate.weight = weight + if obj.mlp_bias: + obj.scalar_gate.bias = np.asarray( + variables["scalar_gate.bias"], dtype=prec + ).reshape(obj.scalar_gate.bias.shape) + if obj.op_type == "branch": + obj.grid_op._load_variables( + { + key[len("grid_op.") :]: value + for key, value in variables.items() + if key.startswith("grid_op.") + } + ) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/indexing.py b/deepmd/dpmodel/descriptor/dpa4_nn/indexing.py new file mode 100644 index 0000000000..179dcefd9f --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/indexing.py @@ -0,0 +1,478 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +SO(3) packed-index and projection helpers for DPA4/SeZM. + +This module defines the packed `(l, m)` indexing helpers and the projection +utilities used by the DPA4 equivariant operators. It is the dpmodel port of +``deepmd.pt.model.descriptor.sezm_nn.indexing``. + +The index-table builders run at module-init time on static index data and are +implemented in plain numpy by design (not array-API); they return ``np.int64`` +arrays. The torch-specific ``device``/``dtype`` keyword parameters of the pt +versions are dropped for those builders. Only ``project_D_to_m`` and +``project_Dt_from_m`` operate on runtime tensors and are array-API compatible. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + + +def get_so3_dim_of_lmax(lmax: int) -> int: + """ + Return SO(3) representation dimension for given lmax. + + The dimension equals:: + + sum_{l<=lmax} (2l+1) = (lmax+1)^2 + + which is the number of spherical harmonics basis functions. + + Parameters + ---------- + lmax + Maximum spherical harmonic degree. + + Returns + ------- + int + The SO(3) dimension D = (lmax+1)^2. + """ + return int((int(lmax) + 1) ** 2) + + +def map_degree_idx(lmax: int) -> np.ndarray: + """ + Build degree (l) index for each position in the packed (l, m) layout. + + For each spherical harmonic coefficient position in the packed tensor, + returns the corresponding angular momentum quantum number l. + + The torch version's ``device`` parameter is dropped: the output is a static + numpy table. + + Examples + -------- + For lmax=2, the packed layout has D=9 positions: + - Position 0: l=0, m=0 + - Positions 1-3: l=1, m=-1,0,+1 + - Positions 4-8: l=2, m=-2,-1,0,+1,+2 + + Returns: [0, 1,1,1, 2,2,2,2,2] + + Parameters + ---------- + lmax + Maximum angular momentum degree. + + Returns + ------- + np.ndarray + ``np.int64`` array with shape (D,), where D=(lmax+1)^2. + Each element is the l value for that position. + """ + lmax = int(lmax) + counts = np.array([2 * degree + 1 for degree in range(lmax + 1)], dtype=np.int64) + return np.repeat(np.arange(lmax + 1, dtype=np.int64), counts) + + +def build_gie_zonal_index(lmax: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Build node-level packed indices for GIE zonal coupling. + + The returned arrays are aligned row-wise for every non-scalar packed + coefficient in the node representation. They select the local ``m=0`` column + of the matching degree from ``Dt_full`` or an equivalent zonal coupling table. + + The torch version's ``device`` parameter is dropped: the output is a static + numpy table. + + Parameters + ---------- + lmax + Maximum node degree used by the geometric initial embedding. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + ``(node_row_index, node_zonal_m0_col_index, node_radial_l_index)``. + The first two index packed SO(3) rows/columns; the last one indexes + radial features with degree slots ``l=1..lmax`` stored as ``0..lmax-1``. + All are ``np.int64`` arrays. + """ + lmax_i = int(lmax) + ebed_dim = get_so3_dim_of_lmax(lmax_i) + if lmax_i == 0: + empty = np.empty(0, dtype=np.int64) + return empty, empty, empty + + packed_degree_by_row = map_degree_idx(lmax_i) + node_row_index = np.arange(1, ebed_dim, dtype=np.int64) + node_degree_by_row = packed_degree_by_row[1:] + node_zonal_m0_col_index = node_degree_by_row * (node_degree_by_row + 1) + node_radial_l_index = node_degree_by_row - 1 + return node_row_index, node_zonal_m0_col_index, node_radial_l_index + + +def project_D_to_m( + D_full: Any, + coeff_index_m: Any, + ebed_dim_full: int, + cache: dict[str, Any] | None, + key_lmax: int, + key_mmax: int, +) -> Any: + """ + Row-project block-diagonal Wigner-D to the m-major truncated layout. + + This function operates on runtime tensors and is array-API compatible. + + Parameters + ---------- + D_full + Block-diagonal Wigner-D with shape (E, D, D). + coeff_index_m + Indices for m-major reduced layout with shape (D_m_trunc,). + ebed_dim_full + Full SO(3) dimension D_full = (lmax+1)^2 to slice the block. + cache + Optional cache mapping (lmax, mmax) -> projected matrix. + key_lmax + lmax used to build coeff_index_m (cache key). + key_mmax + mmax used to build coeff_index_m (cache key). + + Returns + ------- + Array + Projected rotation matrix with shape (E, D_m_trunc, D). + + Examples + -------- + For lmax=2, mmax=1 (D=9, D_m_trunc=7), coeff_index_m selects + [0,2,6,1,5,3,7] in packed (l,m) order. The returned array keeps only those + rows of ``D_full`` while retaining all columns, so that rotating and truncating + is done in a single matmul: ``x_local = D_to_m @ x_global``. + """ + cache_key = f"{int(key_lmax)}:{int(key_mmax)}" + if cache is not None: + cached = cache.get(cache_key) + if cached is not None: + return cached + + xp = array_api_compat.array_namespace(D_full) + D_block = D_full[:, :ebed_dim_full, :ebed_dim_full] + index = xp.asarray(coeff_index_m, device=array_api_compat.device(D_full)) + proj = xp.take(D_block, index, axis=1) + if cache is not None: + cache[cache_key] = proj + return proj + + +def project_Dt_from_m( + Dt_full: Any, + coeff_index_m: Any, + ebed_dim_full: int, + cache: dict[str, Any] | None, + key_lmax: int, + key_mmax: int, +) -> Any: + """ + Column-project block-diagonal Wigner-D^T for inverse rotation. + + This function operates on runtime tensors and is array-API compatible. + + Parameters + ---------- + Dt_full + Block-diagonal Wigner-D^T with shape (E, D, D). + coeff_index_m + Indices for m-major reduced layout with shape (D_m_trunc,). + ebed_dim_full + Full SO(3) dimension D_full = (lmax+1)^2 to slice the block. + cache + Optional cache mapping (lmax, mmax) -> projected matrix. + key_lmax + lmax used to build coeff_index_m (cache key). + key_mmax + mmax used to build coeff_index_m (cache key). + + Returns + ------- + Array + Projected inverse rotation matrix with shape (E, D, D_m_trunc). + + Examples + -------- + Continuing lmax=2, mmax=1, the projection selects the same column subset + [0,2,6,1,5,3,7] from ``Dt_full``. This enables inverse rotation with missing + coefficients implicitly zeroed: ``x_global = Dt_from_m @ x_local``. + """ + cache_key = f"{int(key_lmax)}:{int(key_mmax)}" + if cache is not None: + cached = cache.get(cache_key) + if cached is not None: + return cached + + xp = array_api_compat.array_namespace(Dt_full) + Dt_block = Dt_full[:, :ebed_dim_full, :ebed_dim_full] + index = xp.asarray(coeff_index_m, device=array_api_compat.device(Dt_full)) + proj = xp.take(Dt_block, index, axis=2) + if cache is not None: + cache[cache_key] = proj + return proj + + +def so3_packed_index(degree: int, m: int) -> int: + """ + Compute packed (l, m) index for real spherical harmonics layout. + + The packed layout is l-primary with m ordered as ``-l..+l`` inside each l-block. + The index formula is:: + + idx(l, m) = l^2 + l + m + + Parameters + ---------- + degree + Degree l. + m + Order m, must satisfy ``-l <= m <= l``. + + Returns + ------- + int + Packed index. + """ + degree = int(degree) + m = int(m) + return degree * degree + degree + m + + +def build_l_major_index(lmax: int, mmax: int) -> np.ndarray: + """ + Build coefficient indices for l-major layout truncated by mmax. + + The returned indices select coefficients with ``|m| <= min(mmax, l)`` in the + standard packed (l, m) layout. The order is l-major: + + - l = 0..lmax + - within each l, m = -min(mmax, l) .. +min(mmax, l) + + The torch version's ``device`` parameter is dropped: the output is a static + numpy table. + + Parameters + ---------- + lmax + Maximum degree. + mmax + Maximum order (|m|). Must satisfy ``0 <= mmax <= lmax``. + + Returns + ------- + np.ndarray + ``np.int64`` array of indices with shape (D_m_trunc,), selecting + coefficients from the full packed layout with D=(lmax+1)^2, where + D_m_trunc is the number of coefficients kept under ``|m| <= min(mmax, l)``. + + Examples + -------- + For lmax=2, mmax=1: + - Full packed layout: l=0(0), l=1(1-3), l=2(4-8) + - Truncated by mmax=1: skip (l=2, m=±2) at indices 4,8 + - Returns: [0, 1, 2, 3, 5, 6, 7] + """ + lmax_i = int(lmax) + mmax_i = int(mmax) + if lmax_i < 0: + raise ValueError("`lmax` must be non-negative") + if mmax_i < 0: + raise ValueError("`mmax` must be non-negative") + if mmax_i > lmax_i: + raise ValueError("`mmax` must be <= `lmax`") + + indices: list[int] = [] + for degree in range(lmax_i + 1): + m_keep = min(mmax_i, degree) + for m in range(-m_keep, m_keep + 1): + indices.append(so3_packed_index(degree, m)) + return np.asarray(indices, dtype=np.int64) + + +def build_m_major_index(lmax: int, mmax: int) -> np.ndarray: + """ + Build coefficient indices for m-major layout truncated by mmax. + + This layout minimizes rotation cost and avoids gather-heavy indexing: + + - m = 0: l = 0..lmax (single coefficient per l) + - for each m = 1..mmax: + - negative part: l = m..lmax, coefficient (l, -m) + - positive part: l = m..lmax, coefficient (l, +m) + + The torch version's ``device`` parameter is dropped: the output is a static + numpy table. + + Parameters + ---------- + lmax + Maximum degree. + mmax + Maximum order (|m|). Must satisfy ``0 <= mmax <= lmax``. + + Returns + ------- + np.ndarray + ``np.int64`` array of indices with shape (D_m_trunc,), selecting + coefficients from the full packed layout with D=(lmax+1)^2, where + D_m_trunc is the number of coefficients kept under ``|m| <= min(mmax, l)``. + + Examples + -------- + For lmax=2, mmax=1: + - m=0 group: (l=0,m=0)→0, (l=1,m=0)→2, (l=2,m=0)→6 + - m=1 neg group: (l=1,m=-1)→1, (l=2,m=-1)→5 + - m=1 pos group: (l=1,m=+1)→3, (l=2,m=+1)→7 + - Returns: [0, 2, 6, 1, 5, 3, 7] + """ + lmax_i = int(lmax) + mmax_i = int(mmax) + if lmax_i < 0: + raise ValueError("`lmax` must be non-negative") + if mmax_i < 0: + raise ValueError("`mmax` must be non-negative") + if mmax_i > lmax_i: + raise ValueError("`mmax` must be <= `lmax`") + + indices: list[int] = [] + # === Step 1. m = 0 group (l = 0..lmax) === + for degree in range(lmax_i + 1): + indices.append(so3_packed_index(degree, 0)) + + # === Step 2. m > 0 groups (neg then pos) === + for m in range(1, mmax_i + 1): + for degree in range(m, lmax_i + 1): + indices.append(so3_packed_index(degree, -m)) + for degree in range(m, lmax_i + 1): + indices.append(so3_packed_index(degree, m)) + + return np.asarray(indices, dtype=np.int64) + + +def build_m_major_l_index(lmax: int, mmax: int) -> np.ndarray: + """ + Build degree (l) index aligned with `build_m_major_index`. + + The torch version's ``device`` parameter is dropped: the output is a static + numpy table. + + Parameters + ---------- + lmax + Maximum degree. + mmax + Maximum order (|m|). Must satisfy ``0 <= mmax <= lmax``. + + Returns + ------- + np.ndarray + ``np.int64`` array of degrees with shape (D_m_trunc,). Entry i is the + degree l for the i-th coefficient in the m-major layout. + + Examples + -------- + For lmax=2, mmax=1: + - m=0 group: l=0,1,2 + - m=1 neg group: l=1,2 + - m=1 pos group: l=1,2 + - Returns: [0, 1, 2, 1, 2, 1, 2] + """ + lmax_i = int(lmax) + mmax_i = int(mmax) + if lmax_i < 0: + raise ValueError("`lmax` must be non-negative") + if mmax_i < 0: + raise ValueError("`mmax` must be non-negative") + if mmax_i > lmax_i: + raise ValueError("`mmax` must be <= `lmax`") + + degrees: list[int] = [] + # === Step 1. m = 0 group === + for degree in range(lmax_i + 1): + degrees.append(degree) + + # === Step 2. m > 0 groups (neg then pos) === + for m in range(1, mmax_i + 1): + for degree in range(m, lmax_i + 1): + degrees.append(degree) + for degree in range(m, lmax_i + 1): + degrees.append(degree) + + return np.asarray(degrees, dtype=np.int64) + + +def build_rotate_inv_rescale( + lmax: int, + mmax: int, + degree_index: np.ndarray, + *, + dtype: Any = np.float64, +) -> np.ndarray: + """ + Build reduced-layout inverse-rotation rescale factors. + + When ``mmax < lmax``, the reduced local layout keeps only ``2*mmax+1`` orders + for each degree ``l > mmax``. The inverse rotation rescales those truncated + degrees by ``sqrt((2*l+1)/(2*mmax+1))`` so the reduced representation matches + the amplitude expected by the full SO(3) basis. + + The torch version's ``device`` parameter is dropped: the output is a static + numpy table. ``dtype`` is kept (as a numpy dtype) since the floating-point + precision of the rescale vector is meaningful. + + Parameters + ---------- + lmax + Maximum degree. + mmax + Maximum order (|m|). Must satisfy ``0 <= mmax <= lmax``. + degree_index + Degree index aligned with the reduced coefficient layout, typically + returned by ``build_m_major_l_index``. + dtype + Floating-point numpy dtype for the returned array. + + Returns + ------- + np.ndarray + Rescale vector with shape (D_m_trunc,), aligned with the reduced + coefficient layout. + """ + lmax_i = int(lmax) + mmax_i = int(mmax) + if lmax_i < 0: + raise ValueError("`lmax` must be non-negative") + if mmax_i < 0: + raise ValueError("`mmax` must be non-negative") + if mmax_i > lmax_i: + raise ValueError("`mmax` must be <= `lmax`") + + degrees = np.asarray(degree_index, dtype=np.int64) + rescale = np.ones(degrees.shape[0], dtype=dtype) + if mmax_i == lmax_i: + return rescale + + mask = degrees > mmax_i + if mask.any(): + denom = float(2 * mmax_i + 1) + degree_values = degrees[mask].astype(dtype) + rescale[mask] = np.sqrt((2.0 * degree_values + 1.0) / denom) + return rescale diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/norm.py b/deepmd/dpmodel/descriptor/dpa4_nn/norm.py new file mode 100644 index 0000000000..420f6d469b --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/norm.py @@ -0,0 +1,653 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Normalization layers for the DPA4/SeZM descriptor. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.norm``. +All four pt norm classes are ported: ``RMSNorm`` (used by ``radial.RadialMLP``), +``EquivariantRMSNorm`` (used by ``block``), ``ReducedEquivariantRMSNorm`` and +``ScalarRMSNorm`` (used by ``so2``). + +Serialization contract: the ``@variables`` keys of each class match the +``state_dict`` key names of its pt counterpart, so pt ``serialize()`` output +deserializes directly into the dpmodel classes (and vice versa). +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .indexing import ( + map_degree_idx, +) + + +class RMSNorm(NativeOP): + """ + Generic RMSNorm on tensors with shape `(..., C)`. + + This is the plain channel-wise RMS normalization used for non-equivariant + branches whose last axis stores feature channels. A learnable affine scale + is applied on the channel axis only, while all leading axes are treated as + batch dimensions. + + Parameters + ---------- + channels : int + Feature dimension of the last axis. + eps : float + Small epsilon for numerical stability. + precision : str + Parameter and computation precision. Caller should pass a compute + precision (fp32+) for numerical stability. + trainable : bool + Whether parameters are trainable. + """ + + def __init__( + self, + *, + channels: int, + eps: float = 1e-7, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + ) -> None: + self.channels = int(channels) + self.eps = float(eps) + self.precision = precision + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + # adam_ prefix routes this to Adam (no weight decay) in HybridMuon. + self.adam_scale = np.ones((self.channels,), dtype=prec) + + def call(self, x: Any) -> Any: + """ + Apply RMS normalization. + + Parameters + ---------- + x : Array + Input array with shape `(..., C)`. + + Returns + ------- + Array + Normalized array with shape `(..., C)`, same dtype as input. + """ + xp = array_api_compat.array_namespace(x) + scale = xp.asarray(self.adam_scale[...], device=array_api_compat.device(x)) + in_dtype = x.dtype + if in_dtype != scale.dtype: + x = xp.astype(x, scale.dtype) + inv_rms = 1.0 / xp.sqrt(xp.mean(x * x, axis=-1, keepdims=True) + self.eps) + out = x * inv_rms * scale + if out.dtype != in_dtype: + out = xp.astype(out, in_dtype) + return out + + def serialize(self) -> dict[str, Any]: + """Serialize the RMSNorm to a dict.""" + return { + "@class": "RMSNorm", + "@version": 1, + "config": { + "channels": self.channels, + "eps": self.eps, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + }, + "@variables": {"adam_scale": to_numpy_array(self.adam_scale)}, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> RMSNorm: + """Deserialize an RMSNorm from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "RMSNorm": + raise ValueError(f"Invalid class for RMSNorm: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + channels=int(config["channels"]), + eps=float(config["eps"]), + precision=str(config["precision"]), + trainable=bool(config["trainable"]), + ) + prec = PRECISION_DICT[obj.precision.lower()] + adam_scale = np.asarray(variables["adam_scale"], dtype=prec).reshape(-1) + if adam_scale.shape != obj.adam_scale.shape: + raise ValueError( + f"adam_scale shape {adam_scale.shape} does not match " + f"channels {obj.channels}" + ) + obj.adam_scale = adam_scale + return obj + + +class EquivariantRMSNorm(NativeOP): + """ + Degree-balanced equivariant RMS normalization on packed `(l, m)` layout. + + The scalar slice `l=0` is mean-centered across channels before the shared + RMS is evaluated. All coefficients, including the centered scalar slice, + contribute to the same per-sample and per-focus RMS. Degree balancing + assigns each coefficient from degree `l` the weight + `1 / ((2 * l + 1) * (lmax + 1))`, so each degree contributes equally + regardless of its multiplicity. A learnable per-focus, per-degree scale is + then expanded to all `m` coefficients, and a learnable bias is added only + to the scalar slice. + + Parameters + ---------- + lmax : int + Maximum spherical harmonic degree. + channels : int + Channels per `(l, m)` coefficient in each focus stream. + n_focus : int + Number of focus streams. Affine parameters are independent per focus. + eps : float + Small epsilon for numerical stability. + precision : str + Parameter and computation precision. Caller should pass a compute + precision (fp32+) for numerical stability. + trainable : bool + Whether parameters are trainable. + """ + + def __init__( + self, + lmax: int, + channels: int, + n_focus: int = 1, + *, + eps: float = 1e-5, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + ) -> None: + self.lmax = int(lmax) + self.channels = int(channels) + self.n_focus = int(n_focus) + self.eps = float(eps) + self.precision = precision + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + + # === Step 1. Learnable Parameters === + # Store affine scales in degree-major layout (L, F, C). This matches the + # packed output layout after degree expansion. + # adam_ prefix routes this to Adam (no weight decay) in HybridMuon. + self.adam_scale = np.ones( + (self.lmax + 1, self.n_focus, self.channels), dtype=prec + ) + # Bias only for l=0, independent per focus. + self.bias = np.zeros((self.n_focus, self.channels), dtype=prec) + + # === Step 2. Index and Weight Buffers === + self.expand_index = map_degree_idx(self.lmax) + + # Pre-fuse degree balancing and channel averaging into a single weight: + # w_d = 1 / ((2l+1) * (lmax+1) * C) + # so that the shared RMS statistic is a single weighted sum without + # allocating an intermediate (N, D, F, C) buffer beyond x^2 itself. + weights_list = [] + scale = 1.0 / ((self.lmax + 1) * self.channels) + for l in range(self.lmax + 1): + w = scale / (2 * l + 1) + weights_list.extend([w] * (2 * l + 1)) + self.balance_weight = np.asarray(weights_list, dtype=prec) + + def call(self, x: Any) -> Any: + """ + Apply degree-balanced equivariant RMS normalization. + + Parameters + ---------- + x : Array + Features with shape `(N, D, F, C)` where `D = (lmax + 1)^2`. + + Returns + ------- + Array + Normalized features with shape `(N, D, F, C)`, same dtype as input. + """ + xp = array_api_compat.array_namespace(x) + device = array_api_compat.device(x) + scale = xp.asarray(self.adam_scale[...], device=device) + bias = xp.asarray(self.bias[...], device=device) + balance_weight = xp.asarray( + self.balance_weight, device=array_api_compat.device(x) + ) + in_dtype = x.dtype + if in_dtype != scale.dtype: + x = xp.astype(x, scale.dtype) + x0 = x[:, :1, :, :] # (N, 1, F, C) + xt = x[:, 1:, :, :] # (N, D-1, F, C) + + # === Step 1. Center the scalar slice === + x0 = x0 - xp.mean(x0, axis=-1, keepdims=True) + + # === Step 2. Compute a shared degree-balanced RMS === + mean_variance = xp.sum(x0 * x0, axis=(1, 3)) * balance_weight[0] + if self.lmax > 0: + mean_variance = mean_variance + xp.sum( + (xt * xt) * balance_weight[1:][None, :, None, None], axis=(1, 3) + ) + inv_rms = 1.0 / xp.sqrt(mean_variance + self.eps) + inv_rms = inv_rms[:, None, :, None] # (N, 1, F, 1) + + x0 = x0 * inv_rms + if self.lmax > 0: + xt = xt * inv_rms + + # === Step 3. Apply per-degree affine parameters === + expand_index = xp.asarray(self.expand_index, device=array_api_compat.device(x)) + expanded_scale = xp.take(scale, expand_index, axis=0) + expanded_scale = expanded_scale[None, ...] # (1, D, F, C) + x0 = x0 * expanded_scale[:, :1, :, :] + if self.lmax > 0: + xt = xt * expanded_scale[:, 1:, :, :] + + # === Step 4. Add scalar bias and restore layout === + bias0 = xp.reshape(bias, (1, 1, self.n_focus, -1)) # (1, 1, F, C) + x0 = x0 + bias0 + + out = x0 if self.lmax == 0 else xp.concat([x0, xt], axis=1) + if out.dtype != in_dtype: + out = xp.astype(out, in_dtype) + return out + + def serialize(self) -> dict[str, Any]: + """Serialize the EquivariantRMSNorm to a dict.""" + return { + "@class": "EquivariantRMSNorm", + "@version": 1, + "config": { + "lmax": self.lmax, + "channels": self.channels, + "n_focus": self.n_focus, + "eps": self.eps, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + }, + "@variables": { + "adam_scale": to_numpy_array(self.adam_scale), + "bias": to_numpy_array(self.bias), + "expand_index": to_numpy_array(self.expand_index), + "balance_weight": to_numpy_array(self.balance_weight), + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> EquivariantRMSNorm: + """Deserialize an EquivariantRMSNorm from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "EquivariantRMSNorm": + raise ValueError(f"Invalid class for EquivariantRMSNorm: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + lmax=int(config["lmax"]), + channels=int(config["channels"]), + n_focus=int(config["n_focus"]), + eps=float(config["eps"]), + precision=str(config["precision"]), + trainable=bool(config["trainable"]), + ) + prec = PRECISION_DICT[obj.precision.lower()] + expand_index = np.asarray(variables["expand_index"], dtype=np.int64) + if not np.array_equal(expand_index, obj.expand_index): + raise ValueError("expand_index does not match the lmax-derived table") + for name in ("adam_scale", "bias", "balance_weight"): + value = np.asarray(variables[name], dtype=prec) + if value.shape != getattr(obj, name).shape: + raise ValueError( + f"{name} shape {value.shape} does not match " + f"the expected shape {getattr(obj, name).shape}" + ) + setattr(obj, name, value) + return obj + + +class ReducedEquivariantRMSNorm(NativeOP): + """ + Degree-balanced equivariant RMS normalization on reduced m-major layout. + + The scalar slice `l=0` is mean-centered across channels before the shared + RMS is evaluated. All retained coefficients, including the centered scalar + slice, contribute to the same per-edge and per-focus RMS. Degree balancing + assigns each retained coefficient from degree `l` the weight + `1 / (n_coeff_l * (lmax + 1))`, where + `n_coeff_l = 2 * min(l, mmax) + 1` is the number of retained coefficients + for that degree in the reduced layout. A learnable per-focus, per-degree + scale is expanded with `degree_index_m`, and a learnable bias is added only + to the scalar slice. + + Parameters + ---------- + lmax : int + Maximum spherical harmonic degree. + mmax : int + Maximum order kept in the truncated layout. + channels : int + Number of channels per retained coefficient. + degree_index_m : np.ndarray + Degree index per coefficient in m-major truncated layout, with shape + `(D_m_trunc,)`. + n_focus : int + Number of focus streams. + eps : float + Epsilon for numerical stability. + precision : str + Parameter and computation precision. Caller should pass a compute + precision (fp32+) for numerical stability. + trainable : bool + Whether parameters are trainable. + """ + + def __init__( + self, + *, + lmax: int, + mmax: int, + channels: int, + degree_index_m: np.ndarray, + n_focus: int = 1, + eps: float = 1e-5, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + ) -> None: + self.lmax = int(lmax) + self.mmax = int(mmax) + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.channels = int(channels) + self.n_focus = int(n_focus) + self.eps = float(eps) + self.precision = precision + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + + self.degree_index_m = np.asarray(degree_index_m, dtype=np.int64) + + # Pre-fuse degree balancing and channel averaging into a single weight: + # w_d = 1 / (n_coeff_l * (lmax+1) * C) + # where n_coeff_l is the number of retained coefficients for degree l in + # the reduced layout. + weights = np.zeros(self.degree_index_m.size, dtype=prec) + scale = 1.0 / ((self.lmax + 1) * self.channels) + for l in range(self.lmax + 1): + n_coeff_l = 2 * min(l, self.mmax) + 1 + w_l = scale / float(n_coeff_l) + weights[self.degree_index_m == l] = w_l + if np.any(weights == 0): + raise ValueError( + "ReducedEquivariantRMSNorm: balance_weight has zeros; " + "degree_index_m may be invalid." + ) + self.balance_weight = weights + + # adam_ prefix routes this to Adam (no weight decay) in HybridMuon. + self.adam_scale = np.ones( + (self.n_focus, self.lmax + 1, self.channels), dtype=prec + ) + self.bias0 = np.zeros((self.n_focus, self.channels), dtype=prec) + + def call(self, x: Any) -> Any: + """ + Apply degree-balanced reduced-layout RMS normalization. + + Parameters + ---------- + x : Array + Input array with shape (E, F, D_m_trunc, C). + + Returns + ------- + Array + Normalized array with shape `(E, F, D_m_trunc, C)`, same dtype as + input. + """ + xp = array_api_compat.array_namespace(x) + device = array_api_compat.device(x) + scale = xp.asarray(self.adam_scale[...], device=device) + bias0_w = xp.asarray(self.bias0[...], device=device) + balance_weight = xp.asarray( + self.balance_weight, device=array_api_compat.device(x) + ) + in_dtype = x.dtype + if in_dtype != scale.dtype: + x = xp.astype(x, scale.dtype) + has_xt = self.degree_index_m.size > 1 + x0 = x[:, :, :1, :] # (E, F, 1, C) + xt = x[:, :, 1:, :] # (E, F, D_m_trunc-1, C) + + # === Step 1. Center the scalar slice === + x0 = x0 - xp.mean(x0, axis=-1, keepdims=True) + + # === Step 2. Compute a shared degree-balanced RMS === + mean_variance = xp.sum(x0 * x0, axis=(2, 3)) * balance_weight[0] + if has_xt: + mean_variance = mean_variance + xp.sum( + (xt * xt) * balance_weight[1:][None, None, :, None], axis=(2, 3) + ) + inv_rms = 1.0 / xp.sqrt(mean_variance + self.eps) + inv_rms = inv_rms[:, :, None, None] # (E, F, 1, 1) + + x0 = x0 * inv_rms + if has_xt: + xt = xt * inv_rms + + # === Step 3. Apply per-degree affine parameters === + degree_index_m = xp.asarray( + self.degree_index_m, device=array_api_compat.device(x) + ) + expanded_scale = xp.take(scale, degree_index_m, axis=1) + expanded_scale = expanded_scale[None, ...] # (1, F, D_m_trunc, C) + x0 = x0 * expanded_scale[:, :, :1, :] + if has_xt: + xt = xt * expanded_scale[:, :, 1:, :] + + # === Step 4. Add scalar bias and restore layout === + bias0 = xp.reshape(bias0_w, (1, self.n_focus, 1, -1)) # (1, F, 1, C) + x0 = x0 + bias0 + + out = xp.concat([x0, xt], axis=2) if has_xt else x0 + if out.dtype != in_dtype: + out = xp.astype(out, in_dtype) + return out + + def serialize(self) -> dict[str, Any]: + """Serialize the ReducedEquivariantRMSNorm to a dict.""" + return { + "@class": "ReducedEquivariantRMSNorm", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.mmax, + "channels": self.channels, + "degree_index_m": to_numpy_array(self.degree_index_m), + "n_focus": self.n_focus, + "eps": self.eps, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + }, + "@variables": { + "degree_index_m": to_numpy_array(self.degree_index_m), + "balance_weight": to_numpy_array(self.balance_weight), + "adam_scale": to_numpy_array(self.adam_scale), + "bias0": to_numpy_array(self.bias0), + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> ReducedEquivariantRMSNorm: + """Deserialize a ReducedEquivariantRMSNorm from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "ReducedEquivariantRMSNorm": + raise ValueError(f"Invalid class for ReducedEquivariantRMSNorm: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + lmax=int(config["lmax"]), + mmax=int(config["mmax"]), + channels=int(config["channels"]), + degree_index_m=np.asarray(config["degree_index_m"], dtype=np.int64), + n_focus=int(config["n_focus"]), + eps=float(config["eps"]), + precision=str(config["precision"]), + trainable=bool(config["trainable"]), + ) + prec = PRECISION_DICT[obj.precision.lower()] + degree_index_m = np.asarray(variables["degree_index_m"], dtype=np.int64) + if not np.array_equal(degree_index_m, obj.degree_index_m): + raise ValueError("degree_index_m variable does not match the config") + for name in ("balance_weight", "adam_scale", "bias0"): + value = np.asarray(variables[name], dtype=prec) + if value.shape != getattr(obj, name).shape: + raise ValueError( + f"{name} shape {value.shape} does not match " + f"the expected shape {getattr(obj, name).shape}" + ) + setattr(obj, name, value) + return obj + + +class ScalarRMSNorm(NativeOP): + """ + Lightweight per-focus RMSNorm for scalar branches. + + This is the unified scalar norm used by SeZM: + - `n_focus=1` naturally degenerates to the single-stream behavior. + - `n_focus>1` uses independent learnable scales per focus stream. + Bias is intentionally omitted to keep the gate paths minimal. + + Parameters + ---------- + channels : int + Feature dimension of the last axis. + n_focus : int + Number of focus streams. + eps : float + Small epsilon for numerical stability. + precision : str + Parameter and computation precision. Caller should pass a compute + precision (fp32+) for numerical stability. + trainable : bool + Whether parameters are trainable. + """ + + def __init__( + self, + *, + channels: int, + n_focus: int = 1, + eps: float = 1e-7, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + ) -> None: + self.channels = int(channels) + self.n_focus = int(n_focus) + self.eps = float(eps) + self.precision = precision + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + # adam_ prefix routes this to Adam (no weight decay) in HybridMuon. + self.adam_scale = np.ones((self.n_focus, self.channels), dtype=prec) + + def call(self, x: Any) -> Any: + """ + Apply per-focus RMS normalization. + + Parameters + ---------- + x : Array + Input array with shape (B, F, C) or (B, C) when `n_focus=1`. + + Returns + ------- + Array + Normalized array with the same shape as input and same dtype. + """ + xp = array_api_compat.array_namespace(x) + scale = xp.asarray(self.adam_scale[...], device=array_api_compat.device(x)) + in_dtype = x.dtype + if in_dtype != scale.dtype: + x = xp.astype(x, scale.dtype) + + inv_rms = 1.0 / xp.sqrt(xp.mean(x * x, axis=-1, keepdims=True) + self.eps) + x = x * inv_rms + if x.ndim == 2: + x = x * scale[0, :] + else: + x = x * scale[None, ...] + if x.dtype != in_dtype: + x = xp.astype(x, in_dtype) + return x + + def serialize(self) -> dict[str, Any]: + """Serialize the ScalarRMSNorm to a dict.""" + return { + "@class": "ScalarRMSNorm", + "@version": 1, + "config": { + "channels": self.channels, + "n_focus": self.n_focus, + "eps": self.eps, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + }, + "@variables": {"adam_scale": to_numpy_array(self.adam_scale)}, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> ScalarRMSNorm: + """Deserialize a ScalarRMSNorm from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "ScalarRMSNorm": + raise ValueError(f"Invalid class for ScalarRMSNorm: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + channels=int(config["channels"]), + n_focus=int(config["n_focus"]), + eps=float(config["eps"]), + precision=str(config["precision"]), + trainable=bool(config["trainable"]), + ) + prec = PRECISION_DICT[obj.precision.lower()] + adam_scale = np.asarray(variables["adam_scale"], dtype=prec) + adam_scale = adam_scale.reshape(obj.adam_scale.shape) + obj.adam_scale = adam_scale + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/projection.py b/deepmd/dpmodel/descriptor/dpa4_nn/projection.py new file mode 100644 index 0000000000..bc1ce5add7 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/projection.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +S2 grid projection helpers for DPA4/SeZM function-space nonlinearities. + +This module is the dpmodel port of +``deepmd.pt.model.descriptor.sezm_nn.projection``, restricted to the Lebedev +S2 quadrature path used by the core DPA4 configuration +(``lebedev_quadrature=True``). The projectors only handle basis transforms: +a projector maps coefficient tensors to a fixed quadrature grid, and maps +grid fields back to coefficients with the matching quadrature rule. + +Ported names: ``BaseGridProjector``, ``S2GridProjector`` (Lebedev branch), +``resolve_s2_grid_resolution`` (as-is, both methods — pure arithmetic), and +``_normalize_s2_grid_resolution``. + +Skipped names (SO(3) Wigner-D grid machinery; consumed only by +``SO3GridNet`` in pt ``grid_net.py``, which backs the ``node_wise_so3``, +``message_node_so3``, and ``ffn_so3_grid`` paths — all disabled in the core +DPA4 config): ``SO3GridProjector``, ``resolve_so3_grid``, +``_build_so3_frame_set``. + +Not ported (guarded): the e3nn product-grid branch of ``S2GridProjector`` +(``grid_method="e3nn"``, i.e. ``lebedev_quadrature=False``) raises +``NotImplementedError`` at construction. Only the Lebedev path reproduces +to-grid/from-grid roundtrip identities at machine precision. + +The Lebedev projection matrices are assembled at init time with pure numpy: +``load_lebedev_rule`` replaces the pt Lebedev loader (same packaged data) and +``real_spherical_harmonics`` exactly replaces the e3nn call +``spherical_harmonics(list(range(lmax+1)), points, normalize=True, +normalization="norm")``, so the buffers match the pt float64 buffers to +machine precision. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.utils.lebedev import ( + LEBEDEV_PRECISION_TO_NPOINTS, + load_lebedev_rule, +) +from deepmd.dpmodel.utils.spherical_harmonics import ( + real_spherical_harmonics, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .indexing import ( + build_l_major_index, + build_m_major_index, +) + + +class BaseGridProjector(NativeOP): + """ + Base class for fixed coefficient-to-grid projection matrices. + + Subclasses build ``to_grid_mat`` with shape ``(G, J)`` and + ``from_grid_mat`` with shape ``(J, G)``, where ``G`` is the number of grid + samples and ``J`` is the flattened coefficient axis consumed by the grid + net. For ordinary S2 projections, ``J`` is the SO(3) feature coefficient + axis: ``D = (lmax + 1)^2`` in packed layout, or the retained ``D_m`` axis + in m-major layout. + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None, + precision: str = DEFAULT_PRECISION, + n_frames: int, + coefficient_layout: str, + ) -> None: + self.lmax = int(lmax) + self.mmax = int(self.lmax if mmax is None else mmax) + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.coefficient_layout = str(coefficient_layout).lower() + if self.coefficient_layout not in {"packed", "m_major"}: + raise ValueError( + "`coefficient_layout` must be either 'packed' or 'm_major'" + ) + self.precision = precision + self.n_frames = int(n_frames) + self.packed_dim = int((self.lmax + 1) ** 2) + + coeff_index = self._build_coefficient_index() + to_grid_mat, from_grid_mat = self._build_projection_mats(coeff_index) + self.coeff_dim = int(to_grid_mat.shape[1]) + self.grid_size = int(to_grid_mat.shape[0]) + if self.coeff_dim != int(from_grid_mat.shape[0]): + raise ValueError("Projection matrix coefficient axes `J` do not match") + if self.grid_size != int(from_grid_mat.shape[1]): + raise ValueError("Projection matrix grid axes `G` do not match") + prec = PRECISION_DICT[self.precision.lower()] + self.to_grid_mat = np.ascontiguousarray(to_grid_mat).astype(prec) + self.from_grid_mat = np.ascontiguousarray(from_grid_mat).astype(prec) + + def call(self, *args: Any, **kwargs: Any) -> Any: + """Projectors expose ``to_grid``/``from_grid``; there is no forward.""" + raise NotImplementedError( + "BaseGridProjector has no forward; use `to_grid` or `from_grid`" + ) + + def to_grid(self, embedding: Any) -> Any: + """Project flattened coefficients ``(N, J, C)`` to grid fields ``(N, G, C)``.""" + xp = array_api_compat.array_namespace(embedding) + to_grid_mat = xp.asarray( + self.to_grid_mat[...], device=array_api_compat.device(embedding) + ) + if to_grid_mat.dtype != embedding.dtype: + to_grid_mat = xp.astype(to_grid_mat, embedding.dtype) + # einsum "gj,njc->ngc" as a broadcast batched matmul + return xp.matmul(to_grid_mat[None, ...], embedding) + + def from_grid(self, grid: Any) -> Any: + """Project grid fields ``(N, G, C)`` back to flattened coefficients ``(N, J, C)``.""" + xp = array_api_compat.array_namespace(grid) + from_grid_mat = xp.asarray( + self.from_grid_mat[...], device=array_api_compat.device(grid) + ) + if from_grid_mat.dtype != grid.dtype: + from_grid_mat = xp.astype(from_grid_mat, grid.dtype) + # einsum "jg,ngc->njc" as a broadcast batched matmul + return xp.matmul(from_grid_mat[None, ...], grid) + + def _build_coefficient_index(self) -> np.ndarray: + """Build the coefficient subset consumed by the projector matrices.""" + if self.coefficient_layout == "m_major": + return build_m_major_index(self.lmax, self.mmax) + if self.mmax == self.lmax: + return np.arange((self.lmax + 1) ** 2, dtype=np.int64) + return build_l_major_index(self.lmax, self.mmax) + + def _build_projection_mats( + self, + coeff_index: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + """Build ``to_grid_mat (G, J)`` and ``from_grid_mat (J, G)``.""" + raise NotImplementedError + + +class S2GridProjector(BaseGridProjector): + """ + Project SO(3) coefficients to/from a flattened S2 grid (Lebedev only). + + Parameters + ---------- + lmax + Maximum spherical harmonic degree. + mmax + Maximum order kept in the coefficient layout. If None, use ``lmax``. + precision + Buffer precision used by the projection matrices. + grid_resolution_list + Two-element resolution list ``[precision, n_points]`` for + ``grid_method='lebedev'``. If None, resolved automatically. + coefficient_layout + Coefficient ordering expected by the caller: + - ``"packed"``: packed ``(l, m)`` order, optionally truncated by ``mmax``. + - ``"m_major"``: reduced m-major order used inside ``SO2Convolution``. + grid_method + S2 quadrature backend. Must be ``"e3nn"`` or ``"lebedev"``; only + ``"lebedev"`` (``lebedev_quadrature=True``) is ported to dpmodel. + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None = None, + precision: str = DEFAULT_PRECISION, + grid_resolution_list: list[int] | None = None, + coefficient_layout: str = "packed", + grid_method: str = "e3nn", + ) -> None: + lmax_i = int(lmax) + mmax_i = int(lmax_i if mmax is None else mmax) + self.grid_method = str(grid_method).lower() + if self.grid_method not in {"e3nn", "lebedev"}: + raise ValueError("`grid_method` must be either 'e3nn' or 'lebedev'") + if self.grid_method == "e3nn": + raise NotImplementedError( + "grid_method='e3nn' (lebedev_quadrature=False) is not ported " + "to dpmodel; use lebedev_quadrature=True" + ) + + self.grid_resolution_list = _normalize_s2_grid_resolution( + lmax_i, + mmax_i, + grid_resolution_list, + method=self.grid_method, + ) + self.phi_resolution = 0 + self.theta_resolution = 0 + self.lebedev_precision, self.lebedev_npoints = self.grid_resolution_list + + super().__init__( + lmax=lmax_i, + mmax=mmax_i, + precision=precision, + n_frames=1, + coefficient_layout=coefficient_layout, + ) + + def _rescale_truncated_matrix(self, mat: np.ndarray) -> None: + if self.lmax == self.mmax: + return + for degree in range(self.lmax + 1): + if degree <= self.mmax: + continue + start_idx = degree * degree + length = 2 * degree + 1 + rescale = math.sqrt(length / float(2 * self.mmax + 1)) + mat[:, start_idx : start_idx + length] *= rescale + + def _build_projection_mats( + self, + coeff_index: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + points, weights = load_lebedev_rule(self.lebedev_precision) + points = np.asarray(points, dtype=np.float64) + weights = np.asarray(weights, dtype=np.float64) + # exact numpy replacement for e3nn spherical_harmonics( + # list(range(lmax + 1)), points, normalize=True, normalization="norm") + harmonics = real_spherical_harmonics(points, self.lmax) + # Match the component-normalized product-grid convention used by + # e3nn's ToS2Grid/FromS2Grid pair so both S2 backends are drop-in + # replacements for the same grid net. + scale = math.sqrt(float(self.lmax + 1)) + degree_factors = np.asarray( + [ + float(2 * degree + 1) + for degree in range(self.lmax + 1) + for _ in range(2 * degree + 1) + ], + dtype=np.float64, + ) + to_grid_mat = harmonics / scale + from_grid_mat = harmonics * (weights[:, None] * scale * degree_factors[None, :]) + self._rescale_truncated_matrix(to_grid_mat) + self._rescale_truncated_matrix(from_grid_mat) + + to_grid_mat = to_grid_mat[:, coeff_index] + from_grid_mat = from_grid_mat[:, coeff_index].T + return to_grid_mat, from_grid_mat + + def serialize(self) -> dict[str, Any]: + """Serialize the S2GridProjector to a dict (pt-compatible format).""" + return { + "@class": "S2GridProjector", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.mmax, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "grid_resolution_list": self.grid_resolution_list, + "coefficient_layout": self.coefficient_layout, + "grid_method": self.grid_method, + }, + "@variables": {}, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> S2GridProjector: + """Deserialize an S2GridProjector from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "S2GridProjector": + raise ValueError(f"Invalid class for S2GridProjector: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + data.pop("@variables", None) + return cls( + lmax=int(config["lmax"]), + mmax=int(config["mmax"]), + precision=str(config["precision"]), + grid_resolution_list=config["grid_resolution_list"], + coefficient_layout=str(config["coefficient_layout"]), + grid_method=str(config["grid_method"]), + ) + + +def resolve_s2_grid_resolution( + lmax: int, + mmax: int, + *, + method: str = "e3nn", +) -> list[int]: + """ + Resolve the default S2 grid resolution. + + For ``method='e3nn'``, the automatic default uses even azimuthal sampling + ``R_phi = 2 * mmax + 4`` and even polar sampling + ``R_theta = ceil_even(3 * lmax + 2)``. + + For ``method='lebedev'``, the automatic default picks the smallest packaged + Lebedev rule whose algebraic precision is at least ``3 * lmax`` and returns + ``[precision, n_points]``. + """ + method = str(method).lower() + if method not in {"e3nn", "lebedev"}: + raise ValueError("`method` must be either 'e3nn' or 'lebedev'") + if method == "lebedev": + required_precision = 3 * int(lmax) + for precision, n_points in LEBEDEV_PRECISION_TO_NPOINTS.items(): + if precision >= required_precision: + return [precision, n_points] + raise ValueError( + f"No packaged Lebedev rule has precision >= {required_precision}" + ) + + phi_resolution = 2 * int(mmax) + 4 + theta_resolution = 3 * int(lmax) + 2 + theta_resolution += theta_resolution % 2 + return [phi_resolution, theta_resolution] + + +def _normalize_s2_grid_resolution( + lmax: int, + mmax: int, + grid_resolution_list: list[int] | None, + *, + method: str, +) -> list[int]: + """Resolve default grids or validate already-resolved low-level grids.""" + method = str(method).lower() + if grid_resolution_list is None: + return resolve_s2_grid_resolution(lmax, mmax, method=method) + if method == "lebedev": + if len(grid_resolution_list) != 2: + raise ValueError( + "Lebedev `grid_resolution_list` must be [precision, n_points]" + ) + precision = int(grid_resolution_list[0]) + n_points = int(grid_resolution_list[1]) + expected_n_points = LEBEDEV_PRECISION_TO_NPOINTS.get(precision) + if expected_n_points != n_points: + raise ValueError( + "Lebedev `grid_resolution_list` must match a packaged " + f"[precision, n_points] pair; got [{precision}, {n_points}]" + ) + return [precision, n_points] + + if len(grid_resolution_list) != 2: + raise ValueError("`grid_resolution_list` must contain two integers") + resolution = [int(grid_resolution_list[0]), int(grid_resolution_list[1])] + if resolution[0] < 1 or resolution[1] < 1: + raise ValueError("grid resolutions must be positive") + return resolution diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py new file mode 100644 index 0000000000..b93836ea8e --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Radial building blocks for the DPA4/SeZM descriptor. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.radial``. +It defines the cutoff envelope, radial basis, and radial multilayer perceptron +used by the DPA4 descriptor. ``InnerClamp`` and ``BridgingSwitch`` are ported +by later tasks together with the modules that consume them. + +Serialization contract: the ``@variables`` keys of each class match the +``state_dict`` key names of its pt counterpart, so pt ``serialize()`` output +deserializes directly into the dpmodel classes (and vice versa). +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.network import ( + NativeLayer, + get_activation_fn, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .norm import ( + RMSNorm, +) + + +class RadialMLP(NativeOP): + """ + Radial MLP with channel RMSNorm and configurable activation. + + Parameters + ---------- + mlp_layers : list[int] + Layer sizes including input and output dimensions. + E.g., [in_dim, hidden1, hidden2, out_dim]. + activation_function : str + Activation function name (e.g., "silu", "tanh", "gelu"). + precision : str + Floating point precision for the linear layers. + trainable : bool + Whether the parameters are trainable. + seed : int | list[int] | None + Random seed for the layer initialization. + + Architecture + ------------ + Linear → RMSNorm → Activation for all hidden layers, + with the final layer being a plain Linear (no norm, no activation). + + Notes + ----- + All bias terms are disabled (Linear bias=False, RMSNorm bias-free) to + guarantee ``RadialMLP(0) = 0``. This is required because the compile path + pads masked edges with zero ``edge_rbf``; any non-zero bias would leak + spurious features into GIE scatter, causing energy divergence between + compile and non-compile paths. + """ + + def __init__( + self, + mlp_layers: list[int], + *, + activation_function: str = "silu", + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + if len(mlp_layers) < 2: + raise ValueError("`mlp_layers` must have at least 2 elements") + self.mlp_layers = [int(d) for d in mlp_layers] + self.activation_function = str(activation_function) + self.precision = precision + self.trainable = bool(trainable) + + n_layers = len(self.mlp_layers) + self.layers: list[NativeLayer] = [] + self.norms: list[RMSNorm] = [] + for i in range(n_layers - 1): + self.layers.append( + NativeLayer( + self.mlp_layers[i], + self.mlp_layers[i + 1], + bias=False, + activation_function=None, + precision=self.precision, + seed=child_seed(seed, i), + trainable=self.trainable, + ) + ) + # Last layer: no RMSNorm/activation + if i < n_layers - 2: + self.norms.append( + RMSNorm( + channels=self.mlp_layers[i + 1], + precision=self.precision, + trainable=self.trainable, + ) + ) + + def call(self, x: Any) -> Any: + """ + Forward pass. + + Parameters + ---------- + x : Array + Input array with shape (..., mlp_layers[0]). + + Returns + ------- + Array + Output array with shape (..., mlp_layers[-1]). + """ + n_hidden = len(self.norms) + for i, layer in enumerate(self.layers): + x = layer.call(x) + if i < n_hidden: + x = self.norms[i].call(x) + fn = get_activation_fn(self.activation_function) + x = fn(x) + return x + + def serialize(self) -> dict[str, Any]: + """Serialize the RadialMLP to a dict. + + The ``@variables`` keys follow the pt ``net.state_dict()`` naming: + ``{3*i}.matrix`` for the i-th linear layer and ``{3*i+1}.adam_scale`` + for the i-th RMSNorm (activation modules are parameter-free). + """ + variables: dict[str, np.ndarray] = {} + for i, layer in enumerate(self.layers): + variables[f"{3 * i}.matrix"] = to_numpy_array(layer.w) + if i < len(self.norms): + variables[f"{3 * i + 1}.adam_scale"] = to_numpy_array( + self.norms[i].adam_scale + ) + return { + "@class": "RadialMLP", + "@version": 1, + "mlp_layers": self.mlp_layers.copy(), + "activation_function": self.activation_function, + "dtype": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> RadialMLP: + """Deserialize a RadialMLP from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "RadialMLP": + raise ValueError(f"Invalid class for RadialMLP: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + variables = data.pop("@variables") + precision = str(data.pop("dtype")) + obj = cls( + data.pop("mlp_layers"), + activation_function=str(data.pop("activation_function")), + precision=precision, + trainable=bool(data.pop("trainable")), + ) + prec = PRECISION_DICT[precision.lower()] + expected_keys = {f"{3 * i}.matrix" for i in range(len(obj.layers))} | { + f"{3 * i + 1}.adam_scale" for i in range(len(obj.norms)) + } + if set(variables) != expected_keys: + raise ValueError( + f"variable keys {sorted(variables)} do not match the expected " + f"keys {sorted(expected_keys)}" + ) + for key, value in variables.items(): + idx_s, _, name = key.partition(".") + idx = int(idx_s) + value = np.asarray(value, dtype=prec) + if name == "matrix": + layer = obj.layers[idx // 3] + if value.shape != layer.w.shape: + raise ValueError( + f"shape of {key} {value.shape} does not match " + f"the layer shape {layer.w.shape}" + ) + layer.w = value + else: + norm = obj.norms[idx // 3] + norm.adam_scale = value.reshape(norm.adam_scale.shape) + return obj + + +class C3CutoffEnvelope(NativeOP): + """ + C^3-continuous polynomial cutoff envelope function. + + This envelope provides a smooth transition to zero at the cutoff radius, + ensuring continuity of the function value and the first three derivatives. + + Notes + ----- + The envelope function is defined for scaled distance ``x = r / rcut`` as:: + + E(x) = 1 + x^p * (a + b*x + c*x^2 + d*x^3), for x < 1 + E(x) = 0, for x >= 1 + + where the coefficients are chosen to satisfy:: + + E(0) = 1, E(1) = 0 + E'(1) = 0, E''(1) = 0, E'''(1) = 0 + + This ensures C^3 continuity at the cutoff boundary. The coefficients are:: + + a = -(p + 1)(p + 2)(p + 3) / 6 + b = p(p + 2)(p + 3) / 2 + c = -p(p + 1)(p + 3) / 2 + d = p(p + 1)(p + 2) / 6 + + For the default exponent p=5, the coefficients are a=-56, b=140, c=-120, + d=35:: + + E(x) = 1 + x^5 * (-56 + 140*x - 120*x^2 + 35*x^3) + = 1 - 56*x^5 + 140*x^6 - 120*x^7 + 35*x^8 + + Parameters + ---------- + rcut : float + Cutoff radius in Å. + exponent : int, optional + Polynomial exponent (p), must be positive. Default is 5. + precision : str + Floating point precision label (kept for config parity with pt; the + computation follows the input dtype). + """ + + def __init__( + self, + rcut: float, + exponent: int = 5, + *, + precision: str = DEFAULT_PRECISION, + ) -> None: + if rcut <= 0.0: + raise ValueError("`rcut` must be positive") + if exponent <= 0: + raise ValueError("`exponent` must be positive") + self.rcut = float(rcut) + self.p = int(exponent) + self.precision = precision + self.coeff_a = -((self.p + 1) * (self.p + 2) * (self.p + 3)) / 6.0 + self.coeff_b = (self.p * (self.p + 2) * (self.p + 3)) / 2.0 + self.coeff_c = -(self.p * (self.p + 1) * (self.p + 3)) / 2.0 + self.coeff_d = (self.p * (self.p + 1) * (self.p + 2)) / 6.0 + + def call(self, dst: Any) -> Any: + """Compute the envelope value for given distances.""" + xp = array_api_compat.array_namespace(dst) + d_scaled = xp.clip(dst / self.rcut, min=0.0, max=1.0) + poly = self.coeff_a + d_scaled * ( + self.coeff_b + d_scaled * (self.coeff_c + d_scaled * self.coeff_d) + ) + env_val = 1 + d_scaled**self.p * poly + return env_val * xp.astype(d_scaled < 1.0, dst.dtype) + + def serialize(self) -> dict[str, Any]: + """Serialize the C3CutoffEnvelope to a dict (config only, no state).""" + return { + "@class": "C3CutoffEnvelope", + "@version": 1, + "config": { + "rcut": self.rcut, + "exponent": self.p, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> C3CutoffEnvelope: + """Deserialize a C3CutoffEnvelope from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "C3CutoffEnvelope": + raise ValueError(f"Invalid class for C3CutoffEnvelope: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + return cls( + rcut=float(config["rcut"]), + exponent=int(config["exponent"]), + precision=str(config["precision"]), + ) + + +class RadialBasis(NativeOP): + """ + Radial basis with C^3 cutoff envelope. + + The trainable radial parameters are stored in ``adam_freqs`` so HybridMuon + routes them to Adam without weight decay. + + Notes + ----- + The Bessel basis uses the normalized sinc function for numerical + stability:: + + phi_n(r) = w_n * sinc(w_n * r / π) + + where ``sinc(z) = sin(π*z) / (π*z)`` with ``sinc(0) = 1`` (same convention + as ``torch.sinc`` and ``np.sinc``). This is mathematically equivalent to + the standard form ``sin(w_n * r) / r``, but sinc handles the r->0 limit, + providing continuous gradients without explicit epsilon clamping. + + The ``r -> 0`` limit is finite:: + + lim_{r->0} w_n * sinc(w_n * r / π) = w_n + + The initial Bessel frequencies follow a common spacing:: + + w_n = n * π / rcut, for n = 1..n_radial (in 1/Å) + + The C^3 cutoff envelope is multiplied directly into the output to ensure + strict smoothness at ``rcut``. + + Parameters + ---------- + rcut : float + Cutoff radius in Å. + basis_type : str, optional + Radial basis type. Supported values are ``"bessel"`` and ``"gaussian"``. + n_radial : int + Number of basis functions. + precision : str + Floating-point precision for the radial basis frequencies and outputs. + exponent : int, optional + Exponent for the C^3 cutoff envelope polynomial. Default is 7. + """ + + def __init__( + self, + rcut: float, + basis_type: str = "bessel", + n_radial: int = 10, + precision: str = DEFAULT_PRECISION, + exponent: int = 7, + ) -> None: + self.rcut = float(rcut) + if self.rcut <= 0.0: + raise ValueError("`rcut` must be positive") + self.n_radial = int(n_radial) + if self.n_radial <= 0: + raise ValueError("`n_radial` must be positive") + self.basis_type = str(basis_type).lower() + if self.basis_type not in ("bessel", "gaussian"): + raise ValueError("`basis_type` must be either 'bessel' or 'gaussian'") + self.precision = precision + self.exponent = int(exponent) + prec = PRECISION_DICT[self.precision.lower()] + + # Frequencies: n*π/rcut, n=1..n_radial + # Shape: (1, n_radial), stored as a trainable array. + if self.basis_type == "bessel": + freqs = np.arange(1, self.n_radial + 1, dtype=prec) * (math.pi / self.rcut) + else: + freqs = np.linspace(0.0, self.rcut, self.n_radial, dtype=prec) + self.adam_freqs = np.reshape(freqs.astype(prec), (1, self.n_radial)) + gaussian_width = self.rcut / max(self.n_radial - 1, 1) + self.gaussian_coeff = -0.5 / (gaussian_width * gaussian_width) + + self.envelope = C3CutoffEnvelope( + rcut=self.rcut, + exponent=self.exponent, + precision=self.precision, + ) + + def call(self, r: Any) -> Any: + """ + Compute radial basis functions. + + Parameters + ---------- + r : Array + Pair distances with shape (N, 1) in Å, where N is the number of + pairs. + + Returns + ------- + Array + Radial basis multiplied by C^3 cutoff envelope with shape + (N, n_rbf). The output is smoothly truncated to zero at r = rcut. + """ + xp = array_api_compat.array_namespace(r) + freqs = xp.asarray(self.adam_freqs, device=array_api_compat.device(r)) + # === Step 1. Radial basis === + # Shape: (N, 1) * (1, n_radial) -> (N, n_radial) + if self.basis_type == "bessel": + # phi_n(r) = w_n * sinc(w_n * r / π) + x = r * freqs # (N, n_rbf) + # normalized sinc, mirroring torch.sinc(x / π): + # sinc(z) = sin(π*z) / (π*z), with sinc(0) = 1. + # The zero branch is selected through a safe denominator so that + # gradients stay finite at r = 0. + z = x / math.pi + pz = math.pi * z + zero = z == 0.0 + safe_pz = xp.where(zero, xp.ones_like(pz), pz) + sinc = xp.where(zero, xp.ones_like(pz), xp.sin(safe_pz) / safe_pz) + raw = freqs * sinc # (N, n_rbf) + else: + dr = r - freqs # (N, n_rbf) + raw = xp.exp(dr * dr * self.gaussian_coeff) # (N, n_rbf) + + # === Step 2. Apply C^3 envelope for smooth cutoff === + envelope = self.envelope.call(r) # (N, 1) + return raw * envelope + + def serialize(self) -> dict[str, Any]: + """Serialize RadialBasis including trainable frequencies.""" + return { + "@class": "RadialBasis", + "@version": 1, + "config": { + "rcut": self.rcut, + "basis_type": self.basis_type, + "n_radial": self.n_radial, + "exponent": self.exponent, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + }, + "@variables": {"adam_freqs": to_numpy_array(self.adam_freqs)}, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> RadialBasis: + """Deserialize RadialBasis including trainable frequencies.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "RadialBasis": + raise ValueError(f"Invalid class for RadialBasis: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config", data) + variables = data.pop("@variables", None) + precision = str(config["precision"]) + obj = cls( + rcut=float(config["rcut"]), + basis_type=str(config.get("basis_type", "bessel")), + n_radial=int(config["n_radial"]), + precision=precision, + exponent=int(config.get("exponent", 7)), + ) + if variables is not None: + prec = PRECISION_DICT[precision.lower()] + adam_freqs = np.asarray(variables["adam_freqs"], dtype=prec) + if adam_freqs.shape != obj.adam_freqs.shape: + raise ValueError( + f"adam_freqs shape {adam_freqs.shape} does not match " + f"the expected shape {obj.adam_freqs.shape}" + ) + obj.adam_freqs = adam_freqs + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py new file mode 100644 index 0000000000..eca0209e93 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -0,0 +1,1737 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +SO(2)-equivariant message-passing layers for DPA4/SeZM. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.so2``. +It defines the reduced-layout SO(2) linear operator, the edge-conditioned +radial degree mixer, and the edge convolution used inside SeZM interaction +blocks. + +Padded-edge adaptation +---------------------- +The pt ``SO2Convolution`` consumes a flat *sparse* edge list and aggregates +per destination node with ``index_add_``. The dpmodel port uses the padded, +frame-explicit edge layout documented in ``edge_cache.EdgeCache`` +(``E = nf * nloc * nnei`` with invalid slots marked by ``edge_mask == 0``), +so every destination aggregation becomes a masked sum over the ``nnei`` axis +and the destination-wise softmax becomes a masked softmax over ``nnei`` +(see ``attention.segment_envelope_gated_softmax``). Per-edge math (the +SO(2) linear application, the Wigner rotations via the ``D_to_m`` +projections, and the radial modulation) is identical to pt, just evaluated +over the padded edge axis. + +Branches guarded with ``NotImplementedError`` (flags unused by the core DPA4 +config): ``so2_attn_res != "none"``, ``layer_scale``, ``s2_activation``, +``atten_f_mix``, ``atten_v_proj``, ``atten_o_proj``, ``node_wise_s2``, +``node_wise_so3``, ``message_node_s2``, ``message_node_so3``. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + xp_sigmoid, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .activation import ( + GatedActivation, +) +from .attention import ( + segment_envelope_gated_softmax, +) +from .indexing import ( + build_m_major_index, + build_m_major_l_index, + build_rotate_inv_rescale, + get_so3_dim_of_lmax, + map_degree_idx, + project_D_to_m, + project_Dt_from_m, +) +from .norm import ( + ReducedEquivariantRMSNorm, + ScalarRMSNorm, +) +from .projection import ( + resolve_s2_grid_resolution, +) +from .so3 import ( + ChannelLinear, + FocusLinear, + SO3Linear, +) +from .utils import ( + ATTN_RES_MODES, + init_trunc_normal_fan_in_out, +) + +if TYPE_CHECKING: + from .edge_cache import ( + EdgeCache, + ) + + +def _compute_precision(precision: str) -> str: + """Promote fp16/bf16 to fp32 (dpmodel analog of pt ``get_promoted_dtype``).""" + name = np.dtype(PRECISION_DICT[precision.lower()]).name + if "float16" in name: # matches float16 and bfloat16 + return "float32" + return precision + + +def _check_shape_assign(obj: Any, attr: str, value: Any, dtype: Any, key: str) -> None: + """Assign ``value`` (cast to ``dtype``) to ``obj.attr`` with a shape check.""" + expected = getattr(obj, attr) + arr = np.asarray(value, dtype=dtype) + if arr.shape != expected.shape: + raise ValueError( + f"{key} shape {arr.shape} does not match the expected shape " + f"{expected.shape}" + ) + setattr(obj, attr, arr) + + +def _check_index_table(expected: np.ndarray, value: Any, key: str) -> None: + """Validate that a serialized integer index table matches the rebuilt one.""" + arr = np.asarray(value, dtype=np.int64) + if not np.array_equal(arr.reshape(-1), np.asarray(expected).reshape(-1)): + raise ValueError(f"{key} does not match the table derived from the config") + + +class SO2Linear(NativeOP): + """ + SO(2)-equivariant linear mixing in the edge-aligned local frame. + + Coefficient layout (m-major, truncated by mmax) + ------------------------------------------------ + The coefficient axis D_m_trunc is ordered by |m| groups:: + + [ m=0: l=0..lmax | m=1: (l,-1) then (l,+1) | ... | m=mmax: ... ] + |___ lmax+1 ____| |_______ 2*(lmax) ________| + + Each |m| group is contiguous, enabling per-group block matmuls. + + Block-diagonal weight structure + ------------------------------- + The conceptual full weight matrix is block-diagonal over |m| groups:: + + W = diag[W_m0, B_m1, B_m2, ..., B_mmax] + + - ``W_m0``: unconstrained ``(num_l*Cin, num_l*Cout)`` block for m=0. + Cross-l mixing is allowed since m=0 coefficients are real scalars. + + - ``B_m`` (|m|>0): SO(2)-constrained 2x2 block coupling (-m, +m) pairs:: + + B_m = [ W_u^T , -W_v^T ] where W_u, W_v are learnable + [ W_v^T , W_u^T ] (num_l*Cin, num_l*Cout) each. + + This structure is the real-valued form of complex multiplication + ``(u + iv)(a + ib) = (ua - vb) + i(va + ub)``, which guarantees + SO(2) equivariance. + + Unlike pt (which assembles the dense block-diagonal matrix and applies a + single ``einsum``), the dpmodel forward contracts the diagonal blocks + directly with slicing + matmul + concat, which is array-API friendly and + numerically equivalent (the off-block entries are exact zeros). + + Parameters + ---------- + lmax + Maximum spherical harmonic degree. + mmax + Maximum SO(2) order (|m|) to mix. If None, defaults to ``lmax``. + in_channels + Number of input channels per (l, m) coefficient. + out_channels + Number of output channels per (l, m) coefficient. + n_focus + Number of independent focus streams. Each stream has its own + weight matrices. + precision + Parameter precision. + mlp_bias + Whether to use bias for l=0 (scalar) components. + seed + Random seed for weight initialization. + trainable + Whether parameters are trainable. + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None = None, + in_channels: int, + out_channels: int, + n_focus: int = 1, + precision: str = DEFAULT_PRECISION, + mlp_bias: bool = False, + seed: int | list[int] | None = None, + trainable: bool = True, + ) -> None: + self.lmax = int(lmax) + self.mmax = int(self.lmax if mmax is None else mmax) + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.in_channels = int(in_channels) + self.out_channels = int(out_channels) + self.n_focus = int(n_focus) + self.precision = precision + self.mlp_bias = bool(mlp_bias) + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + + # === Step 1. Build m-major coefficient layout === + # Map each |m| group to contiguous index ranges in the flattened axis. + # Example for lmax=2, mmax=2: + # m=0 : indices [0, 1, 2] (l=0,1,2) + # m=1-: indices [3, 4] (l=1,2 with -m) + # m=1+: indices [5, 6] (l=1,2 with +m) + # m=2-: index [7] (l=2 with -m) + # m=2+: index [8] (l=2 with +m) + # => reduced_dim = 9 + m0_size = self.lmax + 1 + self.m0_idx = np.arange(m0_size, dtype=np.int64) + + pos_indices_list: list[np.ndarray] = [] + neg_indices_list: list[np.ndarray] = [] + # Each entry: (neg_start, pos_start, num_l) for a fixed |m|. + # These ranges are contiguous in m-major layout. + m_ranges: list[tuple[int, int, int]] = [] + + offset = m0_size + for m in range(1, self.mmax + 1): + num_l = self.lmax - m + 1 + neg_start = offset + pos_start = offset + num_l + neg_indices_list.append( + np.arange(neg_start, neg_start + num_l, dtype=np.int64) + ) + pos_indices_list.append( + np.arange(pos_start, pos_start + num_l, dtype=np.int64) + ) + m_ranges.append((neg_start, pos_start, num_l)) + offset += 2 * num_l + + self.reduced_dim = int(offset) + + if len(pos_indices_list) > 0: + self.pos_indices = np.concatenate(pos_indices_list) + self.neg_indices = np.concatenate(neg_indices_list) + else: + self.pos_indices = np.empty(0, dtype=np.int64) + self.neg_indices = np.empty(0, dtype=np.int64) + self._m_ranges = m_ranges + + # === Step 2. Learnable weight parameters === + # weight_m0: folded (num_l*Cin, F*num_l*Cout) storage — (in, out) convention. + # Runtime view: (num_l*Cin, F, num_l*Cout). + # Cross-l mixing is allowed because m=0 coefficients are real. + num_m0 = self.lmax + 1 + num_in_m0 = num_m0 * self.in_channels + num_out_m0 = num_m0 * self.out_channels + weight_m0 = np.empty((num_in_m0, self.n_focus * num_out_m0), dtype=prec) + weight_m0_view = weight_m0.reshape(num_in_m0, self.n_focus, num_out_m0) + for focus_idx in range(self.n_focus): + init_trunc_normal_fan_in_out( + weight_m0_view[:, focus_idx, :], child_seed(seed, 1000 + focus_idx) + ) + self.weight_m0 = weight_m0 + if self.mlp_bias: + self.bias0: np.ndarray | None = np.zeros( + (self.n_focus * self.out_channels,), dtype=prec + ) + else: + self.bias0 = None + + # weight_m[i]: folded (num_l*Cin, F*2*num_l*Cout) storage — (in, out) + # convention. Runtime view: (num_l*Cin, F, 2*num_l*Cout). + # The factor of 2 comes from storing W_u and W_v concatenated along the + # output axis. Scaling by 1/sqrt(2) compensates for the doubled + # parameter count. + self.weight_m: list[np.ndarray] = [] + for m in range(1, self.mmax + 1): + num_l = self.lmax - m + 1 + num_in = num_l * self.in_channels + num_out = 2 * num_l * self.out_channels + weight = np.empty((num_in, self.n_focus * num_out), dtype=prec) + weight_view = weight.reshape(num_in, self.n_focus, num_out) + for focus_idx in range(self.n_focus): + init_trunc_normal_fan_in_out( + weight_view[:, focus_idx, :], + child_seed(seed, 2000 + m * 100 + focus_idx), + ) + # Apply scaling for SO(2) equivariance + weight *= 1.0 / math.sqrt(2.0) + self.weight_m.append(weight) + + # === Step 3. Precompute flattened slice ranges for the block matmuls === + # Each |m|>0 group occupies two sub-blocks (neg, pos) in the flattened + # coefficient*channel axis. + # Tuple layout: (neg_i0, neg_i1, pos_i0, pos_i1, <- input ranges + # neg_o0, neg_o1, pos_o0, pos_o1) <- output ranges + self._m0_in = (self.lmax + 1) * self.in_channels + self._m0_out = (self.lmax + 1) * self.out_channels + self._block_slices: list[tuple[int, int, int, int, int, int, int, int]] = [] + for neg_start, pos_start, num_l in self._m_ranges: + ib = num_l * self.in_channels + ob = num_l * self.out_channels + self._block_slices.append( + ( + neg_start * self.in_channels, + neg_start * self.in_channels + ib, + pos_start * self.in_channels, + pos_start * self.in_channels + ib, + neg_start * self.out_channels, + neg_start * self.out_channels + ob, + pos_start * self.out_channels, + pos_start * self.out_channels + ob, + ) + ) + + @staticmethod + def _focus_matmul(xp: Any, x: Any, w: Any) -> Any: + """Per-focus matmul: einsum("efi,fio->efo") via broadcast batched matmul. + + Parameters + ---------- + x + Input with shape (E, F, in_blk). + w + Weight with shape (F, in_blk, out_blk). + """ + return xp.matmul(x[:, :, None, :], w[None, ...])[..., 0, :] + + def call(self, x: Any) -> Any: + """ + Parameters + ---------- + x + Input with shape (E, F, D_m_trunc, Cin), where D_m_trunc is the + coefficient dimension of the m-major layout truncated by `mmax`. + + Returns + ------- + Array + Output with shape (E, F, D_m_trunc, Cout). + """ + xp = array_api_compat.array_namespace(x) + # === Step 1. Flatten coefficient + channel axes for matmul === + # (E, F, D_m, Cin) -> (E, F, D_m*Cin) + n_edge = x.shape[0] + in_dim_total = self.reduced_dim * self.in_channels + x_flat = xp.reshape(x, (n_edge, self.n_focus, in_dim_total)) + + # === Step 2. Contract the diagonal |m| blocks === + # m=0 block: unconstrained (num_l*Cin, num_l*Cout) per focus. + num_m0 = self.lmax + 1 + device = array_api_compat.device(x) + weight_m0 = xp.reshape( + xp.asarray(self.weight_m0[...], device=device), + (num_m0 * self.in_channels, self.n_focus, num_m0 * self.out_channels), + ) + weight_m0 = xp.permute_dims(weight_m0, (1, 0, 2)) # (F, in, out) + out_blocks = [self._focus_matmul(xp, x_flat[:, :, : self._m0_in], weight_m0)] + + # |m|>0 blocks: real-valued complex multiplication on (-m, +m) pairs. + for m_idx, w in enumerate(self.weight_m): + ni0, ni1, pi0, pi1, no0, no1, po0, po1 = self._block_slices[m_idx] + ib = ni1 - ni0 # in_block size + ob = no1 - no0 # out_block size + w = xp.reshape( + xp.asarray(w[...], device=device), (ib, self.n_focus, 2 * ob) + ) + w = xp.permute_dims(w, (1, 0, 2)) # (F, in_blk, 2*out_blk) + w_u = w[:, :, :ob] # (F, in_blk, out_blk) + w_v = w[:, :, ob:] # (F, in_blk, out_blk) + x_neg = x_flat[:, :, ni0:ni1] + x_pos = x_flat[:, :, pi0:pi1] + # 2x2 coupling: neg_out = x_neg @ W_u - x_pos @ W_v + # pos_out = x_neg @ W_v + x_pos @ W_u + out_blocks.append( + self._focus_matmul(xp, x_neg, w_u) - self._focus_matmul(xp, x_pos, w_v) + ) + out_blocks.append( + self._focus_matmul(xp, x_neg, w_v) + self._focus_matmul(xp, x_pos, w_u) + ) + + out_flat = ( + xp.concat(out_blocks, axis=-1) if len(out_blocks) > 1 else out_blocks[0] + ) + out = xp.reshape( + out_flat, (n_edge, self.n_focus, self.reduced_dim, self.out_channels) + ) + + # === Step 3. Bias on l=0 scalar index === + if self.mlp_bias: + bias0 = xp.reshape( + xp.asarray(self.bias0[...], device=device), + (self.n_focus, self.out_channels), + ) + out0 = out[:, :, :1, :] + bias0[None, :, None, :] + out = xp.concat([out0, out[:, :, 1:, :]], axis=2) + return out + + def _variables(self) -> dict[str, np.ndarray]: + """Variables keyed by the pt ``state_dict`` key names.""" + variables = { + "m0_idx": to_numpy_array(self.m0_idx), + "pos_indices": to_numpy_array(self.pos_indices), + "neg_indices": to_numpy_array(self.neg_indices), + "weight_m0": to_numpy_array(self.weight_m0), + } + if self.mlp_bias: + variables["bias0"] = to_numpy_array(self.bias0) + for m_idx, w in enumerate(self.weight_m): + variables[f"weight_m.{m_idx}"] = to_numpy_array(w) + return variables + + def _load_variables(self, variables: dict[str, Any]) -> None: + """Load variables keyed by the pt ``state_dict`` key names.""" + prec = PRECISION_DICT[self.precision.lower()] + _check_index_table(self.m0_idx, variables["m0_idx"], "m0_idx") + _check_index_table(self.pos_indices, variables["pos_indices"], "pos_indices") + _check_index_table(self.neg_indices, variables["neg_indices"], "neg_indices") + _check_shape_assign( + self, "weight_m0", variables["weight_m0"], prec, "weight_m0" + ) + if self.mlp_bias: + self.bias0 = np.asarray(variables["bias0"], dtype=prec).reshape( + self.bias0.shape + ) + for m_idx in range(len(self.weight_m)): + key = f"weight_m.{m_idx}" + value = np.asarray(variables[key], dtype=prec) + if value.shape != self.weight_m[m_idx].shape: + raise ValueError( + f"{key} shape {value.shape} does not match the expected " + f"shape {self.weight_m[m_idx].shape}" + ) + self.weight_m[m_idx] = value + + def serialize(self) -> dict[str, Any]: + """Serialize the SO2Linear to a dict (pt-compatible format).""" + return { + "@class": "SO2Linear", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.mmax, + "in_channels": self.in_channels, + "out_channels": self.out_channels, + "n_focus": self.n_focus, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "mlp_bias": self.mlp_bias, + "trainable": self.trainable, + "seed": None, + }, + "@variables": self._variables(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SO2Linear: + """Deserialize an SO2Linear from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "SO2Linear": + raise ValueError(f"Invalid class for SO2Linear: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + lmax=int(config["lmax"]), + mmax=int(config["mmax"]), + in_channels=int(config["in_channels"]), + out_channels=int(config["out_channels"]), + n_focus=int(config["n_focus"]), + precision=str(config["precision"]), + mlp_bias=bool(config["mlp_bias"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + obj._load_variables(variables) + return obj + + +class DynamicRadialDegreeMixer(NativeOP): + """ + Edge-conditioned degree mixer in the SO(2) reduced local layout. + + The mixer replaces per-degree scalar radial modulation by an edge-conditioned + degree kernel without channel output mixing: + + degree: + y[e, l_out, m, c] = sum_l_in W[e, l_in, l_out, |m|] x[e, l_in, m, c] + degree_channel: + y[e, l_out, m, c] = sum_l_in W[e, l_in, l_out, |m|, c] x[e, l_in, m, c] + + `mode="degree"` shares W across channels. `mode="degree_channel"` gives each + channel its own W, optionally with a low-rank channel factorization. + + The pt ``index_copy_`` scatter of the compact kernel into the dense + ``(D_m, D_m)`` layout is replaced by a precomputed gather index + mask + (functionally identical, array-API friendly). + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None = None, + channels: int, + mode: str, + rank: int = 0, + precision: str = DEFAULT_PRECISION, + seed: int | list[int] | None = None, + trainable: bool = True, + ) -> None: + self.lmax = int(lmax) + self.mmax = int(self.lmax if mmax is None else mmax) + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.channels = int(channels) + if self.channels < 1: + raise ValueError("`channels` must be positive") + self.mode = str(mode).lower() + if self.mode not in {"degree", "degree_channel"}: + raise ValueError("`mode` must be one of 'degree' or 'degree_channel'") + self.rank = int(rank) + if self.rank < 0: + raise ValueError("`rank` must be non-negative") + self.precision = precision + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + + # m-major reduced layout: m=0 block followed by (-m, +m) blocks. + self.reduced_dim = (self.lmax + 1) + sum( + 2 * (self.lmax - m + 1) for m in range(1, self.mmax + 1) + ) + self.degree_kernel_size = sum( + (self.lmax - m + 1) ** 2 for m in range(self.mmax + 1) + ) + self.input_dim = (self.lmax + 1) * self.channels + if self.mode == "degree": + self.proj_out_dim = self.degree_kernel_size + elif self.rank > 0: + self.proj_out_dim = self.degree_kernel_size * self.rank + else: + self.proj_out_dim = self.degree_kernel_size * self.channels + + weight = np.empty((self.input_dim, self.proj_out_dim), dtype=prec) + init_trunc_normal_fan_in_out(weight, child_seed(seed, 0)) + self.weight = weight + + if self.mode == "degree_channel" and self.rank > 0: + channel_basis = np.empty((self.rank, self.channels), dtype=prec) + init_trunc_normal_fan_in_out(channel_basis, child_seed(seed, 1)) + self.channel_basis: np.ndarray | None = channel_basis + else: + self.channel_basis = None + + compact_idx, dense_idx = self._build_dense_scatter_indices() + self.kernel_compact_index = compact_idx + self.kernel_dense_index = dense_idx + # Gather-form of pt's index_copy_ scatter: + # dense[:, dense_idx[j]] = compact[:, compact_idx[j]] + # becomes + # dense = take(compact, gather_index, axis=1) * scatter_mask + dense_size = self.reduced_dim * self.reduced_dim + gather_index = np.zeros(dense_size, dtype=np.int64) + scatter_mask = np.zeros(dense_size, dtype=prec) + gather_index[dense_idx] = compact_idx + scatter_mask[dense_idx] = 1.0 + self._dense_gather_index = gather_index + self._dense_scatter_mask = scatter_mask + + def _build_dense_scatter_indices(self) -> tuple[np.ndarray, np.ndarray]: + compact_indices: list[int] = [] + dense_indices: list[int] = [] + compact_offset = 0 + reduced_dim = self.reduced_dim + + def append_block(start_in: int, start_out: int, num_l: int) -> None: + for l_in in range(num_l): + for l_out in range(num_l): + compact_indices.append(compact_offset + l_in * num_l + l_out) + # Store dense kernels in matmul layout (out, in) so forward + # can use a batched matmul without transposing. + dense_indices.append( + (start_out + l_out) * reduced_dim + start_in + l_in + ) + + # m=0: single real block. + num_l0 = self.lmax + 1 + append_block(0, 0, num_l0) + compact_offset += num_l0 * num_l0 + + # |m|>0: same degree kernel is applied to the negative and positive + # signed-m blocks. No cross signed-m mixing is introduced. + offset = num_l0 + for m in range(1, self.mmax + 1): + num_l = self.lmax - m + 1 + neg_start = offset + pos_start = offset + num_l + append_block(neg_start, neg_start, num_l) + append_block(pos_start, pos_start, num_l) + compact_offset += num_l * num_l + offset += 2 * num_l + + return ( + np.asarray(compact_indices, dtype=np.int64), + np.asarray(dense_indices, dtype=np.int64), + ) + + def _project_radial(self, xp: Any, radial_feat: Any) -> Any: + radial_m0 = xp.reshape( + radial_feat[:, : self.lmax + 1, :], + (radial_feat.shape[0], self.input_dim), + ) + weight = xp.asarray( + self.weight[...], device=array_api_compat.device(radial_feat) + ) + return xp.matmul(radial_m0, weight) + + def _scatter_dense(self, xp: Any, compact: Any, device: Any) -> Any: + """Scatter the compact per-block kernel into the dense (D_m*D_m, ...) layout.""" + gather_index = xp.asarray(self._dense_gather_index, device=device) + scatter_mask = xp.astype( + xp.asarray(self._dense_scatter_mask, device=device), compact.dtype + ) + dense = xp.take(compact, gather_index, axis=1) + if compact.ndim == 2: + return dense * scatter_mask[None, :] + return dense * scatter_mask[None, :, None] + + def call(self, x_local: Any, radial_feat: Any) -> Any: + """ + Parameters + ---------- + x_local + Local reduced features with shape (E, D_m, C_wide). + radial_feat + Invariant radial/type features with shape (E, D_m, C_wide). + """ + if x_local.shape != radial_feat.shape: + raise ValueError("`x_local` and `radial_feat` must have the same shape") + if x_local.shape[1] != self.reduced_dim or x_local.shape[2] != self.channels: + raise ValueError("Input shape is incompatible with this mixer") + + xp = array_api_compat.array_namespace(x_local) + device = array_api_compat.device(x_local) + n_edge = x_local.shape[0] + kernel_flat = self._project_radial(xp, radial_feat) + if self.mode == "degree": + kernel = xp.reshape( + self._scatter_dense(xp, kernel_flat, device), + (n_edge, self.reduced_dim, self.reduced_dim), + ) + return xp.matmul(kernel, x_local) + + if self.rank > 0: + compact = xp.reshape( + kernel_flat, (n_edge, self.degree_kernel_size, self.rank) + ) + kernel = xp.reshape( + self._scatter_dense(xp, compact, device), + (n_edge, self.reduced_dim, self.reduced_dim, self.rank), + ) + # einsum "eoir,eic->eorc" as a broadcast batched matmul: + # (E, o, r, i) @ (E, 1, i, c) -> (E, o, r, c) + kernel = xp.permute_dims(kernel, (0, 1, 3, 2)) + mixed = xp.matmul(kernel, x_local[:, None, :, :]) + channel_basis = xp.reshape( + xp.asarray(self.channel_basis[...], device=device), + (1, 1, self.rank, self.channels), + ) + return xp.sum(mixed * channel_basis, axis=2) + + compact = xp.reshape( + kernel_flat, (n_edge, self.degree_kernel_size, self.channels) + ) + kernel = xp.reshape( + self._scatter_dense(xp, compact, device), + (n_edge, self.reduced_dim, self.reduced_dim, self.channels), + ) + # einsum "eoic,eic->eoc" + return xp.sum(kernel * x_local[:, None, :, :], axis=2) + + def _variables(self) -> dict[str, np.ndarray]: + """Variables keyed by the pt ``state_dict`` key names.""" + variables = {"weight": to_numpy_array(self.weight)} + if self.channel_basis is not None: + variables["channel_basis"] = to_numpy_array(self.channel_basis) + variables["kernel_compact_index"] = to_numpy_array(self.kernel_compact_index) + variables["kernel_dense_index"] = to_numpy_array(self.kernel_dense_index) + return variables + + def _load_variables(self, variables: dict[str, Any]) -> None: + """Load variables keyed by the pt ``state_dict`` key names.""" + prec = PRECISION_DICT[self.precision.lower()] + _check_index_table( + self.kernel_compact_index, + variables["kernel_compact_index"], + "kernel_compact_index", + ) + _check_index_table( + self.kernel_dense_index, + variables["kernel_dense_index"], + "kernel_dense_index", + ) + _check_shape_assign(self, "weight", variables["weight"], prec, "weight") + if self.channel_basis is not None: + _check_shape_assign( + self, "channel_basis", variables["channel_basis"], prec, "channel_basis" + ) + + def serialize(self) -> dict[str, Any]: + """Serialize the DynamicRadialDegreeMixer to a dict. + + The pt class has no ``serialize()``; the ``@variables`` keys here + match the pt ``state_dict`` key names. + """ + return { + "@class": "DynamicRadialDegreeMixer", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.mmax, + "channels": self.channels, + "mode": self.mode, + "rank": self.rank, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "seed": None, + }, + "@variables": self._variables(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> DynamicRadialDegreeMixer: + """Deserialize a DynamicRadialDegreeMixer from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "DynamicRadialDegreeMixer": + raise ValueError(f"Invalid class for DynamicRadialDegreeMixer: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + lmax=int(config["lmax"]), + mmax=int(config["mmax"]), + channels=int(config["channels"]), + mode=str(config["mode"]), + rank=int(config["rank"]), + precision=str(config["precision"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + obj._load_variables(variables) + return obj + + +class SO2Convolution(NativeOP): + """ + SO(2)-equivariant edge convolution with cached geometry and rotations. + + This module consumes node features in packed SO(3) layout `(N, D, C)` and + performs edge message passing in the reduced m-major local layout. The + operation pipeline is: + + 1. `pre_focus_mix`: project node features `(N, D, C)` to the SO(2) hidden width. + 2. rotate global -> local reduced basis with cached `D_to_m`. + 3. radial modulation in reduced layout. + 4. `so2_layers` stacked local mixers: + `inter_norm -> SO2Linear -> non_linearity -> residual`. + 5. rotate local -> global with cached `Dt_from_m`. + 6. edge aggregation (plain envelope masked sum or envelope-aware masked + softmax attention with output-side head gate); see the module + docstring for the padded-edge adaptation. + 7. `post_focus_mix`: project aggregated hidden messages back to `(N, D, C)`. + + See the pt ``SO2Convolution`` docstring for the full parameter + documentation; this port keeps the same constructor parameters with + ``dtype`` replaced by ``precision``. Flags unused by the core DPA4 config + raise ``NotImplementedError`` (listed in the module docstring). + """ + + def __init__( + self, + *, + lmax: int, + mmax: int | None = None, + kmax: int = 1, + channels: int, + n_focus: int = 1, + focus_dim: int = 0, + focus_compete: bool = True, + so2_norm: bool = False, + so2_layers: int = 4, + so2_attn_res: str = "none", + layer_scale: bool = False, + n_atten_head: int = 1, + atten_f_mix: bool = False, + atten_v_proj: bool = False, + atten_o_proj: bool = False, + s2_activation: bool = False, + node_wise_grid_mlp: bool = False, + node_wise_grid_branch: int = 0, + message_node_grid_mlp: bool = False, + message_node_grid_branch: int = 0, + node_wise_s2: bool = False, + node_wise_so3: bool = False, + message_node_s2: bool = False, + message_node_so3: bool = False, + lebedev_quadrature: bool = False, + activation_function: str = "silu", + mlp_bias: bool = False, + radial_so2_mode: str = "none", + radial_so2_rank: int = 0, + eps: float = 1e-7, + precision: str = DEFAULT_PRECISION, + seed: int | list[int] | None = None, + trainable: bool = True, + ) -> None: + self.lmax = int(lmax) + self.mmax = int(self.lmax if mmax is None else mmax) + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") + self.kmax = int(kmax) + if self.kmax < 0: + raise ValueError("`kmax` must be non-negative") + self.channels = int(channels) + self.n_focus = int(n_focus) + if self.n_focus < 1: + raise ValueError("`n_focus` must be >= 1") + self.focus_dim = int(focus_dim) + if self.focus_dim < 0: + raise ValueError("`focus_dim` must be >= 0") + self.so2_focus_dim = self.channels if self.focus_dim == 0 else self.focus_dim + self.hidden_channels = int(self.n_focus * self.so2_focus_dim) + self.use_hidden_projection = self.hidden_channels != self.channels + self.focus_compete = bool(focus_compete) + self.focus_softmax_tau = 1.0 + self.focus_label_smoothing = 0.02 + self.so2_norm = bool(so2_norm) + self.so2_layers = int(so2_layers) + if self.so2_layers < 1: + raise ValueError("`so2_layers` must be >= 1") + self.so2_attn_res_mode = str(so2_attn_res).lower() + if self.so2_attn_res_mode not in ATTN_RES_MODES: + raise ValueError( + "`so2_attn_res` must be one of 'none', 'independent', or 'dependent'" + ) + if self.so2_attn_res_mode != "none": + raise NotImplementedError( + "so2_attn_res != 'none' (DepthAttnRes) is not ported to dpmodel" + ) + self.layer_scale = bool(layer_scale) + if self.layer_scale: + raise NotImplementedError("layer_scale=True is not ported to dpmodel") + self.n_atten_head = int(n_atten_head) + if self.n_atten_head < 0: + raise ValueError("`n_atten_head` must be non-negative") + self.atten_f_mix = bool(atten_f_mix) + if self.atten_f_mix: + raise NotImplementedError("atten_f_mix=True is not ported to dpmodel") + self.use_atten_v_proj = bool(atten_v_proj) + if self.use_atten_v_proj: + raise NotImplementedError("atten_v_proj=True is not ported to dpmodel") + self.use_atten_o_proj = bool(atten_o_proj) + if self.use_atten_o_proj: + raise NotImplementedError("atten_o_proj=True is not ported to dpmodel") + self.s2_activation = bool(s2_activation) + if self.s2_activation: + raise NotImplementedError( + "s2_activation=True (so2_s2_activation) is not ported to dpmodel" + ) + self.node_wise_grid_mlp = bool(node_wise_grid_mlp) + self.node_wise_grid_branch = int(node_wise_grid_branch) + self.message_node_grid_mlp = bool(message_node_grid_mlp) + self.message_node_grid_branch = int(message_node_grid_branch) + if min(self.node_wise_grid_branch, self.message_node_grid_branch) < 0: + raise ValueError("grid branch counts must be non-negative") + self.node_wise_s2 = bool(node_wise_s2) + self.node_wise_so3 = bool(node_wise_so3) + self.message_node_s2 = bool(message_node_s2) + self.message_node_so3 = bool(message_node_so3) + if self.node_wise_s2 or self.node_wise_so3: + raise NotImplementedError( + "node_wise_s2/node_wise_so3 grid products are not ported to dpmodel" + ) + if self.message_node_s2 or self.message_node_so3: + raise NotImplementedError( + "message_node_s2/message_node_so3 grid products are not ported " + "to dpmodel" + ) + self.lebedev_quadrature = bool(lebedev_quadrature) + self.s2_grid_method = "lebedev" if self.lebedev_quadrature else "e3nn" + self.s2_grid_resolution = resolve_s2_grid_resolution( + self.lmax, + self.mmax, + method=self.s2_grid_method, + ) + self.activation_function = str(activation_function) + self.attn_n_focus = self.n_focus + self.attn_focus_dim = self.so2_focus_dim + if self.n_atten_head > 0 and self.attn_focus_dim % self.n_atten_head != 0: + raise ValueError( + "`n_atten_head` must divide the attention width " + "(`focus_dim` or `n_focus * focus_dim` when `atten_f_mix=True`)" + ) + self.head_dim = ( + None + if self.n_atten_head == 0 + else int(self.attn_focus_dim // self.n_atten_head) + ) + self.mlp_bias = bool(mlp_bias) + self.radial_so2_mode = str(radial_so2_mode).lower() + if self.radial_so2_mode not in {"none", "degree", "degree_channel"}: + raise ValueError( + "`radial_so2_mode` must be one of 'none', 'degree', or 'degree_channel'" + ) + self.radial_so2_rank = int(radial_so2_rank) + if self.radial_so2_rank < 0: + raise ValueError("`radial_so2_rank` must be non-negative") + self.eps = float(eps) + self.ebed_dim_full = get_so3_dim_of_lmax(self.lmax) + self.precision = precision + self.compute_precision = _compute_precision(precision) + self.trainable = bool(trainable) + prec = PRECISION_DICT[self.precision.lower()] + + # === Step 1. Precompute coefficient indices for m-major reduced layout === + self.coeff_index_m = build_m_major_index(self.lmax, self.mmax) + self.degree_index_m = build_m_major_l_index(self.lmax, self.mmax) + degree_index_full = map_degree_idx(self.lmax) + self.rotate_inv_rescale_full = build_rotate_inv_rescale( + self.lmax, + self.mmax, + degree_index_full, + dtype=prec, + ) + self.reduced_dim = int(self.coeff_index_m.shape[0]) + + # === Step 2. Split deterministic seeds at the module top-level === + seed_so2_stack = child_seed(seed, 0) + seed_non_linearities = child_seed(seed, 1) + seed_so3_pre = child_seed(seed, 2) + seed_so3_post = child_seed(seed, 3) + seed_gate = child_seed(seed, 4) + seed_radial_hidden = child_seed(seed, 6) + seed_radial_degree = child_seed(seed, 7) + + # === Step 3. Multiple SO2Linear layers === + # (s2_activation is guarded above, so out_channels == so2_focus_dim.) + self.so2_linears = [ + SO2Linear( + lmax=self.lmax, + mmax=self.mmax, + in_channels=self.so2_focus_dim, + out_channels=self.so2_focus_dim, + n_focus=self.n_focus, + precision=self.precision, + mlp_bias=self.mlp_bias, + seed=child_seed(seed_so2_stack, i), + trainable=trainable, + ) + for i in range(self.so2_layers) + ] + + # === Step 4. Intermediate norms (Optional) === + # pt appends nn.Identity() entries; dpmodel uses None for Identity. + inter_norms: list[ReducedEquivariantRMSNorm | None] = [] + if self.so2_norm: + for _ in range(max(0, self.so2_layers - 1)): + inter_norms.append( + ReducedEquivariantRMSNorm( + lmax=self.lmax, + mmax=self.mmax, + channels=self.so2_focus_dim, + degree_index_m=self.degree_index_m, + n_focus=self.n_focus, + precision=self.compute_precision, + trainable=trainable, + ) + ) + else: + for _ in range(max(0, self.so2_layers - 1)): + inter_norms.append(None) + inter_norms.append(None) + self.so2_inter_norms = inter_norms + + # === Step 5. Intermediate non-linearity === + # pt appends nn.Identity() as the last entry; dpmodel uses None. + non_linearities: list[GatedActivation | None] = [] + for i in range(max(0, self.so2_layers - 1)): + non_linearities.append( + GatedActivation( + lmax=self.lmax, + mmax=self.mmax, + channels=self.so2_focus_dim, + n_focus=self.n_focus, + precision=self.compute_precision, + activation_function=self.activation_function, + mlp_bias=self.mlp_bias, + layout="nfdc", + trainable=trainable, + seed=child_seed(seed_non_linearities, i), + ) + ) + non_linearities.append(None) + self.non_linearities = non_linearities + + # === Step 7. Optional attention projections (n_atten_head > 0) === + self.attn_qk_norm: ScalarRMSNorm | None = None + self.attn_q_proj: FocusLinear | None = None + self.attn_k_proj: FocusLinear | None = None + self.adamw_attn_logit_w: np.ndarray | None = None + self.adamw_attn_z_bias_raw: np.ndarray | None = None + self.attn_output_gate_norm: ScalarRMSNorm | None = None + self.adamw_attn_gate_w: np.ndarray | None = None + cprec = PRECISION_DICT[self.compute_precision.lower()] + if self.n_atten_head > 0: + self.attn_qk_norm = ScalarRMSNorm( + channels=self.attn_focus_dim, + n_focus=self.attn_n_focus, + eps=self.eps, + precision=self.compute_precision, + trainable=trainable, + ) + self.attn_q_proj = FocusLinear( + in_channels=self.attn_focus_dim, + out_channels=self.attn_focus_dim, + n_focus=self.attn_n_focus, + precision=self.compute_precision, + bias=False, + seed=child_seed(seed_gate, 0), + trainable=trainable, + ) + self.attn_k_proj = FocusLinear( + in_channels=self.attn_focus_dim, + out_channels=self.attn_focus_dim, + n_focus=self.attn_n_focus, + precision=self.compute_precision, + bias=False, + seed=child_seed(seed_gate, 1), + trainable=trainable, + ) + rng = np.random.default_rng(child_seed(seed_gate, 2)) + self.adamw_attn_logit_w = rng.normal( + 0.0, + 0.01, + size=(self.attn_focus_dim, self.attn_n_focus, self.n_atten_head), + ).astype(cprec) + # softplus(0.5413) ~= 1.0 provides balanced initial competition. + self.adamw_attn_z_bias_raw = np.full( + (self.attn_n_focus, self.n_atten_head), 0.5413, dtype=cprec + ) + self.attn_output_gate_norm = ScalarRMSNorm( + channels=self.attn_focus_dim, + n_focus=self.attn_n_focus, + eps=self.eps, + precision=self.compute_precision, + trainable=trainable, + ) + rng = np.random.default_rng(child_seed(seed_gate, 3)) + self.adamw_attn_gate_w = rng.normal( + 0.0, + 0.01, + size=(self.attn_focus_dim, self.attn_n_focus, self.n_atten_head), + ).astype(cprec) + + # === Step 7.5. Optional cross-focus competition === + self.focus_compete_norm: ScalarRMSNorm | None = None + self.adamw_focus_compete_w: np.ndarray | None = None + self.focus_compete_bias: np.ndarray | None = None + if self.focus_compete and self.n_focus > 1: + self.focus_compete_norm = ScalarRMSNorm( + channels=self.so2_focus_dim, + n_focus=self.n_focus, + eps=self.eps, + precision=self.compute_precision, + trainable=trainable, + ) + rng = np.random.default_rng(child_seed(seed_gate, 4)) + self.adamw_focus_compete_w = rng.normal( + 0.0, 0.01, size=(self.so2_focus_dim, self.n_focus) + ).astype(cprec) + if self.mlp_bias: + self.focus_compete_bias = np.zeros((self.n_focus,), dtype=cprec) + + # === Step 8. Optional radial hidden projection === + self.radial_hidden_proj: ChannelLinear | None = None + if self.use_hidden_projection: + self.radial_hidden_proj = ChannelLinear( + in_channels=self.channels, + out_channels=self.hidden_channels, + precision=self.precision, + bias=False, + seed=seed_radial_hidden, + trainable=trainable, + ) + self.radial_degree_mixer: DynamicRadialDegreeMixer | None = None + if self.radial_so2_mode != "none": + self.radial_degree_mixer = DynamicRadialDegreeMixer( + lmax=self.lmax, + mmax=self.mmax, + channels=self.hidden_channels, + mode=self.radial_so2_mode, + rank=self.radial_so2_rank, + precision=self.precision, + seed=seed_radial_degree, + trainable=trainable, + ) + + # === Step 9. Pre-focus channel mixing === + # This projects the full channel width before the SO(2) focus split. + self.pre_focus_mix = SO3Linear( + lmax=self.lmax, + in_channels=self.channels, + out_channels=self.hidden_channels, + n_focus=1, + precision=self.precision, + mlp_bias=self.mlp_bias, + trainable=trainable, + seed=seed_so3_pre, + ) + + # === Step 10. Post-focus channel mixing === + self.post_focus_mix = SO3Linear( + lmax=self.lmax, + in_channels=self.hidden_channels, + out_channels=self.channels, + n_focus=1, + precision=self.precision, + mlp_bias=self.mlp_bias, + trainable=trainable, + seed=seed_so3_post, + init_std=0.0, + ) + + def call( + self, + x: Any, + edge_cache: EdgeCache, + radial_feat: Any, + ) -> Any: + """ + Parameters + ---------- + x + Node features with shape (N, D, C), where D=(lmax+1)^2 is the + SO(3) coefficient dimension and N = nf * nloc is the local node + axis. + edge_cache + Precomputed edge cache in the padded-edge layout + (``E = N * nnei``; see ``edge_cache.EdgeCache``). Must be + compatible with this block's lmax. + radial_feat + Per-edge radial features with shape (E, lmax+1, C), already fused + with edge type features. + + Returns + ------- + Array + Message updates with shape (N, D, C). + """ + xp = array_api_compat.array_namespace(x) + device = array_api_compat.device(x) + src, dst = edge_cache.src, edge_cache.dst + n_node = x.shape[0] + n_edge = int(src.shape[0]) + if n_node <= 0 or n_edge % n_node != 0: + raise ValueError( + "padded-edge layout requires E to be a multiple of N; " + f"got E={n_edge}, N={n_node}" + ) + nnei = n_edge // n_node + # Validity mask for the padded-edge layout (1 on real edges). + edge_mask = edge_cache.edge_mask + if edge_mask is not None: + mask_f = xp.astype(xp.reshape(edge_mask, (n_edge,)), x.dtype) + else: + mask_f = xp.ones((n_edge,), dtype=x.dtype, device=device) + + # === Step 1. Pre-focus channel mixing on full width === + # (N, D, C_wide), C_wide = F * Cf + x = self.pre_focus_mix(x[:, :, None, :])[:, :, 0, :] + + # === Step 2. Rotate to edge-aligned local frame === + D_full = edge_cache.D_full + D_m_prime = project_D_to_m( + D_full=D_full, + coeff_index_m=self.coeff_index_m, + ebed_dim_full=self.ebed_dim_full, + cache=edge_cache.D_to_m_cache, + key_lmax=self.lmax, + key_mmax=self.mmax, + ) + src_idx = xp.astype(xp.reshape(src, (n_edge,)), xp.int64) + x_src = xp.take(x, src_idx, axis=0) # (E, D, C_wide) + x_local = xp.matmul(D_m_prime, x_src) # (E, D_m, C_wide) + + # === Step 3. Select radial/type features for reduced layout === + degree_index_m = xp.asarray(self.degree_index_m, device=device) + rad_feat = xp.take(radial_feat, degree_index_m, axis=1) # (E, D_m, C) + if self.radial_hidden_proj is not None: + rad_feat = self.radial_hidden_proj(rad_feat) + if self.radial_degree_mixer is None: + x_local = x_local * rad_feat + else: + x_local = self.radial_degree_mixer(x_local, rad_feat) + rad_feat_l0_focus = xp.reshape( + rad_feat[:, 0, :], (n_edge, self.n_focus, self.so2_focus_dim) + ) # (E, F, Cf) + + # === Step 4. Convert to SO(2) internal focus layout === + focus_gate_src: Any = None + x_local = xp.permute_dims( + xp.reshape( + x_local, (n_edge, self.reduced_dim, self.n_focus, self.so2_focus_dim) + ), + (0, 2, 1, 3), + ) # (E, F, D_m, Cf) + if self.focus_compete and self.n_focus > 1: + focus_gate_src = x_local[:, :, 0, :] + + # === Step 5. Multi-layer SO(2) mixing (pre-norm + residual) === + def apply_bias_correction( + x_local: Any, + so2_linear: SO2Linear, + layer_idx: int, + ) -> Any: + if layer_idx != 0 or so2_linear.bias0 is None: + return x_local + bias0 = xp.reshape( + xp.asarray(so2_linear.bias0[...], device=device), + (1, self.n_focus, so2_linear.out_channels), + ) + if so2_linear.out_channels == self.so2_focus_dim: + radial_factor = rad_feat_l0_focus + else: + raise RuntimeError( + "Unexpected SO2Linear output width in bias correction" + ) + edge_env = xp.reshape( + xp.astype(edge_cache.edge_env, x_local.dtype), (n_edge, 1, 1) + ) + bias_correction = bias0 * (radial_factor * edge_env - 1.0) + x0 = x_local[:, :, :1, :] + bias_correction[:, :, None, :] + return xp.concat([x0, x_local[:, :, 1:, :]], axis=2) + + for layer_idx, (so2_linear, inter_norm, non_linear) in enumerate( + zip( + self.so2_linears, + self.so2_inter_norms, + self.non_linearities, + strict=True, + ) + ): + residual = x_local + if inter_norm is not None: + x_local = inter_norm(x_local) + x_local = so2_linear(x_local) + x_local = apply_bias_correction(x_local, so2_linear, layer_idx) + + if non_linear is not None: + x_local = non_linear(x_local) + + x_local = residual + x_local + + # === Step 6. Cross-focus softmax competition === + if self.focus_compete and self.n_focus > 1: + compete_w = xp.asarray(self.adamw_focus_compete_w[...], device=device) + gate_in = xp.astype(focus_gate_src, compete_w.dtype) + gate_normed = self.focus_compete_norm(gate_in) # (E, F, Cf) + # einsum "efi,if->ef" + focus_logits = xp.sum( + gate_normed * xp.permute_dims(compete_w, (1, 0))[None, ...], + axis=-1, + ) + if self.mlp_bias: + focus_logits = ( + focus_logits + + xp.asarray(self.focus_compete_bias[...], device=device)[None, :] + ) + focus_logits = focus_logits / self.focus_softmax_tau + logits_max = xp.max(focus_logits, axis=1, keepdims=True) + exp_logits = xp.exp(focus_logits - logits_max) + alpha = exp_logits / xp.sum(exp_logits, axis=1, keepdims=True) + alpha = xp.astype(alpha, x_local.dtype) + alpha = alpha * (1.0 - self.focus_label_smoothing) + ( + self.focus_label_smoothing / float(self.n_focus) + ) + x_local = x_local * alpha[:, :, None, None] + + # === Step 7. Rotate back to global frame === + Dt_full = edge_cache.Dt_full + # Restore reduced global layout (E, D_m, C_wide) for inverse rotation. + x_local = xp.reshape( + xp.permute_dims(x_local, (0, 2, 1, 3)), + (n_edge, self.reduced_dim, self.hidden_channels), + ) + Dt_from_m = project_Dt_from_m( + Dt_full=Dt_full, + coeff_index_m=self.coeff_index_m, + ebed_dim_full=self.ebed_dim_full, + cache=edge_cache.Dt_from_m_cache, + key_lmax=self.lmax, + key_mmax=self.mmax, + ) + x_message = xp.matmul(Dt_from_m, x_local) # (E, D, C_wide) + # Reduced layouts keep only 2*mmax+1 orders for l>mmax. Applying the + # inverse-rotation degree rescale after the global lift restores the + # full-basis amplitude expected by the block output contract. + rescale = xp.astype( + xp.asarray(self.rotate_inv_rescale_full, device=device), x_message.dtype + ) + x_message = x_message * xp.reshape(rescale, (1, -1, 1)) + + # === Step 8. Aggregate with optional head-wise gating === + # Source Freeze Propagation Gate: broadcast the per-edge scalar + # eta[src] to the edge message before destination aggregation. + edge_src_gate = edge_cache.edge_src_gate + if self.n_atten_head == 0: + # Baseline path: envelope-weighted masked sum -> degree norm. + edge_weight = xp.astype(edge_cache.edge_env, x_message.dtype) # (E, 1) + edge_weight = xp.reshape(edge_weight, (n_edge, 1)) + if edge_src_gate is not None: + edge_weight = edge_weight * xp.astype( + xp.reshape(edge_src_gate, (n_edge, 1)), edge_weight.dtype + ) + x_message = x_message * edge_weight[:, :, None] + # pt: out.index_add_(0, dst, x_message) — padded-edge masked sum + # over the nnei axis (dst is slot-implicit). + x_message = x_message * mask_f[:, None, None] + out = xp.sum( + xp.reshape( + x_message, + (n_node, nnei, self.ebed_dim_full, self.hidden_channels), + ), + axis=1, + ) + inv_sqrt_deg = xp.astype(edge_cache.inv_sqrt_deg, out.dtype) + out = out * inv_sqrt_deg # (N, D, C_wide) + else: + # === Step 8.1. Build attention logits from scalar channels === + qk_w = xp.asarray(self.attn_q_proj.weight[...], device=device) + x_l0_node = xp.reshape( + x[:, 0, :], (n_node, self.attn_n_focus, self.attn_focus_dim) + ) # (N, Fa, Ca) + x_l0_node = xp.astype(x_l0_node, qk_w.dtype) + qk_input = self.attn_qk_norm(x_l0_node) + q_node = self.attn_q_proj(qk_input) # (N, Fa, Ca) + k_node = self.attn_k_proj(qk_input) # (N, Fa, Ca) + dst_idx = xp.astype(xp.reshape(dst, (n_edge,)), xp.int64) + q_edge = xp.reshape( + xp.take(q_node, dst_idx, axis=0), + (n_edge, self.attn_n_focus, self.n_atten_head, self.head_dim), + ) # (E, Fa, H, Ch), Ca = H * Ch + k_edge = xp.reshape( + xp.take(k_node, src_idx, axis=0), + (n_edge, self.attn_n_focus, self.n_atten_head, self.head_dim), + ) # (E, Fa, H, Ch) + radial_l0 = xp.reshape( + rad_feat[:, 0, :], (n_edge, self.attn_n_focus, self.attn_focus_dim) + ) # (E, Fa, Ca) + radial_l0 = xp.astype(radial_l0, qk_w.dtype) + # einsum "efi,ifo->efo" as a broadcast batched matmul. + logit_w = xp.permute_dims( + xp.asarray(self.adamw_attn_logit_w[...], device=device), (1, 0, 2) + ) # (Fa, Ca, H) + radial_bias = xp.matmul(radial_l0[:, :, None, :], logit_w[None, ...])[ + ..., 0, : + ] # (E, Fa, H) + attn_logits = xp.sum(q_edge * k_edge, axis=-1) * (self.head_dim**-0.5) + attn_logits = attn_logits + radial_bias + + # === Step 8.2. Destination-wise stable envelope-gated softmax === + # pt: scatter-based segment softmax keyed by dst — padded-edge + # masked softmax over the nnei axis. ``src_weight=edge_src_gate`` + # folds SFPG into both the numerator and the denominator so a + # muted source drops out of the normalization entirely. + attn_alpha = segment_envelope_gated_softmax( + logits=attn_logits, + edge_env=xp.astype(edge_cache.edge_env, attn_logits.dtype), + n_nodes=n_node, + z_bias_raw=xp.asarray(self.adamw_attn_z_bias_raw, device=device), + eps=self.eps, + src_weight=( + None + if edge_src_gate is None + else xp.astype(edge_src_gate, attn_logits.dtype) + ), + edge_mask=mask_f, + ) # (E, F, H) + + # === Step 8.3. Value projection and head-wise aggregation === + value_heads = xp.reshape( + xp.astype(x_message, qk_w.dtype), + ( + n_edge, + self.ebed_dim_full, + self.attn_n_focus, + self.n_atten_head, + self.head_dim, + ), + ) # (E, D, Fa, H, Ch) + weighted_value = value_heads * xp.reshape( + attn_alpha, (n_edge, 1, self.attn_n_focus, self.n_atten_head, 1) + ) + # pt: out_heads.index_add_(0, dst, weighted_value) — padded-edge + # masked sum over the nnei axis (dst is slot-implicit). + weighted_value = ( + weighted_value + * xp.astype(mask_f, weighted_value.dtype)[:, None, None, None, None] + ) + out_heads = xp.sum( + xp.reshape( + weighted_value, + ( + n_node, + nnei, + self.ebed_dim_full, + self.attn_n_focus, + self.n_atten_head, + self.head_dim, + ), + ), + axis=1, + ) # (N, D, Fa, H, Ch) + + # === Step 8.4. Output-side head gate === + gate_w = xp.permute_dims( + xp.asarray(self.adamw_attn_gate_w[...], device=device), (1, 0, 2) + ) # (Fa, Ca, H) + gate_in = self.attn_output_gate_norm(x_l0_node) + attn_output_gate = xp_sigmoid( + xp.matmul(gate_in[:, :, None, :], gate_w[None, ...])[..., 0, :] + ) # (N, F, H) + out_heads = out_heads * xp.reshape( + attn_output_gate, + (n_node, 1, self.attn_n_focus, self.n_atten_head, 1), + ) # (N, D, Fa, H, Ch) + + # === Step 8.5. Merge heads === + out = xp.astype( + xp.reshape( + out_heads, (n_node, self.ebed_dim_full, self.hidden_channels) + ), + x.dtype, + ) # (N, D, C_wide) + + # === Step 10. Final channel mixing === + out = self.post_focus_mix(out[:, :, None, :])[:, :, 0, :] + return out # (N, D, C) + + def _variables(self) -> dict[str, np.ndarray]: + """Variables keyed by the pt ``state_dict`` key names.""" + variables: dict[str, np.ndarray] = {} + for i, so2_linear in enumerate(self.so2_linears): + for key, value in so2_linear._variables().items(): + variables[f"so2_linears.{i}.{key}"] = value + for i, inter_norm in enumerate(self.so2_inter_norms): + if inter_norm is not None: + for key, value in inter_norm.serialize()["@variables"].items(): + variables[f"so2_inter_norms.{i}.{key}"] = value + for i, non_linear in enumerate(self.non_linearities): + if non_linear is not None: + for key, value in non_linear.serialize()["@variables"].items(): + variables[f"non_linearities.{i}.{key}"] = value + if self.n_atten_head > 0: + variables["adamw_attn_logit_w"] = to_numpy_array(self.adamw_attn_logit_w) + variables["adamw_attn_z_bias_raw"] = to_numpy_array( + self.adamw_attn_z_bias_raw + ) + variables["adamw_attn_gate_w"] = to_numpy_array(self.adamw_attn_gate_w) + variables["attn_qk_norm.adam_scale"] = to_numpy_array( + self.attn_qk_norm.adam_scale + ) + variables["attn_q_proj.weight"] = to_numpy_array(self.attn_q_proj.weight) + variables["attn_k_proj.weight"] = to_numpy_array(self.attn_k_proj.weight) + variables["attn_output_gate_norm.adam_scale"] = to_numpy_array( + self.attn_output_gate_norm.adam_scale + ) + if self.focus_compete_norm is not None: + variables["adamw_focus_compete_w"] = to_numpy_array( + self.adamw_focus_compete_w + ) + variables["focus_compete_norm.adam_scale"] = to_numpy_array( + self.focus_compete_norm.adam_scale + ) + if self.mlp_bias: + variables["focus_compete_bias"] = to_numpy_array( + self.focus_compete_bias + ) + if self.radial_hidden_proj is not None: + variables["radial_hidden_proj.weight"] = to_numpy_array( + self.radial_hidden_proj.weight + ) + if self.radial_degree_mixer is not None: + for key, value in self.radial_degree_mixer._variables().items(): + variables[f"radial_degree_mixer.{key}"] = value + for name, mix in ( + ("pre_focus_mix", self.pre_focus_mix), + ("post_focus_mix", self.post_focus_mix), + ): + for key, value in mix.serialize()["@variables"].items(): + variables[f"{name}.{key}"] = value + variables["coeff_index_m"] = to_numpy_array(self.coeff_index_m) + variables["degree_index_m"] = to_numpy_array(self.degree_index_m) + variables["rotate_inv_rescale_full"] = to_numpy_array( + self.rotate_inv_rescale_full + ) + return variables + + def _load_variables(self, variables: dict[str, Any]) -> None: + """Load variables keyed by the pt ``state_dict`` key names.""" + variables = dict(variables) + prec = PRECISION_DICT[self.precision.lower()] + cprec = PRECISION_DICT[self.compute_precision.lower()] + + def pop(key: str) -> Any: + try: + return variables.pop(key) + except KeyError: + raise KeyError(f"Missing variable: {key}") from None + + def sub_vars(prefix: str) -> dict[str, Any]: + full = f"{prefix}." + out = { + key[len(full) :]: value + for key, value in variables.items() + if key.startswith(full) + } + for key in list(variables): + if key.startswith(full): + del variables[key] + if not out: + raise KeyError(f"Missing variables with prefix: {full}") + return out + + # Top-level index buffers: validate against the config-derived tables. + _check_index_table(self.coeff_index_m, pop("coeff_index_m"), "coeff_index_m") + _check_index_table(self.degree_index_m, pop("degree_index_m"), "degree_index_m") + _check_shape_assign( + self, + "rotate_inv_rescale_full", + pop("rotate_inv_rescale_full"), + prec, + "rotate_inv_rescale_full", + ) + + for i, so2_linear in enumerate(self.so2_linears): + so2_linear._load_variables(sub_vars(f"so2_linears.{i}")) + for i, inter_norm in enumerate(self.so2_inter_norms): + if inter_norm is not None: + sv = sub_vars(f"so2_inter_norms.{i}") + _check_index_table( + inter_norm.degree_index_m, + sv["degree_index_m"], + f"so2_inter_norms.{i}.degree_index_m", + ) + for name in ("balance_weight", "adam_scale", "bias0"): + _check_shape_assign( + inter_norm, + name, + sv[name], + cprec, + f"so2_inter_norms.{i}.{name}", + ) + for i, non_linear in enumerate(self.non_linearities): + if non_linear is not None: + sv = sub_vars(f"non_linearities.{i}") + _check_index_table( + non_linear.expand_index, + sv["expand_index"], + f"non_linearities.{i}.expand_index", + ) + _check_shape_assign( + non_linear.gate_linear, + "weight", + sv["gate_linear.weight"], + cprec, + f"non_linearities.{i}.gate_linear.weight", + ) + if self.mlp_bias: + _check_shape_assign( + non_linear.gate_linear, + "bias", + sv["gate_linear.bias"], + cprec, + f"non_linearities.{i}.gate_linear.bias", + ) + if self.n_atten_head > 0: + for name in ( + "adamw_attn_logit_w", + "adamw_attn_z_bias_raw", + "adamw_attn_gate_w", + ): + _check_shape_assign(self, name, pop(name), cprec, name) + _check_shape_assign( + self.attn_qk_norm, + "adam_scale", + pop("attn_qk_norm.adam_scale"), + cprec, + "attn_qk_norm.adam_scale", + ) + _check_shape_assign( + self.attn_q_proj, + "weight", + pop("attn_q_proj.weight"), + cprec, + "attn_q_proj.weight", + ) + _check_shape_assign( + self.attn_k_proj, + "weight", + pop("attn_k_proj.weight"), + cprec, + "attn_k_proj.weight", + ) + _check_shape_assign( + self.attn_output_gate_norm, + "adam_scale", + pop("attn_output_gate_norm.adam_scale"), + cprec, + "attn_output_gate_norm.adam_scale", + ) + if self.focus_compete_norm is not None: + _check_shape_assign( + self, + "adamw_focus_compete_w", + pop("adamw_focus_compete_w"), + cprec, + "adamw_focus_compete_w", + ) + _check_shape_assign( + self.focus_compete_norm, + "adam_scale", + pop("focus_compete_norm.adam_scale"), + cprec, + "focus_compete_norm.adam_scale", + ) + if self.mlp_bias: + _check_shape_assign( + self, + "focus_compete_bias", + pop("focus_compete_bias"), + cprec, + "focus_compete_bias", + ) + if self.radial_hidden_proj is not None: + _check_shape_assign( + self.radial_hidden_proj, + "weight", + pop("radial_hidden_proj.weight"), + prec, + "radial_hidden_proj.weight", + ) + if self.radial_degree_mixer is not None: + self.radial_degree_mixer._load_variables(sub_vars("radial_degree_mixer")) + for name, mix in ( + ("pre_focus_mix", self.pre_focus_mix), + ("post_focus_mix", self.post_focus_mix), + ): + sv = sub_vars(name) + _check_index_table( + mix.expand_index, sv["expand_index"], f"{name}.expand_index" + ) + _check_shape_assign(mix, "weight", sv["weight"], prec, f"{name}.weight") + if self.mlp_bias: + _check_shape_assign(mix, "bias", sv["bias"], prec, f"{name}.bias") + + if variables: + raise KeyError(f"Unknown variables: {sorted(variables)}") + + def serialize(self) -> dict[str, Any]: + """Serialize the SO2Convolution to a dict (pt-compatible format).""" + return { + "@class": "SO2Convolution", + "@version": 1, + "config": { + "lmax": self.lmax, + "mmax": self.mmax, + "kmax": self.kmax, + "channels": self.channels, + "n_focus": self.n_focus, + "focus_dim": self.focus_dim, + "focus_compete": self.focus_compete, + "so2_norm": self.so2_norm, + "so2_layers": self.so2_layers, + "so2_attn_res": self.so2_attn_res_mode, + "layer_scale": self.layer_scale, + "n_atten_head": self.n_atten_head, + "atten_f_mix": self.atten_f_mix, + "atten_v_proj": self.use_atten_v_proj, + "atten_o_proj": self.use_atten_o_proj, + "s2_activation": self.s2_activation, + "node_wise_grid_mlp": self.node_wise_grid_mlp, + "node_wise_grid_branch": self.node_wise_grid_branch, + "message_node_grid_mlp": self.message_node_grid_mlp, + "message_node_grid_branch": self.message_node_grid_branch, + "node_wise_s2": self.node_wise_s2, + "node_wise_so3": self.node_wise_so3, + "message_node_s2": self.message_node_s2, + "message_node_so3": self.message_node_so3, + "lebedev_quadrature": self.lebedev_quadrature, + "activation_function": self.activation_function, + "mlp_bias": self.mlp_bias, + "radial_so2_mode": self.radial_so2_mode, + "radial_so2_rank": self.radial_so2_rank, + "eps": self.eps, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "seed": None, + }, + "@variables": self._variables(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SO2Convolution: + """Deserialize an SO2Convolution from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "SO2Convolution": + raise ValueError(f"Invalid class for SO2Convolution: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = dict(data.pop("config")) + variables = data.pop("@variables") + config["precision"] = str(config.pop("precision")) + obj = cls(**config) + obj._load_variables(variables) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py new file mode 100644 index 0000000000..60e8c2bceb --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py @@ -0,0 +1,540 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +SO(3)-equivariant linear layers for DPA4/SeZM. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.so3``. +It defines the channel-only and focus-aware linear maps used by the DPA4 +SO(3) feature transformations. All three pt classes are ported: +``FocusLinear`` (used by ``so2``, ``grid_net``, ``activation``), +``ChannelLinear`` (used by ``so2``, ``grid_net``), and ``SO3Linear`` +(used by ``so2``, ``ffn``). + +Serialization contract: ``SO3Linear`` mirrors the pt ``serialize()`` format +exactly (same config and ``@variables`` keys), so pt ``serialize()`` output +deserializes directly. The pt ``FocusLinear`` and ``ChannelLinear`` define no +``serialize()`` (they only appear nested inside larger modules' state_dicts); +their dpmodel ``serialize()``/``deserialize()`` use ``@variables`` keys equal +to the pt ``state_dict`` key names (``weight``, ``bias``) so that pt +state-dict fragments load directly. + +Weight initialization is distribution-equivalent to the pt version (drawn +from ``np.random.default_rng`` instead of the torch generator stream), the +same convention as ``utils.init_trunc_normal_fan_in_out``. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .indexing import ( + get_so3_dim_of_lmax, + map_degree_idx, +) +from .utils import ( + init_trunc_normal_fan_in_out, +) + + +class FocusLinear(NativeOP): + """ + Per-focus linear projection on the last feature axis. + + Notes + ----- + Parameters are stored in (in, out) convention to match Muon's rectangular + correction assumption (rows=fan_in, cols=fan_out): + - weight: (in_channels, n_focus * out_channels) + - bias: (n_focus * out_channels,) + + Parameters + ---------- + in_channels : int + Input feature dimension. + out_channels : int + Output feature dimension. + n_focus : int + Number of focus streams. + precision : str + Parameter precision. + bias : bool + Whether to use bias. + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for initialization. + init_std : float | None + If given, use normal(0, init_std) instead of default uniform init. + Useful for gate projections where small initial logits are desired. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: int, + n_focus: int, + precision: str = DEFAULT_PRECISION, + bias: bool = True, + trainable: bool = True, + seed: int | list[int] | None = None, + init_std: float | None = None, + ) -> None: + self.in_channels = int(in_channels) + self.out_channels = int(out_channels) + self.n_focus = int(n_focus) + self.precision = precision + self.trainable = bool(trainable) + self.use_bias = bool(bias) + prec = PRECISION_DICT[self.precision.lower()] + rng = np.random.default_rng(seed) + shape = (self.in_channels, self.n_focus * self.out_channels) + if init_std is not None: + weight = rng.normal(0.0, float(init_std), size=shape) + else: + bound = 1.0 / math.sqrt(self.in_channels) + weight = rng.uniform(-bound, bound, size=shape) + self.weight = weight.astype(prec) + if self.use_bias: + self.bias: np.ndarray | None = np.zeros( + (self.n_focus * self.out_channels,), dtype=prec + ) + else: + self.bias = None + + def call(self, x: Any) -> Any: + """ + Apply the per-focus linear projection. + + Parameters + ---------- + x : Array + Input array with shape (B, F, Cin). + + Returns + ------- + Array + Projected array with shape (B, F, Cout). + """ + xp = array_api_compat.array_namespace(x) + weight = xp.asarray(self.weight[...], device=array_api_compat.device(x)) + weight = xp.reshape(weight, (self.in_channels, self.n_focus, self.out_channels)) + # einsum "bfi,ifo->bfo" as a broadcast batched matmul: + # (B, F, 1, Cin) @ (1, F, Cin, Cout) -> (B, F, 1, Cout) + weight = xp.permute_dims(weight, (1, 0, 2)) # (F, Cin, Cout) + out = xp.matmul(x[:, :, None, :], weight[None, ...])[..., 0, :] + if self.use_bias: + bias = xp.asarray(self.bias[...], device=array_api_compat.device(x)) + bias = xp.reshape(bias, (self.n_focus, self.out_channels)) + out = out + bias[None, ...] + return out + + def serialize(self) -> dict[str, Any]: + """Serialize the FocusLinear to a dict. + + The pt ``FocusLinear`` has no ``serialize()``; the ``@variables`` keys + here match the pt ``state_dict`` key names (``weight``, ``bias``). + """ + variables = {"weight": to_numpy_array(self.weight)} + if self.use_bias: + variables["bias"] = to_numpy_array(self.bias) + return { + "@class": "FocusLinear", + "@version": 1, + "config": { + "in_channels": self.in_channels, + "out_channels": self.out_channels, + "n_focus": self.n_focus, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "bias": self.use_bias, + "trainable": self.trainable, + "seed": None, + }, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> FocusLinear: + """Deserialize a FocusLinear from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "FocusLinear": + raise ValueError(f"Invalid class for FocusLinear: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + in_channels=int(config["in_channels"]), + out_channels=int(config["out_channels"]), + n_focus=int(config["n_focus"]), + precision=str(config["precision"]), + bias=bool(config["bias"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + prec = PRECISION_DICT[obj.precision.lower()] + weight = np.asarray(variables["weight"], dtype=prec) + if weight.shape != obj.weight.shape: + raise ValueError( + f"weight shape {weight.shape} does not match " + f"the expected shape {obj.weight.shape}" + ) + obj.weight = weight + if obj.use_bias: + obj.bias = np.asarray(variables["bias"], dtype=prec).reshape(obj.bias.shape) + return obj + + +class ChannelLinear(NativeOP): + """ + Channel-only linear projection on the last feature axis. + + Notes + ----- + Parameters are stored in (in, out) convention to match Muon's rectangular + correction assumption (rows=fan_in, cols=fan_out): + - weight: (in_channels, out_channels) + - bias: (out_channels,) + + Parameters + ---------- + in_channels : int + Input feature dimension. + out_channels : int + Output feature dimension. + precision : str + Parameter precision. + bias : bool + Whether to use bias. + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for initialization. + init_std : float | None + If given, use normal(0, init_std) instead of default uniform init. + Useful for gate projections where small initial logits are desired. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: int, + precision: str = DEFAULT_PRECISION, + bias: bool = True, + trainable: bool = True, + seed: int | list[int] | None = None, + init_std: float | None = None, + ) -> None: + self.in_channels = int(in_channels) + self.out_channels = int(out_channels) + self.precision = precision + self.trainable = bool(trainable) + self.use_bias = bool(bias) + prec = PRECISION_DICT[self.precision.lower()] + rng = np.random.default_rng(seed) + shape = (self.in_channels, self.out_channels) + if init_std is not None: + weight = rng.normal(0.0, float(init_std), size=shape) + else: + bound = 1.0 / math.sqrt(self.in_channels) + weight = rng.uniform(-bound, bound, size=shape) + self.weight = weight.astype(prec) + if self.use_bias: + self.bias: np.ndarray | None = np.zeros((self.out_channels,), dtype=prec) + else: + self.bias = None + + def call(self, x: Any) -> Any: + """ + Apply the channel-only linear projection. + + Parameters + ---------- + x : Array + Input array with shape ``(..., C_in)``. + + Returns + ------- + Array + Projected array with shape ``(..., C_out)``. + """ + xp = array_api_compat.array_namespace(x) + # einsum "...i,io->...o" is a plain matmul on the last axis + device = array_api_compat.device(x) + out = xp.matmul(x, xp.asarray(self.weight[...], device=device)) + if self.use_bias: + out = out + xp.asarray(self.bias[...], device=device) + return out + + def serialize(self) -> dict[str, Any]: + """Serialize the ChannelLinear to a dict. + + The pt ``ChannelLinear`` has no ``serialize()``; the ``@variables`` + keys here match the pt ``state_dict`` key names (``weight``, ``bias``). + """ + variables = {"weight": to_numpy_array(self.weight)} + if self.use_bias: + variables["bias"] = to_numpy_array(self.bias) + return { + "@class": "ChannelLinear", + "@version": 1, + "config": { + "in_channels": self.in_channels, + "out_channels": self.out_channels, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "bias": self.use_bias, + "trainable": self.trainable, + "seed": None, + }, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> ChannelLinear: + """Deserialize a ChannelLinear from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "ChannelLinear": + raise ValueError(f"Invalid class for ChannelLinear: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + in_channels=int(config["in_channels"]), + out_channels=int(config["out_channels"]), + precision=str(config["precision"]), + bias=bool(config["bias"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + prec = PRECISION_DICT[obj.precision.lower()] + weight = np.asarray(variables["weight"], dtype=prec) + if weight.shape != obj.weight.shape: + raise ValueError( + f"weight shape {weight.shape} does not match " + f"the expected shape {obj.weight.shape}" + ) + obj.weight = weight + if obj.use_bias: + obj.bias = np.asarray(variables["bias"], dtype=prec).reshape(obj.bias.shape) + return obj + + +class SO3Linear(NativeOP): + """ + Focus-aware degree-wise linear self-interaction. + + The key insight is that weights are shared across all ``m`` components + within each ``l`` block. + + Notes + ----- + - Weight storage: ``(lmax+1, C_in, F*C_out)``. + - Bias storage: ``(F*C_out,)``, only applied to ``l=0`` scalar components. + - Runtime view restores weights to ``(lmax+1, C_in, F, C_out)`` via reshape. + - ``expand_index`` maps each packed ``(l,m)`` position to its ``l`` value. + - The pt einsum ``ndfi,difo->ndfo`` is expressed as a broadcast batched + matmul, which keeps the whole multi-focus path vectorized. + + Parameters + ---------- + lmax : int + Maximum spherical harmonic degree. + in_channels : int + Number of input channels per (l, m) coefficient. + out_channels : int + Number of output channels per (l, m) coefficient. + n_focus : int + Number of focus streams. + precision : str + Parameter precision. + mlp_bias : bool + Whether to use bias for l=0 (scalar) components. + trainable : bool + Whether parameters are trainable. + seed : int | list[int] | None + Random seed for weight initialization. + init_std : float | None + If given, use normal(0, init_std) for all weights instead of default + trunc-normal fan-in/fan-out init. Use 0.0 for zero initialization. + """ + + def __init__( + self, + *, + lmax: int, + in_channels: int, + out_channels: int, + n_focus: int = 1, + precision: str = DEFAULT_PRECISION, + mlp_bias: bool = False, + trainable: bool = True, + seed: int | list[int] | None = None, + init_std: float | None = None, + ) -> None: + self.lmax = int(lmax) + self.in_channels = int(in_channels) + self.out_channels = int(out_channels) + self.n_focus = int(n_focus) + self.precision = precision + self.trainable = bool(trainable) + self.ebed_dim = get_so3_dim_of_lmax(self.lmax) + self.mlp_bias = bool(mlp_bias) + prec = PRECISION_DICT[self.precision.lower()] + + # === Step 1. Per-l weight matrix with focus folded on output axis === + # Storage: (lmax+1, C_in, F*C_out); runtime view: (lmax+1, C_in, F, C_out). + num_l = self.lmax + 1 + weight = np.empty( + (num_l, self.in_channels, self.n_focus * self.out_channels), + dtype=prec, + ) + if init_std is not None: + if init_std == 0.0: + weight[...] = 0.0 + else: + rng = np.random.default_rng(seed) + weight[...] = rng.normal(0.0, float(init_std), size=weight.shape) + else: + for l_idx in range(num_l): + init_trunc_normal_fan_in_out( + weight[l_idx], + child_seed(seed, 1000 + l_idx), + ) + self.weight = weight + + # === Step 2. Bias only for l=0 (scalar components) === + if self.mlp_bias: + self.bias: np.ndarray | None = np.zeros( + (self.n_focus * self.out_channels,), dtype=prec + ) + else: + self.bias = None + + # === Step 3. Precompute expand_index for weight lookup === + self.expand_index = map_degree_idx(self.lmax) + + def call(self, x: Any) -> Any: + """ + Apply the degree-wise linear self-interaction. + + Parameters + ---------- + x : Array + Input features with shape (N, D, F, C_in) where D=(lmax+1)^2. + + Returns + ------- + Array + Order-wise mixed features with shape (N, D, F, C_out). + """ + xp = array_api_compat.array_namespace(x) + # === Step 1. Expand per-l weights to packed coefficient layout === + # (L, Cin, F*Cout) -> (L, Cin, F, Cout) + weight = xp.reshape( + xp.asarray(self.weight[...], device=array_api_compat.device(x)), + (self.lmax + 1, self.in_channels, self.n_focus, self.out_channels), + ) # (L, Cin, F, Cout) + # (L, Cin, F, Cout) -> (D, Cin, F, Cout) + expand_index = xp.asarray(self.expand_index, device=array_api_compat.device(x)) + weight_expanded = xp.take(weight, expand_index, axis=0) + + # === Step 2. Per-focus, per-degree channel mixing === + # einsum "ndfi,difo->ndfo" as a broadcast batched matmul: + # (N, D, F, 1, Cin) @ (1, D, F, Cin, Cout) -> (N, D, F, 1, Cout) + weight_expanded = xp.permute_dims( + weight_expanded, (0, 2, 1, 3) + ) # (D, F, Cin, Cout) + out = xp.matmul(x[:, :, :, None, :], weight_expanded[None, ...])[..., 0, :] + + # === Step 3. Add l=0 bias === + if self.mlp_bias: + bias = xp.asarray(self.bias[...], device=array_api_compat.device(x)) + bias = xp.reshape(bias, (self.n_focus, self.out_channels)) + out0 = out[:, :1, :, :] + bias[None, None, ...] + out = xp.concat([out0, out[:, 1:, :, :]], axis=1) if self.lmax > 0 else out0 + + return out + + def serialize(self) -> dict[str, Any]: + """Serialize the SO3Linear to a dict (pt-compatible format).""" + variables = {"weight": to_numpy_array(self.weight)} + if self.mlp_bias: + variables["bias"] = to_numpy_array(self.bias) + variables["expand_index"] = to_numpy_array(self.expand_index) + return { + "@class": "SO3Linear", + "@version": 1, + "config": { + "lmax": self.lmax, + "in_channels": self.in_channels, + "out_channels": self.out_channels, + "n_focus": self.n_focus, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "mlp_bias": self.mlp_bias, + "trainable": self.trainable, + "seed": None, + }, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SO3Linear: + """Deserialize an SO3Linear from a dict.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "SO3Linear": + raise ValueError(f"Invalid class for SO3Linear: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + config = data.pop("config") + variables = data.pop("@variables") + obj = cls( + lmax=int(config["lmax"]), + in_channels=int(config["in_channels"]), + out_channels=int(config["out_channels"]), + n_focus=int(config["n_focus"]), + precision=str(config["precision"]), + mlp_bias=bool(config["mlp_bias"]), + trainable=bool(config["trainable"]), + seed=config.get("seed"), + ) + prec = PRECISION_DICT[obj.precision.lower()] + expand_index = np.asarray(variables["expand_index"], dtype=np.int64) + if not np.array_equal(expand_index, obj.expand_index): + raise ValueError("expand_index does not match the lmax-derived table") + weight = np.asarray(variables["weight"], dtype=prec) + if weight.shape != obj.weight.shape: + raise ValueError( + f"weight shape {weight.shape} does not match " + f"the expected shape {obj.weight.shape}" + ) + obj.weight = weight + if obj.mlp_bias: + obj.bias = np.asarray(variables["bias"], dtype=prec).reshape(obj.bias.shape) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/utils.py b/deepmd/dpmodel/descriptor/dpa4_nn/utils.py new file mode 100644 index 0000000000..5a016f6e94 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/utils.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Utility helpers for the DPA4/SeZM descriptor package. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.utils``. +It provides the small numeric helpers shared across the DPA4 descriptor +implementation. + +Init-time helpers (``init_trunc_normal_fan_in_out``) operate on static numpy +data and are plain numpy by design (not array-API). ``safe_norm`` operates on +runtime tensors and is array-API compatible. + +Helpers from the pt version intentionally NOT ported: + +- ``nvtx_range``: CUDA profiling, torch-only. +- ``use_triton_infer``: Triton inference kernels, torch-only. +- ``safe_numpy_to_tensor``: numpy -> torch conversion glue; dpmodel code uses + ``xp.asarray`` directly. +- ``np_safe``: torch -> numpy conversion glue; dpmodel code uses + ``deepmd.dpmodel.common.to_numpy_array`` instead. + +``get_promoted_dtype`` IS ported (numpy equivalent) because core modules use it +to pick a stable computation/storage dtype. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +ATTN_RES_MODES = ("none", "independent", "dependent") + + +def init_trunc_normal_fan_in_out( + weight: np.ndarray, + seed: int | list[int] | None, + scale: float = 1.0, +) -> None: + """Initialize weight with truncated normal distribution. + + Uses Xavier-like variance scaling: std = scale / sqrt(fan_in + fan_out). + Truncation at +/-3*std prevents extreme outliers. + + NumPy equivalent of the pt version: the weight is filled in place from a + ``np.random.default_rng(seed)`` stream (distribution-equivalent to the + torch version, not RNG-stream-identical). + + Parameters + ---------- + weight : np.ndarray + Weight array with shape (out_features, in_features), modified in place. + seed : int | list[int] | None + Random seed for reproducibility. + scale : float, default=1.0 + Multiplicative scale factor in the standard deviation numerator. + """ + if weight.ndim != 2: + raise ValueError("`weight` must be a 2D tensor") + if scale <= 0: + raise ValueError("`scale` must be positive") + fan_out, fan_in = weight.shape + std = float(scale) / math.sqrt(fan_in + fan_out) + rng = np.random.default_rng(seed) + # rejection sampling: exact truncated normal on [-3*std, 3*std] + values = rng.normal(0.0, std, size=weight.shape) + out_of_bounds = np.abs(values) > 3.0 * std + while out_of_bounds.any(): + values[out_of_bounds] = rng.normal( + 0.0, std, size=int(np.count_nonzero(out_of_bounds)) + ) + out_of_bounds = np.abs(values) > 3.0 * std + weight[...] = values.astype(weight.dtype, copy=False) + + +def safe_norm(x: Any, eps: float = 1e-7) -> Any: + """ + Compute vector norm with smooth epsilon regularization. + + Uses float32 for computation when input is fp16/bf16. This function + operates on runtime tensors and is array-API compatible. + + Parameters + ---------- + x : Array + Input array with shape (N, 3), where N is the number of vectors. + eps : float + Lower bound for the norm. + + Returns + ------- + Array + Norm with shape (N, 1). + """ + xp = array_api_compat.array_namespace(x) + in_dtype = x.dtype + # matches "float16" and "bfloat16" dtype names across namespaces + promote = "float16" in str(in_dtype) + if promote: + x = xp.astype(x, xp.float32) + norm = xp.sqrt(xp.sum(x * x, axis=-1, keepdims=True) + float(eps) * float(eps)) + if promote: + norm = xp.astype(norm, in_dtype) + return norm + + +def get_promoted_dtype(dtype: Any) -> Any: + """ + Get promoted dtype for numerical stability. + + For bf16/fp16, use float32 to ensure numerical stability + in computation and storage compatibility. + + NumPy equivalent of the pt version; accepts a numpy dtype (including + ``ml_dtypes.bfloat16``) and returns a numpy dtype. + """ + name = getattr(dtype, "name", None) or str(dtype) + if "float16" in name: # matches float16 and bfloat16 + return np.dtype(np.float32) + return dtype diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/wignerd.py b/deepmd/dpmodel/descriptor/dpa4_nn/wignerd.py new file mode 100644 index 0000000000..0612d233be --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/wignerd.py @@ -0,0 +1,963 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Quaternion-based Wigner-D and edge-frame utilities for the DPA4/SeZM descriptor. + +This module is the dpmodel port of ``deepmd.pt.model.descriptor.sezm_nn.wignerd``. +It defines the quaternion helpers and the Wigner-D evaluator used to construct +edge-aligned SO(3) rotation blocks. + +Port notes +---------- +- The pt reference evaluates the ``l=2..10`` blocks with monomial kernels whose + coefficients are solved at init time by ``torch.linalg.lstsq`` against the + generic closed-form quaternion polynomial path (seeded ``torch.randn`` fit + points). That fit is a performance optimization and is not bit-reproducible + without torch. The dpmodel port instead evaluates the generic closed-form + path (the very reference the pt kernels are fitted to) for every ``l >= 2``. + Outputs agree with pt within the fp64 round-off of the pt fit (validated by + the parity tests at ``rtol=1e-12, atol=1e-14``). +- All coefficient tables are plain numpy arrays computed at ``__init__`` time; + ``call`` is pure array-API. The block-diagonal matrices are assembled + functionally with a precomputed ``xp.take`` gather index (no ``__setitem__`` + on traced values), so the path is safe for later torch.export + functionalization. +- Random-gamma gauge randomization is NOT part of this module in pt either: + it lives in ``edge_cache`` and only consumes the deterministic helpers + ``quaternion_z_rotation`` / ``quaternion_multiply`` ported here. + +Serialization contract: pt ``WignerDCalculator.serialize()`` emits only +``{"@class", "@version"}`` (all buffers are derived constants rebuilt from +``lmax``/``dtype`` by the parent). The dpmodel port mirrors that contract. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + NativeOP, +) +from deepmd.dpmodel.common import ( + get_xp_precision, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .utils import ( + safe_norm, +) + + +def quaternion_normalize(q: Any, eps: float = 1e-7) -> Any: + """Normalize quaternions with a differentiable epsilon floor.""" + # safe_norm is the array-API port of pt's _safe_norm_nd (same formula) + return q / safe_norm(q, eps) + + +def quaternion_multiply(q1: Any, q2: Any) -> Any: + """Hamilton product for batched quaternions in ``(w, x, y, z)`` order.""" + xp = array_api_compat.array_namespace(q1, q2) + w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3] + w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3] + return xp.stack( + [ + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ], + axis=-1, + ) + + +def quaternion_to_rotation_matrix(q: Any) -> Any: + """ + Convert unit quaternions to 3x3 rotation matrices. + + The returned matrix is the active rotation represented by ``q``. In SeZM + this is the global->local edge rotation, so multiplying the edge direction + by this matrix sends it to local ``+Z``. + """ + xp = array_api_compat.array_namespace(q) + w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3] + x2 = x * x + y2 = y * y + z2 = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + return xp.stack( + [ + xp.stack( + [1.0 - 2.0 * (y2 + z2), 2.0 * (xy - wz), 2.0 * (xz + wy)], + axis=-1, + ), + xp.stack( + [2.0 * (xy + wz), 1.0 - 2.0 * (x2 + z2), 2.0 * (yz - wx)], + axis=-1, + ), + xp.stack( + [2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (x2 + y2)], + axis=-1, + ), + ], + axis=-2, + ) + + +def quaternion_z_rotation(gamma: Any) -> Any: + """ + Create quaternions for a rotation about the local ``+Z`` axis. + + Parameters + ---------- + gamma + Roll angles in radians with shape ``(E,)``. + + Returns + ------- + Array + Quaternions with shape ``(E, 4)`` in ``(w, x, y, z)`` order. + """ + xp = array_api_compat.array_namespace(gamma) + half_gamma = 0.5 * gamma + w = xp.cos(half_gamma) + x = xp.zeros_like(gamma) + y = xp.zeros_like(gamma) + z = xp.sin(half_gamma) + return xp.stack([w, x, y, z], axis=-1) + + +def _smooth_step_cinf(x: Any) -> Any: + """ + Smooth ``C^inf`` step on ``[0, 1]``. + + This function equals exactly 0 and 1 at the endpoints, and transitions with + all derivatives vanishing there. It is used only to blend the two valid + quaternion charts; the geometric constraint itself is still enforced by the + charts. The interior denominator ``left + right`` is bounded below by + ``exp(-2)`` on the clamped domain, so the dead branches of the ``where`` + never divide by zero (gradient-safe). + """ + xp = array_api_compat.array_namespace(x) + x_clamped = xp.clip(x, min=0.0, max=1.0) + eps = float(xp.finfo(x_clamped.dtype).eps) + left = xp.exp(-1.0 / xp.clip(x_clamped, min=eps)) + right = xp.exp(-1.0 / xp.clip(1.0 - x_clamped, min=eps)) + interior = left / (left + right) + return xp.where( + x_clamped <= 0.0, + xp.zeros_like(x_clamped), + xp.where(x_clamped >= 1.0, xp.ones_like(x_clamped), interior), + ) + + +def quaternion_nlerp( + q0: Any, + q1: Any, + weight: Any, + *, + eps: float = 1e-7, +) -> Any: + """ + Normalized linear interpolation on the shortest quaternion arc. + + ``q`` and ``-q`` represent the same spatial rotation. Aligning signs before + the interpolation guarantees that the blended chart stays on the shorter + great-circle segment in ``S^3``. + """ + xp = array_api_compat.array_namespace(q0, q1, weight) + dot = xp.sum(q0 * q1, axis=-1, keepdims=True) + q1_aligned = xp.where(dot < 0.0, -q1, q1) + blended = (1.0 - weight[..., None]) * q0 + weight[..., None] * q1_aligned + return quaternion_normalize(blended, eps) + + +def _build_edge_quaternion_chart_pos_z(edge_unit: Any, eps: float) -> Any: + """Quaternion chart that is exact away from the ``-Z`` pole.""" + xp = array_api_compat.array_namespace(edge_unit) + x = edge_unit[..., 0] + y = edge_unit[..., 1] + z = edge_unit[..., 2] + q = xp.stack([1.0 + z, y, -x, xp.zeros_like(x)], axis=-1) + return quaternion_normalize(q, eps) + + +def _build_edge_quaternion_chart_neg_z(edge_unit: Any, eps: float) -> Any: + """Quaternion chart that is exact away from the ``+Z`` pole.""" + xp = array_api_compat.array_namespace(edge_unit) + x = edge_unit[..., 0] + y = edge_unit[..., 1] + z = edge_unit[..., 2] + q = xp.stack([-x, xp.zeros_like(x), 1.0 - z, y], axis=-1) + return quaternion_normalize(q, eps) + + +def build_edge_quaternion( + edge_vec: Any, + *, + edge_len: Any = None, + eps: float = 1e-7, +) -> Any: + """ + Build stable edge quaternions for the SeZM local ``+Z`` convention. + + The returned quaternion represents the global->local edge rotation, so + applying its rotation matrix to the unit edge direction yields exactly + ``(0, 0, 1)``. Two exact quaternion charts are used: + + - a ``+Z`` chart that is regular everywhere except the antipodal ``-Z`` pole; + - a ``-Z`` chart that is regular everywhere except the antipodal ``+Z`` pole. + + Both charts encode the same edge-aligned local frame. A smooth ``C^inf`` + blend in the overlap region removes the hard pole switch while keeping the + represented rotation on the correct quaternion branch. + + Parameters + ---------- + edge_vec + Edge vectors with shape ``(E, 3)``. + edge_len + Optional edge lengths with shape ``(E, 1)``. When omitted, lengths are + recomputed from ``edge_vec``. + eps + Numerical floor used in vector and quaternion normalization. + + Returns + ------- + Array + Unit quaternions with shape ``(E, 4)`` in ``(w, x, y, z)`` order. + """ + xp = array_api_compat.array_namespace(edge_vec) + if edge_len is None: + edge_len = safe_norm(edge_vec, eps) + else: + edge_len = xp.sqrt(edge_len * edge_len + eps * eps) + edge_unit = edge_vec / edge_len + q_pos = _build_edge_quaternion_chart_pos_z(edge_unit, eps) + q_neg = _build_edge_quaternion_chart_neg_z(edge_unit, eps) + blend = _smooth_step_cinf(0.5 * (edge_unit[..., 2] + 1.0)) + return quaternion_nlerp(q_neg, q_pos, blend, eps=eps) + + +def _factorial_table(n: int) -> np.ndarray: + """Return ``[0!, 1!, ..., n!]`` in fp64 (iterative, matching pt bit-exactly).""" + table = np.zeros(n + 1, dtype=np.float64) + table[0] = 1.0 + for i in range(1, n + 1): + table[i] = table[i - 1] * i + return table + + +def _binomial(n: int, k: int, factorial: np.ndarray) -> float: + """Evaluate ``C(n, k)`` from a precomputed factorial table.""" + if k < 0 or k > n: + return 0.0 + return float(factorial[n] / (factorial[k] * factorial[n - k])) + + +class _CaseTables: + """ + Plain numpy tables for one magnitude-ordered branch of the quaternion Wigner path. + + Mirrors pt ``CaseCoefficients`` (init-time constants only). + """ + + def __init__(self, n_primary: int, max_poly_len: int) -> None: + self.coeff = np.zeros(n_primary, dtype=np.float64) + self.horner = np.zeros((n_primary, max_poly_len), dtype=np.float64) + self.poly_len = np.zeros(n_primary, dtype=np.int64) + self.ra_exp = np.zeros(n_primary, dtype=np.float64) + self.rb_exp = np.zeros(n_primary, dtype=np.float64) + self.sign = np.zeros(n_primary, dtype=np.float64) + # filled by _finalize_case_tables + self.valid_mask: np.ndarray | None = None + self.horner_step_mask: np.ndarray | None = None + self.signed_coeff: np.ndarray | None = None + + +def _compute_case_coefficients( + case: _CaseTables, + idx: int, + ell: int, + mp: int, + m: int, + sqrt_factor: float, + factorial: np.ndarray, + *, + is_case1: bool, +) -> None: + """ + Fill one Horner branch for a fixed ``(ell, mp, m)`` entry. + + The closed-form quaternion Wigner formula is reorganized so that only the + ratio ``-(|Rb|/|Ra|)^2`` or ``-(|Ra|/|Rb|)^2`` enters the Horner chain. + """ + if is_case1: + rho_min = max(0, mp - m) + rho_max = min(ell + mp, ell - m) + else: + rho_min = max(0, -(mp + m)) + rho_max = min(ell - m, ell - mp) + + if rho_min > rho_max: + return + + if is_case1: + binom1 = _binomial(ell + mp, rho_min, factorial) + binom2 = _binomial(ell - mp, ell - m - rho_min, factorial) + else: + binom1 = _binomial(ell + mp, ell - m - rho_min, factorial) + binom2 = _binomial(ell - mp, rho_min, factorial) + case.coeff[idx] = sqrt_factor * binom1 * binom2 + + poly_len = rho_max - rho_min + 1 + case.poly_len[idx] = poly_len + for i, rho in enumerate(range(rho_max, rho_min, -1)): + if is_case1: + n1 = ell + mp - rho + 1 + n2 = ell - m - rho + 1 + d1 = rho + d2 = m - mp + rho + else: + n1 = ell - m - rho + 1 + n2 = ell - mp - rho + 1 + d1 = rho + d2 = mp + m + rho + if d1 != 0 and d2 != 0: + case.horner[idx, i] = (n1 * n2) / (d1 * d2) + + if is_case1: + case.ra_exp[idx] = 2 * ell + mp - m - 2 * rho_min + case.rb_exp[idx] = m - mp + 2 * rho_min + case.sign[idx] = (-1) ** rho_min + else: + case.ra_exp[idx] = mp + m + 2 * rho_min + case.rb_exp[idx] = 2 * ell - mp - m - 2 * rho_min + case.sign[idx] = ((-1) ** (ell - m)) * ((-1) ** rho_min) + + +def _finalize_case_tables(case: _CaseTables, max_poly_len: int) -> None: + """Attach runtime-ready masks and fused coefficients for one Horner branch.""" + step_count = np.clip(case.poly_len - 1, 0, None) + if max_poly_len > 1: + horner_step_mask = ( + np.arange(max_poly_len - 1, dtype=np.int64)[None, :] < step_count[:, None] + ) + else: + horner_step_mask = np.zeros((case.poly_len.shape[0], 0), dtype=np.bool_) + case.valid_mask = case.poly_len > 0 + case.horner_step_mask = horner_step_mask + case.signed_coeff = case.sign * case.coeff + + +class _PolyTables: + """ + Precomputed coefficient tables for the generic quaternion Wigner evaluator. + + Mirrors pt ``WignerPolynomialCoefficients`` (init-time numpy constants only). + Only one half of each real block is stored explicitly. The remaining + entries are reconstructed from the exact symmetry + ``D^l_{-m',-m} = (-1)^(m' - m) * conj(D^l_{m',m})``. + """ + + def __init__(self, lmin: int, lmax: int) -> None: + if lmin < 0: + raise ValueError("`lmin` must be non-negative") + if lmax < lmin: + raise ValueError("`lmax` must be >= `lmin`") + + factorial = _factorial_table(2 * lmax + 1) + n_total = sum((2 * ell + 1) ** 2 for ell in range(lmin, lmax + 1)) + n_primary = sum( + 1 + for ell in range(lmin, lmax + 1) + for mp in range(-ell, ell + 1) + for m in range(-ell, ell + 1) + if mp + m > 0 or (mp + m == 0 and mp >= 0) + ) + n_derived = n_total - n_primary + max_poly_len = lmax + 1 + size = (lmax + 1) ** 2 - lmin * lmin + + self.lmin = lmin + self.lmax = lmax + self.size = size + self.max_poly_len = max_poly_len + self.n_primary = n_primary + self.n_derived = n_derived + + self.primary_row = np.zeros(n_primary, dtype=np.int64) + self.primary_col = np.zeros(n_primary, dtype=np.int64) + self.mp_plus_m = np.zeros(n_primary, dtype=np.float64) + self.m_minus_mp = np.zeros(n_primary, dtype=np.float64) + self.diagonal_mask = np.zeros(n_primary, dtype=np.bool_) + self.anti_diagonal_mask = np.zeros(n_primary, dtype=np.bool_) + self.special_2m = np.zeros(n_primary, dtype=np.float64) + self.anti_diag_sign = np.zeros(n_primary, dtype=np.float64) + self.case1 = _CaseTables(n_primary, max_poly_len) + self.case2 = _CaseTables(n_primary, max_poly_len) + self.derived_row = np.zeros(n_derived, dtype=np.int64) + self.derived_col = np.zeros(n_derived, dtype=np.int64) + self.derived_primary_idx = np.zeros(n_derived, dtype=np.int64) + self.derived_sign = np.zeros(n_derived, dtype=np.float64) + + primary_map: dict[tuple[int, int], int] = {} + primary_idx = 0 + block_start = 0 + for ell in range(lmin, lmax + 1): + block_size = 2 * ell + 1 + for mp_local in range(block_size): + mp = mp_local - ell + for m_local in range(block_size): + m = m_local - ell + row = block_start + mp_local + col = block_start + m_local + is_primary = (mp + m > 0) or (mp + m == 0 and mp >= 0) + if not is_primary: + continue + + primary_map[(row, col)] = primary_idx + self.primary_row[primary_idx] = row + self.primary_col[primary_idx] = col + self.mp_plus_m[primary_idx] = mp + m + self.m_minus_mp[primary_idx] = m - mp + self.diagonal_mask[primary_idx] = mp == m + self.anti_diagonal_mask[primary_idx] = mp == -m + self.special_2m[primary_idx] = 2 * m + self.anti_diag_sign[primary_idx] = (-1) ** (ell - m) + + sqrt_factor = math.sqrt( + float(factorial[ell + m] * factorial[ell - m]) + / float(factorial[ell + mp] * factorial[ell - mp]) + ) + _compute_case_coefficients( + self.case1, + primary_idx, + ell, + mp, + m, + sqrt_factor, + factorial, + is_case1=True, + ) + _compute_case_coefficients( + self.case2, + primary_idx, + ell, + mp, + m, + sqrt_factor, + factorial, + is_case1=False, + ) + primary_idx += 1 + block_start += block_size + + derived_idx = 0 + block_start = 0 + for ell in range(lmin, lmax + 1): + block_size = 2 * ell + 1 + for mp_local in range(block_size): + mp = mp_local - ell + for m_local in range(block_size): + m = m_local - ell + row = block_start + mp_local + col = block_start + m_local + is_primary = (mp + m > 0) or (mp + m == 0 and mp >= 0) + if is_primary: + continue + + self.derived_row[derived_idx] = row + self.derived_col[derived_idx] = col + self.derived_primary_idx[derived_idx] = primary_map[ + (block_start + (-mp + ell), block_start + (-m + ell)) + ] + self.derived_sign[derived_idx] = (-1) ** (mp - m) + derived_idx += 1 + block_start += block_size + + _finalize_case_tables(self.case1, max_poly_len) + _finalize_case_tables(self.case2, max_poly_len) + + # Functional scatter replacement: gather index mapping each flat + # (row, col) of the packed (size, size) matrix to its source slot in + # ``concat([primary, derived, zero_slot])``. Entries outside the + # diagonal blocks point at the trailing zero slot. + flat_to_src = np.full(size * size, n_primary + n_derived, dtype=np.int64) + flat_to_src[self.primary_row * size + self.primary_col] = np.arange( + n_primary, dtype=np.int64 + ) + flat_to_src[self.derived_row * size + self.derived_col] = n_primary + np.arange( + n_derived, dtype=np.int64 + ) + self.flat_gather_idx = flat_to_src + + +def _build_complex_to_real_sh_block(ell: int) -> np.ndarray: + """ + Build the complex-to-real basis transform for one ``ell`` block. + + The packed real basis follows the SeZM convention ``m = -ell, ..., +ell`` + inside each block. This unitary transform defines the real tesseral basis + used by the packed ``D_full`` layout. + """ + size = 2 * ell + 1 + inv_sqrt2 = 1.0 / math.sqrt(2.0) + U = np.zeros((size, size), dtype=np.complex128) + for m in range(-ell, ell + 1): + row = m + ell + if m == 0: + U[row, ell] = 1.0 + elif m > 0: + U[row, m + ell] = inv_sqrt2 + U[row, -m + ell] = ((-1) ** m) * inv_sqrt2 + else: + U[row, -m + ell] = -1j * inv_sqrt2 + U[row, m + ell] = ((-1) ** m) * 1j * inv_sqrt2 + return U + + +def _assemble_block_diagonal_real_basis( + lmin: int, lmax: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Assemble per-``ell`` real-basis blocks into one block-diagonal transform.""" + size = sum(2 * ell + 1 for ell in range(lmin, lmax + 1)) + U_re_full = np.zeros((size, size), dtype=np.float64) + U_im_full = np.zeros((size, size), dtype=np.float64) + offset = 0 + for ell in range(lmin, lmax + 1): + U = _build_complex_to_real_sh_block(ell) + block_size = 2 * ell + 1 + block_end = offset + block_size + U_re_full[offset:block_end, offset:block_end] = U.real + U_im_full[offset:block_end, offset:block_end] = U.imag + offset = block_end + return ( + U_re_full, + U_im_full, + np.ascontiguousarray(U_re_full.T), + np.ascontiguousarray(U_im_full.T), + ) + + +def _vectorized_horner( + xp: Any, + ratio: Any, + horner_coeffs: Any, + horner_step_mask: Any, +) -> Any: + """Evaluate many varying-length Horner chains in one batched loop.""" + n_batch = ratio.shape[0] + n_elements = horner_coeffs.shape[0] + result = xp.ones( + (n_batch, n_elements), + dtype=ratio.dtype, + device=array_api_compat.device(ratio), + ) + if horner_step_mask.shape[1] == 0: + return result + ratio = ratio[:, None] + for i in range(horner_step_mask.shape[1]): + new_result = 1.0 + result * (ratio * horner_coeffs[None, :, i]) + result = xp.where(horner_step_mask[None, :, i], new_result, result) + return result + + +class WignerDCalculator(NativeOP): + """ + Quaternion-driven Wigner-D blocks for the SeZM packed real spherical basis. + + Input quaternions represent the global->local edge rotation that sends the + edge direction to local ``+Z``. The returned block-diagonal matrix keeps + the packed SeZM real spherical-harmonics layout, so downstream code + consumes ``D_full`` and ``Dt_full`` directly. + + Runtime structure: + + - ``l=0``: scalar identity block; + - ``l=1``: direct quaternion -> Cartesian rotation -> real ``l=1`` block; + - ``l>=2``: generic quaternion polynomial path with precomputed coefficient + tables (the pt reference path that the pt monomial kernels are fitted to; + see the module docstring). + + Parameters + ---------- + lmax : int + Maximum spherical-harmonics degree. + eps : float + Numerical floor used in quaternion normalization. + precision : str + Working floating-point precision of the returned blocks. The internal + polynomial algebra is evaluated in fp64 (as in pt) before the result is + cast back. + """ + + def __init__( + self, + lmax: int, + *, + eps: float = 1e-7, + precision: str = DEFAULT_PRECISION, + ) -> None: + self.lmax = int(lmax) + if self.lmax < 0: + raise ValueError("`lmax` must be non-negative") + self.precision = precision + self.eps = float(eps) + self.dim_full = (self.lmax + 1) ** 2 + + # l=1 block constants: permutation [1, 2, 0] is applied structurally in + # _compute_l1_block; the sign pattern is a plain numpy constant. + l1_sign = np.array([-1.0, -1.0, 1.0], dtype=np.float64) + self.l1_sign_outer = np.outer(l1_sign, l1_sign) + + if self.lmax >= 2: + self.poly_tables = _PolyTables(lmin=2, lmax=self.lmax) + ( + self.poly_u_re, + self.poly_u_im, + self.poly_u_re_t, + self.poly_u_im_t, + ) = _assemble_block_diagonal_real_basis(2, self.lmax) + + # Functional block-diagonal assembly: gather index mapping each flat + # (row, col) of D_full to its source slot in the concatenated value + # vector [l0 ones (1), l1 block (9), packed l>=2 block (size^2), + # trailing zero slot]. No __setitem__ on traced values is needed. + n_l1 = 9 if self.lmax >= 1 else 0 + n_packed = (self.dim_full - 4) ** 2 if self.lmax >= 2 else 0 + zero_slot = 1 + n_l1 + n_packed + full_idx = np.full(self.dim_full * self.dim_full, zero_slot, dtype=np.int64) + full_idx[0] = 0 # D_full[:, 0, 0] = 1 + if self.lmax >= 1: + for i in range(3): + for j in range(3): + full_idx[(1 + i) * self.dim_full + (1 + j)] = 1 + 3 * i + j + if self.lmax >= 2: + packed_size = self.dim_full - 4 + for i in range(packed_size): + for j in range(packed_size): + full_idx[(4 + i) * self.dim_full + (4 + j)] = ( + 1 + n_l1 + packed_size * i + j + ) + self.full_gather_idx = full_idx + + def call(self, edge_quaternion: Any) -> tuple[Any, Any]: + """ + Build packed block-diagonal Wigner-D matrices from edge quaternions. + + Parameters + ---------- + edge_quaternion : Array + Unit quaternions with shape ``(E, 4)`` representing the + global->local edge rotation. + + Returns + ------- + tuple[Array, Array] + ``(D_full, Dt_full)`` with shape ``(E, (lmax+1)^2, (lmax+1)^2)``. + """ + xp = array_api_compat.array_namespace(edge_quaternion) + dtype = get_xp_precision(xp, self.precision) + device = array_api_compat.device(edge_quaternion) + q = quaternion_normalize( + xp.astype(edge_quaternion, dtype), + eps=self.eps, + ) + n_edge = q.shape[0] + + segments = [xp.ones((n_edge, 1), dtype=dtype, device=device)] + if self.lmax >= 1: + segments.append( + xp.reshape(self._compute_l1_block(q, xp, dtype, device), (n_edge, 9)) + ) + if self.lmax >= 2: + packed = self._compute_packed_blocks(q, xp, dtype, device) + packed_size = self.dim_full - 4 + segments.append(xp.reshape(packed, (n_edge, packed_size * packed_size))) + segments.append(xp.zeros((n_edge, 1), dtype=dtype, device=device)) + values = xp.concat(segments, axis=1) + idx = xp.asarray(self.full_gather_idx, device=device) + D_full = xp.reshape( + xp.take(values, idx, axis=1), + (n_edge, self.dim_full, self.dim_full), + ) + Dt_full = xp.matrix_transpose(D_full) + return D_full, Dt_full + + def forward_zonal(self, edge_quaternion: Any, lmin: int = 1) -> Any: + """ + Build local ``m=0`` to global coupling for GIE. + + The returned layout matches the packed node rows for degrees + ``lmin..lmax``: each degree contributes ``2l+1`` values in packed + ``m=-l..l`` order. These values are equivalent to gathering + ``Dt_full[:, row(l, m), col(l, 0)]`` from :meth:`call` over the same + degree range. + + Parameters + ---------- + edge_quaternion : Array + Unit quaternions with shape ``(E, 4)`` representing the + global->local edge rotation. + lmin : int + First degree to return. + + Returns + ------- + Array + Zonal coupling with shape ``(E, (lmax + 1) ** 2 - lmin ** 2)``. + """ + lmin = int(lmin) + if lmin < 1: + raise ValueError("`lmin` must be >= 1") + xp = array_api_compat.array_namespace(edge_quaternion) + dtype = get_xp_precision(xp, self.precision) + device = array_api_compat.device(edge_quaternion) + n_edge = edge_quaternion.shape[0] + if self.lmax < lmin: + return xp.zeros((n_edge, 0), dtype=dtype, device=device) + q = quaternion_normalize( + xp.astype(edge_quaternion, dtype), + eps=self.eps, + ) + + zonal_blocks = [] + if lmin <= 1 <= self.lmax: + zonal_blocks.append(self._compute_l1_block(q, xp, dtype, device)[:, 1, :]) + if self.lmax >= 2: + packed = self._compute_packed_blocks(q, xp, dtype, device) + offset = 0 + for degree in range(2, self.lmax + 1): + block_size = 2 * degree + 1 + block_end = offset + block_size + if degree >= lmin: + zonal_blocks.append(packed[:, offset + degree, offset:block_end]) + offset = block_end + return xp.concat(zonal_blocks, axis=1) + + def _compute_l1_block(self, q: Any, xp: Any, dtype: Any, device: Any) -> Any: + """Compute the vector block directly from the Cartesian rotation matrix.""" + rot = quaternion_to_rotation_matrix(q) + # row/column permutation [1, 2, 0], applied structurally (no gather) + rot = xp.stack([rot[..., 1, :], rot[..., 2, :], rot[..., 0, :]], axis=-2) + rot = xp.stack([rot[..., 1], rot[..., 2], rot[..., 0]], axis=-1) + sign = xp.asarray(self.l1_sign_outer, dtype=dtype, device=device) + return rot * sign + + def _compute_packed_blocks(self, q: Any, xp: Any, dtype: Any, device: Any) -> Any: + """Evaluate the packed real Wigner blocks for ``l = 2..lmax``.""" + # Cayley-Klein pair: Ra = w - i z, Rb = y - i x (SeZM convention) + ra_re = q[..., 0] + ra_im = -q[..., 3] + rb_re = q[..., 2] + rb_im = -q[..., 1] + D_re, D_im = self._wigner_d_matrix_realpair( + ra_re, ra_im, rb_re, rb_im, xp, dtype, device + ) + u_re = xp.asarray(self.poly_u_re, dtype=dtype, device=device) + u_im = xp.asarray(self.poly_u_im, dtype=dtype, device=device) + u_re_t = xp.asarray(self.poly_u_re_t, dtype=dtype, device=device) + u_im_t = xp.asarray(self.poly_u_im_t, dtype=dtype, device=device) + temp_re = xp.matmul(D_re, u_re_t) + xp.matmul(D_im, u_im_t) + temp_im = xp.matmul(D_im, u_re_t) - xp.matmul(D_re, u_im_t) + return xp.matmul(u_re, temp_re) - xp.matmul(u_im, temp_im) + + def _wigner_d_matrix_realpair( + self, + ra_re: Any, + ra_im: Any, + rb_re: Any, + rb_im: Any, + xp: Any, + out_dtype: Any, + device: Any, + ) -> tuple[Any, Any]: + """ + Evaluate the complex Wigner blocks in real/imaginary form. + + The runtime path uses only real arithmetic. The complex phase is + represented by two real tensors, while the polynomial and magnitude + algebra is evaluated in fp64 before the result is cast back to the + requested output dtype. All denominators are eps-floored before any + division (gradient-safe masked-denominator idiom, as in pt). + """ + coeffs = self.poly_tables + n_batch = ra_re.shape[0] + f64 = xp.float64 + ra_re = xp.astype(ra_re, f64) + ra_im = xp.astype(ra_im, f64) + rb_re = xp.astype(rb_re, f64) + rb_im = xp.astype(rb_im, f64) + + def cv(arr: np.ndarray) -> Any: # constant table -> xp on input device + return xp.asarray(arr, device=device) + + eps = float(np.finfo(np.float64).eps) + eps_sq = eps * eps + ra_sq = ra_re * ra_re + ra_im * ra_im + rb_sq = rb_re * rb_re + rb_im * rb_im + ra_small = ra_sq <= eps_sq + rb_small = rb_sq <= eps_sq + ra = xp.sqrt(xp.clip(ra_sq, min=eps_sq)) + rb = xp.sqrt(xp.clip(rb_sq, min=eps_sq)) + general_mask = ~ra_small & ~rb_small + use_case1 = (ra >= rb) & general_mask + use_case2 = (ra < rb) & general_mask + + safe_ra_re = xp.where(ra_small, xp.ones_like(ra_re), ra_re) + safe_ra_im = xp.where(ra_small, xp.zeros_like(ra_im), ra_im) + safe_rb_re = xp.where(rb_small, xp.ones_like(rb_re), rb_re) + safe_rb_im = xp.where(rb_small, xp.zeros_like(rb_im), rb_im) + phia = xp.atan2(safe_ra_im, safe_ra_re) + phib = xp.atan2(safe_rb_im, safe_rb_re) + + phase = ( + phia[:, None] * cv(coeffs.mp_plus_m)[None, :] + + phib[:, None] * cv(coeffs.m_minus_mp)[None, :] + ) + exp_phase_re = xp.cos(phase) + exp_phase_im = xp.sin(phase) + + safe_ra = xp.clip(ra, min=eps) + safe_rb = xp.clip(rb, min=eps) + log_ra = xp.log(safe_ra) + log_rb = xp.log(safe_rb) + + result_re = xp.zeros((n_batch, coeffs.n_primary), dtype=f64, device=device) + result_im = xp.zeros((n_batch, coeffs.n_primary), dtype=f64, device=device) + + special_2m = cv(coeffs.special_2m) + anti_rows = ra_small + anti_log_rb = xp.where(anti_rows, log_rb, xp.zeros_like(log_rb)) + anti_phib = xp.where(anti_rows, phib, xp.zeros_like(phib)) + rb_power_mag = xp.exp(anti_log_rb[:, None] * special_2m[None, :]) + rb_power_phase = anti_phib[:, None] * special_2m[None, :] + anti_diag_sign = cv(coeffs.anti_diag_sign) + anti_re = anti_diag_sign[None, :] * rb_power_mag * xp.cos(rb_power_phase) + anti_im = anti_diag_sign[None, :] * rb_power_mag * xp.sin(rb_power_phase) + anti_mask = ra_small[:, None] & cv(coeffs.anti_diagonal_mask)[None, :] + result_re = xp.where(anti_mask, anti_re, result_re) + result_im = xp.where(anti_mask, anti_im, result_im) + + diag_rows = rb_small & ~ra_small + diag_log_ra = xp.where(diag_rows, log_ra, xp.zeros_like(log_ra)) + diag_phia = xp.where(diag_rows, phia, xp.zeros_like(phia)) + ra_power_mag = xp.exp(diag_log_ra[:, None] * special_2m[None, :]) + ra_power_phase = diag_phia[:, None] * special_2m[None, :] + diag_re = ra_power_mag * xp.cos(ra_power_phase) + diag_im = ra_power_mag * xp.sin(ra_power_phase) + diag_mask = diag_rows[:, None] & cv(coeffs.diagonal_mask)[None, :] + result_re = xp.where(diag_mask, diag_re, result_re) + result_im = xp.where(diag_mask, diag_im, result_im) + + for case, case_rows, ratio in ( + ( + coeffs.case1, + use_case1, + -(rb * rb) / (safe_ra * safe_ra), + ), + ( + coeffs.case2, + use_case2, + -(ra * ra) / (safe_rb * safe_rb), + ), + ): + magnitude = self._compute_case_magnitude( + xp, + xp.where(case_rows, log_ra, xp.zeros_like(log_ra)), + xp.where(case_rows, log_rb, xp.zeros_like(log_rb)), + xp.where(case_rows, ratio, xp.zeros_like(ratio)), + case, + device, + ) + val_re = magnitude * exp_phase_re + val_im = magnitude * exp_phase_im + mask = case_rows[:, None] & cv(case.valid_mask)[None, :] + result_re = xp.where(mask, val_re, result_re) + result_im = xp.where(mask, val_im, result_im) + + # Functional scatter into the dense packed matrix: derive the + # symmetry-completed entries by gather, then place primary + derived + # values with one precomputed take index (zero slot for off-block). + derived_idx = cv(coeffs.derived_primary_idx) + derived_sign = cv(coeffs.derived_sign) + primary_re = xp.take(result_re, derived_idx, axis=1) + primary_im = xp.take(result_im, derived_idx, axis=1) + derived_re = derived_sign[None, :] * primary_re + derived_im = -derived_sign[None, :] * primary_im + zero_col = xp.zeros((n_batch, 1), dtype=f64, device=device) + flat_idx = cv(coeffs.flat_gather_idx) + D_re = xp.reshape( + xp.take( + xp.concat([result_re, derived_re, zero_col], axis=1), flat_idx, axis=1 + ), + (n_batch, coeffs.size, coeffs.size), + ) + D_im = xp.reshape( + xp.take( + xp.concat([result_im, derived_im, zero_col], axis=1), flat_idx, axis=1 + ), + (n_batch, coeffs.size, coeffs.size), + ) + return xp.astype(D_re, out_dtype), xp.astype(D_im, out_dtype) + + @staticmethod + def _compute_case_magnitude( + xp: Any, + log_ra: Any, + log_rb: Any, + ratio: Any, + case: _CaseTables, + device: Any, + ) -> Any: + """Compute the real magnitude factor for one stable Horner branch.""" + horner_sum = _vectorized_horner( + xp, + ratio, + xp.asarray(case.horner, device=device), + xp.asarray(case.horner_step_mask, device=device), + ) + ra_powers = xp.exp( + log_ra[:, None] * xp.asarray(case.ra_exp, device=device)[None, :] + ) + rb_powers = xp.exp( + log_rb[:, None] * xp.asarray(case.rb_exp, device=device)[None, :] + ) + signed_coeff = xp.asarray(case.signed_coeff, device=device) + magnitude = signed_coeff[None, :] * ra_powers * rb_powers + return magnitude * horner_sum + + def serialize(self) -> dict[str, Any]: + """Serialize WignerDCalculator (lmax and precision are stored by parent).""" + return { + "@class": "WignerDCalculator", + "@version": 1, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> WignerDCalculator: + """Deserialize WignerDCalculator - parent handles lmax/precision reconstruction.""" + data = data.copy() + data_cls = data.pop("@class") + if data_cls != "WignerDCalculator": + raise ValueError(f"Invalid class for WignerDCalculator: {data_cls}") + version = int(data.pop("@version")) + check_version_compatibility(version, 1, 1) + raise NotImplementedError( + "WignerDCalculator.deserialize should be called by parent with lmax/precision" + ) diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 5bdfff2571..37b0a1b731 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -5,6 +5,9 @@ from .dos_fitting import ( DOSFittingNet, ) +from .dpa4_ener import ( + SeZMEnergyFittingNet, +) from .ener_fitting import ( EnergyFittingNet, ) @@ -28,5 +31,6 @@ "InvarFitting", "PolarFitting", "PropertyFittingNet", + "SeZMEnergyFittingNet", "make_base_fitting", ] diff --git a/deepmd/dpmodel/fitting/dpa4_ener.py b/deepmd/dpmodel/fitting/dpa4_ener.py new file mode 100644 index 0000000000..b8f9507014 --- /dev/null +++ b/deepmd/dpmodel/fitting/dpa4_ener.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""SeZM (DPA4) GLU energy fitting network, dpmodel implementation. + +Mirrors ``deepmd.pt.model.task.sezm_ener`` with the array-API ``call`` +convention and the pt-state_dict-key serialization contract. +""" + +import math +from typing import ( + Any, + ClassVar, +) + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.utils.network import ( + NativeLayer, + get_activation_fn, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .invar_fitting import ( + InvarFitting, +) + + +class GLUFittingNet(NativeOP): + """ + GLU-based fitting network for SeZM. + + Parameters + ---------- + in_dim + Input dimension. + out_dim + Output dimension. + neuron + Hidden layer sizes. Empty list means direct linear projection. + activation_function + Activation function used for GLU gating. + resnet_dt + Reserved for compatibility; not used in GLU layers. + precision + Numerical precision. + bias_out + Whether the output layer uses bias. + seed + Random seed. + trainable + Whether parameters are trainable. + descriptor_dim + Descriptor feature width. Kept for serialization compatibility + with the case-FiLM path (not implemented here). + dim_case_embd + Case one-hot width. + case_film_embd + Whether to use case FiLM instead of input concatenation. + Not implemented in the dpmodel backend. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + neuron: list[int] | None = None, + activation_function: str = "silu", + resnet_dt: bool = False, + precision: str = DEFAULT_PRECISION, + bias_out: bool = False, + seed: int | list[int] | None = None, + trainable: bool | list[bool] = True, + descriptor_dim: int | None = None, + dim_case_embd: int = 0, + case_film_embd: bool = False, + ) -> None: + if case_film_embd and int(dim_case_embd) > 0: + raise NotImplementedError( + "case_film_embd is not implemented in the dpmodel backend" + ) + if neuron is None: + neuron = [] + if isinstance(trainable, list): + trainable = all(trainable) + self.in_dim = int(in_dim) + self.out_dim = int(out_dim) + self.neuron = [int(nn_dim) for nn_dim in neuron] + self.activation_function = activation_function + self.resnet_dt = bool(resnet_dt) + self.precision = precision + self.bias_out = bool(bias_out) + self.descriptor_dim = ( + self.in_dim if descriptor_dim is None else int(descriptor_dim) + ) + self.dim_case_embd = int(dim_case_embd) + self.case_film_embd = bool(case_film_embd and self.dim_case_embd > 0) + + # === Step 1. Build GLU hidden layers === + # Each hidden layer is a linear map to 2*hidden_dim, split into + # value and gate halves: out = val * act(gate). + hidden_layers = [] + dim_in = self.in_dim + for layer_idx, hidden_dim in enumerate(self.neuron): + hidden_layers.append( + NativeLayer( + dim_in, + 2 * hidden_dim, + bias=True, + use_timestep=False, + activation_function=None, + resnet=False, + precision=self.precision, + seed=child_seed(seed, layer_idx), + trainable=trainable, + ) + ) + dim_in = hidden_dim + self.hidden_layers = hidden_layers + + # === Step 2. Build output projection === + self.output_layer = NativeLayer( + dim_in, + self.out_dim, + bias=self.bias_out, + use_timestep=False, + activation_function=None, + resnet=False, + precision=self.precision, + seed=child_seed(seed, len(self.neuron) + int(self.case_film_embd)), + trainable=trainable, + ) + + def call_until_last(self, xx: Array) -> Array: + """Return activations before the output projection.""" + act = get_activation_fn(self.activation_function) + for hidden_dim, layer in zip(self.neuron, self.hidden_layers, strict=True): + yy = layer(xx) + val, gate = yy[..., :hidden_dim], yy[..., hidden_dim:] + xx = val * act(gate) + return xx + + def call(self, xx: Array) -> Array: + """Forward pass for the GLU fitting net.""" + return self.output_layer(self.call_until_last(xx)) + + def serialize(self) -> dict[str, Any]: + """Serialize the network to a dict (pt state_dict key contract).""" + variables: dict[str, Any] = {} + for layer_idx, layer in enumerate(self.hidden_layers): + variables[f"hidden_layers.{layer_idx}.linear.matrix"] = layer.w + variables[f"hidden_layers.{layer_idx}.linear.bias"] = layer.b + variables["output_layer.matrix"] = self.output_layer.w + if self.bias_out: + variables["output_layer.bias"] = self.output_layer.b + return { + "@class": "GLUFittingNet", + "@version": 1, + "in_dim": self.in_dim, + "out_dim": self.out_dim, + "neuron": self.neuron.copy(), + "activation_function": self.activation_function, + "resnet_dt": self.resnet_dt, + "precision": self.precision, + "bias_out": self.bias_out, + "descriptor_dim": self.descriptor_dim, + "dim_case_embd": self.dim_case_embd, + "case_film_embd": self.case_film_embd, + "@variables": variables, + } + + @classmethod + def deserialize(cls, data: dict) -> "GLUFittingNet": + """Deserialize the network from a dict.""" + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) + variables = data.pop("@variables", {}) + obj = cls(**data) + for layer_idx, layer in enumerate(obj.hidden_layers): + layer["matrix"] = variables[f"hidden_layers.{layer_idx}.linear.matrix"] + layer["bias"] = variables[f"hidden_layers.{layer_idx}.linear.bias"] + obj.output_layer["matrix"] = variables["output_layer.matrix"] + if obj.bias_out: + obj.output_layer["bias"] = variables["output_layer.bias"] + return obj + + +class SeZMNetworkCollection: + """ + Network collection for SeZM fitting networks. + + Parameters + ---------- + ndim + The number of type dimensions. + ntypes + Number of atom types. + network_type + The network type name. Only "sezm_fitting_network" is supported. + networks + The networks to initialize with. + """ + + NETWORK_TYPE_MAP: ClassVar[dict[str, type]] = { + "sezm_fitting_network": GLUFittingNet, + } + + def __init__( + self, + ndim: int, + ntypes: int, + network_type: str = "sezm_fitting_network", + networks: list[Any] | None = None, + ) -> None: + self.ndim = int(ndim) + self.ntypes = int(ntypes) + if network_type not in self.NETWORK_TYPE_MAP: + raise ValueError(f"Unknown network_type: {network_type}") + self.network_type = self.NETWORK_TYPE_MAP[network_type] + if networks is None: + networks = [] + + total = self.ntypes**self.ndim + self._networks: list[GLUFittingNet | None] = [None for _ in range(total)] + for idx, network in enumerate(networks): + self[idx] = network + if any(net is None for net in self._networks): + raise RuntimeError("SeZMNetworkCollection is incomplete.") + self.networks = self._networks + + def _convert_key(self, key: int | tuple | str) -> int: + if isinstance(key, int): + idx = key + else: + if isinstance(key, tuple): + pass + elif isinstance(key, str): + key = tuple([int(tt) for tt in key.split("_")[1:]]) + else: + raise TypeError(key) + if len(key) != self.ndim: + raise KeyError( + f"key {key} has length {len(key)}, expected ndim {self.ndim}" + ) + if any(not (0 <= int(tt) < self.ntypes) for tt in key): + raise KeyError( + f"key {key} contains type indices outside [0, {self.ntypes})" + ) + idx = sum([tt * self.ntypes**ii for ii, tt in enumerate(key)]) + if not (0 <= idx < self.ntypes**self.ndim): + raise KeyError( + f"key {key} maps to index {idx}, outside [0, {self.ntypes**self.ndim})" + ) + return idx + + def __getitem__(self, key: int | tuple | str) -> GLUFittingNet: + idx = self._convert_key(key) + nn = self._networks[idx] + if nn is None: + raise KeyError(f"network for key {key} is not set") + return nn + + def __setitem__(self, key: int | tuple | str, value: Any) -> None: + if isinstance(value, self.network_type): + network = value + elif isinstance(value, dict): + network = self.network_type.deserialize(value) + else: + raise TypeError(value) + idx = self._convert_key(key) + self._networks[idx] = network + + def serialize(self) -> dict[str, Any]: + """Serialize the networks to a dict.""" + network_type_map_inv = {v: k for k, v in self.NETWORK_TYPE_MAP.items()} + return { + "@class": "NetworkCollection", + "@version": 1, + "ndim": self.ndim, + "ntypes": self.ntypes, + "network_type": network_type_map_inv[self.network_type], + "networks": [ + nn.serialize() if nn is not None else None for nn in self._networks + ], + } + + @classmethod + def deserialize(cls, data: dict) -> "SeZMNetworkCollection": + """Deserialize the networks from a dict.""" + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) + return cls(**data) + + +def _resolve_auto_neuron( + neuron: list[int] | None, + *, + dim_descrpt: int, + numb_fparam: int, + numb_aparam: int, + dim_case_embd: int, + case_film_embd: bool, + use_aparam_as_mask: bool, +) -> list[int]: + """Resolve SeZM fitting hidden widths, using 0 as the auto-width marker.""" + resolved_neuron = [0] if neuron is None else [int(width) for width in neuron] + if any(width < 0 for width in resolved_neuron): + raise ValueError("`fitting_net.neuron` entries must be >= 0") + if 0 not in resolved_neuron: + return resolved_neuron + case_dim = 0 if case_film_embd else int(dim_case_embd) + dim_in = ( + int(dim_descrpt) + + int(numb_fparam) + + (0 if use_aparam_as_mask else int(numb_aparam)) + + case_dim + ) + resolved_width = int(32 * math.ceil((8.0 * float(dim_in) / 3.0) / 32.0)) + return [resolved_width if width == 0 else width for width in resolved_neuron] + + +@InvarFitting.register("dpa4_ener") +@InvarFitting.register("sezm_ener") +class SeZMEnergyFittingNet(InvarFitting): + """ + SeZM energy fitting with GLU hidden layers. + + This uses the same configuration keys as the standard energy fitting + but replaces hidden MLP layers with GLU blocks. + """ + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + neuron: list[int] | None = None, + bias_atom_e: Array | None = None, + resnet_dt: bool = False, + numb_fparam: int = 0, + numb_aparam: int = 0, + dim_case_embd: int = 0, + case_film_embd: bool = False, + activation_function: str = "silu", + bias_out: bool = False, + precision: str = "float32", + mixed_types: bool = True, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, + **kwargs: Any, + ) -> None: + if int(dim_case_embd) > 0: + raise NotImplementedError( + "dim_case_embd > 0 is not implemented in the dpmodel backend" + ) + if case_film_embd: + raise NotImplementedError( + "case_film_embd is not implemented in the dpmodel backend" + ) + neuron = _resolve_auto_neuron( + neuron, + dim_descrpt=dim_descrpt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + dim_case_embd=dim_case_embd, + case_film_embd=case_film_embd, + use_aparam_as_mask=bool(kwargs.get("use_aparam_as_mask", False)), + ) + super().__init__( + "energy", + ntypes, + dim_descrpt, + 1, + neuron=neuron, + bias_atom=bias_atom_e, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + dim_case_embd=dim_case_embd, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + seed=seed, + type_map=type_map, + default_fparam=default_fparam, + **kwargs, + ) + self.seed = seed + self.bias_out = bool(bias_out) + self.case_film_embd = bool(case_film_embd and self.dim_case_embd > 0) + self._build_glu_fitting_layers() + + def _build_glu_fitting_layers(self) -> None: + # === Step 1. Derive input/output dimensions === + case_dim = 0 if self.case_film_embd else self.dim_case_embd + in_dim = ( + self.dim_descrpt + + self.numb_fparam + + (0 if self.use_aparam_as_mask else self.numb_aparam) + + case_dim + ) + net_dim_out = self._net_out_dim() + n_networks = self.ntypes if not self.mixed_types else 1 + + # === Step 2. Build GLU fitting networks === + self.nets = SeZMNetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="sezm_fitting_network", + networks=[ + GLUFittingNet( + in_dim, + net_dim_out, + self.neuron, + activation_function=self.activation_function, + resnet_dt=self.resnet_dt, + precision=self.precision, + bias_out=self.bias_out, + seed=child_seed(self.seed, idx), + trainable=self.trainable, + descriptor_dim=self.dim_descrpt, + dim_case_embd=self.dim_case_embd, + case_film_embd=self.case_film_embd, + ) + for idx in range(n_networks) + ], + ) + + @classmethod + def deserialize(cls, data: dict) -> "SeZMEnergyFittingNet": + data = data.copy() + variables = data.pop("@variables") + nets = data.pop("nets") + check_version_compatibility(data.pop("@version", 1), 4, 1) + data.pop("@class", None) + data.pop("type", None) + data.pop("var_name") + data.pop("dim_out") + obj = cls(**data) + for kk in variables.keys(): + obj[kk] = variables[kk] + obj.nets = SeZMNetworkCollection.deserialize(nets) + return obj + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + **super().serialize(), + "type": "sezm_ener", + "case_film_embd": self.case_film_embd, + } diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py index c6e41c08e5..80153129cc 100644 --- a/deepmd/dpmodel/utils/exclude_mask.py +++ b/deepmd/dpmodel/utils/exclude_mask.py @@ -55,7 +55,11 @@ def build_type_exclude_mask( xp = array_api_compat.array_namespace(atype) nf, natom = atype.shape return xp.reshape( - xp.take(self.type_mask[...], xp.reshape(atype, (-1,)), axis=0), + xp.take( + xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)), + xp.reshape(atype, (-1,)), + axis=0, + ), (nf, natom), ) @@ -147,7 +151,10 @@ def build_type_exclude_mask( # (nf * nloc * nnei,) type_ij_flat = xp.reshape(type_ij, (-1,)) mask = xp.reshape( - xp.take(self.type_mask[...], type_ij_flat), + xp.take( + xp.asarray(self.type_mask[...], device=array_api_compat.device(nlist)), + type_ij_flat, + ), (nf, nloc, nnei), ) return mask diff --git a/deepmd/dpmodel/utils/lebedev.py b/deepmd/dpmodel/utils/lebedev.py new file mode 100644 index 0000000000..7d55121e47 --- /dev/null +++ b/deepmd/dpmodel/utils/lebedev.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Lebedev quadrature data loader for S2 projections.""" + +from __future__ import ( + annotations, +) + +from pathlib import ( + Path, +) + +import numpy as np + +# See: https://people.sc.fsu.edu/~jburkardt/datasets/sphere_lebedev_rule/sphere_lebedev_rule.html +LEBEDEV_RULES_FILE = Path(__file__).with_name("lebedev_rules.npz") +LEBEDEV_PRECISION_TO_NPOINTS = { + 3: 6, + 5: 14, + 7: 26, + 9: 38, + 11: 50, + 13: 74, + 15: 86, + 17: 110, + 19: 146, + 21: 170, + 23: 194, + 25: 230, + 27: 266, + 29: 302, + 31: 350, + 35: 434, + 41: 590, + 47: 770, + 53: 974, + 59: 1202, + 65: 1454, + 71: 1730, + 77: 2030, + 83: 2354, + 89: 2702, + 95: 3074, + 101: 3470, + 107: 3890, + 113: 4334, + 119: 4802, + 125: 5294, + 131: 5810, +} + + +def load_lebedev_rule(precision: int) -> tuple[np.ndarray, np.ndarray]: + """ + Load one Lebedev rule from the packaged compressed data file. + + Parameters + ---------- + precision + Algebraic precision of the requested Lebedev rule. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + Cartesian unit points with shape ``(A, 3)`` and normalized weights with + shape ``(A,)``. The weights sum to one, so the sphere integral is + ``4*pi*sum(weights*f(points))``. + """ + if not isinstance(precision, (int, np.integer)) or isinstance(precision, bool): + raise TypeError( + f"`precision` must be an integer, got {type(precision).__name__}" + ) + rule_key = f"{int(precision):03d}" + if not LEBEDEV_RULES_FILE.exists(): + raise FileNotFoundError( + f"Lebedev quadrature data file is missing: {LEBEDEV_RULES_FILE}" + ) + with np.load(LEBEDEV_RULES_FILE) as rules: + point_key = f"points_{rule_key}" + weight_key = f"weights_{rule_key}" + if point_key not in rules or weight_key not in rules: + raise ValueError( + f"Lebedev rule with precision {precision} is not packaged; " + f"available precisions: {sorted(LEBEDEV_PRECISION_TO_NPOINTS)}" + ) + points = rules[point_key] + weights = rules[weight_key] + return points, weights diff --git a/deepmd/pt/model/descriptor/sezm_nn/lebedev_rules.npz b/deepmd/dpmodel/utils/lebedev_rules.npz similarity index 100% rename from deepmd/pt/model/descriptor/sezm_nn/lebedev_rules.npz rename to deepmd/dpmodel/utils/lebedev_rules.npz diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index e93e46782f..c7a12ba3e2 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -282,10 +282,12 @@ def call(self, x): # noqa: ANN001, ANN201 raise ValueError("w, b, and activation_function must be set") xp = array_api_compat.array_namespace(x) fn = get_activation_fn(self.activation_function) + device = array_api_compat.device(x) + w = xp.asarray(self.w[...], device=device) y = ( - xp.matmul(x, self.w[...]) + self.b[...] + xp.matmul(x, w) + xp.asarray(self.b[...], device=device) if self.b is not None - else xp.matmul(x, self.w[...]) + else xp.matmul(x, w) ) if y.dtype != x.dtype: # workaround for bfloat16 @@ -293,7 +295,7 @@ def call(self, x): # noqa: ANN001, ANN201 y = xp.astype(y, x.dtype) y = fn(y) if self.idt is not None: - y = y * self.idt + y = y * xp.asarray(self.idt, device=array_api_compat.device(x)) if self.resnet and self.w.shape[1] == self.w.shape[0]: y = y + x elif self.resnet and self.w.shape[1] == 2 * self.w.shape[0]: diff --git a/deepmd/dpmodel/utils/spherical_harmonics.py b/deepmd/dpmodel/utils/spherical_harmonics.py new file mode 100644 index 0000000000..fc45abf5f7 --- /dev/null +++ b/deepmd/dpmodel/utils/spherical_harmonics.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +""" +Real spherical harmonics matching the e3nn convention used by SeZM. + +This module is a torch-free, pure-numpy replacement for the exact e3nn call +used by the SeZM Lebedev S2-grid projection code +(``deepmd/pt/model/descriptor/sezm_nn/projection.py``):: + + e3nn.o3.spherical_harmonics( + list(range(lmax + 1)), vecs, normalize=True, normalization="norm" + ) + +Empirically verified e3nn convention facts (probed on basis vectors and +random points, matching to machine precision): + +- Output layout: for each degree ``l`` from 0 to ``lmax``, a block of + ``2l + 1`` components ordered ``m = -l, ..., +l``; component ``l*l + l + m`` + is the order-``m`` harmonic. +- Axis convention: e3nn uses ``y`` as the polar axis. The ``l = 1`` block + under ``normalization="norm"`` is exactly the unit input vector + ``(x, y, z)`` in ``m = (-1, 0, +1)`` order, i.e. the standard z-polar real + spherical harmonics evaluated with axes ``(x_s, y_s, z_s) = (z, x, y)``. +- Phase convention: no Condon-Shortley phase. Negative orders carry + ``sin(m * phi)``, positive orders ``cos(m * phi)`` with + ``phi = atan2(x, z)``. +- Normalization ``"norm"``: ``Y_lm = sqrt(4*pi / (2l+1)) * Y_lm^{orthonormal}`` + so that ``sum_m Y_lm(v)^2 = 1`` for every unit vector ``v``. +- ``normalize=True``: input vectors are normalized internally before + evaluation, making the output invariant to the input vector length. +- Zero vectors: e3nn clamps the norm instead of dividing by zero, so a zero + input vector yields ``[Y00, 0, 0, ...] = [1, 0, 0, ...]``. This module + reproduces that exactly (no NaN, no runtime warning). +""" + +import numpy as np + +__all__ = ["real_spherical_harmonics"] + + +def real_spherical_harmonics(vecs: np.ndarray, lmax: int) -> np.ndarray: + """ + Evaluate real spherical harmonics in the e3nn ``"norm"`` convention. + + Exactly matches ``e3nn.o3.spherical_harmonics(list(range(lmax + 1)), + vecs, normalize=True, normalization="norm")`` (see module docstring). + + Parameters + ---------- + vecs + Input vectors with shape ``(..., 3)``. They do not need to be + normalized; vectors are normalized internally (e3nn + ``normalize=True``). Zero vectors map to ``[1, 0, 0, ...]`` + exactly as in e3nn (which clamps the norm before dividing). + lmax + Maximum angular degree, any non-negative integer. + + Returns + ------- + np.ndarray + Real spherical harmonics with shape ``(..., (lmax + 1) ** 2)`` in + float64, with each degree-``l`` block ordered ``m = -l, ..., +l``. + """ + if lmax < 0: + raise ValueError(f"lmax must be non-negative, got {lmax}") + vecs = np.asarray(vecs, dtype=np.float64) + if vecs.ndim == 0 or vecs.shape[-1] != 3: + raise ValueError(f"vecs must have shape (..., 3), got {vecs.shape}") + lead_shape = vecs.shape[:-1] + vecs = vecs.reshape(-1, 3) + # normalize=True: e3nn normalizes the input vectors internally. + # e3nn clamps the norm, so zero vectors stay zero instead of becoming + # NaN; the corresponding l >= 1 output components are zeroed below. + norm = np.linalg.norm(vecs, axis=-1, keepdims=True) + zero_mask = norm == 0.0 + vecs = vecs / np.where(zero_mask, 1.0, norm) + x, y, z = vecs[:, 0], vecs[:, 1], vecs[:, 2] + # e3nn polar axis is y: standard axes (x_s, y_s, z_s) = (z, x, y) + cos_theta = np.clip(y, -1.0, 1.0) + sin_theta = np.hypot(x, z) + phi = np.arctan2(x, z) + + nbatch = vecs.shape[0] + # Fully normalized associated Legendre functions (Condon-Shortley-free): + # pbar[l][m] = sqrt((2l+1)/(4*pi) * (l-m)!/(l+m)!) * P_l^m(cos_theta) + # computed with the standard stable recurrences on the normalized + # functions (no factorials, no overflow for large lmax). See DLMF + # §14.10 (https://dlmf.nist.gov/14.10) and Holmes & Featherstone, + # J. Geodesy 76, 279-299 (2002), doi:10.1007/s00190-002-0216-2. + pbar = [[None] * (ll + 1) for ll in range(lmax + 1)] + pbar[0][0] = np.full(nbatch, np.sqrt(1.0 / (4.0 * np.pi)), dtype=np.float64) + # diagonal: pbar[m][m] + for m in range(1, lmax + 1): + pbar[m][m] = ( + np.sqrt((2.0 * m + 1.0) / (2.0 * m)) * sin_theta * pbar[m - 1][m - 1] + ) + # first off-diagonal: pbar[m+1][m] + for m in range(lmax): + pbar[m + 1][m] = np.sqrt(2.0 * m + 3.0) * cos_theta * pbar[m][m] + # remaining: three-term recurrence in l + for m in range(lmax + 1): + for ll in range(m + 2, lmax + 1): + a_lm = np.sqrt((4.0 * ll * ll - 1.0) / (ll * ll - m * m)) + b_lm = np.sqrt(((ll - 1.0) ** 2 - m * m) / (4.0 * (ll - 1.0) ** 2 - 1.0)) + pbar[ll][m] = a_lm * (cos_theta * pbar[ll - 1][m] - b_lm * pbar[ll - 2][m]) + + out = np.zeros((nbatch, (lmax + 1) ** 2), dtype=np.float64) + sqrt2 = np.sqrt(2.0) + # sin(m*phi) / cos(m*phi): computed once per order m, reused across l + sin_mphi = [None] + [np.sin(m * phi) for m in range(1, lmax + 1)] + cos_mphi = [None] + [np.cos(m * phi) for m in range(1, lmax + 1)] + for ll in range(lmax + 1): + # normalization="norm": scale orthonormal SH by sqrt(4*pi/(2l+1)) + scale = np.sqrt(4.0 * np.pi / (2.0 * ll + 1.0)) + center = ll * ll + ll + out[:, center] = scale * pbar[ll][0] + for m in range(1, ll + 1): + base = sqrt2 * scale * pbar[ll][m] + out[:, center - m] = base * sin_mphi[m] + out[:, center + m] = base * cos_mphi[m] + # e3nn zero-vector behavior: the clamped-norm zero vector evaluates the + # degree-l polynomials at the origin, giving exactly [1, 0, 0, ...]. + # The theta/phi recurrences above are ill-defined there, so enforce it. + zero_out = np.zeros((lmax + 1) ** 2, dtype=np.float64) + zero_out[0] = 1.0 + out = np.where(zero_mask, zero_out, out) + return out.reshape(*lead_shape, (lmax + 1) ** 2) diff --git a/deepmd/pt/model/descriptor/sezm_nn/lebedev.py b/deepmd/pt/model/descriptor/sezm_nn/lebedev.py index 6e7105677e..a5a06f2f53 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/lebedev.py +++ b/deepmd/pt/model/descriptor/sezm_nn/lebedev.py @@ -5,49 +5,17 @@ annotations, ) -from pathlib import ( - Path, -) - -import numpy as np import torch -# See: https://people.sc.fsu.edu/~jburkardt/datasets/sphere_lebedev_rule/sphere_lebedev_rule.html -LEBEDEV_RULES_FILE = Path(__file__).with_name("lebedev_rules.npz") -LEBEDEV_PRECISION_TO_NPOINTS = { - 3: 6, - 5: 14, - 7: 26, - 9: 38, - 11: 50, - 13: 74, - 15: 86, - 17: 110, - 19: 146, - 21: 170, - 23: 194, - 25: 230, - 27: 266, - 29: 302, - 31: 350, - 35: 434, - 41: 590, - 47: 770, - 53: 974, - 59: 1202, - 65: 1454, - 71: 1730, - 77: 2030, - 83: 2354, - 89: 2702, - 95: 3074, - 101: 3470, - 107: 3890, - 113: 4334, - 119: 4802, - 125: 5294, - 131: 5810, -} +from deepmd.dpmodel.utils.lebedev import ( + LEBEDEV_PRECISION_TO_NPOINTS, +) +from deepmd.dpmodel.utils.lebedev import load_lebedev_rule as load_lebedev_rule_np + +__all__ = [ + "LEBEDEV_PRECISION_TO_NPOINTS", + "load_lebedev_rule", +] def load_lebedev_rule( @@ -75,18 +43,7 @@ def load_lebedev_rule( shape ``(A,)``. The weights sum to one, so the sphere integral is ``4*pi*sum(weights*f(points))``. """ - rule_key = f"{int(precision):03d}" - if not LEBEDEV_RULES_FILE.exists(): - raise FileNotFoundError( - f"Lebedev quadrature data file is missing: {LEBEDEV_RULES_FILE}" - ) - with np.load(LEBEDEV_RULES_FILE) as rules: - point_key = f"points_{rule_key}" - weight_key = f"weights_{rule_key}" - if point_key not in rules or weight_key not in rules: - raise ValueError(f"Lebedev rule with precision {precision} is not packaged") - points_np = rules[point_key] - weights_np = rules[weight_key] + points_np, weights_np = load_lebedev_rule_np(precision) points = torch.as_tensor(points_np, dtype=dtype, device=device) weights = torch.as_tensor(weights_np, dtype=dtype, device=device) return points, weights diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py new file mode 100644 index 0000000000..43c882864a --- /dev/null +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Torch-free unit tests for the dpmodel DPA4 (SeZM) descriptor.""" + +import numpy as np +import pytest + +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, +) + + +def build_neighbor_list_np(coord, rcut, nnei): + """Build a padded, distance-sorted gas-phase neighbor list. + + Parameters + ---------- + coord + Coordinates with shape (nf, nloc, 3); no PBC. + rcut + Cutoff radius. + nnei + Number of neighbor slots; pads with -1. + + Returns + ------- + np.ndarray + Neighbor list with shape (nf, nloc, nnei). + """ + nf, nloc, _ = coord.shape + nlist = -np.ones((nf, nloc, nnei), dtype=np.int64) + for f in range(nf): + dist = np.linalg.norm(coord[f][:, None, :] - coord[f][None, :, :], axis=-1) + for i in range(nloc): + neighbors = [ + (dist[i, j], j) for j in range(nloc) if j != i and dist[i, j] < rcut + ] + neighbors.sort() + for slot, (_, j) in enumerate(neighbors[:nnei]): + nlist[f, i, slot] = j + return nlist + + +def make_descriptor(**overrides) -> DescrptDPA4: + kwargs = { + "ntypes": 3, + "sel": 8, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 3, + "mmax": 1, + "n_blocks": 2, + "grid_branch": [1, 1, 1], + "s2_activation": [False, True], + "random_gamma": False, + "exclude_types": [(0, 0)], + "precision": "float64", + "seed": 42, + } + kwargs.update(overrides) + return DescrptDPA4(**kwargs) + + +def make_inputs(seed=5, nf=2, nloc=6, rcut=4.0, nnei=8, ntypes=3): + rng = np.random.default_rng(seed) + coord = rng.uniform(0.0, 3.5, size=(nf, nloc, 3)) + atype = rng.integers(0, ntypes, size=(nf, nloc)) + nlist = build_neighbor_list_np(coord, rcut, nnei) + return coord, atype, nlist + + +class TestDescrptDPA4: + def test_shapes_and_interface(self) -> None: + dd = make_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + out = dd.call(coord.reshape(nf, -1), atype, nlist, mapping=None) + assert out[0].shape == (nf, nloc, dd.get_dim_out()) + assert out[1:] == (None, None, None, None) + assert np.isfinite(np.asarray(out[0])).all() + # standard descriptor surface + assert dd.get_rcut() == 4.0 + assert dd.get_rcut_smth() == 4.0 + assert dd.get_sel() == [8] + assert dd.get_nsel() == 8 + assert dd.get_ntypes() == 3 + assert dd.get_type_map() == [] + assert dd.get_dim_out() == 16 + assert dd.get_dim_emb() == 16 + assert dd.mixed_types() is True + assert dd.has_message_passing() is True + assert dd.need_sorted_nlist_for_lower() is False + assert dd.get_env_protection() == dd.eps + + def test_has_message_passing_false(self) -> None: + # scalar-only model: lmax=0 carries no directional messages + dd = make_descriptor(lmax=0, mmax=0, kmax=0, n_blocks=1) + assert dd.has_message_passing() is False + + def test_serialize_roundtrip_exact(self) -> None: + dd = make_descriptor() + data = dd.serialize() + assert data["type"] == "SeZM" + dd2 = DescrptDPA4.deserialize(data) + coord, atype, nlist = make_inputs() + nf = atype.shape[0] + out1 = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + out2 = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0]) + np.testing.assert_array_equal(out1, out2) + + def test_permutation_equivariance(self) -> None: + dd = make_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + out = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + rng = np.random.default_rng(11) + perm = rng.permutation(nloc) + inv = np.argsort(perm) + coord2 = coord[:, perm, :] + atype2 = atype[:, perm] + nlist_p = nlist[:, perm, :] + nlist2 = np.where(nlist_p >= 0, inv[np.where(nlist_p >= 0, nlist_p, 0)], -1) + out2 = np.asarray(dd.call(coord2.reshape(nf, -1), atype2, nlist2)[0]) + np.testing.assert_allclose(out2, out[:, perm, :], rtol=1e-10, atol=1e-12) + + def test_rotation_invariance(self) -> None: + dd = make_descriptor() + coord, atype, nlist = make_inputs() + nf = atype.shape[0] + out = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + # a random proper rotation (QR with det fix) + rng = np.random.default_rng(13) + q, _ = np.linalg.qr(rng.normal(size=(3, 3))) + if np.linalg.det(q) < 0: + q[:, 0] = -q[:, 0] + coord_rot = coord @ q.T # distances (and the nlist) are unchanged + out_rot = np.asarray(dd.call(coord_rot.reshape(nf, -1), atype, nlist)[0]) + np.testing.assert_allclose(out_rot, out, rtol=1e-10, atol=1e-12) + + def test_masked_edge_inertness(self) -> None: + # an extra all-(-1) neighbor column must not change the descriptor + dd = make_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + out = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + pad = -np.ones((nf, nloc, 1), dtype=nlist.dtype) + nlist2 = np.concatenate([nlist, pad], axis=-1) + out2 = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist2)[0]) + np.testing.assert_allclose(out2, out, rtol=1e-12, atol=1e-14) + + @pytest.mark.parametrize( + "flag,value", + [ + # guarded at the descriptor level + ("lebedev_quadrature", False), # tensor-product S2 grid + ("lebedev_quadrature", [False, True]), # tensor-product S2 grid (so2) + ("lebedev_quadrature", [True, False]), # tensor-product S2 grid (ffn) + ("add_chg_spin_ebd", True), # ChargeSpinEmbedding + ("inner_clamp_r_inner", 0.5), # zone bridging + ("inner_clamp_r_outer", 1.0), # zone bridging + # delegated to the owning submodules + ("layer_scale", True), # block LayerScale + ("full_attn_res", "independent"), # DepthAttnRes + ("block_attn_res", "dependent"), # DepthAttnRes + ("so2_attn_res", "independent"), # SO(2) DepthAttnRes + ("s2_activation", [True, True]), # so2-side S2 activation + ("node_wise_s2", True), # SO(2) cross-grid product + ("message_node_so3", True), # SO(2) cross-grid product + ("ffn_so3_grid", True), # SO(3) Wigner-D FFN grid + ("atten_f_mix", True), # SO(2) attention focus mix + ("atten_v_proj", True), # SO(2) attention value projection + ("atten_o_proj", True), # SO(2) attention output projection + ], + ) + def test_not_implemented_guards(self, flag, value) -> None: + with pytest.raises(NotImplementedError): + make_descriptor(**{flag: value}) + + @pytest.mark.parametrize( + "flag,value", + [ + ("lebedev_quadrature", True), # supported branch of every guard + ("add_chg_spin_ebd", False), + ("inner_clamp_r_inner", None), + ("layer_scale", False), + ("full_attn_res", "none"), + ("s2_activation", [False, True]), + ("node_wise_s2", False), + ("ffn_so3_grid", False), + ("use_amp", True), # pt-runtime-only switch: accepted and ignored + ("use_amp", False), + ], + ) + def test_supported_branches_construct(self, flag, value) -> None: + dd = make_descriptor(**{flag: value}) + assert isinstance(dd, DescrptDPA4) + + def test_value_errors(self) -> None: + with pytest.raises(ValueError): # kmax must be <= lmax + make_descriptor(kmax=4, lmax=3) + with pytest.raises(ValueError): # m_schedule entries must be <= l_schedule + make_descriptor(l_schedule=[2, 2], m_schedule=[3, 1]) + with pytest.raises(ValueError): # l_schedule must be non-increasing + make_descriptor(l_schedule=[2, 3]) + with pytest.raises(ValueError): # sandwich_norm must have length 4 + make_descriptor(sandwich_norm=[True, False]) + with pytest.raises(ValueError): # env_exp must have length 2 + make_descriptor(env_exp=[7]) + with pytest.raises(ValueError): # attn res mode token + make_descriptor(full_attn_res="depth") + with pytest.raises(ValueError): # wrong class tag + DescrptDPA4.deserialize({"@class": "NotDescriptor", "type": "SeZM"}) + with pytest.raises(ValueError): # wrong type tag + DescrptDPA4.deserialize({"@class": "Descriptor", "type": "se_e2_a"}) diff --git a/source/tests/common/dpmodel/test_dpa4_ener.py b/source/tests/common/dpmodel/test_dpa4_ener.py new file mode 100644 index 0000000000..5f96f56ac2 --- /dev/null +++ b/source/tests/common/dpmodel/test_dpa4_ener.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Torch-free unit tests for the dpmodel DPA4 (SeZM) GLU energy fitting.""" + +import subprocess +import sys + +import numpy as np +import pytest + +from deepmd.dpmodel.fitting.dpa4_ener import ( + GLUFittingNet, + SeZMEnergyFittingNet, + _resolve_auto_neuron, +) + + +def make_fitting(**overrides) -> SeZMEnergyFittingNet: + kwargs = { + "ntypes": 2, + "dim_descrpt": 12, + "neuron": [16], + "precision": "float64", + "seed": 7, + } + kwargs.update(overrides) + return SeZMEnergyFittingNet(**kwargs) + + +def make_inputs(nf=2, nloc=5, dim=12, ntypes=2, seed=0): + rng = np.random.default_rng(seed) + descriptor = rng.normal(size=(nf, nloc, dim)) + atype = rng.integers(0, ntypes, size=(nf, nloc)) + atype[0, 0], atype[0, 1] = 0, 1 + return descriptor, atype + + +class TestGuards: + def test_dim_case_embd_not_implemented(self) -> None: + with pytest.raises(NotImplementedError, match="dim_case_embd"): + make_fitting(dim_case_embd=2) + + def test_case_film_embd_not_implemented(self) -> None: + with pytest.raises(NotImplementedError, match="case_film_embd"): + make_fitting(case_film_embd=True) + + def test_glu_net_case_film_not_implemented(self) -> None: + with pytest.raises(NotImplementedError, match="case_film_embd"): + GLUFittingNet(8, 1, [16], dim_case_embd=2, case_film_embd=True) + + def test_negative_neuron_raises(self) -> None: + with pytest.raises(ValueError, match="neuron"): + make_fitting(neuron=[-1]) + + @pytest.mark.parametrize( + "key", + [ + (0, 0), # wrong length (ndim=1) + (-1,), # negative type index + (2,), # type index >= ntypes + 5, # int index out of bounds + ], + ) # invalid network-collection keys + def test_network_collection_invalid_key(self, key) -> None: + fitting = make_fitting() + # __getitem__ raises KeyError per the lookup-method convention + with pytest.raises(KeyError): + fitting.nets[key] + + +class TestAutoNeuron: + def test_auto_width_marker(self) -> None: + # dim_in=12 -> 32*ceil(8*12/3/32) = 32 + assert _resolve_auto_neuron( + [0], + dim_descrpt=12, + numb_fparam=0, + numb_aparam=0, + dim_case_embd=0, + case_film_embd=False, + use_aparam_as_mask=False, + ) == [32] + + def test_no_marker_passthrough(self) -> None: + assert _resolve_auto_neuron( + [16, 16], + dim_descrpt=12, + numb_fparam=0, + numb_aparam=0, + dim_case_embd=0, + case_film_embd=False, + use_aparam_as_mask=False, + ) == [16, 16] + + def test_none_means_single_auto(self) -> None: + fit = make_fitting(neuron=None) + assert fit.neuron == [32] + + +class TestRoundtrip: + @pytest.mark.parametrize("bias_out", [False, True]) # output-layer bias + @pytest.mark.parametrize("resnet_dt", [False, True]) # serialized flag only + @pytest.mark.parametrize( + "neuron", [[], [16], [16, 16]] + ) # direct linear / shallow / deep + def test_serialize_roundtrip_exact(self, neuron, bias_out, resnet_dt) -> None: + fit = make_fitting(neuron=neuron, bias_out=bias_out, resnet_dt=resnet_dt) + data = fit.serialize() + assert data["type"] == "sezm_ener" + assert data["@class"] == "Fitting" + fit2 = SeZMEnergyFittingNet.deserialize(data) + assert set(fit2.serialize().keys()) == set(data.keys()) + descriptor, atype = make_inputs() + out1 = fit.call(descriptor, atype)["energy"] + out2 = fit2.call(descriptor, atype)["energy"] + np.testing.assert_array_equal(out1, out2) + + def test_glu_net_roundtrip_exact(self) -> None: + net = GLUFittingNet(8, 1, [16], precision="float64", bias_out=True, seed=3) + net2 = GLUFittingNet.deserialize(net.serialize()) + x = np.random.default_rng(1).normal(size=(4, 8)) + np.testing.assert_array_equal(net.call(x), net2.call(x)) + # state keys follow the pt state_dict contract + assert set(net.serialize()["@variables"]) == { + "hidden_layers.0.linear.matrix", + "hidden_layers.0.linear.bias", + "output_layer.matrix", + "output_layer.bias", + } + + def test_glu_net_no_bias_out_keys(self) -> None: + net = GLUFittingNet(8, 1, [16], precision="float64", bias_out=False, seed=3) + assert "output_layer.bias" not in net.serialize()["@variables"] + + +class TestForward: + @pytest.mark.parametrize("mixed_types", [True, False]) # shared vs per-type nets + def test_output_shape(self, mixed_types) -> None: + fit = make_fitting(mixed_types=mixed_types) + descriptor, atype = make_inputs() + out = fit.call(descriptor, atype)["energy"] + assert out.shape == (2, 5, 1) + + def test_bias_atom_e_added(self) -> None: + fit = make_fitting() + bias = np.array([[1.5], [-2.5]]) + fit["bias_atom_e"] = bias + descriptor, atype = make_inputs() + out0 = make_fitting().call(descriptor, atype)["energy"] + out1 = fit.call(descriptor, atype)["energy"] + np.testing.assert_allclose(out1 - out0, bias[atype], rtol=1e-12, atol=1e-14) + + def test_exclude_types_zeroed(self) -> None: + fit = make_fitting(exclude_types=[0]) + descriptor, atype = make_inputs() + out = fit.call(descriptor, atype)["energy"] + assert np.all(out[atype == 0] == 0.0) + assert np.all(out[atype == 1] != 0.0) + + def test_default_fparam_matches_explicit(self) -> None: + fit = make_fitting(numb_fparam=2, default_fparam=[0.25, -0.75]) + descriptor, atype = make_inputs() + out_default = fit.call(descriptor, atype)["energy"] + fparam = np.tile(np.array([[0.25, -0.75]]), (2, 1)) + out_explicit = fit.call(descriptor, atype, fparam=fparam)["energy"] + np.testing.assert_array_equal(out_default, out_explicit) + + def test_trainable_list_accepted(self) -> None: + fit = make_fitting(trainable=[True, False]) + descriptor, atype = make_inputs() + assert fit.call(descriptor, atype)["energy"].shape == (2, 5, 1) + + +class TestNoTorchImport: + def test_dpa4_ener_does_not_import_torch(self) -> None: + code = ( + "import sys; " + "import deepmd.dpmodel.fitting.dpa4_ener; " + "print('torch' in sys.modules)" + ) + out = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ) + assert out.stdout.strip() == "False" diff --git a/source/tests/common/dpmodel/test_lebedev_sh.py b/source/tests/common/dpmodel/test_lebedev_sh.py new file mode 100644 index 0000000000..5b7a448ed1 --- /dev/null +++ b/source/tests/common/dpmodel/test_lebedev_sh.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest + +import deepmd.dpmodel.utils.lebedev as lebedev_module +from deepmd.dpmodel.utils.spherical_harmonics import ( + real_spherical_harmonics, +) + + +class TestLebedevRules: + @pytest.mark.parametrize("precision", [3, 9, 11, 29]) # quadrature precision order + def test_rule_basic(self, precision): + pts, wts = lebedev_module.load_lebedev_rule(precision) + assert pts.shape[1] == 3 and wts.shape == (pts.shape[0],) + assert pts.shape[0] == lebedev_module.LEBEDEV_PRECISION_TO_NPOINTS[precision] + np.testing.assert_allclose(np.linalg.norm(pts, axis=1), 1.0, rtol=1e-12) + np.testing.assert_allclose(wts.sum(), 1.0, rtol=1e-12) + + def test_unpackaged_precision_raises(self): + with pytest.raises(ValueError, match="not packaged"): + lebedev_module.load_lebedev_rule(4) + + @pytest.mark.parametrize("precision", [3.5, 11.0, "11", None]) # non-integers + def test_non_integer_precision_raises(self, precision): + with pytest.raises(TypeError, match="integer"): + lebedev_module.load_lebedev_rule(precision) + + def test_missing_data_file_raises(self, monkeypatch, tmp_path): + monkeypatch.setattr(lebedev_module, "LEBEDEV_RULES_FILE", tmp_path / "nope.npz") + with pytest.raises(FileNotFoundError, match="missing"): + lebedev_module.load_lebedev_rule(11) + + def test_pt_loader_matches(self): + torch = pytest.importorskip("torch") + from deepmd.pt.model.descriptor.sezm_nn.lebedev import ( + load_lebedev_rule as pt_rule, + ) + + pts, wts = lebedev_module.load_lebedev_rule(11) + tpts, twts = pt_rule(11, dtype=torch.float64, device="cpu") + np.testing.assert_allclose(pts, tpts.numpy(), rtol=0, atol=0) + np.testing.assert_allclose(wts, twts.numpy(), rtol=0, atol=0) + + +class TestRealSphericalHarmonics: + @pytest.mark.parametrize("lmax", [0, 1, 2, 3, 4, 6]) # maximum angular degree + def test_matches_e3nn(self, lmax): + pytest.importorskip("e3nn") + import torch + from e3nn import ( + o3, + ) + + rng = np.random.default_rng(0) + v = rng.standard_normal((64, 3)) + v /= np.linalg.norm(v, axis=1, keepdims=True) + # exact call convention used by the SeZM Lebedev projection path + # (deepmd/pt/model/descriptor/sezm_nn/projection.py) + ref = o3.spherical_harmonics( + list(range(lmax + 1)), + torch.from_numpy(v), + normalize=True, + normalization="norm", + ).numpy() + out = real_spherical_harmonics(v, lmax) + assert out.shape == (64, (lmax + 1) ** 2) + assert out.dtype == np.float64 + np.testing.assert_allclose(out, ref, rtol=1e-12, atol=1e-13) + + @pytest.mark.parametrize("lmax", [2, 4]) # maximum angular degree + def test_scale_invariance(self, lmax): + # normalize=True in the e3nn call: input vectors are normalized + # internally, so non-unit inputs must give identical output. + rng = np.random.default_rng(1) + v = rng.standard_normal((32, 3)) + v /= np.linalg.norm(v, axis=1, keepdims=True) + scale = rng.uniform(0.1, 10.0, size=(32, 1)) + np.testing.assert_allclose( + real_spherical_harmonics(v * scale, lmax), + real_spherical_harmonics(v, lmax), + rtol=1e-12, + atol=1e-14, + ) + + def test_batched_leading_dims(self): + rng = np.random.default_rng(2) + v = rng.standard_normal((4, 5, 3)) + out = real_spherical_harmonics(v, 3) + assert out.shape == (4, 5, 16) + flat = real_spherical_harmonics(v.reshape(-1, 3), 3) + np.testing.assert_allclose(out.reshape(-1, 16), flat, rtol=0, atol=0) + + # e3nn-free convention pin: the l=1 block under normalization="norm" + # is exactly the unit input vector in (x, y, z) order (m = -1, 0, +1). + def test_l1_block_is_unit_vector(self): + rng = np.random.default_rng(3) + v = rng.standard_normal((128, 3)) + v /= np.linalg.norm(v, axis=1, keepdims=True) + out = real_spherical_harmonics(v, 1) + np.testing.assert_allclose(out[:, 1:4], v, rtol=1e-12, atol=1e-14) + + def test_basis_vectors_lmax2(self): + # e3nn-free convention pin: analytic SH values at lmax=2 for the + # Cartesian basis vectors. Cross-checked against + # e3nn.o3.spherical_harmonics([0, 1, 2], v, normalize=True, + # normalization="norm"): the analytic values below equal the e3nn + # output to the last bit (max |diff| = 0.0). + h = np.sqrt(3.0) / 2.0 + cases = [ + ((1.0, 0.0, 0.0), [1.0, 1.0, 0.0, 0.0, 0.0, 0.0, -0.5, 0.0, -h]), + ((0.0, 1.0, 0.0), [1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), + ((0.0, 0.0, 1.0), [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, -0.5, 0.0, h]), + ] + for vec, ref in cases: + out = real_spherical_harmonics(np.array([vec]), 2) + np.testing.assert_allclose(out[0], ref, rtol=1e-12, atol=1e-15) + + @pytest.mark.filterwarnings("error") + @pytest.mark.parametrize("vecs", [1.0, np.float64(2.0), [1.0, 0.0]]) # bad shapes + def test_invalid_input_shape_raises(self, vecs): + # scalar/0-d inputs and wrong last-axis sizes must raise, not crash + with pytest.raises(ValueError, match="shape"): + real_spherical_harmonics(vecs, 2) + + def test_negative_lmax_raises(self): + with pytest.raises(ValueError, match="lmax"): + real_spherical_harmonics(np.zeros((4, 3)), -1) + + def test_zero_vector(self): + # e3nn's normalize=True clamps the norm, so a zero vector maps to + # [Y00, 0, 0, ...] = [1, 0, ...]. Verified against + # e3nn.o3.spherical_harmonics([0, 1, 2], zeros, normalize=True, + # normalization="norm") -> [1, 0, 0, 0, 0, 0, 0, 0, 0]. + expected = np.zeros(9) + expected[0] = 1.0 + with np.errstate(invalid="raise", divide="raise"): + out = real_spherical_harmonics(np.zeros((1, 3)), 2) + np.testing.assert_allclose(out[0], expected, rtol=0, atol=0) + + @pytest.mark.filterwarnings("error") + def test_zero_vector_mixed_batch(self): + # batch mixing zero and unit vectors: zero rows give [1, 0, ...], + # nonzero rows are unaffected by the zero-vector guard + rng = np.random.default_rng(4) + v = rng.standard_normal((6, 3)) + v /= np.linalg.norm(v, axis=1, keepdims=True) + v[1] = 0.0 + v[4] = 0.0 + with np.errstate(invalid="raise", divide="raise"): + out = real_spherical_harmonics(v, 2) + expected_zero = np.zeros(9) + expected_zero[0] = 1.0 + for i in (1, 4): + np.testing.assert_allclose(out[i], expected_zero, rtol=0, atol=0) + nonzero = [0, 2, 3, 5] + np.testing.assert_allclose( + out[nonzero], + real_spherical_harmonics(v[nonzero], 2), + rtol=0, + atol=0, + ) + + def test_quadrature_orthogonality(self): + lmax = 3 + pts, wts = lebedev_module.load_lebedev_rule(2 * lmax + 1) + sh = real_spherical_harmonics(pts, lmax) + gram = (sh[:, :, None] * sh[:, None, :] * wts[:, None, None]).sum(axis=0) + # The implemented convention is e3nn normalization="norm": + # Y_lm = sqrt(4*pi/(2l+1)) * Y_lm^{orthonormal} + # so int Y_lm Y_l'm' dOmega = (4*pi/(2l+1)) * delta_ll' delta_mm'. + # Lebedev weights sum to 1 (absorbing the 1/(4*pi) surface factor), + # hence gram = blockdiag over l of I_{2l+1} / (2l+1). + expected = np.zeros_like(gram) + for ll in range(lmax + 1): + for mm in range(ll * ll, (ll + 1) ** 2): + expected[mm, mm] = 1.0 / (2 * ll + 1) + # products of degree <= 2*lmax are integrated exactly by the + # precision-(2*lmax+1) Lebedev rule -> machine precision + np.testing.assert_allclose(gram, expected, atol=1e-12, rtol=0) diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py new file mode 100644 index 0000000000..ced71078f3 --- /dev/null +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, + ClassVar, +) + +import numpy as np +from dargs import ( + Argument, +) + +from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4DP +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.utils.argcheck import ( + descrpt_se_zm_args, +) + +from ..common import ( + INSTALLED_PT, + CommonTest, + parameterized_cases, +) +from .common import ( + DescriptorTest, +) + +if INSTALLED_PT: + from deepmd.pt.model.descriptor.sezm import DescrptSeZM as DescrptDPA4PT +else: + DescrptDPA4PT = None + +# not implemented +DescrptDPA4TF = None + + +DPA4_CASE_FIELDS = ( + "precision", + "grid_branch", + "s2_activation", + "basis_type", +) + +DPA4_BASELINE_CASE = { + "precision": "float64", + "grid_branch": [1, 1, 1], + "s2_activation": [False, True], + "basis_type": "bessel", +} + + +def dpa4_case(**overrides: Any) -> tuple: + unknown = set(overrides) - set(DPA4_BASELINE_CASE) + if unknown: + raise KeyError(f"Unknown DPA4 case override(s): {sorted(unknown)}") + case = dict(DPA4_BASELINE_CASE) + case.update(overrides) + return tuple(case[field] for field in DPA4_CASE_FIELDS) + + +# curated cases (one-factor-at-a-time, dpa3 precedent) instead of full +# cross product to keep CI runtime sane +DPA4_CURATED_CASES = ( + # baseline coverage + dpa4_case(), + # grid branch disabled + dpa4_case(grid_branch=[0, 0, 0]), + # no S2 activation in any FFN + dpa4_case(s2_activation=[False, False]), + # gaussian radial basis + dpa4_case(basis_type="gaussian"), + # float32 baseline + dpa4_case(precision="float32"), + # float32 mixed high-risk path + dpa4_case( + precision="float32", + grid_branch=[0, 0, 0], + s2_activation=[False, False], + basis_type="gaussian", + ), +) + + +@parameterized_cases(*DPA4_CURATED_CASES) +class TestDPA4(CommonTest, DescriptorTest, unittest.TestCase): + @property + def data(self) -> dict: + ( + precision, + grid_branch, + s2_activation, + basis_type, + ) = self.param + return { + "ntypes": self.ntypes, + "sel": 10, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "basis_type": basis_type, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "grid_branch": grid_branch, + "s2_activation": s2_activation, + "random_gamma": False, + "precision": precision, + "trainable": False, + "seed": 20251208, + } + + @property + def skip_pt(self) -> bool: + return CommonTest.skip_pt + + skip_dp = False + skip_tf = True + skip_jax = True + skip_pd = True + skip_pt_expt = True + skip_array_api_strict = True + + tf_class = DescrptDPA4TF + dp_class = DescrptDPA4DP + pt_class = DescrptDPA4PT + pt_expt_class = None + jax_class = None + pd_class = None + array_api_strict_class = None + args: ClassVar[list] = [ + *descrpt_se_zm_args(), + Argument("ntypes", int, optional=False), + ] + + def setUp(self) -> None: + CommonTest.setUp(self) + + self.ntypes = 2 + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ) + self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + raise NotImplementedError("DPA4 is not implemented in TensorFlow") + + def eval_dp(self, dp_obj: Any) -> Any: + return self.eval_dp_descriptor( + dp_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + + def eval_pt(self, pt_obj: Any) -> Any: + return self.eval_pt_descriptor( + pt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + + def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: + return (ret[0],) + + @property + def rtol(self) -> float: + """Relative tolerance for comparing the return value.""" + ( + precision, + _grid_branch, + _s2_activation, + _basis_type, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") + + @property + def atol(self) -> float: + """Absolute tolerance for comparing the return value.""" + ( + precision, + _grid_branch, + _s2_activation, + _basis_type, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py new file mode 100644 index 0000000000..666594bf64 --- /dev/null +++ b/source/tests/consistent/fitting/test_dpa4_ener.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.fitting.dpa4_ener import SeZMEnergyFittingNet as SeZMEnerFittingDP +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.utils.argcheck import ( + fitting_sezm_ener, +) + +from ..common import ( + INSTALLED_PT, + CommonTest, + parameterized, +) +from .common import ( + FittingTest, +) + +if INSTALLED_PT: + import torch + + from deepmd.pt.model.task.sezm_ener import SeZMEnergyFittingNet as SeZMEnerFittingPT + from deepmd.pt.utils.env import DEVICE as PT_DEVICE +else: + SeZMEnerFittingPT = None + +# not implemented +SeZMEnerFittingTF = None + + +@parameterized( + ("float64", "float32"), # precision + ([0], [16, 16]), # neuron ([0] = auto-width placeholder) +) +class TestDPA4Ener(CommonTest, FittingTest, unittest.TestCase): + @property + def data(self) -> dict: + ( + precision, + neuron, + ) = self.param + return { + "neuron": neuron, + "precision": precision, + "seed": 20251208, + "activation_function": "silu", + } + + @property + def skip_pt(self) -> bool: + return CommonTest.skip_pt + + skip_dp = False + skip_tf = True + skip_jax = True + skip_pd = True + skip_pt_expt = True + skip_array_api_strict = True + + tf_class = SeZMEnerFittingTF + dp_class = SeZMEnerFittingDP + pt_class = SeZMEnerFittingPT + pt_expt_class = None + jax_class = None + pd_class = None + array_api_strict_class = None + args = fitting_sezm_ener() + + def setUp(self) -> None: + CommonTest.setUp(self) + + self.ntypes = 2 + self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + rng = np.random.default_rng(20251208) + self.inputs = rng.normal(size=(1, 6, 20)).astype(GLOBAL_NP_FLOAT_PRECISION) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) + # inconsistent if not sorted + self.atype.sort() + + @property + def additional_data(self) -> dict: + return { + "ntypes": self.ntypes, + "dim_descrpt": self.inputs.shape[-1], + "mixed_types": True, + } + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + raise NotImplementedError("dpa4_ener is not implemented in TensorFlow") + + def eval_dp(self, dp_obj: Any) -> Any: + return dp_obj( + self.inputs, + self.atype.reshape(1, -1), + )["energy"] + + def eval_pt(self, pt_obj: Any) -> Any: + return ( + pt_obj( + torch.from_numpy(self.inputs).to(device=PT_DEVICE), + torch.from_numpy(self.atype.reshape(1, -1)).to(device=PT_DEVICE), + )["energy"] + .detach() + .cpu() + .numpy() + ) + + def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: + return (ret,) + + @property + def rtol(self) -> float: + """Relative tolerance for comparing the return value.""" + ( + precision, + _neuron, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") + + @property + def atol(self) -> float: + """Absolute tolerance for comparing the return value.""" + ( + precision, + _neuron, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py new file mode 100644 index 0000000000..08d78e7320 --- /dev/null +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -0,0 +1,3837 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Parity tests: dpmodel ``dpa4_nn`` modules vs the reference pt ``sezm_nn`` modules. + +This file is extended task-by-task as the DPA4 port progresses. Index-table and +numeric-helper parity is exact (``assert_array_equal`` / tight rtol). Module-level +weight-copy parity tests (build the pt module in float64, copy its ``state_dict`` +into the dpmodel module via ``pt_state_to_numpy``, compare forwards with +``assert_parity``) are added by the tasks that port each module. +""" + +import subprocess +import sys + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.descriptor.dpa4_nn import indexing as dp_indexing +from deepmd.dpmodel.descriptor.dpa4_nn import utils as dp_utils +from deepmd.pt.model.descriptor.sezm_nn import indexing as pt_indexing +from deepmd.pt.model.descriptor.sezm_nn import utils as pt_utils +from deepmd.pt.utils import env as pt_env + +# pt reference modules run on their native device (house convention). +# On CPU the pt and numpy fp64 math is identical to ~1 ulp, so the parity +# gate is near-bit (rtol 1e-12). On CUDA, fp64 kernels differ from CPU +# numpy at ULP level per op and index_add_ uses nondeterministic atomics, +# so the gate is relaxed to rtol 1e-10 — still orders of magnitude below +# any logic bug. +PT_DEVICE = pt_env.DEVICE +_ON_CPU = PT_DEVICE.type == "cpu" +PT_RTOL, PT_ATOL = (1e-12, 1e-14) if _ON_CPU else (1e-10, 1e-12) + + +def to_pt(x: np.ndarray) -> torch.Tensor: + """Move a numpy array onto the pt reference device.""" + return torch.from_numpy(np.ascontiguousarray(x)).to(PT_DEVICE) + + +def pt_state_to_numpy(module: torch.nn.Module) -> dict[str, np.ndarray]: + return {k: v.detach().cpu().numpy() for k, v in module.state_dict().items()} + + +def assert_parity(a, t, rtol=PT_RTOL, atol=PT_ATOL): + np.testing.assert_allclose( + np.asarray(a), t.detach().cpu().numpy(), rtol=rtol, atol=atol + ) + + +# The pt indexing helpers below are pure functions taking an explicit +# ``device``; their integer index tables are device-independent, so they +# stay CPU-pinned to allow direct ``.numpy()`` comparison. +CPU = torch.device("cpu") + + +class TestIndexingParity: + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + def test_get_so3_dim_of_lmax(self, lmax) -> None: + assert dp_indexing.get_so3_dim_of_lmax(lmax) == pt_indexing.get_so3_dim_of_lmax( + lmax + ) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + def test_map_degree_idx(self, lmax) -> None: + res = dp_indexing.map_degree_idx(lmax) + ref = pt_indexing.map_degree_idx(lmax, device=CPU) + assert res.dtype == np.int64 + np.testing.assert_array_equal(res, ref.numpy()) + + @pytest.mark.parametrize("lmax", [0, 1, 2, 3, 4]) # incl. lmax=0 empty branch + def test_build_gie_zonal_index(self, lmax) -> None: + res = dp_indexing.build_gie_zonal_index(lmax) + ref = pt_indexing.build_gie_zonal_index(lmax, device=CPU) + assert len(res) == len(ref) == 3 + for r, t in zip(res, ref, strict=True): + assert r.dtype == np.int64 + np.testing.assert_array_equal(r, t.numpy()) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + def test_so3_packed_index(self, lmax) -> None: + for degree in range(lmax + 1): + for m in range(-degree, degree + 1): + assert dp_indexing.so3_packed_index( + degree, m + ) == pt_indexing.so3_packed_index(degree, m) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + @pytest.mark.parametrize("mmax", [1, 2]) # max order |m| + def test_build_l_major_index(self, lmax, mmax) -> None: + if mmax > lmax: + pytest.skip("mmax must be <= lmax") + res = dp_indexing.build_l_major_index(lmax, mmax) + ref = pt_indexing.build_l_major_index(lmax, mmax, device=CPU) + assert res.dtype == np.int64 + np.testing.assert_array_equal(res, ref.numpy()) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + @pytest.mark.parametrize("mmax", [1, 2]) # max order |m| + def test_build_m_major_index(self, lmax, mmax) -> None: + if mmax > lmax: + pytest.skip("mmax must be <= lmax") + res = dp_indexing.build_m_major_index(lmax, mmax) + ref = pt_indexing.build_m_major_index(lmax, mmax, device=CPU) + assert res.dtype == np.int64 + np.testing.assert_array_equal(res, ref.numpy()) + + def test_m_major_index_literal(self) -> None: + # layout contract anchor, cross-checked with sezm_nn docs: + # lmax=2, mmax=1: m=0 block (l=0..2), then m=-1, then m=+1 + np.testing.assert_array_equal( + dp_indexing.build_m_major_index(2, 1), [0, 2, 6, 1, 5, 3, 7] + ) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + @pytest.mark.parametrize("mmax", [1, 2]) # max order |m| + def test_build_m_major_l_index(self, lmax, mmax) -> None: + if mmax > lmax: + pytest.skip("mmax must be <= lmax") + res = dp_indexing.build_m_major_l_index(lmax, mmax) + ref = pt_indexing.build_m_major_l_index(lmax, mmax, device=CPU) + assert res.dtype == np.int64 + np.testing.assert_array_equal(res, ref.numpy()) + + @pytest.mark.parametrize( + "builder", + ["build_l_major_index", "build_m_major_index", "build_m_major_l_index"], + ) # index builder under test + @pytest.mark.parametrize( + "lmax,mmax", [(-1, 0), (1, -1), (1, 2)] + ) # lmax<0, mmax<0, mmax>lmax error branches + def test_index_builder_errors(self, builder, lmax, mmax) -> None: + with pytest.raises(ValueError): + getattr(dp_indexing, builder)(lmax, mmax) + with pytest.raises(ValueError): + getattr(pt_indexing, builder)(lmax, mmax, device=CPU) + + @pytest.mark.parametrize( + "lmax", [1, 2, 3, 4] + ) # max degree; lmax==mmax hits the all-ones branch + @pytest.mark.parametrize("mmax", [1, 2]) # max order |m| + def test_build_rotate_inv_rescale(self, lmax, mmax) -> None: + if mmax > lmax: + pytest.skip("mmax must be <= lmax") + degree_index_np = dp_indexing.build_m_major_l_index(lmax, mmax) + degree_index_pt = pt_indexing.build_m_major_l_index(lmax, mmax, device=CPU) + res = dp_indexing.build_rotate_inv_rescale( + lmax, mmax, degree_index_np, dtype=np.float64 + ) + ref = pt_indexing.build_rotate_inv_rescale( + lmax, mmax, degree_index_pt, device=CPU, dtype=torch.float64 + ) + assert res.dtype == np.float64 + np.testing.assert_allclose(res, ref.numpy(), rtol=1e-15, atol=0.0) + + @pytest.mark.parametrize( + "lmax,mmax", [(-1, 0), (1, -1), (1, 2)] + ) # lmax<0, mmax<0, mmax>lmax error branches + def test_build_rotate_inv_rescale_errors(self, lmax, mmax) -> None: + degree_index = np.zeros(1, dtype=np.int64) + with pytest.raises(ValueError): + dp_indexing.build_rotate_inv_rescale(lmax, mmax, degree_index) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + @pytest.mark.parametrize("mmax", [1, 2]) # max order |m| + def test_project_D_to_m(self, lmax, mmax) -> None: + if mmax > lmax: + pytest.skip("mmax must be <= lmax") + rng = np.random.default_rng(2026) + nfull = dp_indexing.get_so3_dim_of_lmax(4) + d_full_np = rng.normal(size=(5, nfull, nfull)) + d_full_pt = torch.from_numpy(d_full_np) # CPU: pure fn, CPU index table + idx_np = dp_indexing.build_m_major_index(lmax, mmax) + idx_pt = pt_indexing.build_m_major_index(lmax, mmax, device=CPU) + ebed = dp_indexing.get_so3_dim_of_lmax(lmax) + # cache=None branch + res = dp_indexing.project_D_to_m(d_full_np, idx_np, ebed, None, lmax, mmax) + ref = pt_indexing.project_D_to_m(d_full_pt, idx_pt, ebed, None, lmax, mmax) + assert res.shape == (5, idx_np.shape[0], ebed) + np.testing.assert_array_equal(np.asarray(res), ref.numpy()) + # cache branch: miss then hit (returned object identical) + cache: dict = {} + first = dp_indexing.project_D_to_m(d_full_np, idx_np, ebed, cache, lmax, mmax) + second = dp_indexing.project_D_to_m(d_full_np, idx_np, ebed, cache, lmax, mmax) + assert second is first + np.testing.assert_array_equal(np.asarray(first), ref.numpy()) + + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) # max spherical harmonic degree + @pytest.mark.parametrize("mmax", [1, 2]) # max order |m| + def test_project_Dt_from_m(self, lmax, mmax) -> None: + if mmax > lmax: + pytest.skip("mmax must be <= lmax") + rng = np.random.default_rng(2027) + nfull = dp_indexing.get_so3_dim_of_lmax(4) + dt_full_np = rng.normal(size=(5, nfull, nfull)) + dt_full_pt = torch.from_numpy(dt_full_np) # CPU: pure fn, CPU index table + idx_np = dp_indexing.build_m_major_index(lmax, mmax) + idx_pt = pt_indexing.build_m_major_index(lmax, mmax, device=CPU) + ebed = dp_indexing.get_so3_dim_of_lmax(lmax) + # cache=None branch + res = dp_indexing.project_Dt_from_m(dt_full_np, idx_np, ebed, None, lmax, mmax) + ref = pt_indexing.project_Dt_from_m(dt_full_pt, idx_pt, ebed, None, lmax, mmax) + assert res.shape == (5, ebed, idx_np.shape[0]) + np.testing.assert_array_equal(np.asarray(res), ref.numpy()) + # cache branch: miss then hit (returned object identical) + cache: dict = {} + first = dp_indexing.project_Dt_from_m( + dt_full_np, idx_np, ebed, cache, lmax, mmax + ) + second = dp_indexing.project_Dt_from_m( + dt_full_np, idx_np, ebed, cache, lmax, mmax + ) + assert second is first + np.testing.assert_array_equal(np.asarray(first), ref.numpy()) + + def test_project_works_on_torch_tensors(self) -> None: + # dpmodel project_* are array-API: must accept torch tensors at runtime + lmax, mmax = 2, 1 + rng = np.random.default_rng(2028) + nfull = dp_indexing.get_so3_dim_of_lmax(4) + d_full_np = rng.normal(size=(5, nfull, nfull)) + # CPU on purpose: pins the dp class's torch-namespace behavior + d_full_pt = torch.from_numpy(d_full_np) + idx_np = dp_indexing.build_m_major_index(lmax, mmax) + ebed = dp_indexing.get_so3_dim_of_lmax(lmax) + res = dp_indexing.project_D_to_m(d_full_pt, idx_np, ebed, None, lmax, mmax) + assert isinstance(res, torch.Tensor) + ref = dp_indexing.project_D_to_m(d_full_np, idx_np, ebed, None, lmax, mmax) + np.testing.assert_array_equal(res.numpy(), np.asarray(ref)) + rest = dp_indexing.project_Dt_from_m(d_full_pt, idx_np, ebed, None, lmax, mmax) + assert isinstance(rest, torch.Tensor) + reft = dp_indexing.project_Dt_from_m(d_full_np, idx_np, ebed, None, lmax, mmax) + np.testing.assert_array_equal(rest.numpy(), np.asarray(reft)) + + +class TestUtilsParity: + @pytest.mark.parametrize("dtype", ["float64", "float32"]) # input precision + def test_safe_norm(self, dtype) -> None: + rng = np.random.default_rng(1234) + x = rng.normal(size=(8, 3)).astype(getattr(np, dtype)) + # include zero vectors: exercises the eps regularization path + x[2, :] = 0.0 + x[5, :] = 0.0 + res = dp_utils.safe_norm(x) + ref = pt_utils.safe_norm(to_pt(x)) + assert res.shape == (8, 1) + if dtype == "float64": + # fp64: ~1 ulp on CPU; device-conditional gate on CUDA + np.testing.assert_allclose( + res, ref.cpu().numpy(), rtol=1e-15 if _ON_CPU else PT_RTOL, atol=0.0 + ) + else: + # fp32: numpy and torch may differ by ~1 ulp depending on the + # runner's BLAS/SIMD codegen; CUDA fp32 kernels diverge further + # from CPU numpy, so widen the gate there only. + fp32_tol = 2e-7 if _ON_CPU else 1e-5 + np.testing.assert_allclose( + res, ref.cpu().numpy(), rtol=fp32_tol, atol=fp32_tol + ) + + def test_safe_norm_all_zero(self) -> None: + # pure eps branch: norm of zero vector equals eps exactly + x = np.zeros((4, 3), dtype=np.float64) + res = dp_utils.safe_norm(x, eps=1e-7) + ref = pt_utils.safe_norm(to_pt(x), eps=1e-7) + # pure-eps branch is exact on any device + np.testing.assert_allclose(res, ref.cpu().numpy(), rtol=1e-15, atol=0.0) + np.testing.assert_allclose(np.asarray(res), 1e-7, rtol=1e-15) + + def test_safe_norm_float16_promotion(self) -> None: + # fp16 input: both implementations compute in fp32, cast back to fp16 + rng = np.random.default_rng(4321) + x = rng.normal(size=(8, 3)).astype(np.float16) + x[3, :] = 0.0 + res = dp_utils.safe_norm(x) + ref = pt_utils.safe_norm(to_pt(x)) + assert np.asarray(res).dtype == np.float16 + # the internal fp32 math may differ by ~1 ulp across runners, which + # can flip the final fp16 rounding; compare at ~1 ulp fp16 instead + # of bit-exact equality. 1e-3 is already ulp-of-fp16, so it is + # device-tolerant (CPU and CUDA alike). + np.testing.assert_allclose( + np.asarray(res), ref.cpu().numpy(), rtol=1e-3, atol=1e-3 + ) + + def test_safe_norm_torch_input(self) -> None: + # dpmodel safe_norm is array-API: must accept torch tensors + rng = np.random.default_rng(999) + x = rng.normal(size=(8, 3)) + x[0, :] = 0.0 + # CPU on purpose: pins the dp function's torch-namespace behavior; + # both sides see identical CPU tensors, so the compare stays exact. + res = dp_utils.safe_norm(torch.from_numpy(x)) + assert isinstance(res, torch.Tensor) + ref = pt_utils.safe_norm(torch.from_numpy(x)) + np.testing.assert_allclose(res.numpy(), ref.numpy(), rtol=1e-15, atol=0.0) + + def test_attn_res_modes(self) -> None: + assert dp_utils.ATTN_RES_MODES == pt_utils.ATTN_RES_MODES + + @pytest.mark.parametrize( + "in_dtype,out_dtype", + [ + (np.float16, np.float32), # promoted branch + (np.float32, np.float32), # unchanged branch + (np.float64, np.float64), # unchanged branch + ], + ) # (input dtype, expected promoted dtype) + def test_get_promoted_dtype(self, in_dtype, out_dtype) -> None: + assert np.dtype(dp_utils.get_promoted_dtype(np.dtype(in_dtype))) == np.dtype( + out_dtype + ) + + def test_get_promoted_dtype_bfloat16(self) -> None: + ml_dtypes = pytest.importorskip("ml_dtypes") + assert np.dtype( + dp_utils.get_promoted_dtype(np.dtype(ml_dtypes.bfloat16)) + ) == np.dtype(np.float32) + + def test_init_trunc_normal_fan_in_out(self) -> None: + fan_out, fan_in = 256, 128 + w = np.empty((fan_out, fan_in), dtype=np.float64) + dp_utils.init_trunc_normal_fan_in_out(w, seed=7) + std = 1.0 / np.sqrt(fan_in + fan_out) + # truncation bound respected + assert np.abs(w).max() <= 3.0 * std + # statistics close to the (truncated) normal target + assert abs(w.mean()) < 5.0 * std / np.sqrt(w.size) + assert 0.8 * std < w.std() < 1.05 * std + # reproducible for identical seed + w2 = np.empty_like(w) + dp_utils.init_trunc_normal_fan_in_out(w2, seed=7) + np.testing.assert_array_equal(w, w2) + # scale parameter rescales std + w3 = np.empty_like(w) + dp_utils.init_trunc_normal_fan_in_out(w3, seed=7, scale=2.0) + assert np.abs(w3).max() <= 6.0 * std + assert w3.std() > w.std() + + def test_init_trunc_normal_fan_in_out_errors(self) -> None: + with pytest.raises(ValueError): + dp_utils.init_trunc_normal_fan_in_out( + np.empty((2, 3, 4), dtype=np.float64), seed=0 + ) + with pytest.raises(ValueError): + dp_utils.init_trunc_normal_fan_in_out( + np.empty((2, 3), dtype=np.float64), seed=0, scale=0.0 + ) + + +class TestRadialParity: + rcut = 6.0 + + def _r_grid(self) -> np.ndarray: + # r=0 (sinc/envelope zero-distance branch), r=rcut (envelope boundary), + # r>rcut (envelope-zero branch), plus a dense inside/outside sweep + return np.concatenate( + [[0.0, self.rcut, self.rcut + 0.5], np.linspace(0.05, 6.5, 200)] + )[:, None] + + @pytest.mark.parametrize("exponent", [5, 7]) # envelope polynomial exponent + def test_envelope(self, exponent) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as DPEnvelope, + ) + from deepmd.pt.model.descriptor.sezm_nn.radial import ( + C3CutoffEnvelope as PTEnvelope, + ) + + pt_mod = PTEnvelope(rcut=self.rcut, exponent=exponent, dtype=torch.float64) + dp_mod = DPEnvelope(rcut=self.rcut, exponent=exponent, precision="float64") + r = self._r_grid() + res = dp_mod.call(r) + assert_parity(res, pt_mod(to_pt(r))) + # boundary contract: E(0)=1, E(r>=rcut)=0 exactly + np.testing.assert_array_equal(np.asarray(res)[0], 1.0) + np.testing.assert_array_equal(np.asarray(res)[r[:, 0] >= self.rcut], 0.0) + + def test_envelope_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as DPEnvelope, + ) + + dp_mod = DPEnvelope(rcut=self.rcut, exponent=5, precision="float64") + dp_mod2 = DPEnvelope.deserialize(dp_mod.serialize()) + r = self._r_grid() + np.testing.assert_array_equal( + np.asarray(dp_mod.call(r)), np.asarray(dp_mod2.call(r)) + ) + + @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"]) # both bases + @pytest.mark.parametrize("exponent", [5, 7]) # envelope exponent + def test_radial_basis(self, basis_type, exponent) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + RadialBasis as DPRadialBasis, + ) + from deepmd.pt.model.descriptor.sezm_nn.radial import ( + RadialBasis as PTRadialBasis, + ) + + n_radial = 16 + pt_mod = PTRadialBasis( + rcut=self.rcut, + basis_type=basis_type, + n_radial=n_radial, + dtype=torch.float64, + exponent=exponent, + ) + # perturb the trained frequencies so parity exercises copied weights, + # not just identical deterministic init + rng = np.random.default_rng(2030) + with torch.no_grad(): + pt_mod.adam_freqs += to_pt(0.05 * rng.normal(size=(1, n_radial))) + serialized = pt_mod.serialize() + # pt state_dict key contract: only the trainable frequencies + assert list(serialized["@variables"]) == ["adam_freqs"] + dp_mod = DPRadialBasis.deserialize(serialized) + r = self._r_grid() + assert_parity(dp_mod.call(r), pt_mod(to_pt(r))) + + @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"]) # both bases + def test_radial_basis_roundtrip(self, basis_type) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + RadialBasis as DPRadialBasis, + ) + + dp_mod = DPRadialBasis( + rcut=self.rcut, + basis_type=basis_type, + n_radial=12, + precision="float64", + exponent=7, + ) + dp_mod2 = DPRadialBasis.deserialize(dp_mod.serialize()) + r = self._r_grid() + np.testing.assert_array_equal( + np.asarray(dp_mod.call(r)), np.asarray(dp_mod2.call(r)) + ) + + @pytest.mark.parametrize( + "mlp_layers", + [[16, 32, 24], [16, 24]], + ) # with hidden layers (Linear+RMSNorm+act) and pure-linear (no hidden) branch + @pytest.mark.parametrize("activation", ["silu", "tanh"]) # activation mapping + def test_radial_mlp(self, mlp_layers, activation) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as DPRadialMLP + from deepmd.pt.model.descriptor.sezm_nn.radial import RadialMLP as PTRadialMLP + + pt_mod = PTRadialMLP( + mlp_layers, + activation_function=activation, + dtype=torch.float64, + seed=11, + ) + # perturb all parameters (RMSNorm scale inits to ones, which would + # otherwise make the scale copy untested) + rng = np.random.default_rng(2031) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + serialized = pt_mod.serialize() + # pt state_dict key contract: Sequential index 3*i for linear `matrix`, + # 3*i+1 for RMSNorm `adam_scale` (activation modules are parameter-free) + n_lin = len(mlp_layers) - 1 + expected_keys = {f"{3 * i}.matrix" for i in range(n_lin)} | { + f"{3 * i + 1}.adam_scale" for i in range(n_lin - 1) + } + assert set(serialized["@variables"]) == expected_keys + dp_mod = DPRadialMLP.deserialize(serialized) + x = rng.normal(size=(50, mlp_layers[0])) + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + + def test_radial_mlp_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as DPRadialMLP + + dp_mod = DPRadialMLP( + [16, 32, 24], + activation_function="silu", + precision="float64", + seed=5, + ) + dp_mod2 = DPRadialMLP.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2032) + x = rng.normal(size=(50, 16)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_radial_mlp_zero_input_is_zero(self) -> None: + # bias-free design contract: RadialMLP(0) = 0 + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as DPRadialMLP + + dp_mod = DPRadialMLP([8, 16, 4], precision="float64", seed=3) + out = dp_mod.call(np.zeros((5, 8), dtype=np.float64)) + np.testing.assert_array_equal(np.asarray(out), 0.0) + + def test_radial_mlp_unsupported_activation(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as DPRadialMLP + + dp_mod = DPRadialMLP([4, 8, 4], activation_function="nope", seed=0) + with pytest.raises(NotImplementedError): + dp_mod.call(np.zeros((2, 4), dtype=np.float64)) + + def test_rmsnorm_parity_and_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import RMSNorm as DPRMSNorm + from deepmd.pt.model.descriptor.sezm_nn.norm import RMSNorm as PTRMSNorm + + channels = 24 + pt_mod = PTRMSNorm(channels=channels, dtype=torch.float64, trainable=True) + rng = np.random.default_rng(2033) + with torch.no_grad(): + pt_mod.adam_scale += to_pt(0.1 * rng.normal(size=(channels,))) + serialized = pt_mod.serialize() + assert list(serialized["@variables"]) == ["adam_scale"] + dp_mod = DPRMSNorm.deserialize(serialized) + x64 = rng.normal(size=(50, channels)) + assert_parity(dp_mod.call(x64), pt_mod(to_pt(x64))) + # input-dtype promotion branch: fp32 input with fp64 params, + # output cast back to fp32 in both implementations + x32 = x64.astype(np.float32) + res32 = dp_mod.call(x32) + ref32 = pt_mod(to_pt(x32)) + assert np.asarray(res32).dtype == np.float32 + if _ON_CPU: + # identical CPU fp32 truncation points: bit-exact + np.testing.assert_array_equal( + np.asarray(res32), ref32.detach().cpu().numpy() + ) + else: + # CUDA fp32 kernels differ from CPU numpy at ulp level + np.testing.assert_allclose( + np.asarray(res32), ref32.detach().cpu().numpy(), rtol=1e-5, atol=1e-6 + ) + # serialize roundtrip is exact + dp_mod2 = DPRMSNorm.deserialize(dp_mod.serialize()) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x64)), np.asarray(dp_mod2.call(x64)) + ) + + def test_constructor_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as DPEnvelope, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + RadialBasis as DPRadialBasis, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as DPRadialMLP + + with pytest.raises(ValueError): # rcut <= 0 + DPEnvelope(rcut=0.0) + with pytest.raises(ValueError): # exponent <= 0 + DPEnvelope(rcut=6.0, exponent=0) + with pytest.raises(ValueError): # rcut <= 0 + DPRadialBasis(rcut=-1.0) + with pytest.raises(ValueError): # n_radial <= 0 + DPRadialBasis(rcut=6.0, n_radial=0) + with pytest.raises(ValueError): # unknown basis_type + DPRadialBasis(rcut=6.0, basis_type="chebyshev") + with pytest.raises(ValueError): # mlp_layers too short + DPRadialMLP([16]) + + def test_deserialize_wrong_class(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import RMSNorm as DPRMSNorm + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as DPEnvelope, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + RadialBasis as DPRadialBasis, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as DPRadialMLP + + for klass in (DPEnvelope, DPRadialBasis, DPRadialMLP, DPRMSNorm): + with pytest.raises(ValueError): + klass.deserialize({"@class": "Nope", "@version": 1}) + + +def _make_edge_vectors() -> np.ndarray: + """Random edge vectors plus the polar/eps corner cases of the quaternion charts.""" + rng = np.random.default_rng(1) + vec = rng.standard_normal((128, 3)) + vec[0] = [0.0, 0.0, 1.0] # +z axis (rb_small branch in the Wigner path) + vec[1] = [0.0, 0.0, -1.0] # -z axis (antiparallel pole, ra_small branch) + vec[2] = [1e-9, 0.0, 1.0] # near +z (eps branch of the +z chart) + vec[3] = [0.0, 1.0, 0.0] # +y (e3nn polar axis) + vec[4] = [0.0, -1.0, 0.0] # -y + vec[5] = [1e-9, 0.0, -1.0] # near -z (eps branch of the -z chart) + vec[6] = [0.0, 0.0, 0.0] # zero-length edge (eps-floored normalization) + vec[7] = [0.3, -0.4, 0.0] # equator (chart blend midpoint region) + return vec + + +class TestWignerDParity: + def test_build_edge_quaternion(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + build_edge_quaternion as dp_build_edge_quaternion, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + build_edge_quaternion as pt_build_edge_quaternion, + ) + + vec = _make_edge_vectors() + vec_t = torch.tensor(vec, dtype=torch.float64, device=PT_DEVICE) + # edge_len omitted branch + quat_dp = dp_build_edge_quaternion(vec) + quat_pt = pt_build_edge_quaternion(vec_t) + assert_parity(quat_dp, quat_pt) + # edge_len provided branch + edge_len = np.linalg.norm(vec, axis=-1, keepdims=True) + quat_dp = dp_build_edge_quaternion(vec, edge_len=edge_len) + quat_pt = pt_build_edge_quaternion( + vec_t, + edge_len=torch.tensor(edge_len, dtype=torch.float64, device=PT_DEVICE), + ) + assert_parity(quat_dp, quat_pt) + # the quaternion rotates the unit edge direction onto local +z + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + quaternion_to_rotation_matrix as dp_quaternion_to_rotation_matrix, + ) + + rot = dp_quaternion_to_rotation_matrix(quat_dp) + unit = vec / np.sqrt(np.sum(vec * vec, axis=-1, keepdims=True) + 1e-14) + local = np.einsum("eij,ej->ei", rot, unit) + scale = np.linalg.norm(unit, axis=-1) # ~0 for the zero-length edge row + np.testing.assert_allclose(local[:, 2], scale, atol=1e-10) + np.testing.assert_allclose(local[:, :2], 0.0, atol=1e-10) + + def test_quaternion_helpers(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn import wignerd as dp_w + from deepmd.pt.model.descriptor.sezm_nn import wignerd as pt_w + + rng = np.random.default_rng(2) + q1 = rng.standard_normal((16, 4)) + q2 = rng.standard_normal((16, 4)) + gamma = rng.standard_normal((16,)) + weight = rng.uniform(0.0, 1.0, (16,)) + q1_t = torch.tensor(q1, dtype=torch.float64, device=PT_DEVICE) + q2_t = torch.tensor(q2, dtype=torch.float64, device=PT_DEVICE) + assert_parity( + dp_w.quaternion_multiply(q1, q2), pt_w.quaternion_multiply(q1_t, q2_t) + ) + assert_parity( + dp_w.quaternion_z_rotation(gamma), + pt_w.quaternion_z_rotation( + torch.tensor(gamma, dtype=torch.float64, device=PT_DEVICE) + ), + ) + assert_parity(dp_w.quaternion_normalize(q1), pt_w.quaternion_normalize(q1_t)) + assert_parity( + dp_w.quaternion_to_rotation_matrix(dp_w.quaternion_normalize(q1)), + pt_w.quaternion_to_rotation_matrix(pt_w.quaternion_normalize(q1_t)), + ) + assert_parity( + dp_w.quaternion_nlerp(q1, q2, weight), + pt_w.quaternion_nlerp( + q1_t, q2_t, torch.tensor(weight, dtype=torch.float64, device=PT_DEVICE) + ), + ) + + @pytest.mark.parametrize("lmax", [0, 1, 2, 3, 4]) # degree range incl. beyond-core + def test_quat_and_d(self, lmax) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as DPWignerDCalculator, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + build_edge_quaternion as dp_build_edge_quaternion, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator as PTWignerDCalculator, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + build_edge_quaternion as pt_build_edge_quaternion, + ) + + vec = _make_edge_vectors() + quat_dp = dp_build_edge_quaternion(vec) + quat_pt = pt_build_edge_quaternion( + torch.tensor(vec, dtype=torch.float64, device=PT_DEVICE) + ) + assert_parity(quat_dp, quat_pt) + + calc_dp = DPWignerDCalculator(lmax, precision="float64") + calc_pt = PTWignerDCalculator(lmax, dtype=torch.float64) + D_dp, Dt_dp = calc_dp(quat_dp) + D_pt, Dt_pt = calc_pt(quat_pt) + dim = (lmax + 1) ** 2 + assert D_dp.shape == (vec.shape[0], dim, dim) + assert_parity(D_dp, D_pt) + assert_parity(Dt_dp, Dt_pt) + # rotation property: D @ Dt == I + eye = np.broadcast_to(np.eye(dim), D_dp.shape) + np.testing.assert_allclose(D_dp @ Dt_dp, eye, atol=1e-11) + + @pytest.mark.parametrize("lmax", [2, 4]) # calculator degree + @pytest.mark.parametrize("lmin", [1, 2, 3, 4, 5]) # zonal start degree + def test_forward_zonal(self, lmax, lmin) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as DPWignerDCalculator, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + build_edge_quaternion as dp_build_edge_quaternion, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator as PTWignerDCalculator, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + build_edge_quaternion as pt_build_edge_quaternion, + ) + + vec = _make_edge_vectors() + quat_dp = dp_build_edge_quaternion(vec) + quat_pt = pt_build_edge_quaternion( + torch.tensor(vec, dtype=torch.float64, device=PT_DEVICE) + ) + calc_dp = DPWignerDCalculator(lmax, precision="float64") + calc_pt = PTWignerDCalculator(lmax, dtype=torch.float64) + z_dp = calc_dp.forward_zonal(quat_dp, lmin=lmin) + z_pt = calc_pt.forward_zonal(quat_pt, lmin=lmin) + n_expected = max((lmax + 1) ** 2 - lmin * lmin, 0) + assert z_dp.shape == (vec.shape[0], n_expected) + assert tuple(z_pt.shape) == (vec.shape[0], n_expected) + assert_parity(z_dp, z_pt) + + def test_call_works_on_torch_tensors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as DPWignerDCalculator, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + build_edge_quaternion as dp_build_edge_quaternion, + ) + + vec = _make_edge_vectors() + quat_np = dp_build_edge_quaternion(vec) + # CPU on purpose: pins the dp function's torch-namespace behavior + quat_t = dp_build_edge_quaternion( + torch.tensor(vec, dtype=torch.float64, device=CPU) + ) + assert isinstance(quat_t, torch.Tensor) + assert_parity(quat_np, quat_t) + calc_dp = DPWignerDCalculator(3, precision="float64") + D_np, Dt_np = calc_dp(quat_np) + D_t, Dt_t = calc_dp(quat_t) + assert isinstance(D_t, torch.Tensor) + assert_parity(D_np, D_t) + assert_parity(Dt_np, Dt_t) + z_np = calc_dp.forward_zonal(quat_np, lmin=2) + z_t = calc_dp.forward_zonal(quat_t, lmin=2) + assert_parity(z_np, z_t) + + def test_serialize(self) -> None: + # pt WignerDCalculator has buffers, but they are all derived constants: + # its serialize() emits only {"@class", "@version"} and deserialize() + # is delegated to the parent (raises NotImplementedError). The dpmodel + # port mirrors that contract exactly; no @variables roundtrip exists. + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as DPWignerDCalculator, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator as PTWignerDCalculator, + ) + + calc_dp = DPWignerDCalculator(2, precision="float64") + calc_pt = PTWignerDCalculator(2, dtype=torch.float64) + assert calc_dp.serialize() == calc_pt.serialize() + with pytest.raises(NotImplementedError): + DPWignerDCalculator.deserialize(calc_dp.serialize()) + with pytest.raises(ValueError): + DPWignerDCalculator.deserialize({"@class": "Nope", "@version": 1}) + + def test_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as DPWignerDCalculator, + ) + + with pytest.raises(ValueError): # negative lmax + DPWignerDCalculator(-1, precision="float64") + calc = DPWignerDCalculator(2, precision="float64") + with pytest.raises(ValueError): # lmin < 1 + calc.forward_zonal(np.zeros((4, 4)), lmin=0) + + +class TestNormParity: + channels = 8 + + def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: + # perturb all parameters (scales init to ones / biases to zeros, which + # would otherwise make the parameter copy untested) + rng = np.random.default_rng(seed) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + + @pytest.mark.parametrize("lmax", [0, 2, 3]) # 0 covers the scalar-only branch + @pytest.mark.parametrize("n_focus", [1, 2]) # focus streams + def test_equivariant_rmsnorm(self, lmax, n_focus) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + EquivariantRMSNorm as DPEquivariantRMSNorm, + ) + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm as PTEquivariantRMSNorm, + ) + + pt_mod = PTEquivariantRMSNorm( + lmax, self.channels, n_focus, dtype=torch.float64, trainable=True + ) + self._perturb(pt_mod, 2040) + serialized = pt_mod.serialize() + # pt state_dict key contract: 2 parameters + 2 persistent buffers + assert set(serialized["@variables"]) == { + "adam_scale", + "bias", + "expand_index", + "balance_weight", + } + dp_mod = DPEquivariantRMSNorm.deserialize(serialized) + rng = np.random.default_rng(2041) + x = rng.normal(size=(17, (lmax + 1) ** 2, n_focus, self.channels)) + x[0] = 0.0 # all-zeros row exercises the eps path + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + + def test_equivariant_rmsnorm_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + EquivariantRMSNorm as DPEquivariantRMSNorm, + ) + + dp_mod = DPEquivariantRMSNorm(2, self.channels, 2, precision="float64") + dp_mod2 = DPEquivariantRMSNorm.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2042) + x = rng.normal(size=(17, 9, 2, self.channels)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + @pytest.mark.parametrize( + "lmax,mmax", [(0, 0), (2, 1), (2, 2), (3, 2)] + ) # (0,0) covers the scalar-only branch; mmax None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ReducedEquivariantRMSNorm as DPReducedEquivariantRMSNorm, + ) + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + ReducedEquivariantRMSNorm as PTReducedEquivariantRMSNorm, + ) + + degree_index_m = dp_indexing.build_m_major_l_index(lmax, mmax) + pt_mod = PTReducedEquivariantRMSNorm( + lmax=lmax, + mmax=mmax, + channels=self.channels, + degree_index_m=torch.tensor( + degree_index_m, dtype=torch.long, device=PT_DEVICE + ), + n_focus=n_focus, + dtype=torch.float64, + trainable=True, + ) + self._perturb(pt_mod, 2043) + serialized = pt_mod.serialize() + # pt state_dict key contract: 2 parameters + 2 persistent buffers + assert set(serialized["@variables"]) == { + "degree_index_m", + "balance_weight", + "adam_scale", + "bias0", + } + dp_mod = DPReducedEquivariantRMSNorm.deserialize(serialized) + rng = np.random.default_rng(2044) + x = rng.normal(size=(17, n_focus, degree_index_m.size, self.channels)) + x[0] = 0.0 # all-zeros row exercises the eps path + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + + def test_reduced_equivariant_rmsnorm_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ReducedEquivariantRMSNorm as DPReducedEquivariantRMSNorm, + ) + + degree_index_m = dp_indexing.build_m_major_l_index(2, 1) + dp_mod = DPReducedEquivariantRMSNorm( + lmax=2, + mmax=1, + channels=self.channels, + degree_index_m=degree_index_m, + n_focus=2, + precision="float64", + ) + dp_mod2 = DPReducedEquivariantRMSNorm.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2045) + x = rng.normal(size=(17, 2, degree_index_m.size, self.channels)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_reduced_equivariant_rmsnorm_invalid_degree_index(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ReducedEquivariantRMSNorm as DPReducedEquivariantRMSNorm, + ) + + with pytest.raises(ValueError): # degree 5 > lmax leaves zero weights + DPReducedEquivariantRMSNorm( + lmax=2, + mmax=1, + channels=4, + degree_index_m=np.array([0, 1, 5], dtype=np.int64), + precision="float64", + ) + + @pytest.mark.parametrize("mmax", [-1, 3]) # below 0 / above lmax + def test_reduced_equivariant_rmsnorm_invalid_mmax(self, mmax) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ReducedEquivariantRMSNorm as DPReducedEquivariantRMSNorm, + ) + + with pytest.raises(ValueError, match="mmax"): + DPReducedEquivariantRMSNorm( + lmax=2, + mmax=mmax, + channels=4, + degree_index_m=np.array([0, 1, 2], dtype=np.int64), + precision="float64", + ) + + @pytest.mark.parametrize("ndim", [2, 3]) # (B, C) and (B, F, C) branches + def test_scalar_rmsnorm(self, ndim) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ScalarRMSNorm as DPScalarRMSNorm, + ) + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + ScalarRMSNorm as PTScalarRMSNorm, + ) + + n_focus = 1 if ndim == 2 else 2 + pt_mod = PTScalarRMSNorm( + channels=self.channels, + n_focus=n_focus, + dtype=torch.float64, + trainable=True, + ) + self._perturb(pt_mod, 2046) + serialized = pt_mod.serialize() + assert list(serialized["@variables"]) == ["adam_scale"] + dp_mod = DPScalarRMSNorm.deserialize(serialized) + rng = np.random.default_rng(2047) + shape = (17, self.channels) if ndim == 2 else (17, n_focus, self.channels) + x = rng.normal(size=shape) + x[0] = 0.0 # all-zeros row exercises the eps path + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + # serialize roundtrip is exact + dp_mod2 = DPScalarRMSNorm.deserialize(dp_mod.serialize()) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_norm_fp32_input_branch(self) -> None: + # input-dtype promotion branch: fp32 input with fp64 params, output is + # cast back to fp32. Compared at a few ulp fp32: truncation/downcast + # points may differ across BLAS/SIMD codegen and environments, so + # bit-exact equality would be brittle. + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + EquivariantRMSNorm as DPEquivariantRMSNorm, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ScalarRMSNorm as DPScalarRMSNorm, + ) + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm as PTEquivariantRMSNorm, + ) + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + ScalarRMSNorm as PTScalarRMSNorm, + ) + + rng = np.random.default_rng(2048) + pt_eq = PTEquivariantRMSNorm( + 2, self.channels, 1, dtype=torch.float64, trainable=True + ) + dp_eq = DPEquivariantRMSNorm.deserialize(pt_eq.serialize()) + x32 = rng.normal(size=(17, 9, 1, self.channels)).astype(np.float32) + res = dp_eq.call(x32) + ref = pt_eq(to_pt(x32)) + assert np.asarray(res).dtype == np.float32 + # fp32: a few ulp on CPU; CUDA fp32 kernels diverge further from + # CPU numpy, so widen under CUDA only + fp32_tol = 2e-7 if _ON_CPU else 1e-5 + np.testing.assert_allclose( + np.asarray(res), ref.detach().cpu().numpy(), rtol=fp32_tol, atol=fp32_tol + ) + + pt_sc = PTScalarRMSNorm( + channels=self.channels, dtype=torch.float64, trainable=True + ) + dp_sc = DPScalarRMSNorm.deserialize(pt_sc.serialize()) + x32 = rng.normal(size=(17, self.channels)).astype(np.float32) + res = dp_sc.call(x32) + ref = pt_sc(to_pt(x32)) + assert np.asarray(res).dtype == np.float32 + np.testing.assert_allclose( + np.asarray(res), ref.detach().cpu().numpy(), rtol=fp32_tol, atol=fp32_tol + ) + + def test_deserialize_wrong_class(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + EquivariantRMSNorm as DPEquivariantRMSNorm, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ReducedEquivariantRMSNorm as DPReducedEquivariantRMSNorm, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( + ScalarRMSNorm as DPScalarRMSNorm, + ) + + for klass in ( + DPEquivariantRMSNorm, + DPReducedEquivariantRMSNorm, + DPScalarRMSNorm, + ): + with pytest.raises(ValueError): + klass.deserialize({"@class": "Nope", "@version": 1}) + + +class TestSO3LinearParity: + in_channels = 8 + out_channels = 6 + + @pytest.mark.parametrize("lmax", [0, 2, 3]) # 0 covers the scalar-only branch + @pytest.mark.parametrize("mlp_bias", [False, True]) # l=0 bias branch + @pytest.mark.parametrize("n_focus", [1, 2]) # focus streams + def test_so3_linear(self, lmax, mlp_bias, n_focus) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import SO3Linear as DPSO3Linear + from deepmd.pt.model.descriptor.sezm_nn.so3 import SO3Linear as PTSO3Linear + + pt_mod = PTSO3Linear( + lmax=lmax, + in_channels=self.in_channels, + out_channels=self.out_channels, + n_focus=n_focus, + dtype=torch.float64, + mlp_bias=mlp_bias, + trainable=True, + seed=21, + ) + # bias inits to zeros; perturb so the bias copy is exercised + rng = np.random.default_rng(2050) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + serialized = pt_mod.serialize() + expected_keys = {"weight", "expand_index"} | ({"bias"} if mlp_bias else set()) + assert set(serialized["@variables"]) == expected_keys + dp_mod = DPSO3Linear.deserialize(serialized) + x = rng.normal(size=(17, (lmax + 1) ** 2, n_focus, self.in_channels)) + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + + @pytest.mark.parametrize("mlp_bias", [False, True]) # l=0 bias branch + def test_so3_linear_roundtrip(self, mlp_bias) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import SO3Linear as DPSO3Linear + + dp_mod = DPSO3Linear( + lmax=2, + in_channels=self.in_channels, + out_channels=self.out_channels, + n_focus=2, + precision="float64", + mlp_bias=mlp_bias, + seed=7, + ) + dp_mod2 = DPSO3Linear.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2051) + x = rng.normal(size=(17, 9, 2, self.in_channels)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_so3_linear_init_std_branches(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import SO3Linear as DPSO3Linear + + # init_std=0.0 -> exact zero init + dp_zero = DPSO3Linear( + lmax=2, + in_channels=self.in_channels, + out_channels=self.out_channels, + precision="float64", + init_std=0.0, + ) + np.testing.assert_array_equal(dp_zero.weight, 0.0) + # init_std>0 -> normal init (nonzero) + dp_norm = DPSO3Linear( + lmax=2, + in_channels=self.in_channels, + out_channels=self.out_channels, + precision="float64", + seed=3, + init_std=0.5, + ) + assert np.any(dp_norm.weight != 0.0) + + @pytest.mark.parametrize("bias", [False, True]) # bias branch + def test_focus_linear(self, bias) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import FocusLinear as DPFocusLinear + from deepmd.pt.model.descriptor.sezm_nn.so3 import FocusLinear as PTFocusLinear + + n_focus = 2 + pt_mod = PTFocusLinear( + in_channels=self.in_channels, + out_channels=self.out_channels, + n_focus=n_focus, + dtype=torch.float64, + bias=bias, + trainable=True, + seed=5, + ) + rng = np.random.default_rng(2052) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + # pt FocusLinear has no serialize(); copy the state_dict fragment + # (keys "weight"/"bias") directly, the contract used by nested modules + state = pt_state_to_numpy(pt_mod) + assert set(state) == ({"weight", "bias"} if bias else {"weight"}) + dp_mod = DPFocusLinear( + in_channels=self.in_channels, + out_channels=self.out_channels, + n_focus=n_focus, + precision="float64", + bias=bias, + seed=5, + ) + dp_mod.weight = state["weight"] + if bias: + dp_mod.bias = state["bias"] + x = rng.normal(size=(17, n_focus, self.in_channels)) + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + # serialize roundtrip is exact + dp_mod2 = DPFocusLinear.deserialize(dp_mod.serialize()) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + @pytest.mark.parametrize("bias", [False, True]) # bias branch + def test_channel_linear(self, bias) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import ( + ChannelLinear as DPChannelLinear, + ) + from deepmd.pt.model.descriptor.sezm_nn.so3 import ( + ChannelLinear as PTChannelLinear, + ) + + pt_mod = PTChannelLinear( + in_channels=self.in_channels, + out_channels=self.out_channels, + dtype=torch.float64, + bias=bias, + trainable=True, + seed=6, + ) + rng = np.random.default_rng(2053) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + # pt ChannelLinear has no serialize(); copy the state_dict fragment + state = pt_state_to_numpy(pt_mod) + assert set(state) == ({"weight", "bias"} if bias else {"weight"}) + dp_mod = DPChannelLinear( + in_channels=self.in_channels, + out_channels=self.out_channels, + precision="float64", + bias=bias, + seed=6, + ) + dp_mod.weight = state["weight"] + if bias: + dp_mod.bias = state["bias"] + # leading axes are batch: exercise a 3D input + x = rng.normal(size=(17, 4, self.in_channels)) + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + # serialize roundtrip is exact + dp_mod2 = DPChannelLinear.deserialize(dp_mod.serialize()) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_focus_channel_linear_init_std_branch(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import ( + ChannelLinear as DPChannelLinear, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import FocusLinear as DPFocusLinear + + # init_std branch: normal(0, init_std) instead of uniform + dp_f = DPFocusLinear( + in_channels=64, + out_channels=64, + n_focus=1, + precision="float64", + seed=8, + init_std=0.01, + ) + dp_c = DPChannelLinear( + in_channels=64, + out_channels=64, + precision="float64", + seed=8, + init_std=0.01, + ) + for w in (dp_f.weight, dp_c.weight): + # uniform init would have std ~ bound/sqrt(3) = 0.072; normal 0.01 + assert np.std(w) < 0.02 + + def test_deserialize_wrong_class(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import ( + ChannelLinear as DPChannelLinear, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import FocusLinear as DPFocusLinear + from deepmd.dpmodel.descriptor.dpa4_nn.so3 import SO3Linear as DPSO3Linear + + for klass in (DPSO3Linear, DPFocusLinear, DPChannelLinear): + with pytest.raises(ValueError): + klass.deserialize({"@class": "Nope", "@version": 1}) + + +class TestGatedActivationParity: + channels = 8 + + def _build_pair(self, *, lmax, mmax, n_focus, mlp_bias, layout, activation): + from deepmd.dpmodel.descriptor.dpa4_nn.activation import ( + GatedActivation as DPGatedActivation, + ) + from deepmd.pt.model.descriptor.sezm_nn.activation import ( + GatedActivation as PTGatedActivation, + ) + + pt_mod = PTGatedActivation( + lmax=lmax, + mmax=mmax, + channels=self.channels, + n_focus=n_focus, + dtype=torch.float64, + activation_function=activation, + mlp_bias=mlp_bias, + layout=layout, + trainable=True, + seed=31, + ) + # perturb all parameters (gate bias inits to zeros) + rng = np.random.default_rng(2060) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.05 * rng.normal(size=tuple(p.shape))) + serialized = pt_mod.serialize() + expected_keys = {"expand_index"} + if lmax > 0: + expected_keys |= {"gate_linear.weight"} + if mlp_bias: + expected_keys |= {"gate_linear.bias"} + assert set(serialized["@variables"]) == expected_keys + dp_mod = DPGatedActivation.deserialize(serialized) + return dp_mod, pt_mod + + def _shape(self, lmax, mmax, n_focus, layout): + if mmax is None: + ncoeff = (lmax + 1) ** 2 + else: + ncoeff = dp_indexing.build_m_major_l_index(lmax, mmax).size + if layout == "nfdc": + return (17, n_focus, ncoeff, self.channels) + return (17, ncoeff, n_focus, self.channels) + + @pytest.mark.parametrize("layout", ["nfdc", "ndfc"]) # tensor layout + @pytest.mark.parametrize("use_gate", [False, True]) # standard vs GLU mode + def test_gated_activation(self, layout, use_gate) -> None: + lmax, n_focus = 2, 2 + dp_mod, pt_mod = self._build_pair( + lmax=lmax, + mmax=None, + n_focus=n_focus, + mlp_bias=False, + layout=layout, + activation="silu", + ) + rng = np.random.default_rng(2061) + shape = self._shape(lmax, None, n_focus, layout) + x = rng.normal(size=shape) + if use_gate: + gate = rng.normal(size=shape) + res = dp_mod.call(x, gate=gate) + ref = pt_mod(to_pt(x), gate=to_pt(gate)) + else: + res = dp_mod.call(x) + ref = pt_mod(to_pt(x)) + assert_parity(res, ref) + + @pytest.mark.parametrize("mlp_bias", [False, True]) # gate-linear bias branch + def test_gated_activation_mmax_reduced(self, mlp_bias) -> None: + # m-major reduced layout branch (mmax provided) + tanh activation + lmax, mmax, n_focus = 3, 1, 1 + dp_mod, pt_mod = self._build_pair( + lmax=lmax, + mmax=mmax, + n_focus=n_focus, + mlp_bias=mlp_bias, + layout="ndfc", + activation="tanh", + ) + rng = np.random.default_rng(2062) + x = rng.normal(size=self._shape(lmax, mmax, n_focus, "ndfc")) + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + + @pytest.mark.parametrize("use_gate", [False, True]) # standard vs GLU mode + def test_gated_activation_lmax0(self, use_gate) -> None: + # lmax=0 branch: scalar-only output, no gate_linear + dp_mod, pt_mod = self._build_pair( + lmax=0, + mmax=None, + n_focus=1, + mlp_bias=False, + layout="nfdc", + activation="silu", + ) + assert dp_mod.gate_linear is None + rng = np.random.default_rng(2063) + shape = self._shape(0, None, 1, "nfdc") + x = rng.normal(size=shape) + if use_gate: + gate = rng.normal(size=shape) + res = dp_mod.call(x, gate=gate) + ref = pt_mod(to_pt(x), gate=to_pt(gate)) + else: + res = dp_mod.call(x) + ref = pt_mod(to_pt(x)) + assert_parity(res, ref) + + def test_gated_activation_fp32_input_branch(self) -> None: + # input-dtype promotion branch: fp32 input with fp64 gate params; + # downcast happens at different points in the two implementations, + # and truncation points vary across BLAS/SIMD codegen, so compare + # with fp32 round-off headroom rather than a single-machine ulp. + dp_mod, pt_mod = self._build_pair( + lmax=2, + mmax=None, + n_focus=1, + mlp_bias=False, + layout="nfdc", + activation="silu", + ) + rng = np.random.default_rng(2064) + x32 = rng.normal(size=self._shape(2, None, 1, "nfdc")).astype(np.float32) + res = dp_mod.call(x32) + ref = pt_mod(to_pt(x32)) + assert np.asarray(res).dtype == np.float32 + # fp32 round-off headroom on CPU; wider under CUDA (kernel/codegen + # truncation points differ from CPU numpy) + np.testing.assert_allclose( + np.asarray(res), + ref.detach().cpu().numpy(), + rtol=2e-6 if _ON_CPU else 1e-5, + atol=2e-7 if _ON_CPU else 1e-6, + ) + + def test_gated_activation_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.activation import ( + GatedActivation as DPGatedActivation, + ) + + dp_mod = DPGatedActivation( + lmax=2, + channels=self.channels, + n_focus=2, + precision="float64", + mlp_bias=True, + layout="ndfc", + seed=13, + ) + dp_mod2 = DPGatedActivation.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2065) + x = rng.normal(size=(17, 9, 2, self.channels)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_gated_activation_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.activation import ( + GatedActivation as DPGatedActivation, + ) + + with pytest.raises(ValueError): # mmax < 0 + DPGatedActivation(lmax=2, mmax=-1, channels=4) + with pytest.raises(ValueError): # mmax > lmax + DPGatedActivation(lmax=2, mmax=3, channels=4) + with pytest.raises(ValueError): # invalid layout + DPGatedActivation(lmax=2, channels=4, layout="cfdn") + with pytest.raises(ValueError): # wrong class + DPGatedActivation.deserialize({"@class": "Nope", "@version": 1}) + + def test_swiglu(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.activation import SwiGLU as DPSwiGLU + from deepmd.pt.model.descriptor.sezm_nn.activation import SwiGLU as PTSwiGLU + + rng = np.random.default_rng(2066) + x = rng.normal(size=(17, 3, 2 * self.channels)) + assert_parity(DPSwiGLU().call(x), PTSwiGLU()(to_pt(x))) + + +class TestS2GridParity: + channels = 8 + + # ---------------------------------------------------------------- helpers + def _build_projectors(self, lmax, mmax, coefficient_layout): + from deepmd.dpmodel.descriptor.dpa4_nn.projection import ( + S2GridProjector as DPS2GridProjector, + ) + from deepmd.pt.model.descriptor.sezm_nn.projection import ( + S2GridProjector as PTS2GridProjector, + ) + + pt_proj = PTS2GridProjector( + lmax=lmax, + mmax=mmax, + dtype=torch.float64, + coefficient_layout=coefficient_layout, + grid_method="lebedev", + ) + dp_proj = DPS2GridProjector( + lmax=lmax, + mmax=mmax, + precision="float64", + coefficient_layout=coefficient_layout, + grid_method="lebedev", + ) + return pt_proj, dp_proj + + def _build_grid_nets( + self, + *, + lmax, + op_type, + layout, + mlp_bias=False, + n_focus=1, + mmax=None, + coefficient_layout="packed", + grid_branches=1, + seed=7, + ): + """Build a pt S2GridNet, perturb its params, and copy them into dp.""" + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import S2GridNet as DPS2GridNet + from deepmd.pt.model.descriptor.sezm_nn.grid_net import S2GridNet as PTS2GridNet + + pt_net = PTS2GridNet( + lmax=lmax, + mmax=mmax, + channels=self.channels, + n_focus=n_focus, + mode="self", + op_type=op_type, + dtype=torch.float64, + layout=layout, + coefficient_layout=coefficient_layout, + grid_method="lebedev", + grid_branches=grid_branches, + mlp_bias=mlp_bias, + trainable=True, + seed=seed, + ) + rng = np.random.default_rng(2100) + with torch.no_grad(): + for p in pt_net.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + dp_net = DPS2GridNet( + lmax=lmax, + mmax=mmax, + channels=self.channels, + n_focus=n_focus, + mode="self", + op_type=op_type, + precision="float64", + layout=layout, + coefficient_layout=coefficient_layout, + grid_method="lebedev", + grid_branches=grid_branches, + mlp_bias=mlp_bias, + trainable=True, + seed=seed, + ) + # pt S2GridNet has no serialize(); copy the state_dict fragment with + # the pt key names (the contract used by the dp serialize format) + state = pt_state_to_numpy(pt_net) + expected_keys = {"scalar_gate.weight"} + if mlp_bias: + expected_keys.add("scalar_gate.bias") + if op_type == "branch": + expected_keys |= { + "grid_op.left_proj.weight", + "grid_op.right_proj.weight", + "grid_op.router.weight", + "grid_op.out_proj.weight", + } + assert set(state) == expected_keys + dp_net.scalar_gate.weight = state["scalar_gate.weight"] + if mlp_bias: + dp_net.scalar_gate.bias = state["scalar_gate.bias"] + if op_type == "branch": + for name in ("left_proj", "right_proj", "router", "out_proj"): + getattr(dp_net.grid_op, name).weight = state[f"grid_op.{name}.weight"] + return pt_net, dp_net + + # ------------------------------------------------- (a) projector constants + @pytest.mark.parametrize("lmax,mmax", [(2, 2), (3, 3), (3, 2)]) # degree/order + @pytest.mark.parametrize( + "coefficient_layout", ["packed", "m_major"] + ) # coefficient ordering + def test_projector_constants(self, lmax, mmax, coefficient_layout) -> None: + pt_proj, dp_proj = self._build_projectors(lmax, mmax, coefficient_layout) + assert dp_proj.grid_resolution_list == pt_proj.grid_resolution_list + assert dp_proj.grid_size == pt_proj.grid_size + assert dp_proj.coeff_dim == pt_proj.coeff_dim + assert dp_proj.packed_dim == pt_proj.packed_dim + # validates the numpy-SH replacement of e3nn end-to-end + assert_parity(dp_proj.to_grid_mat, pt_proj.to_grid_mat) + assert_parity(dp_proj.from_grid_mat, pt_proj.from_grid_mat) + # to_grid / from_grid forwards + rng = np.random.default_rng(2080) + x = rng.normal(size=(7, dp_proj.coeff_dim, 5)) + assert_parity(dp_proj.to_grid(x), pt_proj.to_grid(to_pt(x))) + g = rng.normal(size=(7, dp_proj.grid_size, 5)) + assert_parity(dp_proj.from_grid(g), pt_proj.from_grid(to_pt(g))) + + @pytest.mark.parametrize("lmax", [2, 3]) # max degree + def test_projector_quadrature_identity(self, lmax) -> None: + # Lebedev path: from_grid o to_grid is the identity at machine + # precision (full mmax == lmax, packed layout) + from deepmd.dpmodel.descriptor.dpa4_nn.projection import ( + S2GridProjector as DPS2GridProjector, + ) + + dp_proj = DPS2GridProjector( + lmax=lmax, precision="float64", grid_method="lebedev" + ) + prod = np.matmul(dp_proj.from_grid_mat, dp_proj.to_grid_mat) + np.testing.assert_allclose(prod, np.eye(dp_proj.coeff_dim), atol=1e-13) + + def test_projector_serialize(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.projection import ( + S2GridProjector as DPS2GridProjector, + ) + + pt_proj, dp_proj = self._build_projectors(3, 2, "m_major") + # the serialize contracts are identical + assert dp_proj.serialize() == pt_proj.serialize() + # dp deserializes pt's real serialize() output + dp_from_pt = DPS2GridProjector.deserialize(pt_proj.serialize()) + np.testing.assert_array_equal(dp_from_pt.to_grid_mat, dp_proj.to_grid_mat) + np.testing.assert_array_equal(dp_from_pt.from_grid_mat, dp_proj.from_grid_mat) + # dp roundtrip is exact + dp_proj2 = DPS2GridProjector.deserialize(dp_proj.serialize()) + np.testing.assert_array_equal(dp_proj2.to_grid_mat, dp_proj.to_grid_mat) + np.testing.assert_array_equal(dp_proj2.from_grid_mat, dp_proj.from_grid_mat) + with pytest.raises(ValueError): # wrong class + DPS2GridProjector.deserialize({"@class": "Nope", "@version": 1}) + + @pytest.mark.parametrize("method", ["lebedev", "e3nn"]) # quadrature backend + @pytest.mark.parametrize("lmax,mmax", [(2, 2), (3, 2), (4, 4)]) # degree/order + def test_resolve_s2_grid_resolution(self, method, lmax, mmax) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.projection import ( + resolve_s2_grid_resolution as dp_resolve, + ) + from deepmd.pt.model.descriptor.sezm_nn.projection import ( + resolve_s2_grid_resolution as pt_resolve, + ) + + assert dp_resolve(lmax, mmax, method=method) == pt_resolve( + lmax, mmax, method=method + ) + with pytest.raises(ValueError): # invalid method + dp_resolve(lmax, mmax, method="cartesian") + + # ------------------------------------------ (b) S2GridNet forward parity + @pytest.mark.parametrize("lmax", [2, 3]) # max degree + @pytest.mark.parametrize("op_type", ["glu", "branch"]) # grid operation + def test_s2_grid_net(self, lmax, op_type) -> None: + # ffn-style core usage: mode="self", layout="ndfc", packed, n_focus=1 + pt_net, dp_net = self._build_grid_nets( + lmax=lmax, op_type=op_type, layout="ndfc" + ) + rng = np.random.default_rng(2081) + n_coeff = (lmax + 1) ** 2 + x = rng.normal(size=(11, n_coeff, 1, 2 * self.channels)) + assert_parity(dp_net.call(x), pt_net(to_pt(x))) + + @pytest.mark.parametrize("mlp_bias", [False, True]) # scalar gate bias + def test_s2_grid_net_nfdc_m_major(self, mlp_bias) -> None: + # so2-style usage: mode="self", op_type="glu", layout="nfdc", + # m-major coefficients truncated at mmax < lmax, multiple foci + lmax, mmax, n_focus = 3, 2, 2 + pt_net, dp_net = self._build_grid_nets( + lmax=lmax, + mmax=mmax, + op_type="glu", + layout="nfdc", + mlp_bias=mlp_bias, + n_focus=n_focus, + coefficient_layout="m_major", + ) + n_coeff = dp_net.projector.coeff_dim + assert n_coeff == pt_net.projector.coeff_dim + rng = np.random.default_rng(2082) + x = rng.normal(size=(11, n_focus, n_coeff, 2 * self.channels)) + assert_parity(dp_net.call(x), pt_net(to_pt(x))) + + def test_s2_grid_net_fp32_input(self) -> None: + # fp32 input through a float64-precision net exercises the cast + # branches; the output dtype must match the input dtype as in pt + pt_net, dp_net = self._build_grid_nets(lmax=2, op_type="branch", layout="ndfc") + rng = np.random.default_rng(2083) + x = rng.normal(size=(11, 9, 1, 2 * self.channels)).astype(np.float32) + dp_out = dp_net.call(x) + assert dp_out.dtype == np.float32 + pt_out = pt_net(to_pt(x)) + assert pt_out.dtype == torch.float32 + np.testing.assert_allclose( + np.asarray(dp_out), pt_out.detach().cpu().numpy(), rtol=1e-6, atol=1e-6 + ) + + # ------------------------------------------- (c) GridBranch forward parity + @pytest.mark.parametrize("n_branches", [1, 2]) # router branches + def test_grid_branch(self, n_branches) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import ( + GridBranch as DPGridBranch, + ) + from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( + GridBranch as PTGridBranch, + ) + + pt_mod = PTGridBranch( + channels=self.channels, + n_branches=n_branches, + dtype=torch.float64, + trainable=True, + seed=9, + ) + rng = np.random.default_rng(2084) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + state = pt_state_to_numpy(pt_mod) + assert set(state) == { + "left_proj.weight", + "right_proj.weight", + "router.weight", + "out_proj.weight", + } + dp_mod = DPGridBranch( + channels=self.channels, + n_branches=n_branches, + precision="float64", + seed=9, + ) + for name in ("left_proj", "right_proj", "router", "out_proj"): + getattr(dp_mod, name).weight = state[f"{name}.weight"] + n_batch, n_grid, n_focus = 5, 26, 2 + query = rng.normal(size=(n_batch, n_grid, n_focus, self.channels)) + context = rng.normal(size=(n_batch, n_grid, n_focus, self.channels)) + scalar = rng.normal(size=(n_batch, n_focus, 2 * self.channels)) + assert_parity( + dp_mod.call(query, context, scalar), + pt_mod( + to_pt(query), + to_pt(context), + to_pt(scalar), + ), + ) + # serialize roundtrip is exact; @variables keys match the pt state dict + ser = dp_mod.serialize() + assert set(ser["@variables"]) == set(state) + dp_mod2 = DPGridBranch.deserialize(ser) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(query, context, scalar)), + np.asarray(dp_mod2.call(query, context, scalar)), + ) + + # ------------------------------------------------------ (e) serialization + @pytest.mark.parametrize("op_type", ["glu", "branch"]) # grid operation + @pytest.mark.parametrize("mlp_bias", [False, True]) # scalar gate bias + def test_s2_grid_net_serialize_roundtrip(self, op_type, mlp_bias) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import S2GridNet as DPS2GridNet + + pt_net, dp_net = self._build_grid_nets( + lmax=2, op_type=op_type, layout="ndfc", mlp_bias=mlp_bias + ) + ser = dp_net.serialize() + # @variables key set equals the pt state_dict key set exactly + assert set(ser["@variables"]) == set(pt_state_to_numpy(pt_net)) + dp_net2 = DPS2GridNet.deserialize(ser) + rng = np.random.default_rng(2085) + x = rng.normal(size=(11, 9, 1, 2 * self.channels)) + np.testing.assert_array_equal( + np.asarray(dp_net.call(x)), np.asarray(dp_net2.call(x)) + ) + # loading pt's real state_dict values through deserialize also works + ser_pt = dict(ser) + ser_pt["@variables"] = pt_state_to_numpy(pt_net) + dp_net3 = DPS2GridNet.deserialize(ser_pt) + assert_parity(dp_net3.call(x), pt_net(to_pt(x))) + with pytest.raises(ValueError): # wrong class + DPS2GridNet.deserialize({"@class": "Nope", "@version": 1}) + + def test_grid_branch_deserialize_wrong_class(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import ( + GridBranch as DPGridBranch, + ) + + with pytest.raises(ValueError): + DPGridBranch.deserialize({"@class": "Nope", "@version": 1}) + + # ------------------------------------------------ (d) not-ported guards + def test_not_ported_guards(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import S2GridNet as DPS2GridNet + from deepmd.dpmodel.descriptor.dpa4_nn.projection import ( + S2GridProjector as DPS2GridProjector, + ) + + common = { + "lmax": 2, + "channels": 4, + "mode": "self", + "op_type": "glu", + "precision": "float64", + "layout": "ndfc", + "grid_method": "lebedev", + } + with pytest.raises(NotImplementedError, match="lebedev_quadrature"): + # e3nn product grid (lebedev_quadrature=False) is not ported + DPS2GridProjector(lmax=2, precision="float64", grid_method="e3nn") + with pytest.raises(NotImplementedError, match="lebedev_quadrature"): + DPS2GridNet(**{**common, "grid_method": "e3nn"}) + # default grid_method is "lebedev" (deliberate divergence from pt's + # "e3nn" default, which dp rejects): default construction works + net = DPS2GridNet(**{k: v for k, v in common.items() if k != "grid_method"}) + assert net.grid_method == "lebedev" + with pytest.raises(NotImplementedError, match="grid_mlp"): + # GridMLP (grid_mlp=True) is not ported + DPS2GridNet(**{**common, "op_type": "mlp"}) + with pytest.raises(NotImplementedError, match="node_wise_s2"): + # cross mode backs node_wise_s2/message_node_s2 only + DPS2GridNet(**{**common, "mode": "cross"}) + with pytest.raises(NotImplementedError, match="residual_scale_init"): + DPS2GridNet(**common, residual_scale_init=1e-3) + + def test_value_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import ( + GridBranch as DPGridBranch, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import S2GridNet as DPS2GridNet + from deepmd.dpmodel.descriptor.dpa4_nn.projection import ( + S2GridProjector as DPS2GridProjector, + ) + + common = { + "lmax": 2, + "channels": 4, + "mode": "self", + "op_type": "glu", + "precision": "float64", + "layout": "ndfc", + "grid_method": "lebedev", + } + with pytest.raises(ValueError): # unknown grid method + DPS2GridProjector(lmax=2, grid_method="cartesian") + with pytest.raises(ValueError): # negative mmax + DPS2GridProjector(lmax=2, mmax=-1, grid_method="lebedev") + with pytest.raises(ValueError): # mmax > lmax + DPS2GridProjector(lmax=2, mmax=3, grid_method="lebedev") + with pytest.raises(ValueError): # bad coefficient layout + DPS2GridProjector( + lmax=2, grid_method="lebedev", coefficient_layout="l_major" + ) + with pytest.raises(ValueError): # non-packaged [precision, n_points] + DPS2GridProjector( + lmax=2, grid_method="lebedev", grid_resolution_list=[7, 10] + ) + with pytest.raises(ValueError): # wrong resolution list length + DPS2GridProjector(lmax=2, grid_method="lebedev", grid_resolution_list=[7]) + with pytest.raises(ValueError): # unknown mode + DPS2GridNet(**{**common, "mode": "pair"}) + with pytest.raises(ValueError): # unknown op_type + DPS2GridNet(**{**common, "op_type": "attention"}) + with pytest.raises(ValueError): # unknown layout + DPS2GridNet(**{**common, "layout": "cdfn"}) + with pytest.raises(ValueError): # flat layout is cross-only + DPS2GridNet(**{**common, "layout": "flat"}) + with pytest.raises(ValueError): # n_branches must be positive + DPGridBranch(channels=4, n_branches=0, precision="float64") + dp_net = DPS2GridNet(**common) + rng = np.random.default_rng(2086) + with pytest.raises(ValueError): # wrong query channel count + dp_net.call(rng.normal(size=(3, 9, 1, 5))) + + +def _build_so2_edge_data( + rng, + *, + nloc, + nnei, + lmax, + channels, + masked="none", + with_gate=False, + n_radial=None, +): + """Build matching pt (sparse) and dp (padded) edge caches. + + The dp cache uses the padded layout (E = nloc * nnei with ``edge_mask``); + the pt cache keeps only the valid slots (flat sparse edges in the same + row-major slot order pt's ``torch.nonzero`` would produce). Both sides + share identical Wigner-D blocks built from the (parity-proven) dpmodel + ``WignerDCalculator``. Invalid slots intentionally keep garbage (nonzero) + envelope/feature values so a missing mask shows up as a parity failure. + + ``n_radial``: when not None, ``edge_rbf`` is filled with random values of + width ``n_radial`` (garbage in masked slots too); otherwise it stays the + zero (E, 1) placeholder used by consumers that ignore ``edge_rbf``. + + ``masked`` is one of: + - ``"none"``: all slots valid; + - ``"slots"``: a few scattered invalid slots; + - ``"node"``: node 2 fully masked (no incoming edges) plus one extra slot. + """ + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + EdgeCache, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator, + build_edge_quaternion, + ) + from deepmd.pt.model.descriptor.sezm_nn.edge_cache import ( + EdgeFeatureCache, + ) + + n_edge = nloc * nnei + dim_full = (lmax + 1) ** 2 + src = np.array( + [(i + 1 + k) % nloc for i in range(nloc) for k in range(nnei)], + dtype=np.int64, + ) + dst = np.repeat(np.arange(nloc, dtype=np.int64), nnei) + mask = np.ones(n_edge, dtype=np.float64) + if masked == "slots": + mask[3] = 0.0 + mask[nnei + 1] = 0.0 + mask[-1] = 0.0 + elif masked == "node": + mask[2 * nnei : 3 * nnei] = 0.0 # node 2: no incoming edges at all + mask[3] = 0.0 + elif masked != "none": + raise ValueError(f"unknown masked mode {masked}") + valid = mask > 0.5 + n_valid = int(valid.sum()) + + edge_vec = rng.normal(size=(n_edge, 3)) + edge_vec /= np.linalg.norm(edge_vec, axis=-1, keepdims=True) + quat = build_edge_quaternion(edge_vec) + D_full, Dt_full = WignerDCalculator(lmax, precision="float64").call(quat) + D_full = np.asarray(D_full) + Dt_full = np.asarray(Dt_full) + if n_radial is None: + edge_rbf = np.zeros((n_edge, 1)) + else: + edge_rbf = rng.normal(size=(n_edge, n_radial)) + edge_env = rng.uniform(0.2, 1.0, size=(n_edge, 1)) + deg = ((edge_env[:, 0] ** 2) * mask).reshape(nloc, nnei).sum(axis=1) + inv_sqrt_deg = (1.0 / np.sqrt(deg + 1.0)).reshape(nloc, 1, 1) + edge_src_gate = rng.uniform(0.1, 1.0, size=(n_edge, 1)) if with_gate else None + radial = rng.normal(size=(n_edge, lmax + 1, channels)) + x = rng.normal(size=(nloc, dim_full, channels)) + + t = to_pt + pt_cache = EdgeFeatureCache( + src=t(src[valid]), + dst=t(dst[valid]), + edge_type_feat=t(np.zeros((n_valid, channels))), + edge_vec=t(edge_vec[valid]), + edge_rbf=t(edge_rbf[valid]), + edge_env=t(edge_env[valid]), + deg=t(deg), + inv_sqrt_deg=t(inv_sqrt_deg), + D_full=t(D_full[valid]), + Dt_full=t(Dt_full[valid]), + edge_src_gate=None if edge_src_gate is None else t(edge_src_gate[valid]), + ) + dp_cache = EdgeCache( + src=src, + dst=dst, + edge_type_feat=np.zeros((n_edge, channels)), + edge_vec=edge_vec, + edge_rbf=edge_rbf, + edge_env=edge_env, + deg=deg, + inv_sqrt_deg=inv_sqrt_deg, + D_full=D_full, + Dt_full=Dt_full, + edge_src_gate=edge_src_gate, + edge_mask=mask, + ) + return pt_cache, dp_cache, radial, radial[valid], x, valid + + +class TestSO2Parity: + nloc = 5 + nnei = 4 + + def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: + rng = np.random.default_rng(seed) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + + # ---------- SO2Linear ---------- + @pytest.mark.parametrize( + "lmax,mmax", [(2, 0), (2, 1), (2, 2), (3, 1), (3, 2)] + ) # degree/order truncations (mmax=0 covers the empty weight_m branch) + @pytest.mark.parametrize("mlp_bias", [False, True]) # l=0 bias branch + @pytest.mark.parametrize("n_focus", [1, 2]) # focus streams + def test_so2_linear(self, lmax, mmax, mlp_bias, n_focus) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Linear as DPSO2Linear + from deepmd.pt.model.descriptor.sezm_nn.so2 import SO2Linear as PTSO2Linear + + pt_mod = PTSO2Linear( + lmax=lmax, + mmax=mmax, + in_channels=5, + out_channels=3, + n_focus=n_focus, + dtype=torch.float64, + mlp_bias=mlp_bias, + seed=11, + trainable=True, + ) + self._perturb(pt_mod, 2052) + serialized = pt_mod.serialize() + dp_mod = DPSO2Linear.deserialize(serialized) + rng = np.random.default_rng(2053) + x = rng.normal(size=(13, n_focus, dp_mod.reduced_dim, 5)) + assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + + def test_so2_linear_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Linear as DPSO2Linear + + dp_mod = DPSO2Linear( + lmax=3, + mmax=1, + in_channels=4, + out_channels=4, + n_focus=2, + precision="float64", + mlp_bias=True, + seed=4, + ) + dp_mod2 = DPSO2Linear.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2054) + x = rng.normal(size=(9, 2, dp_mod.reduced_dim, 4)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_so2_linear_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Linear as DPSO2Linear + + with pytest.raises(ValueError): # mmax > lmax + DPSO2Linear(lmax=2, mmax=3, in_channels=2, out_channels=2) + with pytest.raises(ValueError): # negative mmax + DPSO2Linear(lmax=2, mmax=-1, in_channels=2, out_channels=2) + with pytest.raises(ValueError): # wrong class tag + DPSO2Linear.deserialize({"@class": "NotSO2Linear", "@version": 1}) + + # ---------- DynamicRadialDegreeMixer ---------- + @pytest.mark.parametrize( + "mode,rank", + [ + ("degree", 0), # channel-shared degree kernel + ("degree_channel", 0), # full per-channel kernel + ("degree_channel", 1), # low-rank factorization (core) + ("degree_channel", 2), # low-rank, rank > 1 + ], + ) + def test_radial_degree_mixer(self, mode, rank) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import ( + DynamicRadialDegreeMixer as DPMixer, + ) + from deepmd.pt.model.descriptor.sezm_nn.so2 import ( + DynamicRadialDegreeMixer as PTMixer, + ) + + pt_mod = PTMixer( + lmax=3, + mmax=1, + channels=4, + mode=mode, + rank=rank, + dtype=torch.float64, + seed=5, + trainable=True, + ) + self._perturb(pt_mod, 2055) + dp_mod = DPMixer( + lmax=3, + mmax=1, + channels=4, + mode=mode, + rank=rank, + precision="float64", + seed=5, + ) + # pt has no standalone serialize(); load the pt state_dict fragment + dp_mod._load_variables(pt_state_to_numpy(pt_mod)) + rng = np.random.default_rng(2056) + x_local = rng.normal(size=(17, dp_mod.reduced_dim, 4)) + radial = rng.normal(size=(17, dp_mod.reduced_dim, 4)) + assert_parity( + dp_mod.call(x_local, radial), + pt_mod(to_pt(x_local), to_pt(radial)), + ) + + def test_radial_degree_mixer_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import ( + DynamicRadialDegreeMixer as DPMixer, + ) + + dp_mod = DPMixer( + lmax=3, + mmax=1, + channels=4, + mode="degree_channel", + rank=1, + precision="float64", + seed=6, + ) + dp_mod2 = DPMixer.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2057) + x_local = rng.normal(size=(7, dp_mod.reduced_dim, 4)) + radial = rng.normal(size=(7, dp_mod.reduced_dim, 4)) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x_local, radial)), + np.asarray(dp_mod2.call(x_local, radial)), + ) + + def test_radial_degree_mixer_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import ( + DynamicRadialDegreeMixer as DPMixer, + ) + + common = {"lmax": 2, "mmax": 1, "channels": 4, "precision": "float64"} + with pytest.raises(ValueError): # unknown mode + DPMixer(mode="channel", **common) + with pytest.raises(ValueError): # negative rank + DPMixer(mode="degree_channel", rank=-1, **common) + with pytest.raises(ValueError): # non-positive channels + DPMixer(lmax=2, mmax=1, channels=0, mode="degree") + with pytest.raises(ValueError): # mmax > lmax + DPMixer(lmax=2, mmax=3, channels=4, mode="degree") + dp_mod = DPMixer(mode="degree", **common) + rng = np.random.default_rng(2058) + good = rng.normal(size=(3, dp_mod.reduced_dim, 4)) + with pytest.raises(ValueError): # shape mismatch between inputs + dp_mod.call(good, good[:, :, :2]) + with pytest.raises(ValueError): # incompatible reduced layout + dp_mod.call(good[:, :3, :], good[:, :3, :]) + with pytest.raises(ValueError): # wrong class tag + DPMixer.deserialize({"@class": "NotMixer", "@version": 1}) + + # ---------- segment softmax ---------- + @pytest.mark.parametrize( + "masked", ["none", "slots", "node"] + ) # padded-slot patterns (node = one all-masked destination) + @pytest.mark.parametrize("use_src_weight", [False, True]) # SFPG gate branch + def test_segment_envelope_gated_softmax(self, masked, use_src_weight) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( + segment_envelope_gated_softmax as dp_softmax, + ) + from deepmd.pt.model.descriptor.sezm_nn.attention import ( + segment_envelope_gated_softmax as pt_softmax, + ) + + rng = np.random.default_rng(2059) + nloc, nnei, n_focus, n_head = self.nloc, self.nnei, 2, 3 + pt_cache, dp_cache, _, _, _, valid = _build_so2_edge_data( + rng, + nloc=nloc, + nnei=nnei, + lmax=2, + channels=4, + masked=masked, + with_gate=use_src_weight, + ) + n_edge = nloc * nnei + logits = rng.normal(size=(n_edge, n_focus, n_head)) + # mixed signs exercise both stable-softplus branches for zeta + z_bias_raw = rng.normal(size=(n_focus, n_head)) + alpha_dp = dp_softmax( + logits=logits, + edge_env=dp_cache.edge_env, + n_nodes=nloc, + z_bias_raw=z_bias_raw, + eps=1e-7, + src_weight=dp_cache.edge_src_gate, + edge_mask=dp_cache.edge_mask, + ) + alpha_pt = pt_softmax( + logits=to_pt(logits[valid]), + edge_env=pt_cache.edge_env, + dst=pt_cache.dst, + n_nodes=nloc, + z_bias_raw=to_pt(z_bias_raw), + eps=1e-7, + src_weight=pt_cache.edge_src_gate, + ) + alpha_dp = np.asarray(alpha_dp) + np.testing.assert_allclose( + alpha_dp[valid], + alpha_pt.detach().cpu().numpy(), + rtol=PT_RTOL, + atol=PT_ATOL, + ) + # invalid slots must produce exactly zero attention weights + np.testing.assert_array_equal(alpha_dp[~valid], 0.0) + assert np.all(np.isfinite(alpha_dp)) + + def test_segment_softmax_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( + segment_envelope_gated_softmax as dp_softmax, + ) + + rng = np.random.default_rng(2062) + with pytest.raises(ValueError): # E not a multiple of n_nodes + dp_softmax( + logits=rng.normal(size=(7, 1, 1)), + edge_env=rng.uniform(size=(7, 1)), + n_nodes=3, + z_bias_raw=np.zeros((1, 1)), + eps=1e-7, + ) + + # ---------- SO2Convolution ---------- + def _conv_kwargs(self, **overrides): + kwargs = { + "lmax": 3, + "mmax": 1, + "kmax": 1, + "channels": 4, + "n_focus": 1, + "focus_dim": 0, + "focus_compete": True, + "so2_norm": False, + "so2_layers": 2, + "so2_attn_res": "none", + "layer_scale": False, + "n_atten_head": 1, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "lebedev_quadrature": True, + "activation_function": "silu", + "mlp_bias": False, + "eps": 1e-7, + } + kwargs.update(overrides) + return kwargs + + def _build_conv_pair(self, seed=17, perturb_seed=2060, **overrides): + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Convolution as DPSO2Conv + from deepmd.pt.model.descriptor.sezm_nn.so2 import SO2Convolution as PTSO2Conv + + kwargs = self._conv_kwargs(**overrides) + pt_mod = PTSO2Conv(**kwargs, dtype=torch.float64, seed=seed, trainable=True) + # post_focus_mix is zero-initialized; perturb so the output is nonzero + self._perturb(pt_mod, perturb_seed) + dp_mod = DPSO2Conv.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod, kwargs + + def _assert_conv_parity( + self, pt_mod, dp_mod, kwargs, *, masked="slots", with_gate=False + ) -> None: + rng = np.random.default_rng(2061) + pt_cache, dp_cache, radial, radial_valid, x, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=kwargs["lmax"], + channels=kwargs["channels"], + masked=masked, + with_gate=with_gate, + ) + out_dp = dp_mod.call(x, dp_cache, radial) + out_pt = pt_mod(to_pt(x), pt_cache, to_pt(radial_valid)) + assert_parity(out_dp, out_pt) + + @pytest.mark.parametrize("masked", ["none", "slots"]) # padded-slot pattern + @pytest.mark.parametrize("so2_layers", [2, 4]) # SO(2) layer loop depth (core=4) + def test_so2_convolution(self, masked, so2_layers) -> None: + pt_mod, dp_mod, kwargs = self._build_conv_pair(so2_layers=so2_layers) + self._assert_conv_parity(pt_mod, dp_mod, kwargs, masked=masked) + + def test_so2_convolution_all_masked_node(self) -> None: + # one destination with zero valid incoming edges + pt_mod, dp_mod, kwargs = self._build_conv_pair() + self._assert_conv_parity(pt_mod, dp_mod, kwargs, masked="node") + + @pytest.mark.parametrize( + "radial_so2_mode,radial_so2_rank", + [ + ("none", 0), # elementwise radial modulation + ("degree", 0), # channel-shared dynamic degree kernel + ("degree_channel", 0), # full per-channel dynamic kernel + ], + ) + def test_so2_convolution_radial_modes( + self, radial_so2_mode, radial_so2_rank + ) -> None: + pt_mod, dp_mod, kwargs = self._build_conv_pair( + radial_so2_mode=radial_so2_mode, radial_so2_rank=radial_so2_rank + ) + self._assert_conv_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize( + "n_atten_head", [0, 2] + ) # 0 = plain envelope sum, 2 = multi-head attention + def test_so2_convolution_atten_heads(self, n_atten_head) -> None: + pt_mod, dp_mod, kwargs = self._build_conv_pair(n_atten_head=n_atten_head) + self._assert_conv_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize("focus_compete", [False, True]) # competition branch + def test_so2_convolution_multi_focus(self, focus_compete) -> None: + # n_focus=2 also activates the hidden-width ChannelLinear projection + pt_mod, dp_mod, kwargs = self._build_conv_pair( + n_focus=2, focus_compete=focus_compete + ) + self._assert_conv_parity(pt_mod, dp_mod, kwargs) + + def test_so2_convolution_so2_norm(self) -> None: + pt_mod, dp_mod, kwargs = self._build_conv_pair(so2_norm=True, so2_layers=3) + self._assert_conv_parity(pt_mod, dp_mod, kwargs) + + def test_so2_convolution_mlp_bias(self) -> None: + # exercises bias0 + the layer-0 envelope bias correction + pt_mod, dp_mod, kwargs = self._build_conv_pair(mlp_bias=True) + self._assert_conv_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize("n_atten_head", [0, 1]) # gate enters both paths + def test_so2_convolution_src_gate(self, n_atten_head) -> None: + pt_mod, dp_mod, kwargs = self._build_conv_pair(n_atten_head=n_atten_head) + self._assert_conv_parity(pt_mod, dp_mod, kwargs, with_gate=True) + + def test_so2_convolution_real_edge_cache(self) -> None: + # end-to-end: REAL pt build_edge_cache vs REAL dp build_edge_cache + # feeding the same weight-copied SO2Convolution (no synthetic cache) + pt_mod, dp_mod, kwargs = self._build_conv_pair() + rng = np.random.default_rng(2096) + nf, nloc, nall, nnei = 1, self.nloc, self.nloc + 3, self.nnei + inputs = _build_real_edge_inputs( + rng, + nf=nf, + nloc=nloc, + nall=nall, + nnei=nnei, + channels=kwargs["channels"], + ) + pt_cache, dp_cache = _build_real_edge_caches(inputs, lmax=kwargs["lmax"]) + valid = inputs["valid"] + dim_full = (kwargs["lmax"] + 1) ** 2 + radial = rng.normal(size=(nf * nloc * nnei, kwargs["lmax"] + 1, 4)) + x = rng.normal(size=(nf * nloc, dim_full, kwargs["channels"])) + out_dp = dp_mod.call(x, dp_cache, radial) + out_pt = pt_mod(to_pt(x), pt_cache, to_pt(radial[valid])) + assert_parity(out_dp, out_pt) + + def test_so2_convolution_full_mmax(self) -> None: + # mmax == lmax: rotate_inv_rescale is all ones + pt_mod, dp_mod, kwargs = self._build_conv_pair(lmax=2, mmax=2) + self._assert_conv_parity(pt_mod, dp_mod, kwargs) + + def test_so2_convolution_roundtrip(self) -> None: + _, dp_mod, kwargs = self._build_conv_pair( + n_focus=2, so2_norm=True, mlp_bias=True + ) + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Convolution as DPSO2Conv + + dp_mod2 = DPSO2Conv.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2063) + _, dp_cache, radial, _, x, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=kwargs["lmax"], + channels=kwargs["channels"], + masked="slots", + ) + out1 = np.asarray(dp_mod.call(x, dp_cache, radial)) + # the D_to_m projections are cached in the EdgeCache dicts; reuse is exact + out2 = np.asarray(dp_mod2.call(x, dp_cache, radial)) + np.testing.assert_array_equal(out1, out2) + + @pytest.mark.parametrize( + "flag,value", + [ + ("so2_attn_res", "independent"), # DepthAttnRes + ("so2_attn_res", "dependent"), # DepthAttnRes + ("layer_scale", True), # per-layer LayerScale + ("n_atten_head", -1), # ValueError, not NotImplementedError + ("atten_f_mix", True), # focus-merged attention + ("atten_v_proj", True), # value projection + ("atten_o_proj", True), # output projection + ("s2_activation", True), # S2-grid SwiGLU non-linearity + ("node_wise_s2", True), # edge-local S2 grid product + ("node_wise_so3", True), # edge-local SO(3) grid product + ("message_node_s2", True), # post-aggregation S2 grid product + ("message_node_so3", True), # post-aggregation SO(3) grid product + ], + ) + def test_so2_convolution_guards(self, flag, value) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Convolution as DPSO2Conv + + kwargs = self._conv_kwargs(**{flag: value}) + if flag == "n_atten_head": + with pytest.raises(ValueError): + DPSO2Conv(**kwargs, precision="float64") + else: + with pytest.raises(NotImplementedError, match=flag): + DPSO2Conv(**kwargs, precision="float64") + + def test_so2_convolution_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Convolution as DPSO2Conv + + with pytest.raises(ValueError): # head count must divide focus width + DPSO2Conv(**self._conv_kwargs(n_atten_head=3), precision="float64") + with pytest.raises(ValueError): # so2_layers must be >= 1 + DPSO2Conv(**self._conv_kwargs(so2_layers=0), precision="float64") + with pytest.raises(ValueError): # n_focus must be >= 1 + DPSO2Conv(**self._conv_kwargs(n_focus=0), precision="float64") + with pytest.raises(ValueError): # unknown radial mode + DPSO2Conv( + **self._conv_kwargs(radial_so2_mode="degree_rank"), + precision="float64", + ) + with pytest.raises(ValueError): # mmax > lmax + DPSO2Conv(**self._conv_kwargs(mmax=4), precision="float64") + with pytest.raises(ValueError): # unknown so2_attn_res token + DPSO2Conv(**self._conv_kwargs(so2_attn_res="depth"), precision="float64") + dp_mod = DPSO2Conv(**self._conv_kwargs(), precision="float64", seed=1) + rng = np.random.default_rng(2064) + _, dp_cache, radial, _, x, _ = _build_so2_edge_data( + rng, nloc=self.nloc, nnei=self.nnei, lmax=3, channels=4 + ) + with pytest.raises(ValueError): # E not a multiple of N + dp_mod.call(x[:3], dp_cache, radial) + with pytest.raises(ValueError): # wrong class tag + DPSO2Conv.deserialize({"@class": "NotConv", "@version": 1}) + + +class TestEmbeddingParity: + nloc = 5 + nnei = 4 + + def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: + rng = np.random.default_rng(seed) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + + # ---------- SeZMTypeEmbedding ---------- + @pytest.mark.parametrize("padding", [False, True]) # zero padding row branch + def test_type_embedding(self, padding) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + SeZMTypeEmbedding as DPTypeEmbed, + ) + from deepmd.pt.model.descriptor.sezm_nn.embedding import ( + SeZMTypeEmbedding as PTTypeEmbed, + ) + + ntypes, embed_dim = 4, 6 + pt_mod = PTTypeEmbed( + ntypes=ntypes, + embed_dim=embed_dim, + dtype=torch.float64, + seed=21, + trainable=True, + padding=padding, + ) + dp_mod = DPTypeEmbed( + ntypes=ntypes, + embed_dim=embed_dim, + precision="float64", + seed=21, + padding=padding, + ) + state = pt_state_to_numpy(pt_mod) + # pt has no serialize(); the @variables key set must equal the pt + # state_dict key set so the weights map one-to-one. + assert set(dp_mod.serialize()["@variables"]) == set(state) + assert state["adam_type_embedding"].shape == dp_mod.adam_type_embedding.shape + dp_mod.adam_type_embedding = state["adam_type_embedding"] + rng = np.random.default_rng(2070) + # include the padding row index ntypes when padding=True + atype = rng.integers(0, ntypes + 1 if padding else ntypes, size=(3, 5)) + assert_parity(dp_mod.call(atype), pt_mod(to_pt(atype))) + if padding: + pad_out = np.asarray(dp_mod.call(np.full((2,), ntypes, dtype=np.int64))) + np.testing.assert_array_equal(pad_out, 0.0) + + @pytest.mark.parametrize("padding", [False, True]) # zero padding row branch + def test_type_embedding_roundtrip(self, padding) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + SeZMTypeEmbedding as DPTypeEmbed, + ) + + dp_mod = DPTypeEmbed( + ntypes=3, embed_dim=4, precision="float64", seed=22, padding=padding + ) + dp_mod2 = DPTypeEmbed.deserialize(dp_mod.serialize()) + atype = np.array([0, 2, 1, 1], dtype=np.int64) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(atype)), np.asarray(dp_mod2.call(atype)) + ) + + def test_type_embedding_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + SeZMTypeEmbedding as DPTypeEmbed, + ) + + with pytest.raises(ValueError): # non-positive ntypes + DPTypeEmbed(ntypes=0, embed_dim=4) + with pytest.raises(ValueError): # non-positive embed_dim + DPTypeEmbed(ntypes=3, embed_dim=0) + with pytest.raises(ValueError): # wrong class tag + DPTypeEmbed.deserialize({"@class": "NotTypeEmbed", "@version": 1}) + dp_mod = DPTypeEmbed(ntypes=3, embed_dim=4, precision="float64") + data = dp_mod.serialize() + data["@variables"]["adam_type_embedding"] = np.zeros((2, 4)) + with pytest.raises(ValueError): # table shape mismatch + DPTypeEmbed.deserialize(data) + + # ---------- GeometricInitialEmbedding ---------- + def _build_gie_pair(self, lmax, channels): + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + GeometricInitialEmbedding as DPGIE, + ) + from deepmd.pt.model.descriptor.sezm_nn.embedding import ( + GeometricInitialEmbedding as PTGIE, + ) + + pt_mod = PTGIE(lmax=lmax, channels=channels, dtype=torch.float64) + # pt serialize() is config-only; the dp module is weight-free. + dp_mod = DPGIE.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod + + @pytest.mark.parametrize( + "masked", ["none", "slots", "node"] + ) # padded-slot patterns (node = one all-masked destination) + @pytest.mark.parametrize( + "zonal_provided", [False, True] + ) # zonal_coupling: None (gather from Dt_full) vs provided-zonal input + # path with D_node == D_cache only + @pytest.mark.parametrize("with_gate", [False, True]) # SFPG gate branch + def test_gie(self, masked, zonal_provided, with_gate) -> None: + lmax, channels = 2, 4 + pt_mod, dp_mod = self._build_gie_pair(lmax, channels) + rng = np.random.default_rng(2071) + pt_cache, dp_cache, radial, radial_valid, _, valid = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=lmax, + channels=channels, + masked=masked, + with_gate=with_gate, + ) + if zonal_provided: + # Scope: provided-zonal here uses D_node == D_cache; the genuine + # lmax_node > lmax_mp (dim_full != ebed_dim) path is exercised by + # the descriptor-level tests in a later task. + rows = dp_mod.non_scalar_row_index + cols = dp_mod.zonal_m0_col_index_for_row + dp_zonal = np.asarray(dp_cache.Dt_full)[:, rows, cols] + pt_zonal = pt_cache.Dt_full[:, to_pt(rows), to_pt(cols)] + else: + dp_zonal = pt_zonal = None + out_dp = dp_mod.call( + n_nodes=self.nloc, + edge_cache=dp_cache, + radial_feat=radial[:, 1:, :], + zonal_coupling=dp_zonal, + ) + out_pt = pt_mod( + n_nodes=self.nloc, + edge_cache=pt_cache, + radial_feat=to_pt(radial_valid[:, 1:, :]), + zonal_coupling=pt_zonal, + ) + assert_parity(out_dp, out_pt) + # l=0 row must be exactly zero (comes from type embedding instead) + np.testing.assert_array_equal(np.asarray(out_dp)[:, 0, :], 0.0) + + def test_gie_lmax0(self) -> None: + # lmax=0 short-circuit: all-zero (N, 1, C) on both sides + pt_mod, dp_mod = self._build_gie_pair(0, 3) + rng = np.random.default_rng(2072) + pt_cache, dp_cache, _, _, _, _ = _build_so2_edge_data( + rng, nloc=self.nloc, nnei=self.nnei, lmax=1, channels=3 + ) + out_dp = np.asarray( + dp_mod.call(n_nodes=self.nloc, edge_cache=dp_cache, radial_feat=None) + ) + out_pt = pt_mod(n_nodes=self.nloc, edge_cache=pt_cache, radial_feat=None) + assert out_dp.shape == (self.nloc, 1, 3) + np.testing.assert_array_equal(out_dp, out_pt.detach().cpu().numpy()) + np.testing.assert_array_equal(out_dp, 0.0) + + def test_gie_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + GeometricInitialEmbedding as DPGIE, + ) + + dp_mod = DPGIE(lmax=3, channels=4, precision="float64") + dp_mod2 = DPGIE.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2073) + _, dp_cache, radial, _, _, _ = _build_so2_edge_data( + rng, nloc=self.nloc, nnei=self.nnei, lmax=3, channels=4, masked="slots" + ) + out1 = dp_mod.call( + n_nodes=self.nloc, edge_cache=dp_cache, radial_feat=radial[:, 1:, :] + ) + out2 = dp_mod2.call( + n_nodes=self.nloc, edge_cache=dp_cache, radial_feat=radial[:, 1:, :] + ) + np.testing.assert_array_equal(np.asarray(out1), np.asarray(out2)) + + def test_gie_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + GeometricInitialEmbedding as DPGIE, + ) + + with pytest.raises(ValueError): # wrong class tag + DPGIE.deserialize({"@class": "NotGIE", "@version": 1}) + dp_mod = DPGIE(lmax=2, channels=4, precision="float64") + rng = np.random.default_rng(2074) + _, dp_cache, radial, _, _, _ = _build_so2_edge_data( + rng, nloc=self.nloc, nnei=self.nnei, lmax=2, channels=4 + ) + with pytest.raises(ValueError): # E not a multiple of N + dp_mod.call(n_nodes=3, edge_cache=dp_cache, radial_feat=radial[:, 1:, :]) + + # ---------- EnvironmentInitialEmbedding ---------- + n_radial = 5 + ntypes = 3 + + def _env_kwargs(self, **overrides): + kwargs = { + "ntypes": self.ntypes, + "n_radial": self.n_radial, + "channels": 4, + "embed_dim": 12, + "axis_dim": 3, + "type_dim": 4, + "hidden_dim": 8, + "mlp_bias": False, + "activation_function": "silu", + "eps": 1e-7, + } + kwargs.update(overrides) + return kwargs + + def _build_env_pair(self, seed=23, perturb_seed=2075, **overrides): + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + EnvironmentInitialEmbedding as DPEnv, + ) + from deepmd.pt.model.descriptor.sezm_nn.embedding import ( + EnvironmentInitialEmbedding as PTEnv, + ) + + kwargs = self._env_kwargs(**overrides) + pt_mod = PTEnv(**kwargs, dtype=torch.float64, seed=seed, trainable=True) + # output_proj is zero-initialized; perturb so the output is nonzero + self._perturb(pt_mod, perturb_seed) + dp_mod = DPEnv.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod + + def _assert_env_parity( + self, pt_mod, dp_mod, *, masked="slots", with_gate=False + ) -> None: + rng = np.random.default_rng(2076) + pt_cache, dp_cache, _, _, _, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=1, + channels=4, + masked=masked, + with_gate=with_gate, + n_radial=self.n_radial, + ) + atype = rng.integers(0, self.ntypes, size=(self.nloc,)) + out_dp = dp_mod.call(edge_cache=dp_cache, atype_flat=atype, n_nodes=self.nloc) + out_pt = pt_mod( + edge_cache=pt_cache, + atype_flat=to_pt(atype), + n_nodes=self.nloc, + ) + assert_parity(out_dp, out_pt) + + @pytest.mark.parametrize( + "masked", ["none", "slots", "node"] + ) # padded-slot patterns (node = one all-masked destination) + @pytest.mark.parametrize("mlp_bias", [False, True]) # MLP bias branch + def test_env_embedding(self, masked, mlp_bias) -> None: + pt_mod, dp_mod = self._build_env_pair(mlp_bias=mlp_bias) + self._assert_env_parity(pt_mod, dp_mod, masked=masked) + + def test_env_embedding_src_gate(self) -> None: + pt_mod, dp_mod = self._build_env_pair() + self._assert_env_parity(pt_mod, dp_mod, with_gate=True) + + def test_env_embedding_wide_rbf(self) -> None: + # embed_dim - 2*type_dim > 32 exercises the non-clamped rbf_out_dim + pt_mod, dp_mod = self._build_env_pair(embed_dim=42, axis_dim=3) + self._assert_env_parity(pt_mod, dp_mod) + + @pytest.mark.parametrize("mlp_bias", [False, True]) # MLP bias branch + def test_env_embedding_serialize_keys(self, mlp_bias) -> None: + pt_mod, dp_mod = self._build_env_pair(mlp_bias=mlp_bias) + assert set(dp_mod.serialize()["@variables"]) == set(pt_mod.state_dict()) + + def test_env_embedding_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + EnvironmentInitialEmbedding as DPEnv, + ) + + _, dp_mod = self._build_env_pair(mlp_bias=True) + dp_mod2 = DPEnv.deserialize(dp_mod.serialize()) + rng = np.random.default_rng(2077) + _, dp_cache, _, _, _, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=1, + channels=4, + masked="slots", + n_radial=self.n_radial, + ) + atype = rng.integers(0, self.ntypes, size=(self.nloc,)) + out1 = dp_mod.call(edge_cache=dp_cache, atype_flat=atype, n_nodes=self.nloc) + out2 = dp_mod2.call(edge_cache=dp_cache, atype_flat=atype, n_nodes=self.nloc) + np.testing.assert_array_equal(np.asarray(out1), np.asarray(out2)) + + def test_env_embedding_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + EnvironmentInitialEmbedding as DPEnv, + ) + + with pytest.raises(ValueError): # axis_dim must be < embed_dim + DPEnv(**self._env_kwargs(axis_dim=12), precision="float64") + with pytest.raises(ValueError): # wrong class tag + DPEnv.deserialize({"@class": "NotEnv", "@version": 1}) + dp_mod = DPEnv(**self._env_kwargs(), precision="float64", seed=2) + data = dp_mod.serialize() + data["@variables"].pop("output_proj.matrix") + with pytest.raises(ValueError): # variable key set mismatch + DPEnv.deserialize(data) + data = dp_mod.serialize() + data["@variables"]["output_proj.matrix"] = np.zeros((2, 2)) + with pytest.raises(ValueError): # variable shape mismatch + DPEnv.deserialize(data) + rng = np.random.default_rng(2078) + _, dp_cache, _, _, _, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=1, + channels=4, + n_radial=self.n_radial, + ) + atype = rng.integers(0, self.ntypes, size=(self.nloc,)) + with pytest.raises(ValueError): # E not a multiple of N + dp_mod.call(edge_cache=dp_cache, atype_flat=atype, n_nodes=3) + + +def _build_real_edge_inputs( + rng, + *, + nf, + nloc, + nall, + nnei, + channels, + local_nlist=False, +): + """Build numpy inputs for a real ``build_edge_cache`` run. + + Includes -1 padding slots, ghosts (``nall > nloc``) mapping back to their + owners, one broken mapping entry (``mapping == -1``, exercising pt's + ``src_ok`` drop), and a non-empty type-pair exclusion ``(0, 1)``. + When ``local_nlist`` is True, neighbor indices are drawn directly from the + local range and ``mapping`` is ``None`` (pt's mapping-free branch). + """ + ntypes = 3 + coord = rng.uniform(0.0, 4.0, size=(nf, nall, 3)) + atype = rng.integers(0, ntypes, size=(nf, nloc)) + if local_nlist: + mapping = None + atype_ext = atype + hi = nloc + else: + n_ghost = nall - nloc + mapping = np.concatenate( + [ + np.tile(np.arange(nloc, dtype=np.int64), (nf, 1)), + rng.integers(0, nloc, size=(nf, n_ghost)), + ], + axis=1, + ) + mapping[:, -1] = -1 # broken ghost: pt drops via src_ok, dp masks + atype_ext = np.take_along_axis(atype, np.clip(mapping, 0, nloc - 1), axis=1) + hi = nall + # neighbors over the extended axis, excluding the center itself + nlist = rng.integers(0, hi, size=(nf, nloc, nnei)) + center = np.arange(nloc)[None, :, None] + nlist = np.where(nlist == center, (nlist + 1) % hi, nlist) + nlist[rng.uniform(size=nlist.shape) < 0.25] = -1 # padding slots + # pair_keep_mask from exclude pair (0, 1), built on extended types + nl_safe = np.where(nlist >= 0, nlist, 0) + nb_type = np.take_along_axis(atype_ext, nl_safe.reshape(nf, -1), axis=1).reshape( + nf, nloc, nnei + ) + ct = atype[:, :, None] + pair_keep_mask = ~(((ct == 0) & (nb_type == 1)) | ((ct == 1) & (nb_type == 0))) + type_ebed = rng.normal(size=(nf * nloc, channels)) + # expected validity mask, computed independently in numpy + if local_nlist: + src_local = nl_safe + else: + src_local = np.take_along_axis( + mapping, nl_safe.reshape(nf, -1), axis=1 + ).reshape(nf, nloc, nnei) + valid = ( + (nlist >= 0) & pair_keep_mask & (src_local >= 0) & (src_local < nloc) + ).reshape(-1) + return { + "coord": coord, + "nlist": nlist, + "mapping": mapping, + "pair_keep_mask": pair_keep_mask, + "type_ebed": type_ebed, + "valid": valid, + } + + +def _build_real_edge_caches( + inputs, + *, + lmax, + rcut=6.0, + n_radial=8, + deg_norm_floor=1e-12, + eps=1e-7, + random_gamma=False, + gamma=None, + seed=2090, +): + """Run the REAL pt and dp ``build_edge_cache`` on identical inputs. + + The pt ``RadialBasis`` frequencies are perturbed and weight-copied into + the dp side via ``deserialize`` so parity exercises copied weights. + Returns ``(pt_cache, dp_cache)``. + """ + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + build_edge_cache as dp_build_edge_cache, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import C3CutoffEnvelope as DPEnvelope + from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialBasis as DPRadialBasis + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import WignerDCalculator as DPWigner + from deepmd.pt.model.descriptor.sezm_nn.edge_cache import ( + build_edge_cache as pt_build_edge_cache, + ) + from deepmd.pt.model.descriptor.sezm_nn.radial import C3CutoffEnvelope as PTEnvelope + from deepmd.pt.model.descriptor.sezm_nn.radial import RadialBasis as PTRadialBasis + from deepmd.pt.model.descriptor.sezm_nn.wignerd import WignerDCalculator as PTWigner + + pt_rb = PTRadialBasis(rcut=rcut, n_radial=n_radial, dtype=torch.float64) + rng = np.random.default_rng(seed) + with torch.no_grad(): + pt_rb.adam_freqs += to_pt(0.05 * rng.normal(size=(1, n_radial))) + dp_rb = DPRadialBasis.deserialize(pt_rb.serialize()) + pt_env = PTEnvelope(rcut=rcut, dtype=torch.float64) + dp_env = DPEnvelope(rcut=rcut, precision="float64") + pt_wig = PTWigner(lmax, dtype=torch.float64) + dp_wig = DPWigner(lmax, precision="float64") + + t = to_pt + mapping = inputs["mapping"] + pt_cache = pt_build_edge_cache( + type_ebed=t(inputs["type_ebed"]), + extended_coord=t(inputs["coord"]), + nlist=t(inputs["nlist"]), + mapping=None if mapping is None else t(mapping), + pair_keep_mask=t(inputs["pair_keep_mask"]), + eps=eps, + deg_norm_floor=deg_norm_floor, + edge_envelope=pt_env, + radial_basis=pt_rb, + n_radial=n_radial, + random_gamma=random_gamma, + wigner_calc=pt_wig, + ) + dp_cache = dp_build_edge_cache( + type_ebed=inputs["type_ebed"], + extended_coord=inputs["coord"], + nlist=inputs["nlist"], + mapping=mapping, + pair_keep_mask=inputs["pair_keep_mask"], + eps=eps, + deg_norm_floor=deg_norm_floor, + edge_envelope=dp_env, + radial_basis=dp_rb, + n_radial=n_radial, + random_gamma=random_gamma, + wigner_calc=dp_wig, + gamma=gamma, + ) + return pt_cache, dp_cache + + +class TestEdgeCacheParity: + nf = 2 + nloc = 6 + nall = 10 + nnei = 12 + channels = 4 + lmax = 2 + + def _inputs(self, seed=2086, **overrides): + rng = np.random.default_rng(seed) + kwargs = { + "nf": self.nf, + "nloc": self.nloc, + "nall": self.nall, + "nnei": self.nnei, + "channels": self.channels, + } + kwargs.update(overrides) + return _build_real_edge_inputs(rng, **kwargs) + + @pytest.mark.parametrize("local_nlist", [False, True]) # mapping None branch + @pytest.mark.parametrize( + "deg_norm_floor", [1e-12, 1.0] + ) # legacy-eps floor vs O(1) floor + def test_real_build_parity(self, local_nlist, deg_norm_floor) -> None: + inputs = self._inputs( + local_nlist=local_nlist, + nall=self.nloc if local_nlist else self.nall, + ) + valid = inputs["valid"] + pt_cache, dp_cache = _build_real_edge_caches( + inputs, lmax=self.lmax, deg_norm_floor=deg_norm_floor + ) + # the dp validity mask matches the independently computed mask, and + # pt's sparse edges occupy exactly those slots (in row-major order) + np.testing.assert_array_equal(np.asarray(dp_cache.edge_mask), valid) + assert pt_cache.src.shape[0] == int(valid.sum()) + # padded dst contract: node-contiguous repeat of arange(nf * nloc) + np.testing.assert_array_equal( + np.asarray(dp_cache.dst), + np.repeat(np.arange(self.nf * self.nloc), self.nnei), + ) + # indices on valid slots are exactly pt's sparse indices + np.testing.assert_array_equal( + np.asarray(dp_cache.src)[valid], pt_cache.src.cpu().numpy() + ) + np.testing.assert_array_equal( + np.asarray(dp_cache.dst)[valid], pt_cache.dst.cpu().numpy() + ) + # per-edge fields: compare masked entries against pt's sparse outputs + for name in ( + "edge_vec", + "edge_rbf", + "edge_env", + "edge_type_feat", + "edge_quat", + "D_full", + "Dt_full", + ): + assert_parity( + np.asarray(getattr(dp_cache, name))[valid], getattr(pt_cache, name) + ) + # node-level normalization compares directly + assert_parity(dp_cache.deg, pt_cache.deg) + assert_parity(dp_cache.inv_sqrt_deg, pt_cache.inv_sqrt_deg) + # everything is finite, including masked slots + for name in ( + "edge_vec", + "edge_rbf", + "edge_env", + "edge_type_feat", + "edge_quat", + "D_full", + "Dt_full", + "deg", + "inv_sqrt_deg", + ): + assert np.isfinite(np.asarray(getattr(dp_cache, name))).all(), name + # standard path carries no source gate + assert dp_cache.edge_src_gate is None + assert dp_cache.D_to_m_cache == {} + assert dp_cache.Dt_from_m_cache == {} + + def test_out_of_range_local_index_masked(self) -> None: + # a local nlist entry >= nloc with mapping=None must be masked out + # and must not break the coordinate gather (nlist_safe is re-zeroed + # after the final src_ok mask update) + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + build_edge_cache as dp_build_edge_cache, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as DPEnvelope, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + RadialBasis as DPRadialBasis, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as DPWigner, + ) + + inputs = self._inputs(local_nlist=True, nall=self.nloc) + nlist = inputs["nlist"].copy() + nlist[0, 0, 0] = self.nloc # out of [0, nloc), would gather OOB + n_radial = 8 + cache = dp_build_edge_cache( + type_ebed=inputs["type_ebed"], + extended_coord=inputs["coord"], + nlist=nlist, + mapping=None, + pair_keep_mask=inputs["pair_keep_mask"], + eps=1e-7, + deg_norm_floor=1e-12, + edge_envelope=DPEnvelope(rcut=6.0, precision="float64"), + radial_basis=DPRadialBasis( + rcut=6.0, n_radial=n_radial, precision="float64" + ), + n_radial=n_radial, + random_gamma=False, + wigner_calc=DPWigner(self.lmax, precision="float64"), + ) + mask = np.asarray(cache.edge_mask).reshape(self.nf, self.nloc, self.nnei) + assert not mask[0, 0, 0] + assert np.isfinite(np.asarray(cache.edge_vec)).all() + + def test_masked_edge_inertness(self) -> None: + # an extra all-(-1) neighbor column must not change the masked-view + # fields or the degree normalization + inputs = self._inputs() + pad = -np.ones((self.nf, self.nloc, 1), dtype=inputs["nlist"].dtype) + inputs2 = dict(inputs) + inputs2["nlist"] = np.concatenate([inputs["nlist"], pad], axis=-1) + inputs2["pair_keep_mask"] = np.concatenate( + [ + inputs["pair_keep_mask"], + np.ones((self.nf, self.nloc, 1), dtype=bool), + ], + axis=-1, + ) + _, cache = _build_real_edge_caches(inputs, lmax=self.lmax) + _, cache2 = _build_real_edge_caches(inputs2, lmax=self.lmax) + n_nodes = self.nf * self.nloc + mask = np.asarray(cache.edge_mask).reshape(n_nodes, self.nnei) + mask2 = np.asarray(cache2.edge_mask).reshape(n_nodes, self.nnei + 1) + np.testing.assert_array_equal(mask2[:, : self.nnei], mask) + np.testing.assert_array_equal(mask2[:, self.nnei], False) + for name in ("edge_vec", "edge_rbf", "edge_env", "edge_quat", "D_full"): + a = np.asarray(getattr(cache, name)) + b = np.asarray(getattr(cache2, name)) + a = a.reshape(n_nodes, self.nnei, -1)[mask.astype(bool)] + b = b.reshape(n_nodes, self.nnei + 1, -1)[mask2.astype(bool)] + np.testing.assert_array_equal(a, b, err_msg=name) + np.testing.assert_array_equal(np.asarray(cache.deg), np.asarray(cache2.deg)) + np.testing.assert_array_equal( + np.asarray(cache.inv_sqrt_deg), np.asarray(cache2.inv_sqrt_deg) + ) + + def test_random_gamma(self) -> None: + # pt draws gamma internally with torch.rand, so the draw cannot be + # injected identically into both sides; the dp branch is verified by + # determinism (injected gamma) and gauge properties instead. + inputs = self._inputs() + n_edge = self.nf * self.nloc * self.nnei + gamma = np.random.default_rng(7).uniform(0.0, 2.0 * np.pi, n_edge) + _, base = _build_real_edge_caches(inputs, lmax=self.lmax) + _, c1 = _build_real_edge_caches( + inputs, lmax=self.lmax, random_gamma=True, gamma=gamma + ) + _, c2 = _build_real_edge_caches( + inputs, lmax=self.lmax, random_gamma=True, gamma=gamma + ) + # determinism with an injected gamma + np.testing.assert_array_equal(np.asarray(c1.D_full), np.asarray(c2.D_full)) + np.testing.assert_array_equal( + np.asarray(c1.edge_quat), np.asarray(c2.edge_quat) + ) + # the roll is a gauge choice: D stays orthogonal ... + d = np.asarray(c1.D_full) + dt = np.asarray(c1.Dt_full) + eye = np.broadcast_to(np.eye(d.shape[-1]), d.shape) + np.testing.assert_allclose(d @ dt, eye, rtol=1e-12, atol=1e-12) + # ... the l=0 block is unchanged ... + np.testing.assert_allclose( + d[:, 0, 0], np.asarray(base.D_full)[:, 0, 0], rtol=1e-12, atol=1e-14 + ) + # ... and the rotation-independent fields are bit-identical + for name in ("edge_vec", "edge_env", "edge_rbf", "deg", "inv_sqrt_deg"): + np.testing.assert_array_equal( + np.asarray(getattr(c1, name)), + np.asarray(getattr(base, name)), + err_msg=name, + ) + # internal-draw branch (gamma=None) runs and stays finite + _, c3 = _build_real_edge_caches(inputs, lmax=self.lmax, random_gamma=True) + assert np.isfinite(np.asarray(c3.D_full)).all() + np.testing.assert_allclose( + np.asarray(c3.D_full)[:, 0, 0], + np.asarray(base.D_full)[:, 0, 0], + rtol=1e-12, + atol=1e-14, + ) + + +class TestFFNParity: + n_node = 11 + channels = 8 + + def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: + rng = np.random.default_rng(seed) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + + def _ffn_kwargs(self, **overrides): + kwargs = { + "lmax": 3, + "channels": self.channels, + "hidden_channels": self.channels, + "kmax": 1, + "grid_mlp": False, + "grid_branch": 0, + "s2_activation": False, + "ffn_so3_grid": False, + "lebedev_quadrature": True, + "activation_function": "silu", + "glu_activation": True, + "mlp_bias": False, + } + kwargs.update(overrides) + return kwargs + + def _build_ffn_pair(self, seed=29, perturb_seed=2110, **overrides): + from deepmd.dpmodel.descriptor.dpa4_nn.ffn import EquivariantFFN as DPFFN + from deepmd.pt.model.descriptor.sezm_nn.ffn import EquivariantFFN as PTFFN + + kwargs = self._ffn_kwargs(**overrides) + pt_mod = PTFFN(**kwargs, dtype=torch.float64, seed=seed, trainable=True) + # so3_linear_2 is zero-initialized; perturb so the output is nonzero + self._perturb(pt_mod, perturb_seed) + dp_mod = DPFFN.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod, kwargs + + def _assert_ffn_parity(self, pt_mod, dp_mod, kwargs, seed=2111) -> None: + rng = np.random.default_rng(seed) + dim = (kwargs["lmax"] + 1) ** 2 + x = rng.normal(size=(self.n_node, dim, 1, kwargs["channels"])) + out_dp = dp_mod.call(x) + out_pt = pt_mod(to_pt(x)) + assert out_dp.shape == tuple(out_pt.shape) + assert_parity(out_dp, out_pt) + + @pytest.mark.parametrize("lmax", [2, 3]) # degree truncation (core=3) + @pytest.mark.parametrize("s2_activation", [False, True]) # S2 grid path (core=True) + @pytest.mark.parametrize("glu_activation", [False, True]) # GLU gating (core=True) + def test_ffn(self, lmax, s2_activation, glu_activation) -> None: + pt_mod, dp_mod, kwargs = self._build_ffn_pair( + lmax=lmax, s2_activation=s2_activation, glu_activation=glu_activation + ) + self._assert_ffn_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize("grid_branch", [0, 1]) # branch mixer off/on (core=1) + def test_ffn_grid_branch(self, grid_branch) -> None: + pt_mod, dp_mod, kwargs = self._build_ffn_pair( + s2_activation=True, grid_branch=grid_branch + ) + self._assert_ffn_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize("mlp_bias", [False, True]) # l=0 / gate bias branch + def test_ffn_mlp_bias(self, mlp_bias) -> None: + pt_mod, dp_mod, kwargs = self._build_ffn_pair(mlp_bias=mlp_bias) + self._assert_ffn_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize("s2_activation", [False, True]) # both act sub-modules + def test_ffn_roundtrip(self, s2_activation) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.ffn import EquivariantFFN as DPFFN + + pt_mod, dp_mod, kwargs = self._build_ffn_pair( + s2_activation=s2_activation, grid_branch=1 if s2_activation else 0 + ) + data = dp_mod.serialize() + # exact pt state_dict key-set match + assert set(data["@variables"]) == set(pt_state_to_numpy(pt_mod)) + dp_mod2 = DPFFN.deserialize(data) + rng = np.random.default_rng(2112) + dim = (kwargs["lmax"] + 1) ** 2 + x = rng.normal(size=(self.n_node, dim, 1, kwargs["channels"])) + np.testing.assert_array_equal( + np.asarray(dp_mod.call(x)), np.asarray(dp_mod2.call(x)) + ) + + def test_ffn_guards(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.ffn import EquivariantFFN as DPFFN + + with pytest.raises(NotImplementedError, match="ffn_so3_grid"): + DPFFN(**self._ffn_kwargs(ffn_so3_grid=True), precision="float64") + # grid_mlp guard is delegated to S2GridNet's op_type='mlp' NIE + with pytest.raises(NotImplementedError, match="mlp"): + DPFFN( + **self._ffn_kwargs(s2_activation=True, grid_mlp=True), + precision="float64", + ) + + def test_ffn_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.ffn import EquivariantFFN as DPFFN + + with pytest.raises(ValueError): # kmax must be non-negative + DPFFN(**self._ffn_kwargs(kmax=-1), precision="float64") + with pytest.raises(ValueError): # grid_branch must be non-negative + DPFFN(**self._ffn_kwargs(grid_branch=-1), precision="float64") + with pytest.raises(ValueError): # wrong class tag + DPFFN.deserialize({"@class": "NotFFN", "@version": 1}) + dp_mod = DPFFN(**self._ffn_kwargs(), precision="float64", seed=3) + with pytest.raises(KeyError): # missing sub-module variables + dp_mod._load_variables({"so3_linear_1.weight": dp_mod.so3_linear_1.weight}) + with pytest.raises(KeyError): # unknown variables rejected + dp_mod._load_variables({**dp_mod._variables(), "extra.weight": 0.0}) + + +class TestBlockParity: + nloc = 5 + nnei = 4 + channels = 4 + + def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: + rng = np.random.default_rng(seed) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.1 * rng.normal(size=tuple(p.shape))) + + def _block_kwargs(self, **overrides): + # core DPA4 config (sandwich_norm=[F,T,T,F], so2 s2 off, ffn s2 on) + kwargs = { + "lmax": 3, + "node_lmax": None, + "mmax": 1, + "kmax": 1, + "channels": self.channels, + "n_focus": 1, + "focus_dim": 0, + "focus_compete": True, + "so2_norm": False, + "so2_layers": 4, + "so2_attn_res": "none", + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_atten_head": 1, + "so2_pre_norm": False, + "so2_post_norm": True, + "ffn_pre_norm": True, + "ffn_post_norm": False, + "ffn_neurons": self.channels, + "ffn_grid_branch": 1, + "ffn_blocks": 1, + "ffn_s2_activation": True, + "so2_lebedev_quadrature": True, + "ffn_lebedev_quadrature": True, + "so2_activation_function": "silu", + "ffn_activation_function": "silu", + "ffn_glu_activation": True, + "mlp_bias": False, + "eps": 1e-7, + } + kwargs.update(overrides) + return kwargs + + def _build_block_pair(self, seed=31, perturb_seed=2120, **overrides): + from deepmd.dpmodel.descriptor.dpa4_nn.block import ( + SeZMInteractionBlock as DPBlock, + ) + from deepmd.pt.model.descriptor.sezm_nn.block import ( + SeZMInteractionBlock as PTBlock, + ) + + kwargs = self._block_kwargs(**overrides) + pt_mod = PTBlock(**kwargs, dtype=torch.float64, seed=seed, trainable=True) + # zero-initialized residual projections; perturb so the output is nonzero + self._perturb(pt_mod, perturb_seed) + dp_mod = DPBlock.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod, kwargs + + def _node_dim(self, kwargs): + node_lmax = kwargs["node_lmax"] + if node_lmax is None: + node_lmax = kwargs["lmax"] + return (node_lmax + 1) ** 2 + + def _assert_block_parity(self, pt_mod, dp_mod, kwargs, *, masked="slots") -> None: + rng = np.random.default_rng(2121) + pt_cache, dp_cache, radial, radial_valid, _, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=kwargs["lmax"], + channels=kwargs["channels"], + masked=masked, + ) + node_dim = self._node_dim(kwargs) + x = rng.normal(size=(self.nloc, node_dim, 1, kwargs["channels"])) + out_dp = dp_mod.call(x, dp_cache, radial) + out_pt = pt_mod(to_pt(x), pt_cache, to_pt(radial_valid)) + assert out_dp[1:] == (None, None, None) + assert out_pt[1] is None and out_pt[2] is None and out_pt[3] is None + assert_parity(out_dp[0], out_pt[0]) + + @pytest.mark.parametrize("so2_layers", [2, 4]) # SO(2) layer depth (core=4) + def test_block(self, so2_layers) -> None: + pt_mod, dp_mod, kwargs = self._build_block_pair(so2_layers=so2_layers) + self._assert_block_parity(pt_mod, dp_mod, kwargs) + + @pytest.mark.parametrize( + "sandwich", + [ + (False, True, True, False), # core [so2_pre, so2_post, ffn_pre, ffn_post] + (True, False, False, True), # flips every norm flag's branch + ], + ) + def test_block_sandwich_norm(self, sandwich) -> None: + so2_pre, so2_post, ffn_pre, ffn_post = sandwich + pt_mod, dp_mod, kwargs = self._build_block_pair( + so2_pre_norm=so2_pre, + so2_post_norm=so2_post, + ffn_pre_norm=ffn_pre, + ffn_post_norm=ffn_post, + ) + self._assert_block_parity(pt_mod, dp_mod, kwargs) + + def test_block_ffn_blocks(self) -> None: + # multiple FFN subblocks exercise the per-subblock loop and seeds + pt_mod, dp_mod, kwargs = self._build_block_pair(ffn_blocks=2) + self._assert_block_parity(pt_mod, dp_mod, kwargs) + + def test_block_node_lmax(self) -> None: + # node_lmax > lmax: SO(2) acts on the truncated slice, zero-pads above + pt_mod, dp_mod, kwargs = self._build_block_pair(lmax=2, node_lmax=3, mmax=1) + self._assert_block_parity(pt_mod, dp_mod, kwargs) + + def test_block_mlp_bias(self) -> None: + pt_mod, dp_mod, kwargs = self._build_block_pair(mlp_bias=True) + self._assert_block_parity(pt_mod, dp_mod, kwargs) + + def test_block_plain_ffn_act(self) -> None: + # ffn_s2_activation=False: FFN uses the GatedActivation path + pt_mod, dp_mod, kwargs = self._build_block_pair( + ffn_s2_activation=False, ffn_grid_branch=0 + ) + self._assert_block_parity(pt_mod, dp_mod, kwargs) + + def test_block_real_edge_cache(self) -> None: + # end-to-end: REAL pt build_edge_cache vs REAL dp build_edge_cache + # feeding the same weight-copied block (no synthetic cache) + pt_mod, dp_mod, kwargs = self._build_block_pair() + rng = np.random.default_rng(2122) + nf, nloc, nall, nnei = 1, self.nloc, self.nloc + 3, self.nnei + inputs = _build_real_edge_inputs( + rng, + nf=nf, + nloc=nloc, + nall=nall, + nnei=nnei, + channels=kwargs["channels"], + ) + pt_cache, dp_cache = _build_real_edge_caches(inputs, lmax=kwargs["lmax"]) + valid = inputs["valid"] + node_dim = self._node_dim(kwargs) + radial = rng.normal( + size=(nf * nloc * nnei, kwargs["lmax"] + 1, kwargs["channels"]) + ) + x = rng.normal(size=(nf * nloc, node_dim, 1, kwargs["channels"])) + out_dp = dp_mod.call(x, dp_cache, radial) + out_pt = pt_mod(to_pt(x), pt_cache, to_pt(radial[valid])) + assert_parity(out_dp[0], out_pt[0]) + # masked-slot garbage is inert end-to-end: scribble into invalid slots + # (finite O(10) garbage: the padded layout computes exp() on masked + # attention logits before zero-weighting them, so the garbage must + # not overflow exp; see attention.segment_envelope_gated_softmax) + radial2 = radial.copy() + radial2[~valid] = 10.0 + out_dp2 = dp_mod.call(x, dp_cache, radial2) + np.testing.assert_array_equal(np.asarray(out_dp[0]), np.asarray(out_dp2[0])) + + def test_block_roundtrip(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.block import ( + SeZMInteractionBlock as DPBlock, + ) + + pt_mod, dp_mod, kwargs = self._build_block_pair(ffn_blocks=2) + data = dp_mod.serialize() + # exact pt state_dict key-set match + assert set(data["@variables"]) == set(pt_state_to_numpy(pt_mod)) + dp_mod2 = DPBlock.deserialize(data) + rng = np.random.default_rng(2123) + _, dp_cache, radial, _, _, _ = _build_so2_edge_data( + rng, + nloc=self.nloc, + nnei=self.nnei, + lmax=kwargs["lmax"], + channels=kwargs["channels"], + masked="slots", + ) + x = rng.normal(size=(self.nloc, self._node_dim(kwargs), 1, self.channels)) + out1 = np.asarray(dp_mod.call(x, dp_cache, radial)[0]) + out2 = np.asarray(dp_mod2.call(x, dp_cache, radial)[0]) + np.testing.assert_array_equal(out1, out2) + + @pytest.mark.parametrize( + "flag,value", + [ + ("full_attn_res", "independent"), # block-level DepthAttnRes + ("full_attn_res", "dependent"), # block-level DepthAttnRes + ("block_attn_res", "independent"), # block-level DepthAttnRes + ("block_attn_res", "dependent"), # block-level DepthAttnRes + ("layer_scale", True), # block-level FFN LayerScale + ("so2_s2_activation", True), # delegated to SO2Convolution + ("node_wise_s2", True), # delegated to SO2Convolution + ("message_node_so3", True), # delegated to SO2Convolution + ("ffn_so3_grid", True), # delegated to EquivariantFFN + ], + ) + def test_block_guards(self, flag, value) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.block import ( + SeZMInteractionBlock as DPBlock, + ) + + match = "s2_activation" if flag == "so2_s2_activation" else flag + with pytest.raises(NotImplementedError, match=match): + DPBlock(**self._block_kwargs(**{flag: value}), precision="float64") + + def test_block_errors(self) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.block import ( + SeZMInteractionBlock as DPBlock, + ) + + with pytest.raises(ValueError): # node_lmax must be >= lmax + DPBlock(**self._block_kwargs(node_lmax=2), precision="float64") + with pytest.raises(ValueError): # mmax must be <= lmax + DPBlock(**self._block_kwargs(mmax=4), precision="float64") + with pytest.raises(ValueError): # ffn_blocks must be >= 1 + DPBlock(**self._block_kwargs(ffn_blocks=0), precision="float64") + with pytest.raises(ValueError): # unknown full_attn_res token + DPBlock(**self._block_kwargs(full_attn_res="depth"), precision="float64") + with pytest.raises(ValueError): # unknown block_attn_res token + DPBlock(**self._block_kwargs(block_attn_res="depth"), precision="float64") + with pytest.raises(ValueError): # negative grid branch count + DPBlock(**self._block_kwargs(ffn_grid_branch=-1), precision="float64") + with pytest.raises(ValueError): # wrong class tag + DPBlock.deserialize({"@class": "NotBlock", "@version": 1}) + + +def _build_descriptor_inputs(rng, *, nf, nloc, nall, nnei, ntypes=3): + """Build a real two-frame descriptor fixture with ghosts and mapping. + + Extends the Task-10 ``_build_real_edge_inputs`` fixture with consistent + extended atom types (``atype_ext`` derived from the local types via the + ghost mapping); includes -1 padding slots and one broken mapping entry. + """ + inputs = _build_real_edge_inputs( + rng, nf=nf, nloc=nloc, nall=nall, nnei=nnei, channels=1 + ) + mapping = inputs["mapping"] + atype_loc = rng.integers(0, ntypes, size=(nf, nloc)) + atype_ext = np.take_along_axis(atype_loc, np.clip(mapping, 0, nloc - 1), axis=1) + return { + "coord": inputs["coord"], # (nf, nall, 3) + "atype_ext": atype_ext, # (nf, nall) + "nlist": inputs["nlist"], # (nf, nloc, nnei), -1 padded + "mapping": mapping, # (nf, nall), one -1 entry + } + + +class TestDescriptorParity: + nf = 2 + nloc = 6 + nall = 10 + nnei = 12 + + def _descr_kwargs(self, **overrides): + # small core DPA4 config (see task spec) + kwargs = { + "ntypes": 3, + "sel": self.nnei, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 3, + "mmax": 1, + "n_blocks": 2, + "grid_branch": [1, 1, 1], + "s2_activation": [False, True], + "random_gamma": False, + "exclude_types": [(0, 0)], + "precision": "float64", + "seed": 42, + } + kwargs.update(overrides) + return kwargs + + def _build_descr_pair(self, perturb_seed=2130, **overrides): + from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, + ) + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + kwargs = self._descr_kwargs(**overrides) + pt_mod = DescrptSeZM(**kwargs).double().eval() + # several projections are zero-initialized; perturb for nonzero output + rng = np.random.default_rng(perturb_seed) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += to_pt(0.05 * rng.normal(size=tuple(p.shape))) + dp_mod = DescrptDPA4.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod, kwargs + + def _inputs(self, seed=2131): + rng = np.random.default_rng(seed) + return _build_descriptor_inputs( + rng, nf=self.nf, nloc=self.nloc, nall=self.nall, nnei=self.nnei + ) + + def _assert_descr_parity(self, pt_mod, dp_mod, *, mapping=True) -> None: + inp = self._inputs() + coord, atype_ext, nlist, mp = ( + inp["coord"], + inp["atype_ext"], + inp["nlist"], + inp["mapping"], + ) + if not mapping: + # mapping-free path: neighbor indices already local (no ghosts) + rng = np.random.default_rng(2132) + inp_local = _build_real_edge_inputs( + rng, + nf=self.nf, + nloc=self.nloc, + nall=self.nloc, + nnei=self.nnei, + channels=1, + local_nlist=True, + ) + coord, nlist, mp = inp_local["coord"], inp_local["nlist"], None + atype_ext = rng.integers(0, 3, size=(self.nf, self.nloc)) + nf = coord.shape[0] + out_dp = dp_mod.call( + coord.reshape(nf, -1), + atype_ext, + nlist, + mapping=mp, + ) + out_pt = pt_mod( + to_pt(coord), + to_pt(atype_ext), + to_pt(nlist), + mapping=None if mp is None else to_pt(mp), + ) + assert out_dp[0].shape == tuple(out_pt[0].shape) + # descriptor-level tolerance: rtol 1e-10 / atol 1e-12 + assert_parity(out_dp[0], out_pt[0], rtol=1e-10, atol=1e-12) + # unused returns are None on the dp side (pt returns empty tensors) + assert out_dp[1:] == (None, None, None, None) + + @pytest.mark.parametrize("use_env_seed", [False, True]) # env FiLM + GIE seeding + @pytest.mark.parametrize("n_blocks", [1, 2]) # interaction block stack depth + def test_descriptor(self, use_env_seed, n_blocks) -> None: + pt_mod, dp_mod, _ = self._build_descr_pair( + use_env_seed=use_env_seed, n_blocks=n_blocks + ) + self._assert_descr_parity(pt_mod, dp_mod) + + @pytest.mark.parametrize( + "exclude_types", [[], [(0, 0)]] + ) # pair-exclusion off vs on + def test_descriptor_exclude_types(self, exclude_types) -> None: + pt_mod, dp_mod, _ = self._build_descr_pair(exclude_types=exclude_types) + self._assert_descr_parity(pt_mod, dp_mod) + + def test_descriptor_no_mapping(self) -> None: + # pt forward accepts mapping=None when neighbor indices are local; + # mapping is NOT required by either backend + pt_mod, dp_mod, _ = self._build_descr_pair() + self._assert_descr_parity(pt_mod, dp_mod, mapping=False) + + def test_descriptor_extra_node_l(self) -> None: + # node degrees above message-passing degrees (GIE zonal wigner path) + pt_mod, dp_mod, _ = self._build_descr_pair(extra_node_l=1) + self._assert_descr_parity(pt_mod, dp_mod) + + def test_descriptor_torch_namespace(self) -> None: + # the dp descriptor must run under the torch array namespace as well: + # feeding torch tensors must yield a torch tensor matching the numpy + # result (catches raw numpy attributes mixed into xp arithmetic) + _, dp_mod, _ = self._build_descr_pair(use_env_seed=True) + inp = self._inputs() + nf = inp["coord"].shape[0] + coord = inp["coord"].reshape(nf, -1) + atype_ext, nlist, mp = inp["atype_ext"], inp["nlist"], inp["mapping"] + out_np = dp_mod.call(coord, atype_ext, nlist, mapping=mp)[0] + # CPU on purpose: this pins the dp class's torch-namespace + # behavior (not device placement); CPU keeps the dp-vs-dp compare + # at the strict device-independent gate. + out_t = dp_mod.call( + torch.from_numpy(coord).to(device="cpu"), + torch.from_numpy(atype_ext.astype(np.int64)).to(device="cpu"), + torch.from_numpy(nlist.astype(np.int64)).to(device="cpu"), + mapping=torch.from_numpy(mp.astype(np.int64)).to(device="cpu"), + )[0] + assert isinstance(out_t, torch.Tensor) + np.testing.assert_allclose( + out_t.numpy(), np.asarray(out_np), rtol=1e-12, atol=1e-14 + ) + + def test_descriptor_cross_deserialize(self) -> None: + from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, + ) + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + pt_mod, dp_mod, _ = self._build_descr_pair() + # dp serialize emits exactly the pt state_dict key set + data = dp_mod.serialize() + assert set(data["@variables"]) == set(pt_state_to_numpy(pt_mod)) + assert data["type"] == "SeZM" + # pt <- dp: load the dp serialization into a fresh pt descriptor + pt_mod2 = DescrptSeZM.deserialize(data).double().eval() + self._assert_descr_parity(pt_mod2, dp_mod) + # dp <- dp roundtrip is bit-exact + dp_mod2 = DescrptDPA4.deserialize(data) + inp = self._inputs() + nf = inp["coord"].shape[0] + args = ( + inp["coord"].reshape(nf, -1), + inp["atype_ext"], + inp["nlist"], + ) + out1 = np.asarray(dp_mod.call(*args, mapping=inp["mapping"])[0]) + out2 = np.asarray(dp_mod2.call(*args, mapping=inp["mapping"])[0]) + np.testing.assert_array_equal(out1, out2) + + +class TestNoTorchImport: + def test_dpa4_nn_does_not_import_torch(self) -> None: + code = ( + "import sys; " + "import deepmd.dpmodel.descriptor.dpa4_nn.indexing, " + "deepmd.dpmodel.descriptor.dpa4_nn.utils, " + "deepmd.dpmodel.descriptor.dpa4_nn.norm, " + "deepmd.dpmodel.descriptor.dpa4_nn.radial, " + "deepmd.dpmodel.descriptor.dpa4_nn.so3, " + "deepmd.dpmodel.descriptor.dpa4_nn.activation, " + "deepmd.dpmodel.descriptor.dpa4_nn.wignerd, " + "deepmd.dpmodel.descriptor.dpa4_nn.projection, " + "deepmd.dpmodel.descriptor.dpa4_nn.grid_net, " + "deepmd.dpmodel.descriptor.dpa4_nn.so2, " + "deepmd.dpmodel.descriptor.dpa4_nn.attention, " + "deepmd.dpmodel.descriptor.dpa4_nn.edge_cache, " + "deepmd.dpmodel.descriptor.dpa4_nn.embedding, " + "deepmd.dpmodel.descriptor.dpa4_nn.ffn, " + "deepmd.dpmodel.descriptor.dpa4_nn.block, " + "deepmd.dpmodel.descriptor.dpa4; " + "print('torch' in sys.modules)" + ) + out = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ) + assert out.stdout.strip() == "False" + + +class TestFittingParity: + nf = 2 + nloc = 6 + in_dim = 12 + ntypes = 2 + + def _build_pair(self, **overrides): + from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetDP, + ) + from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetPT, + ) + + kwargs = { + "ntypes": self.ntypes, + "dim_descrpt": self.in_dim, + "neuron": [0], + "precision": "float64", + "seed": 5, + } + kwargs.update(overrides) + pt_mod = SeZMEnergyFittingNetPT(**kwargs).eval() + # bias_atom_e is zero-initialized; perturb for a nontrivial bias path + rng = np.random.default_rng(2140) + with torch.no_grad(): + pt_mod.bias_atom_e += to_pt( + rng.normal(size=tuple(pt_mod.bias_atom_e.shape)) + ) + dp_mod = SeZMEnergyFittingNetDP.deserialize(pt_mod.serialize()) + return pt_mod, dp_mod + + def _inputs(self, seed=2141): + rng = np.random.default_rng(seed) + descriptor = rng.normal(size=(self.nf, self.nloc, self.in_dim)) + # cover both atom types + atype = rng.integers(0, self.ntypes, size=(self.nf, self.nloc)) + atype[0, 0], atype[0, 1] = 0, 1 + return descriptor, atype + + def _assert_fitting_parity(self, pt_mod, dp_mod, fparam=None, aparam=None): + descriptor, atype = self._inputs() + out_dp = dp_mod.call(descriptor, atype, fparam=fparam, aparam=aparam)["energy"] + out_pt = pt_mod( + to_pt(descriptor), + to_pt(atype), + fparam=None if fparam is None else to_pt(fparam), + aparam=None if aparam is None else to_pt(aparam), + )["energy"] + assert out_dp.shape == tuple(out_pt.shape) + assert_parity(out_dp, out_pt) + + @pytest.mark.parametrize("bias_out", [False, True]) # output-layer bias + @pytest.mark.parametrize( + "neuron", [[0], [32], [16, 16], []] + ) # auto-width / fixed / deep / direct linear + def test_fitting(self, neuron, bias_out) -> None: + pt_mod, dp_mod = self._build_pair(neuron=neuron, bias_out=bias_out) + self._assert_fitting_parity(pt_mod, dp_mod) + + def test_fitting_fparam_aparam(self) -> None: + pt_mod, dp_mod = self._build_pair(numb_fparam=2, numb_aparam=3) + rng = np.random.default_rng(2142) + fparam = rng.normal(size=(self.nf, 2)) + aparam = rng.normal(size=(self.nf, self.nloc, 3)) + self._assert_fitting_parity(pt_mod, dp_mod, fparam=fparam, aparam=aparam) + + def test_fitting_default_fparam(self) -> None: + # fparam=None falls back to the default frame parameter on both sides + pt_mod, dp_mod = self._build_pair(numb_fparam=2, default_fparam=[0.5, -1.5]) + self._assert_fitting_parity(pt_mod, dp_mod) + + def test_fitting_not_mixed_types(self) -> None: + # one GLU net per atom type + pt_mod, dp_mod = self._build_pair(mixed_types=False) + self._assert_fitting_parity(pt_mod, dp_mod) + + def test_fitting_exclude_types(self) -> None: + pt_mod, dp_mod = self._build_pair(exclude_types=[0]) + self._assert_fitting_parity(pt_mod, dp_mod) + + @pytest.mark.parametrize("bias_out", [False, True]) # output-layer bias + def test_fitting_cross_deserialize(self, bias_out) -> None: + from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetDP, + ) + from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetPT, + ) + + pt_mod, dp_mod = self._build_pair(neuron=[16], bias_out=bias_out) + data = dp_mod.serialize() + assert data["type"] == "sezm_ener" + # the dp serialization carries exactly the pt state_dict key set + flat = {k: v for k, v in data["@variables"].items() if v is not None} + for ii, net in enumerate(data["nets"]["networks"]): + for kk, vv in net["@variables"].items(): + flat[f"filter_layers.networks.{ii}.{kk}"] = vv + assert set(flat) == set(pt_state_to_numpy(pt_mod)) + # serialized dict key sets match between backends + assert set(data) == set(pt_mod.serialize()) + # pt <- dp + pt_mod2 = SeZMEnergyFittingNetPT.deserialize(data).eval() + self._assert_fitting_parity(pt_mod2, dp_mod) + # dp <- dp roundtrip is bit-exact + dp_mod2 = SeZMEnergyFittingNetDP.deserialize(data) + descriptor, atype = self._inputs() + out1 = np.asarray(dp_mod.call(descriptor, atype)["energy"]) + out2 = np.asarray(dp_mod2.call(descriptor, atype)["energy"]) + np.testing.assert_array_equal(out1, out2) + + +class TestEndToEndParity: + """Chain descriptor and fitting: full dpmodel DPA4 atomic-energy math.""" + + def test_descriptor_fitting_chain(self) -> None: + from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetDP, + ) + from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetPT, + ) + + helper = TestDescriptorParity() + pt_descr, dp_descr, _ = helper._build_descr_pair() + in_dim = dp_descr.get_dim_out() + pt_fit = SeZMEnergyFittingNetPT( + ntypes=3, + dim_descrpt=in_dim, + neuron=[0], + precision="float64", + seed=11, + ).eval() + rng = np.random.default_rng(2143) + with torch.no_grad(): + pt_fit.bias_atom_e += to_pt( + rng.normal(size=tuple(pt_fit.bias_atom_e.shape)) + ) + dp_fit = SeZMEnergyFittingNetDP.deserialize(pt_fit.serialize()) + + inp = helper._inputs() + coord, atype_ext, nlist, mp = ( + inp["coord"], + inp["atype_ext"], + inp["nlist"], + inp["mapping"], + ) + nf, nloc = nlist.shape[:2] + atype_loc = atype_ext[:, :nloc] + d_dp = dp_descr.call(coord.reshape(nf, -1), atype_ext, nlist, mapping=mp)[0] + e_dp = dp_fit.call(d_dp, atype_loc)["energy"] + d_pt = pt_descr( + to_pt(coord), + to_pt(atype_ext), + to_pt(nlist), + mapping=to_pt(mp), + )[0] + e_pt = pt_fit(d_pt, to_pt(atype_loc))["energy"] + assert e_dp.shape == tuple(e_pt.shape) + # end-to-end tolerance: rtol 1e-10 / atol 1e-12 + assert_parity(e_dp, e_pt, rtol=1e-10, atol=1e-12) + + +class TestModelDefCompat: + """Pin the pt SeZM energy model serialized-dict contract for dpmodel. + + Full dpmodel model assembly (SeZMModel / sezm_atomic in dpmodel) is + PR-2/PR-3 scope; until then these tests pin the contract: the pt model's + serialized descriptor/fitting sub-dicts must deserialize via the dpmodel + classes, and the out-of-scope top-level fields must stay disabled for + the core config. + """ + + # top-level keys of SeZMModel.serialize() the dpmodel port relies on + KNOWN_TOP_LEVEL_KEYS = frozenset( + { + "@class", + "@version", + "type", + "atomic_model", + "bridging_method", + "bridging_r_inner", + "bridging_r_outer", + "lora", + } + ) + + def _build_model(self): + from deepmd.pt.model.model import ( + get_model, + ) + + cfg = { + "type": "dpa4", + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4", + "sel": 10, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 42, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [0], + "precision": "float64", + "seed": 42, + }, + } + return get_model(cfg) + + def test_serialized_subdicts_deserialize_via_dpmodel(self) -> None: + from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, + ) + from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, + ) + + model = self._build_model() + data = model.serialize() + atomic = data["atomic_model"] + assert "descriptor" in atomic, sorted(atomic) + assert "fitting" in atomic, sorted(atomic) + dp_descr = DescrptDPA4.deserialize(atomic["descriptor"]) + dp_fit = SeZMEnergyFittingNet.deserialize(atomic["fitting"]) + assert dp_fit.dim_descrpt == dp_descr.get_dim_out() + + def test_top_level_fields_pinned(self) -> None: + model = self._build_model() + data = model.serialize() + unknown = set(data) - self.KNOWN_TOP_LEVEL_KEYS + assert not unknown, ( + f"pt SeZMModel.serialize() gained new top-level field(s) {sorted(unknown)}; " + "the dpmodel DPA4 model port (PR-2/PR-3) must be updated to handle them " + "before this contract test is extended." + ) + # out-of-scope features must be disabled for the core config + assert str(data["bridging_method"]).lower() == "none", data["bridging_method"] + assert data["lora"] is None, data["lora"] + # core-config atomic_model must have no density fitting attached + assert data["atomic_model"].get("dens_fitting") is None