From d8b27ce67e3141a677f69deb19f03fe24b41f8f5 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 17 Sep 2026 06:59:24 +0000 Subject: [PATCH 01/12] feat(pt-expt): add DPA4C-LR model with non-periodic LES/SOG long-range correction Add model type `dpa4c_lr` (pt_expt backend only): the compact DPA4C descriptor plus a non-periodic long-range term evaluated from per-atom latent charges produced by the `dpa4c_lr` fitting. Two kernels are available: LES (erf(alpha*r)/r, default) and a trainable sum-of-Gaussians (SOG) with multi-channel latent charges. The long-range all-pairs graph is built outside the compiled region and passed into the graph lower, so `torch.compile` sees no data-dependent neighbor construction. Only non-periodic systems (box=None) are supported. Periodic misuse is rejected with a clear error at every entry point: training statistics (fail-fast before descriptor stats), eager/ragged forward, and DeepPot evaluation. `.pt2`/`.pte` freeze is not yet implemented. Includes dp-test (DeepPot.eval) support, user documentation (doc/model/dpa4c_lr.md), and unit tests covering LES/SOG force-virial finite differences, compile parity, serialization, checkpoint inference, and periodic-input rejection. --- AGENTS.md | 1 + deepmd/dpmodel/descriptor/dpa4c.py | 11 +- deepmd/dpmodel/fitting/__init__.py | 4 + deepmd/dpmodel/fitting/dpa4c_lr.py | 426 ++++++++++ deepmd/dpmodel/model/make_model.py | 12 +- deepmd/dpmodel/utils/default_neighbor_list.py | 9 +- deepmd/pt_expt/fitting/__init__.py | 4 + deepmd/pt_expt/fitting/dpa4c_lr.py | 78 ++ deepmd/pt_expt/infer/deep_eval.py | 65 +- deepmd/pt_expt/model/__init__.py | 4 + deepmd/pt_expt/model/dpa4c_lr_model.py | 793 ++++++++++++++++++ deepmd/pt_expt/model/get_model.py | 51 ++ deepmd/pt_expt/model/make_model.py | 11 +- deepmd/pt_expt/train/training.py | 82 ++ deepmd/pt_expt/utils/graph_builder.py | 2 + deepmd/utils/argcheck.py | 178 ++++ doc/model/dpa4c_lr.md | 80 ++ doc/model/index.rst | 1 + source/tests/pt_expt/model/test_dpa4c_lr.py | 488 +++++++++++ 19 files changed, 2284 insertions(+), 16 deletions(-) create mode 100644 deepmd/dpmodel/fitting/dpa4c_lr.py create mode 100644 deepmd/pt_expt/fitting/dpa4c_lr.py create mode 100644 deepmd/pt_expt/model/dpa4c_lr_model.py create mode 100644 doc/model/dpa4c_lr.md create mode 100644 source/tests/pt_expt/model/test_dpa4c_lr.py diff --git a/AGENTS.md b/AGENTS.md index 617c904f11..5ca920a19b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -165,6 +165,7 @@ source/ # C++ source code and tests - **Configuration**: Use `input_torch.json` format typically - **Training**: `dp --pt train input_torch.json` - **Requirements**: `torch` package +- **DPA4C-LR long-range variant**: model type `dpa4c_lr` with fitting type `dpa4c_lr` adds a non-periodic LES long-range term on top of the compact DPA4C descriptor (`deepmd/pt_expt/model/dpa4c_lr_model.py`, `deepmd/pt_expt/fitting/dpa4c_lr.py`). Only `box=None` systems are supported and only the `pt_expt` backend. The long-range all-pairs graph is built outside the compiled region and passed into the graph lower alongside the short-range graph, so `torch.compile` sees no data-dependent neighbor construction and memory scales as `Σ_f n_f^2` rather than the full-batch `N^2`. `.pt2`/`.pte` freeze is not yet implemented. Tests: `pytest source/tests/pt_expt/model/test_dpa4c_lr.py`. ### JAX Backend diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index e38d8374da..a5b4adfd1e 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -1452,11 +1452,12 @@ def compute_input_stats( for frame in self._calibration_frames(system): coord = xp.asarray(frame["coord"], dtype=dtype, device=device) atype = xp.asarray(frame["atype"], device=device) - box = ( - None - if frame["box"] is None - else xp.asarray(frame["box"], dtype=dtype, device=device) - ) + if frame["box"] is None: + box = None + else: + box = xp.asarray(frame["box"], dtype=dtype, device=device) + if xp.all(box == 0): + box = None spin = ( None if frame["spin"] is None diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 37b0a1b731..7cbaa0f931 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -8,6 +8,9 @@ from .dpa4_ener import ( SeZMEnergyFittingNet, ) +from .dpa4c_lr import ( + DPA4CLRFitting, +) from .ener_fitting import ( EnergyFittingNet, ) @@ -26,6 +29,7 @@ __all__ = [ "DOSFittingNet", + "DPA4CLRFitting", "DipoleFitting", "EnergyFittingNet", "InvarFitting", diff --git a/deepmd/dpmodel/fitting/dpa4c_lr.py b/deepmd/dpmodel/fitting/dpa4c_lr.py new file mode 100644 index 0000000000..41c58d3ef3 --- /dev/null +++ b/deepmd/dpmodel/fitting/dpa4c_lr.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range latent-charge fitting for DPA4C (dpmodel backend). + +This fitting extends the GLU energy fitting of ``dpa4_ener`` with a second, +lightweight GLU branch that predicts per-atom latent charges. The actual +long-range energy/force is evaluated by the model layer +(:class:`deepmd.pt_expt.model.dpa4c_lr_model.DPA4CLREnergyModel`). + +Only non-periodic systems are supported by the model layer; this fitting +class itself is backend-agnostic and can run on any array-API backend. +""" + +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.array_api import ( + array_api_compat, +) +from deepmd.dpmodel.common import ( + get_xp_precision, + to_numpy_array, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) + +from .dpa4_ener import ( + GLUFittingNet, + SeZMEnergyFittingNet, + SeZMNetworkCollection, +) +from .invar_fitting import ( + InvarFitting, +) + +SOG_DEFAULT_B = 1.62976708826776469 +SOG_DEFAULT_SIGMA = 2.180230445405648 +SOG_DEFAULT_M = 12 + + +@InvarFitting.register("dpa4c_lr") +@fitting_check_output +class DPA4CLRFitting(SeZMEnergyFittingNet): + """DPA4C energy fitting with an independent latent-charge branch. + + Parameters + ---------- + dim_out_lr : int + Latent charge dimension. Must be 1 for the LES kernel; SOG supports + multiple latent channels. + neuron_lr : list[int], optional + Hidden widths of the latent-charge GLU branch. + bias_atom_q : list[float] | None + Per-element latent-charge bias. If ``None`` a trainable array is + allocated. + les_alpha : float + Initial width of the LES kernel ``erf(alpha * r) / r``. + amp : float | list[float] | None + Initial amplitudes of the SOG Gaussian levels. + bandwidth : float | list[float] | None + Initial widths of the SOG Gaussian levels. + b, sigma, M + Hyperparameters used to initialize geometric SOG widths when + ``bandwidth`` is not provided. + use_charge_constraint : bool + Whether to enforce that the sum of latent charges equals the total + charge. In the graph path this is applied by the model layer, because + the per-frame sum is not available inside ``call_graph``. + lr_kernel : str + Long-range kernel selector. DPA4C supports ``"les"`` and ``"sog"``; + it has no equivariant latent, so ``"dipole"`` is impossible. + """ + + def __init__( + self, + *args: Any, + dim_out_lr: int = 1, + neuron_lr: list[int] | None = None, + bias_atom_q: list[float] | None = None, + amp: float | list[float] | None = None, + bandwidth: float | list[float] | None = None, + les_alpha: float | None = None, + b: float | None = None, + sigma: float | None = None, + M: int | None = None, + use_charge_constraint: bool = False, + lr_kernel: str = "les", + **kwargs: Any, + ) -> None: + # Attributes consulted by output_def() during fitting_check_output must + # be set before super().__init__(). + self.dim_out_lr = int(dim_out_lr) + if self.dim_out_lr < 1: + raise ValueError(f"`dim_out_lr` must be positive, got {dim_out_lr}.") + self.neuron_lr = [64, 64] if neuron_lr is None else [int(x) for x in neuron_lr] + self.use_charge_constraint = bool(use_charge_constraint) + self.bias_atom_q_bound = 3.0 + lr_kernel = str(lr_kernel).lower() + if lr_kernel not in ("les", "sog", "dipole"): + raise ValueError( + f"`lr_kernel` should be 'les', 'sog' or 'dipole', got {lr_kernel!r}." + ) + if lr_kernel == "dipole": + raise NotImplementedError( + "DPA4C has no equivariant latent output, so `lr_kernel='dipole'` " + "is not supported." + ) + if lr_kernel == "les" and self.dim_out_lr != 1: + raise ValueError( + "`lr_kernel='les'` requires `dim_out_lr=1`, got " + f"{self.dim_out_lr}. Use `lr_kernel='sog'` for multiple latent " + "channels." + ) + self.lr_kernel = lr_kernel + + super().__init__(*args, **kwargs) + + # --- latent-charge GLU branch --- + case_dim = 0 if self.case_film_embd else self.dim_case_embd + in_dim = ( + self.dim_descrpt + + self.numb_fparam + + (0 if self.use_aparam_as_mask else self.numb_aparam) + + case_dim + ) + n_networks = self.ntypes if not self.mixed_types else 1 + self.filter_layers_lr = SeZMNetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="sezm_fitting_network", + networks=[ + GLUFittingNet( + in_dim, + self.dim_out_lr, + self.neuron_lr, + activation_function=self.activation_function, + resnet_dt=self.resnet_dt, + precision=self.precision, + bias_out=self.bias_out, + seed=child_seed(self.seed, 10_000 + idx), + trainable=self.trainable, + descriptor_dim=self.dim_descrpt, + dim_case_embd=self.dim_case_embd, + case_film_embd=self.case_film_embd, + ) + for idx in range(n_networks) + ], + ) + + # --- per-element latent-charge bias --- + if bias_atom_q is None: + self.bias_atom_q = np.zeros((self.ntypes, self.dim_out_lr), dtype=self.prec) + else: + bias_atom_q_t = np.asarray(bias_atom_q, dtype=self.prec).reshape( + (self.ntypes, self.dim_out_lr) + ) + self.bias_atom_q = bias_atom_q_t + + # --- LES width --- + les_alpha_value = float(les_alpha) if les_alpha is not None else 1.0 + if les_alpha_value <= 0.0: + raise ValueError("`les_alpha` should be positive.") + self.les_alpha = np.array(les_alpha_value, dtype=self.prec) + + # --- SOG Gaussian levels --- + self.b = SOG_DEFAULT_B if b is None else float(b) + self.sigma = SOG_DEFAULT_SIGMA if sigma is None else float(sigma) + self.M = max(1, SOG_DEFAULT_M if M is None else int(M)) + if self.lr_kernel == "sog": + if not np.isfinite(self.b) or self.b <= 0.0: + raise ValueError("`b` should be finite and positive.") + if not np.isfinite(self.sigma) or self.sigma <= 0.0: + raise ValueError("`sigma` should be finite and positive.") + + if bandwidth is None: + bandwidth_array = self.sigma * np.power( + self.b, np.arange(self.M, dtype=self.prec) + ) + else: + bandwidth_array = np.asarray(bandwidth, dtype=self.prec).reshape(-1) + if bandwidth_array.size < 1 or not np.all(np.isfinite(bandwidth_array)): + raise ValueError("`bandwidth` should contain finite values.") + if not np.all(bandwidth_array > 0.0): + raise ValueError("`bandwidth` should contain positive values.") + + if amp is None: + amp_array = np.full_like(bandwidth_array, 4.0 * np.pi * np.log(self.b)) + else: + amp_array = np.asarray(amp, dtype=self.prec).reshape(-1) + if amp_array.size == 1 and bandwidth_array.size > 1: + amp_array = np.full_like(bandwidth_array, amp_array[0]) + elif amp_array.size != bandwidth_array.size: + raise ValueError( + "`amp` should be scalar or have the same length as `bandwidth`." + ) + if not np.all(np.isfinite(amp_array)): + raise ValueError("`amp` should contain finite values.") + + self.bandwidth = bandwidth_array + self.amp = amp_array + self.M = int(bandwidth_array.size) + else: + # LES checkpoints created before SOG existed have no such keys. + # Keep the unused SOG attributes out of the module state entirely. + self.bandwidth = None + self.amp = None + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + "energy", + [1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "latent_charge", + [self.dim_out_lr], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def call( + self, + descriptor: Any, + atype: Any, + gr: Any | None = None, + g2: Any | None = None, + h2: Any | None = None, + fparam: Any | None = None, + aparam: Any | None = None, + ) -> dict[str, Any]: + """Return energy and latent charges.""" + energy_ret = super().call( + descriptor, atype, gr=gr, g2=g2, h2=h2, fparam=fparam, aparam=aparam + ) + xp = array_api_compat.array_namespace(descriptor, atype) + nf = descriptor.shape[0] + if self.numb_fparam > 0 and fparam is None: + assert self.default_fparam_tensor is not None + default_fparam_tensor = xp.asarray( + self.default_fparam_tensor, + dtype=descriptor.dtype, + device=array_api_compat.device(descriptor), + ) + fparam = xp.tile( + xp.reshape(default_fparam_tensor, (1, self.numb_fparam)), (nf, 1) + ) + latent_charge = self._forward_lr( + descriptor, atype, fparam=fparam, aparam=aparam + ) + return { + **energy_ret, + "latent_charge": latent_charge, + } + + def _forward_lr( + self, + descriptor: Any, + atype: Any, + fparam: Any | None = None, + aparam: Any | None = None, + ) -> Any: + """Evaluate the latent-charge branch.""" + xp = array_api_compat.array_namespace(descriptor, atype) + nf, nloc, nd = descriptor.shape + if nd != self.dim_descrpt: + raise ValueError( + f"descriptor dim {nd} inconsistent with {self.dim_descrpt}" + ) + + xx = descriptor + if self.numb_fparam > 0: + assert fparam is not None + fparam = xp.reshape(fparam, (nf, self.numb_fparam)) + fparam_avg = xp.asarray( + self.fparam_avg, + dtype=fparam.dtype, + device=array_api_compat.device(fparam), + ) + fparam_inv_std = xp.asarray( + self.fparam_inv_std, + dtype=fparam.dtype, + device=array_api_compat.device(fparam), + ) + fparam = (fparam - fparam_avg) * fparam_inv_std + fparam = xp.tile( + xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1) + ) + xx = xp.concat([xx, fparam], axis=-1) + + if self.numb_aparam > 0 and not self.use_aparam_as_mask: + assert aparam is not None + aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam)) + aparam_avg = xp.asarray( + self.aparam_avg, + dtype=aparam.dtype, + device=array_api_compat.device(aparam), + ) + aparam_inv_std = xp.asarray( + self.aparam_inv_std, + dtype=aparam.dtype, + device=array_api_compat.device(aparam), + ) + aparam = (aparam - aparam_avg) * aparam_inv_std + xx = xp.concat([xx, aparam], axis=-1) + + if self.dim_case_embd > 0: + assert self.case_embd is not None + case_embd = xp.asarray( + self.case_embd, + dtype=descriptor.dtype, + device=array_api_compat.device(descriptor), + ) + case_embd = xp.tile(xp.reshape(case_embd, (1, 1, -1)), (nf, nloc, 1)) + xx = xp.concat([xx, case_embd], axis=-1) + + net_dtype = get_xp_precision(xp, self.precision) + xx = xp.astype(xx, net_dtype) + lr_out = xp.zeros( + (nf, nloc, self.dim_out_lr), + dtype=net_dtype, + device=array_api_compat.device(descriptor), + ) + if self.mixed_types: + lr_out = lr_out + self.filter_layers_lr.networks[0](xx) + else: + for type_i, ll in enumerate(self.filter_layers_lr.networks): + mask = xp.astype(atype == type_i, descriptor.dtype) + atom_property = ll(xx) + atom_property = xp.where( + xp.astype(mask[:, :, None], xp.bool), + atom_property, + xp.zeros_like(atom_property), + ) + lr_out = lr_out + atom_property + + lr_out = lr_out + self._get_lr_bias(atype, xp) + + # Total-charge constraint is applied by the model layer + # (:class:`~deepmd.pt_expt.model.dpa4c_lr_model.DPA4CLREnergyModel`), + # which has access to the frame-level ``charge_spin`` target. + exclude_mask = self.emask.build_type_exclude_mask(atype) + exclude_mask = xp.astype(exclude_mask, xp.bool) + lr_out = xp.where(exclude_mask[:, :, None], lr_out, xp.zeros_like(lr_out)) + return lr_out + + def _get_lr_bias(self, atype: Any, xp: Any) -> Any: + """Return per-atom latent-charge bias.""" + atype_flat = xp.reshape(xp.astype(atype, xp.int64), (-1,)) + target_dtype = get_xp_precision(xp, self.precision) + target_device = array_api_compat.device(atype) + bias_atom_q = self.bias_atom_q + if array_api_compat.is_torch_array(bias_atom_q): + # Keep the autograd graph: ``.to`` preserves grad_fn, while + # ``xp.asarray`` on a Parameter may return a detached view in some + # array_api_compat versions. + bias_atom_q = bias_atom_q.to(dtype=target_dtype, device=target_device) + else: + bias_atom_q = xp.asarray( + bias_atom_q, dtype=target_dtype, device=target_device + ) + bias = xp.take(bias_atom_q, atype_flat, axis=0).reshape( + *atype.shape, self.dim_out_lr + ) + return self.bias_atom_q_bound * xp.tanh(bias / self.bias_atom_q_bound) + + def serialize(self) -> dict: + """Serialize the fitting.""" + data = super().serialize() + data["type"] = "dpa4c_lr" + data["dim_out_lr"] = self.dim_out_lr + data["neuron_lr"] = self.neuron_lr + data["nets_lr"] = self.filter_layers_lr.serialize() + data["use_charge_constraint"] = self.use_charge_constraint + data["lr_kernel"] = self.lr_kernel + data["b"] = self.b + data["sigma"] = self.sigma + data["M"] = self.M + variables = data.setdefault("@variables", {}) + variables["bias_atom_q"] = to_numpy_array(self.bias_atom_q) + variables["les_alpha"] = to_numpy_array(self.les_alpha) + if self.lr_kernel == "sog": + variables["amp"] = to_numpy_array(self.amp) + variables["bandwidth"] = to_numpy_array(self.bandwidth) + return data + + @classmethod + def deserialize(cls, data: dict) -> "DPA4CLRFitting": + """Deserialize the fitting.""" + data = data.copy() + variables = data.get("@variables", {}).copy() + bias_atom_q = variables.pop("bias_atom_q", None) + les_alpha = variables.pop("les_alpha", None) + amp = variables.pop("amp", None) + bandwidth = variables.pop("bandwidth", None) + nets_lr = data.pop("nets_lr", None) + # Feed LR variables through the constructor. This preserves their + # Parameter status when ``cls`` is the pt_expt torch wrapper; assigning + # numpy arrays after construction would replace Parameters with buffers. + if bias_atom_q is not None: + data["bias_atom_q"] = np.asarray(bias_atom_q).tolist() + if les_alpha is not None: + data["les_alpha"] = float(np.asarray(les_alpha).reshape(-1)[0]) + if amp is not None: + data["amp"] = np.asarray(amp).reshape(-1).tolist() + if bandwidth is not None: + data["bandwidth"] = np.asarray(bandwidth).reshape(-1).tolist() + data["@variables"] = variables + obj = super().deserialize(data) + if nets_lr is not None: + obj.filter_layers_lr = SeZMNetworkCollection.deserialize(nets_lr) + return obj diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b74fb6b697..143304b329 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -131,6 +131,11 @@ def model_call_from_call_lower( nframes, nloc = atype.shape[:2] cc, bb, fp, ap = coord, box, fparam, aparam del coord, box, fparam, aparam + # Non-periodic datasets may carry an all-zero cell; treat it as no box. + if bb is not None: + xp_bb = array_api_compat.array_namespace(bb) + if bool(xp_bb.all(bb == 0)): + bb = None builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() # Model-level pair exclusion is a nlist-BUILD transform (decision #18/A4): # the BUILDER owns it (mirroring build_neighbor_graph on the graph path), so @@ -1071,10 +1076,15 @@ def format_nlist( """ n_nf, n_nloc, n_nnei = nlist.shape mixed_types = self.mixed_types() + # A frame can never contribute more than nall neighbors; cap the + # requested neighbor count so sentinel capacities (e.g. DPA4C's + # effectively-unbounded sel) do not allocate absurd padding. + nall = extended_atype.shape[1] + nnei = min(sum(self.get_sel()), nall) ret = self._format_nlist( extended_coord, nlist, - sum(self.get_sel()), + nnei, extra_nlist_sort=extra_nlist_sort, ) if not mixed_types: diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index 2316ab0545..6bd1703ed4 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -729,6 +729,11 @@ def build( extended_coord, extended_atype, mapping = extend_coord_with_ghosts( coord_normalized, atype, box, rcut ) + # A frame can never contribute more than nall neighbors; cap the + # requested neighbor count so sentinel capacities (e.g. DPA4C's + # effectively-unbounded sel) do not allocate absurd padding. + nall = extended_atype.shape[1] + nsel = min(sum(sel), nall) # Types are distinguished in the lower interface, so keep them merged # here. The dense path remains faster for small systems; the spatial # path avoids the nloc*nall distance matrix for larger supported arrays. @@ -738,7 +743,7 @@ def build( extended_atype, nloc, rcut, - sum(sel), + nsel, pair_excl=pair_excl, ) else: @@ -747,7 +752,7 @@ def build( extended_atype, nloc, rcut, - sel, + nsel, distinguish_types=False, pair_excl=pair_excl, ) diff --git a/deepmd/pt_expt/fitting/__init__.py b/deepmd/pt_expt/fitting/__init__.py index 63bcc6c2ce..41b79b8c03 100644 --- a/deepmd/pt_expt/fitting/__init__.py +++ b/deepmd/pt_expt/fitting/__init__.py @@ -11,6 +11,9 @@ from .dpa4_ener import ( SeZMEnergyFittingNet, ) +from .dpa4c_lr import ( + DPA4CLRFitting, +) from .ener_fitting import ( EnergyFittingNet, ) @@ -30,6 +33,7 @@ __all__ = [ "BaseFitting", "DOSFittingNet", + "DPA4CLRFitting", "DipoleFitting", "EnergyFittingNet", "InvarFitting", diff --git a/deepmd/pt_expt/fitting/dpa4c_lr.py b/deepmd/pt_expt/fitting/dpa4c_lr.py new file mode 100644 index 0000000000..93a988c011 --- /dev/null +++ b/deepmd/pt_expt/fitting/dpa4c_lr.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""PyTorch-exportable wrapper for the DPA4C long-range fitting.""" + +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.fitting.dpa4c_lr import ( + DPA4CLRFitting as DPA4CLRFittingDP, +) +from deepmd.pt_expt.common import ( + register_dpmodel_mapping, + torch_module, +) + +from .base_fitting import ( + BaseFitting, +) + + +@BaseFitting.register("dpa4c_lr") +@torch_module +class DPA4CLRFitting(DPA4CLRFittingDP): + """DPA4C long-range fitting net for the pt_expt backend.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + # In the dpmodel these are plain numpy arrays, which the wrapper stores + # as non-trainable buffers. Promote only the parameters used by the + # selected kernel so every registered Parameter participates in loss + # backpropagation. + if self.lr_kernel == "les": + self._promote_to_parameter("les_alpha") + elif self.lr_kernel == "sog": + self._promote_to_parameter("amp") + self._promote_to_parameter("bandwidth") + self._promote_to_parameter("bias_atom_q") + + def _promote_to_parameter(self, name: str) -> None: + value = getattr(self, name, None) + if value is None or isinstance(value, torch.nn.Parameter): + return + tensor = ( + value.detach().clone() + if isinstance(value, torch.Tensor) + else torch.as_tensor(value) + ) + if name in self._buffers: + del self._buffers[name] + self.register_parameter( + name, torch.nn.Parameter(tensor, requires_grad=bool(self.trainable)) + ) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + return self.call(*args, **kwargs) + + def _load_from_state_dict( + self, + state_dict: dict, + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + if self.lr_kernel == "les": + # A short-lived SOG implementation registered these unused buffers + # on LES models. Ignore them when reading such a checkpoint, while + # older LES checkpoints naturally have neither key. + state_dict.pop(prefix + "amp", None) + state_dict.pop(prefix + "bandwidth", None) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + +register_dpmodel_mapping( + DPA4CLRFittingDP, + lambda v: DPA4CLRFitting.deserialize(v.serialize()), +) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..981835bbc9 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -664,6 +664,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ) use_graph_lower = model_uses_graph_lower(model) + needs_long_range_edges = bool(getattr(model, "_needs_long_range_edges", False)) self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -717,6 +718,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ), "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", + "needs_long_range_edges": needs_long_range_edges, } if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( @@ -805,7 +807,17 @@ def _eager_runner_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None = None, + lr_edge_index: torch.Tensor | None = None, + lr_edge_vec: torch.Tensor | None = None, + lr_edge_mask: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: + long_range_kwargs = {} + if needs_long_range_edges: + long_range_kwargs = { + "lr_edge_index": lr_edge_index, + "lr_edge_vec": lr_edge_vec, + "lr_edge_mask": lr_edge_mask, + } model_ret = model.forward_common_lower_graph( atype, n_node, @@ -822,6 +834,7 @@ def _eager_runner_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + **long_range_kwargs, ) return _translate_energy_keys( model_ret, @@ -2410,6 +2423,19 @@ def _eval_model_graph( # Build the carry-all graph at its exact edge count; the exported edge # axis is dynamic. graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) + lr_graph = None + if self.metadata.get("needs_long_range_edges", False): + if box_input is not None and bool(np.any(box_input != 0)): + raise NotImplementedError( + "DPA4C-LR supports only non-periodic systems (box=None)." + ) + lr_graph = self._build_eval_graph( + coord_input, + atom_types, + None, + DEVICE, + long_range=True, + ) atype_t = torch.tensor( np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE @@ -2528,6 +2554,24 @@ def _eval_model_graph( aparam_t, charge_spin_t, ) + if lr_graph is not None: + model_inputs += ( + torch.as_tensor( + lr_graph.edge_index, + dtype=torch.int64, + device=DEVICE, + ), + torch.as_tensor( + lr_graph.edge_vec, + dtype=edge_dtype, + device=DEVICE, + ), + torch.as_tensor( + lr_graph.edge_mask, + dtype=torch.bool, + device=DEVICE, + ), + ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: @@ -2552,6 +2596,8 @@ def _build_eval_graph( atom_types: np.ndarray, box_input: np.ndarray | None, device: "torch.device", + *, + long_range: bool = False, ) -> "NeighborGraph": """Build the carry-all NeighborGraph for graph-lower inference. @@ -2567,7 +2613,11 @@ def _build_eval_graph( destination-major graph-form ``.pt2`` ABI after construction; the caller transfers its fields to the model device. """ - method = self._neighbor_graph_method + # The LR graph is all-pairs by definition. Do not send its artificial + # 1e6 cutoff through the spatial NV/vesin builders: their capacity/cell + # estimates are intended for finite local cutoffs and can explode for + # this sentinel radius. + method = "dense" if long_range else self._neighbor_graph_method if method == "auto": coord_arr = np.asarray(coord_input) nf = int(coord_arr.shape[0]) if coord_arr.ndim >= 2 else 1 @@ -2576,7 +2626,8 @@ def _build_eval_graph( # (decision #18): apply it here so the exported ``.pt2`` lower consumes a # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++ # ``applyPairExclusion`` and the eager dpmodel/pt_expt build path). - pair_excl = self._model_pair_excl() + pair_excl = None if long_range else self._model_pair_excl() + rcut = 1.0e6 if long_range else self._rcut builder_device = torch.device("cpu") if method == "cell" else device # The fused builder writes the whole destination-major payload from one # search. It applies only where nothing has to be filtered or masked @@ -2616,7 +2667,7 @@ def _build_eval_graph( ) if box_input is not None else None, - self._rcut, + rcut, edge_dtype=edge_dtype, ) if method == "dense": @@ -2628,7 +2679,7 @@ def _build_eval_graph( coord_input, atom_types, box_input, - self._rcut, + rcut, canonicalize=True, pair_excl=pair_excl, ) @@ -2641,7 +2692,7 @@ def _build_eval_graph( coord_input, atom_types, box_input, - self._rcut, + rcut, canonicalize=True, pair_excl=pair_excl, ) @@ -2679,7 +2730,7 @@ def _build_eval_graph( cc, aa, bb, - self._rcut, + rcut, canonicalize=True, pair_excl=pair_excl, ) @@ -2691,7 +2742,7 @@ def _build_eval_graph( cc, aa, bb, - self._rcut, + rcut, canonicalize=True, pair_excl=pair_excl, ) diff --git a/deepmd/pt_expt/model/__init__.py b/deepmd/pt_expt/model/__init__.py index 22a1c2e7db..cdcdb9ff6c 100644 --- a/deepmd/pt_expt/model/__init__.py +++ b/deepmd/pt_expt/model/__init__.py @@ -18,6 +18,9 @@ from .dpa4_model import ( DPA4EnergyModel, ) +from .dpa4c_lr_model import ( + DPA4CLREnergyModel, +) from .ener_model import ( EnergyModel, ) @@ -49,6 +52,7 @@ __all__ = [ "BaseModel", "DOSModel", + "DPA4CLREnergyModel", "DPA4EnergyModel", "DPZBLModel", "DipoleModel", diff --git a/deepmd/pt_expt/model/dpa4c_lr_model.py b/deepmd/pt_expt/model/dpa4c_lr_model.py new file mode 100644 index 0000000000..faaadf3406 --- /dev/null +++ b/deepmd/pt_expt/model/dpa4c_lr_model.py @@ -0,0 +1,793 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""DPA4C energy model with a non-periodic LES/SOG term (pt_expt). + +The short-range part is the standard DPA4C graph-native energy model. The +long-range correction is evaluated from per-atom latent charges produced by +:class:`deepmd.pt_expt.fitting.dpa4c_lr.DPA4CLRFitting`. + +The graph lower consumes two edge sets: + +* the short-range edges (within ``rcut``) used by the descriptor; +* the long-range edges (all pairs within each frame) used by the LR kernel. + +Both edge sets are built eagerly outside the compiled region and passed in as +inputs, so the compiled graph lower contains no data-dependent Python ints. +The force is computed via two gradient paths: + +* ``grad(E_SR + E_LR, edge_vec)`` gives the short-range force plus the + charge-response term ``∂E_LR/∂q · dq/dR``; +* ``grad(E_LR, lr_edge_vec)`` gives the geometry term ``∂E_LR/∂r|_q``, + assembled with the ``edge_force_virial`` conventions. + +Only non-periodic systems (``box=None``) are supported. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np +import torch +import torch.nn as nn + +from deepmd.dpmodel import ( + get_hessian_name, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + compact_nodes, + edge_force_virial, + expand_node_values, + frame_id_from_n_node, + segment_sum, +) +from deepmd.pt_expt.kernels.utils import ( + cuda_infer_level, +) +from deepmd.pt_expt.fitting.dpa4c_lr import ( + DPA4CLRFitting, +) +from deepmd.pt_expt.model.ener_model import ( + EnergyModel, +) +from deepmd.pt_expt.model.make_model import ( + _cal_hessian_ext_graph, +) +from deepmd.pt_expt.model.model import ( + BaseModel, +) +from deepmd.pt_expt.utils.graph_builder import ( + build_neighbor_graph_for_method, + build_ragged_neighbor_graph, + resolve_neighbor_graph_method, +) + +from .edge_transform_output import ( + fit_output_to_model_output_graph, +) +from .transform_output import ( + fit_output_to_model_output, +) + +E2_PER_ANGSTROM_TO_EV = 14.3996454784255 + + +def _les_kernel_from_squared_distance( + r_sq: torch.Tensor, + les_alpha: torch.Tensor, +) -> torch.Tensor: + """Evaluate ``erf(alpha * r) / r`` with a finite derivative at ``r=0``.""" + r = torch.sqrt(r_sq + 1e-12) + return torch.erf(les_alpha * r) / r + + +def _disable_pointless_cumsum_pattern() -> None: + """Remove Inductor's ``pointless_cumsum_replacement`` pattern. + + The pattern rewrites ``cumsum(full(shape, fill_value))`` into + ``arange * fill_value``; it crashes when ``fill_value`` is a symbolic + shape (``Node``) instead of a scalar. DPA4C's local-frame construction + hits exactly that case under symbolic tracing, so we drop the pattern. + The pattern is a pure optimization; removing it is semantics-preserving. + """ + try: + from torch._inductor.fx_passes import ( + post_grad, + ) + except Exception: + return + for pm in post_grad.pass_patterns: + for key in list(pm.patterns.keys()): + pm.patterns[key] = [ + p + for p in pm.patterns[key] + if getattr(getattr(p, "handler", None), "__name__", "") + != "pointless_cumsum_replacement" + ] + + +class _LESKernel(nn.Module): + """Batched non-periodic LES charge-charge kernel.""" + + def __init__(self, norm_factor: float = E2_PER_ANGSTROM_TO_EV) -> None: + super().__init__() + self.norm_factor = float(norm_factor) + + def forward( + self, + positions: torch.Tensor, + latent_charges: torch.Tensor, + les_alpha: torch.Tensor, + ) -> torch.Tensor: + """Return the per-frame LES energy. + + Parameters + ---------- + positions + Local atom coordinates, shape ``(nf, nloc, 3)``. + latent_charges + Per-atom latent charges, shape ``(nf, nloc, 1)``. + les_alpha + Scalar LES width parameter. + + Returns + ------- + energy + Per-frame energy, shape ``(nf,)``. + """ + nloc = positions.shape[1] + q = latent_charges.to(positions.dtype) + r_ij = positions.unsqueeze(1) - positions.unsqueeze(2) # (nf,nloc,nloc,3) + r_sq = torch.sum(r_ij * r_ij, dim=-1) + eye = 1.0 - torch.eye( + nloc, dtype=positions.dtype, device=positions.device + ).unsqueeze(0) + pair_q = q.unsqueeze(2) * q.unsqueeze(1) # (nf,nloc,nloc,1) + + alpha_t = les_alpha.to(positions.dtype) + kernel = _les_kernel_from_squared_distance(r_sq, alpha_t) + kernel = kernel * eye + pot = 0.5 * torch.sum(pair_q.squeeze(-1) * kernel, dim=(1, 2)) # (nf,) + return pot * self.norm_factor + + +class _SOGKernel(nn.Module): + """Batched non-periodic SOG kernel over distinct atom pairs ``i != j``.""" + + def __init__(self, norm_factor: float = E2_PER_ANGSTROM_TO_EV) -> None: + super().__init__() + self.norm_factor = float(norm_factor) + + def forward( + self, + positions: torch.Tensor, + latent_charges: torch.Tensor, + amp: torch.Tensor, + bandwidth: torch.Tensor, + ) -> torch.Tensor: + """Return the per-frame SOG energy.""" + nloc = positions.shape[1] + q = latent_charges.to(dtype=positions.dtype, device=positions.device) + r_ij = positions.unsqueeze(1) - positions.unsqueeze(2) + r_sq = torch.sum(r_ij * r_ij, dim=-1, keepdim=True) + eye = 1.0 - torch.eye( + nloc, dtype=positions.dtype, device=positions.device + ).unsqueeze(0) + pair_q = q.unsqueeze(2) * q.unsqueeze(1) + + amp_t = amp.to(dtype=positions.dtype, device=positions.device).view(1, 1, 1, -1) + bandwidth_sq = ( + bandwidth.to(dtype=positions.dtype, device=positions.device) + .square() + .view(1, 1, 1, -1) + ) + kernel = torch.sum(amp_t * torch.exp(-0.5 * r_sq / bandwidth_sq), dim=-1) + kernel = kernel * eye + pot = 0.5 * torch.sum(pair_q * kernel.unsqueeze(-1), dim=(1, 2, 3)) + + return pot * self.norm_factor + + +@BaseModel.register("dpa4c_lr") +class DPA4CLREnergyModel(EnergyModel): + """DPA4C model with a non-periodic LES or SOG correction.""" + + model_type = "dpa4c_lr" + _needs_long_range_edges = True + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + fitting = self.get_fitting_net() + if not isinstance(fitting, DPA4CLRFitting): + raise TypeError( + "DPA4CLREnergyModel requires a DPA4CLRFitting fitting net " + f"(type 'dpa4c_lr'), got {type(fitting).__name__}." + ) + self.lr_kernel = fitting.lr_kernel + if self.lr_kernel not in ("les", "sog"): + raise NotImplementedError( + f"Unsupported DPA4C-LR kernel {self.lr_kernel!r}." + ) + self.les_kernel = _LESKernel(norm_factor=E2_PER_ANGSTROM_TO_EV) + self.sog_kernel = _SOGKernel(norm_factor=E2_PER_ANGSTROM_TO_EV) + _disable_pointless_cumsum_pattern() + + def compute_or_load_stat( + self, + sampled_func: Any, + stat_file_path: Any = None, + preset_observed_type: list[str] | None = None, + ) -> None: + """Reject periodic training data before statistics are computed. + + Both LR kernels are non-periodic, so fail fast with a clear message + instead of letting the descriptor statistics crash deep inside the + graph construction on ghost atoms. + """ + sampled = sampled_func() + for frame in sampled: + box = frame.get("box") + if box is not None and bool(np.any(np.asarray(box) != 0)): + raise NotImplementedError( + "DPA4C-LR supports only non-periodic systems (box=None); " + "periodic long-range summation is not implemented. " + "Remove the box data (i.e. use nopbc systems) to train " + "this model." + ) + super().compute_or_load_stat( + lambda: sampled, + stat_file_path, + preset_observed_type=preset_observed_type, + ) + + def call_common( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + **kwargs: Any, + ) -> dict[str, torch.Tensor]: + """Reject periodic cells; both LR kernels are non-periodic.""" + if box is not None: + xp = array_api_compat.array_namespace(box) + if bool(xp.any(box != 0)): + raise NotImplementedError( + "DPA4C-LR supports only non-periodic systems (box=None)." + ) + return super().call_common(coord, atype, box=box, **kwargs) + + # ------------------------------------------------------------------ + # Graph lower with LES/SOG + # ------------------------------------------------------------------ + + def forward_common_lower_graph( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor | None = None, + destination_row_ptr: torch.Tensor | None = None, + source_order: torch.Tensor | None = None, + source_row_ptr: torch.Tensor | None = None, + destination_sorted: bool = False, + do_atomic_virial: bool = False, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + comm_dict: dict | None = None, + lr_edge_index: torch.Tensor | None = None, + lr_edge_vec: torch.Tensor | None = None, + lr_edge_mask: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Graph lower with the long-range correction. + + The short-range ``edge_*`` tensors and the long-range ``lr_edge_*`` + tensors are built eagerly by the caller and consumed as inputs, so the + compiled graph lower contains no data-dependent neighbor-list + construction. The force is computed via two gradient paths: + + * ``grad(E_SR + E_LR, edge_vec)`` gives the short-range force plus the + charge-response term ``∂E_LR/∂q · dq/dR``; + * ``grad(E_LR, lr_edge_vec)`` gives the geometry term + ``∂E_LR/∂r|_q``. + """ + if lr_edge_index is None or lr_edge_vec is None or lr_edge_mask is None: + raise ValueError( + "DPA4C-LR requires long-range edges in the graph lower; the caller " + "must pass lr_edge_index, lr_edge_vec, and lr_edge_mask." + ) + + # Short-range graph (edge_vec is the autograd leaf). + edge_vec = edge_vec.detach().requires_grad_(True) + if spin is not None: + spin = spin.detach().requires_grad_(True) + graph = NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + n_local=n_local, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=destination_sorted, + ) + + # Descriptor + fitting on the short-range graph. + atomic_ret = self.atomic_model.forward_common_atomic_graph( + graph, + atype, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, + comm_dict=comm_dict, + ) + + # LR energy on the long-range edge set. + lr_edge_vec = lr_edge_vec.detach().requires_grad_(True) + atomic_ret, e_lr = self._apply_lr_correction_graph( + atomic_ret, + lr_edge_index, + lr_edge_vec, + lr_edge_mask, + n_node, + fparam=fparam, + charge_spin=charge_spin, + ) + + # Short-range force + charge-response LR force via the short-range graph. + model_ret = fit_output_to_model_output_graph( + atomic_ret, + self.atomic_output_def(), + graph, + do_atomic_virial=do_atomic_virial, + create_graph=self.training, + mask=atomic_ret.get("mask"), + node_capacity=atype.shape[0], + n_local=n_local, + force_precision=edge_vec.dtype, + ) + + # Geometry LR force via the long-range graph. + if e_lr is not None and self.do_grad_r("energy"): + (g_lr,) = torch.autograd.grad( + e_lr.sum(), + lr_edge_vec, + create_graph=self.training, + retain_graph=True, + ) + lr_force, lr_atom_virial, lr_virial = edge_force_virial( + g_lr, + lr_edge_vec, + lr_edge_index, + lr_edge_mask, + n_node, + node_capacity=atype.shape[0], + ) + model_ret["energy_derv_r"] = model_ret["energy_derv_r"] + lr_force.reshape( + -1, 1, 3 + ).to(model_ret["energy_derv_r"].dtype) + if self.do_grad_c("energy"): + nf = n_node.shape[0] + model_ret["energy_derv_c_redu"] = model_ret[ + "energy_derv_c_redu" + ] + lr_virial.reshape(nf, 1, 9).to( + model_ret["energy_derv_c_redu"].dtype + ) + if do_atomic_virial: + model_ret["energy_derv_c"] = model_ret[ + "energy_derv_c" + ] + lr_atom_virial.reshape(-1, 1, 9).to( + model_ret["energy_derv_c"].dtype + ) + + return model_ret + + def _apply_lr_correction_graph( + self, + atomic_ret: dict[str, torch.Tensor], + lr_edge_index: torch.Tensor, + lr_edge_vec: torch.Tensor, + lr_edge_mask: torch.Tensor, + n_node: torch.Tensor, + fparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor | None]: + """Compute the LES or SOG energy on the long-range edge set. + + The long-range edges are built eagerly by the caller and passed in as + inputs. They contain only intra-frame all-pairs, so the storage scales + with ``Σ_f n_f^2`` rather than the full batch ``N^2``. + """ + if "latent_charge" not in atomic_ret: + return atomic_ret, None + + fitting = self.get_fitting_net() + nf = n_node.shape[0] + N = atomic_ret["latent_charge"].shape[0] + dtype = lr_edge_vec.dtype + + q = atomic_ret["latent_charge"] # (N, dim) + frame_id = frame_id_from_n_node(n_node, n_total=N) # (N,) + + # Charge constraint (flat, vectorized). + if fitting.use_charge_constraint: + q = self._constrain_charges_flat(q, frame_id, n_node, fparam, charge_spin) + + q_src = q[lr_edge_index[0]] # (E_lr, dim) + q_dst = q[lr_edge_index[1]] # (E_lr, dim) + r_sq = torch.sum(lr_edge_vec * lr_edge_vec, dim=-1) # (E_lr,) + if self.lr_kernel == "les": + les_alpha = fitting.les_alpha + if not isinstance(les_alpha, torch.Tensor): + les_alpha = torch.tensor( + les_alpha, dtype=dtype, device=lr_edge_vec.device + ) + else: + les_alpha = les_alpha.to(dtype=dtype, device=lr_edge_vec.device) + K = _les_kernel_from_squared_distance(r_sq, les_alpha) + else: + amp = fitting.amp.to(dtype=dtype, device=lr_edge_vec.device) + bandwidth_sq = fitting.bandwidth.to( + dtype=dtype, device=lr_edge_vec.device + ).square() + K = torch.sum( + amp.unsqueeze(0) + * torch.exp(-0.5 * r_sq.unsqueeze(-1) / bandwidth_sq.unsqueeze(0)), + dim=-1, + ) + e_edge = (q_src * q_dst).sum(-1) * K * lr_edge_mask.to(dtype) # (E_lr,) + edge_frame = frame_id[lr_edge_index[1]] # (E_lr,) frame of dst + e_lr_raw = segment_sum(e_edge, edge_frame, nf) + e_lr = e_lr_raw * (0.5 * E2_PER_ANGSTROM_TO_EV) + + # Spread the frame LR energy evenly over the real atoms. + nreal = n_node.to(e_lr.dtype).clamp_min(1) + e_lr_atom = (e_lr / nreal)[frame_id].unsqueeze(-1) # (N, 1) + atomic_ret["energy"] = atomic_ret["energy"] + e_lr_atom.to( + atomic_ret["energy"].dtype + ) + return atomic_ret, e_lr + + def _constrain_charges_flat( + self, + q: torch.Tensor, + frame_id: torch.Tensor, + n_node: torch.Tensor, + fparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> torch.Tensor: + """Shift the first latent-charge channel to match the target total charge. + + Flat (ragged) version: ``q`` is ``(N, dim)``, ``frame_id`` is ``(N,)``. + """ + nf = n_node.shape[0] + q_sum = segment_sum(q[:, 0], frame_id, nf) # (nf,) + nreal = n_node.to(q.dtype).clamp_min(1) + if charge_spin is not None and charge_spin.shape[-1] > 0: + target = charge_spin[:, 0].to(q.dtype) + elif fparam is not None and self.get_dim_fparam() > 0: + target = fparam[:, 0].to(q.dtype) + else: + target = torch.zeros(nf, dtype=q.dtype, device=q.device) + correction = (q_sum - target) / nreal # (nf,) + q_new = q.clone() + q_new[:, 0] = q[:, 0] - correction[frame_id] + return q_new + + # ------------------------------------------------------------------ + # Eager path: pass coord into the graph lower + # ------------------------------------------------------------------ + + def call_common_ragged( + self, + coord: torch.Tensor, + atype: torch.Tensor, + n_node: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Ragged graph forward with the LR correction.""" + if not (self.mixed_types() and self.atomic_model.uses_graph_lower()): + raise NotImplementedError( + "a flat node axis requires a mixed_types descriptor with a graph lower" + ) + if box is not None and torch.all(box == 0): + box = None + if box is not None: + raise NotImplementedError( + "DPA4C-LR supports only non-periodic systems (box=None)." + ) + method = getattr(self, "neighbor_graph_method", None) + if method is None: + method = resolve_neighbor_graph_method("auto", coord.device) + graph = build_ragged_neighbor_graph( + method, + coord, + atype, + n_node, + box, + self.get_rcut(), + self.atomic_model.pair_excl, + ) + # Long-range all-pairs graph, built eagerly outside the graph lower. + lr_graph = build_ragged_neighbor_graph( + method, coord, atype, n_node, None, 1e6, None + ) + predict = self.forward_common_lower_graph( + atype, + graph.n_node, + graph.n_node, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + destination_sorted=graph.destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, + lr_edge_index=lr_graph.edge_index, + lr_edge_vec=lr_graph.edge_vec, + lr_edge_mask=lr_graph.edge_mask, + ) + predict["n_node"] = graph.n_node + return predict + + def _call_common_graph( + self, + cc: torch.Tensor, + atype: torch.Tensor, + bb: torch.Tensor | None, + fp: torch.Tensor | None, + ap: torch.Tensor | None, + method: str, + do_atomic_virial: bool = False, + spin: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Carry-all graph forward with the LR correction. + + Mirrors the base :meth:`_call_common_graph` but builds the long-range + all-pairs graph eagerly and passes it into + :meth:`forward_common_lower_graph`. + """ + if not (self.mixed_types() and self.atomic_model.uses_graph_lower()): + raise NotImplementedError( + "neighbor_graph_method requires a mixed_types descriptor with a " + "graph lower." + ) + if bb is not None and torch.all(bb == 0): + bb = None + if bb is not None: + raise NotImplementedError( + "DPA4C-LR supports only non-periodic systems (box=None)." + ) + rcut = self.get_rcut() + _desc = getattr(self.atomic_model, "descriptor", None) + with_csr = ( + not self.training + and cuda_infer_level() >= 1 + and _desc is not None + and _desc.get_geo_compress() + ) + pair_excl = self.atomic_model.pair_excl + ng = build_neighbor_graph_for_method( + method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr + ) + nf, nloc = atype.shape[:2] + n_padded = nf * nloc + atype_flat = atype.reshape(n_padded) + ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_flat = atype_flat[node_index] + ap_flat = ( + ap.reshape(n_padded, ap.shape[-1])[node_index] if ap is not None else None + ) + spin_flat = spin.reshape(n_padded, 3)[node_index] if spin is not None else None + coord_flat = cc.reshape(n_padded, 3)[node_index] + + # Long-range all-pairs graph on the compact node axis. + lr_graph = build_ragged_neighbor_graph( + method, coord_flat, atype_flat, ng.n_node, None, 1e6, None + ) + + model_predict = self.forward_common_lower_graph( + atype_flat, + ng.n_node, + ng.n_node, + ng.edge_index, + ng.edge_vec, + ng.edge_mask, + ng.destination_order, + ng.destination_row_ptr, + ng.source_order, + ng.source_row_ptr, + destination_sorted=ng.destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fp, + aparam=ap_flat, + spin=spin_flat, + charge_spin=charge_spin, + lr_edge_index=lr_graph.edge_index, + lr_edge_vec=lr_graph.edge_vec, + lr_edge_mask=lr_graph.edge_mask, + ) + + # Unravel flat node-axis outputs back to the rectangular I/O boundary. + N = node_index.shape[0] + for k in list(model_predict.keys()): + v = model_predict[k] + if ( + v is not None + and not k.endswith("_redu") + and v.shape[:1] == torch.Size([N]) + ): + model_predict[k] = expand_node_values(v, node_index, n_padded).reshape( + nf, nloc, *v.shape[1:] + ) + + # Graph-native Hessian, if enabled. + aod = self.atomic_output_def() + for kk in aod.keys(): + vdef = aod[kk] + if vdef.reducible and vdef.r_hessian: + model_predict[get_hessian_name(kk)] = _cal_hessian_ext_graph( + model=self, + kk=kk, + vdef=vdef, + coord=cc, + atype=atype, + box=bb, + fparam=fp, + aparam=ap, + spin=spin, + charge_spin=charge_spin, + method=method, + pair_excl=pair_excl, + rcut=rcut, + create_graph=self.training, + ) + return model_predict + + # ------------------------------------------------------------------ + # Dense lower fallback (neighbor_graph_method="legacy") + # ------------------------------------------------------------------ + + def forward_common_atomic( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + extended_coord_corr: torch.Tensor | None = None, + comm_dict: dict | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Dense lower with the LR correction folded in.""" + extended_coord = extended_coord.detach().requires_grad_(True) + atomic_ret = self.atomic_model.forward_common_atomic( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + comm_dict=comm_dict, + charge_spin=charge_spin, + ) + atomic_ret = self._apply_lr_correction( + atomic_ret, + extended_coord, + fparam=fparam, + charge_spin=charge_spin, + ) + model_ret = fit_output_to_model_output( + atomic_ret, + self.atomic_output_def(), + extended_coord, + do_atomic_virial=do_atomic_virial, + create_graph=self.training, + mask=atomic_ret.get("mask"), + extended_coord_corr=extended_coord_corr, + ) + return model_ret + + def _apply_lr_correction( + self, + atomic_ret: dict[str, torch.Tensor], + extended_coord: torch.Tensor, + fparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + """Add the LR energy to the per-atom short-range energy (dense path).""" + if "latent_charge" not in atomic_ret: + return atomic_ret + + fitting = self.get_fitting_net() + nf, nloc = atomic_ret["energy"].shape[:2] + cc = extended_coord[:, :nloc, :] + q = atomic_ret["latent_charge"] + + mask = atomic_ret.get("mask") + if mask is None: + mask_bool = torch.ones(nf, nloc, dtype=torch.bool, device=cc.device) + else: + mask_bool = mask.to(torch.bool) + + if fitting.use_charge_constraint: + q = self._constrain_charges(q, mask_bool, fparam, charge_spin) + + q = q * mask_bool.unsqueeze(-1).to(q.dtype) + if self.lr_kernel == "les": + les_alpha = fitting.les_alpha + if not isinstance(les_alpha, torch.Tensor): + les_alpha = torch.tensor( + les_alpha, + dtype=cc.dtype, + device=cc.device, + ) + else: + les_alpha = les_alpha.to(dtype=cc.dtype, device=cc.device) + e_lr = self.les_kernel(cc, q, les_alpha) + else: + e_lr = self.sog_kernel(cc, q, fitting.amp, fitting.bandwidth) + + # Evenly spread the frame LR energy over the real atoms so that the + # per-atom sum matches the corrected total. + nreal = mask_bool.sum(dim=1, keepdim=True).clamp_min(1).to(e_lr.dtype) + e_lr_atom = (e_lr.view(nf, 1, 1) / nreal.view(nf, 1, 1)) * mask_bool.unsqueeze( + -1 + ).to(e_lr.dtype) + atomic_ret["energy"] = atomic_ret["energy"] + e_lr_atom.to( + atomic_ret["energy"].dtype + ) + return atomic_ret + + def _constrain_charges( + self, + q_rect: torch.Tensor, + mask_rect: torch.Tensor, + fparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> torch.Tensor: + """Shift the first latent-charge channel to match the target total charge.""" + nf, nloc, dim_out_lr = q_rect.shape + if charge_spin is not None and charge_spin.shape[-1] > 0: + target = charge_spin[:, 0].to(q_rect.dtype) + elif fparam is not None and self.get_dim_fparam() > 0: + target = fparam[:, 0].to(q_rect.dtype) + else: + target = torch.zeros(nf, dtype=q_rect.dtype, device=q_rect.device) + + nreal = mask_rect.sum(dim=1, keepdim=True).clamp_min(1).to(q_rect.dtype) + q0 = q_rect[:, :, 0] + q_sum = (q0 * mask_rect.to(q_rect.dtype)).sum(dim=1, keepdim=True) + q_target_per_atom = target.view(nf, 1) / nreal + q0 = q0 - (q_sum / nreal - q_target_per_atom) * mask_rect.to(q_rect.dtype) + if dim_out_lr == 1: + q_rect = q0.unsqueeze(-1) + else: + q_rect = torch.cat([q0.unsqueeze(-1), q_rect[:, :, 1:]], dim=-1) + return q_rect diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index c87ff016b5..574a0e5c9d 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -37,6 +37,9 @@ from deepmd.pt_expt.model.dpa4_model import ( DPA4EnergyModel, ) +from deepmd.pt_expt.model.dpa4c_lr_model import ( + DPA4CLREnergyModel, +) from deepmd.pt_expt.model.model import ( BaseModel, ) @@ -79,6 +82,53 @@ get_zbl_model = _model_factory.get_zbl_model +def get_dpa4c_lr_model(data: dict) -> DPA4CLREnergyModel: + """Build a pt_expt DPA4C-LR energy model from config. + + Mirrors :func:`get_sezm_model`: the descriptor must be DPA4/SeZM and the + fitting net must be ``dpa4c_lr``. + """ + data = copy.deepcopy(data) + data.pop("type", None) + data.setdefault("descriptor", {}) + data.setdefault("fitting_net", {}) + data["descriptor"].setdefault("type", "dpa4c") + if data["descriptor"]["type"] not in ("dpa4", "dpa4c", "DPA4", "sezm", "SeZM"): + raise ValueError( + "Model type 'dpa4c_lr' requires a DPA4/SeZM descriptor, but got " + f"descriptor type '{data['descriptor']['type']}'." + ) + if data["fitting_net"].get("type", "dpa4c_lr") != "dpa4c_lr": + raise ValueError( + "Model type 'dpa4c_lr' requires fitting type 'dpa4c_lr', but got " + f"'{data['fitting_net'].get('type')}'." + ) + + descriptor_exclude_types = [ + list(pair) for pair in (data["descriptor"].get("exclude_types") or []) + ] + if "pair_exclude_types" in data: + pair_exclude_types = [list(pair) for pair in (data["pair_exclude_types"] or [])] + if descriptor_exclude_types and descriptor_exclude_types != pair_exclude_types: + raise ValueError( + "DPA4C-LR `pair_exclude_types` and `descriptor.exclude_types` must match " + "when both are provided." + ) + else: + pair_exclude_types = descriptor_exclude_types + data["pair_exclude_types"] = pair_exclude_types + data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) + + descriptor, fitting, _ = _model_factory.get_model_components(data) + return DPA4CLREnergyModel( + descriptor=descriptor, + fitting=fitting, + type_map=data["type_map"], + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=pair_exclude_types, + ) + + def get_sezm_model(data: dict) -> BaseModel: """Build a pt_expt energy model from a DPA4/SeZM model config. @@ -410,5 +460,6 @@ def get_model(data: dict) -> BaseModel: "DPA4": get_sezm_model, "sezm": get_sezm_model, "SeZM": get_sezm_model, + "dpa4c_lr": get_dpa4c_lr_model, }, ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index c61a5edca0..73fca2fc3e 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -81,6 +81,11 @@ def _translate_energy_keys( ].squeeze(-2) if "mask" in model_ret: out["mask"] = model_ret["mask"] + # Preserve auxiliary outputs needed by post-compiled correction steps + # (e.g. DPA4C-LR latent charges) so a compiled graph can hand them to + # an eager post-processing hook. + if "latent_charge" in model_ret: + out["latent_charge"] = model_ret["latent_charge"] return out @@ -571,6 +576,7 @@ def forward_common_lower_graph( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, comm_dict: dict | None = None, + **kwargs: Any, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -775,7 +781,10 @@ def _resolve_graph_method( if "energy" not in self.atomic_output_def().keys(): return None if self.mixed_types() and self.atomic_model.uses_graph_lower(): - return getattr(self, "neighbor_graph_method", "dense") + method = getattr(self, "neighbor_graph_method", "dense") + # Honor an attribute-level opt-out as well: "legacy" forces + # the dense (nlist) path just like the explicit argument. + return None if method == "legacy" else method return None def call_common_ragged( diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..c45f74cb9b 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -768,6 +768,12 @@ def _trace_and_compile_graph( Representative optional inputs (or ``None``) so the traced branch matches what :meth:`_CompiledModel.forward` passes at run time. ``spin`` is the flat ``(N, 3)`` per-node moment. + + Long-range edges (``lr_edge_index`` / ``lr_edge_vec`` / ``lr_edge_mask``) + are always part of the compiled lower's input schema. Models that do not + use them receive empty tensors at run time; models with a long-range term + (e.g. DPA4C-LR) receive a real all-pairs graph built outside the compiled + region. compile_opts User-supplied inductor options (merged over the built-in defaults). task_buffers @@ -895,6 +901,20 @@ def _trace_and_compile_graph( s_spin = s_conditioning.pop(0) if spin is not None else None s_fparam, s_aparam, s_charge_spin = s_conditioning + # Build a synthetic long-range all-pairs graph for models that consume one + # (e.g. DPA4C-LR). The exact coordinates do not matter for shape tracing; + # we only need edge tensors whose dynamic axes are distinct from nf and N. + s_lr_coord = torch.randn( + trace_N, 3, dtype=GLOBAL_PT_FLOAT_PRECISION, device=_trace_device + ) + from deepmd.pt_expt.utils.graph_builder import ( + build_ragged_neighbor_graph, + ) + + s_lr_graph = build_ragged_neighbor_graph( + "dense", s_lr_coord, s_atype, s_n_node, None, 1e6, None + ) + def fn( atype: torch.Tensor, n_node: torch.Tensor, @@ -910,6 +930,9 @@ def fn( aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, spin: torch.Tensor | None, + lr_edge_index: torch.Tensor, + lr_edge_vec: torch.Tensor, + lr_edge_mask: torch.Tensor, *task_buf_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: # Patch task-specific buffers with the proxy tensors so make_fx records @@ -945,6 +968,9 @@ def fn( aparam=aparam, charge_spin=charge_spin, spin=spin, + lr_edge_index=lr_edge_index, + lr_edge_vec=lr_edge_vec, + lr_edge_mask=lr_edge_mask, ) if spin is None: return _translate_energy_keys( @@ -991,6 +1017,9 @@ def fn( s_aparam, s_charge_spin, s_spin, + s_lr_graph.edge_index, + s_lr_graph.edge_vec, + s_lr_graph.edge_mask, *task_buf_vals_trace, ) @@ -1239,6 +1268,8 @@ def forward( coord_3d = coord.detach().reshape(nframes, nloc, 3) box_flat = box.detach().reshape(nframes, 9) if box is not None else None + if box_flat is not None and torch.all(box_flat == 0): + box_flat = None if box_flat is not None: coord_norm = normalize_coord(coord_3d, box_flat.reshape(nframes, 3, 3)) else: @@ -1247,6 +1278,10 @@ def forward( ext_coord, ext_atype, mapping = extend_coord_with_ghosts( coord_norm, atype, box_flat, rcut ) + # A frame can never contribute more than nall neighbors; cap the + # requested neighbor count so sentinel capacities (e.g. DPA4C's + # effectively-unbounded sel) do not allocate absurd padding. + nsel = min(nsel, ext_atype.shape[1]) nlist = build_neighbor_list( ext_coord, ext_atype, @@ -1634,6 +1669,32 @@ def _forward_graph( if spin is not None: spin = spin[node_index] + # Long-range all-pairs graph (non-periodic). Built eagerly outside the + # compiled region and consumed as an extra edge set by models that have + # a long-range term (e.g. DPA4C-LR). Models without such a term still + # feed the compiled lower empty placeholders so the input signature is + # fixed across all graph models. + needs_lr = getattr(_model, "_needs_long_range_edges", False) + if needs_lr: + if ragged: + lr_graph = build_ragged_neighbor_graph( + method, coord_3d, atype, n_node, None, 1e6, None + ) + else: + coord_flat = coord_3d.reshape(n_padded, 3)[node_index] + lr_graph = build_ragged_neighbor_graph( + method, coord_flat, atype_flat, ng.n_node, None, 1e6, None + ) + lr_edge_index = lr_graph.edge_index + lr_edge_vec = lr_graph.edge_vec + lr_edge_mask = lr_graph.edge_mask + else: + _dev = coord_3d.device + _dt = coord_3d.dtype + lr_edge_index = torch.empty((2, 0), dtype=torch.int64, device=_dev) + lr_edge_vec = torch.empty((0, 3), dtype=_dt, device=_dev) + lr_edge_mask = torch.empty((0,), dtype=torch.bool, device=_dev) + # Lazy compile of the GRAPH lower (cached per structure key and mode). compiled_lower = self._compiled_lower_by_mode.get(self.training) if compiled_lower is None: @@ -1699,10 +1760,31 @@ def _forward_graph( aparam, charge_spin, spin, + lr_edge_index, + lr_edge_vec, + lr_edge_mask, *task_buf_vals, ) self._report_pending_compile() + # Optional model-specific post-processing that needs coordinates and is + # kept eager (e.g. DPA4C-LR's LES correction). The hook runs on the + # flat-node compiled output before the rectangular unravel below. + postprocess = getattr(self.original_model, "_postprocess_compiled_output", None) + if postprocess is not None: + result = postprocess( + coord=coord_3d, + edge_vec=edge_vec, + graph=ng, + result=result, + node_index=node_index, + n_padded=n_padded, + ragged=ragged, + fparam=fparam, + charge_spin=charge_spin, + create_graph=self.training, + ) + # The compiled graph lower emits PUBLIC keys on the FLAT node axis # (``atom_energy`` / ``force`` are (N, *); ``energy`` / ``virial`` are # (nf, *)). A ragged caller reads that axis directly. A rectangular one diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index 3b117732a6..8c2b37593b 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -271,6 +271,8 @@ def build_neighbor_graph_for_method( with_csr: bool = False, ) -> "NeighborGraph": """Build a carry-all graph with one concrete pt_expt backend.""" + if box is not None and torch.all(box == 0): + box = None if method == "dense": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 982c5d2a26..16823017a5 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3139,6 +3139,126 @@ def fitting_sezm_ener() -> list[Argument]: ] +@fitting_args_plugin.register( + "dpa4c_lr", + doc=supported_backends("pt_expt") + + "Fit a DPA4C latent-charge energy model (SR + non-periodic LES/SOG).", +) +def fitting_dpa4c_lr() -> list[Argument]: + doc_dim_out_lr = ( + "The output dimension of the long-range (latent charge) fitting branch. " + "Must be 1 for the LES kernel; SOG supports multiple channels." + ) + doc_neuron_lr = ( + "The number of neurons in each hidden layer of the latent-charge fitting net." + ) + doc_bias_atom_q = ( + "Initial bias values for the latent charge (q) of each atomic type. " + "If not set, a trainable per-type array is used." + ) + doc_les_alpha = ( + "Trainable width parameter for the non-periodic LES correction. " + "The LES kernel is erf(alpha * r) / r, which behaves as 1/r at long range." + ) + doc_amp = "Trainable amplitude of each Gaussian level in the SOG kernel." + doc_bandwidth = ( + "Trainable width of each Gaussian level in the SOG kernel. The stored " + "value is squared inside the kernel." + ) + doc_b = "Geometric ratio used to initialize the SOG bandwidths." + doc_sigma = "Base width used to initialize the SOG bandwidths." + doc_M = "Number of Gaussian levels used to initialize the SOG kernel." + doc_lr_kernel = ( + "Long-range kernel used for non-periodic systems. " + "'les' (default) applies only the LES erf(alpha*r)/r charge-charge kernel. " + "'sog' applies a trainable sum-of-Gaussians kernel over distinct atom " + "pairs (i != j) and supports multi-channel latent charges." + ) + doc_use_charge_constraint = ( + "Whether to enforce the sum of latent charges equals the target total " + "charge (from charge_spin[:, 0] or fparam[:, 0])." + ) + return [ + *fitting_sezm_ener(), + Argument( + "dim_out_lr", + int, + optional=True, + default=1, + doc=supported_backends("pt_expt") + doc_dim_out_lr, + ), + Argument( + "neuron_lr", + list[int], + optional=True, + default=[64, 64], + doc=supported_backends("pt_expt") + doc_neuron_lr, + ), + Argument( + "bias_atom_q", + list[float], + optional=True, + doc=supported_backends("pt_expt") + doc_bias_atom_q, + ), + Argument( + "amp", + [float, list[float], type(None)], + optional=True, + default=None, + doc=supported_backends("pt_expt") + doc_amp, + ), + Argument( + "bandwidth", + [float, list[float], type(None)], + optional=True, + default=None, + doc=supported_backends("pt_expt") + doc_bandwidth, + ), + Argument( + "les_alpha", + [float, type(None)], + optional=True, + default=1.0, + doc=supported_backends("pt_expt") + doc_les_alpha, + ), + Argument( + "b", + float, + optional=True, + default=1.6297670882677647, + doc=supported_backends("pt_expt") + doc_b, + ), + Argument( + "sigma", + float, + optional=True, + default=2.180230445405648, + doc=supported_backends("pt_expt") + doc_sigma, + ), + Argument( + "M", + int, + optional=True, + default=12, + doc=supported_backends("pt_expt") + doc_M, + ), + Argument( + "lr_kernel", + str, + optional=True, + default="les", + doc=supported_backends("pt_expt") + doc_lr_kernel, + ), + Argument( + "use_charge_constraint", + bool, + optional=True, + default=False, + doc=supported_backends("pt_expt") + doc_use_charge_constraint, + ), + ] + + @fitting_args_plugin.register( "dos", doc=supported_backends("tf", "pt", "pt_expt", "tf2") + doc_dos ) @@ -4058,6 +4178,64 @@ def sezm_model_args() -> Argument: return ca +@model_args_plugin.register("dpa4c_lr") +def dpa4c_lr_model_args() -> Argument: + doc_descrpt = ( + "Descriptor configuration for atomic environments. DPA4C-LR uses the " + "compact DPA4C descriptor." + ) + doc_fitting = ( + "Fitting network configuration. DPA4C-LR requires the `dpa4c_lr` " + "fitting, which augments the `dpa4_ener` GLU energy fitting with a " + "long-range latent-charge branch." + ) + doc_model = ( + "DPA4C-LR model scaffold: the compact DPA4C descriptor plus a " + "non-periodic LES or SOG term evaluated from the latent charges " + "produced by the `dpa4c_lr` fitting. Only the pt_expt backend and " + "non-periodic systems (no box) are supported." + ) + ca = Argument( + "dpa4c_lr", + dict, + [ + Argument( + "descriptor", + dict, + [], + [ + Variant( + "type", + [descrpt_args_plugin.get_argument("dpa4c")], + optional=True, + default_tag="dpa4c", + doc="The type of the descriptor.", + ) + ], + doc=supported_backends("pt_expt") + doc_descrpt, + ), + Argument( + "fitting_net", + dict, + [], + [ + Variant( + "type", + [fitting_args_plugin.get_argument("dpa4c_lr")], + optional=True, + default_tag="dpa4c_lr", + doc="The type of the fitting.", + ) + ], + doc=supported_backends("pt_expt") + doc_fitting, + ), + *_bridging_method_args(), + ], + doc=supported_backends("pt_expt") + doc_model, + ) + return ca + + @hybrid_model_args_plugin.register("pairwise_dprc") def pairwise_dprc() -> Argument: qm_model_args = model_args(exclude_hybrid=True) diff --git a/doc/model/dpa4c_lr.md b/doc/model/dpa4c_lr.md new file mode 100644 index 0000000000..e43e2db219 --- /dev/null +++ b/doc/model/dpa4c_lr.md @@ -0,0 +1,80 @@ +# DPA4C-LR: DPA4C with a non-periodic long-range term {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch Exportable {{ pytorch_icon }} (`dp --pt-expt`) + +DPA4C-LR augments the compact [DPA4C](./dpa4c.md) energy model with a +long-range correction evaluated from per-atom latent charges predicted by a +dedicated fitting branch. The total energy is + +```math +E = E_{\text{SR}} + E_{\text{LR}}(\{q_i\}, \{\boldsymbol r_i\}), +``` + +where $E_{\text{SR}}$ is the standard short-range DPA4C energy and +$E_{\text{LR}}$ is a non-periodic charge-charge interaction. Two kernels are +available for the long-range term, selected by +{ref}`lr_kernel `: + +- `les` (default): the LES kernel $\mathrm{erf}(\alpha r)/r$, which behaves as + $1/r$ at long range. The width $\alpha$ is trainable + ({ref}`les_alpha `) and a + single latent charge per atom is used. +- `sog`: a trainable sum-of-Gaussians kernel over distinct atom pairs + ($i \neq j$), initialized from the hyper-parameters + {ref}`b `, + {ref}`sigma `, and + {ref}`M `. Multiple latent-charge + channels are supported via + {ref}`dim_out_lr `. + +Forces and virials are obtained by differentiating the total energy, including +the response of the latent charges to the atomic positions, so the model +remains conservative. + +:::{warning} +Only **non-periodic systems** (no simulation box) are supported. Periodic +long-range summation (Ewald/SPME-style) is not implemented. Freezing to +`.pt2`/`.pte` is not yet implemented for this model type. +::: + +## Training input + +DPA4C-LR is selected as a model scaffold, `model.type: "dpa4c_lr"`, with the +matching fitting type `dpa4c_lr`: + +```json +"model": { + "type": "dpa4c_lr", + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0, + "channels": 64, + "lmax": 2, + "n_radial": 8 + }, + "fitting_net": { + "type": "dpa4c_lr", + "neuron": [128, 128], + "dim_out_lr": 1, + "neuron_lr": [64, 64], + "lr_kernel": "les", + "les_alpha": 1.0, + "use_charge_constraint": false + } +} +``` + +For the SOG kernel, set `lr_kernel` to `"sog"`; the kernel amplitudes and +bandwidths can be initialized explicitly with +{ref}`amp ` and +{ref}`bandwidth `, and +multiple latent-charge channels can be used by setting `dim_out_lr` larger +than 1. + +The latent-charge branch follows the standard energy fitting options (e.g. +`neuron`, `precision`, `seed`); see +{ref}`the fitting_net section ` for the +full argument list. See [training energy models](train-energy.md) for the +general workflow shared by all energy models. diff --git a/doc/model/index.rst b/doc/model/index.rst index c6db98d169..9a405f8401 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -14,6 +14,7 @@ Model dpa4 dpa4c unimol + dpa4c_lr train-hybrid sel train-energy diff --git a/source/tests/pt_expt/model/test_dpa4c_lr.py b/source/tests/pt_expt/model/test_dpa4c_lr.py new file mode 100644 index 0000000000..1f104c01f7 --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4c_lr.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the DPA4C-LR non-periodic latent-charge model.""" + +import copy + +import numpy as np +import pytest +import torch + +from deepmd.infer import ( + DeepPot, +) +from deepmd.pt_expt.fitting.dpa4c_lr import ( + DPA4CLRFitting, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.model.dpa4c_lr_model import ( + E2_PER_ANGSTROM_TO_EV, + _SOGKernel, + _les_kernel_from_squared_distance, +) +from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ..compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + + +def _config(lr_kernel: str = "les") -> dict: + """Minimal float64 DPA4C-LR config for unit tests.""" + config = { + "type": "dpa4c_lr", + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 17, + }, + "fitting_net": { + "type": "dpa4c_lr", + "neuron": [16, 16], + "precision": "float64", + "seed": 19, + "dim_out_lr": 1, + "neuron_lr": [8, 8], + "use_charge_constraint": True, + "lr_kernel": "les", + "les_alpha": 1.0, + }, + } + if lr_kernel == "sog": + config["fitting_net"].update( + { + "dim_out_lr": 3, + "use_charge_constraint": False, + "lr_kernel": "sog", + "amp": [0.7, -0.2, 0.1], + "bandwidth": [0.8, 1.7, 3.2], + } + ) + return config + + +def _build_model(train: bool = False, lr_kernel: str = "les") -> torch.nn.Module: + """Build a small DPA4C-LR model on the default test device.""" + model = get_model(_config(lr_kernel)).to(env.DEVICE) + if train: + model.train() + else: + model.eval() + return model + + +def _make_batch( + nlocs: tuple[int, ...], + seed: int = 0, + box: np.ndarray | None = None, +) -> tuple[np.ndarray, np.ndarray]: + """Build a padded rectangular batch.""" + rng = np.random.default_rng(seed) + pad_nloc = max(nlocs) + padded_coord = np.zeros((len(nlocs), pad_nloc, 3)) + padded_atype = np.full((len(nlocs), pad_nloc), -1, dtype=np.int64) + for index, nloc in enumerate(nlocs): + padded_coord[index, :nloc] = rng.uniform(0.0, 6.0, (nloc, 3)) + padded_atype[index, :nloc] = rng.integers(0, 2, nloc) + return padded_coord, padded_atype + + +def test_sog_kernel_matches_closed_form() -> None: + """The dense SOG kernel contains only distinct-atom pairs.""" + positions = torch.tensor( + [[[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]]], + dtype=torch.float64, + device=env.DEVICE, + ) + charges = torch.tensor( + [[[1.0, 2.0], [-0.5, 1.5]]], dtype=torch.float64, device=env.DEVICE + ) + amp = torch.tensor([0.7, -0.2], dtype=torch.float64, device=env.DEVICE) + bandwidth = torch.tensor([0.8, 1.7], dtype=torch.float64, device=env.DEVICE) + + kernel_at_r = torch.sum(amp * torch.exp(-0.5 * 4.0 / bandwidth.square())) + pair = torch.dot(charges[0, 0], charges[0, 1]) * kernel_at_r + energy = _SOGKernel()(positions, charges, amp, bandwidth) + torch.testing.assert_close(energy[0], pair * E2_PER_ANGSTROM_TO_EV) + + +def test_les_masked_zero_edge_has_finite_gradient() -> None: + """A zero-length guard edge must not inject NaN into the LES force.""" + edge_vec = torch.zeros( + (1, 3), dtype=torch.float64, device=env.DEVICE, requires_grad=True + ) + alpha = torch.tensor(1.0, dtype=torch.float64, device=env.DEVICE) + edge_mask = torch.zeros(1, dtype=torch.float64, device=env.DEVICE) + r_sq = torch.sum(edge_vec * edge_vec, dim=-1) + energy = (_les_kernel_from_squared_distance(r_sq, alpha) * edge_mask).sum() + (gradient,) = torch.autograd.grad(energy, edge_vec) + assert torch.isfinite(gradient).all() + assert torch.count_nonzero(gradient) == 0 + + +def test_energy_force_are_finite_and_phantom_atoms_are_force_free() -> None: + """A padded batch produces finite outputs and zero force on phantom atoms.""" + torch.manual_seed(1234) + model = _build_model() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((4, 7, 3)) + result = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + ) + for key in ("energy", "force", "atom_energy"): + assert torch.isfinite(result[key]).all(), key + phantom = torch.tensor(atype, device=env.DEVICE) < 0 + assert bool(phantom.any()), "fixture must exercise padding" + assert bool(torch.all(result["force"][phantom] == 0.0)) + + +def test_frames_do_not_interact_across_batch() -> None: + """The dense LR kernel must not couple atoms in different frames.""" + torch.manual_seed(1234) + model = _build_model() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + nlocs = (4, 7, 3) + coord, atype = _make_batch(nlocs) + + def run( + batch_coord: np.ndarray, batch_atype: np.ndarray + ) -> dict[str, torch.Tensor]: + return model( + torch.tensor(batch_coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(batch_atype, dtype=torch.long, device=env.DEVICE), + ) + + batched = run(coord, atype) + for index, nloc in enumerate(nlocs): + alone = run( + coord[index : index + 1, :nloc], + atype[index : index + 1, :nloc], + ) + torch.testing.assert_close( + batched["energy"].reshape(-1)[index], + alone["energy"].reshape(-1)[0], + atol=1e-12, + rtol=1e-12, + ) + torch.testing.assert_close( + batched["force"][index, :nloc], + alone["force"][0], + atol=1e-12, + rtol=1e-12, + ) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_force_matches_finite_difference(lr_kernel: str) -> None: + """The public model force agrees with a central finite-difference stencil.""" + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + rng = np.random.default_rng(1) + coord = rng.uniform(0.0, 6.0, (1, 5, 3)) + atype = rng.integers(0, 2, (1, 5)).astype(np.int64) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + + def energy(x: torch.Tensor) -> float: + with torch.enable_grad(): + x = x.detach().requires_grad_(True) + ret = model(x, atype_t) + return ret["energy"].item() + + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + force = model(coord_t, atype_t)["force"].detach().cpu().numpy().reshape(-1, 3) + + eps = 1e-5 + force_fd = np.zeros_like(force) + for i in range(coord.shape[1]): + for d in range(3): + cp = coord.copy() + cp[0, i, d] += eps + cm = coord.copy() + cm[0, i, d] -= eps + Ep = energy(torch.tensor(cp, dtype=torch.float64, device=env.DEVICE)) + Em = energy(torch.tensor(cm, dtype=torch.float64, device=env.DEVICE)) + force_fd[i, d] = -(Ep - Em) / (2 * eps) + + np.testing.assert_allclose(force, force_fd, atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_atom_virial_sums_to_total_virial(lr_kernel: str) -> None: + """Per-atom virials include the explicit LR geometry contribution.""" + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((4, 7, 3), seed=31) + result = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + do_atomic_virial=True, + ) + torch.testing.assert_close( + result["atom_virial"].sum(dim=1), + result["virial"], + rtol=1e-12, + atol=1e-12, + ) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_ragged_matches_padded(lr_kernel: str) -> None: + """Ragged and padded layouts must describe the same batch.""" + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + nlocs = (4, 7, 3) + coord, atype = _make_batch(nlocs) + + padded = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + ) + flat_coord = torch.tensor( + np.concatenate([coord[i, :nloc] for i, nloc in enumerate(nlocs)]), + dtype=torch.float64, + device=env.DEVICE, + ) + flat_atype = torch.tensor( + np.concatenate([atype[i, :nloc] for i, nloc in enumerate(nlocs)]), + dtype=torch.long, + device=env.DEVICE, + ) + n_node = torch.tensor(nlocs, dtype=torch.long, device=env.DEVICE) + ragged = model.forward_ragged(flat_coord, flat_atype, n_node) + + torch.testing.assert_close(padded["energy"], ragged["energy"]) + offset = 0 + for index, nloc in enumerate(nlocs): + torch.testing.assert_close( + padded["force"][index, :nloc], + ragged["force"][offset : offset + nloc], + atol=1e-12, + rtol=1e-12, + ) + offset += nloc + assert offset == ragged["force"].shape[0] + + +def test_charge_constraint_zero_target_preserves_total_charge() -> None: + """With no charge_spin or fparam, constraint drives total latent charge to 0.""" + torch.manual_seed(1234) + model = _build_model() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + fitting = model.get_fitting_net() + nf, nloc = 2, 5 + q = torch.randn(nf * nloc, 1, dtype=torch.float64, device=env.DEVICE) + n_node = torch.tensor([nloc] * nf, dtype=torch.long, device=env.DEVICE) + from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + ) + + frame_id = frame_id_from_n_node(n_node, n_total=q.shape[0]) + q_out = model._constrain_charges_flat( + q, frame_id, n_node, fparam=None, charge_spin=None + ) + for f in range(nf): + assert torch.isclose( + q_out[frame_id == f, 0].sum(), + torch.zeros((), dtype=q_out.dtype, device=q_out.device), + ) + + +def test_periodic_box_is_rejected() -> None: + """DPA4C-LR must reject non-zero simulation cells.""" + model = _build_model() + coord = torch.zeros((1, 3, 3), dtype=torch.float64, device=env.DEVICE) + atype = torch.zeros((1, 3), dtype=torch.long, device=env.DEVICE) + box = torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 9) + with pytest.raises(NotImplementedError): + model(coord, atype, box=box) + + +def test_periodic_training_data_is_rejected_before_stats() -> None: + """Periodic sampled frames fail fast with a clear error before stats.""" + model = _build_model() + + def sampled() -> list[dict]: + return [ + { + "coord": np.zeros((1, 3, 3)), + "atype": np.zeros((1, 3), dtype=np.int64), + "box": np.eye(3).reshape(1, 9), + } + ] + + with pytest.raises(NotImplementedError, match="non-periodic"): + model.compute_or_load_stat(sampled) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_pt_checkpoint_deep_eval_matches_direct_forward( + tmp_path, lr_kernel: str +) -> None: + """Checkpoint inference builds and passes the required long-range graph.""" + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + checkpoint = tmp_path / f"dpa4c_lr_{lr_kernel}.pt" + wrapper = ModelWrapper(model, model_params=copy.deepcopy(_config(lr_kernel))) + torch.save({"model": wrapper.state_dict()}, checkpoint) + + coord, atype = _make_batch((4, 4), seed=29) + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + expected = model(coord_t, atype_t, do_atomic_virial=True) + + dp = DeepPot(str(checkpoint), auto_batch_size=False) + output_names = ("energy", "force", "virial", "atom_energy", "atom_virial") + actual = dict( + zip( + output_names, + dp.eval(coord, None, atype, atomic=True, mixed_type=True), + strict=True, + ) + ) + for name in output_names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_graph_training_gradient_reaches_lr_parameters(lr_kernel: str) -> None: + """Backprop through the graph lower reaches descriptor and LR fitting params.""" + torch.manual_seed(1234) + model = _build_model(train=True, lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((4, 5)) + result = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + ) + loss = result["energy"].square().mean() + result["force"].square().mean() + loss.backward() + + fitting = model.get_fitting_net() + for name, parameter in fitting.named_parameters(): + assert parameter.grad is not None, f"fitting parameter {name} has no grad" + assert torch.isfinite(parameter.grad).all(), f"{name} grad non-finite" + descriptor = model.get_descriptor() + grads = {name: parameter.grad for name, parameter in descriptor.named_parameters()} + assert grads, "descriptor has parameters" + for name, gradient in grads.items(): + assert gradient is not None, name + assert torch.isfinite(gradient).all(), name + + +def test_sog_parameters_serialize_and_remain_trainable() -> None: + """SOG amplitudes and widths survive fitting serialization as Parameters.""" + fitting = _build_model(lr_kernel="sog").get_fitting_net() + named_parameters = dict(fitting.named_parameters()) + assert "amp" in named_parameters + assert "bandwidth" in named_parameters + assert "les_alpha" not in named_parameters + + restored = DPA4CLRFitting.deserialize(fitting.serialize()).to(env.DEVICE) + assert restored.lr_kernel == "sog" + assert restored.dim_out_lr == 3 + restored_parameters = dict(restored.named_parameters()) + assert "amp" in restored_parameters + assert "bandwidth" in restored_parameters + torch.testing.assert_close(restored.amp, fitting.amp) + torch.testing.assert_close(restored.bandwidth, fitting.bandwidth) + + +def test_les_checkpoint_does_not_require_sog_parameters() -> None: + """LES state loading accepts checkpoints from before and during SOG rollout.""" + fitting = _build_model(lr_kernel="les").get_fitting_net() + state_dict = fitting.state_dict() + assert "amp" not in state_dict + assert "bandwidth" not in state_dict + + # Simulate the short-lived format that stored unused SOG buffers in LES. + transitional_state = state_dict.copy() + transitional_state["amp"] = torch.ones(12, device=env.DEVICE) + transitional_state["bandwidth"] = torch.ones(12, device=env.DEVICE) + + restored = _build_model(lr_kernel="les").get_fitting_net() + restored.load_state_dict(state_dict, strict=True) + restored.load_state_dict(transitional_state, strict=True) + + +def test_sog_graph_matches_dense_legacy_path() -> None: + """SOG graph and rectangular dense paths produce the same observables.""" + torch.manual_seed(1234) + model = _build_model(lr_kernel="sog") + coord, atype = _make_batch((5, 5), seed=37) + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + + graph = model(coord_t, atype_t) + # The dense (nlist) fallback must agree with the graph lower. + model.neighbor_graph_method = "legacy" + dense = model(coord_t, atype_t) + for name in ("energy", "force", "virial"): + torch.testing.assert_close(graph[name], dense[name], rtol=1e-10, atol=1e-10) + + +@REQUIRES_SUPPORTED_COMPILE +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_compiled_lower_matches_eager(lr_kernel: str) -> None: + """The compiled graph lower must reproduce the eager energy and force.""" + from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + ) + + torch.manual_seed(0) + model = _build_model(lr_kernel=lr_kernel) + compiled = _CompiledModel(model, _get_model_structure_key(model)) + compiled.eval() + + nlocs = (4, 7, 3) + coord, atype = _make_batch(nlocs) + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + + expected = model(coord_t, atype_t) + got = compiled(coord_t, atype_t, None) + + torch.testing.assert_close(got["energy"], expected["energy"]) + torch.testing.assert_close(got["force"], expected["force"]) From 017d7690ef4abc3ecf71edef6d6f51e8f0211bcc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 07:22:46 +0000 Subject: [PATCH 02/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/fitting/dpa4c_lr.py | 4 +--- deepmd/pt_expt/model/dpa4c_lr_model.py | 6 +++--- source/tests/pt_expt/model/test_dpa4c_lr.py | 8 ++++---- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/deepmd/pt_expt/fitting/dpa4c_lr.py b/deepmd/pt_expt/fitting/dpa4c_lr.py index 93a988c011..5123f53387 100644 --- a/deepmd/pt_expt/fitting/dpa4c_lr.py +++ b/deepmd/pt_expt/fitting/dpa4c_lr.py @@ -7,9 +7,7 @@ import torch -from deepmd.dpmodel.fitting.dpa4c_lr import ( - DPA4CLRFitting as DPA4CLRFittingDP, -) +from deepmd.dpmodel.fitting.dpa4c_lr import DPA4CLRFitting as DPA4CLRFittingDP from deepmd.pt_expt.common import ( register_dpmodel_mapping, torch_module, diff --git a/deepmd/pt_expt/model/dpa4c_lr_model.py b/deepmd/pt_expt/model/dpa4c_lr_model.py index faaadf3406..688c4452bc 100644 --- a/deepmd/pt_expt/model/dpa4c_lr_model.py +++ b/deepmd/pt_expt/model/dpa4c_lr_model.py @@ -46,12 +46,12 @@ frame_id_from_n_node, segment_sum, ) -from deepmd.pt_expt.kernels.utils import ( - cuda_infer_level, -) from deepmd.pt_expt.fitting.dpa4c_lr import ( DPA4CLRFitting, ) +from deepmd.pt_expt.kernels.utils import ( + cuda_infer_level, +) from deepmd.pt_expt.model.ener_model import ( EnergyModel, ) diff --git a/source/tests/pt_expt/model/test_dpa4c_lr.py b/source/tests/pt_expt/model/test_dpa4c_lr.py index 1f104c01f7..7d1bf9d9d6 100644 --- a/source/tests/pt_expt/model/test_dpa4c_lr.py +++ b/source/tests/pt_expt/model/test_dpa4c_lr.py @@ -13,13 +13,13 @@ from deepmd.pt_expt.fitting.dpa4c_lr import ( DPA4CLRFitting, ) -from deepmd.pt_expt.model.get_model import ( - get_model, -) from deepmd.pt_expt.model.dpa4c_lr_model import ( E2_PER_ANGSTROM_TO_EV, - _SOGKernel, _les_kernel_from_squared_distance, + _SOGKernel, +) +from deepmd.pt_expt.model.get_model import ( + get_model, ) from deepmd.pt_expt.train.wrapper import ( ModelWrapper, From b30b119847ef5116854eb613099f8baa862e9ec5 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 17 Sep 2026 08:19:22 +0000 Subject: [PATCH 03/12] fix(dpmodel): skip neighbor-count cap for symbolic shapes The nsel/nnei cap introduced with DPA4C-LR used Python's min() on the frame atom count, which raises InconclusiveDimensionOperation under jax2tf export where dimensions are symbolic. Only apply the cap when nall is a concrete int, and make the all-zero-box check tracer-safe. --- deepmd/dpmodel/model/make_model.py | 15 +++++++++++++-- deepmd/dpmodel/utils/default_neighbor_list.py | 7 ++++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 143304b329..d31fd7b74a 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -134,7 +134,13 @@ def model_call_from_call_lower( # Non-periodic datasets may carry an all-zero cell; treat it as no box. if bb is not None: xp_bb = array_api_compat.array_namespace(bb) - if bool(xp_bb.all(bb == 0)): + try: + zero_box = bool(xp_bb.all(bb == 0)) + except Exception: + # Traced/symbolic arrays (e.g. JAX export) cannot be evaluated to + # a Python bool; keep the box as-is in that case. + zero_box = False + if zero_box: bb = None builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() # Model-level pair exclusion is a nlist-BUILD transform (decision #18/A4): @@ -1079,8 +1085,13 @@ def format_nlist( # A frame can never contribute more than nall neighbors; cap the # requested neighbor count so sentinel capacities (e.g. DPA4C's # effectively-unbounded sel) do not allocate absurd padding. + # Skip the cap when nall is symbolic (e.g. jax2tf export), where + # the comparison would be inconclusive; sentinel capacities only + # occur on backends with concrete shapes. nall = extended_atype.shape[1] - nnei = min(sum(self.get_sel()), nall) + nnei = sum(self.get_sel()) + if isinstance(nall, int) and nnei > nall: + nnei = nall ret = self._format_nlist( extended_coord, nlist, diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index 6bd1703ed4..c7da1a6d80 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -732,8 +732,13 @@ def build( # A frame can never contribute more than nall neighbors; cap the # requested neighbor count so sentinel capacities (e.g. DPA4C's # effectively-unbounded sel) do not allocate absurd padding. + # Skip the cap when nall is symbolic (e.g. jax2tf export), where + # the comparison would be inconclusive; sentinel capacities only + # occur on backends with concrete shapes. nall = extended_atype.shape[1] - nsel = min(sum(sel), nall) + nsel = sum(sel) + if isinstance(nall, int) and nsel > nall: + nsel = nall # Types are distinguished in the lower interface, so keep them merged # here. The dense path remains faster for small systems; the spatial # path avoids the nloc*nall distance matrix for larger supported arrays. From 5a755391d773f1998c23b97d6f1541066605d805 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Fri, 18 Sep 2026 02:27:05 +0000 Subject: [PATCH 04/12] fix(dpmodel): only cap neighbor count for mixed-types descriptors Type-distinguished descriptors (e.g. se_e2_a) require exactly sum(sel) nlist columns: their davg/dstd statistics have sum(sel) x 4 shape and nlist_distinguish_types depends on the column count. The unconditional nnei = min(nnei, nall) cap broke them whenever nall < sum(sel), raising "cannot reshape array of size ..." in env_mat._apply_stat (source/tests/consistent/io/test_io.py::TestDeepPot::test_deep_eval). Carry-all (mixed_types) descriptors such as DPA4C keep the cap, which still prevents absurd padding for their effectively-unbounded sentinel sel. --- deepmd/dpmodel/model/make_model.py | 18 +++++++++++------- deepmd/dpmodel/utils/default_neighbor_list.py | 11 +++++++---- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index d31fd7b74a..0afc8b21a7 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1082,15 +1082,19 @@ def format_nlist( """ n_nf, n_nloc, n_nnei = nlist.shape mixed_types = self.mixed_types() - # A frame can never contribute more than nall neighbors; cap the - # requested neighbor count so sentinel capacities (e.g. DPA4C's - # effectively-unbounded sel) do not allocate absurd padding. - # Skip the cap when nall is symbolic (e.g. jax2tf export), where - # the comparison would be inconclusive; sentinel capacities only - # occur on backends with concrete shapes. + # A frame can never contribute more than nall neighbors, so for + # carry-all (mixed-types) descriptors cap the requested neighbor + # count: sentinel capacities (e.g. DPA4C's effectively-unbounded + # sel) would otherwise allocate absurd padding. Type- + # distinguished descriptors instead REQUIRE exactly sum(sel) + # columns (their statistics and ``nlist_distinguish_types`` depend + # on it), so they must never be capped. Skip the cap when nall is + # symbolic (e.g. jax2tf export), where the comparison would be + # inconclusive; sentinel capacities only occur on backends with + # concrete shapes. nall = extended_atype.shape[1] nnei = sum(self.get_sel()) - if isinstance(nall, int) and nnei > nall: + if mixed_types and isinstance(nall, int) and nnei > nall: nnei = nall ret = self._format_nlist( extended_coord, diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index c7da1a6d80..79ecfc580d 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -731,10 +731,13 @@ def build( ) # A frame can never contribute more than nall neighbors; cap the # requested neighbor count so sentinel capacities (e.g. DPA4C's - # effectively-unbounded sel) do not allocate absurd padding. - # Skip the cap when nall is symbolic (e.g. jax2tf export), where - # the comparison would be inconclusive; sentinel capacities only - # occur on backends with concrete shapes. + # effectively-unbounded sel) do not allocate absurd padding. This is + # always safe here: real neighbors number at most nall, and + # ``format_nlist`` re-pads the neighbor axis to ``sum(sel)`` for + # type-distinguished descriptors downstream. Skip the cap when nall + # is symbolic (e.g. jax2tf export), where the comparison would be + # inconclusive; sentinel capacities only occur on backends with + # concrete shapes. nall = extended_atype.shape[1] nsel = sum(sel) if isinstance(nall, int) and nsel > nall: From 1430ce834b1678ae19cda2ee4205ba5532b9348e Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Fri, 18 Sep 2026 06:50:58 +0000 Subject: [PATCH 05/12] fix(dpmodel): restrict neighbor-count cap to sentinel capacities The nnei = min(nnei, nall) cap introduced with DPA4C-LR broke any model whose sum(sel) exceeds the frame atom count (e.g. spin models, where virtual atoms double nloc): the neighbor list was capped to nall columns while descriptor statistics and serialized lowers still require exactly sum(sel) columns, raising "shape ... is invalid for input of size ..." in env-mat application (source/tests/pt_expt/infer/test_deep_eval_spin.py .pt2/.pte inference). Only cap sentinel (effectively-unbounded) capacities such as DPA4C's; normal capacities must keep exactly sum(sel) columns. --- deepmd/dpmodel/model/make_model.py | 24 ++++++++++--------- deepmd/dpmodel/utils/default_neighbor_list.py | 22 +++++++++-------- deepmd/dpmodel/utils/nlist.py | 11 +++++++++ 3 files changed, 36 insertions(+), 21 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 0afc8b21a7..c30b6ec9e3 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -44,6 +44,9 @@ format_nlist, nlist_distinguish_types, ) +from deepmd.dpmodel.utils.nlist import ( + UNBOUNDED_NSEL, +) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, build_neighbor_graph, @@ -1082,19 +1085,18 @@ def format_nlist( """ n_nf, n_nloc, n_nnei = nlist.shape mixed_types = self.mixed_types() - # A frame can never contribute more than nall neighbors, so for - # carry-all (mixed-types) descriptors cap the requested neighbor - # count: sentinel capacities (e.g. DPA4C's effectively-unbounded - # sel) would otherwise allocate absurd padding. Type- - # distinguished descriptors instead REQUIRE exactly sum(sel) - # columns (their statistics and ``nlist_distinguish_types`` depend - # on it), so they must never be capped. Skip the cap when nall is - # symbolic (e.g. jax2tf export), where the comparison would be - # inconclusive; sentinel capacities only occur on backends with - # concrete shapes. + # Sentinel capacities (e.g. DPA4C's effectively-unbounded sel) + # would allocate absurd padding; a frame can never contribute + # more than nall neighbors, so cap them. Normal capacities must + # NOT be capped: type-distinguished descriptors REQUIRE exactly + # sum(sel) columns (their statistics and + # ``nlist_distinguish_types`` depend on it), so they must never + # be capped. Skip the cap when nall is symbolic (e.g. jax2tf + # export), where the comparison would be inconclusive; sentinel + # capacities only occur on backends with concrete shapes. nall = extended_atype.shape[1] nnei = sum(self.get_sel()) - if mixed_types and isinstance(nall, int) and nnei > nall: + if nnei >= UNBOUNDED_NSEL and isinstance(nall, int) and nnei > nall: nnei = nall ret = self._format_nlist( extended_coord, diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index 79ecfc580d..a7381e706b 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -22,6 +22,7 @@ NeighborList, ) from .nlist import ( + UNBOUNDED_NSEL, apply_pair_exclusion_nlist, build_neighbor_list, extend_coord_with_ghosts, @@ -729,18 +730,19 @@ def build( extended_coord, extended_atype, mapping = extend_coord_with_ghosts( coord_normalized, atype, box, rcut ) - # A frame can never contribute more than nall neighbors; cap the - # requested neighbor count so sentinel capacities (e.g. DPA4C's - # effectively-unbounded sel) do not allocate absurd padding. This is - # always safe here: real neighbors number at most nall, and - # ``format_nlist`` re-pads the neighbor axis to ``sum(sel)`` for - # type-distinguished descriptors downstream. Skip the cap when nall - # is symbolic (e.g. jax2tf export), where the comparison would be - # inconclusive; sentinel capacities only occur on backends with - # concrete shapes. + # Sentinel capacities (e.g. DPA4C's effectively-unbounded sel) would + # allocate absurd padding; a frame can never contribute more than nall + # neighbors, so cap them. Normal capacities must NOT be capped: + # type-distinguished descriptors and serialized lowers (e.g. the jax + # StableHLO export, whose polymorphic shape bakes the neighbor axis + # as the sum(sel) constant) require exactly sum(sel) columns, and + # ``format_nlist`` only re-pads on the eager model path. Skip the cap + # when nall is symbolic (e.g. jax2tf export), where the comparison + # would be inconclusive; sentinel capacities only occur on backends + # with concrete shapes. nall = extended_atype.shape[1] nsel = sum(sel) - if isinstance(nall, int) and nsel > nall: + if nsel >= UNBOUNDED_NSEL and isinstance(nall, int) and nsel > nall: nsel = nall # Types are distinguished in the lower interface, so keep them merged # here. The dense path remains faster for small systems; the spatial diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index cff741dec2..1b9d5bf40a 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -24,6 +24,17 @@ ) +#: Neighbor capacities at or beyond this value are sentinels for +#: "effectively unbounded" descriptors (DPA4C reports sel 999999 per type) +#: that carry every in-cutoff neighbor. Dense neighbor-list builders may +#: cap such capacities at the frame atom count instead of allocating the +#: full sentinel width. Realistic sel values are orders of magnitude +#: smaller, so ordinary descriptors keep their exact ``sum(sel)`` column +#: count, which type-distinguished statistics and serialized lowers +#: require. +UNBOUNDED_NSEL: int = 100_000 + + def _is_ndtensorflow_namespace(xp: Any) -> bool: return getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow" From 29713bd4542f6481fe4079f8a9b6c0e057edae38 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 10:14:19 +0000 Subject: [PATCH 06/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/make_model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index c30b6ec9e3..8d4f459f2f 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -44,9 +44,6 @@ format_nlist, nlist_distinguish_types, ) -from deepmd.dpmodel.utils.nlist import ( - UNBOUNDED_NSEL, -) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, build_neighbor_graph, @@ -54,6 +51,9 @@ compact_nodes, expand_node_values, ) +from deepmd.dpmodel.utils.nlist import ( + UNBOUNDED_NSEL, +) from deepmd.utils.path import ( DPPath, ) From 1080806d16a94f9a42577c96f02d0089659a938a Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Tue, 22 Sep 2026 13:40:49 +0000 Subject: [PATCH 07/12] test(pt-expt): pass empty LR edge placeholders to compiled graph lower in DPA2 tests The merge_dpa4C branch made lr_edge_index/lr_edge_vec/lr_edge_mask unconditional positional inputs of the _trace_and_compile_graph output. Update the three compiled-training tests to feed the empty placeholders (same convention as _CompiledModel at run time), fixing: TypeError: forward() missing 3 required positional arguments: 'arg14_1', 'arg15_1', and 'arg16_1' --- .../pt_expt/model/test_dpa2_graph_lower.py | 81 ++++++++++++++++++- 1 file changed, 78 insertions(+), 3 deletions(-) diff --git a/source/tests/pt_expt/model/test_dpa2_graph_lower.py b/source/tests/pt_expt/model/test_dpa2_graph_lower.py index d8dc85a43e..5d0bd8e957 100644 --- a/source/tests/pt_expt/model/test_dpa2_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py @@ -146,6 +146,24 @@ def _make_dpa2_descriptor( ) +def _empty_lr_edges( + *, dtype: torch.dtype, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Empty long-range edge placeholders for models without an LR term. + + ``_trace_and_compile_graph`` makes ``lr_edge_index`` / ``lr_edge_vec`` / + ``lr_edge_mask`` unconditional positional inputs of the compiled graph + lower; a DPA2 model has no long-range term, so the runtime caller + (``_CompiledModel`` in ``deepmd/pt_expt/train/training.py``) feeds these + empty tensors. The tests must do the same. + """ + return ( + torch.empty((2, 0), dtype=torch.int64, device=device), + torch.empty((0, 3), dtype=dtype, device=device), + torch.empty((0,), dtype=torch.bool, device=device), + ) + + class TestDpa2GraphLower: def setup_method(self) -> None: self.device = env.DEVICE @@ -689,8 +707,27 @@ def test_compiled_training_graph_smoke(self) -> None: ) atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + lr_ei, lr_ev, lr_em = _empty_lr_edges( + dtype=torch.float64, device=torch.device("cpu") + ) compiled_out = compiled_lower( - atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + fp, + ap, + cs, + None, + lr_ei, + lr_ev, + lr_em, ) eager = model.forward_common_lower_graph( atype, @@ -977,8 +1014,27 @@ def test_compiled_training_graph_small_sel(self) -> None: want_charge_spin=False, ) atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + lr_ei, lr_ev, lr_em = _empty_lr_edges( + dtype=torch.float64, device=torch.device("cpu") + ) compiled_out = compiled_lower( - atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + fp, + ap, + cs, + None, + lr_ei, + lr_ev, + lr_em, ) eager = model.forward_common_lower_graph( atype, @@ -1083,8 +1139,27 @@ def test_graph_lower_fparam_symbolic_trace_and_compile(self) -> None: # (b) compiled-training path (fparam threaded through the compile) compiled_lower, _ = _trace_and_compile_graph(model, fp, None, None, None) + lr_ei, lr_ev, lr_em = _empty_lr_edges( + dtype=torch.float64, device=torch.device("cpu") + ) compiled_out = compiled_lower( - atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + fp, + ap, + cs, + None, + lr_ei, + lr_ev, + lr_em, ) ctol = {"rtol": 1e-10, "atol": 1e-10} torch.testing.assert_close(compiled_out["energy"], ref["energy_redu"], **ctol) From 063d2cc1edd71bcaf0185fbeb8e120de8165a01d Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 15 Oct 2026 02:43:44 +0000 Subject: [PATCH 08/12] fix(pt-expt): address DPA4C-LR review comments - Freeze the promoted LES/SOG parameters when fitting_net.trainable is false: GeneralFitting normalizes the flag to a per-layer list, so a plain bool() on it was always True. Normalize list-valued trainable with all() before setting requires_grad, and add regression tests that compare parameter values across a real Adam step for LES and SOG. - Build every long-range all-pairs graph with the dense builder, decoupled from the short-range search policy: the hardcoded 1e6 sentinel cutoff explodes the NV capacity estimate (~1e18 neighbors, ~12.6 EB allocation) on CUDA with nvalchemiops installed. - Normalize an all-zero box to None before building eval graphs so the native periodic search never sees a singular lattice. - Default fitting_net.type to dpa4c_lr in get_dpa4c_lr_model so a config omitting it no longer builds an ener fitting. - Match DPA4CLRFitting.call / DPA4CLREnergyModel.call_common to their base-class signatures (vacuum_descriptor passthrough; explicit call_common parameters). - Drop bridging_method from the dpa4c_lr schema: expand_bridging_method only accepts standard/dpa4/sezm, so documenting it here passed validation and then failed at construction. - Document that SOG is supported and that box=None or an all-zero box is accepted, and note that exported graphs bake the periodic topology at trace time. - Remove unused local variables flagged by CodeQL. --- AGENTS.md | 2 +- deepmd/dpmodel/fitting/dpa4c_lr.py | 10 ++- deepmd/dpmodel/model/make_model.py | 9 ++- deepmd/pt_expt/fitting/dpa4c_lr.py | 8 ++- deepmd/pt_expt/infer/deep_eval.py | 8 ++- deepmd/pt_expt/model/dpa4c_lr_model.py | 33 ++++++++-- deepmd/pt_expt/model/get_model.py | 7 ++- deepmd/pt_expt/train/training.py | 7 ++- deepmd/utils/argcheck.py | 4 +- doc/model/dpa4c_lr.md | 8 ++- source/tests/pt_expt/model/test_dpa4c_lr.py | 69 ++++++++++++++++++++- 11 files changed, 144 insertions(+), 21 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5ca920a19b..ebb487db0c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -165,7 +165,7 @@ source/ # C++ source code and tests - **Configuration**: Use `input_torch.json` format typically - **Training**: `dp --pt train input_torch.json` - **Requirements**: `torch` package -- **DPA4C-LR long-range variant**: model type `dpa4c_lr` with fitting type `dpa4c_lr` adds a non-periodic LES long-range term on top of the compact DPA4C descriptor (`deepmd/pt_expt/model/dpa4c_lr_model.py`, `deepmd/pt_expt/fitting/dpa4c_lr.py`). Only `box=None` systems are supported and only the `pt_expt` backend. The long-range all-pairs graph is built outside the compiled region and passed into the graph lower alongside the short-range graph, so `torch.compile` sees no data-dependent neighbor construction and memory scales as `Σ_f n_f^2` rather than the full-batch `N^2`. `.pt2`/`.pte` freeze is not yet implemented. Tests: `pytest source/tests/pt_expt/model/test_dpa4c_lr.py`. +- **DPA4C-LR long-range variant**: model type `dpa4c_lr` with fitting type `dpa4c_lr` adds a non-periodic long-range term on top of the compact DPA4C descriptor (`deepmd/pt_expt/model/dpa4c_lr_model.py`, `deepmd/pt_expt/fitting/dpa4c_lr.py`). Two kernels are available: `les` (`erf(alpha*r)/r`, the default) and a trainable sum-of-Gaussians `sog` with multi-channel latent charges. Only non-periodic systems are supported — `box=None` or an all-zero box is accepted, while nonzero periodic cells are rejected — and only the `pt_expt` backend. The long-range all-pairs graph is built outside the compiled region and passed into the graph lower alongside the short-range graph, so `torch.compile` sees no data-dependent neighbor construction and memory scales as `Σ_f n_f^2` rather than the full-batch `N^2`. `.pt2`/`.pte` freeze is not yet implemented. Tests: `pytest source/tests/pt_expt/model/test_dpa4c_lr.py`. ### JAX Backend diff --git a/deepmd/dpmodel/fitting/dpa4c_lr.py b/deepmd/dpmodel/fitting/dpa4c_lr.py index 41c58d3ef3..59f97a5fdb 100644 --- a/deepmd/dpmodel/fitting/dpa4c_lr.py +++ b/deepmd/dpmodel/fitting/dpa4c_lr.py @@ -242,10 +242,18 @@ def call( h2: Any | None = None, fparam: Any | None = None, aparam: Any | None = None, + vacuum_descriptor: Any | None = None, ) -> dict[str, Any]: """Return energy and latent charges.""" energy_ret = super().call( - descriptor, atype, gr=gr, g2=g2, h2=h2, fparam=fparam, aparam=aparam + descriptor, + atype, + gr=gr, + g2=g2, + h2=h2, + fparam=fparam, + aparam=aparam, + vacuum_descriptor=vacuum_descriptor, ) xp = array_api_compat.array_namespace(descriptor, atype) nf = descriptor.shape[0] diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 8d4f459f2f..dfa64ce290 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -140,8 +140,12 @@ def model_call_from_call_lower( try: zero_box = bool(xp_bb.all(bb == 0)) except Exception: - # Traced/symbolic arrays (e.g. JAX export) cannot be evaluated to - # a Python bool; keep the box as-is in that case. + # Traced/symbolic arrays (e.g. jax2tf export) cannot be evaluated + # to a Python bool, so this rewrite cannot be baked into an + # exported graph: the periodic topology is fixed at trace time. + # Callers exporting a non-periodic model must pass ``box=None`` + # at export time; a dynamic all-zero box at runtime is not + # representable and keeps the periodic path. zero_box = False if zero_box: bb = None @@ -1083,7 +1087,6 @@ def format_nlist( the formatted nlist. """ - n_nf, n_nloc, n_nnei = nlist.shape mixed_types = self.mixed_types() # Sentinel capacities (e.g. DPA4C's effectively-unbounded sel) # would allocate absurd padding; a frame can never contribute diff --git a/deepmd/pt_expt/fitting/dpa4c_lr.py b/deepmd/pt_expt/fitting/dpa4c_lr.py index 5123f53387..b8be34323f 100644 --- a/deepmd/pt_expt/fitting/dpa4c_lr.py +++ b/deepmd/pt_expt/fitting/dpa4c_lr.py @@ -47,8 +47,14 @@ def _promote_to_parameter(self, name: str) -> None: ) if name in self._buffers: del self._buffers[name] + # GeneralFitting normalizes ``trainable`` to a per-layer list, so a + # plain bool() on it is always True; freeze only when every layer is + # frozen. + trainable = self.trainable + if isinstance(trainable, list): + trainable = all(trainable) self.register_parameter( - name, torch.nn.Parameter(tensor, requires_grad=bool(self.trainable)) + name, torch.nn.Parameter(tensor, requires_grad=bool(trainable)) ) def forward(self, *args: Any, **kwargs: Any) -> Any: diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 981835bbc9..cc60b0a221 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -2420,12 +2420,18 @@ def _eval_model_graph( coord_input = coords.reshape(nframes, natoms, 3) box_input = cells.reshape(nframes, 9) if cells is not None else None + # An all-zero cell marks a non-periodic frame. Normalize it to None + # before any graph is built: the fused-cell path would otherwise hand + # the singular lattice to the native periodic neighbor search, which + # rejects it. + if box_input is not None and not np.any(box_input != 0): + box_input = None # Build the carry-all graph at its exact edge count; the exported edge # axis is dynamic. graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) lr_graph = None if self.metadata.get("needs_long_range_edges", False): - if box_input is not None and bool(np.any(box_input != 0)): + if box_input is not None: raise NotImplementedError( "DPA4C-LR supports only non-periodic systems (box=None)." ) diff --git a/deepmd/pt_expt/model/dpa4c_lr_model.py b/deepmd/pt_expt/model/dpa4c_lr_model.py index 688c4452bc..04e54df606 100644 --- a/deepmd/pt_expt/model/dpa4c_lr_model.py +++ b/deepmd/pt_expt/model/dpa4c_lr_model.py @@ -250,7 +250,14 @@ def call_common( coord: torch.Tensor, atype: torch.Tensor, box: torch.Tensor | None = None, - **kwargs: Any, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + coord_corr_for_virial: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + neighbor_list: Any | None = None, + neighbor_graph_method: str | None = None, + spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Reject periodic cells; both LR kernels are non-periodic.""" if box is not None: @@ -259,7 +266,19 @@ def call_common( raise NotImplementedError( "DPA4C-LR supports only non-periodic systems (box=None)." ) - return super().call_common(coord, atype, box=box, **kwargs) + return super().call_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + coord_corr_for_virial=coord_corr_for_virial, + charge_spin=charge_spin, + neighbor_list=neighbor_list, + neighbor_graph_method=neighbor_graph_method, + spin=spin, + ) # ------------------------------------------------------------------ # Graph lower with LES/SOG @@ -526,8 +545,10 @@ def call_common_ragged( self.atomic_model.pair_excl, ) # Long-range all-pairs graph, built eagerly outside the graph lower. + # Force the dense builder: the 1e6 sentinel cutoff would explode the + # capacity/cell estimates of the spatial NV/vesin builders. lr_graph = build_ragged_neighbor_graph( - method, coord, atype, n_node, None, 1e6, None + "dense", coord, atype, n_node, None, 1e6, None ) predict = self.forward_common_lower_graph( atype, @@ -605,9 +626,11 @@ def _call_common_graph( spin_flat = spin.reshape(n_padded, 3)[node_index] if spin is not None else None coord_flat = cc.reshape(n_padded, 3)[node_index] - # Long-range all-pairs graph on the compact node axis. + # Long-range all-pairs graph on the compact node axis. Force the + # dense builder: the 1e6 sentinel cutoff would explode the + # capacity/cell estimates of the spatial NV/vesin builders. lr_graph = build_ragged_neighbor_graph( - method, coord_flat, atype_flat, ng.n_node, None, 1e6, None + "dense", coord_flat, atype_flat, ng.n_node, None, 1e6, None ) model_predict = self.forward_common_lower_graph( diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 574a0e5c9d..37bdc714f8 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -93,15 +93,18 @@ def get_dpa4c_lr_model(data: dict) -> DPA4CLREnergyModel: data.setdefault("descriptor", {}) data.setdefault("fitting_net", {}) data["descriptor"].setdefault("type", "dpa4c") + # Default the fitting type as well; otherwise the factory below would + # build the standard ``ener`` fitting for a config that omits it. + data["fitting_net"].setdefault("type", "dpa4c_lr") if data["descriptor"]["type"] not in ("dpa4", "dpa4c", "DPA4", "sezm", "SeZM"): raise ValueError( "Model type 'dpa4c_lr' requires a DPA4/SeZM descriptor, but got " f"descriptor type '{data['descriptor']['type']}'." ) - if data["fitting_net"].get("type", "dpa4c_lr") != "dpa4c_lr": + if data["fitting_net"]["type"] != "dpa4c_lr": raise ValueError( "Model type 'dpa4c_lr' requires fitting type 'dpa4c_lr', but got " - f"'{data['fitting_net'].get('type')}'." + f"'{data['fitting_net']['type']}'." ) descriptor_exclude_types = [ diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index c45f74cb9b..3fc81092fc 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1676,14 +1676,17 @@ def _forward_graph( # fixed across all graph models. needs_lr = getattr(_model, "_needs_long_range_edges", False) if needs_lr: + # Force the dense builder for the all-pairs graph: the 1e6 + # sentinel cutoff would explode the capacity/cell estimates of + # the spatial NV/vesin builders. if ragged: lr_graph = build_ragged_neighbor_graph( - method, coord_3d, atype, n_node, None, 1e6, None + "dense", coord_3d, atype, n_node, None, 1e6, None ) else: coord_flat = coord_3d.reshape(n_padded, 3)[node_index] lr_graph = build_ragged_neighbor_graph( - method, coord_flat, atype_flat, ng.n_node, None, 1e6, None + "dense", coord_flat, atype_flat, ng.n_node, None, 1e6, None ) lr_edge_index = lr_graph.edge_index lr_edge_vec = lr_graph.edge_vec diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 16823017a5..7c351f3e20 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -4229,7 +4229,9 @@ def dpa4c_lr_model_args() -> Argument: ], doc=supported_backends("pt_expt") + doc_fitting, ), - *_bridging_method_args(), + # NOTE: no `bridging_method` here. `expand_bridging_method` only + # accepts `standard`/`dpa4`/`sezm`, so advertising the flag for + # `dpa4c_lr` would pass validation and then fail at construction. ], doc=supported_backends("pt_expt") + doc_model, ) diff --git a/doc/model/dpa4c_lr.md b/doc/model/dpa4c_lr.md index e43e2db219..7f13850ae3 100644 --- a/doc/model/dpa4c_lr.md +++ b/doc/model/dpa4c_lr.md @@ -33,9 +33,11 @@ the response of the latent charges to the atomic positions, so the model remains conservative. :::{warning} -Only **non-periodic systems** (no simulation box) are supported. Periodic -long-range summation (Ewald/SPME-style) is not implemented. Freezing to -`.pt2`/`.pte` is not yet implemented for this model type. +Only **non-periodic systems** are supported: input frames may carry no box +(`box=None`) or an all-zero box, while nonzero periodic cells are rejected at +training, forward, and inference entry points. Periodic long-range summation +(Ewald/SPME-style) is not implemented. Freezing to `.pt2`/`.pte` is not yet +implemented for this model type. ::: ## Training input diff --git a/source/tests/pt_expt/model/test_dpa4c_lr.py b/source/tests/pt_expt/model/test_dpa4c_lr.py index 7d1bf9d9d6..f0812aa8c7 100644 --- a/source/tests/pt_expt/model/test_dpa4c_lr.py +++ b/source/tests/pt_expt/model/test_dpa4c_lr.py @@ -301,7 +301,6 @@ def test_charge_constraint_zero_target_preserves_total_charge() -> None: for parameter in model.parameters(): parameter.copy_(torch.randn_like(parameter) * 0.1) - fitting = model.get_fitting_net() nf, nloc = 2, 5 q = torch.randn(nf * nloc, 1, dtype=torch.float64, device=env.DEVICE) n_node = torch.tensor([nloc] * nf, dtype=torch.long, device=env.DEVICE) @@ -411,6 +410,74 @@ def test_graph_training_gradient_reaches_lr_parameters(lr_kernel: str) -> None: assert torch.isfinite(gradient).all(), name +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_untrainable_fitting_freezes_long_range_parameters(lr_kernel: str) -> None: + """`trainable=False` must freeze the promoted long-range Parameters. + + ``GeneralFitting`` normalizes ``trainable`` to a per-layer list, so a + plain ``bool(self.trainable)`` would always be True and the freeze would + silently not apply. + """ + config = _config(lr_kernel) + config["fitting_net"]["trainable"] = False + fitting = get_model(config).to(env.DEVICE).get_fitting_net() + names = ( + ("les_alpha", "bias_atom_q") + if lr_kernel == "les" + else ( + "amp", + "bandwidth", + "bias_atom_q", + ) + ) + for name in names: + param = getattr(fitting, name) + assert isinstance(param, torch.nn.Parameter), name + assert not param.requires_grad, name + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_untrainable_fitting_survives_an_optimizer_step(lr_kernel: str) -> None: + """A real forward/backward/optimizer step must not move frozen LR params. + + Regression test for the review reproduction: with + ``fitting_net.trainable: false`` the LES/SOG parameters changed after one + Adam step because ``bool(self.trainable)`` on the normalized per-layer + list is always True. + """ + config = _config(lr_kernel) + config["fitting_net"]["trainable"] = False + model = get_model(config).to(env.DEVICE) + fitting = model.get_fitting_net() + names = ( + ("les_alpha", "bias_atom_q") + if lr_kernel == "les" + else ( + "amp", + "bandwidth", + "bias_atom_q", + ) + ) + before = {name: getattr(fitting, name).detach().clone() for name in names} + + coord, atype = _make_batch((5,), seed=7) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + result = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + ) + loss = result["energy"].square().sum() + result["force"].square().sum() + loss.backward() + optimizer.step() + + for name in names: + after = getattr(fitting, name).detach() + assert torch.equal(before[name], after), ( + f"{name} moved despite trainable=False: " + f"{before[name].tolist()} -> {after.tolist()}" + ) + + def test_sog_parameters_serialize_and_remain_trainable() -> None: """SOG amplitudes and widths survive fitting serialization as Parameters.""" fitting = _build_model(lr_kernel="sog").get_fitting_net() From dca5758a5159b9401e9db12f1f512026eb95efd3 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 15 Oct 2026 04:14:01 +0000 Subject: [PATCH 09/12] perf(pt-expt): ragged-native unordered all-pairs graph for DPA4C-LR Build the long-range all-pairs graph directly from the per-frame node counts instead of running a cutoff-based neighbor search with a 1e6 sentinel radius: - New build_all_pairs_unordered_graph emits each intra-frame pair once (i < j) via vectorized triangular indexing -- no distance search, no padding to the batch max width, and half the directed pair storage (sum_f n_f(n_f-1)/2). - Drop the 0.5 double-counting factor from the graph-path LR energy. - New _unordered_pair_force_virial scatters each pair force to both endpoints with opposite signs and splits the pair virial evenly between them; reusing the directed full-to-src attribution would make the per-atom virial depend on the arbitrary endpoint order. - One shared builder now serves the eager rectangular/ragged paths, the compiled training/evaluation wrapper, and DeepEval, independent of the short-range builder policy. Numerics are unchanged: the unordered assembly reproduces the directed edge_force_virial semantics to machine precision, and the graph path matches the legacy dense kernels (the old directed implementation) to ~1e-15 in float64, so existing checkpoints load and run identically. Benchmark (CPU): skewed batch (2000+5x100 atoms) builds 5.2x faster; a single 20k-atom frame halves edge storage (400M -> 200M edges) and builds in 34s instead of 77s. Tests: brute-force pair enumeration, directed-assembly equivalence, legacy dense reference, atom-index permutation invariance, zero/single atom frames, force-loss second-order gradients, and pair-exclusion path consistency. --- AGENTS.md | 2 +- deepmd/pt_expt/infer/deep_eval.py | 32 ++- deepmd/pt_expt/model/dpa4c_lr_model.py | 81 +++++-- deepmd/pt_expt/train/training.py | 20 +- deepmd/pt_expt/utils/graph_builder.py | 114 +++++++++ source/tests/pt_expt/model/test_dpa4c_lr.py | 243 ++++++++++++++++++++ 6 files changed, 454 insertions(+), 38 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ebb487db0c..9ffd103aed 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -165,7 +165,7 @@ source/ # C++ source code and tests - **Configuration**: Use `input_torch.json` format typically - **Training**: `dp --pt train input_torch.json` - **Requirements**: `torch` package -- **DPA4C-LR long-range variant**: model type `dpa4c_lr` with fitting type `dpa4c_lr` adds a non-periodic long-range term on top of the compact DPA4C descriptor (`deepmd/pt_expt/model/dpa4c_lr_model.py`, `deepmd/pt_expt/fitting/dpa4c_lr.py`). Two kernels are available: `les` (`erf(alpha*r)/r`, the default) and a trainable sum-of-Gaussians `sog` with multi-channel latent charges. Only non-periodic systems are supported — `box=None` or an all-zero box is accepted, while nonzero periodic cells are rejected — and only the `pt_expt` backend. The long-range all-pairs graph is built outside the compiled region and passed into the graph lower alongside the short-range graph, so `torch.compile` sees no data-dependent neighbor construction and memory scales as `Σ_f n_f^2` rather than the full-batch `N^2`. `.pt2`/`.pte` freeze is not yet implemented. Tests: `pytest source/tests/pt_expt/model/test_dpa4c_lr.py`. +- **DPA4C-LR long-range variant**: model type `dpa4c_lr` with fitting type `dpa4c_lr` adds a non-periodic long-range term on top of the compact DPA4C descriptor (`deepmd/pt_expt/model/dpa4c_lr_model.py`, `deepmd/pt_expt/fitting/dpa4c_lr.py`). Two kernels are available: `les` (`erf(alpha*r)/r`, the default) and a trainable sum-of-Gaussians `sog` with multi-channel latent charges. Only non-periodic systems are supported — `box=None` or an all-zero box is accepted, while nonzero periodic cells are rejected — and only the `pt_expt` backend. The long-range all-pairs graph is built outside the compiled region and passed into the graph lower alongside the short-range graph, so `torch.compile` sees no data-dependent neighbor construction; it is built ragged-native with unordered pairs (`i < j`, `build_all_pairs_unordered_graph` in `deepmd/pt_expt/utils/graph_builder.py`) straight from the per-frame node counts — no cutoff search, no padding to the batch max width, and half the directed pair storage (`Σ_f n_f(n_f-1)/2`). `.pt2`/`.pte` freeze is not yet implemented. Tests: `pytest source/tests/pt_expt/model/test_dpa4c_lr.py`. ### JAX Backend diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index cc60b0a221..8cfc6f1955 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -2607,7 +2607,10 @@ def _build_eval_graph( ) -> "NeighborGraph": """Build the carry-all NeighborGraph for graph-lower inference. - Dispatches on ``self._neighbor_graph_method``: ``auto`` is resolved + ``long_range=True`` bypasses the cutoff-based builders entirely and + returns the ragged-native unordered all-pairs graph used by models + with a non-periodic long-range term. Otherwise dispatches on + ``self._neighbor_graph_method``: ``auto`` is resolved call-time via :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder` using the batch frame count (vesin only when ``nf == 1``); @@ -2619,11 +2622,24 @@ def _build_eval_graph( destination-major graph-form ``.pt2`` ABI after construction; the caller transfers its fields to the model device. """ - # The LR graph is all-pairs by definition. Do not send its artificial - # 1e6 cutoff through the spatial NV/vesin builders: their capacity/cell - # estimates are intended for finite local cutoffs and can explode for - # this sentinel radius. - method = "dense" if long_range else self._neighbor_graph_method + # The LR graph is all-pairs by definition: build it ragged-native and + # unordered straight from the frame sizes -- no cutoff search, no + # padding to a common width, and no dependence on the short-range + # builder policy (the 1e6 sentinel cutoff used to explode the NV + # capacity estimate). + if long_range: + from deepmd.pt_expt.utils.graph_builder import ( + build_all_pairs_unordered_graph, + ) + + coord_arr = np.asarray(coord_input) + nf, natoms = coord_arr.shape[0], coord_arr.shape[1] + coord_t = torch.as_tensor( + coord_arr.reshape(-1, 3), dtype=torch.float64, device=device + ) + n_node_t = torch.full((nf,), natoms, dtype=torch.int64, device=device) + return build_all_pairs_unordered_graph(coord_t, n_node_t) + method = self._neighbor_graph_method if method == "auto": coord_arr = np.asarray(coord_input) nf = int(coord_arr.shape[0]) if coord_arr.ndim >= 2 else 1 @@ -2632,8 +2648,8 @@ def _build_eval_graph( # (decision #18): apply it here so the exported ``.pt2`` lower consumes a # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++ # ``applyPairExclusion`` and the eager dpmodel/pt_expt build path). - pair_excl = None if long_range else self._model_pair_excl() - rcut = 1.0e6 if long_range else self._rcut + pair_excl = self._model_pair_excl() + rcut = self._rcut builder_device = torch.device("cpu") if method == "cell" else device # The fused builder writes the whole destination-major payload from one # search. It applies only where nothing has to be filtered or masked diff --git a/deepmd/pt_expt/model/dpa4c_lr_model.py b/deepmd/pt_expt/model/dpa4c_lr_model.py index 04e54df606..117d831c85 100644 --- a/deepmd/pt_expt/model/dpa4c_lr_model.py +++ b/deepmd/pt_expt/model/dpa4c_lr_model.py @@ -17,7 +17,8 @@ * ``grad(E_SR + E_LR, edge_vec)`` gives the short-range force plus the charge-response term ``∂E_LR/∂q · dq/dR``; * ``grad(E_LR, lr_edge_vec)`` gives the geometry term ``∂E_LR/∂r|_q``, - assembled with the ``edge_force_virial`` conventions. + assembled with the unordered-pair conventions of + :func:`_unordered_pair_force_virial`. Only non-periodic systems (``box=None``) are supported. """ @@ -41,7 +42,6 @@ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, compact_nodes, - edge_force_virial, expand_node_values, frame_id_from_n_node, segment_sum, @@ -62,6 +62,7 @@ BaseModel, ) from deepmd.pt_expt.utils.graph_builder import ( + build_all_pairs_unordered_graph, build_neighbor_graph_for_method, build_ragged_neighbor_graph, resolve_neighbor_graph_method, @@ -77,6 +78,56 @@ E2_PER_ANGSTROM_TO_EV = 14.3996454784255 +def _unordered_pair_force_virial( + g_e: torch.Tensor, + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + n_node: torch.Tensor, + node_capacity: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Assemble force/virial for an UNORDERED (``i < j``) pair graph. + + Each edge is a distinct atom pair counted once, which differs from the + directed ``edge_force_virial`` contract in two ways: + + * the force still scatters the per-edge gradient ``g_e = dE/d(edge_vec)`` + to both endpoints with opposite signs (``dst += g_e``, ``src -= g_e``); + * the pair virial ``-g_e (x) edge_vec`` is split EVENLY between the two + endpoints. The directed convention attributes the full edge virial to + ``src``, which is correct there only because the two opposite + half-weighted edges of a pair split it implicitly; reusing it here + would make the per-atom virial depend on the arbitrary endpoint order. + + The per-frame virial is unaffected by the split: it reduces the per-atom + virial over each frame's nodes, and the two halves of a pair live in the + same frame. + """ + nf = n_node.shape[0] + n_out = node_capacity + dtype = g_e.dtype + device = g_e.device + if isinstance(n_out, int) and n_out == 0: + return ( + torch.zeros((0, 3), dtype=dtype, device=device), + torch.zeros((0, 3, 3), dtype=dtype, device=device), + torch.zeros((nf, 3, 3), dtype=dtype, device=device), + ) + g = g_e * edge_mask.to(dtype).unsqueeze(-1) + src = edge_index[0] % n_out + dst = edge_index[1] % n_out + force = segment_sum(g, dst, n_out) - segment_sum(g, src, n_out) + # Half-weight pair virial to each endpoint (broadcast, no einsum). + w_half = -0.5 * (g[:, :, None] * edge_vec[:, None, :]) # (E, 3, 3) + atom_virial = segment_sum(w_half, src, n_out) + segment_sum(w_half, dst, n_out) + # Reduce the per-atom virial over nodes (N << E) rather than per-edge + # over pairs; both endpoints of a pair share a frame, so the per-frame + # total equals the full pair virial sum. + node_frame = frame_id_from_n_node(n_node, n_total=n_out) + virial = segment_sum(atom_virial, node_frame, nf) + return force, atom_virial, virial + + def _les_kernel_from_squared_distance( r_sq: torch.Tensor, les_alpha: torch.Tensor, @@ -386,7 +437,7 @@ def forward_common_lower_graph( create_graph=self.training, retain_graph=True, ) - lr_force, lr_atom_virial, lr_virial = edge_force_virial( + lr_force, lr_atom_virial, lr_virial = _unordered_pair_force_virial( g_lr, lr_edge_vec, lr_edge_index, @@ -427,7 +478,9 @@ def _apply_lr_correction_graph( The long-range edges are built eagerly by the caller and passed in as inputs. They contain only intra-frame all-pairs, so the storage scales - with ``Σ_f n_f^2`` rather than the full batch ``N^2``. + with ``Σ_f n_f^2`` rather than the full batch ``N^2``. The pairs are + UNORDERED (``i < j``, each pair counted once), so the energy carries + no ``0.5`` double-counting factor. """ if "latent_charge" not in atomic_ret: return atomic_ret, None @@ -469,7 +522,8 @@ def _apply_lr_correction_graph( e_edge = (q_src * q_dst).sum(-1) * K * lr_edge_mask.to(dtype) # (E_lr,) edge_frame = frame_id[lr_edge_index[1]] # (E_lr,) frame of dst e_lr_raw = segment_sum(e_edge, edge_frame, nf) - e_lr = e_lr_raw * (0.5 * E2_PER_ANGSTROM_TO_EV) + # Unordered pairs: each pair counted once, no 0.5 factor. + e_lr = e_lr_raw * E2_PER_ANGSTROM_TO_EV # Spread the frame LR energy evenly over the real atoms. nreal = n_node.to(e_lr.dtype).clamp_min(1) @@ -545,11 +599,9 @@ def call_common_ragged( self.atomic_model.pair_excl, ) # Long-range all-pairs graph, built eagerly outside the graph lower. - # Force the dense builder: the 1e6 sentinel cutoff would explode the - # capacity/cell estimates of the spatial NV/vesin builders. - lr_graph = build_ragged_neighbor_graph( - "dense", coord, atype, n_node, None, 1e6, None - ) + # Ragged-native unordered pairs: no cutoff search, no padding to the + # batch max width, and half the directed pair storage. + lr_graph = build_all_pairs_unordered_graph(coord, n_node) predict = self.forward_common_lower_graph( atype, graph.n_node, @@ -626,12 +678,9 @@ def _call_common_graph( spin_flat = spin.reshape(n_padded, 3)[node_index] if spin is not None else None coord_flat = cc.reshape(n_padded, 3)[node_index] - # Long-range all-pairs graph on the compact node axis. Force the - # dense builder: the 1e6 sentinel cutoff would explode the - # capacity/cell estimates of the spatial NV/vesin builders. - lr_graph = build_ragged_neighbor_graph( - "dense", coord_flat, atype_flat, ng.n_node, None, 1e6, None - ) + # Long-range all-pairs graph on the compact node axis: ragged-native + # unordered pairs, independent of the short-range search policy. + lr_graph = build_all_pairs_unordered_graph(coord_flat, ng.n_node) model_predict = self.forward_common_lower_graph( atype_flat, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 3fc81092fc..836fe7ea36 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -908,12 +908,10 @@ def _trace_and_compile_graph( trace_N, 3, dtype=GLOBAL_PT_FLOAT_PRECISION, device=_trace_device ) from deepmd.pt_expt.utils.graph_builder import ( - build_ragged_neighbor_graph, + build_all_pairs_unordered_graph, ) - s_lr_graph = build_ragged_neighbor_graph( - "dense", s_lr_coord, s_atype, s_n_node, None, 1e6, None - ) + s_lr_graph = build_all_pairs_unordered_graph(s_lr_coord, s_n_node) def fn( atype: torch.Tensor, @@ -1582,6 +1580,7 @@ def _forward_graph( expand_node_values, ) from deepmd.pt_expt.utils.graph_builder import ( + build_all_pairs_unordered_graph, build_neighbor_graph_for_method, build_ragged_neighbor_graph, ) @@ -1676,18 +1675,13 @@ def _forward_graph( # fixed across all graph models. needs_lr = getattr(_model, "_needs_long_range_edges", False) if needs_lr: - # Force the dense builder for the all-pairs graph: the 1e6 - # sentinel cutoff would explode the capacity/cell estimates of - # the spatial NV/vesin builders. + # Ragged-native unordered all-pairs: no cutoff search and no + # dependence on the short-range builder policy. if ragged: - lr_graph = build_ragged_neighbor_graph( - "dense", coord_3d, atype, n_node, None, 1e6, None - ) + lr_graph = build_all_pairs_unordered_graph(coord_3d, n_node) else: coord_flat = coord_3d.reshape(n_padded, 3)[node_index] - lr_graph = build_ragged_neighbor_graph( - "dense", coord_flat, atype_flat, ng.n_node, None, 1e6, None - ) + lr_graph = build_all_pairs_unordered_graph(coord_flat, ng.n_node) lr_edge_index = lr_graph.edge_index lr_edge_vec = lr_graph.edge_vec lr_edge_mask = lr_graph.edge_mask diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index 8c2b37593b..f071dfd3a1 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -342,3 +342,117 @@ def build_neighbor_graph_for_method( f"unknown neighbor_graph_method {method!r}; use 'dense', 'ase', " "'cell', 'vesin', or 'nv'" ) + + +class AllPairsGraph: + """Unordered intra-frame all-pairs graph for non-periodic LR terms. + + Only the fields consumed by the long-range consumers are populated: + ``n_node``, ``edge_index`` (``[src, dst]``), ``edge_vec`` + (``r_src - r_dst``), and an all-true ``edge_mask``. Each atom pair + appears exactly once (``i < j``), so the energy contract is + ``sum_{i < j}`` with NO ``0.5`` factor, and the force/virial assembly + must scatter each pair to BOTH endpoints. + + Attributes + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + edge_index + ``[src, dst]`` node endpoints with shape ``(2, E)``, int64. + edge_vec + Per-edge displacement ``r_src - r_dst`` with shape ``(E, 3)``. + edge_mask + Valid-edge mask with shape ``(E,)``; always all-true (the builder + emits exactly the real pairs, there is no padding). + """ + + def __init__( + self, + n_node: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + ) -> None: + self.n_node = n_node + self.edge_index = edge_index + self.edge_vec = edge_vec + self.edge_mask = edge_mask + + +def build_all_pairs_unordered_graph( + coord: torch.Tensor, + n_node: torch.Tensor, +) -> AllPairsGraph: + """Build the unordered (``i < j``) intra-frame all-pairs graph directly. + + Every pair of atoms that shares a frame participates in a non-periodic + all-pairs interaction, so no distance search, padding, or cutoff is + involved: the pair indices are generated straight from the per-frame + node counts and the displacements are gathered from them. Storage is + proportional to the true pair count ``sum_f n_f (n_f - 1) / 2`` -- half + the directed all-pairs count, and without widening every frame to the + batch maximum width. + + The builder runs eagerly (outside any compiled region), so the + data-dependent pair count as a Python int is fine here. + + Parameters + ---------- + coord + Flat local coordinates with shape ``(N, 3)``, frame-major over + ``n_node``. + n_node + Atoms per frame with shape ``(nf,)``. + + Returns + ------- + AllPairsGraph + The unordered all-pairs graph on the flat node axis. + """ + nf = n_node.shape[0] + device = n_node.device + n_pairs = (n_node * (n_node - 1)) // 2 # (nf,) + total_pairs = int(n_pairs.sum()) + if total_pairs == 0: + return AllPairsGraph( + n_node=n_node, + edge_index=torch.zeros((2, 0), dtype=torch.int64, device=device), + edge_vec=torch.zeros((0, 3), dtype=coord.dtype, device=coord.device), + edge_mask=torch.zeros((0,), dtype=torch.bool, device=device), + ) + + pair_frame = torch.repeat_interleave(torch.arange(nf, device=device), n_pairs) + pair_start = torch.cumsum(n_pairs, 0) - n_pairs # (nf,) + node_start = torch.cumsum(n_node, 0) - n_node # (nf,) + # Pair rank within its frame, in row-major strict-upper-triangle order + # (0,1), (0,2), ..., (0,n-1), (1,2), ... + k = (torch.arange(total_pairs, device=device) - pair_start[pair_frame]).to( + torch.float64 + ) + n = n_node[pair_frame].to(torch.float64) + # Row i starts at S_i = i * (2n - 1 - i) / 2; solve S_i <= k < S_{i+1}. + # float64 is exact for the sizes involved (n^2 << 2^53); one correction + # step below guards the sqrt rounding regardless. + discr = (2 * n - 1).square() - 8 * k + ii = torch.floor(((2 * n - 1) - torch.sqrt(discr.clamp_min(0.0))) / 2) + ii = ii.clamp_min(0.0) + + def row_start(i: torch.Tensor) -> torch.Tensor: + return i * (2 * n - 1 - i) / 2 + + ii = torch.where(k < row_start(ii), ii - 1, ii) + ii = torch.where(k >= row_start(ii + 1), ii + 1, ii) + jj = ii + 1 + (k - row_start(ii)) + + src = node_start[pair_frame] + ii.to(torch.int64) + dst = node_start[pair_frame] + jj.to(torch.int64) + edge_index = torch.stack([src, dst]) + edge_vec = coord[src] - coord[dst] + edge_mask = torch.ones(total_pairs, dtype=torch.bool, device=device) + return AllPairsGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) diff --git a/source/tests/pt_expt/model/test_dpa4c_lr.py b/source/tests/pt_expt/model/test_dpa4c_lr.py index f0812aa8c7..5ee4250725 100644 --- a/source/tests/pt_expt/model/test_dpa4c_lr.py +++ b/source/tests/pt_expt/model/test_dpa4c_lr.py @@ -17,6 +17,13 @@ E2_PER_ANGSTROM_TO_EV, _les_kernel_from_squared_distance, _SOGKernel, + _unordered_pair_force_virial, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + edge_force_virial, +) +from deepmd.pt_expt.utils.graph_builder import ( + build_all_pairs_unordered_graph, ) from deepmd.pt_expt.model.get_model import ( get_model, @@ -553,3 +560,239 @@ def test_compiled_lower_matches_eager(lr_kernel: str) -> None: torch.testing.assert_close(got["energy"], expected["energy"]) torch.testing.assert_close(got["force"], expected["force"]) + + +def test_all_pairs_builder_matches_brute_force() -> None: + """The ragged-native builder emits exactly the unordered intra-frame pairs.""" + torch.manual_seed(0) + for nlocs in [(3,), (1, 4, 2), (0, 5), (7, 1, 0, 3)]: + n_node = torch.tensor(nlocs, dtype=torch.long, device=env.DEVICE) + coord = torch.randn( + int(n_node.sum()), 3, dtype=torch.float64, device=env.DEVICE + ) + graph = build_all_pairs_unordered_graph(coord, n_node) + expected = [] + offset = 0 + for n in nlocs: + expected.extend( + (offset + i, offset + j) for i in range(n) for j in range(i + 1, n) + ) + offset += n + got = [tuple(pair) for pair in graph.edge_index.t().tolist()] + assert got == expected + src, dst = graph.edge_index + assert torch.equal(graph.edge_vec, coord[src] - coord[dst]) + assert bool(graph.edge_mask.all()) + + +def test_unordered_pair_assembly_matches_directed_reference() -> None: + """Unordered force/virial assembly reproduces the directed convention. + + The directed reference represents each pair as two opposite edges, each + carrying half the pair energy, and attributes the full edge virial to + ``src`` -- the two half-weighted edges then split the pair virial evenly + between the endpoints. The unordered assembly must reproduce exactly + that split while storing each pair once. + """ + torch.manual_seed(0) + nlocs = (2, 1, 5, 3) + n_node = torch.tensor(nlocs, dtype=torch.long, device=env.DEVICE) + n_total = int(n_node.sum()) + coord = torch.randn(n_total, 3, dtype=torch.float64, device=env.DEVICE) + + graph = build_all_pairs_unordered_graph(coord, n_node) + n_pair = graph.edge_index.shape[1] + g_pair = torch.randn(n_pair, 3, dtype=torch.float64, device=env.DEVICE) + + force_u, atom_virial_u, virial_u = _unordered_pair_force_virial( + g_pair, + graph.edge_vec, + graph.edge_index, + graph.edge_mask, + n_node, + node_capacity=n_total, + ) + + index_dir = torch.cat([graph.edge_index, graph.edge_index.flip(0)], dim=1) + vec_dir = torch.cat([graph.edge_vec, -graph.edge_vec], dim=0) + g_dir = torch.cat([0.5 * g_pair, -0.5 * g_pair], dim=0) + mask_dir = torch.ones(2 * n_pair, dtype=torch.bool, device=env.DEVICE) + force_d, atom_virial_d, virial_d = edge_force_virial( + g_dir, vec_dir, index_dir, mask_dir, n_node, node_capacity=n_total + ) + + torch.testing.assert_close(force_u, force_d) + torch.testing.assert_close(atom_virial_u, atom_virial_d) + torch.testing.assert_close(virial_u, virial_d) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_graph_matches_legacy_dense_reference(lr_kernel: str) -> None: + """The unordered graph path matches the directed dense legacy path. + + The dense fallback (``neighbor_graph_method='legacy'``) evaluates the LR + kernels as ``0.5 * sum_{i != j}`` over the padded batch -- the directed + reference. Energy, force, and the total virial must match exactly. The + per-atom virial follows a different (pre-existing) decomposition + convention in the dense path, so it is covered separately by the + directed-assembly and permutation tests. + """ + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((4, 7, 3)) + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + + graph_ret = model.call_common(coord_t, atype_t, do_atomic_virial=True) + dense_ret = model.call_common( + coord_t, atype_t, do_atomic_virial=True, neighbor_graph_method="legacy" + ) + for key in ("energy_redu", "energy_derv_r", "energy_derv_c_redu"): + torch.testing.assert_close( + graph_ret[key], dense_ret[key], atol=1e-10, rtol=1e-10 + ) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_atom_index_permutation_invariance(lr_kernel: str) -> None: + """Shuffling the atoms of a frame must not change any output. + + In particular this guards the even pair-virial split: an assembly that + attributes the pair virial to an arbitrary endpoint would change the + per-atom virial under a permutation. + """ + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((6,), seed=5) + perm = np.random.default_rng(0).permutation(6) + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + coord_p = torch.tensor(coord[:, perm], dtype=torch.float64, device=env.DEVICE) + atype_p = torch.tensor(atype[:, perm], dtype=torch.long, device=env.DEVICE) + + base = model(coord_t, atype_t, do_atomic_virial=True) + permuted = model(coord_p, atype_p, do_atomic_virial=True) + + perm_t = torch.tensor(perm, dtype=torch.long, device=env.DEVICE) + torch.testing.assert_close(permuted["energy"], base["energy"]) + torch.testing.assert_close(permuted["virial"], base["virial"]) + torch.testing.assert_close(permuted["force"][0], base["force"][0][perm_t]) + torch.testing.assert_close( + permuted["atom_virial"][0], base["atom_virial"][0][perm_t] + ) + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_zero_and_single_atom_frames(lr_kernel: str) -> None: + """Frames with zero or one atom contribute no LR pairs and stay finite.""" + torch.manual_seed(1234) + model = _build_model(lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + # Ragged layout: a zero-atom frame cannot be represented in a padded + # batch, so exercise it here. + nlocs = (1, 0, 4) + rng = np.random.default_rng(9) + flat_coord = torch.tensor( + rng.uniform(0.0, 6.0, (sum(nlocs), 3)), dtype=torch.float64, device=env.DEVICE + ) + flat_atype = torch.tensor( + rng.integers(0, 2, sum(nlocs)), dtype=torch.long, device=env.DEVICE + ) + n_node = torch.tensor(nlocs, dtype=torch.long, device=env.DEVICE) + ragged = model.forward_ragged(flat_coord, flat_atype, n_node) + for key in ("energy", "force"): + assert torch.isfinite(ragged[key]).all(), key + + # A single-atom frame has no pairs: ragged and padded layouts must agree. + nlocs = (1, 5, 2) + coord, atype = _make_batch(nlocs, seed=13) + padded = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + ) + flat_coord = torch.tensor( + np.concatenate([coord[i, :nloc] for i, nloc in enumerate(nlocs)]), + dtype=torch.float64, + device=env.DEVICE, + ) + flat_atype = torch.tensor( + np.concatenate([atype[i, :nloc] for i, nloc in enumerate(nlocs)]), + dtype=torch.long, + device=env.DEVICE, + ) + ragged = model.forward_ragged( + flat_coord, flat_atype, torch.tensor(nlocs, dtype=torch.long, device=env.DEVICE) + ) + torch.testing.assert_close(padded["energy"], ragged["energy"]) + offset = 0 + for index, nloc in enumerate(nlocs): + torch.testing.assert_close( + padded["force"][index, :nloc], + ragged["force"][offset : offset + nloc], + ) + offset += nloc + + +@pytest.mark.parametrize("lr_kernel", ["les", "sog"]) +def test_force_loss_backward_through_long_range(lr_kernel: str) -> None: + """A training-mode force loss backpropagates into the LR kernel parameters.""" + torch.manual_seed(1234) + model = _build_model(train=True, lr_kernel=lr_kernel) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((5,), seed=11) + result = model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + ) + loss = result["energy"].square().sum() + result["force"].square().sum() + loss.backward() + + fitting = model.get_fitting_net() + names = ("les_alpha",) if lr_kernel == "les" else ("amp", "bandwidth") + for name in names: + grad = getattr(fitting, name).grad + assert grad is not None, name + assert torch.isfinite(grad).all(), name + assert torch.count_nonzero(grad) > 0, name + + +def test_pair_exclusion_consistent_between_paths() -> None: + """Model-level pair exclusion must not change the LR term between paths. + + Both paths apply ``pair_exclude_types`` to the short-range graph only; + the all-pairs LR term deliberately ignores it. The graph and dense + paths must therefore still agree. + """ + config = _config("les") + config["pair_exclude_types"] = [[0, 1]] + torch.manual_seed(1234) + model = get_model(config).to(env.DEVICE) + model.eval() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + + coord, atype = _make_batch((5, 6), seed=17) + coord_t = torch.tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + + graph_ret = model.call_common(coord_t, atype_t) + dense_ret = model.call_common(coord_t, atype_t, neighbor_graph_method="legacy") + for key in ("energy_redu", "energy_derv_r", "energy_derv_c_redu"): + torch.testing.assert_close( + graph_ret[key], dense_ret[key], atol=1e-10, rtol=1e-10 + ) From 71b39b93e9f8fb804d1dda136f862957d9fb0fe1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 04:15:12 +0000 Subject: [PATCH 10/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/model/test_dpa4c_lr.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/source/tests/pt_expt/model/test_dpa4c_lr.py b/source/tests/pt_expt/model/test_dpa4c_lr.py index 5ee4250725..bf5509ddd2 100644 --- a/source/tests/pt_expt/model/test_dpa4c_lr.py +++ b/source/tests/pt_expt/model/test_dpa4c_lr.py @@ -7,6 +7,9 @@ import pytest import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + edge_force_virial, +) from deepmd.infer import ( DeepPot, ) @@ -19,12 +22,6 @@ _SOGKernel, _unordered_pair_force_virial, ) -from deepmd.dpmodel.utils.neighbor_graph import ( - edge_force_virial, -) -from deepmd.pt_expt.utils.graph_builder import ( - build_all_pairs_unordered_graph, -) from deepmd.pt_expt.model.get_model import ( get_model, ) @@ -34,6 +31,9 @@ from deepmd.pt_expt.utils import ( env, ) +from deepmd.pt_expt.utils.graph_builder import ( + build_all_pairs_unordered_graph, +) from ..compile_utils import ( REQUIRES_SUPPORTED_COMPILE, From 69e8b794b77c3f1cda858a2f09b4ffc6a29b6554 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 15 Oct 2026 04:31:47 +0000 Subject: [PATCH 11/12] style(pt-expt): explicit dtype in all-pairs graph builder arange calls Fixes the pre-commit.ci pylint failure (E8002 no-explicit-dtype). --- deepmd/pt_expt/utils/graph_builder.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index f071dfd3a1..08a17ab45c 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -422,14 +422,17 @@ def build_all_pairs_unordered_graph( edge_mask=torch.zeros((0,), dtype=torch.bool, device=device), ) - pair_frame = torch.repeat_interleave(torch.arange(nf, device=device), n_pairs) + pair_frame = torch.repeat_interleave( + torch.arange(nf, dtype=torch.int64, device=device), n_pairs + ) pair_start = torch.cumsum(n_pairs, 0) - n_pairs # (nf,) node_start = torch.cumsum(n_node, 0) - n_node # (nf,) # Pair rank within its frame, in row-major strict-upper-triangle order # (0,1), (0,2), ..., (0,n-1), (1,2), ... - k = (torch.arange(total_pairs, device=device) - pair_start[pair_frame]).to( - torch.float64 - ) + k = ( + torch.arange(total_pairs, dtype=torch.int64, device=device) + - pair_start[pair_frame] + ).to(torch.float64) n = n_node[pair_frame].to(torch.float64) # Row i starts at S_i = i * (2n - 1 - i) / 2; solve S_i <= k < S_{i+1}. # float64 is exact for the sizes involved (n^2 << 2^53); one correction From 013c057d81316b0ae38cafa05f4778731753c0f3 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 15 Oct 2026 10:17:46 +0000 Subject: [PATCH 12/12] ci: retrigger wheel builds after cibuildwheel download outage