diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index d7b4287199..8d59a5caba 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -176,6 +176,15 @@ class DescrptDPA4(NativeOP, BaseDescriptor): radial_mlp Hidden layer sizes for radial networks. An output layer of size `(l_schedule[0]+extra_node_l+1)*channels` will be automatically appended. + edge_norm + Whether to apply channel RMSNorm on the descriptor's cutoff-vanishing + branches: the radial network hidden layers, the environment-seed FiLM + scale/shift logits, the cross-focus competition scalars, and the + post-SO(2) residual messages. ``False`` replaces the first three norms + with identity and changes only the post-SO(2) norm to unit-floor residual + scaling. The unit floor uses ``sqrt(1 + variance)`` so small messages + retain their cutoff envelope instead of receiving the standard + ``1/sqrt(eps)`` small-signal gain. use_env_seed If True, seed the initial node state with local-environment information: apply environment matrix FiLM conditioning on l=0 features using 4D @@ -457,6 +466,7 @@ def __init__( basis_type: str = "bessel", n_radial: int = 16, radial_mlp: list[int] | None = None, + edge_norm: bool = True, use_env_seed: bool = True, random_gamma: bool = True, edge_cartesian: bool = False, @@ -551,6 +561,7 @@ def __init__( if radial_mlp is None: radial_mlp = [0] self.radial_mlp = [self.channels if x == 0 else int(x) for x in radial_mlp] + self.edge_norm = bool(edge_norm) if sandwich_norm is None: sandwich_norm = [False, True, True, False] if not isinstance(sandwich_norm, (list, tuple)) or len(sandwich_norm) != 4: @@ -849,20 +860,28 @@ def __init__( seed=seed_env_seed, ) ) - self.film_scale_norm = ScalarRMSNorm( - channels=self.channels, - n_focus=1, - eps=self.eps, - precision=self.compute_precision, - trainable=self.trainable, - ) - self.film_shift_norm = ScalarRMSNorm( - channels=self.channels, - n_focus=1, - eps=self.eps, - precision=self.compute_precision, - trainable=self.trainable, - ) + # The FiLM logits derive from the env-seed matrix D = envᵀenv, which + # vanishes at rcut; normalizing them shares the radial network's + # cutoff-smoothness issue, so ``edge_norm=False`` also drops these + # norms (identity pass-through) to keep the FiLM scale/shift smooth. + if self.edge_norm: + self.film_scale_norm = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + precision=self.compute_precision, + trainable=self.trainable, + ) + self.film_shift_norm = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + precision=self.compute_precision, + trainable=self.trainable, + ) + else: + self.film_scale_norm = None + self.film_shift_norm = None film_strength_init = 0.01 # Use 1D tensor (not scalar) for FSDP2 compatibility self.film_scale_strength_log = np.full( @@ -902,6 +921,7 @@ def __init__( activation_function=self.activation_function, precision=self.compute_precision, # force fp32+ trainable=self.trainable, + radial_norm=self.edge_norm, seed=seed_radial_embedding, ) @@ -964,6 +984,7 @@ def __init__( channels=self.channels, n_focus=self.n_focus, focus_dim=self.focus_dim, + focus_norm=self.edge_norm, so2_norm=self.so2_norm, mixing_layers=self.mixing_layers, so2_attn_res=self.so2_attn_res_mode, @@ -998,6 +1019,7 @@ def __init__( atten_o_proj=self.use_atten_o_proj, so2_pre_norm=self.so2_pre_norm, so2_post_norm=self.so2_post_norm, + so2_post_norm_eps=1.0e-5 if self.edge_norm else 1.0, so2_activation_function=self.so2_activation_function, ffn_pre_norm=self.ffn_pre_norm, ffn_post_norm=self.ffn_post_norm, @@ -1267,8 +1289,12 @@ def call( ) # (N, 2*C) scale_logits = film[:, : self.channels] # (N, C) shift_logits = film[:, self.channels :] # (N, C) - scale_hat = self.film_scale_norm(scale_logits) # (N, C) - shift_hat = self.film_shift_norm(shift_logits) # (N, C) + scale_hat = ( + self.film_scale_norm(scale_logits) if self.edge_norm else scale_logits + ) # (N, C) + shift_hat = ( + self.film_shift_norm(shift_logits) if self.edge_norm else shift_logits + ) # (N, C) scale_strength = xp.exp( xp_asarray_nodetach( xp, self.film_scale_strength_log[...], device=device @@ -1522,8 +1548,12 @@ def call_with_edges( ) # (N, 2*C) scale_logits = film[:, : self.channels] # (N, C) shift_logits = film[:, self.channels :] # (N, C) - scale_hat = self.film_scale_norm(scale_logits) # (N, C) - shift_hat = self.film_shift_norm(shift_logits) # (N, C) + scale_hat = ( + self.film_scale_norm(scale_logits) if self.edge_norm else scale_logits + ) # (N, C) + shift_hat = ( + self.film_shift_norm(shift_logits) if self.edge_norm else shift_logits + ) # (N, C) scale_strength = xp.exp( xp_asarray_nodetach( xp, self.film_scale_strength_log[...], device=device @@ -2394,10 +2424,15 @@ def _variables(self) -> dict[str, np.ndarray]: if self.use_env_seed: for key, value in self.env_seed_embedding.serialize()["@variables"].items(): variables[f"env_seed_embedding.{key}"] = value - for key, value in self.film_scale_norm.serialize()["@variables"].items(): - variables[f"film_scale_norm.{key}"] = value - for key, value in self.film_shift_norm.serialize()["@variables"].items(): - variables[f"film_shift_norm.{key}"] = value + if self.edge_norm: + for key, value in self.film_scale_norm.serialize()[ + "@variables" + ].items(): + variables[f"film_scale_norm.{key}"] = value + for key, value in self.film_shift_norm.serialize()[ + "@variables" + ].items(): + variables[f"film_shift_norm.{key}"] = value variables["film_scale_strength_log"] = to_numpy_array( self.film_scale_strength_log ) @@ -2496,8 +2531,9 @@ def load(module: Any, prefix: str) -> Any: self.env_seed_embedding = load( self.env_seed_embedding, "env_seed_embedding." ) - self.film_scale_norm = load(self.film_scale_norm, "film_scale_norm.") - self.film_shift_norm = load(self.film_shift_norm, "film_shift_norm.") + if self.edge_norm: + self.film_scale_norm = load(self.film_scale_norm, "film_scale_norm.") + self.film_shift_norm = load(self.film_shift_norm, "film_shift_norm.") self.film_scale_strength_log = np.asarray( variables["film_scale_strength_log"], dtype=compute_prec ) @@ -2547,6 +2583,7 @@ def serialize(self) -> dict[str, Any]: "basis_type": self.basis_type, "n_radial": self.n_radial, "radial_mlp": self.radial_mlp, + "edge_norm": self.edge_norm, "use_env_seed": self.use_env_seed, "random_gamma": self.random_gamma, "edge_cartesian": self.edge_cartesian, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/block.py b/deepmd/dpmodel/descriptor/dpa4_nn/block.py index c13acf2617..751efe3f95 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/block.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/block.py @@ -145,6 +145,10 @@ class SeZMInteractionBlock(NativeOP): ``focus_dim=0`` means using ``channels``. focus_compete If True, enable cross-focus softmax competition in SO(2) convolution. + focus_norm + If True, RMS-normalize the cross-focus competition scalars before the + softmax. The competition input is envelope-gated and vanishes at the + cutoff, so ``False`` drops the norm to keep the competition smooth there. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm between SO(2) mixing layers. When False (default), no normalization is applied between layers. @@ -192,6 +196,10 @@ class SeZMInteractionBlock(NativeOP): If True, apply pre-norm before SO(2) convolution. so2_post_norm If True, apply post-norm on SO(2) output before the residual add. + so2_post_norm_eps + Variance floor for the SO(2) post-norm. A value of ``1`` preserves small + residual messages instead of amplifying them by ``1/sqrt(eps)``. Other + normalization sites retain their own RMSNorm floors. ffn_pre_norm If True, apply pre-norm before each FFN subblock. ffn_post_norm @@ -293,6 +301,7 @@ def __init__( n_focus: int = 1, focus_dim: int = 0, focus_compete: bool = True, + focus_norm: bool = True, so2_norm: bool = False, mixing_layers: int = 4, so2_attn_res: str = "none", @@ -306,6 +315,7 @@ def __init__( atten_o_proj: bool = False, so2_pre_norm: bool = True, so2_post_norm: bool = False, + so2_post_norm_eps: float = 1e-5, ffn_pre_norm: bool = True, ffn_post_norm: bool = False, ffn_neurons: int = 96, @@ -359,6 +369,7 @@ def __init__( if self.focus_dim < 0: raise ValueError("`focus_dim` must be >= 0") self.focus_compete = bool(focus_compete) + self.focus_norm = bool(focus_norm) self.so2_norm = bool(so2_norm) self.mixing_layers = int(mixing_layers) self.so2_attn_res_mode = str(so2_attn_res).lower() @@ -376,6 +387,7 @@ def __init__( 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.so2_post_norm_eps = float(so2_post_norm_eps) self.ffn_pre_norm = bool(ffn_pre_norm) self.ffn_post_norm = bool(ffn_post_norm) self.ffn_neurons = int(ffn_neurons) @@ -456,6 +468,7 @@ def __init__( self.lmax, self.channels, n_focus=1, + eps=self.so2_post_norm_eps, precision=self.compute_precision, trainable=trainable, ) @@ -470,6 +483,7 @@ def __init__( n_focus=self.n_focus, focus_dim=self.focus_dim, focus_compete=self.focus_compete, + focus_norm=self.focus_norm, so2_norm=self.so2_norm, mixing_layers=self.mixing_layers, so2_attn_res=self.so2_attn_res_mode, @@ -1065,6 +1079,7 @@ def serialize(self) -> dict[str, Any]: "n_focus": self.n_focus, "focus_dim": self.focus_dim, "focus_compete": self.focus_compete, + "focus_norm": self.focus_norm, "so2_norm": self.so2_norm, "mixing_layers": self.mixing_layers, "so2_attn_res": self.so2_attn_res_mode, @@ -1078,6 +1093,7 @@ def serialize(self) -> dict[str, Any]: "atten_o_proj": self.use_atten_o_proj, "so2_pre_norm": self.so2_pre_norm, "so2_post_norm": self.so2_post_norm, + "so2_post_norm_eps": self.so2_post_norm_eps, "ffn_pre_norm": self.ffn_pre_norm, "ffn_post_norm": self.ffn_post_norm, "ffn_neurons": self.ffn_neurons, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py index 3766588de9..6012cc4256 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py @@ -50,7 +50,7 @@ class RadialMLP(NativeOP): """ - Radial MLP with channel RMSNorm and configurable activation. + Radial MLP with optional channel RMSNorm and configurable activation. Parameters ---------- @@ -63,11 +63,14 @@ class RadialMLP(NativeOP): Floating point precision for the linear layers. trainable : bool Whether the parameters are trainable. + radial_norm : bool + Whether to insert a channel RMSNorm in each hidden layer. Architecture ------------ - Linear → RMSNorm → Activation for all hidden layers, - with the final layer being a plain Linear (no norm, no activation). + ``radial_norm=True`` : Linear → RMSNorm → Activation for each hidden layer. + ``radial_norm=False`` : Linear → Activation for each hidden layer. + The final layer is always a plain Linear (no norm, no activation). Notes ----- @@ -76,6 +79,15 @@ class RadialMLP(NativeOP): 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. + + The hidden RMSNorm normalizes each edge's radial features by their own RMS. + The input ``edge_rbf`` carries the C^3 cutoff envelope and therefore + vanishes at ``rcut``; the RMSNorm divides that envelope out, and its ``eps`` + floor is crossed as the edge approaches ``rcut``. On a sparse neighborhood + (e.g. a dimer) this floor-crossing produces a sharp kink in the potential + energy surface just inside the cutoff. Setting ``radial_norm=False`` drops + the RMSNorm so the radial features vanish smoothly with the envelope, which + restores C^3 smoothness at the cutoff. """ def __init__( @@ -85,6 +97,7 @@ def __init__( activation_function: str = "silu", precision: str = DEFAULT_PRECISION, trainable: bool = True, + radial_norm: bool = True, seed: int | list[int] | None = None, ) -> None: if len(mlp_layers) < 2: @@ -93,6 +106,7 @@ def __init__( self.activation_function = str(activation_function) self.precision = precision self.trainable = bool(trainable) + self.radial_norm = bool(radial_norm) modules: list = [] n_layers = len(mlp_layers) @@ -109,13 +123,14 @@ def __init__( modules.append(linear) # Last layer: no RMSNorm/activation if i < n_layers - 2: - modules.append( - RMSNorm( - channels=mlp_layers[i + 1], - precision=self.precision, - trainable=trainable, + if self.radial_norm: + modules.append( + RMSNorm( + channels=mlp_layers[i + 1], + precision=self.precision, + trainable=trainable, + ) ) - ) modules.append(get_activation_fn(self.activation_function)) self.net = modules @@ -153,6 +168,7 @@ def serialize(self) -> dict[str, Any]: "activation_function": self.activation_function, "dtype": np.dtype(PRECISION_DICT[self.precision]).name, "trainable": self.trainable, + "radial_norm": self.radial_norm, "@variables": variables, } diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index c0c4c0b262..4c8cd80e9a 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -941,6 +941,10 @@ class SO2Convolution(NativeOP): If True, apply cross-focus softmax competition in SO(2) local layout. Competition logits are constructed only from l=0 scalar channels and the resulting invariant weights are broadcast to all (l, m) components. + focus_norm + If True, RMS-normalize the competition l=0 scalars before the softmax. + Those scalars are envelope-gated and vanish at the cutoff, so ``False`` + drops the norm to keep the competition smooth there. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm as pre-norm before each SO(2) mixing layer. The last SO(2) layer always uses Identity. @@ -1056,6 +1060,7 @@ def __init__( n_focus: int = 1, focus_dim: int = 0, focus_compete: bool = True, + focus_norm: bool = True, so2_norm: bool = False, mixing_layers: int = 4, so2_attn_res: str = "none", @@ -1105,6 +1110,7 @@ def __init__( 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_norm = bool(focus_norm) self.focus_softmax_tau = 1.0 self.focus_label_smoothing = 0.02 self.so2_norm = bool(so2_norm) @@ -1367,13 +1373,17 @@ def __init__( 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, - ) + # ``focus_norm=False`` drops the competition-input RMSNorm (which + # would cross its eps floor as the envelope-gated scalars vanish at + # rcut); the competition weights then decay smoothly to uniform. + if self.focus_norm: + self.focus_compete_norm = ScalarRMSNorm( + channels=self.so2_focus_dim, + n_focus=self.n_focus, + eps=self.eps, + precision=self.compute_precision, + trainable=trainable, + ) self.adamw_focus_compete_w = ( np.random.default_rng(child_seed(seed_gate, 4)) .normal( @@ -2381,10 +2391,13 @@ def _focus_alpha(self, focus_gate_src: Array) -> Array: """ xp = array_api_compat.array_namespace(focus_gate_src) device = array_api_compat.device(focus_gate_src) + focus_in = xp.astype( + focus_gate_src, get_xp_precision(xp, self.compute_precision) + ) + if self.focus_norm: + focus_in = self.focus_compete_norm(focus_in) focus_logits = xp.sum( - self.focus_compete_norm( - xp.astype(focus_gate_src, get_xp_precision(xp, self.compute_precision)) - ) + focus_in * xp.permute_dims( xp_asarray_nodetach(xp, self.adamw_focus_compete_w[...], device=device), (1, 0), @@ -2706,6 +2719,7 @@ def serialize(self) -> dict[str, Any]: "n_focus": self.n_focus, "focus_dim": self.focus_dim, "focus_compete": self.focus_compete, + "focus_norm": self.focus_norm, "so2_norm": self.so2_norm, "mixing_layers": self.mixing_layers, "so2_attn_res": self.so2_attn_res_mode, diff --git a/deepmd/jax/atomic_model/linear_atomic_model.py b/deepmd/jax/atomic_model/linear_atomic_model.py index ae9bae6c4a..e25f8a8658 100644 --- a/deepmd/jax/atomic_model/linear_atomic_model.py +++ b/deepmd/jax/atomic_model/linear_atomic_model.py @@ -24,7 +24,7 @@ def __setattr__(self, name: str, value: Any) -> None: if name == "zbl_weight": # discard since it's only used in tests # to fix flax.errors.TraceContextError: Cannot mutate 'FlaxModule' from different trace level - return + return None return super().__setattr__(name, value) def forward_common_atomic( diff --git a/deepmd/jax/entrypoints/freeze.py b/deepmd/jax/entrypoints/freeze.py index fbc126ffc7..03536d4031 100644 --- a/deepmd/jax/entrypoints/freeze.py +++ b/deepmd/jax/entrypoints/freeze.py @@ -18,6 +18,7 @@ def freeze( *, checkpoint_folder: str, output: str, + hessian: bool = False, **kwargs: object, ) -> None: """Freeze a JAX checkpoint into a serialized model file. @@ -30,6 +31,8 @@ def freeze( output : str Output model filename or prefix. The JAX model suffix is added when the filename has no supported backend suffix. + hessian : bool, default=False + Whether to include the Hessian in the frozen model outputs. **kwargs Other CLI arguments accepted for backend entry-point compatibility. """ @@ -46,4 +49,4 @@ def freeze( strict_prefer=True, ) data = serialize_from_file(checkpoint_folder) - deserialize_to_file(output, data) + deserialize_to_file(output, data, hessian=hessian) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 8148f1e4fe..0e6c11ede6 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -314,9 +314,9 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: The requested output definitions. """ if atomic: - return list(self.output_def.var_defs.values()) + output_defs = list(self.output_def.var_defs.values()) else: - return [ + output_defs = [ x for x in self.output_def.var_defs.values() if x.category @@ -324,8 +324,18 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: OutputVariableCategory.REDU, OutputVariableCategory.DERV_R, OutputVariableCategory.DERV_C_REDU, + OutputVariableCategory.DERV_R_DERV_R, ) ] + # Avoid allocating the quadratic Hessian placeholder when the frozen + # model does not provide that output. + if not self.get_has_hessian(): + output_defs = [ + x + for x in output_defs + if x.category != OutputVariableCategory.DERV_R_DERV_R + ] + return output_defs def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: """Wrapper method with auto batch size. @@ -494,6 +504,10 @@ def get_model_def_script(self) -> dict: """Get model definition script.""" return json.loads(self.dp.get_model_def_script()) + def get_has_hessian(self) -> bool: + """Check if the model has Hessian output.""" + return self.get_model_def_script().get("hessian_mode", False) + def get_model(self) -> Any: """Get the JAX model as BaseModel. diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index a23733ce23..24ea448843 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -29,16 +29,22 @@ BaseModel, ) from deepmd.jax.utils.serialization import ( + _prepare_hessian_model_def_script, _set_model_min_nbor_dist_from_data, ) -def deserialize_to_file(model_file: str, data: dict) -> None: +def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None: """Deserialize the dictionary to a JAX/jax2tf SavedModel.""" if model_file.endswith(".savedmodel"): model = BaseModel.deserialize(data["model"]) _set_model_min_nbor_dist_from_data(model, data) - model_def_script = data["model_def_script"] + model_def_script, hessian = _prepare_hessian_model_def_script( + data["model_def_script"], + hessian, + ) + if hessian: + model.enable_hessian() call_lower = model.call_common_lower dim_chg_spin = model.get_dim_chg_spin() has_chg_spin = dim_chg_spin > 0 diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 8ebc3ae00a..90f939de7d 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import json from typing import ( Any, ) @@ -30,6 +31,14 @@ r_differentiable=True, c_differentiable=True, ), + "energy_hessian": OutputVariableDef( + "energy", + shape=[1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + r_hessian=True, + ), "dos": OutputVariableDef( "dos", shape=[-1], @@ -215,6 +224,10 @@ def model_output_def(self) -> ModelOutputDef: ) def _output_var_def(self, name: str) -> OutputVariableDef: + if name == "energy" and json.loads(self.model_def_script).get( + "hessian_mode", False + ): + return OUTPUT_DEFS["energy_hessian"] if name in OUTPUT_DEFS: return OUTPUT_DEFS[name] # property models carry a user-defined output name (``var_name``) that diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 39354cc1fe..a347c2dff8 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -185,7 +185,19 @@ def _check_compressed_hlo_exportable(data: dict) -> None: ) -def deserialize_to_file(model_file: str, data: dict) -> None: +def _prepare_hessian_model_def_script( + model_def_script: dict, + hessian: bool, +) -> tuple[dict, bool]: + """Return a copied model definition and whether Hessian should be enabled.""" + model_def_script = model_def_script.copy() + hessian = hessian or model_def_script.get("hessian_mode", False) + if hessian: + model_def_script["hessian_mode"] = True + return model_def_script, hessian + + +def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None: """Deserialize the dictionary to a model file. Parameters @@ -194,9 +206,14 @@ def deserialize_to_file(model_file: str, data: dict) -> None: The model file to be saved. data : dict The dictionary to be deserialized. + hessian : bool, default=False + Whether to include the Hessian in the model outputs. """ if model_file.endswith(".jax"): - model_def_script = data["model_def_script"].copy() + model_def_script, hessian = _prepare_hessian_model_def_script( + data["model_def_script"], + hessian, + ) min_nbor_dist = _to_optional_float(data.get("min_nbor_dist")) if min_nbor_dist is None: min_nbor_dist = _to_optional_float( @@ -209,6 +226,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None: model_key: BaseModel.deserialize(data["model"]["model_dict"][model_key]) for model_key in model_def_script["model_dict"] } + if hessian: + for model in models.values(): + model.enable_hessian() state = { "models": { model_key: nnx.split(model)[1].to_pure_dict() @@ -217,6 +237,8 @@ def deserialize_to_file(model_file: str, data: dict) -> None: } else: model = BaseModel.deserialize(data["model"]) + if hessian: + model.enable_hessian() _, state = nnx.split(model) state = state.to_pure_dict() with ocp.Checkpointer( @@ -233,7 +255,12 @@ def deserialize_to_file(model_file: str, data: dict) -> None: _check_compressed_hlo_exportable(data) model = BaseModel.deserialize(data["model"]) _set_model_min_nbor_dist_from_data(model, data) - model_def_script = data["model_def_script"] + model_def_script, hessian = _prepare_hessian_model_def_script( + data["model_def_script"], + hessian, + ) + if hessian: + model.enable_hessian() call_lower = model.call_common_lower nf, nloc, nghost = jax_export.symbolic_shape("nf, nloc, nghost") @@ -298,6 +325,7 @@ def call_lower_with_fixed_do_atomic_virial( serialized_atomic_virial_no_ghost = exported_atomic_virial_no_ghost.serialize() data = data.copy() + data["model_def_script"] = model_def_script data.setdefault("@variables", {}) data["@variables"]["stablehlo"] = np.void(serialized) data["@variables"]["stablehlo_atomic_virial"] = np.void( @@ -344,7 +372,7 @@ def call_lower_with_fixed_do_atomic_virial( deserialize_to_file as deserialize_to_savedmodel, ) - return deserialize_to_savedmodel(model_file, data) + deserialize_to_savedmodel(model_file, data, hessian=hessian) else: raise ValueError("Unsupported file extension") diff --git a/deepmd/kernels/triton/sezm/so2_rotation.py b/deepmd/kernels/triton/sezm/so2_rotation.py index 87b7792121..b69eaef1aa 100644 --- a/deepmd/kernels/triton/sezm/so2_rotation.py +++ b/deepmd/kernels/triton/sezm/so2_rotation.py @@ -271,7 +271,7 @@ def _to_local_fwd_kernel( coeff_rows = tl.load(idx_ptr + row, mask=row_mask, other=0).to(tl.int64) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k0 in range(0, tl.cdiv(dim_full, BLOCK_K)): + for k0 in range(tl.cdiv(dim_full, BLOCK_K)): kk = k0 * BLOCK_K + tl.arange(0, BLOCK_K) # over D k_mask = kk < dim_full w_tile = tl.load( @@ -333,7 +333,7 @@ def _to_local_bwd_dx_kernel( src_idx = tl.load(src_ptr + edge).to(tl.int64) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k0 in range(0, tl.cdiv(reduced_dim, BLOCK_K)): + for k0 in range(tl.cdiv(reduced_dim, BLOCK_K)): mm = k0 * BLOCK_K + tl.arange(0, BLOCK_K) # over Dm m_mask = mm < reduced_dim coeff = tl.load(idx_ptr + mm, mask=m_mask, other=0).to(tl.int64) @@ -395,7 +395,7 @@ def _to_local_bwd_dw_kernel( src_idx = tl.load(src_ptr + edge).to(tl.int64) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k0 in range(0, tl.cdiv(channels, BLOCK_K)): + for k0 in range(tl.cdiv(channels, BLOCK_K)): cc = k0 * BLOCK_K + tl.arange(0, BLOCK_K) # over C c_mask = cc < channels go_tile = tl.load( @@ -454,7 +454,7 @@ def _back_fwd_kernel( chan_mask = chan < channels acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k0 in range(0, tl.cdiv(dim_full, BLOCK_K)): + for k0 in range(tl.cdiv(dim_full, BLOCK_K)): kk = k0 * BLOCK_K + tl.arange(0, BLOCK_K) # over D (contraction) k_mask = kk < dim_full inv_k = tl.load(inv_ptr + kk, mask=k_mask, other=-1).to(tl.int64) @@ -517,7 +517,7 @@ def _back_bwd_dx_kernel( keep = inv_k >= 0 acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k0 in range(0, tl.cdiv(dim_full, BLOCK_K)): + for k0 in range(tl.cdiv(dim_full, BLOCK_K)): dd = k0 * BLOCK_K + tl.arange(0, BLOCK_K) # over D (contraction) d_mask = dd < dim_full w_tile = tl.load( @@ -578,7 +578,7 @@ def _back_bwd_dw_kernel( keep = inv_k >= 0 acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k0 in range(0, tl.cdiv(channels, BLOCK_K)): + for k0 in range(tl.cdiv(channels, BLOCK_K)): cc = k0 * BLOCK_K + tl.arange(0, BLOCK_K) # over C (contraction) c_mask = cc < channels go_tile = tl.load( diff --git a/deepmd/main.py b/deepmd/main.py index 42a03bc3f0..03d1a2ec91 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -350,6 +350,12 @@ def main_parser() -> argparse.ArgumentParser: type=str, help="(Supported backend: PyTorch) Task head (alias: model branch) to freeze if in multi-task mode.", ) + parser_frz.add_argument( + "--hessian", + action="store_true", + default=False, + help="(Supported backend: JAX) Add the Hessian to the frozen model output.", + ) parser_frz.add_argument( "--lower-kind", default="nlist", diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index f397bc358b..acd49c589c 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -103,7 +103,7 @@ def prepare_trainer_input_single( seed: int | None = None, ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: training_dataset_params = data_dict_single["training_data"] - validation_dataset_params = data_dict_single.get("validation_data", None) + validation_dataset_params = data_dict_single.get("validation_data") validation_systems = ( validation_dataset_params["systems"] if validation_dataset_params else None ) @@ -115,7 +115,7 @@ def prepare_trainer_input_single( validation_systems = process_systems(validation_systems, val_patterns) # stat files - stat_file_path_single = data_dict_single.get("stat_file", None) + stat_file_path_single = data_dict_single.get("stat_file") if rank != 0: stat_file_path_single = None elif stat_file_path_single is not None: diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py index f61e9867ab..05817920da 100644 --- a/deepmd/pd/train/wrapper.py +++ b/deepmd/pd/train/wrapper.py @@ -207,7 +207,6 @@ def state_dict(self) -> dict[str, Any]: def set_extra_state(self, extra_state: dict[str, Any]) -> None: self.model_params = extra_state["model_params"] self.train_infos = extra_state["train_infos"] - return None def get_extra_state(self) -> dict: extra_state = { diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index acaadb67aa..2773e27240 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -286,7 +286,7 @@ def __init__( Thread.__init__(self) self._queue = queue self._source = source # Main DL iterator - self._max_len = max_len # + self._max_len = max_len def run(self) -> None: for item in self._source: diff --git a/deepmd/pd/utils/utils.py b/deepmd/pd/utils/utils.py index 0f7b1e7987..158b76f0df 100644 --- a/deepmd/pd/utils/utils.py +++ b/deepmd/pd/utils/utils.py @@ -263,7 +263,7 @@ def to_numpy_array( # Create a reverse mapping of PD_PRECISION_DICT reverse_precision_dict = {v: k for k, v in PD_PRECISION_DICT.items()} # Use the reverse mapping to find keys with the desired value - prec = reverse_precision_dict.get(xx.dtype, None) + prec = reverse_precision_dict.get(xx.dtype) prec = NP_PRECISION_DICT.get(prec, np.float64) if prec is None: raise ValueError(f"unknown precision {xx.dtype}") @@ -293,7 +293,7 @@ def to_paddle_tensor( # Create a reverse mapping of NP_PRECISION_DICT reverse_precision_dict = {v: k for k, v in NP_PRECISION_DICT.items()} # Use the reverse mapping to find keys with the desired value - prec = reverse_precision_dict.get(xx.dtype.type, None) + prec = reverse_precision_dict.get(xx.dtype.type) prec = PD_PRECISION_DICT.get(prec, None) if prec is None: raise ValueError(f"unknown precision {xx.dtype}") diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 560ea5a1ba..9da161580b 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -154,19 +154,19 @@ def prepare_trainer_input_single( ]: # get data modifier modifier = None - modifier_params = model_params_single.get("modifier", None) + modifier_params = model_params_single.get("modifier") if modifier_params is not None: modifier = get_data_modifier(modifier_params).to(DEVICE) training_dataset_params = data_dict_single["training_data"] - validation_dataset_params = data_dict_single.get("validation_data", None) + validation_dataset_params = data_dict_single.get("validation_data") validation_systems = ( validation_dataset_params["systems"] if validation_dataset_params else None ) training_systems = training_dataset_params["systems"] # stat files - stat_file_path_single = data_dict_single.get("stat_file", None) + stat_file_path_single = data_dict_single.get("stat_file") if rank != 0: stat_file_path_single = None elif stat_file_path_single is not None: @@ -185,7 +185,7 @@ def _make_dp_loader_set( dataset_params: dict[str, Any], ) -> DpLoaderSet: """Create a DpLoaderSet from systems with pattern expansion.""" - patterns = dataset_params.get("rglob_patterns", None) + patterns = dataset_params.get("rglob_patterns") systems = process_systems(systems, patterns=patterns) return DpLoaderSet( systems, diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 5ec9a1da19..d39f7a1028 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -159,6 +159,15 @@ class DescrptSeZM(BaseDescriptor, nn.Module): radial_mlp Hidden layer sizes for radial networks. An output layer of size `(l_schedule[0]+extra_node_l+1)*channels` will be automatically appended. + edge_norm + Whether to apply channel RMSNorm on the descriptor's cutoff-vanishing + branches: the radial network hidden layers, the environment-seed FiLM + scale/shift logits, the cross-focus competition scalars, and the + post-SO(2) residual messages. ``False`` replaces the first three norms + with identity and changes only the post-SO(2) norm to unit-floor residual + scaling. The unit floor uses ``sqrt(1 + variance)`` so small messages + retain their cutoff envelope instead of receiving the standard + ``1/sqrt(eps)`` small-signal gain. use_env_seed If True, seed the initial node state with local-environment information: apply environment matrix FiLM conditioning on l=0 features using 4D @@ -441,6 +450,7 @@ def __init__( basis_type: str = "bessel", n_radial: int = 16, radial_mlp: list[int] | None = None, + edge_norm: bool = True, use_env_seed: bool = True, random_gamma: bool = True, edge_cartesian: bool = False, @@ -542,6 +552,7 @@ def __init__( if radial_mlp is None: radial_mlp = [0] self.radial_mlp = [self.channels if x == 0 else int(x) for x in radial_mlp] + self.edge_norm = bool(edge_norm) if sandwich_norm is None: sandwich_norm = [False, True, True, False] if not isinstance(sandwich_norm, (list, tuple)) or len(sandwich_norm) != 4: @@ -845,20 +856,28 @@ def __init__( seed=seed_env_seed, ) ) - self.film_scale_norm = ScalarRMSNorm( - channels=self.channels, - n_focus=1, - eps=self.eps, - dtype=self.compute_dtype, - trainable=self.trainable, - ) - self.film_shift_norm = ScalarRMSNorm( - channels=self.channels, - n_focus=1, - eps=self.eps, - dtype=self.compute_dtype, - trainable=self.trainable, - ) + # The FiLM logits derive from the env-seed matrix D = envᵀenv, which + # vanishes at rcut; normalizing them shares the radial network's + # cutoff-smoothness issue, so ``edge_norm=False`` also drops these + # norms (identity pass-through) to keep the FiLM scale/shift smooth. + if self.edge_norm: + self.film_scale_norm: nn.Module = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + dtype=self.compute_dtype, + trainable=self.trainable, + ) + self.film_shift_norm: nn.Module = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + dtype=self.compute_dtype, + trainable=self.trainable, + ) + else: + self.film_scale_norm = nn.Identity() + self.film_shift_norm = nn.Identity() film_strength_init = 0.01 # Use 1D tensor (not scalar) for FSDP2 compatibility self.film_scale_strength_log = nn.Parameter( @@ -906,6 +925,7 @@ def __init__( activation_function=self.activation_function, dtype=self.compute_dtype, # force fp32+ trainable=self.trainable, + radial_norm=self.edge_norm, seed=seed_radial_embedding, ) @@ -968,6 +988,7 @@ def __init__( channels=self.channels, n_focus=self.n_focus, focus_dim=self.focus_dim, + focus_norm=self.edge_norm, so2_norm=self.so2_norm, mixing_layers=self.mixing_layers, so2_attn_res=self.so2_attn_res_mode, @@ -1002,6 +1023,7 @@ def __init__( atten_o_proj=self.use_atten_o_proj, so2_pre_norm=self.so2_pre_norm, so2_post_norm=self.so2_post_norm, + so2_post_norm_eps=1.0e-5 if self.edge_norm else 1.0, so2_activation_function=self.so2_activation_function, ffn_pre_norm=self.ffn_pre_norm, ffn_post_norm=self.ffn_post_norm, @@ -2439,6 +2461,7 @@ def serialize(self) -> dict[str, Any]: "basis_type": self.basis_type, "n_radial": self.n_radial, "radial_mlp": self.radial_mlp, + "edge_norm": self.edge_norm, "use_env_seed": self.use_env_seed, "random_gamma": self.random_gamma, "edge_cartesian": self.edge_cartesian, diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py index 825ff2a5e3..6b170a8935 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/block.py +++ b/deepmd/pt/model/descriptor/sezm_nn/block.py @@ -151,6 +151,12 @@ class SeZMInteractionBlock(nn.Module): ``focus_dim=0`` means using ``channels``. focus_compete If True, enable cross-focus softmax competition in SO(2) convolution. + focus_norm + If True, RMS-normalize the cross-focus competition scalars before the + softmax. The competition input is envelope-gated (via radial modulation) + and vanishes at the cutoff, so normalizing it crosses the norm ``eps`` + floor near ``rcut``; ``False`` drops the norm so the competition decays + smoothly to uniform weights at the cutoff. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm between SO(2) mixing layers. When False (default), no normalization is applied between layers. @@ -198,6 +204,10 @@ class SeZMInteractionBlock(nn.Module): If True, apply pre-norm before SO(2) convolution. so2_post_norm If True, apply post-norm on SO(2) output before the residual add. + so2_post_norm_eps + Variance floor for the SO(2) post-norm. A value of ``1`` preserves small + residual messages instead of amplifying them by ``1/sqrt(eps)``. Other + normalization sites retain their own RMSNorm floors. ffn_pre_norm If True, apply pre-norm before each FFN subblock. ffn_post_norm @@ -299,6 +309,7 @@ def __init__( n_focus: int = 1, focus_dim: int = 0, focus_compete: bool = True, + focus_norm: bool = True, so2_norm: bool = False, mixing_layers: int = 4, so2_attn_res: str = "none", @@ -312,6 +323,7 @@ def __init__( atten_o_proj: bool = False, so2_pre_norm: bool = True, so2_post_norm: bool = False, + so2_post_norm_eps: float = 1e-5, ffn_pre_norm: bool = True, ffn_post_norm: bool = False, ffn_neurons: int = 96, @@ -366,6 +378,7 @@ def __init__( if self.focus_dim < 0: raise ValueError("`focus_dim` must be >= 0") self.focus_compete = bool(focus_compete) + self.focus_norm = bool(focus_norm) self.so2_norm = bool(so2_norm) self.mixing_layers = int(mixing_layers) self.so2_attn_res_mode = str(so2_attn_res).lower() @@ -383,6 +396,7 @@ def __init__( 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.so2_post_norm_eps = float(so2_post_norm_eps) self.ffn_pre_norm = bool(ffn_pre_norm) self.ffn_post_norm = bool(ffn_post_norm) self.ffn_neurons = int(ffn_neurons) @@ -463,6 +477,7 @@ def __init__( self.lmax, self.channels, n_focus=1, + eps=self.so2_post_norm_eps, dtype=self.compute_dtype, trainable=trainable, ) @@ -477,6 +492,7 @@ def __init__( n_focus=self.n_focus, focus_dim=self.focus_dim, focus_compete=self.focus_compete, + focus_norm=self.focus_norm, so2_norm=self.so2_norm, mixing_layers=self.mixing_layers, so2_attn_res=self.so2_attn_res_mode, @@ -1037,6 +1053,7 @@ def serialize(self) -> dict[str, Any]: "n_focus": self.n_focus, "focus_dim": self.focus_dim, "focus_compete": self.focus_compete, + "focus_norm": self.focus_norm, "so2_norm": self.so2_norm, "mixing_layers": self.mixing_layers, "so2_attn_res": self.so2_attn_res_mode, @@ -1050,6 +1067,7 @@ def serialize(self) -> dict[str, Any]: "atten_o_proj": self.use_atten_o_proj, "so2_pre_norm": self.so2_pre_norm, "so2_post_norm": self.so2_post_norm, + "so2_post_norm_eps": self.so2_post_norm_eps, "ffn_pre_norm": self.ffn_pre_norm, "ffn_post_norm": self.ffn_post_norm, "ffn_neurons": self.ffn_neurons, diff --git a/deepmd/pt/model/descriptor/sezm_nn/radial.py b/deepmd/pt/model/descriptor/sezm_nn/radial.py index 4f4e4d888c..cef4e402c3 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/radial.py +++ b/deepmd/pt/model/descriptor/sezm_nn/radial.py @@ -52,7 +52,7 @@ class RadialMLP(nn.Module): """ - Radial MLP with channel RMSNorm and configurable activation. + Radial MLP with optional channel RMSNorm and configurable activation. Parameters ---------- @@ -65,11 +65,14 @@ class RadialMLP(nn.Module): Floating point dtype for the linear layers. trainable : bool Whether the parameters are trainable. + radial_norm : bool + Whether to insert a channel RMSNorm in each hidden layer. Architecture ------------ - Linear → RMSNorm → Activation for all hidden layers, - with the final layer being a plain Linear (no norm, no activation). + ``radial_norm=True`` : Linear → RMSNorm → Activation for each hidden layer. + ``radial_norm=False`` : Linear → Activation for each hidden layer. + The final layer is always a plain Linear (no norm, no activation). Notes ----- @@ -78,6 +81,15 @@ class RadialMLP(nn.Module): 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. + + The hidden RMSNorm normalizes each edge's radial features by their own RMS. + The input ``edge_rbf`` carries the C^3 cutoff envelope and therefore + vanishes at ``rcut``; the RMSNorm divides that envelope out, and its ``eps`` + floor is crossed as the edge approaches ``rcut``. On a sparse neighborhood + (e.g. a dimer) this floor-crossing produces a sharp kink in the potential + energy surface just inside the cutoff. Setting ``radial_norm=False`` drops + the RMSNorm so the radial features vanish smoothly with the envelope, which + restores C^3 smoothness at the cutoff. """ def __init__( @@ -87,6 +99,7 @@ def __init__( activation_function: str = "silu", dtype: torch.dtype = torch.float32, trainable: bool = True, + radial_norm: bool = True, seed: int | list[int] | None = None, ) -> None: super().__init__() @@ -98,6 +111,7 @@ def __init__( self.device = env.DEVICE self.precision = RESERVED_PRECISION_DICT[self.dtype] self.trainable = bool(trainable) + self.radial_norm = bool(radial_norm) modules: list[nn.Module] = [] n_layers = len(mlp_layers) @@ -114,13 +128,14 @@ def __init__( modules.append(linear) # Last layer: no RMSNorm/activation if i < n_layers - 2: - modules.append( - RMSNorm( - channels=mlp_layers[i + 1], - dtype=self.dtype, - trainable=trainable, + if self.radial_norm: + modules.append( + RMSNorm( + channels=mlp_layers[i + 1], + dtype=self.dtype, + trainable=trainable, + ) ) - ) modules.append(ActivationFn(self.activation_function)) self.net = nn.Sequential(*modules) @@ -151,6 +166,7 @@ def serialize(self) -> dict[str, Any]: "activation_function": self.activation_function, "dtype": RESERVED_PRECISION_DICT[self.dtype], "trainable": self.trainable, + "radial_norm": self.radial_norm, "@variables": {k: np_safe(v) for k, v in state.items()}, } diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py index fd194c94d7..b7548507f3 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/so2.py +++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py @@ -903,6 +903,11 @@ class SO2Convolution(nn.Module): If True, apply cross-focus softmax competition in SO(2) local layout. Competition logits are constructed only from l=0 scalar channels and the resulting invariant weights are broadcast to all (l, m) components. + focus_norm + If True, RMS-normalize the competition l=0 scalars before the softmax. + Those scalars are envelope-gated (radial modulation) and vanish at the + cutoff, so the norm crosses its eps floor near ``rcut``; ``False`` uses an + identity pass-through and lets the competition decay smoothly to uniform. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm as pre-norm before each SO(2) mixing layer. The last SO(2) layer always uses Identity. @@ -1018,6 +1023,7 @@ def __init__( n_focus: int = 1, focus_dim: int = 0, focus_compete: bool = True, + focus_norm: bool = True, so2_norm: bool = False, mixing_layers: int = 4, so2_attn_res: str = "none", @@ -1068,6 +1074,7 @@ def __init__( 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_norm = bool(focus_norm) self.focus_softmax_tau = 1.0 self.focus_label_smoothing = 0.02 self.so2_norm = bool(so2_norm) @@ -1370,16 +1377,24 @@ def __init__( ) # === Step 7.5. Optional cross-focus competition === - self.focus_compete_norm: ScalarRMSNorm | None = None + self.focus_compete_norm: nn.Module | None = None self.adamw_focus_compete_w: nn.Parameter | None = None self.focus_compete_bias: nn.Parameter | 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, - dtype=self.compute_dtype, - trainable=trainable, + # The competition scalars are envelope-gated (radial modulation) and + # vanish at rcut; normalizing them crosses the eps floor near the + # cutoff, so ``focus_norm=False`` uses an identity pass-through and + # lets the softmax decay smoothly to uniform weights there. + self.focus_compete_norm = ( + ScalarRMSNorm( + channels=self.so2_focus_dim, + n_focus=self.n_focus, + eps=self.eps, + dtype=self.compute_dtype, + trainable=trainable, + ) + if self.focus_norm + else nn.Identity() ) self.adamw_focus_compete_w = nn.Parameter( torch.empty( @@ -2488,6 +2503,7 @@ def serialize(self) -> dict[str, Any]: "n_focus": self.n_focus, "focus_dim": self.focus_dim, "focus_compete": self.focus_compete, + "focus_norm": self.focus_norm, "so2_norm": self.so2_norm, "mixing_layers": self.mixing_layers, "so2_attn_res": self.so2_attn_res_mode, diff --git a/deepmd/pt/optimizer/hybrid_muon.py b/deepmd/pt/optimizer/hybrid_muon.py index d2d96ccdb8..24caaadc41 100644 --- a/deepmd/pt/optimizer/hybrid_muon.py +++ b/deepmd/pt/optimizer/hybrid_muon.py @@ -281,7 +281,7 @@ def _mmt_kernel( b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + for k in range(tl.cdiv(K, BLOCK_SIZE_K)): a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) b = tl.load(b_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 76a2e9866a..8dd68c8cdb 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -329,7 +329,7 @@ def get_dataloader_and_iter_lmdb( rank=self.rank, world_size=self.world_size, shuffle=True, - seed=_training_params.get("seed", None), + seed=_training_params.get("seed"), block_targets=_block_targets, ) else: diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index da710f4fdf..b2e4d4f229 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -256,7 +256,6 @@ def _forward_without_loss( def set_extra_state(self, state: dict) -> None: self.model_params = state["model_params"] self.train_infos = state["train_infos"] - return None def get_extra_state(self) -> dict: state = { diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index a3ba852d7e..9f95c59adc 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -247,7 +247,7 @@ def to_numpy_array( # Create a reverse mapping of PT_PRECISION_DICT reverse_precision_dict = {v: k for k, v in PT_PRECISION_DICT.items()} # Use the reverse mapping to find keys with the desired value - prec = reverse_precision_dict.get(xx.dtype, None) + prec = reverse_precision_dict.get(xx.dtype) prec = NP_PRECISION_DICT.get(prec, None) if prec is None: raise ValueError(f"unknown precision {xx.dtype}") @@ -277,7 +277,7 @@ def to_torch_tensor( # Create a reverse mapping of NP_PRECISION_DICT reverse_precision_dict = {v: k for k, v in NP_PRECISION_DICT.items()} # Use the reverse mapping to find keys with the desired value - prec = reverse_precision_dict.get(xx.dtype.type, None) + prec = reverse_precision_dict.get(xx.dtype.type) prec = PT_PRECISION_DICT.get(prec, None) if prec is None: raise ValueError(f"unknown precision {xx.dtype}") diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 2e4f747ccb..3367ee4579 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -151,7 +151,7 @@ def _build_data_system( ) systems = process_systems( systems_raw, - patterns=dataset_params.get("rglob_patterns", None), + patterns=dataset_params.get("rglob_patterns"), ) return DeepmdDataSystem( systems=systems, @@ -159,7 +159,7 @@ def _build_data_system( test_size=1, type_map=type_map, trn_all_set=True, - sys_probs=dataset_params.get("sys_probs", None), + sys_probs=dataset_params.get("sys_probs"), auto_prob_style=dataset_params.get("auto_prob", "prob_sys_size"), ) diff --git a/deepmd/pt_expt/utils/network.py b/deepmd/pt_expt/utils/network.py index adef443de9..004ba94401 100644 --- a/deepmd/pt_expt/utils/network.py +++ b/deepmd/pt_expt/utils/network.py @@ -94,10 +94,10 @@ def __setattr__(self, name: str, value: Any) -> None: if val is None: if name in self._parameters: self._parameters[name] = None - return + return None if name in self._buffers: self._buffers[name] = None - return + return None return super().__setattr__(name, None) if getattr(self, "trainable", False): param = ( @@ -107,14 +107,14 @@ def __setattr__(self, name: str, value: Any) -> None: ) if name in self._parameters: self._parameters[name] = param - return + return None return super().__setattr__(name, param) if name in self._buffers: self._buffers[name] = val - return + return None # Register on first assignment so tensors are in state_dict and moved by .to(). self.register_buffer(name, val) - return + return None return super().__setattr__(name, value) def call(self, x: torch.Tensor) -> torch.Tensor: diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index 3e1c9b127e..acecaf57f7 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -794,7 +794,7 @@ def _pass_filter( trainable: bool = True, ) -> tuple[tf.Tensor, tf.Tensor]: if input_dict is not None: - type_embedding = input_dict.get("type_embedding", None) + type_embedding = input_dict.get("type_embedding") if type_embedding is not None: self.use_tebd = True else: diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 1bbb0a5595..e9feac8d74 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -732,10 +732,9 @@ def _pass_filter( trainable: bool = True, ) -> tuple[tf.Tensor, None]: assert ( - input_dict is not None - and input_dict.get("type_embedding", None) is not None + input_dict is not None and input_dict.get("type_embedding") is not None ), "se_atten descriptor must use type_embedding" - type_embedding = input_dict.get("type_embedding", None) + type_embedding = input_dict.get("type_embedding") inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) output = [] output_qmat = [] @@ -1961,7 +1960,7 @@ def serialize(self, suffix: str = "") -> dict: raise RuntimeError( "The implementation for smooth_type_embedding is inconsistent with other backends" ) - # todo support serialization when tebd_input_mode=='strip' and type_one_side is True + # TODO support serialization when tebd_input_mode=='strip' and type_one_side is True if self.stripped_type_embedding and self.type_one_side: raise NotImplementedError( "serialization is unsupported when tebd_input_mode=='strip' and type_one_side is True" diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 16bec59bf0..5ea019b2eb 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -810,7 +810,7 @@ def clear_ij(type_i: int, type_j: int) -> None: clear_ij(i, j) clear_ij(j, i) for i in range(ntypes): - for j in range(0, i): + for j in range(i): clear_ij(i, j) if suffix != "": diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 917cfe70e1..990c0d41d5 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -187,7 +187,7 @@ def _build_lower( rot_mat_i = tf.slice(rot_mat, [0, start_index, 0], [-1, natoms, -1]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) layer = inputs_i - for ii in range(0, len(self.n_neuron)): + for ii in range(len(self.n_neuron)): if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]: layer += one_layer( layer, @@ -282,8 +282,8 @@ def build( """ if input_dict is None: input_dict = {} - type_embedding = input_dict.get("type_embedding", None) - atype = input_dict.get("atype", None) + type_embedding = input_dict.get("type_embedding") + atype = input_dict.get("atype") nframes = input_dict.get("nframes") start_index = 0 inputs = tf.reshape(input_d, [-1, natoms[0], self.dim_descrpt]) diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 166ed2e355..3d99bcb64d 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -352,7 +352,7 @@ def _build_lower( one_layer = one_layer_nvnmd else: one_layer = one_layer_deepmd - for ii in range(0, len(self.n_neuron)): + for ii in range(len(self.n_neuron)): if self.layer_name is not None and self.layer_name[ii] is not None: layer_suffix = "share_" + self.layer_name[ii] + type_suffix layer_reuse = tf.AUTO_REUSE @@ -452,8 +452,8 @@ def build( if input_dict is None: input_dict = {} bias_dos = self.bias_dos - type_embedding = input_dict.get("type_embedding", None) - atype = input_dict.get("atype", None) + type_embedding = input_dict.get("type_embedding") + atype = input_dict.get("atype") if self.numb_fparam > 0: if self.fparam_avg is None: self.fparam_avg = 0.0 diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index e8accb2087..a787e5249c 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -419,7 +419,7 @@ def _build_lower( one_layer = one_layer_nvnmd else: one_layer = one_layer_deepmd - for ii in range(0, len(self.n_neuron)): + for ii in range(len(self.n_neuron)): if self.layer_name is not None and self.layer_name[ii] is not None: layer_suffix = "share_" + self.layer_name[ii] + type_suffix layer_reuse = tf.AUTO_REUSE @@ -519,8 +519,8 @@ def build( if input_dict is None: input_dict = {} bias_atom_e = self.bias_atom_e - type_embedding = input_dict.get("type_embedding", None) - atype = input_dict.get("atype", None) + type_embedding = input_dict.get("type_embedding") + atype = input_dict.get("atype") if self.numb_fparam > 0: if self.fparam_avg is None: self.fparam_avg = 0.0 diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 5e4e48e96f..3143f89c8b 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -336,7 +336,7 @@ def _build_lower( ) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) layer = inputs_i - for ii in range(0, len(self.n_neuron)): + for ii in range(len(self.n_neuron)): if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]: layer += one_layer( layer, @@ -470,8 +470,8 @@ def build( """ if input_dict is None: input_dict = {} - type_embedding = input_dict.get("type_embedding", None) - atype = input_dict.get("atype", None) + type_embedding = input_dict.get("type_embedding") + atype = input_dict.get("atype") nframes = input_dict.get("nframes") start_index = 0 diff --git a/deepmd/tf/loss/tensor.py b/deepmd/tf/loss/tensor.py index d63987e908..624195c69b 100644 --- a/deepmd/tf/loss/tensor.py +++ b/deepmd/tf/loss/tensor.py @@ -26,7 +26,7 @@ class TensorLoss(Loss): """Loss function for tensorial properties.""" def __init__(self, jdata: dict | None, **kwarg: Any) -> None: - model = kwarg.get("model", None) + model = kwarg.get("model") if model is not None: self.type_sel = model.get_sel_type() else: diff --git a/deepmd/tf/nvnmd/descriptor/se_a.py b/deepmd/tf/nvnmd/descriptor/se_a.py index 96b89cffa7..0d5c3dd5cc 100644 --- a/deepmd/tf/nvnmd/descriptor/se_a.py +++ b/deepmd/tf/nvnmd/descriptor/se_a.py @@ -15,8 +15,6 @@ op_module, tf, ) - -# from deepmd.tf.nvnmd.utils.config import ( nvnmd_cfg, ) @@ -44,7 +42,6 @@ def build_davg_dstd() -> tuple[Any, Any]: def check_switch_range(davg: np.ndarray, dstd: np.ndarray) -> None: r"""Check the range of switch, let it in range [-2, 14].""" rmin = nvnmd_cfg.dscp["rcut_smth"] - # namelist = [n.name for n in tf.get_default_graph().as_graph_def().node] if "train_attr/min_nbor_dist" in namelist: min_dist = get_tensor_by_name_from_graph( @@ -307,7 +304,6 @@ def filter_GR2D(xyz_scatter_1: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: result = tf.ensure_shape(result, [None, M1, M1]) # D': natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, M1 * M1]) - # index_subset = [] for ii in range(M1): for jj in range(ii, ii + M2): @@ -337,7 +333,6 @@ def filter_GR2D(xyz_scatter_1: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: # natom x (outputs_size x outputs_size_2) # result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) result = tf.reshape(result, [-1, M1 * M1]) - # index_subset = [] for ii in range(M1): for jj in range(ii, ii + M2): diff --git a/deepmd/tf/nvnmd/descriptor/se_atten.py b/deepmd/tf/nvnmd/descriptor/se_atten.py index f8f3085ae8..aa15a8ca7d 100644 --- a/deepmd/tf/nvnmd/descriptor/se_atten.py +++ b/deepmd/tf/nvnmd/descriptor/se_atten.py @@ -14,8 +14,6 @@ op_module, tf, ) - -# from deepmd.tf.nvnmd.utils.config import ( nvnmd_cfg, ) @@ -43,7 +41,6 @@ def check_switch_range(davg: np.ndarray, dstd: np.ndarray) -> None: ntype = nvnmd_cfg.dscp["ntype"] NIDP = nvnmd_cfg.dscp["NIDP"] ndescrpt = NIDP * 4 - # namelist = [n.name for n in tf.get_default_graph().as_graph_def().node] if "train_attr/min_nbor_dist" in namelist: min_dist = get_tensor_by_name_from_graph( @@ -271,7 +268,6 @@ def filter_GR2D(xyz_scatter_1: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: result = tf.ensure_shape(result, [None, M1, M1]) # D': natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, M1 * M1]) - # index_subset = [] for ii in range(M1): for jj in range(ii, ii + M2): @@ -301,7 +297,6 @@ def filter_GR2D(xyz_scatter_1: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: # natom x (outputs_size x outputs_size_2) # result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) result = tf.reshape(result, [-1, M1 * M1]) - # index_subset = [] for ii in range(M1): for jj in range(ii, ii + M2): diff --git a/deepmd/tf/nvnmd/entrypoints/mapt.py b/deepmd/tf/nvnmd/entrypoints/mapt.py index 5da708b311..efd72f03da 100644 --- a/deepmd/tf/nvnmd/entrypoints/mapt.py +++ b/deepmd/tf/nvnmd/entrypoints/mapt.py @@ -108,7 +108,6 @@ def build_map(self) -> dict: if self.Gs_Gt_mode == 1: self.shift_Gs = 0 self.shift_Gt = 1 - # M = nvnmd_cfg.dscp["M1"] if nvnmd_cfg.version == 0: ndim = nvnmd_cfg.dscp["ntype"] @@ -183,7 +182,6 @@ def build_map(self) -> dict: if nvnmd_cfg.version == 1: self.map.update(dic_map3) self.map.update(dic_map4) - # FioDic().save(self.map_file, self.map) log.info("NVNMD: finish building mapping table") return self.map @@ -231,7 +229,6 @@ def mapping2(self, x: np.ndarray, dic_map: dict, cfgs: dict) -> dict: t_table_info = tf.placeholder(tf.float64, [None], "t_table_info") t_y = op_module.map_flt_nvnmd(t_x, t_table, t_table_grad, t_table_info) sess = get_sess() - # n = len(x) dic_val = {} for key in dic_map.keys(): @@ -293,7 +290,6 @@ def build_map_coef( y_i = ys grad_i = grads grad_grad_i = grad_grads - # coef_i = [] coef_grad_i = [] for jj in range(Nc): @@ -339,7 +335,6 @@ def cal_coef4( y1 = y1[:Nd] dy0 = dy0[:Nd] dy1 = dy1[:Nd] - # a = (dx * dy1 - 2 * y1 + dx * dy0 + 2 * y0) / dx**3 b = (3 * y1 - dx * dy1 - 2 * dx * dy0 - 3 * y0) / dx**2 c = dy0 @@ -386,7 +381,6 @@ def build_u2s(self, r2: tf.Tensor) -> tf.Tensor: if dmin > 1e-6: min_dist = dmin min_dist = 0.5 if (min_dist > 0.5) else (min_dist - 0.1) - # r = tf.sqrt(r2) r_ = tf.clip_by_value(r, rmin, rmax) r__ = tf.clip_by_value(r, min_dist, rmax) @@ -423,7 +417,6 @@ def build_u2s_grad(self) -> dict: ndim = nvnmd_cfg.dscp["ntype"] if nvnmd_cfg.version == 1: ndim = 1 - # dic_ph = {} dic_ph["u"] = tf.placeholder(tf.float64, [None, 1], "t_u") dic_ph["s"], dic_ph["h"] = self.build_u2s(dic_ph["u"]) @@ -473,7 +466,6 @@ def run_u2s(self) -> tuple[dict, dict]: res_dic["h"][tt][0] = 0 res_dic["h_grad"][tt][0] = 0 res_dic["h_grad_grad"][tt][0] = 0 - # res_dic2["s"][tt][0] = -avg[tt, 0] / std[tt, 0] res_dic2["s_grad"][tt][0] = 0 res_dic2["s_grad_grad"][tt][0] = 0 @@ -493,7 +485,6 @@ def build_s2g(self, s: tf.Tensor) -> tf.Tensor: ntype = nvnmd_cfg.dscp["ntype"] if nvnmd_cfg.version == 1: ntype = 1 - # xyz_scatters = [] for tt2 in range(ntype): wbs = [get_filter_weight(nvnmd_cfg.weight, tt2, ll) for ll in range(1, 5)] @@ -504,7 +495,6 @@ def build_s2g(self, s: tf.Tensor) -> tf.Tensor: def build_s2g_grad(self) -> dict: r"""Build gradient of G with respect to s.""" M1 = nvnmd_cfg.dscp["M1"] - # if nvnmd_cfg.version == 0: ntypex = nvnmd_cfg.dscp["ntypex"] ntype = nvnmd_cfg.dscp["ntype"] @@ -513,7 +503,6 @@ def build_s2g_grad(self) -> dict: if nvnmd_cfg.version == 1: ndim = 1 shift = self.shift_Gs - # dic_ph = {} dic_ph["s"] = tf.placeholder(tf.float64, [None, 1], "t_s") dic_ph["g"] = [g + shift for g in self.build_s2g(dic_ph["s"])] @@ -550,7 +539,6 @@ def run_s2g(self) -> tuple[dict, dict]: smin_ = np.floor(smin * prec - 1) / prec if nvnmd_cfg.version == 1: smin_ = 0 - # keys = list(dic_ph.keys()) vals = list(dic_ph.values()) @@ -616,10 +604,8 @@ def run_t2g(self) -> dict: tf.reset_default_graph() dic_ph = self.build_t2g() sess = get_sess() - # keys = list(dic_ph.keys()) vals = list(dic_ph.values()) - # res_lst = run_sess(sess, vals, feed_dict={}) res_dic = dict(zip(keys, res_lst, strict=True)) @@ -659,7 +645,6 @@ def build_embedding_net( def build_davg_dstd(self) -> dict: ntype = nvnmd_cfg.dscp["ntype"] davg, dstd = get_normalize(nvnmd_cfg.weight) - # res_dic = {} res_dic["davg_opp"] = np.array([-davg[tt, 0:4] for tt in range(ntype)]) res_dic["dstd_inv"] = np.array([1.0 / dstd[tt, 0:4] for tt in range(ntype)]) diff --git a/deepmd/tf/nvnmd/entrypoints/train.py b/deepmd/tf/nvnmd/entrypoints/train.py index d83362a705..63583ebffa 100644 --- a/deepmd/tf/nvnmd/entrypoints/train.py +++ b/deepmd/tf/nvnmd/entrypoints/train.py @@ -88,7 +88,6 @@ def normalized_input(fn: str, PATH_CNN: str, CONFIG_CNN: str) -> str: jdata_train["save_ckpt"] = os.path.join( PATH_CNN, os.path.split(jdata_train["save_ckpt"])[1] ) - # jdata["model"] = nvnmd_cfg.get_model_jdata() jdata["nvnmd"] = nvnmd_cfg.get_nvnmd_jdata() return jdata @@ -98,7 +97,6 @@ def normalized_input_qnn( jdata: dict, PATH_QNN: str, CONFIG_CNN: str, WEIGHT_CNN: str, MAP_CNN: str ) -> str: r"""Normalize a input script file for quantize neural network.""" - # jdata_nvnmd = jdata_deepmd_input_v0["nvnmd"] jdata_nvnmd["enable"] = True jdata_nvnmd["version"] = nvnmd_cfg.version diff --git a/deepmd/tf/nvnmd/entrypoints/wrap.py b/deepmd/tf/nvnmd/entrypoints/wrap.py index e946a74a5d..74cd420410 100755 --- a/deepmd/tf/nvnmd/entrypoints/wrap.py +++ b/deepmd/tf/nvnmd/entrypoints/wrap.py @@ -131,7 +131,6 @@ def wrap(self) -> None: w4 = w * 4 # nbit nhs.append(h) nws.append(w) - # w_full = np.ceil(w4 / nbit) * nbit d = e.extend_hex(d, w_full) # DEVELOP_DEBUG @@ -319,7 +318,6 @@ def wrap_dscp(self) -> str: ) sGSs = "".join(GSs[::-1]) bs = sGSs + bs - # NIX = dscp["NIX"] ln2_NIX = -int(np.log2(NIX)) bs = e.dec2bin(ln2_NIX, NBIT_FLTE, signed=True)[0] + bs @@ -389,7 +387,6 @@ def wrap_fitn(self) -> tuple[list[str], list[str]]: bdc.append(bdct) bwr.append(bwrt) bwc.append(bwct) - # bfps, bbps = [], [] for ss in range(NSEL): tt = ss // NSTDM @@ -537,7 +534,6 @@ def wrap_map(self) -> tuple[list[str], list[str], list[str], list[str]]: d1 = d[:, :, 0:2] d2 = d[:, :, 2:4] d = np.concatenate([d1, d2]) - # bs = e.flt2bin(d, NBIT_FLTE, NBIT_FLTF) bs = e.reverse_bin(bs, nmerges[ii]) bs = e.merge_bin(bs, nmerges[ii]) @@ -552,7 +548,6 @@ def wrap_map(self) -> tuple[list[str], list[str], list[str], list[str]]: d1 = np.reshape(d[:, :, 0:2], [-1, nd * 2]) d2 = np.reshape(d[:, :, 2:4], [-1, nd * 2]) d = np.concatenate([d1, d2], axis=1) - # bs = e.flt2bin(d, NBIT_FLTE, NBIT_FLTF) bss.append(bs) bswt, bdsw, bfea, bgra = bss diff --git a/deepmd/tf/nvnmd/utils/config.py b/deepmd/tf/nvnmd/utils/config.py index 99fb640aae..c88077ae94 100644 --- a/deepmd/tf/nvnmd/utils/config.py +++ b/deepmd/tf/nvnmd/utils/config.py @@ -118,7 +118,6 @@ def init_from_config(self, jdata: dict) -> None: self.init_config_by_version( jdata["ctrl"]["VERSION"], jdata["ctrl"]["MAX_NNEI"] ) - # self.config = FioDic().update(jdata, self.config) self.config["dscp"] = self.init_dscp(self.config["dscp"], self.config) self.config["fitn"] = self.init_fitn(self.config["fitn"], self.config) @@ -167,7 +166,6 @@ def init_from_deepmd_input(self, jdata: dict) -> None: self.config["fitn"] = self.init_fitn(self.config["fitn"], self.config) dp_in = {"type_map": fioObj.get(jdata, "type_map", [])} self.config["dpin"] = fioObj.update(dp_in, self.config["dpin"]) - # self.init_net_size() self.init_value() @@ -290,7 +288,6 @@ def get_s_range(self, davg: np.ndarray, dstd: np.ndarray) -> None: rmax = nvnmd_cfg.dscp["rcut"] ntype = self.dscp["ntype"] dmin = self.dscp["dmin"] - # s0 = r2s(dmin, rmin, rmax) smin_ = -davg[:ntype, 0] / dstd[:ntype, 0] smax_ = (s0 - davg[:ntype, 0]) / dstd[:ntype, 0] diff --git a/deepmd/tf/nvnmd/utils/encode.py b/deepmd/tf/nvnmd/utils/encode.py index 53a860080f..039c61ab9e 100644 --- a/deepmd/tf/nvnmd/utils/encode.py +++ b/deepmd/tf/nvnmd/utils/encode.py @@ -68,7 +68,6 @@ def flt2bin_one(self, v: float, nbit_expo: int, nbit_frac: int) -> str: if h[ii] == "p": ed = ii + 1 is_zero = h[st] == "0" - # if is_zero: return "0" * (1 + nbit_expo + nbit_frac) else: diff --git a/deepmd/tf/nvnmd/utils/fio.py b/deepmd/tf/nvnmd/utils/fio.py index 0994a8995a..8b452a19c9 100644 --- a/deepmd/tf/nvnmd/utils/fio.py +++ b/deepmd/tf/nvnmd/utils/fio.py @@ -178,7 +178,6 @@ def save(self, file_name: str, data: list[str]) -> None: buff = [] for si in data: buff.extend(list(bytearray.fromhex(si))[::-1]) - # with open(file_name, "wb") as fp: fp.write(struct.pack(f"{len(buff)}B", *buff)) diff --git a/deepmd/tf/nvnmd/utils/network.py b/deepmd/tf/nvnmd/utils/network.py index c0a6bf5248..0486ac9840 100644 --- a/deepmd/tf/nvnmd/utils/network.py +++ b/deepmd/tf/nvnmd/utils/network.py @@ -213,7 +213,6 @@ def one_layer( uniform_seed, name, ) - # NTAVC = nvnmd_cfg.fitn["NTAVC"] nd = inputs.get_shape().as_list()[1] - NTAVC inputs2 = tf.slice(inputs, [0, nd], [-1, NTAVC]) diff --git a/deepmd/tf/op/_map_flt_nvnmd_grad.py b/deepmd/tf/op/_map_flt_nvnmd_grad.py index 5443e3286b..a8bc6c56cc 100644 --- a/deepmd/tf/op/_map_flt_nvnmd_grad.py +++ b/deepmd/tf/op/_map_flt_nvnmd_grad.py @@ -25,7 +25,6 @@ def _MapFltNvnmdGrad(op: tf.Operation, grad: tf.Tensor) -> list[tf.Tensor | None N = shx[0] D = shx[1] M = shw[1] // 4 - # dydx = op_module.map_flt_nvnmd(x, table_grad, tf.zeros_like(table_grad), table_info) dydx = tf.ensure_shape(dydx, [N, D, M]) # calculate diff --git a/deepmd/tf/op/_tanh4_flt_nvnmd_grad.py b/deepmd/tf/op/_tanh4_flt_nvnmd_grad.py index f3582b194b..4d0a0fcdc0 100644 --- a/deepmd/tf/op/_tanh4_flt_nvnmd_grad.py +++ b/deepmd/tf/op/_tanh4_flt_nvnmd_grad.py @@ -22,12 +22,10 @@ def _Tanh4FltNvnmdGrad(op: tf.Operation, grad: tf.Tensor) -> list[tf.Tensor]: xx = xhi * xlo xxhi = xx + tf.stop_gradient(tf.floor(xx * prechi) / prechi - xx) xxlo = xx + tf.stop_gradient(tf.floor(xx * preclo) / preclo - xx) - # dydx = xxlo * (xhi / 4 - 3 / 4) + 1 # dydx = xxhi * (xlo/4 - 3/4) + 1 dydxhi = dydx + tf.stop_gradient(tf.floor(dydx * prechi) / prechi - dydx) dydxlo = dydx + tf.stop_gradient(tf.floor(dydx * preclo) / preclo - dydx) - # gradhi = grad + tf.stop_gradient(tf.floor(grad * prechi) / prechi - grad) dx = dydxlo * gradhi dx = dx + tf.stop_gradient(tf.floor(dx * prechi) / prechi - dx) diff --git a/deepmd/tf/utils/tabulate.py b/deepmd/tf/utils/tabulate.py index 614d18d9d8..53137cab19 100644 --- a/deepmd/tf/utils/tabulate.py +++ b/deepmd/tf/utils/tabulate.py @@ -194,7 +194,7 @@ def _get_bias(self) -> dict[str, list[np.ndarray]]: bias["layer_" + str(layer)].append(tf.make_ndarray(node)) elif isinstance(self.descrpt, deepmd.tf.descriptor.DescrptSeA): if self.type_one_side: - for ii in range(0, self.ntypes): + for ii in range(self.ntypes): if not self._all_excluded(ii): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/bias_{layer}_{ii}" @@ -203,7 +203,7 @@ def _get_bias(self) -> dict[str, list[np.ndarray]]: else: bias["layer_" + str(layer)].append(np.array([])) else: - for ii in range(0, self.ntypes * self.ntypes): + for ii in range(self.ntypes * self.ntypes): if ( ii // self.ntypes, ii % self.ntypes, @@ -223,7 +223,7 @@ def _get_bias(self) -> dict[str, list[np.ndarray]]: bias["layer_" + str(layer)].append(tf.make_ndarray(node)) elif isinstance(self.descrpt, deepmd.tf.descriptor.DescrptSeR): if self.type_one_side: - for ii in range(0, self.ntypes): + for ii in range(self.ntypes): if not self._all_excluded(ii): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/bias_{layer}_{ii}" @@ -232,7 +232,7 @@ def _get_bias(self) -> dict[str, list[np.ndarray]]: else: bias["layer_" + str(layer)].append(np.array([])) else: - for ii in range(0, self.ntypes * self.ntypes): + for ii in range(self.ntypes * self.ntypes): if ( ii // self.ntypes, ii % self.ntypes, @@ -260,7 +260,7 @@ def _get_matrix(self) -> dict[str, list[np.ndarray]]: matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) elif isinstance(self.descrpt, deepmd.tf.descriptor.DescrptSeA): if self.type_one_side: - for ii in range(0, self.ntypes): + for ii in range(self.ntypes): if not self._all_excluded(ii): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/matrix_{layer}_{ii}" @@ -269,7 +269,7 @@ def _get_matrix(self) -> dict[str, list[np.ndarray]]: else: matrix["layer_" + str(layer)].append(np.array([])) else: - for ii in range(0, self.ntypes * self.ntypes): + for ii in range(self.ntypes * self.ntypes): if ( ii // self.ntypes, ii % self.ntypes, @@ -289,7 +289,7 @@ def _get_matrix(self) -> dict[str, list[np.ndarray]]: matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) elif isinstance(self.descrpt, deepmd.tf.descriptor.DescrptSeR): if self.type_one_side: - for ii in range(0, self.ntypes): + for ii in range(self.ntypes): if not self._all_excluded(ii): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/matrix_{layer}_{ii}" @@ -298,7 +298,7 @@ def _get_matrix(self) -> dict[str, list[np.ndarray]]: else: matrix["layer_" + str(layer)].append(np.array([])) else: - for ii in range(0, self.ntypes * self.ntypes): + for ii in range(self.ntypes * self.ntypes): if ( ii // self.ntypes, ii % self.ntypes, @@ -508,7 +508,7 @@ def _get_layer_size(self) -> int: @cached_property def _n_all_excluded(self) -> int: """Then number of types excluding all types.""" - return sum(int(self._all_excluded(ii)) for ii in range(0, self.ntypes)) + return sum(int(self._all_excluded(ii)) for ii in range(self.ntypes)) def _convert_numpy_to_tensor(self) -> None: """Convert self.data from np.ndarray to tf.Tensor.""" diff --git a/deepmd/tf2/atomic_model/linear_atomic_model.py b/deepmd/tf2/atomic_model/linear_atomic_model.py index f6030291f1..2c810ee7dd 100644 --- a/deepmd/tf2/atomic_model/linear_atomic_model.py +++ b/deepmd/tf2/atomic_model/linear_atomic_model.py @@ -24,7 +24,7 @@ def __setattr__(self, name: str, value: Any) -> None: if name == "zbl_weight": # discard since it's only used in tests # to fix TensorFlow tracing mutation error: Cannot mutate 'FlaxModule' from different trace level - return + return None return super().__setattr__(name, value) def forward_common_atomic( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index fe5f542933..ab14e59bb1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -416,6 +416,7 @@ def descrpt_se_zm_args() -> list[Argument]: doc_basis_type = "Radial basis type. Supported values are `bessel` and `gaussian`." doc_n_radial = "Number of radial basis functions." doc_radial_mlp = "Hidden layer sizes for radial networks. An output layer of size (l_schedule[0]+extra_node_l+1)*channels will be automatically appended. Use 0 as a placeholder to be replaced by channels." + doc_edge_norm = "Whether to apply standard channel RMSNorm on cutoff-vanishing feature branches. Setting to `false` removes RMSNorm from the radial network, environment-seed FiLM, and cross-focus competition, and uses unit-floor residual scaling for post-SO(2) messages. Setting to `false` is recommended." doc_use_env_seed = ( "If True, seed the initial node state with local-environment information: " "apply environment matrix FiLM conditioning on l=0 features using 4D " @@ -734,6 +735,13 @@ def descrpt_se_zm_args() -> list[Argument]: default=[0], doc=doc_radial_mlp, ), + Argument( + "edge_norm", + bool, + optional=True, + default=True, + doc=doc_edge_norm, + ), Argument( "use_env_seed", bool, diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index c6b839a8f6..d7609f9897 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -119,7 +119,7 @@ def _smth_descriptor(jdata: dict[str, Any]) -> dict[str, Any]: dict with descriptor parameters """ descriptor = {} - seed = jdata.get("seed", None) + seed = jdata.get("seed") if seed is not None: descriptor["seed"] = seed descriptor["type"] = "se_a" @@ -150,7 +150,7 @@ def _fitting_net(jdata: dict[str, Any]) -> dict[str, Any]: """ fitting_net = {} - seed = jdata.get("seed", None) + seed = jdata.get("seed") if seed is not None: fitting_net["seed"] = seed fitting_net["neuron"] = j_deprecated(jdata, "fitting_neuron", ["n_neuron"]) @@ -228,7 +228,7 @@ def _training(jdata: dict[str, Any]) -> dict[str, Any]: dict with training parameters """ training = {} - seed = jdata.get("seed", None) + seed = jdata.get("seed") if seed is not None: training["seed"] = seed diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 9d13cb4699..9713a43d26 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -897,11 +897,11 @@ def get_data( The data system """ systems = jdata["systems"] - rglob_patterns = jdata.get("rglob_patterns", None) + rglob_patterns = jdata.get("rglob_patterns") systems = process_systems(systems, patterns=rglob_patterns) batch_size = jdata["batch_size"] - sys_probs = jdata.get("sys_probs", None) + sys_probs = jdata.get("sys_probs") auto_prob = jdata.get("auto_prob", "prob_sys_size") optional_type_map = not multi_task_mode diff --git a/deepmd/utils/tabulate_math.py b/deepmd/utils/tabulate_math.py index 93fe903e12..401840ecd5 100644 --- a/deepmd/utils/tabulate_math.py +++ b/deepmd/utils/tabulate_math.py @@ -536,7 +536,7 @@ def _get_network_variable(self, var_name: str) -> dict: result["layer_" + str(layer)].append(node) elif self.descrpt_type == "A": if self.type_one_side: - for ii in range(0, self.ntypes): + for ii in range(self.ntypes): if not self._all_excluded(ii): node = self.embedding_net_nodes[ii]["layers"][layer - 1][ "@variables" @@ -545,7 +545,7 @@ def _get_network_variable(self, var_name: str) -> dict: else: result["layer_" + str(layer)].append(np.array([])) else: - for ii in range(0, self.ntypes * self.ntypes): + for ii in range(self.ntypes * self.ntypes): if ( ii // self.ntypes, ii % self.ntypes, @@ -570,7 +570,7 @@ def _get_network_variable(self, var_name: str) -> dict: result["layer_" + str(layer)].append(node) elif self.descrpt_type == "R": if self.type_one_side: - for ii in range(0, self.ntypes): + for ii in range(self.ntypes): if not self._all_excluded(ii): node = self.embedding_net_nodes[ii]["layers"][layer - 1][ "@variables" @@ -579,7 +579,7 @@ def _get_network_variable(self, var_name: str) -> dict: else: result["layer_" + str(layer)].append(np.array([])) else: - for ii in range(0, self.ntypes * self.ntypes): + for ii in range(self.ntypes * self.ntypes): if ( ii // self.ntypes, ii % self.ntypes, @@ -607,4 +607,4 @@ def _convert_numpy_to_tensor(self) -> None: @cached_property def _n_all_excluded(self) -> int: """The number of types excluding all types.""" - return sum(int(self._all_excluded(ii)) for ii in range(0, self.ntypes)) + return sum(int(self._all_excluded(ii)) for ii in range(self.ntypes)) diff --git a/doc/conf.py b/doc/conf.py index c58073d5c0..d7ad2c5673 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -128,7 +128,6 @@ # Tell sphinx what the pygments highlight language should be. # highlight_language = 'cpp' -# myst_heading_anchors = 4 nb_execution_mode = "off" diff --git a/examples/water/dpa4/input.json b/examples/water/dpa4/input.json index 1126c7107a..1819a3afad 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -10,6 +10,7 @@ "rcut": 6.0, "channels": 32, "n_radial": 16, + "edge_norm": false, "use_env_seed": true, "edge_cartesian": false, "node_cartesian": "none", diff --git a/pyproject.toml b/pyproject.toml index 779cce37cf..56aa1edd1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -412,6 +412,14 @@ select = [ "PYI", # flake8-pyi "ANN", # type annotations "B905", # zip-without-explicit-strict + "PIE808", # unnecessary-range-start + "PT022", # pytest-useless-yield-fixture + "RET502", # implicit-return-value + "SIM910", # dict-get-with-none-default + "TD006", # invalid-todo-capitalization + "PLR1711", # useless-return + "PLR1733", # unnecessary-dict-index-lookup + "PLR2044", # empty-comment ] ignore = [ diff --git a/source/tests/dpa_adapt/test_conditions.py b/source/tests/dpa_adapt/test_conditions.py index f2aefad714..182588a36c 100644 --- a/source/tests/dpa_adapt/test_conditions.py +++ b/source/tests/dpa_adapt/test_conditions.py @@ -77,7 +77,6 @@ def _mock_extract_features(self, systems): def _mock_load_descriptor_model(self): self._checkpoint_type_map = ["Cu", "O"] - return None # ====================================================================== diff --git a/source/tests/dpa_adapt/test_finetuner_strategies.py b/source/tests/dpa_adapt/test_finetuner_strategies.py index 6f280ffa90..d332153060 100644 --- a/source/tests/dpa_adapt/test_finetuner_strategies.py +++ b/source/tests/dpa_adapt/test_finetuner_strategies.py @@ -347,7 +347,6 @@ def test_fit_dispatch_calls_training_path(self, tmp_path): def _mock_load_descriptor_model_cache_test(self): self._checkpoint_type_map = ["H", "O"] - return None class TestFitDescriptorCache: diff --git a/source/tests/dpa_adapt/test_predictor.py b/source/tests/dpa_adapt/test_predictor.py index 5f0a8135ba..685183eb59 100644 --- a/source/tests/dpa_adapt/test_predictor.py +++ b/source/tests/dpa_adapt/test_predictor.py @@ -83,7 +83,6 @@ def _mock_extract_features(self, systems): def _mock_load_descriptor_model(self): self._checkpoint_type_map = ["Cu", "O"] - return None # --------------------------------------------------------------------------- diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 72e8a47ec7..b720a71789 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -27,6 +27,9 @@ import numpy as np import optax +from deepmd.dpmodel.output_def import ( + OutputVariableCategory, +) from deepmd.dpmodel.train import ( DEFAULT_TASK_KEY, RankContext, @@ -46,6 +49,12 @@ jnp, nnx, ) +from deepmd.jax.infer.deep_eval import ( + DeepEval, +) +from deepmd.jax.model.hlo import ( + HLO, +) from deepmd.jax.train.trainer import ( DPTrainer, _copy_matching_state_tree, @@ -852,17 +861,19 @@ def test_update_sel_supports_multitask(self, get_nbor_stat, get_data) -> None: def test_freeze_entrypoint_uses_checkpoint_pointer( self, serialize_from_file, deserialize_to_file ) -> None: - """Freeze resolves the stable checkpoint pointer without Hessian options.""" + """Freeze resolves the stable checkpoint pointer and forwards Hessian.""" checkpoint_dir = self.work_dir / "ckpt" checkpoint_dir.mkdir() (checkpoint_dir / "checkpoint").write_text("model-1.jax") serialize_from_file.return_value = {"model": {}, "model_def_script": {}} - freeze(checkpoint_folder=str(checkpoint_dir), output="frozen_model") + freeze( + checkpoint_folder=str(checkpoint_dir), output="frozen_model", hessian=True + ) serialize_from_file.assert_called_once_with(str(checkpoint_dir / "model-1.jax")) deserialize_to_file.assert_called_once_with( - "frozen_model.hlo", serialize_from_file.return_value + "frozen_model.hlo", serialize_from_file.return_value, hessian=True ) @patch("deepmd.jax.entrypoints.main.freeze") @@ -874,11 +885,63 @@ def test_main_dispatches_freeze(self, freeze_entrypoint) -> None: log_path=None, checkpoint_folder=".", output="frozen_model", + hessian=False, ) main(args) freeze_entrypoint.assert_called_once() + self.assertIn("hessian", freeze_entrypoint.call_args.kwargs) + self.assertFalse(freeze_entrypoint.call_args.kwargs["hessian"]) + + def test_hlo_hessian_mode_updates_output_def(self) -> None: + """HLO output definition should expose Hessian when requested.""" + hlo = object.__new__(HLO) + hlo._model_output_type = ["energy"] + hlo.model_def_script = json.dumps({"hessian_mode": True}) + + output_def = hlo.model_output_def() + + self.assertTrue(output_def["energy"].r_hessian) + self.assertIn("energy_derv_r_derv_r", output_def.keys()) + + def test_deep_eval_requests_hessian_for_hessian_model(self) -> None: + """Non-atomic JAX evaluation should request Hessian outputs.""" + hlo = object.__new__(HLO) + hlo._model_output_type = ["energy"] + hlo.model_def_script = json.dumps({"hessian_mode": True}) + deep_eval = object.__new__(DeepEval) + deep_eval.output_def = hlo.model_output_def() + deep_eval.dp = SimpleNamespace( + get_model_def_script=lambda: json.dumps({"hessian_mode": True}) + ) + + request_defs = deep_eval._get_request_defs(atomic=False) + + self.assertTrue(deep_eval.get_has_hessian()) + self.assertIn( + OutputVariableCategory.DERV_R_DERV_R, + {odef.category for odef in request_defs}, + ) + + def test_deep_eval_skips_hessian_for_standard_model(self) -> None: + """Standard JAX evaluation should not request Hessian outputs.""" + hlo = object.__new__(HLO) + hlo._model_output_type = ["energy"] + hlo.model_def_script = json.dumps({"hessian_mode": False}) + deep_eval = object.__new__(DeepEval) + deep_eval.output_def = hlo.model_output_def() + deep_eval.dp = SimpleNamespace( + get_model_def_script=lambda: json.dumps({"hessian_mode": False}) + ) + + request_defs = deep_eval._get_request_defs(atomic=False) + + self.assertFalse(deep_eval.get_has_hessian()) + self.assertNotIn( + OutputVariableCategory.DERV_R_DERV_R, + {odef.category for odef in request_defs}, + ) def test_jax_finetune_state_copy_preserves_random_fitting_target_leaves() -> None: diff --git a/source/tests/pd/model/test_force_grad.py b/source/tests/pd/model/test_force_grad.py index eb6975afec..d1387771f1 100644 --- a/source/tests/pd/model/test_force_grad.py +++ b/source/tests/pd/model/test_force_grad.py @@ -37,7 +37,7 @@ def __init__( def get_disturb(self, index, atom_index, axis_index, delta): for i in range( - 0, len(self.dirs) + 1 + len(self.dirs) + 1 ): # note: if different sets can be merged, prefix sum is unused to calculate if index < self.prefix_sum[i]: break diff --git a/source/tests/pd/model/test_rotation.py b/source/tests/pd/model/test_rotation.py index 48f5ae8983..f47f1d8201 100644 --- a/source/tests/pd/model/test_rotation.py +++ b/source/tests/pd/model/test_rotation.py @@ -35,7 +35,7 @@ def __init__( def get_rotation(self, index, rotation_matrix): for i in range( - 0, len(self.dirs) + 1 + len(self.dirs) + 1 ): # note: if different sets can be merged, prefix sum is unused to calculate if index < self.prefix_sum[i]: break diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index 161e8a87e6..835d7c6a5d 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -24,6 +24,8 @@ ForceEmbedding, InnerClamp, NodeCartesianTensorProduct, + RadialBasis, + RadialMLP, SeZMDirectForceHead, SO2Linear, SpinEmbedding, @@ -1826,6 +1828,131 @@ def test_invalid_params(self) -> None: InnerClamp(1.0, 1.0) +class TestEdgeNorm(_SeZMTestCase): + """The ``edge_norm`` switch and its effect on cutoff smoothness. + + The descriptor exposes a single ``edge_norm`` flag; internally it drives the + ``RadialMLP.radial_norm`` hidden RMSNorm, the FiLM scale/shift norms, and the + cross-focus competition norm, and selects the post-SO(2) residual scaling + floor. The RadialMLP-level tests exercise the radial mechanism directly; the + descriptor test checks the umbrella propagation. + """ + + def setUp(self) -> None: + super().setUp() + self.dtype = torch.float64 + self.rcut = 6.0 + + def _radial_feature_curve(self, *, radial_norm: bool, seed: int) -> torch.Tensor: + """Radial features over a near-cutoff distance sweep. + + ``RadialBasis`` bakes in the C^3 envelope, so ``edge_rbf`` vanishes at + ``rcut``. With ``radial_norm=True`` the hidden RMSNorm divides that + envelope out and its ``eps`` floor is crossed near ``rcut``, injecting a + localized curvature spike; ``radial_norm=False`` drops the RMSNorm so the + feature stays smooth. Both variants share the same linear weights (same + ``seed``), isolating the effect of the norm. + + Parameters + ---------- + radial_norm : bool + Whether the RadialMLP keeps its hidden RMSNorm. + seed : int + Seed shared by both variants so the linear weights match. + + Returns + ------- + torch.Tensor + Radial features with shape (N, out_dim) over the distance sweep. + """ + basis = RadialBasis(rcut=self.rcut, n_radial=8, exponent=7, dtype=self.dtype) + mlp = RadialMLP( + [8, 12, 8], radial_norm=radial_norm, dtype=self.dtype, seed=seed + ) + r = torch.linspace( + 0.5 * self.rcut, + 0.9995 * self.rcut, + 4000, + dtype=self.dtype, + device=self.device, + ).view(-1, 1) + with torch.no_grad(): + return mlp(basis(r)) + + @staticmethod + def _peak_curvature(feat: torch.Tensor) -> float: + """Peak absolute second finite difference of the feature L2 norm.""" + y = feat.norm(dim=1) + return float((y[2:] - 2.0 * y[1:-1] + y[:-2]).abs().max()) + + def test_radial_norm_false_removes_cutoff_curvature_spike(self) -> None: + """``radial_norm=False`` suppresses the eps-crossing kink near ``rcut``.""" + feat_norm = self._radial_feature_curve(radial_norm=True, seed=3) + feat_smooth = self._radial_feature_curve(radial_norm=False, seed=3) + # Both vanish at rcut: edge_rbf -> 0 and RadialMLP(0) = 0 (bias=False). + self.assertLess(feat_smooth[-1].abs().max().item(), 1.0e-8) + # The normalized variant floor-crosses just inside rcut; dropping the + # RMSNorm removes that localized curvature spike by a wide margin. + self.assertLess( + self._peak_curvature(feat_smooth) * 5.0, + self._peak_curvature(feat_norm), + ) + + def test_radial_norm_structure_and_serialization(self) -> None: + """The flag toggles the hidden RMSNorm and round-trips through serialize.""" + for radial_norm in (True, False): + with self.subTest(radial_norm=radial_norm): + mlp = RadialMLP( + [8, 12, 8], radial_norm=radial_norm, dtype=self.dtype, seed=5 + ) + has_norm = any(type(m).__name__ == "RMSNorm" for m in mlp.net) + self.assertEqual(has_norm, radial_norm) + + restored = RadialMLP.deserialize(mlp.serialize()) + self.assertEqual(restored.radial_norm, radial_norm) + x = torch.rand(16, 8, dtype=self.dtype, device=self.device) + with torch.no_grad(): + torch.testing.assert_close(mlp(x), restored(x)) + + def test_edge_norm_gates_all_cutoff_vanishing_norms(self) -> None: + """``edge_norm`` controls every cutoff-vanishing normalization path.""" + for edge_norm in (True, False): + with self.subTest(edge_norm=edge_norm): + desc = DescrptSeZM( + **_descriptor_kwargs( + edge_norm=edge_norm, + use_env_seed=True, + n_focus=2, + sandwich_norm=[True, True, True, True], + precision="float64", + ) + ) + # radial MLP hidden RMSNorm + radial_has_norm = any( + type(m).__name__ == "RMSNorm" for m in desc.radial_embedding.net + ) + self.assertEqual(radial_has_norm, edge_norm) + # env-seed FiLM scale/shift norms + self.assertEqual( + type(desc.film_scale_norm).__name__ == "ScalarRMSNorm", edge_norm + ) + self.assertEqual( + type(desc.film_shift_norm).__name__ == "ScalarRMSNorm", edge_norm + ) + # cross-focus competition norm (n_focus>1 -> competition active) + focus_norm_mod = desc.blocks[0].so2_conv.focus_compete_norm + self.assertEqual( + type(focus_norm_mod).__name__ == "ScalarRMSNorm", edge_norm + ) + # Only the post-SO(2) residual branch uses unit-floor scaling. + expected_eps = 1.0e-5 if edge_norm else 1.0 + self.assertEqual(desc.blocks[0].post_so2_norm.eps, expected_eps) + self.assertEqual(desc.blocks[0].pre_so2_norm.eps, 1.0e-5) + self.assertEqual(desc.blocks[0].pre_ffn_norms[0].eps, 1.0e-5) + self.assertEqual(desc.blocks[0].post_ffn_norms[0].eps, 1.0e-5) + self.assertEqual(desc.serialize()["config"]["edge_norm"], edge_norm) + + class TestDescriptorEnergyCurveSmoothness(_SeZMTestCase): """Test PES smoothness from scaled symmetric eight-atom probes.""" diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index 345583b092..c286f0fca7 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -806,7 +806,8 @@ def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: @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: + @pytest.mark.parametrize("eps", [1.0e-5, 1.0]) + def test_equivariant_rmsnorm(self, lmax, n_focus, eps) -> None: from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( EquivariantRMSNorm as DPEquivariantRMSNorm, ) @@ -815,10 +816,16 @@ def test_equivariant_rmsnorm(self, lmax, n_focus) -> None: ) pt_mod = PTEquivariantRMSNorm( - lmax, self.channels, n_focus, dtype=torch.float64, trainable=True + lmax, + self.channels, + n_focus, + eps=eps, + dtype=torch.float64, + trainable=True, ) self._perturb(pt_mod, 2040) serialized = pt_mod.serialize() + assert serialized["config"]["eps"] == eps # pt state_dict key contract: 2 parameters + 2 persistent buffers assert set(serialized["@variables"]) == { "adam_scale", @@ -832,6 +839,40 @@ def test_equivariant_rmsnorm(self, lmax, n_focus) -> None: 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_eps_one_reparameterization(self) -> None: + """An epsilon of one preserves the legacy branch function class.""" + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm as PTEquivariantRMSNorm, + ) + + legacy = PTEquivariantRMSNorm( + 2, + self.channels, + 1, + eps=1.0e-5, + dtype=torch.float64, + trainable=True, + ) + unit_scale = PTEquivariantRMSNorm( + 2, + self.channels, + 1, + eps=1.0, + dtype=torch.float64, + trainable=True, + ) + self._perturb(legacy, 2042) + unit_scale.load_state_dict(legacy.state_dict()) + + rng = np.random.default_rng(2043) + x = to_pt(rng.normal(size=(17, 9, 1, self.channels))) + torch.testing.assert_close( + legacy(x), + unit_scale(x / np.sqrt(legacy.eps)), + rtol=PT_RTOL, + atol=PT_ATOL, + ) + def test_equivariant_rmsnorm_roundtrip(self) -> None: from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( EquivariantRMSNorm as DPEquivariantRMSNorm, @@ -3293,6 +3334,17 @@ def test_block_sandwich_norm(self, sandwich) -> None: ) self._assert_block_parity(pt_mod, dp_mod, kwargs) + def test_block_post_so2_eps_one(self) -> None: + pt_mod, dp_mod, kwargs = self._build_block_pair( + so2_post_norm=True, + so2_post_norm_eps=1.0, + ) + assert pt_mod.post_so2_norm.eps == 1.0 + assert dp_mod.post_so2_norm.eps == 1.0 + assert pt_mod.pre_ffn_norms[0].eps == 1.0e-5 + assert dp_mod.pre_ffn_norms[0].eps == 1.0e-5 + 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) @@ -3521,6 +3573,22 @@ def test_descriptor(self, use_env_seed, n_blocks) -> None: ) self._assert_descr_parity(pt_mod, dp_mod) + @pytest.mark.parametrize( + "edge_norm", [False, True] + ) # cutoff-vanishing normalization modes + def test_descriptor_edge_norm(self, edge_norm) -> None: + # edge_norm=False drops the radial MLP RMSNorm, turns the FiLM scale/shift + # norms into identity pass-throughs, drops the focus-compete norm, and + # selects unit-floor post-SO(2) residual scaling in both backends. + pt_mod, dp_mod, _ = self._build_descr_pair( + edge_norm=edge_norm, use_env_seed=True, n_focus=2 + ) + assert dp_mod.edge_norm == edge_norm + expected_eps = 1.0e-5 if edge_norm else 1.0 + assert pt_mod.blocks[0].post_so2_norm.eps == expected_eps + assert dp_mod.blocks[0].post_so2_norm.eps == expected_eps + self._assert_descr_parity(pt_mod, dp_mod) + @pytest.mark.parametrize( "exclude_types", [[], [(0, 0)]] ) # pair-exclusion off vs on diff --git a/source/tests/pt/model/test_force_grad.py b/source/tests/pt/model/test_force_grad.py index 27bf660241..599f0d1830 100644 --- a/source/tests/pt/model/test_force_grad.py +++ b/source/tests/pt/model/test_force_grad.py @@ -37,7 +37,7 @@ def __init__( def get_disturb(self, index, atom_index, axis_index, delta): for i in range( - 0, len(self.dirs) + 1 + len(self.dirs) + 1 ): # note: if different sets can be merged, prefix sum is unused to calculate if index < self.prefix_sum[i]: break diff --git a/source/tests/pt/model/test_rotation.py b/source/tests/pt/model/test_rotation.py index f2abc6f9fd..e940177cb8 100644 --- a/source/tests/pt/model/test_rotation.py +++ b/source/tests/pt/model/test_rotation.py @@ -35,7 +35,7 @@ def __init__( def get_rotation(self, index, rotation_matrix): for i in range( - 0, len(self.dirs) + 1 + len(self.dirs) + 1 ): # note: if different sets can be merged, prefix sum is unused to calculate if index < self.prefix_sum[i]: break diff --git a/source/tests/pt_expt/conftest.py b/source/tests/pt_expt/conftest.py index 228c6104ae..ff9347b132 100644 --- a/source/tests/pt_expt/conftest.py +++ b/source/tests/pt_expt/conftest.py @@ -43,11 +43,9 @@ def _pop_device_contexts() -> list: def _clear_leaked_device_context_session(): """Pop any stale DeviceContext once at session start.""" _pop_device_contexts() - yield @pytest.fixture(autouse=True) def _clear_leaked_device_context(): """Pop any stale ``DeviceContext`` before each test (safety net).""" _pop_device_contexts() - yield diff --git a/source/tests/tf/test_compat_input.py b/source/tests/tf/test_compat_input.py index 4e60b6bf3e..6fa82a4cc9 100644 --- a/source/tests/tf/test_compat_input.py +++ b/source/tests/tf/test_compat_input.py @@ -35,9 +35,9 @@ def assertDictAlmostEqual(self, d1, d2, msg=None, places=7) -> None: self.assertEqual(d1.keys(), d2.keys()) for kk, vv in d1.items(): if isinstance(vv, dict): - self.assertDictAlmostEqual(d1[kk], d2[kk], msg=msg) + self.assertDictAlmostEqual(vv, d2[kk], msg=msg) else: - self.assertAlmostEqual(d1[kk], d2[kk], places=places, msg=msg) + self.assertAlmostEqual(vv, d2[kk], places=places, msg=msg) def test_json_yaml_equal(self) -> None: inputs = ("water_v1", "water_se_a_v1") diff --git a/source/tests/tf/test_descrpt_sea_ef_rot.py b/source/tests/tf/test_descrpt_sea_ef_rot.py index 356e9dd5bc..ff48f8c52b 100644 --- a/source/tests/tf/test_descrpt_sea_ef_rot.py +++ b/source/tests/tf/test_descrpt_sea_ef_rot.py @@ -205,7 +205,7 @@ def test_rot_axis(self, suffix="") -> None: # print(v_ae0) # print(ae0) - for kk in range(0, self.natoms[0]): + for kk in range(self.natoms[0]): # print(f0) theta = 45.0 / 180.0 * np.pi rr0 = self.rotate_mat(defield[0][kk * 3 : kk * 3 + 3], theta) @@ -256,8 +256,8 @@ def test_rot_axis(self, suffix="") -> None: self.tnatoms: self.natoms, }, ) - for ii in range(0, self.natoms[0]): - for jj in range(0, self.natoms[0]): + for ii in range(self.natoms[0]): + for jj in range(self.natoms[0]): diff = ( dcoord[0][3 * jj : 3 * jj + 3] - dcoord[0][3 * ii : 3 * ii + 3] ) @@ -394,7 +394,7 @@ def test_rot_diff_axis(self, suffix="") -> None: }, ) - for ii in range(0, self.natoms[0]): + for ii in range(self.natoms[0]): self.assertNotAlmostEqual(p_ae0[ii], p_ae1[ii]) self.assertNotAlmostEqual(v_ae0[ii], v_ae1[ii]) diff --git a/source/tests/tf/test_nvnmd_entrypoints.py b/source/tests/tf/test_nvnmd_entrypoints.py index 4a6877761e..531b30748c 100644 --- a/source/tests/tf/test_nvnmd_entrypoints.py +++ b/source/tests/tf/test_nvnmd_entrypoints.py @@ -55,7 +55,6 @@ def test_mapt_cnn_v0(self) -> None: # mapt mapObj = MapTable(config_file, weight_file, map_file) mapt = mapObj.build_map() - # N = 32 x = np.reshape(np.arange(N) / N * (8.0**2), [-1, 1]) pred = mapObj.mapping2(x, {"s": mapt["s"]}, mapt["cfg_u2s"]) @@ -127,7 +126,6 @@ def test_mapt_cnn_v0(self) -> None: -0.37758207, ] np.testing.assert_almost_equal(pred, ref_dout, 8) - # N = 4 x = np.reshape(np.arange(N) / N * 16, [-1, 1]) pred = mapObj.mapping2(x, {"g": mapt["g"]}, mapt["cfg_s2g"]) @@ -455,7 +453,6 @@ def test_model_qnn_v0(self) -> None: dic_ph["natoms_vec"]: natoms_vec_dat, dic_ph["default_mesh"]: mesh_dat, } - # sess = self.cached_session().__enter__() sess.run(tf.global_variables_initializer()) # get tensordic @@ -528,7 +525,6 @@ def test_mapt_cnn_v1(self) -> None: mapObj = MapTable(config_file, weight_file, map_file) mapObj.Gs_Gt_mode = 0 mapt = mapObj.build_map() - # N = 32 x = np.reshape(np.arange(N) / N * (8.0**2), [-1, 1]) pred = mapObj.mapping2(x, {"s": mapt["s"]}, mapt["cfg_u2s"]) @@ -568,7 +564,6 @@ def test_mapt_cnn_v1(self) -> None: 0.00000000e00, ] np.testing.assert_almost_equal(pred, ref_dout, 8) - # N = 4 x = np.reshape(np.arange(N) / N * 16, [-1, 1]) pred = mapObj.mapping2(x, {"g": mapt["g"]}, mapt["cfg_s2g"]) @@ -772,7 +767,6 @@ def test_model_qnn_v1(self) -> None: dic_ph["natoms_vec"]: natoms_vec_dat, dic_ph["default_mesh"]: mesh_dat, } - # sess = self.cached_session().__enter__() sess.run(tf.global_variables_initializer()) # get tensordic