From d8b27ce67e3141a677f69deb19f03fe24b41f8f5 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 17 Sep 2026 06:59:24 +0000 Subject: [PATCH 01/30] 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/30] [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/30] 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/30] 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/30] 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/30] [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/30] 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/30] 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/30] 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/30] [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/30] 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/30] ci: retrigger wheel builds after cibuildwheel download outage From 4d5795c56a5ef82fe45e9c5bc8898544e616e8ed Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 18 Oct 2026 18:24:30 +0800 Subject: [PATCH 13/30] refactor(pt_expt): make LES/SOG a long-range module of the atomic energy model Move the latent-charge long-range energy into DPAtomicModel, so that every consumer of the atomic model evaluates the complete potential: training and compiled training, output-bias statistics and change-bias, the graph Hessian, linear compositions, fine-tuning and checkpoint inference. - A standard energy model takes a `long_range` section (type `latent_charge`, kernels `les` and `sog`, optional total-charge constraint). The charge network reads the descriptor shared with the energy fitting and takes its conditioning inputs; the energy is the direct pair sum over complete non-periodic frames, each pair energy attributed half to each atom. - Components: LatentChargeFitting (charge head), interaction kernels with their own trainable parameters (widths stored as logarithms), and the LatentChargeLongRange module in deepmd/dpmodel/long_range with pt_expt wrappers in deepmd/pt_expt/long_range. - Atomic models declare the topologies they read (`required_topologies()`); callers build the `all_pairs` topology outside the compiled region with `build_topologies`, which honours model-level pair exclusion and rejects periodic cells. The graph lower makes the pair vectors further autograd leaves and pulls them back with `pair_force_virial`, so force, virial, atomic virial, Hessian and force-loss double backward include the long-range term and the charge response. - Compiled training receives the topology as additional dynamic inputs; in a graph traced for a symbolic pair count the kernel coefficients reach the pairs through the node axis, which keeps their gradient reduction parallel for any batch. Ordinary models keep their compiled input signature and serialization record (version 2; models with a long-range module write version 3). - Only DPA4C is admitted (`check_long_range_support`). The dense lower, periodic cells, domain decomposition, compression, export, spin models, tabulated ZBL, bridging, atom exclusion and other backends are rejected explicitly. - Remove the `dpa4c_lr` model and fitting types, their routing, the long-range edge inputs of every compiled graph, and the neighbor-list and box changes the previous implementation made to ordinary models. --- AGENTS.md | 1 - .../dpmodel/atomic_model/base_atomic_model.py | 48 +- .../dpmodel/atomic_model/dp_atomic_model.py | 205 ++++- .../dpmodel/atomic_model/inner_potential.py | 9 +- .../atomic_model/linear_atomic_model.py | 56 +- deepmd/dpmodel/descriptor/dpa4c.py | 11 +- deepmd/dpmodel/fitting/__init__.py | 8 +- deepmd/dpmodel/fitting/dpa4c_lr.py | 434 --------- deepmd/dpmodel/fitting/latent_charge.py | 45 + deepmd/dpmodel/long_range/__init__.py | 34 + deepmd/dpmodel/long_range/base_long_range.py | 156 ++++ deepmd/dpmodel/long_range/kernel.py | 274 ++++++ deepmd/dpmodel/long_range/latent_charge.py | 354 +++++++ deepmd/dpmodel/model/dp_model.py | 2 + deepmd/dpmodel/model/make_model.py | 67 +- deepmd/dpmodel/model/model.py | 4 + deepmd/dpmodel/model/model_factory.py | 115 +++ deepmd/dpmodel/model/native_spin_model.py | 4 + deepmd/dpmodel/model/spin_model.py | 5 + deepmd/dpmodel/utils/default_neighbor_list.py | 19 +- .../dpmodel/utils/neighbor_graph/__init__.py | 20 +- .../utils/neighbor_graph/derivatives.py | 68 ++ .../utils/neighbor_graph/pair_graph.py | 231 +++++ deepmd/dpmodel/utils/nlist.py | 11 - deepmd/jax/atomic_model/dp_atomic_model.py | 2 + deepmd/pd/model/model/__init__.py | 5 + deepmd/pt/model/model/__init__.py | 10 + deepmd/pt_expt/fitting/__init__.py | 8 +- deepmd/pt_expt/fitting/dpa4c_lr.py | 82 -- deepmd/pt_expt/fitting/latent_charge.py | 12 + deepmd/pt_expt/infer/deep_eval.py | 106 +-- deepmd/pt_expt/long_range/__init__.py | 18 + deepmd/pt_expt/long_range/base_long_range.py | 9 + deepmd/pt_expt/long_range/kernel.py | 55 ++ deepmd/pt_expt/long_range/latent_charge.py | 42 + deepmd/pt_expt/model/__init__.py | 4 - deepmd/pt_expt/model/dpa4c_lr_model.py | 865 ------------------ deepmd/pt_expt/model/edge_transform_output.py | 45 +- deepmd/pt_expt/model/get_model.py | 60 +- deepmd/pt_expt/model/make_model.py | 72 +- deepmd/pt_expt/train/training.py | 157 ++-- deepmd/pt_expt/utils/graph_builder.py | 119 --- deepmd/pt_expt/utils/serialization.py | 5 + deepmd/tf/model/model.py | 5 + deepmd/tf2/atomic_model/dp_atomic_model.py | 2 + deepmd/utils/argcheck.py | 380 ++++---- deepmd/utils/finetune.py | 12 + doc/model/dpa4c_lr.md | 82 -- doc/model/index.rst | 2 +- doc/model/long-range.md | 128 +++ .../tests/common/dpmodel/test_long_range.py | 192 ++++ .../tests/common/dpmodel/test_zbl_bridging.py | 3 + source/tests/common/test_finetune_utils.py | 22 + .../pt_expt/model/test_dpa2_graph_lower.py | 81 +- source/tests/pt_expt/model/test_dpa4c_lr.py | 798 ---------------- source/tests/pt_expt/model/test_long_range.py | 673 ++++++++++++++ source/tests/pt_expt/test_change_bias.py | 2 + 57 files changed, 3271 insertions(+), 2968 deletions(-) delete mode 100644 deepmd/dpmodel/fitting/dpa4c_lr.py create mode 100644 deepmd/dpmodel/fitting/latent_charge.py create mode 100644 deepmd/dpmodel/long_range/__init__.py create mode 100644 deepmd/dpmodel/long_range/base_long_range.py create mode 100644 deepmd/dpmodel/long_range/kernel.py create mode 100644 deepmd/dpmodel/long_range/latent_charge.py create mode 100644 deepmd/dpmodel/utils/neighbor_graph/pair_graph.py delete mode 100644 deepmd/pt_expt/fitting/dpa4c_lr.py create mode 100644 deepmd/pt_expt/fitting/latent_charge.py create mode 100644 deepmd/pt_expt/long_range/__init__.py create mode 100644 deepmd/pt_expt/long_range/base_long_range.py create mode 100644 deepmd/pt_expt/long_range/kernel.py create mode 100644 deepmd/pt_expt/long_range/latent_charge.py delete mode 100644 deepmd/pt_expt/model/dpa4c_lr_model.py delete mode 100644 doc/model/dpa4c_lr.md create mode 100644 doc/model/long-range.md create mode 100644 source/tests/common/dpmodel/test_long_range.py delete mode 100644 source/tests/pt_expt/model/test_dpa4c_lr.py create mode 100644 source/tests/pt_expt/model/test_long_range.py diff --git a/AGENTS.md b/AGENTS.md index 9ffd103aed..617c904f11 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -165,7 +165,6 @@ 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; 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/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 25b770a5fb..df334852de 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -221,6 +221,27 @@ def fused_decomposition(self) -> "tuple[Any, Any] | None": """ return None + def has_long_range(self) -> bool: + """Whether the atomic energy carries a long-range module beside a fitting. + + Generic capability (concrete default ``False``): a model wrapping the + atomic model asks it whether the energy holds a long-range interaction, + e.g. before combining it with a scheme defined for a learned local + energy alone. + """ + return False + + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the graph forward reads beside the neighbor graph. + + Generic capability (concrete default: none) in the sense of + :meth:`uses_graph_lower`: a long-range module, whose interaction has no + cutoff, declares the topology it is evaluated on (see + :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`), and + every caller of the graph forward builds what is declared. + """ + return () + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. @@ -493,6 +514,8 @@ def forward_common_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict: """Graph analogue of :meth:`forward_common_atomic` on the flat node axis. @@ -529,6 +552,13 @@ def forward_common_atomic_graph( (and, from there, the descriptor's ``call_graph``). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_common_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name, built on the + node axis of ``graph``; ``None`` when none is required. + node_mask + flat mask of the nodes an enclosing model counts (real, owned and + not excluded there), combined with this model's own; ``None`` + imposes nothing further. N Returns ------- @@ -537,6 +567,8 @@ def forward_common_atomic_graph( """ graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype) + if node_mask is not None: + output_mask = output_mask & node_mask ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -545,6 +577,8 @@ def forward_common_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=output_mask, ) return self._finalize_atomic_ret(ret_dict, output_mask, atype) @@ -940,6 +974,7 @@ def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]: from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, + build_topologies, ) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -989,9 +1024,19 @@ def model_forward( self.get_rcut(), pair_excl=self.pair_excl, ) + atype_flat = xp.reshape(atype, (-1,)) + topologies = build_topologies( + self.required_topologies(), + xp.reshape(coord, (-1, 3)), + graph.n_node, + box, + node_mask=atype_flat >= 0, + atype=atype_flat, + pair_excl=self.pair_excl, + ) atomic_ret = self.forward_common_atomic_graph( graph, - xp.reshape(atype, (-1,)), + atype_flat, fparam=fparam, aparam=( xp.reshape( @@ -1003,6 +1048,7 @@ def model_forward( ), charge_spin=charge_spin, spin=None if spin is None else xp.reshape(spin, (-1, 3)), + topologies=topologies, ) # The graph route works on a flat node axis; restore the # per-frame layout the dense route returns. diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index b9f20ed8c7..0122c29f8a 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -24,6 +24,10 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, + check_long_range_support, +) from deepmd.dpmodel.output_def import ( FittingOutputDef, ) @@ -102,6 +106,12 @@ class DPAtomicModel(BaseAtomicModel): :math:`\mathcal{F}` is the fitting network, and :math:`\mathbf{y}^i` is the predicted atomic property (energy, dipole, etc.). + An energy model may carry a long-range module beside its fitting: it reads + the same descriptor and adds its own atomic energy to the fitting output, + so the atomic energy :math:`E^i` is the complete potential, including the + interactions without a cutoff. The module is evaluated on the + NeighborGraph route. + Parameters ---------- descriptor @@ -111,6 +121,8 @@ class DPAtomicModel(BaseAtomicModel): type_map Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. + long_range + Long-range module added to the energy output of the fitting. """ @@ -121,12 +133,14 @@ def __init__( descriptor: BaseDescriptor, fitting: BaseFitting, type_map: list[str], + long_range: BaseLongRange | None = None, **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) self.descriptor = descriptor self.fitting_net = fitting self.fitting_net.reinit_exclude(self.atom_exclude_types) + self.long_range = long_range self.type_map = type_map self.add_chg_spin_ebd: bool = self.descriptor.has_chg_spin_ebd() # Structural capability: only descriptors with a native spin @@ -148,6 +162,8 @@ def __init__( self.add_spin_ebd: bool = ( self._supports_native_spin and self.descriptor.use_spin is not None ) + if self.long_range is not None: + self._check_long_range() super().init_out_stat() # === Reference of the isolated atoms === # The fitting references its output only when the preset fixes the @@ -155,6 +171,31 @@ def __init__( resolve_vacuum_ref(self.fitting_net, self.preset_out_bias) self.init_vacuum_conditions() + def _check_long_range(self) -> None: + """Validate that the long-range module can join this descriptor and fitting.""" + check_long_range_support(self.descriptor) + if not self.descriptor.uses_graph_lower(): + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower, " + "which this descriptor does not provide" + ) + if "energy" not in self.fitting_output_def().keys(): + raise ValueError( + "a long-range module adds to the energy output, which this " + "fitting does not have" + ) + if self.atom_exclude_types: + raise NotImplementedError( + "`atom_exclude_types` is not supported with a long-range module, " + "which attributes each pair energy to both atoms of the pair" + ) + if self.long_range.uses_charge_spin() and not self.add_chg_spin_ebd: + raise ValueError( + "the long-range module reads the total charge from " + "`charge_spin`, but the descriptor is not configured with the " + "charge/spin condition (`add_chg_spin_ebd`)" + ) + def init_vacuum_conditions(self) -> None: """Build the conditioning inputs of the reference atoms for the type map. @@ -213,21 +254,45 @@ def uses_graph_lower(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.uses_graph_lower()) - def fused_decomposition(self) -> "tuple[DPAtomicModel, None]": - """A descriptor-fitting model is its own learned part.""" + def fused_decomposition(self) -> "tuple[DPAtomicModel, None] | None": + """A descriptor-fitting model is its own learned part. + + A model with a long-range module is more than a learned part and a + pair potential, so no fused pipeline evaluates it. + """ + if self.long_range is not None: + return None return self, None + def has_long_range(self) -> bool: + """Whether this model carries a long-range module.""" + return self.long_range is not None + + def required_topologies(self) -> tuple[str, ...]: + """The topologies of the long-range module.""" + if self.long_range is None: + return () + return self.long_range.required_topologies() + def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.has_message_passing_across_ranks()) def supports_edge_parallel(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.supports_edge_parallel()) + """Delegates to this model's own descriptor. + + A long-range module sums over complete systems, so a model carrying + one does not run under domain decomposition. + """ + return self.long_range is None and bool( + self.descriptor.supports_edge_parallel() + ) def dense_lower_supports_comm(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.dense_lower_supports_comm()) + """Delegates to this model's own descriptor; a long-range module has no dense lower.""" + return self.long_range is None and bool( + self.descriptor.dense_lower_supports_comm() + ) def uses_compact_edge_pairs(self) -> bool: """Delegates to this model's own descriptor.""" @@ -267,6 +332,8 @@ def set_case_embd(self, case_idx: int) -> None: typically concatenated with the output of the descriptor and fed into the fitting net. """ self.fitting_net.set_case_embd(case_idx) + if self.long_range is not None: + self.long_range.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model @@ -311,6 +378,10 @@ def enable_compression( check_frequency The overflow check frequency """ + if self.long_range is not None: + raise NotImplementedError( + "compression of a model with a long-range module is not supported" + ) self.descriptor.enable_compression( min_nbor_dist, table_extrapolate, @@ -358,6 +429,11 @@ def forward_atomic( the result dict, defined by the `FittingOutputDef`. """ + if self.long_range is not None: + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower " + "only; the dense (nlist) lower cannot evaluate this model" + ) nframes, nloc, nnei = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) @@ -575,13 +651,16 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph analogue of :meth:`forward_atomic` on the flat node axis. Runs the descriptor ``call_graph`` then the fitting ``call_graph`` PER NODE and returns the raw fitting dict on the flat ``(N, *)`` axis (no reduction or masking; the wrapper handles those). ``fparam`` is gathered to nodes by - ``frame_id`` so each node sees its frame's parameter. + ``frame_id`` so each node sees its frame's parameter. The atomic energy + of the long-range module is added to the energy output of the fitting. Parameters ---------- @@ -607,6 +686,11 @@ def forward_atomic_graph( ``call_graph`` (the message-passing part). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name. + node_mask + flat mask of the real and owned nodes, which carry the fields of + the long-range module. N Returns ------- @@ -638,6 +722,13 @@ def forward_atomic_graph( graph.n_local, atype.shape[0], ) + # The long-range module reads the frame conditions of the real frames, + # with the defaults the descriptor and fitting fall back on. + n_node = graph.n_node + if self.long_range is not None: + fparam, charge_spin_frame = self._frame_conditions( + fparam, charge_spin, n_node.shape[0], graph.edge_vec + ) # === Step 2. Vacuum reference nodes === # The reference atoms are appended as single-node frames, so the same # descriptor call yields the vacuum descriptor of every type. @@ -683,7 +774,7 @@ def forward_atomic_graph( # its type, so along the closing of a pair the learned energy follows # the switch linearly while the analytical pair term takes over. gate_kwargs = {} if node_gate is None else {"node_gate": node_gate} - return self.fitting_net.call_graph( + ret = self.fitting_net.call_graph( gg, atype, gr=rot_mat, @@ -694,6 +785,71 @@ def forward_atomic_graph( **vacuum_kwargs, **gate_kwargs, ) + # === Step 5. Long-range module === + # The module reads the descriptor of the real nodes; the reference + # nodes appended above take no part in it. + if self.long_range is not None: + if topologies is None or node_mask is None: + raise ValueError( + "the long-range module of this model reads the topologies " + f"{self.required_topologies()} and the node mask, which were " + "not given" + ) + ret["energy"] = ret["energy"] + self.long_range.call_graph( + gg, + atype, + node_mask, + n_node, + topologies, + fparam=fparam, + aparam=aparam_node, + charge_spin=charge_spin_frame, + ) + return ret + + def _frame_conditions( + self, + fparam: Array | None, + charge_spin: Array | None, + nf: int, + reference: Array, + ) -> tuple[Array | None, Array | None]: + """Frame parameters and charge/spin condition with their defaults resolved. + + Parameters + ---------- + fparam + Frame parameters with shape (nf, numb_fparam), or None. + charge_spin + Charge/spin condition with shape (nf, 2) or (1, 2), or None. + nf + The number of frames. + reference + An array giving the namespace, dtype and device of the defaults. + + Returns + ------- + tuple + The frame parameters with shape (nf, numb_fparam) and the + charge/spin condition with shape (nf, 2), each None when the model + takes no such input or it has no default. + """ + xp = array_api_compat.array_namespace(reference) + device = array_api_compat.device(reference) + if fparam is None and self.has_default_fparam(): + default_fparam = xp.asarray( + self.get_default_fparam(), dtype=reference.dtype, device=device + ) + fparam = xp.tile(xp.reshape(default_fparam, (1, -1)), (nf, 1)) + if not self.add_chg_spin_ebd: + return fparam, None + if charge_spin is None: + default_cs = self.descriptor.get_default_chg_spin() + if default_cs is None: + return fparam, None + charge_spin = xp.asarray(default_cs, dtype=reference.dtype, device=device) + charge_spin = xp.broadcast_to(xp.reshape(charge_spin, (-1, 2)), (nf, 2)) + return fparam, charge_spin def compute_or_load_stat( self, @@ -724,6 +880,10 @@ def compute_or_load_stat( self.fitting_net.compute_input_stats( wrapped_sampler, stat_file_path=stat_file_path ) + if self.long_range is not None: + self.long_range.compute_input_stats( + wrapped_sampler, stat_file_path=stat_file_path + ) self._collect_and_set_observed_type( wrapped_sampler, stat_file_path, preset_observed_type ) @@ -747,6 +907,8 @@ def change_type_map( else None, ) self.fitting_net.change_type_map(type_map=type_map) + if self.long_range is not None: + self.long_range.change_type_map(type_map=type_map) self.init_vacuum_conditions() def compute_fitting_input_stat( @@ -756,6 +918,9 @@ def compute_fitting_input_stat( ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + The networks of the long-range module take the inputs of the fitting + and receive their statistics as well. + Parameters ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] @@ -772,6 +937,12 @@ def compute_fitting_input_stat( protection=self.data_stat_protect, stat_file_path=stat_file_path, ) + if self.long_range is not None: + self.long_range.compute_input_stats( + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def serialize(self) -> dict: dd = super().serialize() @@ -785,6 +956,12 @@ def serialize(self) -> dict: "fitting": self.fitting_net.serialize(), } ) + # A record with a long-range module is version 3, which a reader that + # cannot evaluate the module rejects; every other record keeps + # version 2. + if self.long_range is not None: + dd["@version"] = 3 + dd["long_range"] = self.long_range.serialize() return dd # for subclass overridden @@ -792,17 +969,27 @@ def serialize(self) -> dict: """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls: type | None = BaseLongRange + """The base long-range module class; ``None`` for a backend without one.""" @classmethod def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 2, 2) + check_version_compatibility(data.pop("@version", 1), 3, 2) data.pop("@class") data.pop("type") descriptor_obj = cls.base_descriptor_cls.deserialize(data.pop("descriptor")) fitting_obj = cls.base_fitting_cls.deserialize(data.pop("fitting")) data["descriptor"] = descriptor_obj data["fitting"] = fitting_obj + long_range = data.pop("long_range", None) + if long_range is not None: + if cls.base_long_range_cls is None: + raise NotImplementedError( + "this backend cannot evaluate the long-range module of the " + "model record" + ) + data["long_range"] = cls.base_long_range_cls.deserialize(long_range) obj = super().deserialize(data) return obj diff --git a/deepmd/dpmodel/atomic_model/inner_potential.py b/deepmd/dpmodel/atomic_model/inner_potential.py index 4d3c51f2d7..954c34ab96 100644 --- a/deepmd/dpmodel/atomic_model/inner_potential.py +++ b/deepmd/dpmodel/atomic_model/inner_potential.py @@ -606,12 +606,15 @@ def forward_atomic_graph( charge_spin: Any = None, spin: Any = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Any = None, ) -> dict: """Evaluate the analytical per-atom energy on the flat node axis. - ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict`` are - accepted for pipeline-signature compatibility and ignored (the - analytical term conditions on geometry and types only). + ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict``/ + ``topologies``/``node_mask`` are accepted for pipeline-signature + compatibility and ignored (the analytical term conditions on + geometry and types only; the caller masks its outputs). Parameters ---------- diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index e224cbdc2b..1dd43e0d45 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -34,6 +34,9 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from .inner_potential import ( + InnerPotentialAtomicModel, +) from .pairtab_atomic_model import ( PairTabAtomicModel, ) @@ -112,6 +115,23 @@ def __init__( + ", ".join(f"{type(m).__name__}={dim_of(m)}" for m in models) ) + # A long-range module attributes each pair energy to both atoms of + # the pair. It is not combined with a bridging term, which defines the + # close-pair limit of a learned local energy, nor with atom exclusion, + # which removes atomic outputs. + if any(isinstance(m, InnerPotentialAtomicModel) for m in models) and any( + m.has_long_range() for m in models + ): + raise NotImplementedError( + "an inner_potential bridging term cannot be combined with a " + "sub-model carrying a long-range module" + ) + if self.atom_exclude_types and any(m.has_long_range() for m in models): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a sub-model carrying " + "a long-range module" + ) + self.models = models self.type_map = type_map self._rebuild_mapping_state() @@ -377,6 +397,8 @@ def compute_fitting_input_stats( for model in self.models: if isinstance(model, DPAtomicModel): model.fitting_net.compute_input_stats(sampled_func) + if model.long_range is not None: + model.long_range.compute_input_stats(sampled_func) def supports_native_spin(self) -> bool: """Spin-capable when ANY child consumes the spin input. @@ -389,6 +411,18 @@ def supports_native_spin(self) -> bool: """ return any(m.supports_native_spin() for m in self.models) + def has_long_range(self) -> bool: + """ANY child carrying a long-range module makes the composition carry one.""" + return any(m.has_long_range() for m in self.models) + + def required_topologies(self) -> tuple[str, ...]: + """Every topology a child reads, each named once.""" + return tuple( + dict.fromkeys( + name for model in self.models for name in model.required_topologies() + ) + ) + def forward_atomic_graph( self, graph: Any, @@ -398,14 +432,18 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph-route linear combination on the flat node axis. - Every child consumes the SAME graph, so on autograd backends the - shared ``graph.edge_vec`` leaf makes the summed energy's force and - virial exactly the sum of the children's -- one edge backward - covers the whole composition (this is what makes analytical - bridging terms compose with the learned model for free). + Every child consumes the SAME graph and topologies, so on autograd + backends the shared geometry leaves make the summed energy's force + and virial exactly the sum of the children's -- one backward covers + the whole composition (this is what makes analytical bridging terms + compose with the learned model for free). Each child evaluates its + complete energy, including its long-range module, before its weight + applies. Only constant weights are supported here (``"sum"``/``"mean"`` or a per-child float list); the distance-switched ZBL-interpolation @@ -429,6 +467,12 @@ def forward_atomic_graph( flat (N, 3) per-node spin, forwarded to every child. comm_dict MPI communication metadata, forwarded to every child. + topologies + the topologies of :meth:`required_topologies`, forwarded to every + child. + node_mask + flat mask of the nodes whose outputs count, forwarded to every + child, which combines it with its own exclusions. Returns ------- @@ -472,6 +516,8 @@ def forward_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=node_mask, ) contrib = ret["energy"] * ww energy = contrib if energy is None else energy + contrib diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index d15c739ff1..7dbd9c7f3d 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -1541,12 +1541,11 @@ 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) - if frame["box"] is None: - box = None - else: - box = xp.asarray(frame["box"], dtype=dtype, device=device) - if xp.all(box == 0): - box = None + box = ( + None + if frame["box"] is None + else xp.asarray(frame["box"], dtype=dtype, device=device) + ) spin = ( None if frame["spin"] is None diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 7cbaa0f931..126b51e721 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -8,15 +8,15 @@ from .dpa4_ener import ( SeZMEnergyFittingNet, ) -from .dpa4c_lr import ( - DPA4CLRFitting, -) from .ener_fitting import ( EnergyFittingNet, ) from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .make_base_fitting import ( make_base_fitting, ) @@ -29,10 +29,10 @@ __all__ = [ "DOSFittingNet", - "DPA4CLRFitting", "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/dpmodel/fitting/dpa4c_lr.py b/deepmd/dpmodel/fitting/dpa4c_lr.py deleted file mode 100644 index 59f97a5fdb..0000000000 --- a/deepmd/dpmodel/fitting/dpa4c_lr.py +++ /dev/null @@ -1,434 +0,0 @@ -# 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, - 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, - vacuum_descriptor=vacuum_descriptor, - ) - 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/fitting/latent_charge.py b/deepmd/dpmodel/fitting/latent_charge.py new file mode 100644 index 0000000000..4967879f1f --- /dev/null +++ b/deepmd/dpmodel/fitting/latent_charge.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from .invar_fitting import ( + InvarFitting, +) + + +@fitting_check_output +class LatentChargeFitting(InvarFitting): + r"""Per-atom latent charges predicted from the invariant descriptor. + + The network maps the descriptor of an atom, together with the frame + parameters, atomic parameters and case embedding of the energy fitting it + accompanies, to ``dim_out`` latent charge channels in units of the + elementary charge :math:`e`. The charges are an intermediate field of a + long-range module rather than a model output: the module constrains and + pairs them, and the potential energy is the only quantity they are + trained through. + + The constructor arguments are those of :class:`InvarFitting` with the + output variable fixed to ``latent_charge``. + """ + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def serialize(self) -> dict: + data = super().serialize() + data["type"] = "latent_charge" + return data diff --git a/deepmd/dpmodel/long_range/__init__.py b/deepmd/dpmodel/long_range/__init__.py new file mode 100644 index 0000000000..4a4aad0d7c --- /dev/null +++ b/deepmd/dpmodel/long_range/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range modules of the atomic energy of a descriptor-fitting model. + +A long-range module is a component of the atomic energy model beside its +energy fitting: it reads the descriptor features shared with that fitting and +the topologies it declares, and contributes an energy per atom. Every consumer +of the atomic model therefore evaluates the complete potential. +""" + +from .base_long_range import ( + BaseLongRange, + check_long_range_support, + make_base_long_range, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "COULOMB_CONSTANT", + "BaseInteractionKernel", + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", + "check_long_range_support", + "make_base_long_range", +] diff --git a/deepmd/dpmodel/long_range/base_long_range.py b/deepmd/dpmodel/long_range/base_long_range.py new file mode 100644 index 0000000000..9cf76eda93 --- /dev/null +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, + abstractmethod, +) +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.plugin import ( + PluginVariant, + make_plugin_registry, +) + + +def check_long_range_support(descriptor: Any) -> None: + """Admit only the descriptors whose combination with long-range modules is validated. + + A long-range module reads nothing but the invariant node features every + graph descriptor returns, so this gate is a statement of validation, not + of implementation: a further descriptor is admitted by adding it here once + its combination has been tested. + + Parameters + ---------- + descriptor + The descriptor of the atomic model receiving the long-range module. + + Raises + ------ + NotImplementedError + If the descriptor is not among the validated ones. + """ + from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, + ) + + if not isinstance(descriptor, DescrptDPA4C): + raise NotImplementedError( + "long-range modules are currently validated with the DPA4C " + f"descriptor only, got {type(descriptor).__name__}" + ) + + +def make_base_long_range(t_tensor: Any) -> type: + """Make the base class of the long-range modules of one backend. + + Each backend owns its registry, so a configuration builds the long-range + module native to the backend that reads it. + + Parameters + ---------- + t_tensor + The type of the tensor, used in the type hints. + """ + + class BLR(ABC, PluginVariant, make_plugin_registry("long range")): + """Long-range contribution to the atomic energy of a descriptor-fitting model. + + A long-range module reads the invariant node features of the + descriptor and the topologies it declares, and returns an energy per + node in eV. The atomic model adds it to the output of its energy + fitting before the output statistics and masks, so every consumer of + the atomic model sees one complete potential. The module owns the + attribution of its energy to the nodes; the per-frame energy is their + sum. + """ + + @abstractmethod + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the module reads beside the neighbor graph.""" + + @abstractmethod + def uses_charge_spin(self) -> bool: + """Whether the module reads the frame-level charge/spin condition.""" + + @abstractmethod + def call_graph( + self, + descriptor: t_tensor, + atype: t_tensor, + node_mask: t_tensor, + n_node: t_tensor, + topologies: dict[str, Any], + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, + charge_spin: t_tensor | None = None, + ) -> t_tensor: + """Evaluate the atomic energy of the module on the flat node axis. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any valid + type. + node_mask + Nodes carrying fields with shape ``(N,)``: the real and owned + nodes. + n_node + Per-frame node counts with shape ``(nf,)``. + topologies + The topologies of :meth:`required_topologies` by name. + fparam + Frame parameters with shape ``(nf, numb_fparam)``, defaults + resolved. + aparam + Atomic parameters with shape ``(N, numb_aparam)``. + charge_spin + Frame charge/spin condition with shape ``(nf, 2)``, defaults + resolved, when the model takes one. + + Returns + ------- + t_tensor + The atomic energy with shape ``(N, 1)``, in eV. + """ + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the networks of the module.""" + + def change_type_map(self, type_map: list[str]) -> None: + """Remap the type-dependent state of the module to a new type map.""" + + def set_case_embd(self, case_idx: int) -> None: + """Select the case embedding of the networks of the module.""" + + @abstractmethod + def serialize(self) -> dict: + """Serialize the module to a dict.""" + + @classmethod + def deserialize(cls, data: dict) -> "BLR": + """Deserialize a module, dispatching on its registered type.""" + if cls is BLR: + return BLR.get_class_by_type(data["type"]).deserialize(data) + raise NotImplementedError(f"Not implemented in class {cls.__name__}") + + return BLR + + +BaseLongRange = make_base_long_range(np.ndarray) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py new file mode 100644 index 0000000000..e6227aa629 --- /dev/null +++ b/deepmd/dpmodel/long_range/kernel.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Pair interaction kernels of the latent-charge long-range modules. + +A kernel is a function :math:`K(r)` of the distance between two charges, in +Å⁻¹, so that charges in units of :math:`e` interact with the energy +:math:`k_e q_i q_j K(r_{ij})` in eV, where :math:`k_e` is the Coulomb +constant. A kernel owns its parameters; widths are stored through their +logarithm, which keeps them positive however training moves them. +""" + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_asarray_nodetach, + xp_erf, +) +from deepmd.dpmodel.common import ( + PRECISION_DICT, + NativeOP, + to_numpy_array, +) +from deepmd.utils.plugin import ( + make_plugin_registry, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +COULOMB_CONSTANT = 14.3996454784255 +r"""Coulomb constant :math:`e^2 / (4 \pi \varepsilon_0)` in eV Å.""" + + +class BaseInteractionKernel(NativeOP, make_plugin_registry("interaction kernel")): + """Pair interaction kernel :math:`K(r)` in Å⁻¹ with its own parameters. + + Parameters + ---------- + trainable + Whether the parameters of the kernel are trained. + precision + The precision in which the parameters are stored; the kernel is + evaluated in the precision of its input. + """ + + PARAMETER_NAMES: tuple[str, ...] = () + """Attributes holding the trainable parameters of the kernel.""" + + def __init__(self, trainable: bool = True, precision: str = "default") -> None: + self.trainable = bool(trainable) + self.precision = precision + self.prec = PRECISION_DICT[precision.lower()] + + def call(self, r_sq: Array) -> Array: + """Evaluate the kernel on squared pair distances. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.evaluate(r_sq, self.coefficients(r_sq)) + + def coefficients(self, like: Array) -> dict[str, Array]: + """The coefficients of the kernel formula, derived from its parameters. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the coefficients. + + Returns + ------- + dict[str, Array] + Every coefficient with its own shape, differentiable with respect + to the parameters. + """ + raise NotImplementedError + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + """Evaluate the kernel formula with given coefficients. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + coefficients + The coefficients of :meth:`coefficients`, each with its own shape + or with a leading pair axis of length ``P``. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + raise NotImplementedError + + def _parameter(self, name: str, like: Array) -> Array: + """A parameter in the namespace, dtype and device of ``like``, autograd kept.""" + xp = array_api_compat.array_namespace(like) + return xp_asarray_nodetach( + xp, + getattr(self, name), + dtype=like.dtype, + device=array_api_compat.device(like), + ) + + def serialize(self) -> dict: + return { + "@class": "InteractionKernel", + "type": self.TYPE, + "@version": 1, + "trainable": self.trainable, + "precision": self.precision, + "@variables": { + name: to_numpy_array(getattr(self, name)) + for name in self.PARAMETER_NAMES + }, + } + + @classmethod + def deserialize(cls, data: dict) -> "BaseInteractionKernel": + if cls is BaseInteractionKernel: + return cls.get_class_by_type(data["type"]).deserialize(data) + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + data.pop("type") + variables = data.pop("@variables") + obj = cls(**data, **cls._shape_kwargs(variables)) + for name in cls.PARAMETER_NAMES: + setattr(obj, name, np.asarray(variables[name], dtype=obj.prec)) + return obj + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + """Constructor arguments that size the parameters stored in ``variables``.""" + return {} + + +@BaseInteractionKernel.register("les") +class LESKernel(BaseInteractionKernel): + r"""Latent Ewald summation kernel :math:`K(r) = \mathrm{erf}(\alpha r) / r`. + + The kernel is the interaction of two Gaussian charge clouds of width + :math:`1/(\sqrt{2}\alpha)`: it reaches :math:`1/r` beyond a few + :math:`1/\alpha` and stays finite, :math:`2\alpha/\sqrt{\pi}`, at contact. + + Parameters + ---------- + alpha + The inverse width :math:`\alpha > 0`, in Å⁻¹. + trainable + Whether :math:`\alpha` is trained. + precision + The precision in which :math:`\log \alpha` is stored. + """ + + TYPE = "les" + PARAMETER_NAMES = ("log_alpha",) + + def __init__( + self, alpha: float = 1.0, trainable: bool = True, precision: str = "default" + ) -> None: + super().__init__(trainable, precision) + if not np.isfinite(alpha) or alpha <= 0.0: + raise ValueError(f"`alpha` must be finite and positive, got {alpha}") + self.log_alpha = np.array(np.log(alpha), dtype=self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return {"alpha": xp.exp(self._parameter("log_alpha", like))} + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + # The softening keeps the derivative finite on a coincident pair. + r = xp.sqrt(r_sq + 1e-12) + return xp_erf(coefficients["alpha"] * r) / r + + +@BaseInteractionKernel.register("sog") +class SOGKernel(BaseInteractionKernel): + r"""Sum-of-Gaussians kernel :math:`K(r) = \sum_k a_k \exp(-r^2 / 2\sigma_k^2)`. + + Parameters + ---------- + amplitude + The amplitudes :math:`a_k`, in Å⁻¹, of either sign: a scalar shared by + every Gaussian or one value per Gaussian. Required with ``width``; + with the default widths it defaults to :math:`4\pi \ln b` for every + Gaussian. + width + The widths :math:`\sigma_k > 0`, in Å. Defaults to the geometric + sequence :math:`\sigma_k = \sigma b^k`, :math:`k = 0, \dots, M - 1`. + n_gaussian + The number :math:`M` of Gaussians of the default widths. + base_width + The width :math:`\sigma` of the first default Gaussian, in Å. + width_ratio + The ratio :math:`b > 1` of consecutive default widths. + trainable + Whether the amplitudes and widths are trained. + precision + The precision in which the amplitudes and log-widths are stored. + """ + + TYPE = "sog" + PARAMETER_NAMES = ("amplitude", "log_width") + + def __init__( + self, + amplitude: float | list[float] | None = None, + width: list[float] | None = None, + n_gaussian: int = 12, + base_width: float = 2.180230445405648, + width_ratio: float = 1.62976708826776469, + trainable: bool = True, + precision: str = "default", + ) -> None: + super().__init__(trainable, precision) + if width is None: + if n_gaussian < 1: + raise ValueError(f"`n_gaussian` must be positive, got {n_gaussian}") + if not np.isfinite(base_width) or base_width <= 0.0: + raise ValueError(f"`base_width` must be positive, got {base_width}") + if not np.isfinite(width_ratio) or width_ratio <= 1.0: + raise ValueError( + f"`width_ratio` must be greater than one, got {width_ratio}" + ) + width = base_width * np.power( + width_ratio, np.arange(n_gaussian, dtype=np.float64) + ) + if amplitude is None: + amplitude = 4.0 * np.pi * np.log(width_ratio) + elif amplitude is None: + raise ValueError("`amplitude` is required with explicit `width`") + width = np.asarray(width, dtype=np.float64).reshape(-1) + if width.size < 1 or not np.all(np.isfinite(width)) or np.any(width <= 0.0): + raise ValueError("`width` must hold finite positive values") + amplitude = np.asarray(amplitude, dtype=np.float64).reshape(-1) + if amplitude.size == 1: + amplitude = np.full_like(width, amplitude[0]) + if amplitude.shape != width.shape or not np.all(np.isfinite(amplitude)): + raise ValueError( + "`amplitude` must be a finite scalar or hold one value per Gaussian" + ) + self.amplitude = amplitude.astype(self.prec) + self.log_width = np.log(width).astype(self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return { + "amplitude": self._parameter("amplitude", like), + "inv_two_width_sq": 0.5 * xp.exp(-2.0 * self._parameter("log_width", like)), + } + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + gaussian = xp.exp(-r_sq[:, None] * coefficients["inv_two_width_sq"]) + return xp.sum(coefficients["amplitude"] * gaussian, axis=-1) + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + return {"n_gaussian": int(np.asarray(variables["log_width"]).size)} diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py new file mode 100644 index 0000000000..3e3ea6b99d --- /dev/null +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + frame_id_from_n_node, + segment_sum, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) + +CHARGE_TARGET_SOURCES = ("neutral", "charge_spin", "fparam") +"""Sources of the total charge a constrained channel sums to.""" + +CHARGE_NET_NEURON = [64, 64] +"""Default hidden widths of the charge network.""" + +CHARGE_NET_SEED_INDEX = 10_000 +"""Child index of the fitting seed that seeds a charge network configured without a seed. + +The networks of the fitting and of the descriptor draw from small child +indices of their seeds, so the charge network draws a stream of its own even +when the fitting and descriptor seeds coincide. +""" + + +@BaseLongRange.register("latent_charge") +class LatentChargeLongRange(NativeOP, BaseLongRange): + r"""Energy of interacting latent charges in complete non-periodic systems. + + A charge network predicts :math:`C` latent charge channels + :math:`q_i^c` per atom, in units of :math:`e`, from the descriptor shared + with the energy fitting; padding and atoms owned by another rank carry no + charge. The energy is the direct sum over the pairs of atoms of each frame + listed by the :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` + topology, each unordered pair once, + + .. math:: + E = k_e \sum_{i None: + if charge_net.vacuum_ref: + raise ValueError( + "the charge network reads the descriptor of the real atoms only " + "and cannot reference isolated atoms (`vacuum_ref`)" + ) + self.charge_net = charge_net + self.kernel = kernel + self.charge_constraint = self._normalize_constraint( + charge_constraint, charge_net + ) + + @staticmethod + def _normalize_constraint( + constraint: dict | None, charge_net: LatentChargeFitting + ) -> dict | None: + if constraint is None: + return None + source = str(constraint.get("source", "neutral")) + channel = int(constraint.get("channel", 0)) + fparam_index = int(constraint.get("fparam_index", 0)) + if source not in CHARGE_TARGET_SOURCES: + raise ValueError( + f"charge constraint source must be one of {CHARGE_TARGET_SOURCES}, " + f"got {source!r}" + ) + if not 0 <= channel < charge_net.dim_out: + raise ValueError( + f"constrained channel {channel} is out of range for " + f"{charge_net.dim_out} latent charge channel(s)" + ) + if source == "fparam" and not 0 <= fparam_index < charge_net.numb_fparam: + raise ValueError( + f"`fparam_index` {fparam_index} is out of range for " + f"{charge_net.numb_fparam} frame parameter(s)" + ) + return {"source": source, "channel": channel, "fparam_index": fparam_index} + + @classmethod + def from_config( + cls, + config: dict, + descriptor: Any, + fitting: Any, + fitting_seed: int | list[int] | None = None, + ) -> "LatentChargeLongRange": + """Build the module from its configuration. + + The charge network reads the descriptor and takes the conditioning + inputs of the energy fitting: its frame and atomic parameters, their + defaults and its case embedding. It follows the activation function, + skip connection and precision of the fitting unless its configuration + sets them, and its seed defaults to a child of the fitting seed. The + kernel parameters are stored in the precision of the charge network. + + Parameters + ---------- + config + The module configuration: ``channels``, ``charge_net``, ``kernel`` + and ``charge_constraint``. + descriptor + The descriptor of the model. + fitting + The energy fitting of the model. + fitting_seed + The seed of the energy fitting. + """ + charge_net_config = { + "neuron": CHARGE_NET_NEURON, + "activation_function": fitting.activation_function, + "resnet_dt": fitting.resnet_dt, + "precision": fitting.precision, + "seed": child_seed(fitting_seed, CHARGE_NET_SEED_INDEX), + **config.get("charge_net", {}), + } + kernel_config = { + "precision": charge_net_config["precision"], + **config["kernel"], + } + kernel_cls = BaseInteractionKernel.get_class_by_type(kernel_config.pop("type")) + charge_net = LatentChargeFitting( + var_name="latent_charge", + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + dim_out=int(config.get("channels", 1)), + numb_fparam=fitting.get_dim_fparam(), + numb_aparam=fitting.get_dim_aparam(), + dim_case_embd=fitting.dim_case_embd, + default_fparam=fitting.get_default_fparam(), + use_aparam_as_mask=fitting.use_aparam_as_mask, + mixed_types=descriptor.mixed_types(), + type_map=fitting.get_type_map(), + **charge_net_config, + ) + return cls( + charge_net, + kernel_cls(**kernel_config), + config.get("charge_constraint"), + ) + + def required_topologies(self) -> tuple[str, ...]: + return (ALL_PAIRS,) + + def uses_charge_spin(self) -> bool: + return ( + self.charge_constraint is not None + and self.charge_constraint["source"] == "charge_spin" + ) + + def call(self, *args: Any, **kwargs: Any) -> Array: + return self.call_graph(*args, **kwargs) + + def call_graph( + self, + descriptor: Array, + atype: Array, + node_mask: Array, + n_node: Array, + topologies: dict[str, Any], + fparam: Array | None = None, + aparam: Array | None = None, + charge_spin: Array | None = None, + ) -> Array: + xp = array_api_compat.array_namespace(descriptor) + n = descriptor.shape[0] + pairs = topologies[ALL_PAIRS] + frame_id = frame_id_from_n_node(n_node, n_total=n) + fparam_node = None if fparam is None else xp.take(fparam, frame_id, axis=0) + charge = self.charge_net.call_graph( + descriptor, atype, fparam=fparam_node, aparam=aparam + )["latent_charge"] + # The pair sum cancels between charges of opposite sign, so it runs in + # the precision of the geometry. + charge = xp.astype(charge, pairs.vec.dtype) + charge = xp.where(node_mask[:, None], charge, xp.zeros_like(charge)) + if self.charge_constraint is not None: + charge = self._constrain( + charge, node_mask, frame_id, n_node, fparam, charge_spin + ) + src = pairs.index[0] + dst = pairs.index[1] + pair_charge = xp.sum( + xp.take(charge, src, axis=0) * xp.take(charge, dst, axis=0), axis=-1 + ) + r_sq = xp.sum(pairs.vec * pairs.vec, axis=-1) + pair_energy = pair_charge * self._pair_kernel(r_sq, src, n) + half_energy = (0.5 * COULOMB_CONSTANT) * xp.where( + pairs.mask, pair_energy, xp.zeros_like(r_sq) + ) + energy = segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) + return xp.reshape(energy, (n, 1)) + + def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: + """Kernel values of the pairs. + + The pairs are located on the node axis for a backend that brings the + kernel coefficients to the pairs through it. + + Parameters + ---------- + r_sq + Squared pair distances with shape ``(P,)``, in Ų. + src + First node of every pair with shape ``(P,)``. + n + The number of nodes. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.kernel(r_sq) + + def _constrain( + self, + charge: Array, + node_mask: Array, + frame_id: Array, + n_node: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """Shift the constrained channel so that it sums to the frame total charge.""" + xp = array_api_compat.array_namespace(charge) + nf = n_node.shape[0] + source = self.charge_constraint["source"] + channel = self.charge_constraint["channel"] + if source == "charge_spin": + if charge_spin is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`charge_spin`, which was not given" + ) + target = charge_spin[:, 0] + elif source == "fparam": + if fparam is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`fparam`, which was not given" + ) + target = fparam[:, self.charge_constraint["fparam_index"]] + else: + target = xp.zeros( + (nf,), dtype=charge.dtype, device=array_api_compat.device(charge) + ) + target = xp.astype(target, charge.dtype) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + total = segment_sum(charge[:, channel], frame_id, nf) + # A frame without a charged atom has nothing to shift. + shift = (total - target) / xp.maximum(count, xp.ones_like(count)) + constrained = charge[:, channel] - xp.take(shift, frame_id, axis=0) * weight + return xp.concat( + [charge[:, :channel], constrained[:, None], charge[:, channel + 1 :]], + axis=-1, + ) + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + self.charge_net.compute_input_stats( + merged, protection=protection, stat_file_path=stat_file_path + ) + + def change_type_map(self, type_map: list[str]) -> None: + self.charge_net.change_type_map(type_map=type_map) + + def set_case_embd(self, case_idx: int) -> None: + self.charge_net.set_case_embd(case_idx) + + def serialize(self) -> dict: + return { + "@class": "LongRange", + "type": "latent_charge", + "@version": 1, + "charge_net": self.charge_net.serialize(), + "kernel": self.kernel.serialize(), + "charge_constraint": self.charge_constraint, + } + + @classmethod + def deserialize(cls, data: dict) -> "LatentChargeLongRange": + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + return cls( + LatentChargeFitting.deserialize(data["charge_net"]), + BaseInteractionKernel.deserialize(data["kernel"]), + data["charge_constraint"], + ) diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 7f3f8b9705..b5549d4adb 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -73,3 +73,5 @@ def compute_fitting_input_stats( The lazy sampled function to get data frames from different data systems. """ self.get_fitting_net().compute_input_stats(sampled_func) + if self.atomic_model.has_long_range(): + self.atomic_model.long_range.compute_input_stats(sampled_func) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index dfa64ce290..f05d1d9e09 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -46,14 +46,13 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, + PairGraph, build_neighbor_graph, build_neighbor_graph_ase, + build_topologies, compact_nodes, expand_node_values, ) -from deepmd.dpmodel.utils.nlist import ( - UNBOUNDED_NSEL, -) from deepmd.utils.path import ( DPPath, ) @@ -134,21 +133,6 @@ 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) - try: - zero_box = bool(xp_bb.all(bb == 0)) - except Exception: - # 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 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 @@ -582,12 +566,21 @@ def _call_common_graph( # atom count the mask is all true and this is a renumbering by the # identity. ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_real = xp.take(atype_flat, node_index, axis=0) + topologies = build_topologies( + self.atomic_model.required_topologies(), + xp.take(xp.reshape(cc, (n_padded, 3)), node_index, axis=0), + ng.n_node, + bb, + atype=atype_real, + pair_excl=pair_excl, + ) # OUTPUT-AGNOSTIC standard model dict (````, ``_redu``, # derivative name-holders ``None``, plus int ``mask``), like the # dense ``call_common``. ``call_lower_graph`` masks virtual atoms # (atype<0) and sets the real int mask. model_predict = self.call_lower_graph( - atype=xp.take(atype_flat, node_index, axis=0), + atype=atype_real, n_node=ng.n_node, edge_index=ng.edge_index, edge_vec=ng.edge_vec, @@ -607,6 +600,7 @@ def _call_common_graph( else None ), charge_spin=charge_spin, + topologies=topologies, ) # Public ABI is rectangular (nf, nloc, *); the lower is flat over # the real atoms. Scatter per-atom keys back onto the padded width @@ -749,6 +743,7 @@ def forward_common_atomic_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Model-level graph forward (no type cast). Analogue of the dense :meth:`forward_common_atomic`. @@ -791,6 +786,9 @@ def forward_common_atomic_graph( Per-node spin vectors, flat (N, 3), or ``None``. Forwarded unchanged to the atomic model's ``forward_common_atomic_graph`` (and, from there, the descriptor's ``call_graph``). + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -813,6 +811,7 @@ def forward_common_atomic_graph( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) return fit_output_to_model_output_graph( atomic_ret, @@ -835,6 +834,7 @@ def call_common_lower_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Graph-native PUBLIC lower (dpa1/se_atten concat-tebd, attention included). @@ -876,6 +876,9 @@ def call_common_lower_graph( Per-node spin vectors, flat (N, 3), or ``None``. Cast to the model precision alongside the other node inputs and forwarded unchanged to :meth:`forward_common_atomic_graph`. + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -889,6 +892,15 @@ def call_common_lower_graph( charge_spin=charge_spin, spin=spin, ) + if topologies: + # The pair vectors are geometry like the edge vectors. + xp = array_api_compat.array_namespace(edge_vec) + topologies = { + name: PairGraph( + pairs.index, xp.astype(pairs.vec, edge_vec.dtype), pairs.mask + ) + for name, pairs in topologies.items() + } model_predict = self.forward_common_atomic_graph( atype, n_node, @@ -901,6 +913,7 @@ def call_common_lower_graph( comm_dict=comm_dict, charge_spin=cs, spin=sp, + topologies=topologies, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -1087,24 +1100,12 @@ 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 - # 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 nnei >= UNBOUNDED_NSEL and isinstance(nall, int) and nnei > nall: - nnei = nall ret = self._format_nlist( extended_coord, nlist, - nnei, + sum(self.get_sel()), extra_nlist_sort=extra_nlist_sort, ) if not mixed_types: diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 66d32f8f8b..21c3be799b 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -47,6 +50,7 @@ atomic_model=DPAtomicModel, pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 38a4f4b9e4..da96ca32b9 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -65,6 +65,83 @@ def get_model_components( return descriptor, fitting, fitting_type +def get_long_range( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> Any | None: + """Construct the long-range module a model configuration adds to its fitting. + + The module reads the descriptor of the model, and its networks take the + conditioning inputs of the energy fitting they accompany. + + Parameters + ---------- + data : dict + The model configuration, with the optional ``long_range`` section. + descriptor : Any + The descriptor of the model. + fitting : Any + The energy fitting of the model. + long_range_base : type or None + Backend long-range registry base class; ``None`` for a backend + without long-range modules. + backend_name : str + Backend name used in error messages. + + Returns + ------- + Any or None + The backend-native long-range module, or ``None`` when none is + configured. + + Raises + ------ + NotImplementedError + If a long-range module is configured for a backend without them. + """ + config = data.get("long_range") + if config is None: + return None + if long_range_base is None: + raise NotImplementedError( + f"long-range modules are not supported in the {backend_name} backend" + ) + config = copy.deepcopy(config) + long_range_cls = long_range_base.get_class_by_type(config.pop("type")) + return long_range_cls.from_config( + config, + descriptor, + fitting, + fitting_seed=data.get("fitting_net", {}).get("seed"), + ) + + +def _long_range_kwargs( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> dict[str, Any]: + """The ``long_range`` constructor argument, absent when none is configured. + + Model classes without a long-range module keep their constructor contract. + """ + long_range = get_long_range( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ) + return {} if long_range is None else {"long_range": long_range} + + def get_standard_model( data: dict, *, @@ -72,6 +149,7 @@ def get_standard_model( fitting_base: type, model_base: type, backend_name: str, + long_range_base: type | None = None, ) -> Any: """Construct a standard model through backend registries.""" descriptor, fitting, fitting_type = get_model_components( @@ -89,6 +167,13 @@ def get_standard_model( atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), preset_out_bias=data.get("preset_out_bias"), + **_long_range_kwargs( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) @@ -103,6 +188,11 @@ def get_zbl_model( backend_name: str, ) -> Any: """Construct a ZBL model from backend-native atomic model classes.""" + if data.get("long_range") is not None: + raise NotImplementedError( + "a long-range module is not supported with the tabulated ZBL " + "potential (`use_srtab`)" + ) descriptor, fitting, fitting_type = get_model_components( data, descriptor_base=descriptor_base, @@ -141,6 +231,7 @@ def get_linear_atomic_model( pairtab_model: type, linear_atomic_model: type | None = None, descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, + long_range_base: type | None = None, ) -> Any: """Build the ``LinearEnergyAtomicModel`` composition from a config. @@ -181,6 +272,9 @@ def get_linear_atomic_model( back to the generic registry build. Backends use it to route family-specific model types (e.g. DPA4/SeZM) through their validated builders. + long_range_base : type, optional + Backend long-range registry base class, which builds the long-range + module of a child; ``None`` for a backend without long-range modules. Raises ------ @@ -325,6 +419,13 @@ def get_linear_atomic_model( fitting, type_map=sub["type_map"], preset_out_bias=preset_out_bias, + **_long_range_kwargs( + sub, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) built[i] = child else: @@ -467,6 +568,7 @@ def __init__( pairtab_model: type | None = None, zbl_model: type | None = None, linear_atomic_model: type | None = None, + long_range_base: type | None = None, ) -> None: """Store backend-native classes used by all model construction paths.""" self.descriptor_base = descriptor_base @@ -477,6 +579,7 @@ def __init__( self.pairtab_model = pairtab_model self.zbl_model = zbl_model self.linear_atomic_model = linear_atomic_model + self.long_range_base = long_range_base def get_model_components(self, data: dict) -> tuple[Any, Any, str]: """Construct descriptor and fitting objects for this backend.""" @@ -495,6 +598,17 @@ def get_standard_model(self, data: dict) -> Any: fitting_base=self.fitting_base, model_base=self.model_base, backend_name=self.backend_name, + long_range_base=self.long_range_base, + ) + + def get_long_range(self, data: dict, descriptor: Any, fitting: Any) -> Any | None: + """Construct the long-range module of a model config for this backend.""" + return get_long_range( + data, + descriptor, + fitting, + long_range_base=self.long_range_base, + backend_name=self.backend_name, ) def get_linear_atomic_model( @@ -515,6 +629,7 @@ def get_linear_atomic_model( pairtab_model=self.pairtab_model, linear_atomic_model=self.linear_atomic_model, descriptor_child_builder=descriptor_child_builder, + long_range_base=self.long_range_base, ) def get_zbl_model(self, data: dict) -> Any: diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index 5df83d91df..6fdd5744d6 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -79,6 +79,10 @@ class NSM(T_Model, NativeSpinModelKind): def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + if self.atomic_model.has_long_range(): + raise NotImplementedError( + "native spin is not supported for a model with a long-range module" + ) self.spin = spin self.ntypes_real = self.spin.ntypes_real # Per-real-type 0/1 spin gate, derived from ``use_spin`` and hence diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 8678e51398..2d86dd0698 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -118,6 +118,11 @@ def __init__( spin: Spin, ) -> None: super().__init__() + if backbone_model.atomic_model.has_long_range(): + raise NotImplementedError( + "the virtual-atom spin scheme is not supported for a model with " + "a long-range module" + ) self.backbone_model = backbone_model self.spin = spin # Spin graph-lower unsupported: carry-all routing diverges on diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index a7381e706b..2316ab0545 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -22,7 +22,6 @@ NeighborList, ) from .nlist import ( - UNBOUNDED_NSEL, apply_pair_exclusion_nlist, build_neighbor_list, extend_coord_with_ghosts, @@ -730,20 +729,6 @@ def build( extended_coord, extended_atype, mapping = extend_coord_with_ghosts( coord_normalized, atype, box, rcut ) - # 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 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 # path avoids the nloc*nall distance matrix for larger supported arrays. @@ -753,7 +738,7 @@ def build( extended_atype, nloc, rcut, - nsel, + sum(sel), pair_excl=pair_excl, ) else: @@ -762,7 +747,7 @@ def build( extended_atype, nloc, rcut, - nsel, + sel, distinguish_types=False, pair_excl=pair_excl, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 6895ce3e1d..dce69ed20b 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -5,8 +5,10 @@ ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity + edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the -``from_dense_quartet`` legacy converter), ``segment`` (mask-aware -segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly). +``from_dense_quartet`` legacy converter), ``pair_graph`` (the complete +intra-frame pair topology of interactions without a cutoff), ``segment`` +(mask-aware segment-reduction toolkit), and ``derivatives`` (edge and pair +force/virial assembly). See the design discussion wanghan-iapcm/deepmd-kit#4. """ @@ -33,6 +35,7 @@ ) from .derivatives import ( edge_force_virial, + pair_force_virial, ) from .env import ( edge_env_mat, @@ -53,6 +56,13 @@ pad_and_guard_angles, pad_and_guard_edges, ) +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + build_topologies, + unordered_pair_index, +) from .pairs import ( center_edge_pairs, ) @@ -64,8 +74,10 @@ ) __all__ = [ + "ALL_PAIRS", "GraphLayout", "NeighborGraph", + "PairGraph", "angle_padding_fraction", "angle_to_edge_sum", "angle_to_node_sum", @@ -77,6 +89,8 @@ "build_edge_csr", "build_neighbor_graph", "build_neighbor_graph_ase", + "build_pair_graph", + "build_topologies", "canonicalize_neighbor_graph", "center_edge_pairs", "compact_nodes", @@ -92,8 +106,10 @@ "node_validity_mask", "pad_and_guard_angles", "pad_and_guard_edges", + "pair_force_virial", "segment_max", "segment_mean", "segment_softmax", "segment_sum", + "unordered_pair_index", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 09111fa5d5..5e3f34ec04 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -118,3 +118,71 @@ def edge_force_virial( node_frame = frame_id_from_n_node(n_node, n_total=n_out) # (N,) frame per node virial = segment_sum(atom_virial, node_frame, nf) # (nf, 3, 3) return force, atom_virial, virial + + +def pair_force_virial( + g_p: Array, + pair_vec: Array, + pair_index: Array, + pair_mask: Array, + n_node: Array, + node_capacity: int | None = None, +) -> tuple[Array, Array, Array]: + """Assemble per-node force/atom-virial and PER-FRAME virial of a pair set. + + A pair set lists every unordered pair once (see + :class:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.PairGraph`), so its + per-pair gradient ``g_p = dE/d(pair_vec)`` acts on both endpoints: + ``pair_vec = r_src - r_dst`` gives ``F_dst += g_p`` and ``F_src -= g_p``, + as for an edge. The pair virial ``w = -g_p (x) pair_vec`` is split evenly + between the two endpoints, which keeps the atomic virial independent of + the order in which a pair is listed; its per-frame sum is the sum of + ``w`` over the pairs of that frame. + + Parameters + ---------- + g_p + (P, 3) per-pair gradient ``dE/d(pair_vec)``. + pair_vec + (P, 3) per-pair displacement ``r_src - r_dst``. + pair_index + (2, P) ``[src, dst]`` node endpoints of each pair. + pair_mask + (P,) boolean real-pair mask; guard pairs (``False``) are zeroed + before any scatter. + n_node + (nf,) per-frame node counts. + node_capacity + Size of the node axis ``N``; ``None`` takes ``sum(n_node)``. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, half of each pair virial per endpoint. + virial + (nf, 3, 3) per-frame virial. + """ + xp = array_api_compat.array_namespace(g_p) + n_out = node_capacity if node_capacity is not None else int(xp.sum(n_node)) + nf = n_node.shape[0] + if isinstance(n_out, int) and n_out == 0: + device = array_api_compat.device(g_p) + return ( + xp.zeros((0, 3), dtype=g_p.dtype, device=device), + xp.zeros((0, 3, 3), dtype=g_p.dtype, device=device), + xp.zeros((nf, 3, 3), dtype=g_p.dtype, device=device), + ) + g = g_p * xp.astype(pair_mask[:, None], g_p.dtype) + # As for edges, the modulo bounds the guard endpoints under export. + src = pair_index[0] % n_out + dst = pair_index[1] % n_out + force = segment_sum(g, dst, n_out) - segment_sum(g, src, n_out) + w_half = -0.5 * (g[:, :, None] * pair_vec[:, None, :]) # (P, 3, 3) + atom_virial = segment_sum(w_half, src, n_out) + segment_sum(w_half, dst, n_out) + # Both endpoints of a pair share a frame, so reducing the per-node virial + # over nodes equals reducing the pair virial over pairs. + 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 diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py new file mode 100644 index 0000000000..2e33cc6141 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Complete intra-frame pair topology for interactions without a cutoff. + +A module whose interaction acts between every two atoms of a system, such as +a non-periodic long-range module, cannot read its geometry from the +short-range :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph`. It +declares the :data:`ALL_PAIRS` topology instead: every unordered pair of valid +nodes sharing a frame, listed once with ``src < dst``. The pairs are generated +from the per-frame node counts alone, so their storage is the true pair count +``sum_f n_f (n_f - 1) / 2`` with neither a distance search nor padding to the +widest frame. Pairs of types excluded by the model-level pair exclusion are +left out, as they are left out of the neighbor graph. Like the edge axis of +the neighbor graph, the pair axis ends in masked guard pairs, which keep its +length a dynamic dimension of its own however few real pairs a batch holds. + +Geometry enters only through :attr:`PairGraph.vec`, which an autograd backend +makes a leaf next to the short-range ``edge_vec``; its gradient is pulled back +to the nodes by :func:`~deepmd.dpmodel.utils.neighbor_graph.derivatives.pair_force_virial`. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import array_api_compat + +from .graph import ( + frame_id_from_n_node, + pad_and_guard_edges, +) +from .segment import ( + segment_sum, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + +ALL_PAIRS = "all_pairs" +"""Name of the topology holding every unordered pair of valid nodes of a frame.""" + + +@dataclass +class PairGraph: + """Unordered pairs of nodes on the flat node axis of a NeighborGraph.""" + + index: Array + """(2, P) int64 [src, dst] with src < dst; both endpoints in one frame.""" + vec: Array + """(P, 3) float r_src - r_dst; the geometry leaf of the pair interaction.""" + mask: Array + """(P,) bool real (1) vs guard (0) pairs; guards join node 0 to itself.""" + + +def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: + """Enumerate the unordered pairs of valid nodes within every frame. + + Within a frame of ``n`` valid nodes the pairs follow the row-major order + of the strict upper triangle, ``(0, 1), (0, 2), ..., (0, n-1), (1, 2), ...``. + Row ``i`` starts at rank ``S(i) = i (2n - 1 - i) / 2``; the row of rank + ``k`` is the floor of the smaller root of ``S(i) = k``, evaluated in + float64 and corrected by one integer step against rounding. + + The pair count is data dependent, so the enumeration runs eagerly. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + node_mask + Validity of the flat nodes with shape ``(N,)``; ``None`` takes every + node as valid. Invalid nodes appear in no pair. + + Returns + ------- + Array + The ``[src, dst]`` node indices with shape ``(2, P)``, int64. + """ + xp = array_api_compat.array_namespace(n_node) + device = array_api_compat.device(n_node) + n_node = xp.astype(n_node, xp.int64) + if node_mask is None: + count = n_node + else: + frame_id = frame_id_from_n_node(n_node, n_total=node_mask.shape[0]) + count = segment_sum(xp.astype(node_mask, xp.int64), frame_id, n_node.shape[0]) + n_pair = count * (count - 1) // 2 + total = int(xp.sum(n_pair)) + if total == 0: + return xp.zeros((2, 0), dtype=xp.int64, device=device) + pair_end = xp.cumulative_sum(n_pair) + rank = xp.arange(total, dtype=xp.int64, device=device) + pair_frame = xp.astype(xp.searchsorted(pair_end, rank, side="right"), xp.int64) + local = rank - xp.take(pair_end - n_pair, pair_frame, axis=0) + n = xp.take(count, pair_frame, axis=0) + + def row_start(row: Array) -> Array: + return row * (2 * n - 1 - row) // 2 + + b = xp.astype(2 * n - 1, xp.float64) + discriminant = b * b - 8.0 * xp.astype(local, xp.float64) + discriminant = xp.where( + discriminant > 0.0, discriminant, xp.zeros_like(discriminant) + ) + row = xp.astype(xp.floor((b - xp.sqrt(discriminant)) / 2.0), xp.int64) + row = xp.where(local < row_start(row), row - 1, row) + row = xp.where(local >= row_start(row + 1), row + 1, row) + col = row + 1 + (local - row_start(row)) + first = xp.take(xp.cumulative_sum(count) - count, pair_frame, axis=0) + src = first + row + dst = first + col + if node_mask is not None: + # The ranks above count valid nodes only; map them to node positions. + position = xp.reshape(xp.nonzero(node_mask)[0], (-1,)) + position = xp.astype(position, xp.int64) + src = xp.take(position, src, axis=0) + dst = xp.take(position, dst, axis=0) + return xp.stack([src, dst]) + + +def build_pair_graph( + coord: Array, + n_node: Array, + atype: Array, + node_mask: Array | None = None, + pair_excl: PairExcludeMask | None = None, +) -> PairGraph: + """Build the :data:`ALL_PAIRS` topology of flat node coordinates. + + Parameters + ---------- + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + A differentiable input keeps :attr:`PairGraph.vec` differentiable. + n_node + Per-frame node counts with shape ``(nf,)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + pair_excl + The model-level pair exclusion, or ``None``; pairs of excluded types + are left out. + + Returns + ------- + PairGraph + The unordered pairs of valid nodes and their displacements, followed + by two masked guard pairs. + """ + xp = array_api_compat.array_namespace(coord) + index = unordered_pair_index(n_node, node_mask) + if pair_excl is not None: + keep = pair_excl.build_edge_exclude_mask(index, atype) + index = xp.take(index, xp.reshape(xp.nonzero(keep)[0], (-1,)), axis=1) + vec = xp.take(coord, index[0], axis=0) - xp.take(coord, index[1], axis=0) + index, vec, mask = pad_and_guard_edges(index, vec, None) + return PairGraph(index=index, vec=vec, mask=mask) + + +def build_topologies( + names: Iterable[str], + coord: Array, + n_node: Array, + box: Array | None = None, + node_mask: Array | None = None, + *, + atype: Array, + pair_excl: PairExcludeMask | None = None, +) -> dict[str, PairGraph]: + """Build the topologies an atomic model declares, beside its neighbor graph. + + Parameters + ---------- + names + Topology names from ``required_topologies()`` of the atomic model. + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + n_node + Per-frame node counts with shape ``(nf,)``. + box + Simulation cells, or ``None`` for non-periodic systems. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + pair_excl + The model-level pair exclusion, or ``None``. + + Returns + ------- + dict[str, PairGraph] + Each requested topology under its name; empty when none is required. + + Raises + ------ + NotImplementedError + If a cell is given: the pair topology describes complete non-periodic + systems only. + ValueError + If a name is unknown. + """ + topologies = {} + for name in names: + if name != ALL_PAIRS: + raise ValueError(f"unknown topology {name!r}; expected {ALL_PAIRS!r}") + if box is not None: + raise NotImplementedError( + "the all-pairs topology of the long-range modules describes " + "complete non-periodic systems only; periodic cells are not " + "supported (pass box=None for a non-periodic system)" + ) + topologies[name] = build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + return topologies diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 1b9d5bf40a..cff741dec2 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -24,17 +24,6 @@ ) -#: 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" diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index 8db31d6d5d..2ef1dc5eb9 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -40,6 +40,8 @@ class jax_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" def forward_common_atomic( self, diff --git a/deepmd/pd/model/model/__init__.py b/deepmd/pd/model/model/__init__.py index 59dce4f89b..94897bcfd1 100644 --- a/deepmd/pd/model/model/__init__.py +++ b/deepmd/pd/model/model/__init__.py @@ -64,6 +64,11 @@ def _get_standard_model_components( def get_standard_model(model_params: dict) -> BaseModel: + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index d2e0a09e96..b65604c40c 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -208,6 +208,11 @@ def get_linear_model(model_params: dict) -> BaseModel: "sub-model: add an `inner_potential` sub-model to the " "composition instead." ) + if sub.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) if any( sub.get("type") == "inner_potential" for sub in model_params.get("models", []) ): @@ -414,6 +419,11 @@ def get_standard_model(model_params: dict) -> BaseModel: "expands the flag), or spell the composition explicitly with " '`type: "linear_ener"` and an `inner_potential` sub-model.' ) + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt_expt/fitting/__init__.py b/deepmd/pt_expt/fitting/__init__.py index 41b79b8c03..2f42a57d48 100644 --- a/deepmd/pt_expt/fitting/__init__.py +++ b/deepmd/pt_expt/fitting/__init__.py @@ -11,15 +11,15 @@ from .dpa4_ener import ( SeZMEnergyFittingNet, ) -from .dpa4c_lr import ( - DPA4CLRFitting, -) from .ener_fitting import ( EnergyFittingNet, ) from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .polarizability_fitting import ( PolarFitting, ) @@ -33,10 +33,10 @@ __all__ = [ "BaseFitting", "DOSFittingNet", - "DPA4CLRFitting", "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/pt_expt/fitting/dpa4c_lr.py b/deepmd/pt_expt/fitting/dpa4c_lr.py deleted file mode 100644 index b8be34323f..0000000000 --- a/deepmd/pt_expt/fitting/dpa4c_lr.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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] - # 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(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/fitting/latent_charge.py b/deepmd/pt_expt/fitting/latent_charge.py new file mode 100644 index 0000000000..f4f38d2f0a --- /dev/null +++ b/deepmd/pt_expt/fitting/latent_charge.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting as LatentChargeFittingDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) + + +@torch_module +class LatentChargeFitting(LatentChargeFittingDP): + """Per-atom latent charges predicted from the invariant descriptor (pt_expt).""" diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 8cfc6f1955..b5c04eaadd 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -22,6 +22,10 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + build_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -223,6 +227,10 @@ def __init__( # Only a compressed ``.pt2`` folds its charge state into constants; a # model that reads the condition as an ordinary input needs no rebuild. self._charge_state_fold: ChargeStateFold | None = None + # Topologies the graph lower reads beside the neighbor graph; only a + # checkpoint can declare them, since no deployable archive holds a + # model that does. + self._required_topologies: tuple[str, ...] = () if self._is_pt2: self._load_pt2(model_file) @@ -664,7 +672,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._required_topologies = model.atomic_model.required_topologies() self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -718,7 +726,6 @@ 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 ( @@ -807,17 +814,8 @@ 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, + topologies: dict[str, PairGraph] | 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, @@ -834,7 +832,7 @@ def _eager_runner_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, - **long_range_kwargs, + topologies=topologies, ) return _translate_energy_keys( model_ret, @@ -2420,28 +2418,9 @@ 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: - 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 @@ -2560,22 +2539,20 @@ def _eval_model_graph( aparam_t, charge_spin_t, ) - if lr_graph is not None: + if self._required_topologies: + # Padding slots of a mixed-type batch are no atoms and join + # no pair. 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, + build_topologies( + self._required_topologies, + torch.as_tensor( + coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE + ), + n_node_t, + box_input, + node_mask=atype_t >= 0, + atype=atype_t, + pair_excl=self._model_pair_excl(), ), ) if self._is_pt2: @@ -2602,15 +2579,10 @@ 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. - ``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 + 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``); @@ -2622,23 +2594,6 @@ 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: 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) @@ -2649,7 +2604,6 @@ def _build_eval_graph( # 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() - 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 @@ -2689,7 +2643,7 @@ def _build_eval_graph( ) if box_input is not None else None, - rcut, + self._rcut, edge_dtype=edge_dtype, ) if method == "dense": @@ -2701,7 +2655,7 @@ def _build_eval_graph( coord_input, atom_types, box_input, - rcut, + self._rcut, canonicalize=True, pair_excl=pair_excl, ) @@ -2714,7 +2668,7 @@ def _build_eval_graph( coord_input, atom_types, box_input, - rcut, + self._rcut, canonicalize=True, pair_excl=pair_excl, ) @@ -2752,7 +2706,7 @@ def _build_eval_graph( cc, aa, bb, - rcut, + self._rcut, canonicalize=True, pair_excl=pair_excl, ) @@ -2764,7 +2718,7 @@ def _build_eval_graph( cc, aa, bb, - rcut, + self._rcut, canonicalize=True, pair_excl=pair_excl, ) diff --git a/deepmd/pt_expt/long_range/__init__.py b/deepmd/pt_expt/long_range/__init__.py new file mode 100644 index 0000000000..9c230575e7 --- /dev/null +++ b/deepmd/pt_expt/long_range/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", +] diff --git a/deepmd/pt_expt/long_range/base_long_range.py b/deepmd/pt_expt/long_range/base_long_range.py new file mode 100644 index 0000000000..66135372f9 --- /dev/null +++ b/deepmd/pt_expt/long_range/base_long_range.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import torch + +from deepmd.dpmodel.long_range import ( + make_base_long_range, +) + +BaseLongRange = make_base_long_range(torch.Tensor) diff --git a/deepmd/pt_expt/long_range/kernel.py b/deepmd/pt_expt/long_range/kernel.py new file mode 100644 index 0000000000..3675e6b988 --- /dev/null +++ b/deepmd/pt_expt/long_range/kernel.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP +from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP +from deepmd.pt_expt.common import ( + register_dpmodel_mapping, + to_torch_array, +) + + +def _torch_kernel(kernel_dp: type) -> type: + """Wrap a dpmodel interaction kernel as a module whose parameters are trained. + + The attributes listed in ``PARAMETER_NAMES`` are stored as + ``torch.nn.Parameter`` that require a gradient exactly when the kernel is + trainable, both at construction and when a deserialized value is assigned. + """ + + class TorchKernel(kernel_dp, torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + torch.nn.Module.__init__(self) + kernel_dp.__init__(self, *args, **kwargs) + + def __call__(self, *args: Any, **kwargs: Any) -> torch.Tensor: + return torch.nn.Module.__call__(self, *args, **kwargs) + + def forward(self, r_sq: torch.Tensor) -> torch.Tensor: + return self.call(r_sq) + + def __setattr__(self, name: str, value: Any) -> None: + if name in self.PARAMETER_NAMES and "_parameters" in self.__dict__: + if not isinstance(value, torch.nn.Parameter): + value = torch.nn.Parameter( + to_torch_array(value).detach().clone(), + requires_grad=self.trainable, + ) + if name in self._parameters: + self._parameters[name] = value + return + super().__setattr__(name, value) + + TorchKernel.__name__ = TorchKernel.__qualname__ = kernel_dp.__name__ + register_dpmodel_mapping( + kernel_dp, lambda v, cls=TorchKernel: cls.deserialize(v.serialize()) + ) + return TorchKernel + + +LESKernel = _torch_kernel(LESKernelDP) +SOGKernel = _torch_kernel(SOGKernelDP) diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py new file mode 100644 index 0000000000..2c612c03de --- /dev/null +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import torch + +from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange as LatentChargeLongRangeDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) + +# The charge network and the kernels built by the dpmodel class are converted +# to their pt_expt modules on assignment; importing them registers the +# conversions. +from deepmd.pt_expt.fitting import latent_charge as _latent_charge_fitting # noqa: F401 + +from . import kernel as _kernel # noqa: F401 +from .base_long_range import ( + BaseLongRange, +) + + +@BaseLongRange.register("latent_charge") +@torch_module +class LatentChargeLongRange(LatentChargeLongRangeDP): + """Energy of interacting latent charges in complete non-periodic systems (pt_expt).""" + + def _pair_kernel( + self, r_sq: torch.Tensor, src: torch.Tensor, n: int + ) -> torch.Tensor: + # A graph traced for a symbolic pair count plans the reduction of the + # coefficient gradients over the pairs before it sees a batch. In + # training the coefficients therefore reach the pairs through the node + # axis, which turns that reduction into a scatter over the pairs and a + # sum over the nodes, parallel in the pair count of every batch. A + # concrete pair count is reduced directly. + if not (self.training and isinstance(r_sq.shape[0], torch.SymInt)): + return self.kernel(r_sq) + coefficients = { + name: value.expand(n, *value.shape)[src] + for name, value in self.kernel.coefficients(r_sq).items() + } + return self.kernel.evaluate(r_sq, coefficients) diff --git a/deepmd/pt_expt/model/__init__.py b/deepmd/pt_expt/model/__init__.py index cdcdb9ff6c..22a1c2e7db 100644 --- a/deepmd/pt_expt/model/__init__.py +++ b/deepmd/pt_expt/model/__init__.py @@ -18,9 +18,6 @@ from .dpa4_model import ( DPA4EnergyModel, ) -from .dpa4c_lr_model import ( - DPA4CLREnergyModel, -) from .ener_model import ( EnergyModel, ) @@ -52,7 +49,6 @@ __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 deleted file mode 100644 index 117d831c85..0000000000 --- a/deepmd/pt_expt/model/dpa4c_lr_model.py +++ /dev/null @@ -1,865 +0,0 @@ -# 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 unordered-pair conventions of - :func:`_unordered_pair_force_virial`. - -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, - expand_node_values, - frame_id_from_n_node, - segment_sum, -) -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, -) -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_all_pairs_unordered_graph, - 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 _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, -) -> 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, - 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: - 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, - 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 - # ------------------------------------------------------------------ - - 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 = _unordered_pair_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``. 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 - - 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) - # 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) - 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. - # 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, - 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: 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, - 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/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index a4a497d453..3e95d33a3a 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -1,10 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Autograd assembly: graph energy -> force/virial/atom_virial via grad(E, edge_vec). -torch-only. The pure-array scatter (edge_force_virial) is shared with dpmodel; -this module supplies the single backward pass that produces g_e = dE/d(edge_vec). +torch-only. The pure-array scatters (edge_force_virial, pair_force_virial) are +shared with dpmodel; this module supplies the single backward pass that +produces the gradient of the energy with respect to every geometry leaf: the +edge vectors of the neighbor graph and the pair vectors of the topologies. """ +from collections.abc import ( + Sequence, +) + import torch from deepmd.dpmodel import ( @@ -20,8 +26,10 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, + PairGraph, edge_force_virial, frame_id_from_n_node, + pair_force_virial, segment_sum, ) from deepmd.pt.utils import ( @@ -60,6 +68,7 @@ def edge_energy_deriv( do_atomic_virial: bool = False, create_graph: bool = False, force_precision: torch.dtype | None = None, + pair_graphs: Sequence[PairGraph] = (), ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -79,6 +88,12 @@ def edge_energy_deriv( reduction. A retained graph (training / double backward) keeps the fp64 leaf dtype so training numerics are unchanged. + The same backward differentiates the energy with respect to the vectors + of every pair topology in ``pair_graphs``. Their gradients are pulled back + by the shared pair_force_virial scatter and summed in: the energy reaches + the coordinates through all its geometry leaves, so its complete force and + virial are the sum of their pullbacks. + Parameters ---------- energy @@ -109,6 +124,9 @@ def edge_energy_deriv( whether to materialize the per-atom virial (else ``None`` is returned). create_graph whether the backward retains a graph (training, for second-order grad). + pair_graphs + Pair topologies whose ``vec`` are further autograd leaves of + ``energy``, on the same node axis. Returns ------- @@ -119,9 +137,9 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - (g_e,) = torch.autograd.grad( + g_e, *g_pairs = torch.autograd.grad( energy.sum() if energy.dim() else energy, - edge_vec, + (edge_vec, *(pairs.vec for pairs in pair_graphs)), create_graph=create_graph, retain_graph=True, ) @@ -197,6 +215,19 @@ def edge_energy_deriv( force, atom_virial, virial = edge_force_virial( g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity ) + for g_p, pairs in zip(g_pairs, pair_graphs, strict=True): + pair_force, pair_atom_virial, pair_virial = pair_force_virial( + g_p, + pairs.vec, + pairs.index, + pairs.mask, + n_node, + node_capacity=force.shape[0], + ) + force = force + pair_force.to(force.dtype) + virial = virial + pair_virial.to(virial.dtype) + if do_atomic_virial: + atom_virial = atom_virial + pair_atom_virial.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -211,6 +242,7 @@ def fit_output_to_model_output_graph( n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, spin_leaf: torch.Tensor | None = None, + topologies: dict[str, PairGraph] | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -282,6 +314,10 @@ def fit_output_to_model_output_graph( ``grad(E, [edge_vec, spin])`` jointly in one call; here the two grads are separate so :func:`edge_energy_deriv`'s signature stays untouched. ``None`` (default): unchanged, no mag output. + topologies + Pair topologies read by the atomic model, by name. Their ``vec`` are + autograd leaves of ``fit_ret`` next to ``graph.edge_vec``; force and + virial include their pullbacks (see :func:`edge_energy_deriv`). Returns ------- @@ -393,6 +429,7 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, force_precision=force_precision if not create_graph else None, + pair_graphs=() if topologies is None else tuple(topologies.values()), ) # force (N, 3) -> (N, 1, 3) [flat; caller unravels at I/O boundary] ff_list.append(force.reshape(N, 1, 3)) diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 94a6f3277f..5bd6e68717 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -31,15 +31,15 @@ from deepmd.pt_expt.fitting import ( BaseFitting, ) +from deepmd.pt_expt.long_range import ( + BaseLongRange, +) from deepmd.pt_expt.model.dp_zbl_model import ( DPZBLModel, ) 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, ) @@ -78,60 +78,11 @@ pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, linear_atomic_model=LinearEnergyAtomicModel, + long_range_base=BaseLongRange, ) 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") - # 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"]["type"] != "dpa4c_lr": - raise ValueError( - "Model type 'dpa4c_lr' requires fitting type 'dpa4c_lr', but got " - f"'{data['fitting_net']['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. @@ -221,6 +172,7 @@ def get_sezm_model(data: dict) -> BaseModel: data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) descriptor, fitting, _ = _model_factory.get_model_components(data) + long_range = _model_factory.get_long_range(data, descriptor, fitting) return DPA4EnergyModel( descriptor=descriptor, fitting=fitting, @@ -228,6 +180,7 @@ def get_sezm_model(data: dict) -> BaseModel: atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, preset_out_bias=data.get("preset_out_bias"), + **({} if long_range is None else {"long_range": long_range}), ) @@ -462,6 +415,5 @@ 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 01dbb29a15..ba49568ce8 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -21,6 +21,8 @@ OutputVariableDef, ) from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + build_topologies, compact_nodes, expand_node_values, ) @@ -81,11 +83,6 @@ 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 @@ -361,7 +358,8 @@ class _WrapperForwardEnergyGraph: """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian. Given flattened LOCAL coordinates for one frame, rebuilds the carry-all - ``NeighborGraph`` (so ``edge_vec`` tracks the coordinates through + ``NeighborGraph`` and the topologies the atomic model declares (so + ``edge_vec`` and the pair vectors track the coordinates through ``autograd.functional.hessian``'s double-backward) and returns the scalar reduced-output component. Unlike the dense wrapper this differentiates w.r.t. the ``(nloc, 3)`` LOCAL coordinates directly -- the carry-all graph @@ -404,6 +402,16 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: ng = build_neighbor_graph_for_method( self.method, cc, self.atype, self.box, self.rcut, self.pair_excl ) + # The topologies follow the same coordinates, so the Hessian holds + # every geometry path of the energy and the blocks coupling them. + topologies = build_topologies( + self.model.atomic_model.required_topologies(), + coord_flat.reshape(self.nloc, 3), + ng.n_node, + self.box, + atype=self.atype.reshape(-1), + pair_excl=self.pair_excl, + ) atomic_ret = self.model.atomic_model.forward_common_atomic_graph( ng, self.atype.reshape(-1), @@ -411,6 +419,7 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: aparam=self.aparam, spin=self.spin, charge_spin=self.charge_spin, + topologies=topologies, ) # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all # graph (all nodes owned); reduced output = sum over the node axis. @@ -613,7 +622,7 @@ def forward_common_lower_graph( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, comm_dict: dict | None = None, - **kwargs: Any, + topologies: dict[str, PairGraph] | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -702,6 +711,12 @@ def forward_common_lower_graph( DescrptBlockRepformers._exchange_ghosts_graph`). Only meaningful for descriptors whose ``has_message_passing_across_ranks()`` is ``True``. + topologies + The topologies the atomic model declares by + ``required_topologies()``, built on the node axis of the + graph from complete systems. Each ``vec`` becomes an autograd + leaf next to ``edge_vec``, and force and virial include its + pullback. ``None`` (default) for models declaring none. Returns ------- @@ -718,6 +733,21 @@ def forward_common_lower_graph( # make edge_vec the autograd leaf for the energy backward edge_vec = edge_vec.detach().requires_grad_(True) + if topologies: + if comm_dict is not None: + raise NotImplementedError( + "the long-range topologies describe complete systems; " + "a domain-decomposed system is not supported" + ) + # Every topology carries a geometry leaf of its own. + topologies = { + name: PairGraph( + index=pairs.index, + vec=pairs.vec.detach().requires_grad_(True), + mask=pairs.mask, + ) + for name, pairs in topologies.items() + } if spin is not None: # second autograd leaf: force_mag = -dE/dspin (native spin). # Deliberately a SEPARATE leaf/backward from edge_vec rather @@ -757,6 +787,7 @@ def forward_common_lower_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, ) # ``forward_common_atomic_graph`` returns flat ``(N, *)`` output. # Pass directly to the flat-N transform; no rectangular reshape needed. @@ -778,6 +809,7 @@ def forward_common_lower_graph( # dynamic-edge torch.export. See fit_output_to_model_output_graph. node_capacity=atype.shape[0], n_local=n_local, + topologies=topologies, ) def _resolve_graph_method( @@ -818,10 +850,7 @@ 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(): - 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 getattr(self, "neighbor_graph_method", "dense") return None def call_common_ragged( @@ -901,6 +930,7 @@ def call_common_ragged( method = getattr(self, "neighbor_graph_method", None) if method is None: method = resolve_neighbor_graph_method("auto", coord.device) + pair_excl = getattr(self.atomic_model, "pair_excl", None) graph = build_ragged_neighbor_graph( method, coord, @@ -908,7 +938,15 @@ def call_common_ragged( n_node, box, self.get_rcut(), - getattr(self.atomic_model, "pair_excl", None), + pair_excl, + ) + topologies = build_topologies( + self.atomic_model.required_topologies(), + coord.detach(), + graph.n_node, + box, + atype=atype, + pair_excl=pair_excl, ) predict = self.forward_common_lower_graph( atype, @@ -927,6 +965,7 @@ def call_common_ragged( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) # The per-atom mask a rectangular batch carries is what tells the # loss each frame's real atom count. Nothing here is padded, so the @@ -1031,6 +1070,14 @@ def _call_common_graph( spin_flat = ( spin.reshape(n_padded, 3)[node_index] if spin is not None else None ) + topologies = build_topologies( + self.atomic_model.required_topologies(), + cc.detach().reshape(n_padded, 3)[node_index], + ng.n_node, + bb, + atype=atype_flat, + pair_excl=pair_excl, + ) model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -1048,6 +1095,7 @@ def _call_common_graph( aparam=ap_flat, spin=spin_flat, charge_spin=charge_spin, + topologies=topologies, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom # outputs over the real atoms. Scatter them back onto the padded diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9e400896cf..476a582ec0 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -72,6 +72,10 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + build_topologies, +) from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) @@ -297,6 +301,10 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: atomic model and fitting. Shared weights do not make those branches interchangeable: tasks without a preset bias disable the reference, while a folded reference no longer appends isolated atoms to the graph. + + A long-range module is part of the traced computation and of its input + schema, so its identity (that of its first parameter) closes the key; + tasks share a graph only when they share the module as well. """ descriptor_id: int = 0 try: @@ -322,7 +330,17 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: break except AttributeError: pass - return (int(model.has_spin()), descriptor_id, fitting_id, *vacuum_state) + long_range_id = 0 + long_range = getattr(getattr(model, "atomic_model", None), "long_range", None) + if long_range is not None: + long_range_id = next((id(p) for p in long_range.parameters()), id(long_range)) + return ( + int(model.has_spin()), + descriptor_id, + fitting_id, + *vacuum_state, + long_range_id, + ) # --------------------------------------------------------------------------- @@ -765,6 +783,11 @@ def _trace_and_compile_graph( autograd leaf. Their compiled lower returns ``force_mag`` and ``mask_mag`` through the same model-owned translation used by eager execution. + A model declaring topologies (``atomic_model.required_topologies()``) + receives each of them as three further inputs, ``(index, vec, mask)`` in + the declared order, ahead of the task buffers; its pair axis is a fourth + dynamic dimension, kept a distinct symbol like the other three. + Parameters ---------- model @@ -773,12 +796,6 @@ 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 @@ -905,18 +922,23 @@ def _trace_and_compile_graph( # tensors, while forward_common_lower_graph takes it after them. 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_all_pairs_unordered_graph, - ) - - s_lr_graph = build_all_pairs_unordered_graph(s_lr_coord, s_n_node) + # The trace reads only the shapes of the pair topologies, so the synthetic + # pairs join fixed nodes over a pair count distinct from the other axes. + topology_names = model.atomic_model.required_topologies() + s_topology_vals: list[torch.Tensor] = [] + if topology_names: + p_trace = _next_safe_prime(e_max + 2, _forbidden | {trace_nf, trace_N, e_max}) + for _ in topology_names: + s_topology_vals += [ + torch.arange(2 * p_trace, dtype=torch.int64, device=_trace_device) + .reshape(2, p_trace) + .remainder(trace_N), + torch.ones( + p_trace, 3, dtype=GLOBAL_PT_FLOAT_PRECISION, device=_trace_device + ), + torch.ones(p_trace, dtype=torch.bool, device=_trace_device), + ] + n_topology_vals = len(s_topology_vals) def fn( atype: torch.Tensor, @@ -933,11 +955,14 @@ 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, + *extra_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: + topology_vals = extra_vals[:n_topology_vals] + task_buf_vals = extra_vals[n_topology_vals:] + topologies = { + name: PairGraph(*topology_vals[3 * ii : 3 * ii + 3]) + for ii, name in enumerate(topology_names) + } # Patch task-specific buffers with the proxy tensors so make_fx records # them as FX placeholders (mirrors the dense ``_trace_and_compile``). originals: dict[str, torch.Tensor | None] = {} @@ -971,9 +996,7 @@ 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, + topologies=topologies, ) if spin is None: return _translate_energy_keys( @@ -1020,9 +1043,7 @@ def fn( s_aparam, s_charge_spin, s_spin, - s_lr_graph.edge_index, - s_lr_graph.edge_vec, - s_lr_graph.edge_mask, + *s_topology_vals, *task_buf_vals_trace, ) @@ -1271,8 +1292,6 @@ 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: @@ -1281,10 +1300,6 @@ 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, @@ -1587,7 +1602,6 @@ 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, ) @@ -1674,30 +1688,16 @@ def _forward_graph( aparam = aparam[node_index] 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: - # Ragged-native unordered all-pairs: no cutoff search and no - # dependence on the short-range builder policy. - if ragged: - 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_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 - 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) + # The topologies the atomic model declares are built here as well, + # outside the compiled region, on the node axis of the graph. + topologies = build_topologies( + _model.atomic_model.required_topologies(), + coord_3d if ragged else coord_3d.reshape(n_padded, 3)[node_index], + ng.n_node, + box_flat, + atype=atype_flat, + pair_excl=pair_excl, + ) # Lazy compile of the GRAPH lower (cached per structure key and mode). compiled_lower = self._compiled_lower_by_mode.get(self.training) @@ -1725,6 +1725,17 @@ def _forward_graph( edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) + # Each pair topology enters as (index, vec, mask), its vec a further + # leaf. + topology_vals = tuple( + tensor + for pairs in topologies.values() + for tensor in ( + pairs.index, + pairs.vec.detach().requires_grad_(True), + pairs.mask, + ) + ) if self._task_buf_order: try: @@ -1766,31 +1777,11 @@ def _forward_graph( aparam, charge_spin, spin, - lr_edge_index, - lr_edge_vec, - lr_edge_mask, + *topology_vals, *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 @@ -3639,8 +3630,8 @@ def _learned_part_path(model: torch.nn.Module) -> str | None: def _move_learned_key(key: str, path: str, source_path: str) -> str: - """Rewrite a descriptor or fitting-net key from one learned-part path to another.""" - for part in ("descriptor.", "fitting_net."): + """Rewrite a key of the learned part from one learned-part path to another.""" + for part in ("descriptor.", "fitting_net.", "long_range."): if f".{path}{part}" in key: return key.replace(f".{path}{part}", f".{source_path}{part}", 1) return key diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index 08a17ab45c..3b117732a6 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -271,8 +271,6 @@ 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, @@ -342,120 +340,3 @@ 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, 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, 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 - # 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/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index b3edf999ad..e5252d6825 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1832,6 +1832,11 @@ def _trace_and_export_impl( # Registry-dispatched (incl. native spin, type "native_spin"): the # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) + if model.atomic_model.has_long_range(): + raise NotImplementedError( + "a model with a long-range module cannot be exported to a " + "deployable archive; evaluate its training checkpoint (.pt)" + ) model.eval() # The vacuum reference is resolved on the export object, folded into the # fitting bias or stored as a per-type table, so the exported graph diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 068bef876f..e88c69fcf1 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -670,6 +670,11 @@ def __init__( type_map: list[str] | None = None, **kwargs: Any, ) -> None: + if kwargs.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) super().__init__( descriptor=descriptor, fitting=fitting_net, type_map=type_map, **kwargs ) diff --git a/deepmd/tf2/atomic_model/dp_atomic_model.py b/deepmd/tf2/atomic_model/dp_atomic_model.py index 2d68ba81ef..966c005b10 100644 --- a/deepmd/tf2/atomic_model/dp_atomic_model.py +++ b/deepmd/tf2/atomic_model/dp_atomic_model.py @@ -41,6 +41,8 @@ class tf2_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" @tf.autograph.experimental.do_not_convert def make_atom_mask( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index d6d723d593..88d249c221 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3143,126 +3143,6 @@ 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 ) @@ -3666,6 +3546,205 @@ def fitting_variant_type_args() -> Variant: ) +# --- Long-range module configurations: --- # +long_range_args_plugin = ArgsPlugin() +interaction_kernel_args_plugin = ArgsPlugin() + + +@interaction_kernel_args_plugin.register( + "les", + doc="The latent Ewald summation kernel erf(alpha r) / r, the interaction of " + "two Gaussian charge clouds: it reaches 1/r beyond a few 1/alpha and stays " + "finite at contact.", +) +def interaction_kernel_les() -> list[Argument]: + doc_alpha = "The inverse width alpha of the kernel, in Å^-1. Must be positive." + doc_trainable = "Whether alpha is trained." + return [ + Argument("alpha", float, optional=True, default=1.0, doc=doc_alpha), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +@interaction_kernel_args_plugin.register( + "sog", + doc="The sum-of-Gaussians kernel sum_k a_k exp(-r^2 / (2 sigma_k^2)).", +) +def interaction_kernel_sog() -> list[Argument]: + doc_amplitude = ( + "The amplitudes a_k of the Gaussians, in Å^-1, of either sign: one value " + "shared by every Gaussian or one value per Gaussian. Required with " + "`width`; with the default widths it defaults to 4 pi ln(`width_ratio`) " + "for every Gaussian." + ) + doc_width = ( + "The widths sigma_k of the Gaussians, in Å, each positive. Defaults to " + "the geometric sequence `base_width` * `width_ratio`^k for " + "k = 0, ..., `n_gaussian` - 1." + ) + doc_n_gaussian = "The number of Gaussians of the default widths." + doc_base_width = "The width of the first default Gaussian, in Å." + doc_width_ratio = "The ratio of consecutive default widths. Must exceed one." + doc_trainable = "Whether the amplitudes and widths are trained." + return [ + Argument( + "amplitude", + [float, list[float], type(None)], + optional=True, + default=None, + doc=doc_amplitude, + ), + Argument( + "width", + [list[float], type(None)], + optional=True, + default=None, + doc=doc_width, + ), + Argument("n_gaussian", int, optional=True, default=12, doc=doc_n_gaussian), + Argument( + "base_width", + float, + optional=True, + default=2.180230445405648, + doc=doc_base_width, + ), + Argument( + "width_ratio", + float, + optional=True, + default=1.6297670882677647, + doc=doc_width_ratio, + ), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +def latent_charge_net_args() -> list[Argument]: + doc_neuron = "The number of neurons in each hidden layer of the charge network." + doc_activation_function = f"The activation function in the charge network. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Defaults to the activation function of the energy fitting." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection. Defaults to the setting of the energy fitting.' + doc_precision = f"The precision of the charge network and kernel parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Defaults to the precision of the energy fitting." + doc_trainable = "Whether the parameters of the charge network are trained. A list sets every hidden layer and the output layer separately." + doc_seed = "Random seed for parameter initialization of the charge network. Defaults to a seed derived from the seed of the energy fitting." + return [ + Argument("neuron", list[int], optional=True, default=[64, 64], doc=doc_neuron), + Argument( + "activation_function", str, optional=True, doc=doc_activation_function + ), + Argument("resnet_dt", bool, optional=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, doc=doc_precision), + Argument( + "trainable", + [list[bool], bool], + optional=True, + default=True, + doc=doc_trainable, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] + + +def charge_constraint_args() -> list[Argument]: + doc_source = ( + "The source of the total charge of a frame: `neutral` for zero, " + "`charge_spin` for the charge of the frame charge/spin condition, which " + "requires a descriptor with `add_chg_spin_ebd`, or `fparam` for the frame " + "parameter at `fparam_index`." + ) + doc_fparam_index = "The column of the frame parameters holding the total charge, used when `source` is `fparam`." + doc_channel = "The latent charge channel that is constrained to the total charge; the other channels stay unconstrained." + return [ + Argument( + "source", + str, + optional=True, + default="neutral", + doc=doc_source, + extra_check=lambda x: x in ("neutral", "charge_spin", "fparam"), + extra_check_errmsg="must be one of 'neutral', 'charge_spin' or 'fparam'", + ), + Argument("fparam_index", int, optional=True, default=0, doc=doc_fparam_index), + Argument("channel", int, optional=True, default=0, doc=doc_channel), + ] + + +@long_range_args_plugin.register( + "latent_charge", + doc="Interaction of latent charges predicted from the descriptor, summed " + "directly over every pair of atoms of a complete non-periodic system. " + "Every channel interacts through the same kernel; half of each pair energy " + "is attributed to each atom.", +) +def long_range_latent_charge() -> list[Argument]: + doc_channels = "The number of latent charge channels predicted per atom." + doc_charge_net = ( + "The network predicting the latent charges, in units of the elementary " + "charge, from the descriptor; it takes the frame and atomic parameters " + "of the energy fitting." + ) + doc_kernel = "The pair interaction kernel of the latent charges, in Å^-1." + doc_charge_constraint = ( + "Shift one latent charge channel uniformly over the atoms of every frame " + "so that it sums to the total charge of the frame. Unconstrained if not " + "set." + ) + return [ + Argument("channels", int, optional=True, default=1, doc=doc_channels), + Argument( + "charge_net", + dict, + latent_charge_net_args(), + optional=True, + doc=doc_charge_net, + ), + Argument( + "kernel", + dict, + [], + [ + Variant( + "type", + interaction_kernel_args_plugin.get_all_argument(), + doc="The type of the kernel.", + ) + ], + doc=doc_kernel, + ), + Argument( + "charge_constraint", + dict, + charge_constraint_args(), + optional=True, + default=None, + doc=doc_charge_constraint, + ), + ] + + +def long_range_args() -> Argument: + doc_long_range = ( + "A long-range module added to the atomic energy of the energy fitting. " + "It reads the same descriptor and is evaluated on complete " + "non-periodic systems; a periodic cell is rejected. Currently validated " + "with the DPA4C descriptor." + ) + return Argument( + "long_range", + dict, + [], + [ + Variant( + "type", + long_range_args_plugin.get_all_argument(), + doc="The type of the long-range module.", + ) + ], + optional=True, + doc=supported_backends("pt_expt") + doc_long_range, + ) + + # --- Modifier configurations: --- # modifier_args_plugin = ArgsPlugin() @@ -3963,6 +4042,7 @@ def standard_model_args() -> Argument: default={}, doc=supported_backends("pt", "jax", "pd", "pt_expt", "tf2") + doc_info, ), + long_range_args(), *_bridging_method_args(), ], doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") @@ -4215,66 +4295,6 @@ 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, - ), - # 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, - ) - return ca - - @hybrid_model_args_plugin.register("pairwise_dprc") def pairwise_dprc() -> Argument: qm_model_args = model_args(exclude_hybrid=True) diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 5dd3623cba..71f2e85834 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -490,6 +490,18 @@ def _apply_pretrained_model_params( single_config["descriptor"] = deepcopy(pretrained_config["descriptor"]) if not random_fitting: single_config["fitting_net"] = deepcopy(pretrained_config["fitting_net"]) + # A long-range module reads the descriptor beside the fitting and + # is carried over with it; the trainability set in the input + # applies to its parts as to the descriptor and the fitting. + if "long_range" in pretrained_config: + long_range = deepcopy(pretrained_config["long_range"]) + for part in ("charge_net", "kernel"): + long_range.setdefault(part, {})["trainable"] = ( + single_config.get("long_range", {}) + .get(part, {}) + .get("trainable", True) + ) + single_config["long_range"] = long_range log.info( "Change the '%s' model configurations according to pretrained branch '%s'.", model_branch, diff --git a/doc/model/dpa4c_lr.md b/doc/model/dpa4c_lr.md deleted file mode 100644 index 7f13850ae3..0000000000 --- a/doc/model/dpa4c_lr.md +++ /dev/null @@ -1,82 +0,0 @@ -# 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** 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 - -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 9a405f8401..b3069c6e7b 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -13,8 +13,8 @@ Model dpa3 dpa4 dpa4c + long-range unimol - dpa4c_lr train-hybrid sel train-energy diff --git a/doc/model/long-range.md b/doc/model/long-range.md new file mode 100644 index 0000000000..cbb7b96875 --- /dev/null +++ b/doc/model/long-range.md @@ -0,0 +1,128 @@ +# Long-range latent charges {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch Exportable {{ pytorch_icon }} (`dp --pt-expt`) + +A short-range model sees each atom only through the neighbors within its +cutoff. Interactions that decay slowly with distance, such as those between +charged groups of a molecule, therefore escape it. The long-range module adds +such an interaction to the energy of a standard model. A second network reads +the same descriptor as the energy fitting and predicts per-atom _latent +charges_, and every pair of atoms of a system interacts through those charges: + +```math +E = \sum_i E_i^{\text{SR}} + k_e \sum_{i [!WARNING] +> The pair sum covers complete non-periodic systems: frames without a cell, +> such as isolated molecules and clusters. A periodic cell is rejected. The +> cost of the sum grows with the square of the number of atoms of a frame. + +## Configuration + +The module is a `long_range` section of a standard energy model, beside the +descriptor and the fitting. It is currently available with the +[DPA4C](./dpa4c.md) descriptor: + +```json +"model": { + "type_map": ["C", "H", "O"], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0 + }, + "fitting_net": { + "neuron": [128, 128, 128] + }, + "long_range": { + "type": "latent_charge", + "channels": 1, + "kernel": { + "type": "les", + "alpha": 1.0 + }, + "charge_constraint": { + "source": "neutral" + } + } +} +``` + +The charge network, +{ref}`charge_net `, +takes the frame and atomic parameters of the energy fitting and follows its +activation function, skip connection and precision unless they are set; its +seed defaults to one derived from the seed of the fitting. +Its width is set by +{ref}`neuron `. + +### Kernels + +{ref}`kernel ` selects the +interaction of the charges. Both kernels are trained with the model unless +their `trainable` option is false, and their widths stay positive throughout +training. + +- `les`: $K(r) = \mathrm{erf}(\alpha r)/r$, the interaction of two Gaussian + charge clouds. It approaches $1/r$ beyond a few $1/\alpha$ and stays finite, + $2\alpha/\sqrt{\pi}$, at contact. The inverse width $\alpha$ is set by + {ref}`alpha `. +- `sog`: $K(r) = \sum_k a_k \exp(-r^2/2\sigma_k^2)$, a sum of Gaussians whose + amplitudes $a_k$ may take either sign. The widths are given by + {ref}`width ` + or default to a geometric sequence, and the amplitudes by + {ref}`amplitude `, + which explicit widths require. + +Several charge channels, set by +{ref}`channels `, let +the model represent interactions that one charge per atom cannot. + +### Total charge + +{ref}`charge_constraint ` +shifts one channel uniformly over the atoms of every frame so that it sums to +the total charge of the frame. The total charge is zero (`neutral`), the +charge of the frame charge/spin condition of the descriptor (`charge_spin`), +or a column of the frame parameters (`fparam`); a frame without the input +takes its default. The constrained channel then carries the total charge, +while the other channels remain free. Without the section, no channel is +constrained. + +### Exclusions + +Pairs of the types in `pair_exclude_types` see each other neither in the +descriptor nor in the long-range interaction, while the `exclude_types` of +the descriptor act on the descriptor alone. `atom_exclude_types` is not +supported with a long-range module. + +## Training and inference + +Training uses the usual energy loss; see +[training energy models](train-energy.md). With `enable_compile`, the +long-range interaction is part of the compiled model: the pairs of each frame +are listed outside the compiled region and enter it as inputs of dynamic +length. + +A trained checkpoint (`.pt`) is evaluated by `dp test` and the Python +inference interface. The following are not supported for a model with a +long-range module: freezing to `.pt2` or `.pte`, compression, multi-rank +(domain-decomposed) inference, spin models, the tabulated ZBL potential, +and short-range bridging. diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py new file mode 100644 index 0000000000..9a7ee1b770 --- /dev/null +++ b/source/tests/common/dpmodel/test_long_range.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Primitives of the long-range modules: pair topology, kernels and constraint.""" + +import itertools + +import numpy as np +import pytest +from scipy.special import ( + erf, +) + +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.long_range import ( + BaseInteractionKernel, + LatentChargeLongRange, + LESKernel, + SOGKernel, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + build_pair_graph, + build_topologies, + pair_force_virial, + unordered_pair_index, +) + + +def _brute_force_pairs(counts: list[int], mask: np.ndarray | None) -> np.ndarray: + pairs = [] + start = 0 + for count in counts: + nodes = [ + node for node in range(start, start + count) if mask is None or mask[node] + ] + pairs += itertools.combinations(nodes, 2) + start += count + return np.asarray(pairs, dtype=np.int64).reshape(-1, 2).T + + +@pytest.mark.parametrize( + "counts", [[0], [1], [2], [3, 0, 1, 5], [7, 200, 1, 0, 33], [1000]] +) +@pytest.mark.parametrize("masked", [False, True]) +def test_pair_index_matches_brute_force(counts: list[int], masked: bool) -> None: + """Every unordered pair of valid nodes of a frame appears once, in row-major order.""" + mask = np.random.default_rng(0).random(sum(counts)) > 0.3 if masked else None + index = unordered_pair_index(np.asarray(counts), mask) + np.testing.assert_array_equal(index, _brute_force_pairs(counts, mask)) + + +def test_pair_graph_guards_and_periodic_rejection() -> None: + """The pair axis ends in masked guards; a cell is rejected.""" + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + pairs = build_pair_graph(coord, np.array([4]), atype) + assert pairs.index.shape == (2, 8) + assert pairs.mask.tolist() == [True] * 6 + [False] * 2 + np.testing.assert_allclose( + pairs.vec[:6], coord[pairs.index[0, :6]] - coord[pairs.index[1, :6]] + ) + assert build_topologies((), coord, np.array([4]), atype=atype) == {} + with pytest.raises(NotImplementedError, match="non-periodic"): + build_topologies( + ("all_pairs",), coord, np.array([4]), np.eye(3)[None], atype=atype + ) + + +def test_pair_force_virial_of_a_central_potential() -> None: + """For E = sum 1/r the assembled force is Newtonian and the virial sums.""" + coord = np.array([[0.0, 0.0, 0.0], [1.0, 0.2, 0.0], [0.0, 1.5, 0.3]]) + pairs = build_pair_graph(coord, np.array([3]), np.array([0, 1, 0])) + # d(1/r)/d(vec) on the real pairs; the zero-length guards carry no gradient. + r = np.where(pairs.mask, np.linalg.norm(pairs.vec, axis=-1), 1.0) + grad = np.where(pairs.mask[:, None], -pairs.vec / r[:, None] ** 3, 0.0) + force, atom_virial, virial = pair_force_virial( + grad, pairs.vec, pairs.index, pairs.mask, np.array([3]) + ) + expected = np.zeros((3, 3)) + for i, j in itertools.permutations(range(3), 2): + d = coord[i] - coord[j] + expected[i] += d / np.linalg.norm(d) ** 3 + np.testing.assert_allclose(force, expected, rtol=1e-12) + np.testing.assert_allclose(atom_virial.sum(0), virial[0], rtol=1e-12) + + +def test_les_kernel_closed_form() -> None: + """LES is erf(alpha r)/r, finite at contact.""" + kernel = LESKernel(alpha=0.7) + r = np.array([0.0, 0.5, 2.0, 9.0]) + values = kernel(r**2) + np.testing.assert_allclose(values[1:], erf(0.7 * r[1:]) / r[1:], rtol=1e-10) + np.testing.assert_allclose(values[0], 2 * 0.7 / np.sqrt(np.pi), rtol=1e-10) + + +def test_sog_kernel_closed_form_and_defaults() -> None: + """SOG is sum_k a_k exp(-r^2 / 2 sigma_k^2); the defaults follow the geometric widths.""" + kernel = SOGKernel(amplitude=[0.7, -0.2], width=[0.8, 1.7]) + r = np.array([0.0, 1.1, 3.0]) + expected = 0.7 * np.exp(-(r**2) / (2 * 0.8**2)) - 0.2 * np.exp( + -(r**2) / (2 * 1.7**2) + ) + np.testing.assert_allclose(kernel(r**2), expected, rtol=1e-12) + default = SOGKernel(n_gaussian=3, base_width=2.0, width_ratio=1.5) + np.testing.assert_allclose(np.exp(default.log_width), [2.0, 3.0, 4.5]) + np.testing.assert_allclose(default.amplitude, 4 * np.pi * np.log(1.5)) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"type": "les", "alpha": 0.0}, + {"type": "sog", "width": [1.0, -2.0]}, + {"type": "sog", "width": [1.0, 2.0], "amplitude": [1.0, 2.0, 3.0]}, + {"type": "sog", "width": [1.0, 2.0]}, + {"type": "sog", "width_ratio": 1.0}, + {"type": "sog", "amplitude": 1.0, "width_ratio": 0.5}, + ], +) +def test_kernel_parameter_domain(kwargs: dict) -> None: + """Widths must be positive and increase by default; amplitudes match the Gaussians.""" + kernel_cls = BaseInteractionKernel.get_class_by_type(kwargs.pop("type")) + with pytest.raises(ValueError): + kernel_cls(**kwargs) + + +@pytest.mark.parametrize("kernel", [LESKernel(alpha=1.3), SOGKernel(n_gaussian=4)]) +def test_kernel_round_trip(kernel: BaseInteractionKernel) -> None: + """A kernel record restores the same parameters.""" + restored = BaseInteractionKernel.deserialize(kernel.serialize()) + r_sq = np.linspace(0.0, 30.0, 7) + np.testing.assert_allclose(restored(r_sq), kernel(r_sq), rtol=0, atol=0) + + +def _long_range(constraint: dict | None, channels: int = 2) -> LatentChargeLongRange: + charge_net = LatentChargeFitting( + "latent_charge", 2, 4, channels, neuron=[4], numb_fparam=1, seed=1 + ) + return LatentChargeLongRange(charge_net, LESKernel(), constraint) + + +def test_constraint_shifts_one_channel_over_the_charged_atoms() -> None: + """The constrained channel sums to the target over the charged atoms only.""" + long_range = _long_range({"source": "fparam", "channel": 1, "fparam_index": 0}) + rng = np.random.default_rng(2) + charge = rng.normal(size=(7, 2)) + node_mask = np.array([True, True, False, True, True, True, False]) + charge[~node_mask] = 0.0 + n_node = np.array([4, 3]) + frame_id = np.array([0, 0, 0, 0, 1, 1, 1]) + fparam = np.array([[1.5], [-2.0]]) + constrained = long_range._constrain( + charge, node_mask, frame_id, n_node, fparam, None + ) + np.testing.assert_array_equal(constrained[:, 0], charge[:, 0]) + np.testing.assert_array_equal(constrained[~node_mask], 0.0) + totals = np.bincount(frame_id, weights=constrained[:, 1], minlength=2) + np.testing.assert_allclose(totals, [1.5, -2.0], rtol=1e-12) + # A frame without charged atoms is left as it is. + empty = long_range._constrain( + charge * 0.0, + np.zeros(7, dtype=bool), + frame_id, + n_node, + fparam, + None, + ) + np.testing.assert_array_equal(empty, 0.0) + + +@pytest.mark.parametrize( + "constraint", + [ + {"source": "total"}, + {"source": "neutral", "channel": 2}, + {"source": "fparam", "fparam_index": 1}, + ], +) +def test_constraint_configuration_is_validated(constraint: dict) -> None: + """Unknown sources and out-of-range channels or columns are rejected.""" + with pytest.raises(ValueError): + _long_range(constraint) + + +def test_long_range_round_trip() -> None: + """The module record holds the charge network, kernel and constraint.""" + long_range = _long_range({"source": "neutral"}) + record = long_range.serialize() + restored = LatentChargeLongRange.deserialize(record) + assert restored.serialize()["charge_constraint"] == record["charge_constraint"] + assert restored.required_topologies() == ("all_pairs",) + assert not restored.uses_charge_spin() diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 9b61ee26cf..061ada901a 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -606,6 +606,9 @@ def has_default_chg_spin(self) -> bool: def get_default_chg_spin(self): return None + def has_long_range(self) -> bool: + return False + def _compose(self, children): return LinearEnergyAtomicModel(children, type_map=["Ni", "O"], weights="sum") diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py index 33547b34d0..1f2bd45116 100644 --- a/source/tests/common/test_finetune_utils.py +++ b/source/tests/common/test_finetune_utils.py @@ -156,6 +156,28 @@ def test_finetune_rule_builder_updates_single_task_config_preserving_trainable() assert rule.get_has_new_type() +def test_finetune_rule_builder_carries_long_range_preserving_trainable(): + long_range = {"type": "latent_charge", "kernel": {"type": "les", "alpha": 0.7}} + pretrained = _model_config(["O", "H"]) + pretrained["long_range"] = deepcopy(long_range) + target = _model_config(["O", "H"], trainable=False) + target["long_range"] = { + "type": "latent_charge", + "charge_net": {"trainable": False}, + "kernel": {"type": "les", "alpha": 2.0, "trainable": False}, + } + + updated, _ = finetune.FinetuneRuleBuilder( + pretrained, + target, + change_model_params=True, + ).build() + + assert updated["long_range"]["kernel"]["alpha"] == 0.7 + assert updated["long_range"]["kernel"]["trainable"] is False + assert updated["long_range"]["charge_net"]["trainable"] is False + + def test_finetune_rule_builder_random_fitting_keeps_target_fitting_net(): pretrained = _model_config(["O", "H"], descriptor_sel=[8, 16], fitting_neuron=[32]) target = _model_config(["O", "H"], descriptor_sel=[1, 1], fitting_neuron=[2]) 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 5d0bd8e957..d8dc85a43e 100644 --- a/source/tests/pt_expt/model/test_dpa2_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py @@ -146,24 +146,6 @@ 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 @@ -707,27 +689,8 @@ 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, - lr_ei, - lr_ev, - lr_em, + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None ) eager = model.forward_common_lower_graph( atype, @@ -1014,27 +977,8 @@ 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, - lr_ei, - lr_ev, - lr_em, + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None ) eager = model.forward_common_lower_graph( atype, @@ -1139,27 +1083,8 @@ 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, - lr_ei, - lr_ev, - lr_em, + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None ) ctol = {"rtol": 1e-10, "atol": 1e-10} torch.testing.assert_close(compiled_out["energy"], ref["energy_redu"], **ctol) diff --git a/source/tests/pt_expt/model/test_dpa4c_lr.py b/source/tests/pt_expt/model/test_dpa4c_lr.py deleted file mode 100644 index bf5509ddd2..0000000000 --- a/source/tests/pt_expt/model/test_dpa4c_lr.py +++ /dev/null @@ -1,798 +0,0 @@ -# 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.dpmodel.utils.neighbor_graph import ( - edge_force_virial, -) -from deepmd.infer import ( - DeepPot, -) -from deepmd.pt_expt.fitting.dpa4c_lr import ( - DPA4CLRFitting, -) -from deepmd.pt_expt.model.dpa4c_lr_model import ( - E2_PER_ANGSTROM_TO_EV, - _les_kernel_from_squared_distance, - _SOGKernel, - _unordered_pair_force_virial, -) -from deepmd.pt_expt.model.get_model import ( - get_model, -) -from deepmd.pt_expt.train.wrapper import ( - ModelWrapper, -) -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, -) - - -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) - - 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 - - -@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() - 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"]) - - -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 - ) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py new file mode 100644 index 0000000000..697e3bb68f --- /dev/null +++ b/source/tests/pt_expt/model/test_long_range.py @@ -0,0 +1,673 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range latent-charge modules on the DPA4C graph lower (pt_expt).""" + +import copy + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + build_pair_graph, +) +from deepmd.infer import ( + DeepPot, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + model_change_out_bias, +) +from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ..compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + +KERNELS = { + "les": {"type": "les", "alpha": 0.8}, + "sog": {"type": "sog", "amplitude": [0.7, -0.2, 0.1], "width": [0.8, 1.7, 3.2]}, +} +CHANNELS = {"les": 1, "sog": 3} +NEUTRAL = {"source": "neutral"} + + +def _config(kernel: str = "les", constraint: dict | None = NEUTRAL) -> dict: + """Small float64 DPA4C model with a latent-charge long-range module.""" + return { + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 17, + }, + "fitting_net": { + "type": "ener", + "neuron": [16, 16], + "precision": "float64", + "seed": 19, + }, + "long_range": { + "type": "latent_charge", + "channels": CHANNELS[kernel], + "charge_net": {"neuron": [8, 8], "precision": "float64", "seed": 23}, + "kernel": copy.deepcopy(KERNELS[kernel]), + "charge_constraint": copy.deepcopy(constraint), + }, + } + + +def _model(kernel: str = "les", **kwargs: object) -> torch.nn.Module: + return get_model(_config(kernel, **kwargs)).to(env.DEVICE).eval() + + +def _tensor(x: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tensor: + return torch.tensor(x, dtype=dtype, device=env.DEVICE) + + +def _batch(seed: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Three frames of 6, 3 and 5 real atoms, padded to a common width.""" + rng = np.random.default_rng(seed) + coord = rng.uniform(0.0, 6.0, (3, 6, 3)) + atype = np.array( + [[0, 1, 0, 1, 1, 0], [1, 0, 0, -1, -1, -1], [0, 0, 1, 1, 0, -1]], + dtype=np.int64, + ) + return _tensor(coord), _tensor(atype, torch.int64) + + +def _set_constant_charges(model: torch.nn.Module, charges: list[float]) -> None: + """Zero every network weight and give the atoms of each type a fixed charge. + + The energy is then the long-range pair sum alone, of charges that do not + depend on the geometry; the kernel keeps its configured parameters. + """ + with torch.no_grad(): + for name, parameter in model.atomic_model.named_parameters(): + if ".kernel." not in f".{name}": + parameter.zero_() + bias = model.atomic_model.long_range.charge_net.bias_atom_e + bias.zero_() + bias[:, 0] = torch.as_tensor(charges, dtype=bias.dtype) + + +def _reference_energy(coord: np.ndarray, charge: np.ndarray, kernel: str) -> float: + """Direct pair sum of the latent-charge energy of one frame, in eV.""" + energy = 0.0 + for i in range(len(coord)): + for j in range(i + 1, len(coord)): + r = np.linalg.norm(coord[i] - coord[j]) + if kernel == "les": + from scipy.special import ( + erf, + ) + + k = erf(0.8 * r) / r + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k = np.sum(amplitude * np.exp(-(r**2) / (2 * width**2))) + energy += COULOMB_CONSTANT * float(charge[i] @ charge[j]) * k + return energy + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_energy_matches_direct_pair_sum(kernel: str) -> None: + """With fixed charges the energy is the documented pair sum, attributed half per atom.""" + model = _model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord = np.array([[[0.0, 0.0, 0.0], [1.3, 0.0, 0.0], [0.0, 2.1, 0.4]]]) + atype = np.array([[0, 1, 1]]) + out = model(_tensor(coord), _tensor(atype, torch.int64)) + charge = np.zeros((3, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, -0.5] + expected = _reference_energy(coord[0], charge, kernel) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12) + np.testing.assert_allclose( + out["atom_energy"].sum().item(), expected, rtol=1e-12, atol=1e-14 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_force_virial_hessian_match_finite_differences(kernel: str) -> None: + """Force, affine-strain virial and Hessian are derivatives of the complete energy.""" + model = _model(kernel) + # Atoms within the cutoff of each other make the charges depend on the + # geometry, so the derivatives hold the response of the charges. + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 2.5, (1, 4, 3)) + atype = np.array([[0, 1, 0, 1]]) + coord_t, atype_t = _tensor(coord), _tensor(atype, torch.int64) + out = model(coord_t, atype_t) + + def energy(c: np.ndarray) -> float: + return model(_tensor(c), atype_t)["energy"].item() + + h = 1e-5 + force_fd = np.zeros((4, 3)) + for i in range(4): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus) - energy(minus)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-8 + ) + + # The virial is -dE/d(strain) at zero strain, for x -> (1 + eps) x. + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain) + minus = coord @ (np.eye(3) - strain) + virial_fd[a, b] = -(energy(plus) - energy(minus)) / (2 * h) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-8 + ) + + hessian_model = _model(kernel) + hessian_model.enable_hessian() + hessian = hessian_model(coord_t, atype_t)["hessian"].detach().cpu().numpy() + hessian = hessian.reshape(12, 12) + hessian_fd = np.zeros((12, 12)) + for k in range(12): + plus, minus = coord.copy(), coord.copy() + plus.reshape(-1)[k] += h + minus.reshape(-1)[k] -= h + f_plus = model(_tensor(plus), atype_t)["force"].detach().cpu().numpy() + f_minus = model(_tensor(minus), atype_t)["force"].detach().cpu().numpy() + hessian_fd[:, k] = -(f_plus - f_minus).reshape(-1) / (2 * h) + np.testing.assert_allclose(hessian, hessian_fd, atol=1e-6) + np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_atomic_reductions_and_padding(kernel: str) -> None: + """Atomic energy and virial sum to the frame values; padding atoms are inert.""" + model = _model(kernel) + coord, atype = _batch() + out = model(coord, atype, do_atomic_virial=True) + torch.testing.assert_close(out["atom_energy"].sum(1), out["energy"]) + torch.testing.assert_close(out["atom_virial"].sum(1), out["virial"]) + assert out["force"][1, 3:].abs().max() == 0 + moved = coord.clone() + moved[1, 3:] += 50.0 + moved_out = model(moved, atype) + torch.testing.assert_close(moved_out["energy"], out["energy"], rtol=0, atol=0) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_frames_are_independent(kernel: str) -> None: + """A frame keeps its energy and force whatever other frames are batched with it.""" + model = _model(kernel) + coord, atype = _batch() + batched = model(coord, atype) + for frame, nloc in enumerate((6, 3, 5)): + alone = model(coord[frame : frame + 1, :nloc], atype[frame : frame + 1, :nloc]) + torch.testing.assert_close(alone["energy"][0], batched["energy"][frame]) + torch.testing.assert_close( + alone["force"][0], batched["force"][frame, :nloc], atol=1e-12, rtol=0 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_ragged_matches_rectangular(kernel: str) -> None: + """The flat ragged entry evaluates the same model as the padded one.""" + model = _model(kernel) + coord, atype = _batch() + padded = model(coord, atype) + keep = atype.reshape(-1) >= 0 + ragged = model.forward_ragged( + coord.reshape(-1, 3)[keep], + atype.reshape(-1)[keep], + torch.tensor([6, 3, 5], device=env.DEVICE), + ) + torch.testing.assert_close(ragged["energy"], padded["energy"]) + torch.testing.assert_close(ragged["force"], padded["force"].reshape(-1, 3)[keep]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_empty_and_single_atom_frames(kernel: str) -> None: + """Frames without pairs contribute no long-range energy and stay finite.""" + model = _model(kernel) + coord = _tensor(np.random.default_rng(5).uniform(0.0, 3.0, (3, 2, 3))) + atype = _tensor(np.array([[0, -1], [1, 0], [-1, -1]]), torch.int64) + out = model(coord, atype) + assert torch.isfinite(out["energy"]).all() + assert torch.isfinite(out["force"]).all() + single = model(coord[:1, :1], atype[:1, :1]) + torch.testing.assert_close(single["energy"][0], out["energy"][0]) + + +def test_permutation_and_rigid_motion_invariance() -> None: + """The energy is invariant under atom permutation, translation and rotation.""" + model = _model("sog") + coord, atype = _batch() + coord, atype = coord[:1], atype[:1] + out = model(coord, atype) + perm = torch.tensor([3, 0, 5, 1, 4, 2], device=env.DEVICE) + permuted = model(coord[:, perm], atype[:, perm]) + torch.testing.assert_close(permuted["energy"], out["energy"]) + torch.testing.assert_close(permuted["force"], out["force"][:, perm]) + rotation, _ = torch.linalg.qr( + _tensor(np.random.default_rng(13).normal(size=(3, 3))) + ) + moved = model(coord @ rotation.T + 3.7, atype) + torch.testing.assert_close(moved["energy"], out["energy"]) + torch.testing.assert_close(moved["force"], out["force"] @ rotation.T) + + +@pytest.mark.parametrize("source", ["charge_spin", "fparam"]) +def test_constraint_default_matches_explicit_target(source: str) -> None: + """An omitted frame condition takes its default for the constraint as well.""" + config = _config("les", constraint={"source": source, "fparam_index": 0}) + if source == "charge_spin": + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [2.0, 1.0] + explicit = {"charge_spin": _tensor(np.array([[2.0, 1.0], [2.0, 1.0]]))} + other = {"charge_spin": _tensor(np.array([[0.0, 1.0], [0.0, 1.0]]))} + else: + config["fitting_net"]["numb_fparam"] = 1 + config["fitting_net"]["default_fparam"] = [2.0] + explicit = {"fparam": _tensor(np.array([[2.0], [2.0]]))} + other = {"fparam": _tensor(np.array([[0.0], [0.0]]))} + model = get_model(config).to(env.DEVICE).eval() + coord, atype = _batch() + coord, atype = coord[:2], atype[:2] + implicit_out = model(coord, atype) + explicit_out = model(coord, atype, **explicit) + torch.testing.assert_close(implicit_out["energy"], explicit_out["energy"]) + torch.testing.assert_close(implicit_out["force"], explicit_out["force"]) + other_out = model(coord, atype, **other) + assert (other_out["energy"] - implicit_out["energy"]).abs().max() > 1e-6 + + +@pytest.mark.parametrize("constraint", [None, NEUTRAL]) +def test_pair_exclusion_removes_the_interaction(constraint: dict | None) -> None: + """Pairs of excluded types do not interact; their atoms keep their charges.""" + config = _config("les", constraint=constraint) + config["pair_exclude_types"] = [[0, 1]] + model = get_model(config).to(env.DEVICE).eval() + _set_constant_charges(model, [1.0, -1.0]) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + out = model(coord, atype) + # Only the two type-0 atoms interact. The neutral constraint shifts the + # charges (1, 1, -1) of the three atoms to (2/3, 2/3, -4/3). + charge = 1.0 if constraint is None else 2.0 / 3.0 + erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item() + expected = COULOMB_CONSTANT * charge**2 * erf / 4.0 + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12) + assert out["force"][0, 2].abs().max() == 0 + + +@pytest.mark.parametrize("compiled", [False, True]) +def test_vacuum_reference_with_long_range(compiled: bool) -> None: + """An isolated atom keeps its reference energy; reference rows join no pair.""" + config = _config("les", constraint={"source": "charge_spin"}) + config["type_map"] = ["O", "H"] + config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}} + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [0.0, 1.0] + config["fitting_net"]["vacuum_ref"] = True + model = get_model(config).to(env.DEVICE).eval() + assert model.get_fitting_net().needs_vacuum_descriptor() + rng = np.random.default_rng(11) + coord = _tensor(rng.uniform(0.0, 2.5, (2, 4, 3))) + atype = _tensor(np.array([[0, 1, 1, 0], [1, 0, -1, -1]]), torch.int64) + if compiled: + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model(coord, atype) + out = runner(coord, atype) + torch.testing.assert_close(out["energy"], reference["energy"]) + torch.testing.assert_close(out["force"], reference["force"]) + return + out = model(coord, atype) + for frame in range(2): + alone = model(coord[frame : frame + 1], atype[frame : frame + 1]) + torch.testing.assert_close(alone["energy"][0], out["energy"][frame]) + # The output statistics set the bias; an atom without neighbors and in + # its reference state then has exactly the bias of its type. + bias = model.atomic_model.out_bias.clone() + bias[0, :, 0] = torch.tensor([-3.0, 0.5]) + model.atomic_model.set_out_bias(bias) + isolated = model( + coord[:1, :1], + atype[:1, :1], + charge_spin=_tensor(np.array([[0.0, 3.0]])), + ) + torch.testing.assert_close( + isolated["energy"].reshape(()), + torch.tensor(-3.0, dtype=torch.float64, device=env.DEVICE), + rtol=0, + atol=1e-12, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_self_label_bias_adjustment_is_invariant(kernel: str) -> None: + """Labelling a sample with the model's own energies leaves its bias unchanged.""" + model = _model(kernel) + with torch.no_grad(): + model.atomic_model.long_range.charge_net.bias_atom_e[:, 0] = torch.tensor( + [0.8, -0.6] + ) + coord, atype = _batch() + coord, atype = coord[:1], atype[:1] + energy = model(coord, atype)["energy"].detach().cpu().numpy() + natoms = np.array([[6, 6, 3, 3]], dtype=np.int32) + sample = [ + { + "coord": coord.cpu().numpy(), + "atype": atype.cpu().numpy(), + "box": None, + "natoms": natoms, + "energy": energy.reshape(1, 1), + "find_energy": np.float32(1.0), + } + ] + bias_before = model.atomic_model.out_bias.detach().clone() + model_change_out_bias(model, lambda: sample, "change-by-statistic") + torch.testing.assert_close( + model.atomic_model.out_bias, bias_before, atol=1e-10, rtol=0 + ) + after = model(coord, atype)["energy"].detach().cpu().numpy() + np.testing.assert_allclose(after, energy, atol=1e-10) + + +@pytest.mark.parametrize("weights", ["sum", "mean", [0.3, 1.4]]) +def test_linear_composition_combines_complete_energies(weights: object) -> None: + """A linear composition weights the complete energy of each child. + + The second child is an ordinary model, which reads no topology. + """ + child = _config("les") + child["type"] = "standard" + other = copy.deepcopy(child) + other.pop("long_range") + other["descriptor"]["seed"] = 31 + composition = ( + get_model( + { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": weights, + "models": [child, other], + } + ) + .to(env.DEVICE) + .eval() + ) + standalone = [_model("les"), get_model(other).to(env.DEVICE).eval()] + for index, model in enumerate(standalone): + model.atomic_model.load_state_dict( + composition.atomic_model.models[index].state_dict() + ) + coord, atype = _batch() + out = composition(coord, atype) + if weights == "sum": + w = [1.0, 1.0] + elif weights == "mean": + w = [0.5, 0.5] + else: + w = weights + energy = sum( + ww * m(coord, atype)["energy"] for ww, m in zip(w, standalone, strict=True) + ) + force = sum( + ww * m(coord, atype)["force"] for ww, m in zip(w, standalone, strict=True) + ) + torch.testing.assert_close(out["energy"], energy) + torch.testing.assert_close(out["force"], force) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_parameter_gradients_and_freezing(kernel: str) -> None: + """Energy and force losses reach the charge network and the kernel; frozen parts stay out.""" + model = _model(kernel).train() + coord, atype = _batch() + out = model(coord, atype) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + long_range = model.atomic_model.long_range + for name, parameter in long_range.named_parameters(): + assert parameter.grad is not None and parameter.grad.abs().max() > 0, name + + config = _config(kernel) + config["long_range"]["kernel"]["trainable"] = False + config["long_range"]["charge_net"]["trainable"] = False + frozen = get_model(config).to(env.DEVICE).train() + names = [ + name + for name, parameter in frozen.named_parameters() + if parameter.requires_grad and "long_range" in name + ] + assert names == [] + frozen(coord, atype)["energy"].sum().backward() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_serialization_round_trip(kernel: str) -> None: + """The portable record and the state dict restore the same model.""" + model = _model(kernel) + rng = np.random.default_rng(5) + with torch.no_grad(): + for parameter in model.atomic_model.long_range.parameters(): + parameter.add_(0.01 * _tensor(rng.normal(size=tuple(parameter.shape)))) + coord, atype = _batch() + expected = model(coord, atype) + record = model.serialize() + assert record["@version"] == 3 + restored = type(model).deserialize(copy.deepcopy(record)).to(env.DEVICE).eval() + torch.testing.assert_close(restored(coord, atype)["energy"], expected["energy"]) + rebuilt = _model(kernel) + rebuilt.load_state_dict(model.state_dict()) + torch.testing.assert_close(rebuilt(coord, atype)["force"], expected["force"]) + + +def test_construction_is_reproducible() -> None: + """Without a seed of its own the charge network follows the fitting seed.""" + config = _config("les") + del config["long_range"]["charge_net"]["seed"] + first, second = (get_model(copy.deepcopy(config)) for _ in range(2)) + for (name, a), (_, b) in zip( + first.named_parameters(), second.named_parameters(), strict=True + ): + torch.testing.assert_close(a, b, rtol=0, atol=0, msg=name) + + +def test_ordinary_record_keeps_version() -> None: + """A model without a long-range module keeps its record and constructor.""" + config = _config("les") + config.pop("long_range") + model = get_model(config) + record = model.serialize() + assert record["@version"] == 2 + assert "long_range" not in record + assert not model.atomic_model.has_long_range() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_checkpoint_inference_matches_direct_forward(tmp_path, kernel: str) -> None: + """DeepPot on a checkpoint builds the topology over the real atoms only.""" + model = _model(kernel) + checkpoint = tmp_path / f"long_range_{kernel}.pt" + wrapper = ModelWrapper(model, model_params=_config(kernel)) + torch.save({"model": wrapper.state_dict()}, checkpoint) + coord, atype = _batch(seed=7) + expected = model(coord, atype, do_atomic_virial=True) + dp = DeepPot(str(checkpoint), auto_batch_size=False) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + coord_np = coord.cpu().numpy() + actual = dict( + zip( + names, + dp.eval(coord_np, None, atype.cpu().numpy(), atomic=True, mixed_type=True), + strict=True, + ) + ) + for name in names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + coord_np[1, 3:] += 40.0 + moved = dp.eval(coord_np, None, atype.cpu().numpy(), mixed_type=True) + np.testing.assert_allclose(moved[0], actual["energy"], rtol=0, atol=1e-12) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the complete model, train and eval. + + In training the gradient of an energy and force loss, whose force part is + a second-order derivative of the energy, reaches the long-range + parameters as in eager execution. + """ + model = _model(kernel) + coord, atype = _batch() + parameters = list(model.atomic_model.long_range.parameters()) + + def loss(out: dict) -> torch.Tensor: + return out["energy"].square().sum() + out["force"].square().sum() + + for training in (True, False): + model.train(training) + compiled = _CompiledModel( + model, _get_model_structure_key(model), compile_eval=True + ) + compiled.train(training) + reference = model(coord, atype) + result = compiled(coord, atype) + assert compiled._compiled_lower_by_mode, "the compiled graph was not used" + torch.testing.assert_close(result["energy"], reference["energy"]) + torch.testing.assert_close(result["force"], reference["force"]) + torch.testing.assert_close(result["virial"], reference["virial"]) + if training: + expected = torch.autograd.grad(loss(reference), parameters) + actual = torch.autograd.grad(loss(result), parameters) + for expected_grad, actual_grad in zip(expected, actual, strict=True): + torch.testing.assert_close(actual_grad, expected_grad) + + +def test_unsupported_combinations_fail() -> None: + """Unsupported inputs, routes, combinations and descriptors are rejected.""" + model = _model("les") + coord, atype = _batch() + box = 10.0 * torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 9) + with pytest.raises(NotImplementedError, match="non-periodic"): + model(coord[:1], atype[:1], box=box) + nlist = torch.full((1, 6, 4), -1, dtype=torch.int64, device=env.DEVICE) + with pytest.raises(NotImplementedError, match="NeighborGraph"): + model.atomic_model.forward_atomic(coord[:1], atype[:1], nlist) + with pytest.raises(NotImplementedError, match="compression"): + model.atomic_model.enable_compression(0.5) + config = _config("les") + config["spin"] = {"use_spin": [True, False], "scheme": "native"} + with pytest.raises(NotImplementedError): + get_model(config) + config = _config("les") + config["use_srtab"] = "table.txt" + config["sw_rmin"], config["sw_rmax"] = 0.5, 1.0 + with pytest.raises(NotImplementedError, match="ZBL"): + get_model(config) + config = _config("les") + config["type_map"] = ["O", "H"] + config["bridging_method"] = "ZBL" + with pytest.raises(NotImplementedError, match="bridging"): + get_model(config) + config = _config("les") + config["atom_exclude_types"] = [1] + with pytest.raises(NotImplementedError, match="atom_exclude_types"): + get_model(config) + child = _config("les") + child["type"] = "standard" + with pytest.raises(NotImplementedError, match="atom_exclude_types"): + get_model( + { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": "sum", + "models": [child], + "atom_exclude_types": [1], + } + ) + config = _config("les") + config["descriptor"] = { + "type": "se_e2_a", + "rcut": 3.0, + "rcut_smth": 2.0, + "sel": [8, 8], + } + with pytest.raises(NotImplementedError, match="DPA4C"): + get_model(config) + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + with pytest.raises(NotImplementedError, match="deployable"): + _trace_and_export({"model": model.serialize()}) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_dpmodel_reference_matches(kernel: str) -> None: + """The dpmodel reference evaluates the same energy from the portable record.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + model = _model(kernel) + coord, atype = _batch() + expected = model(coord, atype) + reference = BaseModelDP.deserialize(model.serialize()) + ret = reference.call_common( + coord.cpu().numpy(), atype.cpu().numpy(), neighbor_graph_method="dense" + ) + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + ret["energy"].reshape(3, 6), + expected["atom_energy"].detach().cpu().numpy().reshape(3, 6), + rtol=1e-12, + atol=1e-12, + ) + + +def test_pair_graph_skips_padding() -> None: + """The pair topology joins real atoms of one frame only and ends in masked guards.""" + coord = torch.arange(15, dtype=torch.float64).reshape(5, 3) + pairs = build_pair_graph( + coord, + torch.tensor([3, 2]), + torch.tensor([0, -1, 1, 0, 1]), + torch.tensor([True, False, True, True, True]), + ) + real = pairs.index[:, pairs.mask] + assert real.tolist() == [[0, 3], [2, 4]] + assert (~pairs.mask).sum() == 2 diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 2401aa59c5..aad817ad2d 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -472,6 +472,8 @@ def _make_mock_model(self): fitting_net = MagicMock() class FakeModel(DPModelCommon): + atomic_model = MagicMock(**{"has_long_range.return_value": False}) + def get_out_bias(self): return np.array([[0.0, 0.0]]) From 043c253d12efa02d839534aeb7e1792208ad3bc4 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 18 Oct 2026 19:56:19 +0800 Subject: [PATCH 14/30] feat(pt_expt): freeze models with a long-range module Export a model with a latent-charge long-range module to .pt2 and .pte through the graph lower, and evaluate the frozen model with the Python inference interface. - The graph lower of such a model takes the `all_pairs` topology after its other inputs, as `pair_index`, `pair_vec` and `pair_mask`, each with a dynamic pair axis of its own. The metadata names the topologies the lower reads (`topologies`), and the caller builds them from the coordinates of each frame, as checkpoint inference already does. - `flatten_topologies` and `unflatten_topologies` define the positional layout once; graph export, compiled training and DeepEval share it, and `build_synthetic_topologies` the trace-time samples of export and compiled training. - The C++ interface rejects such an archive when it is loaded, with a message pointing to the Python interface. --- .../dpmodel/utils/neighbor_graph/__init__.py | 4 + .../utils/neighbor_graph/pair_graph.py | 56 +++++++++++ deepmd/pt_expt/infer/deep_eval.py | 34 +++---- deepmd/pt_expt/model/ener_model.py | 84 +++++++++++++++++ deepmd/pt_expt/model/make_model.py | 70 +++++++++++++- deepmd/pt_expt/train/training.py | 39 ++++---- deepmd/pt_expt/utils/serialization.py | 92 +++++++++++++++++-- doc/model/long-range.md | 14 ++- source/api_cc/src/DeepPotPTExpt.cc | 7 ++ source/tests/pt_expt/model/test_long_range.py | 70 ++++++++++++-- 10 files changed, 413 insertions(+), 57 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index dce69ed20b..09617ded48 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -61,6 +61,8 @@ PairGraph, build_pair_graph, build_topologies, + flatten_topologies, + unflatten_topologies, unordered_pair_index, ) from .pairs import ( @@ -97,6 +99,7 @@ "edge_env_mat", "edge_force_virial", "expand_node_values", + "flatten_topologies", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", @@ -111,5 +114,6 @@ "segment_mean", "segment_softmax", "segment_sum", + "unflatten_topologies", "unordered_pair_index", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index 2e33cc6141..4a78487a69 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -42,6 +42,7 @@ if TYPE_CHECKING: from collections.abc import ( Iterable, + Sequence, ) from deepmd.dpmodel.array_api import ( @@ -229,3 +230,58 @@ def build_topologies( ) topologies[name] = build_pair_graph(coord, n_node, atype, node_mask, pair_excl) return topologies + + +def flatten_topologies(topologies: dict[str, PairGraph]) -> tuple[Array, ...]: + """Lay topologies out as positional inputs of a graph lower. + + A compiled or exported graph lower receives every topology as three + inputs, ``(index, vec, mask)``, in the order of the dict, after its other + inputs; :func:`unflatten_topologies` restores them. + + Parameters + ---------- + topologies + The topologies by name. + + Returns + ------- + tuple + ``(index, vec, mask)`` of every topology, concatenated. + """ + return tuple( + value + for pairs in topologies.values() + for value in (pairs.index, pairs.vec, pairs.mask) + ) + + +def unflatten_topologies( + names: Sequence[str], values: Sequence[Array] +) -> dict[str, PairGraph]: + """Restore topologies laid out by :func:`flatten_topologies`. + + Parameters + ---------- + names + The topology names, in the order they were laid out. + values + ``(index, vec, mask)`` of every topology, concatenated. + + Returns + ------- + dict[str, PairGraph] + The topologies by name. + + Raises + ------ + ValueError + If the number of values is not three per name. + """ + if len(values) != 3 * len(names): + raise ValueError( + f"expected three inputs per topology for {tuple(names)}, got {len(values)}" + ) + return { + name: PairGraph(*values[3 * ii : 3 * ii + 3]) for ii, name in enumerate(names) + } diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index b5c04eaadd..7ccee2cf3c 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -23,8 +23,9 @@ OutputVariableDef, ) from deepmd.dpmodel.utils.neighbor_graph import ( - PairGraph, build_topologies, + flatten_topologies, + unflatten_topologies, ) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, @@ -227,10 +228,6 @@ def __init__( # Only a compressed ``.pt2`` folds its charge state into constants; a # model that reads the condition as an ordinary input needs no rebuild. self._charge_state_fold: ChargeStateFold | None = None - # Topologies the graph lower reads beside the neighbor graph; only a - # checkpoint can declare them, since no deployable archive holds a - # model that does. - self._required_topologies: tuple[str, ...] = () if self._is_pt2: self._load_pt2(model_file) @@ -672,7 +669,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ) use_graph_lower = model_uses_graph_lower(model) - self._required_topologies = model.atomic_model.required_topologies() + topology_names = model.atomic_model.required_topologies() self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -727,6 +724,8 @@ 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", } + if topology_names: + self.metadata["topologies"] = list(topology_names) if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, @@ -744,7 +743,8 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # ``communicate_extended_output``, while the graph runner emits local # public keys consumed directly by ``_eval_model_graph``. # - # Graph: (..., source_row_ptr, fparam, aparam, charge_spin) + # Graph: (..., source_row_ptr, fparam, aparam, charge_spin, + # *topology inputs) # Graph spin: (..., source_row_ptr, spin, fparam, aparam, charge_spin) # Nlist: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) # Nlist spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) @@ -814,7 +814,7 @@ def _eager_runner_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None = None, - topologies: dict[str, PairGraph] | None = None, + *topology_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: model_ret = model.forward_common_lower_graph( atype, @@ -832,7 +832,7 @@ def _eager_runner_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, - topologies=topologies, + topologies=unflatten_topologies(topology_names, topology_vals), ) return _translate_energy_keys( model_ret, @@ -2398,9 +2398,10 @@ def _eval_model_graph( positional schema ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, - fparam, aparam, charge_spin)`` to the graph forward. Exported AOTI - artifacts use a dynamic edge axis, so no ``edge_capacity`` padding is - needed. The ``graph_edge_dtype`` metadata selects float32 geometry for + fparam, aparam, charge_spin)`` to the graph forward, followed by the + ``(index, vec, mask)`` of every topology the ``topologies`` metadata + names. Exported AOTI artifacts use dynamic edge and pair axes, so no + ``edge_capacity`` padding is needed. The ``graph_edge_dtype`` metadata selects float32 geometry for compressed DPA1 and float64 for generic graph descriptors. The forward returns local public keys directly, so results are reshaped without ``communicate_extended_output``. @@ -2539,12 +2540,13 @@ def _eval_model_graph( aparam_t, charge_spin_t, ) - if self._required_topologies: + topology_names = self.metadata.get("topologies", ()) + if topology_names: # Padding slots of a mixed-type batch are no atoms and join # no pair. - model_inputs += ( + model_inputs += flatten_topologies( build_topologies( - self._required_topologies, + topology_names, torch.as_tensor( coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE ), @@ -2553,7 +2555,7 @@ def _eval_model_graph( node_mask=atype_t >= 0, atype=atype_t, pair_excl=self._model_pair_excl(), - ), + ) ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index f7e021a684..fe74f20de7 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -18,6 +18,9 @@ from deepmd.dpmodel.model.make_hessian_model import ( make_hessian_model, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, +) from deepmd.dpmodel.utils.neighbor_list import ( NeighborList, ) @@ -613,6 +616,7 @@ def forward_lower_graph_exportable( do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, PairGraph] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -646,6 +650,11 @@ def forward_lower_graph_exportable( and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_lower``. + topologies + The sample pair topology the atomic model declares, or ``None``; + its ``index``, ``vec`` and ``mask`` become the inputs + ``pair_index``, ``pair_vec`` and ``pair_mask`` after + ``charge_spin``. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -680,6 +689,7 @@ def forward_lower_graph_exportable( charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, destination_sorted=destination_sorted, + topologies=topologies, **make_fx_kwargs, ) @@ -688,6 +698,80 @@ def forward_lower_graph_exportable( do_grad_r = self.do_grad_r("energy") do_grad_c = self.do_grad_c("energy") + if topologies: + # The inputs of the deployed lower keep their names, which the + # export reads, so the traced signature spells out the one pair + # topology a long-range module declares. + if len(topologies) != 1: + raise NotImplementedError( + "the deployed graph lower takes one pair topology, got " + f"{list(topologies)}" + ) + (pairs,) = topologies.values() + + def fn_pairs( + 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, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + pair_index: torch.Tensor, + pair_vec: torch.Tensor, + pair_mask: torch.Tensor, + ) -> dict[str, torch.Tensor]: + model_ret = traced( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + pair_index, + pair_vec, + pair_mask, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn_pairs, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + pairs.index, + pairs.vec, + pairs.mask, + ) + def fn( atype: torch.Tensor, n_node: torch.Tensor, diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index ba49568ce8..73f5fc9eb4 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -25,6 +25,8 @@ build_topologies, compact_nodes, expand_node_values, + flatten_topologies, + unflatten_topologies, ) from deepmd.pt_expt.common import ( auto_wrapped_class, @@ -1306,6 +1308,7 @@ def forward_common_lower_graph_exportable( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, PairGraph] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` @@ -1346,6 +1349,13 @@ def forward_common_lower_graph_exportable( output. ``None`` (default) is the existing, unconditioned trace with no mag output -- used by every non-spin caller of this generic (descriptor-agnostic) exportable. + topologies + Sample topologies the atomic model declares by + ``required_topologies()``, or ``None``. When given, they are + traced inputs laid out by + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies` + after ``charge_spin``; a model reading topologies carries no + native spin. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -1356,8 +1366,9 @@ def forward_common_lower_graph_exportable( A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, fparam, aparam, charge_spin, spin)`` and - returns a dict with the same internal keys as + source_row_ptr, fparam, aparam, charge_spin, spin)`` -- or, + with topologies, ``spin`` replaced by the topology inputs -- + and returns a dict with the same internal keys as ``forward_common_lower_graph``. """ validate_graph_csr_for_export( @@ -1372,6 +1383,61 @@ def forward_common_lower_graph_exportable( ) model = self + if topologies: + topology_names = tuple(topologies) + + def fn_topologies( + 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, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + *topology_vals: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), + ) + + return make_fx(fn_topologies, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *flatten_topologies(topologies), + ) + # ``spin`` is a traced INPUT only when a real spin tensor is # given (native-spin models). For every non-spin caller # (``spin is None``) the traced ``fn`` must have the SAME arity as diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 476a582ec0..f4e24bc824 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -73,8 +73,9 @@ make_learning_rate_schedule, ) from deepmd.dpmodel.utils.neighbor_graph import ( - PairGraph, build_topologies, + flatten_topologies, + unflatten_topologies, ) from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, @@ -865,6 +866,7 @@ def _trace_and_compile_graph( # its graph builders omit the CSR metadata carried by deployment inputs. from deepmd.pt_expt.utils.serialization import ( build_synthetic_graph_inputs, + build_synthetic_topologies, check_graph_trace_torch_version, count_synthetic_graph_edges, ) @@ -922,22 +924,19 @@ def _trace_and_compile_graph( # tensors, while forward_common_lower_graph takes it after them. s_spin = s_conditioning.pop(0) if spin is not None else None s_fparam, s_aparam, s_charge_spin = s_conditioning - # The trace reads only the shapes of the pair topologies, so the synthetic - # pairs join fixed nodes over a pair count distinct from the other axes. - topology_names = model.atomic_model.required_topologies() - s_topology_vals: list[torch.Tensor] = [] - if topology_names: - p_trace = _next_safe_prime(e_max + 2, _forbidden | {trace_nf, trace_N, e_max}) - for _ in topology_names: - s_topology_vals += [ - torch.arange(2 * p_trace, dtype=torch.int64, device=_trace_device) - .reshape(2, p_trace) - .remainder(trace_N), - torch.ones( - p_trace, 3, dtype=GLOBAL_PT_FLOAT_PRECISION, device=_trace_device - ), - torch.ones(p_trace, dtype=torch.bool, device=_trace_device), - ] + # The pair axis of every topology is a fourth dynamic axis, its sample + # length distinct from the other three and from the CSR offsets. + s_topologies = build_synthetic_topologies( + model, + trace_N, + _next_safe_prime( + e_max + 2, _forbidden | {trace_nf, trace_N, trace_N + 1, e_max} + ), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=_trace_device, + ) + topology_names = tuple(s_topologies) + s_topology_vals = flatten_topologies(s_topologies) n_topology_vals = len(s_topology_vals) def fn( @@ -957,12 +956,8 @@ def fn( spin: torch.Tensor | None, *extra_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: - topology_vals = extra_vals[:n_topology_vals] + topologies = unflatten_topologies(topology_names, extra_vals[:n_topology_vals]) task_buf_vals = extra_vals[n_topology_vals:] - topologies = { - name: PairGraph(*topology_vals[3 * ii : 3 * ii + 3]) - for ii, name in enumerate(topology_names) - } # Patch task-specific buffers with the proxy tensors so make_fx records # them as FX placeholders (mirrors the dense ``_trace_and_compile``). originals: dict[str, torch.Tensor | None] = {} diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index e5252d6825..31d11891f8 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -14,6 +14,10 @@ import numpy as np import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + flatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -674,6 +678,53 @@ def build_synthetic_graph_inputs( ) +def build_synthetic_topologies( + model: torch.nn.Module, + n_node_total: int, + n_pair: int, + *, + dtype: torch.dtype, + device: torch.device, +) -> dict[str, PairGraph]: + """Build trace-time samples of the topologies a model declares. + + Single source of the trace-time topology inputs, shared by graph export + and compiled training. Tracing reads only their shapes, so every topology + joins fixed nodes over ``n_pair`` pairs, a length the caller keeps + distinct from the other dynamic axes. + + Parameters + ---------- + model : torch.nn.Module + The model; its atomic model declares the topologies. + n_node_total : int + Length of the node axis the pairs index. + n_pair : int + Length of the pair axis of every sample. + dtype : torch.dtype + Precision of the pair vectors. + device : torch.device + Target device. + + Returns + ------- + dict[str, PairGraph] + A sample of every declared topology; empty when none is declared. + """ + index = ( + torch.arange(2 * n_pair, dtype=torch.int64, device=device).reshape(2, n_pair) + % n_node_total + ) + return { + name: PairGraph( + index=index, + vec=torch.ones(n_pair, 3, dtype=dtype, device=device), + mask=torch.ones(n_pair, dtype=torch.bool, device=device), + ) + for name in model.atomic_model.required_topologies() + } + + def build_synthetic_canonical_graph_inputs( model: torch.nn.Module, e_max: int, @@ -830,6 +881,7 @@ def count_synthetic_graph_edges( def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, is_native_spin: bool = False, + topology_names: tuple[str, ...] = (), ) -> tuple: """Build dynamic-shape specifications for the graph-form forward_lower export. @@ -857,6 +909,10 @@ def _build_graph_dynamic_shapes( Whether ``sample_inputs`` follows the native-spin positional ABI (spin at slot 10) instead of the regular energy ABI (charge_spin at slot 12). + topology_names : tuple[str, ...] + Topologies whose ``(index, vec, mask)`` follow the inputs above, each + with a dynamic pair axis of its own; the guard pairs of a topology + bound it below by two. """ nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) @@ -873,6 +929,10 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) ) + topologies = [] + for name in topology_names: + npair_dim = torch.export.Dim(f"npair_{name}", min=2) + topologies += [{1: npair_dim}, {0: npair_dim}, {0: npair_dim}] if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -886,6 +946,7 @@ def _build_graph_dynamic_shapes( {0: nframes_dim} if fparam is not None else None, # fparam {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) fparam = sample_inputs[10] aparam = sample_inputs[11] @@ -898,6 +959,7 @@ def _build_graph_dynamic_shapes( # dim would make torch.export prove/reject the equality). {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) @@ -1286,6 +1348,13 @@ def _probe_has_message_passing(obj: object) -> bool | None: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break meta["pair_exclude_types"] = pair_exclude_types + # Topologies the graph lower reads beside the neighbor graph, laid out + # after its other inputs (see ``flatten_topologies``); the caller of the + # lower builds them from the coordinates of each frame. Absent when the + # model reads none. + topologies = model.atomic_model.required_topologies() + if topologies: + meta["topologies"] = list(topologies) return meta @@ -1832,11 +1901,6 @@ def _trace_and_export_impl( # Registry-dispatched (incl. native spin, type "native_spin"): the # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) - if model.atomic_model.has_long_range(): - raise NotImplementedError( - "a model with a long-range module cannot be exported to a " - "deployable archive; evaluate its training checkpoint (.pt)" - ) model.eval() # The vacuum reference is resolved on the export object, folded into the # fitting bias or stored as a per-type table, so the exported graph @@ -2090,6 +2154,18 @@ def _trace_and_export_impl( device=torch.device("cpu"), want_spin=is_native_spin, ) + # The pair axis of every topology is dynamic; its sample length + # is kept distinct from the other axes and from the CSR offsets. + topologies = build_synthetic_topologies( + model, + n_sample, + next_safe_prime( + e_sample + 2, + _forbidden | {nframes_sample, n_sample, n_sample + 1, e_sample}, + ), + dtype=edge_dtype, + device=torch.device("cpu"), + ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): # slot 10 is ``spin`` (mandatory), slots 11/12 are @@ -2115,12 +2191,16 @@ def _trace_and_export_impl( do_atomic_virial=do_atomic_virial, charge_spin=sample_inputs[12], destination_sorted=True, + topologies=topologies, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) dynamic_shapes = _build_graph_dynamic_shapes( - *sample_inputs, is_native_spin=is_native_spin + *sample_inputs, + is_native_spin=is_native_spin, + topology_names=tuple(topologies), ) + sample_inputs = (*sample_inputs, *flatten_topologies(topologies)) output_keys = traced_output_keys(traced) exported = torch.export.export( traced, diff --git a/doc/model/long-range.md b/doc/model/long-range.md index cbb7b96875..295f251662 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -121,8 +121,12 @@ long-range interaction is part of the compiled model: the pairs of each frame are listed outside the compiled region and enter it as inputs of dynamic length. -A trained checkpoint (`.pt`) is evaluated by `dp test` and the Python -inference interface. The following are not supported for a model with a -long-range module: freezing to `.pt2` or `.pte`, compression, multi-rank -(domain-decomposed) inference, spin models, the tabulated ZBL potential, -and short-range bridging. +`dp --pt-expt freeze` exports the model, long-range module included, to +`.pt2` or `.pte`. A trained checkpoint (`.pt`) and the frozen model are +evaluated by `dp test` and the Python inference interface, which list the +pairs of each frame before calling the model. The C++ interface, and with it +LAMMPS, does not evaluate a model with a long-range module. + +The following are not supported for a model with a long-range module: +compression, multi-rank (domain-decomposed) inference, spin models, the +tabulated ZBL potential, and short-range bridging. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 005a18e722..1bbdc9b9bc 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -256,6 +256,13 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_graph_ = false; lower_input_is_canonical_ = false; } + // A model with a long-range module reads pair topologies after the neighbor + // graph, which this interface does not build. + if (metadata.obj_val.count("topologies")) { + throw deepmd::deepmd_exception( + "this .pt2 holds a model with a long-range module, which the C++ " + "interface does not evaluate; use the Python inference interface."); + } if (lower_input_is_edge_) { std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " "schema (pt-backend SeZM/DPA4 freeze). Support will be " diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 697e3bb68f..751451373e 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -30,6 +30,9 @@ from deepmd.pt_expt.utils import ( env, ) +from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, +) from ..compile_utils import ( REQUIRES_SUPPORTED_COMPILE, @@ -539,6 +542,67 @@ def test_checkpoint_inference_matches_direct_forward(tmp_path, kernel: str) -> N np.testing.assert_allclose(moved[0], actual["energy"], rtol=0, atol=1e-12) +@pytest.mark.parametrize( + ("kernel", "source", "suffix"), + [("les", "neutral", ".pt2"), ("sog", "neutral", ".pte"), ("les", "fparam", ".pte")], +) +def test_frozen_model_matches_direct_forward( + tmp_path, kernel: str, source: str, suffix: str +) -> None: + """A frozen model reads the pair topology as named inputs of any length.""" + config = _config(kernel, constraint={"source": source, "fparam_index": 0}) + frame = {} + if source == "fparam": + config["fitting_net"]["numb_fparam"] = 1 + frame["fparam"] = np.array([[1.0], [-0.5], [2.0]]) + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / f"long_range_{kernel}{suffix}" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="graph", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + assert dp.deep_eval.metadata["topologies"] == ["all_pairs"] + if suffix == ".pte": + names = [ + node.name + for node in torch.export.load(str(frozen)).graph_module.graph.nodes + if node.op == "placeholder" + ] + for name, expected_name in zip( + names[-3:], ("pair_index", "pair_vec", "pair_mask"), strict=True + ): + assert name.startswith(expected_name), names + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + for seed, nloc in ((7, 6), (8, 4)): + coord, atype = _batch(seed=seed) + coord, atype = coord[:, :nloc], atype[:, :nloc] + expected = model( + coord, + atype, + do_atomic_virial=True, + **{key: _tensor(value) for key, value in frame.items()}, + ) + actual = dp.eval( + coord.cpu().numpy(), + None, + atype.cpu().numpy(), + atomic=True, + mixed_type=True, + **frame, + ) + for name, value in zip(names, actual, strict=True): + np.testing.assert_allclose( + value, + expected[name].detach().cpu().numpy().reshape(value.shape), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + @pytest.mark.parametrize("kernel", ["les", "sog"]) @REQUIRES_SUPPORTED_COMPILE def test_compiled_matches_eager(kernel: str) -> None: @@ -625,12 +689,6 @@ def test_unsupported_combinations_fail() -> None: } with pytest.raises(NotImplementedError, match="DPA4C"): get_model(config) - from deepmd.pt_expt.utils.serialization import ( - _trace_and_export, - ) - - with pytest.raises(NotImplementedError, match="deployable"): - _trace_and_export({"model": model.serialize()}) @pytest.mark.parametrize("kernel", ["les", "sog"]) From c60df4a7a3cf3681b8a9f9e0f995b64192ca0913 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 01:03:34 +0800 Subject: [PATCH 15/30] feat(pt_expt): compressed DPA4C inference with a long-range module Compress a DPA4C model that carries a latent-charge long-range module, and evaluate it on a CUDA device through the compact canonical lower, with a direct-sum operator whose memory is linear in the atom count. - `deepmd::latent_charge_direct_sum` takes the charges the charge network predicts for the nodes of complete frames, applies the total-charge constraint, and returns the atomic energy, its derivative with respect to those charges, and the force, per-atom virial and per-frame virial at fixed charges, without listing the pairs. Each pair is evaluated by both of its nodes in float32 from double-float positions; the partial sums, the constraint and the per-frame sums are added in a fixed order in float64, so the result is deterministic. The CPU kernel is the float64 reference. - The compact canonical lower of such a model takes the float64 node positions as a ninth input. The descriptor operator runs once; the charge gradient, pulled back through the charge network, joins the descriptor cotangent of the fitting, so one descriptor backward yields the edge gradient of both. - Compressing a descriptor conditioned on `charge_spin` folds the charge state into its tables; compressing the atomic model folds a constraint that reads its total charge from `charge_spin` to the same state (`fold_charge_state`), and the charge-state fold of a .pt2 archive rewrites both together. A descriptor compressed alone, whose constraint still reads the condition, is rejected. - A model the operator does not serve (pair exclusion, frame or atomic parameters of the charge network, more than 8 channels or 32 Gaussians), like any CPU export, keeps the graph lower and its pair lists. --- .../dpmodel/atomic_model/dp_atomic_model.py | 30 +- deepmd/dpmodel/long_range/base_long_range.py | 35 + deepmd/dpmodel/long_range/latent_charge.py | 106 +- .../dpmodel/utils/neighbor_graph/__init__.py | 2 + .../utils/neighbor_graph/pair_graph.py | 28 +- deepmd/pt_expt/infer/deep_eval.py | 10 + deepmd/pt_expt/kernels/dpa4c/canonical.py | 15 + .../kernels/dpa4c/canonical_long_range.py | 149 +++ deepmd/pt_expt/kernels/long_range.py | 291 +++++ deepmd/pt_expt/long_range/latent_charge.py | 109 ++ deepmd/pt_expt/model/ener_model.py | 76 +- deepmd/pt_expt/model/make_model.py | 3 +- deepmd/pt_expt/utils/serialization.py | 143 ++- doc/model/long-range.md | 28 +- source/op/pt/CMakeLists.txt | 9 +- source/op/pt/long_range/direct_sum.cu | 1058 +++++++++++++++++ source/op/pt/long_range/ops.cc | 63 + .../kernels/test_long_range_direct_sum.py | 217 ++++ source/tests/pt_expt/model/test_long_range.py | 153 ++- 19 files changed, 2467 insertions(+), 58 deletions(-) create mode 100644 deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py create mode 100644 deepmd/pt_expt/kernels/long_range.py create mode 100644 source/op/pt/long_range/direct_sum.cu create mode 100644 source/op/pt/long_range/ops.cc create mode 100644 source/tests/pt_expt/kernels/test_long_range_direct_sum.py diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 0122c29f8a..c52ed7aadc 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -254,14 +254,13 @@ def uses_graph_lower(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.uses_graph_lower()) - def fused_decomposition(self) -> "tuple[DPAtomicModel, None] | None": + def fused_decomposition(self) -> "tuple[DPAtomicModel, None]": """A descriptor-fitting model is its own learned part. - A model with a long-range module is more than a learned part and a - pair potential, so no fused pipeline evaluates it. + The learned part carries the long-range module of the model, if any; + a fused pipeline that does not evaluate long-range modules declines a + learned part with one. """ - if self.long_range is not None: - return None return self, None def has_long_range(self) -> bool: @@ -365,6 +364,12 @@ def enable_compression( ) -> None: """Call descriptor enable_compression(). + A long-range module keeps its networks: it reads the output of the + compressed descriptor like that of the uncompressed one. A compressed + descriptor folds the charge/spin condition into its tables, so a + constraint that reads the total charge from the condition is folded + to the same state, the default condition of the descriptor. + Parameters ---------- min_nbor_dist @@ -378,10 +383,6 @@ def enable_compression( check_frequency The overflow check frequency """ - if self.long_range is not None: - raise NotImplementedError( - "compression of a model with a long-range module is not supported" - ) self.descriptor.enable_compression( min_nbor_dist, table_extrapolate, @@ -389,6 +390,8 @@ def enable_compression( table_stride_2, check_frequency, ) + if self.long_range is not None and self.long_range.uses_charge_spin(): + self.long_range.fold_charge_state(self.descriptor.get_default_chg_spin()) def forward_atomic( self, @@ -726,6 +729,15 @@ def forward_atomic_graph( # with the defaults the descriptor and fitting fall back on. n_node = graph.n_node if self.long_range is not None: + # The constraint and a compressed descriptor serve one charge + # state; compressing the atomic model folds both to it. + if self.long_range.uses_charge_spin() and self.get_dim_chg_spin() == 0: + raise ValueError( + "the descriptor folds the charge/spin condition into its " + "compressed tables, while the long-range charge constraint " + "still reads it; compress the atomic model, which folds " + "both to one state" + ) fparam, charge_spin_frame = self._frame_conditions( fparam, charge_spin, n_node.shape[0], graph.edge_vec ) diff --git a/deepmd/dpmodel/long_range/base_long_range.py b/deepmd/dpmodel/long_range/base_long_range.py index 9cf76eda93..f37b821333 100644 --- a/deepmd/dpmodel/long_range/base_long_range.py +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -139,6 +139,41 @@ def change_type_map(self, type_map: list[str]) -> None: def set_case_embd(self, case_idx: int) -> None: """Select the case embedding of the networks of the module.""" + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the module to one frame charge/spin condition. + + Called when the descriptor of the model is compressed, which folds + the condition into its tables, for a module whose + :meth:`uses_charge_spin` holds. + + Raises + ------ + NotImplementedError + If the module cannot be specialized. + """ + raise NotImplementedError( + f"{type(self).__name__} cannot be specialized to a charge/spin " + "condition, which the compression of the model requires" + ) + + def charge_state_constants(self) -> list[t_tensor]: + """The tensors that carry the folded charge/spin condition. + + A deployment that serves another condition writes the outputs of + :meth:`rebuild_charge_state` over them, in this order. + """ + return [] + + def rebuild_charge_state(self, charge_spin: t_tensor) -> list[t_tensor]: + """Build the tensors of :meth:`charge_state_constants` for a condition. + + Parameters + ---------- + charge_spin + The condition with shape ``(1, 2)``. + """ + return [] + @abstractmethod def serialize(self) -> dict: """Serialize the module to a dict.""" diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index 3e3ea6b99d..4cc6f3ccd1 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -7,12 +7,15 @@ ) import array_api_compat +import numpy as np from deepmd.dpmodel.array_api import ( Array, ) from deepmd.dpmodel.common import ( + GLOBAL_NP_FLOAT_PRECISION, NativeOP, + to_numpy_array, ) from deepmd.dpmodel.fitting.latent_charge import ( LatentChargeFitting, @@ -110,6 +113,9 @@ def __init__( self.charge_constraint = self._normalize_constraint( charge_constraint, charge_net ) + # The frame charge/spin condition a compressed model is specialized + # to; see :meth:`fold_charge_state`. + self.charge_state = None @staticmethod def _normalize_constraint( @@ -203,11 +209,41 @@ def required_topologies(self) -> tuple[str, ...]: return (ALL_PAIRS,) def uses_charge_spin(self) -> bool: + # A folded state stands in for the condition. return ( self.charge_constraint is not None and self.charge_constraint["source"] == "charge_spin" + and self.charge_state is None ) + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the constraint to one frame charge/spin condition. + + A compressed descriptor folds the condition into its tables and reads + none at run time, so a constraint that reads its total charge from the + condition takes it from this state instead. The state is a constant + of the deployed model that a deployment rewrites, together with the + tables of the descriptor, when it serves another condition. + + Parameters + ---------- + charge_spin + The condition ``[charge, spin multiplicity]``. + """ + self.charge_state = np.asarray( + charge_spin, dtype=GLOBAL_NP_FLOAT_PRECISION + ).reshape(1, 2) + + def charge_state_constants(self) -> list[Array]: + return [] if self.charge_state is None else [self.charge_state] + + def rebuild_charge_state(self, charge_spin: Array) -> list[Array]: + if self.charge_state is None: + return [] + xp = array_api_compat.array_namespace(charge_spin) + state = xp.astype(charge_spin, self.charge_state.dtype) + return [xp.reshape(state, (1, 2))] + def call(self, *args: Any, **kwargs: Any) -> Array: return self.call_graph(*args, **kwargs) @@ -285,9 +321,49 @@ def _constrain( """Shift the constrained channel so that it sums to the frame total charge.""" xp = array_api_compat.array_namespace(charge) nf = n_node.shape[0] - source = self.charge_constraint["source"] channel = self.charge_constraint["channel"] - if source == "charge_spin": + target = self._total_charge(nf, charge, fparam, charge_spin) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + total = segment_sum(charge[:, channel], frame_id, nf) + # A frame without a charged atom has nothing to shift. + shift = (total - target) / xp.maximum(count, xp.ones_like(count)) + constrained = charge[:, channel] - xp.take(shift, frame_id, axis=0) * weight + return xp.concat( + [charge[:, :channel], constrained[:, None], charge[:, channel + 1 :]], + axis=-1, + ) + + def _total_charge( + self, + nf: int, + like: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """The total charge of every frame, which the constrained channel sums to. + + Parameters + ---------- + nf + The number of frames. + like + An array giving the namespace, dtype and device of the result. + fparam, charge_spin + As in :meth:`call_graph`. + + Returns + ------- + Array + The total charges with shape ``(nf,)``, in units of :math:`e`. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + source = self.charge_constraint["source"] + if source == "charge_spin" and self.charge_state is not None: + state = xp.asarray(self.charge_state, device=device) + target = xp.broadcast_to(state[:, 0], (nf,)) + elif source == "charge_spin": if charge_spin is None: raise ValueError( "the charge constraint reads the total charge from " @@ -302,20 +378,8 @@ def _constrain( ) target = fparam[:, self.charge_constraint["fparam_index"]] else: - target = xp.zeros( - (nf,), dtype=charge.dtype, device=array_api_compat.device(charge) - ) - target = xp.astype(target, charge.dtype) - weight = xp.astype(node_mask, charge.dtype) - count = segment_sum(weight, frame_id, nf) - total = segment_sum(charge[:, channel], frame_id, nf) - # A frame without a charged atom has nothing to shift. - shift = (total - target) / xp.maximum(count, xp.ones_like(count)) - constrained = charge[:, channel] - xp.take(shift, frame_id, axis=0) * weight - return xp.concat( - [charge[:, :channel], constrained[:, None], charge[:, channel + 1 :]], - axis=-1, - ) + target = xp.zeros((nf,), dtype=like.dtype, device=device) + return xp.astype(target, like.dtype) def compute_input_stats( self, @@ -334,7 +398,7 @@ def set_case_embd(self, case_idx: int) -> None: self.charge_net.set_case_embd(case_idx) def serialize(self) -> dict: - return { + data = { "@class": "LongRange", "type": "latent_charge", "@version": 1, @@ -342,13 +406,19 @@ def serialize(self) -> dict: "kernel": self.kernel.serialize(), "charge_constraint": self.charge_constraint, } + if self.charge_state is not None: + data["charge_state"] = to_numpy_array(self.charge_state).tolist() + return data @classmethod def deserialize(cls, data: dict) -> "LatentChargeLongRange": data = data.copy() check_version_compatibility(data.pop("@version"), 1, 1) - return cls( + obj = cls( LatentChargeFitting.deserialize(data["charge_net"]), BaseInteractionKernel.deserialize(data["kernel"]), data["charge_constraint"], ) + if data.get("charge_state") is not None: + obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) + return obj diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 09617ded48..8779648f15 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -61,6 +61,7 @@ PairGraph, build_pair_graph, build_topologies, + check_complete_frames, flatten_topologies, unflatten_topologies, unordered_pair_index, @@ -95,6 +96,7 @@ "build_topologies", "canonicalize_neighbor_graph", "center_edge_pairs", + "check_complete_frames", "compact_nodes", "edge_env_mat", "edge_force_virial", diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index 4a78487a69..4e9c653128 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -222,16 +222,32 @@ def build_topologies( for name in names: if name != ALL_PAIRS: raise ValueError(f"unknown topology {name!r}; expected {ALL_PAIRS!r}") - if box is not None: - raise NotImplementedError( - "the all-pairs topology of the long-range modules describes " - "complete non-periodic systems only; periodic cells are not " - "supported (pass box=None for a non-periodic system)" - ) + check_complete_frames(box) topologies[name] = build_pair_graph(coord, n_node, atype, node_mask, pair_excl) return topologies +def check_complete_frames(box: Array | None) -> None: + """Reject a cell, since the all-pairs topology covers non-periodic frames. + + Parameters + ---------- + box + Simulation cells, or ``None`` for non-periodic systems. + + Raises + ------ + NotImplementedError + If a cell is given. + """ + if box is not None: + raise NotImplementedError( + "the all-pairs topology of the long-range modules describes " + "complete non-periodic systems only; periodic cells are not " + "supported (pass box=None for a non-periodic system)" + ) + + def flatten_topologies(topologies: dict[str, PairGraph]) -> tuple[Array, ...]: """Lay topologies out as positional inputs of a graph lower. diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 7ccee2cf3c..2f9eef288c 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -24,6 +24,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( build_topologies, + check_complete_frames, flatten_topologies, unflatten_topologies, ) @@ -2515,6 +2516,15 @@ def _eval_model_graph( compact.source_row_ptr, compact.source_order, ) + if self.metadata.get("topologies"): + # The compact lower of a model with a long-range module reads + # the positions of complete non-periodic frames. + check_complete_frames(box_input) + model_inputs += ( + torch.as_tensor( + coord_input.reshape(-1, 3), dtype=torch.float64, device=DEVICE + ), + ) else: fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 750bb31b25..1aaf2e1010 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -58,6 +58,21 @@ def canonical_model_eligible(model: Any) -> bool: return False if getattr(atomic_model, "atom_excl", None) is not None: return False + # A latent-charge long-range module joins the pipeline through the node + # positions, which its direct-sum operator reads in place of a pair list; + # that pipeline evaluates no pair potential. + if parts[0].has_long_range(): + from deepmd.pt_expt.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + long_range = parts[0].long_range + if not ( + isinstance(long_range, LatentChargeLongRange) + and long_range.direct_sum_eligible() + and parts[1] is None + ): + return False from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( mega_eligible, ) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py new file mode 100644 index 0000000000..2f1fb6c6b0 --- /dev/null +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact canonical deployment of compressed DPA4C with a long-range module.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +if TYPE_CHECKING: + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) + + +def dpa4c_canonical_long_range_energy_force( + descriptor: Any, + fitting: Any, + long_range: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + coord: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA4C and its long-range module from a compact graph. + + The descriptor operator runs once, and its output feeds both the fused + energy fitting and the charge network of the long-range module. The + derivative of the energy with respect to the descriptor collects the + fitting term and the response of the charges, so a single descriptor + backward yields the edge gradient of both; the long-range force at fixed + charges is added to the assembled force. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor without a spin condition. + fitting + Eligible pt_expt energy fitting network. + long_range + Long-range module whose ``direct_sum_eligible()`` holds. + graph + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and + whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. + atype + Flat node atom types with shape ``(N,)``. + coord + Node positions with shape ``(N, 3)`` in Å; every frame is a complete + non-periodic system. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + do_atomic_virial + Whether to return per-node virials. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp32. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + ensure_registered, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compressed_operator_arguments, + fitting_energy_and_gradient, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + canonical_edge_force_virial, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + ensure_registered as ensure_force_registered, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + frame_scalar_sum, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + operator_args = ( + graph.source, + graph.destination_row_ptr, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + ) + node_descriptor, state = torch.ops.deepmd.dpa4c_canonical_compress( + graph.edge_vec, *operator_args + ) + # The charge network reads the descriptor before the fused fitting + # overwrites it with its own cotangent. + ( + long_range_energy, + long_range_gradient, + long_range_force, + long_range_atom_virial, + long_range_virial, + ) = long_range.energy_force_from_coord( + node_descriptor, atype, ownership, graph.n_node, coord + ) + atom_energy, descriptor_gradient = fitting_energy_and_gradient( + fitting, node_descriptor, atype, ownership, atom_bias + ) + edge_gradient, _, _ = torch.ops.deepmd.dpa4c_canonical_compress_backward_inplace( + descriptor_gradient + long_range_gradient, + state, + graph.edge_vec, + *operator_args, + ) + force, atom_virial, virial, _ = canonical_edge_force_virial( + edge_gradient, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.edge_vec.new_empty(0), + atype.shape[0], + do_atomic_virial, + ) + # Nodes outside the ownership carry no charge and hence no long-range + # energy. + atom_energy = atom_energy + long_range_energy + energy = frame_scalar_sum(atom_energy, graph.n_node) + virial = virial + long_range_virial.reshape(-1, 3, 3) + if do_atomic_virial: + atom_virial = atom_virial + long_range_atom_virial.reshape(-1, 3, 3) + return energy, atom_energy, force + long_range_force, virial, atom_virial diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py new file mode 100644 index 0000000000..d63ebfb344 --- /dev/null +++ b/deepmd/pt_expt/kernels/long_range.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the direct-sum operator of the latent-charge long range. + +The operator ``deepmd::latent_charge_direct_sum`` evaluates, for the nodes of +complete frames, the potential that the charges of every other node of the +frame create at each node and the force between the charges, without listing +the pairs: its memory is linear in the node count. The CUDA kernel is +``source/op/pt/long_range/direct_sum.cu``; the CPU kernel registered here is +the reference it is tested against. + +The operator serves inference. It applies the total-charge constraint of the +module, and returns the atomic energy, its derivative with respect to the +charges the charge network predicts, and the force and virial at fixed +charges; the caller adds the response of the charges by back-propagating the +charge gradient through the charge network. +""" + +import math +from functools import ( + cache, +) +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + unordered_pair_index, +) +from deepmd.pt_expt.kernels.utils import ( + operator_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +__all__ = [ + "MAX_CHANNELS", + "MAX_GAUSSIANS", + "kernel_eligible", + "kernel_operator_arguments", + "latent_charge_direct_sum", + "op_available", +] + +#: Charge channels and Gaussians the CUDA kernel is instantiated for. +MAX_CHANNELS = 8 +MAX_GAUSSIANS = 32 + +_LES = 0 +_SOG = 1 + + +def op_available() -> bool: + """Whether the backend device carries the direct-sum operator. + + The CPU kernel is the reference of this module, which is registered here + so that the answer does not depend on earlier calls. + """ + # Importing the module loads the operator library. + import deepmd.pt.cxx_op # noqa: F401 + + ensure_registered() + return operator_available("latent_charge_direct_sum") + + +def kernel_eligible(kernel: Any) -> bool: + """Whether the operator evaluates this interaction kernel.""" + if isinstance(kernel, LESKernel): + return True + return isinstance(kernel, SOGKernel) and kernel.amplitude.numel() <= MAX_GAUSSIANS + + +def kernel_operator_arguments(kernel: Any) -> tuple[int, torch.Tensor]: + """Convert an interaction kernel into the code and coefficients of the operator. + + Parameters + ---------- + kernel + A pt_expt interaction kernel accepted by :func:`kernel_eligible`. + + Returns + ------- + tuple[int, torch.Tensor] + The kernel code and its float32 coefficients: ``alpha`` for the erf + kernel; the amplitudes ``a_k`` and then the rates ``c_k = log2(e) / + (2 sigma_k^2)`` of ``a_k 2^(-c_k r^2)`` for the sum of Gaussians. + + Raises + ------ + ValueError + If the operator does not evaluate the kernel. + """ + if not kernel_eligible(kernel): + raise ValueError( + f"the direct-sum operator does not evaluate {type(kernel).__name__}" + ) + if isinstance(kernel, LESKernel): + coefficient = torch.exp(kernel.log_alpha).reshape(1) + return _LES, coefficient.detach().to(torch.float32).contiguous() + rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width) + coefficient = torch.cat((kernel.amplitude, rate)) + return _SOG, coefficient.detach().to(torch.float32).contiguous() + + +def _kernel_and_slope( + r_sq: torch.Tensor, coefficient: torch.Tensor, kernel: int +) -> tuple[torch.Tensor, torch.Tensor]: + """K(r) and K'(r) / r at the squared distances ``r_sq``.""" + coefficient = coefficient.to(dtype=r_sq.dtype, device=r_sq.device) + if kernel == _LES: + alpha = coefficient[0] + # Coincident nodes take the limit at the shortest distance the CUDA + # kernel evaluates. + r_sq = torch.clamp(r_sq, min=1e-12 / alpha**2) + r = torch.sqrt(r_sq) + value = torch.erf(alpha * r) / r + contact = (2.0 / math.sqrt(math.pi)) * alpha + slope = (contact * torch.exp(-(alpha**2) * r_sq) - value) / r_sq + return value, slope + amplitude, rate = coefficient.reshape(2, -1) + weight = torch.exp2(-r_sq[:, None] * rate) + return weight @ amplitude, weight @ (-2.0 * math.log(2.0) * rate * amplitude) + + +def _reference( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the operator over the explicit pairs of every frame, in float64.""" + n = coord.shape[0] + frame_count = n_node_per_frame.shape[0] + # Padding nodes belong to no frame and carry no charge. + frame = frame_id_from_n_node(n_node_per_frame, n_total=n) + real = torch.arange(n, dtype=torch.int64, device=coord.device) < torch.sum( + n_node_per_frame + ) + weight = (node_mask & real).to(torch.float64) + q = charge.to(torch.float64) * weight[:, None] + if channel >= 0: + columns = torch.arange(q.shape[1], dtype=torch.int64, device=coord.device) + selector = (columns == channel).to(torch.float64) + count = coord.new_zeros(frame_count).index_add(0, frame, weight) + count = torch.clamp(count, min=1.0) + total = coord.new_zeros(frame_count).index_add(0, frame, q[:, channel]) + shift = (total - total_charge.to(torch.float64)) / count + q = q - (shift[frame] * weight)[:, None] * selector + pairs = unordered_pair_index(n_node_per_frame) + src, dst = pairs[0], pairs[1] + d = coord[src] - coord[dst] + value, slope = _kernel_and_slope(torch.sum(d * d, dim=-1), coefficient, kernel) + potential = coord.new_zeros(q.shape) + potential = potential.index_add(0, src, q[dst] * value[:, None]) + potential = prefactor * potential.index_add(0, dst, q[src] * value[:, None]) + strength = torch.sum(q[src] * q[dst], dim=-1) * slope + pair_force = -strength[:, None] * d + force = torch.zeros_like(coord).index_add(0, src, pair_force) + force = force.index_add(0, dst, -pair_force) + half = (-0.5 * strength)[:, None, None] * d[:, :, None] * d[:, None, :] + atom_virial = coord.new_zeros(n, 3, 3).index_add(0, src, half) + atom_virial = prefactor * atom_virial.index_add(0, dst, half).reshape(n, 9) + virial = coord.new_zeros(frame_count, 9).index_add(0, frame, atom_virial) + atom_energy = 0.5 * torch.sum(q * potential, dim=-1, keepdim=True) + gradient = potential * weight[:, None] + if channel >= 0: + mean = coord.new_zeros(frame_count).index_add(0, frame, gradient[:, channel]) + gradient = gradient - ((mean / count)[frame] * weight)[:, None] * selector + return ( + atom_energy, + gradient.to(charge.dtype), + (prefactor * force).to(charge.dtype), + atom_virial.to(charge.dtype), + virial.to(charge.dtype), + ) + + +def _fake( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + n = coord.shape[0] + return ( + coord.new_empty(n, 1), + torch.empty_like(charge), + charge.new_empty(n, 3), + charge.new_empty(n, 9), + charge.new_empty(n_node_per_frame.shape[0], 9), + ) + + +@cache +def _register() -> None: + torch.library.register_fake("deepmd::latent_charge_direct_sum")(_fake) + library = torch.library.Library("deepmd", "IMPL") + library.impl("latent_charge_direct_sum", _reference, "CPU") + # The registration lives as long as the library object. + _register.library = library + + +def ensure_registered() -> None: + """Register the fake and CPU implementations when the schema is loaded.""" + if isinstance( + getattr(torch.ops.deepmd, "latent_charge_direct_sum", None), + torch._ops.OpOverloadPacket, + ): + _register() + + +def latent_charge_direct_sum( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + kernel: Any, + channel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Energy, charge gradient, force and virial of the charges of complete frames. + + Parameters + ---------- + coord + Node positions with shape ``(N, 3)`` in Å, float64. The frames occupy + contiguous spans of the node axis; nodes past ``sum(n_node_per_frame)`` + are padding and join no frame. + charge + Raw node charges of the charge network with shape ``(N, C)``, + float32, with ``C`` at most :data:`MAX_CHANNELS`. + node_mask + Nodes that carry charges with shape ``(N,)``, bool; the others carry + none. + n_node_per_frame + Node count of every frame with shape ``(F,)``. + total_charge + Total charge of every frame with shape ``(F,)``, float64, which the + constrained channel sums to. + kernel + The interaction kernel, accepted by :func:`kernel_eligible`. + channel + The constrained channel, which is shifted uniformly over the charged + nodes of every frame to sum to its total charge, or -1 for none. + prefactor + Factor of every pair term, in eV Å. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64, each pair + energy split half to each of its nodes. + charge_gradient + Derivative of the energy with respect to ``charge`` with shape + ``(N, C)``, in eV per unit charge, through the constraint. + force + Force at fixed charges with shape ``(N, 3)``, in eV/Å. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV, each + pair term split half to each of its nodes. + virial + Per-frame virial at fixed charges with shape ``(F, 9)`` in eV. + """ + ensure_registered() + code, coefficient = kernel_operator_arguments(kernel) + return torch.ops.deepmd.latent_charge_direct_sum( + coord, + charge, + node_mask, + n_node_per_frame, + total_charge, + coefficient.to(coord.device), + code, + int(channel), + float(prefactor), + ) diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py index 2c612c03de..badb2aa26c 100644 --- a/deepmd/pt_expt/long_range/latent_charge.py +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import torch +from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, +) from deepmd.dpmodel.long_range.latent_charge import ( LatentChargeLongRange as LatentChargeLongRangeDP, ) @@ -40,3 +43,109 @@ def _pair_kernel( for name, value in self.kernel.coefficients(r_sq).items() } return self.kernel.evaluate(r_sq, coefficients) + + def direct_sum_eligible(self) -> bool: + """Whether :meth:`energy_force_from_coord` evaluates this module. + + The direct-sum operator evaluates the kernels and channel counts it is + instantiated for. The compact lower that calls it reads no frame + input: the charges depend on the descriptor alone, and a constrained + channel sums to zero or to the total charge of a folded charge state. + """ + from deepmd.pt_expt.kernels.long_range import ( + MAX_CHANNELS, + kernel_eligible, + op_available, + ) + + return ( + op_available() + and kernel_eligible(self.kernel) + and self.charge_net.dim_out <= MAX_CHANNELS + and self.charge_net.numb_fparam == 0 + and self.charge_net.numb_aparam == 0 + and not self.uses_charge_spin() + ) + + def energy_force_from_coord( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + node_mask: torch.Tensor, + n_node: torch.Tensor, + coord: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the module for inference from the positions of the nodes. + + The direct-sum operator enumerates the pairs of every frame itself, + so the memory is linear in the node count. It applies the constraint + and returns the force at fixed charges and the derivative of the + energy with respect to the charges the network predicts; pulled back + through the charge network, that derivative is the derivative of the + energy with respect to the descriptor, through which the caller adds + the response of the charges to the positions. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``. + node_mask + Nodes carrying charges with shape ``(N,)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + coord + Node positions with shape ``(N, 3)`` in Å. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64. + descriptor_gradient + Derivative of the energy with respect to the descriptor through + the charges, with the shape and precision of ``descriptor``. + force + Force at fixed charges with shape ``(N, 3)`` in eV/Å, float32. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV, + float32, each pair term split half to each of its nodes. + virial + Per-frame virial at fixed charges with shape ``(nf, 9)`` in eV, + float32. + """ + from deepmd.pt_expt.kernels.long_range import ( + latent_charge_direct_sum, + ) + + leaf = descriptor.detach().requires_grad_(True) + with torch.enable_grad(): + charge = self.charge_net.call_graph(leaf, atype)["latent_charge"] + total_charge = coord.new_zeros(n_node.shape[0], dtype=torch.float64) + channel = -1 + if self.charge_constraint is not None: + channel = self.charge_constraint["channel"] + total_charge = self._total_charge(n_node.shape[0], total_charge, None, None) + atom_energy, charge_gradient, force, atom_virial, virial = ( + latent_charge_direct_sum( + coord.to(torch.float64).contiguous(), + charge.detach().to(torch.float32).contiguous(), + node_mask.contiguous(), + n_node, + total_charge.contiguous(), + self.kernel, + channel, + COULOMB_CONSTANT, + ) + ) + (descriptor_gradient,) = torch.autograd.grad( + charge, leaf, charge_gradient.to(charge.dtype) + ) + return ( + atom_energy, + descriptor_gradient.to(descriptor.dtype), + force, + atom_virial, + virial, + ) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index fe74f20de7..960c501e8d 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -141,6 +141,7 @@ def forward_lower_canonical_graph( *, do_atomic_virial: bool, spin: torch.Tensor | None = None, + coord: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -168,6 +169,10 @@ def forward_lower_canonical_graph( spin Per-node magnetic moments with shape ``(N, 3)`` for a native-spin model, or ``None``. Ghost rows carry their owner's moment. + coord + Node positions with shape ``(N, 3)`` in Å, float64, for a model + with a long-range module, whose frames are complete non-periodic + systems; ``None`` otherwise. Returns ------- @@ -214,7 +219,30 @@ def forward_lower_canonical_graph( descriptor = learned.descriptor fitting = learned.fitting_net atom_bias = fused_atom_bias(self.atomic_model, learned) - if use_dpa4c: + if learned.has_long_range(): + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + dpa4c_canonical_long_range_energy_force, + ) + + if coord is None: + raise ValueError( + "a model with a long-range module reads the node positions" + ) + force_mag = None + energy, atom_energy, force, virial, atom_virial = ( + dpa4c_canonical_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + graph, + atype, + coord, + output_mask, + atom_bias, + do_atomic_virial, + ) + ) + elif use_dpa4c: energy, atom_energy, force, virial, atom_virial, force_mag = ( dpa4c_canonical_compress_energy_force( descriptor, @@ -274,6 +302,7 @@ def forward_lower_canonical_graph_exportable( destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, + coord: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -286,6 +315,10 @@ def forward_lower_canonical_graph_exportable( Compact graph node and edge tensors. destination_row_ptr, source_row_ptr, source_order Compact dual-CSR topology. + coord + Sample node positions with shape ``(N, 3)``, float64, for a model + with a long-range module, which then takes them as a ninth input; + ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -294,10 +327,49 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced eight-input compact deployment module. + Traced compact deployment module of eight inputs, or nine with + the node positions. """ model = self + if coord is not None: + + def fn_coord( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + coord: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_lower_canonical_graph( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + do_atomic_virial=do_atomic_virial, + coord=coord, + ) + + return make_fx(fn_coord, **make_fx_kwargs)( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + coord, + ) + def fn( atype: torch.Tensor, n_node: torch.Tensor, diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 73f5fc9eb4..b2d1e2c4ef 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -141,7 +141,8 @@ def _fused_energy_force_graph( """ am = model.atomic_model parts = am.fused_decomposition() - if parts is None: + if parts is None or parts[0].has_long_range(): + # The descriptor pipelines evaluate no long-range module. return None learned, pair_potential = parts desc, fit = learned.descriptor, learned.fitting_net diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 31d11891f8..a1e94de1e1 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -731,12 +731,15 @@ def build_synthetic_canonical_graph_inputs( *, device: torch.device, want_spin: bool = False, + want_coord: bool = False, ) -> tuple[torch.Tensor, ...]: """Build compact canonical trace inputs for compressed CUDA descriptors. ``want_spin`` appends the per-node moment at slot 8, the last slot of the compact ABI, matching :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. + ``want_coord`` appends the float64 node positions at the same slot, which + the compact lower of a model with a long-range module reads. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -793,7 +796,16 @@ def build_synthetic_canonical_graph_inputs( compact.source_row_ptr, compact.source_order, ) - return compact_sample if spin is None else (*compact_sample, spin) + if spin is not None: + return (*compact_sample, spin) + if want_coord: + # Tracing reads the shape of the positions alone; distinct values keep + # the sample nodes apart. + coord = 0.7 * torch.arange( + 3 * atype.shape[0], dtype=torch.float64, device=device + ).reshape(-1, 3) + return (*compact_sample, coord) + return compact_sample def _build_canonical_graph_dynamic_shapes( @@ -801,8 +813,9 @@ def _build_canonical_graph_dynamic_shapes( ) -> tuple: """Build dynamic shapes for the compact deployment ABI. - The trailing spin slot is present only for a native-spin model, so the - sample length selects the shape tuple. + The trailing node slot, the moments of a native-spin model or the + positions of a model with a long-range module, is present only for such + models, so the sample length selects the shape tuple. """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, @@ -1348,10 +1361,12 @@ def _probe_has_message_passing(obj: object) -> bool | None: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break meta["pair_exclude_types"] = pair_exclude_types - # Topologies the graph lower reads beside the neighbor graph, laid out - # after its other inputs (see ``flatten_topologies``); the caller of the - # lower builds them from the coordinates of each frame. Absent when the - # model reads none. + # Topologies the lower reads beside the neighbor graph; absent when the + # model reads none. The graph lower takes each as its ``(index, vec, + # mask)`` after its other inputs (see ``flatten_topologies``), built by the + # caller from the coordinates of each frame. The compact canonical lower + # takes the node positions instead, from which its operators enumerate the + # pairs. topologies = model.atomic_model.required_topologies() if topologies: meta["topologies"] = list(topologies) @@ -2068,11 +2083,15 @@ def _trace_and_export_impl( while e_sample in (_forbidden | {nframes_sample, n_sample}): e_sample += 1 if canonical: + # The compact lower of a model with a long-range module reads the + # node positions, from which its operators enumerate the pairs. + want_coord = bool(model.atomic_model.required_topologies()) sample_inputs = build_synthetic_canonical_graph_inputs( model, e_sample, device=torch.device("cpu"), want_spin=is_native_spin, + want_coord=want_coord, ) traced = model.forward_lower_canonical_graph_exportable( *sample_inputs, @@ -2531,7 +2550,81 @@ def _charge_state_descriptor(data: dict, metadata: dict) -> Any | None: return descriptor -def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: +def _charge_state_long_range(data: dict, metadata: dict) -> Any | None: + """Return the long-range module folded to the charge state, if any. + + Compressing the atomic model folds a long-range charge constraint that + reads its total charge from the charge state to the condition of the + descriptor, so the fold rebuilds the state of the module with the + artifacts of the descriptor. Only a model with a long-range module, which + the collected metadata names by its topologies, pays for rebuilding the + model. + + Parameters + ---------- + data + Serialized model dictionary, as passed to the export entry point. + metadata + Archive metadata collected from that same model. + + Returns + ------- + Any or None + The evaluated long-range module, or ``None`` when the model folds no + long-range state. + """ + if not metadata.get("topologies"): + return None + + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = BaseModel.deserialize(data["model"]) + parts = model.atomic_model.fused_decomposition() + long_range = None if parts is None else parts[0].long_range + if long_range is None or not long_range.charge_state_constants(): + return None + model.to("cpu") + model.eval() + return long_range + + +class _ChargeStateRebuild(torch.nn.Module): + """Rebuild the charge state of a descriptor and of a long-range module. + + The outputs are those of the descriptor fold followed by the folded state + of the long-range module. + + Parameters + ---------- + descriptor + Compressed charge-conditioned descriptor, as returned by + :func:`_charge_state_descriptor`. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`. + """ + + def __init__(self, descriptor: Any, long_range: Any) -> None: + super().__init__() + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + ChargeStateFold, + ) + + self.descriptor = ChargeStateFold(descriptor) + self.long_range = long_range + + def forward(self, charge_spin: torch.Tensor) -> tuple[torch.Tensor, ...]: + return ( + *self.descriptor(charge_spin), + *self.long_range.rebuild_charge_state(charge_spin), + ) + + +def _match_charge_state_constants( + descriptor: Any, exported: Any, long_range: Any | None = None +) -> list[str]: """Name, per fold output, the constant of ``exported`` it replaces. A compressed charge-conditioned descriptor carries its frame condition in @@ -2564,6 +2657,10 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: :func:`_charge_state_descriptor`. exported Exported lower whose constants are to be named. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state is matched after the + artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2584,9 +2681,18 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: for name, tensor in exported.state_dict.items() if isinstance(tensor, torch.Tensor) } + artifacts = [ + (name, getattr(descriptor, f"compress_{name}")) + for name in CHARGE_STATE_ARTIFACTS + ] + if long_range is not None: + artifacts += [ + ("long-range charge state", state) + for state in long_range.charge_state_constants() + ] constants: list[str] = [] - for name in CHARGE_STATE_ARTIFACTS: - artifact = getattr(descriptor, f"compress_{name}").detach().cpu() + for name, tensor in artifacts: + artifact = tensor.detach().cpu() if artifact.numel() == 0: constants.append("") continue @@ -2611,6 +2717,7 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: def _compile_charge_state_fold( descriptor: Any, aoti_configs: dict, + long_range: Any | None = None, ) -> bytes: """Compile the rebuild of the charge-state artifacts as its own archive. @@ -2626,6 +2733,10 @@ def _compile_charge_state_fold( :func:`_charge_state_descriptor`. aoti_configs Inductor options the inference lower was compiled with. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state the fold rebuilds after + the artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2652,7 +2763,12 @@ def _compile_charge_state_fold( dtype=torch.float32, device="cpu", ) - fold = torch.export.export(ChargeStateFold(descriptor), (sample,)) + module = ( + ChargeStateFold(descriptor) + if long_range is None + else _ChargeStateRebuild(descriptor, long_range) + ) + fold = torch.export.export(module, (sample,)) if _env.DEVICE.type != "cpu": from torch.export.passes import ( move_to_device_pass, @@ -2785,11 +2901,12 @@ def _deserialize_to_file_pt2( # state becomes known. charge_state_bytes: bytes | None = None if charge_state_descriptor is not None: + long_range = _charge_state_long_range(data, metadata) metadata["charge_state_constants"] = _match_charge_state_constants( - charge_state_descriptor, exported + charge_state_descriptor, exported, long_range ) charge_state_bytes = _compile_charge_state_fold( - charge_state_descriptor, aoti_configs + charge_state_descriptor, aoti_configs, long_range ) # Second artifact: with-comm. Only for descriptors whose message diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 295f251662..6610f19530 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -122,11 +122,27 @@ are listed outside the compiled region and enter it as inputs of dynamic length. `dp --pt-expt freeze` exports the model, long-range module included, to -`.pt2` or `.pte`. A trained checkpoint (`.pt`) and the frozen model are -evaluated by `dp test` and the Python inference interface, which list the -pairs of each frame before calling the model. The C++ interface, and with it -LAMMPS, does not evaluate a model with a long-range module. +`.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a +`.pt2` archive while the networks of the long-range module stay as they are. +A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` +and the Python inference interface. The C++ interface, and with it LAMMPS, +does not evaluate a model with a long-range module. + +Compressing a descriptor conditioned on `charge_spin` specializes the model +to one charge state at a time, and a constraint that reads its total charge +from `charge_spin` follows the same state. The compressed archive serves any +charge state, which the Python interface applies at each call. + +A compressed DPA4C model frozen for a CUDA device evaluates the pair sum with +a dedicated operator, which reads the positions of the atoms and enumerates +their pairs itself: its memory grows linearly with the number of atoms, and +its cost with their square. The operator serves the `les` kernel and the +`sog` kernel of up to 32 Gaussians, with up to 8 charge channels, a charge +network without frame or atomic parameters, and no `pair_exclude_types`. Any +other model, like an uncompressed one, lists the pairs of each frame before +calling the model, so its memory grows with the square of the number of +atoms. The following are not supported for a model with a long-range module: -compression, multi-rank (domain-decomposed) inference, spin models, the -tabulated ZBL potential, and short-range bridging. +multi-rank (domain-decomposed) inference, spin models, the tabulated ZBL +potential, and short-range bridging. diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 5702e6c18c..52b4921d05 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -23,7 +23,8 @@ list( cpu/neighbor_search_cpu.cc dpa4c/ops.cc dpa4c/graph_compress_cpu.cc - ${DPA4C_CPU_ISA_SRC}) + ${DPA4C_CPU_ISA_SRC} + long_range/ops.cc) # The architecture level is pinned per unit rather than added to it: the project # may be configured with ENABLE_NATIVE_OPTIMIZATION, and a trailing -march wins, # so this is what keeps the fallback units free of instructions the dispatcher @@ -179,6 +180,7 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC} graph_fitting.cu edge_force_virial.cu + long_range/direct_sum.cu dpa1_graph_energy_force.cu dpa4/so2_conv.cu ${DPA4_SO2_CONV_KERNEL_SRC} @@ -245,6 +247,11 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) set_source_files_properties( ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC} PROPERTIES COMPILE_OPTIONS "--use_fast_math") + # The direct sum of the long range evaluates every pair with the special + # function unit; without denormal handling each of its three calls is a single + # instruction. + set_source_files_properties(long_range/direct_sum.cu + PROPERTIES COMPILE_OPTIONS "--use_fast_math") endif() if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI}) target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD}) diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu new file mode 100644 index 0000000000..47c5a0e330 --- /dev/null +++ b/source/op/pt/long_range/direct_sum.cu @@ -0,0 +1,1058 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernel of the latent-charge long-range direct sum (schema and +// formulas in ops.cc). +// +// Every node reads every node of its frame, so the memory is linear in the +// node count while the work is quadratic. One thread owns one node and walks +// the nodes of its frame in tiles staged in shared memory; the partner of a +// thread at a given step is the same for the whole block, so each tile entry +// is a broadcast read. Each pair is evaluated by both of its nodes, which +// keeps every sum in the registers of its owner: no atomic operation is +// needed, and the result does not depend on scheduling. +// +// Positions are carried relative to the first node of their frame as a +// double-float pair: a float32 head and the float32 tail of the remainder. +// The displacement of two nodes, formed as the difference of the heads plus +// the difference of the tails, keeps the float32 precision of the +// displacement itself wherever the frame lies, while every pair evaluates in +// float32; the far field of the erf kernel forms it from the heads alone (see +// evaluate_kernel). For the erf kernel the positions are scaled by alpha, so +// that the pair loop evaluates erf(x) / x and its slope without a +// coefficient, and the reduction restores the units. +// +// A thread sums one tile in float32 and adds the tile sums with compensated +// float32 arithmetic. The partner axis is split into contiguous chunks, so +// that a launch fills the device several times over; every chunk writes its +// sums, and a second pass adds the chunks in a fixed order in float64. +// +// The charge constraint and the per-frame virial reduce over the nodes of +// every frame, in slices whose float64 sums are added in a fixed order, so +// that every output is independent of scheduling. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +// Nodes per block, which is also the length of a tile of partners. +constexpr int kThreads = 128; +// Partners evaluated together, which share every coefficient load. +constexpr int kBatch = 4; +constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; +// Squared distances are evaluated no shorter than this, in the units of the +// staged positions, which keeps the kernels finite for a node paired with +// itself and for coincident padding nodes; both carry a zero weight. +constexpr float kMinSquaredDistance = 1e-12f; +// Fills of the resident block slots of the device that one launch aims for. +constexpr int kWaves = 8; +// Shortest chunk of the partner axis, in tiles. +constexpr int kMinChunkTiles = 2; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// Sums of one node: the potential of every channel, the force direction and +// the six independent components of the virial. +template +constexpr int sum_width() { + return kChannels + 3 + 6; +} + +// Running sum with Kahan compensation. +struct CompensatedSum { + float sum = 0.f; + float error = 0.f; + + __device__ __forceinline__ void add(float value) { + const float corrected = value - error; + const float total = sum + corrected; + error = (total - sum) - corrected; + sum = total; + } +}; + +// Index of the frame holding ``node``, or -1 for a padding node past every +// frame. ``frame_ptr`` holds ``frame_count + 1`` offsets. +__device__ __forceinline__ int frame_of(const int64_t* frame_ptr, + int frame_count, + int64_t node) { + if (node >= frame_ptr[frame_count]) { + return -1; + } + int low = 0; + int high = frame_count; + while (high - low > 1) { + const int middle = (low + high) / 2; + if (frame_ptr[middle] <= node) { + low = middle; + } else { + high = middle; + } + } + return low; +} + +// Writes the offsets of the frames on the node axis: frame_ptr[f] counts the +// nodes of the frames before f, for f in [0, frame_count]. One block scans +// the counts, each thread a contiguous run of frames. +constexpr int kOffsetThreads = 256; + +__global__ void __launch_bounds__(kOffsetThreads) + frame_offsets(const int64_t* __restrict__ n_node_per_frame, + int frame_count, + int64_t* __restrict__ frame_ptr) { + __shared__ int64_t run_total[kOffsetThreads]; + const int run = (frame_count + kOffsetThreads - 1) / kOffsetThreads; + const int begin = min(frame_count, static_cast(threadIdx.x) * run); + const int end = min(frame_count, begin + run); + int64_t total = 0; + for (int f = begin; f < end; ++f) { + total += n_node_per_frame[f]; + } + run_total[threadIdx.x] = total; + __syncthreads(); + // Inclusive scan of the run totals. + for (int stride = 1; stride < kOffsetThreads; stride *= 2) { + const int64_t addend = + threadIdx.x >= stride ? run_total[threadIdx.x - stride] : 0; + __syncthreads(); + run_total[threadIdx.x] += addend; + __syncthreads(); + } + int64_t offset = threadIdx.x == 0 ? 0 : run_total[threadIdx.x - 1]; + for (int f = begin; f < end; ++f) { + frame_ptr[f] = offset; + offset += n_node_per_frame[f]; + } + if (threadIdx.x == kOffsetThreads - 1) { + frame_ptr[frame_count] = run_total[kOffsetThreads - 1]; + } +} + +// Per-frame sums. The nodes of a frame are split into ``gridDim.y`` +// contiguous slices, one block of kFrameThreads each, and a finalizing kernel +// adds the float64 sums of the slices in slice order. +constexpr int kFrameThreads = 256; +// Nodes per slice of a frame that a slice pass aims for. +constexpr int kSliceNodes = 2048; +constexpr int kMaxSlices = 256; + +struct NodeRange { + int64_t begin; + int64_t end; +}; + +// The nodes of the slice of the frame that this block sums. +__device__ __forceinline__ NodeRange +slice_of_frame(const int64_t* __restrict__ frame_ptr) { + const int64_t frame_begin = frame_ptr[blockIdx.x]; + const int64_t frame_end = frame_ptr[blockIdx.x + 1]; + const int64_t length = (frame_end - frame_begin + gridDim.y - 1) / gridDim.y; + const int64_t begin = frame_begin + blockIdx.y * length; + return {begin, min(frame_end, begin + length)}; +} + +// Adds ``value`` over the threads of the block in a fixed order and writes the +// totals to the entry of this slice in ``slice_sum``. +template +__device__ __forceinline__ void write_slice_sum( + const double (&value)[kComponents], double* __restrict__ slice_sum) { + __shared__ double partial[kComponents][kFrameThreads]; +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] = value[k]; + } + __syncthreads(); + for (int stride = kFrameThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] += partial[k][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kComponents) { + slice_sum[(static_cast(blockIdx.x) * gridDim.y + blockIdx.y) * + kComponents + + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Per-frame statistics of the charge constraint: the shift of the +// constrained channel, the number of charged nodes, and the mean potential of +// the charged nodes in that channel. +constexpr int kStatShift = 0; +constexpr int kStatCount = 1; +constexpr int kStatMean = 2; +constexpr int kStatWidth = 3; + +// Sums, over one slice of one frame, the number of charged nodes and their raw +// charges in the constrained channel. +__global__ void __launch_bounds__(kFrameThreads) + charge_slices(const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + const NodeRange slice = slice_of_frame(frame_ptr); + double value[2] = {0.0, 0.0}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { + if (node_mask[node]) { + value[0] += 1.0; + value[1] += charge[node * channels + channel]; + } + } + write_slice_sum<2>(value, slice_sum); +} + +// The shift s_f = (sum_f c - Q_f) / max(n_f, 1) of the constrained channel of +// every frame, from the slice sums of charge_slices. +__global__ void charge_shift(const double* __restrict__ slice_sum, + int slice_count, + int frame_count, + const double* __restrict__ total_charge, + double* __restrict__ frame_stat) { + const int f = blockIdx.x * blockDim.x + threadIdx.x; + if (f >= frame_count) { + return; + } + double count = 0.0; + double sum = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + const double* entry = + slice_sum + (static_cast(f) * slice_count + slice) * 2; + count += entry[0]; + sum += entry[1]; + } + // A frame without a charged node has nothing to shift. + count = fmax(count, 1.0); + frame_stat[f * kStatWidth + kStatShift] = (sum - total_charge[f]) / count; + frame_stat[f * kStatWidth + kStatCount] = count; +} + +// Writes the double-float positions relative to the first node of each +// frame, scaled by alpha for the erf kernel, the frame of every node, and the +// charges of the pairs: the raw charges of the charged nodes, shifted in the +// constrained channel, and zero for every other node. ``alpha`` is a device +// scalar, or null for a kernel read in plain units; ``frame_stat`` is null +// without a constraint. +__global__ void prepare_nodes(const double* __restrict__ coord, + const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int frame_count, + int node_count, + int channels, + int channel, + const float* __restrict__ alpha, + const double* __restrict__ frame_stat, + float4* __restrict__ head, + float4* __restrict__ tail, + int* __restrict__ frame, + float* __restrict__ pair_charge) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= node_count) { + return; + } + const double scale = alpha == nullptr ? 1.0 : static_cast(*alpha); + const int f = frame_of(frame_ptr, frame_count, node); + double3 offset = make_double3(0.0, 0.0, 0.0); + if (f >= 0) { + const int64_t first = frame_ptr[f]; + offset.x = scale * (coord[3 * node + 0] - coord[3 * first + 0]); + offset.y = scale * (coord[3 * node + 1] - coord[3 * first + 1]); + offset.z = scale * (coord[3 * node + 2] - coord[3 * first + 2]); + } + const float hx = __double2float_rn(offset.x); + const float hy = __double2float_rn(offset.y); + const float hz = __double2float_rn(offset.z); + head[node] = make_float4(hx, hy, hz, 0.f); + tail[node] = make_float4(__double2float_rn(offset.x - hx), + __double2float_rn(offset.y - hy), + __double2float_rn(offset.z - hz), 0.f); + frame[node] = f; + const bool charged = f >= 0 && node_mask[node]; + const double shift = charged && frame_stat != nullptr + ? frame_stat[f * kStatWidth + kStatShift] + : 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t index = static_cast(node) * channels + c; + const double value = charged ? static_cast(charge[index]) : 0.0; + pair_charge[index] = + static_cast(c == channel ? value - shift : value); + } +} + +// K and K'(r) / r of a batch of pairs at the squared distances ``r2``. +// +// The erf kernel reads positions scaled by alpha, x = alpha r, and evaluates +// erf(x) / x and its slope (2 / sqrt(pi) exp(-x^2) - erf(x) / x) / x^2, with +// erf(x) = 1 - t P(t) exp(-x^2), t = 1 / (1 + p x), and the coefficients of +// Abramowitz and Stegun 7.1.26 (absolute error below 1.5e-7), so that one +// exponential serves the kernel and its slope. +// +// Beyond x^2 = kFarSquared the float32 kernel equals 1 / x and its slope +// -1 / x^3 bit for bit: exp(-x^2) and erfc(x) fall below half a unit in the +// last place of the terms they would correct. A warp whose batch lies there +// entirely skips the evaluation of erf, which most pairs of a large system +// do, and forms the displacements of the batch from the heads alone. The +// omitted tails move a displacement by at most 2^-24 (|x_i| + |x_j|), with x +// the alpha-scaled offset of a node from the first node of its frame: the +// rounding of float32 positions over the extent of the frame, against a +// distance of at least sqrt(kFarSquared). For a cluster of 100,000 carbon +// atoms at alpha = 0.5 per angstrom, with offsets of up to 26 units, the +// potential and the force keep the error of the exact displacements. +// +// The sum of Gaussians evaluates sum_k a_k 2^(-c_k r^2) and its slope +// sum_k s_k 2^(-c_k r^2), s_k = -2 ln(2) c_k a_k, from the staged +// coefficients ``{a_k, -c_k, s_k}``. +constexpr float kFarSquared = 20.25f; + +// Whether every pair of the batch of every thread of the warp lies in the far +// field of the erf kernel, judged from the head displacements, whose squared +// lengths it writes to ``r2``. The sum of Gaussians has no far field. +template +__device__ __forceinline__ bool far_batch(const float (&dx)[kBatch], + const float (&dy)[kBatch], + const float (&dz)[kBatch], + float (&r2)[kBatch]) { + if constexpr (kKernel == kSog) { + return false; + } else { + float nearest = std::numeric_limits::infinity(); +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])); + nearest = fminf(nearest, r2[b]); + } + return __all_sync(0xffffffffu, nearest >= kFarSquared); + } +} + +template +__device__ __forceinline__ void evaluate_kernel(const float (&r2)[kBatch], + const float4* gaussians, + int gaussian_count, + float (&value)[kBatch], + float (&slope)[kBatch]) { + if constexpr (kKernel == kLes) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + const float x = r2[b] * inverse; + const float t = __fdividef(1.f, fmaf(0.3275911f, x, 1.f)); + const float poly = + t * + fmaf(t, + fmaf(t, + fmaf(t, fmaf(t, 1.061405429f, -1.453152027f), 1.421413741f), + -0.284496736f), + 0.254829592f); + const float gaussian = exp2f(-1.44269504088896341f * r2[b]); + value[b] = fmaf(-poly, gaussian, 1.f) * inverse; + slope[b] = + fmaf(1.12837916709551257f, gaussian, -value[b]) * (inverse * inverse); + } + } else { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + value[b] = 0.f; + slope[b] = 0.f; + } + for (int k = 0; k < gaussian_count; ++k) { + const float4 coefficient = gaussians[k]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float weight = exp2f(coefficient.y * r2[b]); + value[b] = fmaf(coefficient.x, weight, value[b]); + slope[b] = fmaf(coefficient.z, weight, slope[b]); + } + } + } +} + +// Partners a tile holds besides those of the frame of the node. +enum Exclusion : int { + // Every partner belongs to the frame and differs from the node. + kNoExclusion = 0, + // The node itself is among the partners. + kExcludeSelf = 1, + // The partners may also lie in other frames. + kExcludeOutsideFrame = 2, +}; + +// The state of the node a thread owns. +template +struct OwnNode { + int index; + int low; // first node of its frame + int high; // past the last node of its frame + float4 head; + float4 tail; + float charge[kChannels]; +}; + +// Adds the unscaled potential, force direction and virial that the partners +// of one staged tile give the node to ``sum``. +template +__device__ __forceinline__ void accumulate_tile( + const float4* __restrict__ tile_head, + const float4* __restrict__ tile_tail, + const float* __restrict__ tile_charge, + const float4* __restrict__ gaussians, + int gaussian_count, + int start, + const OwnNode& own, + float (&sum)[sum_width()]) { +#pragma unroll 4 + for (int base = 0; base < kThreads; base += kBatch) { + float dx[kBatch]; + float dy[kBatch]; + float dz[kBatch]; + float r2[kBatch]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float4 other_head = tile_head[base + b]; + dx[b] = own.head.x - other_head.x; + dy[b] = own.head.y - other_head.y; + dz[b] = own.head.z - other_head.z; + } + float value[kBatch]; + float slope[kBatch]; + if (far_batch(dx, dy, dz, r2)) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + value[b] = inverse; + slope[b] = -inverse * (inverse * inverse); + } + } else { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float4 other_tail = tile_tail[base + b]; + dx[b] += own.tail.x - other_tail.x; + dy[b] += own.tail.y - other_tail.y; + dz[b] += own.tail.z - other_tail.z; + r2[b] = fmaxf(fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])), + kMinSquaredDistance); + } + evaluate_kernel(r2, gaussians, gaussian_count, value, slope); + } +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const int other = start + base + b; + float weight = 1.f; + if constexpr (kExclusion == kExcludeSelf) { + weight = other != own.index ? 1.f : 0.f; + } else if constexpr (kExclusion == kExcludeOutsideFrame) { + weight = (other != own.index && + static_cast(other - own.low) < + static_cast(own.high - own.low)) + ? 1.f + : 0.f; + } + float product = 0.f; +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + const float other_charge = + kExclusion == kNoExclusion + ? tile_charge[(base + b) * kChannels + c] + : weight * tile_charge[(base + b) * kChannels + c]; + sum[c] = fmaf(other_charge, value[b], sum[c]); + product = fmaf(own.charge[c], other_charge, product); + } + const float strength = product * slope[b]; + const float sx = strength * dx[b]; + const float sy = strength * dy[b]; + const float sz = strength * dz[b]; + sum[kChannels + 0] += sx; + sum[kChannels + 1] += sy; + sum[kChannels + 2] += sz; + sum[kChannels + 3] = fmaf(sx, dx[b], sum[kChannels + 3]); + sum[kChannels + 4] = fmaf(sy, dy[b], sum[kChannels + 4]); + sum[kChannels + 5] = fmaf(sz, dz[b], sum[kChannels + 5]); + sum[kChannels + 6] = fmaf(sx, dy[b], sum[kChannels + 6]); + sum[kChannels + 7] = fmaf(sx, dz[b], sum[kChannels + 7]); + sum[kChannels + 8] = fmaf(sy, dz[b], sum[kChannels + 8]); + } + } +} + +// Sums, for every node of one block and every partner of one chunk, the +// unscaled potential, force direction and virial of the node, and writes them +// to ``partial[chunk][component][node]``. +template +__global__ void __launch_bounds__(kThreads) + direct_sum_kernel(const float4* __restrict__ head, + const float4* __restrict__ tail, + const float* __restrict__ charge, + const int* __restrict__ frame, + const int64_t* __restrict__ frame_ptr, + int node_count, + const float* __restrict__ coefficient, + int gaussian_count, + float* __restrict__ partial) { + constexpr int kWidth = sum_width(); + __shared__ float4 tile_head[kThreads]; + __shared__ float4 tile_tail[kThreads]; + __shared__ float tile_charge[kThreads * kChannels]; + __shared__ float4 gaussians[kMaxGaussians]; + __shared__ int range[2]; + __shared__ bool single_frame; + + const int first = blockIdx.x * kThreads; + const int node = first + threadIdx.x; + const int own_frame = node < node_count ? frame[node] : -1; + + if constexpr (kKernel == kSog) { + for (int k = threadIdx.x; k < gaussian_count; k += kThreads) { + const float amplitude = coefficient[k]; + const float rate = coefficient[gaussian_count + k]; + gaussians[k] = make_float4(amplitude, -rate, + -1.38629436111989061f * rate * amplitude, 0.f); + } + } + // The partners of the block are the nodes of the frames its nodes belong + // to: a contiguous range, since frames are contiguous on the node axis and + // padding nodes trail them. + if (threadIdx.x == 0) { + int last = min(first + kThreads, node_count) - 1; + while (last > first && frame[last] < 0) { + --last; + } + const int low_frame = frame[first]; + const int high_frame = frame[last]; + range[0] = low_frame < 0 ? 0 : static_cast(frame_ptr[low_frame]); + range[1] = low_frame < 0 ? 0 : static_cast(frame_ptr[high_frame + 1]); + // A block whose nodes all lie in one frame finds every partner of its + // range in that frame, which spares the frame test of every pair. + single_frame = low_frame >= 0 && low_frame == high_frame && + min(first + kThreads, node_count) - 1 == last; + } + __syncthreads(); + + // The chunk of the partner range this block sums, in whole tiles. + const int64_t tile_count = (range[1] - range[0] + kThreads - 1) / kThreads; + const int begin = + range[0] + + static_cast(tile_count * blockIdx.y / gridDim.y) * kThreads; + const int end = min( + range[1], + range[0] + static_cast(tile_count * (blockIdx.y + 1) / gridDim.y) * + kThreads); + + OwnNode own; + own.index = node; + own.low = 0; + own.high = 0; + own.head = make_float4(0.f, 0.f, 0.f, 0.f); + own.tail = own.head; + if (own_frame >= 0) { + own.head = head[node]; + own.tail = tail[node]; + own.low = static_cast(frame_ptr[own_frame]); + own.high = static_cast(frame_ptr[own_frame + 1]); + } +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + own.charge[c] = own_frame >= 0 ? charge[node * kChannels + c] : 0.f; + } + + CompensatedSum total[kWidth]; + for (int start = begin; start < end; start += kThreads) { + const int partner = start + threadIdx.x; + const bool staged = partner < end; + // A slot past the chunk carries a zero charge and contributes nothing. + tile_head[threadIdx.x] = + staged ? head[partner] : make_float4(0.f, 0.f, 0.f, 0.f); + tile_tail[threadIdx.x] = + staged ? tail[partner] : make_float4(0.f, 0.f, 0.f, 0.f); +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + tile_charge[threadIdx.x * kChannels + c] = + staged ? charge[partner * kChannels + c] : 0.f; + } + __syncthreads(); + + float sum[kWidth]; +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + sum[k] = 0.f; + } + // In a single-frame block only the tiles that overlap the block's own + // nodes hold a node itself, so the tests are confined to them and to the + // blocks that span frames; the choice is uniform over the block. + if (!single_frame) { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } else if (start < first + kThreads && first < start + kThreads) { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } else { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + total[k].add(sum[k]); + } + __syncthreads(); + } + + if (node < node_count) { + float* out = + partial + static_cast(blockIdx.y) * kWidth * node_count + node; +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + out[static_cast(k) * node_count] = total[k].sum; + } + } +} + +// Adds the chunk sums of one component of one node in chunk order, one +// thread per component and node, and scales the result into the potential, +// the force or the per-node virial. With positions scaled by alpha, +// K(r) = alpha f(x), K'(r) / r = alpha^3 f'(x) / x and d = d' / alpha, so the +// potential carries alpha, the force alpha^2 and the virial alpha. +template +__global__ void reduce_chunks(const float* __restrict__ partial, + int node_count, + int chunk_count, + double prefactor, + const float* __restrict__ alpha, + float* __restrict__ potential, + float* __restrict__ force, + float* __restrict__ atom_virial) { + constexpr int kWidth = sum_width(); + const int64_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= static_cast(kWidth) * node_count) { + return; + } + const int component = static_cast(index / node_count); + const int node = static_cast(index - component * int64_t(node_count)); + double value = 0.0; + for (int chunk = 0; chunk < chunk_count; ++chunk) { + value += partial[static_cast(chunk) * kWidth * node_count + index]; + } + const double unit = alpha == nullptr ? 1.0 : static_cast(*alpha); + if (component < kChannels) { + potential[node * kChannels + component] = + static_cast(prefactor * unit * value); + } else if (component < kChannels + 3) { + force[node * 3 + component - kChannels] = + static_cast(-prefactor * unit * unit * value); + } else { + // Each pair term is split half to each of its two nodes. The six sums + // are xx, yy, zz, xy, xz, yz. + constexpr int kRow[6] = {0, 1, 2, 0, 0, 1}; + constexpr int kColumn[6] = {0, 1, 2, 1, 2, 2}; + const int pair = component - kChannels - 3; + const float entry = static_cast(-0.5 * prefactor * unit * value); + float* out = atom_virial + static_cast(node) * 9; + out[3 * kRow[pair] + kColumn[pair]] = entry; + out[3 * kColumn[pair] + kRow[pair]] = entry; + } +} + +// Sums, over one slice of one frame, the per-node virial and, with a +// constraint, the potential of the charged nodes in the constrained channel. +template +__global__ void __launch_bounds__(kFrameThreads) + virial_slices(const float* __restrict__ atom_virial, + const float* __restrict__ potential, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + constexpr int kComponents = kConstrained ? 10 : 9; + const NodeRange slice = slice_of_frame(frame_ptr); + double value[kComponents] = {}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { +#pragma unroll + for (int k = 0; k < 9; ++k) { + value[k] += atom_virial[node * 9 + k]; + } + if constexpr (kConstrained) { + if (node_mask[node]) { + value[9] += potential[node * channels + channel]; + } + } + } + write_slice_sum(value, slice_sum); +} + +// Adds the slice sums of virial_slices into the virial of every frame and, +// with a constraint, the mean potential of its charged nodes. +template +__global__ void virial_total(const double* __restrict__ slice_sum, + int slice_count, + float* __restrict__ virial, + double* __restrict__ frame_stat) { + constexpr int kComponents = kConstrained ? 10 : 9; + if (threadIdx.x >= kComponents) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * + kComponents + + threadIdx.x]; + } + if (threadIdx.x < 9) { + virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value); + } else { + double* stat = frame_stat + blockIdx.x * kStatWidth; + stat[kStatMean] = value / stat[kStatCount]; + } +} + +// Writes the atomic energy q_i . phi_i / 2 and the derivative of the energy +// with respect to the raw charges: w_i (phi_i - mean_f) in the constrained +// channel and w_i phi_i in the others, w_i marking the charged nodes. +__global__ void project_charges(const float* __restrict__ potential, + const float* __restrict__ pair_charge, + const bool* __restrict__ node_mask, + const int* __restrict__ frame, + const double* __restrict__ frame_stat, + int node_count, + int channels, + int channel, + double* __restrict__ atom_energy, + float* __restrict__ charge_gradient) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= node_count) { + return; + } + const int f = frame[node]; + const bool charged = f >= 0 && node_mask[node]; + const double mean = charged && frame_stat != nullptr + ? frame_stat[f * kStatWidth + kStatMean] + : 0.0; + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t index = static_cast(node) * channels + c; + const double phi = potential[index]; + energy += static_cast(pair_charge[index]) * phi; + const double gradient = charged ? phi : 0.0; + charge_gradient[index] = + static_cast(c == channel ? gradient - mean : gradient); + } + atom_energy[node] = 0.5 * energy; +} + +// Chunks of the partner axis for the given node count. The blocks of one +// launch all carry the same work, so the launch fills the resident block +// slots of the device several times over, which bounds the idle share of the +// last wave; no chunk is shorter than kMinChunkTiles tiles of a frame that +// spans every node. +int chunk_count_for(int node_count, int resident_blocks_per_multiprocessor) { + const int multiprocessors = + at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int node_blocks = (node_count + kThreads - 1) / kThreads; + const int64_t wanted_blocks = static_cast(kWaves) * + resident_blocks_per_multiprocessor * + multiprocessors; + const int wanted = + static_cast((wanted_blocks + node_blocks - 1) / node_blocks); + return std::max(1, std::min(wanted, node_blocks / kMinChunkTiles)); +} + +struct DirectSumBuffers { + const torch::Tensor& head; + const torch::Tensor& tail; + const torch::Tensor& charge; + const torch::Tensor& frame; + const torch::Tensor& frame_ptr; + const torch::Tensor& coefficient; + torch::Tensor& potential; + torch::Tensor& force; + torch::Tensor& atom_virial; +}; + +template +int resident_blocks_per_multiprocessor() { + // The occupancy is a property of the compiled kernel and of the device; a + // process drives one device architecture. + static const int blocks = [] { + int value = 0; + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &value, direct_sum_kernel, kThreads, 0)); + return std::max(1, value); + }(); + return blocks; +} + +template +void launch(const DirectSumBuffers& buffers, + int node_count, + double prefactor, + cudaStream_t stream) { + constexpr int kWidth = sum_width(); + const int chunk_count = chunk_count_for( + node_count, resident_blocks_per_multiprocessor()); + auto partial = + torch::empty({static_cast(chunk_count) * node_count * kWidth}, + buffers.charge.options()); + const dim3 grid((node_count + kThreads - 1) / kThreads, chunk_count); + direct_sum_kernel<<>>( + reinterpret_cast(buffers.head.data_ptr()), + reinterpret_cast(buffers.tail.data_ptr()), + buffers.charge.data_ptr(), buffers.frame.data_ptr(), + buffers.frame_ptr.data_ptr(), node_count, + buffers.coefficient.data_ptr(), + static_cast(buffers.coefficient.numel() / 2), + partial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t reduce_threads = static_cast(kWidth) * node_count; + reduce_chunks + <<((reduce_threads + 255) / 256), 256, 0, stream>>>( + partial.data_ptr(), node_count, chunk_count, prefactor, + kKernel == kLes ? buffers.coefficient.data_ptr() : nullptr, + buffers.potential.data_ptr(), buffers.force.data_ptr(), + buffers.atom_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void dispatch_channels(int channels, + const DirectSumBuffers& buffers, + int node_count, + double prefactor, + cudaStream_t stream) { + switch (channels) { + case 1: + return launch(buffers, node_count, prefactor, stream); + case 2: + return launch(buffers, node_count, prefactor, stream); + case 3: + return launch(buffers, node_count, prefactor, stream); + case 4: + return launch(buffers, node_count, prefactor, stream); + case 5: + return launch(buffers, node_count, prefactor, stream); + case 6: + return launch(buffers, node_count, prefactor, stream); + case 7: + return launch(buffers, node_count, prefactor, stream); + case 8: + return launch(buffers, node_count, prefactor, stream); + default: + TORCH_CHECK(false, "latent_charge_direct_sum: unsupported channel count ", + channels); + } +} + +std::tuple +latent_charge_direct_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + torch::Tensor coefficient, + int64_t kernel, + int64_t channel, + double prefactor) { + TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && + n_node_per_frame.is_cuda() && total_charge.is_cuda() && + coefficient.is_cuda(), + "latent_charge_direct_sum: inputs must be CUDA tensors"); + TORCH_CHECK(charge.device() == coord.device() && + node_mask.device() == coord.device() && + n_node_per_frame.device() == coord.device() && + total_charge.device() == coord.device() && + coefficient.device() == coord.device(), + "latent_charge_direct_sum: inputs must share one device"); + TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && + coord.scalar_type() == torch::kFloat64, + "latent_charge_direct_sum: coord must be float64 with shape " + "(N, 3)"); + TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && + charge.scalar_type() == torch::kFloat32, + "latent_charge_direct_sum: charge must be float32 with shape " + "(N, C)"); + TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && + node_mask.scalar_type() == torch::kBool, + "latent_charge_direct_sum: node_mask must be bool with shape " + "(N,)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "latent_charge_direct_sum: n_node_per_frame must be " + "one-dimensional"); + TORCH_CHECK(total_charge.dim() == 1 && + total_charge.size(0) == n_node_per_frame.size(0) && + total_charge.scalar_type() == torch::kFloat64, + "latent_charge_direct_sum: total_charge must be float64 with " + "one entry per frame"); + TORCH_CHECK( + coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat32, + "latent_charge_direct_sum: coefficient must be a float32 vector"); + const int64_t channels = charge.size(1); + TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, + "latent_charge_direct_sum: the channel count must lie in [1, ", + kMaxChannels, "], got ", channels); + TORCH_CHECK(channel >= -1 && channel < channels, + "latent_charge_direct_sum: the constrained channel must be -1 " + "or lie in [0, ", + channels, "), got ", channel); + TORCH_CHECK(coord.size(0) < std::numeric_limits::max() / kThreads, + "latent_charge_direct_sum: too many nodes"); + if (kernel == kLes) { + TORCH_CHECK(coefficient.numel() == 1, + "latent_charge_direct_sum: the erf kernel takes one " + "coefficient"); + } else if (kernel == kSog) { + TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && + coefficient.numel() <= 2 * kMaxGaussians, + "latent_charge_direct_sum: the Gaussian kernel takes two " + "coefficients per Gaussian, for at most ", + kMaxGaussians, " Gaussians"); + } else { + TORCH_CHECK(false, "latent_charge_direct_sum: unknown kernel ", kernel); + } + + const c10::cuda::CUDAGuard guard(coord.device()); + const int node_count = static_cast(coord.size(0)); + const int frame_count = static_cast(n_node_per_frame.size(0)); + const auto float_options = charge.options(); + const auto double_options = coord.options(); + if (node_count == 0 || frame_count == 0) { + return {torch::zeros({node_count, 1}, double_options), + torch::zeros({node_count, channels}, float_options), + torch::zeros({node_count, 3}, float_options), + torch::zeros({node_count, 9}, float_options), + torch::zeros({frame_count, 9}, float_options)}; + } + const bool constrained = channel >= 0; + const int constrained_channel = static_cast(channel); + // The kernels write every entry of the outputs. + auto atom_energy = torch::empty({node_count, 1}, double_options); + auto charge_gradient = torch::empty({node_count, channels}, float_options); + auto force = torch::empty({node_count, 3}, float_options); + auto atom_virial = torch::empty({node_count, 9}, float_options); + auto virial = torch::empty({frame_count, 9}, float_options); + + coord = coord.contiguous(); + charge = charge.contiguous(); + node_mask = node_mask.contiguous(); + total_charge = total_charge.contiguous(); + coefficient = coefficient.contiguous(); + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + auto frame_ptr = torch::empty({frame_count + 1}, counts.options()); + auto head = torch::empty({node_count, 4}, float_options); + auto tail = torch::empty({node_count, 4}, float_options); + auto frame = torch::empty({node_count}, coord.options().dtype(torch::kInt32)); + auto pair_charge = torch::empty({node_count, channels}, float_options); + auto potential = torch::empty({node_count, channels}, float_options); + auto frame_stat = torch::empty({constrained ? frame_count * kStatWidth : 0}, + double_options); + const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; + const int slice_count = static_cast(std::min( + kMaxSlices, + std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); + const dim3 slice_grid(frame_count, slice_count); + const auto stream = at::cuda::getCurrentCUDAStream(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>( + counts.data_ptr(), frame_count, frame_ptr.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained) { + auto slice_charge = torch::empty( + {static_cast(frame_count) * slice_count * 2}, double_options); + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), + frame_ptr.data_ptr(), static_cast(channels), + constrained_channel, slice_charge.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge.data_ptr(), slice_count, frame_count, + total_charge.data_ptr(), frame_stat.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // Alpha stays on the device, so that a call never waits for it. + prepare_nodes<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr.data_ptr(), frame_count, + node_count, static_cast(channels), constrained_channel, + kernel == kLes ? coefficient.data_ptr() : nullptr, + constrained ? frame_stat.data_ptr() : nullptr, + reinterpret_cast(head.data_ptr()), + reinterpret_cast(tail.data_ptr()), frame.data_ptr(), + pair_charge.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + const DirectSumBuffers buffers{head, tail, pair_charge, + frame, frame_ptr, coefficient, + potential, force, atom_virial}; + if (kernel == kLes) { + dispatch_channels(static_cast(channels), buffers, node_count, + prefactor, stream.stream()); + } else { + dispatch_channels(static_cast(channels), buffers, node_count, + prefactor, stream.stream()); + } + const int components = constrained ? 10 : 9; + auto slice_virial = torch::empty( + {static_cast(frame_count) * slice_count * components}, + double_options); + if (constrained) { + virial_slices<<>>( + atom_virial.data_ptr(), potential.data_ptr(), + node_mask.data_ptr(), frame_ptr.data_ptr(), + static_cast(channels), constrained_channel, + slice_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial.data_ptr(), slice_count, virial.data_ptr(), + frame_stat.data_ptr()); + } else { + virial_slices<<>>( + atom_virial.data_ptr(), potential.data_ptr(), + node_mask.data_ptr(), frame_ptr.data_ptr(), + static_cast(channels), constrained_channel, + slice_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial.data_ptr(), slice_count, virial.data_ptr(), + nullptr); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + project_charges<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + potential.data_ptr(), pair_charge.data_ptr(), + node_mask.data_ptr(), frame.data_ptr(), + constrained ? frame_stat.data_ptr() : nullptr, node_count, + static_cast(channels), constrained_channel, + atom_energy.data_ptr(), charge_gradient.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {atom_energy, charge_gradient, force, atom_virial, virial}; +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.impl("latent_charge_direct_sum", torch::kCUDA, + &latent_charge_direct_sum); +} diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc new file mode 100644 index 0000000000..7f98709fe1 --- /dev/null +++ b/source/op/pt/long_range/ops.cc @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Operator schema of the latent-charge long-range direct sum. +// +// The schema is declared here, unconditionally, while the CUDA kernel +// registers itself only against a CUDA-enabled PyTorch (direct_sum.cu), so +// that the Python front end can tell "library not loaded" from "this device +// has no kernel". +// +// For the nodes of one frame, the operator takes the raw charges c_i of the +// charge network in C float32 channels, the nodes that carry them (w_i = 1, +// and 0 elsewhere) and float64 positions r_i. The pairs carry the charges +// q_i = w_i c_i; when `channel` is a channel index rather than -1, that +// channel is shifted uniformly over the charged nodes of every frame so that +// it sums to the total charge Q_f of the frame, +// +// q_i^ch = w_i (c_i^ch - s_f), s_f = (sum_{j in f} w_j c_j^ch - Q_f) / +// max(n_f, 1), +// +// with n_f charged nodes in f. With the potential of the charges +// +// phi_i = prefactor * sum_{j != i} q_j K(r_ij), +// +// the operator returns +// +// atom_energy_i = q_i . phi_i / 2 (float64) +// charge_gradient_i = dE / dc_i = w_i (phi_i - [c = ch] m_f), m_f the +// mean of phi^ch over the charged nodes of f +// force_i = -prefactor * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij +// * d_ij +// W_i = -prefactor / 2 * sum_{j != i} (q_i . q_j) +// K'(r_ij) / r_ij * d_ij (x) d_ij +// W_f = sum_{i in f} W_i +// +// in float32 unless noted, with d_ij = r_i - r_j and the sums running over +// the nodes of the frame of i. The frames occupy the leading +// sum(n_node_per_frame) <= N nodes, which the caller ensures; the nodes past +// them are padding and join none. The frame energy is the sum of the atomic +// energies. `kernel` selects K: 0 is +// erf(alpha r) / r with `coefficient` holding alpha; 1 is a sum of Gaussians +// a_k 2^(-c_k r^2), c_k = log2(e) / (2 sigma_k^2), with `coefficient` holding +// the amplitudes a_k and then the rates c_k. `total_charge` holds Q_f in +// float64 and is read only with a constrained channel. The force and the +// virial hold the charges fixed; the response of the charges is the +// caller's, through the charge gradient. + +#include + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.def( + "latent_charge_direct_sum(Tensor coord, Tensor charge, " + "Tensor node_mask, Tensor n_node_per_frame, Tensor total_charge, " + "Tensor coefficient, int kernel, int channel, float prefactor) -> " + "(Tensor atom_energy, Tensor charge_gradient, Tensor force, " + "Tensor atom_virial, Tensor virial)"); +} + +// The operator serves inference, whose caller differentiates nothing through +// it; its outputs carry no gradient. +TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { + library.impl("latent_charge_direct_sum", + torch::CppFunction::makeFallthrough()); +} diff --git a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py new file mode 100644 index 0000000000..b02ce0758d --- /dev/null +++ b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Correctness of the direct-sum operator of the latent-charge long range. + +The CUDA kernel evaluates every pair in float32 from positions relative to +the first node of each frame, carried as float32 head and tail, and sums in +compensated float32 before a float64 reduction; the total-charge constraint +and the per-frame reductions sum in float64 in a fixed order. It is held to +the float64 evaluation over the explicit pairs, at a tolerance relative to +the scale of each output, and to the exact invariants of a pair sum: padding +nodes and frames do not interact, the forces of a frame cancel, and the +energy, force, charge gradient and virial follow from one another. +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.pt_expt.kernels.long_range import ( + _reference, + kernel_operator_arguments, + latent_charge_direct_sum, + op_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not op_available(), + reason="the direct-sum operator is CUDA only", +) + +KERNELS = { + "les": lambda: LESKernel(alpha=0.8), + "sog": lambda: SOGKernel( + amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2], precision="float32" + ), +} + + +def _system( + n_node: list[int], channels: int, padding: int, seed: int, offset: float +) -> tuple[torch.Tensor, ...]: + """Frames of random nodes far from the origin, with trailing padding nodes. + + Every seventh node carries no charge, and every frame a total charge. + """ + rng = np.random.default_rng(seed) + total = sum(n_node) + padding + coord = offset + rng.uniform(0.0, 4.0 * max(n_node) ** (1 / 3), (total, 3)) + charge = rng.normal(0.0, 0.5, (total, channels)) + device = torch.device("cuda") + node_mask = torch.ones(total, dtype=torch.bool, device=device) + node_mask[::7] = False + return ( + torch.tensor(coord, dtype=torch.float64, device=device), + torch.tensor(charge, dtype=torch.float32, device=device), + node_mask, + torch.tensor(n_node, dtype=torch.int64, device=device), + torch.tensor( + rng.integers(-2, 3, len(n_node)), dtype=torch.float64, device=device + ), + ) + + +def _evaluate( + system: tuple[torch.Tensor, ...], kernel_name: str, channel: int +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + coord, charge, node_mask, n_node, total_charge = system + kernel = KERNELS[kernel_name]() + actual = latent_charge_direct_sum( + coord, + charge, + node_mask, + n_node, + total_charge, + kernel, + channel, + COULOMB_CONSTANT, + ) + code, coefficient = kernel_operator_arguments(kernel) + expected = _reference( + coord, + charge.double(), + node_mask, + n_node, + total_charge, + coefficient.double(), + code, + channel, + COULOMB_CONSTANT, + ) + return actual, expected + + +NAMES = ("atom_energy", "charge_gradient", "force", "atom_virial", "virial") + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel"), + [ + ([37, 1, 0, 300], 1, 5, 0), + ([37, 1, 0, 300], 1, 5, -1), + ([2048], 3, 0, -1), + ([5, 129], 8, 3, 2), + ], +) +def test_matches_float64_pair_sum( + kernel_name: str, n_node: list[int], channels: int, padding: int, channel: int +) -> None: + """Every output agrees with the float64 pair sum.""" + system = _system(n_node, channels, padding, 7, 1.0e3) + actual, expected = _evaluate(system, kernel_name, channel) + for name, value, reference in zip(NAMES, actual, expected, strict=True): + assert value.dtype == ( + torch.float64 if name == "atom_energy" else torch.float32 + ) + scale = reference.abs().max().item() + error = (value.double() - reference.double()).abs().max().item() + assert error <= 2e-5 * scale, (name, error, scale) + # Padding nodes join no frame. + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_invariants(kernel_name: str) -> None: + """Frames are independent, forces cancel, and the virial is the pair virial.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + together = latent_charge_direct_sum( + coord, + charge, + node_mask, + n_node_t, + total_charge, + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + _, _, force, atom_virial, virial = together + # Each frame alone gives the same result. + first = n_node[0] + alone = latent_charge_direct_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + # Internal forces sum to zero per frame, and the virial of a frame is + # sum_i r_i (x) f_i. + for frame, (begin, end) in enumerate(((0, first), (first, sum(n_node)))): + frame_force = force[begin:end].double() + assert frame_force.sum(0).abs().max() <= 1e-5 * frame_force.abs().max() + position = coord[begin:end] - coord[begin] + expected = torch.einsum("ni,nj->ij", position, frame_force) + scale = expected.abs().max().item() + for value in ( + atom_virial[begin:end].double().sum(0).reshape(3, 3), + virial[frame].double().reshape(3, 3), + ): + torch.testing.assert_close(value, expected, rtol=0, atol=2e-5 * scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize("channel", [-1, 1]) +def test_reference_outputs_are_energy_derivatives( + kernel_name: str, channel: int +) -> None: + """The float64 force and charge gradient differentiate the energy.""" + coord, charge, node_mask, n_node, total_charge = _system([9, 6], 2, 2, 3, 0.0) + coord = coord.clone().requires_grad_(True) + charge = charge.double().requires_grad_(True) + kernel = KERNELS[kernel_name]() + code, coefficient = kernel_operator_arguments(kernel) + atom_energy, charge_gradient, force, _, _ = _reference( + coord, + charge, + node_mask, + n_node, + total_charge, + coefficient.double(), + code, + channel, + COULOMB_CONSTANT, + ) + grad_coord, grad_charge = torch.autograd.grad(atom_energy.sum(), (coord, charge)) + torch.testing.assert_close(force, -grad_coord, rtol=1e-9, atol=1e-9) + torch.testing.assert_close( + charge_gradient.detach(), grad_charge, rtol=1e-9, atol=1e-9 + ) + + +def test_bitwise_reproducible() -> None: + """Two evaluations of the same input agree bit for bit.""" + system = _system([3000], 1, 0, 5, 0.0) + first = latent_charge_direct_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT) + second = latent_charge_direct_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT) + for a, b in zip(first, second, strict=True): + assert torch.equal(a, b) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 751451373e..02f5ed6587 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -11,6 +11,7 @@ COULOMB_CONSTANT, ) from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, build_pair_graph, ) from deepmd.infer import ( @@ -603,6 +604,139 @@ def test_frozen_model_matches_direct_forward( ) +@pytest.mark.parametrize( + ("kernel", "source"), [("les", None), ("les", "neutral"), ("sog", "charge_spin")] +) +def test_coordinate_evaluation_matches_pair_graph( + kernel: str, source: str | None +) -> None: + """The evaluation from node positions matches the pair-graph evaluation. + + It returns the atomic energy, the derivative of the energy with respect to + the descriptor through the constrained charges, and the force at a fixed + descriptor, all of which the pair graph yields by automatic + differentiation. Two frames, an uncharged node and trailing padding nodes + exercise the frame sums of the constraint. + """ + from deepmd.pt_expt.kernels.long_range import ( + op_available, + ) + + if not op_available(): + pytest.skip("the direct-sum operator is not available") + config = _config(kernel, constraint=None if source is None else {"source": source}) + if source == "charge_spin": + config["descriptor"]["add_chg_spin_ebd"] = True + long_range = get_model(config).atomic_model.long_range.eval() + if source == "charge_spin": + long_range.fold_charge_state([2.0, 1.0]) + rng = np.random.default_rng(7) + n_node = torch.tensor([9, 6], device=env.DEVICE) + node_mask = torch.ones(17, dtype=torch.bool, device=env.DEVICE) + node_mask[4] = False + node_mask[15:] = False + coord = _tensor(rng.uniform(0.0, 6.0, (17, 3))) + descriptor = _tensor(rng.normal(size=(17, long_range.charge_net.dim_descrpt))) + atype = _tensor(rng.integers(0, 2, 17), torch.int64) + + atom_energy, descriptor_gradient, force, _, _ = long_range.energy_force_from_coord( + descriptor, atype, node_mask, n_node, coord + ) + + leaf_descriptor = descriptor.clone().requires_grad_(True) + leaf_coord = coord.clone().requires_grad_(True) + pairs = build_pair_graph(leaf_coord, n_node, atype, node_mask) + expected = long_range.call_graph( + leaf_descriptor, atype, node_mask, n_node, {ALL_PAIRS: pairs} + ) + expected_descriptor_gradient, expected_coord_gradient = torch.autograd.grad( + expected.sum(), (leaf_descriptor, leaf_coord) + ) + for actual, target in ( + (atom_energy, expected), + (descriptor_gradient, expected_descriptor_gradient), + (force, -expected_coord_gradient), + ): + target = target.detach() + scale = target.abs().max().item() + assert scale > 0.0 + np.testing.assert_allclose( + actual.detach().to(torch.float64).cpu().numpy(), + target.cpu().numpy(), + rtol=0.0, + atol=2e-5 * scale, + ) + + +@pytest.mark.parametrize( + ("kernel", "source"), [("les", "charge_spin"), ("sog", "neutral")] +) +def test_compressed_frozen_model_matches_uncompressed( + tmp_path, kernel: str, source: str +) -> None: + """A compressed model keeps its long-range module through the freeze. + + On a CUDA device the freeze selects the compact canonical lower, whose + direct-sum operator reads the positions of the atoms; on the CPU the + graph lower reads the pair lists. A constraint that reads the total + charge from the charge state is folded with the descriptor and follows + the state the deployment serves. Both match the uncompressed model. + """ + config = _config(kernel, constraint={"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + if source == "charge_spin": + config["descriptor"].update(add_chg_spin_ebd=True, default_chg_spin=[0.0, 1.0]) + config["fitting_net"].update( + neuron=[32, 32], activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23} + model = get_model(config).to(env.DEVICE).eval() + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + reference = copy.deepcopy(model) + model.atomic_model.enable_compression(0.5) + frozen = tmp_path / f"long_range_compressed_{kernel}.pt2" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="auto", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" + assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 + atype = rng.integers(0, 2, (1, 40)) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + # The tables of the compressed descriptor approximate its networks. On the + # graph lower the charges of this model amplify that error at a charged + # state, to about 1e-4 of the scale at [1, 2]; the compact canonical lower + # stays within a few 1e-6. + tolerance = 2e-5 if expected_lower == "dpa4c_canonical" else 5e-4 + states = [[0.0, 1.0], [1.0, 2.0], [0.0, 1.0]] if source == "charge_spin" else [None] + for state in states: + condition = {} if state is None else {"charge_spin": np.array([state])} + expected = reference( + _tensor(coord), + _tensor(atype, torch.int64), + do_atomic_virial=True, + **{key: _tensor(value) for key, value in condition.items()}, + ) + actual = dp.eval(coord, None, atype, atomic=True, **condition) + for name, value in zip(names, actual, strict=True): + target = expected[name].detach().cpu().numpy().reshape(value.shape) + scale = np.abs(target).max() + assert np.abs(value - target).max() <= tolerance * scale + 1e-5, ( + name, + state, + ) + + @pytest.mark.parametrize("kernel", ["les", "sog"]) @REQUIRES_SUPPORTED_COMPILE def test_compiled_matches_eager(kernel: str) -> None: @@ -648,8 +782,6 @@ def test_unsupported_combinations_fail() -> None: nlist = torch.full((1, 6, 4), -1, dtype=torch.int64, device=env.DEVICE) with pytest.raises(NotImplementedError, match="NeighborGraph"): model.atomic_model.forward_atomic(coord[:1], atype[:1], nlist) - with pytest.raises(NotImplementedError, match="compression"): - model.atomic_model.enable_compression(0.5) config = _config("les") config["spin"] = {"use_spin": [True, False], "scheme": "native"} with pytest.raises(NotImplementedError): @@ -691,6 +823,23 @@ def test_unsupported_combinations_fail() -> None: get_model(config) +def test_compressing_the_descriptor_alone_is_rejected() -> None: + """A folded descriptor and a constraint reading the condition are rejected.""" + config = _config("les", constraint={"source": "charge_spin"}) + config["descriptor"].update( + n_radial=8, + precision="float32", + add_chg_spin_ebd=True, + default_chg_spin=[0.0, 1.0], + ) + config["fitting_net"]["precision"] = "float32" + model = get_model(config).to(env.DEVICE).eval() + model.atomic_model.descriptor.enable_compression(0.5) + coord, atype = _batch() + with pytest.raises(ValueError, match="folds the charge/spin condition"): + model(coord[:1], atype[:1]) + + @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_dpmodel_reference_matches(kernel: str) -> None: """The dpmodel reference evaluates the same energy from the portable record.""" From c8ba94fe37a71825b20c7973d214026908c5d948 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 01:04:15 +0800 Subject: [PATCH 16/30] feat(api_cc): evaluate long-range models in C++ and LAMMPS Serve a .pt2 model with a latent-charge long-range module in the C++ interface and in `pair_style deepmd` and `deepmd/kk`, for complete non-periodic systems on a single MPI rank. - DeepPotPTExpt builds the `all_pairs` topology of the graph lower from the coordinates of each frame for an uncompressed model, and passes the node positions to the compact canonical lower of a compressed one. A periodic box, ghost atoms and multiple ranks are rejected. - `DeepPot::has_long_range()` and a `compute_canonical_graph_gpu` overload that takes the device coordinates; C API version 31 adds `DP_DeepPotHasLongRange` and `DP_DeepPotComputeCanonicalGraphGPU2`. The backend virtual signature changes, so the DeepPot backend plugin symbols become `deepmd_create_deeppot_backend_v2` and `deepmd_delete_deeppot_backend_v2`. - `PairDeepMD::init_style`, shared by `deepmd/kk`, rejects a long-range model with a periodic dimension or more than one rank; the Kokkos path hands the model node positions to the canonical lower. - Tests: gen_dpa4c_long_range.py writes an uncompressed fixture and, for the CUDA variant, a compressed one; C++, C API and LAMMPS cases compare energies and forces with the Python interface, follow `set_charge_spin`, and check the rejections. --- doc/model/long-range.md | 10 +- source/api_c/include/c_api.h | 52 +++- source/api_c/include/deepmd.hpp | 34 +++ source/api_c/src/c_api.cc | 30 +++ .../tests/test_deeppot_long_range_hpp.cc | 64 +++++ source/api_cc/include/BackendPlugin.h | 10 +- source/api_cc/include/DeepPot.h | 40 +++ source/api_cc/include/DeepPotPTExpt.h | 22 +- source/api_cc/src/DeepPot.cc | 32 ++- source/api_cc/src/DeepPotJAXPlugin.cc | 4 +- source/api_cc/src/DeepPotPDPlugin.cc | 4 +- source/api_cc/src/DeepPotPTExpt.cc | 112 ++++++-- source/api_cc/src/DeepPotPTExptPlugin.cc | 4 +- source/api_cc/src/DeepPotPTPlugin.cc | 4 +- source/api_cc/src/DeepPotTFPlugin.cc | 4 +- source/api_cc/src/longRangePTExpt.h | 59 ++++ .../test_deeppot_dpa4c_long_range_ptexpt.cc | 253 ++++++++++++++++++ source/install/test_cc_local.sh | 7 + source/lmp/pair_deepmd.cpp | 13 + source/lmp/pair_deepmd.h | 1 + source/lmp/pair_deepmd_kokkos.cpp | 46 +++- source/lmp/pair_deepmd_kokkos.h | 7 + .../tests/test_lammps_dpa4c_long_range_pt2.py | 131 +++++++++ source/tests/infer/gen_dpa4c_long_range.py | 238 ++++++++++++++++ 24 files changed, 1133 insertions(+), 48 deletions(-) create mode 100644 source/api_c/tests/test_deeppot_long_range_hpp.cc create mode 100644 source/api_cc/src/longRangePTExpt.h create mode 100644 source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc create mode 100644 source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py create mode 100644 source/tests/infer/gen_dpa4c_long_range.py diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 6610f19530..67d31e079e 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -125,13 +125,17 @@ length. `.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a `.pt2` archive while the networks of the long-range module stay as they are. A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` -and the Python inference interface. The C++ interface, and with it LAMMPS, -does not evaluate a model with a long-range module. +and the Python inference interface; the C++ interface and LAMMPS, with +`pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. +Since the pair sum covers the complete system, LAMMPS runs such a model with +non-periodic boundaries (`f`, `s` or `m`) in every dimension on a single MPI +rank, and rejects any other setting. Compressing a descriptor conditioned on `charge_spin` specializes the model to one charge state at a time, and a constraint that reads its total charge from `charge_spin` follows the same state. The compressed archive serves any -charge state, which the Python interface applies at each call. +charge state: the Python interface applies the state of each call, and LAMMPS +the one given by the `charge_spin` keyword of `pair_style deepmd`. A compressed DPA4C model frozen for a CUDA device evaluates the pair sum with a dedicated operator, which reads the positions of the atoms and enumerates diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index d98d744387..44f70b9318 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -13,7 +13,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 30 +#define DP_C_API_VERSION 31 /** * @brief Neighbor list. @@ -479,6 +479,56 @@ extern bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp); */ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); +/** + * @brief Evaluate a compact canonical graph and the node positions on the + * model device. + * + * A model with a long-range module reads the positions of the nodes, which + * must form a complete non-periodic system in a single domain + * (``nall_nodes == nloc``); any other model ignores them. + * + * @param[in] dp The DP to use. + * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``. + * @param[out] d_force Per-node force, shape ``(nall_nodes, 3)``. + * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``. + * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``. + * @param[in] d_coord Node positions in model-node order, shape + * ``(nall_nodes, 3)``, or NULL for a model without a long-range module. + * @param[in] d_source Source node per edge storage slot. + * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``. + * @param[in] d_destination_row_ptr Destination CSR offsets. + * @param[in] d_source_row_ptr Source CSR offsets. + * @param[in] d_source_order Source-grouped edge storage positions. + * @param[in] nloc Number of owned local nodes. + * @param[in] nall_nodes Total local-plus-halo node count. + * @param[in] edge_storage Number of edge storage slots. + * @note Supersedes DP_DeepPotComputeCanonicalGraphGPU, adding the node + * positions. + * @since API version 31 + */ +extern void DP_DeepPotComputeCanonicalGraphGPU2( + DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const double* d_coord, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + int nloc, + int nall_nodes, + int64_t edge_storage); + +/** + * @brief Query whether the model carries a long-range module, which + * evaluates complete non-periodic systems in a single domain only. + * @since API version 31 + */ +extern bool DP_DeepPotHasLongRange(DP_DeepPot* dp); + /** * @brief Evaluate the energy, force and virial by using a DP. (double version) * @version 2 diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 681b197b0e..9aee9e8857 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1406,6 +1406,40 @@ class DeepPot : public DeepBaseModel { return result; } + /** + * @brief Evaluate a compact canonical graph and the node positions on the + * model device; see DP_DeepPotComputeCanonicalGraphGPU2. + */ + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const double* d_coord, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_DeepPotComputeCanonicalGraphGPU2( + dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, + nloc, nall_nodes, edge_storage); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + + /** + * @brief Query whether the model carries a long-range module, which + * evaluates complete non-periodic systems in a single domain only. + */ + bool has_long_range() const { + const bool result = DP_DeepPotHasLongRange(dp); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + return result; + } + /** * @brief Evaluate the energy, force and virial by using this DP. * @param[out] ener The system energy. diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index c35a4a1763..0391429e61 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1864,6 +1864,36 @@ bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp) { } } +void DP_DeepPotComputeCanonicalGraphGPU2(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const double* d_coord, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_REQUIRES_OK( + dp, dp->dp.compute_canonical_graph_gpu( + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage)); +} + +bool DP_DeepPotHasLongRange(DP_DeepPot* dp) { + try { + return dp->dp.has_long_range(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + // multiple frames void DP_DeepPotCompute2(DP_DeepPot* dp, const int nframes, diff --git a/source/api_c/tests/test_deeppot_long_range_hpp.cc b/source/api_c/tests/test_deeppot_long_range_hpp.cc new file mode 100644 index 0000000000..1f16483156 --- /dev/null +++ b/source/api_c/tests/test_deeppot_long_range_hpp.cc @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test the C API with a pt_expt (.pt2) DPA4C model that carries a +// latent-charge long-range module (source/tests/infer/gen_dpa4c_long_range.py). +#include + +#include +#include + +#include "../../api_cc/tests/expected_ref.h" +#include "DeepPotPTExpt.h" +#include "deepmd.hpp" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4c_long_range.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4c_long_range.expected"; +} // namespace + +template +class TestInferDeepPotLongRangeHPP : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector expected_f; + double expected_tot_e = 0.; + deepmd::hpp::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + dp.init(kModelPath); + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_f = ref.get("default", "expected_f"); + for (const VALUETYPE value : ref.get("default", "expected_e")) { + expected_tot_e += value; + } + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotLongRangeHPP, ValueTypes); + +TYPED_TEST(TestInferDeepPotLongRangeHPP, has_long_range) { + EXPECT_TRUE(this->dp.has_long_range()); +} + +TYPED_TEST(TestInferDeepPotLongRangeHPP, nopbc) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector()); + ASSERT_EQ(force.size(), this->expected_f.size()); + EXPECT_LT(fabs(ener - this->expected_tot_e), epsilon); + for (std::size_t ii = 0; ii < force.size(); ++ii) { + EXPECT_LT(fabs(force[ii] - this->expected_f[ii]), epsilon); + } +} diff --git a/source/api_cc/include/BackendPlugin.h b/source/api_cc/include/BackendPlugin.h index 4149917414..272ec8e679 100644 --- a/source/api_cc/include/BackendPlugin.h +++ b/source/api_cc/include/BackendPlugin.h @@ -15,11 +15,17 @@ class DipoleChargeModifierBase; /** * @brief C ABI symbol names exported by backend plugin libraries. + * + * The suffix of a symbol is the version of the interface it serves: that of + * the DeepPot and DeepSpin symbols changes with the virtual interface of + * DeepPotBackend and DeepSpinBackend, so that a plugin built against another + * interface is refused when it is loaded rather than called through a table + * of another layout. **/ constexpr const char* DEEPMD_DEEPPOT_PLUGIN_CREATE_SYMBOL = - "deepmd_create_deeppot_backend_v1"; + "deepmd_create_deeppot_backend_v2"; constexpr const char* DEEPMD_DEEPPOT_PLUGIN_DELETE_SYMBOL = - "deepmd_delete_deeppot_backend_v1"; + "deepmd_delete_deeppot_backend_v2"; constexpr const char* DEEPMD_DEEPSPIN_PLUGIN_CREATE_SYMBOL = "deepmd_create_deepspin_backend_v1"; constexpr const char* DEEPMD_DEEPSPIN_PLUGIN_DELETE_SYMBOL = diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 575c7924c9..c39466ee52 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -426,6 +426,7 @@ class DeepPotBackend : public DeepBaseModelBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -437,6 +438,7 @@ class DeepPotBackend : public DeepBaseModelBackend { virtual bool supports_device_edge_inference() const; virtual bool uses_fp32_edge_vectors() const; virtual bool uses_canonical_graph_inference() const; + virtual bool has_long_range() const; }; /** @@ -894,10 +896,48 @@ class DeepPot : public DeepBaseModel { */ bool uses_canonical_graph_inference() const; + /** + * @brief Whether the model carries a long-range module. + * + * Such a model sums its long-range interaction over every pair of atoms of + * a complete non-periodic system, so it evaluates a system without a + * periodic cell in a single domain only, and the device entries take the + * positions of its atoms. + */ + bool has_long_range() const; + + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage); + + /** + * @brief Evaluate a compact canonical graph together with the positions of + * its nodes. + * + * Every pointer resides on the model device. A model with a long-range + * module reads the positions; any other model ignores them. + * + * @param[in] d_coord Node positions in model-node order, shape + * ``(nall_nodes, 3)``, in the length unit of the model. A model with a + * long-range module requires them and a single-domain node set + * (``nall_nodes == nloc``) of a system without a periodic cell. + * + * The other parameters are those of the overload without positions. + */ void compute_canonical_graph_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 4707ee9ce6..05cf7b2767 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -377,6 +377,7 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -413,6 +414,7 @@ class DeepPotPTExpt : public DeepPotBackend { */ bool uses_fp32_edge_vectors() const override; bool uses_canonical_graph_inference() const override; + bool has_long_range() const override; private: template @@ -434,6 +436,7 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -473,6 +476,10 @@ class DeepPotPTExpt : public DeepPotBackend { bool lower_input_is_graph_ = false; bool lower_input_is_canonical_ = false; bool graph_edge_fp32_ = false; + // The model carries a long-range module, whose lower reads the all-pairs + // topology of the frame after its other inputs: the pair lists for the graph + // lower, the node positions for the compact canonical lower. + bool long_range_ = false; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) std::vector mapping_; // cached mapping vector (LAMMPS path) @@ -615,6 +622,9 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] destination_row_ptr Destination CSR offsets ``(N + 1,)``. * @param[in] source_order Source-grouped edge permutation ``(E,)``. * @param[in] source_row_ptr Source CSR offsets ``(N + 1,)``. + * @param[in] coord Positions of the nodes of the single frame, shape + * ``(N, 3)`` float64, from which the all-pairs topology of a model with a + * long-range module is built; read only for such a model. */ std::vector run_model_graph( const torch::Tensor& atype, @@ -629,8 +639,15 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& source_row_ptr, const torch::Tensor& fparam, const torch::Tensor& aparam, - const torch::Tensor& charge_spin); + const torch::Tensor& charge_spin, + const torch::Tensor& coord); + /** + * @brief Run a compact canonical ``.pt2`` artifact. + * + * @param[in] coord Node positions, shape ``(N, 3)`` float64, appended for a + * model with a long-range module and ignored otherwise. + */ std::vector run_model_canonical_graph( const torch::Tensor& atype, const torch::Tensor& n_node, @@ -639,7 +656,8 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& edge_vec, const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, - const torch::Tensor& source_order); + const torch::Tensor& source_order, + const torch::Tensor& coord); /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 3373d6fe71..66ed2f05c6 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -661,6 +661,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -673,6 +674,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( (void)d_force; (void)d_atom_virial; (void)d_atype; + (void)d_coord; (void)d_source; (void)d_edge_vec; (void)d_destination_row_ptr; @@ -692,6 +694,8 @@ bool DeepPotBackend::supports_device_edge_inference() const { return false; } bool DeepPotBackend::uses_canonical_graph_inference() const { return false; } +bool DeepPotBackend::has_long_range() const { return false; } + void DeepPot::compute_edges_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, @@ -759,9 +763,29 @@ void DeepPot::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { dp->compute_canonical_graph_gpu( - d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec, - d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, - edge_storage); + d_atom_energy, d_force, d_atom_virial, d_atype, /*d_coord=*/nullptr, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage); +} + +void DeepPot::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const double* d_coord, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + dp->compute_canonical_graph_gpu( + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, + nall_nodes, edge_storage); } bool DeepPot::uses_fp32_edge_vectors() const { @@ -776,6 +800,8 @@ bool DeepPot::uses_canonical_graph_inference() const { return dp->uses_canonical_graph_inference(); } +bool DeepPot::has_long_range() const { return dp->has_long_range(); } + int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } void DeepPot::set_charge_spin(const std::vector& charge_spin) { diff --git a/source/api_cc/src/DeepPotJAXPlugin.cc b/source/api_cc/src/DeepPotJAXPlugin.cc index ac2d5485a1..2522148833 100644 --- a/source/api_cc/src/DeepPotJAXPlugin.cc +++ b/source/api_cc/src/DeepPotJAXPlugin.cc @@ -4,7 +4,7 @@ #include "BackendPluginFactory.h" #include "DeepPotJAX.h" -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -13,7 +13,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotPDPlugin.cc b/source/api_cc/src/DeepPotPDPlugin.cc index fbc920cfb5..3bbd6842de 100644 --- a/source/api_cc/src/DeepPotPDPlugin.cc +++ b/source/api_cc/src/DeepPotPDPlugin.cc @@ -4,7 +4,7 @@ #include "BackendPluginFactory.h" #include "DeepPotPD.h" -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -13,7 +13,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 1bbdc9b9bc..2c7d14a2a2 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -22,6 +22,7 @@ #include "commonPTExpt.h" #include "device.h" #include "errors.h" +#include "longRangePTExpt.h" #include "neighbor_list.h" using deepmd::ptexpt::check_call_charge_spin; @@ -256,12 +257,22 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_graph_ = false; lower_input_is_canonical_ = false; } - // A model with a long-range module reads pair topologies after the neighbor - // graph, which this interface does not build. - if (metadata.obj_val.count("topologies")) { - throw deepmd::deepmd_exception( - "this .pt2 holds a model with a long-range module, which the C++ " - "interface does not evaluate; use the Python inference interface."); + // A model with a long-range module reads the all-pairs topology of the + // frame after its other inputs. + long_range_ = metadata.obj_val.count("topologies") > 0; + if (long_range_) { + for (const auto& name : metadata["topologies"].as_array()) { + if (name.as_string() != "all_pairs") { + throw deepmd::deepmd_exception("unknown topology '" + name.as_string() + + "' in the .pt2 metadata; expected " + "'all_pairs'."); + } + } + if (!lower_input_is_graph_ && !lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "a model with a long-range module requires the graph or the " + "compact canonical lower."); + } } if (lower_input_is_edge_) { std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " @@ -520,7 +531,8 @@ std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& source_row_ptr, const torch::Tensor& fparam, const torch::Tensor& aparam, - const torch::Tensor& charge_spin) { + const torch::Tensor& charge_spin, + const torch::Tensor& coord) { // Graph-input ABI: original edge payload plus destination/source CSR views. check_graph_aparam_flat(aparam, daparam, "run_model_graph"); std::vector inputs = { @@ -536,6 +548,12 @@ std::vector DeepPotPTExpt::run_model_graph( if (dchgspin > 0) { inputs.push_back(charge_spin); } + if (long_range_) { + const auto pairs = deepmd::ptexpt::buildAllPairsTopology( + coord, atype, pair_exclude_table_, ntypes, + graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64); + inputs.insert(inputs.end(), pairs.begin(), pairs.end()); + } return loader->run(inputs); } @@ -547,9 +565,16 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( const torch::Tensor& edge_vec, const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, - const torch::Tensor& source_order) { - return loader->run({atype, n_node, n_local, source, edge_vec, - destination_row_ptr, source_row_ptr, source_order}); + const torch::Tensor& source_order, + const torch::Tensor& coord) { + std::vector inputs = {atype, n_node, + n_local, source, + edge_vec, destination_row_ptr, + source_row_ptr, source_order}; + if (long_range_) { + inputs.push_back(coord); + } + return loader->run(inputs); } std::vector DeepPotPTExpt::run_model_with_comm( @@ -806,6 +831,17 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; + // A long-range module sums over every pair of atoms of the system, which + // therefore has to be complete in this call: without periodic images and + // without other domains, both of which appear as ghost atoms. + if (long_range_ && (multi_rank || nghost_real > 0)) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a complete non-periodic " + "system in a single domain, but this call carries " + + std::to_string(nghost_real) + " ghost atoms on " + + std::to_string(lmp_list.nprocs) + + " domain(s); use a non-periodic box on a single MPI rank."); + } // Whether the edge topology folds ghost neighbours onto their local owners, // which reads an owner for every extended atom out of ``mapping_``. The // edge lower folds unless the with-comm artifact carries ghost features @@ -1275,19 +1311,21 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, graph_pack.atype = node_atype; graph_pack.n_node = n_node_tensor; graph_pack.n_local = n_local_tensor; + const at::Tensor node_coord = + coord_Tensor.slice(1, 0, n_node_count).reshape({n_node_count, 3}); if (lower_input_is_canonical_) { const auto compact = compactCanonicalGraph(graph_pack); flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, - compact.source_row_ptr, compact.source_order); + compact.source_row_ptr, compact.source_order, node_coord); } else { flat_outputs = run_model_graph( node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor); + graph_aparam, charge_spin_tensor, node_coord); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -1463,6 +1501,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, throw deepmd::deepmd_exception( "The coordinate-based inference interface requires at least one atom."); } + // The long-range module sums over the atoms of the system alone. + if (long_range_ && !box.empty()) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates complete non-periodic " + "systems only; pass an empty box."); + } const std::size_t frame_coord_size = static_cast(natoms) * 3; if (coord.empty() || coord.size() % frame_coord_size != 0) { throw deepmd::deepmd_exception( @@ -1655,12 +1699,14 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (graph_edge_fp32_) { graph_tensors.edge_vec = graph_tensors.edge_vec.to(torch::kFloat32); } + const at::Tensor node_coord = + coord_Tensor.slice(1, 0, nloc).reshape({nloc, 3}); if (lower_input_is_canonical_) { const auto compact = compactCanonicalGraph(graph_tensors); flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, - compact.source_order); + compact.source_order, node_coord); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -1672,7 +1718,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // normalizes the rectangular runtime aparam to the flat // (N, daparam) node axis of the graph ABI. extend_graph_aparam(aparam_tensor, natoms, natoms, daparam), - charge_spin_tensor); + charge_spin_tensor, node_coord); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -2228,6 +2274,12 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, const int nnode = (nall_nodes > 0) ? nall_nodes : nloc; const bool extended = nnode != nloc; const bool use_with_comm = has_comm_artifact_ && comm_nlist; + if (long_range_ && (extended || comm_nlist != nullptr)) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: a model with a long-range module evaluates a " + "complete non-periodic system in a single domain (nall_nodes == nloc " + "and no comm_nlist)."); + } if (lower_input_is_graph_ && extended && has_message_passing_) { throw deepmd::deepmd_exception( "compute_edges_gpu: message-passing graph-input inference requires an " @@ -2455,7 +2507,7 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor)); + graph_aparam, charge_spin_tensor, coord_t.reshape({nloc, 3}))); ae = out["atom_energy"].reshape({nnode}).slice(0, 0, nloc).contiguous(); force_t = out["force"].reshape({nnode, 3}).contiguous(); av = out["atom_virial"].reshape({nnode, 9}).contiguous(); @@ -2544,6 +2596,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2566,6 +2619,12 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "invalid compact canonical graph dimensions."); } + if (long_range_ && (d_coord == nullptr || nall_nodes != nloc)) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module reads " + "the node positions of a complete non-periodic system in a single " + "domain (d_coord given and nall_nodes == nloc)."); + } translate_error([&] { const torch::Device device(torch::kCUDA, gpu_id); @@ -2592,12 +2651,18 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( const_cast(d_source_order), {edge_storage}, opt_u32); auto n_node = torch::full({1}, nall_nodes, opt_i64); auto n_local = torch::full({1}, nloc, opt_i64); + const at::Tensor coord = + d_coord != nullptr + ? torch::from_blob( + const_cast(d_coord), {nall_nodes, 3}, + torch::TensorOptions().dtype(torch::kFloat64).device(device)) + : at::Tensor(); std::map output; - extract_outputs(output, - run_model_canonical_graph(atype, n_node, n_local, source, - edge_vec, destination_row_ptr, - source_row_ptr, source_order)); + extract_outputs( + output, run_model_canonical_graph(atype, n_node, n_local, source, + edge_vec, destination_row_ptr, + source_row_ptr, source_order, coord)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2679,6 +2744,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2688,9 +2754,9 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { compute_canonical_graph_gpu_impl( - d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec, - d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, - edge_storage); + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, + nall_nodes, edge_storage); } bool DeepPotPTExpt::uses_fp32_edge_vectors() const { @@ -2707,4 +2773,6 @@ bool DeepPotPTExpt::uses_canonical_graph_inference() const { return lower_input_is_canonical_; } +bool DeepPotPTExpt::has_long_range() const { return long_range_; } + #endif diff --git a/source/api_cc/src/DeepPotPTExptPlugin.cc b/source/api_cc/src/DeepPotPTExptPlugin.cc index 3499c62dcb..6843dc86fc 100644 --- a/source/api_cc/src/DeepPotPTExptPlugin.cc +++ b/source/api_cc/src/DeepPotPTExptPlugin.cc @@ -57,7 +57,7 @@ bool native_spin_backend_serves(const char* model) { } // namespace #endif -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -74,7 +74,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, #endif } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotPTPlugin.cc b/source/api_cc/src/DeepPotPTPlugin.cc index ae584b12b6..80050bb2d2 100644 --- a/source/api_cc/src/DeepPotPTPlugin.cc +++ b/source/api_cc/src/DeepPotPTPlugin.cc @@ -6,7 +6,7 @@ #include "DeepSpinPT.h" #include "DeepTensorPT.h" -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -15,7 +15,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotTFPlugin.cc b/source/api_cc/src/DeepPotTFPlugin.cc index 5e20b6fcf7..16c3a1599c 100644 --- a/source/api_cc/src/DeepPotTFPlugin.cc +++ b/source/api_cc/src/DeepPotTFPlugin.cc @@ -46,7 +46,7 @@ void convert_pbtxt_to_pb_impl(const char* pbtxt, const char* pb) { } // namespace -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -55,7 +55,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h new file mode 100644 index 0000000000..cbb37f49f9 --- /dev/null +++ b/source/api_cc/src/longRangePTExpt.h @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include + +#include + +#include "commonPT.h" + +namespace deepmd { +namespace ptexpt { + +/** + * @brief Build the all-pairs topology a long-range module of the graph lower + * reads. + * + * Inference-path twin of Python ``build_pair_graph`` in + * ``deepmd/dpmodel/utils/neighbor_graph/pair_graph.py`` for a single frame: + * every unordered pair of nodes once, ``[src, dst]`` with ``src < dst``, its + * vector ``r_src - r_dst``, and two masked guard pairs that keep the pair axis + * at the lower bound of the exported lower. The model-level pair exclusion is + * applied to the mask, as it is to the neighbor graph. + * + * @param coord Node positions, shape ``(N, 3)``, float64, on the model device. + * @param atype Node types, shape ``(N,)``, int64. + * @param type_mask_table Device keep table of the pair exclusion, or an + * undefined tensor when the model excludes no pair. + * @param ntypes Number of atom types. + * @param vec_dtype Precision of the pair vectors, that of the edge vectors of + * the artifact. + * @return ``{pair_index (2, P + 2) int64, pair_vec (P + 2, 3), pair_mask + * (P + 2,) bool}``. + */ +inline std::vector buildAllPairsTopology( + const torch::Tensor& coord, + const torch::Tensor& atype, + const torch::Tensor& type_mask_table, + const int ntypes, + const torch::ScalarType vec_dtype) { + const std::int64_t node_count = coord.size(0); + const auto index_options = + torch::TensorOptions().dtype(torch::kInt64).device(coord.device()); + const auto mask_options = + torch::TensorOptions().dtype(torch::kBool).device(coord.device()); + const torch::Tensor index = + torch::triu_indices(node_count, node_count, 1, index_options); + const torch::Tensor vec = (coord.index_select(0, index.select(0, 0)) - + coord.index_select(0, index.select(0, 1))) + .to(vec_dtype); + const torch::Tensor mask = + applyPairExclusion(index, torch::ones({index.size(1)}, mask_options), + atype, type_mask_table, ntypes); + return {torch::cat({index, torch::zeros({2, 2}, index_options)}, 1), + torch::cat({vec, torch::zeros({2, 3}, vec.options())}, 0), + torch::cat({mask, torch::zeros({2}, mask_options)})}; +} + +} // namespace ptexpt +} // namespace deepmd diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc new file mode 100644 index 0000000000..fa4df07ad4 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test C++ inference for pt_expt (.pt2) DPA4C models with a latent-charge +// long-range module, which sums over every pair of atoms of a complete +// non-periodic system. Models generated by +// source/tests/infer/gen_dpa4c_long_range.py: an uncompressed archive on the +// graph lower, which receives the pair list of the atoms, and a compressed +// archive on the compact canonical lower, which reads their positions. +#include + +#include +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4c_long_range.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4c_long_range.expected"; +constexpr const char* kCompressedModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_compressed.pt2"; +constexpr const char* kCompressedRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_compressed.expected"; + +// The full neighbor list of a non-periodic system within ``rc``: every atom +// is local, so the list holds no ghost. +template +std::vector> nopbc_nlist(const std::vector& coord, + const double rc) { + const int natoms = static_cast(coord.size() / 3); + std::vector> nlist(natoms); + for (int ii = 0; ii < natoms; ++ii) { + for (int jj = 0; jj < natoms; ++jj) { + double r2 = 0.; + for (int dd = 0; dd < 3; ++dd) { + const double diff = coord[ii * 3 + dd] - coord[jj * 3 + dd]; + r2 += diff * diff; + } + if (jj != ii && r2 < rc * rc) { + nlist[ii].push_back(jj); + } + } + } + return nlist; +} +} // namespace + +// One reference case of an archive: per-atom energies, forces and virials, +// with the totals the interface returns. +template +struct LongRangeRef { + std::vector expected_f; + double expected_tot_e = 0.; + std::vector expected_tot_v = std::vector(9, 0.); + + void load(const deepmd_test::ExpectedRef& ref, const char* section) { + const std::vector atom_e = + ref.template get(section, "expected_e"); + const std::vector atom_v = + ref.template get(section, "expected_v"); + expected_f = ref.template get(section, "expected_f"); + for (std::size_t ii = 0; ii < atom_e.size(); ++ii) { + expected_tot_e += atom_e[ii]; + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += atom_v[ii * 9 + dd]; + } + } + } + + void check(double ener, + const std::vector& force, + const std::vector& virial, + double epsilon) const { + ASSERT_EQ(force.size(), expected_f.size()); + ASSERT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), epsilon); + for (std::size_t ii = 0; ii < force.size(); ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), epsilon); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), epsilon); + } + } +}; + +template +class TestInferDeepPotDpa4cLongRangePtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector charge_spin_explicit = {1.0, 2.0}; + LongRangeRef ref_default; + LongRangeRef ref_explicit; + + static deepmd::DeepPot dp; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp.init(kModelPath); +#endif + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + ref_default.load(ref, "default"); + ref_explicit.load(ref, "explicit"); + } + + static void TearDownTestSuite() { dp = deepmd::DeepPot(); } + + // Computes through the LAMMPS-style entry point, with the full neighbor + // list of the system and no ghost atom. + void compute_lmp(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& charge_spin) { + std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff()); + const int nloc = static_cast(atype.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, std::vector(), 0, + inlist, 0, {}, {}, charge_spin); + } +}; + +template +deepmd::DeepPot TestInferDeepPotDpa4cLongRangePtExpt::dp; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, has_long_range) { + EXPECT_TRUE(this->dp.has_long_range()); + EXPECT_EQ(this->dp.dim_chg_spin(), 2); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_default) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector()); + this->ref_default.check(ener, force, virial, epsilon); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_explicit) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector(), {}, {}, + this->charge_spin_explicit); + this->ref_explicit.check(ener, force, virial, epsilon); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, lmp_nlist_explicit) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->compute_lmp(ener, force, virial, this->charge_spin_explicit); + this->ref_explicit.check(ener, force, virial, epsilon); +} + +// A cell would leave the pair sum incomplete: its periodic images are no +// part of it. +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, periodic_box_is_rejected) { + using VALUETYPE = TypeParam; + const std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + double ener; + std::vector force, virial; + EXPECT_THROW( + this->dp.compute(ener, force, virial, this->coord, this->atype, box), + deepmd::deepmd_exception); +} + +// The compressed archive folds the charge state into constants, which +// set_charge_spin rewrites: the tables of the descriptor and the total charge +// of the constraint together. +template +class TestInferDeepPotDpa4cLongRangeCompressedPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + LongRangeRef ref_default; + LongRangeRef ref_explicit; + deepmd::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kCompressedModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kCompressedModelPath + << " was not generated (gen_dpa4c_long_range.py writes it " + "on a CUDA device only)."; + } + dp.init(kCompressedModelPath); + deepmd_test::ExpectedRef ref; + ref.load(kCompressedRefPath); + ref_default.load(ref, "default"); + ref_explicit.load(ref, "explicit"); + } + + void compute_lmp(double& ener, + std::vector& force, + std::vector& virial) { + std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff()); + const int nloc = static_cast(atype.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, std::vector(), 0, + inlist, 0); + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, + set_charge_spin_follows_the_state) { + using VALUETYPE = TypeParam; + constexpr double epsilon = 1e-4; + EXPECT_TRUE(this->dp.has_long_range()); + double ener; + std::vector force, virial; + this->compute_lmp(ener, force, virial); + this->ref_default.check(ener, force, virial, epsilon); + this->dp.set_charge_spin({1.0, 2.0}); + this->compute_lmp(ener, force, virial); + this->ref_explicit.check(ener, force, virial, epsilon); + this->dp.set_charge_spin({0.0, 1.0}); + this->compute_lmp(ener, force, virial); + this->ref_default.check(ener, force, virial, epsilon); +} diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 17b4f2f593..4a154bc28a 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -145,6 +145,12 @@ else: env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_spin.py & PID15=$! fi + # DPA4C with a latent-charge long-range module: the uncompressed archive + # takes the all-pairs topology of the graph lower; for the CUDA variant the + # script also writes the compressed archive, whose compact canonical lower + # reads the node positions. + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_long_range.py & + PID16=$! wait $PID11 wait $PID12 wait $PID13 @@ -152,6 +158,7 @@ else: if [ "${DP_VARIANT}" = "cuda" ]; then wait $PID15 fi + wait $PID16 fi if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index a03482a90a..090d8dc597 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -966,6 +966,19 @@ void PairDeepMD::settings(int narg, char** arg) { all_force.resize(numb_models); } +void PairDeepMD::init_style() { + PairDeepBaseModel::init_style(); + // A long-range module sums over every pair of atoms of the system, which a + // periodic image or another domain would leave incomplete. + if (deep_pot.has_long_range() && (domain->xperiodic || domain->yperiodic || + domain->zperiodic || comm->nprocs > 1)) { + error->all(FLERR, + "pair style deepmd: a model with a long-range module evaluates " + "a complete non-periodic system; use a non-periodic boundary " + "(f, s or m) in every dimension on a single MPI rank."); + } +} + /* ---------------------------------------------------------------------- set coeffs for one or more type pairs ------------------------------------------------------------------------- */ diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index c7c40d9b18..ac8881b553 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -46,6 +46,7 @@ class PairDeepMD : public PairDeepBaseModel { ~PairDeepMD() override; void settings(int, char**) override; void coeff(int, char**) override; + void init_style() override; void compute(int, int) override; int pack_reverse_comm(int, int, double*) override; void unpack_reverse_comm(int, int*, double*) override; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index f54873dbe1..db412a9ae1 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -34,6 +34,7 @@ PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) edge_capacity(0), edge_vec_fp32(false), canonical_graph(false), + long_range(false), device_path_ok(false), reverse_virial(false), reverse_used_host(false) { @@ -212,6 +213,9 @@ void PairDeepMDKokkos::init_style() { request->enable_full(); edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors(); canonical_graph = deep_pot.uses_canonical_graph_inference(); + // PairDeepMD::init_style() has confined such a model to a complete + // non-periodic system on one rank. + long_range = deep_pot.has_long_range(); // Force exchange transfers three values per atom. Centroid per-atom virial // uses nine values even though the Kokkos full-list request runs newton off, // so comm_reverse_off reserves the classic host buffer for the wider mode. @@ -408,6 +412,35 @@ int PairDeepMDKokkos::build_edges_device() { return nedge; } +template +const double* PairDeepMDKokkos::model_node_positions() { + atomKK->sync(execution_space, X_MASK); + auto x = atomKK->k_x.template view(); + // Without virtual atoms the model nodes are the leading local atoms, already + // in the length unit of the model. + if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0) { + return x.data(); + } + const int nnode_model = compact_graph.nnode_model; + if ((int)d_node_position.extent(0) < 3 * nnode_model) { + d_node_position = Kokkos::View( + "deepmd/kk:node_position", 3 * nnode_model); + } + auto node_position = d_node_position; + auto model2loc = compact_graph.d_model2loc; + const double inv_dist = 1.0 / dist_unit_cvt_factor; + Kokkos::parallel_for( + "deepmd/kk:node_position", + Kokkos::RangePolicy(0, nnode_model), + KOKKOS_LAMBDA(const int m) { + const int i = model2loc(m); + node_position(3 * m + 0) = x(i, 0) * inv_dist; + node_position(3 * m + 1) = x(i, 1) * inv_dist; + node_position(3 * m + 2) = x(i, 2) * inv_dist; + }); + return node_position.data(); +} + template void PairDeepMDKokkos::compute(int eflag, int vflag) { if (!device_path_ok) { @@ -528,18 +561,21 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // the two runtimes are bracketed by explicit synchronization: fence the // Kokkos work before the model reads the edges, and synchronize the device // after so the scatter sees the finished model outputs. - Kokkos::fence(); // Coordinates are model-node order: the compacted buffer when virtual atoms // are present, else the local coordinates directly (the graph lower ignores - // them; edge-input models consume them). + // them; edge-input models consume them). A long-range module reads the + // positions of its nodes in the length unit of the model. const double* coord_ptr = - compact_graph.has_null_types ? d_coord_model.data() : x.data(); + long_range + ? model_node_positions() + : (compact_graph.has_null_types ? d_coord_model.data() : x.data()); + Kokkos::fence(); try { if (canonical_graph) { deep_pot.compute_canonical_graph_gpu( d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), - compact_graph.d_model_type.data(), compact_graph.d_source.data(), - compact_graph.d_edge_vec.data(), + compact_graph.d_model_type.data(), long_range ? coord_ptr : nullptr, + compact_graph.d_source.data(), compact_graph.d_edge_vec.data(), compact_graph.d_destination_row_ptr.data(), compact_graph.d_source_row_ptr.data(), compact_graph.d_source_order.data(), nloc_m, nnode_m, diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index bd3689aedf..65a1ce315f 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -82,6 +82,10 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); + // Positions of the model nodes in the length unit of the model, which a + // long-range module reads; public for the same reason. + const double* model_node_positions(); + protected: // Model node set and, for a compact canonical artifact, the graph itself. CompactCanonicalGraphKokkos compact_graph; @@ -90,6 +94,8 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { bool multi_rank; // domain-decomposed run -> extended (local+ghost) node set Kokkos::View d_coord_model; // (3 * nnode_model), NULL case + Kokkos::View + d_node_position; // (3 * nnode_model), long-range model positions // Compact edge schema: edge_index is [2 * nedge] (src rows then dst rows), // edge_vec is [3 * nedge]; offsets is the per-atom exclusive edge prefix. @@ -121,6 +127,7 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { int edge_capacity; // allocated edges in d_edge_index / d_edge_vec bool edge_vec_fp32; // model graph ABI consumes edge vectors in fp32 bool canonical_graph; // compact source-only graph artifact + bool long_range; // model reads the positions of every node bool device_path_ok; // resolved once in init_style bool reverse_virial; // reverse communication operates on centroid virial bool reverse_used_host; // force reverse communication selected host staging diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py new file mode 100644 index 0000000000..71fb10a7be --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test LAMMPS with a .pt2 DPA4C model that carries a long-range module. + +The latent-charge module sums over every pair of atoms of a complete +non-periodic system, so ``pair_style deepmd`` runs it with non-periodic +boundaries on a single rank and rejects a periodic one. The model and the +reference values, without a cell, at the default charge state and at +``charge_spin 1.0 2.0``, are generated by +source/tests/infer/gen_dpa4c_long_range.py. +""" + +import os +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps import ( + PyLammps, +) +from lammps_test_utils import ( + make_atomic_lammps, +) +from write_lmp_data import ( + write_lmp_data, +) + +pt2_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4c_long_range.pt2" +) +ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4c_long_range.expected" +) +data_file = Path(__file__).parent / "data_dpa4c_long_range_pt2.lmp" + +# Guarded with try/except because the generator only runs when PyTorch is +# built; matrices that disable PyTorch still load this file at collection time. +try: + _ref = read_expected_ref(ref_file) + expected_e_default = float(np.sum(_ref["default"]["expected_e"])) + expected_f_default = _ref["default"]["expected_f"].reshape(6, 3) + expected_e_explicit = float(np.sum(_ref["explicit"]["expected_e"])) + expected_f_explicit = _ref["explicit"]["expected_f"].reshape(6, 3) +except FileNotFoundError: + expected_e_default = expected_f_default = None + expected_e_explicit = expected_f_explicit = None + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2] +type_OH = np.array([1, 2, 2, 1, 2, 2]) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("Skip test because PyTorch support is not enabled.") + write_lmp_data(box, coord, type_OH, data_file) + + +def teardown_module() -> None: + if data_file.exists(): + os.remove(data_file) + + +def _lammps(boundary: str) -> PyLammps: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + return make_atomic_lammps(data_file, boundary=boundary, atom_map="yes") + + +@pytest.fixture +def lammps(): + lmp = _lammps("f f f") + yield lmp + lmp.close() + + +@pytest.fixture +def lammps_periodic(): + lmp = _lammps("p p p") + yield lmp + lmp.close() + + +def test_pair_deepmd(lammps) -> None: + lammps.pair_style(f"deepmd {pt2_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_default) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_default[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +def test_pair_deepmd_charge_spin(lammps) -> None: + lammps.pair_style(f"deepmd {pt2_file.resolve()} charge_spin 1.0 2.0") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_explicit) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_explicit[lammps.atoms[ii].id - 1] + ) + + +def test_periodic_boundary_is_rejected(lammps_periodic) -> None: + lammps_periodic.pair_style(f"deepmd {pt2_file.resolve()}") + lammps_periodic.pair_coeff("* *") + with pytest.raises(Exception, match="long-range module"): + lammps_periodic.run(0) diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py new file mode 100644 index 0000000000..554704c7ae --- /dev/null +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the DPA4C long-range .pt2 test models and reference values. + +A latent-charge long-range module sums over every pair of atoms of a complete +non-periodic system, so the C++ interface hands the model what the pair sum +needs beyond the neighbor list: the graph lower of an uncompressed model +receives the pair list the interface builds from the positions of the atoms, +and the compact canonical lower of a compressed model reads the positions +themselves. One archive covers each lower: + +- ``deeppot_dpa4c_long_range.pt2``: uncompressed, fp64, graph lower, which + compiles on the CPU; +- ``deeppot_dpa4c_long_range_compressed.pt2``: compressed, fp32, compact + canonical lower, whose operators are CUDA only; it is generated for the + CUDA variant of the C++ suite (``DP_VARIANT=cuda``). + +Both models condition the descriptor on the charge state and constrain the +latent charges to the total charge that state names, so a state other than +the default changes the long-range energy as well as the descriptor. The +references are those of the Python interface on each archive, without a cell, +at the default state and at ``[1, 2]``; each ``.expected`` sidecar holds the +sections ``default`` and ``explicit``. The six atoms form two groups about +9 Å apart, beyond the cutoff of the descriptor, which only the long-range +module couples. +""" + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +# The gen scripts run standalone, outside pytest's package machinery, so the +# relative imports the test suite uses do not apply here. +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +LONG_RANGE = { + "type": "latent_charge", + "channels": 1, + "charge_net": {"neuron": [8, 8], "seed": 9}, + "kernel": {"type": "les", "alpha": 0.8}, + "charge_constraint": {"source": "charge_spin"}, +} + +# The uncompressed model is fp64, so that its references hold to the +# precision of the double C++ interface. +CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 7, + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + }, + "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 7}, + "long_range": LONG_RANGE, +} + +# Compression takes an fp32 descriptor and, for the compact canonical lower, +# a fitting that the fused fitting operator evaluates. +COMPRESSED_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 8, + "precision": "float32", + "seed": 7, + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + }, + "fitting_net": { + "neuron": [32, 32], + "activation_function": "silu", + "resnet_dt": False, + "precision": "float32", + "seed": 7, + }, + "long_range": LONG_RANGE, +} + +# Jitters the zero-initialized arrays, among them the charge-state head of the +# descriptor, so that the charge state reaches the descriptor as well as the +# constraint. +_JITTER_SEED = 20261018 + +_ATYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) +_COORDS = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ], + dtype=np.float64, +).reshape(1, 6, 3) +_OTHER_CHG_SPIN = np.array([[1.0, 2.0]], dtype=np.float64) + + +def _model_dict(config: dict) -> dict: + """Build the dpmodel of a configuration and jitter it in place of its zeros.""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(config)) + return jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED)) + + +def _check_metadata(pt2_path: str, lower: str) -> None: + """Assert the archive declares its lower and the long-range input it needs.""" + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + assert md["lower_input_kind"] == lower, md["lower_input_kind"] + if lower == "graph": + assert md["topologies"] == ["all_pairs"] + assert md["dim_chg_spin"] == 2 + else: + # The compressed descriptor folds the charge state; the archive + # carries the rebuild of every constant that holds it. + assert md["dim_chg_spin"] == 0 + assert len(md["charge_state_constants"]) == 5 + + +def _write_ref(pt2_path: str, ref_path: str) -> None: + """Evaluate the reference cases on the archive and write the C++ sidecar.""" + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + sections = {} + energies = {} + for name, charge_spin in (("default", None), ("explicit", _OTHER_CHG_SPIN)): + condition = {} if charge_spin is None else {"charge_spin": charge_spin} + e, f, v, ae, av = dp.eval(_COORDS, None, _ATYPES, atomic=True, **condition) + assert np.all(np.isfinite(e)) and np.all(np.isfinite(f)) + energies[name] = float(e[0, 0]) + print(f"// {name} total energy: {e[0, 0]:.18e}") # noqa: T201 + sections[name] = { + "expected_e": ae[0, :, 0], + "expected_f": f[0], + "expected_v": av[0], + } + # The C++ cases compare the two states, which needs a real difference. + assert abs(energies["explicit"] - energies["default"]) > 1e-3, ( + "the two charge states must differ, or a call that ignored the " + "requested state would still match" + ) + write_expected_ref( + ref_path, + sections=sections, + source_script="source/tests/infer/gen_dpa4c_long_range.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + + +def _export_compressed(pt2_path: str) -> None: + """Compress the fp32 model and export it for the CUDA device.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + model = BaseModel.deserialize(_model_dict(COMPRESSED_CONFIG)) + model.atomic_model.enable_compression(0.5) + data = { + "model": model.serialize(), + "model_def_script": COMPRESSED_CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {pt2_path} ...") # noqa: T201 + deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="auto") + + +def main() -> None: + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + pt2_path = os.path.join(base_dir, "deeppot_dpa4c_long_range.pt2") + data = { + "model": _model_dict(CONFIG), + "model_def_script": CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {pt2_path} (lower_kind='graph') ...") # noqa: T201 + deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="graph") + _check_metadata(pt2_path, "graph") + _write_ref(pt2_path, os.path.join(base_dir, "deeppot_dpa4c_long_range.expected")) + + if os.environ.get("DP_VARIANT") == "cuda": + compressed_path = os.path.join( + base_dir, "deeppot_dpa4c_long_range_compressed.pt2" + ) + _export_compressed(compressed_path) + _check_metadata(compressed_path, "dpa4c_canonical") + _write_ref( + compressed_path, + os.path.join(base_dir, "deeppot_dpa4c_long_range_compressed.expected"), + ) + + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() From 16f40a2b383830453339383c8faf34ab2ac056b3 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 08:59:56 +0800 Subject: [PATCH 17/30] fix(dpmodel): default the SOG amplitudes to a discretized 1/r The default amplitude of the sum-of-Gaussians kernel was the constant 4 pi ln b for every Gaussian. In the real-space kernel sum_k a_k exp(-r^2 / 2 sigma_k^2) that constant gives a kernel of logarithmic shape, 74 per Angstrom at contact and about 500 times 1/r at 10 Angstrom. The default amplitudes are now a_k = sqrt(2 / pi) ln b / sigma_k, the trapezoidal discretization in ln s of 1/r = sqrt(2 / pi) int_0^inf exp(-r^2 / 2 s^2) s^-2 ds at the default widths. With the default widths the kernel is 0.46 per Angstrom at contact, and its force stays within 0.4% of the Coulomb force from 6.5 to 150 Angstrom. It differs from 1/r there by a constant, which exerts no force. --- deepmd/dpmodel/long_range/kernel.py | 11 ++++++++--- deepmd/utils/argcheck.py | 7 +++++-- doc/model/long-range.md | 7 ++++++- source/tests/common/dpmodel/test_long_range.py | 13 +++++++++++-- 4 files changed, 30 insertions(+), 8 deletions(-) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py index e6227aa629..90a75cbee5 100644 --- a/deepmd/dpmodel/long_range/kernel.py +++ b/deepmd/dpmodel/long_range/kernel.py @@ -197,8 +197,13 @@ class SOGKernel(BaseInteractionKernel): amplitude The amplitudes :math:`a_k`, in Å⁻¹, of either sign: a scalar shared by every Gaussian or one value per Gaussian. Required with ``width``; - with the default widths it defaults to :math:`4\pi \ln b` for every - Gaussian. + with the default widths it defaults to + :math:`a_k = \sqrt{2/\pi} \ln b / \sigma_k`, the trapezoidal + discretization in :math:`\ln s` of + :math:`1/r = \sqrt{2/\pi} \int_0^\infty e^{-r^2/2s^2} s^{-2} \, ds`. + The default kernel then follows :math:`1/r`, less a constant that + exerts no force, from a few :math:`\sigma` to a fraction of + :math:`\sigma b^{M-1}`, and stays finite at contact. width The widths :math:`\sigma_k > 0`, in Å. Defaults to the geometric sequence :math:`\sigma_k = \sigma b^k`, :math:`k = 0, \dots, M - 1`. @@ -241,7 +246,7 @@ def __init__( width_ratio, np.arange(n_gaussian, dtype=np.float64) ) if amplitude is None: - amplitude = 4.0 * np.pi * np.log(width_ratio) + amplitude = np.sqrt(2.0 / np.pi) * np.log(width_ratio) / width elif amplitude is None: raise ValueError("`amplitude` is required with explicit `width`") width = np.asarray(width, dtype=np.float64).reshape(-1) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 88d249c221..a99c8d60b4 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3574,8 +3574,11 @@ def interaction_kernel_sog() -> list[Argument]: doc_amplitude = ( "The amplitudes a_k of the Gaussians, in Å^-1, of either sign: one value " "shared by every Gaussian or one value per Gaussian. Required with " - "`width`; with the default widths it defaults to 4 pi ln(`width_ratio`) " - "for every Gaussian." + "`width`; with the default widths it defaults to " + "sqrt(2 / pi) ln(`width_ratio`) / sigma_k, the weights of a discretized " + "1 / r, so that the default kernel follows 1 / r, less a constant that " + "exerts no force, from a few `base_width` to a fraction of the widest " + "Gaussian." ) doc_width = ( "The widths sigma_k of the Gaussians, in Å, each positive. Defaults to " diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 67d31e079e..815607f6a8 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -89,7 +89,12 @@ training. {ref}`width ` or default to a geometric sequence, and the amplitudes by {ref}`amplitude `, - which explicit widths require. + which explicit widths require. The default amplitudes, + $a_k = \sqrt{2/\pi}\,\ln b/\sigma_k$ for the ratio $b$ of consecutive + widths, discretize $1/r$: with the default widths the kernel stays finite, + $0.46$ Å⁻¹, at contact, and its force stays within 0.4% of that of $1/r$ + from 6.5 Å to 150 Å. It differs from $1/r$ there by a constant, which + exerts no force. Several charge channels, set by {ref}`channels `, let diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 9a7ee1b770..108fd7082c 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -94,7 +94,7 @@ def test_les_kernel_closed_form() -> None: def test_sog_kernel_closed_form_and_defaults() -> None: - """SOG is sum_k a_k exp(-r^2 / 2 sigma_k^2); the defaults follow the geometric widths.""" + """SOG is sum_k a_k exp(-r^2 / 2 sigma_k^2); the defaults discretize 1/r.""" kernel = SOGKernel(amplitude=[0.7, -0.2], width=[0.8, 1.7]) r = np.array([0.0, 1.1, 3.0]) expected = 0.7 * np.exp(-(r**2) / (2 * 0.8**2)) - 0.2 * np.exp( @@ -103,7 +103,16 @@ def test_sog_kernel_closed_form_and_defaults() -> None: np.testing.assert_allclose(kernel(r**2), expected, rtol=1e-12) default = SOGKernel(n_gaussian=3, base_width=2.0, width_ratio=1.5) np.testing.assert_allclose(np.exp(default.log_width), [2.0, 3.0, 4.5]) - np.testing.assert_allclose(default.amplitude, 4 * np.pi * np.log(1.5)) + np.testing.assert_allclose( + default.amplitude, np.sqrt(2 / np.pi) * np.log(1.5) / np.array([2.0, 3.0, 4.5]) + ) + # The default kernel exerts the Coulomb force, r^2 K'(r) = -1, between a + # few times its first width and a fraction of its widest one. + default = SOGKernel() + r = np.geomspace(6.5, 150.0, 200) + step = 1e-4 * r + slope = (default((r + step) ** 2) - default((r - step) ** 2)) / (2 * step) + np.testing.assert_allclose(r**2 * slope, -1.0, rtol=0, atol=5e-3) @pytest.mark.parametrize( From 91e6b1576bc13e58f33c46244d34567b7b149c91 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 11:16:18 +0800 Subject: [PATCH 18/30] feat(lmp): evaluate compressed long-range models on several MPI ranks The latent-charge direct sum of a compressed DPA4C model now spans the domains of a decomposed system. Every rank keeps the atoms it owns: the operator gathers the positions and charges of the owned atoms of all ranks over MPI, in rank order, and evaluates its own atoms against all of them, so that the ranks divide the quadratic work. Every rank shifts the constrained channel by the same sum of the gathered charges, the mean potential of the constraint adds the sums of the ranks in rank order, the forces stay with their owners, and the per-frame virial is the share of the rank. Periodic boundaries stay unsupported. - op: deepmd::latent_charge_direct_sum takes a communicator tensor, the bytes of an MPI_Comm handle as for border_op, or zero for a system held whole; the exchanges live in long_range/domain_exchange.cc and stage through pinned host memory unless MPI reads device memory. - pt_expt: the compact canonical lower of a long-range model takes the communicator as its tenth input. - api_cc, api_c: compute_canonical_graph_gpu and DP_DeepPotComputeCanonicalGraphGPU2 take the communicator (world). A rank without nodes runs the lower on one phantom halo node, so that it joins the exchanges. DP_DeepPotModelDeviHasLongRange and DeepPotModelDevi::has_long_range report a long-range module in any model of an ensemble. - lmp: pair_style deepmd and deepmd/kk accept several ranks for such a model and reject periodic boundaries for any model of a model deviation; deepmd/kk passes the communicator and evaluates the model on every rank at every step. Several ranks require the compact canonical lower, which a compressed DPA4C model frozen for a CUDA device uses when the direct-sum operator serves its long-range module; any other long-range model, like an uncompressed one, is rejected on several ranks. DP_DeepPotComputeCanonicalGraphGPU2 (API version 31) has not been released, so its argument list changes in place; a binary built against the previous commit must be rebuilt together with the library. --- deepmd/pt_expt/infer/deep_eval.py | 4 +- .../kernels/dpa4c/canonical_long_range.py | 8 +- deepmd/pt_expt/kernels/long_range.py | 22 + deepmd/pt_expt/long_range/latent_charge.py | 9 +- deepmd/pt_expt/model/ener_model.py | 33 +- deepmd/pt_expt/utils/serialization.py | 31 +- doc/model/long-range.md | 14 +- source/api_c/include/c_api.h | 23 +- source/api_c/include/deepmd.hpp | 19 +- source/api_c/src/c_api.cc | 14 +- .../tests/test_deeppot_long_range_hpp.cc | 17 + source/api_cc/include/DeepPot.h | 33 +- source/api_cc/include/DeepPotPTExpt.h | 20 +- source/api_cc/src/DeepPot.cc | 9 +- source/api_cc/src/DeepPotPTExpt.cc | 119 +++- source/lmp/pair_deepmd.cpp | 11 +- source/lmp/pair_deepmd_kokkos.cpp | 9 +- .../lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py | 9 +- .../tests/test_lammps_dpa4c_long_range_pt2.py | 180 ++++++- source/op/pt/CMakeLists.txt | 1 + source/op/pt/long_range/direct_sum.cu | 510 ++++++++++++------ source/op/pt/long_range/domain_exchange.cc | 169 ++++++ source/op/pt/long_range/domain_exchange.h | 45 ++ source/op/pt/long_range/ops.cc | 11 +- .../kernels/test_long_range_direct_sum.py | 2 + 25 files changed, 1056 insertions(+), 266 deletions(-) create mode 100644 source/op/pt/long_range/domain_exchange.cc create mode 100644 source/op/pt/long_range/domain_exchange.h diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 2f9eef288c..abee40c07c 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -2518,12 +2518,14 @@ def _eval_model_graph( ) if self.metadata.get("topologies"): # The compact lower of a model with a long-range module reads - # the positions of complete non-periodic frames. + # the positions of complete non-periodic frames, and the zero + # communicator of a frame held whole. check_complete_frames(box_input) model_inputs += ( torch.as_tensor( coord_input.reshape(-1, 3), dtype=torch.float64, device=DEVICE ), + torch.zeros(1, dtype=torch.int64, device="cpu"), ) else: fparam_t, aparam_t = self._prepare_optional_lower_inputs( diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py index 2f1fb6c6b0..f78b5c2db1 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -25,6 +25,7 @@ def dpa4c_canonical_long_range_energy_force( graph: CanonicalGraph, atype: torch.Tensor, coord: torch.Tensor, + communicator: torch.Tensor | None, ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, @@ -54,6 +55,11 @@ def dpa4c_canonical_long_range_energy_force( coord Node positions with shape ``(N, 3)`` in Å; every frame is a complete non-periodic system. + communicator + The bytes of the MPI_Comm handle of a frame decomposed into domains, + one per rank, as an int64 CPU tensor of one element, or zero or + ``None`` for a single domain. Across domains the owned nodes of every + rank make up the frame. ownership Boolean mask selecting energy-contributing nodes with shape ``(N,)``. atom_bias @@ -117,7 +123,7 @@ def dpa4c_canonical_long_range_energy_force( long_range_atom_virial, long_range_virial, ) = long_range.energy_force_from_coord( - node_descriptor, atype, ownership, graph.n_node, coord + node_descriptor, atype, ownership, graph.n_node, coord, communicator ) atom_energy, descriptor_gradient = fitting_energy_and_gradient( fitting, node_descriptor, atype, ownership, atom_bias diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py index d63ebfb344..302320b6fc 100644 --- a/deepmd/pt_expt/kernels/long_range.py +++ b/deepmd/pt_expt/kernels/long_range.py @@ -13,6 +13,10 @@ charges the charge network predicts, and the force and virial at fixed charges; the caller adds the response of the charges by back-propagating the charge gradient through the charge network. + +Given the MPI communicator of a system decomposed into domains, one per rank, +the CUDA kernel evaluates the charged nodes of its rank against those of +every rank; the CPU reference evaluates a single domain. """ import math @@ -133,11 +137,17 @@ def _reference( n_node_per_frame: torch.Tensor, total_charge: torch.Tensor, coefficient: torch.Tensor, + communicator: torch.Tensor, kernel: int, channel: int, prefactor: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Evaluate the operator over the explicit pairs of every frame, in float64.""" + if int(communicator.reshape(-1)[0]) != 0: + raise NotImplementedError( + "the CPU reference of latent_charge_direct_sum evaluates a single " + "domain; a communicator of several ranks requires the CUDA kernel" + ) n = coord.shape[0] frame_count = n_node_per_frame.shape[0] # Padding nodes belong to no frame and carry no charge. @@ -191,6 +201,7 @@ def _fake( n_node_per_frame: torch.Tensor, total_charge: torch.Tensor, coefficient: torch.Tensor, + communicator: torch.Tensor, kernel: int, channel: int, prefactor: float, @@ -232,6 +243,7 @@ def latent_charge_direct_sum( kernel: Any, channel: int, prefactor: float, + communicator: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Energy, charge gradient, force and virial of the charges of complete frames. @@ -259,6 +271,13 @@ def latent_charge_direct_sum( nodes of every frame to sum to its total charge, or -1 for none. prefactor Factor of every pair term, in eV Å. + communicator + The bytes of the MPI_Comm handle of a system decomposed into domains, + one per rank, as an int64 CPU tensor of one element, or zero, the + default, for a single domain. Across domains the call takes one + frame, whose charged nodes are those of every rank; the outputs are + those of the charged nodes of this rank, the other nodes carry no + charge, and ``virial`` is the share of this rank. Returns ------- @@ -278,6 +297,8 @@ def latent_charge_direct_sum( """ ensure_registered() code, coefficient = kernel_operator_arguments(kernel) + if communicator is None: + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") return torch.ops.deepmd.latent_charge_direct_sum( coord, charge, @@ -285,6 +306,7 @@ def latent_charge_direct_sum( n_node_per_frame, total_charge, coefficient.to(coord.device), + communicator, code, int(channel), float(prefactor), diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py index badb2aa26c..ef207f44d3 100644 --- a/deepmd/pt_expt/long_range/latent_charge.py +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -74,6 +74,7 @@ def energy_force_from_coord( node_mask: torch.Tensor, n_node: torch.Tensor, coord: torch.Tensor, + communicator: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Evaluate the module for inference from the positions of the nodes. @@ -98,6 +99,11 @@ def energy_force_from_coord( contiguous spans of the node axis. coord Node positions with shape ``(N, 3)`` in Å. + communicator + The MPI communicator of a frame decomposed into domains, one per + rank, whose charged nodes are those of every rank, in the form of + :func:`~deepmd.pt_expt.kernels.long_range.latent_charge_direct_sum`; + ``None`` for a single domain. Returns ------- @@ -113,7 +119,7 @@ def energy_force_from_coord( float32, each pair term split half to each of its nodes. virial Per-frame virial at fixed charges with shape ``(nf, 9)`` in eV, - float32. + float32; across domains, the share of the nodes of this rank. """ from deepmd.pt_expt.kernels.long_range import ( latent_charge_direct_sum, @@ -137,6 +143,7 @@ def energy_force_from_coord( self.kernel, channel, COULOMB_CONSTANT, + communicator, ) ) (descriptor_gradient,) = torch.autograd.grad( diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 960c501e8d..ff9b3e77b7 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -142,6 +142,7 @@ def forward_lower_canonical_graph( do_atomic_virial: bool, spin: torch.Tensor | None = None, coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -173,6 +174,11 @@ def forward_lower_canonical_graph( Node positions with shape ``(N, 3)`` in Å, float64, for a model with a long-range module, whose frames are complete non-periodic systems; ``None`` otherwise. + communicator + For a model with a long-range module, the bytes of the MPI_Comm + handle of a frame decomposed into domains, one per rank, as an + int64 CPU tensor of one element, whose owned nodes are then those + of one domain; zero or ``None`` for a complete frame. Returns ------- @@ -237,6 +243,7 @@ def forward_lower_canonical_graph( graph, atype, coord, + communicator, output_mask, atom_bias, do_atomic_virial, @@ -303,6 +310,7 @@ def forward_lower_canonical_graph_exportable( source_row_ptr: torch.Tensor, source_order: torch.Tensor, coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -315,10 +323,11 @@ def forward_lower_canonical_graph_exportable( Compact graph node and edge tensors. destination_row_ptr, source_row_ptr, source_order Compact dual-CSR topology. - coord - Sample node positions with shape ``(N, 3)``, float64, for a model - with a long-range module, which then takes them as a ninth input; - ``None`` otherwise. + coord, communicator + For a model with a long-range module, sample node positions with + shape ``(N, 3)``, float64, and a sample communicator, an int64 CPU + tensor of one element, which the traced module then takes as its + ninth and tenth inputs; ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -327,14 +336,19 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced compact deployment module of eight inputs, or nine with - the node positions. + Traced compact deployment module of eight inputs, or ten with the + node positions and the communicator. """ model = self + if (coord is None) != (communicator is None): + raise ValueError( + "the compact lower of a model with a long-range module takes " + "the node positions and the communicator together" + ) if coord is not None: - def fn_coord( + def fn_long_range( atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, @@ -344,6 +358,7 @@ def fn_coord( source_row_ptr: torch.Tensor, source_order: torch.Tensor, coord: torch.Tensor, + communicator: torch.Tensor, ) -> dict[str, torch.Tensor]: return model.forward_lower_canonical_graph( atype, @@ -356,9 +371,10 @@ def fn_coord( source_order, do_atomic_virial=do_atomic_virial, coord=coord, + communicator=communicator, ) - return make_fx(fn_coord, **make_fx_kwargs)( + return make_fx(fn_long_range, **make_fx_kwargs)( atype, n_node, n_local, @@ -368,6 +384,7 @@ def fn_coord( source_row_ptr, source_order, coord, + communicator, ) def fn( diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index a1e94de1e1..ebba8ca4a3 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -731,15 +731,16 @@ def build_synthetic_canonical_graph_inputs( *, device: torch.device, want_spin: bool = False, - want_coord: bool = False, + want_long_range: bool = False, ) -> tuple[torch.Tensor, ...]: """Build compact canonical trace inputs for compressed CUDA descriptors. ``want_spin`` appends the per-node moment at slot 8, the last slot of the compact ABI, matching :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. - ``want_coord`` appends the float64 node positions at the same slot, which - the compact lower of a model with a long-range module reads. + ``want_long_range`` appends the inputs that the compact lower of a model + with a long-range module reads: the float64 node positions at slot 8 and + the int64 communicator, a CPU tensor of one element, at slot 9. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -798,13 +799,14 @@ def build_synthetic_canonical_graph_inputs( ) if spin is not None: return (*compact_sample, spin) - if want_coord: + if want_long_range: # Tracing reads the shape of the positions alone; distinct values keep - # the sample nodes apart. + # the sample nodes apart. The communicator of a single domain is zero. coord = 0.7 * torch.arange( 3 * atype.shape[0], dtype=torch.float64, device=device ).reshape(-1, 3) - return (*compact_sample, coord) + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + return (*compact_sample, coord, communicator) return compact_sample @@ -813,9 +815,10 @@ def _build_canonical_graph_dynamic_shapes( ) -> tuple: """Build dynamic shapes for the compact deployment ABI. - The trailing node slot, the moments of a native-spin model or the - positions of a model with a long-range module, is present only for such - models, so the sample length selects the shape tuple. + The trailing slots, the moments of a native-spin model or the positions + and the communicator of a model with a long-range module, are present + only for such models, so the sample length selects the shape tuple. The + communicator holds one handle and is static. """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, @@ -838,7 +841,8 @@ def _build_canonical_graph_dynamic_shapes( {0: node_dim + 1}, {0: edge_storage_dim}, ) - return shapes if len(sample_inputs) == len(shapes) else (*shapes, {0: node_dim}) + trailing = ({0: node_dim}, None)[: len(sample_inputs) - len(shapes)] + return (*shapes, *trailing) def count_synthetic_graph_edges( @@ -2084,14 +2088,15 @@ def _trace_and_export_impl( e_sample += 1 if canonical: # The compact lower of a model with a long-range module reads the - # node positions, from which its operators enumerate the pairs. - want_coord = bool(model.atomic_model.required_topologies()) + # node positions, from which its operators enumerate the pairs, + # and the communicator of the domains of a decomposed frame. + want_long_range = bool(model.atomic_model.required_topologies()) sample_inputs = build_synthetic_canonical_graph_inputs( model, e_sample, device=torch.device("cpu"), want_spin=is_native_spin, - want_coord=want_coord, + want_long_range=want_long_range, ) traced = model.forward_lower_canonical_graph_exportable( *sample_inputs, diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 815607f6a8..066bdef194 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -133,8 +133,13 @@ A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` and the Python inference interface; the C++ interface and LAMMPS, with `pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. Since the pair sum covers the complete system, LAMMPS runs such a model with -non-periodic boundaries (`f`, `s` or `m`) in every dimension on a single MPI -rank, and rejects any other setting. +non-periodic boundaries (`f`, `s` or `m`) in every dimension, and rejects a +periodic one. A model that the dedicated operator described below evaluates, +a compressed DPA4C model frozen for a CUDA device, also runs on several MPI +ranks: every rank gathers the positions and charges of the atoms of all ranks +over MPI and evaluates the pairs of its own atoms, so that the ranks share +the cost of the pair sum. This requires DeePMD-kit built with MPI. Any other +model, like an uncompressed one, runs on a single rank. Compressing a descriptor conditioned on `charge_spin` specializes the model to one charge state at a time, and a constraint that reads its total charge @@ -153,5 +158,6 @@ calling the model, so its memory grows with the square of the number of atoms. The following are not supported for a model with a long-range module: -multi-rank (domain-decomposed) inference, spin models, the tabulated ZBL -potential, and short-range bridging. +multi-rank (domain-decomposed) inference of a model that the dedicated +operator does not evaluate, such as an uncompressed one, spin models, the +tabulated ZBL potential, and short-range bridging. diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 44f70b9318..a535e0c516 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -484,8 +484,12 @@ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); * model device. * * A model with a long-range module reads the positions of the nodes, which - * must form a complete non-periodic system in a single domain - * (``nall_nodes == nloc``); any other model ignores them. + * must form a complete non-periodic system: held whole by one rank + * (``nall_nodes == nloc``), or decomposed over the ranks of ``world``, whose + * owned nodes ``[0, nloc)`` are the share of this rank. With such a model on + * several ranks, every rank calls this function on every evaluation, also + * when it holds no node. Any other model ignores the positions and the + * communicator. * * @param[in] dp The DP to use. * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``. @@ -494,6 +498,8 @@ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``. * @param[in] d_coord Node positions in model-node order, shape * ``(nall_nodes, 3)``, or NULL for a model without a long-range module. + * @param[in] world The MPI communicator (a host pointer to an ``MPI_Comm``) + * of a system decomposed over several ranks, or NULL on a single rank. * @param[in] d_source Source node per edge storage slot. * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``. * @param[in] d_destination_row_ptr Destination CSR offsets. @@ -503,7 +509,7 @@ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); * @param[in] nall_nodes Total local-plus-halo node count. * @param[in] edge_storage Number of edge storage slots. * @note Supersedes DP_DeepPotComputeCanonicalGraphGPU, adding the node - * positions. + * positions and the communicator. * @since API version 31 */ extern void DP_DeepPotComputeCanonicalGraphGPU2( @@ -513,6 +519,7 @@ extern void DP_DeepPotComputeCanonicalGraphGPU2( double* d_atom_virial, const int64_t* d_atype, const double* d_coord, + void* world, const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, @@ -524,7 +531,7 @@ extern void DP_DeepPotComputeCanonicalGraphGPU2( /** * @brief Query whether the model carries a long-range module, which - * evaluates complete non-periodic systems in a single domain only. + * evaluates complete non-periodic systems. * @since API version 31 */ extern bool DP_DeepPotHasLongRange(DP_DeepPot* dp); @@ -2733,6 +2740,14 @@ int DP_DeepPotModelDeviGetDimAParam(DP_DeepPotModelDevi* dp); */ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp); +/** + * @brief Query whether any model of a DP Model Deviation carries a long-range + * module, which evaluates complete non-periodic systems. + * @param[in] dp The DP Model Deviation to use. + * @since API version 31 + */ +extern bool DP_DeepPotModelDeviHasLongRange(DP_DeepPotModelDevi* dp); + /** * @brief Fix the charge/spin condition served for the rest of the run. * Applied to every model, so that the deviation is taken between models under diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 9aee9e8857..c5153087f3 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1415,6 +1415,7 @@ class DeepPot : public DeepBaseModel { double* d_atom_virial, const int64_t* d_atype, const double* d_coord, + void* world, const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, @@ -1424,15 +1425,15 @@ class DeepPot : public DeepBaseModel { const int nall_nodes, const int64_t edge_storage) { DP_DeepPotComputeCanonicalGraphGPU2( - dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, - d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, - nloc, nall_nodes, edge_storage); + dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage); DP_CHECK_OK(DP_DeepPotCheckOK, dp); } /** * @brief Query whether the model carries a long-range module, which - * evaluates complete non-periodic systems in a single domain only. + * evaluates complete non-periodic systems. */ bool has_long_range() const { const bool result = DP_DeepPotHasLongRange(dp); @@ -2642,6 +2643,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return dchgspin; } + /** + * @brief Query whether any of the models carries a long-range module, which + * evaluates complete non-periodic systems. + */ + bool has_long_range() const { + const bool result = DP_DeepPotModelDeviHasLongRange(dp); + DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); + return result; + } + /** * @brief Fix the charge/spin condition served for the rest of the run. * Applied to every model, so that the deviation is taken between models diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 0391429e61..e6e2904295 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1870,6 +1870,7 @@ void DP_DeepPotComputeCanonicalGraphGPU2(DP_DeepPot* dp, double* d_atom_virial, const int64_t* d_atype, const double* d_coord, + void* world, const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, @@ -1880,8 +1881,8 @@ void DP_DeepPotComputeCanonicalGraphGPU2(DP_DeepPot* dp, const int64_t edge_storage) { DP_REQUIRES_OK( dp, dp->dp.compute_canonical_graph_gpu( - d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, - d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage)); } @@ -2878,6 +2879,15 @@ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp) { return dp->dp.dim_chg_spin(); } +bool DP_DeepPotModelDeviHasLongRange(DP_DeepPotModelDevi* dp) { + try { + return dp->dp.has_long_range(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + void DP_DeepPotModelDeviSetChargeSpin(DP_DeepPotModelDevi* dp, const double* charge_spin, const int numb_chg_spin) { diff --git a/source/api_c/tests/test_deeppot_long_range_hpp.cc b/source/api_c/tests/test_deeppot_long_range_hpp.cc index 1f16483156..9e6804e056 100644 --- a/source/api_c/tests/test_deeppot_long_range_hpp.cc +++ b/source/api_c/tests/test_deeppot_long_range_hpp.cc @@ -4,6 +4,7 @@ #include #include +#include #include #include "../../api_cc/tests/expected_ref.h" @@ -16,6 +17,8 @@ constexpr const char* kModelPath = "../../tests/infer/deeppot_dpa4c_long_range.pt2"; constexpr const char* kRefPath = "../../tests/infer/deeppot_dpa4c_long_range.expected"; +constexpr const char* kShortRangeModelPath = + "../../tests/infer/deeppot_sea.pth"; } // namespace template @@ -49,6 +52,20 @@ TYPED_TEST(TestInferDeepPotLongRangeHPP, has_long_range) { EXPECT_TRUE(this->dp.has_long_range()); } +TEST(TestInferDeepPotModelDeviLongRangeHPP, has_long_range) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + // A long-range model anywhere in the ensemble confines all of them. + deepmd::hpp::DeepPotModelDevi mixed; + mixed.init(std::vector{kShortRangeModelPath, kModelPath}); + EXPECT_TRUE(mixed.has_long_range()); + deepmd::hpp::DeepPotModelDevi short_range; + short_range.init( + std::vector{kShortRangeModelPath, kShortRangeModelPath}); + EXPECT_FALSE(short_range.has_long_range()); +} + TYPED_TEST(TestInferDeepPotLongRangeHPP, nopbc) { using VALUETYPE = TypeParam; const double epsilon = std::is_same::value ? 1e-8 : 1e-4; diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index c39466ee52..41b03b4572 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -427,6 +427,7 @@ class DeepPotBackend : public DeepBaseModelBackend { double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -901,8 +902,9 @@ class DeepPot : public DeepBaseModel { * * Such a model sums its long-range interaction over every pair of atoms of * a complete non-periodic system, so it evaluates a system without a - * periodic cell in a single domain only, and the device entries take the - * positions of its atoms. + * periodic cell, held whole by one rank or, through the compact canonical + * lower, decomposed over the ranks of an MPI communicator; the device + * entries take the positions of its atoms. */ bool has_long_range() const; @@ -923,13 +925,20 @@ class DeepPot : public DeepBaseModel { * @brief Evaluate a compact canonical graph together with the positions of * its nodes. * - * Every pointer resides on the model device. A model with a long-range - * module reads the positions; any other model ignores them. + * Every pointer except ``world`` resides on the model device; ``world`` + * points to the caller's ``MPI_Comm`` in host memory. A model with a + * long-range module reads the positions and the communicator; any other + * model ignores them. * * @param[in] d_coord Node positions in model-node order, shape * ``(nall_nodes, 3)``, in the length unit of the model. A model with a - * long-range module requires them and a single-domain node set - * (``nall_nodes == nloc``) of a system without a periodic cell. + * long-range module requires them, for a system without a periodic cell. + * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of a + * system decomposed over several ranks, of which the owned nodes + * ``[0, nloc)`` are the share of this rank, or NULL on a single rank, + * whose node set then holds the complete system (``nall_nodes == nloc``). + * With a long-range module and several ranks, every rank calls this + * function on every evaluation, also when it holds no node. * * The other parameters are those of the overload without positions. */ @@ -938,6 +947,7 @@ class DeepPot : public DeepBaseModel { double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -999,6 +1009,17 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return numb_models > 0 ? dps[0]->dim_chg_spin() : 0; }; + /** + * @brief Whether any of the models carries a long-range module, which + * confines all of them to complete non-periodic systems. + **/ + bool has_long_range() const { + return std::any_of(dps.begin(), dps.end(), + [](const std::shared_ptr& model) { + return model->has_long_range(); + }); + }; + /** * @brief Fix the charge/spin condition served for the rest of the run. * Applied to every model, so that the deviation is taken between models diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 05cf7b2767..cfa0a551ca 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -378,6 +378,7 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -437,6 +438,7 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -478,7 +480,8 @@ class DeepPotPTExpt : public DeepPotBackend { bool graph_edge_fp32_ = false; // The model carries a long-range module, whose lower reads the all-pairs // topology of the frame after its other inputs: the pair lists for the graph - // lower, the node positions for the compact canonical lower. + // lower, the node positions and the communicator of the ranks for the + // compact canonical lower. bool long_range_ = false; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) @@ -647,6 +650,9 @@ class DeepPotPTExpt : public DeepPotBackend { * * @param[in] coord Node positions, shape ``(N, 3)`` float64, appended for a * model with a long-range module and ignored otherwise. + * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of the + * ranks a decomposed system spans, or NULL for a system that one rank holds + * whole; a model with a long-range module receives it after the positions. */ std::vector run_model_canonical_graph( const torch::Tensor& atype, @@ -657,7 +663,17 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, const torch::Tensor& source_order, - const torch::Tensor& coord); + const torch::Tensor& coord, + const void* world); + + /** + * @brief Evaluate the long-range module on a rank that holds no node. + * + * The compact lower runs on one halo node without edges, which carries no + * charge and contributes nothing, so that the rank joins the exchanges of + * the module with the other ranks of ``world``. + */ + void join_long_range_exchange(const void* world); /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 66ed2f05c6..b8491393ee 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -662,6 +662,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -675,6 +676,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( (void)d_atom_virial; (void)d_atype; (void)d_coord; + (void)world; (void)d_source; (void)d_edge_vec; (void)d_destination_row_ptr; @@ -764,8 +766,8 @@ void DeepPot::compute_canonical_graph_gpu( const std::int64_t edge_storage) { dp->compute_canonical_graph_gpu( d_atom_energy, d_force, d_atom_virial, d_atype, /*d_coord=*/nullptr, - d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, - d_source_order, nloc, nall_nodes, edge_storage); + /*world=*/nullptr, d_source, d_edge_vec, d_destination_row_ptr, + d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage); } void DeepPot::compute_canonical_graph_gpu( @@ -774,6 +776,7 @@ void DeepPot::compute_canonical_graph_gpu( double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -783,7 +786,7 @@ void DeepPot::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { dp->compute_canonical_graph_gpu( - d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage); } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 2c7d14a2a2..3ac1f112c7 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -111,6 +111,18 @@ FrameParameterLayout resolve_frame_parameter_layout(const char* name, " for frame-specific values."); } +// The communicator input of the compact lower of a long-range model: an int64 +// CPU tensor holding the bytes of the ``MPI_Comm`` that ``world`` points to, +// read as for the comm tensors of a message-passing model, or zero for a +// system that one rank holds whole. +at::Tensor communicator_tensor(const void* world) { + const auto options = torch::TensorOptions().dtype(torch::kInt64); + if (world == nullptr) { + return torch::zeros({1}, options); + } + return torch::from_blob(const_cast(world), {1}, options).clone(); +} + } // namespace void DeepPotPTExpt::translate_error(std::function f) { @@ -566,17 +578,33 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, const torch::Tensor& source_order, - const torch::Tensor& coord) { + const torch::Tensor& coord, + const void* world) { std::vector inputs = {atype, n_node, n_local, source, edge_vec, destination_row_ptr, source_row_ptr, source_order}; if (long_range_) { inputs.push_back(coord); + inputs.push_back(communicator_tensor(world)); } return loader->run(inputs); } +void DeepPotPTExpt::join_long_range_exchange(const void* world) { + const torch::Device device(torch::kCUDA, gpu_id); + const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto index_type = deepmd::canonicalGraphIndexType(); + // The compact lower stores at least two edge slots; without edges both are + // padding outside every CSR row. + run_model_canonical_graph( + torch::zeros({1}, i64), torch::ones({1}, i64), torch::zeros({1}, i64), + torch::zeros({2}, i64.dtype(index_type)), + torch::zeros({2, 3}, i64.dtype(torch::kFloat32)), torch::zeros({2}, i64), + torch::zeros({2}, i64), torch::arange(2, i64).to(index_type), + torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world); +} + std::vector DeepPotPTExpt::run_model_with_comm( const torch::Tensor& coord, const torch::Tensor& atype, @@ -831,17 +859,34 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; - // A long-range module sums over every pair of atoms of the system, which - // therefore has to be complete in this call: without periodic images and - // without other domains, both of which appear as ghost atoms. - if (long_range_ && (multi_rank || nghost_real > 0)) { + // A long-range module sums over every pair of atoms of the system. A single + // rank holds all of them, so a ghost atom there is a periodic image. Across + // ranks, the compact lower gathers the owned atoms of every rank through + // the communicator; the pairs that the graph lower lists stay within one. + if (long_range_ && !multi_rank && nghost_real > 0) { throw deepmd::deepmd_exception( "a model with a long-range module evaluates a complete non-periodic " - "system in a single domain, but this call carries " + - std::to_string(nghost_real) + " ghost atoms on " + + "system, but this single-rank call carries " + + std::to_string(nghost_real) + + " ghost atoms, the periodic images of a periodic cell; use a " + "non-periodic box."); + } + if (long_range_ && multi_rank && !lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a system decomposed over " + std::to_string(lmp_list.nprocs) + - " domain(s); use a non-periodic box on a single MPI rank."); + " MPI ranks only through the compact canonical lower, which a " + "compressed DPA4C model frozen for a CUDA device uses when the " + "direct-sum operator serves its long-range module; this archive uses " + "another lower, so run it on a single rank."); + } + if (long_range_ && multi_rank && lmp_list.world == nullptr) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a system decomposed over " + "several MPI ranks through their communicator, but the neighbor list " + "carries none (InputNlist::world)."); } + const void* const world = multi_rank ? lmp_list.world : nullptr; // Whether the edge topology folds ghost neighbours onto their local owners, // which reads an owner for every extended atom out of ``mapping_``. The // edge lower folds unless the with-comm artifact carries ghost features @@ -1270,7 +1315,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // ``Dim("n_node_total", min=1)``. Such a rank contributes nothing, so // fill zero outputs and return instead of running the model. (The // tested ``nloc == 0`` empty-subdomain case has ``nall_real > 0`` -- - // ghosts within rcut -- so it still runs the model normally.) + // ghosts within rcut -- so it still runs the model normally.) A + // long-range module across ranks still joins the exchanges of the + // others. + if (long_range_ && multi_rank) { + join_long_range_exchange(world); + } ener.assign(nframes, static_cast(0)); force.assign(static_cast(nframes) * fwd_map.size() * 3, static_cast(0)); @@ -1318,7 +1368,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, - compact.source_row_ptr, compact.source_order, node_coord); + compact.source_row_ptr, compact.source_order, node_coord, world); } else { flat_outputs = run_model_graph( node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, @@ -1706,7 +1756,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, - compact.source_order, node_coord); + compact.source_order, node_coord, /*world=*/nullptr); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -2277,8 +2327,10 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, if (long_range_ && (extended || comm_nlist != nullptr)) { throw deepmd::deepmd_exception( "compute_edges_gpu: a model with a long-range module evaluates a " - "complete non-periodic system in a single domain (nall_nodes == nloc " - "and no comm_nlist)."); + "complete non-periodic system on a single rank (nall_nodes == nloc and " + "no comm_nlist); several ranks require the compact canonical lower, " + "which a compressed DPA4C model frozen for a CUDA device uses when " + "the direct-sum operator serves its long-range module."); } if (lower_input_is_graph_ && extended && has_message_passing_) { throw deepmd::deepmd_exception( @@ -2597,6 +2649,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2613,18 +2666,35 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a CUDA device."); } + // A long-range module reads the positions of a complete non-periodic + // system: the node set of a single rank, or the owned nodes of every rank + // of ``world``, which a rank without nodes still joins. + if (long_range_ && d_coord == nullptr && nall_nodes > 0) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module reads " + "the node positions (d_coord)."); + } + if (long_range_ && world == nullptr && nall_nodes != nloc) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module on a " + "single rank evaluates a complete non-periodic system " + "(nall_nodes == nloc); a system decomposed over several ranks takes " + "their communicator."); + } + if (long_range_ && world != nullptr && nall_nodes == 0) { + translate_error([&] { + const c10::DeviceGuard device_guard(torch::Device(torch::kCUDA, gpu_id)); + join_long_range_exchange(world); + synchronize_current_accelerator_stream(); + }); + return; + } if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2 || static_cast(edge_storage) > std::numeric_limits::max()) { throw deepmd::deepmd_exception( "invalid compact canonical graph dimensions."); } - if (long_range_ && (d_coord == nullptr || nall_nodes != nloc)) { - throw deepmd::deepmd_exception( - "compute_canonical_graph_gpu: a model with a long-range module reads " - "the node positions of a complete non-periodic system in a single " - "domain (d_coord given and nall_nodes == nloc)."); - } translate_error([&] { const torch::Device device(torch::kCUDA, gpu_id); @@ -2659,10 +2729,10 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( : at::Tensor(); std::map output; - extract_outputs( - output, run_model_canonical_graph(atype, n_node, n_local, source, - edge_vec, destination_row_ptr, - source_row_ptr, source_order, coord)); + extract_outputs(output, run_model_canonical_graph( + atype, n_node, n_local, source, edge_vec, + destination_row_ptr, source_row_ptr, + source_order, coord, world)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2745,6 +2815,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( double* d_atom_virial, const std::int64_t* d_atype, const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2754,7 +2825,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { compute_canonical_graph_gpu_impl( - d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, d_source, + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage); } diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 090d8dc597..7dd715f48c 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -969,13 +969,16 @@ void PairDeepMD::settings(int narg, char** arg) { void PairDeepMD::init_style() { PairDeepBaseModel::init_style(); // A long-range module sums over every pair of atoms of the system, which a - // periodic image or another domain would leave incomplete. - if (deep_pot.has_long_range() && (domain->xperiodic || domain->yperiodic || - domain->zperiodic || comm->nprocs > 1)) { + // periodic image would leave incomplete; the atoms of the other ranks reach + // it through the communicator. + const bool long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() + : deep_pot.has_long_range(); + if (long_range && + (domain->xperiodic || domain->yperiodic || domain->zperiodic)) { error->all(FLERR, "pair style deepmd: a model with a long-range module evaluates " "a complete non-periodic system; use a non-periodic boundary " - "(f, s or m) in every dimension on a single MPI rank."); + "(f, s or m) in every dimension."); } } diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index db412a9ae1..9c24685e1a 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -214,7 +214,8 @@ void PairDeepMDKokkos::init_style() { edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors(); canonical_graph = deep_pot.uses_canonical_graph_inference(); // PairDeepMD::init_style() has confined such a model to a complete - // non-periodic system on one rank. + // non-periodic system, whose atoms on other ranks it reaches through the + // communicator. long_range = deep_pot.has_long_range(); // Force exchange transfers three values per atom. Centroid per-atom virial // uses nine values even though the Kokkos full-list request runs newton off, @@ -553,7 +554,10 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { comm_ptr = &comm_list; } - if ((canonical_graph && nnode_m > 0) || + // A long-range module across ranks exchanges the charges of every rank, so + // that every rank evaluates the model, also without nodes. + const bool long_range_across_ranks = long_range && multi_rank; + if ((canonical_graph && (nnode_m > 0 || long_range_across_ranks)) || (!canonical_graph && (nloc_m > 0 || comm_ptr != nullptr))) { // Fully device-resident inference: raw device pointers in and out. The // edge buffers are produced on the Kokkos stream and consumed by the model @@ -575,6 +579,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { deep_pot.compute_canonical_graph_gpu( d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), compact_graph.d_model_type.data(), long_range ? coord_ptr : nullptr, + long_range_across_ranks ? static_cast(&world) : nullptr, compact_graph.d_source.data(), compact_graph.d_edge_vec.data(), compact_graph.d_destination_row_ptr.data(), compact_graph.d_source_row_ptr.data(), diff --git a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py index a574cef821..a73908d6f3 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py @@ -51,6 +51,13 @@ "domain decomposition (nswap>0). Pass '1 1 1' for a single-rank " "reference run on the same archive (single-artifact dispatch).", ) +parser.add_argument( + "--boundary", + type=str, + default="p p p", + help="LAMMPS boundary style. Default 'p p p'; a model with a long-range " + "module evaluates complete non-periodic systems and passes 'f f f'.", +) parser.add_argument( "--pair-coeff", type=str, @@ -137,7 +144,7 @@ # (the regular artifact handles nswap==0). lammps.processors(args.processors) lammps.units("metal") -lammps.boundary("p p p") +lammps.boundary(args.boundary) lammps.atom_style("atomic") # ``atom_modify map yes`` is required when single-rank dispatch goes # through the regular artifact of a use_loc_mapping=False .pt2: the diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py index 71fb10a7be..5241bd8e5e 100644 --- a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -3,17 +3,22 @@ The latent-charge module sums over every pair of atoms of a complete non-periodic system, so ``pair_style deepmd`` runs it with non-periodic -boundaries on a single rank and rejects a periodic one. The model and the -reference values, without a cell, at the default charge state and at -``charge_spin 1.0 2.0``, are generated by -source/tests/infer/gen_dpa4c_long_range.py. +boundaries and rejects a periodic one. A compressed model also runs on +several ranks, whose atoms it gathers over MPI; an uncompressed one runs on a +single rank. The models and the reference values, without a cell, at the +default charge state and at ``charge_spin 1.0 2.0``, are generated by +source/tests/infer/gen_dpa4c_long_range.py, the compressed model for CUDA +only. """ +import importlib.util import os +import shutil from pathlib import ( Path, ) +import constants import numpy as np import pytest from expected_ref import ( @@ -24,6 +29,7 @@ ) from lammps_test_utils import ( make_atomic_lammps, + run_mpi_pair_runner, ) from write_lmp_data import ( write_lmp_data, @@ -41,7 +47,13 @@ / "infer" / "deeppot_dpa4c_long_range.expected" ) +compressed_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_compressed.pt2") +# A model without a long-range module, for a model deviation beside one. +short_range_file = pt2_file.with_name("deeppot_sea.pth") +compressed_ref_file = ref_file.with_name("deeppot_dpa4c_long_range_compressed.expected") data_file = Path(__file__).parent / "data_dpa4c_long_range_pt2.lmp" +replicated_data_file = Path(__file__).parent / "data_dpa4c_long_range_grid_pt2.lmp" +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" # Guarded with try/except because the generator only runs when PyTorch is # built; matrices that disable PyTorch still load this file at collection time. @@ -54,6 +66,10 @@ except FileNotFoundError: expected_e_default = expected_f_default = None expected_e_explicit = expected_f_explicit = None +try: + compressed_ref = read_expected_ref(compressed_ref_file) +except FileNotFoundError: + compressed_ref = None box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) coord = np.array( @@ -68,6 +84,10 @@ ) # atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2] type_OH = np.array([1, 2, 2, 1, 2, 2]) +# LAMMPS centroid/stress/atom reports the per-atom stress, the negative of the +# per-atom virial, with the components xx, yy, zz, xy, xz, yz, yx, zx, zy; these +# columns of it give the nine components of the reference virial. +expected_v_to_lammps = [0, 6, 7, 3, 1, 8, 4, 5, 2] def setup_module() -> None: @@ -77,8 +97,9 @@ def setup_module() -> None: def teardown_module() -> None: - if data_file.exists(): - os.remove(data_file) + for path in (data_file, replicated_data_file): + if path.exists(): + os.remove(path) def _lammps(boundary: str) -> PyLammps: @@ -129,3 +150,150 @@ def test_periodic_boundary_is_rejected(lammps_periodic) -> None: lammps_periodic.pair_coeff("* *") with pytest.raises(Exception, match="long-range module"): lammps_periodic.run(0) + + +@pytest.mark.parametrize("long_range_first", [True, False]) +def test_model_deviation_periodic_boundary_is_rejected( + lammps_periodic, long_range_first: bool, tmp_path: Path +) -> None: + """Any model of a model deviation with a long-range module confines it.""" + if not short_range_file.is_file(): + pytest.skip(f"Skip because {short_range_file} was not generated.") + models = [pt2_file, short_range_file] + if not long_range_first: + models.reverse() + lammps_periodic.pair_style( + "deepmd " + + " ".join(str(model.resolve()) for model in models) + + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}" + ) + lammps_periodic.pair_coeff("* *") + with pytest.raises(Exception, match="long-range module"): + lammps_periodic.run(0) + + +mpi_available = pytest.mark.skipif( + shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, + reason="MPI or mpi4py is not installed", +) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors", "section", "pair_extra"), + [ + # Both ranks own atoms. + (2, "2 1 1", "default", ""), + # Every atom lies in the first slab along z: two ranks hold ghost atoms + # alone and the last none at all, and all of them join the exchange. + (4, "1 1 4", "explicit", "charge_spin 1.0 2.0"), + ], +) +def test_pair_deepmd_mpi_compressed( + nprocs: int, processors: str, section: str, pair_extra: str +) -> None: + if compressed_ref is None or not compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") + runner_args = ["--boundary", "f f f"] + if pair_extra: + runner_args += ["--pair-extra", pair_extra] + result = run_mpi_pair_runner( + mpi_runner, + data_file, + compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=runner_args, + ) + expected_e = float(np.sum(compressed_ref[section]["expected_e"])) + expected_f = compressed_ref[section]["expected_f"].reshape(6, 3) + expected_v = -compressed_ref[section]["expected_v"].reshape(6, 9) + assert result["pe"] == pytest.approx(expected_e, rel=1e-6) + np.testing.assert_allclose(result["forces"], expected_f, rtol=0, atol=1e-5) + np.testing.assert_allclose( + result["virials"][:, expected_v_to_lammps] / constants.nktv2p, + expected_v, + rtol=0, + atol=1e-5, + ) + + +@pytest.fixture(scope="module") +def replicated_reference() -> dict: + """A grid of 168 copies of the six atoms, and its single-rank result. + + The copies lie farther apart than the short-range cutoff, so only the + long-range module couples them; the 1008 atoms fill several blocks and + chunks of the direct-sum operator. + """ + if compressed_ref is None or not compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") + shifts = np.stack( + np.meshgrid( + 18.0 * np.arange(4), 6.0 * np.arange(7), 6.0 * np.arange(6), indexing="ij" + ), + axis=-1, + ).reshape(-1, 1, 3) + grid_coord = (coord[None] + shifts).reshape(-1, 3) + grid_type = np.tile(type_OH, shifts.shape[0]) + # The box reaches beyond the atoms along z, so that the last of four slabs + # holds no atom, owned or ghost. + grid_box = np.array([0, 72, 0, 42, 0, 60, 0, 0, 0]) + write_lmp_data(grid_box, grid_coord, grid_type, replicated_data_file) + return run_mpi_pair_runner( + mpi_runner, + replicated_data_file, + compressed_pt2_file, + nprocs=1, + runner_args=["--boundary", "f f f"], + ) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors"), + [ + # Uneven slabs that cut through copies of the six atoms. + (3, "3 1 1"), + # The last slab holds no atom at all. + (4, "1 1 4"), + ], +) +def test_pair_deepmd_mpi_compressed_matches_single_rank( + replicated_reference: dict, nprocs: int, processors: str +) -> None: + result = run_mpi_pair_runner( + mpi_runner, + replicated_data_file, + compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=["--boundary", "f f f"], + ) + assert result["pe"] == pytest.approx(replicated_reference["pe"], rel=1e-6) + np.testing.assert_allclose( + result["forces"], replicated_reference["forces"], rtol=0, atol=1e-5 + ) + np.testing.assert_allclose( + result["virials"] / constants.nktv2p, + replicated_reference["virials"] / constants.nktv2p, + rtol=0, + atol=1e-5, + ) + + +@mpi_available +def test_pair_deepmd_mpi_uncompressed_is_rejected() -> None: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + result = run_mpi_pair_runner( + mpi_runner, + data_file, + pt2_file, + nprocs=2, + processors="2 1 1", + runner_args=["--boundary", "f f f"], + capture=True, + ) + assert result["returncode"] != 0 + assert "compact canonical lower" in result["stdout"] + result["stderr"] diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 52b4921d05..933f0fe7c9 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -181,6 +181,7 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) graph_fitting.cu edge_force_virial.cu long_range/direct_sum.cu + long_range/domain_exchange.cc dpa1_graph_energy_force.cu dpa4/so2_conv.cu ${DPA4_SO2_CONV_KERNEL_SRC} diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu index 47c5a0e330..a44667b002 100644 --- a/source/op/pt/long_range/direct_sum.cu +++ b/source/op/pt/long_range/direct_sum.cu @@ -29,6 +29,16 @@ // The charge constraint and the per-frame virial reduce over the nodes of // every frame, in slices whose float64 sums are added in a fixed order, so // that every output is independent of scheduling. +// +// The nodes whose sums a call evaluates, its owners, are either every node +// or, across the domains of a decomposed system, the charged nodes of one +// MPI rank. The ranks then gather the positions and charges of the charged +// nodes of every domain into one frame, in rank order (domain_exchange.h), +// and each rank evaluates its own nodes against the whole frame, so that the +// quadratic work divides among the ranks. Every rank shifts the constrained +// channel by the same sum of the gathered charges; the mean potential of the +// constraint adds the sums of the ranks in rank order, and the per-frame +// virial is the share of the owners of the rank. #include #include @@ -41,6 +51,8 @@ #include #include +#include "domain_exchange.h" + namespace { // Nodes per block, which is also the length of a tile of partners. @@ -189,11 +201,11 @@ __device__ __forceinline__ void write_slice_sum( } // Per-frame statistics of the charge constraint: the shift of the -// constrained channel, the number of charged nodes, and the mean potential of -// the charged nodes in that channel. +// constrained channel, the number of charged nodes, and the potential of the +// charged nodes in that channel summed over them. constexpr int kStatShift = 0; constexpr int kStatCount = 1; -constexpr int kStatMean = 2; +constexpr int kStatPotential = 2; constexpr int kStatWidth = 3; // Sums, over one slice of one frame, the number of charged nodes and their raw @@ -490,9 +502,10 @@ __device__ __forceinline__ void accumulate_tile( } } -// Sums, for every node of one block and every partner of one chunk, the +// Sums, for every owner node of one block and every partner of one chunk, the // unscaled potential, force direction and virial of the node, and writes them -// to ``partial[chunk][component][node]``. +// to ``partial[chunk][component][node - owner_begin]``. The owners are the +// nodes [owner_begin, owner_begin + owner_count). template __global__ void __launch_bounds__(kThreads) direct_sum_kernel(const float4* __restrict__ head, @@ -500,7 +513,8 @@ __global__ void __launch_bounds__(kThreads) const float* __restrict__ charge, const int* __restrict__ frame, const int64_t* __restrict__ frame_ptr, - int node_count, + int owner_begin, + int owner_count, const float* __restrict__ coefficient, int gaussian_count, float* __restrict__ partial) { @@ -512,9 +526,10 @@ __global__ void __launch_bounds__(kThreads) __shared__ int range[2]; __shared__ bool single_frame; - const int first = blockIdx.x * kThreads; + const int owner_end = owner_begin + owner_count; + const int first = owner_begin + blockIdx.x * kThreads; const int node = first + threadIdx.x; - const int own_frame = node < node_count ? frame[node] : -1; + const int own_frame = node < owner_end ? frame[node] : -1; if constexpr (kKernel == kSog) { for (int k = threadIdx.x; k < gaussian_count; k += kThreads) { @@ -528,7 +543,7 @@ __global__ void __launch_bounds__(kThreads) // to: a contiguous range, since frames are contiguous on the node axis and // padding nodes trail them. if (threadIdx.x == 0) { - int last = min(first + kThreads, node_count) - 1; + int last = min(first + kThreads, owner_end) - 1; while (last > first && frame[last] < 0) { --last; } @@ -539,7 +554,7 @@ __global__ void __launch_bounds__(kThreads) // A block whose nodes all lie in one frame finds every partner of its // range in that frame, which spares the frame test of every pair. single_frame = low_frame >= 0 && low_frame == high_frame && - min(first + kThreads, node_count) - 1 == last; + min(first + kThreads, owner_end) - 1 == last; } __syncthreads(); @@ -614,24 +629,25 @@ __global__ void __launch_bounds__(kThreads) __syncthreads(); } - if (node < node_count) { - float* out = - partial + static_cast(blockIdx.y) * kWidth * node_count + node; + if (node < owner_end) { + float* out = partial + + static_cast(blockIdx.y) * kWidth * owner_count + + (node - owner_begin); #pragma unroll for (int k = 0; k < kWidth; ++k) { - out[static_cast(k) * node_count] = total[k].sum; + out[static_cast(k) * owner_count] = total[k].sum; } } } -// Adds the chunk sums of one component of one node in chunk order, one -// thread per component and node, and scales the result into the potential, -// the force or the per-node virial. With positions scaled by alpha, -// K(r) = alpha f(x), K'(r) / r = alpha^3 f'(x) / x and d = d' / alpha, so the -// potential carries alpha, the force alpha^2 and the virial alpha. +// Adds the chunk sums of one component of one owner node in chunk order, one +// thread per component and owner, and scales the result into the potential, +// the force or the per-node virial of the owner. With positions scaled by +// alpha, K(r) = alpha f(x), K'(r) / r = alpha^3 f'(x) / x and d = d' / alpha, +// so the potential carries alpha, the force alpha^2 and the virial alpha. template __global__ void reduce_chunks(const float* __restrict__ partial, - int node_count, + int owner_count, int chunk_count, double prefactor, const float* __restrict__ alpha, @@ -641,14 +657,15 @@ __global__ void reduce_chunks(const float* __restrict__ partial, constexpr int kWidth = sum_width(); const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (index >= static_cast(kWidth) * node_count) { + if (index >= static_cast(kWidth) * owner_count) { return; } - const int component = static_cast(index / node_count); - const int node = static_cast(index - component * int64_t(node_count)); + const int component = static_cast(index / owner_count); + const int node = static_cast(index - component * int64_t(owner_count)); double value = 0.0; for (int chunk = 0; chunk < chunk_count; ++chunk) { - value += partial[static_cast(chunk) * kWidth * node_count + index]; + value += + partial[static_cast(chunk) * kWidth * owner_count + index]; } const double unit = alpha == nullptr ? 1.0 : static_cast(*alpha); if (component < kChannels) { @@ -700,7 +717,7 @@ __global__ void __launch_bounds__(kFrameThreads) } // Adds the slice sums of virial_slices into the virial of every frame and, -// with a constraint, the mean potential of its charged nodes. +// with a constraint, the potential summed over its charged nodes. template __global__ void virial_total(const double* __restrict__ slice_sum, int slice_count, @@ -720,32 +737,34 @@ __global__ void virial_total(const double* __restrict__ slice_sum, if (threadIdx.x < 9) { virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value); } else { - double* stat = frame_stat + blockIdx.x * kStatWidth; - stat[kStatMean] = value / stat[kStatCount]; + frame_stat[blockIdx.x * kStatWidth + kStatPotential] = value; } } -// Writes the atomic energy q_i . phi_i / 2 and the derivative of the energy -// with respect to the raw charges: w_i (phi_i - mean_f) in the constrained -// channel and w_i phi_i in the others, w_i marking the charged nodes. +// Writes the atomic energy q_i . phi_i / 2 of every owner node and the +// derivative of the energy with respect to its raw charges: w_i (phi_i - +// mean_f) in the constrained channel and w_i phi_i in the others, w_i marking +// the charged nodes. ``pair_charge``, ``node_mask`` and ``frame`` start at the +// first owner. __global__ void project_charges(const float* __restrict__ potential, const float* __restrict__ pair_charge, const bool* __restrict__ node_mask, const int* __restrict__ frame, const double* __restrict__ frame_stat, - int node_count, + int owner_count, int channels, int channel, double* __restrict__ atom_energy, float* __restrict__ charge_gradient) { const int node = blockIdx.x * blockDim.x + threadIdx.x; - if (node >= node_count) { + if (node >= owner_count) { return; } const int f = frame[node]; const bool charged = f >= 0 && node_mask[node]; const double mean = charged && frame_stat != nullptr - ? frame_stat[f * kStatWidth + kStatMean] + ? frame_stat[f * kStatWidth + kStatPotential] / + frame_stat[f * kStatWidth + kStatCount] : 0.0; double energy = 0.0; for (int c = 0; c < channels; ++c) { @@ -759,23 +778,27 @@ __global__ void project_charges(const float* __restrict__ potential, atom_energy[node] = 0.5 * energy; } -// Chunks of the partner axis for the given node count. The blocks of one -// launch all carry the same work, so the launch fills the resident block -// slots of the device several times over, which bounds the idle share of the -// last wave; no chunk is shorter than kMinChunkTiles tiles of a frame that -// spans every node. -int chunk_count_for(int node_count, int resident_blocks_per_multiprocessor) { +// Chunks of the partner axis for the given counts of owners and partners. The +// blocks of one launch all carry the same work, so the launch fills the +// resident block slots of the device several times over, which bounds the +// idle share of the last wave; no chunk is shorter than kMinChunkTiles tiles +// of a frame that spans every partner. +int chunk_count_for(int owner_count, + int partner_count, + int resident_blocks_per_multiprocessor) { const int multiprocessors = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; - const int node_blocks = (node_count + kThreads - 1) / kThreads; + const int owner_blocks = (owner_count + kThreads - 1) / kThreads; + const int partner_tiles = (partner_count + kThreads - 1) / kThreads; const int64_t wanted_blocks = static_cast(kWaves) * resident_blocks_per_multiprocessor * multiprocessors; const int wanted = - static_cast((wanted_blocks + node_blocks - 1) / node_blocks); - return std::max(1, std::min(wanted, node_blocks / kMinChunkTiles)); + static_cast((wanted_blocks + owner_blocks - 1) / owner_blocks); + return std::max(1, std::min(wanted, partner_tiles / kMinChunkTiles)); } +// The staged partners, one row per node, and the sums of the owners. struct DirectSumBuffers { const torch::Tensor& head; const torch::Tensor& tail; @@ -803,29 +826,31 @@ int resident_blocks_per_multiprocessor() { template void launch(const DirectSumBuffers& buffers, - int node_count, + int owner_begin, + int owner_count, double prefactor, cudaStream_t stream) { constexpr int kWidth = sum_width(); - const int chunk_count = chunk_count_for( - node_count, resident_blocks_per_multiprocessor()); + const int chunk_count = + chunk_count_for(owner_count, static_cast(buffers.head.size(0)), + resident_blocks_per_multiprocessor()); auto partial = - torch::empty({static_cast(chunk_count) * node_count * kWidth}, + torch::empty({static_cast(chunk_count) * owner_count * kWidth}, buffers.charge.options()); - const dim3 grid((node_count + kThreads - 1) / kThreads, chunk_count); + const dim3 grid((owner_count + kThreads - 1) / kThreads, chunk_count); direct_sum_kernel<<>>( reinterpret_cast(buffers.head.data_ptr()), reinterpret_cast(buffers.tail.data_ptr()), buffers.charge.data_ptr(), buffers.frame.data_ptr(), - buffers.frame_ptr.data_ptr(), node_count, + buffers.frame_ptr.data_ptr(), owner_begin, owner_count, buffers.coefficient.data_ptr(), static_cast(buffers.coefficient.numel() / 2), partial.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); - const int64_t reduce_threads = static_cast(kWidth) * node_count; + const int64_t reduce_threads = static_cast(kWidth) * owner_count; reduce_chunks <<((reduce_threads + 255) / 256), 256, 0, stream>>>( - partial.data_ptr(), node_count, chunk_count, prefactor, + partial.data_ptr(), owner_count, chunk_count, prefactor, kKernel == kLes ? buffers.coefficient.data_ptr() : nullptr, buffers.potential.data_ptr(), buffers.force.data_ptr(), buffers.atom_virial.data_ptr()); @@ -835,142 +860,108 @@ void launch(const DirectSumBuffers& buffers, template void dispatch_channels(int channels, const DirectSumBuffers& buffers, - int node_count, + int owner_begin, + int owner_count, double prefactor, cudaStream_t stream) { switch (channels) { case 1: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 2: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 3: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 4: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 5: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 6: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 7: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); case 8: - return launch(buffers, node_count, prefactor, stream); + return launch(buffers, owner_begin, owner_count, prefactor, + stream); default: TORCH_CHECK(false, "latent_charge_direct_sum: unsupported channel count ", channels); } } -std::tuple -latent_charge_direct_sum(torch::Tensor coord, - torch::Tensor charge, - torch::Tensor node_mask, - torch::Tensor n_node_per_frame, - torch::Tensor total_charge, - torch::Tensor coefficient, - int64_t kernel, - int64_t channel, - double prefactor) { - TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && - n_node_per_frame.is_cuda() && total_charge.is_cuda() && - coefficient.is_cuda(), - "latent_charge_direct_sum: inputs must be CUDA tensors"); - TORCH_CHECK(charge.device() == coord.device() && - node_mask.device() == coord.device() && - n_node_per_frame.device() == coord.device() && - total_charge.device() == coord.device() && - coefficient.device() == coord.device(), - "latent_charge_direct_sum: inputs must share one device"); - TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && - coord.scalar_type() == torch::kFloat64, - "latent_charge_direct_sum: coord must be float64 with shape " - "(N, 3)"); - TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && - charge.scalar_type() == torch::kFloat32, - "latent_charge_direct_sum: charge must be float32 with shape " - "(N, C)"); - TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && - node_mask.scalar_type() == torch::kBool, - "latent_charge_direct_sum: node_mask must be bool with shape " - "(N,)"); - TORCH_CHECK(n_node_per_frame.dim() == 1, - "latent_charge_direct_sum: n_node_per_frame must be " - "one-dimensional"); - TORCH_CHECK(total_charge.dim() == 1 && - total_charge.size(0) == n_node_per_frame.size(0) && - total_charge.scalar_type() == torch::kFloat64, - "latent_charge_direct_sum: total_charge must be float64 with " - "one entry per frame"); - TORCH_CHECK( - coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat32, - "latent_charge_direct_sum: coefficient must be a float32 vector"); - const int64_t channels = charge.size(1); - TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, - "latent_charge_direct_sum: the channel count must lie in [1, ", - kMaxChannels, "], got ", channels); - TORCH_CHECK(channel >= -1 && channel < channels, - "latent_charge_direct_sum: the constrained channel must be -1 " - "or lie in [0, ", - channels, "), got ", channel); +using DirectSumOutputs = std::tuple; + +// Evaluates the owner nodes [owner_begin, owner_begin + owner_count) against +// the nodes of their frames. Returns the outputs of the owners, row r for +// owner owner_begin + r, and the per-frame virial of the owners. Owners other +// than every node lie in a single frame. A non-null ``communicator`` spans +// the domains whose owners share that frame, over which the potential of the +// constraint is summed. +DirectSumOutputs evaluate(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) { TORCH_CHECK(coord.size(0) < std::numeric_limits::max() / kThreads, "latent_charge_direct_sum: too many nodes"); - if (kernel == kLes) { - TORCH_CHECK(coefficient.numel() == 1, - "latent_charge_direct_sum: the erf kernel takes one " - "coefficient"); - } else if (kernel == kSog) { - TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && - coefficient.numel() <= 2 * kMaxGaussians, - "latent_charge_direct_sum: the Gaussian kernel takes two " - "coefficients per Gaussian, for at most ", - kMaxGaussians, " Gaussians"); - } else { - TORCH_CHECK(false, "latent_charge_direct_sum: unknown kernel ", kernel); - } - - const c10::cuda::CUDAGuard guard(coord.device()); const int node_count = static_cast(coord.size(0)); const int frame_count = static_cast(n_node_per_frame.size(0)); + const int64_t channels = charge.size(1); const auto float_options = charge.options(); const auto double_options = coord.options(); if (node_count == 0 || frame_count == 0) { - return {torch::zeros({node_count, 1}, double_options), - torch::zeros({node_count, channels}, float_options), - torch::zeros({node_count, 3}, float_options), - torch::zeros({node_count, 9}, float_options), + return {torch::zeros({owner_count, 1}, double_options), + torch::zeros({owner_count, channels}, float_options), + torch::zeros({owner_count, 3}, float_options), + torch::zeros({owner_count, 9}, float_options), torch::zeros({frame_count, 9}, float_options)}; } + const bool every_node = owner_begin == 0 && owner_count == node_count; + TORCH_CHECK(every_node || frame_count == 1, + "latent_charge_direct_sum: owners other than every node lie in " + "one frame"); const bool constrained = channel >= 0; const int constrained_channel = static_cast(channel); // The kernels write every entry of the outputs. - auto atom_energy = torch::empty({node_count, 1}, double_options); - auto charge_gradient = torch::empty({node_count, channels}, float_options); - auto force = torch::empty({node_count, 3}, float_options); - auto atom_virial = torch::empty({node_count, 9}, float_options); + auto atom_energy = torch::empty({owner_count, 1}, double_options); + auto charge_gradient = torch::empty({owner_count, channels}, float_options); + auto force = torch::empty({owner_count, 3}, float_options); + auto atom_virial = torch::empty({owner_count, 9}, float_options); auto virial = torch::empty({frame_count, 9}, float_options); - coord = coord.contiguous(); - charge = charge.contiguous(); - node_mask = node_mask.contiguous(); - total_charge = total_charge.contiguous(); - coefficient = coefficient.contiguous(); const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); auto frame_ptr = torch::empty({frame_count + 1}, counts.options()); auto head = torch::empty({node_count, 4}, float_options); auto tail = torch::empty({node_count, 4}, float_options); auto frame = torch::empty({node_count}, coord.options().dtype(torch::kInt32)); auto pair_charge = torch::empty({node_count, channels}, float_options); - auto potential = torch::empty({node_count, channels}, float_options); + auto potential = torch::empty({owner_count, channels}, float_options); auto frame_stat = torch::empty({constrained ? frame_count * kStatWidth : 0}, double_options); - const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; - const int slice_count = static_cast(std::min( - kMaxSlices, - std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); + const auto slices_for = [frame_count](int64_t nodes) { + const int64_t average_nodes = (nodes + frame_count - 1) / frame_count; + return static_cast(std::min( + kMaxSlices, + std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); + }; + const int slice_count = slices_for(node_count); const dim3 slice_grid(frame_count, slice_count); const auto stream = at::cuda::getCurrentCUDAStream(); frame_offsets<<<1, kOffsetThreads, 0, stream>>>( @@ -1003,53 +994,214 @@ latent_charge_direct_sum(torch::Tensor coord, pair_charge.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); - const DirectSumBuffers buffers{head, tail, pair_charge, - frame, frame_ptr, coefficient, - potential, force, atom_virial}; - if (kernel == kLes) { - dispatch_channels(static_cast(channels), buffers, node_count, - prefactor, stream.stream()); - } else { - dispatch_channels(static_cast(channels), buffers, node_count, - prefactor, stream.stream()); + if (owner_count > 0) { + const DirectSumBuffers buffers{head, tail, pair_charge, + frame, frame_ptr, coefficient, + potential, force, atom_virial}; + if (kernel == kLes) { + dispatch_channels(static_cast(channels), buffers, owner_begin, + owner_count, prefactor, stream.stream()); + } else { + dispatch_channels(static_cast(channels), buffers, owner_begin, + owner_count, prefactor, stream.stream()); + } + } + // The frames of the owners, on the axis of the owners: the frames + // themselves when every node is an owner, else the one frame the owners + // share, written on the device. + torch::Tensor owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_ptr = + torch::full({2}, int64_t{owner_count}, frame_ptr.options()); + owner_frame_ptr.select(0, 0).zero_(); } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const int owner_slice_count = slices_for(owner_count); + const dim3 owner_slice_grid(frame_count, owner_slice_count); const int components = constrained ? 10 : 9; auto slice_virial = torch::empty( - {static_cast(frame_count) * slice_count * components}, + {static_cast(frame_count) * owner_slice_count * components}, double_options); if (constrained) { - virial_slices<<>>( - atom_virial.data_ptr(), potential.data_ptr(), - node_mask.data_ptr(), frame_ptr.data_ptr(), - static_cast(channels), constrained_channel, - slice_virial.data_ptr()); + virial_slices<<>>( + atom_virial.data_ptr(), potential.data_ptr(), owner_mask, + owner_frame_ptr.data_ptr(), static_cast(channels), + constrained_channel, slice_virial.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); virial_total<<>>( - slice_virial.data_ptr(), slice_count, virial.data_ptr(), - frame_stat.data_ptr()); + slice_virial.data_ptr(), owner_slice_count, + virial.data_ptr(), frame_stat.data_ptr()); } else { - virial_slices<<>>( - atom_virial.data_ptr(), potential.data_ptr(), - node_mask.data_ptr(), frame_ptr.data_ptr(), - static_cast(channels), constrained_channel, - slice_virial.data_ptr()); + virial_slices<<>>( + atom_virial.data_ptr(), potential.data_ptr(), owner_mask, + owner_frame_ptr.data_ptr(), static_cast(channels), + constrained_channel, slice_virial.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); virial_total<<>>( - slice_virial.data_ptr(), slice_count, virial.data_ptr(), - nullptr); + slice_virial.data_ptr(), owner_slice_count, + virial.data_ptr(), nullptr); } C10_CUDA_KERNEL_LAUNCH_CHECK(); - project_charges<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, - stream>>>( - potential.data_ptr(), pair_charge.data_ptr(), - node_mask.data_ptr(), frame.data_ptr(), - constrained ? frame_stat.data_ptr() : nullptr, node_count, - static_cast(channels), constrained_channel, - atom_energy.data_ptr(), charge_gradient.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained && communicator != nullptr) { + // The mean potential of the constraint runs over the charged nodes of + // every domain. + auto potential_sum = + frame_stat.view({frame_count, kStatWidth}).select(1, kStatPotential); + potential_sum.copy_(deepmd::long_range::sum_over_domains( + *communicator, potential_sum.contiguous())); + } + if (owner_count > 0) { + project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + potential.data_ptr(), + pair_charge.data_ptr() + int64_t{owner_begin} * channels, + owner_mask, frame.data_ptr() + owner_begin, + constrained ? frame_stat.data_ptr() : nullptr, owner_count, + static_cast(channels), constrained_channel, + atom_energy.data_ptr(), charge_gradient.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } return {atom_energy, charge_gradient, force, atom_virial, virial}; } +// Evaluates the frame of a system decomposed into the domains of +// ``communicator``, of which this rank holds one: its charged nodes are its +// owners, the charged nodes of every domain its partners. The nodes the rank +// does not own carry no charge, and the per-frame virial is the share of its +// owners. +DirectSumOutputs evaluate_across_domains(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + const torch::Tensor& communicator) { + TORCH_CHECK(n_node_per_frame.size(0) == 1, + "latent_charge_direct_sum: an evaluation across MPI ranks takes " + "one frame, got ", + n_node_per_frame.size(0)); + const int64_t node_count = coord.size(0); + // The charged nodes of the frame; the padding nodes past it carry none. + // Their count is the one value the host has to wait for before the + // exchange. + const auto in_frame = + torch::arange(node_count, n_node_per_frame.options()) < n_node_per_frame; + const auto owned = torch::nonzero(node_mask & in_frame).squeeze(1); + const auto rows = + torch::cat({coord.index_select(0, owned), + charge.index_select(0, owned).to(torch::kFloat64)}, + 1) + .contiguous(); + const auto [partners, owner_begin] = + deepmd::long_range::gather_domain_rows(communicator, rows); + const int64_t partner_count = partners.size(0); + const auto partner_coord = partners.slice(1, 0, 3).contiguous(); + const auto partner_charge = + partners.slice(1, 3).to(torch::kFloat32).contiguous(); + const auto partner_mask = torch::ones({partner_count}, node_mask.options()); + const auto partner_frame = + torch::full({1}, partner_count, n_node_per_frame.options()); + const auto [owner_energy, owner_gradient, owner_force, owner_atom_virial, + virial] = + evaluate(partner_coord, partner_charge, partner_mask, partner_frame, + total_charge, coefficient, kernel, channel, prefactor, + static_cast(owner_begin), static_cast(owned.size(0)), + &communicator); + return {torch::zeros({node_count, 1}, coord.options()) + .index_copy_(0, owned, owner_energy), + torch::zeros({node_count, charge.size(1)}, charge.options()) + .index_copy_(0, owned, owner_gradient), + torch::zeros({node_count, 3}, charge.options()) + .index_copy_(0, owned, owner_force), + torch::zeros({node_count, 9}, charge.options()) + .index_copy_(0, owned, owner_atom_virial), + virial}; +} + +DirectSumOutputs latent_charge_direct_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor) { + TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && + n_node_per_frame.is_cuda() && total_charge.is_cuda() && + coefficient.is_cuda(), + "latent_charge_direct_sum: inputs must be CUDA tensors"); + TORCH_CHECK(charge.device() == coord.device() && + node_mask.device() == coord.device() && + n_node_per_frame.device() == coord.device() && + total_charge.device() == coord.device() && + coefficient.device() == coord.device(), + "latent_charge_direct_sum: inputs must share one device"); + TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && + coord.scalar_type() == torch::kFloat64, + "latent_charge_direct_sum: coord must be float64 with shape " + "(N, 3)"); + TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && + charge.scalar_type() == torch::kFloat32, + "latent_charge_direct_sum: charge must be float32 with shape " + "(N, C)"); + TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && + node_mask.scalar_type() == torch::kBool, + "latent_charge_direct_sum: node_mask must be bool with shape " + "(N,)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "latent_charge_direct_sum: n_node_per_frame must be " + "one-dimensional"); + TORCH_CHECK(total_charge.dim() == 1 && + total_charge.size(0) == n_node_per_frame.size(0) && + total_charge.scalar_type() == torch::kFloat64, + "latent_charge_direct_sum: total_charge must be float64 with " + "one entry per frame"); + TORCH_CHECK( + coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat32, + "latent_charge_direct_sum: coefficient must be a float32 vector"); + const int64_t channels = charge.size(1); + TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, + "latent_charge_direct_sum: the channel count must lie in [1, ", + kMaxChannels, "], got ", channels); + TORCH_CHECK(channel >= -1 && channel < channels, + "latent_charge_direct_sum: the constrained channel must be -1 " + "or lie in [0, ", + channels, "), got ", channel); + if (kernel == kLes) { + TORCH_CHECK(coefficient.numel() == 1, + "latent_charge_direct_sum: the erf kernel takes one " + "coefficient"); + } else if (kernel == kSog) { + TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && + coefficient.numel() <= 2 * kMaxGaussians, + "latent_charge_direct_sum: the Gaussian kernel takes two " + "coefficients per Gaussian, for at most ", + kMaxGaussians, " Gaussians"); + } else { + TORCH_CHECK(false, "latent_charge_direct_sum: unknown kernel ", kernel); + } + + const c10::cuda::CUDAGuard guard(coord.device()); + coord = coord.contiguous(); + charge = charge.contiguous(); + node_mask = node_mask.contiguous(); + total_charge = total_charge.contiguous(); + coefficient = coefficient.contiguous(); + if (deepmd::long_range::domain_of(communicator).size == 1) { + return evaluate(coord, charge, node_mask, n_node_per_frame, total_charge, + coefficient, kernel, channel, prefactor, 0, + static_cast(coord.size(0)), nullptr); + } + return evaluate_across_domains(coord, charge, node_mask, n_node_per_frame, + total_charge, coefficient, kernel, channel, + prefactor, communicator); +} + } // namespace TORCH_LIBRARY_FRAGMENT(deepmd, library) { diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc new file mode 100644 index 0000000000..71acbc6957 --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.cc @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// MPI exchanges of the long-range direct sum (contract in domain_exchange.h). +// +// With an MPI library that reads device memory (MPIX_Query_cuda_support), the +// rows travel between device buffers once the current stream has finished +// writing them; otherwise they are staged through pinned host buffers. Either +// way the host waits for the device once per exchange, before MPI reads the +// data; the results return to the device without a further wait. + +#include "domain_exchange.h" + +#ifdef USE_MPI +#include +#ifdef OMPI_MPI_H +#include +#endif +#endif +#include +#include + +#include +#include +#include + +namespace deepmd { +namespace long_range { +namespace { + +std::int64_t handle_of(const torch::Tensor& communicator) { + TORCH_CHECK(communicator.device().is_cpu() && + communicator.scalar_type() == torch::kInt64 && + communicator.numel() == 1, + "latent_charge_direct_sum: communicator must be an int64 CPU " + "tensor with one element"); + return communicator.contiguous().data_ptr()[0]; +} + +#ifdef USE_MPI +// The tensor holds the bytes of the handle from its first byte on, whatever +// the width of MPI_Comm in the MPI library. +MPI_Comm communicator_of(const torch::Tensor& communicator) { + const torch::Tensor bytes = communicator.contiguous(); + MPI_Comm comm; + static_assert(sizeof(MPI_Comm) <= sizeof(std::int64_t), + "an MPI_Comm handle must fit in an int64 tensor"); + std::memcpy(&comm, bytes.data_ptr(), sizeof(MPI_Comm)); + return comm; +} + +bool reads_device_memory() { +#ifdef NO_CUDA_AWARE + return false; +#else + return MPIX_Query_cuda_support() != 0; +#endif +} +#endif + +} // namespace + +Domain domain_of(const torch::Tensor& communicator) { + if (handle_of(communicator) == 0) { + return {0, 1}; + } +#ifdef USE_MPI + int initialized = 0; + MPI_Initialized(&initialized); + TORCH_CHECK(initialized != 0, + "latent_charge_direct_sum: a communicator was given, but MPI is " + "not initialized"); + const MPI_Comm comm = communicator_of(communicator); + Domain domain{0, 1}; + MPI_Comm_rank(comm, &domain.rank); + MPI_Comm_size(comm, &domain.size); + return domain; +#else + TORCH_CHECK(false, + "latent_charge_direct_sum: an evaluation across MPI ranks " + "requires DeePMD-kit built with MPI"); +#endif +} + +std::pair gather_domain_rows( + const torch::Tensor& communicator, const torch::Tensor& rows) { +#ifdef USE_MPI + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && + rows.scalar_type() == torch::kFloat64 && rows.is_contiguous(), + "gather_domain_rows: rows must be a contiguous float64 CUDA " + "matrix"); + const MPI_Comm comm = communicator_of(communicator); + int rank = 0; + int size = 1; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + const std::int64_t width = rows.size(1); + std::int64_t local = rows.size(0); + std::vector counts(size); + MPI_Allgather(&local, 1, MPI_INT64_T, counts.data(), 1, MPI_INT64_T, comm); + std::vector receive_counts(size); + std::vector displacements(size); + std::int64_t total = 0; + std::int64_t offset = 0; + for (int r = 0; r < size; ++r) { + if (r == rank) { + offset = total; + } + TORCH_CHECK((total + counts[r]) * width <= std::numeric_limits::max(), + "gather_domain_rows: the rows of all ranks exceed the int " + "range of an MPI count"); + receive_counts[r] = static_cast(counts[r] * width); + displacements[r] = static_cast(total * width); + total += counts[r]; + } + const int send_count = static_cast(local * width); + auto gathered = torch::empty({total, width}, rows.options()); + if (reads_device_memory()) { + at::cuda::getCurrentCUDAStream().synchronize(); + MPI_Allgatherv(rows.data_ptr(), send_count, MPI_DOUBLE, + gathered.data_ptr(), receive_counts.data(), + displacements.data(), MPI_DOUBLE, comm); + } else { + const auto host = rows.options().device(torch::kCPU).pinned_memory(true); + auto host_rows = torch::empty({local, width}, host); + host_rows.copy_(rows); + auto host_gathered = torch::empty({total, width}, host); + MPI_Allgatherv(host_rows.data_ptr(), send_count, MPI_DOUBLE, + host_gathered.data_ptr(), receive_counts.data(), + displacements.data(), MPI_DOUBLE, comm); + gathered.copy_(host_gathered, /*non_blocking=*/true); + } + return {gathered, offset}; +#else + TORCH_CHECK(false, + "gather_domain_rows: an evaluation across MPI ranks requires " + "DeePMD-kit built with MPI"); +#endif +} + +torch::Tensor sum_over_domains(const torch::Tensor& communicator, + const torch::Tensor& values) { +#ifdef USE_MPI + TORCH_CHECK(values.dim() == 1 && values.scalar_type() == torch::kFloat64, + "sum_over_domains: values must be a float64 vector"); + const MPI_Comm comm = communicator_of(communicator); + int size = 1; + MPI_Comm_size(comm, &size); + const torch::Tensor local = values.to(torch::kCPU).contiguous(); + const int count = static_cast(local.numel()); + std::vector every(static_cast(size) * count); + MPI_Allgather(local.data_ptr(), count, MPI_DOUBLE, every.data(), + count, MPI_DOUBLE, comm); + auto total = torch::zeros({count}, local.options().pinned_memory(true)); + double* sum = total.data_ptr(); + for (int r = 0; r < size; ++r) { + for (int k = 0; k < count; ++k) { + sum[k] += every[static_cast(r) * count + k]; + } + } + return total.to(values.device(), /*non_blocking=*/true); +#else + TORCH_CHECK(false, + "sum_over_domains: an evaluation across MPI ranks requires " + "DeePMD-kit built with MPI"); +#endif +} + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h new file mode 100644 index 0000000000..323e60d92d --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.h @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Exchange of the long-range direct sum between the domains of a decomposed +// system. +// +// A domain is one MPI rank of the communicator that a caller hands to +// deepmd::latent_charge_direct_sum as an int64 tensor holding the bytes of its +// MPI_Comm handle, the convention of deepmd_export::border_op. A zero handle +// stands for a single domain. Both exchanges follow the rank order, so every +// rank receives the same rows and the same sums, independent of the +// reduction algorithm of the MPI library. + +#pragma once + +#include + +#include +#include + +namespace deepmd { +namespace long_range { + +struct Domain { + int rank; + int size; +}; + +// The rank of this process and the number of ranks of ``communicator``; one +// rank for a zero handle. +Domain domain_of(const torch::Tensor& communicator); + +// Concatenates, in rank order, the rows of every rank. ``rows`` is a +// contiguous float64 matrix of shape (M_r, W) on the CUDA device, with the +// same W on every rank. Returns the rows of all ranks, of shape (sum_r M_r, +// W) on the device of ``rows``, and the index of the first row of this rank. +std::pair gather_domain_rows( + const torch::Tensor& communicator, const torch::Tensor& rows); + +// Sums the float64 vector ``values`` over the ranks, adding the vectors of +// the ranks in rank order. Returns the sum on the device of ``values``. +torch::Tensor sum_over_domains(const torch::Tensor& communicator, + const torch::Tensor& values); + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc index 7f98709fe1..d1bf4b4aa6 100644 --- a/source/op/pt/long_range/ops.cc +++ b/source/op/pt/long_range/ops.cc @@ -43,6 +43,14 @@ // float64 and is read only with a constrained channel. The force and the // virial hold the charges fixed; the response of the charges is the // caller's, through the charge gradient. +// +// `communicator` is an int64 CPU tensor of one element holding the bytes of +// an MPI_Comm handle, as for deepmd_export::border_op, or zero. A handle of +// several ranks decomposes one frame into domains, one per rank: the charged +// nodes of every rank are the charged nodes of the frame, and the sums run +// over all of them. The outputs of a rank are those of its charged nodes, and +// W_f is their share of the virial of the frame; the other nodes of the rank +// carry no charge. #include @@ -50,7 +58,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { library.def( "latent_charge_direct_sum(Tensor coord, Tensor charge, " "Tensor node_mask, Tensor n_node_per_frame, Tensor total_charge, " - "Tensor coefficient, int kernel, int channel, float prefactor) -> " + "Tensor coefficient, Tensor communicator, int kernel, int channel, " + "float prefactor) -> " "(Tensor atom_energy, Tensor charge_gradient, Tensor force, " "Tensor atom_virial, Tensor virial)"); } diff --git a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py index b02ce0758d..475562488c 100644 --- a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py +++ b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py @@ -94,6 +94,7 @@ def _evaluate( n_node, total_charge, coefficient.double(), + torch.zeros(1, dtype=torch.int64), code, channel, COULOMB_CONSTANT, @@ -197,6 +198,7 @@ def test_reference_outputs_are_energy_derivatives( n_node, total_charge, coefficient.double(), + torch.zeros(1, dtype=torch.int64), code, channel, COULOMB_CONSTANT, From c34d04f1bdb2afc22836b4a50d19e9a54bb9b756 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 15:52:15 +0800 Subject: [PATCH 19/30] fix(dpmodel): form the DPA4C angular moments by broadcasting DPA4C formed the non-scalar block of its moment payload with two index gathers, of the amplitude channels and of the harmonics, so that the backward of the payload scatter-adds through the index arrays it gathers with in the same loop. When the C++ backend of Inductor generates code for one thread, it declares such an index array after its first use, and the kernel does not compile ("'tmpN' was not declared in this scope"). The CPU freeze of a DPA4C model to .pt2 therefore failed with PyTorch 2.11 and 2.13 whenever Inductor ran with one thread, as on the hosted CI, with or without spin or a long-range module. Degree l of the payload is the outer product of its 2l + 1 harmonics with the leading C_l channels of the amplitude, which is now formed by broadcasting, harmonic-major as before; its backward reduces over the broadcast axes. The payload and the energies are unchanged bit for bit, and forces and virials agree to the rounding. The index tables and build_moment_indices are gone; the pt_expt descriptor drops the tables from a checkpoint written by DeePMD-kit 3.2 when it loads one. --- deepmd/dpmodel/descriptor/dpa4c.py | 42 +++++++++---------- .../dpmodel/descriptor/dpa4c_nn/__init__.py | 2 - .../dpmodel/descriptor/dpa4c_nn/geometry.py | 34 --------------- deepmd/pt_expt/descriptor/dpa4c.py | 14 +++++++ source/tests/pt_expt/descriptor/test_dpa4c.py | 16 +++++++ 5 files changed, 49 insertions(+), 59 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index 7dbd9c7f3d..b9a87db721 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -90,7 +90,6 @@ OrderedPairFiLM, SpinChannels, build_angular_basis, - build_moment_indices, canonicalize_charge_spin, degree_offsets, derive_bispectrum_ranks, @@ -461,12 +460,9 @@ def __init__( # === Step 4. Lay out the flat moment payload === # Degree zero owns the leading `channels` entries of the flat layout, - # so the non-scalar block is exactly its complement and needs no - # separate degree index. The spin families, when present, are appended - # after the geometric degrees, leaving every geometric offset intact. - channel_index, harmonic_index = build_moment_indices(self.degree_channels) - self.angular_channel_index = channel_index[self.channels :] - self.angular_harmonic_index = harmonic_index[self.channels :] + # and every higher degree follows in its own harmonic-major block. The + # spin families, when present, are appended after the geometric + # degrees, leaving every geometric offset intact. self.degree_offsets = degree_offsets(self.degree_channels) # === Step 5. Initialize the output calibration state === @@ -1146,30 +1142,30 @@ def aggregate_moments( ) xp = array_api_compat.array_namespace(amplitude) - device = array_api_compat.device(amplitude) - channel_index = xp_asarray_nodetach( - xp, - self.angular_channel_index, - device=device, - ) - harmonic_index = xp_asarray_nodetach( - xp, - self.angular_harmonic_index, - device=device, - ) + n_edge = amplitude.shape[0] # Payload layout: [chi^2, chi^4, degree zero (C_0), degrees one and # above (S - C_0)]. Degree zero is the whole edge amplitude under its - # single envelope, while the non-scalar block gathers an amplitude and - # a harmonic per flat moment coordinate and carries a second envelope. + # single envelope. Degree l is the outer product of its 2l + 1 + # harmonics with the leading C_l channels of the amplitude, laid out + # harmonic-major and carrying a second envelope. The products are + # formed by broadcasting, so their backward reduces over the + # broadcast axes. + angular = [ + xp.reshape( + basis[:, degree * degree : (degree + 1) * (degree + 1), None] + * amplitude[:, None, :width], + (n_edge, (2 * degree + 1) * width), + ) + for degree, width in enumerate(self.degree_channels) + if degree > 0 + ] envelope_squared = envelope * envelope parts = [ envelope_squared[:, None], (envelope_squared * envelope_squared)[:, None], amplitude, - xp.take(amplitude, channel_index, axis=1) - * xp.take(basis, harmonic_index, axis=1) - * envelope[:, None], + xp.concat(angular, axis=1) * envelope[:, None], ] if spin_payload is not None: parts.append(spin_payload) diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 460c7970e6..8bd43d4cc0 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -14,7 +14,6 @@ from .geometry import ( MAX_ANGULAR_DEGREE, build_angular_basis, - build_moment_indices, degree_offsets, derive_degree_channels, packed_l2_to_stf, @@ -41,7 +40,6 @@ "SpinChannels", "build_angular_basis", "build_bispectrum_layout", - "build_moment_indices", "canonicalize_charge_spin", "degree_offsets", "derive_bispectrum_ranks", diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py index 766a9d060c..a9378f4ea1 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py @@ -11,7 +11,6 @@ ) import array_api_compat -import numpy as np if TYPE_CHECKING: from collections.abc import ( @@ -236,36 +235,3 @@ def degree_offsets(degree_channels: Sequence[int]) -> tuple[int, ...]: for degree, width in enumerate(degree_channels): offsets.append(offsets[-1] + (2 * degree + 1) * int(width)) return tuple(offsets) - - -def build_moment_indices( - degree_channels: Sequence[int], -) -> tuple[np.ndarray, np.ndarray]: - """Build channel and harmonic indices for one flat edge payload. - - Every degree reads the leading channels of the shared radial map, so the - channel index of a degree is simply ``range(degree_channels[degree])``. - - Parameters - ---------- - degree_channels - Number of channels for degrees zero through ``lmax``. - - Returns - ------- - channel_index - Edge-amplitude channel indices with shape ``(S,)``. - harmonic_index - Packed harmonic indices with shape ``(S,)``, where - ``S = sum((2 * l + 1) * degree_channels[l])``. - """ - channel_index = [] - harmonic_index = [] - for degree, width in enumerate(degree_channels): - for component in range(2 * degree + 1): - channel_index.extend(range(int(width))) - harmonic_index.extend([degree * degree + component] * int(width)) - return ( - np.asarray(channel_index, dtype=np.int64), - np.asarray(harmonic_index, dtype=np.int64), - ) diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index a6482fbaa6..fb4495ab11 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -158,6 +158,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.use_amp_infer = use_amp_infer() self._apply_autocast_policy() + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written by DeePMD-kit 3.2 hold the gather + # tables of the angular moments, which are formed by broadcasting; + # the stale entries are dropped so that the strict load succeeds. + for name in ("angular_channel_index", "angular_harmonic_index"): + state_dict.pop(prefix + name, None) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def call_graph( self, graph: Any, diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index 4eef728456..eaa3b539a6 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -265,6 +265,22 @@ def test_serialization_preserves_parameters(self) -> None: original_parameters[name], ) + def test_state_dict_of_version_3_2_loads(self) -> None: + """A checkpoint holding the gather tables of DeePMD-kit 3.2 loads strictly.""" + state = self.descriptor.state_dict() + width = 32 + state["angular_channel_index"] = torch.arange(width) + state["angular_harmonic_index"] = torch.arange(width) + restored = self.build() + restored.load_state_dict(state) + torch.testing.assert_close( + self._evaluate(restored, self.coord), + self._evaluate(self.descriptor, self.coord), + ) + state["unknown"] = torch.zeros(1) + with pytest.raises(RuntimeError, match="unknown"): + restored.load_state_dict(state) + def test_deserialization_preserves_version_one_raw_radial_basis(self) -> None: data = self.descriptor.serialize() radial_basis = data["radial_basis"] From b55b8a01ef1f4d21a99b8755980269eb1e635c7a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 15:52:58 +0800 Subject: [PATCH 20/30] fix(pt_expt): address the review of the long-range module - LES near contact: LESKernel evaluates erf(x) / x from nine terms of its Taylor series in x^2 below x^2 = 0.0625, where the closed form loses its derivatives to cancellation: in float32 its radial derivative deviated from the 40-digit value by up to 37% at 1e-3 Angstrom, and its second derivative by more than three orders of magnitude. Together they keep the relative error of the radial derivative within 2e-14 in float64 and 1e-5 in float32. The CUDA direct sum does the same with six float32 terms below x^2 = 0.16, which brings the force of a small alpha at atomic distances to float32 precision (0.7% error in the slope at alpha r = 0.03 before) and needs no distance clamp anymore. A warp evaluates the series only for a batch with a pair below the threshold: a choice per pair cost 11-15% of the sum on a carbon cluster, this one about 3% at 100,000 atoms. - The communicator input of the compact lower holds the address of the caller's MPI_Comm, which the operator reads with the width of MPI_Comm in its MPI library, in place of eight bytes copied from the handle, which read past a four-byte MPICH handle. - A frame that one domain holds whole owns every node: the graph lower and the compact lower reject n_local < n_node without a communicator, whose halo nodes would otherwise carry no charge and drop their interactions. - linear_ener: a sub-model with a long-range module rejects its own pair_exclude_types, which the composition owns: it shapes the graph and the pair topology its sub-models share, and gives such a sub-model the semantics of the same model alone. Its atom_exclude_types is rejected as for the model alone. - bridging: long_range is routed to the learned child, beside which the composition rejects a bridging term. - The common charge bias of a constrained channel cancels in the constraint, so its gradient is exactly zero: the gradient test requires a nonzero gradient of every other parameter, and a new test checks the invariance. - The SOG documentation qualifies the force of the default kernel as that between fixed charges, whose nearly constant offset contributes through the response of the charges, and states its size and range. - The pt_expt latent-charge module names the classes of the charge network and the kernel it converts to, in place of imports kept only for their registration, which CodeQL reported as unused. --- deepmd/dpmodel/long_range/kernel.py | 45 +++- deepmd/dpmodel/model/model_factory.py | 23 +- .../kernels/dpa4c/canonical_long_range.py | 8 +- deepmd/pt_expt/kernels/long_range.py | 10 +- deepmd/pt_expt/long_range/latent_charge.py | 19 +- deepmd/pt_expt/model/ener_model.py | 4 +- deepmd/pt_expt/model/make_model.py | 46 ++++ deepmd/pt_expt/utils/serialization.py | 2 +- deepmd/utils/argcheck.py | 7 +- deepmd/utils/bridging.py | 3 + doc/model/long-range.md | 18 +- source/api_cc/src/DeepPotPTExpt.cc | 14 +- source/op/pt/long_range/direct_sum.cu | 104 +++++++-- source/op/pt/long_range/domain_exchange.cc | 19 +- source/op/pt/long_range/domain_exchange.h | 11 +- source/op/pt/long_range/ops.cc | 13 +- .../kernels/test_long_range_direct_sum.py | 39 +++- source/tests/pt_expt/model/test_long_range.py | 202 +++++++++++++++++- 18 files changed, 487 insertions(+), 100 deletions(-) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py index 90a75cbee5..3486b663ce 100644 --- a/deepmd/dpmodel/long_range/kernel.py +++ b/deepmd/dpmodel/long_range/kernel.py @@ -148,6 +148,25 @@ def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: return {} +# Taylor coefficients of erf(x) / x in powers of x^2, 2 / sqrt(pi) (-1)^n / +# (n! (2n + 1)), and the square of x below which LESKernel evaluates them: there +# the nine terms reach the float64 rounding of erf(x) / x. Above it the closed +# form keeps the relative error of its radial derivative within 2e-14 in +# float64 and 1e-5 in float32. +_ERF_OVER_X_SERIES = ( + 1.1283791670955126, + -0.37612638903183754, + 0.11283791670955126, + -0.026866170645131252, + 0.005223977625442188, + -0.0008548327023450852, + 0.00012055332981789664, + -1.492565035840625e-05, + 1.6462114365889246e-06, +) +_ERF_SERIES_X_SQ = 0.0625 + + @BaseInteractionKernel.register("les") class LESKernel(BaseInteractionKernel): r"""Latent Ewald summation kernel :math:`K(r) = \mathrm{erf}(\alpha r) / r`. @@ -183,9 +202,21 @@ def coefficients(self, like: Array) -> dict[str, Array]: def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: xp = array_api_compat.array_namespace(r_sq) - # The softening keeps the derivative finite on a coincident pair. - r = xp.sqrt(r_sq + 1e-12) - return xp_erf(coefficients["alpha"] * r) / r + alpha = coefficients["alpha"] + x_sq = alpha**2 * r_sq + # K = alpha erf(x) / x with x = alpha r. The derivative of the closed + # form subtracts two terms that agree to second order in x, so near + # contact erf(x) / x comes from its Taylor series in x^2. Both are + # evaluated on every pair, which keeps the shapes static; each reads a + # value of its own domain where the other is selected, which keeps the + # derivatives of both finite. + near = x_sq < _ERF_SERIES_X_SQ + x_far = xp.sqrt(xp.where(near, xp.ones_like(x_sq), x_sq)) + x_sq_near = xp.where(near, x_sq, xp.zeros_like(x_sq)) + series = xp.full_like(x_sq, _ERF_OVER_X_SERIES[-1]) + for coefficient in reversed(_ERF_OVER_X_SERIES[:-1]): + series = series * x_sq_near + coefficient + return alpha * xp.where(near, series, xp_erf(x_far) / x_far) @BaseInteractionKernel.register("sog") @@ -201,9 +232,11 @@ class SOGKernel(BaseInteractionKernel): :math:`a_k = \sqrt{2/\pi} \ln b / \sigma_k`, the trapezoidal discretization in :math:`\ln s` of :math:`1/r = \sqrt{2/\pi} \int_0^\infty e^{-r^2/2s^2} s^{-2} \, ds`. - The default kernel then follows :math:`1/r`, less a constant that - exerts no force, from a few :math:`\sigma` to a fraction of - :math:`\sigma b^{M-1}`, and stays finite at contact. + The default kernel then follows :math:`1/r`, less a nearly constant + offset, from a few :math:`\sigma` to a fraction of + :math:`\sigma b^{M-1}`, and stays finite at contact. The offset + exerts no force between fixed charges; through charges that respond + to the positions it does. width The widths :math:`\sigma_k > 0`, in Å. Defaults to the geometric sequence :math:`\sigma_k = \sigma b^k`, :math:`k = 0, \dots, M - 1`. diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index da96ca32b9..118b09a75f 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -282,8 +282,11 @@ def get_linear_atomic_model( If more than one ``inner_potential`` child is given, if an ``inner_potential`` child has no unique learned sibling, if a bridged composition does not use ``weights: "sum"``, if a child - carries a ``bridging_method`` flag, or if a child is of an + carries a ``bridging_method`` flag, if a child with a long-range + module carries its own ``pair_exclude_types``, or if a child is of an unsupported kind. + NotImplementedError + If a child with a long-range module carries ``atom_exclude_types``. """ from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, @@ -405,6 +408,24 @@ def get_linear_atomic_model( if descriptor_child_builder is not None: child = descriptor_child_builder(sub) if child is None: + # A long-range child is checked as the same model alone: the + # pair exclusion of the composition shapes the graph and the + # pair topology the children share, and atom exclusion is + # unsupported with a long-range module. + if sub.get("long_range") is not None: + if sub.get("pair_exclude_types"): + raise ValueError( + "a linear_ener sub-model with a long-range module " + "takes `pair_exclude_types` from the composition, " + "whose sub-models share one neighbor graph and one " + "pair topology; set it on the composition" + ) + if sub.get("atom_exclude_types"): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a " + "long-range module, which attributes each pair " + "energy to both atoms of the pair" + ) descriptor, fitting, _ = get_model_components( sub, descriptor_base=descriptor_base, diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py index f78b5c2db1..69e5edb332 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -56,10 +56,10 @@ def dpa4c_canonical_long_range_energy_force( Node positions with shape ``(N, 3)`` in Å; every frame is a complete non-periodic system. communicator - The bytes of the MPI_Comm handle of a frame decomposed into domains, - one per rank, as an int64 CPU tensor of one element, or zero or - ``None`` for a single domain. Across domains the owned nodes of every - rank make up the frame. + The address of the MPI_Comm handle of a frame decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero + or ``None`` for a single domain. Across domains the owned nodes of + every rank make up the frame. ownership Boolean mask selecting energy-contributing nodes with shape ``(N,)``. atom_bias diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py index 302320b6fc..994b78134f 100644 --- a/deepmd/pt_expt/kernels/long_range.py +++ b/deepmd/pt_expt/kernels/long_range.py @@ -117,8 +117,8 @@ def _kernel_and_slope( coefficient = coefficient.to(dtype=r_sq.dtype, device=r_sq.device) if kernel == _LES: alpha = coefficient[0] - # Coincident nodes take the limit at the shortest distance the CUDA - # kernel evaluates. + # Coincident nodes, which carry no weight, are evaluated at a short + # distance, where the closed form stays finite. r_sq = torch.clamp(r_sq, min=1e-12 / alpha**2) r = torch.sqrt(r_sq) value = torch.erf(alpha * r) / r @@ -272,9 +272,9 @@ def latent_charge_direct_sum( prefactor Factor of every pair term, in eV Å. communicator - The bytes of the MPI_Comm handle of a system decomposed into domains, - one per rank, as an int64 CPU tensor of one element, or zero, the - default, for a single domain. Across domains the call takes one + The address of the MPI_Comm handle of a system decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero, + the default, for a single domain. Across domains the call takes one frame, whose charged nodes are those of every rank; the outputs are those of the charged nodes of this rank, the other nodes carry no charge, and ``virial`` is the share of this rank. diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py index ef207f44d3..2cf9bf009f 100644 --- a/deepmd/pt_expt/long_range/latent_charge.py +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -10,16 +10,17 @@ from deepmd.pt_expt.common import ( torch_module, ) +from deepmd.pt_expt.fitting.latent_charge import ( + LatentChargeFitting, +) -# The charge network and the kernels built by the dpmodel class are converted -# to their pt_expt modules on assignment; importing them registers the -# conversions. -from deepmd.pt_expt.fitting import latent_charge as _latent_charge_fitting # noqa: F401 - -from . import kernel as _kernel # noqa: F401 from .base_long_range import ( BaseLongRange, ) +from .kernel import ( + LESKernel, + SOGKernel, +) @BaseLongRange.register("latent_charge") @@ -27,6 +28,12 @@ class LatentChargeLongRange(LatentChargeLongRangeDP): """Energy of interacting latent charges in complete non-periodic systems (pt_expt).""" + # The dpmodel class builds the charge network and the kernel, which the + # assignment converts to these pt_expt modules; importing their classes + # registers the conversions. + charge_net: LatentChargeFitting + kernel: LESKernel | SOGKernel + def _pair_kernel( self, r_sq: torch.Tensor, src: torch.Tensor, n: int ) -> torch.Tensor: diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index ff9b3e77b7..0523084381 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -27,6 +27,7 @@ from .make_model import ( _translate_energy_keys, + check_owned_frames, fused_atom_bias, make_model, ) @@ -175,7 +176,7 @@ def forward_lower_canonical_graph( with a long-range module, whose frames are complete non-periodic systems; ``None`` otherwise. communicator - For a model with a long-range module, the bytes of the MPI_Comm + For a model with a long-range module, the address of the MPI_Comm handle of a frame decomposed into domains, one per rank, as an int64 CPU tensor of one element, whose owned nodes are then those of one domain; zero or ``None`` for a complete frame. @@ -234,6 +235,7 @@ def forward_lower_canonical_graph( raise ValueError( "a model with a long-range module reads the node positions" ) + check_owned_frames(n_node, n_local, communicator) force_mag = None energy, atom_energy, force, virial, atom_virial = ( dpa4c_canonical_long_range_energy_force( diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index b2d1e2c4ef..9b132b2a53 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -196,6 +196,51 @@ def _fused_energy_force_graph( return ret +def check_owned_frames( + n_node: torch.Tensor, + n_local: torch.Tensor | None, + communicator: torch.Tensor | None = None, +) -> None: + """Reject halo nodes in the frames of a single domain with a long-range module. + + The module sums over every pair of atoms of a frame, so a frame that one + domain holds whole owns all its nodes; halo nodes would carry no charge + and remove their interactions from the sum. Across the domains of a + communicator, every domain owns its share of the frame instead. The + counts are read eagerly only: a traced lower receives them from callers + that check them. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + n_local + Per-frame owned node counts with shape ``(nf,)``, or ``None``. + communicator + The communicator input of the compact lower, a CPU tensor holding the + address of an MPI_Comm handle or zero, or ``None`` for a single domain. + + Raises + ------ + ValueError + If a frame of a single domain carries halo nodes. + """ + from torch._subclasses.fake_tensor import ( + is_fake, + ) + + if n_local is None or n_local is n_node or is_fake(n_local): + return + if communicator is not None and int(communicator.reshape(-1)[0]) != 0: + return + if not torch.equal(n_local, n_node): + raise ValueError( + "a model with a long-range module evaluates complete frames on a " + "single domain, whose graph owns every node (n_local == n_node), " + f"but got n_local {n_local.tolist()} for n_node {n_node.tolist()}" + ) + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -742,6 +787,7 @@ def forward_common_lower_graph( "the long-range topologies describe complete systems; " "a domain-decomposed system is not supported" ) + check_owned_frames(n_node, n_local) # Every topology carries a geometry leaf of its own. topologies = { name: PairGraph( diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index ebba8ca4a3..154423aa0e 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -818,7 +818,7 @@ def _build_canonical_graph_dynamic_shapes( The trailing slots, the moments of a native-spin model or the positions and the communicator of a model with a long-range module, are present only for such models, so the sample length selects the shape tuple. The - communicator holds one handle and is static. + communicator holds one address and is static. """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a99c8d60b4..b8a729b908 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3576,9 +3576,10 @@ def interaction_kernel_sog() -> list[Argument]: "shared by every Gaussian or one value per Gaussian. Required with " "`width`; with the default widths it defaults to " "sqrt(2 / pi) ln(`width_ratio`) / sigma_k, the weights of a discretized " - "1 / r, so that the default kernel follows 1 / r, less a constant that " - "exerts no force, from a few `base_width` to a fraction of the widest " - "Gaussian." + "1 / r, so that the default kernel follows 1 / r, less a nearly " + "constant offset, from a few `base_width` to a fraction of the widest " + "Gaussian. The offset exerts no force between fixed charges; through " + "charges that respond to the positions it does." ) doc_width = ( "The widths sigma_k of the Gaussians, in Å, each positive. Defaults to " diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index b5bdb2ea94..2fa225cc64 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -142,6 +142,9 @@ def _is_dpa4_family(sub: dict) -> bool: _LEARNED_CHILD_KEYS = ( "descriptor", "fitting_net", + # A long-range module belongs to the learned energy; the composition then + # rejects it beside the bridging term. + "long_range", "model_branch_alias", "info", "use_compile", diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 066bdef194..95e198caba 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -92,9 +92,15 @@ training. which explicit widths require. The default amplitudes, $a_k = \sqrt{2/\pi}\,\ln b/\sigma_k$ for the ratio $b$ of consecutive widths, discretize $1/r$: with the default widths the kernel stays finite, - $0.46$ Å⁻¹, at contact, and its force stays within 0.4% of that of $1/r$ - from 6.5 Å to 150 Å. It differs from $1/r$ there by a constant, which - exerts no force. + $0.46$ Å⁻¹, at contact, and its radial force between fixed charges stays + within 0.4% of that of $1/r$ from 6.5 Å to 150 Å. It differs from $1/r$ + there by a nearly constant offset $C \approx 1.3 \times 10^{-3}$ Å⁻¹, which + exerts no force between fixed charges. The charges respond to the positions, + however, and through them the offset adds + $-k_e C \sum_c (Q_c^2 - \sum_i (q_i^c)^2)/2$ to the energy of a frame, with + $Q_c$ the total charge of channel $c$, which the response of the charges + turns into a force. Below about 5 Å the kernel falls short of $1/r$ by more + than $C$, by 0.12 Å⁻¹ at 2 Å. Several charge channels, set by {ref}`channels `, let @@ -115,8 +121,10 @@ constrained. Pairs of the types in `pair_exclude_types` see each other neither in the descriptor nor in the long-range interaction, while the `exclude_types` of -the descriptor act on the descriptor alone. `atom_exclude_types` is not -supported with a long-range module. +the descriptor act on the descriptor alone. In a `linear_ener` composition, +`pair_exclude_types` belongs to the composition, whose sub-models share its +neighbor graph and pair topology. `atom_exclude_types` is not supported with +a long-range module. ## Training and inference diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 3ac1f112c7..a76cd96fa5 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -112,15 +112,13 @@ FrameParameterLayout resolve_frame_parameter_layout(const char* name, } // The communicator input of the compact lower of a long-range model: an int64 -// CPU tensor holding the bytes of the ``MPI_Comm`` that ``world`` points to, -// read as for the comm tensors of a message-passing model, or zero for a -// system that one rank holds whole. +// CPU tensor holding the address of the ``MPI_Comm`` that ``world`` points to, +// or zero for a system that one rank holds whole. The handle itself is read +// only by the operator, with the width of ``MPI_Comm`` in its MPI library. at::Tensor communicator_tensor(const void* world) { - const auto options = torch::TensorOptions().dtype(torch::kInt64); - if (world == nullptr) { - return torch::zeros({1}, options); - } - return torch::from_blob(const_cast(world), {1}, options).clone(); + return torch::full( + {1}, static_cast(reinterpret_cast(world)), + torch::TensorOptions().dtype(torch::kInt64)); } } // namespace diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu index a44667b002..c12e144537 100644 --- a/source/op/pt/long_range/direct_sum.cu +++ b/source/op/pt/long_range/direct_sum.cu @@ -61,10 +61,6 @@ constexpr int kThreads = 128; constexpr int kBatch = 4; constexpr int kMaxChannels = 8; constexpr int kMaxGaussians = 32; -// Squared distances are evaluated no shorter than this, in the units of the -// staged positions, which keeps the kernels finite for a node paired with -// itself and for coincident padding nodes; both carry a zero weight. -constexpr float kMinSquaredDistance = 1e-12f; // Fills of the resident block slots of the device that one launch aims for. constexpr int kWaves = 8; // Shortest chunk of the partner axis, in tiles. @@ -313,7 +309,22 @@ __global__ void prepare_nodes(const double* __restrict__ coord, // erf(x) / x and its slope (2 / sqrt(pi) exp(-x^2) - erf(x) / x) / x^2, with // erf(x) = 1 - t P(t) exp(-x^2), t = 1 / (1 + p x), and the coefficients of // Abramowitz and Stegun 7.1.26 (absolute error below 1.5e-7), so that one -// exponential serves the kernel and its slope. +// exponential serves the kernel and its slope. The two terms of the slope +// agree to second order in x, which amplifies the error of the approximation +// by 1 / x^2; below x^2 = kSeriesSquared the kernel and its slope therefore +// come from their Taylor series in x^2, +// +// erf(x) / x = 2 / sqrt(pi) sum_n (-1)^n x^(2n) / (n! (2n + 1)), +// slope = 2 / sqrt(pi) sum_{n >= 1} (-1)^n 2n x^(2n - 2) / +// (n! (2n + 1)), +// +// whose first six terms reach float32 precision there. The relative error of +// the kernel stays below 1e-6 for every x, that of its slope below 7e-6, the +// largest just above the threshold. A warp evaluates the series only for a +// batch with a pair below the threshold: a batch that pairs a node with +// itself, or a contact shorter than 0.4 / alpha, which for an alpha of 0.5 per +// angstrom or more is shorter than every bond but that of H2. Every other +// batch evaluates the closed form without a branch. // // Beyond x^2 = kFarSquared the float32 kernel equals 1 / x and its slope // -1 / x^3 bit for bit: exp(-x^2) and erfc(x) fall below half a unit in the @@ -331,6 +342,20 @@ __global__ void prepare_nodes(const double* __restrict__ coord, // sum_k s_k 2^(-c_k r^2), s_k = -2 ln(2) c_k a_k, from the staged // coefficients ``{a_k, -c_k, s_k}``. constexpr float kFarSquared = 20.25f; +constexpr float kSeriesSquared = 0.16f; + +// A series in x^2 with the coefficients ``c``, in increasing powers, by +// Horner's rule. +template +__device__ __forceinline__ float even_series(const float (&c)[kTerms], + float x2) { + float sum = c[kTerms - 1]; +#pragma unroll + for (int n = kTerms - 2; n >= 0; --n) { + sum = fmaf(sum, x2, c[n]); + } + return sum; +} // Whether every pair of the batch of every thread of the warp lies in the far // field of the erf kernel, judged from the head displacements, whose squared @@ -353,6 +378,30 @@ __device__ __forceinline__ bool far_batch(const float (&dx)[kBatch], } } +// erf(x) / x and its slope by the approximation of Abramowitz and Stegun at +// the squared distances ``r2``. +__device__ __forceinline__ void les_closed_form(const float (&r2)[kBatch], + float (&value)[kBatch], + float (&slope)[kBatch]) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + const float x = r2[b] * inverse; + const float t = __fdividef(1.f, fmaf(0.3275911f, x, 1.f)); + const float poly = + t * + fmaf( + t, + fmaf(t, fmaf(t, fmaf(t, 1.061405429f, -1.453152027f), 1.421413741f), + -0.284496736f), + 0.254829592f); + const float gaussian = exp2f(-1.44269504088896341f * r2[b]); + value[b] = fmaf(-poly, gaussian, 1.f) * inverse; + slope[b] = + fmaf(1.12837916709551257f, gaussian, -value[b]) * (inverse * inverse); + } +} + template __device__ __forceinline__ void evaluate_kernel(const float (&r2)[kBatch], const float4* gaussians, @@ -360,22 +409,34 @@ __device__ __forceinline__ void evaluate_kernel(const float (&r2)[kBatch], float (&value)[kBatch], float (&slope)[kBatch]) { if constexpr (kKernel == kLes) { + float nearest = r2[0]; +#pragma unroll + for (int b = 1; b < kBatch; ++b) { + nearest = fminf(nearest, r2[b]); + } + if (!__any_sync(0xffffffffu, nearest < kSeriesSquared)) { + les_closed_form(r2, value, slope); + return; + } + constexpr float kValueSeries[6] = {1.128379167f, -0.376126389f, + 0.1128379167f, -0.02686617065f, + 0.005223977625f, -0.0008548327023f}; + constexpr float kSlopeSeries[6] = {-0.7522527781f, 0.4513516668f, + -0.1611970239f, 0.041791821f, + -0.008548327023f, 0.001446639958f}; + // The closed form reads the threshold where the series is selected. + float clamped[kBatch]; #pragma unroll for (int b = 0; b < kBatch; ++b) { - const float inverse = rsqrtf(r2[b]); - const float x = r2[b] * inverse; - const float t = __fdividef(1.f, fmaf(0.3275911f, x, 1.f)); - const float poly = - t * - fmaf(t, - fmaf(t, - fmaf(t, fmaf(t, 1.061405429f, -1.453152027f), 1.421413741f), - -0.284496736f), - 0.254829592f); - const float gaussian = exp2f(-1.44269504088896341f * r2[b]); - value[b] = fmaf(-poly, gaussian, 1.f) * inverse; - slope[b] = - fmaf(1.12837916709551257f, gaussian, -value[b]) * (inverse * inverse); + clamped[b] = fmaxf(r2[b], kSeriesSquared); + } + les_closed_form(clamped, value, slope); +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + if (r2[b] < kSeriesSquared) { + value[b] = even_series(kValueSeries, r2[b]); + slope[b] = even_series(kSlopeSeries, r2[b]); + } } } else { #pragma unroll @@ -457,8 +518,9 @@ __device__ __forceinline__ void accumulate_tile( dx[b] += own.tail.x - other_tail.x; dy[b] += own.tail.y - other_tail.y; dz[b] += own.tail.z - other_tail.z; - r2[b] = fmaxf(fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])), - kMinSquaredDistance); + // Both kernels are finite at r2 = 0, which a node paired with itself + // and coincident padding nodes reach with a zero weight. + r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])); } evaluate_kernel(r2, gaussians, gaussian_count, value, slope); } diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc index 71acbc6957..24c393c5ee 100644 --- a/source/op/pt/long_range/domain_exchange.cc +++ b/source/op/pt/long_range/domain_exchange.cc @@ -19,7 +19,6 @@ #include #include -#include #include #include @@ -27,7 +26,7 @@ namespace deepmd { namespace long_range { namespace { -std::int64_t handle_of(const torch::Tensor& communicator) { +std::int64_t address_of(const torch::Tensor& communicator) { TORCH_CHECK(communicator.device().is_cpu() && communicator.scalar_type() == torch::kInt64 && communicator.numel() == 1, @@ -37,15 +36,13 @@ std::int64_t handle_of(const torch::Tensor& communicator) { } #ifdef USE_MPI -// The tensor holds the bytes of the handle from its first byte on, whatever -// the width of MPI_Comm in the MPI library. +// The handle of the caller, of the width of MPI_Comm in this MPI library: an +// int in MPICH, a pointer in Open MPI. MPI_Comm communicator_of(const torch::Tensor& communicator) { - const torch::Tensor bytes = communicator.contiguous(); - MPI_Comm comm; - static_assert(sizeof(MPI_Comm) <= sizeof(std::int64_t), - "an MPI_Comm handle must fit in an int64 tensor"); - std::memcpy(&comm, bytes.data_ptr(), sizeof(MPI_Comm)); - return comm; + static_assert(sizeof(std::uintptr_t) <= sizeof(std::int64_t), + "an address must fit in an int64 tensor"); + return *reinterpret_cast( + static_cast(address_of(communicator))); } bool reads_device_memory() { @@ -60,7 +57,7 @@ bool reads_device_memory() { } // namespace Domain domain_of(const torch::Tensor& communicator) { - if (handle_of(communicator) == 0) { + if (address_of(communicator) == 0) { return {0, 1}; } #ifdef USE_MPI diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h index 323e60d92d..c091589b31 100644 --- a/source/op/pt/long_range/domain_exchange.h +++ b/source/op/pt/long_range/domain_exchange.h @@ -4,11 +4,10 @@ // system. // // A domain is one MPI rank of the communicator that a caller hands to -// deepmd::latent_charge_direct_sum as an int64 tensor holding the bytes of its -// MPI_Comm handle, the convention of deepmd_export::border_op. A zero handle -// stands for a single domain. Both exchanges follow the rank order, so every -// rank receives the same rows and the same sums, independent of the -// reduction algorithm of the MPI library. +// deepmd::latent_charge_direct_sum as an int64 tensor holding the address of +// its MPI_Comm handle. A zero address stands for a single domain. Both +// exchanges follow the rank order, so every rank receives the same rows and the +// same sums, independent of the reduction algorithm of the MPI library. #pragma once @@ -26,7 +25,7 @@ struct Domain { }; // The rank of this process and the number of ranks of ``communicator``; one -// rank for a zero handle. +// rank for a zero address. Domain domain_of(const torch::Tensor& communicator); // Concatenates, in rank order, the rows of every rank. ``rows`` is a diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc index d1bf4b4aa6..dd312bc555 100644 --- a/source/op/pt/long_range/ops.cc +++ b/source/op/pt/long_range/ops.cc @@ -44,13 +44,12 @@ // virial hold the charges fixed; the response of the charges is the // caller's, through the charge gradient. // -// `communicator` is an int64 CPU tensor of one element holding the bytes of -// an MPI_Comm handle, as for deepmd_export::border_op, or zero. A handle of -// several ranks decomposes one frame into domains, one per rank: the charged -// nodes of every rank are the charged nodes of the frame, and the sums run -// over all of them. The outputs of a rank are those of its charged nodes, and -// W_f is their share of the virial of the frame; the other nodes of the rank -// carry no charge. +// `communicator` is an int64 CPU tensor of one element holding the address of +// the caller's MPI_Comm handle, or zero. A handle of several ranks decomposes +// one frame into domains, one per rank: the charged nodes of every rank are +// the charged nodes of the frame, and the sums run over all of them. The +// outputs of a rank are those of its charged nodes, and W_f is their share of +// the virial of the frame; the other nodes of the rank carry no charge. #include diff --git a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py index 475562488c..b5c775e8cc 100644 --- a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py +++ b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py @@ -15,6 +15,10 @@ annotations, ) +from typing import ( + Any, +) + import numpy as np import pytest import torch @@ -72,10 +76,9 @@ def _system( def _evaluate( - system: tuple[torch.Tensor, ...], kernel_name: str, channel: int + system: tuple[torch.Tensor, ...], kernel: Any, channel: int ) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: coord, charge, node_mask, n_node, total_charge = system - kernel = KERNELS[kernel_name]() actual = latent_charge_direct_sum( coord, charge, @@ -120,7 +123,33 @@ def test_matches_float64_pair_sum( ) -> None: """Every output agrees with the float64 pair sum.""" system = _system(n_node, channels, padding, 7, 1.0e3) - actual, expected = _evaluate(system, kernel_name, channel) + actual, expected = _evaluate(system, KERNELS[kernel_name](), channel) + _assert_close(actual, expected) + # Padding nodes join no frame. + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("alpha", [1e-3, 0.03, 0.3]) +@pytest.mark.parametrize("channel", [0, -1]) +def test_les_matches_float64_pair_sum_at_small_alpha_r( + alpha: float, channel: int +) -> None: + """The erf kernel keeps its precision where alpha r is small. + + The two terms of the slope of erf(x) / x agree to second order in x, so + that the slope of an approximate erf loses precision in proportion to + 1 / x^2 at atomic distances and a small alpha. + """ + system = _system([96], 2, 0, 5, 0.0) + _assert_close(*_evaluate(system, LESKernel(alpha=alpha), channel)) + + +def _assert_close( + actual: tuple[torch.Tensor, ...], expected: tuple[torch.Tensor, ...] +) -> None: + """Every output agrees with the float64 pair sum, relative to its largest entry.""" for name, value, reference in zip(NAMES, actual, expected, strict=True): assert value.dtype == ( torch.float64 if name == "atom_energy" else torch.float32 @@ -128,10 +157,6 @@ def test_matches_float64_pair_sum( scale = reference.abs().max().item() error = (value.double() - reference.double()).abs().max().item() assert error <= 2e-5 * scale, (name, error, scale) - # Padding nodes join no frame. - if padding: - for value in actual[:4]: - assert torch.all(value[-padding:] == 0) @pytest.mark.parametrize("kernel_name", ["les", "sog"]) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 02f5ed6587..606e078245 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -323,6 +323,153 @@ def test_pair_exclusion_removes_the_interaction(constraint: dict | None) -> None assert out["force"][0, 2].abs().max() == 0 +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_les_kernel_derivatives_down_to_contact(dtype: torch.dtype) -> None: + """The radial derivatives of LES keep the precision of the dtype near contact. + + The derivatives of erf(alpha r) / r subtract terms that agree to second + order in alpha r, which the series of the kernel avoids; the reference + differentiates in 40-digit arithmetic. + """ + import mpmath + + from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + ) + + alpha = 0.8 + kernel = LESKernel(alpha=alpha) + r = torch.tensor( + np.geomspace(1e-5, 6.0, 80), dtype=dtype, device=env.DEVICE + ).requires_grad_(True) + (first,) = torch.autograd.grad(kernel(r * r).sum(), r, create_graph=True) + (second,) = torch.autograd.grad(first.sum(), r) + + def les(t: mpmath.mpf) -> mpmath.mpf: + return mpmath.erf(alpha * t) / t + + rtol = 2e-5 if dtype == torch.float32 else 1e-11 + for order, actual in ((1, first), (2, second)): + with mpmath.workdps(40): + expected = np.array( + [ + float(mpmath.diff(les, mpmath.mpf(float(t)), order)) + for t in r.detach().cpu() + ] + ) + # The second derivative changes sign near 1.2 Å, where it is the + # difference of terms much larger than itself. + np.testing.assert_allclose( + actual.detach().cpu().numpy(), + expected, + rtol=rtol, + atol=rtol * np.abs(expected).max(), + err_msg=str(order), + ) + + +def test_owned_frames_across_domains() -> None: + """Halo nodes are rejected on a single domain and kept across domains.""" + from deepmd.pt_expt.model.make_model import ( + check_owned_frames, + ) + + n_node = torch.tensor([5, 3], dtype=torch.int64) + n_local = torch.tensor([4, 3], dtype=torch.int64) + check_owned_frames(n_node, n_node.clone()) + check_owned_frames(n_node, n_local, torch.ones(1, dtype=torch.int64)) + for communicator in (None, torch.zeros(1, dtype=torch.int64)): + with pytest.raises(ValueError, match="n_local == n_node"): + check_owned_frames(n_node, n_local, communicator) + + +def test_graph_lower_rejects_halo_nodes() -> None: + """A graph whose frames carry halo nodes is rejected, not summed without them. + + The halo nodes of a graph without communication would carry no charge and + remove their interactions from the sum of a frame held whole. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + from deepmd.pt_expt.utils.graph_builder import ( + build_neighbor_graph_for_method, + ) + + model = _model("les", constraint=None) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + graph = build_neighbor_graph_for_method( + "dense", coord, atype, None, model.get_rcut(), None + ) + topologies = build_topologies( + model.atomic_model.required_topologies(), + coord.reshape(-1, 3), + graph.n_node, + atype=atype.reshape(-1), + ) + inputs = ( + atype.reshape(-1), + graph.n_node, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + ) + + def lower(n_local: torch.Tensor) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + inputs[0], inputs[1], n_local, *inputs[2:], topologies=topologies + ) + + lower(graph.n_node.clone()) + with pytest.raises(ValueError, match="n_local == n_node"): + lower(graph.n_node - 1) + + +def test_linear_composition_takes_the_exclusions_of_the_composition() -> None: + """A long-range child follows the pair exclusion of its composition. + + The children of a composition share its neighbor graph and pair topology, + so the exclusions of the composition give the child the semantics of the + same model alone, and an exclusion of the child itself is rejected. + """ + child = _config("les", constraint=None) + child["type"] = "standard" + composition = { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": "sum", + "models": [child], + "pair_exclude_types": [[0, 1]], + } + model = get_model(copy.deepcopy(composition)).to(env.DEVICE).eval() + with torch.no_grad(): + for name, parameter in model.atomic_model.named_parameters(): + if ".kernel." not in f".{name}": + parameter.zero_() + bias = model.atomic_model.models[0].long_range.charge_net.bias_atom_e + bias[:, 0] = torch.as_tensor([1.0, -1.0], dtype=bias.dtype) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + # As in test_pair_exclusion_removes_the_interaction: only the two type-0 + # atoms interact. + erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item() + np.testing.assert_allclose( + model(coord, atype)["energy"].item(), + COULOMB_CONSTANT * erf / 4.0, + rtol=1e-12, + ) + for key, value, error, match in ( + ("pair_exclude_types", [[0, 1]], ValueError, "from the composition"), + ("atom_exclude_types", [1], NotImplementedError, "atom_exclude_types"), + ): + config = copy.deepcopy(composition) + config.pop("pair_exclude_types") + config["models"][0][key] = value + with pytest.raises(error, match=match): + get_model(config) + + @pytest.mark.parametrize("compiled", [False, True]) def test_vacuum_reference_with_long_range(compiled: bool) -> None: """An isolated atom keeps its reference energy; reference rows join no pair.""" @@ -447,15 +594,29 @@ def test_linear_composition_combines_complete_energies(weights: object) -> None: @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_parameter_gradients_and_freezing(kernel: str) -> None: - """Energy and force losses reach the charge network and the kernel; frozen parts stay out.""" - model = _model(kernel).train() + """Energy and force losses reach the charge network and the kernel; frozen parts stay out. + + Every parameter of the module moves the losses, except under a constraint + the common charge bias of the constrained channel, whose gradient + vanishes (see test_constraint_removes_the_common_charge_bias). + """ coord, atype = _batch() - out = model(coord, atype) - loss = out["energy"].square().sum() + out["force"].square().sum() - loss.backward() - long_range = model.atomic_model.long_range - for name, parameter in long_range.named_parameters(): - assert parameter.grad is not None and parameter.grad.abs().max() > 0, name + for constraint in (None, NEUTRAL): + model = _model(kernel, constraint=constraint).train() + out = model(coord, atype) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + long_range = model.atomic_model.long_range + bias = long_range.charge_net.nets[0].layers[-1].b + for name, parameter in long_range.named_parameters(): + assert parameter.grad is not None, name + assert torch.isfinite(parameter.grad).all(), name + gradient = parameter.grad + if constraint is not None and parameter is bias: + # Channel 0 is the constrained channel. + gradient = gradient[..., 1:] + if gradient.numel(): + assert gradient.abs().max() > 0, name config = _config(kernel) config["long_range"]["kernel"]["trainable"] = False @@ -470,6 +631,31 @@ def test_parameter_gradients_and_freezing(kernel: str) -> None: frozen(coord, atype)["energy"].sum().backward() +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_constraint_removes_the_common_charge_bias(kernel: str) -> None: + """A uniform shift of the raw charges of the constrained channel changes nothing. + + The constraint subtracts the mean shift of every frame, so the output bias + of that channel leaves energy and force unchanged, while that of another + channel does not. + """ + model = _model(kernel) + coord, atype = _batch() + reference = model(coord, atype) + bias = model.atomic_model.long_range.charge_net.nets[0].layers[-1].b + for channel in range(bias.shape[-1]): + with torch.no_grad(): + bias[..., channel] += 0.3 + out = model(coord, atype) + with torch.no_grad(): + bias[..., channel] -= 0.3 + if channel == 0: + torch.testing.assert_close(out["energy"], reference["energy"]) + torch.testing.assert_close(out["force"], reference["force"]) + else: + assert (out["energy"] - reference["energy"]).abs().max() > 1e-6 + + @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_serialization_round_trip(kernel: str) -> None: """The portable record and the state dict restore the same model.""" From f7ab0d8221c914293c40c94eb96981351acfb7f3 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Tue, 20 Oct 2026 07:33:08 +0000 Subject: [PATCH 21/30] feat(long-range): evaluate latent-charge modules on periodic cells in reciprocal space A long-range module with "periodic": true evaluates a frame given a cell as a reciprocal-space lattice sum of the Fourier-transformed kernel, - les: 4 pi exp(-k^2/4 alpha^2) / k^2, - sog: sum_m a_m (2 pi)^{3/2} sigma_m^3 exp(-sigma_m^2 k^2 / 2), with the self term subtracted analytically and the k=0 mode dropped (tin-foil boundary); every latent channel must be neutral. The k grid is enumerated eagerly per batch at the kernel's wave-number cutoff and enters the compiled region as a "k_grid" topology of dynamic length beside "all_pairs", which a periodic batch leaves guard-only. The module keeps serving non-periodic batches. The geometry reaches the graph through two leaves flattened with the grid: the node positions, whose gradient is the force, and a zero per-frame strain, whose gradient is the negative k-space virial. Both re-leaf inside forward_common_lower_graph so the backward roots at graph inputs; a leaf created inside a make_fx trace produces a NaN gradient under inductor. The SOG transform unrolls its modes: an (nk, M) intermediate breaks inductor's double backward of the k-space virial. Freezing/export, compression, the C++ interface and LAMMPS reject a periodic long-range module until they can serve the k grid. Validated against an independent Ewald-split reference and per-mode Poisson checks (verify_periodic_lr.py), by force/virial finite differences and by compiled-vs-eager parity (tests in source/tests/.../test_long_range.py). --- AGENTS.md | 1 + .../dpmodel/atomic_model/base_atomic_model.py | 36 ++- .../dpmodel/atomic_model/dp_atomic_model.py | 29 ++ .../atomic_model/linear_atomic_model.py | 31 ++ deepmd/dpmodel/long_range/kernel.py | 97 ++++++ deepmd/dpmodel/long_range/latent_charge.py | 183 ++++++++++- deepmd/dpmodel/model/make_model.py | 39 ++- .../dpmodel/utils/neighbor_graph/__init__.py | 8 + .../utils/neighbor_graph/pair_graph.py | 281 +++++++++++++++-- deepmd/pt_expt/infer/deep_eval.py | 33 +- deepmd/pt_expt/model/edge_transform_output.py | 63 +++- deepmd/pt_expt/model/make_model.py | 48 ++- deepmd/pt_expt/train/training.py | 38 ++- deepmd/pt_expt/utils/serialization.py | 70 ++++- deepmd/utils/argcheck.py | 30 +- doc/model/long-range.md | 59 +++- periodic_long_range_design.md | 296 ++++++++++++++++++ .../tests/common/dpmodel/test_long_range.py | 118 +++++++ source/tests/pt_expt/model/test_long_range.py | 251 +++++++++++++++ 19 files changed, 1598 insertions(+), 113 deletions(-) create mode 100644 periodic_long_range_design.md diff --git a/AGENTS.md b/AGENTS.md index 617c904f11..b736c6159e 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 +- **Long-range latent-charge module**: a standard energy model (DPA4C descriptor) takes a `long_range` section (`type: latent_charge`, kernels `les`/`sog`, optional `charge_constraint`) that adds an interacting-charge energy evaluated on complete systems — directly over all pairs for non-periodic frames, and in reciprocal space for periodic frames when `periodic: true` (k grid enumerated eagerly per batch from the kernel's `kspace_cutoff`, tin-foil boundary, neutral channels required). Periodic evaluation supports eager and compiled training plus the Python interface; freeze/export, compression, C++ and LAMMPS reject periodic long-range models. Tests: `pytest source/tests/pt_expt/model/test_long_range.py`, `source/tests/common/dpmodel/test_long_range.py`. ### JAX Backend diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index df334852de..752545f4d3 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -242,6 +242,36 @@ def required_topologies(self) -> tuple[str, ...]: """ return () + def build_topologies( + self, + coord: Any, + n_node: Any, + box: Any, + *, + atype: Any, + node_mask: Any = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of :meth:`required_topologies` for one batch. + + The default builds them from their names; a module that owns build + parameters (e.g. the wave-number cutoff of the k grid of a periodic + long-range module) overrides this so they stay with it. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + + return build_topologies( + self.required_topologies(), + coord, + n_node, + box, + node_mask, + atype=atype, + pair_excl=pair_excl, + ) + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. @@ -974,7 +1004,6 @@ def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]: from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, - build_topologies, ) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -1025,13 +1054,12 @@ def model_forward( pair_excl=self.pair_excl, ) atype_flat = xp.reshape(atype, (-1,)) - topologies = build_topologies( - self.required_topologies(), + topologies = self.build_topologies( xp.reshape(coord, (-1, 3)), graph.n_node, box, - node_mask=atype_flat >= 0, atype=atype_flat, + node_mask=atype_flat >= 0, pair_excl=self.pair_excl, ) atomic_ret = self.forward_common_atomic_graph( diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index c52ed7aadc..edb9230592 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -273,6 +273,29 @@ def required_topologies(self) -> tuple[str, ...]: return () return self.long_range.required_topologies() + def build_topologies( + self, + coord: Array, + n_node: Array, + box: Array | None, + *, + atype: Array, + node_mask: Array | None = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of :meth:`required_topologies` for one batch. + + The module builds its own topologies so that its build parameters + (e.g. the wave-number cutoff of the k grid, which follows the kernel) + stay with it. A model without a long-range module declares no + topologies. + """ + if self.long_range is None: + return {} + return self.long_range.build_topologies( + coord, n_node, box, atype=atype, node_mask=node_mask, pair_excl=pair_excl + ) + def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.has_message_passing_across_ranks()) @@ -392,6 +415,12 @@ def enable_compression( ) if self.long_range is not None and self.long_range.uses_charge_spin(): self.long_range.fold_charge_state(self.descriptor.get_default_chg_spin()) + if self.long_range is not None and getattr(self.long_range, "periodic", False): + raise NotImplementedError( + "a periodic long-range module evaluates its reciprocal-space sum " + "with a k grid the compressed path does not serve; compress only " + "non-periodic long-range models" + ) def forward_atomic( self, diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 1dd43e0d45..eefa809590 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -423,6 +423,37 @@ def required_topologies(self) -> tuple[str, ...]: ) ) + def build_topologies( + self, + coord: Array, + n_node: Array, + box: Array | None, + *, + atype: Array, + node_mask: Array | None = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of :meth:`required_topologies` for one batch. + + The first child declaring a name builds it; every child declaring the + same name must build it identically. + """ + topologies: dict[str, Any] = {} + for model in self.models: + if not hasattr(model, "build_topologies"): + continue + built = model.build_topologies( + coord, + n_node, + box, + atype=atype, + node_mask=node_mask, + pair_excl=pair_excl, + ) + for name, topology in built.items(): + topologies.setdefault(name, topology) + return topologies + def forward_atomic_graph( self, graph: Any, diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py index 90a75cbee5..adbb2299dc 100644 --- a/deepmd/dpmodel/long_range/kernel.py +++ b/deepmd/dpmodel/long_range/kernel.py @@ -105,6 +105,53 @@ def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: """ raise NotImplementedError + def kspace(self, k_sq: Array) -> Array: + r"""Fourier transform of the kernel on squared wave numbers. + + :math:`\\hat K(k) = \\int K(r) e^{-i \\mathbf{k} \\cdot \\mathbf{r}} d^3r`, + so that the energy of charges in a periodic cell of volume ``V`` reads + :math:`(k_e / 2V) \\sum_{\\mathbf{k} \\neq 0} \\hat K(k) |\\hat\\rho(\\mathbf{k})|^2`. + + Parameters + ---------- + k_sq + Squared wave numbers with shape ``(nk,)``, in Å⁻²; every entry + must be positive. + + Returns + ------- + Array + The transformed kernel with shape ``(nk,)``, in ų. + """ + raise NotImplementedError + + def contact_value(self, like: Array) -> Array: + r"""The kernel at contact, :math:`K(0)`, in Å⁻¹. + + The k-space lattice sum of a periodic system contains the self terms, + which the evaluation subtracts as :math:`\\frac{1}{2} K(0) \\sum_i q_i^2`. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the result. + """ + raise NotImplementedError + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + """Wave number at which the kernel weight falls below ``tolerance``. + + The Gaussian tail of the Fourier-transformed kernel makes the k-space + truncation exponentially convergent; the k grid of a periodic + evaluation ends at this wave number. + + Parameters + ---------- + tolerance + The relative weight of the slowest-decaying mode at the cutoff. + """ + raise NotImplementedError + def _parameter(self, name: str, like: Array) -> Array: """A parameter in the namespace, dtype and device of ``like``, autograd kept.""" xp = array_api_compat.array_namespace(like) @@ -187,6 +234,24 @@ def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: r = xp.sqrt(r_sq + 1e-12) return xp_erf(coefficients["alpha"] * r) / r + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + # erf(alpha r)/r is the interaction of two Gaussian charge clouds of + # width 1/(sqrt(2) alpha); its transform is the damped Coulomb + # 4 pi exp(-k^2 / 4 alpha^2) / k^2. + alpha = xp.exp(self._parameter("log_alpha", k_sq)) + return 4.0 * np.pi * xp.exp(-k_sq / (4.0 * alpha * alpha)) / k_sq + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + alpha = xp.exp(self._parameter("log_alpha", like)) + return 2.0 * alpha / np.sqrt(np.pi) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # exp(-k^2 / 4 alpha^2) < tolerance + alpha = float(np.exp(to_numpy_array(self.log_alpha).astype(np.float64))) + return 2.0 * alpha * np.sqrt(np.log(1.0 / tolerance)) + @BaseInteractionKernel.register("sog") class SOGKernel(BaseInteractionKernel): @@ -274,6 +339,38 @@ def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: gaussian = xp.exp(-r_sq[:, None] * coefficients["inv_two_width_sq"]) return xp.sum(coefficients["amplitude"] * gaussian, axis=-1) + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + amplitude = self._parameter("amplitude", k_sq) + width = xp.exp(self._parameter("log_width", k_sq)) + # FT of a_k exp(-r^2 / 2 sigma_k^2) is a_k (2 pi)^{3/2} sigma_k^3 + # exp(-sigma_k^2 k^2 / 2); for the default amplitudes + # a_k = sqrt(2/pi) ln b / sigma_k the coefficient is 4 pi ln b sigma_k^2, + # whose ln-sigma trapezoid sum approaches 4 pi / k^2. The mode loop is + # unrolled: a (nk, M) intermediate breaks inductor's double backward of + # the k-space virial (NaN), while per-mode (nk,) terms compile fine. + terms = [ + amplitude[m] + * (2.0 * np.pi) ** 1.5 + * width[m] ** 3 + * xp.exp(-0.5 * width[m] ** 2 * k_sq) + for m in range(width.shape[0]) + ] + return sum(terms[1:], terms[0]) + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + amplitude = self._parameter("amplitude", like) + return xp.sum(amplitude) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # the slowest-decaying mode is the NARROWEST Gaussian: + # exp(-sigma_min^2 k^2 / 2) < tolerance + sigma_min = float( + np.exp(np.min(to_numpy_array(self.log_width).astype(np.float64))) + ) + return np.sqrt(2.0 * np.log(1.0 / tolerance)) / sigma_min + @classmethod def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: return {"n_gaussian": int(np.asarray(variables["log_width"]).size)} diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index 4cc6f3ccd1..c6c00f124a 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -22,6 +22,9 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( ALL_PAIRS, + K_GRID, + KGrid, + build_topologies, frame_id_from_n_node, segment_sum, ) @@ -60,23 +63,33 @@ @BaseLongRange.register("latent_charge") class LatentChargeLongRange(NativeOP, BaseLongRange): - r"""Energy of interacting latent charges in complete non-periodic systems. + r"""Energy of interacting latent charges, periodic or not. A charge network predicts :math:`C` latent charge channels :math:`q_i^c` per atom, in units of :math:`e`, from the descriptor shared with the energy fitting; padding and atoms owned by another rank carry no - charge. The energy is the direct sum over the pairs of atoms of each frame - listed by the :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` - topology, each unordered pair once, + charge. For a non-periodic batch the energy is the direct sum over the + pairs of atoms of each frame listed by the + :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` topology, each + unordered pair once; for a periodic batch (with ``periodic=True``) it is + the reciprocal-space lattice sum over the wave numbers of the + :data:`~deepmd.dpmodel.utils.neighbor_graph.K_GRID` topology, .. math:: - E = k_e \sum_{i None: if charge_net.vacuum_ref: raise ValueError( @@ -113,6 +136,15 @@ def __init__( self.charge_constraint = self._normalize_constraint( charge_constraint, charge_net ) + self.periodic = bool(periodic) + self.kspace_tolerance = float(kspace_tolerance) + if ( + not np.isfinite(self.kspace_tolerance) + or not 0.0 < self.kspace_tolerance < 1.0 + ): + raise ValueError( + f"`kspace_tolerance` must be in (0, 1), got {self.kspace_tolerance}" + ) # The frame charge/spin condition a compressed model is specialized # to; see :meth:`fold_charge_state`. self.charge_state = None @@ -203,9 +235,43 @@ def from_config( charge_net, kernel_cls(**kernel_config), config.get("charge_constraint"), + periodic=config.get("periodic", False), + kspace_tolerance=config.get("kspace_tolerance", 1e-8), + ) + + def build_topologies( + self, + coord: Array, + n_node: Array, + box: Array | None, + *, + atype: Array, + node_mask: Array | None = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of the module for one batch. + + The k grid of a periodic batch is enumerated from the cells with the + wave-number cutoff of the kernel (the trainable parameters are read at + build time, so a drifting kernel re-derives its grid every batch). + """ + params = {} + if self.periodic: + params[K_GRID] = {"k_max": self.kernel.kspace_cutoff(self.kspace_tolerance)} + return build_topologies( + self.required_topologies(), + coord, + n_node, + box, + node_mask, + atype=atype, + pair_excl=pair_excl, + topology_params=params, ) def required_topologies(self) -> tuple[str, ...]: + if self.periodic: + return (ALL_PAIRS, K_GRID) return (ALL_PAIRS,) def uses_charge_spin(self) -> bool: @@ -274,6 +340,16 @@ def call_graph( charge = self._constrain( charge, node_mask, frame_id, n_node, fparam, charge_spin ) + energy = self._pair_branch(charge, pairs, n) + if self.periodic: + energy = energy + self._kspace_branch( + charge, topologies[K_GRID], frame_id, node_mask, n_node + ) + return xp.reshape(energy, (n, 1)) + + def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: + """Direct pair sum over the non-periodic frames, per atom.""" + xp = array_api_compat.array_namespace(charge) src = pairs.index[0] dst = pairs.index[1] pair_charge = xp.sum( @@ -284,8 +360,84 @@ def call_graph( half_energy = (0.5 * COULOMB_CONSTANT) * xp.where( pairs.mask, pair_energy, xp.zeros_like(r_sq) ) - energy = segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) - return xp.reshape(energy, (n, 1)) + return segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) + + def _kspace_branch( + self, + charge: Array, + k_grid: KGrid, + frame_id: Array, + node_mask: Array, + n_node: Array, + ) -> Array: + r"""Reciprocal-space energy of the periodic frames, per atom. + + Evaluates + :math:`E_f = k_e [ \frac{1}{2V_f} \sum_{\mathbf{k}\neq\mathbf{0}} + \hat K(\mathbf{k}) |\hat\rho(\mathbf{k})|^2 + - \frac{1}{2} K(0) \sum_{i \in f} |q_i|^2 ]` + on the strained geometry of :attr:`KGrid.strain` (identity when no + strain leaf is set): coordinates by ``I + strain``, wave numbers by + ``(I + strain)^{-1}`` on their row layout and volumes by + ``det(I + strain)``, so that the gradient of the energy with respect + to the strain leaf is the negative k-space virial of the frame. The + reciprocal energy is attributed to the charged atoms of the frame in + equal shares; the self term is per atom by construction. Guard points + carry a zero mask. + """ + xp = array_api_compat.array_namespace(charge) + nf = n_node.shape[0] + coord = k_grid.coord + if coord is None: + raise ValueError( + "the k-grid topology carries no coordinates; the graph lower " + "re-leafs them before the evaluation" + ) + coord = xp.astype(coord, charge.dtype) + k_vec = xp.astype(k_grid.vec, charge.dtype) + volume = xp.astype(k_grid.volume, charge.dtype) + if k_grid.strain is not None: + eye = xp.eye(3, dtype=charge.dtype, device=array_api_compat.device(charge)) + s = xp.astype(k_grid.strain, charge.dtype) + eye # (nf, 3, 3) + # r' = S r (columns) + s_node = xp.take(s, frame_id, axis=0) # (N, 3, 3) + coord = xp.matmul(s_node, coord[:, :, None])[:, :, 0] + # k' = S^{-T} k (columns), i.e. rows multiplied by S^{-1} + s_inv_t = xp.matrix_transpose(xp.linalg.inv(s)) + s_inv_t_k = xp.take(s_inv_t, k_grid.frame, axis=0) # (nk, 3, 3) + k_vec = xp.matmul(s_inv_t_k, k_vec[:, :, None])[:, :, 0] + volume = volume * xp.linalg.det(s) + phase = xp.matmul(k_vec, xp.matrix_transpose(coord)) # (nk, N) + same_frame = frame_id[None, :] == k_grid.frame[:, None] + cos_phase = xp.cos(phase) * same_frame + sin_phase = xp.sin(phase) * same_frame + rho_real = xp.matmul(cos_phase, charge) # (nk, C) + rho_imag = xp.matmul(sin_phase, charge) # (nk, C) + s_sq = xp.sum(rho_real**2 + rho_imag**2, axis=-1) # (nk,) + k_sq = xp.sum(k_vec * k_vec, axis=-1) + kernel_value = self.kernel.kspace(k_sq) # (nk,) + half_ke = 0.5 * COULOMB_CONSTANT + volume_k = xp.take(volume, k_grid.frame, axis=0) + energy_k = half_ke * kernel_value * s_sq / volume_k + energy_k = xp.where(k_grid.mask, energy_k, xp.zeros_like(energy_k)) + e_rec_frame = segment_sum(energy_k, k_grid.frame, nf) # (nf,) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + count = xp.maximum(count, xp.ones_like(count)) + e_rec = xp.take(e_rec_frame / count, frame_id, axis=0) * weight # (N,) + # The self term belongs to the periodic lattice sum alone: a frame + # without real k points (a non-periodic batch on a periodic-enabled + # model) keeps its non-periodic self-free energy. + periodic_frame = ( + segment_sum(xp.astype(k_grid.mask, charge.dtype), k_grid.frame, nf) > 0 + ) + e_self = ( + -half_ke + * self.kernel.contact_value(charge) + * xp.sum(charge * charge, axis=-1) + * xp.astype(xp.take(periodic_frame, frame_id, axis=0), charge.dtype) + ) + return e_rec + e_self def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: """Kernel values of the pairs. @@ -406,6 +558,9 @@ def serialize(self) -> dict: "kernel": self.kernel.serialize(), "charge_constraint": self.charge_constraint, } + if self.periodic: + data["periodic"] = True + data["kspace_tolerance"] = self.kspace_tolerance if self.charge_state is not None: data["charge_state"] = to_numpy_array(self.charge_state).tolist() return data @@ -418,6 +573,8 @@ def deserialize(cls, data: dict) -> "LatentChargeLongRange": LatentChargeFitting.deserialize(data["charge_net"]), BaseInteractionKernel.deserialize(data["kernel"]), data["charge_constraint"], + periodic=data.get("periodic", False), + kspace_tolerance=data.get("kspace_tolerance", 1e-8), ) if data.get("charge_state") is not None: obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index f05d1d9e09..42b85f07d9 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -49,7 +49,6 @@ PairGraph, build_neighbor_graph, build_neighbor_graph_ase, - build_topologies, compact_nodes, expand_node_values, ) @@ -567,8 +566,7 @@ def _call_common_graph( # identity. ng, node_index = compact_nodes(ng, atype_flat >= 0) atype_real = xp.take(atype_flat, node_index, axis=0) - topologies = build_topologies( - self.atomic_model.required_topologies(), + topologies = self.atomic_model.build_topologies( xp.take(xp.reshape(cc, (n_padded, 3)), node_index, axis=0), ng.n_node, bb, @@ -893,14 +891,35 @@ def call_common_lower_graph( spin=spin, ) if topologies: - # The pair vectors are geometry like the edge vectors. + # The pair vectors and the k-grid geometry are geometry like + # the edge vectors. + from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + ) + xp = array_api_compat.array_namespace(edge_vec) - topologies = { - name: PairGraph( - pairs.index, xp.astype(pairs.vec, edge_vec.dtype), pairs.mask - ) - for name, pairs in topologies.items() - } + cast = {} + for name, topology in topologies.items(): + if isinstance(topology, KGrid): + cast[name] = KGrid( + frame=topology.frame, + vec=xp.astype(topology.vec, edge_vec.dtype), + mask=topology.mask, + volume=xp.astype(topology.volume, edge_vec.dtype), + coord=( + None + if topology.coord is None + else xp.astype(topology.coord, edge_vec.dtype) + ), + strain=topology.strain, + ) + else: + cast[name] = PairGraph( + topology.index, + xp.astype(topology.vec, edge_vec.dtype), + topology.mask, + ) + topologies = cast model_predict = self.forward_common_atomic_graph( atype, n_node, diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 8779648f15..bf503530de 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -58,11 +58,15 @@ ) from .pair_graph import ( ALL_PAIRS, + K_GRID, + KGrid, PairGraph, + build_k_grid, build_pair_graph, build_topologies, check_complete_frames, flatten_topologies, + topology_flat_fields, unflatten_topologies, unordered_pair_index, ) @@ -78,7 +82,9 @@ __all__ = [ "ALL_PAIRS", + "K_GRID", "GraphLayout", + "KGrid", "NeighborGraph", "PairGraph", "angle_padding_fraction", @@ -90,6 +96,7 @@ "attach_edge_csr", "build_angle_index", "build_edge_csr", + "build_k_grid", "build_neighbor_graph", "build_neighbor_graph_ase", "build_pair_graph", @@ -116,6 +123,7 @@ "segment_mean", "segment_softmax", "segment_sum", + "topology_flat_fields", "unflatten_topologies", "unordered_pair_index", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index 4e9c653128..e98cc10d51 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -22,6 +22,7 @@ annotations, ) +import math from dataclasses import ( dataclass, ) @@ -30,6 +31,11 @@ ) import array_api_compat +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) from .graph import ( frame_id_from_n_node, @@ -55,6 +61,14 @@ ALL_PAIRS = "all_pairs" """Name of the topology holding every unordered pair of valid nodes of a frame.""" +K_GRID = "k_grid" +"""Name of the topology holding the reciprocal lattice points of periodic frames. + +A long-range module declares it beside :data:`ALL_PAIRS` when it serves +periodic systems: the reciprocal-space evaluation of the pair kernel needs the +wave numbers, which the pair topology cannot carry. +""" + @dataclass class PairGraph: @@ -68,6 +82,44 @@ class PairGraph: """(P,) bool real (1) vs guard (0) pairs; guards join node 0 to itself.""" +@dataclass +class KGrid: + r"""Reciprocal lattice points of the periodic frames of a batch. + + The wave numbers enumerate :math:`\\mathbf{k} = 2\\pi B^{-1}\\mathbf{n}` for + integer ``n`` with :math:`0 < |\\mathbf{k}| \\le k_{max}`; they are built + eagerly outside any compiled region, concatenated ragged over the frames, + and the k axis ends in masked guard points that keep its length a dynamic + dimension of its own however few points a batch holds. + """ + + frame: Array + """(nk,) int64 the frame of every wave number.""" + vec: Array + """(nk, 3) float the wave numbers, in Å⁻¹.""" + mask: Array + """(nk,) bool real (1) vs guard (0) points.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array | None = None + """(N, 3) float node positions; the geometry leaf of the k-space energy. + + ``None`` when the grid is reconstructed from flattened inputs (export); a + graph lower re-leafs it before the evaluation. + """ + strain: Array | None = None + """(nf, 3, 3) float zero strain leaf for the k-space virial. + + The module evaluates the energy on coordinates strained by ``I + strain``, + wave numbers strained by ``(I + strain)^{-T}`` and volumes scaled by + ``det(I + strain)``; the gradient of the energy with respect to this leaf + is the negative k-space virial. The leaf is built eagerly beside the grid + (all zeros) and re-leafed by the graph lower, so the backward roots at an + input of the traced graph — a leaf created INSIDE a ``make_fx`` trace + produces a NaN gradient under inductor. + """ + + def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: """Enumerate the unordered pairs of valid nodes within every frame. @@ -134,6 +186,113 @@ def row_start(row: Array) -> Array: return xp.stack([src, dst]) +def build_k_grid( + box: Array, + k_max: float, + n_node: Array, + coord: Array | None = None, +) -> KGrid: + r"""Enumerate the reciprocal lattice points of the periodic frames. + + For every frame the wave numbers are + :math:`\\mathbf{k} = 2\\pi \\, B^{-1} \\mathbf{n}` over integer + :math:`\\mathbf{n}` with :math:`0 < |\\mathbf{k}| \\le k_{max}` — note the + cell enters through ``B^{-1}``, NOT ``B^{-T}``: with the lattice vectors as + the ROWS of ``B``, orthogonality reads ``B @ (B^{-1} n) = n``. The count is + data dependent, so the enumeration runs eagerly. Every frame's points come + in lexicographic integer order, and the batch ends in two masked guard + points. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + k_max + The wave-number cutoff, in Å⁻¹. + n_node + Per-frame node counts with shape ``(nf,)``. + coord + Node positions with shape ``(N, 3)``, stored on the grid as the + geometry context (see :attr:`KGrid.coord`). + + Returns + ------- + KGrid + The reciprocal lattice points of every frame, ragged-concatenated. + """ + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + # A flat (nf, 9) layout arrives from callers that carry the cell flattened. + if box.ndim == 2: + box = xp.reshape(box, (box.shape[0], 3, 3)) + nf = box.shape[0] + k_lists = [] + frame_lists = [] + for f in range(nf): + cell = np.asarray(to_numpy_array(box[f]), dtype=np.float64) + rec = 2.0 * math.pi * np.linalg.inv(cell) # columns are g_d + a_len = np.linalg.norm(cell, axis=1) + n_d = np.ceil(k_max * a_len / (2.0 * math.pi)).astype(np.int64) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=np.float64) for n in n_d], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0.0) & (k2 <= k_max * k_max) + k_lists.append(k[keep]) + frame_lists.append(np.full(int(np.sum(keep)), f, dtype=np.int64)) + if k_lists: + vec_np = np.concatenate(k_lists, axis=0) + frame_np = np.concatenate(frame_lists, axis=0) + else: + vec_np = np.zeros((0, 3), dtype=np.float64) + frame_np = np.zeros((0,), dtype=np.int64) + # Guard points: masked out, keep the k axis a dynamic dimension of its own. + guard = np.tile(np.array([[1e-6, 0.0, 0.0]]), (2, 1)) + vec_np = np.concatenate([vec_np, guard], axis=0) + frame_np = np.concatenate([frame_np, np.zeros(2, dtype=np.int64)], axis=0) + mask_np = np.zeros(vec_np.shape[0], dtype=bool) + mask_np[: vec_np.shape[0] - 2] = True + volume = xp.linalg.det(box) + volume = xp.abs(volume) + return KGrid( + frame=xp.asarray(frame_np, device=device), + vec=xp.astype(xp.asarray(vec_np, device=device), box.dtype), + mask=xp.asarray(mask_np, device=device), + volume=xp.astype(volume, box.dtype), + coord=coord, + strain=xp.zeros((nf, 3, 3), dtype=box.dtype, device=device), + ) + + +def _empty_k_grid(nf: int, like: Array, coord: Array | None) -> KGrid: + """A guard-only k grid for a batch without periodic frames.""" + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + return KGrid( + frame=xp.zeros((2,), dtype=xp.int64, device=device), + vec=xp.tile( + xp.asarray([[1e-6, 0.0, 0.0]], dtype=like.dtype, device=device), (2, 1) + ), + mask=xp.zeros((2,), dtype=xp.bool, device=device), + volume=xp.ones((nf,), dtype=like.dtype, device=device), + coord=coord, + strain=xp.zeros((nf, 3, 3), dtype=like.dtype, device=device), + ) + + +def _empty_pair_graph(like: Array, atype: Array) -> PairGraph: + """A guard-only pair topology for a batch without non-periodic frames.""" + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + index = xp.zeros((2, 0), dtype=xp.int64, device=device) + vec = xp.zeros((0, 3), dtype=like.dtype, device=device) + index, vec, mask = pad_and_guard_edges(index, vec, None) + return PairGraph(index=index, vec=vec, mask=mask) + + def build_pair_graph( coord: Array, n_node: Array, @@ -184,7 +343,8 @@ def build_topologies( *, atype: Array, pair_excl: PairExcludeMask | None = None, -) -> dict[str, PairGraph]: + topology_params: dict[str, dict] | None = None, +) -> dict[str, PairGraph | KGrid]: """Build the topologies an atomic model declares, beside its neighbor graph. Parameters @@ -196,7 +356,9 @@ def build_topologies( n_node Per-frame node counts with shape ``(nf,)``. box - Simulation cells, or ``None`` for non-periodic systems. + Simulation cells, or ``None`` for non-periodic systems. With a cell, + the pair topology comes back guard-only and the k grid holds the + reciprocal lattice points of the frames; without one, the reverse. node_mask Validity of the nodes with shape ``(N,)``, or ``None`` when every node is valid. @@ -204,26 +366,55 @@ def build_topologies( Node types with shape ``(N,)``; invalid nodes carry any type. pair_excl The model-level pair exclusion, or ``None``. + topology_params + Per-name build parameters; the k grid takes ``{"k_max": ...}``. Returns ------- - dict[str, PairGraph] + dict[str, PairGraph | KGrid] Each requested topology under its name; empty when none is required. Raises ------ NotImplementedError - If a cell is given: the pair topology describes complete non-periodic - systems only. + If a cell is given for a pair-only declaration, or a pair exclusion is + set for a periodic batch (the k-space sum is global). ValueError - If a name is unknown. + If a name is unknown, or the k grid parameters are missing. """ - topologies = {} + topologies: dict[str, PairGraph | KGrid] = {} for name in names: - if name != ALL_PAIRS: - raise ValueError(f"unknown topology {name!r}; expected {ALL_PAIRS!r}") - check_complete_frames(box) - topologies[name] = build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + if name == ALL_PAIRS: + if box is None: + topologies[name] = build_pair_graph( + coord, n_node, atype, node_mask, pair_excl + ) + elif K_GRID in names: + # A periodic-capable module reads the k grid for these frames. + topologies[name] = _empty_pair_graph(coord, atype) + else: + check_complete_frames(box) + elif name == K_GRID: + if box is None: + topologies[name] = _empty_k_grid(n_node.shape[0], coord, coord) + else: + if pair_excl is not None: + raise NotImplementedError( + "the k-space sum of a periodic system is global; " + "`pair_exclude_types` cannot exclude pairs from it" + ) + params = (topology_params or {}).get(K_GRID, {}) + if "k_max" not in params: + raise ValueError( + f"the {K_GRID!r} topology requires its `k_max` build parameter" + ) + topologies[name] = build_k_grid( + box, params["k_max"], n_node, coord=coord + ) + else: + raise ValueError( + f"unknown topology {name!r}; expected one of ({ALL_PAIRS!r}, {K_GRID!r})" + ) return topologies @@ -248,12 +439,14 @@ def check_complete_frames(box: Array | None) -> None: ) -def flatten_topologies(topologies: dict[str, PairGraph]) -> tuple[Array, ...]: +def flatten_topologies(topologies: dict[str, PairGraph | KGrid]) -> tuple[Array, ...]: """Lay topologies out as positional inputs of a graph lower. - A compiled or exported graph lower receives every topology as three - inputs, ``(index, vec, mask)``, in the order of the dict, after its other - inputs; :func:`unflatten_topologies` restores them. + A compiled or exported graph lower receives every topology positionally, + in the order of the dict, after its other inputs: a pair topology as its + ``(index, vec, mask)``, a k grid as its ``(frame, vec, mask, volume, + coord)``; + :func:`unflatten_topologies` restores them. Parameters ---------- @@ -263,18 +456,37 @@ def flatten_topologies(topologies: dict[str, PairGraph]) -> tuple[Array, ...]: Returns ------- tuple - ``(index, vec, mask)`` of every topology, concatenated. + The flattened tensors of every topology, concatenated. """ return tuple( value - for pairs in topologies.values() - for value in (pairs.index, pairs.vec, pairs.mask) + for topology in topologies.values() + for value in topology_flat_fields(topology) ) +def topology_flat_fields(topology: PairGraph | KGrid) -> tuple[Array, ...]: + """The positional tensors of one topology.""" + if isinstance(topology, KGrid): + if topology.coord is None or topology.strain is None: + raise ValueError( + "a k grid carries no coordinates or strain; it cannot be laid out " + "positionally" + ) + return ( + topology.frame, + topology.vec, + topology.mask, + topology.volume, + topology.coord, + topology.strain, + ) + return (topology.index, topology.vec, topology.mask) + + def unflatten_topologies( names: Sequence[str], values: Sequence[Array] -) -> dict[str, PairGraph]: +) -> dict[str, PairGraph | KGrid]: """Restore topologies laid out by :func:`flatten_topologies`. Parameters @@ -282,22 +494,37 @@ def unflatten_topologies( names The topology names, in the order they were laid out. values - ``(index, vec, mask)`` of every topology, concatenated. + The flattened tensors of every topology, concatenated. Returns ------- - dict[str, PairGraph] + dict[str, PairGraph | KGrid] The topologies by name. Raises ------ ValueError - If the number of values is not three per name. + If the number of values does not match the layout of the names. """ - if len(values) != 3 * len(names): + arity = {ALL_PAIRS: 3, K_GRID: 6} + topologies: dict[str, PairGraph | KGrid] = {} + offset = 0 + for name in names: + if name not in arity: + raise ValueError(f"unknown topology {name!r}") + fields = values[offset : offset + arity[name]] + if len(fields) != arity[name]: + raise ValueError( + f"expected {arity[name]} inputs for topology {name!r}, " + f"got {len(values) - offset} remaining" + ) + if name == ALL_PAIRS: + topologies[name] = PairGraph(*fields) + else: + topologies[name] = KGrid(*fields) + offset += arity[name] + if offset != len(values): raise ValueError( - f"expected three inputs per topology for {tuple(names)}, got {len(values)}" + f"{len(values)} flattened values given, {offset} consumed by {tuple(names)}" ) - return { - name: PairGraph(*values[3 * ii : 3 * ii + 3]) for ii, name in enumerate(names) - } + return topologies diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index abee40c07c..1dd832aad1 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -2556,19 +2556,40 @@ def _eval_model_graph( if topology_names: # Padding slots of a mixed-type batch are no atoms and join # no pair. - model_inputs += flatten_topologies( - build_topologies( + coord_flat_t = torch.as_tensor( + coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE + ) + if self._dpmodel is not None: + # The live model builds its own topologies, so the k-grid + # cutoff follows the kernel parameters. + topologies = self._dpmodel.atomic_model.build_topologies( + coord_flat_t, + n_node_t, + box_input, + atype=atype_t, + node_mask=atype_t >= 0, + pair_excl=self._model_pair_excl(), + ) + else: + from deepmd.dpmodel.utils.neighbor_graph import ( + K_GRID, + ) + + if K_GRID in topology_names: + raise NotImplementedError( + "a periodic long-range model cannot be evaluated from " + "a .pt2/.pte archive yet" + ) + topologies = build_topologies( topology_names, - torch.as_tensor( - coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE - ), + coord_flat_t, n_node_t, box_input, node_mask=atype_t >= 0, atype=atype_t, pair_excl=self._model_pair_excl(), ) - ) + model_inputs += flatten_topologies(topologies) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 3e95d33a3a..a2bd41aaac 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -25,6 +25,7 @@ get_deriv_name_mag, ) from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, NeighborGraph, PairGraph, edge_force_virial, @@ -69,6 +70,8 @@ def edge_energy_deriv( create_graph: bool = False, force_precision: torch.dtype | None = None, pair_graphs: Sequence[PairGraph] = (), + coord_leaf: torch.Tensor | None = None, + strain_leaf: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -94,6 +97,11 @@ def edge_energy_deriv( the coordinates through all its geometry leaves, so its complete force and virial are the sum of their pullbacks. + A k-space long-range term reaches the geometry through ``coord_leaf`` (the + node positions, whose gradient is its force directly) and ``strain_leaf`` + (a zero (nf, 3, 3) strain the module applies to coordinates, wave numbers + and volumes, whose gradient is the negative k-space virial). + Parameters ---------- energy @@ -127,6 +135,10 @@ def edge_energy_deriv( pair_graphs Pair topologies whose ``vec`` are further autograd leaves of ``energy``, on the same node axis. + coord_leaf + (N, 3) node positions of a k-space long-range term, or ``None``. + strain_leaf + (nf, 3, 3) zero strain of a k-space long-range term, or ``None``. Returns ------- @@ -137,12 +149,16 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - g_e, *g_pairs = torch.autograd.grad( + extra_leaves = tuple(leaf for leaf in (coord_leaf, strain_leaf) if leaf is not None) + g_e, *g_rest = torch.autograd.grad( energy.sum() if energy.dim() else energy, - (edge_vec, *(pairs.vec for pairs in pair_graphs)), + (edge_vec, *(pairs.vec for pairs in pair_graphs), *extra_leaves), create_graph=create_graph, retain_graph=True, ) + g_pairs = g_rest[: len(pair_graphs)] + g_coord = g_rest[len(pair_graphs)] if coord_leaf is not None else None + g_strain = g_rest[-1] if strain_leaf is not None else None if ( force_precision is not None and not create_graph @@ -228,6 +244,31 @@ def edge_energy_deriv( virial = virial + pair_virial.to(virial.dtype) if do_atomic_virial: atom_virial = atom_virial + pair_atom_virial.to(atom_virial.dtype) + if coord_leaf is not None or strain_leaf is not None: + # The k-space term: the coordinate gradient is its force, the strain + # gradient its negative virial. The atomic virial convention keeps the + # per-frame sum exact: the r x F part per node, the remainder shared + # evenly over the frame. + frame_id = frame_id_from_n_node(n_node, force.shape[0]) + kspace_vir_frame = None + if strain_leaf is not None: + kspace_vir_frame = -g_strain # (nf, 3, 3) + virial = virial + kspace_vir_frame.to(virial.dtype) + if coord_leaf is not None: + f_coord = -g_coord + force = force + f_coord.to(force.dtype) + if do_atomic_virial: + av = torch.zeros_like(atom_virial) + if coord_leaf is not None: + # w = -g (x) r, force index first as for the edge virial + av = av + torch.einsum("ni,nj->nij", f_coord, coord_leaf) + if kspace_vir_frame is not None: + remainder = kspace_vir_frame - segment_sum( + av, frame_id, n_node.shape[0] + ) + count = n_node.to(av.dtype)[frame_id].clamp(min=1).reshape(-1, 1, 1) + av = av + remainder[frame_id] / count + atom_virial = atom_virial + av.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -409,6 +450,14 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: size *= ii # split the reduced output into ``size`` per-frame scalar components. svv = model_ret[kk_redu].reshape(nf, size) + # The k-space geometry leaves of a periodic long-range module: the + # node positions (its force) and the zero strain (its virial). + coord_leaf = strain_leaf = None + if topologies: + for topology in topologies.values(): + if isinstance(topology, KGrid): + coord_leaf = topology.coord + strain_leaf = topology.strain ff_list: list[torch.Tensor] = [] av_list: list[torch.Tensor] = [] vir_list: list[torch.Tensor] = [] @@ -429,7 +478,15 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, force_precision=force_precision if not create_graph else None, - pair_graphs=() if topologies is None else tuple(topologies.values()), + pair_graphs=( + () + if topologies is None + else tuple( + t for t in topologies.values() if isinstance(t, PairGraph) + ) + ), + coord_leaf=coord_leaf, + strain_leaf=strain_leaf, ) # force (N, 3) -> (N, 1, 3) [flat; caller unravels at I/O boundary] ff_list.append(force.reshape(N, 1, 3)) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index b2d1e2c4ef..3499d0410b 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -21,8 +21,8 @@ OutputVariableDef, ) from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, PairGraph, - build_topologies, compact_nodes, expand_node_values, flatten_topologies, @@ -407,8 +407,7 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: ) # The topologies follow the same coordinates, so the Hessian holds # every geometry path of the energy and the blocks coupling them. - topologies = build_topologies( - self.model.atomic_model.required_topologies(), + topologies = self.model.atomic_model.build_topologies( coord_flat.reshape(self.nloc, 3), ng.n_node, self.box, @@ -743,14 +742,35 @@ def forward_common_lower_graph( "a domain-decomposed system is not supported" ) # Every topology carries a geometry leaf of its own. - topologies = { - name: PairGraph( - index=pairs.index, - vec=pairs.vec.detach().requires_grad_(True), - mask=pairs.mask, - ) - for name, pairs in topologies.items() - } + re_leafed: dict[str, Any] = {} + for name, topology in topologies.items(): + if isinstance(topology, KGrid): + if topology.coord is None or topology.strain is None: + raise ValueError( + "the k-grid topology carries no coordinates or strain" + ) + # The positions are the geometry leaf of the k-space + # energy; the zero strain leaf turns its gradient into + # the k-space virial (see KGrid.strain). Both re-leaf + # here so the backward roots at graph INPUTS — a leaf + # created inside the trace is NaN under inductor. + coord_leaf = topology.coord.detach().requires_grad_(True) + strain_leaf = topology.strain.detach().requires_grad_(True) + re_leafed[name] = KGrid( + frame=topology.frame, + vec=topology.vec, + mask=topology.mask, + volume=topology.volume, + coord=coord_leaf, + strain=strain_leaf, + ) + else: + re_leafed[name] = PairGraph( + index=topology.index, + vec=topology.vec.detach().requires_grad_(True), + mask=topology.mask, + ) + topologies = re_leafed if spin is not None: # second autograd leaf: force_mag = -dE/dspin (native spin). # Deliberately a SEPARATE leaf/backward from edge_vec rather @@ -943,8 +963,7 @@ def call_common_ragged( self.get_rcut(), pair_excl, ) - topologies = build_topologies( - self.atomic_model.required_topologies(), + topologies = self.atomic_model.build_topologies( coord.detach(), graph.n_node, box, @@ -1073,8 +1092,7 @@ def _call_common_graph( spin_flat = ( spin.reshape(n_padded, 3)[node_index] if spin is not None else None ) - topologies = build_topologies( - self.atomic_model.required_topologies(), + topologies = self.atomic_model.build_topologies( cc.detach().reshape(n_padded, 3)[node_index], ng.n_node, bb, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f4e24bc824..d90e67d36e 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -73,7 +73,7 @@ make_learning_rate_schedule, ) from deepmd.dpmodel.utils.neighbor_graph import ( - build_topologies, + KGrid, flatten_topologies, unflatten_topologies, ) @@ -934,6 +934,7 @@ def _trace_and_compile_graph( ), dtype=GLOBAL_PT_FLOAT_PRECISION, device=_trace_device, + nframes=trace_nf, ) topology_names = tuple(s_topologies) s_topology_vals = flatten_topologies(s_topologies) @@ -1685,8 +1686,7 @@ def _forward_graph( spin = spin[node_index] # The topologies the atomic model declares are built here as well, # outside the compiled region, on the node axis of the graph. - topologies = build_topologies( - _model.atomic_model.required_topologies(), + topologies = _model.atomic_model.build_topologies( coord_3d if ragged else coord_3d.reshape(n_padded, 3)[node_index], ng.n_node, box_flat, @@ -1721,16 +1721,28 @@ def _forward_graph( if spin is not None: spin = spin.detach().requires_grad_(True) # Each pair topology enters as (index, vec, mask), its vec a further - # leaf. - topology_vals = tuple( - tensor - for pairs in topologies.values() - for tensor in ( - pairs.index, - pairs.vec.detach().requires_grad_(True), - pairs.mask, - ) - ) + # leaf; a k grid enters as (frame, vec, mask, volume, coord, strain), + # its coord and zero strain further leaves (the k-space force and + # virial root at these inputs). + topology_vals_list: list[torch.Tensor] = [] + for topology in topologies.values(): + if isinstance(topology, KGrid): + assert topology.coord is not None and topology.strain is not None + topology_vals_list += [ + topology.frame, + topology.vec, + topology.mask, + topology.volume, + topology.coord.detach().requires_grad_(True), + topology.strain.detach().requires_grad_(True), + ] + else: + topology_vals_list += [ + topology.index, + topology.vec.detach().requires_grad_(True), + topology.mask, + ] + topology_vals = tuple(topology_vals_list) if self._task_buf_order: try: diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index ebba8ca4a3..cc92d173ee 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -685,13 +685,15 @@ def build_synthetic_topologies( *, dtype: torch.dtype, device: torch.device, + nframes: int | None = None, ) -> dict[str, PairGraph]: """Build trace-time samples of the topologies a model declares. Single source of the trace-time topology inputs, shared by graph export and compiled training. Tracing reads only their shapes, so every topology - joins fixed nodes over ``n_pair`` pairs, a length the caller keeps - distinct from the other dynamic axes. + joins fixed nodes over ``n_pair`` entries, a length the caller keeps + distinct from the other dynamic axes. A k grid additionally carries the + frame volumes and the node coordinates. Parameters ---------- @@ -705,24 +707,45 @@ def build_synthetic_topologies( Precision of the pair vectors. device : torch.device Target device. + nframes : int | None + The frame count, needed for the volume axis of a k grid; ``None`` + for a model declaring none. Returns ------- dict[str, PairGraph] A sample of every declared topology; empty when none is declared. """ + from deepmd.dpmodel.utils.neighbor_graph import ( + K_GRID, + KGrid, + ) + index = ( torch.arange(2 * n_pair, dtype=torch.int64, device=device).reshape(2, n_pair) % n_node_total ) - return { - name: PairGraph( - index=index, - vec=torch.ones(n_pair, 3, dtype=dtype, device=device), - mask=torch.ones(n_pair, dtype=torch.bool, device=device), - ) - for name in model.atomic_model.required_topologies() - } + topologies = {} + for name in model.atomic_model.required_topologies(): + if name == K_GRID: + if nframes is None: + raise ValueError(f"a {K_GRID!r} topology sample needs the frame count") + topologies[name] = KGrid( + frame=torch.zeros(n_pair, dtype=torch.int64, device=device), + # Nonzero wave numbers: the kernel transform carries 1/k^2. + vec=torch.ones(n_pair, 3, dtype=dtype, device=device), + mask=torch.ones(n_pair, dtype=torch.bool, device=device), + volume=torch.ones(nframes, dtype=dtype, device=device), + coord=torch.ones(n_node_total, 3, dtype=dtype, device=device), + strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), + ) + else: + topologies[name] = PairGraph( + index=index, + vec=torch.ones(n_pair, 3, dtype=dtype, device=device), + mask=torch.ones(n_pair, dtype=torch.bool, device=device), + ) + return topologies def build_synthetic_canonical_graph_inputs( @@ -949,7 +972,19 @@ def _build_graph_dynamic_shapes( topologies = [] for name in topology_names: npair_dim = torch.export.Dim(f"npair_{name}", min=2) - topologies += [{1: npair_dim}, {0: npair_dim}, {0: npair_dim}] + if name == "k_grid": + # (frame, vec, mask, volume, coord, strain): the k axis, the + # frame axis and the node axis. + topologies += [ + {0: npair_dim}, + {0: npair_dim}, + {0: npair_dim}, + {0: nframes_dim}, + {0: n_node_total_dim}, + {0: nframes_dim}, + ] + else: + topologies += [{1: npair_dim}, {0: npair_dim}, {0: npair_dim}] if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -1921,6 +1956,18 @@ def _trace_and_export_impl( # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) model.eval() + # The k-space evaluation of a periodic long-range module needs the + # reciprocal grid of the input cells, which no exported lower ABI carries + # yet; freeze and compress reject such models until they do. + from deepmd.dpmodel.utils.neighbor_graph import ( + K_GRID, + ) + + if K_GRID in model.atomic_model.required_topologies(): + raise NotImplementedError( + "a periodic long-range module (the k-grid topology) cannot be " + "frozen or exported yet; serve it through the Python interface" + ) # The vacuum reference is resolved on the export object, folded into the # fitting bias or stored as a per-type table, so the exported graph # carries no reference atoms; the move to the tracing device follows, so @@ -2189,6 +2236,7 @@ def _trace_and_export_impl( ), dtype=edge_dtype, device=torch.device("cpu"), + nframes=nframes_sample, ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a99c8d60b4..4cc56193f1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3674,8 +3674,10 @@ def charge_constraint_args() -> list[Argument]: @long_range_args_plugin.register( "latent_charge", - doc="Interaction of latent charges predicted from the descriptor, summed " - "directly over every pair of atoms of a complete non-periodic system. " + doc="Interaction of latent charges predicted from the descriptor. " + "For a non-periodic system the kernel is summed directly over every pair " + "of atoms; for a periodic system (``periodic: true``) its Fourier " + "transform is summed over the reciprocal lattice points of the cell. " "Every channel interacts through the same kernel; half of each pair energy " "is attributed to each atom.", ) @@ -3692,6 +3694,16 @@ def long_range_latent_charge() -> list[Argument]: "so that it sums to the total charge of the frame. Unconstrained if not " "set." ) + doc_periodic = ( + "Evaluate the module on periodic systems in reciprocal space: the " + "Fourier transform of the kernel is summed over the wave numbers of " + "each cell, with the tin-foil boundary convention, and the self term " + "subtracted. Every latent channel must be neutral." + ) + doc_kspace_tolerance = ( + "The relative weight of the slowest-decaying kernel mode at the " + "wave-number cutoff of the k grid of a periodic batch." + ) return [ Argument("channels", int, optional=True, default=1, doc=doc_channels), Argument( @@ -3722,15 +3734,23 @@ def long_range_latent_charge() -> list[Argument]: default=None, doc=doc_charge_constraint, ), + Argument("periodic", bool, optional=True, default=False, doc=doc_periodic), + Argument( + "kspace_tolerance", + float, + optional=True, + default=1e-8, + doc=doc_kspace_tolerance, + ), ] def long_range_args() -> Argument: doc_long_range = ( "A long-range module added to the atomic energy of the energy fitting. " - "It reads the same descriptor and is evaluated on complete " - "non-periodic systems; a periodic cell is rejected. Currently validated " - "with the DPA4C descriptor." + "It reads the same descriptor and is evaluated on complete systems; " + "non-periodic directly, periodic in reciprocal space when the module " + "enables it. Currently validated with the DPA4C descriptor." ) return Argument( "long_range", diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 066bdef194..00fd65f3f0 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -17,9 +17,12 @@ E = \sum_i E_i^{\text{SR}} + k_e \sum_{i [!WARNING] -> The pair sum covers complete non-periodic systems: frames without a cell, -> such as isolated molecules and clusters. A periodic cell is rejected. The -> cost of the sum grows with the square of the number of atoms of a frame. +> The pair sum covers complete systems. Without `periodic: true`, a periodic +> cell is rejected; the cost of the non-periodic sum grows with the square of +> the number of atoms of a frame. ## Configuration @@ -100,6 +103,50 @@ Several charge channels, set by {ref}`channels `, let the model represent interactions that one charge per atom cannot. +### Periodic systems + +With +{ref}`periodic ` enabled, +a frame given a cell is evaluated in reciprocal space. The Fourier transform +of the kernel, + +- `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, the damped + Coulomb transform (the kernel is the interaction of two Gaussian charge + clouds, so its transform needs no new parameter), or +- `sog`: $\hat K(\mathbf{k}) = \sum_m a_m (2\pi)^{3/2} \sigma_m^3\, e^{-\sigma_m^2 k^2/2}$, + the term-by-term transform of the Gaussians; with the default amplitudes the + coefficient simplifies to $4\pi\ln b\, \sigma_m^2$ and the sum approaches the + Coulomb transform $4\pi/k^2$, + +is summed over the wave numbers $\mathbf{k} = 2\pi B^{-1} \mathbf{n}$, +$\mathbf{n} \in \mathbb{Z}^3 \setminus \{\mathbf{0}\}$, of each cell: + +```math +E = k_e \left[\frac{1}{2V}\sum_{\mathbf{k}\neq\mathbf{0}} \sum_{c=1}^{C} +\hat K(\mathbf{k})\, |\hat\rho_c(\mathbf{k})|^2 +- \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \right], +\qquad \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i\mathbf{k}\cdot\mathbf{r}_i}, +``` + +with the cell volume $V$ and the kernel value at contact $K(0)$. The +wave-number cutoff follows the kernel parameters through +{ref}`kspace_tolerance `: +the grid ends where the weight of the slowest-decaying mode falls below it, +so the truncation is exponentially convergent. The $\mathbf{k}=\mathbf{0}$ +mode is dropped (the tin-foil boundary convention), so every latent channel +must be neutral: use a `neutral` constraint, or a `charge_spin`/`fparam` +constraint holding a neutral total charge. + +The k grid is enumerated outside the compiled region and enters it as inputs +of dynamic length, like the pairs of the non-periodic evaluation; a model +with `periodic: true` still serves non-periodic batches. Forces come from the +positions the phases are evaluated at, and the virial from a strain the +energy is differentiated against, so both hold the complete k-space term. + +The periodic evaluation is supported for training, compiled training and the +Python inference interface. Freezing/export, compression, the C++ interface +and LAMMPS do not serve a periodic long-range module yet: they reject it. + ### Total charge {ref}`charge_constraint ` diff --git a/periodic_long_range_design.md b/periodic_long_range_design.md new file mode 100644 index 0000000000..e9c21d4876 --- /dev/null +++ b/periodic_long_range_design.md @@ -0,0 +1,296 @@ +# 周期性 LES / SOG 长程能量:公式与 k 点实现方案 + +> 分支:`dpa4c+period` +> 目标:给 `deepmd/dpmodel/long_range`(latent_charge 模块)的 LES 和 SOG 两种 +> kernel 增加周期边界支持,第一版只覆盖 **Python 训练路径(graph lower)**; +> 压缩 / CUDA direct-sum / C++ / LAMMPS 路径显式拒绝。 +> +> 设计取向:**纯倒易空间(all-reciprocal)方案**。SOG 高斯核的傅里叶变换是 +> 带本征高斯阻尼的整函数,实空间 erfc 残差恒为零,因此不需要 Ewald β 分离、 +> 不需要实空间项、不需要动 edge 图。这与 `deepmd-kit-charge` 仓库 +> (`sog_model._compute_periodic`)的路线一致,但修正其三处一致性问题: +> ① k 空间振幅缺 `(2π)^{3/2}σ³` 因子;② 周期分支丢失 LES;③ 自能用截断 k 和近似。 + +## 1. 记号 + +- 原子数 $N$,latent 通道数 $C$,电荷 $q_i^c$($c=1..C$,单位 $e$) +- $k_e = 14.3996454784255$ eV·Å(`ElectrostaticConvertion`,`source/lib/include/ewald.h:19`) +- box 矩阵 $B$(行向量为 $\mathbf{a}_1,\mathbf{a}_2,\mathbf{a}_3$),体积 $V = |\det B| > 0$ +- 倒易格点:$\mathbf{k}(\mathbf{n}) = 2\pi B^{-T}\mathbf{n}$,$\mathbf{n}\in\mathbb{Z}^3$ +- 每通道结构因子:$\hat\rho_c(\mathbf{k}) = \sum_i q_i^c\, e^{i\mathbf{k}\cdot\mathbf{r}_i}$ + (整体平移只给 $\hat\rho$ 乘一个模 1 相位,$|\hat\rho|^2$ 不变,坐标原点任意) +- SOG 参数:宽度几何序列 $\sigma_m = \sigma_0 b^m$($m=0..M-1$), + **修正后振幅** $a_m = \sqrt{2/\pi}\,\ln b/\sigma_m$ + (本仓库 `16f40a2b`;旧常数振幅 $4\pi\ln b$ 给出对数形核,已废弃) + +## 2. 非周期公式(现有实现,不变) + +$$E_{LR} = k_e \sum_{i |n_d| <= k_max*|a_d|/(2pi) (Cauchy-Schwarz) + # 倒易基矢(列):g_d = 2*pi*B^{-1}[:, d],使 a_d · g_e = 2pi delta_de + # 注意是 inv(B) 而非 inv(B).T(B 非对称时 transpose 是错的) + rec = 2*pi * inv(box) # rec[:, d] = g_d + n_d = ceil(k_max * norm(box, axis=1) / (2*pi)) # (3,) + # (3) 枚举 + 字典序(确定性顺序!)+ 球形截断 + ks, ns = [], [] + for n0 in range(-n_d[0], n_d[0]+1): # 外层→内层固定序 + for n1 in range(-n_d[1], n_d[1]+1): + for n2 in range(-n_d[2], n_d[2]+1): + if n0 == n1 == n2 == 0: continue + k = rec @ (n0, n1, n2) + if dot(k, k) <= k_max*k_max: keep + return k_vec (nk,3), k2 (nk,) +``` + +- **数量级**:$\sigma_0 = 2.18$ Å、$\varepsilon=10^{-8}$ → $k_{max}\approx 6.2$ Å⁻¹; + 15 Å 立方盒约 $9^3-1 \approx 700$ 个 k 点;$\hat\rho$ 计算量 $O(N_k\cdot N\cdot C)$。 +- 每帧网格不同(box 不同)→ batch 拼接时需要逐帧归属(§5.2)。 +- **compile:ragged 动态轴**(`nk` 成为第四个符号轴,照 `npair` 协议, + flatten/unflatten 为 k_grid 携 `(frame, vec, mask, volume, coord, strain)`)—— + 固定 nk + mask 方案未采用:拓扑协议已原生支持动态轴,零额外复杂度。 + +### 5.2 图内计算(伪代码,全部纯张量运算,无 data-dependent 控制流) + +```python +# 输入:coord (N,3) [leaf], strain (nf,3,3) [全零 leaf], q (N,C) [电荷网络输出], +# k_vec (nk,3), k_frame (nk,), k_mask (nk,), volume (nf,), +# node_frame (N,) +# 应变 leaf 技巧:位力 = -dE/d(strain),不再需解析位力公式 +S = I + strain +coord' = S[node_frame] @ coord # 每节点 (3,3)@(3,) +k' = inv(S)[k_frame]^{-T} @ k_vec # 每 k 点 (3,3)@(3,) +volume' = volume * det(S) +kr = k' @ coord'.T # (nk, N) +same = (node_frame[None, :] == k_frame[:, None]) # (nk, N) 帧归属 +rho_r = (cos(kr) * same) @ q # (nk, C) segment sum +rho_i = (sin(kr) * same) @ q # (nk, C) +s_sq = (rho_r**2 + rho_i**2).sum(-1) # (nk,) + +# 模式权重(alpha 或 sigma 可训练,梯度照常;按 kernel 类型二选一) +if kernel_type == "les": + K_hat = 4*pi * exp(-k2' / (4*alpha**2)) / k2' # (nk,) + K0 = 2*alpha / sqrt(pi) +else: # "sog" + K_hat = sum_m( amp[m]*(2pi)**1.5*sigma[m]**3 # (nk,),模式循环展开: + * exp(-0.5*sigma[m]**2*k2') ) # 避免 (nk,M) 中间张量 + K0 = amp.sum() # (inductor 在其二次反传 + # 里生成 NaN;LES 不受影响) + +# 能量(自能项只作用于有真实 k 点的帧——周期帧) +E_rec = k_e * segment_sum(k_mask * K_hat * s_sq / (2*volume'[k_frame]), + k_frame, nf) +E_self = -0.5 * k_e * K0 * (q**2).sum(-1) * is_periodic_frame[node_frame] +``` + +- 精度策略(对齐仓库惯例):`rho_r/rho_i` 与 k 求和用 float64 累加, + k 点按 §5.1 固定字典序 → 完全确定性。 +- 力:对 `coord` leaf autograd(框架内 `fit_output_to_model_output_graph` + 的 leaf 列表扩展,与 `edge_vec`/`pair_vec` 并列)。 +- 位力:§3.5 解析式,tensor 实现,`kappa(k)` 对可训练 $\alpha,\sigma$ 可导。 + +### 5.3 图 lower 签名影响(本方案真正动架构的一处) + +现有 graph lower **没有坐标输入**。$e^{i\mathbf{k}\cdot\mathbf{r}_i}$ 需要 +绝对坐标相位。最终实现采用:**`coord` 和零应变 `strain` 都作为 k_grid 拓扑 +的 flatten 输入**(`(frame, vec, mask, volume, coord, strain)`),在 +`forward_common_lower_graph` 里 `.detach().requires_grad_(True)` 重 leaf—— +反向传播必须根植在图输入上(make_fx 内部创建的 leaf 在 inductor 下产出 +NaN)。力由 coord leaf 的梯度直接给出;位力 $W = -\partial E/\partial\varepsilon$ +由 strain leaf 的梯度给出(模块内以 $r'=r(I+\varepsilon)^T$、 +$k'=k(I+\varepsilon)^{-1}$、$V'=V\det(I+\varepsilon)$ 求值)。 +原子位力约定:$r\otimes F$ 部分逐节点,帧余量均摊。 + +## 6. 验证计划 + +1. **M1 对账脚本**(`verify_periodic_lr.py`,独立 torch float64,不接框架)—— + 已实现并全部通过: + - **LES**:k 空间公式 vs 独立双参数 Ewald 分解(实空间 `[erf(αr)−erf(βr)]/r` + 镜像和 + k 空间 β 项 + 自能),diff 1e-9; + - **SOG 分模式 Poisson 恒等式**:每个高斯模式(σ ≤ 10 Å)实空间镜像和 vs + k 空间,diff ~1e-10~1e-11(宽模式以同样的数学自然延伸); + - **SOG 复合对账**:全 k 空间 vs 窄模式(σ≤4 Å)实空间镜像和 + 宽模式 k 空间, + diff 1e-8——注意**不要**用 `[K_sog−erf(βr)/r]` 做参考:该差核在 + $r \gtrsim \sigma_{max}\approx 610$ Å 处残余 $-1/r$ 尾巴,壳层求和无法收敛; + - 力(autograd vs FD)~1e-8,位力(解析式 §3.5 vs 应变 FD)~1e-7; + - k_max 收敛表:两种 kernel 均指数收敛; + - 纯 LES 可用 `deepmd/tf/infer/ewald_recp.py` 交叉验证(其 β = 本方案 α, + 只含倒易项)——需要 C++ TF 算子库构建,未在 M1 中执行。 +2. **M2 框架接入**(已完成):`LatentChargeLongRange` 周期分支 + + k 网格 eager 构建 + `coord`/`strain` 双 leaf;`direct_sum_eligible()` + 不变(压缩路径在 `enable_compression` 拒绝周期模块);导出/冻结拒绝。 + 单测:`source/tests/pt_expt/model/test_long_range.py`(周期组,14 个) + 与 `source/tests/common/dpmodel/test_long_range.py`(k 网格/kernel 原语)。 + 实测(12 原子三斜盒):能量分解 vs M1 参考 2e-7(LES)/1e-6(SOG), + 力 FD 3.6e-7/1.5e-8,位力 FD 3e-7/9.8e-9,编译 == eager 至 1e-13, + Hessian(图路径)vs FD 2e-6。 +3. 拒绝面:周期 + 压缩/CUDA/LAMMPS/非中性通道 → 显式报错。 + +## 8. 实现踩坑记录 + +- **倒易格点约定**:box 行向量是晶格向量 $\mathbf{a}_d$ 时,倒易格点 + $\mathbf{k} = 2\pi\,\mathbf{B}^{-1}\mathbf{n}$(**没有转置**)。写错成 + $2\pi B^{-T}\mathbf{n}$ 在正交/立方盒上碰巧正确,三斜盒上 k 点偏离格点 + 达 0.5·$2\pi/|a|$,造成 0.5 eV 量级的能量错误。验证方法: + `B @ rec / (2*pi) == I`(必须逐元素成立)。 +- **SOG 的 (nk, M) 中间张量**:inductor 对该广播的二次反传生成 NaN + (make_fx 单独正确、LES 不受影响);按模式 Python 循环展开为 (nk,) 项 + 求和后正常。 +- **内部创建的 autograd leaf**:在 make_fx 图内 `new_zeros().requires_grad_()` + 的叶子(strain)在 inductor 下梯度为 NaN;叶子必须作为 flatten 输入传入 + 再在图内 re-leaf(`.detach().requires_grad_()`),与 edge_vec 同模式。 +- **自能项的帧归属**:$-\frac{1}{2}K(0)\sum q^2$ 只作用于有真实 k 点的帧 + (周期帧);非周期批在周期模型上跑时 k_grid 全 guard,自能必须跳过, + 否则能量错 $O(K(0)\sum q^2)$。 +- **FD 位力应变的约定**:坐标与 box 必须同族变形($r\to(I+\varepsilon)r$, + $B\to B(I+\varepsilon)^T$),混用(`d @ box`)破坏旋转不变性。 +- **近截断原子对的 FD Hessian**:原子跨越 rcut 时力有拐点,FD Hessian 在 + 任意步长下都停在 O(0.1) 误差——测试体系必须让所有原子远离 cutoff 边界 + (紧致簇 + 大盒)。 + +## 7. 开放项(不阻塞 v1) + +- 非中性通道的背景自能修正(多阻尼宽度下需选定 regularization); +- 真空表面项配置开关; +- box 在大范围内变化的训练集的 ragged k 轴(v2); +- 压缩路径 / CUDA 周期 direct-sum / LAMMPS 周期部署; +- Hessian 与周期 k 项的兼容性检查。 diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 108fd7082c..c9b52f0788 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -19,6 +19,7 @@ SOGKernel, ) from deepmd.dpmodel.utils.neighbor_graph import ( + build_k_grid, build_pair_graph, build_topologies, pair_force_virial, @@ -199,3 +200,120 @@ def test_long_range_round_trip() -> None: assert restored.serialize()["charge_constraint"] == record["charge_constraint"] assert restored.required_topologies() == ("all_pairs",) assert not restored.uses_charge_spin() + + +# --------------------------------------------------------------------------- +# Periodic (k-space) primitives +# --------------------------------------------------------------------------- + + +def _spherical_ft( + kernel: str, k: float, n_r: int = 20000, r_max: float = 60.0 +) -> float: + """Isotropic Fourier transform of the kernel by direct radial quadrature.""" + r = np.linspace(1e-9, r_max, n_r) + if kernel == "les": + k_r = erf(0.8 * r) / r + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_r = np.sum( + amplitude[None, :] * np.exp(-(r[:, None] ** 2) / (2 * width**2)), -1 + ) + return float(np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(k * r / np.pi), r)) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_kernel_kspace_matches_radial_ft(kernel: str) -> None: + """``kernel.kspace`` is the analytic Fourier transform of the real-space kernel.""" + k_sq = np.array([0.5, 1.0, 2.0, 5.0]) ** 2 + if kernel == "les": + mod = LESKernel(alpha=0.8) + else: + mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2]) + got = mod.kspace(k_sq) + for kk, value in zip(np.sqrt(k_sq), got, strict=True): + if kernel == "les": + # erf(alpha r)/r is not integrable; quadrature of the integrable + # complement erfc(alpha r)/r = 1/r - erf(alpha r)/r converges, and + # K_hat(k) = 4 pi / k^2 - FT(erfc(alpha r)/r). + r = np.linspace(1e-9, 60.0, 20000) + k_r = (1.0 - erf(0.8 * r)) / r # erfc(0.8 r)/r + expected = 4.0 * np.pi / kk**2 - float( + np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(kk * r / np.pi), r) + ) + else: + expected = _spherical_ft(kernel, kk) + np.testing.assert_allclose(value, expected, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_kernel_contact_value_matches_real_space_limit(kernel: str) -> None: + """``kernel.contact_value`` is K(r) at contact.""" + if kernel == "les": + mod = LESKernel(alpha=0.8) + else: + mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2]) + r_sq = np.array([1e-12]) + np.testing.assert_allclose( + mod.contact_value(r_sq), mod(r_sq)[0], rtol=1e-6, atol=1e-9 + ) + + +def test_build_k_grid_orthogonality_and_guards() -> None: + """The k points satisfy k . a_d = 2 pi n_d, are ordered and carry guards.""" + box = np.array([[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]])[None] + grid = build_k_grid(box, 2.0, np.array([4]), coord=np.zeros((4, 3))) + # guard points are two masked (1e-6, 0, 0) entries + assert grid.mask[-2:].tolist() == [False, False] + assert grid.mask[:-2].all() + n_integer = grid.vec[:-2] @ box[0].T / (2 * np.pi) + np.testing.assert_allclose(n_integer, np.round(n_integer), atol=1e-10) + k_sq = np.sum(grid.vec[:-2] ** 2, axis=-1) + assert (k_sq > 0).all() and (k_sq <= 4.0 + 1e-12).all() + np.testing.assert_allclose(grid.volume, np.abs(np.linalg.det(box[0]))) + # every integer point inside the sphere is present exactly once + ns = np.round(n_integer).astype(int) + assert len({tuple(row) for row in ns}) == ns.shape[0] + rec = 2 * np.pi * np.linalg.inv(box[0]) + n_d = np.ceil(2.0 * np.linalg.norm(box[0], axis=1) / (2 * np.pi)).astype(int) + count = 0 + for n0 in range(-n_d[0], n_d[0] + 1): + for n1 in range(-n_d[1], n_d[1] + 1): + for n2 in range(-n_d[2], n_d[2] + 1): + if (n0, n1, n2) != (0, 0, 0): + k = rec @ np.array([n0, n1, n2], dtype=float) + count += int(np.sum(k * k) <= 4.0) + assert count == ns.shape[0] + + +def test_build_topologies_periodic_branching() -> None: + """A cell switches the pair topology to guards and the k grid on, and back.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + K_GRID, + ) + + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + n_node = np.array([4]) + box = np.eye(3)[None] * 10.0 + names = (ALL_PAIRS, K_GRID) + params = {K_GRID: {"k_max": 2.0}} + with_box = build_topologies( + names, coord, n_node, box, atype=atype, topology_params=params + ) + assert with_box[ALL_PAIRS].index.shape[1] == 2 # guards only + assert not with_box[ALL_PAIRS].mask.any() + assert with_box[K_GRID].mask.sum() > 0 + without_box = build_topologies( + names, coord, n_node, None, atype=atype, topology_params=params + ) + assert without_box[ALL_PAIRS].mask.any() + assert not without_box[K_GRID].mask.any() # guards only + # a pair-only declaration still rejects a cell loudly + with pytest.raises(NotImplementedError, match="non-periodic"): + build_topologies((ALL_PAIRS,), coord, n_node, box, atype=atype) + # the k grid needs its build parameter + with pytest.raises(ValueError, match="k_max"): + build_topologies(names, coord, n_node, box, atype=atype) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 02f5ed6587..8c1c7440e9 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -878,3 +878,254 @@ def test_pair_graph_skips_padding() -> None: real = pairs.index[:, pairs.mask] assert real.tolist() == [[0, 3], [2, 4]] assert (~pairs.mask).sum() == 2 + + +# --------------------------------------------------------------------------- +# Periodic (k-space) evaluation +# --------------------------------------------------------------------------- + + +def _periodic_model(kernel: str, **kwargs: object) -> torch.nn.Module: + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + return get_model(config).to(env.DEVICE).eval() + + +def _periodic_case() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A compact cluster (all atoms within rcut of each other) in a large triclinic cell.""" + box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]]) + rng = np.random.default_rng(7) + coord = 3.0 + rng.uniform(0.0, 1.4, (6, 3)) + atype = np.array([0, 1, 0, 1, 1, 0], dtype=np.int64) + return coord[None], atype[None], box[None] + + +def _kspace_reference( + coord: np.ndarray, charge: np.ndarray, box: np.ndarray, kernel: str +) -> float: + """Independent k-space evaluation of the periodic long-range energy, in eV. + + E = k_e [ 1/(2V) sum_{k != 0} K_hat(k) |rho(k)|^2 - K(0)/2 sum_i |q_i|^2 ]. + """ + rec = 2.0 * np.pi * np.linalg.inv(box) # columns g_d: k(n) = rec @ n + # the module's default kspace_tolerance=1e-8 sets the truncation + if kernel == "les": + k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + else: + k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=float) for n in n_d], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T # (nk, N) + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge # (nk, C) + s_sq = np.abs(rho) ** 2 + s_sq = s_sq.sum(axis=-1) + if kernel == "les": + k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2 + k0 = 2.0 * 0.8 / np.sqrt(np.pi) + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_hat = np.sum( + amplitude[None, :] + * (2.0 * np.pi) ** 1.5 + * width[None, :] ** 3 + * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2), + axis=-1, + ) + k0 = amplitude.sum() + volume = abs(np.linalg.det(box)) + return float( + COULOMB_CONSTANT + * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()) + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_energy_matches_kspace_reference(kernel: str) -> None: + """With fixed charges the periodic energy is the documented k-space sum.""" + model = _periodic_model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord, atype, box = _periodic_case() + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference(coord[0], charge, box[0], kernel) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) + np.testing.assert_allclose( + out["atom_energy"].sum().item(), expected, rtol=1e-9, atol=1e-9 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_force_virial_hessian_match_finite_differences(kernel: str) -> None: + """Force, strain virial and Hessian derive from the complete periodic energy.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + out = model(coord_t, atype_t, box=box_t) + + def energy(c: np.ndarray, b: np.ndarray) -> float: + return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item() + + h = 1e-5 + force_fd = np.zeros((6, 3)) + for i in range(6): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus, box) - energy(minus, box)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-6 + ) + + # The virial is -dE/d(strain) at zero strain, for r -> (I + eps) r and + # B -> B (I + eps)^T (the strain family of the k-space strain leaf). + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain).T + minus = coord @ (np.eye(3) - strain).T + box_plus = box @ (np.eye(3) + strain).T + box_minus = box @ (np.eye(3) - strain).T + virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / ( + 2 * h + ) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-5 + ) + + hessian_model = _periodic_model(kernel) + hessian_model.load_state_dict(model.state_dict()) + hessian_model.enable_hessian() + hessian = ( + hessian_model(coord_t, atype_t, box=box_t)["hessian"].detach().cpu().numpy() + ) + hessian = hessian.reshape(18, 18) + hessian_fd = np.zeros((18, 18)) + for kk in range(18): + plus, minus = coord.copy(), coord.copy() + plus.reshape(-1)[kk] += h + minus.reshape(-1)[kk] -= h + f_plus = ( + model(_tensor(plus), atype_t, box=box_t)["force"].detach().cpu().numpy() + ) + f_minus = ( + model(_tensor(minus), atype_t, box=box_t)["force"].detach().cpu().numpy() + ) + hessian_fd[:, kk] = -(f_plus - f_minus).reshape(-1) / (2 * h) + np.testing.assert_allclose(hessian, hessian_fd, atol=1e-4) + np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the periodic model like the eager one.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + # the backward reaches the charge network and the kernel through the + # compiled graph + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_model_serves_non_periodic_batches(kernel: str) -> None: + """A periodic-enabled model keeps its non-periodic evaluation.""" + model = _periodic_model(kernel) + plain = _model(kernel) + plain.load_state_dict(model.state_dict()) + coord, atype = _batch() + periodic_enabled = model(coord, atype) + plain_out = plain(coord, atype) + torch.testing.assert_close(periodic_enabled["energy"], plain_out["energy"]) + torch.testing.assert_close(periodic_enabled["force"], plain_out["force"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_frames_are_independent(kernel: str) -> None: + """Each frame's k grid and energy are independent of the rest of the batch.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.1], axis=0) + batched = model(_tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2)) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_serialization_round_trip(kernel: str) -> None: + """The record keeps the periodic flag and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + model = _periodic_model(kernel) + assert model.atomic_model.required_topologies() == ("all_pairs", "k_grid") + record = model.serialize() + assert record["long_range"]["periodic"] is True + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.required_topologies() == ("all_pairs", "k_grid") + coord, atype, box = _periodic_case() + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common( + coord, + atype, + box, + neighbor_graph_method="dense", + ) + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_compression_is_rejected(kernel: str) -> None: + """The compressed path does not serve the k-space evaluation.""" + config = _config(kernel) + config["descriptor"]["precision"] = "float32" + config["fitting_net"]["precision"] = "float32" + config["long_range"]["charge_net"]["precision"] = "float32" + config["long_range"]["periodic"] = True + model = get_model(config).to(env.DEVICE).eval() + with pytest.raises(NotImplementedError, match="reciprocal-space"): + model.atomic_model.enable_compression(0.5) From f6c8396aaa3eb374ff875c66201406fff4306369 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Wed, 21 Oct 2026 11:29:14 +0000 Subject: [PATCH 22/30] update long range --- source/tests/pt_expt/model/test_long_range.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 8c1c7440e9..9a85c5af01 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -454,8 +454,24 @@ def test_parameter_gradients_and_freezing(kernel: str) -> None: loss = out["energy"].square().sum() + out["force"].square().sum() loss.backward() long_range = model.atomic_model.long_range + # The charge constraint fixes the uniform offset of the constrained + # channel, so its entry in the output bias of the charge net is redundant + # by construction and carries an exactly zero gradient. + output_bias = long_range.charge_net.nets._module_networks["0"].layers[-1].b + constrained = ( + long_range.charge_constraint["channel"] + if long_range.charge_constraint is not None + else None + ) for name, parameter in long_range.named_parameters(): - assert parameter.grad is not None and parameter.grad.abs().max() > 0, name + assert parameter.grad is not None, name + grad = parameter.grad.abs() + if parameter is output_bias and constrained is not None: + grad = grad[[c for c in range(grad.shape[0]) if c != constrained]] + if grad.numel() == 0: + # every channel is constrained + continue + assert grad.max() > 0, name config = _config(kernel) config["long_range"]["kernel"]["trainable"] = False From 8fc95cdaeb70c2b547acc52615ed79fceb6f9560 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Thu, 22 Oct 2026 08:19:19 +0000 Subject: [PATCH 23/30] feat(long-range): add fixed k-grid mode and decoupled cutoff option - k_grid_mode: "fixed" shares one integer k grid over the frames of a periodic batch, sized by the largest cell of the data statistics or by k_grid_bounds, and masks the points a frame does not reach; a frame outgrowing the grid is truncated to it with a warning - decouple_k_grid freezes the wave-number cutoff at the initial kernel parameters, so training the kernel no longer moves the k-point selection (the kernel values still follow the live parameters); the frozen cutoff is serialized with the model --- AGENTS.md | 2 +- deepmd/dpmodel/long_range/latent_charge.py | 141 ++++++++- .../dpmodel/utils/neighbor_graph/__init__.py | 6 + .../utils/neighbor_graph/pair_graph.py | 165 ++++++++++- deepmd/utils/argcheck.py | 45 +++ periodic_long_range_design.md | 20 +- .../tests/common/dpmodel/test_long_range.py | 240 ++++++++++++++++ source/tests/pt_expt/model/test_long_range.py | 270 ++++++++++++++++++ 8 files changed, 880 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index b736c6159e..7e21751d7c 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 -- **Long-range latent-charge module**: a standard energy model (DPA4C descriptor) takes a `long_range` section (`type: latent_charge`, kernels `les`/`sog`, optional `charge_constraint`) that adds an interacting-charge energy evaluated on complete systems — directly over all pairs for non-periodic frames, and in reciprocal space for periodic frames when `periodic: true` (k grid enumerated eagerly per batch from the kernel's `kspace_cutoff`, tin-foil boundary, neutral channels required). Periodic evaluation supports eager and compiled training plus the Python interface; freeze/export, compression, C++ and LAMMPS reject periodic long-range models. Tests: `pytest source/tests/pt_expt/model/test_long_range.py`, `source/tests/common/dpmodel/test_long_range.py`. +- **Long-range latent-charge module**: a standard energy model (DPA4C descriptor) takes a `long_range` section (`type: latent_charge`, kernels `les`/`sog`, optional `charge_constraint`) that adds an interacting-charge energy evaluated on complete systems — directly over all pairs for non-periodic frames, and in reciprocal space for periodic frames when `periodic: true`. The k grid has two layouts selected by `k_grid_mode`: `"ragged"` (default; k grid enumerated eagerly per batch from the kernel's `kspace_cutoff`, tin-foil boundary, neutral channels required) and `"fixed"` (one integer grid shared over frames, sized by the largest cell of the data statistics or by `k_grid_bounds`, with a mask selecting each frame's k points; out-of-grid cells are truncated to the grid with a warning). The wave-number cutoff is re-derived from the live (trainable) kernel every batch unless `decouple_k_grid: true` freezes it at the initial parameters (the kernel values still follow the live parameters; the frozen cutoff is serialized). Periodic evaluation supports eager and compiled training plus the Python interface; freeze/export, compression, C++ and LAMMPS reject periodic long-range models. Tests: `pytest source/tests/pt_expt/model/test_long_range.py`, `source/tests/common/dpmodel/test_long_range.py`. ### JAX Backend diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index c6c00f124a..fa5d60d95c 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -25,7 +25,9 @@ K_GRID, KGrid, build_topologies, + enumerate_k_grid_bounds, frame_id_from_n_node, + k_grid_bounds_of_cells, segment_sum, ) from deepmd.dpmodel.utils.seed import ( @@ -116,6 +118,28 @@ class LatentChargeLongRange(NativeOP, BaseLongRange): kspace_tolerance The relative weight of the slowest-decaying kernel mode at the wave-number cutoff of the k grid of a periodic batch. + k_grid_mode + How the k grid of a periodic batch is laid out: ``"ragged"`` + enumerates the wave numbers of every frame eagerly, so the k axis is + a dynamic dimension; ``"fixed"`` shares one integer grid over the + frames — sized by the largest cell of the data statistics, or by + ``k_grid_bounds`` — evaluates every frame on it by a static-shaped + transform, and masks the points a frame does not reach. The fixed + grid is hard: a frame needing an integer index beyond it (a larger + cell, or a kernel that moved the cutoff out) is served by the grid + alone, so its k-sum is truncated to the grid with a warning. + k_grid_bounds + The per-axis integer bounds :math:`n_d^{max}` of the fixed grid, in + place of the data statistics; only meaningful with + ``k_grid_mode="fixed"``. + decouple_k_grid + Decouple the k-point selection of a periodic batch from the kernel + parameters: the wave-number cutoff is derived once from the initial + parameters and kept for the life of the model, so training the + kernel no longer moves the k grid. The kernel values the grid is + evaluated with still follow the live parameters. The frozen cutoff + is serialized with the model, so a restarted run keeps the grid the + model was built with. Requires ``periodic=True``. """ def __init__( @@ -125,6 +149,9 @@ def __init__( charge_constraint: dict | None = None, periodic: bool = False, kspace_tolerance: float = 1e-8, + k_grid_mode: str = "ragged", + k_grid_bounds: list[int] | None = None, + decouple_k_grid: bool = False, ) -> None: if charge_net.vacuum_ref: raise ValueError( @@ -148,6 +175,38 @@ def __init__( # The frame charge/spin condition a compressed model is specialized # to; see :meth:`fold_charge_state`. self.charge_state = None + self.k_grid_mode = str(k_grid_mode) + if self.k_grid_mode not in ("ragged", "fixed"): + raise ValueError( + f"`k_grid_mode` must be 'ragged' or 'fixed', got {self.k_grid_mode!r}" + ) + if self.k_grid_mode == "fixed" and not self.periodic: + raise ValueError("`k_grid_mode='fixed'` requires `periodic=True`") + if k_grid_bounds is not None and self.k_grid_mode != "fixed": + raise ValueError( + "`k_grid_bounds` is only meaningful with `k_grid_mode='fixed'`" + ) + self.decouple_k_grid = bool(decouple_k_grid) + if self.decouple_k_grid and not self.periodic: + raise ValueError("`decouple_k_grid` requires `periodic=True`") + # A decoupled grid reads the kernel once, at construction: the k-point + # selection stays at the initial kernel however training moves it, + # while the kernel values in the graph still use the live parameters. + self.frozen_k_max = ( + float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + if self.decouple_k_grid + else None + ) + # The shared integer grid of the fixed mode, empty until the data + # statistics size it (see :meth:`compute_input_stats`) — an array even + # then, so a backend keeping arrays in its state dict sees a stable + # set of keys and a checkpoint restores the sized grid. + if k_grid_bounds is not None: + self.k_grid_n = enumerate_k_grid_bounds(k_grid_bounds) + elif self.k_grid_mode == "fixed": + self.k_grid_n = np.zeros((0, 3), dtype=np.int64) + else: + self.k_grid_n = None @staticmethod def _normalize_constraint( @@ -237,8 +296,22 @@ def from_config( config.get("charge_constraint"), periodic=config.get("periodic", False), kspace_tolerance=config.get("kspace_tolerance", 1e-8), + k_grid_mode=config.get("k_grid_mode", "ragged"), + k_grid_bounds=config.get("k_grid_bounds"), + decouple_k_grid=config.get("decouple_k_grid", False), ) + def _k_max(self) -> float: + """The wave-number cutoff selecting the k points of a periodic batch. + + A coupled grid reads the current kernel parameters, so a trained + kernel re-derives its grid every batch; a decoupled grid keeps the + cutoff of the initial parameters. + """ + if self.decouple_k_grid: + return self.frozen_k_max + return float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + def build_topologies( self, coord: Array, @@ -253,11 +326,22 @@ def build_topologies( The k grid of a periodic batch is enumerated from the cells with the wave-number cutoff of the kernel (the trainable parameters are read at - build time, so a drifting kernel re-derives its grid every batch). + build time, so a drifting kernel re-derives its grid every batch, + unless the grid is decoupled from them by ``decouple_k_grid``). In + the fixed grid mode the batch is evaluated on the shared integer grid + instead, whose mask the same cutoff derives. """ params = {} if self.periodic: - params[K_GRID] = {"k_max": self.kernel.kspace_cutoff(self.kspace_tolerance)} + params[K_GRID] = {"k_max": self._k_max()} + if self.k_grid_mode == "fixed": + if self.k_grid_n is None or self.k_grid_n.shape[0] == 0: + raise ValueError( + "the fixed k grid is not sized yet: run the data " + "statistics (`compute_input_stats`) on periodic " + "frames, or configure `k_grid_bounds`" + ) + params[K_GRID]["n_grid"] = self.k_grid_n return build_topologies( self.required_topologies(), coord, @@ -539,10 +623,47 @@ def compute_input_stats( protection: float = 1e-2, stat_file_path: DPPath | None = None, ) -> None: + if self.k_grid_mode == "fixed" and ( + self.k_grid_n is None or self.k_grid_n.shape[0] == 0 + ): + # Materialize the samples once for the grid sizing and the + # networks; a lazy sampler is shared by the callers. + sampled = merged() if callable(merged) else merged + self._size_k_grid(sampled) + merged = sampled self.charge_net.compute_input_stats( merged, protection=protection, stat_file_path=stat_file_path ) + def _size_k_grid(self, sampled: list[dict]) -> None: + r"""Size the fixed integer k grid by the largest sampled cell. + + Every sampled frame needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` per + axis at the wave-number cutoff of the kernel; the grid takes the + elementwise maximum of the bounds over the frames. The kernel + parameters are read now, so a grid is sized for the kernel at + statistics time — frames that later outgrow it (larger cells, or a + kernel moving its cutoff out) are truncated to the grid. + """ + k_max = self._k_max() + bounds = None + for frame in sampled: + box = frame.get("box") + if box is None: + continue + box = np.asarray(to_numpy_array(box), dtype=np.float64) + if box.size == 0: + continue + frame_max = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + bounds = frame_max if bounds is None else np.maximum(bounds, frame_max) + if bounds is None: + raise ValueError( + "the fixed k grid takes its size from the data statistics, " + "but the sampled frames hold no simulation cell; provide " + "periodic training data or configure `k_grid_bounds`" + ) + self.k_grid_n = enumerate_k_grid_bounds(bounds) + def change_type_map(self, type_map: list[str]) -> None: self.charge_net.change_type_map(type_map=type_map) @@ -561,6 +682,13 @@ def serialize(self) -> dict: if self.periodic: data["periodic"] = True data["kspace_tolerance"] = self.kspace_tolerance + if self.k_grid_mode != "ragged": + data["k_grid_mode"] = self.k_grid_mode + if self.k_grid_n is not None and self.k_grid_n.shape[0] > 0: + data["k_grid_n"] = to_numpy_array(self.k_grid_n).tolist() + if self.decouple_k_grid: + data["decouple_k_grid"] = True + data["frozen_k_max"] = self.frozen_k_max if self.charge_state is not None: data["charge_state"] = to_numpy_array(self.charge_state).tolist() return data @@ -575,7 +703,16 @@ def deserialize(cls, data: dict) -> "LatentChargeLongRange": data["charge_constraint"], periodic=data.get("periodic", False), kspace_tolerance=data.get("kspace_tolerance", 1e-8), + k_grid_mode=data.get("k_grid_mode", "ragged"), + decouple_k_grid=data.get("decouple_k_grid", False), ) + if data.get("frozen_k_max") is not None: + # The constructor derives the frozen cutoff from the (possibly + # trained) kernel; the record holds the value the grid was + # actually built with. + obj.frozen_k_max = float(data["frozen_k_max"]) + if data.get("k_grid_n") is not None: + obj.k_grid_n = np.asarray(data["k_grid_n"], dtype=np.int64).reshape(-1, 3) if data.get("charge_state") is not None: obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) return obj diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index bf503530de..8bcecaecd6 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -61,11 +61,14 @@ K_GRID, KGrid, PairGraph, + build_fixed_k_grid, build_k_grid, build_pair_graph, build_topologies, check_complete_frames, + enumerate_k_grid_bounds, flatten_topologies, + k_grid_bounds_of_cells, topology_flat_fields, unflatten_topologies, unordered_pair_index, @@ -96,6 +99,7 @@ "attach_edge_csr", "build_angle_index", "build_edge_csr", + "build_fixed_k_grid", "build_k_grid", "build_neighbor_graph", "build_neighbor_graph_ase", @@ -107,12 +111,14 @@ "compact_nodes", "edge_env_mat", "edge_force_virial", + "enumerate_k_grid_bounds", "expand_node_values", "flatten_topologies", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", "graph_from_dense_quartet", + "k_grid_bounds_of_cells", "neighbor_graph_from_ijs", "node_ownership_mask", "node_validity_mask", diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index e98cc10d51..e6be02bfb5 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -23,6 +23,7 @@ ) import math +import warnings from dataclasses import ( dataclass, ) @@ -267,6 +268,155 @@ def build_k_grid( ) +def k_grid_bounds_of_cells(box: Array, k_max: float) -> np.ndarray: + r"""Per-axis integer wave-number bounds the cells need at ``k_max``. + + A wave number :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` with + :math:`|\mathbf{k}| \le k_{max}` needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` + per axis (Cauchy-Schwarz), so the bound of every frame is the ceiling of + the right-hand side. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + + Returns + ------- + np.ndarray + The per-frame, per-axis integer bounds with shape ``(nf, 3)``. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + a_len = np.linalg.norm(cell, axis=-1) + return np.ceil(k_max * a_len / (2.0 * math.pi)).astype(np.int64) + + +def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: + r"""Enumerate the integer wave-number indices within per-axis bounds. + + The grid is the full meshgrid of :math:`\mathbf{n}` with + :math:`|n_d| \le n_d^{max}`, in lexicographic order (the first axis + outermost). The zero vector is left out: its kernel transform is + singular and the tin-foil convention drops the :math:`\mathbf{k} = 0` + mode anyway. + + Parameters + ---------- + bounds + The per-axis bounds :math:`n_d^{max}` with shape ``(3,)``. + + Returns + ------- + np.ndarray + The integer indices with shape ``(nk, 3)``, int64. + """ + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0): + raise ValueError( + f"k-grid bounds must be three non-negative integers, got {bounds}" + ) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=np.int64) for n in bounds], indexing="ij" + ) + n_grid = np.stack(grids, axis=-1).reshape(-1, 3) + return n_grid[np.any(n_grid != 0, axis=-1)] + + +def build_fixed_k_grid( + box: Array, + n_grid: Array, + k_max: float, + n_node: Array, + coord: Array | None = None, +) -> KGrid: + r"""Evaluate a fixed integer k grid on the cells of the periodic frames. + + Every frame shares the integer indices ``n_grid`` (see + :func:`enumerate_k_grid_bounds`), so the k axis has the static length + ``nf * nk``; the wave numbers of a frame are + :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` and the points the frame does + not reach — outside :math:`0 < |\mathbf{k}| \le k_{max}` — are masked out. + The mask is evaluated eagerly with the wave-number cutoff current at + build time, exactly as in :func:`build_k_grid`. The enumeration runs + eagerly beside the ragged one, but no count is data dependent: the grid + comes back without guard points. + + The grid is HARD-FIXED: a frame whose cell would need integer indices + beyond the grid (a larger cell than the grid was sized for, or a kernel + that moved the cutoff out since) is served by the grid alone — its + reciprocal lattice is denser than the grid, so its energy is the + truncated k-sum of the grid, a well-defined approximation whose forces + and virial stay the exact derivatives of that energy. A warning is + issued when this happens. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + n_grid + The shared integer wave-number indices with shape ``(nk, 3)``. + k_max + The wave-number cutoff, in Å⁻¹. + n_node + Per-frame node counts with shape ``(nf,)`` (interface symmetry with + :func:`build_k_grid`; the grid does not read it). + coord + Node positions with shape ``(N, 3)``, stored on the grid as the + geometry context (see :attr:`KGrid.coord`). + + Returns + ------- + KGrid + The wave numbers of every frame, frame-major over the shared grid. + + Raises + ------ + ValueError + Never for an oversized cell (it is truncated, see above); kept for + interface symmetry. + """ + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + if box.ndim == 2: + box = xp.reshape(box, (box.shape[0], 3, 3)) + nf = box.shape[0] + n_grid = np.asarray(to_numpy_array(n_grid), dtype=np.int64).reshape(-1, 3) + nk = n_grid.shape[0] + bounds = np.max(np.abs(n_grid), axis=0) if nk else np.zeros(3, dtype=np.int64) + required = k_grid_bounds_of_cells(box, k_max) + over = np.any(required > bounds[None, :], axis=1) + if np.any(over): + f = int(np.argmax(over)) + warnings.warn( + f"frame {f} needs integer k-grid bounds {required[f].tolist()}, " + f"beyond the fixed grid {bounds.tolist()}: its k-sum is truncated " + f"to the grid (an approximation for this frame); enlarge " + f"`k_grid_bounds` for full accuracy", + stacklevel=2, + ) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + rec = 2.0 * math.pi * np.linalg.inv(cells) # (nf, 3, 3), columns are g_d + k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) + k2 = np.sum(k * k, axis=-1) + mask_np = (k2 > 0.0) & (k2 <= k_max * k_max) # (nf, nk) + vec_np = k.reshape(nf * nk, 3) + frame_np = np.repeat(np.arange(nf, dtype=np.int64), nk) + volume = xp.linalg.det(box) + volume = xp.abs(volume) + return KGrid( + frame=xp.asarray(frame_np, device=device), + vec=xp.astype(xp.asarray(vec_np, device=device), box.dtype), + mask=xp.asarray(mask_np.reshape(-1), device=device), + volume=xp.astype(volume, box.dtype), + coord=coord, + strain=xp.zeros((nf, 3, 3), dtype=box.dtype, device=device), + ) + + def _empty_k_grid(nf: int, like: Array, coord: Array | None) -> KGrid: """A guard-only k grid for a batch without periodic frames.""" xp = array_api_compat.array_namespace(like) @@ -367,7 +517,9 @@ def build_topologies( pair_excl The model-level pair exclusion, or ``None``. topology_params - Per-name build parameters; the k grid takes ``{"k_max": ...}``. + Per-name build parameters; the k grid takes ``{"k_max": ...}`` and, + for a fixed integer grid, additionally ``{"n_grid": ...}`` (see + :func:`build_fixed_k_grid`). Returns ------- @@ -408,9 +560,14 @@ def build_topologies( raise ValueError( f"the {K_GRID!r} topology requires its `k_max` build parameter" ) - topologies[name] = build_k_grid( - box, params["k_max"], n_node, coord=coord - ) + if "n_grid" in params: + topologies[name] = build_fixed_k_grid( + box, params["n_grid"], params["k_max"], n_node, coord=coord + ) + else: + topologies[name] = build_k_grid( + box, params["k_max"], n_node, coord=coord + ) else: raise ValueError( f"unknown topology {name!r}; expected one of ({ALL_PAIRS!r}, {K_GRID!r})" diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 4cc56193f1..79538293ae 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3704,6 +3704,30 @@ def long_range_latent_charge() -> list[Argument]: "The relative weight of the slowest-decaying kernel mode at the " "wave-number cutoff of the k grid of a periodic batch." ) + doc_k_grid_mode = ( + "How the k grid of a periodic batch is laid out: 'ragged' enumerates " + "the wave numbers of every frame eagerly (the k axis is a dynamic " + "dimension); 'fixed' shares one integer grid over the frames, sized " + "by the largest cell of the data statistics or by `k_grid_bounds`, " + "and masks the points a frame does not reach, so the k axis has a " + "static length. The fixed grid is hard: a frame needing an index " + "beyond it (a larger cell, or a kernel that moved the cutoff out) " + "has its k-sum truncated to the grid, with a warning. Requires " + "`periodic: true` for 'fixed'." + ) + doc_k_grid_bounds = ( + "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " + "place of the data statistics; only meaningful with " + "`k_grid_mode: fixed`." + ) + doc_decouple_k_grid = ( + "Decouple the k-point selection of a periodic batch from the kernel " + "parameters: the wave-number cutoff is derived once from the initial " + "parameters and kept for the life of the model, so training the " + "kernel no longer moves the k grid (the kernel values the grid is " + "evaluated with still follow the live parameters). The frozen cutoff " + "is serialized with the model. Requires `periodic: true`." + ) return [ Argument("channels", int, optional=True, default=1, doc=doc_channels), Argument( @@ -3742,6 +3766,27 @@ def long_range_latent_charge() -> list[Argument]: default=1e-8, doc=doc_kspace_tolerance, ), + Argument( + "k_grid_mode", + str, + optional=True, + default="ragged", + doc=doc_k_grid_mode, + ), + Argument( + "k_grid_bounds", + list[int], + optional=True, + default=None, + doc=doc_k_grid_bounds, + ), + Argument( + "decouple_k_grid", + bool, + optional=True, + default=False, + doc=doc_decouple_k_grid, + ), ] diff --git a/periodic_long_range_design.md b/periodic_long_range_design.md index e9c21d4876..1cb7573882 100644 --- a/periodic_long_range_design.md +++ b/periodic_long_range_design.md @@ -186,8 +186,24 @@ def build_k_grid(box, alpha=None, sigma0=None, eps=1e-8): 15 Å 立方盒约 $9^3-1 \approx 700$ 个 k 点;$\hat\rho$ 计算量 $O(N_k\cdot N\cdot C)$。 - 每帧网格不同(box 不同)→ batch 拼接时需要逐帧归属(§5.2)。 - **compile:ragged 动态轴**(`nk` 成为第四个符号轴,照 `npair` 协议, - flatten/unflatten 为 k_grid 携 `(frame, vec, mask, volume, coord, strain)`)—— - 固定 nk + mask 方案未采用:拓扑协议已原生支持动态轴,零额外复杂度。 + flatten/unflatten 为 k_grid 携 `(frame, vec, mask, volume, coord, strain)`)。 +- **固定网格 + mask(已实现,`k_grid_mode: fixed`)**:v1 时未采用,后按需求 + 补上作为可选模式。整数网格 $\mathbf{n} \in [-n_d, n_d]^3$(不含零矢量) + 由数据统计中最大 cell 的逐方向界 $n_d = \lceil k_{max}|a_d|/2\pi \rceil$ + 一次性枚举(`enumerate_k_grid_bounds`,可用 `k_grid_bounds` 配置覆盖), + 存入模型并随其序列化;每个 batch 以静态形状 matmul 求各帧波矢 + $\mathbf{k} = 2\pi B_f^{-1}\mathbf{n}$(k 轴长度固定为 $n_f \times n_k$,无 + guard 点),帧达不到的点($|k| > k_{max}$,$k_{max}$ 取构建时 kernel 参数) + 由 mask 屏蔽(`build_fixed_k_grid`)。网格是硬性的:cell 超出网格容量或 + kernel 漂移使 $k_{max}$ 越界时,该帧仅以网格内 k 点截断求和(倒格点欠采样的 + 近似,力和位力仍是该近似能量的精确导数),并发 warning。零矢量不入网格:LES 的 $\hat K$ 在 + $k=0$ 奇异,masked 的 inf 会在反向传播产生 NaN。开销约为 ragged 球内点数 + 的 ~2 倍(各向同性 cell),换来完全静态的形状。 +- **解耦开关(`decouple_k_grid: true`)**:默认(耦合)每 batch 以当前 kernel + 参数现算 $k_{max}$,kernel 漂移 → 网格跟随,能量定义始终是当前 kernel 的收敛 + 能量;置 true 后 $k_{max}$ 在构造时按初始参数冻结(`frozen_k_max`,随模型 + 序列化),kernel 漂移不再影响 k 点选择,图内 $\hat K(k)$ 仍用活参数求值 + (梯度不受影响)。两种模式共用;fixed 模式下解耦后 mask 也不再随漂移收紧。 ### 5.2 图内计算(伪代码,全部纯张量运算,无 data-dependent 控制流) diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index c9b52f0788..b3c5032558 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -22,6 +22,7 @@ build_k_grid, build_pair_graph, build_topologies, + enumerate_k_grid_bounds, pair_force_virial, unordered_pair_index, ) @@ -317,3 +318,242 @@ def test_build_topologies_periodic_branching() -> None: # the k grid needs its build parameter with pytest.raises(ValueError, match="k_max"): build_topologies(names, coord, n_node, box, atype=atype) + + +# --------------------------------------------------------------------------- +# Fixed k-grid mode +# --------------------------------------------------------------------------- + + +def _periodic_long_range(**kwargs: object) -> LatentChargeLongRange: + charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + return LatentChargeLongRange( + charge_net, LESKernel(alpha=0.3), None, periodic=True, **kwargs + ) + + +def _cells() -> np.ndarray: + """Two cells of different sizes, the second the larger one.""" + return np.array( + [ + [[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]], + [[12.0, 0.5, 0.0], [0.2, 11.0, 0.3], [0.1, 0.4, 10.5]], + ] + ) + + +def test_enumerate_k_grid_bounds() -> None: + """The grid is the lexicographic meshgrid of the bounds, without the zero vector.""" + n_grid = enumerate_k_grid_bounds([2, 1, 0]) + assert n_grid.shape == (5 * 3 * 1 - 1, 3) + assert not np.any(np.all(n_grid == 0, axis=-1)) + expected = [ + (n0, n1, 0) for n0 in range(-2, 3) for n1 in range(-1, 2) if (n0, n1) != (0, 0) + ] + np.testing.assert_array_equal(n_grid, np.asarray(expected, dtype=np.int64)) + with pytest.raises(ValueError, match="bounds"): + enumerate_k_grid_bounds([1, -1, 0]) + + +def test_fixed_k_grid_matches_ragged_enumeration() -> None: + """The masked points of the fixed grid are exactly the ragged k points.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_fixed_k_grid, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + ) + + box = _cells() + k_max = 2.0 + n_grid = enumerate_k_grid_bounds(np.max(k_grid_bounds_of_cells(box, k_max), axis=0)) + nk = n_grid.shape[0] + fixed = build_fixed_k_grid(box, n_grid, k_max, np.array([4, 5])) + # a static, guard-free layout: nf * nk points, frame-major + assert fixed.vec.shape == (2 * nk, 3) + np.testing.assert_array_equal(fixed.frame, np.repeat([0, 1], nk)) + ragged = build_k_grid(box, k_max, np.array([4, 5])) + for frame in range(2): + expected = { + tuple(row) + for row in np.round(ragged.vec[(ragged.frame == frame) & ragged.mask], 8) + } + got = { + tuple(row) + for row in np.round(fixed.vec[(fixed.frame == frame) & fixed.mask], 8) + } + assert got == expected + # the smaller cell masks more points of the shared grid + assert fixed.mask[:nk].sum() < fixed.mask[nk:].sum() + np.testing.assert_allclose(fixed.volume, np.abs(np.linalg.det(box)), rtol=1e-12) + + +def test_fixed_k_grid_truncates_oversized_cells() -> None: + """A cell beyond the grid is served by the grid alone (truncated k-sum), with a warning.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_fixed_k_grid, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + ) + + box = _cells() + k_max = 2.0 + n_grid = enumerate_k_grid_bounds(np.max(k_grid_bounds_of_cells(box, k_max), axis=0)) + nk = n_grid.shape[0] + with pytest.warns(UserWarning, match="truncated"): + bigger = build_fixed_k_grid(box * 2.0, n_grid, k_max, np.array([4, 5])) + # the layout is untouched: static, guard-free, every grid point present + assert bigger.vec.shape == (2 * nk, 3) + np.testing.assert_array_equal(bigger.frame, np.repeat([0, 1], nk)) + # every grid point is still a valid wave number of the larger cell; the + # mask keeps the whole k_max sphere the grid reaches + k2 = np.sum(bigger.vec**2, axis=-1) + np.testing.assert_array_equal(bigger.mask, (k2 > 0) & (k2 <= k_max * k_max)) + # the larger cell activates more of the shared grid than the original did + original = build_fixed_k_grid(box, n_grid, k_max, np.array([4, 5])) + assert bigger.mask.sum() > original.mask.sum() + # a larger wave-number cutoff (a drifting kernel) truncates the same way + with pytest.warns(UserWarning, match="truncated"): + build_fixed_k_grid(box, n_grid, k_max * 2.0, np.array([4, 5])) + + +def test_fixed_k_grid_mode_configuration() -> None: + """The fixed mode requires periodicity, and the bounds the fixed mode.""" + with pytest.raises(ValueError, match="k_grid_mode"): + _periodic_long_range(k_grid_mode="dense") + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + k_grid_mode="fixed", + ) + with pytest.raises(ValueError, match="k_grid_bounds"): + _periodic_long_range(k_grid_bounds=[3, 3, 3]) + # explicit bounds build the grid at construction, no statistics needed + long_range = _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[2, 2, 2]) + assert long_range.k_grid_n.shape == (5**3 - 1, 3) + + +def test_fixed_k_grid_stats_size_the_grid() -> None: + """The data statistics size the grid by the largest sampled cell.""" + long_range = _periodic_long_range(k_grid_mode="fixed") + # the grid is not sized before the statistics run + coord = np.zeros((9, 3)) + with pytest.raises(ValueError, match="not sized yet"): + long_range.build_topologies( + coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) + ) + long_range.compute_input_stats( + [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}] + ) + k_max = long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + from deepmd.dpmodel.utils.neighbor_graph import ( + k_grid_bounds_of_cells, + ) + + expected = np.max(k_grid_bounds_of_cells(_cells(), k_max), axis=0) + np.testing.assert_array_equal(np.max(np.abs(long_range.k_grid_n), axis=0), expected) + topologies = long_range.build_topologies( + coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) + ) + assert topologies["k_grid"].vec.shape == (2 * long_range.k_grid_n.shape[0], 3) + # statistics without a cell cannot size the grid + unsized = _periodic_long_range(k_grid_mode="fixed") + with pytest.raises(ValueError, match="no simulation cell"): + unsized.compute_input_stats([{"coord": np.zeros((1, 3))}]) + + +def test_fixed_k_grid_round_trip() -> None: + """The record keeps the grid mode and the sized grid.""" + long_range = _periodic_long_range(k_grid_mode="fixed") + long_range.compute_input_stats([{"box": _cells()}]) + record = long_range.serialize() + assert record["k_grid_mode"] == "fixed" + restored = LatentChargeLongRange.deserialize(record) + assert restored.k_grid_mode == "fixed" + np.testing.assert_array_equal(restored.k_grid_n, long_range.k_grid_n) + # the default mode serializes no grid fields + plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) + assert plain.k_grid_mode == "ragged" + assert plain.k_grid_n is None + + +# --------------------------------------------------------------------------- +# Decoupled k grid +# --------------------------------------------------------------------------- + + +def test_decoupled_k_grid_requires_periodic() -> None: + """Decoupling is only meaningful for a periodic module.""" + with pytest.raises(ValueError, match="decouple_k_grid"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + decouple_k_grid=True, + ) + + +def test_decoupled_k_grid_ignores_kernel_drift() -> None: + """A decoupled grid keeps the k points of the initial kernel however the + kernel parameters move; a coupled one re-derives them every batch. + """ + coord = np.random.default_rng(3).normal(size=(5, 3)) + atype = np.zeros(5, dtype=np.int64) + n_node = np.array([5]) + box = np.eye(3)[None] * 8.0 + + decoupled = _periodic_long_range(decouple_k_grid=True) + coupled = _periodic_long_range() + # the frozen cutoff is the cutoff of the initial kernel + assert decoupled.frozen_k_max == pytest.approx( + LESKernel(alpha=0.3).kspace_cutoff(1e-8) + ) + initial = decoupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + reference = coupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + np.testing.assert_array_equal(initial.vec, reference.vec) + # a larger alpha moves the cutoff out, growing the coupled grid + for module in (decoupled, coupled): + module.kernel.log_alpha = np.log(1.5) + drifted = decoupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + moved = coupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + np.testing.assert_array_equal(drifted.vec, initial.vec) + np.testing.assert_array_equal(drifted.mask, initial.mask) + assert moved.vec.shape[0] > initial.vec.shape[0] + + +def test_decoupled_fixed_grid_keeps_its_mask() -> None: + """The mask of a decoupled fixed grid ignores the kernel drift.""" + coord = np.zeros((5, 3)) + atype = np.zeros(5, dtype=np.int64) + n_node = np.array([5]) + box = np.eye(3)[None] * 8.0 + + long_range = _periodic_long_range( + k_grid_mode="fixed", k_grid_bounds=[5, 5, 5], decouple_k_grid=True + ) + before = long_range.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + # a smaller alpha moves the cutoff in; a coupled mask would tighten + long_range.kernel.log_alpha = np.log(0.05) + after = long_range.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + np.testing.assert_array_equal(after.mask, before.mask) + np.testing.assert_array_equal(after.vec, before.vec) + + +def test_decoupled_k_grid_round_trip() -> None: + """The frozen cutoff survives serialization, past the kernel drift.""" + long_range = _periodic_long_range(decouple_k_grid=True) + frozen = long_range.frozen_k_max + long_range.kernel.log_alpha = np.log(1.5) + record = long_range.serialize() + assert record["decouple_k_grid"] is True + assert record["frozen_k_max"] == pytest.approx(frozen) + restored = LatentChargeLongRange.deserialize(record) + assert restored.decouple_k_grid + assert restored.frozen_k_max == pytest.approx(frozen) + # ...even though the restored (drifted) kernel would derive a different one + assert restored.kernel.kspace_cutoff(1e-8) != pytest.approx(frozen) + # the default stays coupled and serializes nothing + plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) + assert not plain.decouple_k_grid + assert plain.frozen_k_max is None diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 9a85c5af01..7388f90b6c 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -1145,3 +1145,273 @@ def test_periodic_compression_is_rejected(kernel: str) -> None: model = get_model(config).to(env.DEVICE).eval() with pytest.raises(NotImplementedError, match="reciprocal-space"): model.atomic_model.enable_compression(0.5) + + +# --------------------------------------------------------------------------- +# Periodic fixed k-grid mode +# --------------------------------------------------------------------------- + + +def _periodic_fixed_model( + kernel: str, boxes: np.ndarray | None = None, **kwargs: object +) -> torch.nn.Module: + """A periodic model in the fixed k-grid mode, sized by the given cells.""" + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "fixed" + model = get_model(config).to(env.DEVICE).eval() + if boxes is not None: + boxes = torch.as_tensor(np.asarray(boxes, dtype=np.float64).reshape(-1, 9)) + model.atomic_model.long_range.compute_input_stats([{"box": boxes}]) + return model + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_matches_ragged(kernel: str) -> None: + """The fixed grid with its mask evaluates like the ragged enumeration.""" + coord, atype, box = _periodic_case() + ragged = _periodic_model(kernel) + fixed = _periodic_fixed_model(kernel, box) + # the state dicts of the two grid modes differ by the fixed grid buffer + fixed.load_state_dict(ragged.state_dict(), strict=False) + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + expected = ragged(coord_t, atype_t, box=box_t) + got = fixed(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(got[key], expected[key], rtol=1e-9, atol=1e-12) + # the gradient reaches the charge network and the kernel without NaN, + # though masked points carry a singular kernel transform nowhere + fixed.train() + out = fixed(coord_t, atype_t, box=box_t) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [ + parameter.grad + for parameter in fixed.atomic_model.long_range.parameters() + if parameter.grad is not None + ] + assert grads and all(torch.isfinite(grad).all() for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_mixed_cells(kernel: str) -> None: + """One shared grid serves a batch of different cells, each masked to its sphere.""" + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.4], axis=0) + fixed = _periodic_fixed_model(kernel, box2) + ragged = _periodic_model(kernel) + ragged.load_state_dict(fixed.state_dict(), strict=False) + coord_t, atype_t, box_t = ( + _tensor(coord2), + _tensor(atype2, torch.int64), + _tensor(box2), + ) + batched = fixed(coord_t, atype_t, box=box_t) + expected = ragged(coord_t, atype_t, box=box_t) + torch.testing.assert_close( + batched["energy"], expected["energy"], rtol=1e-9, atol=1e-12 + ) + torch.testing.assert_close( + batched["force"], expected["force"], rtol=1e-9, atol=1e-12 + ) + # the batched energy of a frame equals its standalone evaluation + for frame in range(2): + alone = fixed( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +def _kspace_reference_bounded( + coord: np.ndarray, + charge: np.ndarray, + box: np.ndarray, + kernel: str, + bounds: np.ndarray, +) -> float: + """The k-space reference of :func:`_kspace_reference` restricted to the + integer indices within ``bounds`` — the truncated sum of a fixed grid. + """ + rec = 2.0 * np.pi * np.linalg.inv(box) + if kernel == "les": + k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + else: + k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=float) for n in bounds], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge + s_sq = (np.abs(rho) ** 2).sum(axis=-1) + if kernel == "les": + k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2 + k0 = 2.0 * 0.8 / np.sqrt(np.pi) + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_hat = np.sum( + amplitude[None, :] + * (2.0 * np.pi) ** 1.5 + * width[None, :] ** 3 + * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2), + axis=-1, + ) + k0 = amplitude.sum() + volume = abs(np.linalg.det(box)) + return float( + COULOMB_CONSTANT + * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()) + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_truncates_oversized_cells(kernel: str) -> None: + """A cell beyond the grid is evaluated with the grid alone, not an error.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + big_box = box * 3.0 + with pytest.warns(UserWarning, match="truncated"): + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(big_box)) + grid_n = model.atomic_model.long_range.k_grid_n + grid_n = grid_n.detach().cpu().numpy() if hasattr(grid_n, "detach") else grid_n + bounds = np.abs(np.asarray(grid_n)).max(axis=0) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference_bounded(coord[0], charge, big_box[0], kernel, bounds) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_needs_statistics_or_bounds(kernel: str) -> None: + """Without statistics or explicit bounds the grid is not sized; bounds skip them.""" + coord, atype, box = _periodic_case() + unsized = _periodic_fixed_model(kernel) + with pytest.raises(ValueError, match="not sized yet"): + unsized(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "fixed" + config["long_range"]["k_grid_bounds"] = [12, 12, 12] + model = get_model(config).to(env.DEVICE).eval() + assert model.atomic_model.long_range.k_grid_n.shape == (25**3 - 1, 3) + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + reference = _periodic_model(kernel) + reference.load_state_dict(model.state_dict(), strict=False) + expected = reference(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + torch.testing.assert_close(out["energy"], expected["energy"], rtol=1e-9, atol=1e-12) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_serialization_round_trip(kernel: str) -> None: + """The record keeps the sized grid and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + record = model.serialize() + assert record["long_range"]["k_grid_mode"] == "fixed" + assert "k_grid_n" in record["long_range"] + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.long_range.k_grid_mode == "fixed" + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense") + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_fixed_grid_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the fixed-grid model like the eager one.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + + +# --------------------------------------------------------------------------- +# Periodic decoupled k grid +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: + """A decoupled model selects the k points with the initial kernel but + evaluates them with the live one: the drift moves the energy, not the grid. + """ + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["decouple_k_grid"] = True + model = get_model(config).to(env.DEVICE).eval() + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + long_range = model.atomic_model.long_range + grid_args = ( + coord_t.reshape(-1, 3), + torch.tensor([6], dtype=torch.int64, device=env.DEVICE), + box_t, + ) + before = long_range.build_topologies(*grid_args, atype=atype_t.reshape(-1))[ + "k_grid" + ] + energy_before = model(coord_t, atype_t, box=box_t)["energy"] + # move the kernel far enough that a coupled grid would change + with torch.no_grad(): + if kernel == "les": + long_range.kernel.log_alpha.add_(1.0) + else: + long_range.kernel.log_width.sub_(1.0) + after = long_range.build_topologies(*grid_args, atype=atype_t.reshape(-1))["k_grid"] + torch.testing.assert_close(after.vec, before.vec) + torch.testing.assert_close(after.mask, before.mask) + # the kernel values still follow the live parameters + energy_after = model(coord_t, atype_t, box=box_t)["energy"] + assert not torch.allclose(energy_after, energy_before) + # and the drifted kernel still takes a finite gradient through the graph + model.train() + out = model(coord_t, atype_t, box=box_t) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + kernel_grads = [parameter.grad for parameter in long_range.kernel.parameters()] + assert kernel_grads and all( + grad is not None and torch.isfinite(grad).all() for grad in kernel_grads + ) From 50c108c0910be2182b52ce04818f70c8edfd49f5 Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Fri, 23 Oct 2026 07:25:13 +0000 Subject: [PATCH 24/30] feat(long-range): evaluate periodic systems on a static charge mesh (SPME) - k_grid_mode: "pme" spreads the charges to a static real-space mesh with a B-spline of order pme_order and evaluates the reciprocal sum by an FFT with the SPME deconvolution factor, near-linear in system size; the mesh is sized by the largest cell of the data statistics or by pme_mesh (default spacing pme_spacing = pi/k_max, respecting decouple_k_grid) and serialized with the model - the PME_MESH topology carries (box, volume, coord, strain, mask); the strain family mirrors the k-grid one (fractional coordinates are strain-invariant, the virial comes from the strained reciprocal vectors and volume) - restore the stats-sized mesh lazily after load_state_dict, and feed the graph lower torch cells in DeepEval (fixes periodic dp test for both the mesh and the k-grid paths) --- AGENTS.md | 2 +- deepmd/dpmodel/long_range/latent_charge.py | 408 +++++++++++++++++- deepmd/dpmodel/model/make_model.py | 17 +- .../dpmodel/utils/neighbor_graph/__init__.py | 6 + .../utils/neighbor_graph/pair_graph.py | 157 ++++++- deepmd/pt_expt/infer/deep_eval.py | 20 +- deepmd/pt_expt/model/edge_transform_output.py | 3 +- deepmd/pt_expt/model/make_model.py | 16 + deepmd/pt_expt/train/training.py | 10 + deepmd/pt_expt/utils/serialization.py | 31 +- deepmd/utils/argcheck.py | 54 ++- periodic_long_range_design.md | 18 + .../tests/common/dpmodel/test_long_range.py | 205 +++++++++ source/tests/pt_expt/model/test_long_range.py | 264 ++++++++++++ 14 files changed, 1166 insertions(+), 45 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7e21751d7c..bcf7a9e945 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 -- **Long-range latent-charge module**: a standard energy model (DPA4C descriptor) takes a `long_range` section (`type: latent_charge`, kernels `les`/`sog`, optional `charge_constraint`) that adds an interacting-charge energy evaluated on complete systems — directly over all pairs for non-periodic frames, and in reciprocal space for periodic frames when `periodic: true`. The k grid has two layouts selected by `k_grid_mode`: `"ragged"` (default; k grid enumerated eagerly per batch from the kernel's `kspace_cutoff`, tin-foil boundary, neutral channels required) and `"fixed"` (one integer grid shared over frames, sized by the largest cell of the data statistics or by `k_grid_bounds`, with a mask selecting each frame's k points; out-of-grid cells are truncated to the grid with a warning). The wave-number cutoff is re-derived from the live (trainable) kernel every batch unless `decouple_k_grid: true` freezes it at the initial parameters (the kernel values still follow the live parameters; the frozen cutoff is serialized). Periodic evaluation supports eager and compiled training plus the Python interface; freeze/export, compression, C++ and LAMMPS reject periodic long-range models. Tests: `pytest source/tests/pt_expt/model/test_long_range.py`, `source/tests/common/dpmodel/test_long_range.py`. +- **Long-range latent-charge module**: a standard energy model (DPA4C descriptor) takes a `long_range` section (`type: latent_charge`, kernels `les`/`sog`, optional `charge_constraint`) that adds an interacting-charge energy evaluated on complete systems — directly over all pairs for non-periodic frames, and in reciprocal space for periodic frames when `periodic: true`. The k grid has two layouts selected by `k_grid_mode`: `"ragged"` (default; k grid enumerated eagerly per batch from the kernel's `kspace_cutoff`, tin-foil boundary, neutral channels required) and `"fixed"` (one integer grid shared over frames, sized by the largest cell of the data statistics or by `k_grid_bounds`, with a mask selecting each frame's k points; out-of-grid cells are truncated to the grid with a warning) and `"pme"` (charges spread to a static real-space mesh by a B-spline of order `pme_order`, FFT, kernel transform with the SPME deconvolution factor; mesh sized by the largest cell of the data statistics or by `pme_mesh`, default spacing `pme_spacing` = pi/k_max; near-linear in system size). The wave-number cutoff is re-derived from the live (trainable) kernel every batch unless `decouple_k_grid: true` freezes it at the initial parameters (the kernel values still follow the live parameters; the frozen cutoff is serialized). Periodic evaluation supports eager and compiled training plus the Python interface; freeze/export, compression, C++ and LAMMPS reject periodic long-range models. Tests: `pytest source/tests/pt_expt/model/test_long_range.py`, `source/tests/common/dpmodel/test_long_range.py`. ### JAX Backend diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index fa5d60d95c..1e85ad4440 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import math from collections.abc import ( Callable, ) @@ -11,6 +12,7 @@ from deepmd.dpmodel.array_api import ( Array, + xp_add_at, ) from deepmd.dpmodel.common import ( GLOBAL_NP_FLOAT_PRECISION, @@ -23,7 +25,9 @@ from deepmd.dpmodel.utils.neighbor_graph import ( ALL_PAIRS, K_GRID, + PME_MESH, KGrid, + MeshGrid, build_topologies, enumerate_k_grid_bounds, frame_id_from_n_node, @@ -63,6 +67,73 @@ """ +def _fft_smooth(n: int) -> int: + """The smallest integer >= ``n`` whose prime factors are among 2, 3, 5, 7.""" + + def smooth(m: int) -> bool: + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + return m == 1 + + m = max(1, n) + while not smooth(m): + m += 1 + return m + + +def _bspline_integer_weights(p: int) -> np.ndarray: + r"""The order-``p`` cardinal B-spline at :math:`M_p(l+1)`, ``l = 0..p-2``. + + The weights of the SPME denominator + :math:`D(\theta) = \sum_l M_p(l+1) e^{i \theta l}`. + """ + out = [] + for l in range(p - 1): + x = float(l + 1) + out.append( + sum( + (-1.0) ** j * math.comb(p, j) * max(x - j, 0.0) ** (p - 1) + for j in range(p + 1) + ) + / math.factorial(p - 1) + ) + return np.asarray(out) + + +def _bspline_eval(xp: Any, x: Array, p: int) -> Array: + r"""The order-``p`` cardinal B-spline :math:`M_p(x)` elementwise. + + :math:`M_p(x) = \frac{1}{(p-1)!} \sum_j (-1)^j \binom{p}{j} + \mathrm{relu}(x - j)^{p-1}` — a piecewise polynomial, differentiable + (twice for :math:`p \ge 4`) with respect to ``x``. + """ + w = None + for j in range(p + 1): + term = ((-1.0) ** j * math.comb(p, j)) * xp.maximum( + x - j, xp.zeros_like(x) + ) ** (p - 1) + w = term if w is None else w + term + return w / math.factorial(p - 1) + + +def _inv3(xp: Any, m: Array) -> Array: + """Inverse of a batch of 3x3 matrices, adjugate over determinant. + + Rows a, b, c: the columns of the inverse are (b x c)/det, (c x a)/det, + (a x b)/det. The closed form is cheap and forward-only here — the cell is + not an autograd leaf; the strain leaf keeps ``linalg.inv`` whose + registered backward handles a symbolic frame count. + """ + a, b, c = m[..., 0, :], m[..., 1, :], m[..., 2, :] + bc = xp.linalg.cross(b, c) + ca = xp.linalg.cross(c, a) + ab = xp.linalg.cross(a, b) + det = xp.sum(a * bc, axis=-1) + inv = xp.stack([bc, ca, ab], axis=-1) # columns + return inv / det[..., None, None] + + @BaseLongRange.register("latent_charge") class LatentChargeLongRange(NativeOP, BaseLongRange): r"""Energy of interacting latent charges, periodic or not. @@ -119,15 +190,21 @@ class LatentChargeLongRange(NativeOP, BaseLongRange): The relative weight of the slowest-decaying kernel mode at the wave-number cutoff of the k grid of a periodic batch. k_grid_mode - How the k grid of a periodic batch is laid out: ``"ragged"`` + How the reciprocal sum of a periodic batch is evaluated: ``"ragged"`` enumerates the wave numbers of every frame eagerly, so the k axis is a dynamic dimension; ``"fixed"`` shares one integer grid over the frames — sized by the largest cell of the data statistics, or by ``k_grid_bounds`` — evaluates every frame on it by a static-shaped - transform, and masks the points a frame does not reach. The fixed - grid is hard: a frame needing an integer index beyond it (a larger - cell, or a kernel that moved the cutoff out) is served by the grid - alone, so its k-sum is truncated to the grid with a warning. + transform, and masks the points a frame does not reach; ``"pme"`` + spreads the charges to a static real-space mesh with a B-spline and + evaluates the sum by an FFT (SPME), sized by the largest cell of the + data statistics or by ``pme_mesh``. The fixed grid is hard: a frame + needing an integer index beyond it (a larger cell, or a kernel that + moved the cutoff out) is served by the grid alone, so its k-sum is + truncated to the grid with a warning; on the mesh such a frame is + evaluated with a coarser effective spacing (a smooth accuracy + degradation; no warning — sizing the mesh is the caller's + responsibility). k_grid_bounds The per-axis integer bounds :math:`n_d^{max}` of the fixed grid, in place of the data statistics; only meaningful with @@ -140,6 +217,15 @@ class LatentChargeLongRange(NativeOP, BaseLongRange): evaluated with still follow the live parameters. The frozen cutoff is serialized with the model, so a restarted run keeps the grid the model was built with. Requires ``periodic=True``. + pme_spacing + The target real-space mesh spacing of the ``"pme"`` mode, in Å; the + default is the Nyquist spacing of the kernel's wave-number cutoff, + :math:`\pi / k_{max}`. + pme_order + The B-spline order of the ``"pme"`` charge assignment. + pme_mesh + The per-axis mesh dimensions :math:`(n_1, n_2, n_3)` of the + ``"pme"`` mode, in place of the data statistics. """ def __init__( @@ -152,6 +238,9 @@ def __init__( k_grid_mode: str = "ragged", k_grid_bounds: list[int] | None = None, decouple_k_grid: bool = False, + pme_spacing: float | None = None, + pme_order: int = 4, + pme_mesh: list[int] | None = None, ) -> None: if charge_net.vacuum_ref: raise ValueError( @@ -176,12 +265,14 @@ def __init__( # to; see :meth:`fold_charge_state`. self.charge_state = None self.k_grid_mode = str(k_grid_mode) - if self.k_grid_mode not in ("ragged", "fixed"): + if self.k_grid_mode not in ("ragged", "fixed", "pme"): raise ValueError( - f"`k_grid_mode` must be 'ragged' or 'fixed', got {self.k_grid_mode!r}" + f"`k_grid_mode` must be 'ragged', 'fixed' or 'pme', got {self.k_grid_mode!r}" + ) + if self.k_grid_mode in ("fixed", "pme") and not self.periodic: + raise ValueError( + f"`k_grid_mode='{self.k_grid_mode}'` requires `periodic=True`" ) - if self.k_grid_mode == "fixed" and not self.periodic: - raise ValueError("`k_grid_mode='fixed'` requires `periodic=True`") if k_grid_bounds is not None and self.k_grid_mode != "fixed": raise ValueError( "`k_grid_bounds` is only meaningful with `k_grid_mode='fixed'`" @@ -207,6 +298,58 @@ def __init__( self.k_grid_n = np.zeros((0, 3), dtype=np.int64) else: self.k_grid_n = None + # The PME mesh parameters. The mesh dims are an array even while + # unsized, so a backend keeping arrays in its state dict sees a + # stable set of keys and a checkpoint restores the sized mesh; the + # plain-integer copy is what a trace reads as static shapes. + self.pme_order = int(pme_order) + if self.pme_order < 2: + raise ValueError(f"`pme_order` must be at least 2, got {self.pme_order}") + if pme_spacing is not None: + pme_spacing = float(pme_spacing) + if not np.isfinite(pme_spacing) or pme_spacing <= 0.0: + raise ValueError( + f"`pme_spacing` must be finite and positive, got {pme_spacing}" + ) + self.pme_spacing = pme_spacing + if self.k_grid_mode != "pme": + if pme_spacing is not None or pme_mesh is not None: + raise ValueError( + "`pme_spacing`/`pme_mesh` are only meaningful with " + "`k_grid_mode='pme'`" + ) + self.pme_mesh = None + self._pme_n = None + elif pme_mesh is not None: + mesh = np.asarray(pme_mesh, dtype=np.int64).reshape(-1) + if mesh.shape != (3,) or np.any(mesh < 1): + raise ValueError( + f"`pme_mesh` must be three positive integers, got {pme_mesh}" + ) + self.pme_mesh = mesh + self._pme_n = (int(mesh[0]), int(mesh[1]), int(mesh[2])) + else: + # Unsized sentinel: the data statistics size it (see + # :meth:`compute_input_stats`), or a checkpoint restores the + # sized mesh through the state dict. + self.pme_mesh = np.zeros((3,), dtype=np.int64) + self._pme_n = None + + def _pme_dims(self) -> tuple[int, int, int] | None: + """The mesh dimensions as plain integers, or ``None`` when unsized. + + The plain-integer cache is what a trace reads as static shapes; it is + set at construction and at statistics time. A backend restoring the + sized mesh through its state dict (a checkpoint does not run + :meth:`compute_input_stats`) is picked up lazily from the mesh array. + """ + if self._pme_n is not None: + return self._pme_n + if self.pme_mesh is None or int(self.pme_mesh[0]) < 1: + return None + mesh = self.pme_mesh + self._pme_n = (int(mesh[0]), int(mesh[1]), int(mesh[2])) + return self._pme_n @staticmethod def _normalize_constraint( @@ -299,6 +442,9 @@ def from_config( k_grid_mode=config.get("k_grid_mode", "ragged"), k_grid_bounds=config.get("k_grid_bounds"), decouple_k_grid=config.get("decouple_k_grid", False), + pme_spacing=config.get("pme_spacing"), + pme_order=config.get("pme_order", 4), + pme_mesh=config.get("pme_mesh"), ) def _k_max(self) -> float: @@ -312,6 +458,17 @@ def _k_max(self) -> float: return self.frozen_k_max return float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + def _pme_target_spacing(self) -> float: + r"""The target PME mesh spacing, in Å. + + The explicit ``pme_spacing``, or the Nyquist spacing of the kernel's + wave-number cutoff :math:`\pi / k_{max}` — fine enough that the + kernel's spectral weight at the mesh cutoff is negligible. + """ + if self.pme_spacing is not None: + return self.pme_spacing + return math.pi / self._k_max() + def build_topologies( self, coord: Array, @@ -333,15 +490,24 @@ def build_topologies( """ params = {} if self.periodic: - params[K_GRID] = {"k_max": self._k_max()} - if self.k_grid_mode == "fixed": - if self.k_grid_n is None or self.k_grid_n.shape[0] == 0: + if self.k_grid_mode == "pme": + if self._pme_dims() is None: raise ValueError( - "the fixed k grid is not sized yet: run the data " + "the PME mesh is not sized yet: run the data " "statistics (`compute_input_stats`) on periodic " - "frames, or configure `k_grid_bounds`" + "frames, or configure `pme_mesh`" ) - params[K_GRID]["n_grid"] = self.k_grid_n + params[PME_MESH] = {"n_mesh": self.pme_mesh} + else: + params[K_GRID] = {"k_max": self._k_max()} + if self.k_grid_mode == "fixed": + if self.k_grid_n is None or self.k_grid_n.shape[0] == 0: + raise ValueError( + "the fixed k grid is not sized yet: run the data " + "statistics (`compute_input_stats`) on periodic " + "frames, or configure `k_grid_bounds`" + ) + params[K_GRID]["n_grid"] = self.k_grid_n return build_topologies( self.required_topologies(), coord, @@ -355,6 +521,8 @@ def build_topologies( def required_topologies(self) -> tuple[str, ...]: if self.periodic: + if self.k_grid_mode == "pme": + return (ALL_PAIRS, PME_MESH) return (ALL_PAIRS, K_GRID) return (ALL_PAIRS,) @@ -426,9 +594,14 @@ def call_graph( ) energy = self._pair_branch(charge, pairs, n) if self.periodic: - energy = energy + self._kspace_branch( - charge, topologies[K_GRID], frame_id, node_mask, n_node - ) + if self.k_grid_mode == "pme": + energy = energy + self._pme_branch( + charge, topologies[PME_MESH], frame_id, node_mask, n_node + ) + else: + energy = energy + self._kspace_branch( + charge, topologies[K_GRID], frame_id, node_mask, n_node + ) return xp.reshape(energy, (n, 1)) def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: @@ -523,6 +696,161 @@ def _kspace_branch( ) return e_rec + e_self + def _pme_b_factor(self, xp: Any, dtype: Any, device: Any) -> Array: + r"""The SPME assignment correction :math:`|B(m)|^2` on the mesh. + + Per axis :math:`|D_d(m)|^2 = |\sum_l M_p(l+1) e^{2\pi i m l / n_d}|^2` + (equal to :math:`\mathrm{sinc}(m/n_d)^{2p}`), the full factor the + three-axis outer product, in the ``fftfreq`` index order. + """ + w = _bspline_integer_weights(self.pme_order) + axes = [] + for n_d in self._pme_dims(): + m = xp.arange(n_d, dtype=dtype, device=device) + m = xp.where(m > n_d // 2, m - n_d, m) + theta = 2.0 * math.pi * m / n_d + re = xp.zeros_like(theta) + im = xp.zeros_like(theta) + for l, wl in enumerate(w): + re = re + float(wl) * xp.cos(theta * l) + im = im + float(wl) * xp.sin(theta * l) + axes.append(1.0 / (re * re + im * im)) + b = axes[0][:, None, None] * axes[1][None, :, None] * axes[2][None, None, :] + return xp.reshape(b, (-1,)) + + def _pme_branch( + self, + charge: Array, + mesh: MeshGrid, + frame_id: Array, + node_mask: Array, + n_node: Array, + ) -> Array: + r"""Reciprocal-space energy of the periodic frames on a mesh, per atom. + + Evaluates the same lattice sum as :meth:`_kspace_branch` by the SPME + algorithm: the charges spread to a static real-space mesh with a + B-spline of order :attr:`pme_order`, the mesh transforms by an FFT, + and the energy convolves with the kernel transform corrected by the + assignment factor (:meth:`_pme_b_factor`). The strain leaf enters the + coordinates by ``I + strain`` and the cells by + ``B (I + strain)^T`` — the fractional coordinates are strain- + invariant, so the assignment is unstrained and the virial comes from + the reciprocal vectors and the volume alone. Non-periodic frames + (``mesh.mask`` false) carry no mesh energy and no self term. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + nf = n_node.shape[0] + n_channel = charge.shape[-1] + n1, n2, n3 = self._pme_dims() + p = self.pme_order + m_total = n1 * n2 * n3 + box = xp.astype(mesh.box, charge.dtype) + coord = xp.astype(mesh.coord, charge.dtype) + volume = xp.astype(mesh.volume, charge.dtype) + b_inv = _inv3(xp, box) # (nf, 3, 3); columns are g_d / 2pi + b_inv_node = xp.take(b_inv, frame_id, axis=0) # (N, 3, 3) + # fractional coordinates, wrapped to [0, 1): frac = r B^-1 + frac = xp.matmul(coord[:, None, :], b_inv_node)[:, 0, :] + frac = frac - xp.floor(frac) + # per-axis assignment: weights (N, p) and mesh indices (N, p) + weights = [] + indices = [] + for d, n_d in enumerate((n1, n2, n3)): + u = frac[:, d] * n_d + base = xp.floor(u) + t = (u - base)[:, None] + shift = xp.arange(p, dtype=charge.dtype, device=device)[None, :] + # x_k = t + p - 1 - k lies in the support [0, p] of M_p + weights.append(_bspline_eval(xp, t + (p - 1) - shift, p)) + idx = ( + xp.astype(base, xp.int64)[:, None] + - (p - 1) + + xp.arange(p, dtype=xp.int64, device=device)[None, :] + ) + indices.append(idx % n_d) + w1, w2, w3 = weights + j1, j2, j3 = indices + w = xp.reshape( + w1[:, :, None, None] * w2[:, None, :, None] * w3[:, None, None, :], + (frac.shape[0], p**3), + ) + lin = xp.reshape( + (j1[:, :, None, None] * n2 + j2[:, None, :, None]) * n3 + + j3[:, None, None, :], + (frac.shape[0], p**3), + ) + # spread the charges to the mesh, frame- and channel-major + q_mesh = xp.zeros( + (nf * n_channel * m_total,), dtype=charge.dtype, device=device + ) + for c in range(n_channel): + target = ((frame_id * n_channel + c) * m_total)[:, None] + lin + value = charge[:, c][:, None] * w + q_mesh = xp_add_at( + q_mesh, xp.reshape(target, (-1,)), xp.reshape(value, (-1,)) + ) + q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) + rho = xp.fft.fftn(q_mesh, axes=(-3, -2, -1)) + rho_sq = xp.real(rho) ** 2 + xp.imag(rho) ** 2 + rho_sq = xp.reshape(rho_sq, (nf, n_channel, m_total)) + # the reciprocal vectors of every frame on the mesh: mode m is the + # integer Miller index, k(m) = 2 pi B^-1 m + h_axes = [] + for n_d in (n1, n2, n3): + m = xp.arange(n_d, dtype=charge.dtype, device=device) + m = xp.where(m > n_d // 2, m - n_d, m) + h_axes.append(m) + g1, g2, g3 = xp.meshgrid(*h_axes, indexing="ij") + h = xp.stack( + [xp.reshape(g1, (-1,)), xp.reshape(g2, (-1,)), xp.reshape(g3, (-1,))], + axis=-1, + ) # (M, 3) + k_vec = 2.0 * math.pi * xp.matmul(h[None, :, :], xp.matrix_transpose(b_inv)) + if mesh.strain is not None: + # The strain family of the virial: B' = B S^T with S = I + strain. + # The fractional coordinates are exactly strain-invariant, so the + # assignment above needs no strain; the reciprocal vectors and the + # volume pick it up: k' = S^-T k, V' = V det(S). The 3x3 inverse + # and determinant stay on linalg — its registered backward handles + # a symbolic frame count, unlike hand-written select/cross forms. + eye = xp.eye(3, dtype=charge.dtype, device=device) + s = xp.astype(mesh.strain, charge.dtype) + eye # (nf, 3, 3) + s_inv = xp.linalg.inv(s) + # rows: k'_i = sum_j k_j (S^-1)[j, i] + k_vec = xp.matmul(k_vec, s_inv) + volume = volume * xp.linalg.det(s) + k_sq = xp.reshape(xp.sum(k_vec * k_vec, axis=-1), (-1,)) # (nf * M,) + # the zero mode sits at flat index 0; guard the singular kernel there + safe = xp.where(k_sq == 0.0, xp.ones_like(k_sq), k_sq) + influence = xp.reshape(self.kernel.kspace(safe), (nf, m_total)) + influence = influence * self._pme_b_factor(xp, charge.dtype, device)[None, :] + # the zero mode sits at flat index 0 (tin-foil); keep everything else + keep = xp.concat( + [ + xp.zeros((1,), dtype=xp.bool, device=device), + xp.ones((m_total - 1,), dtype=xp.bool, device=device), + ] + ) + influence = xp.where(keep[None, :], influence, xp.zeros_like(influence)) + influence = influence * xp.astype(mesh.mask, charge.dtype)[:, None] + half_ke = 0.5 * COULOMB_CONSTANT + e_frame = ( + half_ke * xp.sum(influence * xp.sum(rho_sq, axis=1), axis=-1) / volume + ) # (nf,) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + count = xp.maximum(count, xp.ones_like(count)) + e_rec = xp.take(e_frame / count, frame_id, axis=0) * weight # (N,) + e_self = ( + -half_ke + * self.kernel.contact_value(charge) + * xp.sum(charge * charge, axis=-1) + * xp.astype(xp.take(mesh.mask, frame_id, axis=0), charge.dtype) + ) + return e_rec + e_self + def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: """Kernel values of the pairs. @@ -631,6 +959,10 @@ def compute_input_stats( sampled = merged() if callable(merged) else merged self._size_k_grid(sampled) merged = sampled + elif self.k_grid_mode == "pme" and self._pme_dims() is None: + sampled = merged() if callable(merged) else merged + self._size_pme_mesh(sampled) + merged = sampled self.charge_net.compute_input_stats( merged, protection=protection, stat_file_path=stat_file_path ) @@ -664,6 +996,37 @@ def _size_k_grid(self, sampled: list[dict]) -> None: ) self.k_grid_n = enumerate_k_grid_bounds(bounds) + def _size_pme_mesh(self, sampled: list[dict]) -> None: + r"""Size the PME mesh by the largest sampled cell. + + Every frame is evaluated on the shared mesh with the effective + spacing :math:`|a_d|/n_d` per axis; the mesh takes the largest sampled + cell axis down to the target spacing, rounded up to an FFT-smooth + dimension. A frame larger than the mesh (a larger cell than the + statistics saw) is evaluated with a coarser effective spacing — a + smooth accuracy degradation, by design without a warning. + """ + spacing = self._pme_target_spacing() + a_max = None + for frame in sampled: + box = frame.get("box") + if box is None: + continue + box = np.asarray(to_numpy_array(box), dtype=np.float64) + if box.size == 0: + continue + frame_max = np.max(np.linalg.norm(box.reshape(-1, 3, 3), axis=-1), axis=0) + a_max = frame_max if a_max is None else np.maximum(a_max, frame_max) + if a_max is None: + raise ValueError( + "the PME mesh takes its size from the data statistics, " + "but the sampled frames hold no simulation cell; provide " + "periodic training data or configure `pme_mesh`" + ) + dims = [_fft_smooth(int(np.ceil(a / spacing))) for a in a_max] + self.pme_mesh = np.asarray(dims, dtype=np.int64) + self._pme_n = (dims[0], dims[1], dims[2]) + def change_type_map(self, type_map: list[str]) -> None: self.charge_net.change_type_map(type_map=type_map) @@ -689,6 +1052,12 @@ def serialize(self) -> dict: if self.decouple_k_grid: data["decouple_k_grid"] = True data["frozen_k_max"] = self.frozen_k_max + if self.k_grid_mode == "pme": + data["pme_order"] = self.pme_order + if self.pme_spacing is not None: + data["pme_spacing"] = self.pme_spacing + if self._pme_dims() is not None: + data["pme_mesh"] = to_numpy_array(self.pme_mesh).tolist() if self.charge_state is not None: data["charge_state"] = to_numpy_array(self.charge_state).tolist() return data @@ -705,6 +1074,9 @@ def deserialize(cls, data: dict) -> "LatentChargeLongRange": kspace_tolerance=data.get("kspace_tolerance", 1e-8), k_grid_mode=data.get("k_grid_mode", "ragged"), decouple_k_grid=data.get("decouple_k_grid", False), + pme_spacing=data.get("pme_spacing"), + pme_order=data.get("pme_order", 4), + pme_mesh=data.get("pme_mesh"), ) if data.get("frozen_k_max") is not None: # The constructor derives the frozen cutoff from the (possibly diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 42b85f07d9..36e3cd3440 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -891,10 +891,11 @@ def call_common_lower_graph( spin=spin, ) if topologies: - # The pair vectors and the k-grid geometry are geometry like - # the edge vectors. + # The pair vectors and the k-grid/mesh geometry are geometry + # like the edge vectors. from deepmd.dpmodel.utils.neighbor_graph import ( KGrid, + MeshGrid, ) xp = array_api_compat.array_namespace(edge_vec) @@ -913,6 +914,18 @@ def call_common_lower_graph( ), strain=topology.strain, ) + elif isinstance(topology, MeshGrid): + cast[name] = MeshGrid( + box=xp.astype(topology.box, edge_vec.dtype), + volume=xp.astype(topology.volume, edge_vec.dtype), + coord=( + None + if topology.coord is None + else xp.astype(topology.coord, edge_vec.dtype) + ), + strain=topology.strain, + mask=topology.mask, + ) else: cast[name] = PairGraph( topology.index, diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 8bcecaecd6..3d427c0ec4 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -59,10 +59,13 @@ from .pair_graph import ( ALL_PAIRS, K_GRID, + PME_MESH, KGrid, + MeshGrid, PairGraph, build_fixed_k_grid, build_k_grid, + build_mesh_grid, build_pair_graph, build_topologies, check_complete_frames, @@ -86,8 +89,10 @@ __all__ = [ "ALL_PAIRS", "K_GRID", + "PME_MESH", "GraphLayout", "KGrid", + "MeshGrid", "NeighborGraph", "PairGraph", "angle_padding_fraction", @@ -101,6 +106,7 @@ "build_edge_csr", "build_fixed_k_grid", "build_k_grid", + "build_mesh_grid", "build_neighbor_graph", "build_neighbor_graph_ase", "build_pair_graph", diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index e6be02bfb5..64c3ca0f82 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -70,6 +70,17 @@ wave numbers, which the pair topology cannot carry. """ +PME_MESH = "pme_mesh" +"""Name of the topology carrying the cells of periodic frames for a mesh evaluation. + +A long-range module declares it beside :data:`ALL_PAIRS` when it evaluates the +reciprocal sum of a smooth kernel on a real-space mesh (SPME): charges spread +to a static mesh by a B-spline, transformed by an FFT, and convolved with the +kernel transform. The mesh axes are static and live on the module, so the +topology holds no ragged axis of its own — only the per-frame cells and the +geometry leaves. +""" + @dataclass class PairGraph: @@ -121,6 +132,42 @@ class KGrid: """ +@dataclass +class MeshGrid: + r"""Cells and geometry of the periodic frames for a mesh (FFT) evaluation. + + The reciprocal sum is evaluated on a static real-space mesh shared by the + frames, so the topology carries no ragged axis: the cells (for the + reciprocal vectors), the volumes, the geometry leaves, and a per-frame + flag telling periodic frames from non-periodic ones (a batch without a + cell comes back with an all-false mask). The volume is computed eagerly + (as in :class:`KGrid`): the graph lower only rescales it by the strain + determinant, which keeps ``linalg.det`` out of the traced graph. + """ + + box: Array + """(nf, 3, 3) float simulation cells; rows are the lattice vectors.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array | None + """(N, 3) float node positions; the geometry leaf of the mesh energy. + + ``None`` when the grid is reconstructed from flattened inputs (export); a + graph lower re-leafs it before the evaluation. + """ + strain: Array | None + """(nf, 3, 3) float zero strain leaf for the mesh virial. + + Same contract as :attr:`KGrid.strain`: the energy is evaluated with the + coordinates strained by ``I + strain`` and the cells by + ``B' = B (I + strain)^T`` with the lattice vectors as rows, so the + fractional coordinates are strain-invariant and the gradient with respect + to the leaf is the negative virial. + """ + mask: Array + """(nf,) bool periodic (1) vs non-periodic (0) frames.""" + + def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: """Enumerate the unordered pairs of valid nodes within every frame. @@ -433,6 +480,67 @@ def _empty_k_grid(nf: int, like: Array, coord: Array | None) -> KGrid: ) +def _empty_mesh_grid(nf: int, like: Array, coord: Array | None) -> MeshGrid: + """A non-periodic mesh grid: every frame masked out, identity cells.""" + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + return MeshGrid( + box=xp.tile(xp.eye(3, dtype=like.dtype, device=device), (nf, 1, 1)), + volume=xp.ones((nf,), dtype=like.dtype, device=device), + coord=coord, + strain=xp.zeros((nf, 3, 3), dtype=like.dtype, device=device), + mask=xp.zeros((nf,), dtype=xp.bool, device=device), + ) + + +def build_mesh_grid( + box: Array, + n_mesh: Array, + coord: Array | None = None, +) -> MeshGrid: + r"""Pack the cells of the periodic frames for the mesh (FFT) evaluation. + + The mesh is shared by every frame; a frame whose cell is larger than the + mesh was sized for is evaluated with a coarser effective spacing + :math:`|a_d|/n_d` than the target — an undersampled kernel spectrum, a + well-defined approximation whose forces and virial stay the exact + derivatives of that energy. Sizing the mesh (data statistics or + ``pme_mesh``) is the caller's responsibility; no warning is issued here. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + n_mesh + The per-axis mesh dimensions with shape ``(3,)``. + coord + Node positions with shape ``(N, 3)``, stored on the grid as the + geometry context (see :attr:`MeshGrid.coord`). + + Returns + ------- + MeshGrid + The cells and geometry of the frames, every frame flagged periodic. + """ + xp = array_api_compat.array_namespace(box) + if box.ndim == 2: + box = xp.reshape(box, (box.shape[0], 3, 3)) + volume = xp.linalg.det(box) + volume = xp.abs(volume) + return MeshGrid( + box=box, + volume=xp.astype(volume, box.dtype), + coord=coord, + strain=xp.zeros( + (box.shape[0], 3, 3), dtype=box.dtype, device=array_api_compat.device(box) + ), + mask=xp.ones( + (box.shape[0],), dtype=xp.bool, device=array_api_compat.device(box) + ), + ) + + def _empty_pair_graph(like: Array, atype: Array) -> PairGraph: """A guard-only pair topology for a batch without non-periodic frames.""" xp = array_api_compat.array_namespace(like) @@ -534,15 +642,16 @@ def build_topologies( ValueError If a name is unknown, or the k grid parameters are missing. """ - topologies: dict[str, PairGraph | KGrid] = {} + topologies: dict[str, PairGraph | KGrid | MeshGrid] = {} for name in names: if name == ALL_PAIRS: if box is None: topologies[name] = build_pair_graph( coord, n_node, atype, node_mask, pair_excl ) - elif K_GRID in names: - # A periodic-capable module reads the k grid for these frames. + elif K_GRID in names or PME_MESH in names: + # A periodic-capable module reads the k grid or the mesh for + # these frames. topologies[name] = _empty_pair_graph(coord, atype) else: check_complete_frames(box) @@ -568,9 +677,26 @@ def build_topologies( topologies[name] = build_k_grid( box, params["k_max"], n_node, coord=coord ) + elif name == PME_MESH: + if box is None: + topologies[name] = _empty_mesh_grid(n_node.shape[0], coord, coord) + else: + if pair_excl is not None: + raise NotImplementedError( + "the k-space sum of a periodic system is global; " + "`pair_exclude_types` cannot exclude pairs from it" + ) + params = (topology_params or {}).get(PME_MESH, {}) + if "n_mesh" not in params: + raise ValueError( + f"the {PME_MESH!r} topology requires its `n_mesh` " + f"build parameter" + ) + topologies[name] = build_mesh_grid(box, params["n_mesh"], coord=coord) else: raise ValueError( - f"unknown topology {name!r}; expected one of ({ALL_PAIRS!r}, {K_GRID!r})" + f"unknown topology {name!r}; expected one of " + f"({ALL_PAIRS!r}, {K_GRID!r}, {PME_MESH!r})" ) return topologies @@ -622,7 +748,7 @@ def flatten_topologies(topologies: dict[str, PairGraph | KGrid]) -> tuple[Array, ) -def topology_flat_fields(topology: PairGraph | KGrid) -> tuple[Array, ...]: +def topology_flat_fields(topology: PairGraph | KGrid | MeshGrid) -> tuple[Array, ...]: """The positional tensors of one topology.""" if isinstance(topology, KGrid): if topology.coord is None or topology.strain is None: @@ -638,6 +764,19 @@ def topology_flat_fields(topology: PairGraph | KGrid) -> tuple[Array, ...]: topology.coord, topology.strain, ) + if isinstance(topology, MeshGrid): + if topology.coord is None or topology.strain is None: + raise ValueError( + "a mesh grid carries no coordinates or strain; it cannot be laid " + "out positionally" + ) + return ( + topology.box, + topology.volume, + topology.coord, + topology.strain, + topology.mask, + ) return (topology.index, topology.vec, topology.mask) @@ -663,8 +802,8 @@ def unflatten_topologies( ValueError If the number of values does not match the layout of the names. """ - arity = {ALL_PAIRS: 3, K_GRID: 6} - topologies: dict[str, PairGraph | KGrid] = {} + arity = {ALL_PAIRS: 3, K_GRID: 6, PME_MESH: 5} + topologies: dict[str, PairGraph | KGrid | MeshGrid] = {} offset = 0 for name in names: if name not in arity: @@ -677,8 +816,10 @@ def unflatten_topologies( ) if name == ALL_PAIRS: topologies[name] = PairGraph(*fields) - else: + elif name == K_GRID: topologies[name] = KGrid(*fields) + else: + topologies[name] = MeshGrid(*fields) offset += arity[name] if offset != len(values): raise ValueError( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 1dd832aad1..23dc7bf818 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -2555,17 +2555,28 @@ def _eval_model_graph( topology_names = self.metadata.get("topologies", ()) if topology_names: # Padding slots of a mixed-type batch are no atoms and join - # no pair. + # no pair. The cell is torch like the coordinates: the + # topologies it namespaces feed the torch graph lower + # directly (a numpy cell would produce numpy leaves). coord_flat_t = torch.as_tensor( coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE ) + box_t = ( + torch.as_tensor( + np.asarray(box_input).reshape(n_node_t.shape[0], 3, 3), + dtype=edge_dtype, + device=DEVICE, + ) + if box_input is not None + else None + ) if self._dpmodel is not None: # The live model builds its own topologies, so the k-grid # cutoff follows the kernel parameters. topologies = self._dpmodel.atomic_model.build_topologies( coord_flat_t, n_node_t, - box_input, + box_t, atype=atype_t, node_mask=atype_t >= 0, pair_excl=self._model_pair_excl(), @@ -2573,9 +2584,10 @@ def _eval_model_graph( else: from deepmd.dpmodel.utils.neighbor_graph import ( K_GRID, + PME_MESH, ) - if K_GRID in topology_names: + if K_GRID in topology_names or PME_MESH in topology_names: raise NotImplementedError( "a periodic long-range model cannot be evaluated from " "a .pt2/.pte archive yet" @@ -2584,7 +2596,7 @@ def _eval_model_graph( topology_names, coord_flat_t, n_node_t, - box_input, + box_t, node_mask=atype_t >= 0, atype=atype_t, pair_excl=self._model_pair_excl(), diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index a2bd41aaac..6e0f837de9 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -26,6 +26,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( KGrid, + MeshGrid, NeighborGraph, PairGraph, edge_force_virial, @@ -455,7 +456,7 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: coord_leaf = strain_leaf = None if topologies: for topology in topologies.values(): - if isinstance(topology, KGrid): + if isinstance(topology, (KGrid, MeshGrid)): coord_leaf = topology.coord strain_leaf = topology.strain ff_list: list[torch.Tensor] = [] diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 3499d0410b..3bd0c788d0 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -22,6 +22,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( KGrid, + MeshGrid, PairGraph, compact_nodes, expand_node_values, @@ -764,6 +765,21 @@ def forward_common_lower_graph( coord=coord_leaf, strain=strain_leaf, ) + elif isinstance(topology, MeshGrid): + if topology.coord is None or topology.strain is None: + raise ValueError( + "the mesh-grid topology carries no coordinates or strain" + ) + # Same re-leaf contract as the k grid (see above). + coord_leaf = topology.coord.detach().requires_grad_(True) + strain_leaf = topology.strain.detach().requires_grad_(True) + re_leafed[name] = MeshGrid( + box=topology.box, + volume=topology.volume, + coord=coord_leaf, + strain=strain_leaf, + mask=topology.mask, + ) else: re_leafed[name] = PairGraph( index=topology.index, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index d90e67d36e..d1339adc4e 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -74,6 +74,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( KGrid, + MeshGrid, flatten_topologies, unflatten_topologies, ) @@ -1736,6 +1737,15 @@ def _forward_graph( topology.coord.detach().requires_grad_(True), topology.strain.detach().requires_grad_(True), ] + elif isinstance(topology, MeshGrid): + assert topology.coord is not None and topology.strain is not None + topology_vals_list += [ + topology.box, + topology.volume, + topology.coord.detach().requires_grad_(True), + topology.strain.detach().requires_grad_(True), + topology.mask, + ] else: topology_vals_list += [ topology.index, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index cc92d173ee..cf0ce1a5bf 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -718,7 +718,9 @@ def build_synthetic_topologies( """ from deepmd.dpmodel.utils.neighbor_graph import ( K_GRID, + PME_MESH, KGrid, + MeshGrid, ) index = ( @@ -739,6 +741,18 @@ def build_synthetic_topologies( coord=torch.ones(n_node_total, 3, dtype=dtype, device=device), strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), ) + elif name == PME_MESH: + if nframes is None: + raise ValueError( + f"a {PME_MESH!r} topology sample needs the frame count" + ) + topologies[name] = MeshGrid( + box=torch.eye(3, dtype=dtype, device=device).expand(nframes, 3, 3), + volume=torch.ones(nframes, dtype=dtype, device=device), + coord=torch.ones(n_node_total, 3, dtype=dtype, device=device), + strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), + mask=torch.ones(nframes, dtype=torch.bool, device=device), + ) else: topologies[name] = PairGraph( index=index, @@ -983,6 +997,16 @@ def _build_graph_dynamic_shapes( {0: n_node_total_dim}, {0: nframes_dim}, ] + elif name == "pme_mesh": + # (box, volume, coord, strain, mask): the frame axis and the node + # axis; the mesh axes are static and live on the module. + topologies += [ + {0: nframes_dim}, + {0: nframes_dim}, + {0: n_node_total_dim}, + {0: nframes_dim}, + {0: nframes_dim}, + ] else: topologies += [{1: npair_dim}, {0: npair_dim}, {0: npair_dim}] if is_native_spin: @@ -1961,11 +1985,14 @@ def _trace_and_export_impl( # yet; freeze and compress reject such models until they do. from deepmd.dpmodel.utils.neighbor_graph import ( K_GRID, + PME_MESH, ) - if K_GRID in model.atomic_model.required_topologies(): + if K_GRID in model.atomic_model.required_topologies() or ( + PME_MESH in model.atomic_model.required_topologies() + ): raise NotImplementedError( - "a periodic long-range module (the k-grid topology) cannot be " + "a periodic long-range module (the k-grid/mesh topology) cannot be " "frozen or exported yet; serve it through the Python interface" ) # The vacuum reference is resolved on the export object, folded into the diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 79538293ae..d2802db947 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3705,15 +3705,20 @@ def long_range_latent_charge() -> list[Argument]: "wave-number cutoff of the k grid of a periodic batch." ) doc_k_grid_mode = ( - "How the k grid of a periodic batch is laid out: 'ragged' enumerates " - "the wave numbers of every frame eagerly (the k axis is a dynamic " - "dimension); 'fixed' shares one integer grid over the frames, sized " - "by the largest cell of the data statistics or by `k_grid_bounds`, " - "and masks the points a frame does not reach, so the k axis has a " - "static length. The fixed grid is hard: a frame needing an index " - "beyond it (a larger cell, or a kernel that moved the cutoff out) " - "has its k-sum truncated to the grid, with a warning. Requires " - "`periodic: true` for 'fixed'." + "How the reciprocal sum of a periodic batch is evaluated: 'ragged' " + "enumerates the wave numbers of every frame eagerly (the k axis is a " + "dynamic dimension); 'fixed' shares one integer grid over the frames, " + "sized by the largest cell of the data statistics or by " + "`k_grid_bounds`, and masks the points a frame does not reach, so the " + "k axis has a static length; 'pme' spreads the charges to a static " + "real-space mesh with a B-spline and evaluates the sum by an FFT " + "(SPME), sized by the largest cell of the data statistics or by " + "`pme_mesh`. The fixed grid is hard: a frame needing an index beyond " + "it (a larger cell, or a kernel that moved the cutoff out) has its " + "k-sum truncated to the grid, with a warning; on the mesh such a " + "frame is evaluated with a coarser effective spacing (a smooth " + "accuracy degradation, without a warning). Requires `periodic: true` " + "for 'fixed' and 'pme'." ) doc_k_grid_bounds = ( "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " @@ -3728,6 +3733,16 @@ def long_range_latent_charge() -> list[Argument]: "evaluated with still follow the live parameters). The frozen cutoff " "is serialized with the model. Requires `periodic: true`." ) + doc_pme_spacing = ( + "The target real-space mesh spacing of the 'pme' mode, in Angstrom; " + "the default is the Nyquist spacing of the kernel's wave-number " + "cutoff, pi / k_max." + ) + doc_pme_order = "The B-spline order of the 'pme' charge assignment." + doc_pme_mesh = ( + "The per-axis mesh dimensions [n1, n2, n3] of the 'pme' mode, in " + "place of the data statistics." + ) return [ Argument("channels", int, optional=True, default=1, doc=doc_channels), Argument( @@ -3787,6 +3802,27 @@ def long_range_latent_charge() -> list[Argument]: default=False, doc=doc_decouple_k_grid, ), + Argument( + "pme_spacing", + float, + optional=True, + default=None, + doc=doc_pme_spacing, + ), + Argument( + "pme_order", + int, + optional=True, + default=4, + doc=doc_pme_order, + ), + Argument( + "pme_mesh", + list[int], + optional=True, + default=None, + doc=doc_pme_mesh, + ), ] diff --git a/periodic_long_range_design.md b/periodic_long_range_design.md index 1cb7573882..c58e2e0e12 100644 --- a/periodic_long_range_design.md +++ b/periodic_long_range_design.md @@ -204,6 +204,24 @@ def build_k_grid(box, alpha=None, sigma0=None, eps=1e-8): 能量;置 true 后 $k_{max}$ 在构造时按初始参数冻结(`frozen_k_max`,随模型 序列化),kernel 漂移不再影响 k 点选择,图内 $\hat K(k)$ 仍用活参数求值 (梯度不受影响)。两种模式共用;fixed 模式下解耦后 mask 也不再随漂移收紧。 +- **PME 网格模式(`k_grid_mode: "pme"`)**:不直接枚举 k 点,而是把电荷用 + p 阶 B 样条撒到全 batch 共享的静态实空间网格,FFT 后乘影响函数(SPME): + $E = \frac{k_e}{2V}\sum_{m\neq 0}\hat K(|k(m)|)\,|B(m)|^2\,|\hat\rho_{\rm mesh}(m)|^2$, + 其中 $|B(m)|^2 = \prod_d |\sum_l M_p(l+1)e^{2\pi i m_d l/n_d}|^{-2}$ 是样条去卷积 + 修正(构造时按 Essmann 公式计算),$k(m) = 2\pi B^{-1} m$($m$ 为网格模式的整数 + Miller 指标),$m=0$ 置零(tin-foil)。网格尺寸由数据统计的最大 cell 按目标 + 间距 $h$(缺省 $\pi/\_k\_max()$,尊重解耦开关;或 `pme_spacing` 显式)定 + $n_d = \lceil |a_d|/h \rceil$ 并上取 FFT 光滑数(2·3·5·7),或 `pme_mesh` + 显式跳过统计;随模型序列化。大 cell 超出统计 → 有效间距变粗(平滑的精度退化, + 力/位力仍是该近似能量的精确导数,不发 warning——网格尺寸是用户的显式责任)。代价 $O(N p^3 C + M\log M)$ 近线性,无相位矩阵显存墙。 + 拓扑 `PME_MESH` 无 ragged 轴,flatten 为 `(box, volume, coord, strain, mask)` + (arity 5,volume 同 KGrid 一样 eager 计算);应变族与直接求和相同(分数坐标 + 应变不变,撒点不带应变;$k'=S^{-T}k$、$V'=V\det S$——strain leaf 的 + `linalg.inv`/`linalg.det` 反向与 KGrid 路径同族,已验证 nf=1 符号帧数下编译 + 反向正确)。eager/compiled 训练与 Python 接口均支持。 + 实现坑(已踩过):$k(m)$ 里 $m$ 是**整数** Miller 指标(不是 $m/n$——那是 + 把网格当成只跨 $1/n_d$ 的 cell);tin-foil 的 `where` 掩码条件别写反 + (`where(keep, influence, 0)`)。 ### 5.2 图内计算(伪代码,全部纯张量运算,无 data-dependent 控制流) diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index b3c5032558..b8976801cc 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -13,6 +13,7 @@ LatentChargeFitting, ) from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, BaseInteractionKernel, LatentChargeLongRange, LESKernel, @@ -557,3 +558,207 @@ def test_decoupled_k_grid_round_trip() -> None: plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) assert not plain.decouple_k_grid assert plain.frozen_k_max is None + + +# --------------------------------------------------------------------------- +# PME mesh mode +# --------------------------------------------------------------------------- + + +def _pme_long_range(**kwargs: object) -> LatentChargeLongRange: + kwargs.setdefault("k_grid_mode", "pme") + return _periodic_long_range(**kwargs) + + +def test_pme_mesh_grid_build_and_flatten() -> None: + """The mesh grid packs the cells, flags the frames, and flattens to 4 tensors.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + build_topologies, + flatten_topologies, + unflatten_topologies, + ) + + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + box = _cells() + params = {PME_MESH: {"n_mesh": np.array([16, 16, 16]), "spacing": 1.0}} + tops = build_topologies( + ("all_pairs", PME_MESH), + coord, + np.array([2, 2]), + box, + atype=atype, + topology_params=params, + ) + mesh = tops[PME_MESH] + np.testing.assert_array_equal(mesh.box, box) + assert mesh.mask.shape == (2,) and mesh.mask.all() + assert tops["all_pairs"].index.shape[1] == 2 # guards only + flat = flatten_topologies(tops) + assert len(flat) == 3 + 5 + back = unflatten_topologies(("all_pairs", PME_MESH), flat) + np.testing.assert_array_equal(back[PME_MESH].box, box) + np.testing.assert_array_equal(back[PME_MESH].mask, mesh.mask) + # a non-periodic batch: real pairs, masked-out mesh + tops_np = build_topologies( + ("all_pairs", PME_MESH), + coord, + np.array([2, 2]), + None, + atype=atype, + topology_params=params, + ) + assert tops_np["all_pairs"].mask.any() + assert not tops_np[PME_MESH].mask.any() + # missing build parameters are an error + with pytest.raises(ValueError, match="n_mesh"): + build_topologies((PME_MESH,), coord, np.array([4]), box, atype=atype) + + +def test_pme_fft_smooth() -> None: + from deepmd.dpmodel.long_range.latent_charge import ( + _fft_smooth, + ) + + assert _fft_smooth(1) == 1 + assert _fft_smooth(17) == 18 + assert _fft_smooth(19) == 20 + assert _fft_smooth(23) == 24 + assert _fft_smooth(29) == 30 + assert _fft_smooth(31) == 32 + assert _fft_smooth(37) == 40 + + +def _pme_energy(module, charge, coord, box) -> float: + from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + frame_id_from_n_node, + ) + + n_node = np.array([coord.shape[0]]) + tops = module.build_topologies( + coord, n_node, box[None], atype=np.zeros(len(coord), dtype=np.int64) + ) + fid = frame_id_from_n_node(n_node, n_total=coord.shape[0]) + e = module._pme_branch( + charge, tops[PME_MESH], fid, np.ones(len(coord), dtype=bool), n_node + ) + return float(np.sum(e)) + + +def test_pme_energy_converges_to_direct_reference() -> None: + """The PME energy approaches the direct k-space sum as the mesh refines.""" + coord = 3.0 + np.random.default_rng(7).uniform(0.0, 1.4, (6, 3)) + box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]]) + charge = np.array([[1.0], [-0.5], [1.0], [-0.5], [-0.5], [1.0]]) + charge = charge - charge.mean(axis=0, keepdims=True) + + kernel = LESKernel(alpha=0.3) + module = LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + kernel, + None, + periodic=True, + k_grid_mode="pme", + pme_mesh=[40, 40, 40], + ) + # converged direct k-space reference (tol 1e-10) + k_max = kernel.kspace_cutoff(1e-10) + rec = 2.0 * np.pi * np.linalg.inv(box) + n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) + grids = np.meshgrid(*[np.arange(-n, n + 1) for n in n_d], indexing="ij") + n_all = np.stack(grids, -1).reshape(-1, 3) + k = n_all @ rec.T + k2 = (k * k).sum(-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge + s_sq = (np.abs(rho) ** 2).sum(-1) + k_hat = kernel.kspace(k2) + volume = abs(np.linalg.det(box)) + k0 = 2.0 * 0.3 / np.sqrt(np.pi) + e_ref = COULOMB_CONSTANT * ( + (k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum() + ) + errors = [] + for dims in ([16, 16, 16], [40, 40, 40]): + module.pme_mesh = np.asarray(dims) + module._pme_n = tuple(dims) + errors.append(abs(_pme_energy(module, charge, coord, box) - e_ref)) + assert errors[1] < errors[0] # refinement improves the energy + assert errors[1] < 1e-2 + + +def test_pme_mesh_stats_size_the_mesh() -> None: + """The data statistics size the mesh by the largest sampled cell.""" + import math + + long_range = _pme_long_range() + coord = np.zeros((9, 3)) + with pytest.raises(ValueError, match="not sized yet"): + long_range.build_topologies( + coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) + ) + long_range.compute_input_stats( + [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}] + ) + spacing = math.pi / long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + a_max = np.max(np.linalg.norm(_cells(), axis=-1), axis=0) + for d in range(3): + assert long_range.pme_mesh[d] >= math.ceil(a_max[d] / spacing) + # FFT-smooth dimensions + m = int(long_range.pme_mesh[d]) + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + assert m == 1 + topologies = long_range.build_topologies( + coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) + ) + assert topologies["pme_mesh"].mask.all() + # explicit mesh skips the statistics + explicit = _pme_long_range(pme_mesh=[16, 16, 16]) + assert explicit._pme_dims() == (16, 16, 16) + # statistics without a cell cannot size the mesh + unsized = _pme_long_range() + with pytest.raises(ValueError, match="no simulation cell"): + unsized.compute_input_stats([{"coord": np.zeros((1, 3))}]) + + +def test_pme_mode_configuration() -> None: + """The pme mode requires periodicity, and its parameters the pme mode.""" + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + k_grid_mode="pme", + ) + with pytest.raises(ValueError, match="pme_mesh"): + _periodic_long_range(pme_mesh=[8, 8, 8]) + with pytest.raises(ValueError, match="pme_spacing"): + _periodic_long_range(pme_spacing=1.0) + with pytest.raises(ValueError, match="pme_order"): + _pme_long_range(pme_order=1) + with pytest.raises(ValueError, match="pme_mesh"): + _pme_long_range(pme_mesh=[8, 8]) + with pytest.raises(ValueError, match="k_grid_mode"): + _periodic_long_range(k_grid_mode="dense") + + +def test_pme_round_trip() -> None: + """The record keeps the pme mode, the sized mesh and the spacing.""" + long_range = _pme_long_range(pme_spacing=1.0, pme_order=5) + long_range.compute_input_stats([{"box": _cells()}]) + record = long_range.serialize() + assert record["k_grid_mode"] == "pme" + assert record["pme_order"] == 5 + assert record["pme_spacing"] == 1.0 + restored = LatentChargeLongRange.deserialize(record) + assert restored.k_grid_mode == "pme" + assert restored.pme_order == 5 + assert restored.pme_spacing == 1.0 + np.testing.assert_array_equal(restored.pme_mesh, long_range.pme_mesh) + assert restored._pme_dims() == long_range._pme_dims() diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 7388f90b6c..4496725dc8 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -559,6 +559,57 @@ def test_checkpoint_inference_matches_direct_forward(tmp_path, kernel: str) -> N np.testing.assert_allclose(moved[0], actual["energy"], rtol=0, atol=1e-12) +@pytest.mark.parametrize("grid_mode", ["ragged", "pme"]) +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_checkpoint_inference_matches_direct_forward( + tmp_path, kernel: str, grid_mode: str +) -> None: + """DeepPot on a periodic checkpoint builds the periodic topology and matches.""" + if grid_mode == "pme": + model = _periodic_pme_model(kernel) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + config["long_range"]["pme_mesh"] = [32, 32, 32] + else: + model = _periodic_model(kernel) + config = _config(kernel) + config["long_range"]["periodic"] = True + checkpoint = tmp_path / f"periodic_{kernel}_{grid_mode}.pt" + wrapper = ModelWrapper(model, model_params=config) + torch.save({"model": wrapper.state_dict()}, checkpoint) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + expected = model(coord_t, atype_t, box=box_t, do_atomic_virial=True) + dp = DeepPot(str(checkpoint), auto_batch_size=False) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + actual = dict( + zip( + names, + dp.eval( + coord, + box, + atype, + atomic=True, + mixed_type=True, + ), + strict=True, + ) + ) + for name in names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-9, + atol=1e-9, + err_msg=name, + ) + + @pytest.mark.parametrize( ("kernel", "source", "suffix"), [("les", "neutral", ".pt2"), ("sog", "neutral", ".pte"), ("les", "fparam", ".pte")], @@ -1415,3 +1466,216 @@ def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: assert kernel_grads and all( grad is not None and torch.isfinite(grad).all() for grad in kernel_grads ) + + +# --------------------------------------------------------------------------- +# Periodic PME mesh mode +# --------------------------------------------------------------------------- + + +def _periodic_pme_model( + kernel: str, boxes: np.ndarray | None = None, **kwargs: object +) -> torch.nn.Module: + """A periodic model in the PME mesh mode. + + Sized by the given cells through the data statistics, or by an explicit + ``pme_mesh`` in ``kwargs``. + """ + pme_kwargs = { + key: kwargs.pop(key) + for key in ("pme_mesh", "pme_spacing", "pme_order") + if key in kwargs + } + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + config["long_range"].update(pme_kwargs) + if boxes is None: + config["long_range"].setdefault("pme_mesh", [32, 32, 32]) + model = get_model(config).to(env.DEVICE).eval() + if boxes is not None: + boxes_t = torch.as_tensor(np.asarray(boxes, dtype=np.float64).reshape(-1, 9)) + model.atomic_model.long_range.compute_input_stats([{"box": boxes_t}]) + return model + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_matches_kspace_reference(kernel: str) -> None: + """With fixed charges the PME energy approaches the converged k-space sum.""" + model = _periodic_pme_model(kernel, constraint=None, pme_mesh=[48, 48, 48]) + _set_constant_charges(model, [1.0, -0.5]) + coord, atype, box = _periodic_case() + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference(coord[0], charge, box[0], kernel) + # the PME energy is a mesh approximation of the converged lattice sum + np.testing.assert_allclose(out["energy"].item(), expected, atol=1e-2) + np.testing.assert_allclose(out["atom_energy"].sum().item(), expected, atol=1e-2) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_force_virial_match_finite_differences(kernel: str) -> None: + """Force and strain virial are the exact derivatives of the PME energy.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + out = model(coord_t, atype_t, box=box_t) + + def energy(c: np.ndarray, b: np.ndarray) -> float: + return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item() + + h = 1e-5 + force_fd = np.zeros((6, 3)) + for i in range(6): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus, box) - energy(minus, box)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-5 + ) + + # the strain family of the mesh strain leaf: r -> (I + eps) r and + # B -> B (I + eps)^T leave the fractional coordinates invariant + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain).T + minus = coord @ (np.eye(3) - strain).T + box_plus = box @ (np.eye(3) + strain).T + box_minus = box @ (np.eye(3) - strain).T + virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / ( + 2 * h + ) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-4 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_frames_are_independent(kernel: str) -> None: + """Each frame's mesh energy is independent of the rest of the batch.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.1], axis=0) + batched = model(_tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2)) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_serves_non_periodic_batches(kernel: str) -> None: + """A PME model keeps its non-periodic pair evaluation.""" + model = _periodic_pme_model(kernel) + plain = _model(kernel) + plain.load_state_dict(model.state_dict(), strict=False) + coord, atype = _batch() + mesh_out = model(coord, atype) + plain_out = plain(coord, atype) + torch.testing.assert_close(mesh_out["energy"], plain_out["energy"]) + torch.testing.assert_close(mesh_out["force"], plain_out["force"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_needs_statistics_or_bounds(kernel: str) -> None: + """Without statistics or an explicit mesh the mesh is not sized.""" + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + model = get_model(config).to(env.DEVICE).eval() + coord, atype, box = _periodic_case() + with pytest.raises(ValueError, match="not sized yet"): + model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + sized = _periodic_pme_model(kernel, box) + out = sized(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + # the statistics-sized model evaluates exactly like a model given the + # same mesh explicitly + mesh = sized.atomic_model.long_range.pme_mesh.tolist() + reference = _periodic_pme_model(kernel, pme_mesh=mesh) + reference.load_state_dict(sized.state_dict()) + expected = reference(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + torch.testing.assert_close(out["energy"], expected["energy"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_serialization_round_trip(kernel: str) -> None: + """The record keeps the pme mode and the sized mesh, and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + coord, atype, box = _periodic_case() + model = _periodic_pme_model(kernel, box) + record = model.serialize() + assert record["long_range"]["k_grid_mode"] == "pme" + assert "pme_mesh" in record["long_range"] + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.long_range.k_grid_mode == "pme" + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense") + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_pme_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the PME model like the eager one.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + # the backward reaches the charge network and the kernel through the + # compiled graph + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + assert all(torch.isfinite(grad).all() for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_checkpoint_restores_sized_mesh(kernel: str) -> None: + """A checkpoint restores the stats-sized mesh through the state dict alone.""" + coord, atype, box = _periodic_case() + trained = _periodic_pme_model(kernel, box) + state = trained.state_dict() + # a fresh model from the bare config is unsized + fresh = _periodic_pme_model(kernel, pme_mesh=None) + with pytest.raises(ValueError, match="not sized yet"): + fresh(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + fresh.load_state_dict(state) + assert fresh.atomic_model.long_range._pme_dims() is not None + out = fresh(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + expected = trained(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + torch.testing.assert_close(out["energy"], expected["energy"]) From 17a102d01355ad59e6fbd6072782c7d5a282dc7f Mon Sep 17 00:00:00 2001 From: yuzhibytedance Date: Fri, 23 Oct 2026 10:45:23 +0000 Subject: [PATCH 25/30] perf(long-range): eliminate GEMMs from the PME branch and halve the FFT - evaluate |k|^2 as the quadratic form m^T W m with the per-frame reciprocal metric W (six static monomials times six per-frame coefficients, all elementwise) instead of batched 3x3 matmuls, which were 56% of the training step on the profile; the strain folds into the metric through A = S^-T B^-1 - transform the (real) charge mesh with rfftn and sum the half spectrum with the hermitian doubling weight - fractional coordinates by elementwise multiply-sum instead of a batched GEMV f64 40^3 mesh: 34 -> 6 ms forward, 44 -> 16 ms training step --- deepmd/dpmodel/long_range/latent_charge.py | 118 +++++++++++++-------- periodic_long_range_design.md | 4 + 2 files changed, 79 insertions(+), 43 deletions(-) diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index 1e85ad4440..e470dc9c2e 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -696,18 +696,17 @@ def _kspace_branch( ) return e_rec + e_self - def _pme_b_factor(self, xp: Any, dtype: Any, device: Any) -> Array: - r"""The SPME assignment correction :math:`|B(m)|^2` on the mesh. + def _pme_b_factor(self, xp: Any, dtype: Any, device: Any, m_axes: list) -> Array: + r"""The SPME assignment correction :math:`|B(m)|^2` on the mesh modes. Per axis :math:`|D_d(m)|^2 = |\sum_l M_p(l+1) e^{2\pi i m l / n_d}|^2` (equal to :math:`\mathrm{sinc}(m/n_d)^{2p}`), the full factor the - three-axis outer product, in the ``fftfreq`` index order. + three-axis outer product. ``m_axes`` holds the integer mode values of + each axis (the third axis carries only the non-negative rfft modes). """ w = _bspline_integer_weights(self.pme_order) axes = [] - for n_d in self._pme_dims(): - m = xp.arange(n_d, dtype=dtype, device=device) - m = xp.where(m > n_d // 2, m - n_d, m) + for m, n_d in zip(m_axes, self._pme_dims(), strict=True): theta = 2.0 * math.pi * m / n_d re = xp.zeros_like(theta) im = xp.zeros_like(theta) @@ -730,14 +729,21 @@ def _pme_branch( Evaluates the same lattice sum as :meth:`_kspace_branch` by the SPME algorithm: the charges spread to a static real-space mesh with a - B-spline of order :attr:`pme_order`, the mesh transforms by an FFT, - and the energy convolves with the kernel transform corrected by the - assignment factor (:meth:`_pme_b_factor`). The strain leaf enters the - coordinates by ``I + strain`` and the cells by - ``B (I + strain)^T`` — the fractional coordinates are strain- - invariant, so the assignment is unstrained and the virial comes from - the reciprocal vectors and the volume alone. Non-periodic frames - (``mesh.mask`` false) carry no mesh energy and no self term. + B-spline of order :attr:`pme_order`, the mesh transforms by a + real-input FFT (only the non-negative :math:`m_3` half-spectrum), and + the energy convolves with the kernel transform corrected by the + assignment factor (:meth:`_pme_b_factor`) and the hermitian doubling + weight of the interior :math:`m_3` planes. No wave-number vector is + ever formed: the kernel reads :math:`|k|^2 = (2\pi)^2 m^\mathsf{T} W m` + with the per-frame metric :math:`W = (B^{-1})^\mathsf{T} B^{-1}` — + six static monomials of the mode indices times six per-frame metric + entries, all elementwise. The strain leaf enters the metric by + :math:`W' = A^\mathsf{T} A` with :math:`A = S^{-\mathsf{T}} B^{-1}` + (:math:`S = I +` strain) and the volume by :math:`\det S` — the + fractional coordinates are strain-invariant, so the assignment is + unstrained and the virial comes from the metric and the volume alone. + Non-periodic frames (``mesh.mask`` false) carry no mesh energy and no + self term. """ xp = array_api_compat.array_namespace(charge) device = array_api_compat.device(charge) @@ -751,8 +757,9 @@ def _pme_branch( volume = xp.astype(mesh.volume, charge.dtype) b_inv = _inv3(xp, box) # (nf, 3, 3); columns are g_d / 2pi b_inv_node = xp.take(b_inv, frame_id, axis=0) # (N, 3, 3) - # fractional coordinates, wrapped to [0, 1): frac = r B^-1 - frac = xp.matmul(coord[:, None, :], b_inv_node)[:, 0, :] + # fractional coordinates, wrapped to [0, 1): frac_d = sum_j r_j B^-1_jd + # (elementwise; a batched 3x3 GEMV is the worst GEMM shape on GPU) + frac = xp.sum(coord[:, :, None] * b_inv_node, axis=1) frac = frac - xp.floor(frac) # per-axis assignment: weights (N, p) and mesh indices (N, p) weights = [] @@ -792,45 +799,70 @@ def _pme_branch( q_mesh, xp.reshape(target, (-1,)), xp.reshape(value, (-1,)) ) q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) - rho = xp.fft.fftn(q_mesh, axes=(-3, -2, -1)) + # the charge mesh is real: only the non-negative m3 half-spectrum is + # transformed and summed (with the doubling weight below) + rho = xp.fft.rfftn(q_mesh, axes=(-3, -2, -1)) rho_sq = xp.real(rho) ** 2 + xp.imag(rho) ** 2 - rho_sq = xp.reshape(rho_sq, (nf, n_channel, m_total)) - # the reciprocal vectors of every frame on the mesh: mode m is the - # integer Miller index, k(m) = 2 pi B^-1 m - h_axes = [] - for n_d in (n1, n2, n3): - m = xp.arange(n_d, dtype=charge.dtype, device=device) - m = xp.where(m > n_d // 2, m - n_d, m) - h_axes.append(m) - g1, g2, g3 = xp.meshgrid(*h_axes, indexing="ij") - h = xp.stack( - [xp.reshape(g1, (-1,)), xp.reshape(g2, (-1,)), xp.reshape(g3, (-1,))], - axis=-1, - ) # (M, 3) - k_vec = 2.0 * math.pi * xp.matmul(h[None, :, :], xp.matrix_transpose(b_inv)) + m_half = n1 * n2 * (n3 // 2 + 1) + rho_sq = xp.reshape(rho_sq, (nf, n_channel, m_half)) + # the reciprocal-lattice metric: |k|^2 = (2 pi)^2 m^T W m with + # W = (B^-1)^T B^-1; the strain folds in through A = S^-T B^-1 (the + # linalg inverse of S keeps its registered backward, safe under a + # symbolic frame count) if mesh.strain is not None: - # The strain family of the virial: B' = B S^T with S = I + strain. - # The fractional coordinates are exactly strain-invariant, so the - # assignment above needs no strain; the reciprocal vectors and the - # volume pick it up: k' = S^-T k, V' = V det(S). The 3x3 inverse - # and determinant stay on linalg — its registered backward handles - # a symbolic frame count, unlike hand-written select/cross forms. eye = xp.eye(3, dtype=charge.dtype, device=device) s = xp.astype(mesh.strain, charge.dtype) + eye # (nf, 3, 3) s_inv = xp.linalg.inv(s) - # rows: k'_i = sum_j k_j (S^-1)[j, i] - k_vec = xp.matmul(k_vec, s_inv) + a = xp.matmul(xp.matrix_transpose(s_inv), b_inv) + w_metric = xp.matmul(xp.matrix_transpose(a), a) volume = volume * xp.linalg.det(s) - k_sq = xp.reshape(xp.sum(k_vec * k_vec, axis=-1), (-1,)) # (nf * M,) + else: + w_metric = xp.matmul(xp.matrix_transpose(b_inv), b_inv) + # static integer mode grids: full fftfreq order on axes 1, 2; the rfft + # axis keeps m3 = 0 .. n3//2 + m_axes = [] + for d, n_d in enumerate((n1, n2, n3)): + m = xp.arange(n_d, dtype=charge.dtype, device=device) + if d == 2: + m = m[: n_d // 2 + 1] + else: + m = xp.where(m > n_d // 2, m - n_d, m) + m_axes.append(m) + g1, g2, g3 = xp.meshgrid(*m_axes, indexing="ij") + m1f = xp.reshape(g1, (-1,)) + m2f = xp.reshape(g2, (-1,)) + m3f = xp.reshape(g3, (-1,)) + # six static monomials x six per-frame metric entries, no GEMM + two_pi_sq = (2.0 * math.pi) ** 2 + k_sq = two_pi_sq * ( + w_metric[:, 0, 0][:, None] * (m1f * m1f)[None, :] + + w_metric[:, 1, 1][:, None] * (m2f * m2f)[None, :] + + w_metric[:, 2, 2][:, None] * (m3f * m3f)[None, :] + + 2.0 * w_metric[:, 0, 1][:, None] * (m1f * m2f)[None, :] + + 2.0 * w_metric[:, 0, 2][:, None] * (m1f * m3f)[None, :] + + 2.0 * w_metric[:, 1, 2][:, None] * (m2f * m3f)[None, :] + ) # (nf, m_half) # the zero mode sits at flat index 0; guard the singular kernel there safe = xp.where(k_sq == 0.0, xp.ones_like(k_sq), k_sq) - influence = xp.reshape(self.kernel.kspace(safe), (nf, m_total)) - influence = influence * self._pme_b_factor(xp, charge.dtype, device)[None, :] + influence = xp.reshape( + self.kernel.kspace(xp.reshape(safe, (-1,))), (nf, m_half) + ) + influence = ( + influence * self._pme_b_factor(xp, charge.dtype, device, m_axes)[None, :] + ) + # the spectrum of a real mesh is hermitian: interior m3 planes count + # twice, the m3 = 0 and Nyquist planes once + w_rfft = xp.where( + (m3f > 0) & (m3f < n3 / 2.0), + 2.0 * xp.ones_like(m3f), + xp.ones_like(m3f), + ) + influence = influence * w_rfft[None, :] # the zero mode sits at flat index 0 (tin-foil); keep everything else keep = xp.concat( [ xp.zeros((1,), dtype=xp.bool, device=device), - xp.ones((m_total - 1,), dtype=xp.bool, device=device), + xp.ones((m_half - 1,), dtype=xp.bool, device=device), ] ) influence = xp.where(keep[None, :], influence, xp.zeros_like(influence)) diff --git a/periodic_long_range_design.md b/periodic_long_range_design.md index c58e2e0e12..7d18797848 100644 --- a/periodic_long_range_design.md +++ b/periodic_long_range_design.md @@ -214,6 +214,10 @@ def build_k_grid(box, alpha=None, sigma0=None, eps=1e-8): $n_d = \lceil |a_d|/h \rceil$ 并上取 FFT 光滑数(2·3·5·7),或 `pme_mesh` 显式跳过统计;随模型序列化。大 cell 超出统计 → 有效间距变粗(平滑的精度退化, 力/位力仍是该近似能量的精确导数,不发 warning——网格尺寸是用户的显式责任)。代价 $O(N p^3 C + M\log M)$ 近线性,无相位矩阵显存墙。 + 性能实现(profile 驱动):$|k|^2$ 走度规二次型 $(2\pi)^2 m^\mathsf{T}W m$ + (6 个静态单项式 × 6 个逐帧系数,全 elementwise),**不构造 k 矢量**—— + batched 3×3 GEMM 曾占训练步 56%;FFT 用 rfftn 半谱($m_3 \ge 0$,内部面加倍 + 权重 ×2、$m_3=0$ 与 Nyquist 面 ×1);f64 下 40³ 网格前向 34→6ms。 拓扑 `PME_MESH` 无 ragged 轴,flatten 为 `(box, volume, coord, strain, mask)` (arity 5,volume 同 KGrid 一样 eager 计算);应变族与直接求和相同(分数坐标 应变不变,撒点不带应变;$k'=S^{-T}k$、$V'=V\det S$——strain leaf 的 From 56e6732d0a6e533e068e9dfcfa5bbeb7181461a4 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 01:05:59 +0800 Subject: [PATCH 26/30] feat(long-range): serve periodic systems in frozen models, C++ and LAMMPS Periodic long-range modules (`periodic: true`) now freeze, compress and run in DeePMD-kit C++ and LAMMPS, and the reciprocal-space evaluation is reorganized: - Topologies: every topology is a dataclass that declares its geometry fields, its autograd leaves and its dynamic axes, from which casting, re-leafing, the layout of the graph-lower inputs and the export shapes follow (neighbor_graph/topology.py); the pair topology and the reciprocal-space ones live in pair_graph.py and reciprocal_graph.py. A model gives its build parameters through topology_params() instead of overriding build_topologies, and an archive records them as metadata. - The mesh of a periodic frame is sized by its cell at a spacing, per batch, its three mode axes being inputs of dynamic (even) length, in place of a mesh sized once from the data statistics and frozen into the model; `pme_mesh` is gone. The defaults, B-splines of order 8 at three quarters of the Nyquist spacing of the wave-number cutoff, keep the mesh energy and forces within a few 1e-6 of the sum over the wave numbers; order 4 at the Nyquist spacing, the former default, was 1e-3 off in energy and 7e-3 in the forces. - dpmodel/long_range/reciprocal.py evaluates the k-space and mesh sums. An atom carries half of the interaction of its charges with all others, E_i = q_i phi_i / 2, as for the direct sum, in place of an equal share of the frame energy. The reciprocal energy is formed in float64: the lattice sum and the self term cancel to a small fraction of either, which cost a float32 compressed model 1e-3 eV. The cell part of the atomic virial is shared over the owned nodes only, not over padding and halo nodes. - A frozen model evaluates every periodic frame on the mesh of its cell, whatever its training mode, with the spacing fixed at freeze time. - The fixed k grid keeps its per-axis bounds, an array of fixed shape, rather than the enumerated grid, so that a checkpoint of a sized grid loads into a fresh model and DDP broadcasts it. - A pretrained model without a long-range module fine-tunes into one with it: the new module takes the input statistics of the fine-tuning data before the bias adjustment. - Rejected at construction: pair_exclude_types beside a module that serves periodic systems, whose reciprocal-space sum is global, and a linear_ener composition with a periodic module beside another long-range module. Rejected at evaluation: a cell without volume, and a model deviation of a non-periodic system in LAMMPS over models of which only some carry a long-range module. - CUDA: deepmd::latent_charge_direct_sum becomes deepmd::latent_charge_sum, which takes an optional cell per frame. Without one it is the direct sum as before; with one it evaluates the smooth particle-mesh Ewald sum (mesh_sum.cu): the B-spline stencils are evaluated once per node and axis in float64, the charges spread to the mesh in 64-bit fixed point, so that the result does not depend on the order of the atomic additions, the transforms run in float64, and the virial is that of the strain of the cell. It serves the constraint and several MPI ranks as the direct sum does. Its outputs agree with the float64 reference to about 5e-8. - C++: the compact canonical lower takes the cell as an eleventh input, and DP_DeepPotComputeCanonicalGraphGPU2 and the matching C++ entries take the host cell after the communicator; the graph lower receives the mesh of the cell. A cell requires a module that serves periodic systems; with one, the halo nodes of a single rank are images of its atoms. - LAMMPS: pair_style deepmd passes the cell of a system periodic in every dimension and none for a non-periodic one, and rejects a partially periodic box for a long-range model; deepmd/kk passes it to the compact canonical lower, with positions relative to the origin of the box, and rejects a periodic system for an uncompressed long-range model. - The author's design notes and AGENTS.md lines leave the repository. Verified: CPU and CUDA test suites of the module, the operator and fine-tuning; C++ and C API long-range tests, periodic graph and compact lowers on float and double, on CUDA and the graph lower on the CPU; LAMMPS deepmd and deepmd/kk on 1, 2 and 4 ranks, periodic and cluster, against the Python references. --- AGENTS.md | 1 - .../dpmodel/atomic_model/base_atomic_model.py | 16 +- .../dpmodel/atomic_model/dp_atomic_model.py | 66 +- .../atomic_model/linear_atomic_model.py | 55 +- deepmd/dpmodel/long_range/base_long_range.py | 4 + deepmd/dpmodel/long_range/kernel.py | 6 +- deepmd/dpmodel/long_range/latent_charge.py | 677 +++------- deepmd/dpmodel/long_range/reciprocal.py | 379 ++++++ deepmd/dpmodel/model/make_model.py | 49 +- .../dpmodel/utils/neighbor_graph/__init__.py | 42 +- .../utils/neighbor_graph/pair_graph.py | 658 +--------- .../utils/neighbor_graph/reciprocal_graph.py | 527 ++++++++ .../dpmodel/utils/neighbor_graph/topology.py | 270 ++++ deepmd/pt_expt/infer/deep_eval.py | 73 +- deepmd/pt_expt/kernels/dpa4c/canonical.py | 6 +- .../kernels/dpa4c/canonical_long_range.py | 15 +- deepmd/pt_expt/kernels/long_range.py | 262 +++- deepmd/pt_expt/long_range/latent_charge.py | 62 +- deepmd/pt_expt/model/edge_transform_output.py | 222 ++-- deepmd/pt_expt/model/ener_model.py | 74 +- deepmd/pt_expt/model/make_model.py | 113 +- deepmd/pt_expt/train/training.py | 85 +- deepmd/pt_expt/utils/serialization.py | 220 ++-- deepmd/utils/argcheck.py | 93 +- doc/model/long-range.md | 145 ++- periodic_long_range_design.md | 334 ----- source/api_c/include/c_api.h | 30 +- source/api_c/include/deepmd.hpp | 11 +- source/api_c/src/c_api.cc | 11 +- source/api_cc/include/DeepPot.h | 36 +- source/api_cc/include/DeepPotPTExpt.h | 44 +- source/api_cc/src/DeepPot.cc | 35 +- source/api_cc/src/DeepPotPTExpt.cc | 177 ++- source/api_cc/src/longRangePTExpt.h | 100 ++ .../test_deeppot_dpa4c_long_range_ptexpt.cc | 129 +- source/lmp/pair_deepmd.cpp | 30 +- source/lmp/pair_deepmd.h | 3 + source/lmp/pair_deepmd_kokkos.cpp | 47 +- source/lmp/pair_deepmd_kokkos.h | 6 +- .../tests/test_lammps_dpa4c_long_range_pt2.py | 170 ++- source/op/pt/CMakeLists.txt | 2 + source/op/pt/long_range/charge_sum.h | 68 + source/op/pt/long_range/direct_sum.cu | 367 +----- source/op/pt/long_range/domain_exchange.cc | 9 +- source/op/pt/long_range/domain_exchange.h | 12 +- source/op/pt/long_range/frame_sums.cuh | 181 +++ source/op/pt/long_range/latent_charge_sum.cu | 228 ++++ source/op/pt/long_range/mesh_sum.cu | 1113 +++++++++++++++++ source/op/pt/long_range/ops.cc | 80 +- .../tests/common/dpmodel/test_long_range.py | 266 ++-- source/tests/infer/gen_dpa4c_long_range.py | 120 +- .../kernels/test_long_range_direct_sum.py | 244 ---- .../pt_expt/kernels/test_long_range_sum.py | 421 +++++++ source/tests/pt_expt/model/test_long_range.py | 327 +++-- source/tests/pt_expt/test_finetune.py | 73 ++ 55 files changed, 5602 insertions(+), 3192 deletions(-) create mode 100644 deepmd/dpmodel/long_range/reciprocal.py create mode 100644 deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py create mode 100644 deepmd/dpmodel/utils/neighbor_graph/topology.py delete mode 100644 periodic_long_range_design.md create mode 100644 source/op/pt/long_range/charge_sum.h create mode 100644 source/op/pt/long_range/frame_sums.cuh create mode 100644 source/op/pt/long_range/latent_charge_sum.cu create mode 100644 source/op/pt/long_range/mesh_sum.cu delete mode 100644 source/tests/pt_expt/kernels/test_long_range_direct_sum.py create mode 100644 source/tests/pt_expt/kernels/test_long_range_sum.py diff --git a/AGENTS.md b/AGENTS.md index bcf7a9e945..617c904f11 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -165,7 +165,6 @@ source/ # C++ source code and tests - **Configuration**: Use `input_torch.json` format typically - **Training**: `dp --pt train input_torch.json` - **Requirements**: `torch` package -- **Long-range latent-charge module**: a standard energy model (DPA4C descriptor) takes a `long_range` section (`type: latent_charge`, kernels `les`/`sog`, optional `charge_constraint`) that adds an interacting-charge energy evaluated on complete systems — directly over all pairs for non-periodic frames, and in reciprocal space for periodic frames when `periodic: true`. The k grid has two layouts selected by `k_grid_mode`: `"ragged"` (default; k grid enumerated eagerly per batch from the kernel's `kspace_cutoff`, tin-foil boundary, neutral channels required) and `"fixed"` (one integer grid shared over frames, sized by the largest cell of the data statistics or by `k_grid_bounds`, with a mask selecting each frame's k points; out-of-grid cells are truncated to the grid with a warning) and `"pme"` (charges spread to a static real-space mesh by a B-spline of order `pme_order`, FFT, kernel transform with the SPME deconvolution factor; mesh sized by the largest cell of the data statistics or by `pme_mesh`, default spacing `pme_spacing` = pi/k_max; near-linear in system size). The wave-number cutoff is re-derived from the live (trainable) kernel every batch unless `decouple_k_grid: true` freezes it at the initial parameters (the kernel values still follow the live parameters; the frozen cutoff is serialized). Periodic evaluation supports eager and compiled training plus the Python interface; freeze/export, compression, C++ and LAMMPS reject periodic long-range models. Tests: `pytest source/tests/pt_expt/model/test_long_range.py`, `source/tests/common/dpmodel/test_long_range.py`. ### JAX Backend diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 292d5eebf7..90ba87a14c 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -243,6 +243,16 @@ def required_topologies(self) -> tuple[str, ...]: """ return () + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of :meth:`required_topologies`. + + Generic capability (concrete default: none) in the sense of + :meth:`required_topologies`: a module whose topology depends on its + own state, such as the wave-number cutoff of a periodic long-range + module, which follows its kernel, gives it here, by topology name. + """ + return {} + def build_topologies( self, coord: Any, @@ -255,9 +265,8 @@ def build_topologies( ) -> dict[str, Any]: """Build the topologies of :meth:`required_topologies` for one batch. - The default builds them from their names; a module that owns build - parameters (e.g. the wave-number cutoff of the k grid of a periodic - long-range module) overrides this so they stay with it. + See :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`; + the build parameters are those of :meth:`topology_params`. """ from deepmd.dpmodel.utils.neighbor_graph import ( build_topologies, @@ -271,6 +280,7 @@ def build_topologies( node_mask, atype=atype, pair_excl=pair_excl, + topology_params=self.topology_params(), ) def supports_native_spin(self) -> bool: diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index a6b633d0cb..564b6fd56a 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -173,6 +173,10 @@ def __init__( def _check_long_range(self) -> None: """Validate that the long-range module can join this descriptor and fitting.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + ) + check_long_range_support(self.descriptor) if not self.descriptor.uses_graph_lower(): raise NotImplementedError( @@ -189,6 +193,14 @@ def _check_long_range(self) -> None: "`atom_exclude_types` is not supported with a long-range module, " "which attributes each pair energy to both atoms of the pair" ) + if self.pair_exclude_types and any( + name != ALL_PAIRS for name in self.long_range.required_topologies() + ): + raise NotImplementedError( + "`pair_exclude_types` is not supported with a long-range module " + "that serves periodic systems, whose reciprocal-space sum is " + "global and excludes no pair" + ) if self.long_range.uses_charge_spin() and not self.add_chg_spin_ebd: raise ValueError( "the long-range module reads the total charge from " @@ -273,28 +285,11 @@ def required_topologies(self) -> tuple[str, ...]: return () return self.long_range.required_topologies() - def build_topologies( - self, - coord: Array, - n_node: Array, - box: Array | None, - *, - atype: Array, - node_mask: Array | None = None, - pair_excl: Any = None, - ) -> dict[str, Any]: - """Build the topologies of :meth:`required_topologies` for one batch. - - The module builds its own topologies so that its build parameters - (e.g. the wave-number cutoff of the k grid, which follows the kernel) - stay with it. A model without a long-range module declares no - topologies. - """ + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the long-range module.""" if self.long_range is None: return {} - return self.long_range.build_topologies( - coord, n_node, box, atype=atype, node_mask=node_mask, pair_excl=pair_excl - ) + return self.long_range.topology_params() def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" @@ -415,12 +410,6 @@ def enable_compression( ) if self.long_range is not None and self.long_range.uses_charge_spin(): self.long_range.fold_charge_state(self.descriptor.get_default_chg_spin()) - if self.long_range is not None and getattr(self.long_range, "periodic", False): - raise NotImplementedError( - "a periodic long-range module evaluates its reciprocal-space sum " - "with a k grid the compressed path does not serve; compress only " - "non-periodic long-range models" - ) def forward_atomic( self, @@ -985,6 +974,31 @@ def compute_fitting_input_stat( stat_file_path=stat_file_path, ) + def compute_long_range_input_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the long-range module alone. + + A fine-tuning run that adds a long-range module to a pretrained model + calls it for the new module, whose statistics, unlike those of the + inherited descriptor and fitting, the pretrained model cannot provide. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data + systems. + stat_file_path + The path to the stat file. + """ + self.long_range.compute_input_stats( + self._make_wrapped_sampler(sampled_func), + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) + def serialize(self) -> dict: dd = super().serialize() dd.update( diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index eefa809590..1356569b8e 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.array_api import ( Array, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, +) from deepmd.dpmodel.utils.nlist import ( build_multiple_neighbor_list, get_multiple_nlist_key, @@ -131,6 +134,23 @@ def __init__( "`atom_exclude_types` is not supported with a sub-model carrying " "a long-range module" ) + # The composition builds every topology once for all sub-models: the + # pairs of a periodic batch are guards only, which would silence a + # module serving non-periodic systems alone, and a reciprocal-space + # topology is built with the parameters of a single module. + declared = [m.required_topologies() for m in models if m.has_long_range()] + periodic = any(name != ALL_PAIRS for names in declared for name in names) + if len(declared) > 1 and periodic: + raise NotImplementedError( + "a sub-model whose long-range module serves periodic systems " + "takes no other sub-model carrying a long-range module beside it" + ) + if self.pair_exclude_types and periodic: + raise NotImplementedError( + "`pair_exclude_types` is not supported with a sub-model whose " + "long-range module serves periodic systems, whose " + "reciprocal-space sum is global and excludes no pair" + ) self.models = models self.type_map = type_map @@ -423,36 +443,17 @@ def required_topologies(self) -> tuple[str, ...]: ) ) - def build_topologies( - self, - coord: Array, - n_node: Array, - box: Array | None, - *, - atype: Array, - node_mask: Array | None = None, - pair_excl: Any = None, - ) -> dict[str, Any]: - """Build the topologies of :meth:`required_topologies` for one batch. + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the sub-models. - The first child declaring a name builds it; every child declaring the - same name must build it identically. + At most one sub-model declares a reciprocal-space topology, the only + kind with build parameters, so that the parameters of every topology + come from one sub-model. """ - topologies: dict[str, Any] = {} + params: dict[str, dict[str, Any]] = {} for model in self.models: - if not hasattr(model, "build_topologies"): - continue - built = model.build_topologies( - coord, - n_node, - box, - atype=atype, - node_mask=node_mask, - pair_excl=pair_excl, - ) - for name, topology in built.items(): - topologies.setdefault(name, topology) - return topologies + params.update(model.topology_params()) + return params def forward_atomic_graph( self, diff --git a/deepmd/dpmodel/long_range/base_long_range.py b/deepmd/dpmodel/long_range/base_long_range.py index f37b821333..d00617fc1c 100644 --- a/deepmd/dpmodel/long_range/base_long_range.py +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -78,6 +78,10 @@ class BLR(ABC, PluginVariant, make_plugin_registry("long range")): def required_topologies(self) -> tuple[str, ...]: """Names of the topologies the module reads beside the neighbor graph.""" + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies, by name; none by default.""" + return {} + @abstractmethod def uses_charge_spin(self) -> bool: """Whether the module reads the frame-level charge/spin condition.""" diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py index 500256fbbd..67065a6eec 100644 --- a/deepmd/dpmodel/long_range/kernel.py +++ b/deepmd/dpmodel/long_range/kernel.py @@ -108,9 +108,9 @@ def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: def kspace(self, k_sq: Array) -> Array: r"""Fourier transform of the kernel on squared wave numbers. - :math:`\\hat K(k) = \\int K(r) e^{-i \\mathbf{k} \\cdot \\mathbf{r}} d^3r`, + :math:`\hat K(k) = \int K(r) e^{-i \mathbf{k} \cdot \mathbf{r}} d^3r`, so that the energy of charges in a periodic cell of volume ``V`` reads - :math:`(k_e / 2V) \\sum_{\\mathbf{k} \\neq 0} \\hat K(k) |\\hat\\rho(\\mathbf{k})|^2`. + :math:`(k_e / 2V) \sum_{\mathbf{k} \neq 0} \hat K(k) |\hat\rho(\mathbf{k})|^2`. Parameters ---------- @@ -129,7 +129,7 @@ def contact_value(self, like: Array) -> Array: r"""The kernel at contact, :math:`K(0)`, in Å⁻¹. The k-space lattice sum of a periodic system contains the self terms, - which the evaluation subtracts as :math:`\\frac{1}{2} K(0) \\sum_i q_i^2`. + which the evaluation subtracts as :math:`\frac{1}{2} K(0) \sum_i q_i^2`. Parameters ---------- diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index e470dc9c2e..0480c92f06 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -12,7 +12,6 @@ from deepmd.dpmodel.array_api import ( Array, - xp_add_at, ) from deepmd.dpmodel.common import ( GLOBAL_NP_FLOAT_PRECISION, @@ -26,10 +25,6 @@ ALL_PAIRS, K_GRID, PME_MESH, - KGrid, - MeshGrid, - build_topologies, - enumerate_k_grid_bounds, frame_id_from_n_node, k_grid_bounds_of_cells, segment_sum, @@ -51,6 +46,10 @@ COULOMB_CONSTANT, BaseInteractionKernel, ) +from .reciprocal import ( + kspace_energy, + mesh_energy, +) CHARGE_TARGET_SOURCES = ("neutral", "charge_spin", "fparam") """Sources of the total charge a constrained channel sums to.""" @@ -58,6 +57,17 @@ CHARGE_NET_NEURON = [64, 64] """Default hidden widths of the charge network.""" +PME_ORDER = 8 +"""Default order of the B-splines that spread the charges to the mesh.""" + +PME_SPACING_FRACTION = 0.75 +"""Default mesh spacing, as a fraction of the Nyquist spacing of the wave-number cutoff. + +At order 8 the mesh evaluation of a periodic frame then stays within a few +parts in a million of the sum over the wave numbers, in the reciprocal energy +and in the force. +""" + CHARGE_NET_SEED_INDEX = 10_000 """Child index of the fitting seed that seeds a charge network configured without a seed. @@ -67,73 +77,6 @@ """ -def _fft_smooth(n: int) -> int: - """The smallest integer >= ``n`` whose prime factors are among 2, 3, 5, 7.""" - - def smooth(m: int) -> bool: - for p in (2, 3, 5, 7): - while m % p == 0: - m //= p - return m == 1 - - m = max(1, n) - while not smooth(m): - m += 1 - return m - - -def _bspline_integer_weights(p: int) -> np.ndarray: - r"""The order-``p`` cardinal B-spline at :math:`M_p(l+1)`, ``l = 0..p-2``. - - The weights of the SPME denominator - :math:`D(\theta) = \sum_l M_p(l+1) e^{i \theta l}`. - """ - out = [] - for l in range(p - 1): - x = float(l + 1) - out.append( - sum( - (-1.0) ** j * math.comb(p, j) * max(x - j, 0.0) ** (p - 1) - for j in range(p + 1) - ) - / math.factorial(p - 1) - ) - return np.asarray(out) - - -def _bspline_eval(xp: Any, x: Array, p: int) -> Array: - r"""The order-``p`` cardinal B-spline :math:`M_p(x)` elementwise. - - :math:`M_p(x) = \frac{1}{(p-1)!} \sum_j (-1)^j \binom{p}{j} - \mathrm{relu}(x - j)^{p-1}` — a piecewise polynomial, differentiable - (twice for :math:`p \ge 4`) with respect to ``x``. - """ - w = None - for j in range(p + 1): - term = ((-1.0) ** j * math.comb(p, j)) * xp.maximum( - x - j, xp.zeros_like(x) - ) ** (p - 1) - w = term if w is None else w + term - return w / math.factorial(p - 1) - - -def _inv3(xp: Any, m: Array) -> Array: - """Inverse of a batch of 3x3 matrices, adjugate over determinant. - - Rows a, b, c: the columns of the inverse are (b x c)/det, (c x a)/det, - (a x b)/det. The closed form is cheap and forward-only here — the cell is - not an autograd leaf; the strain leaf keeps ``linalg.inv`` whose - registered backward handles a symbolic frame count. - """ - a, b, c = m[..., 0, :], m[..., 1, :], m[..., 2, :] - bc = xp.linalg.cross(b, c) - ca = xp.linalg.cross(c, a) - ab = xp.linalg.cross(a, b) - det = xp.sum(a * bc, axis=-1) - inv = xp.stack([bc, ca, ab], axis=-1) # columns - return inv / det[..., None, None] - - @BaseLongRange.register("latent_charge") class LatentChargeLongRange(NativeOP, BaseLongRange): r"""Energy of interacting latent charges, periodic or not. @@ -145,8 +88,8 @@ class LatentChargeLongRange(NativeOP, BaseLongRange): pairs of atoms of each frame listed by the :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` topology, each unordered pair once; for a periodic batch (with ``periodic=True``) it is - the reciprocal-space lattice sum over the wave numbers of the - :data:`~deepmd.dpmodel.utils.neighbor_graph.K_GRID` topology, + the lattice sum of the kernel over the images of the cell, evaluated in + reciprocal space (see :mod:`~deepmd.dpmodel.long_range.reciprocal`), .. math:: E_{\rm non-periodic} = k_e \sum_{i None: if charge_net.vacuum_ref: raise ValueError( @@ -288,21 +223,29 @@ def __init__( if self.decouple_k_grid else None ) - # The shared integer grid of the fixed mode, empty until the data - # statistics size it (see :meth:`compute_input_stats`) — an array even - # then, so a backend keeping arrays in its state dict sees a stable - # set of keys and a checkpoint restores the sized grid. + # The per-axis bounds of the shared integer grid of the fixed mode, + # zero until the data statistics size it (see + # :meth:`compute_input_stats`): an array of fixed shape, so that a + # backend keeping arrays in its state dict sees the same keys and + # shapes before and after the sizing, and a checkpoint restores the + # sized grid. if k_grid_bounds is not None: - self.k_grid_n = enumerate_k_grid_bounds(k_grid_bounds) + bounds = np.asarray(k_grid_bounds, dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0) or not np.any(bounds): + raise ValueError( + "`k_grid_bounds` must be three non-negative integers, not " + f"all zero, got {list(k_grid_bounds)}" + ) + self.k_grid_bounds = bounds elif self.k_grid_mode == "fixed": - self.k_grid_n = np.zeros((0, 3), dtype=np.int64) + self.k_grid_bounds = np.zeros(3, dtype=np.int64) else: - self.k_grid_n = None - # The PME mesh parameters. The mesh dims are an array even while - # unsized, so a backend keeping arrays in its state dict sees a - # stable set of keys and a checkpoint restores the sized mesh; the - # plain-integer copy is what a trace reads as static shapes. - self.pme_order = int(pme_order) + self.k_grid_bounds = None + # The mesh of a periodic frame, on which training in the "pme" mode + # and every frozen model evaluate the reciprocal sum. + if not self.periodic and (pme_spacing is not None or pme_order is not None): + raise ValueError("`pme_spacing` and `pme_order` require `periodic=True`") + self.pme_order = PME_ORDER if pme_order is None else int(pme_order) if self.pme_order < 2: raise ValueError(f"`pme_order` must be at least 2, got {self.pme_order}") if pme_spacing is not None: @@ -312,44 +255,6 @@ def __init__( f"`pme_spacing` must be finite and positive, got {pme_spacing}" ) self.pme_spacing = pme_spacing - if self.k_grid_mode != "pme": - if pme_spacing is not None or pme_mesh is not None: - raise ValueError( - "`pme_spacing`/`pme_mesh` are only meaningful with " - "`k_grid_mode='pme'`" - ) - self.pme_mesh = None - self._pme_n = None - elif pme_mesh is not None: - mesh = np.asarray(pme_mesh, dtype=np.int64).reshape(-1) - if mesh.shape != (3,) or np.any(mesh < 1): - raise ValueError( - f"`pme_mesh` must be three positive integers, got {pme_mesh}" - ) - self.pme_mesh = mesh - self._pme_n = (int(mesh[0]), int(mesh[1]), int(mesh[2])) - else: - # Unsized sentinel: the data statistics size it (see - # :meth:`compute_input_stats`), or a checkpoint restores the - # sized mesh through the state dict. - self.pme_mesh = np.zeros((3,), dtype=np.int64) - self._pme_n = None - - def _pme_dims(self) -> tuple[int, int, int] | None: - """The mesh dimensions as plain integers, or ``None`` when unsized. - - The plain-integer cache is what a trace reads as static shapes; it is - set at construction and at statistics time. A backend restoring the - sized mesh through its state dict (a checkpoint does not run - :meth:`compute_input_stats`) is picked up lazily from the mesh array. - """ - if self._pme_n is not None: - return self._pme_n - if self.pme_mesh is None or int(self.pme_mesh[0]) < 1: - return None - mesh = self.pme_mesh - self._pme_n = (int(mesh[0]), int(mesh[1]), int(mesh[2])) - return self._pme_n @staticmethod def _normalize_constraint( @@ -443,8 +348,7 @@ def from_config( k_grid_bounds=config.get("k_grid_bounds"), decouple_k_grid=config.get("decouple_k_grid", False), pme_spacing=config.get("pme_spacing"), - pme_order=config.get("pme_order", 4), - pme_mesh=config.get("pme_mesh"), + pme_order=config.get("pme_order"), ) def _k_max(self) -> float: @@ -458,66 +362,59 @@ def _k_max(self) -> float: return self.frozen_k_max return float(self.kernel.kspace_cutoff(self.kspace_tolerance)) - def _pme_target_spacing(self) -> float: - r"""The target PME mesh spacing, in Å. + def mesh_spacing(self) -> float: + r"""The largest spacing of the mesh of a periodic frame, in Å. - The explicit ``pme_spacing``, or the Nyquist spacing of the kernel's - wave-number cutoff :math:`\pi / k_{max}` — fine enough that the - kernel's spectral weight at the mesh cutoff is negligible. + The configured ``pme_spacing``, or :data:`PME_SPACING_FRACTION` of the + Nyquist spacing :math:`\pi / k_{max}` of the wave-number cutoff. """ if self.pme_spacing is not None: return self.pme_spacing - return math.pi / self._k_max() + return PME_SPACING_FRACTION * math.pi / self._k_max() - def build_topologies( - self, - coord: Array, - n_node: Array, - box: Array | None, - *, - atype: Array, - node_mask: Array | None = None, - pair_excl: Any = None, - ) -> dict[str, Any]: - """Build the topologies of the module for one batch. + def use_mesh_evaluation(self) -> None: + """Evaluate the reciprocal sum of every periodic batch on the mesh. - The k grid of a periodic batch is enumerated from the cells with the - wave-number cutoff of the kernel (the trainable parameters are read at - build time, so a drifting kernel re-derives its grid every batch, - unless the grid is decoupled from them by ``decouple_k_grid``). In - the fixed grid mode the batch is evaluated on the shared integer grid - instead, whose mask the same cutoff derives. + A frozen model takes the mesh of every frame as an input of dynamic + size, sized by the cell of the frame, so that it serves any cell at + the accuracy of :meth:`mesh_spacing`; the explicit wave numbers of + the other modes have a count that no input layout fixes. The spacing + is fixed at the current kernel parameters, a constant of the frozen + model. """ - params = {} if self.periodic: - if self.k_grid_mode == "pme": - if self._pme_dims() is None: - raise ValueError( - "the PME mesh is not sized yet: run the data " - "statistics (`compute_input_stats`) on periodic " - "frames, or configure `pme_mesh`" - ) - params[PME_MESH] = {"n_mesh": self.pme_mesh} - else: - params[K_GRID] = {"k_max": self._k_max()} - if self.k_grid_mode == "fixed": - if self.k_grid_n is None or self.k_grid_n.shape[0] == 0: - raise ValueError( - "the fixed k grid is not sized yet: run the data " - "statistics (`compute_input_stats`) on periodic " - "frames, or configure `k_grid_bounds`" - ) - params[K_GRID]["n_grid"] = self.k_grid_n - return build_topologies( - self.required_topologies(), - coord, - n_node, - box, - node_mask, - atype=atype, - pair_excl=pair_excl, - topology_params=params, - ) + self.pme_spacing = self.mesh_spacing() + self.k_grid_mode = "pme" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the reciprocal-space topology of a periodic batch. + + The wave-number cutoff reads the kernel parameters at build time, so + that a trained kernel moves the k grid and the mesh with it, unless + ``decouple_k_grid`` keeps the initial cutoff. The k grid holds the + wave numbers of every frame within the cutoff, or the points of the + fixed integer grid with those beyond it masked; the mesh resolves + every cell of the batch at :meth:`mesh_spacing`. + + Raises + ------ + ValueError + If the fixed grid has not been sized. + """ + if not self.periodic: + return {} + if self.k_grid_mode == "pme": + return {PME_MESH: {"spacing": self.mesh_spacing()}} + params = {"k_max": self._k_max()} + if self.k_grid_mode == "fixed": + if not self._k_grid_sized(): + raise ValueError( + "the fixed k grid is not sized yet: run the data " + "statistics (`compute_input_stats`) on periodic " + "frames, or configure `k_grid_bounds`" + ) + params["bounds"] = to_numpy_array(self.k_grid_bounds) + return {K_GRID: params} def required_topologies(self) -> tuple[str, ...]: if self.periodic: @@ -594,14 +491,23 @@ def call_graph( ) energy = self._pair_branch(charge, pairs, n) if self.periodic: + # The lattice sum of a periodic frame and its self term cancel to + # a small fraction of either, so the reciprocal-space energy is + # formed in float64 whatever the precision of the geometry. + charge_64 = xp.astype(charge, xp.float64) if self.k_grid_mode == "pme": - energy = energy + self._pme_branch( - charge, topologies[PME_MESH], frame_id, node_mask, n_node + reciprocal = mesh_energy( + charge_64, + topologies[PME_MESH], + frame_id, + self.kernel, + self.pme_order, ) else: - energy = energy + self._kspace_branch( - charge, topologies[K_GRID], frame_id, node_mask, n_node + reciprocal = kspace_energy( + charge_64, topologies[K_GRID], frame_id, self.kernel ) + energy = energy + xp.astype(reciprocal, energy.dtype) return xp.reshape(energy, (n, 1)) def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: @@ -619,270 +525,6 @@ def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: ) return segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) - def _kspace_branch( - self, - charge: Array, - k_grid: KGrid, - frame_id: Array, - node_mask: Array, - n_node: Array, - ) -> Array: - r"""Reciprocal-space energy of the periodic frames, per atom. - - Evaluates - :math:`E_f = k_e [ \frac{1}{2V_f} \sum_{\mathbf{k}\neq\mathbf{0}} - \hat K(\mathbf{k}) |\hat\rho(\mathbf{k})|^2 - - \frac{1}{2} K(0) \sum_{i \in f} |q_i|^2 ]` - on the strained geometry of :attr:`KGrid.strain` (identity when no - strain leaf is set): coordinates by ``I + strain``, wave numbers by - ``(I + strain)^{-1}`` on their row layout and volumes by - ``det(I + strain)``, so that the gradient of the energy with respect - to the strain leaf is the negative k-space virial of the frame. The - reciprocal energy is attributed to the charged atoms of the frame in - equal shares; the self term is per atom by construction. Guard points - carry a zero mask. - """ - xp = array_api_compat.array_namespace(charge) - nf = n_node.shape[0] - coord = k_grid.coord - if coord is None: - raise ValueError( - "the k-grid topology carries no coordinates; the graph lower " - "re-leafs them before the evaluation" - ) - coord = xp.astype(coord, charge.dtype) - k_vec = xp.astype(k_grid.vec, charge.dtype) - volume = xp.astype(k_grid.volume, charge.dtype) - if k_grid.strain is not None: - eye = xp.eye(3, dtype=charge.dtype, device=array_api_compat.device(charge)) - s = xp.astype(k_grid.strain, charge.dtype) + eye # (nf, 3, 3) - # r' = S r (columns) - s_node = xp.take(s, frame_id, axis=0) # (N, 3, 3) - coord = xp.matmul(s_node, coord[:, :, None])[:, :, 0] - # k' = S^{-T} k (columns), i.e. rows multiplied by S^{-1} - s_inv_t = xp.matrix_transpose(xp.linalg.inv(s)) - s_inv_t_k = xp.take(s_inv_t, k_grid.frame, axis=0) # (nk, 3, 3) - k_vec = xp.matmul(s_inv_t_k, k_vec[:, :, None])[:, :, 0] - volume = volume * xp.linalg.det(s) - phase = xp.matmul(k_vec, xp.matrix_transpose(coord)) # (nk, N) - same_frame = frame_id[None, :] == k_grid.frame[:, None] - cos_phase = xp.cos(phase) * same_frame - sin_phase = xp.sin(phase) * same_frame - rho_real = xp.matmul(cos_phase, charge) # (nk, C) - rho_imag = xp.matmul(sin_phase, charge) # (nk, C) - s_sq = xp.sum(rho_real**2 + rho_imag**2, axis=-1) # (nk,) - k_sq = xp.sum(k_vec * k_vec, axis=-1) - kernel_value = self.kernel.kspace(k_sq) # (nk,) - half_ke = 0.5 * COULOMB_CONSTANT - volume_k = xp.take(volume, k_grid.frame, axis=0) - energy_k = half_ke * kernel_value * s_sq / volume_k - energy_k = xp.where(k_grid.mask, energy_k, xp.zeros_like(energy_k)) - e_rec_frame = segment_sum(energy_k, k_grid.frame, nf) # (nf,) - weight = xp.astype(node_mask, charge.dtype) - count = segment_sum(weight, frame_id, nf) - count = xp.maximum(count, xp.ones_like(count)) - e_rec = xp.take(e_rec_frame / count, frame_id, axis=0) * weight # (N,) - # The self term belongs to the periodic lattice sum alone: a frame - # without real k points (a non-periodic batch on a periodic-enabled - # model) keeps its non-periodic self-free energy. - periodic_frame = ( - segment_sum(xp.astype(k_grid.mask, charge.dtype), k_grid.frame, nf) > 0 - ) - e_self = ( - -half_ke - * self.kernel.contact_value(charge) - * xp.sum(charge * charge, axis=-1) - * xp.astype(xp.take(periodic_frame, frame_id, axis=0), charge.dtype) - ) - return e_rec + e_self - - def _pme_b_factor(self, xp: Any, dtype: Any, device: Any, m_axes: list) -> Array: - r"""The SPME assignment correction :math:`|B(m)|^2` on the mesh modes. - - Per axis :math:`|D_d(m)|^2 = |\sum_l M_p(l+1) e^{2\pi i m l / n_d}|^2` - (equal to :math:`\mathrm{sinc}(m/n_d)^{2p}`), the full factor the - three-axis outer product. ``m_axes`` holds the integer mode values of - each axis (the third axis carries only the non-negative rfft modes). - """ - w = _bspline_integer_weights(self.pme_order) - axes = [] - for m, n_d in zip(m_axes, self._pme_dims(), strict=True): - theta = 2.0 * math.pi * m / n_d - re = xp.zeros_like(theta) - im = xp.zeros_like(theta) - for l, wl in enumerate(w): - re = re + float(wl) * xp.cos(theta * l) - im = im + float(wl) * xp.sin(theta * l) - axes.append(1.0 / (re * re + im * im)) - b = axes[0][:, None, None] * axes[1][None, :, None] * axes[2][None, None, :] - return xp.reshape(b, (-1,)) - - def _pme_branch( - self, - charge: Array, - mesh: MeshGrid, - frame_id: Array, - node_mask: Array, - n_node: Array, - ) -> Array: - r"""Reciprocal-space energy of the periodic frames on a mesh, per atom. - - Evaluates the same lattice sum as :meth:`_kspace_branch` by the SPME - algorithm: the charges spread to a static real-space mesh with a - B-spline of order :attr:`pme_order`, the mesh transforms by a - real-input FFT (only the non-negative :math:`m_3` half-spectrum), and - the energy convolves with the kernel transform corrected by the - assignment factor (:meth:`_pme_b_factor`) and the hermitian doubling - weight of the interior :math:`m_3` planes. No wave-number vector is - ever formed: the kernel reads :math:`|k|^2 = (2\pi)^2 m^\mathsf{T} W m` - with the per-frame metric :math:`W = (B^{-1})^\mathsf{T} B^{-1}` — - six static monomials of the mode indices times six per-frame metric - entries, all elementwise. The strain leaf enters the metric by - :math:`W' = A^\mathsf{T} A` with :math:`A = S^{-\mathsf{T}} B^{-1}` - (:math:`S = I +` strain) and the volume by :math:`\det S` — the - fractional coordinates are strain-invariant, so the assignment is - unstrained and the virial comes from the metric and the volume alone. - Non-periodic frames (``mesh.mask`` false) carry no mesh energy and no - self term. - """ - xp = array_api_compat.array_namespace(charge) - device = array_api_compat.device(charge) - nf = n_node.shape[0] - n_channel = charge.shape[-1] - n1, n2, n3 = self._pme_dims() - p = self.pme_order - m_total = n1 * n2 * n3 - box = xp.astype(mesh.box, charge.dtype) - coord = xp.astype(mesh.coord, charge.dtype) - volume = xp.astype(mesh.volume, charge.dtype) - b_inv = _inv3(xp, box) # (nf, 3, 3); columns are g_d / 2pi - b_inv_node = xp.take(b_inv, frame_id, axis=0) # (N, 3, 3) - # fractional coordinates, wrapped to [0, 1): frac_d = sum_j r_j B^-1_jd - # (elementwise; a batched 3x3 GEMV is the worst GEMM shape on GPU) - frac = xp.sum(coord[:, :, None] * b_inv_node, axis=1) - frac = frac - xp.floor(frac) - # per-axis assignment: weights (N, p) and mesh indices (N, p) - weights = [] - indices = [] - for d, n_d in enumerate((n1, n2, n3)): - u = frac[:, d] * n_d - base = xp.floor(u) - t = (u - base)[:, None] - shift = xp.arange(p, dtype=charge.dtype, device=device)[None, :] - # x_k = t + p - 1 - k lies in the support [0, p] of M_p - weights.append(_bspline_eval(xp, t + (p - 1) - shift, p)) - idx = ( - xp.astype(base, xp.int64)[:, None] - - (p - 1) - + xp.arange(p, dtype=xp.int64, device=device)[None, :] - ) - indices.append(idx % n_d) - w1, w2, w3 = weights - j1, j2, j3 = indices - w = xp.reshape( - w1[:, :, None, None] * w2[:, None, :, None] * w3[:, None, None, :], - (frac.shape[0], p**3), - ) - lin = xp.reshape( - (j1[:, :, None, None] * n2 + j2[:, None, :, None]) * n3 - + j3[:, None, None, :], - (frac.shape[0], p**3), - ) - # spread the charges to the mesh, frame- and channel-major - q_mesh = xp.zeros( - (nf * n_channel * m_total,), dtype=charge.dtype, device=device - ) - for c in range(n_channel): - target = ((frame_id * n_channel + c) * m_total)[:, None] + lin - value = charge[:, c][:, None] * w - q_mesh = xp_add_at( - q_mesh, xp.reshape(target, (-1,)), xp.reshape(value, (-1,)) - ) - q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) - # the charge mesh is real: only the non-negative m3 half-spectrum is - # transformed and summed (with the doubling weight below) - rho = xp.fft.rfftn(q_mesh, axes=(-3, -2, -1)) - rho_sq = xp.real(rho) ** 2 + xp.imag(rho) ** 2 - m_half = n1 * n2 * (n3 // 2 + 1) - rho_sq = xp.reshape(rho_sq, (nf, n_channel, m_half)) - # the reciprocal-lattice metric: |k|^2 = (2 pi)^2 m^T W m with - # W = (B^-1)^T B^-1; the strain folds in through A = S^-T B^-1 (the - # linalg inverse of S keeps its registered backward, safe under a - # symbolic frame count) - if mesh.strain is not None: - eye = xp.eye(3, dtype=charge.dtype, device=device) - s = xp.astype(mesh.strain, charge.dtype) + eye # (nf, 3, 3) - s_inv = xp.linalg.inv(s) - a = xp.matmul(xp.matrix_transpose(s_inv), b_inv) - w_metric = xp.matmul(xp.matrix_transpose(a), a) - volume = volume * xp.linalg.det(s) - else: - w_metric = xp.matmul(xp.matrix_transpose(b_inv), b_inv) - # static integer mode grids: full fftfreq order on axes 1, 2; the rfft - # axis keeps m3 = 0 .. n3//2 - m_axes = [] - for d, n_d in enumerate((n1, n2, n3)): - m = xp.arange(n_d, dtype=charge.dtype, device=device) - if d == 2: - m = m[: n_d // 2 + 1] - else: - m = xp.where(m > n_d // 2, m - n_d, m) - m_axes.append(m) - g1, g2, g3 = xp.meshgrid(*m_axes, indexing="ij") - m1f = xp.reshape(g1, (-1,)) - m2f = xp.reshape(g2, (-1,)) - m3f = xp.reshape(g3, (-1,)) - # six static monomials x six per-frame metric entries, no GEMM - two_pi_sq = (2.0 * math.pi) ** 2 - k_sq = two_pi_sq * ( - w_metric[:, 0, 0][:, None] * (m1f * m1f)[None, :] - + w_metric[:, 1, 1][:, None] * (m2f * m2f)[None, :] - + w_metric[:, 2, 2][:, None] * (m3f * m3f)[None, :] - + 2.0 * w_metric[:, 0, 1][:, None] * (m1f * m2f)[None, :] - + 2.0 * w_metric[:, 0, 2][:, None] * (m1f * m3f)[None, :] - + 2.0 * w_metric[:, 1, 2][:, None] * (m2f * m3f)[None, :] - ) # (nf, m_half) - # the zero mode sits at flat index 0; guard the singular kernel there - safe = xp.where(k_sq == 0.0, xp.ones_like(k_sq), k_sq) - influence = xp.reshape( - self.kernel.kspace(xp.reshape(safe, (-1,))), (nf, m_half) - ) - influence = ( - influence * self._pme_b_factor(xp, charge.dtype, device, m_axes)[None, :] - ) - # the spectrum of a real mesh is hermitian: interior m3 planes count - # twice, the m3 = 0 and Nyquist planes once - w_rfft = xp.where( - (m3f > 0) & (m3f < n3 / 2.0), - 2.0 * xp.ones_like(m3f), - xp.ones_like(m3f), - ) - influence = influence * w_rfft[None, :] - # the zero mode sits at flat index 0 (tin-foil); keep everything else - keep = xp.concat( - [ - xp.zeros((1,), dtype=xp.bool, device=device), - xp.ones((m_half - 1,), dtype=xp.bool, device=device), - ] - ) - influence = xp.where(keep[None, :], influence, xp.zeros_like(influence)) - influence = influence * xp.astype(mesh.mask, charge.dtype)[:, None] - half_ke = 0.5 * COULOMB_CONSTANT - e_frame = ( - half_ke * xp.sum(influence * xp.sum(rho_sq, axis=1), axis=-1) / volume - ) # (nf,) - weight = xp.astype(node_mask, charge.dtype) - count = segment_sum(weight, frame_id, nf) - count = xp.maximum(count, xp.ones_like(count)) - e_rec = xp.take(e_frame / count, frame_id, axis=0) * weight # (N,) - e_self = ( - -half_ke - * self.kernel.contact_value(charge) - * xp.sum(charge * charge, axis=-1) - * xp.astype(xp.take(mesh.mask, frame_id, axis=0), charge.dtype) - ) - return e_rec + e_self - def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: """Kernel values of the pairs. @@ -983,18 +625,12 @@ def compute_input_stats( protection: float = 1e-2, stat_file_path: DPPath | None = None, ) -> None: - if self.k_grid_mode == "fixed" and ( - self.k_grid_n is None or self.k_grid_n.shape[0] == 0 - ): + if self.k_grid_mode == "fixed" and not self._k_grid_sized(): # Materialize the samples once for the grid sizing and the # networks; a lazy sampler is shared by the callers. sampled = merged() if callable(merged) else merged self._size_k_grid(sampled) merged = sampled - elif self.k_grid_mode == "pme" and self._pme_dims() is None: - sampled = merged() if callable(merged) else merged - self._size_pme_mesh(sampled) - merged = sampled self.charge_net.compute_input_stats( merged, protection=protection, stat_file_path=stat_file_path ) @@ -1005,9 +641,9 @@ def _size_k_grid(self, sampled: list[dict]) -> None: Every sampled frame needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` per axis at the wave-number cutoff of the kernel; the grid takes the elementwise maximum of the bounds over the frames. The kernel - parameters are read now, so a grid is sized for the kernel at - statistics time — frames that later outgrow it (larger cells, or a - kernel moving its cutoff out) are truncated to the grid. + parameters are read at statistics time, so the grid is sized for the + kernel of that moment; frames that outgrow it later, by larger cells + or a kernel whose cutoff grows, are truncated to the grid. """ k_max = self._k_max() bounds = None @@ -1026,38 +662,11 @@ def _size_k_grid(self, sampled: list[dict]) -> None: "but the sampled frames hold no simulation cell; provide " "periodic training data or configure `k_grid_bounds`" ) - self.k_grid_n = enumerate_k_grid_bounds(bounds) - - def _size_pme_mesh(self, sampled: list[dict]) -> None: - r"""Size the PME mesh by the largest sampled cell. + self.k_grid_bounds = bounds.astype(np.int64) - Every frame is evaluated on the shared mesh with the effective - spacing :math:`|a_d|/n_d` per axis; the mesh takes the largest sampled - cell axis down to the target spacing, rounded up to an FFT-smooth - dimension. A frame larger than the mesh (a larger cell than the - statistics saw) is evaluated with a coarser effective spacing — a - smooth accuracy degradation, by design without a warning. - """ - spacing = self._pme_target_spacing() - a_max = None - for frame in sampled: - box = frame.get("box") - if box is None: - continue - box = np.asarray(to_numpy_array(box), dtype=np.float64) - if box.size == 0: - continue - frame_max = np.max(np.linalg.norm(box.reshape(-1, 3, 3), axis=-1), axis=0) - a_max = frame_max if a_max is None else np.maximum(a_max, frame_max) - if a_max is None: - raise ValueError( - "the PME mesh takes its size from the data statistics, " - "but the sampled frames hold no simulation cell; provide " - "periodic training data or configure `pme_mesh`" - ) - dims = [_fft_smooth(int(np.ceil(a / spacing))) for a in a_max] - self.pme_mesh = np.asarray(dims, dtype=np.int64) - self._pme_n = (dims[0], dims[1], dims[2]) + def _k_grid_sized(self) -> bool: + """Whether the fixed integer grid has its bounds.""" + return bool(np.any(to_numpy_array(self.k_grid_bounds))) def change_type_map(self, type_map: list[str]) -> None: self.charge_net.change_type_map(type_map=type_map) @@ -1079,17 +688,15 @@ def serialize(self) -> dict: data["kspace_tolerance"] = self.kspace_tolerance if self.k_grid_mode != "ragged": data["k_grid_mode"] = self.k_grid_mode - if self.k_grid_n is not None and self.k_grid_n.shape[0] > 0: - data["k_grid_n"] = to_numpy_array(self.k_grid_n).tolist() + if self.k_grid_mode == "fixed" and self._k_grid_sized(): + data["k_grid_bounds"] = to_numpy_array(self.k_grid_bounds).tolist() if self.decouple_k_grid: data["decouple_k_grid"] = True data["frozen_k_max"] = self.frozen_k_max - if self.k_grid_mode == "pme": + if self.periodic: data["pme_order"] = self.pme_order if self.pme_spacing is not None: data["pme_spacing"] = self.pme_spacing - if self._pme_dims() is not None: - data["pme_mesh"] = to_numpy_array(self.pme_mesh).tolist() if self.charge_state is not None: data["charge_state"] = to_numpy_array(self.charge_state).tolist() return data @@ -1105,18 +712,16 @@ def deserialize(cls, data: dict) -> "LatentChargeLongRange": periodic=data.get("periodic", False), kspace_tolerance=data.get("kspace_tolerance", 1e-8), k_grid_mode=data.get("k_grid_mode", "ragged"), + k_grid_bounds=data.get("k_grid_bounds"), decouple_k_grid=data.get("decouple_k_grid", False), pme_spacing=data.get("pme_spacing"), - pme_order=data.get("pme_order", 4), - pme_mesh=data.get("pme_mesh"), + pme_order=data.get("pme_order"), ) if data.get("frozen_k_max") is not None: # The constructor derives the frozen cutoff from the (possibly # trained) kernel; the record holds the value the grid was # actually built with. obj.frozen_k_max = float(data["frozen_k_max"]) - if data.get("k_grid_n") is not None: - obj.k_grid_n = np.asarray(data["k_grid_n"], dtype=np.int64).reshape(-1, 3) if data.get("charge_state") is not None: obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) return obj diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py new file mode 100644 index 0000000000..8643ed3473 --- /dev/null +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Reciprocal-space energy of latent charges in periodic cells. + +For the charges :math:`q_i^c` of a periodic cell of volume :math:`V` and a +pair kernel :math:`K` with Fourier transform :math:`\hat K`, the energy of the +lattice of images, + +.. math:: + E = k_e \Big[ \frac{1}{2V} \sum_{\mathbf{k} \neq \mathbf{0}} \sum_c + \hat K(\mathbf{k}) \, |\hat\rho_c(\mathbf{k})|^2 + - \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \Big], + \qquad + \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i \mathbf{k} \cdot \mathbf{r}_i}, + +is the sum over the wave numbers of the cell minus the interaction of every +charge with itself. The :math:`\mathbf{k} = \mathbf{0}` mode is dropped, the +tin-foil convention, which for a channel with a net charge amounts to a +uniform neutralizing background. The energy is attributed to the atoms as +:math:`E_i = \frac{1}{2} \sum_c q_i^c \phi_i^c`, with the potential +:math:`\phi_i^c = \partial E / \partial q_i^c` that the other charges and the +images create at atom :math:`i`, as the pair sum of a non-periodic system +splits every pair term evenly between its atoms. + +Two evaluations serve the topologies of +:mod:`~deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph`: +:func:`kspace_energy` sums over explicit wave numbers, and +:func:`mesh_energy` evaluates the same sum by the smooth particle-mesh Ewald +method on a mesh. +""" + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_add_at, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, + segment_sum, +) + +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) + + +def _self_potential( + kernel: BaseInteractionKernel, charge: Array, periodic_node: Array +) -> Array: + """The potential of every charge at itself, k_e K(0) q, on periodic nodes.""" + xp = array_api_compat.array_namespace(charge) + return ( + COULOMB_CONSTANT + * kernel.contact_value(charge) + * charge + * xp.astype(periodic_node, charge.dtype)[:, None] + ) + + +def kspace_energy( + charge: Array, + k_grid: KGrid, + frame_id: Array, + kernel: BaseInteractionKernel, +) -> Array: + r"""Energy of the periodic frames over explicit wave numbers, per atom. + + The structure factors of the wave numbers of every frame come from the + phases of its nodes, and the potential at a node from the same phases: + :math:`\phi_i = \frac{k_e}{V} \sum_{\mathbf{k}} \hat K(\mathbf{k}) + [\cos(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Re}\,\hat\rho + + \sin(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Im}\,\hat\rho] - k_e K(0) q_i`. + The geometry is strained by :attr:`KGrid.strain` (see :class:`KGrid`). + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + k_grid + The wave numbers of the periodic frames. + frame_id + The frame of every node with shape ``(N,)``. + kernel + The pair kernel. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + without wave numbers. + """ + xp = array_api_compat.array_namespace(charge) + dtype = charge.dtype + nf = k_grid.volume.shape[0] + coord = xp.astype(k_grid.coord, dtype) + k_vec = xp.astype(k_grid.vec, dtype) + volume = xp.astype(k_grid.volume, dtype) + # === Step 1. Strained geometry: r' = S r, k' = S^-T k, V' = V det S === + eye = xp.eye(3, dtype=dtype, device=array_api_compat.device(charge)) + strain = xp.astype(k_grid.strain, dtype) + eye + coord = xp.matmul(xp.take(strain, frame_id, axis=0), coord[:, :, None])[:, :, 0] + strain_inv_t = xp.matrix_transpose(xp.linalg.inv(strain)) + k_vec = xp.matmul(xp.take(strain_inv_t, k_grid.frame, axis=0), k_vec[:, :, None]) + k_vec = k_vec[:, :, 0] + volume = volume * xp.linalg.det(strain) + # === Step 2. Structure factors, (nk, C) === + phase = xp.matmul(k_vec, xp.matrix_transpose(coord)) + same_frame = xp.astype(frame_id[None, :] == k_grid.frame[:, None], dtype) + cos_phase = xp.cos(phase) * same_frame + sin_phase = xp.sin(phase) * same_frame + rho_real = xp.matmul(cos_phase, charge) + rho_imag = xp.matmul(sin_phase, charge) + # === Step 3. Potential at the nodes, (N, C) === + weight = COULOMB_CONSTANT * kernel.kspace(xp.sum(k_vec * k_vec, axis=-1)) + weight = weight / xp.take(volume, k_grid.frame, axis=0) + weight = xp.where(k_grid.mask, weight, xp.zeros_like(weight))[:, None] + potential = xp.matmul(xp.matrix_transpose(cos_phase), weight * rho_real) + potential = potential + xp.matmul(xp.matrix_transpose(sin_phase), weight * rho_imag) + # The self term belongs to the lattice sum of a frame with wave numbers; + # a non-periodic batch evaluated by a periodic module carries none. + has_points = segment_sum(xp.astype(k_grid.mask, dtype), k_grid.frame, nf) > 0 + periodic_node = xp.take(has_points, frame_id, axis=0) + potential = potential - _self_potential(kernel, charge, periodic_node) + return 0.5 * xp.sum(charge * potential, axis=-1) + + +def bspline_integer_weights(order: int) -> np.ndarray: + r"""The cardinal B-spline of order ``order`` at :math:`1, 2, \dots, p - 1`. + + These are the weights of the assignment factor + :math:`D(\theta) = \sum_{l=0}^{p-2} M_p(l + 1) e^{i \theta l}` of the + smooth particle-mesh Ewald method. + """ + return np.asarray( + [ + sum( + (-1.0) ** j * math.comb(order, j) * max(l + 1.0 - j, 0.0) ** (order - 1) + for j in range(order + 1) + ) + / math.factorial(order - 1) + for l in range(order - 1) + ] + ) + + +def _bspline(xp: Any, x: Array, order: int) -> Array: + r"""The cardinal B-spline :math:`M_p(x)` of order :math:`p`, elementwise. + + :math:`M_p(x) = \frac{1}{(p-1)!} \sum_{j=0}^{p} (-1)^j \binom{p}{j} + \max(x - j, 0)^{p-1}`, a piecewise polynomial supported on + :math:`[0, p]` with :math:`p - 2` continuous derivatives. + """ + value = None + for j in range(order + 1): + term = ((-1.0) ** j * math.comb(order, j)) * xp.maximum( + x - j, xp.zeros_like(x) + ) ** (order - 1) + value = term if value is None else value + term + return value / math.factorial(order - 1) + + +def _inverse3(xp: Any, m: Array) -> Array: + """Inverse of a batch of 3x3 matrices, its columns the cross products of the rows over the determinant.""" + a, b, c = m[..., 0, :], m[..., 1, :], m[..., 2, :] + bc = xp.linalg.cross(b, c) + ca = xp.linalg.cross(c, a) + ab = xp.linalg.cross(a, b) + det = xp.sum(a * bc, axis=-1) + return xp.stack([bc, ca, ab], axis=-1) / det[..., None, None] + + +def _stencil( + xp: Any, frac: Array, dims: tuple[Any, Any, Any], order: int +) -> tuple[Array, Array]: + r"""The mesh points every node spreads to and their weights. + + Along axis :math:`d` the scaled coordinate :math:`u = n_d s_d` of the + fractional coordinate :math:`s_d \in [0, 1)` reaches the ``order`` points + :math:`\lfloor u \rfloor - p + 1, \dots, \lfloor u \rfloor` with the + weights :math:`M_p(u - m)`; the stencil of a node is the outer product of + its three axes. + + Returns + ------- + tuple[Array, Array] + The flat mesh indices and the weights, both with shape + ``(N, order**3)``. + """ + device = array_api_compat.device(frac) + n_node = frac.shape[0] + shift = xp.arange(order, dtype=frac.dtype, device=device)[None, :] + offset = xp.arange(order, dtype=xp.int64, device=device)[None, :] + index = None + weight = None + for d, n_d in enumerate(dims): + u = frac[:, d] * n_d + base = xp.floor(u) + # The distance of point base - p + 1 + j from u is t + p - 1 - j. + w = _bspline(xp, (u - base)[:, None] + (order - 1) - shift, order) + m = (xp.astype(base, xp.int64)[:, None] - (order - 1) + offset) % n_d + if index is None: + index, weight = m, w + else: + index = xp.reshape(index[:, :, None] * n_d + m[:, None, :], (n_node, -1)) + weight = xp.reshape(weight[:, :, None] * w[:, None, :], (n_node, -1)) + return index, weight + + +def _assignment_factor( + m_axes: list[Array], dims: tuple[Any, Any, Any], order: int +) -> list[Array]: + r"""The inverse squared assignment factor :math:`|D_d(m)|^{-2}` of every axis. + + The B-spline spreading multiplies every mode of the charge mesh by + :math:`\prod_d D_d(m_d)`, :math:`D_d(m) = \sum_l M_p(l + 1) + e^{2\pi i m l / n_d}`, which the influence function divides out. + """ + xp = array_api_compat.array_namespace(m_axes[0]) + weights = bspline_integer_weights(order) + factors = [] + for m, n_d in zip(m_axes, dims, strict=True): + theta = (2.0 * math.pi) * m / n_d + real = xp.zeros_like(theta) + imag = xp.zeros_like(theta) + for l, w in enumerate(weights): + real = real + float(w) * xp.cos(theta * l) + imag = imag + float(w) * xp.sin(theta * l) + factors.append(1.0 / (real * real + imag * imag)) + return factors + + +def mesh_energy( + charge: Array, + mesh: MeshGrid, + frame_id: Array, + kernel: BaseInteractionKernel, + order: int, +) -> Array: + r"""Energy of the periodic frames by the smooth particle-mesh Ewald method, per atom. + + The charges spread to the mesh of every frame with B-splines of order + ``order`` along the fractional coordinates; a real FFT gives the half + spectrum :math:`Q(\mathbf{m})` of the charge mesh. The influence function + + .. math:: + G(\mathbf{m}) = \frac{k_e}{V} \hat K(\mathbf{k}_\mathbf{m}) + \prod_d |D_d(m_d)|^{-2}, + \qquad \mathbf{k}_\mathbf{m} = 2\pi B^{-1} \mathbf{m}, + + corrects for the spreading, so that the potential on the mesh is the + inverse transform of :math:`G Q` and the potential at a node its + interpolation by the same B-splines. :math:`|\mathbf{k}|^2` is the + quadratic form :math:`(2\pi)^2 \mathbf{m}^\mathsf{T} W \mathbf{m}` of the + reciprocal metric :math:`W = A^\mathsf{T} A`, :math:`A = S^{-\mathsf{T}} + B^{-1}` of the strained cell, and the strain scales the volume by + :math:`\det S`; the fractional coordinates do not depend on it. The zero + mode carries no weight, nor do the modes at the Nyquist frequency of an + axis, which have no unique wave number. + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + mesh + The cells and the mesh of the periodic frames. + frame_id + The frame of every node with shape ``(N,)``. + kernel + The pair kernel. + order + The order of the B-splines. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + masked out by :attr:`MeshGrid.mask`. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + nf = mesh.box.shape[0] + n_node, n_channel = charge.shape + m_full = [xp.astype(m, dtype) for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)] + n1, n2, n3 = (m.shape[0] for m in m_full) + n_mesh = n1 * n2 * n3 + box = xp.astype(mesh.box, dtype) + volume = xp.astype(mesh.volume, dtype) + # === Step 1. Stencils of the nodes on the mesh === + box_inv = _inverse3(xp, box) + # frac_d = sum_j r_j (B^-1)_jd, elementwise rather than a batched product + # of 3x3 matrices. + frac = xp.sum( + xp.astype(mesh.coord, dtype)[:, :, None] * xp.take(box_inv, frame_id, axis=0), + axis=1, + ) + frac = frac - xp.floor(frac) + index, weight = _stencil(xp, frac, (n1, n2, n3), order) + # Flat targets on the frame- and channel-major meshes, (N, C, order**3). + channel = xp.arange(n_channel, dtype=xp.int64, device=device) + first = (frame_id[:, None] * n_channel + channel[None, :]) * n_mesh + target = first[:, :, None] + index[:, None, :] + # === Step 2. Charge mesh and its half spectrum === + q_mesh = xp.zeros((nf * n_channel * n_mesh,), dtype=dtype, device=device) + q_mesh = xp_add_at( + q_mesh, + xp.reshape(target, (-1,)), + xp.reshape(charge[:, :, None] * weight[:, None, :], (-1,)), + ) + q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) + spectrum = xp.fft.rfftn(q_mesh, axes=(-3, -2, -1)) + # === Step 3. Influence function on the half spectrum, (nf, n1, n2, h3) === + m3 = m_full[2][: n3 // 2 + 1] + m_half = [m_full[0], m_full[1], m3] + eye = xp.eye(3, dtype=dtype, device=device) + strain = xp.astype(mesh.strain, dtype) + eye + a = xp.matmul(xp.matrix_transpose(xp.linalg.inv(strain)), box_inv) + metric = xp.matmul(xp.matrix_transpose(a), a) + volume = volume * xp.linalg.det(strain) + g1 = m_half[0][:, None, None] + g2 = m_half[1][None, :, None] + g3 = m_half[2][None, None, :] + # Six monomials of the mode indices times six entries of the metric of + # every frame, elementwise. + monomials = (g1 * g1, g2 * g2, g3 * g3, g1 * g2, g1 * g3, g2 * g3) + entries = ( + metric[:, 0, 0], + metric[:, 1, 1], + metric[:, 2, 2], + 2.0 * metric[:, 0, 1], + 2.0 * metric[:, 0, 2], + 2.0 * metric[:, 1, 2], + ) + k_sq = None + for entry, monomial in zip(entries, monomials, strict=True): + term = entry[:, None, None, None] * monomial[None, ...] + k_sq = term if k_sq is None else k_sq + term + k_sq = (2.0 * math.pi) ** 2 * k_sq + # The zero mode and the Nyquist planes carry no weight. + keep = ( + (2.0 * xp.abs(g1) != n1) & (2.0 * xp.abs(g2) != n2) & (2.0 * xp.abs(g3) != n3) + ) + keep = keep & ((g1 != 0.0) | (g2 != 0.0) | (g3 != 0.0)) + keep = keep[None, ...] & mesh.mask[:, None, None, None] + safe_k_sq = xp.where(keep, k_sq, xp.ones_like(k_sq)) + shape = safe_k_sq.shape + transform = xp.reshape(kernel.kspace(xp.reshape(safe_k_sq, (-1,))), shape) + f1, f2, f3 = _assignment_factor(m_half, (n1, n2, n3), order) + factor = f1[:, None, None] * f2[None, :, None] * f3[None, None, :] + influence = (COULOMB_CONSTANT * transform * factor[None, ...]) / volume[ + :, None, None, None + ] + influence = xp.where(keep, influence, xp.zeros_like(influence)) + # === Step 4. Potential on the mesh and at the nodes === + spectrum = xp.reshape(spectrum, (nf, n_channel, n1, n2, n3 // 2 + 1)) + potential_mesh = xp.fft.irfftn( + spectrum * influence[:, None, ...], + s=(n1, n2, n3), + axes=(-3, -2, -1), + norm="forward", + ) + stencil_potential = xp.take( + xp.reshape(potential_mesh, (-1,)), xp.reshape(target, (-1,)), axis=0 + ) + potential = xp.sum( + xp.reshape(stencil_potential, (n_node, n_channel, -1)) * weight[:, None, :], + axis=-1, + ) + periodic_node = xp.take(mesh.mask, frame_id, axis=0) + potential = potential - _self_potential(kernel, charge, periodic_node) + return 0.5 * xp.sum(charge * potential, axis=-1) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 4456eda91d..4788e27b29 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -46,9 +46,9 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, - PairGraph, build_neighbor_graph, build_neighbor_graph_ase, + cast_topology_geometry, compact_nodes, expand_node_values, ) @@ -891,48 +891,11 @@ def call_common_lower_graph( spin=spin, ) if topologies: - # The pair vectors and the k-grid/mesh geometry are geometry - # like the edge vectors. - from deepmd.dpmodel.utils.neighbor_graph import ( - KGrid, - MeshGrid, - ) - - xp = array_api_compat.array_namespace(edge_vec) - cast = {} - for name, topology in topologies.items(): - if isinstance(topology, KGrid): - cast[name] = KGrid( - frame=topology.frame, - vec=xp.astype(topology.vec, edge_vec.dtype), - mask=topology.mask, - volume=xp.astype(topology.volume, edge_vec.dtype), - coord=( - None - if topology.coord is None - else xp.astype(topology.coord, edge_vec.dtype) - ), - strain=topology.strain, - ) - elif isinstance(topology, MeshGrid): - cast[name] = MeshGrid( - box=xp.astype(topology.box, edge_vec.dtype), - volume=xp.astype(topology.volume, edge_vec.dtype), - coord=( - None - if topology.coord is None - else xp.astype(topology.coord, edge_vec.dtype) - ), - strain=topology.strain, - mask=topology.mask, - ) - else: - cast[name] = PairGraph( - topology.index, - xp.astype(topology.vec, edge_vec.dtype), - topology.mask, - ) - topologies = cast + # The topologies carry geometry like the edge vectors. + topologies = { + name: cast_topology_geometry(topology, edge_vec.dtype) + for name, topology in topologies.items() + } model_predict = self.forward_common_atomic_graph( atype, n_node, diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 3d427c0ec4..f5667a4510 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -6,9 +6,11 @@ + edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the ``from_dense_quartet`` legacy converter), ``pair_graph`` (the complete -intra-frame pair topology of interactions without a cutoff), ``segment`` -(mask-aware segment-reduction toolkit), and ``derivatives`` (edge and pair -force/virial assembly). +intra-frame pair topology of interactions without a cutoff), +``reciprocal_graph`` (the reciprocal-space topologies of periodic frames), +``topology`` (building and laying out the topologies a model declares), +``segment`` (mask-aware segment-reduction toolkit), and ``derivatives`` (edge +and pair force/virial assembly). See the design discussion wanghan-iapcm/deepmd-kit#4. """ @@ -58,26 +60,26 @@ ) from .pair_graph import ( ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + unordered_pair_index, +) +from .pairs import ( + center_edge_pairs, +) +from .reciprocal_graph import ( K_GRID, PME_MESH, KGrid, MeshGrid, - PairGraph, build_fixed_k_grid, build_k_grid, build_mesh_grid, - build_pair_graph, - build_topologies, - check_complete_frames, + check_periodic_cells, enumerate_k_grid_bounds, - flatten_topologies, k_grid_bounds_of_cells, - topology_flat_fields, - unflatten_topologies, - unordered_pair_index, -) -from .pairs import ( - center_edge_pairs, + mesh_dims_of_cells, ) from .segment import ( segment_max, @@ -85,11 +87,20 @@ segment_softmax, segment_sum, ) +from .topology import ( + TOPOLOGY_TYPES, + build_topologies, + cast_topology_geometry, + flatten_topologies, + topology_flat_fields, + unflatten_topologies, +) __all__ = [ "ALL_PAIRS", "K_GRID", "PME_MESH", + "TOPOLOGY_TYPES", "GraphLayout", "KGrid", "MeshGrid", @@ -112,8 +123,10 @@ "build_pair_graph", "build_topologies", "canonicalize_neighbor_graph", + "cast_topology_geometry", "center_edge_pairs", "check_complete_frames", + "check_periodic_cells", "compact_nodes", "edge_env_mat", "edge_force_virial", @@ -125,6 +138,7 @@ "graph_angle_cos", "graph_from_dense_quartet", "k_grid_bounds_of_cells", + "mesh_dims_of_cells", "neighbor_graph_from_ijs", "node_ownership_mask", "node_validity_mask", diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index 64c3ca0f82..8e68625d78 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -22,21 +22,15 @@ annotations, ) -import math -import warnings from dataclasses import ( dataclass, ) from typing import ( TYPE_CHECKING, + ClassVar, ) import array_api_compat -import numpy as np - -from deepmd.dpmodel.common import ( - to_numpy_array, -) from .graph import ( frame_id_from_n_node, @@ -47,11 +41,6 @@ ) if TYPE_CHECKING: - from collections.abc import ( - Iterable, - Sequence, - ) - from deepmd.dpmodel.array_api import ( Array, ) @@ -62,25 +51,6 @@ ALL_PAIRS = "all_pairs" """Name of the topology holding every unordered pair of valid nodes of a frame.""" -K_GRID = "k_grid" -"""Name of the topology holding the reciprocal lattice points of periodic frames. - -A long-range module declares it beside :data:`ALL_PAIRS` when it serves -periodic systems: the reciprocal-space evaluation of the pair kernel needs the -wave numbers, which the pair topology cannot carry. -""" - -PME_MESH = "pme_mesh" -"""Name of the topology carrying the cells of periodic frames for a mesh evaluation. - -A long-range module declares it beside :data:`ALL_PAIRS` when it evaluates the -reciprocal sum of a smooth kernel on a real-space mesh (SPME): charges spread -to a static mesh by a B-spline, transformed by an FFT, and convolved with the -kernel transform. The mesh axes are static and live on the module, so the -topology holds no ragged axis of its own — only the per-frame cells and the -geometry leaves. -""" - @dataclass class PairGraph: @@ -93,79 +63,16 @@ class PairGraph: mask: Array """(P,) bool real (1) vs guard (0) pairs; guards join node 0 to itself.""" - -@dataclass -class KGrid: - r"""Reciprocal lattice points of the periodic frames of a batch. - - The wave numbers enumerate :math:`\\mathbf{k} = 2\\pi B^{-1}\\mathbf{n}` for - integer ``n`` with :math:`0 < |\\mathbf{k}| \\le k_{max}`; they are built - eagerly outside any compiled region, concatenated ragged over the frames, - and the k axis ends in masked guard points that keep its length a dynamic - dimension of its own however few points a batch holds. - """ - - frame: Array - """(nk,) int64 the frame of every wave number.""" - vec: Array - """(nk, 3) float the wave numbers, in Å⁻¹.""" - mask: Array - """(nk,) bool real (1) vs guard (0) points.""" - volume: Array - """(nf,) float the cell volumes, in ų.""" - coord: Array | None = None - """(N, 3) float node positions; the geometry leaf of the k-space energy. - - ``None`` when the grid is reconstructed from flattened inputs (export); a - graph lower re-leafs it before the evaluation. - """ - strain: Array | None = None - """(nf, 3, 3) float zero strain leaf for the k-space virial. - - The module evaluates the energy on coordinates strained by ``I + strain``, - wave numbers strained by ``(I + strain)^{-T}`` and volumes scaled by - ``det(I + strain)``; the gradient of the energy with respect to this leaf - is the negative k-space virial. The leaf is built eagerly beside the grid - (all zeros) and re-leafed by the graph lower, so the backward roots at an - input of the traced graph — a leaf created INSIDE a ``make_fx`` trace - produces a NaN gradient under inductor. - """ - - -@dataclass -class MeshGrid: - r"""Cells and geometry of the periodic frames for a mesh (FFT) evaluation. - - The reciprocal sum is evaluated on a static real-space mesh shared by the - frames, so the topology carries no ragged axis: the cells (for the - reciprocal vectors), the volumes, the geometry leaves, and a per-frame - flag telling periodic frames from non-periodic ones (a batch without a - cell comes back with an all-false mask). The volume is computed eagerly - (as in :class:`KGrid`): the graph lower only rescales it by the strain - determinant, which keeps ``linalg.det`` out of the traced graph. - """ - - box: Array - """(nf, 3, 3) float simulation cells; rows are the lattice vectors.""" - volume: Array - """(nf,) float the cell volumes, in ų.""" - coord: Array | None - """(N, 3) float node positions; the geometry leaf of the mesh energy. - - ``None`` when the grid is reconstructed from flattened inputs (export); a - graph lower re-leafs it before the evaluation. - """ - strain: Array | None - """(nf, 3, 3) float zero strain leaf for the mesh virial. - - Same contract as :attr:`KGrid.strain`: the energy is evaluated with the - coordinates strained by ``I + strain`` and the cells by - ``B' = B (I + strain)^T`` with the lattice vectors as rows, so the - fractional coordinates are strain-invariant and the gradient with respect - to the leaf is the negative virial. - """ - mask: Array - """(nf,) bool periodic (1) vs non-periodic (0) frames.""" + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec",) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("vec",) + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "index": {1: "entry"}, + "vec": {0: "entry"}, + "mask": {0: "entry"}, + } + """Dynamic axes of every field: the pair axis.""" def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: @@ -234,323 +141,6 @@ def row_start(row: Array) -> Array: return xp.stack([src, dst]) -def build_k_grid( - box: Array, - k_max: float, - n_node: Array, - coord: Array | None = None, -) -> KGrid: - r"""Enumerate the reciprocal lattice points of the periodic frames. - - For every frame the wave numbers are - :math:`\\mathbf{k} = 2\\pi \\, B^{-1} \\mathbf{n}` over integer - :math:`\\mathbf{n}` with :math:`0 < |\\mathbf{k}| \\le k_{max}` — note the - cell enters through ``B^{-1}``, NOT ``B^{-T}``: with the lattice vectors as - the ROWS of ``B``, orthogonality reads ``B @ (B^{-1} n) = n``. The count is - data dependent, so the enumeration runs eagerly. Every frame's points come - in lexicographic integer order, and the batch ends in two masked guard - points. - - Parameters - ---------- - box - Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice - vectors. - k_max - The wave-number cutoff, in Å⁻¹. - n_node - Per-frame node counts with shape ``(nf,)``. - coord - Node positions with shape ``(N, 3)``, stored on the grid as the - geometry context (see :attr:`KGrid.coord`). - - Returns - ------- - KGrid - The reciprocal lattice points of every frame, ragged-concatenated. - """ - xp = array_api_compat.array_namespace(box) - device = array_api_compat.device(box) - # A flat (nf, 9) layout arrives from callers that carry the cell flattened. - if box.ndim == 2: - box = xp.reshape(box, (box.shape[0], 3, 3)) - nf = box.shape[0] - k_lists = [] - frame_lists = [] - for f in range(nf): - cell = np.asarray(to_numpy_array(box[f]), dtype=np.float64) - rec = 2.0 * math.pi * np.linalg.inv(cell) # columns are g_d - a_len = np.linalg.norm(cell, axis=1) - n_d = np.ceil(k_max * a_len / (2.0 * math.pi)).astype(np.int64) - grids = np.meshgrid( - *[np.arange(-n, n + 1, dtype=np.float64) for n in n_d], indexing="ij" - ) - n_all = np.stack(grids, axis=-1).reshape(-1, 3) - k = n_all @ rec.T - k2 = np.sum(k * k, axis=-1) - keep = (k2 > 0.0) & (k2 <= k_max * k_max) - k_lists.append(k[keep]) - frame_lists.append(np.full(int(np.sum(keep)), f, dtype=np.int64)) - if k_lists: - vec_np = np.concatenate(k_lists, axis=0) - frame_np = np.concatenate(frame_lists, axis=0) - else: - vec_np = np.zeros((0, 3), dtype=np.float64) - frame_np = np.zeros((0,), dtype=np.int64) - # Guard points: masked out, keep the k axis a dynamic dimension of its own. - guard = np.tile(np.array([[1e-6, 0.0, 0.0]]), (2, 1)) - vec_np = np.concatenate([vec_np, guard], axis=0) - frame_np = np.concatenate([frame_np, np.zeros(2, dtype=np.int64)], axis=0) - mask_np = np.zeros(vec_np.shape[0], dtype=bool) - mask_np[: vec_np.shape[0] - 2] = True - volume = xp.linalg.det(box) - volume = xp.abs(volume) - return KGrid( - frame=xp.asarray(frame_np, device=device), - vec=xp.astype(xp.asarray(vec_np, device=device), box.dtype), - mask=xp.asarray(mask_np, device=device), - volume=xp.astype(volume, box.dtype), - coord=coord, - strain=xp.zeros((nf, 3, 3), dtype=box.dtype, device=device), - ) - - -def k_grid_bounds_of_cells(box: Array, k_max: float) -> np.ndarray: - r"""Per-axis integer wave-number bounds the cells need at ``k_max``. - - A wave number :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` with - :math:`|\mathbf{k}| \le k_{max}` needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` - per axis (Cauchy-Schwarz), so the bound of every frame is the ceiling of - the right-hand side. - - Parameters - ---------- - box - Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are - the lattice vectors. - k_max - The wave-number cutoff, in Å⁻¹. - - Returns - ------- - np.ndarray - The per-frame, per-axis integer bounds with shape ``(nf, 3)``. - """ - cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) - a_len = np.linalg.norm(cell, axis=-1) - return np.ceil(k_max * a_len / (2.0 * math.pi)).astype(np.int64) - - -def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: - r"""Enumerate the integer wave-number indices within per-axis bounds. - - The grid is the full meshgrid of :math:`\mathbf{n}` with - :math:`|n_d| \le n_d^{max}`, in lexicographic order (the first axis - outermost). The zero vector is left out: its kernel transform is - singular and the tin-foil convention drops the :math:`\mathbf{k} = 0` - mode anyway. - - Parameters - ---------- - bounds - The per-axis bounds :math:`n_d^{max}` with shape ``(3,)``. - - Returns - ------- - np.ndarray - The integer indices with shape ``(nk, 3)``, int64. - """ - bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(-1) - if bounds.shape != (3,) or np.any(bounds < 0): - raise ValueError( - f"k-grid bounds must be three non-negative integers, got {bounds}" - ) - grids = np.meshgrid( - *[np.arange(-n, n + 1, dtype=np.int64) for n in bounds], indexing="ij" - ) - n_grid = np.stack(grids, axis=-1).reshape(-1, 3) - return n_grid[np.any(n_grid != 0, axis=-1)] - - -def build_fixed_k_grid( - box: Array, - n_grid: Array, - k_max: float, - n_node: Array, - coord: Array | None = None, -) -> KGrid: - r"""Evaluate a fixed integer k grid on the cells of the periodic frames. - - Every frame shares the integer indices ``n_grid`` (see - :func:`enumerate_k_grid_bounds`), so the k axis has the static length - ``nf * nk``; the wave numbers of a frame are - :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` and the points the frame does - not reach — outside :math:`0 < |\mathbf{k}| \le k_{max}` — are masked out. - The mask is evaluated eagerly with the wave-number cutoff current at - build time, exactly as in :func:`build_k_grid`. The enumeration runs - eagerly beside the ragged one, but no count is data dependent: the grid - comes back without guard points. - - The grid is HARD-FIXED: a frame whose cell would need integer indices - beyond the grid (a larger cell than the grid was sized for, or a kernel - that moved the cutoff out since) is served by the grid alone — its - reciprocal lattice is denser than the grid, so its energy is the - truncated k-sum of the grid, a well-defined approximation whose forces - and virial stay the exact derivatives of that energy. A warning is - issued when this happens. - - Parameters - ---------- - box - Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice - vectors. - n_grid - The shared integer wave-number indices with shape ``(nk, 3)``. - k_max - The wave-number cutoff, in Å⁻¹. - n_node - Per-frame node counts with shape ``(nf,)`` (interface symmetry with - :func:`build_k_grid`; the grid does not read it). - coord - Node positions with shape ``(N, 3)``, stored on the grid as the - geometry context (see :attr:`KGrid.coord`). - - Returns - ------- - KGrid - The wave numbers of every frame, frame-major over the shared grid. - - Raises - ------ - ValueError - Never for an oversized cell (it is truncated, see above); kept for - interface symmetry. - """ - xp = array_api_compat.array_namespace(box) - device = array_api_compat.device(box) - if box.ndim == 2: - box = xp.reshape(box, (box.shape[0], 3, 3)) - nf = box.shape[0] - n_grid = np.asarray(to_numpy_array(n_grid), dtype=np.int64).reshape(-1, 3) - nk = n_grid.shape[0] - bounds = np.max(np.abs(n_grid), axis=0) if nk else np.zeros(3, dtype=np.int64) - required = k_grid_bounds_of_cells(box, k_max) - over = np.any(required > bounds[None, :], axis=1) - if np.any(over): - f = int(np.argmax(over)) - warnings.warn( - f"frame {f} needs integer k-grid bounds {required[f].tolist()}, " - f"beyond the fixed grid {bounds.tolist()}: its k-sum is truncated " - f"to the grid (an approximation for this frame); enlarge " - f"`k_grid_bounds` for full accuracy", - stacklevel=2, - ) - cells = np.asarray(to_numpy_array(box), dtype=np.float64) - rec = 2.0 * math.pi * np.linalg.inv(cells) # (nf, 3, 3), columns are g_d - k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) - k2 = np.sum(k * k, axis=-1) - mask_np = (k2 > 0.0) & (k2 <= k_max * k_max) # (nf, nk) - vec_np = k.reshape(nf * nk, 3) - frame_np = np.repeat(np.arange(nf, dtype=np.int64), nk) - volume = xp.linalg.det(box) - volume = xp.abs(volume) - return KGrid( - frame=xp.asarray(frame_np, device=device), - vec=xp.astype(xp.asarray(vec_np, device=device), box.dtype), - mask=xp.asarray(mask_np.reshape(-1), device=device), - volume=xp.astype(volume, box.dtype), - coord=coord, - strain=xp.zeros((nf, 3, 3), dtype=box.dtype, device=device), - ) - - -def _empty_k_grid(nf: int, like: Array, coord: Array | None) -> KGrid: - """A guard-only k grid for a batch without periodic frames.""" - xp = array_api_compat.array_namespace(like) - device = array_api_compat.device(like) - return KGrid( - frame=xp.zeros((2,), dtype=xp.int64, device=device), - vec=xp.tile( - xp.asarray([[1e-6, 0.0, 0.0]], dtype=like.dtype, device=device), (2, 1) - ), - mask=xp.zeros((2,), dtype=xp.bool, device=device), - volume=xp.ones((nf,), dtype=like.dtype, device=device), - coord=coord, - strain=xp.zeros((nf, 3, 3), dtype=like.dtype, device=device), - ) - - -def _empty_mesh_grid(nf: int, like: Array, coord: Array | None) -> MeshGrid: - """A non-periodic mesh grid: every frame masked out, identity cells.""" - xp = array_api_compat.array_namespace(like) - device = array_api_compat.device(like) - return MeshGrid( - box=xp.tile(xp.eye(3, dtype=like.dtype, device=device), (nf, 1, 1)), - volume=xp.ones((nf,), dtype=like.dtype, device=device), - coord=coord, - strain=xp.zeros((nf, 3, 3), dtype=like.dtype, device=device), - mask=xp.zeros((nf,), dtype=xp.bool, device=device), - ) - - -def build_mesh_grid( - box: Array, - n_mesh: Array, - coord: Array | None = None, -) -> MeshGrid: - r"""Pack the cells of the periodic frames for the mesh (FFT) evaluation. - - The mesh is shared by every frame; a frame whose cell is larger than the - mesh was sized for is evaluated with a coarser effective spacing - :math:`|a_d|/n_d` than the target — an undersampled kernel spectrum, a - well-defined approximation whose forces and virial stay the exact - derivatives of that energy. Sizing the mesh (data statistics or - ``pme_mesh``) is the caller's responsibility; no warning is issued here. - - Parameters - ---------- - box - Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice - vectors. - n_mesh - The per-axis mesh dimensions with shape ``(3,)``. - coord - Node positions with shape ``(N, 3)``, stored on the grid as the - geometry context (see :attr:`MeshGrid.coord`). - - Returns - ------- - MeshGrid - The cells and geometry of the frames, every frame flagged periodic. - """ - xp = array_api_compat.array_namespace(box) - if box.ndim == 2: - box = xp.reshape(box, (box.shape[0], 3, 3)) - volume = xp.linalg.det(box) - volume = xp.abs(volume) - return MeshGrid( - box=box, - volume=xp.astype(volume, box.dtype), - coord=coord, - strain=xp.zeros( - (box.shape[0], 3, 3), dtype=box.dtype, device=array_api_compat.device(box) - ), - mask=xp.ones( - (box.shape[0],), dtype=xp.bool, device=array_api_compat.device(box) - ), - ) - - -def _empty_pair_graph(like: Array, atype: Array) -> PairGraph: - """A guard-only pair topology for a batch without non-periodic frames.""" - xp = array_api_compat.array_namespace(like) - device = array_api_compat.device(like) - index = xp.zeros((2, 0), dtype=xp.int64, device=device) - vec = xp.zeros((0, 3), dtype=like.dtype, device=device) - index, vec, mask = pad_and_guard_edges(index, vec, None) - return PairGraph(index=index, vec=vec, mask=mask) - - def build_pair_graph( coord: Array, n_node: Array, @@ -592,113 +182,14 @@ def build_pair_graph( return PairGraph(index=index, vec=vec, mask=mask) -def build_topologies( - names: Iterable[str], - coord: Array, - n_node: Array, - box: Array | None = None, - node_mask: Array | None = None, - *, - atype: Array, - pair_excl: PairExcludeMask | None = None, - topology_params: dict[str, dict] | None = None, -) -> dict[str, PairGraph | KGrid]: - """Build the topologies an atomic model declares, beside its neighbor graph. - - Parameters - ---------- - names - Topology names from ``required_topologies()`` of the atomic model. - coord - Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. - n_node - Per-frame node counts with shape ``(nf,)``. - box - Simulation cells, or ``None`` for non-periodic systems. With a cell, - the pair topology comes back guard-only and the k grid holds the - reciprocal lattice points of the frames; without one, the reverse. - node_mask - Validity of the nodes with shape ``(N,)``, or ``None`` when every - node is valid. - atype - Node types with shape ``(N,)``; invalid nodes carry any type. - pair_excl - The model-level pair exclusion, or ``None``. - topology_params - Per-name build parameters; the k grid takes ``{"k_max": ...}`` and, - for a fixed integer grid, additionally ``{"n_grid": ...}`` (see - :func:`build_fixed_k_grid`). - - Returns - ------- - dict[str, PairGraph | KGrid] - Each requested topology under its name; empty when none is required. - - Raises - ------ - NotImplementedError - If a cell is given for a pair-only declaration, or a pair exclusion is - set for a periodic batch (the k-space sum is global). - ValueError - If a name is unknown, or the k grid parameters are missing. - """ - topologies: dict[str, PairGraph | KGrid | MeshGrid] = {} - for name in names: - if name == ALL_PAIRS: - if box is None: - topologies[name] = build_pair_graph( - coord, n_node, atype, node_mask, pair_excl - ) - elif K_GRID in names or PME_MESH in names: - # A periodic-capable module reads the k grid or the mesh for - # these frames. - topologies[name] = _empty_pair_graph(coord, atype) - else: - check_complete_frames(box) - elif name == K_GRID: - if box is None: - topologies[name] = _empty_k_grid(n_node.shape[0], coord, coord) - else: - if pair_excl is not None: - raise NotImplementedError( - "the k-space sum of a periodic system is global; " - "`pair_exclude_types` cannot exclude pairs from it" - ) - params = (topology_params or {}).get(K_GRID, {}) - if "k_max" not in params: - raise ValueError( - f"the {K_GRID!r} topology requires its `k_max` build parameter" - ) - if "n_grid" in params: - topologies[name] = build_fixed_k_grid( - box, params["n_grid"], params["k_max"], n_node, coord=coord - ) - else: - topologies[name] = build_k_grid( - box, params["k_max"], n_node, coord=coord - ) - elif name == PME_MESH: - if box is None: - topologies[name] = _empty_mesh_grid(n_node.shape[0], coord, coord) - else: - if pair_excl is not None: - raise NotImplementedError( - "the k-space sum of a periodic system is global; " - "`pair_exclude_types` cannot exclude pairs from it" - ) - params = (topology_params or {}).get(PME_MESH, {}) - if "n_mesh" not in params: - raise ValueError( - f"the {PME_MESH!r} topology requires its `n_mesh` " - f"build parameter" - ) - topologies[name] = build_mesh_grid(box, params["n_mesh"], coord=coord) - else: - raise ValueError( - f"unknown topology {name!r}; expected one of " - f"({ALL_PAIRS!r}, {K_GRID!r}, {PME_MESH!r})" - ) - return topologies +def empty_pair_graph(like: Array) -> PairGraph: + """The pair topology of a batch without non-periodic frames: guard pairs only.""" + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + index = xp.zeros((2, 0), dtype=xp.int64, device=device) + vec = xp.zeros((0, 3), dtype=like.dtype, device=device) + index, vec, mask = pad_and_guard_edges(index, vec, None) + return PairGraph(index=index, vec=vec, mask=mask) def check_complete_frames(box: Array | None) -> None: @@ -717,112 +208,7 @@ def check_complete_frames(box: Array | None) -> None: if box is not None: raise NotImplementedError( "the all-pairs topology of the long-range modules describes " - "complete non-periodic systems only; periodic cells are not " - "supported (pass box=None for a non-periodic system)" - ) - - -def flatten_topologies(topologies: dict[str, PairGraph | KGrid]) -> tuple[Array, ...]: - """Lay topologies out as positional inputs of a graph lower. - - A compiled or exported graph lower receives every topology positionally, - in the order of the dict, after its other inputs: a pair topology as its - ``(index, vec, mask)``, a k grid as its ``(frame, vec, mask, volume, - coord)``; - :func:`unflatten_topologies` restores them. - - Parameters - ---------- - topologies - The topologies by name. - - Returns - ------- - tuple - The flattened tensors of every topology, concatenated. - """ - return tuple( - value - for topology in topologies.values() - for value in topology_flat_fields(topology) - ) - - -def topology_flat_fields(topology: PairGraph | KGrid | MeshGrid) -> tuple[Array, ...]: - """The positional tensors of one topology.""" - if isinstance(topology, KGrid): - if topology.coord is None or topology.strain is None: - raise ValueError( - "a k grid carries no coordinates or strain; it cannot be laid out " - "positionally" - ) - return ( - topology.frame, - topology.vec, - topology.mask, - topology.volume, - topology.coord, - topology.strain, - ) - if isinstance(topology, MeshGrid): - if topology.coord is None or topology.strain is None: - raise ValueError( - "a mesh grid carries no coordinates or strain; it cannot be laid " - "out positionally" - ) - return ( - topology.box, - topology.volume, - topology.coord, - topology.strain, - topology.mask, - ) - return (topology.index, topology.vec, topology.mask) - - -def unflatten_topologies( - names: Sequence[str], values: Sequence[Array] -) -> dict[str, PairGraph | KGrid]: - """Restore topologies laid out by :func:`flatten_topologies`. - - Parameters - ---------- - names - The topology names, in the order they were laid out. - values - The flattened tensors of every topology, concatenated. - - Returns - ------- - dict[str, PairGraph | KGrid] - The topologies by name. - - Raises - ------ - ValueError - If the number of values does not match the layout of the names. - """ - arity = {ALL_PAIRS: 3, K_GRID: 6, PME_MESH: 5} - topologies: dict[str, PairGraph | KGrid | MeshGrid] = {} - offset = 0 - for name in names: - if name not in arity: - raise ValueError(f"unknown topology {name!r}") - fields = values[offset : offset + arity[name]] - if len(fields) != arity[name]: - raise ValueError( - f"expected {arity[name]} inputs for topology {name!r}, " - f"got {len(values) - offset} remaining" - ) - if name == ALL_PAIRS: - topologies[name] = PairGraph(*fields) - elif name == K_GRID: - topologies[name] = KGrid(*fields) - else: - topologies[name] = MeshGrid(*fields) - offset += arity[name] - if offset != len(values): - raise ValueError( - f"{len(values)} flattened values given, {offset} consumed by {tuple(names)}" + "complete non-periodic systems only; a module evaluates periodic " + "cells with `periodic: true` (pass box=None for a non-periodic " + "system)" ) - return topologies diff --git a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py new file mode 100644 index 0000000000..fdcd2c57d0 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reciprocal-space topologies of the periodic frames of a batch. + +A long-range module that serves periodic systems evaluates the lattice sum of +its pair kernel in reciprocal space, which needs the cell of every frame and +the absolute node positions rather than pair displacements. It declares one of +two topologies beside +:data:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.ALL_PAIRS`: + +- :data:`K_GRID` lists the wave numbers of every frame explicitly, ragged over + the frames, for a direct sum over them; +- :data:`PME_MESH` carries the cells and the mode indices of a real-space mesh + shared by the frames, for a smooth particle-mesh evaluation. + +Both carry the node positions and a zero strain as geometry: an autograd +backend makes them leaves, whose gradients are the force and the negative +virial of the reciprocal-space energy. A batch without cells comes back with +masked topologies, so a module serving periodic systems also serves +non-periodic batches through its pair topology. +""" + +from __future__ import ( + annotations, +) + +import math +import warnings +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + +K_GRID = "k_grid" +"""Name of the topology holding the reciprocal lattice points of periodic frames.""" + +PME_MESH = "pme_mesh" +"""Name of the topology carrying the cells and the mesh of periodic frames.""" + + +@dataclass +class KGrid: + r"""Reciprocal lattice points of the periodic frames of a batch. + + The wave numbers :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` of every frame, + for integer :math:`\mathbf{n}`, are built eagerly outside any compiled + region and concatenated over the frames. The k axis ends in masked guard + points, which keep its length a dynamic dimension of its own however few + points a batch holds. + + The energy is evaluated on the strained geometry of :attr:`strain`: + positions by ``S = I + strain``, wave numbers by :math:`S^{-T}` and volumes + by :math:`\det S`, so that the gradient with respect to the strain is the + negative virial. The graph lower makes :attr:`coord` and :attr:`strain` + autograd leaves at its inputs. + """ + + frame: Array + """(nk,) int64 the frame of every wave number.""" + vec: Array + """(nk, 3) float the wave numbers, in Å⁻¹.""" + mask: Array + """(nk,) bool real (1) vs guard (0) points.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec", "volume", "coord", "strain") + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "frame": {0: "entry"}, + "vec": {0: "entry"}, + "mask": {0: "entry"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + } + """Dynamic axes of every field: the k axis, the frame axis, the node axis.""" + + +@dataclass +class MeshGrid: + r"""Cells and mesh of the periodic frames of a batch. + + The frames share a real-space mesh of :math:`n_1 \times n_2 \times n_3` + points spanning each cell, whose integer mode indices along every axis + are the fields :attr:`mode_1`, :attr:`mode_2` and :attr:`mode_3`: their + lengths are the mesh dimensions, so a compiled or exported graph takes the + mesh as dynamic dimensions. Every dimension is even. A batch without a + cell comes back with identity cells, a minimal mesh and an all-false + :attr:`mask`. + + The geometry contract is that of :class:`KGrid`: the energy is evaluated + with the positions strained by ``S = I + strain`` and the cells by + :math:`B S^{\mathsf{T}}`, the lattice vectors being the rows of + :math:`B`, so that the fractional coordinates are invariant and the + gradient with respect to the strain is the negative virial. The volumes + of the unstrained cells are computed eagerly, and the evaluation scales + them by :math:`\det S`. + """ + + box: Array + """(nf, 3, 3) float simulation cells; rows are the lattice vectors, in Å.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + mask: Array + """(nf,) bool periodic (1) vs non-periodic (0) frames.""" + mode_1: Array + """(n1,) float mode indices of the first mesh axis, in FFT order.""" + mode_2: Array + """(n2,) float mode indices of the second mesh axis, in FFT order.""" + mode_3: Array + """(n3,) float mode indices of the third mesh axis, in FFT order.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ( + "box", + "volume", + "coord", + "strain", + "mode_1", + "mode_2", + "mode_3", + ) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "box": {0: "frame"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + "mask": {0: "frame"}, + "mode_1": {0: "mesh_1"}, + "mode_2": {0: "mesh_2"}, + "mode_3": {0: "mesh_3"}, + } + """Dynamic axes of every field: the frame axis, the node axis, the mesh axes.""" + + +def check_periodic_cells(box: Array) -> None: + """Reject a cell that spans no lattice. + + A cell whose volume is a vanishing fraction of the product of the lengths + of its lattice vectors, or that is not finite, spans no lattice; an + all-zero cell is no marker of a non-periodic frame, which takes no cell. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``; rows are + the lattice vectors. + + Raises + ------ + ValueError + If a cell spans no lattice. + """ + cells = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + volume = np.abs(np.linalg.det(cells)) + spanned = volume > 1e-6 * np.prod(np.linalg.norm(cells, axis=-1), axis=-1) + if not np.all(spanned): + f = int(np.argmin(spanned)) + raise ValueError( + f"the cell of periodic frame {f} has no volume; a periodic frame " + "needs three independent lattice vectors, and a non-periodic " + "system takes no cell" + ) + + +def _cells(box: Array) -> Array: + """The cells with shape ``(nf, 3, 3)``, from a ``(nf, 9)`` or ``(nf, 3, 3)`` array.""" + xp = array_api_compat.array_namespace(box) + return xp.reshape(box, (box.shape[0], 3, 3)) + + +def _volume(box: Array) -> Array: + xp = array_api_compat.array_namespace(box) + return xp.abs(xp.linalg.det(box)) + + +def _zero_strain(box: Array) -> Array: + xp = array_api_compat.array_namespace(box) + return xp.zeros( + (box.shape[0], 3, 3), dtype=box.dtype, device=array_api_compat.device(box) + ) + + +def k_grid_bounds_of_cells(box: Array, k_max: float) -> np.ndarray: + r"""Per-axis integer wave-number bounds the cells need at ``k_max``. + + A wave number :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` with + :math:`|\mathbf{k}| \le k_{max}` has :math:`|n_d| \le k_{max} |a_d| / 2\pi` + along every lattice vector :math:`\mathbf{a}_d` (Cauchy-Schwarz), so the + bound of every frame is the ceiling of the right-hand side. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + + Returns + ------- + np.ndarray + The per-frame, per-axis integer bounds with shape ``(nf, 3)``. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + length = np.linalg.norm(cell, axis=-1) + return np.ceil(k_max * length / (2.0 * math.pi)).astype(np.int64) + + +def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: + r"""Enumerate the integer wave-number indices within per-axis bounds. + + The grid holds every :math:`\mathbf{n} \neq \mathbf{0}` with + :math:`|n_d| \le n_d^{max}`, in lexicographic order with the first axis + outermost. The zero vector is left out: the transform of a Coulomb-like + kernel is singular there, and the tin-foil convention drops that mode. + + Parameters + ---------- + bounds + The per-axis bounds :math:`n_d^{max}` with shape ``(3,)``. + + Returns + ------- + np.ndarray + The integer indices with shape ``(nk, 3)``, int64. + + Raises + ------ + ValueError + If the bounds are not three non-negative integers. + """ + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0): + raise ValueError( + f"k-grid bounds must be three non-negative integers, got {bounds}" + ) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=np.int64) for n in bounds], indexing="ij" + ) + n_grid = np.stack(grids, axis=-1).reshape(-1, 3) + return n_grid[np.any(n_grid != 0, axis=-1)] + + +def _k_grid( + box: Array, frame: np.ndarray, vec: np.ndarray, mask: np.ndarray, coord: Array +) -> KGrid: + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + return KGrid( + frame=xp.asarray(frame, device=device), + vec=xp.astype(xp.asarray(vec, device=device), box.dtype), + mask=xp.asarray(mask, device=device), + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + ) + + +def build_k_grid(box: Array, k_max: float, coord: Array) -> KGrid: + r"""Enumerate the reciprocal lattice points of the periodic frames. + + With the lattice vectors as the rows of the cell :math:`B`, the wave + numbers of a frame are :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` over + integer :math:`\mathbf{n}` with :math:`0 < |\mathbf{k}| \le k_{max}`, so + that :math:`\mathbf{a}_d \cdot \mathbf{k} = 2\pi n_d`. Their count is data + dependent, so the enumeration runs eagerly. The points of every frame + follow the lexicographic order of their integer indices, and the batch + ends in two masked guard points. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + + Returns + ------- + KGrid + The reciprocal lattice points of every frame, concatenated. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + bounds = k_grid_bounds_of_cells(cells, k_max) + vec_list = [] + frame_list = [] + for f, cell in enumerate(cells): + n_grid = enumerate_k_grid_bounds(bounds[f]).astype(np.float64) + k = n_grid @ (2.0 * math.pi * np.linalg.inv(cell)).T + k = k[np.sum(k * k, axis=-1) <= k_max * k_max] + vec_list.append(k) + frame_list.append(np.full(k.shape[0], f, dtype=np.int64)) + # Two guard points with a nonzero wave number, whose kernel transform + # stays finite. + vec_list.append(np.array([[1e-6, 0.0, 0.0], [1e-6, 0.0, 0.0]])) + frame_list.append(np.zeros(2, dtype=np.int64)) + vec = np.concatenate(vec_list, axis=0) + mask = np.arange(vec.shape[0], dtype=np.int64) < vec.shape[0] - 2 + return _k_grid(box, np.concatenate(frame_list), vec, mask, coord) + + +def build_fixed_k_grid(box: Array, bounds: Array, k_max: float, coord: Array) -> KGrid: + r"""Evaluate a fixed integer k grid on the cells of the periodic frames. + + Every frame shares the integer indices within ``bounds`` (see + :func:`enumerate_k_grid_bounds`), so the k axis has the static length + ``nf * nk`` without guard points. The wave numbers of a frame are + :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}`, and the points beyond the + cutoff :math:`|\mathbf{k}| \le k_{max}` are masked out. + + A frame whose cell needs indices beyond the grid, a larger cell than the + grid was sized for or a kernel whose cutoff has grown since, is evaluated + on the grid alone: its sum is truncated to the points of the grid, which + is issued as a warning. Its forces and virial remain the derivatives of + that energy. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + bounds + The per-axis bounds :math:`n_d^{max}` of the shared integer grid with + shape ``(3,)``. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + + Returns + ------- + KGrid + The wave numbers of every frame, frame-major over the shared grid. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + n_grid = enumerate_k_grid_bounds(bounds) + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(3) + nf = cells.shape[0] + nk = n_grid.shape[0] + required = k_grid_bounds_of_cells(cells, k_max) + beyond = np.any(required > bounds[None, :], axis=1) + if np.any(beyond): + f = int(np.argmax(beyond)) + warnings.warn( + f"frame {f} needs integer k-grid bounds {required[f].tolist()}, " + f"beyond the fixed grid {bounds.tolist()}: its k-space sum is " + "truncated to the grid; enlarge `k_grid_bounds` for the full sum", + stacklevel=2, + ) + rec = 2.0 * math.pi * np.linalg.inv(cells) + k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) + k_sq = np.sum(k * k, axis=-1) + mask = (k_sq > 0.0) & (k_sq <= k_max * k_max) + frame = np.repeat(np.arange(nf, dtype=np.int64), nk) + return _k_grid(box, frame, k.reshape(nf * nk, 3), mask.reshape(-1), coord) + + +def empty_k_grid(nf: int, coord: Array) -> KGrid: + """The k grid of a batch without periodic frames: two guard points.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + eye = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + return _k_grid( + eye, + np.zeros(2, dtype=np.int64), + np.array([[1e-6, 0.0, 0.0], [1e-6, 0.0, 0.0]]), + np.zeros(2, dtype=bool), + coord, + ) + + +def _even_fft_size(n: int) -> int: + """The smallest even integer of at least ``n`` with no prime factor above 7.""" + m = max(2, n + n % 2) + while True: + rest = m + for p in (2, 3, 5, 7): + while rest % p == 0: + rest //= p + if rest == 1: + return m + m += 2 + + +def mesh_dims_of_cells(box: Array, spacing: float) -> tuple[int, int, int]: + r"""The mesh that resolves every cell at ``spacing``. + + Along each lattice vector :math:`\mathbf{a}_d` the mesh holds the + smallest even FFT-friendly number of points (no prime factor above 7) of + at least :math:`\lceil \max_f |\mathbf{a}_d| / h \rceil`, the maximum + running over the frames, so that no frame is evaluated with a spacing + coarser than :math:`h`. The C++ interface sizes the mesh of a frozen + model by the same rule. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing :math:`h`, in Å. + + Returns + ------- + tuple[int, int, int] + The mesh dimensions. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + # The length is formed term by term, as the C++ interface forms it, so + # that both round the ceiling identically. + length = np.sqrt( + cell[..., 0] * cell[..., 0] + + cell[..., 1] * cell[..., 1] + + cell[..., 2] * cell[..., 2] + ) + longest = np.max(length, axis=0) + n1, n2, n3 = (_even_fft_size(math.ceil(float(a) / spacing)) for a in longest) + return n1, n2, n3 + + +def _mesh_modes(xp: object, n: int, dtype: object, device: object) -> Array: + """The integer mode indices of an axis of ``n`` points, in FFT order.""" + m = xp.arange(n, dtype=dtype, device=device) + return xp.where(m > n // 2, m - n, m) + + +def mesh_grid_of_dims( + box: Array, coord: Array, dims: tuple[int, int, int], mask: Array +) -> MeshGrid: + """Pack cells, node positions and a mesh of the given dimensions. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + coord + Node positions with shape ``(N, 3)``. + dims + The even mesh dimensions. + mask + Periodic (1) vs non-periodic (0) frames with shape ``(nf,)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames. + """ + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + mode_1, mode_2, mode_3 = (_mesh_modes(xp, n, box.dtype, device) for n in dims) + return MeshGrid( + box=box, + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + mask=mask, + mode_1=mode_1, + mode_2=mode_2, + mode_3=mode_3, + ) + + +def build_mesh_grid(box: Array, spacing: float, coord: Array) -> MeshGrid: + """Pack the cells of the periodic frames and their shared mesh. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing, in Å (see :func:`mesh_dims_of_cells`). + coord + Node positions with shape ``(N, 3)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames, every frame + periodic. + """ + xp = array_api_compat.array_namespace(box) + box = _cells(box) + mask = xp.ones((box.shape[0],), dtype=xp.bool, device=array_api_compat.device(box)) + return mesh_grid_of_dims(box, coord, mesh_dims_of_cells(box, spacing), mask) + + +def empty_mesh_grid(nf: int, coord: Array) -> MeshGrid: + """The mesh grid of a batch without periodic frames: every frame masked.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + box = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + mask = xp.zeros((nf,), dtype=xp.bool, device=device) + return mesh_grid_of_dims(box, coord, (2, 2, 2), mask) diff --git a/deepmd/dpmodel/utils/neighbor_graph/topology.py b/deepmd/dpmodel/utils/neighbor_graph/topology.py new file mode 100644 index 0000000000..0fe4127123 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/topology.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Topologies an atomic model declares beside its neighbor graph. + +An interaction that the short-range neighbor graph cannot describe reads a +topology of its own, declared by name (see +:meth:`~deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel.required_topologies`): + +- :data:`~.pair_graph.ALL_PAIRS`, the pairs of the non-periodic frames; +- :data:`~.reciprocal_graph.K_GRID` and :data:`~.reciprocal_graph.PME_MESH`, + the reciprocal space of the periodic frames. + +Every topology is a dataclass of arrays that declares which of its fields +carry geometry (``GEOMETRY``), which an autograd backend makes leaves of the +energy (``LEAVES``), and the dynamic axes of every field (``AXES``), so that +casting, re-leafing and the inputs of a compiled or exported graph lower are +derived from the declaration alone. +""" + +from __future__ import ( + annotations, +) + +import dataclasses +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat + +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + empty_pair_graph, +) +from .reciprocal_graph import ( + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + empty_k_grid, + empty_mesh_grid, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + Topology = PairGraph | KGrid | MeshGrid + +TOPOLOGY_TYPES: dict[str, type] = { + ALL_PAIRS: PairGraph, + K_GRID: KGrid, + PME_MESH: MeshGrid, +} +"""The dataclass of every topology name.""" + + +def _build_reciprocal( + name: str, + box: Array, + coord: Array, + params: dict[str, Any], +) -> KGrid | MeshGrid: + """Build a reciprocal-space topology of periodic frames from its parameters.""" + required = {K_GRID: "k_max", PME_MESH: "spacing"}[name] + if required not in params: + raise ValueError( + f"the {name!r} topology requires its `{required}` build parameter" + ) + check_periodic_cells(box) + if name == PME_MESH: + return build_mesh_grid(box, params["spacing"], coord) + if "bounds" in params: + return build_fixed_k_grid(box, params["bounds"], params["k_max"], coord) + return build_k_grid(box, params["k_max"], coord) + + +def build_topologies( + names: Iterable[str], + coord: Array, + n_node: Array, + box: Array | None = None, + node_mask: Array | None = None, + *, + atype: Array, + pair_excl: PairExcludeMask | None = None, + topology_params: dict[str, dict[str, Any]] | None = None, +) -> dict[str, Topology]: + """Build the topologies an atomic model declares, beside its neighbor graph. + + A batch is periodic when it carries cells. A declaration of the pairs + alone serves non-periodic batches only. Beside a reciprocal-space + topology, the pairs of a periodic batch come back guard-only and the + reciprocal-space topology of a non-periodic batch masked, so that one + model serves both. + + Parameters + ---------- + names + The declared topology names. + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + n_node + Per-frame node counts with shape ``(nf,)``. + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``, or + ``None`` for non-periodic systems. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + pair_excl + The model-level pair exclusion, or ``None``. + topology_params + Build parameters by name: ``{"k_max": ...}`` for the k grid, with + ``"bounds"`` for a fixed integer grid (see + :func:`~.reciprocal_graph.build_fixed_k_grid`), and + ``{"spacing": ...}`` for the mesh. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + Each declared topology under its name; empty when none is declared. + + Raises + ------ + NotImplementedError + If cells are given for a declaration of the pairs alone, or a pair + exclusion for a periodic batch, whose reciprocal-space sum is global. + ValueError + If a name is unknown, or a build parameter is missing. + """ + names = tuple(names) + unknown = [name for name in names if name not in TOPOLOGY_TYPES] + if unknown: + raise ValueError( + f"unknown topology {unknown[0]!r}; expected one of {tuple(TOPOLOGY_TYPES)}" + ) + periodic = any(name in (K_GRID, PME_MESH) for name in names) + if box is None: + nf = n_node.shape[0] + empty = {K_GRID: empty_k_grid, PME_MESH: empty_mesh_grid} + return { + name: ( + build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + if name == ALL_PAIRS + else empty[name](nf, coord) + ) + for name in names + } + if not periodic: + if names: + check_complete_frames(box) + return {} + if pair_excl is not None: + raise NotImplementedError( + "the reciprocal-space sum of a periodic system is global; " + "`pair_exclude_types` cannot exclude pairs from it" + ) + params = topology_params or {} + return { + name: ( + empty_pair_graph(coord) + if name == ALL_PAIRS + else _build_reciprocal(name, box, coord, params.get(name, {})) + ) + for name in names + } + + +def topology_flat_fields(topology: Topology) -> tuple[Array, ...]: + """The fields of a topology in their declaration order.""" + return tuple( + getattr(topology, field.name) for field in dataclasses.fields(topology) + ) + + +def flatten_topologies(topologies: dict[str, Topology]) -> tuple[Array, ...]: + """Lay topologies out as positional inputs of a graph lower. + + A compiled or exported graph lower receives the fields of every topology, + in the order of the dict and of the declaration of each dataclass, after + its other inputs; :func:`unflatten_topologies` restores them. + + Parameters + ---------- + topologies + The topologies by name. + + Returns + ------- + tuple + The fields of every topology, concatenated. + """ + return tuple( + value + for topology in topologies.values() + for value in topology_flat_fields(topology) + ) + + +def unflatten_topologies( + names: Sequence[str], values: Sequence[Array] +) -> dict[str, Topology]: + """Restore topologies laid out by :func:`flatten_topologies`. + + Parameters + ---------- + names + The topology names, in the order they were laid out. + values + The fields of every topology, concatenated. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + The topologies by name. + + Raises + ------ + ValueError + If a name is unknown, or the number of values does not match the + layout of the names. + """ + topologies = {} + offset = 0 + for name in names: + if name not in TOPOLOGY_TYPES: + raise ValueError(f"unknown topology {name!r}") + cls = TOPOLOGY_TYPES[name] + arity = len(dataclasses.fields(cls)) + if offset + arity > len(values): + break + topologies[name] = cls(*values[offset : offset + arity]) + offset += arity + if len(topologies) != len(names) or offset != len(values): + raise ValueError( + f"{len(values)} flattened values given, which do not lay out the " + f"topologies {tuple(names)}" + ) + return topologies + + +def cast_topology_geometry(topology: Topology, dtype: Any) -> Topology: + """A copy of a topology with its geometry fields cast to ``dtype``.""" + xp = array_api_compat.array_namespace(getattr(topology, topology.GEOMETRY[0])) + return dataclasses.replace( + topology, + **{ + name: xp.astype(getattr(topology, name), dtype) + for name in topology.GEOMETRY + }, + ) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 23dc7bf818..c4f0771b01 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -23,8 +23,10 @@ OutputVariableDef, ) from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, build_topologies, check_complete_frames, + check_periodic_cells, flatten_topologies, unflatten_topologies, ) @@ -727,6 +729,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: } if topology_names: self.metadata["topologies"] = list(topology_names) + self.metadata["topology_params"] = model.atomic_model.topology_params() if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, @@ -2400,8 +2403,9 @@ def _eval_model_graph( ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, fparam, aparam, charge_spin)`` to the graph forward, followed by the - ``(index, vec, mask)`` of every topology the ``topologies`` metadata - names. Exported AOTI artifacts use dynamic edge and pair axes, so no + fields of every topology the ``topologies`` metadata names, built with + its ``topology_params``. Exported AOTI artifacts use dynamic edge, + pair and mesh axes, so no ``edge_capacity`` padding is needed. The ``graph_edge_dtype`` metadata selects float32 geometry for compressed DPA1 and float64 for generic graph descriptors. The forward returns local public keys directly, so results are reshaped without @@ -2516,16 +2520,30 @@ def _eval_model_graph( compact.source_row_ptr, compact.source_order, ) - if self.metadata.get("topologies"): + topology_names = self.metadata.get("topologies", ()) + if topology_names: # The compact lower of a model with a long-range module reads - # the positions of complete non-periodic frames, and the zero - # communicator of a frame held whole. - check_complete_frames(box_input) + # the positions and the cells of the frames, and the zero + # communicator of frames held whole; a module that serves no + # periodic system takes no cell, and the zero cell that marks + # a non-periodic frame is never one a caller passes. + if PME_MESH not in topology_names: + check_complete_frames(box_input) + elif box_input is not None: + check_periodic_cells(box_input) + box_t = ( + torch.zeros(nframes, 9, dtype=torch.float64, device="cpu") + if box_input is None + else torch.as_tensor( + np.asarray(box_input, dtype=np.float64).reshape(nframes, 9) + ) + ) model_inputs += ( torch.as_tensor( coord_input.reshape(-1, 3), dtype=torch.float64, device=DEVICE ), torch.zeros(1, dtype=torch.int64, device="cpu"), + box_t, ) else: fparam_t, aparam_t = self._prepare_optional_lower_inputs( @@ -2570,37 +2588,18 @@ def _eval_model_graph( if box_input is not None else None ) - if self._dpmodel is not None: - # The live model builds its own topologies, so the k-grid - # cutoff follows the kernel parameters. - topologies = self._dpmodel.atomic_model.build_topologies( - coord_flat_t, - n_node_t, - box_t, - atype=atype_t, - node_mask=atype_t >= 0, - pair_excl=self._model_pair_excl(), - ) - else: - from deepmd.dpmodel.utils.neighbor_graph import ( - K_GRID, - PME_MESH, - ) - - if K_GRID in topology_names or PME_MESH in topology_names: - raise NotImplementedError( - "a periodic long-range model cannot be evaluated from " - "a .pt2/.pte archive yet" - ) - topologies = build_topologies( - topology_names, - coord_flat_t, - n_node_t, - box_t, - node_mask=atype_t >= 0, - atype=atype_t, - pair_excl=self._model_pair_excl(), - ) + # The parameters are those of the archive, or those of the + # live model of a checkpoint. + topologies = build_topologies( + topology_names, + coord_flat_t, + n_node_t, + box_t, + node_mask=atype_t >= 0, + atype=atype_t, + pair_excl=self._model_pair_excl(), + topology_params=self.metadata.get("topology_params"), + ) model_inputs += flatten_topologies(topologies) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 1aaf2e1010..48eef91e74 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -59,8 +59,8 @@ def canonical_model_eligible(model: Any) -> bool: if getattr(atomic_model, "atom_excl", None) is not None: return False # A latent-charge long-range module joins the pipeline through the node - # positions, which its direct-sum operator reads in place of a pair list; - # that pipeline evaluates no pair potential. + # positions and the cells, which its operator reads in place of a pair + # list or a mesh; that pipeline evaluates no pair potential. if parts[0].has_long_range(): from deepmd.pt_expt.long_range.latent_charge import ( LatentChargeLongRange, @@ -69,7 +69,7 @@ def canonical_model_eligible(model: Any) -> bool: long_range = parts[0].long_range if not ( isinstance(long_range, LatentChargeLongRange) - and long_range.direct_sum_eligible() + and long_range.operator_eligible() and parts[1] is None ): return False diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py index 69e5edb332..3bb786ff58 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -26,6 +26,7 @@ def dpa4c_canonical_long_range_energy_force( atype: torch.Tensor, coord: torch.Tensor, communicator: torch.Tensor | None, + box: torch.Tensor | None, ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, @@ -46,20 +47,23 @@ def dpa4c_canonical_long_range_energy_force( fitting Eligible pt_expt energy fitting network. long_range - Long-range module whose ``direct_sum_eligible()`` holds. + Long-range module whose ``operator_eligible()`` holds. graph Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. atype Flat node atom types with shape ``(N,)``. coord - Node positions with shape ``(N, 3)`` in Å; every frame is a complete - non-periodic system. + Node positions with shape ``(N, 3)`` in Å. communicator The address of the MPI_Comm handle of a frame decomposed into domains, one per rank, as an int64 CPU tensor of one element, or zero or ``None`` for a single domain. Across domains the owned nodes of every rank make up the frame. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the CPU, + a zero row marking a complete non-periodic frame, or ``None`` when no + frame is periodic. ownership Boolean mask selecting energy-contributing nodes with shape ``(N,)``. atom_bias @@ -123,7 +127,7 @@ def dpa4c_canonical_long_range_energy_force( long_range_atom_virial, long_range_virial, ) = long_range.energy_force_from_coord( - node_descriptor, atype, ownership, graph.n_node, coord, communicator + node_descriptor, atype, ownership, graph.n_node, coord, communicator, box ) atom_energy, descriptor_gradient = fitting_energy_and_gradient( fitting, node_descriptor, atype, ownership, atom_bias @@ -146,7 +150,8 @@ def dpa4c_canonical_long_range_energy_force( do_atomic_virial, ) # Nodes outside the ownership carry no charge and hence no long-range - # energy. + # energy; the halo nodes of a periodic frame are images, which the + # reciprocal-space sum holds itself. atom_energy = atom_energy + long_range_energy energy = frame_scalar_sum(atom_energy, graph.n_node) virial = virial + long_range_virial.reshape(-1, 3, 3) diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py index 994b78134f..89db40ad93 100644 --- a/deepmd/pt_expt/kernels/long_range.py +++ b/deepmd/pt_expt/kernels/long_range.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Bindings for the direct-sum operator of the latent-charge long range. +"""Bindings for the operator of the latent-charge long range. -The operator ``deepmd::latent_charge_direct_sum`` evaluates, for the nodes of +The operator ``deepmd::latent_charge_sum`` evaluates, for the nodes of complete frames, the potential that the charges of every other node of the -frame create at each node and the force between the charges, without listing -the pairs: its memory is linear in the node count. The CUDA kernel is -``source/op/pt/long_range/direct_sum.cu``; the CPU kernel registered here is -the reference it is tested against. +frame create at each node and the force between the charges. A non-periodic +frame is summed directly over its pairs without listing them, so that the +memory is linear in the node count; a periodic frame is summed over the images +of its cell by the smooth particle-mesh Ewald method, on a mesh sized by the +cell. The CUDA kernels are in ``source/op/pt/long_range``; the CPU kernel +registered here is the reference they are tested against. The operator serves inference. It applies the total-charge constraint of the module, and returns the atomic energy, its derivative with respect to the @@ -44,22 +46,25 @@ __all__ = [ "MAX_CHANNELS", "MAX_GAUSSIANS", + "MAX_ORDER", "kernel_eligible", "kernel_operator_arguments", - "latent_charge_direct_sum", + "latent_charge_sum", "op_available", ] -#: Charge channels and Gaussians the CUDA kernel is instantiated for. +#: Charge channels and Gaussians the CUDA kernels are instantiated for. MAX_CHANNELS = 8 MAX_GAUSSIANS = 32 +#: The largest order of the B-splines of the mesh. +MAX_ORDER = 12 _LES = 0 _SOG = 1 def op_available() -> bool: - """Whether the backend device carries the direct-sum operator. + """Whether the backend device carries the operator. The CPU kernel is the reference of this module, which is registered here so that the answer does not depend on earlier calls. @@ -68,7 +73,7 @@ def op_available() -> bool: import deepmd.pt.cxx_op # noqa: F401 ensure_registered() - return operator_available("latent_charge_direct_sum") + return operator_available("latent_charge_sum") def kernel_eligible(kernel: Any) -> bool: @@ -89,7 +94,7 @@ def kernel_operator_arguments(kernel: Any) -> tuple[int, torch.Tensor]: Returns ------- tuple[int, torch.Tensor] - The kernel code and its float32 coefficients: ``alpha`` for the erf + The kernel code and its float64 coefficients: ``alpha`` for the erf kernel; the amplitudes ``a_k`` and then the rates ``c_k = log2(e) / (2 sigma_k^2)`` of ``a_k 2^(-c_k r^2)`` for the sum of Gaussians. @@ -100,14 +105,14 @@ def kernel_operator_arguments(kernel: Any) -> tuple[int, torch.Tensor]: """ if not kernel_eligible(kernel): raise ValueError( - f"the direct-sum operator does not evaluate {type(kernel).__name__}" + f"the long-range operator does not evaluate {type(kernel).__name__}" ) if isinstance(kernel, LESKernel): - coefficient = torch.exp(kernel.log_alpha).reshape(1) - return _LES, coefficient.detach().to(torch.float32).contiguous() - rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width) - coefficient = torch.cat((kernel.amplitude, rate)) - return _SOG, coefficient.detach().to(torch.float32).contiguous() + coefficient = torch.exp(kernel.log_alpha.to(torch.float64)).reshape(1) + return _LES, coefficient.detach().contiguous() + rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width.to(torch.float64)) + coefficient = torch.cat((kernel.amplitude.to(torch.float64), rate)) + return _SOG, coefficient.detach().contiguous() def _kernel_and_slope( @@ -130,41 +135,48 @@ def _kernel_and_slope( return weight @ amplitude, weight @ (-2.0 * math.log(2.0) * rate * amplitude) -def _reference( - coord: torch.Tensor, +def _charges_of_frames( charge: torch.Tensor, node_mask: torch.Tensor, n_node_per_frame: torch.Tensor, total_charge: torch.Tensor, - coefficient: torch.Tensor, - communicator: torch.Tensor, - kernel: int, channel: int, - prefactor: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Evaluate the operator over the explicit pairs of every frame, in float64.""" - if int(communicator.reshape(-1)[0]) != 0: - raise NotImplementedError( - "the CPU reference of latent_charge_direct_sum evaluates a single " - "domain; a communicator of several ranks requires the CUDA kernel" - ) - n = coord.shape[0] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The charges of the interactions in float64, shifted in the constrained channel. + + Returns the charges with shape ``(N, C)``, the weights of the charged + nodes and the frame of every node with shape ``(N,)``, and the number of + charged nodes of every frame, at least one, with shape ``(F,)``. + """ + n = charge.shape[0] frame_count = n_node_per_frame.shape[0] # Padding nodes belong to no frame and carry no charge. frame = frame_id_from_n_node(n_node_per_frame, n_total=n) - real = torch.arange(n, dtype=torch.int64, device=coord.device) < torch.sum( + real = torch.arange(n, dtype=torch.int64, device=charge.device) < torch.sum( n_node_per_frame ) weight = (node_mask & real).to(torch.float64) q = charge.to(torch.float64) * weight[:, None] + count = weight.new_zeros(frame_count).index_add(0, frame, weight) + count = torch.clamp(count, min=1.0) if channel >= 0: - columns = torch.arange(q.shape[1], dtype=torch.int64, device=coord.device) - selector = (columns == channel).to(torch.float64) - count = coord.new_zeros(frame_count).index_add(0, frame, weight) - count = torch.clamp(count, min=1.0) - total = coord.new_zeros(frame_count).index_add(0, frame, q[:, channel]) + total = weight.new_zeros(frame_count).index_add(0, frame, q[:, channel]) shift = (total - total_charge.to(torch.float64)) / count - q = q - (shift[frame] * weight)[:, None] * selector + q = q.clone() + q[:, channel] = q[:, channel] - shift[frame] * weight + return q, weight, frame, count + + +def _direct_reference( + coord: torch.Tensor, + q: torch.Tensor, + n_node_per_frame: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial over the explicit pairs of every frame.""" + n = coord.shape[0] pairs = unordered_pair_index(n_node_per_frame) src, dst = pairs[0], pairs[1] d = coord[src] - coord[dst] @@ -175,20 +187,149 @@ def _reference( strength = torch.sum(q[src] * q[dst], dim=-1) * slope pair_force = -strength[:, None] * d force = torch.zeros_like(coord).index_add(0, src, pair_force) - force = force.index_add(0, dst, -pair_force) + force = prefactor * force.index_add(0, dst, -pair_force) half = (-0.5 * strength)[:, None, None] * d[:, :, None] * d[:, None, :] atom_virial = coord.new_zeros(n, 3, 3).index_add(0, src, half) atom_virial = prefactor * atom_virial.index_add(0, dst, half).reshape(n, 9) + return potential, force, atom_virial + + +def _interaction_kernel(coefficient: torch.Tensor, kernel: int) -> Any: + """The dpmodel interaction kernel of the coefficients of the operator.""" + from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP + from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP + + values = coefficient.detach().to("cpu", torch.float64).numpy() + if kernel == _LES: + return LESKernelDP(alpha=float(values[0]), precision="float64") + amplitude, rate = values.reshape(2, -1) + width = (0.5 / (math.log(2.0) * rate)) ** 0.5 + return SOGKernelDP( + amplitude=amplitude.tolist(), width=width.tolist(), precision="float64" + ) + + +def _mesh_reference( + coord: torch.Tensor, + q: torch.Tensor, + weight: torch.Tensor, + n_node_per_frame: torch.Tensor, + box: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial on the mesh of every periodic frame. + + Every frame is evaluated on its own mesh by the evaluation of the training + path, and differentiated with respect to its charges, its positions and + its strain. The atomic virial takes ``F (x) r`` at every node and shares + the rest of the virial of the frame evenly over its charged nodes. + """ + from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, + ) + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + + interaction = _interaction_kernel(coefficient, kernel) + potential = torch.zeros_like(q) + force = torch.zeros_like(coord) + atom_virial = coord.new_zeros(coord.shape[0], 3, 3) + begin = 0 + for f, length in enumerate(n_node_per_frame.tolist()): + end = begin + length + frame_q = q[begin:end].detach().requires_grad_(True) + frame_coord = coord[begin:end].detach().requires_grad_(True) + with torch.enable_grad(): + mesh = build_mesh_grid(box[f].reshape(1, 3, 3), spacing, frame_coord) + mesh.strain.requires_grad_(True) + frame_id = torch.zeros(length, dtype=torch.int64, device=coord.device) + energy = mesh_energy(frame_q, mesh, frame_id, interaction, order).sum() + energy = energy * (prefactor / COULOMB_CONSTANT) + d_q, d_coord, d_strain = torch.autograd.grad( + energy, (frame_q, frame_coord, mesh.strain) + ) + frame_force = -d_coord + moment = frame_force[:, :, None] * coord[begin:end, None, :] + frame_weight = weight[begin:end] + count = torch.clamp(torch.sum(frame_weight), min=1.0) + rest = (-d_strain[0] - torch.sum(moment, dim=0)) / count + potential[begin:end] = d_q + force[begin:end] = frame_force + atom_virial[begin:end] = moment + rest * frame_weight[:, None, None] + begin = end + return potential, force, atom_virial.reshape(-1, 9) + + +def _reference( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the operator in float64, over the explicit pairs or the meshes.""" + if int(communicator.reshape(-1)[0]) != 0: + raise NotImplementedError( + "the CPU reference of latent_charge_sum evaluates a single domain; " + "a communicator of several ranks requires the CUDA kernel" + ) + frame_count = n_node_per_frame.shape[0] + q, weight, frame, count = _charges_of_frames( + charge, node_mask, n_node_per_frame, total_charge, channel + ) + has_cell = ( + torch.zeros(frame_count, dtype=torch.bool, device="cpu") + if box is None + else torch.any(box != 0.0, dim=-1) + ) + if bool(torch.any(has_cell)) != bool(torch.all(has_cell)): + raise ValueError( + "latent_charge_sum: every frame of a call is periodic, or none is" + ) + if bool(torch.any(has_cell)): + potential, force, atom_virial = _mesh_reference( + coord, + q, + weight, + n_node_per_frame, + box.to(coord.device), + coefficient, + kernel, + prefactor, + order, + spacing, + ) + else: + potential, force, atom_virial = _direct_reference( + coord, q, n_node_per_frame, coefficient, kernel, prefactor + ) virial = coord.new_zeros(frame_count, 9).index_add(0, frame, atom_virial) atom_energy = 0.5 * torch.sum(q * potential, dim=-1, keepdim=True) gradient = potential * weight[:, None] if channel >= 0: mean = coord.new_zeros(frame_count).index_add(0, frame, gradient[:, channel]) - gradient = gradient - ((mean / count)[frame] * weight)[:, None] * selector + gradient = gradient.clone() + gradient[:, channel] = gradient[:, channel] - (mean / count)[frame] * weight return ( atom_energy, gradient.to(charge.dtype), - (prefactor * force).to(charge.dtype), + force.to(charge.dtype), atom_virial.to(charge.dtype), virial.to(charge.dtype), ) @@ -200,11 +341,14 @@ def _fake( node_mask: torch.Tensor, n_node_per_frame: torch.Tensor, total_charge: torch.Tensor, + box: torch.Tensor | None, coefficient: torch.Tensor, communicator: torch.Tensor, kernel: int, channel: int, prefactor: float, + order: int, + spacing: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: n = coord.shape[0] return ( @@ -218,9 +362,9 @@ def _fake( @cache def _register() -> None: - torch.library.register_fake("deepmd::latent_charge_direct_sum")(_fake) + torch.library.register_fake("deepmd::latent_charge_sum")(_fake) library = torch.library.Library("deepmd", "IMPL") - library.impl("latent_charge_direct_sum", _reference, "CPU") + library.impl("latent_charge_sum", _reference, "CPU") # The registration lives as long as the library object. _register.library = library @@ -228,13 +372,13 @@ def _register() -> None: def ensure_registered() -> None: """Register the fake and CPU implementations when the schema is loaded.""" if isinstance( - getattr(torch.ops.deepmd, "latent_charge_direct_sum", None), + getattr(torch.ops.deepmd, "latent_charge_sum", None), torch._ops.OpOverloadPacket, ): _register() -def latent_charge_direct_sum( +def latent_charge_sum( coord: torch.Tensor, charge: torch.Tensor, node_mask: torch.Tensor, @@ -244,6 +388,9 @@ def latent_charge_direct_sum( channel: int, prefactor: float, communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Energy, charge gradient, force and virial of the charges of complete frames. @@ -278,20 +425,32 @@ def latent_charge_direct_sum( frame, whose charged nodes are those of every rank; the outputs are those of the charged nodes of this rank, the other nodes carry no charge, and ``virial`` is the share of this rank. + box + The cells of periodic frames with shape ``(F, 9)``, float64 on the + CPU, the rows of every row-major 3x3 cell being its lattice vectors, + or ``None`` for non-periodic frames; a zero row marks a non-periodic + frame. A periodic frame is evaluated on a mesh of spacing at most + ``spacing``, with B-splines of order ``order``. + order + The order of the B-splines of the mesh, at most :data:`MAX_ORDER`. + spacing + The largest mesh spacing, in Å. Returns ------- atom_energy - Per-node energy with shape ``(N, 1)`` in eV, float64, each pair - energy split half to each of its nodes. + Per-node energy with shape ``(N, 1)`` in eV, float64, half of the + interaction of the charges of the node with all others. charge_gradient Derivative of the energy with respect to ``charge`` with shape ``(N, C)``, in eV per unit charge, through the constraint. force Force at fixed charges with shape ``(N, 3)``, in eV/Å. atom_virial - Per-node virial at fixed charges with shape ``(N, 9)`` in eV, each - pair term split half to each of its nodes. + Per-node virial at fixed charges with shape ``(N, 9)`` in eV: half of + every pair term to each of its nodes in a non-periodic frame; in a + periodic frame ``F (x) r`` with the rest of the virial of the cell + shared evenly over the charged nodes. virial Per-frame virial at fixed charges with shape ``(F, 9)`` in eV. """ @@ -299,15 +458,18 @@ def latent_charge_direct_sum( code, coefficient = kernel_operator_arguments(kernel) if communicator is None: communicator = torch.zeros(1, dtype=torch.int64, device="cpu") - return torch.ops.deepmd.latent_charge_direct_sum( + return torch.ops.deepmd.latent_charge_sum( coord, charge, node_mask, n_node_per_frame, total_charge, + box, coefficient.to(coord.device), communicator, code, int(channel), float(prefactor), + int(order), + float(spacing), ) diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py index 2cf9bf009f..b79dffef4c 100644 --- a/deepmd/pt_expt/long_range/latent_charge.py +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -26,7 +26,7 @@ @BaseLongRange.register("latent_charge") @torch_module class LatentChargeLongRange(LatentChargeLongRangeDP): - """Energy of interacting latent charges in complete non-periodic systems (pt_expt).""" + """Energy of interacting latent charges, periodic or not (pt_expt).""" # The dpmodel class builds the charge network and the kernel, which the # assignment converts to these pt_expt modules; importing their classes @@ -51,16 +51,17 @@ def _pair_kernel( } return self.kernel.evaluate(r_sq, coefficients) - def direct_sum_eligible(self) -> bool: + def operator_eligible(self) -> bool: """Whether :meth:`energy_force_from_coord` evaluates this module. - The direct-sum operator evaluates the kernels and channel counts it is - instantiated for. The compact lower that calls it reads no frame + The operator evaluates the kernels, channel counts and B-spline orders + it is instantiated for. The compact lower that calls it reads no frame input: the charges depend on the descriptor alone, and a constrained channel sums to zero or to the total charge of a folded charge state. """ from deepmd.pt_expt.kernels.long_range import ( MAX_CHANNELS, + MAX_ORDER, kernel_eligible, op_available, ) @@ -69,6 +70,7 @@ def direct_sum_eligible(self) -> bool: op_available() and kernel_eligible(self.kernel) and self.charge_net.dim_out <= MAX_CHANNELS + and self.pme_order <= MAX_ORDER and self.charge_net.numb_fparam == 0 and self.charge_net.numb_aparam == 0 and not self.uses_charge_spin() @@ -82,16 +84,18 @@ def energy_force_from_coord( n_node: torch.Tensor, coord: torch.Tensor, communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Evaluate the module for inference from the positions of the nodes. - The direct-sum operator enumerates the pairs of every frame itself, - so the memory is linear in the node count. It applies the constraint - and returns the force at fixed charges and the derivative of the - energy with respect to the charges the network predicts; pulled back - through the charge network, that derivative is the derivative of the - energy with respect to the descriptor, through which the caller adds - the response of the charges to the positions. + The operator sums a non-periodic frame directly over its pairs, which + it enumerates itself, so that the memory is linear in the node count, + and a periodic frame on the mesh of its cell. It applies the + constraint and returns the force at fixed charges and the derivative + of the energy with respect to the charges the network predicts; + pulled back through the charge network, that derivative is the + derivative of the energy with respect to the descriptor, through which + the caller adds the response of the charges to the positions. Parameters ---------- @@ -109,8 +113,13 @@ def energy_force_from_coord( communicator The MPI communicator of a frame decomposed into domains, one per rank, whose charged nodes are those of every rank, in the form of - :func:`~deepmd.pt_expt.kernels.long_range.latent_charge_direct_sum`; + :func:`~deepmd.pt_expt.kernels.long_range.latent_charge_sum`; ``None`` for a single domain. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the + CPU, a zero row marking a non-periodic frame, or ``None`` when no + frame is periodic. A module that does not serve periodic systems + evaluates every frame as non-periodic. Returns ------- @@ -129,7 +138,7 @@ def energy_force_from_coord( float32; across domains, the share of the nodes of this rank. """ from deepmd.pt_expt.kernels.long_range import ( - latent_charge_direct_sum, + latent_charge_sum, ) leaf = descriptor.detach().requires_grad_(True) @@ -140,18 +149,21 @@ def energy_force_from_coord( if self.charge_constraint is not None: channel = self.charge_constraint["channel"] total_charge = self._total_charge(n_node.shape[0], total_charge, None, None) - atom_energy, charge_gradient, force, atom_virial, virial = ( - latent_charge_direct_sum( - coord.to(torch.float64).contiguous(), - charge.detach().to(torch.float32).contiguous(), - node_mask.contiguous(), - n_node, - total_charge.contiguous(), - self.kernel, - channel, - COULOMB_CONSTANT, - communicator, - ) + atom_energy, charge_gradient, force, atom_virial, virial = latent_charge_sum( + coord.to(torch.float64).contiguous(), + charge.detach().to(torch.float32).contiguous(), + node_mask.contiguous(), + n_node, + total_charge.contiguous(), + self.kernel, + channel, + COULOMB_CONSTANT, + communicator, + # A module that does not serve periodic systems sums every frame + # directly. + box if self.periodic else None, + self.pme_order, + self.mesh_spacing() if self.periodic else 1.0, ) (descriptor_gradient,) = torch.autograd.grad( charge, leaf, charge_gradient.to(charge.dtype) diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 6e0f837de9..f149053489 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -4,9 +4,10 @@ torch-only. The pure-array scatters (edge_force_virial, pair_force_virial) are shared with dpmodel; this module supplies the single backward pass that produces the gradient of the energy with respect to every geometry leaf: the -edge vectors of the neighbor graph and the pair vectors of the topologies. +edge vectors of the neighbor graph and the geometry leaves of the topologies. """ +import dataclasses from collections.abc import ( Sequence, ) @@ -31,6 +32,7 @@ PairGraph, edge_force_virial, frame_id_from_n_node, + node_validity_mask, pair_force_virial, segment_sum, ) @@ -54,6 +56,91 @@ ) +def releaf_topologies( + topologies: dict[str, PairGraph | KGrid | MeshGrid], +) -> dict[str, PairGraph | KGrid | MeshGrid]: + """Make the geometry fields of every topology autograd leaves. + + The leaves are made from the inputs of a traced graph lower, so that the + backward roots at its inputs: a leaf created inside a ``make_fx`` trace + has a NaN gradient under inductor. + """ + return { + name: dataclasses.replace( + topology, + **{ + leaf: getattr(topology, leaf).detach().requires_grad_(True) + for leaf in topology.LEAVES + }, + ) + for name, topology in topologies.items() + } + + +def _cell_force_virial( + g_coord: torch.Tensor, + g_strain: torch.Tensor, + coord: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor | None, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """Force and virial of a reciprocal-space energy from its geometry leaves. + + The gradient with respect to the node positions is the negative force, + and that with respect to the zero strain of every frame the negative + virial. The atomic virial takes ``F_i (x) r_i`` at every node, the force + index first as for the edge virial, and shares the rest of the virial of + the frame, the part of the cell, evenly over its owned nodes, so that the + nodes of a frame sum to its virial; the halo nodes, images that carry no + charge, take no share. + + Parameters + ---------- + g_coord + (N, 3) gradient of the energy with respect to the positions. + g_strain + (nf, 3, 3) gradient of the energy with respect to the strain. + coord + (N, 3) node positions. + n_node + (nf,) per-frame node counts. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned. + do_atomic_virial + Whether to return the per-node virial. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, or ``None``. + virial + (nf, 3, 3) per-frame virial. + """ + force = -g_coord + virial = -g_strain + if not do_atomic_virial: + return force, None, virial + n_total = force.shape[0] + frame_id = frame_id_from_n_node(n_node, n_total) + atom_virial = force[:, :, None] * coord.to(force.dtype)[:, None, :] + remainder = virial - segment_sum(atom_virial, frame_id, n_node.shape[0]) + if n_local is None: + owned = node_validity_mask(n_node, n_total) + count = n_node + else: + owned = node_ownership_mask(n_node, n_local, n_total) + count = n_local + share = owned.to(force.dtype) / torch.clamp( + count.to(force.dtype)[frame_id], min=1.0 + ) + atom_virial = atom_virial + remainder[frame_id] * share[:, None, None] + return force, atom_virial, virial + + def edge_energy_deriv( energy: torch.Tensor, edge_vec: torch.Tensor, @@ -70,9 +157,8 @@ def edge_energy_deriv( do_atomic_virial: bool = False, create_graph: bool = False, force_precision: torch.dtype | None = None, - pair_graphs: Sequence[PairGraph] = (), - coord_leaf: torch.Tensor | None = None, - strain_leaf: torch.Tensor | None = None, + topologies: Sequence[PairGraph | KGrid | MeshGrid] = (), + n_local: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -92,16 +178,14 @@ def edge_energy_deriv( reduction. A retained graph (training / double backward) keeps the fp64 leaf dtype so training numerics are unchanged. - The same backward differentiates the energy with respect to the vectors - of every pair topology in ``pair_graphs``. Their gradients are pulled back - by the shared pair_force_virial scatter and summed in: the energy reaches - the coordinates through all its geometry leaves, so its complete force and - virial are the sum of their pullbacks. - - A k-space long-range term reaches the geometry through ``coord_leaf`` (the - node positions, whose gradient is its force directly) and ``strain_leaf`` - (a zero (nf, 3, 3) strain the module applies to coordinates, wave numbers - and volumes, whose gradient is the negative k-space virial). + The same backward differentiates the energy with respect to the geometry + leaves of every topology in ``topologies``, and their pullbacks are summed + in: the energy reaches the coordinates through all its geometry leaves, so + its complete force and virial are the sum of their pullbacks. The pair + vectors of a pair topology are pulled back by the shared + pair_force_virial scatter; a reciprocal-space topology carries the node + positions, whose gradient is its force, and a zero strain, whose gradient + is its negative virial (see :func:`_cell_force_virial`). Parameters ---------- @@ -133,13 +217,12 @@ def edge_energy_deriv( whether to materialize the per-atom virial (else ``None`` is returned). create_graph whether the backward retains a graph (training, for second-order grad). - pair_graphs - Pair topologies whose ``vec`` are further autograd leaves of - ``energy``, on the same node axis. - coord_leaf - (N, 3) node positions of a k-space long-range term, or ``None``. - strain_leaf - (nf, 3, 3) zero strain of a k-space long-range term, or ``None``. + topologies + Topologies on the same node axis whose ``LEAVES`` are further + autograd leaves of ``energy``. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned; the owned nodes share the virial of a cell. Returns ------- @@ -150,16 +233,15 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - extra_leaves = tuple(leaf for leaf in (coord_leaf, strain_leaf) if leaf is not None) - g_e, *g_rest = torch.autograd.grad( + leaves = [ + getattr(topology, name) for topology in topologies for name in topology.LEAVES + ] + g_e, *g_leaves = torch.autograd.grad( energy.sum() if energy.dim() else energy, - (edge_vec, *(pairs.vec for pairs in pair_graphs), *extra_leaves), + (edge_vec, *leaves), create_graph=create_graph, retain_graph=True, ) - g_pairs = g_rest[: len(pair_graphs)] - g_coord = g_rest[len(pair_graphs)] if coord_leaf is not None else None - g_strain = g_rest[-1] if strain_leaf is not None else None if ( force_precision is not None and not create_graph @@ -232,44 +314,29 @@ def edge_energy_deriv( force, atom_virial, virial = edge_force_virial( g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity ) - for g_p, pairs in zip(g_pairs, pair_graphs, strict=True): - pair_force, pair_atom_virial, pair_virial = pair_force_virial( - g_p, - pairs.vec, - pairs.index, - pairs.mask, - n_node, - node_capacity=force.shape[0], - ) - force = force + pair_force.to(force.dtype) - virial = virial + pair_virial.to(virial.dtype) - if do_atomic_virial: - atom_virial = atom_virial + pair_atom_virial.to(atom_virial.dtype) - if coord_leaf is not None or strain_leaf is not None: - # The k-space term: the coordinate gradient is its force, the strain - # gradient its negative virial. The atomic virial convention keeps the - # per-frame sum exact: the r x F part per node, the remainder shared - # evenly over the frame. - frame_id = frame_id_from_n_node(n_node, force.shape[0]) - kspace_vir_frame = None - if strain_leaf is not None: - kspace_vir_frame = -g_strain # (nf, 3, 3) - virial = virial + kspace_vir_frame.to(virial.dtype) - if coord_leaf is not None: - f_coord = -g_coord - force = force + f_coord.to(force.dtype) + offset = 0 + for topology in topologies: + grads = g_leaves[offset : offset + len(topology.LEAVES)] + offset += len(topology.LEAVES) + if isinstance(topology, PairGraph): + (g_p,) = grads + leaf_force, leaf_atom_virial, leaf_virial = pair_force_virial( + g_p, + topology.vec, + topology.index, + topology.mask, + n_node, + node_capacity=force.shape[0], + ) + else: + g_coord, g_strain = grads + leaf_force, leaf_atom_virial, leaf_virial = _cell_force_virial( + g_coord, g_strain, topology.coord, n_node, n_local, do_atomic_virial + ) + force = force + leaf_force.to(force.dtype) + virial = virial + leaf_virial.to(virial.dtype) if do_atomic_virial: - av = torch.zeros_like(atom_virial) - if coord_leaf is not None: - # w = -g (x) r, force index first as for the edge virial - av = av + torch.einsum("ni,nj->nij", f_coord, coord_leaf) - if kspace_vir_frame is not None: - remainder = kspace_vir_frame - segment_sum( - av, frame_id, n_node.shape[0] - ) - count = n_node.to(av.dtype)[frame_id].clamp(min=1).reshape(-1, 1, 1) - av = av + remainder[frame_id] / count - atom_virial = atom_virial + av.to(atom_virial.dtype) + atom_virial = atom_virial + leaf_atom_virial.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -284,7 +351,7 @@ def fit_output_to_model_output_graph( n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, spin_leaf: torch.Tensor | None = None, - topologies: dict[str, PairGraph] | None = None, + topologies: dict[str, PairGraph | KGrid | MeshGrid] | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -357,9 +424,9 @@ def fit_output_to_model_output_graph( are separate so :func:`edge_energy_deriv`'s signature stays untouched. ``None`` (default): unchanged, no mag output. topologies - Pair topologies read by the atomic model, by name. Their ``vec`` are - autograd leaves of ``fit_ret`` next to ``graph.edge_vec``; force and - virial include their pullbacks (see :func:`edge_energy_deriv`). + Topologies read by the atomic model, by name. Their geometry leaves + are autograd leaves of ``fit_ret`` next to ``graph.edge_vec``; force + and virial include their pullbacks (see :func:`edge_energy_deriv`). Returns ------- @@ -451,14 +518,6 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: size *= ii # split the reduced output into ``size`` per-frame scalar components. svv = model_ret[kk_redu].reshape(nf, size) - # The k-space geometry leaves of a periodic long-range module: the - # node positions (its force) and the zero strain (its virial). - coord_leaf = strain_leaf = None - if topologies: - for topology in topologies.values(): - if isinstance(topology, (KGrid, MeshGrid)): - coord_leaf = topology.coord - strain_leaf = topology.strain ff_list: list[torch.Tensor] = [] av_list: list[torch.Tensor] = [] vir_list: list[torch.Tensor] = [] @@ -479,15 +538,8 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, force_precision=force_precision if not create_graph else None, - pair_graphs=( - () - if topologies is None - else tuple( - t for t in topologies.values() if isinstance(t, PairGraph) - ) - ), - coord_leaf=coord_leaf, - strain_leaf=strain_leaf, + topologies=() if topologies is None else tuple(topologies.values()), + n_local=n_local, ) # force (N, 3) -> (N, 1, 3) [flat; caller unravels at I/O boundary] ff_list.append(force.reshape(N, 1, 3)) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 0523084381..5ee8794186 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -19,7 +19,7 @@ make_hessian_model, ) from deepmd.dpmodel.utils.neighbor_graph import ( - PairGraph, + flatten_topologies, ) from deepmd.dpmodel.utils.neighbor_list import ( NeighborList, @@ -144,6 +144,7 @@ def forward_lower_canonical_graph( spin: torch.Tensor | None = None, coord: torch.Tensor | None = None, communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -173,13 +174,19 @@ def forward_lower_canonical_graph( model, or ``None``. Ghost rows carry their owner's moment. coord Node positions with shape ``(N, 3)`` in Å, float64, for a model - with a long-range module, whose frames are complete non-periodic - systems; ``None`` otherwise. + with a long-range module; ``None`` otherwise. communicator For a model with a long-range module, the address of the MPI_Comm handle of a frame decomposed into domains, one per rank, as an int64 CPU tensor of one element, whose owned nodes are then those of one domain; zero or ``None`` for a complete frame. + box + For a model with a long-range module, the cells of the frames + with shape ``(nf, 9)``, float64 on the CPU, the rows of every + row-major 3x3 cell its lattice vectors; a zero row marks a + complete non-periodic frame, and ``None`` non-periodic frames. + The halo nodes of a periodic frame are images of its owned + nodes. Returns ------- @@ -235,7 +242,7 @@ def forward_lower_canonical_graph( raise ValueError( "a model with a long-range module reads the node positions" ) - check_owned_frames(n_node, n_local, communicator) + check_owned_frames(n_node, n_local, communicator, box) force_mag = None energy, atom_energy, force, virial, atom_virial = ( dpa4c_canonical_long_range_energy_force( @@ -246,6 +253,7 @@ def forward_lower_canonical_graph( atype, coord, communicator, + box, output_mask, atom_bias, do_atomic_virial, @@ -313,6 +321,7 @@ def forward_lower_canonical_graph_exportable( source_order: torch.Tensor, coord: torch.Tensor | None = None, communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -325,11 +334,12 @@ def forward_lower_canonical_graph_exportable( Compact graph node and edge tensors. destination_row_ptr, source_row_ptr, source_order Compact dual-CSR topology. - coord, communicator + coord, communicator, box For a model with a long-range module, sample node positions with - shape ``(N, 3)``, float64, and a sample communicator, an int64 CPU - tensor of one element, which the traced module then takes as its - ninth and tenth inputs; ``None`` otherwise. + shape ``(N, 3)``, float64, a sample communicator, an int64 CPU + tensor of one element, and sample cells with shape ``(nf, 9)``, + float64 on the CPU, which the traced module then takes as its + ninth to eleventh inputs; ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -338,14 +348,14 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced compact deployment module of eight inputs, or ten with the - node positions and the communicator. + Traced compact deployment module of eight inputs, or eleven with + the node positions, the communicator and the cells. """ model = self - if (coord is None) != (communicator is None): + if not (coord is None) == (communicator is None) == (box is None): raise ValueError( "the compact lower of a model with a long-range module takes " - "the node positions and the communicator together" + "the node positions, the communicator and the cells together" ) if coord is not None: @@ -361,6 +371,7 @@ def fn_long_range( source_order: torch.Tensor, coord: torch.Tensor, communicator: torch.Tensor, + box: torch.Tensor, ) -> dict[str, torch.Tensor]: return model.forward_lower_canonical_graph( atype, @@ -374,6 +385,7 @@ def fn_long_range( do_atomic_virial=do_atomic_virial, coord=coord, communicator=communicator, + box=box, ) return make_fx(fn_long_range, **make_fx_kwargs)( @@ -387,6 +399,7 @@ def fn_long_range( source_order, coord, communicator, + box, ) def fn( @@ -707,7 +720,7 @@ def forward_lower_graph_exportable( do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, - topologies: dict[str, PairGraph] | None = None, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -742,10 +755,10 @@ def forward_lower_graph_exportable( fparam, aparam, do_atomic_virial, charge_spin As in ``forward_lower``. topologies - The sample pair topology the atomic model declares, or ``None``; - its ``index``, ``vec`` and ``mask`` become the inputs - ``pair_index``, ``pair_vec`` and ``pair_mask`` after - ``charge_spin``. + The sample topologies the atomic model declares, or ``None``; + their fields follow ``charge_spin`` as further inputs, in the + layout of + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies`. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -790,17 +803,10 @@ def forward_lower_graph_exportable( do_grad_c = self.do_grad_c("energy") if topologies: - # The inputs of the deployed lower keep their names, which the - # export reads, so the traced signature spells out the one pair - # topology a long-range module declares. - if len(topologies) != 1: - raise NotImplementedError( - "the deployed graph lower takes one pair topology, got " - f"{list(topologies)}" - ) - (pairs,) = topologies.values() + # The topology inputs follow ``charge_spin``, their fields in the + # layout of flatten_topologies. - def fn_pairs( + def fn_topologies( atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, @@ -814,9 +820,7 @@ def fn_pairs( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, - pair_index: torch.Tensor, - pair_vec: torch.Tensor, - pair_mask: torch.Tensor, + *topology_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: model_ret = traced( atype, @@ -832,9 +836,7 @@ def fn_pairs( fparam, aparam, charge_spin, - pair_index, - pair_vec, - pair_mask, + *topology_vals, ) return _translate_energy_keys( model_ret, @@ -844,7 +846,7 @@ def fn_pairs( local=True, ) - return make_fx(fn_pairs, **make_fx_kwargs)( + return make_fx(fn_topologies, **make_fx_kwargs)( atype, n_node, n_local, @@ -858,9 +860,7 @@ def fn_pairs( fparam, aparam, charge_spin, - pairs.index, - pairs.vec, - pairs.mask, + *flatten_topologies(topologies), ) def fn( diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index d833140bc4..25e88487a7 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -21,9 +21,6 @@ OutputVariableDef, ) from deepmd.dpmodel.utils.neighbor_graph import ( - KGrid, - MeshGrid, - PairGraph, compact_nodes, expand_node_values, flatten_topologies, @@ -48,6 +45,7 @@ from .edge_transform_output import ( fit_output_to_model_output_graph, + releaf_topologies, ) from .transform_output import ( fit_output_to_model_output, @@ -201,15 +199,18 @@ def check_owned_frames( n_node: torch.Tensor, n_local: torch.Tensor | None, communicator: torch.Tensor | None = None, + periodic_marks: torch.Tensor | None = None, ) -> None: - """Reject halo nodes in the frames of a single domain with a long-range module. + """Reject halo nodes in the non-periodic frames of a single domain with a long-range module. - The module sums over every pair of atoms of a frame, so a frame that one - domain holds whole owns all its nodes; halo nodes would carry no charge - and remove their interactions from the sum. Across the domains of a - communicator, every domain owns its share of the frame instead. The - counts are read eagerly only: a traced lower receives them from callers - that check them. + The module sums over every pair of atoms of a non-periodic frame, so a + frame that one domain holds whole owns all its nodes; halo nodes would + carry no charge and remove their interactions from the sum. Across the + domains of a communicator, every domain owns its share of the frame + instead. The halo nodes of a periodic frame are images of its atoms, + which carry no charge: the reciprocal-space sum holds the images itself. + The counts are read eagerly only: a traced lower receives them from + callers that check them. Parameters ---------- @@ -220,11 +221,16 @@ def check_owned_frames( communicator The communicator input of the compact lower, a CPU tensor holding the address of an MPI_Comm handle or zero, or ``None`` for a single domain. + periodic_marks + A tensor with a nonzero entry when the frames are periodic, such as + their cells or the frame mask of a reciprocal-space topology, or + ``None`` for non-periodic frames. It is read only after the fake + check, so that a traced lower computes nothing on it. Raises ------ ValueError - If a frame of a single domain carries halo nodes. + If a non-periodic frame of a single domain carries halo nodes. """ from torch._subclasses.fake_tensor import ( is_fake, @@ -234,14 +240,31 @@ def check_owned_frames( return if communicator is not None and int(communicator.reshape(-1)[0]) != 0: return + if periodic_marks is not None and bool(torch.any(periodic_marks != 0)): + return if not torch.equal(n_local, n_node): raise ValueError( - "a model with a long-range module evaluates complete frames on a " - "single domain, whose graph owns every node (n_local == n_node), " - f"but got n_local {n_local.tolist()} for n_node {n_node.tolist()}" + "a model with a long-range module evaluates complete non-periodic " + "frames on a single domain, whose graph owns every node (n_local " + f"== n_node), but got n_local {n_local.tolist()} for n_node " + f"{n_node.tolist()}" ) +def _reciprocal_masks(topologies: dict[str, Any]) -> torch.Tensor | None: + """The mask of the reciprocal-space topology, true on periodic entries, or ``None``.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + ) + + masks = [ + topology.mask + for topology in topologies.values() + if not isinstance(topology, PairGraph) + ] + return masks[0] if masks else None + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -670,7 +693,7 @@ def forward_common_lower_graph( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, comm_dict: dict | None = None, - topologies: dict[str, PairGraph] | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -762,9 +785,10 @@ def forward_common_lower_graph( topologies The topologies the atomic model declares by ``required_topologies()``, built on the node axis of the - graph from complete systems. Each ``vec`` becomes an autograd - leaf next to ``edge_vec``, and force and virial include its - pullback. ``None`` (default) for models declaring none. + graph. The geometry fields every topology names in its + ``LEAVES`` become autograd leaves next to ``edge_vec``, and + force and virial include their pullbacks. ``None`` (default) + for models declaring none. Returns ------- @@ -787,52 +811,11 @@ def forward_common_lower_graph( "the long-range topologies describe complete systems; " "a domain-decomposed system is not supported" ) - check_owned_frames(n_node, n_local) - # Every topology carries a geometry leaf of its own. - re_leafed: dict[str, Any] = {} - for name, topology in topologies.items(): - if isinstance(topology, KGrid): - if topology.coord is None or topology.strain is None: - raise ValueError( - "the k-grid topology carries no coordinates or strain" - ) - # The positions are the geometry leaf of the k-space - # energy; the zero strain leaf turns its gradient into - # the k-space virial (see KGrid.strain). Both re-leaf - # here so the backward roots at graph INPUTS — a leaf - # created inside the trace is NaN under inductor. - coord_leaf = topology.coord.detach().requires_grad_(True) - strain_leaf = topology.strain.detach().requires_grad_(True) - re_leafed[name] = KGrid( - frame=topology.frame, - vec=topology.vec, - mask=topology.mask, - volume=topology.volume, - coord=coord_leaf, - strain=strain_leaf, - ) - elif isinstance(topology, MeshGrid): - if topology.coord is None or topology.strain is None: - raise ValueError( - "the mesh-grid topology carries no coordinates or strain" - ) - # Same re-leaf contract as the k grid (see above). - coord_leaf = topology.coord.detach().requires_grad_(True) - strain_leaf = topology.strain.detach().requires_grad_(True) - re_leafed[name] = MeshGrid( - box=topology.box, - volume=topology.volume, - coord=coord_leaf, - strain=strain_leaf, - mask=topology.mask, - ) - else: - re_leafed[name] = PairGraph( - index=topology.index, - vec=topology.vec.detach().requires_grad_(True), - mask=topology.mask, - ) - topologies = re_leafed + check_owned_frames( + n_node, n_local, periodic_marks=_reciprocal_masks(topologies) + ) + # Every topology carries geometry leaves of its own. + topologies = releaf_topologies(topologies) if spin is not None: # second autograd leaf: force_mag = -dE/dspin (native spin). # Deliberately a SEPARATE leaf/backward from edge_vec rather @@ -1389,7 +1372,7 @@ def forward_common_lower_graph_exportable( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, destination_sorted: bool = False, - topologies: dict[str, PairGraph] | None = None, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 586d8bb79a..bb327bb3be 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -73,8 +73,6 @@ make_learning_rate_schedule, ) from deepmd.dpmodel.utils.neighbor_graph import ( - KGrid, - MeshGrid, flatten_topologies, unflatten_topologies, ) @@ -118,6 +116,9 @@ from deepmd.pt_expt.model import ( get_model, ) +from deepmd.pt_expt.model.edge_transform_output import ( + releaf_topologies, +) from deepmd.pt_expt.model.graph_lower import ( model_uses_graph_lower, ) @@ -786,9 +787,10 @@ def _trace_and_compile_graph( through the same model-owned translation used by eager execution. A model declaring topologies (``atomic_model.required_topologies()``) - receives each of them as three further inputs, ``(index, vec, mask)`` in - the declared order, ahead of the task buffers; its pair axis is a fourth - dynamic dimension, kept a distinct symbol like the other three. + receives the fields of each, laid out by :func:`flatten_topologies`, as + further inputs ahead of the task buffers; their dynamic axes, the pairs, + the wave numbers or the axes of the mesh, are symbols distinct from the + other three. Parameters ---------- @@ -925,8 +927,8 @@ def _trace_and_compile_graph( # tensors, while forward_common_lower_graph takes it after them. s_spin = s_conditioning.pop(0) if spin is not None else None s_fparam, s_aparam, s_charge_spin = s_conditioning - # The pair axis of every topology is a fourth dynamic axis, its sample - # length distinct from the other three and from the CSR offsets. + # The dynamic axes of the topologies take sample lengths distinct from the + # other three and from the CSR offsets. s_topologies = build_synthetic_topologies( model, trace_N, @@ -936,6 +938,7 @@ def _trace_and_compile_graph( dtype=GLOBAL_PT_FLOAT_PRECISION, device=_trace_device, nframes=trace_nf, + forbidden=_forbidden | {e_max}, ) topology_names = tuple(s_topologies) s_topology_vals = flatten_topologies(s_topologies) @@ -1721,38 +1724,9 @@ def _forward_graph( edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) - # Each pair topology enters as (index, vec, mask), its vec a further - # leaf; a k grid enters as (frame, vec, mask, volume, coord, strain), - # its coord and zero strain further leaves (the k-space force and - # virial root at these inputs). - topology_vals_list: list[torch.Tensor] = [] - for topology in topologies.values(): - if isinstance(topology, KGrid): - assert topology.coord is not None and topology.strain is not None - topology_vals_list += [ - topology.frame, - topology.vec, - topology.mask, - topology.volume, - topology.coord.detach().requires_grad_(True), - topology.strain.detach().requires_grad_(True), - ] - elif isinstance(topology, MeshGrid): - assert topology.coord is not None and topology.strain is not None - topology_vals_list += [ - topology.box, - topology.volume, - topology.coord.detach().requires_grad_(True), - topology.strain.detach().requires_grad_(True), - topology.mask, - ] - else: - topology_vals_list += [ - topology.index, - topology.vec.detach().requires_grad_(True), - topology.mask, - ] - topology_vals = tuple(topology_vals_list) + # The topologies enter as their fields, the geometry of every one a + # further leaf. + topology_vals = flatten_topologies(releaf_topologies(topologies)) if self._task_buf_order: try: @@ -2344,6 +2318,9 @@ def initialize_statistics( for model_key in self.model_keys } new_state = {} + # The atomic models, by module path, whose long-range module + # takes its initialization rather than pretrained weights. + fresh_long_range: dict[str, str] = {} for key in target_state: if key == "_extra_state": new_state[key] = target_state[key] @@ -2370,17 +2347,39 @@ def initialize_statistics( finetune_rule.get_random_fitting() and ".descriptor." not in key ) - if use_random: + if use_random or pretrained_key not in pretrained_state: new_state[key] = target_state[key] - elif pretrained_key in pretrained_state: - new_state[key] = pretrained_state[pretrained_key] + if ".long_range." in key: + owner = key[: key.index(".long_range.")] + fresh_long_range[owner] = model_key else: - new_state[key] = target_state[key] + new_state[key] = pretrained_state[pretrained_key] break if not matched: new_state[key] = target_state[key] self._unwrapped.load_state_dict(new_state) + # A long-range module added to the pretrained model takes the + # input statistics of the fine-tuning data, as in training from + # scratch, before the bias adjustment evaluates the model. + for owner, model_key in fresh_long_range.items(): + atomic_model = self._unwrapped.get_submodule(owner) + + def long_range_stat( + _atomic_model: torch.nn.Module = atomic_model, + _model_key: str = model_key, + ) -> None: + _atomic_model.compute_long_range_input_stat( + self._sample_funcs[_model_key] + ) + + self._run_stat_on_chief( + long_range_stat, + operation=f"long-range statistics for task {model_key!r}", + ) + if self.is_distributed: + self._broadcast_model_stat(atomic_model.long_range) + # Per-branch bias adjustment (rank 0 only, then broadcast) for model_key in self.model_keys: finetune_rule = finetune_links[model_key] diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 9e6ec43662..13b15d63b7 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import contextlib import ctypes +import dataclasses import json import logging import os @@ -15,7 +16,6 @@ import torch from deepmd.dpmodel.utils.neighbor_graph import ( - PairGraph, flatten_topologies, ) from deepmd.dpmodel.utils.nlist import ( @@ -695,83 +695,88 @@ def build_synthetic_graph_inputs( def build_synthetic_topologies( model: torch.nn.Module, n_node_total: int, - n_pair: int, + n_entry: int, *, dtype: torch.dtype, device: torch.device, - nframes: int | None = None, -) -> dict[str, PairGraph]: + nframes: int, + forbidden: set[int] | frozenset[int] = frozenset(), +) -> dict[str, Any]: """Build trace-time samples of the topologies a model declares. Single source of the trace-time topology inputs, shared by graph export - and compiled training. Tracing reads only their shapes, so every topology - joins fixed nodes over ``n_pair`` entries, a length the caller keeps - distinct from the other dynamic axes. A k grid additionally carries the - frame volumes and the node coordinates. + and compiled training. Tracing reads only their shapes. Every topology + lists ``n_entry`` entries on its own axis, the pairs or the wave numbers, + a length the caller keeps distinct from the other dynamic axes; a mesh + takes the three smallest even dimensions of at least four that no other + axis and no entry of ``forbidden`` has. Parameters ---------- model : torch.nn.Module The model; its atomic model declares the topologies. n_node_total : int - Length of the node axis the pairs index. - n_pair : int - Length of the pair axis of every sample. + Length of the node axis. + n_entry : int + Length of the entry axis of every sample. dtype : torch.dtype - Precision of the pair vectors. + Precision of the geometry. device : torch.device Target device. - nframes : int | None - The frame count, needed for the volume axis of a k grid; ``None`` - for a model declaring none. + nframes : int + Length of the frame axis. + forbidden : set[int] + Sizes of the static dimensions of the model, which a mesh avoids. Returns ------- - dict[str, PairGraph] + dict[str, PairGraph | KGrid | MeshGrid] A sample of every declared topology; empty when none is declared. """ from deepmd.dpmodel.utils.neighbor_graph import ( K_GRID, PME_MESH, KGrid, - MeshGrid, + PairGraph, ) - - index = ( - torch.arange(2 * n_pair, dtype=torch.int64, device=device).reshape(2, n_pair) - % n_node_total + from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import ( + mesh_grid_of_dims, ) + + taken = set(forbidden) | {nframes, n_node_total, n_node_total + 1, n_entry} + mesh_dims = [] + for size in range(4, 4 + 2 * len(taken) + 6, 2): + if size not in taken and len(mesh_dims) < 3: + mesh_dims.append(size) + coord = torch.ones(n_node_total, 3, dtype=dtype, device=device) + # Contiguous like the inputs at run time: a compiled graph reads the + # strides of its samples. + box = torch.eye(3, dtype=dtype, device=device).repeat(nframes, 1, 1) topologies = {} for name in model.atomic_model.required_topologies(): if name == K_GRID: - if nframes is None: - raise ValueError(f"a {K_GRID!r} topology sample needs the frame count") topologies[name] = KGrid( - frame=torch.zeros(n_pair, dtype=torch.int64, device=device), + frame=torch.zeros(n_entry, dtype=torch.int64, device=device), # Nonzero wave numbers: the kernel transform carries 1/k^2. - vec=torch.ones(n_pair, 3, dtype=dtype, device=device), - mask=torch.ones(n_pair, dtype=torch.bool, device=device), + vec=torch.ones(n_entry, 3, dtype=dtype, device=device), + mask=torch.ones(n_entry, dtype=torch.bool, device=device), volume=torch.ones(nframes, dtype=dtype, device=device), - coord=torch.ones(n_node_total, 3, dtype=dtype, device=device), + coord=coord, strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), ) elif name == PME_MESH: - if nframes is None: - raise ValueError( - f"a {PME_MESH!r} topology sample needs the frame count" - ) - topologies[name] = MeshGrid( - box=torch.eye(3, dtype=dtype, device=device).expand(nframes, 3, 3), - volume=torch.ones(nframes, dtype=dtype, device=device), - coord=torch.ones(n_node_total, 3, dtype=dtype, device=device), - strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), - mask=torch.ones(nframes, dtype=torch.bool, device=device), + topologies[name] = mesh_grid_of_dims( + box, + coord, + tuple(mesh_dims), + torch.ones(nframes, dtype=torch.bool, device=device), ) else: + index = torch.arange(2 * n_entry, dtype=torch.int64, device=device) topologies[name] = PairGraph( - index=index, - vec=torch.ones(n_pair, 3, dtype=dtype, device=device), - mask=torch.ones(n_pair, dtype=torch.bool, device=device), + index=index.reshape(2, n_entry) % n_node_total, + vec=torch.ones(n_entry, 3, dtype=dtype, device=device), + mask=torch.ones(n_entry, dtype=torch.bool, device=device), ) return topologies @@ -790,8 +795,9 @@ def build_synthetic_canonical_graph_inputs( compact ABI, matching :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. ``want_long_range`` appends the inputs that the compact lower of a model - with a long-range module reads: the float64 node positions at slot 8 and - the int64 communicator, a CPU tensor of one element, at slot 9. + with a long-range module reads: the float64 node positions at slot 8, the + int64 communicator, a CPU tensor of one element, at slot 9, and the + float64 cells of the frames, a CPU tensor of shape ``(nf, 9)``, at slot 10. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -857,7 +863,8 @@ def build_synthetic_canonical_graph_inputs( 3 * atype.shape[0], dtype=torch.float64, device=device ).reshape(-1, 3) communicator = torch.zeros(1, dtype=torch.int64, device="cpu") - return (*compact_sample, coord, communicator) + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64, device="cpu") + return (*compact_sample, coord, communicator, box) return compact_sample @@ -866,10 +873,11 @@ def _build_canonical_graph_dynamic_shapes( ) -> tuple: """Build dynamic shapes for the compact deployment ABI. - The trailing slots, the moments of a native-spin model or the positions - and the communicator of a model with a long-range module, are present - only for such models, so the sample length selects the shape tuple. The - communicator holds one address and is static. + The trailing slots, the moments of a native-spin model or the positions, + the communicator and the cells of a model with a long-range module, are + present only for such models, so the sample length selects the shape + tuple. The communicator holds one address and is static; the cells follow + the frame axis. """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, @@ -892,7 +900,9 @@ def _build_canonical_graph_dynamic_shapes( {0: node_dim + 1}, {0: edge_storage_dim}, ) - trailing = ({0: node_dim}, None)[: len(sample_inputs) - len(shapes)] + trailing = ({0: node_dim}, None, {0: nframes_dim})[ + : len(sample_inputs) - len(shapes) + ] return (*shapes, *trailing) @@ -946,6 +956,47 @@ def count_synthetic_graph_edges( return int(edge_mask.sum()) +def _topology_dynamic_shapes( + topology_names: tuple[str, ...], shared: dict[str, Any] +) -> list[dict[int, Any]]: + """Dynamic-shape specifications of the fields of the topologies of a graph lower. + + Every field follows the ``AXES`` its dataclass declares. The frame and + node axes are those of the other inputs in ``shared``; the entry axis of + every topology is a dimension of its own, of at least the two guard + entries; a mesh axis is twice a dimension, since the mesh dimensions are + even. + + Parameters + ---------- + topology_names : tuple[str, ...] + The topologies, in the order of their inputs. + shared : dict[str, Any] + The dimensions of the ``"frame"`` and ``"node"`` axes. + + Returns + ------- + list[dict[int, Any]] + One specification per field, in the layout of flatten_topologies. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + TOPOLOGY_TYPES, + ) + + specs = [] + for name in topology_names: + cls = TOPOLOGY_TYPES[name] + dims = dict(shared) + dims["entry"] = torch.export.Dim(f"{name}_entries", min=2) + for axis in ("mesh_1", "mesh_2", "mesh_3"): + dims[axis] = 2 * torch.export.Dim(f"{name}_{axis}_half", min=1) + specs += [ + {dim: dims[axis] for dim, axis in cls.AXES[field.name].items()} + for field in dataclasses.fields(cls) + ] + return specs + + def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, is_native_spin: bool = False, @@ -978,9 +1029,8 @@ def _build_graph_dynamic_shapes( (spin at slot 10) instead of the regular energy ABI (charge_spin at slot 12). topology_names : tuple[str, ...] - Topologies whose ``(index, vec, mask)`` follow the inputs above, each - with a dynamic pair axis of its own; the guard pairs of a topology - bound it below by two. + Topologies whose fields follow the inputs above (see + :func:`_topology_dynamic_shapes`). """ nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) @@ -997,32 +1047,9 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) ) - topologies = [] - for name in topology_names: - npair_dim = torch.export.Dim(f"npair_{name}", min=2) - if name == "k_grid": - # (frame, vec, mask, volume, coord, strain): the k axis, the - # frame axis and the node axis. - topologies += [ - {0: npair_dim}, - {0: npair_dim}, - {0: npair_dim}, - {0: nframes_dim}, - {0: n_node_total_dim}, - {0: nframes_dim}, - ] - elif name == "pme_mesh": - # (box, volume, coord, strain, mask): the frame axis and the node - # axis; the mesh axes are static and live on the module. - topologies += [ - {0: nframes_dim}, - {0: nframes_dim}, - {0: n_node_total_dim}, - {0: nframes_dim}, - {0: nframes_dim}, - ] - else: - topologies += [{1: npair_dim}, {0: npair_dim}, {0: npair_dim}] + topologies = _topology_dynamic_shapes( + topology_names, {"frame": nframes_dim, "node": n_node_total_dim} + ) if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -1438,15 +1465,16 @@ def _probe_has_message_passing(obj: object) -> bool | None: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break meta["pair_exclude_types"] = pair_exclude_types - # Topologies the lower reads beside the neighbor graph; absent when the - # model reads none. The graph lower takes each as its ``(index, vec, - # mask)`` after its other inputs (see ``flatten_topologies``), built by the - # caller from the coordinates of each frame. The compact canonical lower - # takes the node positions instead, from which its operators enumerate the - # pairs. + # Topologies the lower reads beside the neighbor graph, and the parameters + # they are built with; absent when the model reads none. The graph lower + # takes the fields of each after its other inputs (see + # ``flatten_topologies``), built by the caller from the coordinates and + # the cell of each frame. The compact canonical lower takes the node + # positions instead, from which its operators evaluate the long range. topologies = model.atomic_model.required_topologies() if topologies: meta["topologies"] = list(topologies) + meta["topology_params"] = model.atomic_model.topology_params() return meta @@ -2058,21 +2086,15 @@ def _trace_and_export_impl( # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) model.eval() - # The k-space evaluation of a periodic long-range module needs the - # reciprocal grid of the input cells, which no exported lower ABI carries - # yet; freeze and compress reject such models until they do. - from deepmd.dpmodel.utils.neighbor_graph import ( - K_GRID, - PME_MESH, + # A frozen model evaluates the reciprocal sum of a periodic frame on the + # mesh of its cell, whatever the mode it was trained in. + from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange, ) - if K_GRID in model.atomic_model.required_topologies() or ( - PME_MESH in model.atomic_model.required_topologies() - ): - raise NotImplementedError( - "a periodic long-range module (the k-grid/mesh topology) cannot be " - "frozen or exported yet; serve it through the Python interface" - ) + for module in model.modules(): + if isinstance(module, LatentChargeLongRange): + module.use_mesh_evaluation() # Fold the vacuum reference on the CPU tracing device. vacuum_descriptor # evaluates an isolated atom of every type through the descriptor; on CUDA # that path can allocate several GiB even when export already holds most of @@ -2243,8 +2265,9 @@ def _trace_and_export_impl( e_sample += 1 if canonical: # The compact lower of a model with a long-range module reads the - # node positions, from which its operators enumerate the pairs, - # and the communicator of the domains of a decomposed frame. + # node positions and the cells, from which its operator evaluates + # the long range, and the communicator of the domains of a + # decomposed frame. want_long_range = bool(model.atomic_model.required_topologies()) sample_inputs = build_synthetic_canonical_graph_inputs( model, @@ -2345,6 +2368,7 @@ def _trace_and_export_impl( dtype=edge_dtype, device=torch.device("cpu"), nframes=nframes_sample, + forbidden=_forbidden | {e_sample}, ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index ed740e9ab5..a9633e425f 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3677,10 +3677,10 @@ def charge_constraint_args() -> list[Argument]: "latent_charge", doc="Interaction of latent charges predicted from the descriptor. " "For a non-periodic system the kernel is summed directly over every pair " - "of atoms; for a periodic system (``periodic: true``) its Fourier " - "transform is summed over the reciprocal lattice points of the cell. " - "Every channel interacts through the same kernel; half of each pair energy " - "is attributed to each atom.", + "of atoms; for a periodic system (``periodic: true``) it is summed over " + "the images of the cell in reciprocal space. Every channel interacts " + "through the same kernel; every atom carries half of the interaction of " + "its charges with all others.", ) def long_range_latent_charge() -> list[Argument]: doc_channels = "The number of latent charge channels predicted per atom." @@ -3696,30 +3696,30 @@ def long_range_latent_charge() -> list[Argument]: "set." ) doc_periodic = ( - "Evaluate the module on periodic systems in reciprocal space: the " - "Fourier transform of the kernel is summed over the wave numbers of " - "each cell, with the tin-foil boundary convention, and the self term " - "subtracted. Every latent channel must be neutral." + "Evaluate the module on periodic systems as well: the kernel is summed " + "over the images of the cell in reciprocal space, with the " + "zero-wave-number mode dropped (the tin-foil convention, which places " + "a uniform neutralizing background under a channel with a net " + "charge), less the interaction of every charge with itself. " + "Non-periodic systems keep the direct sum." ) doc_kspace_tolerance = ( "The relative weight of the slowest-decaying kernel mode at the " - "wave-number cutoff of the k grid of a periodic batch." + "wave-number cutoff k_max of a periodic system." ) doc_k_grid_mode = ( - "How the reciprocal sum of a periodic batch is evaluated: 'ragged' " - "enumerates the wave numbers of every frame eagerly (the k axis is a " - "dynamic dimension); 'fixed' shares one integer grid over the frames, " - "sized by the largest cell of the data statistics or by " - "`k_grid_bounds`, and masks the points a frame does not reach, so the " - "k axis has a static length; 'pme' spreads the charges to a static " - "real-space mesh with a B-spline and evaluates the sum by an FFT " - "(SPME), sized by the largest cell of the data statistics or by " - "`pme_mesh`. The fixed grid is hard: a frame needing an index beyond " - "it (a larger cell, or a kernel that moved the cutoff out) has its " - "k-sum truncated to the grid, with a warning; on the mesh such a " - "frame is evaluated with a coarser effective spacing (a smooth " - "accuracy degradation, without a warning). Requires `periodic: true` " - "for 'fixed' and 'pme'." + "How training evaluates the reciprocal sum of a periodic batch: " + "'ragged' over the wave numbers of every frame within k_max, " + "enumerated per batch, so that the k axis is a dynamic dimension; " + "'fixed' over one integer grid shared by the frames, sized by the " + "largest cell of the data statistics or by `k_grid_bounds`, with the " + "points a frame does not reach masked, so that the k axis has a static " + "length (a frame whose cell needs points beyond the grid is truncated " + "to it, with a warning); 'pme' by the smooth particle-mesh Ewald " + "method on a mesh sized by the cells of the batch at `pme_spacing`, " + "whose cost grows nearly linearly with the system. A frozen model " + "evaluates every periodic frame on the mesh of its cell, whatever the " + "mode. Requires `periodic: true` for 'fixed' and 'pme'." ) doc_k_grid_bounds = ( "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " @@ -3727,22 +3727,21 @@ def long_range_latent_charge() -> list[Argument]: "`k_grid_mode: fixed`." ) doc_decouple_k_grid = ( - "Decouple the k-point selection of a periodic batch from the kernel " - "parameters: the wave-number cutoff is derived once from the initial " - "parameters and kept for the life of the model, so training the " - "kernel no longer moves the k grid (the kernel values the grid is " - "evaluated with still follow the live parameters). The frozen cutoff " - "is serialized with the model. Requires `periodic: true`." + "Decouple the wave-number cutoff of a periodic system from the kernel " + "parameters: it is derived once from the initial parameters and kept " + "for the life of the model, so that the k grid and the mesh stay put " + "while the kernel trains (the kernel values still follow the live " + "parameters). The frozen cutoff is serialized with the model. " + "Requires `periodic: true`." ) doc_pme_spacing = ( - "The target real-space mesh spacing of the 'pme' mode, in Angstrom; " - "the default is the Nyquist spacing of the kernel's wave-number " - "cutoff, pi / k_max." + "The largest spacing of the mesh of a periodic frame, in Angstrom; " + "the default is three quarters of the Nyquist spacing of the " + "wave-number cutoff, 0.75 pi / k_max. Requires `periodic: true`." ) - doc_pme_order = "The B-spline order of the 'pme' charge assignment." - doc_pme_mesh = ( - "The per-axis mesh dimensions [n1, n2, n3] of the 'pme' mode, in " - "place of the data statistics." + doc_pme_order = ( + "The order of the B-splines that spread the charges to the mesh, 8 " + "by default. Requires `periodic: true`." ) return [ Argument("channels", int, optional=True, default=1, doc=doc_channels), @@ -3788,6 +3787,8 @@ def long_range_latent_charge() -> list[Argument]: optional=True, default="ragged", doc=doc_k_grid_mode, + extra_check=lambda x: x in ("ragged", "fixed", "pme"), + extra_check_errmsg="must be 'ragged', 'fixed' or 'pme'", ), Argument( "k_grid_bounds", @@ -3805,24 +3806,17 @@ def long_range_latent_charge() -> list[Argument]: ), Argument( "pme_spacing", - float, + [float, None], optional=True, default=None, doc=doc_pme_spacing, ), Argument( "pme_order", - int, - optional=True, - default=4, - doc=doc_pme_order, - ), - Argument( - "pme_mesh", - list[int], + [int, None], optional=True, default=None, - doc=doc_pme_mesh, + doc=doc_pme_order, ), ] @@ -3830,9 +3824,10 @@ def long_range_latent_charge() -> list[Argument]: def long_range_args() -> Argument: doc_long_range = ( "A long-range module added to the atomic energy of the energy fitting. " - "It reads the same descriptor and is evaluated on complete systems; " - "non-periodic directly, periodic in reciprocal space when the module " - "enables it. Currently validated with the DPA4C descriptor." + "It reads the same descriptor and is evaluated on complete systems: " + "non-periodic ones directly, periodic ones in reciprocal space when " + "the module serves them. Currently validated with the DPA4C " + "descriptor." ) return Argument( "long_range", diff --git a/doc/model/long-range.md b/doc/model/long-range.md index ddeb4c2cb5..1f190dbb9a 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -21,22 +21,22 @@ Coulomb constant. Every channel interacts through the same kernel. For a non-periodic system, the sum runs over the distinct pairs of each frame: it has no self term, and no periodic image. For a periodic system (with -`periodic: true`, see below) the kernel is summed over the reciprocal lattice -points of the cell instead. +`periodic: true`, see below) the kernel is summed over the images of the cell +as well, in reciprocal space. The charges depend on the atomic environments, so forces and virials, which are obtained by differentiating the energy, include the response of the charges to the positions as well as the direct dependence of the kernel on -the distance. Half of each pair energy is attributed to each of its atoms, -so the atomic energies still sum to the energy of the frame. Every quantity +the distance. Every atom carries half of the interaction of its charges with +all others, so the atomic energies still sum to the energy of the frame. Every quantity derived from the model sees the complete energy: training, `dp test`, the output-bias statistics and change-bias, the Hessian, and a linear composition of models. > [!WARNING] -> The pair sum covers complete systems. Without `periodic: true`, a periodic -> cell is rejected; the cost of the non-periodic sum grows with the square of -> the number of atoms of a frame. +> The sum covers complete systems. Without `periodic: true`, a periodic cell +> is rejected. The cost of the non-periodic sum grows with the square of the +> number of atoms of a frame. ## Configuration @@ -113,19 +113,17 @@ the model represent interactions that one charge per atom cannot. With {ref}`periodic ` enabled, -a frame given a cell is evaluated in reciprocal space. The Fourier transform -of the kernel, +a frame given a cell is evaluated over the images of the cell. The sum +converges in reciprocal space, where the Fourier transform of the kernel, -- `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, the damped - Coulomb transform (the kernel is the interaction of two Gaussian charge - clouds, so its transform needs no new parameter), or -- `sog`: $\hat K(\mathbf{k}) = \sum_m a_m (2\pi)^{3/2} \sigma_m^3\, e^{-\sigma_m^2 k^2/2}$, - the term-by-term transform of the Gaussians; with the default amplitudes the - coefficient simplifies to $4\pi\ln b\, \sigma_m^2$ and the sum approaches the - Coulomb transform $4\pi/k^2$, +- `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, or +- `sog`: $\hat K(\mathbf{k}) = \sum_k a_k (2\pi)^{3/2} \sigma_k^3\, e^{-\sigma_k^2 k^2/2}$, + which with the default amplitudes approaches the Coulomb transform + $4\pi/k^2$, -is summed over the wave numbers $\mathbf{k} = 2\pi B^{-1} \mathbf{n}$, -$\mathbf{n} \in \mathbb{Z}^3 \setminus \{\mathbf{0}\}$, of each cell: +decays like a Gaussian. Over the wave numbers $\mathbf{k} = 2\pi B^{-1} +\mathbf{n}$, $\mathbf{n} \in \mathbb{Z}^3 \setminus \{\mathbf{0}\}$, of a cell +$B$ whose rows are the lattice vectors, ```math E = k_e \left[\frac{1}{2V}\sum_{\mathbf{k}\neq\mathbf{0}} \sum_{c=1}^{C} @@ -134,24 +132,45 @@ E = k_e \left[\frac{1}{2V}\sum_{\mathbf{k}\neq\mathbf{0}} \sum_{c=1}^{C} \qquad \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i\mathbf{k}\cdot\mathbf{r}_i}, ``` -with the cell volume $V$ and the kernel value at contact $K(0)$. The -wave-number cutoff follows the kernel parameters through -{ref}`kspace_tolerance `: -the grid ends where the weight of the slowest-decaying mode falls below it, -so the truncation is exponentially convergent. The $\mathbf{k}=\mathbf{0}$ -mode is dropped (the tin-foil boundary convention), so every latent channel -must be neutral: use a `neutral` constraint, or a `charge_spin`/`fparam` -constraint holding a neutral total charge. - -The k grid is enumerated outside the compiled region and enters it as inputs -of dynamic length, like the pairs of the non-periodic evaluation; a model -with `periodic: true` still serves non-periodic batches. Forces come from the -positions the phases are evaluated at, and the virial from a strain the -energy is differentiated against, so both hold the complete k-space term. - -The periodic evaluation is supported for training, compiled training and the -Python inference interface. Freezing/export, compression, the C++ interface -and LAMMPS do not serve a periodic long-range module yet: they reject it. +with the cell volume $V$ and the kernel value at contact $K(0)$, whose term +removes the interaction of every charge with itself. The +$\mathbf{k}=\mathbf{0}$ mode is dropped, the tin-foil convention, which places +a uniform neutralizing background under a channel with a net charge. The +wave-number cutoff $k_{max}$ is where the weight of the slowest-decaying mode +falls below +{ref}`kspace_tolerance `, +so the truncation converges exponentially; it follows the kernel as it trains, +unless +{ref}`decouple_k_grid ` +keeps the cutoff of the initial kernel. + +{ref}`k_grid_mode ` +selects how training evaluates the sum: + +- `ragged`: over the wave numbers of every frame within $k_{max}$, listed for + every batch; their number grows with the volume of the cell. +- `fixed`: over one integer grid shared by the frames of a batch, sized by the + largest cell of the data statistics or by + {ref}`k_grid_bounds `; + a frame whose cell needs a wave number beyond the grid is truncated to it, + with a warning. +- `pme`: by the smooth particle-mesh Ewald method. The charges are spread + with B-splines of order + {ref}`pme_order ` to a + mesh that resolves every cell of the batch at + {ref}`pme_spacing `, + whose transform gives the sum; the cost grows nearly linearly with the + number of atoms. With the default order 8 and spacing $0.75\pi/k_{max}$, the + energy and the forces stay within a few parts in a million of the sum over + the wave numbers. + +A frozen model evaluates every periodic frame on the mesh of its cell, +whatever its training mode, so a model trained in another mode can differ +from its frozen form by the accuracy of the mesh. Forces and virials hold the +complete periodic term, and a model with `periodic: true` still evaluates +non-periodic frames by the direct sum. In a `linear_ener` composition, a +sub-model whose module serves periodic systems takes no other sub-model with +a long-range module beside it. ### Total charge @@ -170,16 +189,18 @@ Pairs of the types in `pair_exclude_types` see each other neither in the descriptor nor in the long-range interaction, while the `exclude_types` of the descriptor act on the descriptor alone. In a `linear_ener` composition, `pair_exclude_types` belongs to the composition, whose sub-models share its -neighbor graph and pair topology. `atom_exclude_types` is not supported with -a long-range module. +neighbor graph and pair topology. The reciprocal-space sum of a periodic +system is global and excludes no pair, so `pair_exclude_types` is not +supported with a module that serves periodic systems, nor +`atom_exclude_types` with any long-range module. ## Training and inference Training uses the usual energy loss; see [training energy models](train-energy.md). With `enable_compile`, the -long-range interaction is part of the compiled model: the pairs of each frame -are listed outside the compiled region and enter it as inputs of dynamic -length. +long-range interaction is part of the compiled model: the pairs, the wave +numbers or the mesh of each batch are built outside the compiled region and +enter it as inputs of dynamic size. `dp --pt-expt freeze` exports the model, long-range module included, to `.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a @@ -187,14 +208,20 @@ length. A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` and the Python inference interface; the C++ interface and LAMMPS, with `pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. -Since the pair sum covers the complete system, LAMMPS runs such a model with -non-periodic boundaries (`f`, `s` or `m`) in every dimension, and rejects a -periodic one. A model that the dedicated operator described below evaluates, -a compressed DPA4C model frozen for a CUDA device, also runs on several MPI -ranks: every rank gathers the positions and charges of the atoms of all ranks -over MPI and evaluates the pairs of its own atoms, so that the ranks share -the cost of the pair sum. This requires DeePMD-kit built with MPI. Any other -model, like an uncompressed one, runs on a single rank. +Since the sum covers the complete system, LAMMPS runs a model whose module +serves no periodic system with non-periodic boundaries (`f`, `s` or `m`) in +every dimension, and a model with `periodic: true` with boundaries periodic in +every dimension or in none. Several models given to `pair_style deepmd` for +the model deviation evaluate a non-periodic system together only when every +one of them carries a long-range module. A model that the dedicated operator described +below evaluates, a compressed DPA4C model frozen for a CUDA device, also runs +on several MPI ranks: every rank gathers the positions and charges of the +atoms of all ranks over MPI and evaluates its own atoms against all of them, +so that the ranks share the cost of the direct sum, while each evaluates the +mesh of a periodic system whole. This requires DeePMD-kit built with MPI. Any +other model, like an uncompressed one, runs on a single rank; +`pair_style deepmd/kk` evaluates an uncompressed model on non-periodic systems +only. Compressing a descriptor conditioned on `charge_spin` specializes the model to one charge state at a time, and a constraint that reads its total charge @@ -202,15 +229,17 @@ from `charge_spin` follows the same state. The compressed archive serves any charge state: the Python interface applies the state of each call, and LAMMPS the one given by the `charge_spin` keyword of `pair_style deepmd`. -A compressed DPA4C model frozen for a CUDA device evaluates the pair sum with -a dedicated operator, which reads the positions of the atoms and enumerates -their pairs itself: its memory grows linearly with the number of atoms, and -its cost with their square. The operator serves the `les` kernel and the -`sog` kernel of up to 32 Gaussians, with up to 8 charge channels, a charge -network without frame or atomic parameters, and no `pair_exclude_types`. Any -other model, like an uncompressed one, lists the pairs of each frame before -calling the model, so its memory grows with the square of the number of -atoms. +A compressed DPA4C model frozen for a CUDA device evaluates the long-range +term with a dedicated operator, which reads the positions of the atoms and +the cell. It enumerates the pairs of a non-periodic system itself, so that its +memory grows linearly with the number of atoms and its cost with their +square, and evaluates a periodic system on the mesh of its cell. The operator +serves the `les` kernel and the `sog` kernel of up to 32 Gaussians, with up to +8 charge channels, B-splines of order up to 12, a charge network without +frame or atomic parameters, and no `pair_exclude_types`. Any other model, +like an uncompressed one, receives the pairs or the mesh of each frame from +the interface, so the memory of a non-periodic evaluation grows with the +square of the number of atoms. The following are not supported for a model with a long-range module: multi-rank (domain-decomposed) inference of a model that the dedicated diff --git a/periodic_long_range_design.md b/periodic_long_range_design.md deleted file mode 100644 index 7d18797848..0000000000 --- a/periodic_long_range_design.md +++ /dev/null @@ -1,334 +0,0 @@ -# 周期性 LES / SOG 长程能量:公式与 k 点实现方案 - -> 分支:`dpa4c+period` -> 目标:给 `deepmd/dpmodel/long_range`(latent_charge 模块)的 LES 和 SOG 两种 -> kernel 增加周期边界支持,第一版只覆盖 **Python 训练路径(graph lower)**; -> 压缩 / CUDA direct-sum / C++ / LAMMPS 路径显式拒绝。 -> -> 设计取向:**纯倒易空间(all-reciprocal)方案**。SOG 高斯核的傅里叶变换是 -> 带本征高斯阻尼的整函数,实空间 erfc 残差恒为零,因此不需要 Ewald β 分离、 -> 不需要实空间项、不需要动 edge 图。这与 `deepmd-kit-charge` 仓库 -> (`sog_model._compute_periodic`)的路线一致,但修正其三处一致性问题: -> ① k 空间振幅缺 `(2π)^{3/2}σ³` 因子;② 周期分支丢失 LES;③ 自能用截断 k 和近似。 - -## 1. 记号 - -- 原子数 $N$,latent 通道数 $C$,电荷 $q_i^c$($c=1..C$,单位 $e$) -- $k_e = 14.3996454784255$ eV·Å(`ElectrostaticConvertion`,`source/lib/include/ewald.h:19`) -- box 矩阵 $B$(行向量为 $\mathbf{a}_1,\mathbf{a}_2,\mathbf{a}_3$),体积 $V = |\det B| > 0$ -- 倒易格点:$\mathbf{k}(\mathbf{n}) = 2\pi B^{-T}\mathbf{n}$,$\mathbf{n}\in\mathbb{Z}^3$ -- 每通道结构因子:$\hat\rho_c(\mathbf{k}) = \sum_i q_i^c\, e^{i\mathbf{k}\cdot\mathbf{r}_i}$ - (整体平移只给 $\hat\rho$ 乘一个模 1 相位,$|\hat\rho|^2$ 不变,坐标原点任意) -- SOG 参数:宽度几何序列 $\sigma_m = \sigma_0 b^m$($m=0..M-1$), - **修正后振幅** $a_m = \sqrt{2/\pi}\,\ln b/\sigma_m$ - (本仓库 `16f40a2b`;旧常数振幅 $4\pi\ln b$ 给出对数形核,已废弃) - -## 2. 非周期公式(现有实现,不变) - -$$E_{LR} = k_e \sum_{i |n_d| <= k_max*|a_d|/(2pi) (Cauchy-Schwarz) - # 倒易基矢(列):g_d = 2*pi*B^{-1}[:, d],使 a_d · g_e = 2pi delta_de - # 注意是 inv(B) 而非 inv(B).T(B 非对称时 transpose 是错的) - rec = 2*pi * inv(box) # rec[:, d] = g_d - n_d = ceil(k_max * norm(box, axis=1) / (2*pi)) # (3,) - # (3) 枚举 + 字典序(确定性顺序!)+ 球形截断 - ks, ns = [], [] - for n0 in range(-n_d[0], n_d[0]+1): # 外层→内层固定序 - for n1 in range(-n_d[1], n_d[1]+1): - for n2 in range(-n_d[2], n_d[2]+1): - if n0 == n1 == n2 == 0: continue - k = rec @ (n0, n1, n2) - if dot(k, k) <= k_max*k_max: keep - return k_vec (nk,3), k2 (nk,) -``` - -- **数量级**:$\sigma_0 = 2.18$ Å、$\varepsilon=10^{-8}$ → $k_{max}\approx 6.2$ Å⁻¹; - 15 Å 立方盒约 $9^3-1 \approx 700$ 个 k 点;$\hat\rho$ 计算量 $O(N_k\cdot N\cdot C)$。 -- 每帧网格不同(box 不同)→ batch 拼接时需要逐帧归属(§5.2)。 -- **compile:ragged 动态轴**(`nk` 成为第四个符号轴,照 `npair` 协议, - flatten/unflatten 为 k_grid 携 `(frame, vec, mask, volume, coord, strain)`)。 -- **固定网格 + mask(已实现,`k_grid_mode: fixed`)**:v1 时未采用,后按需求 - 补上作为可选模式。整数网格 $\mathbf{n} \in [-n_d, n_d]^3$(不含零矢量) - 由数据统计中最大 cell 的逐方向界 $n_d = \lceil k_{max}|a_d|/2\pi \rceil$ - 一次性枚举(`enumerate_k_grid_bounds`,可用 `k_grid_bounds` 配置覆盖), - 存入模型并随其序列化;每个 batch 以静态形状 matmul 求各帧波矢 - $\mathbf{k} = 2\pi B_f^{-1}\mathbf{n}$(k 轴长度固定为 $n_f \times n_k$,无 - guard 点),帧达不到的点($|k| > k_{max}$,$k_{max}$ 取构建时 kernel 参数) - 由 mask 屏蔽(`build_fixed_k_grid`)。网格是硬性的:cell 超出网格容量或 - kernel 漂移使 $k_{max}$ 越界时,该帧仅以网格内 k 点截断求和(倒格点欠采样的 - 近似,力和位力仍是该近似能量的精确导数),并发 warning。零矢量不入网格:LES 的 $\hat K$ 在 - $k=0$ 奇异,masked 的 inf 会在反向传播产生 NaN。开销约为 ragged 球内点数 - 的 ~2 倍(各向同性 cell),换来完全静态的形状。 -- **解耦开关(`decouple_k_grid: true`)**:默认(耦合)每 batch 以当前 kernel - 参数现算 $k_{max}$,kernel 漂移 → 网格跟随,能量定义始终是当前 kernel 的收敛 - 能量;置 true 后 $k_{max}$ 在构造时按初始参数冻结(`frozen_k_max`,随模型 - 序列化),kernel 漂移不再影响 k 点选择,图内 $\hat K(k)$ 仍用活参数求值 - (梯度不受影响)。两种模式共用;fixed 模式下解耦后 mask 也不再随漂移收紧。 -- **PME 网格模式(`k_grid_mode: "pme"`)**:不直接枚举 k 点,而是把电荷用 - p 阶 B 样条撒到全 batch 共享的静态实空间网格,FFT 后乘影响函数(SPME): - $E = \frac{k_e}{2V}\sum_{m\neq 0}\hat K(|k(m)|)\,|B(m)|^2\,|\hat\rho_{\rm mesh}(m)|^2$, - 其中 $|B(m)|^2 = \prod_d |\sum_l M_p(l+1)e^{2\pi i m_d l/n_d}|^{-2}$ 是样条去卷积 - 修正(构造时按 Essmann 公式计算),$k(m) = 2\pi B^{-1} m$($m$ 为网格模式的整数 - Miller 指标),$m=0$ 置零(tin-foil)。网格尺寸由数据统计的最大 cell 按目标 - 间距 $h$(缺省 $\pi/\_k\_max()$,尊重解耦开关;或 `pme_spacing` 显式)定 - $n_d = \lceil |a_d|/h \rceil$ 并上取 FFT 光滑数(2·3·5·7),或 `pme_mesh` - 显式跳过统计;随模型序列化。大 cell 超出统计 → 有效间距变粗(平滑的精度退化, - 力/位力仍是该近似能量的精确导数,不发 warning——网格尺寸是用户的显式责任)。代价 $O(N p^3 C + M\log M)$ 近线性,无相位矩阵显存墙。 - 性能实现(profile 驱动):$|k|^2$ 走度规二次型 $(2\pi)^2 m^\mathsf{T}W m$ - (6 个静态单项式 × 6 个逐帧系数,全 elementwise),**不构造 k 矢量**—— - batched 3×3 GEMM 曾占训练步 56%;FFT 用 rfftn 半谱($m_3 \ge 0$,内部面加倍 - 权重 ×2、$m_3=0$ 与 Nyquist 面 ×1);f64 下 40³ 网格前向 34→6ms。 - 拓扑 `PME_MESH` 无 ragged 轴,flatten 为 `(box, volume, coord, strain, mask)` - (arity 5,volume 同 KGrid 一样 eager 计算);应变族与直接求和相同(分数坐标 - 应变不变,撒点不带应变;$k'=S^{-T}k$、$V'=V\det S$——strain leaf 的 - `linalg.inv`/`linalg.det` 反向与 KGrid 路径同族,已验证 nf=1 符号帧数下编译 - 反向正确)。eager/compiled 训练与 Python 接口均支持。 - 实现坑(已踩过):$k(m)$ 里 $m$ 是**整数** Miller 指标(不是 $m/n$——那是 - 把网格当成只跨 $1/n_d$ 的 cell);tin-foil 的 `where` 掩码条件别写反 - (`where(keep, influence, 0)`)。 - -### 5.2 图内计算(伪代码,全部纯张量运算,无 data-dependent 控制流) - -```python -# 输入:coord (N,3) [leaf], strain (nf,3,3) [全零 leaf], q (N,C) [电荷网络输出], -# k_vec (nk,3), k_frame (nk,), k_mask (nk,), volume (nf,), -# node_frame (N,) -# 应变 leaf 技巧:位力 = -dE/d(strain),不再需解析位力公式 -S = I + strain -coord' = S[node_frame] @ coord # 每节点 (3,3)@(3,) -k' = inv(S)[k_frame]^{-T} @ k_vec # 每 k 点 (3,3)@(3,) -volume' = volume * det(S) -kr = k' @ coord'.T # (nk, N) -same = (node_frame[None, :] == k_frame[:, None]) # (nk, N) 帧归属 -rho_r = (cos(kr) * same) @ q # (nk, C) segment sum -rho_i = (sin(kr) * same) @ q # (nk, C) -s_sq = (rho_r**2 + rho_i**2).sum(-1) # (nk,) - -# 模式权重(alpha 或 sigma 可训练,梯度照常;按 kernel 类型二选一) -if kernel_type == "les": - K_hat = 4*pi * exp(-k2' / (4*alpha**2)) / k2' # (nk,) - K0 = 2*alpha / sqrt(pi) -else: # "sog" - K_hat = sum_m( amp[m]*(2pi)**1.5*sigma[m]**3 # (nk,),模式循环展开: - * exp(-0.5*sigma[m]**2*k2') ) # 避免 (nk,M) 中间张量 - K0 = amp.sum() # (inductor 在其二次反传 - # 里生成 NaN;LES 不受影响) - -# 能量(自能项只作用于有真实 k 点的帧——周期帧) -E_rec = k_e * segment_sum(k_mask * K_hat * s_sq / (2*volume'[k_frame]), - k_frame, nf) -E_self = -0.5 * k_e * K0 * (q**2).sum(-1) * is_periodic_frame[node_frame] -``` - -- 精度策略(对齐仓库惯例):`rho_r/rho_i` 与 k 求和用 float64 累加, - k 点按 §5.1 固定字典序 → 完全确定性。 -- 力:对 `coord` leaf autograd(框架内 `fit_output_to_model_output_graph` - 的 leaf 列表扩展,与 `edge_vec`/`pair_vec` 并列)。 -- 位力:§3.5 解析式,tensor 实现,`kappa(k)` 对可训练 $\alpha,\sigma$ 可导。 - -### 5.3 图 lower 签名影响(本方案真正动架构的一处) - -现有 graph lower **没有坐标输入**。$e^{i\mathbf{k}\cdot\mathbf{r}_i}$ 需要 -绝对坐标相位。最终实现采用:**`coord` 和零应变 `strain` 都作为 k_grid 拓扑 -的 flatten 输入**(`(frame, vec, mask, volume, coord, strain)`),在 -`forward_common_lower_graph` 里 `.detach().requires_grad_(True)` 重 leaf—— -反向传播必须根植在图输入上(make_fx 内部创建的 leaf 在 inductor 下产出 -NaN)。力由 coord leaf 的梯度直接给出;位力 $W = -\partial E/\partial\varepsilon$ -由 strain leaf 的梯度给出(模块内以 $r'=r(I+\varepsilon)^T$、 -$k'=k(I+\varepsilon)^{-1}$、$V'=V\det(I+\varepsilon)$ 求值)。 -原子位力约定:$r\otimes F$ 部分逐节点,帧余量均摊。 - -## 6. 验证计划 - -1. **M1 对账脚本**(`verify_periodic_lr.py`,独立 torch float64,不接框架)—— - 已实现并全部通过: - - **LES**:k 空间公式 vs 独立双参数 Ewald 分解(实空间 `[erf(αr)−erf(βr)]/r` - 镜像和 + k 空间 β 项 + 自能),diff 1e-9; - - **SOG 分模式 Poisson 恒等式**:每个高斯模式(σ ≤ 10 Å)实空间镜像和 vs - k 空间,diff ~1e-10~1e-11(宽模式以同样的数学自然延伸); - - **SOG 复合对账**:全 k 空间 vs 窄模式(σ≤4 Å)实空间镜像和 + 宽模式 k 空间, - diff 1e-8——注意**不要**用 `[K_sog−erf(βr)/r]` 做参考:该差核在 - $r \gtrsim \sigma_{max}\approx 610$ Å 处残余 $-1/r$ 尾巴,壳层求和无法收敛; - - 力(autograd vs FD)~1e-8,位力(解析式 §3.5 vs 应变 FD)~1e-7; - - k_max 收敛表:两种 kernel 均指数收敛; - - 纯 LES 可用 `deepmd/tf/infer/ewald_recp.py` 交叉验证(其 β = 本方案 α, - 只含倒易项)——需要 C++ TF 算子库构建,未在 M1 中执行。 -2. **M2 框架接入**(已完成):`LatentChargeLongRange` 周期分支 + - k 网格 eager 构建 + `coord`/`strain` 双 leaf;`direct_sum_eligible()` - 不变(压缩路径在 `enable_compression` 拒绝周期模块);导出/冻结拒绝。 - 单测:`source/tests/pt_expt/model/test_long_range.py`(周期组,14 个) - 与 `source/tests/common/dpmodel/test_long_range.py`(k 网格/kernel 原语)。 - 实测(12 原子三斜盒):能量分解 vs M1 参考 2e-7(LES)/1e-6(SOG), - 力 FD 3.6e-7/1.5e-8,位力 FD 3e-7/9.8e-9,编译 == eager 至 1e-13, - Hessian(图路径)vs FD 2e-6。 -3. 拒绝面:周期 + 压缩/CUDA/LAMMPS/非中性通道 → 显式报错。 - -## 8. 实现踩坑记录 - -- **倒易格点约定**:box 行向量是晶格向量 $\mathbf{a}_d$ 时,倒易格点 - $\mathbf{k} = 2\pi\,\mathbf{B}^{-1}\mathbf{n}$(**没有转置**)。写错成 - $2\pi B^{-T}\mathbf{n}$ 在正交/立方盒上碰巧正确,三斜盒上 k 点偏离格点 - 达 0.5·$2\pi/|a|$,造成 0.5 eV 量级的能量错误。验证方法: - `B @ rec / (2*pi) == I`(必须逐元素成立)。 -- **SOG 的 (nk, M) 中间张量**:inductor 对该广播的二次反传生成 NaN - (make_fx 单独正确、LES 不受影响);按模式 Python 循环展开为 (nk,) 项 - 求和后正常。 -- **内部创建的 autograd leaf**:在 make_fx 图内 `new_zeros().requires_grad_()` - 的叶子(strain)在 inductor 下梯度为 NaN;叶子必须作为 flatten 输入传入 - 再在图内 re-leaf(`.detach().requires_grad_()`),与 edge_vec 同模式。 -- **自能项的帧归属**:$-\frac{1}{2}K(0)\sum q^2$ 只作用于有真实 k 点的帧 - (周期帧);非周期批在周期模型上跑时 k_grid 全 guard,自能必须跳过, - 否则能量错 $O(K(0)\sum q^2)$。 -- **FD 位力应变的约定**:坐标与 box 必须同族变形($r\to(I+\varepsilon)r$, - $B\to B(I+\varepsilon)^T$),混用(`d @ box`)破坏旋转不变性。 -- **近截断原子对的 FD Hessian**:原子跨越 rcut 时力有拐点,FD Hessian 在 - 任意步长下都停在 O(0.1) 误差——测试体系必须让所有原子远离 cutoff 边界 - (紧致簇 + 大盒)。 - -## 7. 开放项(不阻塞 v1) - -- 非中性通道的背景自能修正(多阻尼宽度下需选定 regularization); -- 真空表面项配置开关; -- box 在大范围内变化的训练集的 ragged k 轴(v2); -- 压缩路径 / CUDA 周期 direct-sum / LAMMPS 周期部署; -- Hessian 与周期 k 项的兼容性检查。 diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index a535e0c516..4a866538ae 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -480,16 +480,20 @@ extern bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp); extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); /** - * @brief Evaluate a compact canonical graph and the node positions on the - * model device. + * @brief Evaluate a compact canonical graph, the node positions and the cell + * on the model device. * * A model with a long-range module reads the positions of the nodes, which - * must form a complete non-periodic system: held whole by one rank - * (``nall_nodes == nloc``), or decomposed over the ranks of ``world``, whose - * owned nodes ``[0, nloc)`` are the share of this rank. With such a model on - * several ranks, every rank calls this function on every evaluation, also - * when it holds no node. Any other model ignores the positions and the - * communicator. + * must form a complete system: held whole by one rank, or decomposed over the + * ranks of ``world``, whose owned nodes ``[0, nloc)`` are the share of this + * rank. A non-periodic system, which every such model serves, comes without a + * cell, and one rank then holds every node (``nall_nodes == nloc``); a + * periodic system, which a model whose module serves periodic systems + * evaluates over the images of its cell, comes with it, and the halo nodes of + * a single rank are images of its owned nodes. With such a model on several + * ranks, every rank calls this function on every evaluation, also when it + * holds no node. Any other model ignores the positions, the communicator and + * the cell. * * @param[in] dp The DP to use. * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``. @@ -500,6 +504,9 @@ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); * ``(nall_nodes, 3)``, or NULL for a model without a long-range module. * @param[in] world The MPI communicator (a host pointer to an ``MPI_Comm``) * of a system decomposed over several ranks, or NULL on a single rank. + * @param[in] box The cell of a periodic system in host memory, nine values in + * the length unit of the model, row-major with the lattice vectors as rows, + * or NULL for a non-periodic system. * @param[in] d_source Source node per edge storage slot. * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``. * @param[in] d_destination_row_ptr Destination CSR offsets. @@ -509,7 +516,7 @@ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); * @param[in] nall_nodes Total local-plus-halo node count. * @param[in] edge_storage Number of edge storage slots. * @note Supersedes DP_DeepPotComputeCanonicalGraphGPU, adding the node - * positions and the communicator. + * positions, the communicator and the cell. * @since API version 31 */ extern void DP_DeepPotComputeCanonicalGraphGPU2( @@ -520,6 +527,7 @@ extern void DP_DeepPotComputeCanonicalGraphGPU2( const int64_t* d_atype, const double* d_coord, void* world, + const double* box, const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, @@ -531,7 +539,7 @@ extern void DP_DeepPotComputeCanonicalGraphGPU2( /** * @brief Query whether the model carries a long-range module, which - * evaluates complete non-periodic systems. + * evaluates complete systems; see DP_DeepPotComputeCanonicalGraphGPU2. * @since API version 31 */ extern bool DP_DeepPotHasLongRange(DP_DeepPot* dp); @@ -2742,7 +2750,7 @@ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp); /** * @brief Query whether any model of a DP Model Deviation carries a long-range - * module, which evaluates complete non-periodic systems. + * module, which evaluates complete systems. * @param[in] dp The DP Model Deviation to use. * @since API version 31 */ diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index c5153087f3..5ff9a7833e 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1407,8 +1407,8 @@ class DeepPot : public DeepBaseModel { } /** - * @brief Evaluate a compact canonical graph and the node positions on the - * model device; see DP_DeepPotComputeCanonicalGraphGPU2. + * @brief Evaluate a compact canonical graph, the node positions and the + * cell on the model device; see DP_DeepPotComputeCanonicalGraphGPU2. */ void compute_canonical_graph_gpu(double* d_atom_energy, double* d_force, @@ -1416,6 +1416,7 @@ class DeepPot : public DeepBaseModel { const int64_t* d_atype, const double* d_coord, void* world, + const double* box, const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, @@ -1425,7 +1426,7 @@ class DeepPot : public DeepBaseModel { const int nall_nodes, const int64_t edge_storage) { DP_DeepPotComputeCanonicalGraphGPU2( - dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, + dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, box, d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage); DP_CHECK_OK(DP_DeepPotCheckOK, dp); @@ -1433,7 +1434,7 @@ class DeepPot : public DeepBaseModel { /** * @brief Query whether the model carries a long-range module, which - * evaluates complete non-periodic systems. + * evaluates complete systems; see DP_DeepPotComputeCanonicalGraphGPU2. */ bool has_long_range() const { const bool result = DP_DeepPotHasLongRange(dp); @@ -2645,7 +2646,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { /** * @brief Query whether any of the models carries a long-range module, which - * evaluates complete non-periodic systems. + * evaluates complete systems. */ bool has_long_range() const { const bool result = DP_DeepPotModelDeviHasLongRange(dp); diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index e6e2904295..6d0736e08a 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1871,6 +1871,7 @@ void DP_DeepPotComputeCanonicalGraphGPU2(DP_DeepPot* dp, const int64_t* d_atype, const double* d_coord, void* world, + const double* box, const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, @@ -1879,11 +1880,11 @@ void DP_DeepPotComputeCanonicalGraphGPU2(DP_DeepPot* dp, const int nloc, const int nall_nodes, const int64_t edge_storage) { - DP_REQUIRES_OK( - dp, dp->dp.compute_canonical_graph_gpu( - d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, - d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, - d_source_order, nloc, nall_nodes, edge_storage)); + DP_REQUIRES_OK(dp, dp->dp.compute_canonical_graph_gpu( + d_atom_energy, d_force, d_atom_virial, d_atype, + d_coord, world, box, d_source, d_edge_vec, + d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage)); } bool DP_DeepPotHasLongRange(DP_DeepPot* dp) { diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 41b03b4572..c02fbe4365 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -428,6 +428,7 @@ class DeepPotBackend : public DeepBaseModelBackend { const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -901,10 +902,12 @@ class DeepPot : public DeepBaseModel { * @brief Whether the model carries a long-range module. * * Such a model sums its long-range interaction over every pair of atoms of - * a complete non-periodic system, so it evaluates a system without a - * periodic cell, held whole by one rank or, through the compact canonical - * lower, decomposed over the ranks of an MPI communicator; the device - * entries take the positions of its atoms. + * a complete non-periodic system, or over the images of the cell of a + * periodic one when its module serves periodic systems; it evaluates a + * system held whole by one rank or, through the compact canonical lower, + * decomposed over the ranks of an MPI communicator. The device entries + * take the positions of its atoms and the cell, which is NULL for a + * non-periodic system. */ bool has_long_range() const; @@ -932,13 +935,18 @@ class DeepPot : public DeepBaseModel { * * @param[in] d_coord Node positions in model-node order, shape * ``(nall_nodes, 3)``, in the length unit of the model. A model with a - * long-range module requires them, for a system without a periodic cell. + * long-range module requires them. * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of a * system decomposed over several ranks, of which the owned nodes * ``[0, nloc)`` are the share of this rank, or NULL on a single rank, - * whose node set then holds the complete system (``nall_nodes == nloc``). - * With a long-range module and several ranks, every rank calls this - * function on every evaluation, also when it holds no node. + * whose node set then holds the complete system (``nall_nodes == nloc`` + * without a cell). With a long-range module and several ranks, every rank + * calls this function on every evaluation, also when it holds no node. + * @param[in] box The cell of a periodic system in host memory, nine values + * in the length unit of the model, row-major with the lattice vectors as + * rows, or NULL for a non-periodic system. A model with a long-range module + * that serves periodic systems reads it; the halo nodes of a single rank + * are then images of the owned nodes. * * The other parameters are those of the overload without positions. */ @@ -948,6 +956,7 @@ class DeepPot : public DeepBaseModel { const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -1011,7 +1020,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { /** * @brief Whether any of the models carries a long-range module, which - * confines all of them to complete non-periodic systems. + * confines all of them to complete systems (see DeepPot::has_long_range). **/ bool has_long_range() const { return std::any_of(dps.begin(), dps.end(), @@ -1195,5 +1204,14 @@ class DeepPotModelDevi : public DeepBaseModelDevi { protected: std::vector> dps; + + private: + /** + * @brief Check that every model accepts the cell handed to the ensemble. + * Without one, the neighbor-list entries serve a model with a long-range + * module beside models that take the cell of the region. + **/ + template + void check_ensemble_box(const std::vector& dbox) const; }; } // namespace deepmd diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index cfa0a551ca..bc9b685021 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -379,6 +379,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -439,6 +440,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -478,11 +480,15 @@ class DeepPotPTExpt : public DeepPotBackend { bool lower_input_is_graph_ = false; bool lower_input_is_canonical_ = false; bool graph_edge_fp32_ = false; - // The model carries a long-range module, whose lower reads the all-pairs - // topology of the frame after its other inputs: the pair lists for the graph - // lower, the node positions and the communicator of the ranks for the - // compact canonical lower. + // The model carries a long-range module, whose lower reads the topologies + // of the frame after its other inputs: the pair lists, and the mesh of the + // cell of a module that serves periodic systems, for the graph lower; the + // node positions, the communicator of the ranks and the cell for the compact + // canonical lower. bool long_range_ = false; + // The module serves periodic systems, on a mesh of at most this spacing. + bool long_range_periodic_ = false; + double mesh_spacing_ = 0.0; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) std::vector mapping_; // cached mapping vector (LAMMPS path) @@ -626,8 +632,10 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] source_order Source-grouped edge permutation ``(E,)``. * @param[in] source_row_ptr Source CSR offsets ``(N + 1,)``. * @param[in] coord Positions of the nodes of the single frame, shape - * ``(N, 3)`` float64, from which the all-pairs topology of a model with a - * long-range module is built; read only for such a model. + * ``(N, 3)`` float64, from which the topologies of a model with a + * long-range module are built; read only for such a model. + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one; read only for such a model. */ std::vector run_model_graph( const torch::Tensor& atype, @@ -643,7 +651,8 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& fparam, const torch::Tensor& aparam, const torch::Tensor& charge_spin, - const torch::Tensor& coord); + const torch::Tensor& coord, + const std::vector& box); /** * @brief Run a compact canonical ``.pt2`` artifact. @@ -653,6 +662,9 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of the * ranks a decomposed system spans, or NULL for a system that one rank holds * whole; a model with a long-range module receives it after the positions. + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one; a model with a long-range module receives + * it after the communicator. */ std::vector run_model_canonical_graph( const torch::Tensor& atype, @@ -664,16 +676,28 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& source_row_ptr, const torch::Tensor& source_order, const torch::Tensor& coord, - const void* world); + const void* world, + const std::vector& box); /** * @brief Evaluate the long-range module on a rank that holds no node. * * The compact lower runs on one halo node without edges, which carries no * charge and contributes nothing, so that the rank joins the exchanges of - * the module with the other ranks of ``world``. + * the module with the other ranks of ``world``, in the cell ``box`` of the + * system, empty for a non-periodic one. + */ + void join_long_range_exchange(const void* world, + const std::vector& box); + + /** + * @brief Check the cell given to a model with a long-range module. + * + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one, which every long-range module serves; a + * cell requires a module that serves periodic systems. */ - void join_long_range_exchange(const void* world); + void check_long_range_box(const std::vector& box) const; /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index b8491393ee..d3ec820acf 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -663,6 +663,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -677,6 +678,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( (void)d_atype; (void)d_coord; (void)world; + (void)box; (void)d_source; (void)d_edge_vec; (void)d_destination_row_ptr; @@ -766,8 +768,9 @@ void DeepPot::compute_canonical_graph_gpu( const std::int64_t edge_storage) { dp->compute_canonical_graph_gpu( d_atom_energy, d_force, d_atom_virial, d_atype, /*d_coord=*/nullptr, - /*world=*/nullptr, d_source, d_edge_vec, d_destination_row_ptr, - d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage); + /*world=*/nullptr, /*box=*/nullptr, d_source, d_edge_vec, + d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, + edge_storage); } void DeepPot::compute_canonical_graph_gpu( @@ -777,6 +780,7 @@ void DeepPot::compute_canonical_graph_gpu( const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -786,9 +790,9 @@ void DeepPot::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { dp->compute_canonical_graph_gpu( - d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, d_source, - d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, - nall_nodes, edge_storage); + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, box, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage); } bool DeepPot::uses_fp32_edge_vectors() const { @@ -950,6 +954,25 @@ template void DeepPotModelDevi::compute( const std::vector& aparam, const std::vector& charge_spin); +// LAMMPS hands a model with a long-range module no cell for a non-periodic +// system, whereas the neighbor-list entries of other models, as of the +// TensorFlow backend, read the cell of the region. +template +void DeepPotModelDevi::check_ensemble_box( + const std::vector& dbox) const { + if (!dbox.empty() || !has_long_range()) { + return; + } + for (const auto& model : dps) { + if (!model->has_long_range()) { + throw deepmd::deepmd_exception( + "an ensemble that holds a model with a long-range module evaluates " + "a non-periodic system only when every model carries one; use " + "models with long-range modules throughout, or a periodic system."); + } + } +} + template void DeepPotModelDevi::compute(std::vector& all_energy, std::vector>& all_force, @@ -966,6 +989,7 @@ void DeepPotModelDevi::compute(std::vector& all_energy, if (numb_models == 0) { return; } + check_ensemble_box(dbox); all_energy.resize(numb_models); all_force.resize(numb_models); all_virial.resize(numb_models); @@ -1028,6 +1052,7 @@ void DeepPotModelDevi::compute( all_virial.resize(numb_models); all_atom_energy.resize(numb_models); all_atom_virial.resize(numb_models); + check_ensemble_box(dbox); for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], dcoord_, datype_, diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index a76cd96fa5..751d793830 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -121,6 +122,18 @@ at::Tensor communicator_tensor(const void* world) { torch::TensorOptions().dtype(torch::kInt64)); } +// The cell input of the compact lower of a long-range model: a float64 CPU +// tensor of shape (1, 9), the rows of the row-major cell its lattice vectors, +// zero for a non-periodic frame. +at::Tensor cell_tensor(const std::vector& box) { + auto cell = + torch::zeros({1, 9}, torch::TensorOptions().dtype(torch::kFloat64)); + if (!box.empty()) { + std::copy(box.begin(), box.begin() + 9, cell.data_ptr()); + } + return cell; +} + } // namespace void DeepPotPTExpt::translate_error(std::function f) { @@ -268,16 +281,24 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_canonical_ = false; } // A model with a long-range module reads the all-pairs topology of the - // frame after its other inputs. + // frame after its other inputs, and, when the module serves periodic + // systems, the mesh of the cell, sized by the spacing of the metadata. long_range_ = metadata.obj_val.count("topologies") > 0; + long_range_periodic_ = false; if (long_range_) { for (const auto& name : metadata["topologies"].as_array()) { - if (name.as_string() != "all_pairs") { + if (name.as_string() == "pme_mesh") { + long_range_periodic_ = true; + } else if (name.as_string() != "all_pairs") { throw deepmd::deepmd_exception("unknown topology '" + name.as_string() + "' in the .pt2 metadata; expected " - "'all_pairs'."); + "'all_pairs' or 'pme_mesh'."); } } + if (long_range_periodic_) { + mesh_spacing_ = + metadata["topology_params"]["pme_mesh"]["spacing"].as_double(); + } if (!lower_input_is_graph_ && !lower_input_is_canonical_) { throw deepmd::deepmd_exception( "a model with a long-range module requires the graph or the " @@ -542,7 +563,8 @@ std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& fparam, const torch::Tensor& aparam, const torch::Tensor& charge_spin, - const torch::Tensor& coord) { + const torch::Tensor& coord, + const std::vector& box) { // Graph-input ABI: original edge payload plus destination/source CSR views. check_graph_aparam_flat(aparam, daparam, "run_model_graph"); std::vector inputs = { @@ -559,10 +581,19 @@ std::vector DeepPotPTExpt::run_model_graph( inputs.push_back(charge_spin); } if (long_range_) { - const auto pairs = deepmd::ptexpt::buildAllPairsTopology( - coord, atype, pair_exclude_table_, ntypes, - graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64); + // A periodic frame holds no pair; the mesh describes its interactions. + const auto vec_dtype = graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64; + const auto pairs = + box.empty() + ? deepmd::ptexpt::buildAllPairsTopology( + coord, atype, pair_exclude_table_, ntypes, vec_dtype) + : deepmd::ptexpt::buildGuardPairTopology(coord.device(), vec_dtype); inputs.insert(inputs.end(), pairs.begin(), pairs.end()); + if (long_range_periodic_) { + const auto mesh = deepmd::ptexpt::buildMeshTopology( + coord, box, mesh_spacing_, vec_dtype); + inputs.insert(inputs.end(), mesh.begin(), mesh.end()); + } } return loader->run(inputs); } @@ -577,7 +608,8 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( const torch::Tensor& source_row_ptr, const torch::Tensor& source_order, const torch::Tensor& coord, - const void* world) { + const void* world, + const std::vector& box) { std::vector inputs = {atype, n_node, n_local, source, edge_vec, destination_row_ptr, @@ -585,11 +617,48 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( if (long_range_) { inputs.push_back(coord); inputs.push_back(communicator_tensor(world)); + inputs.push_back(cell_tensor(box)); } return loader->run(inputs); } -void DeepPotPTExpt::join_long_range_exchange(const void* world) { +void DeepPotPTExpt::check_long_range_box(const std::vector& box) const { + if (!box.empty() && box.size() != 9) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates one frame with a cell of " + "nine values, or without one, but got " + + std::to_string(box.size()) + " values."); + } + if (!box.empty() && !long_range_periodic_) { + throw deepmd::deepmd_exception( + "the long-range module of this model serves complete non-periodic " + "systems only; pass no cell (in LAMMPS, a non-periodic boundary f, s " + "or m in every dimension), or train the module with periodic: true."); + } + if (!box.empty()) { + // A cell whose volume is a vanishing fraction of the product of the + // lengths of its lattice vectors, or not finite, spans no lattice. + const double volume = + std::abs(box[0] * (box[4] * box[8] - box[5] * box[7]) - + box[1] * (box[3] * box[8] - box[5] * box[6]) + + box[2] * (box[3] * box[7] - box[4] * box[6])); + double lengths = 1.0; + for (int d = 0; d < 3; ++d) { + lengths *= + std::sqrt(box[3 * d] * box[3 * d] + box[3 * d + 1] * box[3 * d + 1] + + box[3 * d + 2] * box[3 * d + 2]); + } + if (!(volume > 1e-6 * lengths)) { + throw deepmd::deepmd_exception( + "the cell of a periodic system needs three independent lattice " + "vectors, but its volume is " + + std::to_string(volume) + "; pass no cell for a non-periodic system."); + } + } +} + +void DeepPotPTExpt::join_long_range_exchange(const void* world, + const std::vector& box) { const torch::Device device(torch::kCUDA, gpu_id); const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); const auto index_type = deepmd::canonicalGraphIndexType(); @@ -600,7 +669,7 @@ void DeepPotPTExpt::join_long_range_exchange(const void* world) { torch::zeros({2}, i64.dtype(index_type)), torch::zeros({2, 3}, i64.dtype(torch::kFloat32)), torch::zeros({2}, i64), torch::zeros({2}, i64), torch::arange(2, i64).to(index_type), - torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world); + torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world, box); } std::vector DeepPotPTExpt::run_model_with_comm( @@ -857,17 +926,23 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; - // A long-range module sums over every pair of atoms of the system. A single - // rank holds all of them, so a ghost atom there is a periodic image. Across - // ranks, the compact lower gathers the owned atoms of every rank through - // the communicator; the pairs that the graph lower lists stay within one. - if (long_range_ && !multi_rank && nghost_real > 0) { + // A long-range module sums over every pair of atoms of a non-periodic + // system, which the caller marks with an empty box, and over the images of + // the cell of a periodic one. A single rank holds every atom, so a ghost + // atom there is a periodic image. Across ranks, the compact lower gathers + // the owned atoms of every rank through the communicator; the pairs that + // the graph lower lists stay within one. + const std::vector long_range_box(box.begin(), box.end()); + if (long_range_) { + check_long_range_box(long_range_box); + } + if (long_range_ && long_range_box.empty() && !multi_rank && nghost_real > 0) { throw deepmd::deepmd_exception( "a model with a long-range module evaluates a complete non-periodic " - "system, but this single-rank call carries " + + "system without a box, but this single-rank call carries " + std::to_string(nghost_real) + - " ghost atoms, the periodic images of a periodic cell; use a " - "non-periodic box."); + " ghost atoms, the periodic images of a periodic cell; pass the cell " + "of a periodic system to a model whose module serves one."); } if (long_range_ && multi_rank && !lower_input_is_canonical_) { throw deepmd::deepmd_exception( @@ -875,7 +950,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, std::to_string(lmp_list.nprocs) + " MPI ranks only through the compact canonical lower, which a " "compressed DPA4C model frozen for a CUDA device uses when the " - "direct-sum operator serves its long-range module; this archive uses " + "long-range operator serves its long-range module; this archive uses " "another lower, so run it on a single rank."); } if (long_range_ && multi_rank && lmp_list.world == nullptr) { @@ -1317,7 +1392,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // long-range module across ranks still joins the exchanges of the // others. if (long_range_ && multi_rank) { - join_long_range_exchange(world); + join_long_range_exchange(world, long_range_box); } ener.assign(nframes, static_cast(0)); force.assign(static_cast(nframes) * fwd_map.size() * 3, @@ -1366,14 +1441,15 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, - compact.source_row_ptr, compact.source_order, node_coord, world); + compact.source_row_ptr, compact.source_order, node_coord, world, + long_range_box); } else { flat_outputs = run_model_graph( node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor, node_coord); + graph_aparam, charge_spin_tensor, node_coord, long_range_box); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -1549,12 +1625,6 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, throw deepmd::deepmd_exception( "The coordinate-based inference interface requires at least one atom."); } - // The long-range module sums over the atoms of the system alone. - if (long_range_ && !box.empty()) { - throw deepmd::deepmd_exception( - "a model with a long-range module evaluates complete non-periodic " - "systems only; pass an empty box."); - } const std::size_t frame_coord_size = static_cast(natoms) * 3; if (coord.empty() || coord.size() % frame_coord_size != 0) { throw deepmd::deepmd_exception( @@ -1572,6 +1642,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, check_call_charge_spin(charge_spin, /*nframes=*/1, settable_chgspin, dchgspin > 0, default_chg_spin_, chg_spin_table_ranges_); + // A long-range module sums over the atoms of a non-periodic system, which + // the empty box marks, and over the images of the cell of a periodic one. + const std::vector long_range_box(box.begin(), box.end()); + if (long_range_) { + check_long_range_box(long_range_box); + } // The .pt2 model only contains forward_common_lower, which requires // nlist as input. We must build the nlist in C++ and fold back the // extended-region outputs to local atoms. @@ -1754,7 +1830,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, - compact.source_order, node_coord, /*world=*/nullptr); + compact.source_order, node_coord, /*world=*/nullptr, long_range_box); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -1766,7 +1842,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // normalizes the rectangular runtime aparam to the flat // (N, daparam) node axis of the graph ABI. extend_graph_aparam(aparam_tensor, natoms, natoms, daparam), - charge_spin_tensor, node_coord); + charge_spin_tensor, node_coord, long_range_box); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -2326,9 +2402,10 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, throw deepmd::deepmd_exception( "compute_edges_gpu: a model with a long-range module evaluates a " "complete non-periodic system on a single rank (nall_nodes == nloc and " - "no comm_nlist); several ranks require the compact canonical lower, " - "which a compressed DPA4C model frozen for a CUDA device uses when " - "the direct-sum operator serves its long-range module."); + "no comm_nlist); periodic systems and several ranks require the " + "compact canonical lower, which a compressed DPA4C model frozen for a " + "CUDA device uses when the long-range operator serves its module, or " + "the evaluation from a neighbor list (pair_style deepmd)."); } if (lower_input_is_graph_ && extended && has_message_passing_) { throw deepmd::deepmd_exception( @@ -2557,7 +2634,8 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor, coord_t.reshape({nloc, 3}))); + graph_aparam, charge_spin_tensor, coord_t.reshape({nloc, 3}), + /*box=*/{})); ae = out["atom_energy"].reshape({nnode}).slice(0, 0, nloc).contiguous(); force_t = out["force"].reshape({nnode, 3}).contiguous(); av = out["atom_virial"].reshape({nnode, 9}).contiguous(); @@ -2648,6 +2726,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2664,25 +2743,32 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a CUDA device."); } - // A long-range module reads the positions of a complete non-periodic - // system: the node set of a single rank, or the owned nodes of every rank - // of ``world``, which a rank without nodes still joins. + // A long-range module reads the positions of a complete system and the + // cell of a periodic one: the node set of a single rank, whose halo nodes + // are images of a periodic cell, or the owned nodes of every rank of + // ``world``, which a rank without nodes still joins. + const std::vector long_range_box = + box != nullptr ? std::vector(box, box + 9) + : std::vector(); + if (long_range_) { + check_long_range_box(long_range_box); + } if (long_range_ && d_coord == nullptr && nall_nodes > 0) { throw deepmd::deepmd_exception( "compute_canonical_graph_gpu: a model with a long-range module reads " "the node positions (d_coord)."); } - if (long_range_ && world == nullptr && nall_nodes != nloc) { + if (long_range_ && world == nullptr && box == nullptr && nall_nodes != nloc) { throw deepmd::deepmd_exception( "compute_canonical_graph_gpu: a model with a long-range module on a " "single rank evaluates a complete non-periodic system " - "(nall_nodes == nloc); a system decomposed over several ranks takes " - "their communicator."); + "(nall_nodes == nloc), or a periodic one with its cell; a system " + "decomposed over several ranks takes their communicator."); } if (long_range_ && world != nullptr && nall_nodes == 0) { translate_error([&] { const c10::DeviceGuard device_guard(torch::Device(torch::kCUDA, gpu_id)); - join_long_range_exchange(world); + join_long_range_exchange(world, long_range_box); synchronize_current_accelerator_stream(); }); return; @@ -2730,7 +2816,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( extract_outputs(output, run_model_canonical_graph( atype, n_node, n_local, source, edge_vec, destination_row_ptr, source_row_ptr, - source_order, coord, world)); + source_order, coord, world, long_range_box)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2814,6 +2900,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( const std::int64_t* d_atype, const double* d_coord, void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2823,9 +2910,9 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { compute_canonical_graph_gpu_impl( - d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, d_source, - d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, - nall_nodes, edge_storage); + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, box, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage); } bool DeepPotPTExpt::uses_fp32_edge_vectors() const { diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h index cbb37f49f9..9e5e05faa7 100644 --- a/source/api_cc/src/longRangePTExpt.h +++ b/source/api_cc/src/longRangePTExpt.h @@ -3,6 +3,9 @@ #include +#include +#include +#include #include #include "commonPT.h" @@ -55,5 +58,102 @@ inline std::vector buildAllPairsTopology( torch::cat({mask, torch::zeros({2}, mask_options)})}; } +/** + * @brief Build the all-pairs topology of a periodic frame, which holds no + * pair: the reciprocal-space topology describes the interactions. + * + * Twin of Python ``empty_pair_graph``: two masked guard pairs. + */ +inline std::vector buildGuardPairTopology( + const torch::Device& device, const torch::ScalarType vec_dtype) { + const auto options = torch::TensorOptions().device(device); + return {torch::zeros({2, 2}, options.dtype(torch::kInt64)), + torch::zeros({2, 3}, options.dtype(vec_dtype)), + torch::zeros({2}, options.dtype(torch::kBool))}; +} + +/** + * @brief The smallest even integer of at least ``n`` without a prime factor + * above 7. + */ +inline std::int64_t evenFftSize(std::int64_t n) { + std::int64_t m = std::max(2, n + n % 2); + while (true) { + std::int64_t rest = m; + for (const std::int64_t p : {2, 3, 5, 7}) { + while (rest % p == 0) { + rest /= p; + } + } + if (rest == 1) { + return m; + } + m += 2; + } +} + +/** + * @brief Build the mesh topology a periodic long-range module of the graph + * lower reads. + * + * Twin of Python ``build_mesh_grid`` in + * ``deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py`` for a single + * frame: along every lattice vector ``a_d`` the mesh holds the smallest even + * number of points without a prime factor above 7 of at least + * ``ceil(|a_d| / spacing)``, with the length formed term by term as Python + * forms it, so that both size the mesh identically. Without a cell, the frame + * is non-periodic: the cell is the identity, the mesh minimal and the frame + * masked. + * + * @param coord Node positions, shape ``(N, 3)``, on the model device. + * @param box The row-major cell, the rows its lattice vectors, nine values in + * the length unit of the model, or empty for a non-periodic frame. + * @param spacing The largest mesh spacing, in the length unit of the model. + * @param dtype Precision of the geometry, that of the edge vectors. + * @return ``{box (1, 3, 3), volume (1,), coord (N, 3), strain (1, 3, 3), + * mask (1,) bool, mode_1 (n1,), mode_2 (n2,), mode_3 (n3,)}``. + */ +inline std::vector buildMeshTopology( + const torch::Tensor& coord, + const std::vector& box, + const double spacing, + const torch::ScalarType dtype) { + const auto options = + torch::TensorOptions().dtype(dtype).device(coord.device()); + const bool periodic = !box.empty(); + std::array cell = {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0}; + std::array dims = {2, 2, 2}; + if (periodic) { + for (int k = 0; k < 9; ++k) { + cell[k] = box[k]; + } + for (int d = 0; d < 3; ++d) { + const double* v = cell.data() + 3 * d; + const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + dims[d] = + evenFftSize(static_cast(std::ceil(length / spacing))); + } + } + const double volume = + std::abs(cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) - + cell[1] * (cell[3] * cell[8] - cell[5] * cell[6]) + + cell[2] * (cell[3] * cell[7] - cell[4] * cell[6])); + // The copy detaches the tensor from the local array, which a conversion + // to float64 on the host would alias. + const auto cell_tensor = + torch::from_blob(cell.data(), {1, 3, 3}, torch::kFloat64) + .to(options, /*non_blocking=*/false, /*copy=*/true); + std::vector fields = { + cell_tensor, torch::full({1}, volume, options), coord.to(dtype), + torch::zeros({1, 3, 3}, options), + torch::full({1}, periodic, options.dtype(torch::kBool))}; + for (const std::int64_t n : dims) { + // The integer frequencies of an axis of n points, in FFT order. + const auto m = torch::arange(n, options); + fields.push_back(torch::where(m > n / 2, m - n, m)); + } + return fields; +} + } // namespace ptexpt } // namespace deepmd diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc index fa4df07ad4..7ca2b916f2 100644 --- a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -1,10 +1,12 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // Test C++ inference for pt_expt (.pt2) DPA4C models with a latent-charge // long-range module, which sums over every pair of atoms of a complete -// non-periodic system. Models generated by -// source/tests/infer/gen_dpa4c_long_range.py: an uncompressed archive on the -// graph lower, which receives the pair list of the atoms, and a compressed -// archive on the compact canonical lower, which reads their positions. +// non-periodic system, and over the images of the cell of a periodic one when +// the module serves periodic systems. Models generated by +// source/tests/infer/gen_dpa4c_long_range.py: uncompressed archives on the +// graph lower, which receives the pair list or the mesh of the cell, and +// compressed archives on the compact canonical lower, which reads the +// positions of the atoms and the cell. #include #include @@ -26,6 +28,14 @@ constexpr const char* kCompressedModelPath = "../../tests/infer/deeppot_dpa4c_long_range_compressed.pt2"; constexpr const char* kCompressedRefPath = "../../tests/infer/deeppot_dpa4c_long_range_compressed.expected"; +constexpr const char* kPeriodicModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic.pt2"; +constexpr const char* kPeriodicRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic.expected"; +constexpr const char* kPeriodicCompressedModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.pt2"; +constexpr const char* kPeriodicCompressedRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.expected"; // The full neighbor list of a non-periodic system within ``rc``: every atom // is local, so the list holds no ghost. @@ -176,8 +186,7 @@ TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, lmp_nlist_explicit) { this->ref_explicit.check(ener, force, virial, epsilon); } -// A cell would leave the pair sum incomplete: its periodic images are no -// part of it. +// A module that serves no periodic system rejects a cell. TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, periodic_box_is_rejected) { using VALUETYPE = TypeParam; const std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; @@ -251,3 +260,111 @@ TYPED_TEST(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, this->compute_lmp(ener, force, virial); this->ref_default.check(ener, force, virial, epsilon); } + +// A module that serves periodic systems evaluates a periodic system over the +// images of its cell, through the coordinates and the cell or through the +// neighbor list of a LAMMPS-style call, whose ghost atoms are images, and +// keeps the direct sum of a system without a cell. ``Archive`` selects the +// uncompressed archive on the graph lower or, on a CUDA device, the +// compressed one on the compact canonical lower. +template +class TestInferDeepPotDpa4cLongRangePeriodicPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13.0, 0.0, 0.0, 0.6, 12.5, + 0.0, -0.4, 0.8, 13.5}; + LongRangeRef ref_periodic; + LongRangeRef ref_cluster; + deepmd::DeepPot dp; + + void load(const char* model_path, const char* ref_path) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(model_path); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << model_path + << " was not generated (gen_dpa4c_long_range.py writes the " + "compressed archives on a CUDA device only)."; + } + dp.init(model_path); + deepmd_test::ExpectedRef ref; + ref.load(ref_path); + ref_periodic.load(ref, "periodic"); + ref_cluster.load(ref, "cluster"); + } + + // Computes the periodic system through the LAMMPS-style entry point, with + // the ghost atoms of the cell, and folds their forces back onto their + // owners. + void compute_lmp_periodic(double& ener, + std::vector& force, + std::vector& virial) { + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector> nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, static_cast(dp.cutoff())); + const int nloc = static_cast(atype.size()); + const int nall = static_cast(atype_cpy.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + inlist.mapping = mapping.data(); + std::vector force_all; + dp.compute(ener, force_all, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + _fold_back(force, force_all, mapping, nloc, nall, 3); + } + + void check_all(const double epsilon) { + EXPECT_TRUE(dp.has_long_range()); + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + ref_periodic.check(ener, force, virial, epsilon); + compute_lmp_periodic(ener, force, virial); + ref_periodic.check(ener, force, virial, epsilon); + dp.compute(ener, force, virial, coord, atype, std::vector()); + ref_cluster.check(ener, force, virial, epsilon); + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, graph_lower) { + using VALUETYPE = TypeParam; + this->load(kPeriodicModelPath, kPeriodicRefPath); + if (this->IsSkipped()) { + return; + } + this->check_all(std::is_same::value ? 1e-8 : 1e-4); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, compact_lower) { + this->load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath); + if (this->IsSkipped()) { + return; + } + this->check_all(1e-4); +} + +// A cell without volume marks no system: it is rejected rather than taken as +// a non-periodic one. +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, + degenerate_cell_is_rejected) { + using VALUETYPE = TypeParam; + this->load(kPeriodicModelPath, kPeriodicRefPath); + if (this->IsSkipped()) { + return; + } + double ener; + std::vector force, virial; + EXPECT_THROW(this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector(9, 0.)), + deepmd::deepmd_exception); +} diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 7dd715f48c..2ed86f2da2 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -178,6 +178,11 @@ double PairDeepMD::eval_energy_with_fparam( dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + // A long-range module reads the cell of a periodic system only; the empty + // box marks a non-periodic one. + if (long_range && domain->nonperiodic) { + dbox.clear(); + } // get coord for (int ii = 0; ii < nall; ++ii) { @@ -297,6 +302,11 @@ void PairDeepMD::compute(int eflag, int vflag) { dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + // A long-range module reads the cell of a periodic system only; the empty + // box marks a non-periodic one. + if (long_range && domain->nonperiodic) { + dbox.clear(); + } // get coord for (int ii = 0; ii < nall; ++ii) { @@ -968,17 +978,19 @@ void PairDeepMD::settings(int narg, char** arg) { void PairDeepMD::init_style() { PairDeepBaseModel::init_style(); - // A long-range module sums over every pair of atoms of the system, which a - // periodic image would leave incomplete; the atoms of the other ranks reach - // it through the communicator. - const bool long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() - : deep_pot.has_long_range(); - if (long_range && - (domain->xperiodic || domain->yperiodic || domain->zperiodic)) { + // A long-range module sums over every pair of atoms of a non-periodic + // system, or over the images of the cell of a system periodic in every + // dimension; the atoms of the other ranks reach it through the + // communicator. The model checks which of the two its module serves. + long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() + : deep_pot.has_long_range(); + const int periodic_dimensions = + domain->xperiodic + domain->yperiodic + domain->zperiodic; + if (long_range && periodic_dimensions != 0 && periodic_dimensions != 3) { error->all(FLERR, "pair style deepmd: a model with a long-range module evaluates " - "a complete non-periodic system; use a non-periodic boundary " - "(f, s or m) in every dimension."); + "a system periodic in every dimension or in none; use p in " + "every dimension, or f, s or m in every dimension."); } } diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index ac8881b553..69466f740a 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -59,6 +59,9 @@ class PairDeepMD : public PairDeepBaseModel { // geometry fields are unused) for the device-resident message-passing path, // where ghost features are exchanged across ranks inside the forward pass. deepmd_compat::InputNlist make_comm_nlist(); + // The model carries a long-range module, which reads the cell of a + // periodic system and no cell for a non-periodic one. + bool long_range = false; private: CommBrickDeepMD* commdata_; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 9c24685e1a..cd9f102fff 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -2,6 +2,7 @@ #ifdef LMP_KOKKOS #include "pair_deepmd_kokkos.h" +#include #include #include #include @@ -34,7 +35,6 @@ PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) edge_capacity(0), edge_vec_fp32(false), canonical_graph(false), - long_range(false), device_path_ok(false), reverse_virial(false), reverse_used_host(false) { @@ -213,10 +213,14 @@ void PairDeepMDKokkos::init_style() { request->enable_full(); edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors(); canonical_graph = deep_pot.uses_canonical_graph_inference(); - // PairDeepMD::init_style() has confined such a model to a complete - // non-periodic system, whose atoms on other ranks it reaches through the - // communicator. - long_range = deep_pot.has_long_range(); + // The device edge path passes a long-range module no cell, so a periodic + // system reaches such a module through the compact canonical lower alone. + if (long_range && !canonical_graph && !domain->nonperiodic) { + error->all(FLERR, + "pair style deepmd/kk evaluates a periodic system with a " + "long-range module only through a compressed DPA4C model frozen " + "for a CUDA device; use pair_style deepmd for this model."); + } // Force exchange transfers three values per atom. Centroid per-atom virial // uses nine values even though the Kokkos full-list request runs newton off, // so comm_reverse_off reserves the classic host buffer for the wider mode. @@ -417,11 +421,6 @@ template const double* PairDeepMDKokkos::model_node_positions() { atomKK->sync(execution_space, X_MASK); auto x = atomKK->k_x.template view(); - // Without virtual atoms the model nodes are the leading local atoms, already - // in the length unit of the model. - if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0) { - return x.data(); - } const int nnode_model = compact_graph.nnode_model; if ((int)d_node_position.extent(0) < 3 * nnode_model) { d_node_position = Kokkos::View( @@ -429,15 +428,21 @@ const double* PairDeepMDKokkos::model_node_positions() { } auto node_position = d_node_position; auto model2loc = compact_graph.d_model2loc; + const bool has_null_types = compact_graph.has_null_types; const double inv_dist = 1.0 / dist_unit_cvt_factor; + // Relative to the origin of the box, as the plain pair style passes them. + const double x0 = domain->boxlo[0]; + const double y0 = domain->boxlo[1]; + const double z0 = domain->boxlo[2]; Kokkos::parallel_for( "deepmd/kk:node_position", Kokkos::RangePolicy(0, nnode_model), KOKKOS_LAMBDA(const int m) { - const int i = model2loc(m); - node_position(3 * m + 0) = x(i, 0) * inv_dist; - node_position(3 * m + 1) = x(i, 1) * inv_dist; - node_position(3 * m + 2) = x(i, 2) * inv_dist; + // Without virtual atoms the model nodes are the leading atoms. + const int i = has_null_types ? model2loc(m) : m; + node_position(3 * m + 0) = (x(i, 0) - x0) * inv_dist; + node_position(3 * m + 1) = (x(i, 1) - y0) * inv_dist; + node_position(3 * m + 2) = (x(i, 2) - z0) * inv_dist; }); return node_position.data(); } @@ -573,6 +578,17 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { long_range ? model_node_positions() : (compact_graph.has_null_types ? d_coord_model.data() : x.data()); + // A long-range module reads the cell of a periodic system, row-major with + // the lattice vectors as rows, in the length unit of the model, and no + // cell for a non-periodic one. + const double* h = domain->h; + std::array cell = {h[0], 0.0, 0.0, h[5], h[1], + 0.0, h[4], h[3], h[2]}; + for (double& entry : cell) { + entry /= dist_unit_cvt_factor; + } + const double* cell_ptr = + long_range && !domain->nonperiodic ? cell.data() : nullptr; Kokkos::fence(); try { if (canonical_graph) { @@ -580,7 +596,8 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), compact_graph.d_model_type.data(), long_range ? coord_ptr : nullptr, long_range_across_ranks ? static_cast(&world) : nullptr, - compact_graph.d_source.data(), compact_graph.d_edge_vec.data(), + cell_ptr, compact_graph.d_source.data(), + compact_graph.d_edge_vec.data(), compact_graph.d_destination_row_ptr.data(), compact_graph.d_source_row_ptr.data(), compact_graph.d_source_order.data(), nloc_m, nnode_m, diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index 65a1ce315f..4db0f380e6 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -82,8 +82,9 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); - // Positions of the model nodes in the length unit of the model, which a - // long-range module reads; public for the same reason. + // Positions of the model nodes relative to the origin of the box, in the + // length unit of the model, which a long-range module reads; public for + // the same reason. const double* model_node_positions(); protected: @@ -127,7 +128,6 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { int edge_capacity; // allocated edges in d_edge_index / d_edge_vec bool edge_vec_fp32; // model graph ABI consumes edge vectors in fp32 bool canonical_graph; // compact source-only graph artifact - bool long_range; // model reads the positions of every node bool device_path_ok; // resolved once in init_style bool reverse_virial; // reverse communication operates on centroid virial bool reverse_used_host; // force reverse communication selected host staging diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py index 5241bd8e5e..4829a91a7d 100644 --- a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -2,13 +2,16 @@ """Test LAMMPS with a .pt2 DPA4C model that carries a long-range module. The latent-charge module sums over every pair of atoms of a complete -non-periodic system, so ``pair_style deepmd`` runs it with non-periodic -boundaries and rejects a periodic one. A compressed model also runs on -several ranks, whose atoms it gathers over MPI; an uncompressed one runs on a -single rank. The models and the reference values, without a cell, at the -default charge state and at ``charge_spin 1.0 2.0``, are generated by -source/tests/infer/gen_dpa4c_long_range.py, the compressed model for CUDA -only. +non-periodic system, and over the images of the cell of a periodic one when +it serves periodic systems, so ``pair_style deepmd`` runs a model whose module +serves no periodic system with non-periodic boundaries only, and any +long-range model with boundaries periodic in every dimension or in none. A +compressed model also runs on several ranks, whose atoms it gathers over MPI; +an uncompressed one runs on a single rank. The models and the reference +values are generated by source/tests/infer/gen_dpa4c_long_range.py, the +compressed models for CUDA only: without a cell, at the default charge state +and at ``charge_spin 1.0 2.0``, for a module that serves no periodic system, +and in its cell and without one for a module that does. """ import importlib.util @@ -48,10 +51,15 @@ / "deeppot_dpa4c_long_range.expected" ) compressed_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_compressed.pt2") +periodic_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_periodic.pt2") +periodic_compressed_pt2_file = pt2_file.with_name( + "deeppot_dpa4c_long_range_periodic_compressed.pt2" +) # A model without a long-range module, for a model deviation beside one. short_range_file = pt2_file.with_name("deeppot_sea.pth") compressed_ref_file = ref_file.with_name("deeppot_dpa4c_long_range_compressed.expected") data_file = Path(__file__).parent / "data_dpa4c_long_range_pt2.lmp" +periodic_data_file = Path(__file__).parent / "data_dpa4c_long_range_periodic_pt2.lmp" replicated_data_file = Path(__file__).parent / "data_dpa4c_long_range_grid_pt2.lmp" mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" @@ -66,12 +74,27 @@ except FileNotFoundError: expected_e_default = expected_f_default = None expected_e_explicit = expected_f_explicit = None -try: - compressed_ref = read_expected_ref(compressed_ref_file) -except FileNotFoundError: - compressed_ref = None + + +def _read_ref(path: Path) -> dict | None: + try: + return read_expected_ref(path) + except FileNotFoundError: + return None + + +compressed_ref = _read_ref(compressed_ref_file) +periodic_ref = _read_ref( + ref_file.with_name("deeppot_dpa4c_long_range_periodic.expected") +) +periodic_compressed_ref = _read_ref( + ref_file.with_name("deeppot_dpa4c_long_range_periodic_compressed.expected") +) box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +# The triclinic cell of the periodic references: xlo xhi ylo yhi zlo zhi xy xz +# yz of the lattice vectors (13, 0, 0), (0.6, 12.5, 0) and (-0.4, 0.8, 13.5). +periodic_box = np.array([0, 13, 0, 12.5, 0, 13.5, 0.6, -0.4, 0.8]) coord = np.array( [ [12.83, 2.56, 2.18], @@ -94,18 +117,19 @@ def setup_module() -> None: if os.environ.get("ENABLE_PYTORCH", "1") != "1": pytest.skip("Skip test because PyTorch support is not enabled.") write_lmp_data(box, coord, type_OH, data_file) + write_lmp_data(periodic_box, coord, type_OH, periodic_data_file) def teardown_module() -> None: - for path in (data_file, replicated_data_file): + for path in (data_file, periodic_data_file, replicated_data_file): if path.exists(): os.remove(path) -def _lammps(boundary: str) -> PyLammps: +def _lammps(boundary: str, data: Path = data_file) -> PyLammps: if not pt2_file.is_file(): pytest.skip(f"Skip because {pt2_file} was not generated.") - return make_atomic_lammps(data_file, boundary=boundary, atom_map="yes") + return make_atomic_lammps(data, boundary=boundary, atom_map="yes") @pytest.fixture @@ -172,6 +196,82 @@ def test_model_deviation_periodic_boundary_is_rejected( lammps_periodic.run(0) +@pytest.mark.parametrize("long_range_first", [True, False]) +def test_model_deviation_without_cell_needs_long_range_throughout( + lammps, long_range_first: bool, tmp_path: Path +) -> None: + """A non-periodic model deviation passes no cell, which every model must serve.""" + if not short_range_file.is_file(): + pytest.skip(f"Skip because {short_range_file} was not generated.") + models = [pt2_file, short_range_file] + if not long_range_first: + models.reverse() + lammps.pair_style( + "deepmd " + + " ".join(str(model.resolve()) for model in models) + + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}" + ) + lammps.pair_coeff("* *") + with pytest.raises(Exception, match="every model carries one"): + lammps.run(0) + + +def test_model_deviation_of_long_range_models(lammps, tmp_path: Path) -> None: + """Models that all carry a long-range module serve a non-periodic deviation.""" + md_file = tmp_path / "model_devi.out" + lammps.pair_style( + f"deepmd {pt2_file.resolve()} {pt2_file.resolve()} " + f"out_freq 1 out_file {md_file}" + ) + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_default) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_default[lammps.atoms[ii].id - 1] + ) + md = np.loadtxt(md_file) + assert md[1:] == pytest.approx(np.zeros(md.shape[0] - 1), abs=1e-10) + + +def _check_reference(lmp: PyLammps, ref: dict, section: str) -> None: + """Compare the energy and the forces of a run with a reference section.""" + lmp.run(0) + assert lmp.eval("pe") == pytest.approx(float(np.sum(ref[section]["expected_e"]))) + expected_f = ref[section]["expected_f"].reshape(6, 3) + for ii in range(6): + assert lmp.atoms[ii].force == pytest.approx( + expected_f[lmp.atoms[ii].id - 1], abs=1e-8 + ) + + +@pytest.mark.parametrize( + ("boundary", "section"), [("p p p", "periodic"), ("f f f", "cluster")] +) +def test_pair_deepmd_periodic_module(boundary: str, section: str) -> None: + """A module that serves periodic systems sums over the images of the cell.""" + if periodic_ref is None or not periodic_pt2_file.is_file(): + pytest.skip(f"Skip because {periodic_pt2_file} was not generated.") + lmp = _lammps(boundary, periodic_data_file) + lmp.pair_style(f"deepmd {periodic_pt2_file.resolve()}") + lmp.pair_coeff("* *") + _check_reference(lmp, periodic_ref, section) + lmp.run(1) + lmp.close() + + +@pytest.mark.parametrize("model", [pt2_file, periodic_pt2_file]) +def test_partially_periodic_boundary_is_rejected(model: Path) -> None: + if not model.is_file(): + pytest.skip(f"Skip because {model} was not generated.") + lmp = _lammps("p p f", periodic_data_file) + lmp.pair_style(f"deepmd {model.resolve()}") + lmp.pair_coeff("* *") + with pytest.raises(Exception, match="every dimension"): + lmp.run(0) + lmp.close() + + mpi_available = pytest.mark.skipif( shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, reason="MPI or mpi4py is not installed", @@ -218,13 +318,53 @@ def test_pair_deepmd_mpi_compressed( ) +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors", "boundary", "section"), + [ + (2, "2 1 1", "p p p", "periodic"), + # Every atom lies in the lower half along z; the ranks above hold + # ghost atoms alone. + (4, "1 1 4", "p p p", "periodic"), + (2, "1 2 1", "f f f", "cluster"), + ], +) +def test_pair_deepmd_mpi_periodic_compressed( + nprocs: int, processors: str, boundary: str, section: str +) -> None: + if periodic_compressed_ref is None or not periodic_compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {periodic_compressed_pt2_file} was not generated.") + if not periodic_data_file.is_file(): + write_lmp_data(periodic_box, coord, type_OH, periodic_data_file) + result = run_mpi_pair_runner( + mpi_runner, + periodic_data_file, + periodic_compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=["--boundary", boundary], + ) + ref = periodic_compressed_ref[section] + expected_v = -ref["expected_v"].reshape(6, 9) + assert result["pe"] == pytest.approx(float(np.sum(ref["expected_e"])), rel=1e-6) + np.testing.assert_allclose( + result["forces"], ref["expected_f"].reshape(6, 3), rtol=0, atol=1e-5 + ) + np.testing.assert_allclose( + result["virials"][:, expected_v_to_lammps] / constants.nktv2p, + expected_v, + rtol=0, + atol=1e-5, + ) + + @pytest.fixture(scope="module") def replicated_reference() -> dict: """A grid of 168 copies of the six atoms, and its single-rank result. The copies lie farther apart than the short-range cutoff, so only the long-range module couples them; the 1008 atoms fill several blocks and - chunks of the direct-sum operator. + chunks of the long-range operator. """ if compressed_ref is None or not compressed_pt2_file.is_file(): pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 933f0fe7c9..17ef7c3d8d 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -182,6 +182,8 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) edge_force_virial.cu long_range/direct_sum.cu long_range/domain_exchange.cc + long_range/latent_charge_sum.cu + long_range/mesh_sum.cu dpa1_graph_energy_force.cu dpa4/so2_conv.cu ${DPA4_SO2_CONV_KERNEL_SRC} diff --git a/source/op/pt/long_range/charge_sum.h b/source/op/pt/long_range/charge_sum.h new file mode 100644 index 0000000000..d570804e8c --- /dev/null +++ b/source/op/pt/long_range/charge_sum.h @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// The two evaluations behind deepmd::latent_charge_sum (schema and formulas in +// ops.cc): the direct sum over the pairs of non-periodic frames +// (direct_sum.cu) and the smooth particle-mesh Ewald sum of periodic frames +// (mesh_sum.cu). Both take float64 positions and float32 raw charges of the +// nodes of the frames on one CUDA device, apply the constraint, and evaluate +// the owner nodes [owner_begin, owner_begin + owner_count) against every node +// of their frames. Owners other than every node lie in a single frame, whose +// nodes are then the charged nodes of every domain of ``communicator``, over +// which the per-frame sums of the owners are added. + +#pragma once + +#include + +#include +#include + +#include "domain_exchange.h" + +namespace deepmd { +namespace long_range { + +// The atomic energy (owners, 1) float64, the charge gradient (owners, C), the +// force (owners, 3), the atomic virial (owners, 9) and the virial of the +// owners of every frame (F, 9), float32. +using ChargeSumOutputs = std::tuple; + +// The direct sum. ``coefficient`` holds the float32 coefficients of the +// kernel on the device. +ChargeSumOutputs direct_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int owner_begin, + int owner_count, + const torch::Tensor* communicator); + +// The mesh sum. ``box`` holds the cell of every frame, float64 on the host; +// ``coefficient`` the float64 coefficients of the kernel on the device. +ChargeSumOutputs mesh_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& box, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + int owner_begin, + int owner_count, + const torch::Tensor* communicator); + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu index c12e144537..29d142e513 100644 --- a/source/op/pt/long_range/direct_sum.cu +++ b/source/op/pt/long_range/direct_sum.cu @@ -1,7 +1,8 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // -// CUDA kernel of the latent-charge long-range direct sum (schema and -// formulas in ops.cc). +// CUDA kernels of the latent-charge long range of non-periodic frames, by the +// direct sum over their pairs (schema and formulas in ops.cc, the entry in +// latent_charge_sum.cu). // // Every node reads every node of its frame, so the memory is linear in the // node count while the work is quadratic. One thread owns one node and walks @@ -32,18 +33,15 @@ // // The nodes whose sums a call evaluates, its owners, are either every node // or, across the domains of a decomposed system, the charged nodes of one -// MPI rank. The ranks then gather the positions and charges of the charged -// nodes of every domain into one frame, in rank order (domain_exchange.h), -// and each rank evaluates its own nodes against the whole frame, so that the +// MPI rank, which the entry gathers with those of every rank into one frame; +// each rank evaluates its own nodes against the whole frame, so that the // quadratic work divides among the ranks. Every rank shifts the constrained // channel by the same sum of the gathered charges; the mean potential of the // constraint adds the sums of the ranks in rank order, and the per-frame // virial is the share of the owners of the rank. #include -#include #include -#include #include #include @@ -51,7 +49,8 @@ #include #include -#include "domain_exchange.h" +#include "charge_sum.h" +#include "frame_sums.cuh" namespace { @@ -59,7 +58,6 @@ namespace { constexpr int kThreads = 128; // Partners evaluated together, which share every coefficient load. constexpr int kBatch = 4; -constexpr int kMaxChannels = 8; constexpr int kMaxGaussians = 32; // Fills of the resident block slots of the device that one launch aims for. constexpr int kWaves = 8; @@ -88,168 +86,6 @@ struct CompensatedSum { } }; -// Index of the frame holding ``node``, or -1 for a padding node past every -// frame. ``frame_ptr`` holds ``frame_count + 1`` offsets. -__device__ __forceinline__ int frame_of(const int64_t* frame_ptr, - int frame_count, - int64_t node) { - if (node >= frame_ptr[frame_count]) { - return -1; - } - int low = 0; - int high = frame_count; - while (high - low > 1) { - const int middle = (low + high) / 2; - if (frame_ptr[middle] <= node) { - low = middle; - } else { - high = middle; - } - } - return low; -} - -// Writes the offsets of the frames on the node axis: frame_ptr[f] counts the -// nodes of the frames before f, for f in [0, frame_count]. One block scans -// the counts, each thread a contiguous run of frames. -constexpr int kOffsetThreads = 256; - -__global__ void __launch_bounds__(kOffsetThreads) - frame_offsets(const int64_t* __restrict__ n_node_per_frame, - int frame_count, - int64_t* __restrict__ frame_ptr) { - __shared__ int64_t run_total[kOffsetThreads]; - const int run = (frame_count + kOffsetThreads - 1) / kOffsetThreads; - const int begin = min(frame_count, static_cast(threadIdx.x) * run); - const int end = min(frame_count, begin + run); - int64_t total = 0; - for (int f = begin; f < end; ++f) { - total += n_node_per_frame[f]; - } - run_total[threadIdx.x] = total; - __syncthreads(); - // Inclusive scan of the run totals. - for (int stride = 1; stride < kOffsetThreads; stride *= 2) { - const int64_t addend = - threadIdx.x >= stride ? run_total[threadIdx.x - stride] : 0; - __syncthreads(); - run_total[threadIdx.x] += addend; - __syncthreads(); - } - int64_t offset = threadIdx.x == 0 ? 0 : run_total[threadIdx.x - 1]; - for (int f = begin; f < end; ++f) { - frame_ptr[f] = offset; - offset += n_node_per_frame[f]; - } - if (threadIdx.x == kOffsetThreads - 1) { - frame_ptr[frame_count] = run_total[kOffsetThreads - 1]; - } -} - -// Per-frame sums. The nodes of a frame are split into ``gridDim.y`` -// contiguous slices, one block of kFrameThreads each, and a finalizing kernel -// adds the float64 sums of the slices in slice order. -constexpr int kFrameThreads = 256; -// Nodes per slice of a frame that a slice pass aims for. -constexpr int kSliceNodes = 2048; -constexpr int kMaxSlices = 256; - -struct NodeRange { - int64_t begin; - int64_t end; -}; - -// The nodes of the slice of the frame that this block sums. -__device__ __forceinline__ NodeRange -slice_of_frame(const int64_t* __restrict__ frame_ptr) { - const int64_t frame_begin = frame_ptr[blockIdx.x]; - const int64_t frame_end = frame_ptr[blockIdx.x + 1]; - const int64_t length = (frame_end - frame_begin + gridDim.y - 1) / gridDim.y; - const int64_t begin = frame_begin + blockIdx.y * length; - return {begin, min(frame_end, begin + length)}; -} - -// Adds ``value`` over the threads of the block in a fixed order and writes the -// totals to the entry of this slice in ``slice_sum``. -template -__device__ __forceinline__ void write_slice_sum( - const double (&value)[kComponents], double* __restrict__ slice_sum) { - __shared__ double partial[kComponents][kFrameThreads]; -#pragma unroll - for (int k = 0; k < kComponents; ++k) { - partial[k][threadIdx.x] = value[k]; - } - __syncthreads(); - for (int stride = kFrameThreads / 2; stride > 0; stride /= 2) { - if (threadIdx.x < stride) { -#pragma unroll - for (int k = 0; k < kComponents; ++k) { - partial[k][threadIdx.x] += partial[k][threadIdx.x + stride]; - } - } - __syncthreads(); - } - if (threadIdx.x < kComponents) { - slice_sum[(static_cast(blockIdx.x) * gridDim.y + blockIdx.y) * - kComponents + - threadIdx.x] = partial[threadIdx.x][0]; - } -} - -// Per-frame statistics of the charge constraint: the shift of the -// constrained channel, the number of charged nodes, and the potential of the -// charged nodes in that channel summed over them. -constexpr int kStatShift = 0; -constexpr int kStatCount = 1; -constexpr int kStatPotential = 2; -constexpr int kStatWidth = 3; - -// Sums, over one slice of one frame, the number of charged nodes and their raw -// charges in the constrained channel. -__global__ void __launch_bounds__(kFrameThreads) - charge_slices(const float* __restrict__ charge, - const bool* __restrict__ node_mask, - const int64_t* __restrict__ frame_ptr, - int channels, - int channel, - double* __restrict__ slice_sum) { - const NodeRange slice = slice_of_frame(frame_ptr); - double value[2] = {0.0, 0.0}; - for (int64_t node = slice.begin + threadIdx.x; node < slice.end; - node += kFrameThreads) { - if (node_mask[node]) { - value[0] += 1.0; - value[1] += charge[node * channels + channel]; - } - } - write_slice_sum<2>(value, slice_sum); -} - -// The shift s_f = (sum_f c - Q_f) / max(n_f, 1) of the constrained channel of -// every frame, from the slice sums of charge_slices. -__global__ void charge_shift(const double* __restrict__ slice_sum, - int slice_count, - int frame_count, - const double* __restrict__ total_charge, - double* __restrict__ frame_stat) { - const int f = blockIdx.x * blockDim.x + threadIdx.x; - if (f >= frame_count) { - return; - } - double count = 0.0; - double sum = 0.0; - for (int slice = 0; slice < slice_count; ++slice) { - const double* entry = - slice_sum + (static_cast(f) * slice_count + slice) * 2; - count += entry[0]; - sum += entry[1]; - } - // A frame without a charged node has nothing to shift. - count = fmax(count, 1.0); - frame_stat[f * kStatWidth + kStatShift] = (sum - total_charge[f]) / count; - frame_stat[f * kStatWidth + kStatCount] = count; -} - // Writes the double-float positions relative to the first node of each // frame, scaled by alpha for the erf kernel, the frame of every node, and the // charges of the pairs: the raw charges of the charged nodes, shifted in the @@ -952,16 +788,15 @@ void dispatch_channels(int channels, return launch(buffers, owner_begin, owner_count, prefactor, stream); default: - TORCH_CHECK(false, "latent_charge_direct_sum: unsupported channel count ", + TORCH_CHECK(false, "latent_charge_sum: unsupported channel count ", channels); } } -using DirectSumOutputs = std::tuple; +} // namespace + +namespace deepmd { +namespace long_range { // Evaluates the owner nodes [owner_begin, owner_begin + owner_count) against // the nodes of their frames. Returns the outputs of the owners, row r for @@ -969,20 +804,20 @@ using DirectSumOutputs = std::tuple::max() / kThreads, - "latent_charge_direct_sum: too many nodes"); + "latent_charge_sum: too many nodes"); const int node_count = static_cast(coord.size(0)); const int frame_count = static_cast(n_node_per_frame.size(0)); const int64_t channels = charge.size(1); @@ -997,7 +832,7 @@ DirectSumOutputs evaluate(const torch::Tensor& coord, } const bool every_node = owner_begin == 0 && owner_count == node_count; TORCH_CHECK(every_node || frame_count == 1, - "latent_charge_direct_sum: owners other than every node lie in " + "latent_charge_sum: owners other than every node lie in " "one frame"); const bool constrained = channel >= 0; const int constrained_channel = static_cast(channel); @@ -1109,8 +944,8 @@ DirectSumOutputs evaluate(const torch::Tensor& coord, // every domain. auto potential_sum = frame_stat.view({frame_count, kStatWidth}).select(1, kStatPotential); - potential_sum.copy_(deepmd::long_range::sum_over_domains( - *communicator, potential_sum.contiguous())); + potential_sum.copy_( + sum_over_domains(*communicator, potential_sum.contiguous())); } if (owner_count > 0) { project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0, @@ -1126,147 +961,5 @@ DirectSumOutputs evaluate(const torch::Tensor& coord, return {atom_energy, charge_gradient, force, atom_virial, virial}; } -// Evaluates the frame of a system decomposed into the domains of -// ``communicator``, of which this rank holds one: its charged nodes are its -// owners, the charged nodes of every domain its partners. The nodes the rank -// does not own carry no charge, and the per-frame virial is the share of its -// owners. -DirectSumOutputs evaluate_across_domains(const torch::Tensor& coord, - const torch::Tensor& charge, - const torch::Tensor& node_mask, - const torch::Tensor& n_node_per_frame, - const torch::Tensor& total_charge, - const torch::Tensor& coefficient, - int64_t kernel, - int64_t channel, - double prefactor, - const torch::Tensor& communicator) { - TORCH_CHECK(n_node_per_frame.size(0) == 1, - "latent_charge_direct_sum: an evaluation across MPI ranks takes " - "one frame, got ", - n_node_per_frame.size(0)); - const int64_t node_count = coord.size(0); - // The charged nodes of the frame; the padding nodes past it carry none. - // Their count is the one value the host has to wait for before the - // exchange. - const auto in_frame = - torch::arange(node_count, n_node_per_frame.options()) < n_node_per_frame; - const auto owned = torch::nonzero(node_mask & in_frame).squeeze(1); - const auto rows = - torch::cat({coord.index_select(0, owned), - charge.index_select(0, owned).to(torch::kFloat64)}, - 1) - .contiguous(); - const auto [partners, owner_begin] = - deepmd::long_range::gather_domain_rows(communicator, rows); - const int64_t partner_count = partners.size(0); - const auto partner_coord = partners.slice(1, 0, 3).contiguous(); - const auto partner_charge = - partners.slice(1, 3).to(torch::kFloat32).contiguous(); - const auto partner_mask = torch::ones({partner_count}, node_mask.options()); - const auto partner_frame = - torch::full({1}, partner_count, n_node_per_frame.options()); - const auto [owner_energy, owner_gradient, owner_force, owner_atom_virial, - virial] = - evaluate(partner_coord, partner_charge, partner_mask, partner_frame, - total_charge, coefficient, kernel, channel, prefactor, - static_cast(owner_begin), static_cast(owned.size(0)), - &communicator); - return {torch::zeros({node_count, 1}, coord.options()) - .index_copy_(0, owned, owner_energy), - torch::zeros({node_count, charge.size(1)}, charge.options()) - .index_copy_(0, owned, owner_gradient), - torch::zeros({node_count, 3}, charge.options()) - .index_copy_(0, owned, owner_force), - torch::zeros({node_count, 9}, charge.options()) - .index_copy_(0, owned, owner_atom_virial), - virial}; -} - -DirectSumOutputs latent_charge_direct_sum(torch::Tensor coord, - torch::Tensor charge, - torch::Tensor node_mask, - torch::Tensor n_node_per_frame, - torch::Tensor total_charge, - torch::Tensor coefficient, - torch::Tensor communicator, - int64_t kernel, - int64_t channel, - double prefactor) { - TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && - n_node_per_frame.is_cuda() && total_charge.is_cuda() && - coefficient.is_cuda(), - "latent_charge_direct_sum: inputs must be CUDA tensors"); - TORCH_CHECK(charge.device() == coord.device() && - node_mask.device() == coord.device() && - n_node_per_frame.device() == coord.device() && - total_charge.device() == coord.device() && - coefficient.device() == coord.device(), - "latent_charge_direct_sum: inputs must share one device"); - TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && - coord.scalar_type() == torch::kFloat64, - "latent_charge_direct_sum: coord must be float64 with shape " - "(N, 3)"); - TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && - charge.scalar_type() == torch::kFloat32, - "latent_charge_direct_sum: charge must be float32 with shape " - "(N, C)"); - TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && - node_mask.scalar_type() == torch::kBool, - "latent_charge_direct_sum: node_mask must be bool with shape " - "(N,)"); - TORCH_CHECK(n_node_per_frame.dim() == 1, - "latent_charge_direct_sum: n_node_per_frame must be " - "one-dimensional"); - TORCH_CHECK(total_charge.dim() == 1 && - total_charge.size(0) == n_node_per_frame.size(0) && - total_charge.scalar_type() == torch::kFloat64, - "latent_charge_direct_sum: total_charge must be float64 with " - "one entry per frame"); - TORCH_CHECK( - coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat32, - "latent_charge_direct_sum: coefficient must be a float32 vector"); - const int64_t channels = charge.size(1); - TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, - "latent_charge_direct_sum: the channel count must lie in [1, ", - kMaxChannels, "], got ", channels); - TORCH_CHECK(channel >= -1 && channel < channels, - "latent_charge_direct_sum: the constrained channel must be -1 " - "or lie in [0, ", - channels, "), got ", channel); - if (kernel == kLes) { - TORCH_CHECK(coefficient.numel() == 1, - "latent_charge_direct_sum: the erf kernel takes one " - "coefficient"); - } else if (kernel == kSog) { - TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && - coefficient.numel() <= 2 * kMaxGaussians, - "latent_charge_direct_sum: the Gaussian kernel takes two " - "coefficients per Gaussian, for at most ", - kMaxGaussians, " Gaussians"); - } else { - TORCH_CHECK(false, "latent_charge_direct_sum: unknown kernel ", kernel); - } - - const c10::cuda::CUDAGuard guard(coord.device()); - coord = coord.contiguous(); - charge = charge.contiguous(); - node_mask = node_mask.contiguous(); - total_charge = total_charge.contiguous(); - coefficient = coefficient.contiguous(); - if (deepmd::long_range::domain_of(communicator).size == 1) { - return evaluate(coord, charge, node_mask, n_node_per_frame, total_charge, - coefficient, kernel, channel, prefactor, 0, - static_cast(coord.size(0)), nullptr); - } - return evaluate_across_domains(coord, charge, node_mask, n_node_per_frame, - total_charge, coefficient, kernel, channel, - prefactor, communicator); -} - -} // namespace - -TORCH_LIBRARY_FRAGMENT(deepmd, library) { - library.impl("latent_charge_direct_sum", torch::kCUDA, - &latent_charge_direct_sum); -} +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc index 24c393c5ee..2ae3d06e05 100644 --- a/source/op/pt/long_range/domain_exchange.cc +++ b/source/op/pt/long_range/domain_exchange.cc @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // -// MPI exchanges of the long-range direct sum (contract in domain_exchange.h). +// MPI exchanges of the latent-charge long range (contract in +// domain_exchange.h). // // With an MPI library that reads device memory (MPIX_Query_cuda_support), the // rows travel between device buffers once the current stream has finished @@ -30,7 +31,7 @@ std::int64_t address_of(const torch::Tensor& communicator) { TORCH_CHECK(communicator.device().is_cpu() && communicator.scalar_type() == torch::kInt64 && communicator.numel() == 1, - "latent_charge_direct_sum: communicator must be an int64 CPU " + "latent_charge_sum: communicator must be an int64 CPU " "tensor with one element"); return communicator.contiguous().data_ptr()[0]; } @@ -64,7 +65,7 @@ Domain domain_of(const torch::Tensor& communicator) { int initialized = 0; MPI_Initialized(&initialized); TORCH_CHECK(initialized != 0, - "latent_charge_direct_sum: a communicator was given, but MPI is " + "latent_charge_sum: a communicator was given, but MPI is " "not initialized"); const MPI_Comm comm = communicator_of(communicator); Domain domain{0, 1}; @@ -73,7 +74,7 @@ Domain domain_of(const torch::Tensor& communicator) { return domain; #else TORCH_CHECK(false, - "latent_charge_direct_sum: an evaluation across MPI ranks " + "latent_charge_sum: an evaluation across MPI ranks " "requires DeePMD-kit built with MPI"); #endif } diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h index c091589b31..f423974f26 100644 --- a/source/op/pt/long_range/domain_exchange.h +++ b/source/op/pt/long_range/domain_exchange.h @@ -1,13 +1,13 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // -// Exchange of the long-range direct sum between the domains of a decomposed -// system. +// Exchange of the latent-charge long range between the domains of a +// decomposed system. // // A domain is one MPI rank of the communicator that a caller hands to -// deepmd::latent_charge_direct_sum as an int64 tensor holding the address of -// its MPI_Comm handle. A zero address stands for a single domain. Both -// exchanges follow the rank order, so every rank receives the same rows and the -// same sums, independent of the reduction algorithm of the MPI library. +// deepmd::latent_charge_sum as an int64 tensor holding the address of its +// MPI_Comm handle. A zero address stands for a single domain. Both exchanges +// follow the rank order, so every rank receives the same rows and the same +// sums, independent of the reduction algorithm of the MPI library. #pragma once diff --git a/source/op/pt/long_range/frame_sums.cuh b/source/op/pt/long_range/frame_sums.cuh new file mode 100644 index 0000000000..a773edd63a --- /dev/null +++ b/source/op/pt/long_range/frame_sums.cuh @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Per-frame bookkeeping shared by the CUDA units of the latent-charge long +// range: the offsets of the frames on the node axis, the sums over the nodes +// of every frame in a fixed order, and the statistics of the charge +// constraint. +// +// The kernels are defined in an unnamed namespace, so that every unit that +// includes this header compiles its own copy. + +#pragma once + +#include + +#include + +namespace { + +// Index of the frame holding ``node``, or -1 for a padding node past every +// frame. ``frame_ptr`` holds ``frame_count + 1`` offsets. +__device__ __forceinline__ int frame_of(const int64_t* frame_ptr, + int frame_count, + int64_t node) { + if (node >= frame_ptr[frame_count]) { + return -1; + } + int low = 0; + int high = frame_count; + while (high - low > 1) { + const int middle = (low + high) / 2; + if (frame_ptr[middle] <= node) { + low = middle; + } else { + high = middle; + } + } + return low; +} + +// Writes the offsets of the frames on the node axis: frame_ptr[f] counts the +// nodes of the frames before f, for f in [0, frame_count]. One block scans +// the counts, each thread a contiguous run of frames. +constexpr int kOffsetThreads = 256; + +__global__ void __launch_bounds__(kOffsetThreads) + frame_offsets(const int64_t* __restrict__ n_node_per_frame, + int frame_count, + int64_t* __restrict__ frame_ptr) { + __shared__ int64_t run_total[kOffsetThreads]; + const int run = (frame_count + kOffsetThreads - 1) / kOffsetThreads; + const int begin = min(frame_count, static_cast(threadIdx.x) * run); + const int end = min(frame_count, begin + run); + int64_t total = 0; + for (int f = begin; f < end; ++f) { + total += n_node_per_frame[f]; + } + run_total[threadIdx.x] = total; + __syncthreads(); + // Inclusive scan of the run totals. + for (int stride = 1; stride < kOffsetThreads; stride *= 2) { + const int64_t addend = + threadIdx.x >= stride ? run_total[threadIdx.x - stride] : 0; + __syncthreads(); + run_total[threadIdx.x] += addend; + __syncthreads(); + } + int64_t offset = threadIdx.x == 0 ? 0 : run_total[threadIdx.x - 1]; + for (int f = begin; f < end; ++f) { + frame_ptr[f] = offset; + offset += n_node_per_frame[f]; + } + if (threadIdx.x == kOffsetThreads - 1) { + frame_ptr[frame_count] = run_total[kOffsetThreads - 1]; + } +} + +// Per-frame sums. The nodes of a frame are split into ``gridDim.y`` +// contiguous slices, one block of kFrameThreads each, and a finalizing kernel +// adds the float64 sums of the slices in slice order. +constexpr int kFrameThreads = 256; +// Nodes per slice of a frame that a slice pass aims for. +constexpr int kSliceNodes = 2048; +constexpr int kMaxSlices = 256; + +struct NodeRange { + int64_t begin; + int64_t end; +}; + +// The nodes of the slice of the frame that this block sums. +__device__ __forceinline__ NodeRange +slice_of_frame(const int64_t* __restrict__ frame_ptr) { + const int64_t frame_begin = frame_ptr[blockIdx.x]; + const int64_t frame_end = frame_ptr[blockIdx.x + 1]; + const int64_t length = (frame_end - frame_begin + gridDim.y - 1) / gridDim.y; + const int64_t begin = frame_begin + blockIdx.y * length; + return {begin, min(frame_end, begin + length)}; +} + +// Adds ``value`` over the threads of the block in a fixed order and writes the +// totals to the entry of this slice in ``slice_sum``. +template +__device__ __forceinline__ void write_slice_sum( + const double (&value)[kComponents], double* __restrict__ slice_sum) { + __shared__ double partial[kComponents][kFrameThreads]; +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] = value[k]; + } + __syncthreads(); + for (int stride = kFrameThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] += partial[k][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kComponents) { + slice_sum[(static_cast(blockIdx.x) * gridDim.y + blockIdx.y) * + kComponents + + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Per-frame statistics of the charge constraint: the shift of the +// constrained channel, the number of charged nodes, and the potential of the +// charged nodes in that channel summed over them. +constexpr int kStatShift = 0; +constexpr int kStatCount = 1; +constexpr int kStatPotential = 2; +constexpr int kStatWidth = 3; + +// Sums, over one slice of one frame, the number of charged nodes and their raw +// charges in the constrained channel. +__global__ void __launch_bounds__(kFrameThreads) + charge_slices(const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + const NodeRange slice = slice_of_frame(frame_ptr); + double value[2] = {0.0, 0.0}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { + if (node_mask[node]) { + value[0] += 1.0; + value[1] += charge[node * channels + channel]; + } + } + write_slice_sum<2>(value, slice_sum); +} + +// The shift s_f = (sum_f c - Q_f) / max(n_f, 1) of the constrained channel of +// every frame, from the slice sums of charge_slices. +__global__ void charge_shift(const double* __restrict__ slice_sum, + int slice_count, + int frame_count, + const double* __restrict__ total_charge, + double* __restrict__ frame_stat) { + const int f = blockIdx.x * blockDim.x + threadIdx.x; + if (f >= frame_count) { + return; + } + double count = 0.0; + double sum = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + const double* entry = + slice_sum + (static_cast(f) * slice_count + slice) * 2; + count += entry[0]; + sum += entry[1]; + } + // A frame without a charged node has nothing to shift. + count = fmax(count, 1.0); + frame_stat[f * kStatWidth + kStatShift] = (sum - total_charge[f]) / count; + frame_stat[f * kStatWidth + kStatCount] = count; +} + +} // namespace diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu new file mode 100644 index 0000000000..f787c5cdb8 --- /dev/null +++ b/source/op/pt/long_range/latent_charge_sum.cu @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA entry of deepmd::latent_charge_sum (schema and formulas in ops.cc). +// +// A call evaluates its frames by the direct sum when they carry no cell and +// by the mesh sum when they do (charge_sum.h). Across the domains of a +// decomposed system, the ranks gather the positions and charges of the charged +// nodes of every domain into one frame, in rank order (domain_exchange.h), and +// each rank evaluates its own charged nodes against that frame; the outputs +// return to the node axis of the rank, where the nodes it does not own carry +// no charge. + +#include +#include +#include + +#include +#include + +#include "charge_sum.h" +#include "domain_exchange.h" + +namespace { + +using deepmd::long_range::ChargeSumOutputs; + +constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// The evaluation of the frames, by the direct or the mesh sum, of the owner +// nodes [owner_begin, owner_begin + owner_count). ``box`` is defined for +// periodic frames only. +struct Evaluation { + const torch::Tensor& box; + const torch::Tensor& coefficient; + int64_t kernel; + int64_t channel; + double prefactor; + int64_t order; + double spacing; + bool periodic; + + ChargeSumOutputs operator()(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) const { + if (periodic) { + return deepmd::long_range::mesh_sum( + coord, charge, node_mask, n_node_per_frame, total_charge, box, + coefficient, kernel, channel, prefactor, order, spacing, owner_begin, + owner_count, communicator); + } + return deepmd::long_range::direct_sum( + coord, charge, node_mask, n_node_per_frame, total_charge, + coefficient.to(torch::kFloat32), kernel, channel, prefactor, + owner_begin, owner_count, communicator); + } +}; + +// Evaluates the frame of a system decomposed into the domains of +// ``communicator``, of which this rank holds one: its charged nodes are its +// owners, the charged nodes of every domain its partners. The nodes the rank +// does not own carry no charge, and the per-frame virial is the share of its +// owners. +ChargeSumOutputs evaluate_across_domains(const Evaluation& evaluation, + const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& communicator) { + TORCH_CHECK(n_node_per_frame.size(0) == 1, + "latent_charge_sum: an evaluation across MPI ranks takes one " + "frame, got ", + n_node_per_frame.size(0)); + const int64_t node_count = coord.size(0); + // The charged nodes of the frame; the padding nodes past it carry none. + // Their count is the one value the host has to wait for before the + // exchange. + const auto in_frame = + torch::arange(node_count, n_node_per_frame.options()) < n_node_per_frame; + const auto owned = torch::nonzero(node_mask & in_frame).squeeze(1); + const auto rows = + torch::cat({coord.index_select(0, owned), + charge.index_select(0, owned).to(torch::kFloat64)}, + 1) + .contiguous(); + const auto [partners, owner_begin] = + deepmd::long_range::gather_domain_rows(communicator, rows); + const int64_t partner_count = partners.size(0); + const auto partner_coord = partners.slice(1, 0, 3).contiguous(); + const auto partner_charge = + partners.slice(1, 3).to(torch::kFloat32).contiguous(); + const auto partner_mask = torch::ones({partner_count}, node_mask.options()); + const auto partner_frame = + torch::full({1}, partner_count, n_node_per_frame.options()); + const auto [owner_energy, owner_gradient, owner_force, owner_atom_virial, + virial] = + evaluation(partner_coord, partner_charge, partner_mask, partner_frame, + total_charge, static_cast(owner_begin), + static_cast(owned.size(0)), &communicator); + return {torch::zeros({node_count, 1}, coord.options()) + .index_copy_(0, owned, owner_energy), + torch::zeros({node_count, charge.size(1)}, charge.options()) + .index_copy_(0, owned, owner_gradient), + torch::zeros({node_count, 3}, charge.options()) + .index_copy_(0, owned, owner_force), + torch::zeros({node_count, 9}, charge.options()) + .index_copy_(0, owned, owner_atom_virial), + virial}; +} + +// Whether the frames carry cells: ``box`` is given and no row of it is zero, +// or no frame carries one. +bool periodic_frames(const c10::optional& box) { + if (!box.has_value()) { + return false; + } + const auto has_cell = box->ne(0.0).any(1); + const bool any = has_cell.any().item(); + TORCH_CHECK(!any || has_cell.all().item(), + "latent_charge_sum: every frame of a call is periodic, or none " + "is"); + return any; +} + +ChargeSumOutputs latent_charge_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + c10::optional box, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing) { + TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && + n_node_per_frame.is_cuda() && total_charge.is_cuda() && + coefficient.is_cuda(), + "latent_charge_sum: inputs must be CUDA tensors, except box and " + "communicator"); + TORCH_CHECK(charge.device() == coord.device() && + node_mask.device() == coord.device() && + n_node_per_frame.device() == coord.device() && + total_charge.device() == coord.device() && + coefficient.device() == coord.device(), + "latent_charge_sum: inputs must share one device"); + TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && + coord.scalar_type() == torch::kFloat64, + "latent_charge_sum: coord must be float64 with shape (N, 3)"); + TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && + charge.scalar_type() == torch::kFloat32, + "latent_charge_sum: charge must be float32 with shape (N, C)"); + TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && + node_mask.scalar_type() == torch::kBool, + "latent_charge_sum: node_mask must be bool with shape (N,)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "latent_charge_sum: n_node_per_frame must be one-dimensional"); + TORCH_CHECK(total_charge.dim() == 1 && + total_charge.size(0) == n_node_per_frame.size(0) && + total_charge.scalar_type() == torch::kFloat64, + "latent_charge_sum: total_charge must be float64 with one entry " + "per frame"); + TORCH_CHECK( + !box.has_value() || + (box->device().is_cpu() && box->scalar_type() == torch::kFloat64 && + box->dim() == 2 && box->size(0) == n_node_per_frame.size(0) && + box->size(1) == 9), + "latent_charge_sum: box must be a float64 CPU tensor with shape " + "(F, 9)"); + TORCH_CHECK( + coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat64, + "latent_charge_sum: coefficient must be a float64 vector"); + const int64_t channels = charge.size(1); + TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, + "latent_charge_sum: the channel count must lie in [1, ", + kMaxChannels, "], got ", channels); + TORCH_CHECK(channel >= -1 && channel < channels, + "latent_charge_sum: the constrained channel must be -1 or lie " + "in [0, ", + channels, "), got ", channel); + if (kernel == kLes) { + TORCH_CHECK(coefficient.numel() == 1, + "latent_charge_sum: the erf kernel takes one coefficient"); + } else if (kernel == kSog) { + TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && + coefficient.numel() <= 2 * kMaxGaussians, + "latent_charge_sum: the Gaussian kernel takes two " + "coefficients per Gaussian, for at most ", + kMaxGaussians, " Gaussians"); + } else { + TORCH_CHECK(false, "latent_charge_sum: unknown kernel ", kernel); + } + TORCH_CHECK(coord.size(0) < std::numeric_limits::max(), + "latent_charge_sum: too many nodes"); + + const c10::cuda::CUDAGuard guard(coord.device()); + coord = coord.contiguous(); + charge = charge.contiguous(); + node_mask = node_mask.contiguous(); + total_charge = total_charge.contiguous(); + const bool periodic = periodic_frames(box); + const torch::Tensor cells = periodic ? box->contiguous() : torch::Tensor(); + coefficient = coefficient.contiguous(); + const Evaluation evaluation{cells, coefficient, kernel, channel, + prefactor, order, spacing, periodic}; + if (deepmd::long_range::domain_of(communicator).size == 1) { + return evaluation(coord, charge, node_mask, n_node_per_frame, total_charge, + 0, static_cast(coord.size(0)), nullptr); + } + return evaluate_across_domains(evaluation, coord, charge, node_mask, + n_node_per_frame, total_charge, communicator); +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.impl("latent_charge_sum", torch::kCUDA, &latent_charge_sum); +} diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu new file mode 100644 index 0000000000..f952494955 --- /dev/null +++ b/source/op/pt/long_range/mesh_sum.cu @@ -0,0 +1,1113 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernels of the latent-charge long range of periodic frames, by the +// smooth particle-mesh Ewald method (schema and formulas in ops.cc, the +// evaluation in deepmd/dpmodel/long_range/reciprocal.py). +// +// Every frame owns a mesh of n1 x n2 x n3 points spanning its cell. The +// charges spread to the mesh with cardinal B-splines of order p along the +// fractional coordinates s = r B^-1, the lattice vectors being the rows of the +// cell B: the scaled coordinate u_d = n_d s_d reaches the p points +// floor(u_d) - p + 1, ..., floor(u_d) with the weights M_p(u_d - m). A real +// FFT gives the half spectrum Q(m) of the charge mesh, and the influence +// function +// +// G(m) = prefactor * K^(k_m) / V * prod_d |D_d(m_d)|^-2, k_m = 2 pi B^-1 m, +// +// with the Fourier transform K^ of the kernel and the assignment factor +// D_d(m) = sum_l M_p(l + 1) exp(2 pi i m l / n_d), turns it into the +// potential mesh, the unnormalized inverse transform of G Q. The potential at +// a node is its interpolation by the B-splines of the node, less the potential +// prefactor * K(0) q of the charge at itself, and the force at fixed charges +// the negative gradient of that interpolation times the charges. The zero mode +// and the modes at the Nyquist frequency of an axis, which have no unique wave +// number, carry no weight. +// +// The virial of a frame is that of the strain of its cell, under which the +// fractional coordinates stay put while k_m and the volume change: +// +// W = sum_m e_m (1 + 2 dln K^ / d(k^2) k_m (x) k_m), e_m = G |Q|^2 / 2, +// +// the sum over the full spectrum, every interior plane of the half spectrum +// counting twice. Its per-node share is F_i (x) r_i, the force index first, +// plus the rest of W shared evenly over the charged nodes. +// +// The charges spread in 64-bit fixed point, so that the mesh, and every +// output, does not depend on the order of the atomic additions. The sums over +// the modes and over the nodes of a frame add per-block float64 sums in a +// fixed order. The meshes and the transforms are float64; the stencils, which +// would otherwise be bound by the float64 throughput of the device, weigh and +// read them in float32, the precision of the outputs, each warp adding the +// sums of its lanes in float64. +// +// The nodes whose outputs a call evaluates, its owners, are either every node +// or, across the domains of a decomposed system, the charged nodes of one MPI +// rank, whose frame then holds the charged nodes of every rank. Every rank +// spreads every charge and transforms the whole mesh, and interpolates only at +// its owners; the per-frame sums of the owners are added over the domains. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "charge_sum.h" +#include "frame_sums.cuh" + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxChannels = 8; +constexpr int kMaxOrder = 12; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// Resolution of the fixed-point charge mesh: 2^-40 e, far below the float32 +// precision of the raw charges, for a sum of at most 2^23 e at a point. +constexpr double kFixedScale = 1099511627776.0; // 2^40 +constexpr float kFixedScaleFloat = 1099511627776.f; + +constexpr double kPi = 3.14159265358979323846; + +// The cell of one frame and its mesh. +struct FrameMesh { + // B^-1, row-major: s_d = sum_j r_j inverse[3 j + d]. + double inverse[9]; + double volume; + int n[3]; + // The first point of the meshes of the frame, channel-major, in the mesh + // buffer of the call. + int64_t offset; +}; + +// The squared lengths of the lattice vectors are formed term by term, as the +// Python front end forms them, so that both round the mesh size identically. +int64_t even_fft_size(int64_t n) { + int64_t m = std::max(2, n + n % 2); + while (true) { + int64_t rest = m; + for (const int64_t p : {2, 3, 5, 7}) { + while (rest % p == 0) { + rest /= p; + } + } + if (rest == 1) { + return m; + } + m += 2; + } +} + +FrameMesh frame_mesh(const double* cell, double spacing) { + FrameMesh mesh{}; + const double a = cell[0], b = cell[1], c = cell[2]; + const double d = cell[3], e = cell[4], f = cell[5]; + const double g = cell[6], h = cell[7], i = cell[8]; + const double det = + a * (e * i - f * h) - b * (d * i - f * g) + c * (d * h - e * g); + // A volume that is a vanishing fraction of the product of the lengths of + // the lattice vectors, or not finite, spans no lattice. + const double lengths = std::sqrt(a * a + b * b + c * c) * + std::sqrt(d * d + e * e + f * f) * + std::sqrt(g * g + h * h + i * i); + TORCH_CHECK(std::abs(det) > 1e-6 * lengths, + "latent_charge_sum: a periodic cell is singular"); + const double adjugate[9] = {e * i - f * h, c * h - b * i, b * f - c * e, + f * g - d * i, a * i - c * g, c * d - a * f, + d * h - e * g, b * g - a * h, a * e - b * d}; + for (int k = 0; k < 9; ++k) { + mesh.inverse[k] = adjugate[k] / det; + } + mesh.volume = std::abs(det); + for (int axis = 0; axis < 3; ++axis) { + const double* v = cell + 3 * axis; + const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + const double points = std::ceil(length / spacing); + TORCH_CHECK(points < 1 << 20, + "latent_charge_sum: the mesh of a cell exceeds 2^20 points " + "along a lattice vector"); + mesh.n[axis] = + static_cast(even_fft_size(static_cast(points))); + } + return mesh; +} + +__host__ __device__ int64_t mesh_points(const FrameMesh& mesh) { + return int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2]; +} + +// Nodes per block of the kernels that spread and interpolate, one warp each. +constexpr int kStencilWarps = 4; +constexpr int kWarp = 32; + +// The stencil of one node along one axis of its mesh: the index of the +// first of its p points on the axis, which the others follow cyclically, and +// their B-spline weights and slopes, evaluated in float64 and stored in +// float32, the precision of the outputs: rounding a recursion of float32 steps +// instead would bias the sum of the weights. +template +struct AxisStencil { + int start; + float value[kOrder]; + float slope[kOrder]; +}; + +// The index of the point ``start + k`` of an axis of ``n`` points, for +// ``start`` < n and k < p, which exceeds 2 n only when a stencil wider than +// the axis wraps more than once. +__device__ __forceinline__ int wrap_index(int index, int n) { + if (index < n) { + return index; + } + return index - n < n ? index - n : index % n; +} + +// The stencil of a node along axis ``d``. The scaled coordinate u = n_d s_d +// of the fractional coordinate s_d reaches the points floor(u) - p + 1 + k, +// k = 0, ..., p - 1, with the weights value[k] = M_p(t + p - 1 - k), t the +// fractional part of u, and slope[k] their derivative with respect to u. The +// recursion M_n(t + j) = ((t + j) M_{n-1}(t + j) + (n - t - j) M_{n-1}(t + j - +// 1)) / (n - 1) builds a_j = M_n(t + j) from M_2; the slope follows from +// M_p'(x) = M_{p-1}(x) - M_{p-1}(x - 1). +template +__device__ void axis_stencil(const double* coord, + const FrameMesh& mesh, + int d, + AxisStencil& stencil) { + double s = coord[0] * mesh.inverse[d] + coord[1] * mesh.inverse[3 + d] + + coord[2] * mesh.inverse[6 + d]; + s -= floor(s); + const double u = s * mesh.n[d]; + const double base = floor(u); + const double t = u - base; + double a[kOrder]; + a[0] = t; + a[1] = 1.0 - t; +#pragma unroll + for (int n = 3; n <= kOrder; ++n) { + if (n == kOrder) { + // The weights of order p - 1 give the slopes. +#pragma unroll + for (int j = 0; j < kOrder; ++j) { + const double high = j < kOrder - 1 ? a[j] : 0.0; + const double low = j > 0 ? a[j - 1] : 0.0; + stencil.slope[kOrder - 1 - j] = static_cast(high - low); + } + } + const double scale = 1.0 / (n - 1); + a[n - 1] = scale * (1.0 - t) * a[n - 2]; +#pragma unroll + for (int j = n - 2; j > 0; --j) { + a[j] = scale * ((t + j) * a[j] + (n - t - j) * a[j - 1]); + } + a[0] = scale * t * a[0]; + } + if constexpr (kOrder == 2) { + stencil.slope[0] = -1.f; + stencil.slope[1] = 1.f; + } + const int first = static_cast(base) - kOrder + 1; + stencil.start = ((first % mesh.n[d]) + mesh.n[d]) % mesh.n[d]; +#pragma unroll + for (int k = 0; k < kOrder; ++k) { + stencil.value[k] = static_cast(a[kOrder - 1 - k]); + } +} + +// Writes the stencils of the charged nodes along the three axes of their +// meshes, one thread per node and axis, which the kernels that spread and +// interpolate read: every warp instruction in float64 occupies the float64 +// pipe of the device whatever the number of its active lanes, so the +// recursion runs once per node and axis with every lane busy. +template +__global__ void __launch_bounds__(kThreads) + node_stencils(const double* __restrict__ coord, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + int frame_count, + int64_t node_count, + AxisStencil* __restrict__ stencils) { + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < 3 * node_count; index += int64_t{gridDim.x} * blockDim.x) { + const int64_t node = index / 3; + const int f = frame_of(frame_ptr, frame_count, node); + if (f >= 0 && node_mask[node]) { + axis_stencil(coord + 3 * node, frames[f], + static_cast(index - 3 * node), stencils[index]); + } + } +} + +// The point ``t`` of the p^3 points of a stencil: its offset on the flat mesh +// of a frame and its indices (k1, k2, k3) along the axes. The index of the +// last axis runs fastest, so that the lanes of a warp touch contiguous runs +// of the mesh. +template +struct StencilPoint { + int64_t offset; + int k1; + int k2; + int k3; + + __device__ StencilPoint(int t, + const FrameMesh& mesh, + const AxisStencil* axes) + : k1(t / (kOrder * kOrder)), k2((t / kOrder) % kOrder), k3(t % kOrder) { + offset = (int64_t{wrap_index(axes[0].start + k1, mesh.n[0])} * mesh.n[1] + + wrap_index(axes[1].start + k2, mesh.n[1])) * + mesh.n[2] + + wrap_index(axes[2].start + k3, mesh.n[2]); + } +}; + +// Writes the charges of every node, shifted in the constrained channel and +// zero on the nodes that carry none, and spreads them to the meshes of their +// frames, one warp per node, the lanes taking the points of its stencil in +// turn. +template +__global__ void __launch_bounds__(kStencilWarps* kWarp) + spread_charges(const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const AxisStencil* __restrict__ stencils, + const double* __restrict__ frame_stat, + int frame_count, + int64_t node_count, + int channels, + int channel, + float* __restrict__ node_charge, + unsigned long long* __restrict__ mesh) { + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x % kWarp; + for (int64_t node = int64_t{blockIdx.x} * kStencilWarps + warp; + node < node_count; node += int64_t{gridDim.x} * kStencilWarps) { + const int f = frame_of(frame_ptr, frame_count, node); + const bool charged = f >= 0 && node_mask[node]; + const double shift = + charged && channel >= 0 ? frame_stat[f * kStatWidth + kStatShift] : 0.0; + float q[kMaxChannels]; +#pragma unroll + for (int c = 0; c < kMaxChannels; ++c) { + const double raw = + charged && c < channels ? charge[node * channels + c] : 0.0; + q[c] = static_cast(c == channel ? raw - shift : raw); + } + if (lane < channels) { + node_charge[node * channels + lane] = charged ? q[lane] : 0.f; + } + if (!charged) { + continue; + } + const FrameMesh grid = frames[f]; + const AxisStencil* axes = stencils + 3 * node; + const int64_t points = mesh_points(grid); + unsigned long long* frame_mesh = mesh + grid.offset; + for (int t = lane; t < kOrder * kOrder * kOrder; t += kWarp) { + const StencilPoint point(t, grid, axes); + const float w = axes[0].value[point.k1] * axes[1].value[point.k2] * + axes[2].value[point.k3] * kFixedScaleFloat; +#pragma unroll + for (int c = 0; c < kMaxChannels; ++c) { + if (c < channels) { + atomicAdd(frame_mesh + c * points + point.offset, + static_cast(__float2ll_rn(q[c] * w))); + } + } + } + } +} + +// Converts the fixed-point mesh to float64 in place. +__global__ void fixed_to_double(int64_t count, unsigned long long* mesh) { + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < count; index += int64_t{gridDim.x} * blockDim.x) { + const long long fixed = static_cast(mesh[index]); + reinterpret_cast(mesh)[index] = + static_cast(fixed) / kFixedScale; + } +} + +// The inverse squared assignment factor |D_d(m)|^-2 along every axis of the +// half spectrum of one frame: n1 entries, then n2, then n3 / 2 + 1. +__global__ void assignment_factor(FrameMesh mesh, + int order, + const double* __restrict__ integer_weight, + double* __restrict__ factor) { + const int half = mesh.n[2] / 2 + 1; + const int total = mesh.n[0] + mesh.n[1] + half; + for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < total; + index += gridDim.x * blockDim.x) { + int axis = 0; + int m = index; + if (m >= mesh.n[0]) { + m -= mesh.n[0]; + axis = 1; + if (m >= mesh.n[1]) { + m -= mesh.n[1]; + axis = 2; + } + } + const double theta = 2.0 * kPi * m / mesh.n[axis]; + double re = 0.0; + double im = 0.0; + for (int l = 0; l < order - 1; ++l) { + double s; + double c; + sincos(theta * l, &s, &c); + re += integer_weight[l] * c; + im += integer_weight[l] * s; + } + factor[index] = 1.0 / (re * re + im * im); + } +} + +// The transform of the kernel and its logarithmic derivative with respect to +// k^2. ``coefficient`` holds alpha, or the amplitudes and then the rates +// c_g = log2(e) / (2 sigma_g^2) of the Gaussians. +__device__ void kernel_transform(int kernel, + const double* __restrict__ coefficient, + int gaussians, + double k_sq, + double& value, + double& log_slope) { + if (kernel == kLes) { + const double alpha = coefficient[0]; + const double inv_four_alpha_sq = 0.25 / (alpha * alpha); + value = 4.0 * kPi * exp(-k_sq * inv_four_alpha_sq) / k_sq; + log_slope = -inv_four_alpha_sq - 1.0 / k_sq; + return; + } + // (2 pi)^{3/2} + constexpr double kGaussianNorm = 15.749609945722419; + value = 0.0; + double slope = 0.0; + for (int g = 0; g < gaussians; ++g) { + // sigma^2 = log2(e) / (2 c) = 1 / (2 c ln 2) + const double width_sq = 0.72134752044448170 / coefficient[gaussians + g]; + const double term = coefficient[g] * kGaussianNorm * width_sq * + sqrt(width_sq) * exp(-0.5 * width_sq * k_sq); + value += term; + slope -= 0.5 * width_sq * term; + } + log_slope = value != 0.0 ? slope / value : 0.0; +} + +// The influence function of one frame on its half spectrum: multiplies every +// mode of the spectra of the channels by G in place and sums, per block, the +// energy of the modes and their virial, xx yy zz xy xz yz. +constexpr int kModeSums = 7; + +__global__ void __launch_bounds__(kThreads) + apply_influence(FrameMesh mesh, + int channels, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + const double* __restrict__ factor, + double2* __restrict__ spectrum, + double* __restrict__ block_sum) { + const int n1 = mesh.n[0]; + const int n2 = mesh.n[1]; + const int n3 = mesh.n[2]; + const int half = n3 / 2 + 1; + const int64_t modes = int64_t{n1} * n2 * half; + double sum[kModeSums] = {}; + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < modes; index += int64_t{gridDim.x} * blockDim.x) { + const int i3 = static_cast(index % half); + const int i2 = static_cast((index / half) % n2); + const int i1 = static_cast(index / (int64_t{half} * n2)); + const int m1 = i1 <= n1 / 2 ? i1 : i1 - n1; + const int m2 = i2 <= n2 / 2 ? i2 : i2 - n2; + const int m3 = i3; + const bool weightless = (m1 == 0 && m2 == 0 && m3 == 0) || 2 * m1 == n1 || + 2 * m2 == n2 || 2 * m3 == n3; + double influence = 0.0; + double k[3] = {0.0, 0.0, 0.0}; + double log_slope = 0.0; + if (!weightless) { + for (int a = 0; a < 3; ++a) { + k[a] = 2.0 * kPi * + (mesh.inverse[3 * a] * m1 + mesh.inverse[3 * a + 1] * m2 + + mesh.inverse[3 * a + 2] * m3); + } + const double k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2]; + double transform; + kernel_transform(kernel, coefficient, gaussians, k_sq, transform, + log_slope); + influence = prefactor * transform / mesh.volume * factor[i1] * + factor[n1 + i2] * factor[n1 + n2 + i3]; + } + // An interior plane of the half spectrum stands for its mirror as well. + const double multiplicity = m3 == 0 ? 1.0 : 2.0; + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + double2& value = spectrum[c * modes + index]; + energy += value.x * value.x + value.y * value.y; + value.x *= influence; + value.y *= influence; + } + energy *= 0.5 * multiplicity * influence; + const double twice = 2.0 * log_slope * energy; + sum[0] += energy; + sum[1] += energy + twice * k[0] * k[0]; + sum[2] += energy + twice * k[1] * k[1]; + sum[3] += energy + twice * k[2] * k[2]; + sum[4] += twice * k[0] * k[1]; + sum[5] += twice * k[0] * k[2]; + sum[6] += twice * k[1] * k[2]; + } + __shared__ double partial[kModeSums][kThreads]; + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] = sum[s]; + } + __syncthreads(); + for (int stride = kThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] += partial[s][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kModeSums) { + block_sum[blockIdx.x * kModeSums + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Adds the block sums of apply_influence into the virial of the frame, +// row-major: every thread sums a fixed strided subset of the blocks, and the +// block adds the thread sums in a fixed order. +__global__ void __launch_bounds__(kThreads) + mode_virial(const double* __restrict__ block_sum, + int block_count, + double* __restrict__ virial) { + double value[kModeSums] = {}; + for (int block = threadIdx.x; block < block_count; block += kThreads) { +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + value[s] += block_sum[block * kModeSums + s]; + } + } + __shared__ double partial[kModeSums][kThreads]; +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] = value[s]; + } + __syncthreads(); + for (int stride = kThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] += partial[s][threadIdx.x + stride]; + } + } + __syncthreads(); + } + // Components 1..6 are xx yy zz xy xz yz; the energy is not stored. + constexpr int kEntry[kModeSums][2] = {{-1, -1}, {0, 0}, {4, 4}, {8, 8}, + {1, 3}, {2, 6}, {5, 7}}; + if (threadIdx.x > 0 && threadIdx.x < kModeSums) { + virial[kEntry[threadIdx.x][0]] = partial[threadIdx.x][0]; + virial[kEntry[threadIdx.x][1]] = partial[threadIdx.x][0]; + } +} + +// The contact value K(0) of the kernel. +__device__ double contact_value(int kernel, + const double* __restrict__ coefficient, + int gaussians) { + if (kernel == kLes) { + return 1.1283791670955126 * coefficient[0]; + } + double value = 0.0; + for (int g = 0; g < gaussians; ++g) { + value += coefficient[g]; + } + return value; +} + +// Adds ``value`` over the lanes of a warp in a fixed butterfly order; every +// lane receives the total. +__device__ __forceinline__ double warp_sum(double value) { +#pragma unroll + for (int stride = kWarp / 2; stride > 0; stride /= 2) { + value += __shfl_xor_sync(0xffffffffu, value, stride); + } + return value; +} + +// Interpolates the potential mesh at every owner node, one warp per node: +// its potential, less that of its charge at itself, its force at fixed +// charges and the share F (x) r of its virial. Owners that carry no charge get +// zeros. The lanes take the points of the stencil in turn and the warp adds +// their sums. +template +__global__ void __launch_bounds__(kStencilWarps* kWarp) + interpolate_owners(const double* __restrict__ coord, + const float* __restrict__ node_charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const AxisStencil* __restrict__ stencils, + const float* __restrict__ potential_mesh, + const double* __restrict__ coefficient, + int frame_count, + int owner_begin, + int owner_count, + int channels, + int kernel, + int gaussians, + double prefactor, + double* __restrict__ potential, + double* __restrict__ force, + double* __restrict__ moment, + double* __restrict__ atom_energy) { + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x % kWarp; + const double self = prefactor * contact_value(kernel, coefficient, gaussians); + for (int owner = blockIdx.x * kStencilWarps + warp; owner < owner_count; + owner += gridDim.x * kStencilWarps) { + const int64_t node = int64_t{owner_begin} + owner; + const int f = frame_of(frame_ptr, frame_count, node); + const bool charged = f >= 0 && node_mask[node]; + const float* q = node_charge + node * channels; + double phi[kMaxChannels] = {}; + double f_node[3] = {0.0, 0.0, 0.0}; + if (charged) { + const FrameMesh grid = frames[f]; + const AxisStencil* axes = stencils + 3 * node; + const int64_t points = mesh_points(grid); + const float* frame_mesh = potential_mesh + grid.offset; + // Every lane sums its points in float32, the warp the lane sums in + // float64. ``gradient`` is the derivative of sum_c q_c phi_c with + // respect to u_d. + float lane_phi[kMaxChannels] = {}; + float lane_gradient[3] = {0.f, 0.f, 0.f}; + for (int t = lane; t < kOrder * kOrder * kOrder; t += kWarp) { + const StencilPoint point(t, grid, axes); + const float v1 = axes[0].value[point.k1]; + const float v2 = axes[1].value[point.k2]; + const float v3 = axes[2].value[point.k3]; + float charge_potential = 0.f; +#pragma unroll + for (int c = 0; c < kMaxChannels; ++c) { + if (c < channels) { + const float value = frame_mesh[c * points + point.offset]; + lane_phi[c] += v1 * v2 * v3 * value; + charge_potential += q[c] * value; + } + } + lane_gradient[0] += + axes[0].slope[point.k1] * v2 * v3 * charge_potential; + lane_gradient[1] += + v1 * axes[1].slope[point.k2] * v3 * charge_potential; + lane_gradient[2] += + v1 * v2 * axes[2].slope[point.k3] * charge_potential; + } +#pragma unroll + for (int c = 0; c < kMaxChannels; ++c) { + if (c < channels) { + phi[c] = warp_sum(lane_phi[c]) - self * q[c]; + } + } + double gradient[3]; +#pragma unroll + for (int d = 0; d < 3; ++d) { + gradient[d] = warp_sum(lane_gradient[d]) * grid.n[d]; + } + // d E / d r_j = sum_d d E / d u_d n_d (B^-1)_jd +#pragma unroll + for (int j = 0; j < 3; ++j) { + f_node[j] = -(gradient[0] * grid.inverse[3 * j] + + gradient[1] * grid.inverse[3 * j + 1] + + gradient[2] * grid.inverse[3 * j + 2]); + } + } + if (lane == 0) { + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + potential[int64_t{owner} * channels + c] = phi[c]; + energy += q[c] * phi[c]; + } + atom_energy[owner] = 0.5 * energy; + const double* r = coord + 3 * node; + for (int a = 0; a < 3; ++a) { + force[int64_t{owner} * 3 + a] = f_node[a]; + for (int b = 0; b < 3; ++b) { + moment[int64_t{owner} * 9 + 3 * a + b] = f_node[a] * r[b]; + } + } + } + } +} + +// Per-frame sums over one slice of the owners: the potential of the charged +// owners in the constrained channel, the number of charged owners and the +// nine components of ``moment``. +constexpr int kOwnerSums = 11; + +__global__ void __launch_bounds__(kFrameThreads) + owner_slices(const double* __restrict__ potential, + const double* __restrict__ moment, + const bool* __restrict__ owner_mask, + const int64_t* __restrict__ owner_frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + const NodeRange slice = slice_of_frame(owner_frame_ptr); + double value[kOwnerSums] = {}; + for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end; + owner += kFrameThreads) { + if (!owner_mask[owner]) { + continue; + } + if (channel >= 0) { + value[0] += potential[owner * channels + channel]; + } + value[1] += 1.0; + for (int k = 0; k < 9; ++k) { + value[2 + k] += moment[owner * 9 + k]; + } + } + write_slice_sum(value, slice_sum); +} + +// Adds the slice sums of owner_slices of every frame in slice order. +__global__ void owner_totals(const double* __restrict__ slice_sum, + int slice_count, + double* __restrict__ total) { + if (threadIdx.x >= kOwnerSums) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * + kOwnerSums + + threadIdx.x]; + } + total[blockIdx.x * kOwnerSums + threadIdx.x] = value; +} + +// Writes the outputs of every owner: the charge gradient w (phi - [c = ch] +// mean), the force, and the virial F (x) r plus the share of the rest of the +// virial of the frame over its charged nodes, then the owner sums of that +// virial per frame. +__global__ void finish_owners(const double* __restrict__ potential, + const double* __restrict__ force, + const double* __restrict__ moment, + const bool* __restrict__ owner_mask, + const int64_t* __restrict__ owner_frame_ptr, + const double* __restrict__ frame_virial, + const double* __restrict__ owner_total, + const double* __restrict__ frame_stat, + int frame_count, + int owner_count, + int channels, + int channel, + float* __restrict__ charge_gradient, + float* __restrict__ force_out, + float* __restrict__ atom_virial) { + for (int owner = blockIdx.x * blockDim.x + threadIdx.x; owner < owner_count; + owner += gridDim.x * blockDim.x) { + const int f = frame_of(owner_frame_ptr, frame_count, owner); + const bool charged = f >= 0 && owner_mask[owner]; + // The charged nodes of the frame, of every domain. + const double count = + charged ? frame_stat[f * kStatWidth + kStatCount] : 1.0; + const double mean = + charged && channel >= 0 ? owner_total[f * kOwnerSums] / count : 0.0; + for (int c = 0; c < channels; ++c) { + const double phi = + charged ? potential[int64_t{owner} * channels + c] : 0.0; + charge_gradient[int64_t{owner} * channels + c] = + static_cast(c == channel ? phi - mean : phi); + } + for (int a = 0; a < 3; ++a) { + force_out[int64_t{owner} * 3 + a] = + static_cast(force[int64_t{owner} * 3 + a]); + } + for (int k = 0; k < 9; ++k) { + const double rest = charged ? (frame_virial[f * 9 + k] - + owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + atom_virial[int64_t{owner} * 9 + k] = + static_cast(moment[int64_t{owner} * 9 + k] + rest); + } + } +} + +// Sums, over one slice of the owners of one frame, their virial. +__global__ void __launch_bounds__(kFrameThreads) + virial_slices(const float* __restrict__ atom_virial, + const int64_t* __restrict__ owner_frame_ptr, + double* __restrict__ slice_sum) { + const NodeRange slice = slice_of_frame(owner_frame_ptr); + double value[9] = {}; + for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end; + owner += kFrameThreads) { + for (int k = 0; k < 9; ++k) { + value[k] += atom_virial[owner * 9 + k]; + } + } + write_slice_sum<9>(value, slice_sum); +} + +__global__ void virial_total(const double* __restrict__ slice_sum, + int slice_count, + float* __restrict__ virial) { + if (threadIdx.x >= 9) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * 9 + + threadIdx.x]; + } + virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value); +} + +// The weights M_p(1), ..., M_p(p - 1) of the assignment factor. +std::vector integer_bspline_weights(int order) { + std::vector weight(order - 1); + for (int l = 0; l < order - 1; ++l) { + const double x = l + 1.0; + double value = 0.0; + double binomial = 1.0; + for (int j = 0; j <= order; ++j) { + if (j > 0) { + binomial = binomial * (order - j + 1) / j; + } + const double base = std::max(x - j, 0.0); + value += (j % 2 == 0 ? 1.0 : -1.0) * binomial * std::pow(base, order - 1); + } + double factorial = 1.0; + for (int k = 2; k < order; ++k) { + factorial *= k; + } + weight[l] = value / factorial; + } + return weight; +} + +int blocks_for(int64_t work) { + return static_cast(std::max( + 1, std::min((work + kThreads - 1) / kThreads, 65535))); +} + +// Blocks of a stencil kernel over ``nodes`` nodes, one warp each. +int stencil_blocks_for(int64_t nodes) { + return static_cast(std::max( + 1, std::min((nodes + kStencilWarps - 1) / kStencilWarps, + int64_t{1} << 20))); +} + +// Calls ``launch`` with the B-spline order as a compile-time constant, for +// which the stencil kernels unroll their loops. +template +void dispatch_order(int64_t order, const Launch& launch) { + switch (order) { + case 2: + return launch(std::integral_constant{}); + case 3: + return launch(std::integral_constant{}); + case 4: + return launch(std::integral_constant{}); + case 5: + return launch(std::integral_constant{}); + case 6: + return launch(std::integral_constant{}); + case 7: + return launch(std::integral_constant{}); + case 8: + return launch(std::integral_constant{}); + case 9: + return launch(std::integral_constant{}); + case 10: + return launch(std::integral_constant{}); + case 11: + return launch(std::integral_constant{}); + case 12: + return launch(std::integral_constant{}); + default: + TORCH_CHECK(false, "latent_charge_sum: unsupported B-spline order ", + order); + } +} + +} // namespace + +namespace deepmd { +namespace long_range { + +ChargeSumOutputs mesh_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& box, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) { + TORCH_CHECK(order >= 2 && order <= kMaxOrder, + "latent_charge_sum: the B-spline order must lie in [2, ", + kMaxOrder, "], got ", order); + TORCH_CHECK(std::isfinite(spacing) && spacing > 0.0, + "latent_charge_sum: the mesh spacing must be positive, got ", + spacing); + const int64_t node_count = coord.size(0); + TORCH_CHECK(node_count < std::numeric_limits::max(), + "latent_charge_sum: too many nodes"); + const int frame_count = static_cast(n_node_per_frame.size(0)); + const int channels = static_cast(charge.size(1)); + const int gaussians = + kernel == kSog ? static_cast(coefficient.numel() / 2) : 0; + const auto double_options = coord.options(); + const auto float_options = charge.options(); + if (node_count == 0 || frame_count == 0) { + return {torch::zeros({owner_count, 1}, double_options), + torch::zeros({owner_count, channels}, float_options), + torch::zeros({owner_count, 3}, float_options), + torch::zeros({owner_count, 9}, float_options), + torch::zeros({frame_count, 9}, float_options)}; + } + const bool every_node = owner_begin == 0 && owner_count == node_count; + TORCH_CHECK(every_node || frame_count == 1, + "latent_charge_sum: owners other than every node lie in one " + "frame"); + const auto stream = at::cuda::getCurrentCUDAStream(); + + // === Step 1. The meshes of the frames, sized on the host === + std::vector frames(frame_count); + const double* cells = box.data_ptr(); + int64_t mesh_total = 0; + for (int f = 0; f < frame_count; ++f) { + frames[f] = frame_mesh(cells + 9 * f, spacing); + frames[f].offset = mesh_total; + mesh_total += channels * mesh_points(frames[f]); + } + const auto byte_options = torch::TensorOptions().dtype(torch::kUInt8); + auto host_frames = + torch::empty({static_cast(sizeof(FrameMesh)) * frame_count}, + byte_options.pinned_memory(true)); + std::memcpy(host_frames.data_ptr(), frames.data(), + sizeof(FrameMesh) * frame_count); + const auto device_frames = + host_frames.to(coord.device(), /*non_blocking=*/true); + const auto* frames_ptr = + reinterpret_cast(device_frames.data_ptr()); + const std::vector weight_host = integer_bspline_weights(order); + auto host_weight = torch::from_blob(const_cast(weight_host.data()), + {order - 1}, torch::kFloat64) + .pin_memory(); + const auto integer_weight = + host_weight.to(coord.device(), /*non_blocking=*/true); + + // === Step 2. The constraint and the charged nodes of every frame === + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + auto frame_ptr = torch::empty({frame_count + 1}, counts.options()); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>( + counts.data_ptr(), frame_count, frame_ptr.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; + const auto slices_for = [](int64_t nodes) { + return static_cast(std::min( + kMaxSlices, + std::max(1, (nodes + kSliceNodes - 1) / kSliceNodes))); + }; + const int slice_count = slices_for(average_nodes); + auto slice_charge = torch::empty( + {static_cast(frame_count) * slice_count * 2}, double_options); + auto frame_stat = torch::empty({frame_count * kStatWidth}, double_options); + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), + frame_ptr.data_ptr(), channels, + static_cast(std::max(channel, 0)), + slice_charge.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge.data_ptr(), slice_count, frame_count, + total_charge.data_ptr(), frame_stat.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // === Step 3. The stencils of the nodes and the charge meshes === + auto mesh = torch::zeros({mesh_total}, double_options.dtype(torch::kInt64)); + auto node_charge = torch::empty({node_count, channels}, float_options); + torch::Tensor stencils; + dispatch_order(order, [&](auto order_constant) { + constexpr int kOrder = decltype(order_constant)::value; + stencils = torch::empty( + {3 * node_count * static_cast(sizeof(AxisStencil))}, + byte_options.device(coord.device())); + auto* node_stencil = + reinterpret_cast*>(stencils.data_ptr()); + node_stencils<<>>( + coord.data_ptr(), node_mask.data_ptr(), + frame_ptr.data_ptr(), frames_ptr, frame_count, node_count, + node_stencil); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + spread_charges + <<>>( + charge.data_ptr(), node_mask.data_ptr(), + frame_ptr.data_ptr(), frames_ptr, node_stencil, + frame_stat.data_ptr(), frame_count, node_count, channels, + static_cast(channel), node_charge.data_ptr(), + reinterpret_cast(mesh.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + fixed_to_double<<>>( + mesh_total, + reinterpret_cast(mesh.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const auto charge_mesh = mesh.view(torch::kFloat64); + + // === Step 4. Influence function and potential meshes, frame by frame === + auto potential_mesh = torch::empty({mesh_total}, double_options); + auto frame_virial = torch::zeros({frame_count, 9}, double_options); + for (int f = 0; f < frame_count; ++f) { + const FrameMesh& grid = frames[f]; + const int64_t points = mesh_points(grid); + const std::vector shape = {channels, grid.n[0], grid.n[1], + grid.n[2]}; + auto spectrum = at::fft_rfftn( + charge_mesh.narrow(0, grid.offset, channels * points).view(shape), + c10::nullopt, std::vector{1, 2, 3}); + const int half = grid.n[2] / 2 + 1; + auto factor = torch::empty({grid.n[0] + grid.n[1] + half}, double_options); + assignment_factor<<>>( + grid, static_cast(order), integer_weight.data_ptr(), + factor.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t modes = int64_t{grid.n[0]} * grid.n[1] * half; + // A few blocks per multiprocessor fill the device; every thread then + // walks several modes, which keeps the block sums few. + const int mode_blocks = std::min( + blocks_for(modes), + 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount); + auto block_sum = + torch::empty({int64_t{mode_blocks} * kModeSums}, double_options); + apply_influence<<>>( + grid, channels, static_cast(kernel), + coefficient.data_ptr(), gaussians, prefactor, + factor.data_ptr(), + reinterpret_cast(spectrum.data_ptr>()), + block_sum.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + mode_virial<<<1, kThreads, 0, stream>>>(block_sum.data_ptr(), + mode_blocks, + frame_virial[f].data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto frame_potential = + potential_mesh.narrow(0, grid.offset, channels * points).view(shape); + at::fft_irfftn_out(frame_potential, spectrum, + std::vector{grid.n[0], grid.n[1], grid.n[2]}, + std::vector{1, 2, 3}, "forward"); + } + + // === Step 5. Potential, force and virial of the owners === + // The stencils read the potential meshes in float32, the precision of the + // outputs. + const auto potential_mesh_float = potential_mesh.to(torch::kFloat32); + auto potential = torch::empty({owner_count, channels}, double_options); + auto force = torch::empty({owner_count, 3}, double_options); + auto moment = torch::empty({owner_count, 9}, double_options); + auto atom_energy = torch::empty({owner_count, 1}, double_options); + if (owner_count > 0) { + dispatch_order(order, [&](auto order_constant) { + constexpr int kOrder = decltype(order_constant)::value; + interpolate_owners<<>>( + coord.data_ptr(), node_charge.data_ptr(), + node_mask.data_ptr(), frame_ptr.data_ptr(), frames_ptr, + reinterpret_cast*>(stencils.data_ptr()), + potential_mesh_float.data_ptr(), + coefficient.data_ptr(), frame_count, owner_begin, owner_count, + channels, static_cast(kernel), gaussians, prefactor, + potential.data_ptr(), force.data_ptr(), + moment.data_ptr(), atom_energy.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // The frames of the owners, on the axis of the owners: the frames + // themselves when every node is an owner, else the one frame they share. + torch::Tensor owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_ptr = + torch::full({2}, int64_t{owner_count}, frame_ptr.options()); + owner_frame_ptr.select(0, 0).zero_(); + } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const int owner_slice_count = slices_for( + (int64_t{owner_count} + frame_count - 1) / std::max(frame_count, 1)); + auto owner_slice = torch::empty( + {int64_t{frame_count} * owner_slice_count * kOwnerSums}, double_options); + owner_slices<<>>( + potential.data_ptr(), moment.data_ptr(), owner_mask, + owner_frame_ptr.data_ptr(), channels, static_cast(channel), + owner_slice.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto owner_total = + torch::empty({int64_t{frame_count} * kOwnerSums}, double_options); + owner_totals<<>>(owner_slice.data_ptr(), + owner_slice_count, + owner_total.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (communicator != nullptr) { + // The mean potential of the constraint and the virial share run over the + // charged nodes of every domain. + owner_total.copy_(sum_over_domains(*communicator, owner_total)); + } + auto charge_gradient = torch::empty({owner_count, channels}, float_options); + auto force_out = torch::empty({owner_count, 3}, float_options); + auto atom_virial = torch::empty({owner_count, 9}, float_options); + if (owner_count > 0) { + finish_owners<<>>( + potential.data_ptr(), force.data_ptr(), + moment.data_ptr(), owner_mask, + owner_frame_ptr.data_ptr(), frame_virial.data_ptr(), + owner_total.data_ptr(), frame_stat.data_ptr(), + frame_count, owner_count, channels, static_cast(channel), + charge_gradient.data_ptr(), force_out.data_ptr(), + atom_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + auto virial_slice = torch::empty( + {int64_t{frame_count} * owner_slice_count * 9}, double_options); + virial_slices<<>>(atom_virial.data_ptr(), + owner_frame_ptr.data_ptr(), + virial_slice.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto virial = torch::empty({frame_count, 9}, float_options); + virial_total<<>>(virial_slice.data_ptr(), + owner_slice_count, + virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {atom_energy, charge_gradient, force_out, atom_virial, virial}; +} + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc index dd312bc555..364fa6c2c2 100644 --- a/source/op/pt/long_range/ops.cc +++ b/source/op/pt/long_range/ops.cc @@ -1,48 +1,67 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // -// Operator schema of the latent-charge long-range direct sum. +// Operator schema of the latent-charge long-range energy. // // The schema is declared here, unconditionally, while the CUDA kernel -// registers itself only against a CUDA-enabled PyTorch (direct_sum.cu), so -// that the Python front end can tell "library not loaded" from "this device +// registers itself only against a CUDA-enabled PyTorch (latent_charge_sum.cu), +// so that the Python front end can tell "library not loaded" from "this device // has no kernel". // -// For the nodes of one frame, the operator takes the raw charges c_i of the +// For the nodes of the frames, the operator takes the raw charges c_i of the // charge network in C float32 channels, the nodes that carry them (w_i = 1, -// and 0 elsewhere) and float64 positions r_i. The pairs carry the charges -// q_i = w_i c_i; when `channel` is a channel index rather than -1, that -// channel is shifted uniformly over the charged nodes of every frame so that -// it sums to the total charge Q_f of the frame, +// and 0 elsewhere) and float64 positions r_i. The interactions carry the +// charges q_i = w_i c_i; when `channel` is a channel index rather than -1, +// that channel is shifted uniformly over the charged nodes of every frame so +// that it sums to the total charge Q_f of the frame, // // q_i^ch = w_i (c_i^ch - s_f), s_f = (sum_{j in f} w_j c_j^ch - Q_f) / // max(n_f, 1), // -// with n_f charged nodes in f. With the potential of the charges +// with n_f charged nodes in f. Without `box`, every frame is a complete +// non-periodic system, in which the potential of the charges is the direct +// sum // -// phi_i = prefactor * sum_{j != i} q_j K(r_ij), +// phi_i = prefactor * sum_{j != i} q_j K(r_ij); // +// with `box`, whose row f holds the row-major cell of frame f, the rows of the +// cell being its lattice vectors, phi_i sums over the images of the cell, by +// the smooth particle-mesh Ewald method (mesh_sum.cu), less the potential +// prefactor * K(0) q_i of every charge at itself. An all-zero row marks a +// non-periodic frame, which a call does not mix with periodic ones. Either way // the operator returns // // atom_energy_i = q_i . phi_i / 2 (float64) // charge_gradient_i = dE / dc_i = w_i (phi_i - [c = ch] m_f), m_f the // mean of phi^ch over the charged nodes of f -// force_i = -prefactor * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij -// * d_ij -// W_i = -prefactor / 2 * sum_{j != i} (q_i . q_j) -// K'(r_ij) / r_ij * d_ij (x) d_ij -// W_f = sum_{i in f} W_i +// force_i = -dE / dr_i at fixed charges +// atom_virial_i (9 components, row-major) +// virial_f = sum_{i in f} atom_virial_i // -// in float32 unless noted, with d_ij = r_i - r_j and the sums running over -// the nodes of the frame of i. The frames occupy the leading -// sum(n_node_per_frame) <= N nodes, which the caller ensures; the nodes past -// them are padding and join none. The frame energy is the sum of the atomic -// energies. `kernel` selects K: 0 is -// erf(alpha r) / r with `coefficient` holding alpha; 1 is a sum of Gaussians -// a_k 2^(-c_k r^2), c_k = log2(e) / (2 sigma_k^2), with `coefficient` holding -// the amplitudes a_k and then the rates c_k. `total_charge` holds Q_f in -// float64 and is read only with a constrained channel. The force and the -// virial hold the charges fixed; the response of the charges is the -// caller's, through the charge gradient. +// in float32 unless noted. In a non-periodic frame, with d_ij = r_i - r_j, +// +// force_i = -prefactor * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij * d_ij, +// W_i = -prefactor / 2 * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij +// * d_ij (x) d_ij; +// +// in a periodic frame the virial of the frame is that of the strain of the +// cell, W_f = -dE / d(strain), and W_i = F_i (x) r_i + (W_f - sum_{j in f} +// F_j (x) r_j) / n_f on the charged nodes, the force index first, so that the +// charged nodes of the frame sum to W_f. +// +// The frames occupy the leading sum(n_node_per_frame) <= N nodes, which the +// caller ensures; the nodes past them are padding and join none. The frame +// energy is the sum of the atomic energies. +// +// `kernel` selects K: 0 is erf(alpha r) / r with `coefficient` holding alpha; +// 1 is a sum of Gaussians a_k 2^(-c_k r^2), c_k = log2(e) / (2 sigma_k^2), +// with `coefficient` holding the amplitudes a_k and then the rates c_k. +// `coefficient` is float64. `total_charge` holds Q_f in float64 and is read +// only with a constrained channel. `box` is a float64 CPU tensor of shape +// (F, 9). The mesh of a periodic frame holds along every lattice vector a_d +// the smallest even number of points without a prime factor above 7 of at +// least ceil(|a_d| / spacing), and spreads the charges with B-splines of +// order `order`. The force and the virial hold the charges fixed; the response +// of the charges is the caller's, through the charge gradient. // // `communicator` is an int64 CPU tensor of one element holding the address of // the caller's MPI_Comm handle, or zero. A handle of several ranks decomposes @@ -55,10 +74,10 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { library.def( - "latent_charge_direct_sum(Tensor coord, Tensor charge, " - "Tensor node_mask, Tensor n_node_per_frame, Tensor total_charge, " + "latent_charge_sum(Tensor coord, Tensor charge, Tensor node_mask, " + "Tensor n_node_per_frame, Tensor total_charge, Tensor? box, " "Tensor coefficient, Tensor communicator, int kernel, int channel, " - "float prefactor) -> " + "float prefactor, int order, float spacing) -> " "(Tensor atom_energy, Tensor charge_gradient, Tensor force, " "Tensor atom_virial, Tensor virial)"); } @@ -66,6 +85,5 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { // The operator serves inference, whose caller differentiates nothing through // it; its outputs carry no gradient. TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { - library.impl("latent_charge_direct_sum", - torch::CppFunction::makeFallthrough()); + library.impl("latent_charge_sum", torch::CppFunction::makeFallthrough()); } diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index b8976801cc..2f826a1e3b 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -19,6 +19,9 @@ LESKernel, SOGKernel, ) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( build_k_grid, build_pair_graph, @@ -265,7 +268,7 @@ def test_kernel_contact_value_matches_real_space_limit(kernel: str) -> None: def test_build_k_grid_orthogonality_and_guards() -> None: """The k points satisfy k . a_d = 2 pi n_d, are ordered and carry guards.""" box = np.array([[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]])[None] - grid = build_k_grid(box, 2.0, np.array([4]), coord=np.zeros((4, 3))) + grid = build_k_grid(box, 2.0, np.zeros((4, 3))) # guard points are two masked (1e-6, 0, 0) entries assert grid.mask[-2:].tolist() == [False, False] assert grid.mask[:-2].all() @@ -319,6 +322,27 @@ def test_build_topologies_periodic_branching() -> None: # the k grid needs its build parameter with pytest.raises(ValueError, match="k_max"): build_topologies(names, coord, n_node, box, atype=atype) + # the reciprocal-space sum excludes no pair; a non-periodic batch still does + pair_excl = PairExcludeMask(2, [(0, 1)]) + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + build_topologies( + names, + coord, + n_node, + box, + atype=atype, + pair_excl=pair_excl, + topology_params=params, + ) + excluded = build_topologies( + names, coord, n_node, None, atype=atype, pair_excl=pair_excl + ) + assert 0 < excluded[ALL_PAIRS].mask.sum() < without_box[ALL_PAIRS].mask.sum() + # a cell without volume is no periodic system + with pytest.raises(ValueError, match="no volume"): + build_topologies( + names, coord, n_node, np.zeros((1, 9)), atype=atype, topology_params=params + ) # --------------------------------------------------------------------------- @@ -326,6 +350,24 @@ def test_build_topologies_periodic_branching() -> None: # --------------------------------------------------------------------------- +def _module_topologies( + module: LatentChargeLongRange, + coord: np.ndarray, + n_node: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, +) -> dict: + """The topologies of a module, built with its parameters.""" + return build_topologies( + module.required_topologies(), + coord, + n_node, + box, + atype=atype, + topology_params=module.topology_params(), + ) + + def _periodic_long_range(**kwargs: object) -> LatentChargeLongRange: charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) return LatentChargeLongRange( @@ -366,13 +408,13 @@ def test_fixed_k_grid_matches_ragged_enumeration() -> None: box = _cells() k_max = 2.0 - n_grid = enumerate_k_grid_bounds(np.max(k_grid_bounds_of_cells(box, k_max), axis=0)) - nk = n_grid.shape[0] - fixed = build_fixed_k_grid(box, n_grid, k_max, np.array([4, 5])) + bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + nk = enumerate_k_grid_bounds(bounds).shape[0] + fixed = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3))) # a static, guard-free layout: nf * nk points, frame-major assert fixed.vec.shape == (2 * nk, 3) np.testing.assert_array_equal(fixed.frame, np.repeat([0, 1], nk)) - ragged = build_k_grid(box, k_max, np.array([4, 5])) + ragged = build_k_grid(box, k_max, np.zeros((9, 3))) for frame in range(2): expected = { tuple(row) @@ -398,10 +440,10 @@ def test_fixed_k_grid_truncates_oversized_cells() -> None: box = _cells() k_max = 2.0 - n_grid = enumerate_k_grid_bounds(np.max(k_grid_bounds_of_cells(box, k_max), axis=0)) - nk = n_grid.shape[0] + bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + nk = enumerate_k_grid_bounds(bounds).shape[0] with pytest.warns(UserWarning, match="truncated"): - bigger = build_fixed_k_grid(box * 2.0, n_grid, k_max, np.array([4, 5])) + bigger = build_fixed_k_grid(box * 2.0, bounds, k_max, np.zeros((9, 3))) # the layout is untouched: static, guard-free, every grid point present assert bigger.vec.shape == (2 * nk, 3) np.testing.assert_array_equal(bigger.frame, np.repeat([0, 1], nk)) @@ -410,11 +452,11 @@ def test_fixed_k_grid_truncates_oversized_cells() -> None: k2 = np.sum(bigger.vec**2, axis=-1) np.testing.assert_array_equal(bigger.mask, (k2 > 0) & (k2 <= k_max * k_max)) # the larger cell activates more of the shared grid than the original did - original = build_fixed_k_grid(box, n_grid, k_max, np.array([4, 5])) + original = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3))) assert bigger.mask.sum() > original.mask.sum() # a larger wave-number cutoff (a drifting kernel) truncates the same way with pytest.warns(UserWarning, match="truncated"): - build_fixed_k_grid(box, n_grid, k_max * 2.0, np.array([4, 5])) + build_fixed_k_grid(box, bounds, k_max * 2.0, np.zeros((9, 3))) def test_fixed_k_grid_mode_configuration() -> None: @@ -432,7 +474,9 @@ def test_fixed_k_grid_mode_configuration() -> None: _periodic_long_range(k_grid_bounds=[3, 3, 3]) # explicit bounds build the grid at construction, no statistics needed long_range = _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[2, 2, 2]) - assert long_range.k_grid_n.shape == (5**3 - 1, 3) + np.testing.assert_array_equal(long_range.k_grid_bounds, [2, 2, 2]) + with pytest.raises(ValueError, match="non-negative"): + _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[0, 0, 0]) def test_fixed_k_grid_stats_size_the_grid() -> None: @@ -440,10 +484,9 @@ def test_fixed_k_grid_stats_size_the_grid() -> None: long_range = _periodic_long_range(k_grid_mode="fixed") # the grid is not sized before the statistics run coord = np.zeros((9, 3)) + atype = np.zeros(9, dtype=np.int64) with pytest.raises(ValueError, match="not sized yet"): - long_range.build_topologies( - coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) - ) + _module_topologies(long_range, coord, np.array([4, 5]), _cells(), atype) long_range.compute_input_stats( [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}] ) @@ -453,11 +496,11 @@ def test_fixed_k_grid_stats_size_the_grid() -> None: ) expected = np.max(k_grid_bounds_of_cells(_cells(), k_max), axis=0) - np.testing.assert_array_equal(np.max(np.abs(long_range.k_grid_n), axis=0), expected) - topologies = long_range.build_topologies( - coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) + np.testing.assert_array_equal(long_range.k_grid_bounds, expected) + topologies = _module_topologies( + long_range, coord, np.array([4, 5]), _cells(), atype ) - assert topologies["k_grid"].vec.shape == (2 * long_range.k_grid_n.shape[0], 3) + assert topologies["k_grid"].vec.shape == (2 * (np.prod(2 * expected + 1) - 1), 3) # statistics without a cell cannot size the grid unsized = _periodic_long_range(k_grid_mode="fixed") with pytest.raises(ValueError, match="no simulation cell"): @@ -472,11 +515,11 @@ def test_fixed_k_grid_round_trip() -> None: assert record["k_grid_mode"] == "fixed" restored = LatentChargeLongRange.deserialize(record) assert restored.k_grid_mode == "fixed" - np.testing.assert_array_equal(restored.k_grid_n, long_range.k_grid_n) + np.testing.assert_array_equal(restored.k_grid_bounds, long_range.k_grid_bounds) # the default mode serializes no grid fields plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) assert plain.k_grid_mode == "ragged" - assert plain.k_grid_n is None + assert plain.k_grid_bounds is None # --------------------------------------------------------------------------- @@ -510,14 +553,14 @@ def test_decoupled_k_grid_ignores_kernel_drift() -> None: assert decoupled.frozen_k_max == pytest.approx( LESKernel(alpha=0.3).kspace_cutoff(1e-8) ) - initial = decoupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] - reference = coupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + initial = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"] + reference = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] np.testing.assert_array_equal(initial.vec, reference.vec) # a larger alpha moves the cutoff out, growing the coupled grid for module in (decoupled, coupled): module.kernel.log_alpha = np.log(1.5) - drifted = decoupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] - moved = coupled.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + drifted = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"] + moved = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] np.testing.assert_array_equal(drifted.vec, initial.vec) np.testing.assert_array_equal(drifted.mask, initial.mask) assert moved.vec.shape[0] > initial.vec.shape[0] @@ -533,10 +576,10 @@ def test_decoupled_fixed_grid_keeps_its_mask() -> None: long_range = _periodic_long_range( k_grid_mode="fixed", k_grid_bounds=[5, 5, 5], decouple_k_grid=True ) - before = long_range.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + before = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"] # a smaller alpha moves the cutoff in; a coupled mask would tighten long_range.kernel.log_alpha = np.log(0.05) - after = long_range.build_topologies(coord, n_node, box, atype=atype)["k_grid"] + after = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"] np.testing.assert_array_equal(after.mask, before.mask) np.testing.assert_array_equal(after.vec, before.vec) @@ -571,18 +614,19 @@ def _pme_long_range(**kwargs: object) -> LatentChargeLongRange: def test_pme_mesh_grid_build_and_flatten() -> None: - """The mesh grid packs the cells, flags the frames, and flattens to 4 tensors.""" + """The mesh grid packs the cells and the mesh, flags the frames and flattens.""" from deepmd.dpmodel.utils.neighbor_graph import ( PME_MESH, build_topologies, flatten_topologies, + mesh_dims_of_cells, unflatten_topologies, ) coord = np.random.default_rng(1).normal(size=(4, 3)) atype = np.array([0, 1, 1, 0]) box = _cells() - params = {PME_MESH: {"n_mesh": np.array([16, 16, 16]), "spacing": 1.0}} + params = {PME_MESH: {"spacing": 1.0}} tops = build_topologies( ("all_pairs", PME_MESH), coord, @@ -594,12 +638,21 @@ def test_pme_mesh_grid_build_and_flatten() -> None: mesh = tops[PME_MESH] np.testing.assert_array_equal(mesh.box, box) assert mesh.mask.shape == (2,) and mesh.mask.all() + dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)) + assert dims == mesh_dims_of_cells(box, 1.0) + # the modes of an axis are its integer frequencies in FFT order + n1 = dims[0] + expected = np.arange(n1, dtype=float) + np.testing.assert_array_equal( + mesh.mode_1, np.where(expected > n1 // 2, expected - n1, expected) + ) assert tops["all_pairs"].index.shape[1] == 2 # guards only flat = flatten_topologies(tops) - assert len(flat) == 3 + 5 + assert len(flat) == 3 + 8 back = unflatten_topologies(("all_pairs", PME_MESH), flat) np.testing.assert_array_equal(back[PME_MESH].box, box) np.testing.assert_array_equal(back[PME_MESH].mask, mesh.mask) + np.testing.assert_array_equal(back[PME_MESH].mode_3, mesh.mode_3) # a non-periodic batch: real pairs, masked-out mesh tops_np = build_topologies( ("all_pairs", PME_MESH), @@ -612,39 +665,51 @@ def test_pme_mesh_grid_build_and_flatten() -> None: assert tops_np["all_pairs"].mask.any() assert not tops_np[PME_MESH].mask.any() # missing build parameters are an error - with pytest.raises(ValueError, match="n_mesh"): + with pytest.raises(ValueError, match="spacing"): build_topologies((PME_MESH,), coord, np.array([4]), box, atype=atype) -def test_pme_fft_smooth() -> None: - from deepmd.dpmodel.long_range.latent_charge import ( - _fft_smooth, +def test_pme_mesh_dims_resolve_every_cell() -> None: + """The mesh is the smallest even 7-smooth one resolving every cell.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, ) - assert _fft_smooth(1) == 1 - assert _fft_smooth(17) == 18 - assert _fft_smooth(19) == 20 - assert _fft_smooth(23) == 24 - assert _fft_smooth(29) == 30 - assert _fft_smooth(31) == 32 - assert _fft_smooth(37) == 40 + for spacing in (0.37, 0.5, 1.1): + dims = mesh_dims_of_cells(_cells(), spacing) + longest = np.max(np.linalg.norm(_cells(), axis=-1), axis=0) + for n, a in zip(dims, longest, strict=True): + assert n % 2 == 0 and n >= a / spacing + m = n + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + assert m == 1 + # no smaller even 7-smooth size reaches the spacing + for smaller in range(n - 2, 1, -2): + m = smaller + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + if m == 1: + assert smaller < a / spacing + break + # the larger cell alone sizes the mesh of the batch + assert mesh_dims_of_cells(_cells(), 0.5) == mesh_dims_of_cells(_cells()[1:], 0.5) def _pme_energy(module, charge, coord, box) -> float: + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) from deepmd.dpmodel.utils.neighbor_graph import ( - PME_MESH, - frame_id_from_n_node, + build_mesh_grid, ) - n_node = np.array([coord.shape[0]]) - tops = module.build_topologies( - coord, n_node, box[None], atype=np.zeros(len(coord), dtype=np.int64) - ) - fid = frame_id_from_n_node(n_node, n_total=coord.shape[0]) - e = module._pme_branch( - charge, tops[PME_MESH], fid, np.ones(len(coord), dtype=bool), n_node - ) - return float(np.sum(e)) + mesh = build_mesh_grid(box[None], module.mesh_spacing(), coord) + frame_id = np.zeros(coord.shape[0], dtype=np.int64) + energy = mesh_energy(charge, mesh, frame_id, module.kernel, module.pme_order) + return float(np.sum(energy)) def test_pme_energy_converges_to_direct_reference() -> None: @@ -655,14 +720,6 @@ def test_pme_energy_converges_to_direct_reference() -> None: charge = charge - charge.mean(axis=0, keepdims=True) kernel = LESKernel(alpha=0.3) - module = LatentChargeLongRange( - LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), - kernel, - None, - periodic=True, - k_grid_mode="pme", - pme_mesh=[40, 40, 40], - ) # converged direct k-space reference (tol 1e-10) k_max = kernel.kspace_cutoff(1e-10) rec = 2.0 * np.pi * np.linalg.inv(box) @@ -683,52 +740,47 @@ def test_pme_energy_converges_to_direct_reference() -> None: (k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum() ) errors = [] - for dims in ([16, 16, 16], [40, 40, 40]): - module.pme_mesh = np.asarray(dims) - module._pme_n = tuple(dims) + for spacing in (2.0, None): + module = LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + kernel, + None, + periodic=True, + k_grid_mode="pme", + pme_spacing=spacing, + ) errors.append(abs(_pme_energy(module, charge, coord, box) - e_ref)) assert errors[1] < errors[0] # refinement improves the energy - assert errors[1] < 1e-2 + # the default mesh reaches the accuracy it is chosen for + assert errors[1] < 1e-6 * abs(e_ref) -def test_pme_mesh_stats_size_the_mesh() -> None: - """The data statistics size the mesh by the largest sampled cell.""" - import math +def test_pme_mesh_follows_the_cells_of_the_batch() -> None: + """The mesh of a batch resolves its cells at the spacing, without statistics.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, + ) long_range = _pme_long_range() coord = np.zeros((9, 3)) - with pytest.raises(ValueError, match="not sized yet"): - long_range.build_topologies( - coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) - ) - long_range.compute_input_stats( - [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}] + atype = np.zeros(9, dtype=np.int64) + small = _module_topologies(long_range, coord, np.array([4, 5]), _cells()[:1], atype) + large = _module_topologies( + long_range, coord, np.array([4, 5]), 2.0 * _cells(), atype ) - spacing = math.pi / long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) - a_max = np.max(np.linalg.norm(_cells(), axis=-1), axis=0) - for d in range(3): - assert long_range.pme_mesh[d] >= math.ceil(a_max[d] / spacing) - # FFT-smooth dimensions - m = int(long_range.pme_mesh[d]) - for p in (2, 3, 5, 7): - while m % p == 0: - m //= p - assert m == 1 - topologies = long_range.build_topologies( - coord, np.array([4, 5]), _cells(), atype=np.zeros(9, dtype=np.int64) - ) - assert topologies["pme_mesh"].mask.all() - # explicit mesh skips the statistics - explicit = _pme_long_range(pme_mesh=[16, 16, 16]) - assert explicit._pme_dims() == (16, 16, 16) - # statistics without a cell cannot size the mesh - unsized = _pme_long_range() - with pytest.raises(ValueError, match="no simulation cell"): - unsized.compute_input_stats([{"coord": np.zeros((1, 3))}]) + spacing = long_range.mesh_spacing() + assert spacing == pytest.approx( + 0.75 * np.pi / long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + ) + for topologies, box in ((small, _cells()[:1]), (large, 2.0 * _cells())): + mesh = topologies["pme_mesh"] + dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)) + assert dims == mesh_dims_of_cells(box, spacing) + assert mesh.mask.all() def test_pme_mode_configuration() -> None: - """The pme mode requires periodicity, and its parameters the pme mode.""" + """The pme mode requires periodicity, and so do the mesh parameters.""" with pytest.raises(ValueError, match="periodic"): LatentChargeLongRange( LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), @@ -736,22 +788,28 @@ def test_pme_mode_configuration() -> None: None, k_grid_mode="pme", ) - with pytest.raises(ValueError, match="pme_mesh"): - _periodic_long_range(pme_mesh=[8, 8, 8]) - with pytest.raises(ValueError, match="pme_spacing"): - _periodic_long_range(pme_spacing=1.0) + for kwargs in ({"pme_spacing": 1.0}, {"pme_order": 6}): + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + **kwargs, + ) + # every periodic mode takes them: a frozen model evaluates on the mesh + assert _periodic_long_range(pme_spacing=1.0).mesh_spacing() == 1.0 + assert _periodic_long_range().pme_order == 8 with pytest.raises(ValueError, match="pme_order"): _pme_long_range(pme_order=1) - with pytest.raises(ValueError, match="pme_mesh"): - _pme_long_range(pme_mesh=[8, 8]) + with pytest.raises(ValueError, match="pme_spacing"): + _pme_long_range(pme_spacing=-1.0) with pytest.raises(ValueError, match="k_grid_mode"): _periodic_long_range(k_grid_mode="dense") def test_pme_round_trip() -> None: - """The record keeps the pme mode, the sized mesh and the spacing.""" + """The record keeps the pme mode, the order and the spacing.""" long_range = _pme_long_range(pme_spacing=1.0, pme_order=5) - long_range.compute_input_stats([{"box": _cells()}]) record = long_range.serialize() assert record["k_grid_mode"] == "pme" assert record["pme_order"] == 5 @@ -760,5 +818,5 @@ def test_pme_round_trip() -> None: assert restored.k_grid_mode == "pme" assert restored.pme_order == 5 assert restored.pme_spacing == 1.0 - np.testing.assert_array_equal(restored.pme_mesh, long_range.pme_mesh) - assert restored._pme_dims() == long_range._pme_dims() + # the default order is recorded too, so that it survives a change of default + assert _periodic_long_range().serialize()["pme_order"] == 8 diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py index 554704c7ae..8fa069eb05 100644 --- a/source/tests/infer/gen_dpa4c_long_range.py +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -3,26 +3,33 @@ """Generate the DPA4C long-range .pt2 test models and reference values. A latent-charge long-range module sums over every pair of atoms of a complete -non-periodic system, so the C++ interface hands the model what the pair sum +non-periodic system, and over the images of the cell of a periodic one when +it serves periodic systems, so the C++ interface hands the model what the sum needs beyond the neighbor list: the graph lower of an uncompressed model -receives the pair list the interface builds from the positions of the atoms, -and the compact canonical lower of a compressed model reads the positions -themselves. One archive covers each lower: +receives the pair list, or the mesh of the cell, that the interface builds +from the positions of the atoms, and the compact canonical lower of a +compressed model reads the positions and the cell themselves. One archive +covers each lower and each kind of module: -- ``deeppot_dpa4c_long_range.pt2``: uncompressed, fp64, graph lower, which - compiles on the CPU; -- ``deeppot_dpa4c_long_range_compressed.pt2``: compressed, fp32, compact - canonical lower, whose operators are CUDA only; it is generated for the - CUDA variant of the C++ suite (``DP_VARIANT=cuda``). +- ``deeppot_dpa4c_long_range.pt2`` and + ``deeppot_dpa4c_long_range_periodic.pt2``: uncompressed, fp64, graph lower, + which compiles on the CPU; +- ``deeppot_dpa4c_long_range_compressed.pt2`` and + ``deeppot_dpa4c_long_range_periodic_compressed.pt2``: compressed, fp32, + compact canonical lower, whose operators are CUDA only; they are generated + for the CUDA variant of the C++ suite (``DP_VARIANT=cuda``). -Both models condition the descriptor on the charge state and constrain the +Every model conditions the descriptor on the charge state and constrains the latent charges to the total charge that state names, so a state other than the default changes the long-range energy as well as the descriptor. The -references are those of the Python interface on each archive, without a cell, -at the default state and at ``[1, 2]``; each ``.expected`` sidecar holds the -sections ``default`` and ``explicit``. The six atoms form two groups about +references are those of the Python interface on each archive. For a +non-periodic module they are taken without a cell, at the default state and +at ``[1, 2]`` (sections ``default`` and ``explicit``); for a periodic module, +in the cell :data:`BOX` and without a cell, at the neutral default state +(sections ``periodic`` and ``cluster``). The six atoms form two groups about 9 Å apart, beyond the cutoff of the descriptor, which only the long-range -module couples. +module couples; in the cell, the first group meets its images across the +boundary. """ import copy @@ -116,6 +123,17 @@ dtype=np.float64, ).reshape(1, 6, 3) _OTHER_CHG_SPIN = np.array([[1.0, 2.0]], dtype=np.float64) +# A triclinic cell, row-major with the lattice vectors as rows. +BOX = np.array( + [[13.0, 0.0, 0.0], [0.6, 12.5, 0.0], [-0.4, 0.8, 13.5]], dtype=np.float64 +).reshape(1, 9) + + +def _periodic(config: dict) -> dict: + """The configuration with a long-range module that serves periodic systems.""" + config = copy.deepcopy(config) + config["long_range"]["periodic"] = True + return config def _model_dict(config: dict) -> dict: @@ -128,13 +146,14 @@ def _model_dict(config: dict) -> dict: return jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED)) -def _check_metadata(pt2_path: str, lower: str) -> None: +def _check_metadata(pt2_path: str, lower: str, periodic: bool = False) -> None: """Assert the archive declares its lower and the long-range input it needs.""" with zipfile.ZipFile(pt2_path) as zf: md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) assert md["lower_input_kind"] == lower, md["lower_input_kind"] + expected = ["all_pairs", "pme_mesh"] if periodic else ["all_pairs"] + assert md["topologies"] == expected, md["topologies"] if lower == "graph": - assert md["topologies"] == ["all_pairs"] assert md["dim_chg_spin"] == 2 else: # The compressed descriptor folds the charge state; the archive @@ -143,18 +162,22 @@ def _check_metadata(pt2_path: str, lower: str) -> None: assert len(md["charge_state_constants"]) == 5 -def _write_ref(pt2_path: str, ref_path: str) -> None: +def _write_ref(pt2_path: str, ref_path: str, periodic: bool = False) -> None: """Evaluate the reference cases on the archive and write the C++ sidecar.""" from deepmd.infer import ( DeepPot, ) dp = DeepPot(pt2_path) + if periodic: + cases = (("periodic", BOX, None), ("cluster", None, None)) + else: + cases = (("default", None, None), ("explicit", None, _OTHER_CHG_SPIN)) sections = {} energies = {} - for name, charge_spin in (("default", None), ("explicit", _OTHER_CHG_SPIN)): + for name, box, charge_spin in cases: condition = {} if charge_spin is None else {"charge_spin": charge_spin} - e, f, v, ae, av = dp.eval(_COORDS, None, _ATYPES, atomic=True, **condition) + e, f, v, ae, av = dp.eval(_COORDS, box, _ATYPES, atomic=True, **condition) assert np.all(np.isfinite(e)) and np.all(np.isfinite(f)) energies[name] = float(e[0, 0]) print(f"// {name} total energy: {e[0, 0]:.18e}") # noqa: T201 @@ -163,9 +186,10 @@ def _write_ref(pt2_path: str, ref_path: str) -> None: "expected_f": f[0], "expected_v": av[0], } - # The C++ cases compare the two states, which needs a real difference. - assert abs(energies["explicit"] - energies["default"]) > 1e-3, ( - "the two charge states must differ, or a call that ignored the " + # The C++ cases compare the two cases, which needs a real difference. + first, second = (name for name, _, _ in cases) + assert abs(energies[second] - energies[first]) > 1e-3, ( + "the two cases must differ, or a call that ignored the cell or the " "requested state would still match" ) write_expected_ref( @@ -176,7 +200,7 @@ def _write_ref(pt2_path: str, ref_path: str) -> None: print(f"Wrote {ref_path}") # noqa: T201 -def _export_compressed(pt2_path: str) -> None: +def _export_compressed(pt2_path: str, config: dict) -> None: """Compress the fp32 model and export it for the CUDA device.""" from deepmd.pt_expt.model.model import ( BaseModel, @@ -185,11 +209,11 @@ def _export_compressed(pt2_path: str) -> None: deserialize_to_file, ) - model = BaseModel.deserialize(_model_dict(COMPRESSED_CONFIG)) + model = BaseModel.deserialize(_model_dict(config)) model.atomic_model.enable_compression(0.5) data = { "model": model.serialize(), - "model_def_script": COMPRESSED_CONFIG, + "model_def_script": config, "backend": "dpmodel", "software": "deepmd-kit", "version": "3.0.0", @@ -207,29 +231,31 @@ def main() -> None: load_custom_ops() base_dir = os.path.dirname(__file__) - pt2_path = os.path.join(base_dir, "deeppot_dpa4c_long_range.pt2") - data = { - "model": _model_dict(CONFIG), - "model_def_script": CONFIG, - "backend": "dpmodel", - "software": "deepmd-kit", - "version": "3.0.0", - } - print(f"Exporting to {pt2_path} (lower_kind='graph') ...") # noqa: T201 - deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="graph") - _check_metadata(pt2_path, "graph") - _write_ref(pt2_path, os.path.join(base_dir, "deeppot_dpa4c_long_range.expected")) - - if os.environ.get("DP_VARIANT") == "cuda": - compressed_path = os.path.join( - base_dir, "deeppot_dpa4c_long_range_compressed.pt2" - ) - _export_compressed(compressed_path) - _check_metadata(compressed_path, "dpa4c_canonical") - _write_ref( - compressed_path, - os.path.join(base_dir, "deeppot_dpa4c_long_range_compressed.expected"), + for suffix, periodic in (("", False), ("_periodic", True)): + config = _periodic(CONFIG) if periodic else CONFIG + stem = os.path.join(base_dir, f"deeppot_dpa4c_long_range{suffix}") + data = { + "model": _model_dict(config), + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {stem}.pt2 (lower_kind='graph') ...") # noqa: T201 + deserialize_to_file( + f"{stem}.pt2", data, do_atomic_virial=True, lower_kind="graph" ) + _check_metadata(f"{stem}.pt2", "graph", periodic) + _write_ref(f"{stem}.pt2", f"{stem}.expected", periodic) + + if os.environ.get("DP_VARIANT") == "cuda": + compressed = f"{stem}_compressed" + compressed_config = ( + _periodic(COMPRESSED_CONFIG) if periodic else COMPRESSED_CONFIG + ) + _export_compressed(f"{compressed}.pt2", compressed_config) + _check_metadata(f"{compressed}.pt2", "dpa4c_canonical", periodic) + _write_ref(f"{compressed}.pt2", f"{compressed}.expected", periodic) print("\nDone!") # noqa: T201 diff --git a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py deleted file mode 100644 index b5c775e8cc..0000000000 --- a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py +++ /dev/null @@ -1,244 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Correctness of the direct-sum operator of the latent-charge long range. - -The CUDA kernel evaluates every pair in float32 from positions relative to -the first node of each frame, carried as float32 head and tail, and sums in -compensated float32 before a float64 reduction; the total-charge constraint -and the per-frame reductions sum in float64 in a fixed order. It is held to -the float64 evaluation over the explicit pairs, at a tolerance relative to -the scale of each output, and to the exact invariants of a pair sum: padding -nodes and frames do not interact, the forces of a frame cancel, and the -energy, force, charge gradient and virial follow from one another. -""" - -from __future__ import ( - annotations, -) - -from typing import ( - Any, -) - -import numpy as np -import pytest -import torch - -from deepmd.dpmodel.long_range import ( - COULOMB_CONSTANT, -) -from deepmd.pt_expt.kernels.long_range import ( - _reference, - kernel_operator_arguments, - latent_charge_direct_sum, - op_available, -) -from deepmd.pt_expt.long_range.kernel import ( - LESKernel, - SOGKernel, -) - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or not op_available(), - reason="the direct-sum operator is CUDA only", -) - -KERNELS = { - "les": lambda: LESKernel(alpha=0.8), - "sog": lambda: SOGKernel( - amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2], precision="float32" - ), -} - - -def _system( - n_node: list[int], channels: int, padding: int, seed: int, offset: float -) -> tuple[torch.Tensor, ...]: - """Frames of random nodes far from the origin, with trailing padding nodes. - - Every seventh node carries no charge, and every frame a total charge. - """ - rng = np.random.default_rng(seed) - total = sum(n_node) + padding - coord = offset + rng.uniform(0.0, 4.0 * max(n_node) ** (1 / 3), (total, 3)) - charge = rng.normal(0.0, 0.5, (total, channels)) - device = torch.device("cuda") - node_mask = torch.ones(total, dtype=torch.bool, device=device) - node_mask[::7] = False - return ( - torch.tensor(coord, dtype=torch.float64, device=device), - torch.tensor(charge, dtype=torch.float32, device=device), - node_mask, - torch.tensor(n_node, dtype=torch.int64, device=device), - torch.tensor( - rng.integers(-2, 3, len(n_node)), dtype=torch.float64, device=device - ), - ) - - -def _evaluate( - system: tuple[torch.Tensor, ...], kernel: Any, channel: int -) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: - coord, charge, node_mask, n_node, total_charge = system - actual = latent_charge_direct_sum( - coord, - charge, - node_mask, - n_node, - total_charge, - kernel, - channel, - COULOMB_CONSTANT, - ) - code, coefficient = kernel_operator_arguments(kernel) - expected = _reference( - coord, - charge.double(), - node_mask, - n_node, - total_charge, - coefficient.double(), - torch.zeros(1, dtype=torch.int64), - code, - channel, - COULOMB_CONSTANT, - ) - return actual, expected - - -NAMES = ("atom_energy", "charge_gradient", "force", "atom_virial", "virial") - - -@pytest.mark.parametrize("kernel_name", ["les", "sog"]) -@pytest.mark.parametrize( - ("n_node", "channels", "padding", "channel"), - [ - ([37, 1, 0, 300], 1, 5, 0), - ([37, 1, 0, 300], 1, 5, -1), - ([2048], 3, 0, -1), - ([5, 129], 8, 3, 2), - ], -) -def test_matches_float64_pair_sum( - kernel_name: str, n_node: list[int], channels: int, padding: int, channel: int -) -> None: - """Every output agrees with the float64 pair sum.""" - system = _system(n_node, channels, padding, 7, 1.0e3) - actual, expected = _evaluate(system, KERNELS[kernel_name](), channel) - _assert_close(actual, expected) - # Padding nodes join no frame. - if padding: - for value in actual[:4]: - assert torch.all(value[-padding:] == 0) - - -@pytest.mark.parametrize("alpha", [1e-3, 0.03, 0.3]) -@pytest.mark.parametrize("channel", [0, -1]) -def test_les_matches_float64_pair_sum_at_small_alpha_r( - alpha: float, channel: int -) -> None: - """The erf kernel keeps its precision where alpha r is small. - - The two terms of the slope of erf(x) / x agree to second order in x, so - that the slope of an approximate erf loses precision in proportion to - 1 / x^2 at atomic distances and a small alpha. - """ - system = _system([96], 2, 0, 5, 0.0) - _assert_close(*_evaluate(system, LESKernel(alpha=alpha), channel)) - - -def _assert_close( - actual: tuple[torch.Tensor, ...], expected: tuple[torch.Tensor, ...] -) -> None: - """Every output agrees with the float64 pair sum, relative to its largest entry.""" - for name, value, reference in zip(NAMES, actual, expected, strict=True): - assert value.dtype == ( - torch.float64 if name == "atom_energy" else torch.float32 - ) - scale = reference.abs().max().item() - error = (value.double() - reference.double()).abs().max().item() - assert error <= 2e-5 * scale, (name, error, scale) - - -@pytest.mark.parametrize("kernel_name", ["les", "sog"]) -def test_invariants(kernel_name: str) -> None: - """Frames are independent, forces cancel, and the virial is the pair virial.""" - n_node = [64, 200] - system = _system(n_node, 2, 0, 11, 0.0) - coord, charge, node_mask, n_node_t, total_charge = system - together = latent_charge_direct_sum( - coord, - charge, - node_mask, - n_node_t, - total_charge, - KERNELS[kernel_name](), - 1, - COULOMB_CONSTANT, - ) - _, _, force, atom_virial, virial = together - # Each frame alone gives the same result. - first = n_node[0] - alone = latent_charge_direct_sum( - coord[:first], - charge[:first], - node_mask[:first], - n_node_t[:1], - total_charge[:1], - KERNELS[kernel_name](), - 1, - COULOMB_CONSTANT, - ) - for both, single in zip(together, alone, strict=True): - torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) - # Internal forces sum to zero per frame, and the virial of a frame is - # sum_i r_i (x) f_i. - for frame, (begin, end) in enumerate(((0, first), (first, sum(n_node)))): - frame_force = force[begin:end].double() - assert frame_force.sum(0).abs().max() <= 1e-5 * frame_force.abs().max() - position = coord[begin:end] - coord[begin] - expected = torch.einsum("ni,nj->ij", position, frame_force) - scale = expected.abs().max().item() - for value in ( - atom_virial[begin:end].double().sum(0).reshape(3, 3), - virial[frame].double().reshape(3, 3), - ): - torch.testing.assert_close(value, expected, rtol=0, atol=2e-5 * scale) - - -@pytest.mark.parametrize("kernel_name", ["les", "sog"]) -@pytest.mark.parametrize("channel", [-1, 1]) -def test_reference_outputs_are_energy_derivatives( - kernel_name: str, channel: int -) -> None: - """The float64 force and charge gradient differentiate the energy.""" - coord, charge, node_mask, n_node, total_charge = _system([9, 6], 2, 2, 3, 0.0) - coord = coord.clone().requires_grad_(True) - charge = charge.double().requires_grad_(True) - kernel = KERNELS[kernel_name]() - code, coefficient = kernel_operator_arguments(kernel) - atom_energy, charge_gradient, force, _, _ = _reference( - coord, - charge, - node_mask, - n_node, - total_charge, - coefficient.double(), - torch.zeros(1, dtype=torch.int64), - code, - channel, - COULOMB_CONSTANT, - ) - grad_coord, grad_charge = torch.autograd.grad(atom_energy.sum(), (coord, charge)) - torch.testing.assert_close(force, -grad_coord, rtol=1e-9, atol=1e-9) - torch.testing.assert_close( - charge_gradient.detach(), grad_charge, rtol=1e-9, atol=1e-9 - ) - - -def test_bitwise_reproducible() -> None: - """Two evaluations of the same input agree bit for bit.""" - system = _system([3000], 1, 0, 5, 0.0) - first = latent_charge_direct_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT) - second = latent_charge_direct_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT) - for a, b in zip(first, second, strict=True): - assert torch.equal(a, b) diff --git a/source/tests/pt_expt/kernels/test_long_range_sum.py b/source/tests/pt_expt/kernels/test_long_range_sum.py new file mode 100644 index 0000000000..18552e553f --- /dev/null +++ b/source/tests/pt_expt/kernels/test_long_range_sum.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Correctness of the operator of the latent-charge long range. + +For non-periodic frames the CUDA kernel evaluates every pair in float32 from +positions relative to the first node of each frame, carried as float32 head +and tail, and sums in compensated float32 before a float64 reduction; the +total-charge constraint and the per-frame reductions sum in float64 in a fixed +order. It is held to the float64 evaluation over the explicit pairs, at a +tolerance relative to the scale of each output, and to the exact invariants of +a pair sum: padding nodes and frames do not interact, the forces of a frame +cancel, and the energy, force, charge gradient and virial follow from one +another. + +For periodic frames the CUDA kernels evaluate the mesh sum in float64, and are +held to the mesh evaluation of the training path, differentiated by autograd. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.pt_expt.kernels.long_range import ( + _reference, + kernel_operator_arguments, + latent_charge_sum, + op_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not op_available(), + reason="the long-range operator is CUDA only", +) + +KERNELS = { + "les": lambda: LESKernel(alpha=0.8), + "sog": lambda: SOGKernel( + amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2], precision="float32" + ), +} + + +def _system( + n_node: list[int], channels: int, padding: int, seed: int, offset: float +) -> tuple[torch.Tensor, ...]: + """Frames of random nodes far from the origin, with trailing padding nodes. + + Every seventh node carries no charge, and every frame a total charge. + """ + rng = np.random.default_rng(seed) + total = sum(n_node) + padding + coord = offset + rng.uniform(0.0, 4.0 * max(n_node) ** (1 / 3), (total, 3)) + charge = rng.normal(0.0, 0.5, (total, channels)) + device = torch.device("cuda") + node_mask = torch.ones(total, dtype=torch.bool, device=device) + node_mask[::7] = False + return ( + torch.tensor(coord, dtype=torch.float64, device=device), + torch.tensor(charge, dtype=torch.float32, device=device), + node_mask, + torch.tensor(n_node, dtype=torch.int64, device=device), + torch.tensor( + rng.integers(-2, 3, len(n_node)), dtype=torch.float64, device=device + ), + ) + + +def _reference_of( + system: tuple[torch.Tensor, ...], + kernel: Any, + channel: int, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[torch.Tensor, ...]: + """The float64 reference of the operator.""" + coord, charge, node_mask, n_node, total_charge = system + code, coefficient = kernel_operator_arguments(kernel) + if box is None: + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64) + return _reference( + coord, + charge.double(), + node_mask, + n_node, + total_charge, + box, + coefficient, + torch.zeros(1, dtype=torch.int64), + code, + channel, + COULOMB_CONSTANT, + order, + spacing, + ) + + +def _evaluate( + system: tuple[torch.Tensor, ...], + kernel: Any, + channel: int, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + actual = latent_charge_sum( + *system, + kernel, + channel, + COULOMB_CONSTANT, + box=box, + order=order, + spacing=spacing, + ) + return actual, _reference_of(system, kernel, channel, box, order, spacing) + + +NAMES = ("atom_energy", "charge_gradient", "force", "atom_virial", "virial") + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel"), + [ + ([37, 1, 0, 300], 1, 5, 0), + ([37, 1, 0, 300], 1, 5, -1), + ([2048], 3, 0, -1), + ([5, 129], 8, 3, 2), + ], +) +def test_matches_float64_pair_sum( + kernel_name: str, n_node: list[int], channels: int, padding: int, channel: int +) -> None: + """Every output agrees with the float64 pair sum.""" + system = _system(n_node, channels, padding, 7, 1.0e3) + actual, expected = _evaluate(system, KERNELS[kernel_name](), channel) + _assert_close(actual, expected) + # Padding nodes join no frame. + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("alpha", [1e-3, 0.03, 0.3]) +@pytest.mark.parametrize("channel", [0, -1]) +def test_les_matches_float64_pair_sum_at_small_alpha_r( + alpha: float, channel: int +) -> None: + """The erf kernel keeps its precision where alpha r is small. + + The two terms of the slope of erf(x) / x agree to second order in x, so + that the slope of an approximate erf loses precision in proportion to + 1 / x^2 at atomic distances and a small alpha. + """ + system = _system([96], 2, 0, 5, 0.0) + _assert_close(*_evaluate(system, LESKernel(alpha=alpha), channel)) + + +def _assert_close( + actual: tuple[torch.Tensor, ...], + expected: tuple[torch.Tensor, ...], + tolerance: float = 2e-5, +) -> None: + """Every output agrees with the float64 reference, relative to its largest entry.""" + for name, value, reference in zip(NAMES, actual, expected, strict=True): + assert value.dtype == ( + torch.float64 if name == "atom_energy" else torch.float32 + ) + scale = reference.abs().max().item() + error = (value.double() - reference.double()).abs().max().item() + assert error <= tolerance * scale, (name, error, scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_invariants(kernel_name: str) -> None: + """Frames are independent, forces cancel, and the virial is the pair virial.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + together = latent_charge_sum( + coord, + charge, + node_mask, + n_node_t, + total_charge, + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + _, _, force, atom_virial, virial = together + # Each frame alone gives the same result. + first = n_node[0] + alone = latent_charge_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + # Internal forces sum to zero per frame, and the virial of a frame is + # sum_i r_i (x) f_i. + for frame, (begin, end) in enumerate(((0, first), (first, sum(n_node)))): + frame_force = force[begin:end].double() + assert frame_force.sum(0).abs().max() <= 1e-5 * frame_force.abs().max() + position = coord[begin:end] - coord[begin] + expected = torch.einsum("ni,nj->ij", position, frame_force) + scale = expected.abs().max().item() + for value in ( + atom_virial[begin:end].double().sum(0).reshape(3, 3), + virial[frame].double().reshape(3, 3), + ): + torch.testing.assert_close(value, expected, rtol=0, atol=2e-5 * scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize("channel", [-1, 1]) +def test_reference_outputs_are_energy_derivatives( + kernel_name: str, channel: int +) -> None: + """The float64 force and charge gradient differentiate the energy.""" + coord, charge, node_mask, n_node, total_charge = _system([9, 6], 2, 2, 3, 0.0) + coord = coord.clone().requires_grad_(True) + charge = charge.double().requires_grad_(True) + atom_energy, charge_gradient, force, _, _ = _reference_of( + (coord, charge, node_mask, n_node, total_charge), + KERNELS[kernel_name](), + channel, + ) + grad_coord, grad_charge = torch.autograd.grad(atom_energy.sum(), (coord, charge)) + torch.testing.assert_close(force, -grad_coord, rtol=1e-9, atol=1e-9) + torch.testing.assert_close( + charge_gradient.detach(), grad_charge, rtol=1e-9, atol=1e-9 + ) + + +def test_bitwise_reproducible() -> None: + """Two evaluations of the same input agree bit for bit, with or without cells.""" + system = _system([3000], 1, 0, 5, 0.0) + for box in (None, _cells(1, 40.0)): + first = latent_charge_sum( + *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box + ) + second = latent_charge_sum( + *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box + ) + for a, b in zip(first, second, strict=True): + assert torch.equal(a, b) + + +# --------------------------------------------------------------------------- +# Periodic frames +# --------------------------------------------------------------------------- + + +def _cells(frame_count: int, length: float, seed: int = 3) -> torch.Tensor: + """Triclinic cells of about ``length`` per side, row-major, one row per frame.""" + rng = np.random.default_rng(seed) + cells = [] + for _ in range(frame_count): + cell = np.diag(length * rng.uniform(0.8, 1.2, 3)) + cell[1, 0], cell[2, 0], cell[2, 1] = length * rng.uniform(-0.2, 0.2, 3) + cells.append(cell.reshape(9)) + return torch.tensor(np.asarray(cells), dtype=torch.float64) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel", "order"), + [ + ([37, 1, 300], 1, 5, 0, 8), + ([37, 1, 300], 1, 5, -1, 4), + ([500], 3, 0, -1, 6), + ([5, 129], 8, 3, 2, 8), + ], +) +def test_mesh_matches_float64_reference( + kernel_name: str, + n_node: list[int], + channels: int, + padding: int, + channel: int, + order: int, +) -> None: + """Every output of a periodic frame agrees with the mesh evaluation of training. + + The positions extend beyond the cell, so that the fractional coordinates + wrap. + """ + system = _system(n_node, channels, padding, 7, -5.0) + box = _cells(len(n_node), 12.0) + actual, expected = _evaluate( + system, KERNELS[kernel_name](), channel, box, order, 0.6 + ) + _assert_close(actual, expected, 1e-6) + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_mesh_virial_is_the_strain_derivative(kernel_name: str) -> None: + """The virial of a periodic frame is the strain derivative of its energy. + + Straining the positions and the cell together by ``I + eps`` changes the + energy by ``-W : eps``; the atomic virial of the charged nodes sums to W. + The mesh interpolates in float32, whose rounding the central difference + divides by its step, so that the step stays large against it. + """ + system = _system([60], 2, 0, 13, 0.0) + coord, charge, node_mask, n_node, total_charge = system + box = _cells(1, 10.0) + kernel = KERNELS[kernel_name]() + atom_energy, _, _, atom_virial, virial = latent_charge_sum( + *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + h = 1e-3 + finite = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + energies = [] + for sign in (1.0, -1.0): + strain = torch.eye(3, dtype=torch.float64) + strain[a, b] += sign * h + strained = ( + coord @ strain.T.to(coord.device), + charge, + node_mask, + n_node, + total_charge, + ) + cell = box.reshape(3, 3) @ strain.T + energies.append( + latent_charge_sum( + *strained, + kernel, + 0, + COULOMB_CONSTANT, + box=cell.reshape(1, 9), + spacing=0.5, + )[0].sum() + ) + finite[a, b] = -(energies[0] - energies[1]).item() / (2 * h) + scale = np.abs(finite).max() + np.testing.assert_allclose( + virial[0].double().cpu().numpy().reshape(3, 3), finite, atol=2e-4 * scale + ) + np.testing.assert_allclose( + atom_virial.double().sum(0).cpu().numpy().reshape(3, 3), + finite, + atol=2e-4 * scale, + ) + + +@pytest.mark.parametrize("periodic", [False, True]) +def test_frame_without_charged_nodes(periodic: bool) -> None: + """A frame whose nodes carry no charge contributes zeros, the others are unaffected.""" + n_node = [40, 6, 50] + coord, charge, node_mask, n_node_t, total_charge = _system(n_node, 2, 0, 5, 0.0) + node_mask[40:46] = False + system = (coord, charge, node_mask, n_node_t, total_charge) + box = _cells(3, 10.0) if periodic else None + actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 0.6) + _assert_close(actual, expected, 1e-6) + for value in actual[:4]: + assert torch.all(value[40:46] == 0) + assert torch.all(actual[4][1] == 0) + + +def test_mesh_frames_are_independent() -> None: + """Every periodic frame is evaluated on its own mesh, bit for bit.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + box = _cells(2, 9.0) + kernel = KERNELS["sog"]() + together = latent_charge_sum( + *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + first = n_node[0] + alone = latent_charge_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + kernel, + 1, + COULOMB_CONSTANT, + box=box[:1], + spacing=0.5, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + + +def test_mixed_frames_are_rejected() -> None: + """A call holds periodic frames or non-periodic ones, not both.""" + system = _system([10, 10], 1, 0, 1, 0.0) + box = _cells(2, 9.0) + box[1] = 0.0 + with pytest.raises(RuntimeError, match="periodic"): + latent_charge_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 39b74f46da..62f0d5eb58 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -468,6 +468,27 @@ def test_linear_composition_takes_the_exclusions_of_the_composition() -> None: config["models"][0][key] = value with pytest.raises(error, match=match): get_model(config) + # a periodic module shares its composition with no other long-range module + periodic = copy.deepcopy(child) + periodic["long_range"]["periodic"] = True + for models in ([periodic, child], [periodic, periodic]): + config = copy.deepcopy(composition) + config.pop("pair_exclude_types") + config["models"] = copy.deepcopy(models) + with pytest.raises(NotImplementedError, match="serves periodic systems"): + get_model(config) + config["models"] = [copy.deepcopy(child), copy.deepcopy(child)] + get_model(config) + # a periodic module takes no pair exclusion, alone or in a composition + config = copy.deepcopy(composition) + config["models"] = [copy.deepcopy(periodic)] + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + get_model(config) + config = copy.deepcopy(periodic) + config["type"] = "standard" + config["pair_exclude_types"] = [[0, 1]] + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + get_model(config) @pytest.mark.parametrize("compiled", [False, True]) @@ -740,7 +761,6 @@ def test_periodic_checkpoint_inference_matches_direct_forward( config = _config(kernel) config["long_range"]["periodic"] = True config["long_range"]["k_grid_mode"] = "pme" - config["long_range"]["pme_mesh"] = [32, 32, 32] else: model = _periodic_model(kernel) config = _config(kernel) @@ -804,15 +824,9 @@ def test_frozen_model_matches_direct_forward( dp = DeepPot(str(frozen), auto_batch_size=False) assert dp.deep_eval.metadata["topologies"] == ["all_pairs"] if suffix == ".pte": - names = [ - node.name - for node in torch.export.load(str(frozen)).graph_module.graph.nodes - if node.op == "placeholder" - ] - for name, expected_name in zip( - names[-3:], ("pair_index", "pair_vec", "pair_mask"), strict=True - ): - assert name.startswith(expected_name), names + # The pair topology follows the thirteen inputs of the graph lower. + signature = torch.export.load(str(frozen)).graph_signature + assert len(signature.user_inputs) == 13 + 3 names = ("energy", "force", "virial", "atom_energy", "atom_virial") for seed, nloc in ((7, 6), (8, 4)): coord, atype = _batch(seed=seed) @@ -860,7 +874,7 @@ def test_coordinate_evaluation_matches_pair_graph( ) if not op_available(): - pytest.skip("the direct-sum operator is not available") + pytest.skip("the long-range operator is not available") config = _config(kernel, constraint=None if source is None else {"source": source}) if source == "charge_spin": config["descriptor"]["add_chg_spin_ebd"] = True @@ -914,7 +928,7 @@ def test_compressed_frozen_model_matches_uncompressed( """A compressed model keeps its long-range module through the freeze. On a CUDA device the freeze selects the compact canonical lower, whose - direct-sum operator reads the positions of the atoms; on the CPU the + long-range operator reads the positions of the atoms; on the CPU the graph lower reads the pair lists. A constraint that reads the total charge from the charge state is folded with the descriptor and follows the state the deployment serves. Both match the uncompressed model. @@ -1355,17 +1369,152 @@ def test_periodic_serialization_round_trip(kernel: str) -> None: ) -@pytest.mark.parametrize("kernel", ["les", "sog"]) -def test_periodic_compression_is_rejected(kernel: str) -> None: - """The compressed path does not serve the k-space evaluation.""" +def _periodic_frames() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Two frames of the periodic case, the second strained and shifted.""" + coord, atype, box = _periodic_case() + strain = np.array([[1.04, 0.02, 0.0], [0.0, 0.97, 0.01], [0.03, 0.0, 1.02]]) + coord2 = np.concatenate([coord, coord @ strain.T + 0.3], axis=0) + box2 = np.concatenate([box, box @ strain.T], axis=0) + return coord2, np.concatenate([atype, atype], axis=0), box2 + + +@pytest.mark.parametrize( + ("kernel", "grid_mode", "suffix"), + [("les", "ragged", ".pt2"), ("sog", "pme", ".pte"), ("sog", "fixed", ".pt2")], +) +def test_periodic_frozen_model_matches_mesh_evaluation( + tmp_path, kernel: str, grid_mode: str, suffix: str +) -> None: + """A frozen periodic model evaluates every periodic frame on the mesh of its cell. + + Whatever its training mode, the frozen model reads the mesh of the cells + of a batch as inputs of dynamic size, and matches the model evaluated on + the mesh; a batch without a cell keeps the direct sum. + """ config = _config(kernel) - config["descriptor"]["precision"] = "float32" + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = grid_mode + if grid_mode == "fixed": + # Bounds that hold every wave number of the cells within the cutoff. + config["long_range"]["k_grid_bounds"] = [14, 14, 14] + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / f"periodic_{kernel}_{grid_mode}{suffix}" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="graph", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + assert dp.deep_eval.metadata["topologies"] == ["all_pairs", "pme_mesh"] + mesh_model = copy.deepcopy(model) + mesh_model.atomic_model.long_range.use_mesh_evaluation() + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + coord, atype, box = _periodic_frames() + for frame in (slice(0, 1), slice(0, 2)): + expected = mesh_model( + _tensor(coord[frame]), + _tensor(atype[frame], torch.int64), + box=_tensor(box[frame]), + do_atomic_virial=True, + ) + actual = dp.eval( + coord[frame], box[frame], atype[frame], atomic=True, mixed_type=True + ) + for name, value in zip(names, actual, strict=True): + np.testing.assert_allclose( + value, + expected[name].detach().cpu().numpy().reshape(value.shape), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + # The mesh approximates the sum of the training mode. + trained = model( + _tensor(coord[:1]), _tensor(atype[:1], torch.int64), box=_tensor(box[:1]) + ) + np.testing.assert_allclose( + dp.eval(coord[:1], box[:1], atype[:1], mixed_type=True)[0], + trained["energy"].detach().cpu().numpy(), + rtol=1e-5, + ) + cluster, cluster_type = _batch() + expected = model(cluster, cluster_type) + actual = dp.eval( + cluster.cpu().numpy(), None, cluster_type.cpu().numpy(), mixed_type=True + ) + np.testing.assert_allclose( + actual[0], expected["energy"].detach().cpu().numpy(), rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize(("kernel", "source"), [("les", "neutral"), ("sog", None)]) +def test_periodic_compressed_frozen_model_matches_uncompressed( + tmp_path, kernel: str, source: str | None +) -> None: + """A compressed periodic model keeps the mesh evaluation through the freeze. + + On a CUDA device the compact canonical lower evaluates every periodic + frame on the mesh of its cell with the long-range operator; on the CPU + the graph lower reads the mesh. Both match the uncompressed model on the + mesh, for a periodic frame and for a cluster. + """ + config = _config(kernel, constraint=None if source is None else {"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + config["fitting_net"].update( + neuron=[32, 32], activation_function="silu", resnet_dt=False + ) config["fitting_net"]["precision"] = "float32" - config["long_range"]["charge_net"]["precision"] = "float32" + config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23} config["long_range"]["periodic"] = True model = get_model(config).to(env.DEVICE).eval() - with pytest.raises(NotImplementedError, match="reciprocal-space"): - model.atomic_model.enable_compression(0.5) + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + reference = copy.deepcopy(model) + reference.atomic_model.long_range.use_mesh_evaluation() + model.atomic_model.enable_compression(0.5) + frozen = tmp_path / f"periodic_compressed_{kernel}.pt2" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="auto", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" + assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + rng = np.random.default_rng(3) + box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]]) + periodic = rng.uniform(0.0, 1.0, (1, 40, 3)) @ box[0] + cluster = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 + atype = rng.integers(0, 2, (1, 40)) + for coord, cell in ((periodic, box), (cluster, None)): + expected = reference( + _tensor(coord), + _tensor(atype, torch.int64), + box=None if cell is None else _tensor(cell), + do_atomic_virial=True, + ) + actual = dp.eval(coord, cell, atype, atomic=True) + for name, value in zip(names, actual, strict=True): + target = expected[name].detach().cpu().numpy().reshape(value.shape) + scale = np.abs(target).max() + # The tables of the compressed descriptor approximate its + # networks, to a few 1e-6 of the scale of every output. + assert np.abs(value - target).max() <= 2e-5 * scale + 1e-5, ( + name, + cell is None, + np.abs(value - target).max(), + scale, + ) + # A cell without volume marks no system, on either lower: the graph lower + # meets it in the neighbor search, the compact lower at its input. + with pytest.raises((ValueError, RuntimeError), match=r"no volume|singular"): + dp.eval(periodic, np.zeros((1, 9)), atype) # --------------------------------------------------------------------------- @@ -1509,9 +1658,7 @@ def test_periodic_fixed_grid_truncates_oversized_cells(kernel: str) -> None: big_box = box * 3.0 with pytest.warns(UserWarning, match="truncated"): out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(big_box)) - grid_n = model.atomic_model.long_range.k_grid_n - grid_n = grid_n.detach().cpu().numpy() if hasattr(grid_n, "detach") else grid_n - bounds = np.abs(np.asarray(grid_n)).max(axis=0) + bounds = model.atomic_model.long_range.k_grid_bounds.detach().cpu().numpy() charge = np.zeros((6, CHANNELS[kernel])) charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] expected = _kspace_reference_bounded(coord[0], charge, big_box[0], kernel, bounds) @@ -1530,7 +1677,9 @@ def test_periodic_fixed_grid_needs_statistics_or_bounds(kernel: str) -> None: config["long_range"]["k_grid_mode"] = "fixed" config["long_range"]["k_grid_bounds"] = [12, 12, 12] model = get_model(config).to(env.DEVICE).eval() - assert model.atomic_model.long_range.k_grid_n.shape == (25**3 - 1, 3) + np.testing.assert_array_equal( + model.atomic_model.long_range.k_grid_bounds.cpu().numpy(), [12, 12, 12] + ) out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) reference = _periodic_model(kernel) reference.load_state_dict(model.state_dict(), strict=False) @@ -1538,6 +1687,20 @@ def test_periodic_fixed_grid_needs_statistics_or_bounds(kernel: str) -> None: torch.testing.assert_close(out["energy"], expected["energy"], rtol=1e-9, atol=1e-12) +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_restores_from_state_dict(kernel: str) -> None: + """The state dict of a sized grid loads into an unsized model, as a checkpoint does.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + restored = _periodic_fixed_model(kernel) + restored.load_state_dict(model.state_dict()) + args = (_tensor(coord), _tensor(atype, torch.int64)) + expected = model(*args, box=_tensor(box)) + got = restored(*args, box=_tensor(box)) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(got[key], expected[key], rtol=1e-12, atol=1e-12) + + @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_periodic_fixed_grid_serialization_round_trip(kernel: str) -> None: """The record keeps the sized grid and evaluates identically.""" @@ -1547,7 +1710,7 @@ def test_periodic_fixed_grid_serialization_round_trip(kernel: str) -> None: model = _periodic_fixed_model(kernel, box) record = model.serialize() assert record["long_range"]["k_grid_mode"] == "fixed" - assert "k_grid_n" in record["long_range"] + assert "k_grid_bounds" in record["long_range"] reference = BaseModelDP.deserialize(record) assert reference.atomic_model.long_range.k_grid_mode == "fixed" expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) @@ -1611,9 +1774,8 @@ def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: torch.tensor([6], dtype=torch.int64, device=env.DEVICE), box_t, ) - before = long_range.build_topologies(*grid_args, atype=atype_t.reshape(-1))[ - "k_grid" - ] + build = model.atomic_model.build_topologies + before = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"] energy_before = model(coord_t, atype_t, box=box_t)["energy"] # move the kernel far enough that a coupled grid would change with torch.no_grad(): @@ -1621,7 +1783,7 @@ def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: long_range.kernel.log_alpha.add_(1.0) else: long_range.kernel.log_width.sub_(1.0) - after = long_range.build_topologies(*grid_args, atype=atype_t.reshape(-1))["k_grid"] + after = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"] torch.testing.assert_close(after.vec, before.vec) torch.testing.assert_close(after.mask, before.mask) # the kernel values still follow the live parameters @@ -1643,45 +1805,32 @@ def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: # --------------------------------------------------------------------------- -def _periodic_pme_model( - kernel: str, boxes: np.ndarray | None = None, **kwargs: object -) -> torch.nn.Module: - """A periodic model in the PME mesh mode. - - Sized by the given cells through the data statistics, or by an explicit - ``pme_mesh`` in ``kwargs``. - """ +def _periodic_pme_model(kernel: str, **kwargs: object) -> torch.nn.Module: + """A periodic model in the PME mesh mode, ``pme_spacing`` and ``pme_order`` in ``kwargs``.""" pme_kwargs = { - key: kwargs.pop(key) - for key in ("pme_mesh", "pme_spacing", "pme_order") - if key in kwargs + key: kwargs.pop(key) for key in ("pme_spacing", "pme_order") if key in kwargs } config = _config(kernel, **kwargs) config["long_range"]["periodic"] = True config["long_range"]["k_grid_mode"] = "pme" config["long_range"].update(pme_kwargs) - if boxes is None: - config["long_range"].setdefault("pme_mesh", [32, 32, 32]) - model = get_model(config).to(env.DEVICE).eval() - if boxes is not None: - boxes_t = torch.as_tensor(np.asarray(boxes, dtype=np.float64).reshape(-1, 9)) - model.atomic_model.long_range.compute_input_stats([{"box": boxes_t}]) - return model + return get_model(config).to(env.DEVICE).eval() @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_periodic_pme_matches_kspace_reference(kernel: str) -> None: """With fixed charges the PME energy approaches the converged k-space sum.""" - model = _periodic_pme_model(kernel, constraint=None, pme_mesh=[48, 48, 48]) + model = _periodic_pme_model(kernel, constraint=None) _set_constant_charges(model, [1.0, -0.5]) coord, atype, box = _periodic_case() out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) charge = np.zeros((6, CHANNELS[kernel])) charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] expected = _kspace_reference(coord[0], charge, box[0], kernel) - # the PME energy is a mesh approximation of the converged lattice sum - np.testing.assert_allclose(out["energy"].item(), expected, atol=1e-2) - np.testing.assert_allclose(out["atom_energy"].sum().item(), expected, atol=1e-2) + # The default mesh approximates the reciprocal sum within a few parts in a + # million; the energy, less the self term, is several times smaller. + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-5) + np.testing.assert_allclose(out["atom_energy"].sum().item(), expected, rtol=1e-5) @pytest.mark.parametrize("kernel", ["les", "sog"]) @@ -1732,22 +1881,32 @@ def energy(c: np.ndarray, b: np.ndarray) -> float: @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_periodic_pme_frames_are_independent(kernel: str) -> None: - """Each frame's mesh energy is independent of the rest of the batch.""" + """Frames on one mesh are independent; a larger cell sharing the batch refines it. + + The frames of a batch share the mesh that resolves the largest cell, so a + smaller cell evaluated beside a larger one sees a finer mesh, which moves + its energy by less than the accuracy of the mesh. + """ model = _periodic_pme_model(kernel) coord, atype, box = _periodic_case() - coord2 = np.concatenate([coord, coord + 0.1], axis=0) atype2 = np.concatenate([atype, atype], axis=0) - box2 = np.concatenate([box, box * 1.1], axis=0) - batched = model(_tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2)) - for frame in range(2): - alone = model( - _tensor(coord2[frame : frame + 1]), - _tensor(atype2[frame : frame + 1], torch.int64), - box=_tensor(box2[frame : frame + 1]), - ) - torch.testing.assert_close( - alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + for box2, rtol in ( + (np.concatenate([box, box], axis=0), 1e-10), + (np.concatenate([box, box * 1.3], axis=0), 1e-6), + ): + batched = model( + _tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2) ) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=rtol, atol=1e-10 + ) @pytest.mark.parametrize("kernel", ["les", "sog"]) @@ -1763,37 +1922,16 @@ def test_periodic_pme_serves_non_periodic_batches(kernel: str) -> None: torch.testing.assert_close(mesh_out["force"], plain_out["force"]) -@pytest.mark.parametrize("kernel", ["les", "sog"]) -def test_periodic_pme_needs_statistics_or_bounds(kernel: str) -> None: - """Without statistics or an explicit mesh the mesh is not sized.""" - config = _config(kernel) - config["long_range"]["periodic"] = True - config["long_range"]["k_grid_mode"] = "pme" - model = get_model(config).to(env.DEVICE).eval() - coord, atype, box = _periodic_case() - with pytest.raises(ValueError, match="not sized yet"): - model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) - sized = _periodic_pme_model(kernel, box) - out = sized(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) - # the statistics-sized model evaluates exactly like a model given the - # same mesh explicitly - mesh = sized.atomic_model.long_range.pme_mesh.tolist() - reference = _periodic_pme_model(kernel, pme_mesh=mesh) - reference.load_state_dict(sized.state_dict()) - expected = reference(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) - torch.testing.assert_close(out["energy"], expected["energy"]) - - @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_periodic_pme_serialization_round_trip(kernel: str) -> None: - """The record keeps the pme mode and the sized mesh, and evaluates identically.""" + """The record keeps the pme mode and the order, and evaluates identically.""" from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP coord, atype, box = _periodic_case() - model = _periodic_pme_model(kernel, box) + model = _periodic_pme_model(kernel, pme_order=6) record = model.serialize() assert record["long_range"]["k_grid_mode"] == "pme" - assert "pme_mesh" in record["long_range"] + assert record["long_range"]["pme_order"] == 6 reference = BaseModelDP.deserialize(record) assert reference.atomic_model.long_range.k_grid_mode == "pme" expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) @@ -1832,20 +1970,3 @@ def test_periodic_pme_compiled_matches_eager(kernel: str) -> None: grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] assert grads and all(grad is not None for grad in grads) assert all(torch.isfinite(grad).all() for grad in grads) - - -@pytest.mark.parametrize("kernel", ["les", "sog"]) -def test_periodic_pme_checkpoint_restores_sized_mesh(kernel: str) -> None: - """A checkpoint restores the stats-sized mesh through the state dict alone.""" - coord, atype, box = _periodic_case() - trained = _periodic_pme_model(kernel, box) - state = trained.state_dict() - # a fresh model from the bare config is unsized - fresh = _periodic_pme_model(kernel, pme_mesh=None) - with pytest.raises(ValueError, match="not sized yet"): - fresh(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) - fresh.load_state_dict(state) - assert fresh.atomic_model.long_range._pme_dims() is not None - out = fresh(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) - expected = trained(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) - torch.testing.assert_close(out["energy"], expected["energy"]) diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index 859e1df8c2..87d4fd9f58 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -868,6 +868,79 @@ def test_finetune_random_fitting_into_bridged(self) -> None: """A RANDOM fitting of a bridged target takes the plain descriptor.""" self._finetune_across_bridging(source_bridged=False, model_branch="RANDOM") + def test_finetune_short_range_into_long_range(self) -> None: + """A checkpoint without a long-range module fine-tunes into one with it. + + The descriptor and the fitting net come from the checkpoint, while the + new periodic module starts from its initialization and takes the + input statistics of the fine-tuning data, here the bounds of its fixed + k grid, before the bias adjustment evaluates the complete model. + """ + plain = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "seed": 1, + }, + "fitting_net": {"neuron": [16, 16], "seed": 1}, + "data_stat_nbatch": 1, + } + target = deepcopy(plain) + target["long_range"] = { + "type": "latent_charge", + "charge_net": {"neuron": [8], "seed": 1}, + "kernel": {"type": "les", "alpha": 1.0}, + "charge_constraint": {"source": "neutral"}, + "periodic": True, + "k_grid_mode": "fixed", + } + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_long_range_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = normalize( + update_deepmd_input( + _make_config(self.data_dir, plain, numb_steps=1), warning=False + ) + ) + ckpt_path = self._train_pretrained(config, tmpdir) + ft_config = normalize( + update_deepmd_input( + _make_config(self.data_dir, target, numb_steps=1), warning=False + ) + ) + ft_config["model"], finetune_links = get_finetune_rules( + ckpt_path, ft_config["model"] + ) + trainer_ft = get_trainer( + ft_config, + finetune_model=ckpt_path, + finetune_links=finetune_links, + ) + pretrained_state = torch.load( + ckpt_path, map_location=DEVICE, weights_only=True + )["model"] + ft_state = trainer_ft.wrapper.state_dict() + learned = [ + key + for key in ft_state + if ".descriptor." in key + or (".fitting_net." in key and "bias_atom_e" not in key) + ] + self.assertTrue(learned) + for key in learned: + torch.testing.assert_close(ft_state[key], pretrained_state[key]) + long_range = trainer_ft.wrapper.model["Default"].atomic_model.long_range + self.assertTrue(torch.all(long_range.k_grid_bounds > 0)) + trainer_ft.run() + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + def test_finetune_from_pte(self) -> None: """Train -> freeze to .pte -> finetune from .pte -> verify checkpoint.""" from deepmd.pt_expt.entrypoints.main import ( From 1cc7884dbdcd2c547e868aad0b0cf1454dbc64fc Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 05:05:23 +0800 Subject: [PATCH 27/30] perf(long-range): float32 mesh pipeline and leaner host side of the operator The smooth particle-mesh Ewald sum of deepmd::latent_charge_sum is bound by float64 throughput, which is a small fraction of float32 on many devices, and by host work at small sizes: - The charge and potential meshes, the transforms and the influence function are float32; the B-spline weights are still evaluated in float64, once per node and axis in a separate pass, and rounded to float32. The fixed-point spreading keeps the result independent of the order of the atomic additions, with an integer rounding in place of the float64 conversion instruction. - Every lane of a warp takes a (k2, k3) pair of the stencil and walks the k1 axis, reading precomputed axis offsets; the interpolation forms the potential, its gradient and the force per channel in float32 and the owner sums form F (x) r from the force, so that the kernel writes four values per owner instead of fourteen in float64. The influence function takes rows of the half spectrum, one warp each. - The transforms run on cached cuFFT plans with the work area of the call; the intermediates of the mesh sum live in one workspace allocation, in which the potential mesh of a frame overwrites its charge mesh, and those of the direct sum but its chunk sums likewise. The direct sum reads the float64 coefficients directly, and the periodicity of the cells is checked on the host. On an RTX PRO 6000 (sm_120), 100,000 atoms: the LES mesh sum takes 0.67 ms instead of 3.7 ms, the SOG sum 0.50 ms instead of 3.2 ms; the host time of a call falls from about 300 us to 140 us. The forces stay within about 1e-6 of their scale of the float64 reference, as before. The energy of a frame carries the float32 rounding of the reciprocal sum, which includes the interaction of every charge with itself: for 4,000 neutral random charges it departs from the reference by 1e-7 to 5e-7 eV per atom, against 1e-9 to 5e-8 with float64 transforms, far below the error of a compressed descriptor, the models the operator evaluates. In LAMMPS with Kokkos on the compressed Air preset and the BCN diamond systems, a periodic LES step takes 1.71 ms instead of 2.08 ms at 512 atoms (1.30 ms without long range), 2.91 ms instead of 3.23 ms at 10,296 atoms (2.71 ms) and 18.9 ms instead of 24.0 ms at 100,800 atoms (18.1 ms). The virial test takes its finite difference from the float64 reference, since the float32 outputs carry no differences of 1e-5 strains; further tests cover every B-spline order and a stencil wider than the mesh. --- source/op/pt/CMakeLists.txt | 3 +- source/op/pt/long_range/charge_sum.h | 2 +- source/op/pt/long_range/direct_sum.cu | 182 +-- source/op/pt/long_range/frame_sums.cuh | 39 +- source/op/pt/long_range/latent_charge_sum.cu | 25 +- source/op/pt/long_range/mesh_sum.cu | 1150 ++++++++++------- .../pt_expt/kernels/test_long_range_sum.py | 58 +- 7 files changed, 901 insertions(+), 558 deletions(-) diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 17ef7c3d8d..9bbda9d61d 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -218,7 +218,8 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) set_property(TARGET deepmd_op_pt PROPERTY CUDA_ARCHITECTURES ${DEEPMD_PT_CUDA_ARCHITECTURES}) endif() - target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublas CUDA::cublasLt) + target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublas CUDA::cublasLt + CUDA::cufft) # libtorch headers require at least C++17, and PyTorch 2.14+ requires C++20. # Preserve an explicitly higher CUDA standard requested by the user. set(DEEPMD_PT_CUDA_STANDARD 17) diff --git a/source/op/pt/long_range/charge_sum.h b/source/op/pt/long_range/charge_sum.h index d570804e8c..04245d3287 100644 --- a/source/op/pt/long_range/charge_sum.h +++ b/source/op/pt/long_range/charge_sum.h @@ -31,7 +31,7 @@ using ChargeSumOutputs = std::tuple; -// The direct sum. ``coefficient`` holds the float32 coefficients of the +// The direct sum. ``coefficient`` holds the float64 coefficients of the // kernel on the device. ChargeSumOutputs direct_sum(const torch::Tensor& coord, const torch::Tensor& charge, diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu index 29d142e513..2e0f128d85 100644 --- a/source/op/pt/long_range/direct_sum.cu +++ b/source/op/pt/long_range/direct_sum.cu @@ -100,7 +100,7 @@ __global__ void prepare_nodes(const double* __restrict__ coord, int node_count, int channels, int channel, - const float* __restrict__ alpha, + const double* __restrict__ alpha, const double* __restrict__ frame_stat, float4* __restrict__ head, float4* __restrict__ tail, @@ -110,7 +110,7 @@ __global__ void prepare_nodes(const double* __restrict__ coord, if (node >= node_count) { return; } - const double scale = alpha == nullptr ? 1.0 : static_cast(*alpha); + const double scale = alpha == nullptr ? 1.0 : *alpha; const int f = frame_of(frame_ptr, frame_count, node); double3 offset = make_double3(0.0, 0.0, 0.0); if (f >= 0) { @@ -413,7 +413,7 @@ __global__ void __launch_bounds__(kThreads) const int64_t* __restrict__ frame_ptr, int owner_begin, int owner_count, - const float* __restrict__ coefficient, + const double* __restrict__ coefficient, int gaussian_count, float* __restrict__ partial) { constexpr int kWidth = sum_width(); @@ -431,8 +431,8 @@ __global__ void __launch_bounds__(kThreads) if constexpr (kKernel == kSog) { for (int k = threadIdx.x; k < gaussian_count; k += kThreads) { - const float amplitude = coefficient[k]; - const float rate = coefficient[gaussian_count + k]; + const float amplitude = static_cast(coefficient[k]); + const float rate = static_cast(coefficient[gaussian_count + k]); gaussians[k] = make_float4(amplitude, -rate, -1.38629436111989061f * rate * amplitude, 0.f); } @@ -548,7 +548,7 @@ __global__ void reduce_chunks(const float* __restrict__ partial, int owner_count, int chunk_count, double prefactor, - const float* __restrict__ alpha, + const double* __restrict__ alpha, float* __restrict__ potential, float* __restrict__ force, float* __restrict__ atom_virial) { @@ -565,7 +565,7 @@ __global__ void reduce_chunks(const float* __restrict__ partial, value += partial[static_cast(chunk) * kWidth * owner_count + index]; } - const double unit = alpha == nullptr ? 1.0 : static_cast(*alpha); + const double unit = alpha == nullptr ? 1.0 : *alpha; if (component < kChannels) { potential[node * kChannels + component] = static_cast(prefactor * unit * value); @@ -698,15 +698,17 @@ int chunk_count_for(int owner_count, // The staged partners, one row per node, and the sums of the owners. struct DirectSumBuffers { - const torch::Tensor& head; - const torch::Tensor& tail; - const torch::Tensor& charge; - const torch::Tensor& frame; - const torch::Tensor& frame_ptr; - const torch::Tensor& coefficient; - torch::Tensor& potential; - torch::Tensor& force; - torch::Tensor& atom_virial; + const float4* head; + const float4* tail; + const float* charge; + const int* frame; + const int64_t* frame_ptr; + int node_count; + const double* coefficient; + int gaussian_count; + float* potential; + float* force; + float* atom_virial; }; template @@ -730,28 +732,25 @@ void launch(const DirectSumBuffers& buffers, cudaStream_t stream) { constexpr int kWidth = sum_width(); const int chunk_count = - chunk_count_for(owner_count, static_cast(buffers.head.size(0)), + chunk_count_for(owner_count, buffers.node_count, resident_blocks_per_multiprocessor()); auto partial = torch::empty({static_cast(chunk_count) * owner_count * kWidth}, - buffers.charge.options()); + torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA, c10::cuda::current_device())); const dim3 grid((owner_count + kThreads - 1) / kThreads, chunk_count); direct_sum_kernel<<>>( - reinterpret_cast(buffers.head.data_ptr()), - reinterpret_cast(buffers.tail.data_ptr()), - buffers.charge.data_ptr(), buffers.frame.data_ptr(), - buffers.frame_ptr.data_ptr(), owner_begin, owner_count, - buffers.coefficient.data_ptr(), - static_cast(buffers.coefficient.numel() / 2), - partial.data_ptr()); + buffers.head, buffers.tail, buffers.charge, buffers.frame, + buffers.frame_ptr, owner_begin, owner_count, buffers.coefficient, + buffers.gaussian_count, partial.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); const int64_t reduce_threads = static_cast(kWidth) * owner_count; reduce_chunks <<((reduce_threads + 255) / 256), 256, 0, stream>>>( partial.data_ptr(), owner_count, chunk_count, prefactor, - kKernel == kLes ? buffers.coefficient.data_ptr() : nullptr, - buffers.potential.data_ptr(), buffers.force.data_ptr(), - buffers.atom_virial.data_ptr()); + kKernel == kLes ? buffers.coefficient : nullptr, buffers.potential, + buffers.force, buffers.atom_virial); C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -844,14 +843,6 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, auto virial = torch::empty({frame_count, 9}, float_options); const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); - auto frame_ptr = torch::empty({frame_count + 1}, counts.options()); - auto head = torch::empty({node_count, 4}, float_options); - auto tail = torch::empty({node_count, 4}, float_options); - auto frame = torch::empty({node_count}, coord.options().dtype(torch::kInt32)); - auto pair_charge = torch::empty({node_count, channels}, float_options); - auto potential = torch::empty({owner_count, channels}, float_options); - auto frame_stat = torch::empty({constrained ? frame_count * kStatWidth : 0}, - double_options); const auto slices_for = [frame_count](int64_t nodes) { const int64_t average_nodes = (nodes + frame_count - 1) / frame_count; return static_cast(std::min( @@ -859,42 +850,73 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); }; const int slice_count = slices_for(node_count); - const dim3 slice_grid(frame_count, slice_count); + const int owner_slice_count = slices_for(owner_count); + const int components = constrained ? 10 : 9; + Workspace workspace; + const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); + const int64_t head_at = workspace.reserve(node_count); + const int64_t tail_at = workspace.reserve(node_count); + const int64_t frame_at = workspace.reserve(node_count); + const int64_t pair_charge_at = + workspace.reserve(int64_t{node_count} * channels); + const int64_t potential_at = + workspace.reserve(int64_t{owner_count} * channels); + const int64_t frame_stat_at = + workspace.reserve(int64_t{frame_count} * kStatWidth); + const int64_t slice_charge_at = + workspace.reserve(int64_t{frame_count} * slice_count * 2); + const int64_t slice_virial_at = workspace.reserve( + int64_t{frame_count} * owner_slice_count * components); + workspace.allocate(coord.device()); + auto* frame_ptr = workspace.at(frame_ptr_at); + auto* head = workspace.at(head_at); + auto* tail = workspace.at(tail_at); + auto* frame = workspace.at(frame_at); + auto* pair_charge = workspace.at(pair_charge_at); + auto* potential = workspace.at(potential_at); + auto* frame_stat = + constrained ? workspace.at(frame_stat_at) : nullptr; + auto* slice_charge = workspace.at(slice_charge_at); + auto* slice_virial = workspace.at(slice_virial_at); + const auto stream = at::cuda::getCurrentCUDAStream(); - frame_offsets<<<1, kOffsetThreads, 0, stream>>>( - counts.data_ptr(), frame_count, frame_ptr.data_ptr()); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(), + frame_count, frame_ptr); C10_CUDA_KERNEL_LAUNCH_CHECK(); if (constrained) { - auto slice_charge = torch::empty( - {static_cast(frame_count) * slice_count * 2}, double_options); - charge_slices<<>>( - charge.data_ptr(), node_mask.data_ptr(), - frame_ptr.data_ptr(), static_cast(channels), - constrained_channel, slice_charge.data_ptr()); + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), frame_ptr, + static_cast(channels), constrained_channel, slice_charge); C10_CUDA_KERNEL_LAUNCH_CHECK(); charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, kFrameThreads, 0, stream>>>( - slice_charge.data_ptr(), slice_count, frame_count, - total_charge.data_ptr(), frame_stat.data_ptr()); + slice_charge, slice_count, frame_count, total_charge.data_ptr(), + frame_stat); C10_CUDA_KERNEL_LAUNCH_CHECK(); } // Alpha stays on the device, so that a call never waits for it. + const double* alpha = + kernel == kLes ? coefficient.data_ptr() : nullptr; prepare_nodes<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, - stream>>>( - coord.data_ptr(), charge.data_ptr(), - node_mask.data_ptr(), frame_ptr.data_ptr(), frame_count, - node_count, static_cast(channels), constrained_channel, - kernel == kLes ? coefficient.data_ptr() : nullptr, - constrained ? frame_stat.data_ptr() : nullptr, - reinterpret_cast(head.data_ptr()), - reinterpret_cast(tail.data_ptr()), frame.data_ptr(), - pair_charge.data_ptr()); + stream>>>(coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frame_count, + node_count, static_cast(channels), + constrained_channel, alpha, frame_stat, head, tail, + frame, pair_charge); C10_CUDA_KERNEL_LAUNCH_CHECK(); if (owner_count > 0) { - const DirectSumBuffers buffers{head, tail, pair_charge, - frame, frame_ptr, coefficient, - potential, force, atom_virial}; + const DirectSumBuffers buffers{head, + tail, + pair_charge, + frame, + frame_ptr, + node_count, + coefficient.data_ptr(), + static_cast(coefficient.numel() / 2), + potential, + force.data_ptr(), + atom_virial.data_ptr()}; if (kernel == kLes) { dispatch_channels(static_cast(channels), buffers, owner_begin, owner_count, prefactor, stream.stream()); @@ -906,54 +928,46 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, // The frames of the owners, on the axis of the owners: the frames // themselves when every node is an owner, else the one frame the owners // share, written on the device. - torch::Tensor owner_frame_ptr = frame_ptr; + torch::Tensor owner_frame_tensor; + const int64_t* owner_frame_ptr = frame_ptr; if (!every_node) { - owner_frame_ptr = - torch::full({2}, int64_t{owner_count}, frame_ptr.options()); - owner_frame_ptr.select(0, 0).zero_(); + owner_frame_tensor = + torch::full({2}, int64_t{owner_count}, counts.options()); + owner_frame_tensor.select(0, 0).zero_(); + owner_frame_ptr = owner_frame_tensor.data_ptr(); } const bool* owner_mask = node_mask.data_ptr() + owner_begin; - const int owner_slice_count = slices_for(owner_count); const dim3 owner_slice_grid(frame_count, owner_slice_count); - const int components = constrained ? 10 : 9; - auto slice_virial = torch::empty( - {static_cast(frame_count) * owner_slice_count * components}, - double_options); if (constrained) { virial_slices<<>>( - atom_virial.data_ptr(), potential.data_ptr(), owner_mask, - owner_frame_ptr.data_ptr(), static_cast(channels), - constrained_channel, slice_virial.data_ptr()); + atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr, + static_cast(channels), constrained_channel, slice_virial); C10_CUDA_KERNEL_LAUNCH_CHECK(); virial_total<<>>( - slice_virial.data_ptr(), owner_slice_count, - virial.data_ptr(), frame_stat.data_ptr()); + slice_virial, owner_slice_count, virial.data_ptr(), frame_stat); } else { virial_slices<<>>( - atom_virial.data_ptr(), potential.data_ptr(), owner_mask, - owner_frame_ptr.data_ptr(), static_cast(channels), - constrained_channel, slice_virial.data_ptr()); + atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr, + static_cast(channels), constrained_channel, slice_virial); C10_CUDA_KERNEL_LAUNCH_CHECK(); virial_total<<>>( - slice_virial.data_ptr(), owner_slice_count, - virial.data_ptr(), nullptr); + slice_virial, owner_slice_count, virial.data_ptr(), nullptr); } C10_CUDA_KERNEL_LAUNCH_CHECK(); if (constrained && communicator != nullptr) { // The mean potential of the constraint runs over the charged nodes of // every domain. - auto potential_sum = - frame_stat.view({frame_count, kStatWidth}).select(1, kStatPotential); + auto stat = + torch::from_blob(frame_stat, {frame_count, kStatWidth}, double_options); + auto potential_sum = stat.select(1, kStatPotential); potential_sum.copy_( sum_over_domains(*communicator, potential_sum.contiguous())); } if (owner_count > 0) { project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0, stream>>>( - potential.data_ptr(), - pair_charge.data_ptr() + int64_t{owner_begin} * channels, - owner_mask, frame.data_ptr() + owner_begin, - constrained ? frame_stat.data_ptr() : nullptr, owner_count, + potential, pair_charge + int64_t{owner_begin} * channels, owner_mask, + frame + owner_begin, frame_stat, owner_count, static_cast(channels), constrained_channel, atom_energy.data_ptr(), charge_gradient.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); diff --git a/source/op/pt/long_range/frame_sums.cuh b/source/op/pt/long_range/frame_sums.cuh index a773edd63a..d3f8d40af6 100644 --- a/source/op/pt/long_range/frame_sums.cuh +++ b/source/op/pt/long_range/frame_sums.cuh @@ -1,9 +1,9 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // // Per-frame bookkeeping shared by the CUDA units of the latent-charge long -// range: the offsets of the frames on the node axis, the sums over the nodes -// of every frame in a fixed order, and the statistics of the charge -// constraint. +// range: the workspace of a call, the offsets of the frames on the node axis, +// the sums over the nodes of every frame in a fixed order, and the statistics +// of the charge constraint. // // The kernels are defined in an unnamed namespace, so that every unit that // includes this header compiles its own copy. @@ -11,11 +11,44 @@ #pragma once #include +#include +#include #include namespace { +// The intermediate buffers of one call, carved out of a single device +// allocation: ``reserve`` lays out a buffer of ``count`` elements and returns +// its byte offset, and ``at`` its address once ``allocate`` has made the +// allocation. Every buffer starts on a 256-byte boundary. +class Workspace { + public: + template + int64_t reserve(int64_t count) { + bytes_ = (bytes_ + kAlignment - 1) / kAlignment * kAlignment; + const int64_t offset = bytes_; + bytes_ += count * static_cast(sizeof(T)); + return offset; + } + + void allocate(const torch::Device& device) { + storage_ = torch::empty( + {std::max(bytes_, 1)}, + torch::TensorOptions().dtype(torch::kUInt8).device(device)); + } + + template + T* at(int64_t offset) const { + return reinterpret_cast(storage_.data_ptr() + offset); + } + + private: + static constexpr int64_t kAlignment = 256; + int64_t bytes_ = 0; + torch::Tensor storage_; +}; + // Index of the frame holding ``node``, or -1 for a padding node past every // frame. ``frame_ptr`` holds ``frame_count + 1`` offsets. __device__ __forceinline__ int frame_of(const int64_t* frame_ptr, diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu index f787c5cdb8..4763c72fcc 100644 --- a/source/op/pt/long_range/latent_charge_sum.cu +++ b/source/op/pt/long_range/latent_charge_sum.cu @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -57,9 +58,8 @@ struct Evaluation { owner_count, communicator); } return deepmd::long_range::direct_sum( - coord, charge, node_mask, n_node_per_frame, total_charge, - coefficient.to(torch::kFloat32), kernel, channel, prefactor, - owner_begin, owner_count, communicator); + coord, charge, node_mask, n_node_per_frame, total_charge, coefficient, + kernel, channel, prefactor, owner_begin, owner_count, communicator); } }; @@ -116,18 +116,25 @@ ChargeSumOutputs evaluate_across_domains(const Evaluation& evaluation, virial}; } -// Whether the frames carry cells: ``box`` is given and no row of it is zero, -// or no frame carries one. +// Whether the frames carry cells: true when ``box`` is given and none of its +// rows is zero, false when it is absent or all of its rows are; a call that +// mixes the two is rejected. bool periodic_frames(const c10::optional& box) { if (!box.has_value()) { return false; } - const auto has_cell = box->ne(0.0).any(1); - const bool any = has_cell.any().item(); - TORCH_CHECK(!any || has_cell.all().item(), + const auto cells = box->contiguous(); + const double* value = cells.data_ptr(); + const int64_t frame_count = cells.size(0); + int64_t periodic = 0; + for (int64_t f = 0; f < frame_count; ++f) { + periodic += std::any_of(value + 9 * f, value + 9 * (f + 1), + [](double entry) { return entry != 0.0; }); + } + TORCH_CHECK(periodic == 0 || periodic == frame_count, "latent_charge_sum: every frame of a call is periodic, or none " "is"); - return any; + return periodic > 0; } ChargeSumOutputs latent_charge_sum(torch::Tensor coord, diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu index f952494955..130cbdcb11 100644 --- a/source/op/pt/long_range/mesh_sum.cu +++ b/source/op/pt/long_range/mesh_sum.cu @@ -34,11 +34,14 @@ // // The charges spread in 64-bit fixed point, so that the mesh, and every // output, does not depend on the order of the atomic additions. The sums over -// the modes and over the nodes of a frame add per-block float64 sums in a -// fixed order. The meshes and the transforms are float64; the stencils, which -// would otherwise be bound by the float64 throughput of the device, weigh and -// read them in float32, the precision of the outputs, each warp adding the -// sums of its lanes in float64. +// the modes and over the nodes of a frame add float64 sums in a fixed order. +// The meshes, the transforms and the influence function are float32, against +// a float64 throughput of a small fraction of the float32 one on many +// devices; the B-spline weights are evaluated in float64, whose rounding to +// float32 leaves their sum unbiased. The forces then carry the rounding of +// float32, and the energy of a frame that of the reciprocal sum including the +// interaction of every charge with itself, which the self term cancels: a few +// 1e-7 of the self energy. // // The nodes whose outputs a call evaluates, its owners, are either every node // or, across the domains of a decomposed system, the charged nodes of one MPI @@ -49,6 +52,7 @@ #include #include #include +#include #include #include @@ -57,7 +61,11 @@ #include #include #include +#include +#include +#include #include +#include #include #include "charge_sum.h" @@ -66,15 +74,14 @@ namespace { constexpr int kThreads = 256; -constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; constexpr int kMaxOrder = 12; enum KernelKind : int { kLes = 0, kSog = 1 }; // Resolution of the fixed-point charge mesh: 2^-40 e, far below the float32 // precision of the raw charges, for a sum of at most 2^23 e at a point. -constexpr double kFixedScale = 1099511627776.0; // 2^40 -constexpr float kFixedScaleFloat = 1099511627776.f; +constexpr float kFixedScaleFloat = 1099511627776.f; // 2^40 constexpr double kPi = 3.14159265358979323846; @@ -83,14 +90,16 @@ struct FrameMesh { // B^-1, row-major: s_d = sum_j r_j inverse[3 j + d]. double inverse[9]; double volume; + // -n_d (B^-1)_jd, row-major in (j, d): the force at fixed charges from the + // derivatives of the energy with respect to the scaled coordinates u_d. + float gradient_to_force[9]; int n[3]; // The first point of the meshes of the frame, channel-major, in the mesh // buffer of the call. int64_t offset; }; -// The squared lengths of the lattice vectors are formed term by term, as the -// Python front end forms them, so that both round the mesh size identically. +// The smallest even number of at least ``n`` without a prime factor above 7. int64_t even_fft_size(int64_t n) { int64_t m = std::max(2, n + n % 2); while (true) { @@ -128,6 +137,9 @@ FrameMesh frame_mesh(const double* cell, double spacing) { mesh.inverse[k] = adjugate[k] / det; } mesh.volume = std::abs(det); + // The squared lengths of the lattice vectors are formed term by term, as + // the Python front end forms them, so that both round the mesh size + // identically. for (int axis = 0; axis < 3; ++axis) { const double* v = cell + 3 * axis; const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); @@ -138,6 +150,15 @@ FrameMesh frame_mesh(const double* cell, double spacing) { mesh.n[axis] = static_cast(even_fft_size(static_cast(points))); } + TORCH_CHECK(int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2] < + std::numeric_limits::max(), + "latent_charge_sum: the mesh of a cell exceeds 2^31 points"); + for (int j = 0; j < 3; ++j) { + for (int d = 0; d < 3; ++d) { + mesh.gradient_to_force[3 * j + d] = + static_cast(-mesh.n[d] * mesh.inverse[3 * j + d]); + } + } return mesh; } @@ -149,28 +170,19 @@ __host__ __device__ int64_t mesh_points(const FrameMesh& mesh) { constexpr int kStencilWarps = 4; constexpr int kWarp = 32; -// The stencil of one node along one axis of its mesh: the index of the -// first of its p points on the axis, which the others follow cyclically, and -// their B-spline weights and slopes, evaluated in float64 and stored in -// float32, the precision of the outputs: rounding a recursion of float32 steps -// instead would bias the sum of the weights. +// The stencil of one node along one axis of its mesh: for each of its p +// points on the axis, its offset on the flat mesh of a frame along that axis, +// the index of the point times the stride of the axis, and its B-spline +// weight and slope, evaluated in float64 and stored in float32, the precision +// of the outputs: rounding a recursion of float32 steps instead would bias +// the sum of the weights. template struct AxisStencil { - int start; + int offset[kOrder]; float value[kOrder]; float slope[kOrder]; }; -// The index of the point ``start + k`` of an axis of ``n`` points, for -// ``start`` < n and k < p, which exceeds 2 n only when a stencil wider than -// the axis wraps more than once. -__device__ __forceinline__ int wrap_index(int index, int n) { - if (index < n) { - return index; - } - return index - n < n ? index - n : index % n; -} - // The stencil of a node along axis ``d``. The scaled coordinate u = n_d s_d // of the fractional coordinate s_d reaches the points floor(u) - p + 1 + k, // k = 0, ..., p - 1, with the weights value[k] = M_p(t + p - 1 - k), t the @@ -215,135 +227,198 @@ __device__ void axis_stencil(const double* coord, stencil.slope[0] = -1.f; stencil.slope[1] = 1.f; } - const int first = static_cast(base) - kOrder + 1; - stencil.start = ((first % mesh.n[d]) + mesh.n[d]) % mesh.n[d]; + const int n = mesh.n[d]; + const int stride = d == 0 ? mesh.n[1] * mesh.n[2] : d == 1 ? mesh.n[2] : 1; + // A stencil wider than the axis wraps more than once. + const int first = ((static_cast(base) - kOrder + 1) % n + n) % n; #pragma unroll for (int k = 0; k < kOrder; ++k) { stencil.value[k] = static_cast(a[kOrder - 1 - k]); + stencil.offset[k] = (first + k) % n * stride; } } -// Writes the stencils of the charged nodes along the three axes of their -// meshes, one thread per node and axis, which the kernels that spread and -// interpolate read: every warp instruction in float64 occupies the float64 -// pipe of the device whatever the number of its active lanes, so the -// recursion runs once per node and axis with every lane busy. +// Writes, one thread per node and axis, the stencils of the charged nodes +// along the three axes of their meshes, and, on the thread of the first axis, +// the charges of every node, shifted in the constrained channel and zero on +// the nodes that carry none. Every float64 warp instruction occupies the +// float64 pipe of the device whatever the number of its active lanes, so the +// float64 work of a node runs here with every lane busy rather than in the +// warps of the kernels that spread and interpolate. template __global__ void __launch_bounds__(kThreads) node_stencils(const double* __restrict__ coord, + const float* __restrict__ charge, const bool* __restrict__ node_mask, const int64_t* __restrict__ frame_ptr, const FrameMesh* __restrict__ frames, + const double* __restrict__ frame_stat, int frame_count, int64_t node_count, - AxisStencil* __restrict__ stencils) { + int channels, + int channel, + AxisStencil* __restrict__ stencils, + float* __restrict__ node_charge) { for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; index < 3 * node_count; index += int64_t{gridDim.x} * blockDim.x) { const int64_t node = index / 3; + const int axis = static_cast(index - 3 * node); const int f = frame_of(frame_ptr, frame_count, node); - if (f >= 0 && node_mask[node]) { - axis_stencil(coord + 3 * node, frames[f], - static_cast(index - 3 * node), stencils[index]); + const bool charged = f >= 0 && node_mask[node]; + if (charged) { + axis_stencil(coord + 3 * node, frames[f], axis, stencils[index]); + } + if (axis == 0) { + const double shift = charged && channel >= 0 + ? frame_stat[f * kStatWidth + kStatShift] + : 0.0; + for (int c = 0; c < channels; ++c) { + const double raw = charged ? charge[node * channels + c] : 0.0; + node_charge[node * channels + c] = + static_cast(c == channel ? raw - shift : raw); + } } } } -// The point ``t`` of the p^3 points of a stencil: its offset on the flat mesh -// of a frame and its indices (k1, k2, k3) along the axes. The index of the -// last axis runs fastest, so that the lanes of a warp touch contiguous runs -// of the mesh. +// Rounds ``x`` to the nearest integer, ties to even, as __float2ll_rn does, +// in integer operations: the conversion instruction occupies the float64 pipe. +// A finite x is a 24-bit significand times 2^exponent; a negative exponent +// shifts the significand right with rounding. Magnitudes of 2^63 and beyond, +// infinities and NaN saturate. +__device__ __forceinline__ long long float_to_fixed(float x) { + const unsigned bits = __float_as_uint(x); + const int biased = static_cast((bits >> 23) & 0xffu); + // |x| < 1/2, zero and subnormal numbers included, rounds to zero. + if (biased < 126) { + return 0; + } + const int exponent = biased - 150; + if (exponent >= 40) { + return bits >> 31 ? std::numeric_limits::min() + : std::numeric_limits::max(); + } + const long long significand = (bits & 0x7fffffu) | 0x800000u; + long long magnitude; + if (exponent >= 0) { + magnitude = significand << exponent; + } else { + const int shift = -exponent; + const long long kept = significand >> shift; + const long long rest = significand & ((1ll << shift) - 1); + const long long half = 1ll << (shift - 1); + magnitude = kept + (rest > half || (rest == half && (kept & 1)) ? 1 : 0); + } + return bits >> 31 ? -magnitude : magnitude; +} + +// The assignment of the p^3 points of a stencil to the lanes of a warp: the +// lanes take the (k2, k3) pairs, k3 running fastest, so that a warp touches +// contiguous runs of the mesh, and every lane walks the k1 axis for its +// pairs; while the pairs are fewer than the lanes, groups of lanes split the +// k1 axis. template -struct StencilPoint { - int64_t offset; - int k1; - int k2; - int k3; +struct LaneSlots { + static constexpr int kPairs = kOrder * kOrder; + static constexpr int kGroups = kPairs >= kWarp ? 1 : kWarp / kPairs; + static constexpr int kSlots = kGroups * kPairs; + static constexpr int kRounds = (kSlots + kWarp - 1) / kWarp; +}; - __device__ StencilPoint(int t, - const FrameMesh& mesh, - const AxisStencil* axes) - : k1(t / (kOrder * kOrder)), k2((t / kOrder) % kOrder), k3(t % kOrder) { - offset = (int64_t{wrap_index(axes[0].start + k1, mesh.n[0])} * mesh.n[1] + - wrap_index(axes[1].start + k2, mesh.n[1])) * - mesh.n[2] + - wrap_index(axes[2].start + k3, mesh.n[2]); +// Copies the stencils of a node along its three axes to the shared memory of +// its warp. +template +__device__ __forceinline__ void load_stencil( + const AxisStencil* __restrict__ stencils, + int64_t node, + int lane, + AxisStencil* shared) { + constexpr int kWords = 3 * sizeof(AxisStencil) / sizeof(int); + const int* source = reinterpret_cast(stencils + 3 * node); + int* target = reinterpret_cast(shared); + for (int word = lane; word < kWords; word += kWarp) { + target[word] = source[word]; } -}; + __syncwarp(); +} -// Writes the charges of every node, shifted in the constrained channel and -// zero on the nodes that carry none, and spreads them to the meshes of their -// frames, one warp per node, the lanes taking the points of its stencil in -// turn. +// Spreads the charges of the charged nodes to the meshes of their frames, one +// warp per node, one pass per channel. template __global__ void __launch_bounds__(kStencilWarps* kWarp) - spread_charges(const float* __restrict__ charge, - const bool* __restrict__ node_mask, + spread_charges(const bool* __restrict__ node_mask, const int64_t* __restrict__ frame_ptr, const FrameMesh* __restrict__ frames, const AxisStencil* __restrict__ stencils, - const double* __restrict__ frame_stat, + const float* __restrict__ node_charge, int frame_count, int64_t node_count, int channels, - int channel, - float* __restrict__ node_charge, unsigned long long* __restrict__ mesh) { + using Slots = LaneSlots; + __shared__ AxisStencil shared[kStencilWarps][3]; const int warp = threadIdx.x / kWarp; const int lane = threadIdx.x % kWarp; + const AxisStencil* axes = shared[warp]; for (int64_t node = int64_t{blockIdx.x} * kStencilWarps + warp; node < node_count; node += int64_t{gridDim.x} * kStencilWarps) { const int f = frame_of(frame_ptr, frame_count, node); - const bool charged = f >= 0 && node_mask[node]; - const double shift = - charged && channel >= 0 ? frame_stat[f * kStatWidth + kStatShift] : 0.0; - float q[kMaxChannels]; -#pragma unroll - for (int c = 0; c < kMaxChannels; ++c) { - const double raw = - charged && c < channels ? charge[node * channels + c] : 0.0; - q[c] = static_cast(c == channel ? raw - shift : raw); - } - if (lane < channels) { - node_charge[node * channels + lane] = charged ? q[lane] : 0.f; - } - if (!charged) { + if (f < 0 || !node_mask[node]) { continue; } - const FrameMesh grid = frames[f]; - const AxisStencil* axes = stencils + 3 * node; - const int64_t points = mesh_points(grid); - unsigned long long* frame_mesh = mesh + grid.offset; - for (int t = lane; t < kOrder * kOrder * kOrder; t += kWarp) { - const StencilPoint point(t, grid, axes); - const float w = axes[0].value[point.k1] * axes[1].value[point.k2] * - axes[2].value[point.k3] * kFixedScaleFloat; + load_stencil(stencils, node, lane, shared[warp]); + const FrameMesh& grid = frames[f]; + for (int c = 0; c < channels; ++c) { + const float q = node_charge[node * channels + c] * kFixedScaleFloat; + unsigned long long* channel_mesh = + mesh + grid.offset + c * mesh_points(grid); #pragma unroll - for (int c = 0; c < kMaxChannels; ++c) { - if (c < channels) { - atomicAdd(frame_mesh + c * points + point.offset, - static_cast(__float2ll_rn(q[c] * w))); + for (int round = 0; round < Slots::kRounds; ++round) { + const int slot = lane + kWarp * round; + if (slot < Slots::kSlots) { + const int group = slot / Slots::kPairs; + const int pair = slot - group * Slots::kPairs; + const int k2 = pair / kOrder; + const int k3 = pair - k2 * kOrder; + const float weight = q * axes[1].value[k2] * axes[2].value[k3]; + const int column = axes[1].offset[k2] + axes[2].offset[k3]; +#pragma unroll + for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) { + atomicAdd(channel_mesh + axes[0].offset[k1] + column, + static_cast( + float_to_fixed(weight * axes[0].value[k1]))); + } } } } + // The stencil of the next node of the warp overwrites this one. + __syncwarp(); } } -// Converts the fixed-point mesh to float64 in place. -__global__ void fixed_to_double(int64_t count, unsigned long long* mesh) { +// Converts the fixed-point mesh to the float32 mesh that the transforms read. +__global__ void fixed_to_float(int64_t count, + const unsigned long long* __restrict__ fixed, + float* __restrict__ mesh) { + constexpr float kFixedInverse = 1.f / kFixedScaleFloat; for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; index < count; index += int64_t{gridDim.x} * blockDim.x) { - const long long fixed = static_cast(mesh[index]); - reinterpret_cast(mesh)[index] = - static_cast(fixed) / kFixedScale; + mesh[index] = + __ll2float_rn(static_cast(fixed[index])) * kFixedInverse; } } +// The weights M_p(1), ..., M_p(p - 1) of the assignment factor. +struct IntegerWeights { + double value[kMaxOrder - 1]; +}; + // The inverse squared assignment factor |D_d(m)|^-2 along every axis of the // half spectrum of one frame: n1 entries, then n2, then n3 / 2 + 1. __global__ void assignment_factor(FrameMesh mesh, int order, - const double* __restrict__ integer_weight, - double* __restrict__ factor) { + IntegerWeights integer_weight, + float* __restrict__ factor) { const int half = mesh.n[2] / 2 + 1; const int total = mesh.n[0] + mesh.n[1] + half; for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < total; @@ -365,47 +440,75 @@ __global__ void assignment_factor(FrameMesh mesh, double s; double c; sincos(theta * l, &s, &c); - re += integer_weight[l] * c; - im += integer_weight[l] * s; + re += integer_weight.value[l] * c; + im += integer_weight.value[l] * s; } - factor[index] = 1.0 / (re * re + im * im); + factor[index] = static_cast(1.0 / (re * re + im * im)); } } -// The transform of the kernel and its logarithmic derivative with respect to -// k^2. ``coefficient`` holds alpha, or the amplitudes and then the rates -// c_g = log2(e) / (2 sigma_g^2) of the Gaussians. -__device__ void kernel_transform(int kernel, - const double* __restrict__ coefficient, - int gaussians, - double k_sq, - double& value, - double& log_slope) { +// The terms of the transform of the kernel, in float32, the precision of the +// transforms: 1 / (4 alpha^2) of the les kernel, or, for every Gaussian of the +// sog kernel, its weight a_g (2 pi)^(3/2) sigma_g^3 and half its squared +// width, from ``coefficient``, which holds alpha, or the amplitudes and then +// the rates c_g = log2(e) / (2 sigma_g^2). +struct KernelTerms { + float inv_four_alpha_sq; + float weight[kMaxGaussians]; + float half_width_sq[kMaxGaussians]; +}; + +__device__ void load_kernel_terms(int kernel, + const double* __restrict__ coefficient, + int gaussians, + KernelTerms& terms) { if (kernel == kLes) { - const double alpha = coefficient[0]; - const double inv_four_alpha_sq = 0.25 / (alpha * alpha); - value = 4.0 * kPi * exp(-k_sq * inv_four_alpha_sq) / k_sq; - log_slope = -inv_four_alpha_sq - 1.0 / k_sq; + if (threadIdx.x == 0) { + const double alpha = coefficient[0]; + terms.inv_four_alpha_sq = static_cast(0.25 / (alpha * alpha)); + } return; } // (2 pi)^{3/2} constexpr double kGaussianNorm = 15.749609945722419; - value = 0.0; - double slope = 0.0; - for (int g = 0; g < gaussians; ++g) { + for (int g = threadIdx.x; g < gaussians; g += blockDim.x) { // sigma^2 = log2(e) / (2 c) = 1 / (2 c ln 2) const double width_sq = 0.72134752044448170 / coefficient[gaussians + g]; - const double term = coefficient[g] * kGaussianNorm * width_sq * - sqrt(width_sq) * exp(-0.5 * width_sq * k_sq); + terms.weight[g] = static_cast(coefficient[g] * kGaussianNorm * + width_sq * sqrt(width_sq)); + terms.half_width_sq[g] = static_cast(0.5 * width_sq); + } +} + +// The transform of the kernel at k^2 and its logarithmic derivative with +// respect to k^2. +__device__ __forceinline__ void kernel_transform(int kernel, + const KernelTerms& terms, + int gaussians, + float k_sq, + float& value, + float& log_slope) { + if (kernel == kLes) { + value = 4.f * static_cast(kPi) * + expf(-k_sq * terms.inv_four_alpha_sq) / k_sq; + log_slope = -terms.inv_four_alpha_sq - 1.f / k_sq; + return; + } + value = 0.f; + float slope = 0.f; + for (int g = 0; g < gaussians; ++g) { + const float term = terms.weight[g] * expf(-terms.half_width_sq[g] * k_sq); value += term; - slope -= 0.5 * width_sq * term; + slope -= terms.half_width_sq[g] * term; } - log_slope = value != 0.0 ? slope / value : 0.0; + log_slope = value != 0.f ? slope / value : 0.f; } // The influence function of one frame on its half spectrum: multiplies every // mode of the spectra of the channels by G in place and sums, per block, the -// energy of the modes and their virial, xx yy zz xy xz yz. +// energy of the modes and their virial, xx yy zz xy xz yz. The modes are +// evaluated in float32, the precision of the transforms, summed in float32 +// over the modes of a row that a lane takes, and in float64 over the rows. constexpr int kModeSums = 7; __global__ void __launch_bounds__(kThreads) @@ -415,59 +518,84 @@ __global__ void __launch_bounds__(kThreads) const double* __restrict__ coefficient, int gaussians, double prefactor, - const double* __restrict__ factor, - double2* __restrict__ spectrum, + const float* __restrict__ factor, + float2* __restrict__ spectrum, double* __restrict__ block_sum) { + __shared__ KernelTerms terms; + load_kernel_terms(kernel, coefficient, gaussians, terms); + __syncthreads(); const int n1 = mesh.n[0]; const int n2 = mesh.n[1]; const int n3 = mesh.n[2]; const int half = n3 / 2 + 1; - const int64_t modes = int64_t{n1} * n2 * half; + const int modes = n1 * n2 * half; + // k_a = 2 pi sum_d (B^-1)_ad m_d + float reciprocal[9]; +#pragma unroll + for (int k = 0; k < 9; ++k) { + reciprocal[k] = static_cast(2.0 * kPi * mesh.inverse[k]); + } + const float scale = static_cast(prefactor / mesh.volume); + // Every warp takes rows (i1, i2) of the half spectrum, its lanes the modes + // of a row, which lie contiguously. + const int lane = threadIdx.x % kWarp; + const int warps = gridDim.x * (kThreads / kWarp); double sum[kModeSums] = {}; - for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; - index < modes; index += int64_t{gridDim.x} * blockDim.x) { - const int i3 = static_cast(index % half); - const int i2 = static_cast((index / half) % n2); - const int i1 = static_cast(index / (int64_t{half} * n2)); + for (int row = blockIdx.x * (kThreads / kWarp) + threadIdx.x / kWarp; + row < n1 * n2; row += warps) { + float row_sum[kModeSums] = {}; + const int i1 = row / n2; + const int i2 = row - i1 * n2; const int m1 = i1 <= n1 / 2 ? i1 : i1 - n1; const int m2 = i2 <= n2 / 2 ? i2 : i2 - n2; - const int m3 = i3; - const bool weightless = (m1 == 0 && m2 == 0 && m3 == 0) || 2 * m1 == n1 || - 2 * m2 == n2 || 2 * m3 == n3; - double influence = 0.0; - double k[3] = {0.0, 0.0, 0.0}; - double log_slope = 0.0; - if (!weightless) { - for (int a = 0; a < 3; ++a) { - k[a] = 2.0 * kPi * - (mesh.inverse[3 * a] * m1 + mesh.inverse[3 * a + 1] * m2 + - mesh.inverse[3 * a + 2] * m3); + const bool row_weightless = 2 * m1 == n1 || 2 * m2 == n2; + float k_row[3]; +#pragma unroll + for (int a = 0; a < 3; ++a) { + k_row[a] = reciprocal[3 * a] * m1 + reciprocal[3 * a + 1] * m2; + } + const float row_factor = scale * factor[i1] * factor[n1 + i2]; + for (int i3 = lane; i3 < half; i3 += kWarp) { + const int m3 = i3; + const bool weightless = + row_weightless || 2 * m3 == n3 || (m1 == 0 && m2 == 0 && m3 == 0); + float influence = 0.f; + float k[3] = {0.f, 0.f, 0.f}; + float log_slope = 0.f; + if (!weightless) { +#pragma unroll + for (int a = 0; a < 3; ++a) { + k[a] = k_row[a] + reciprocal[3 * a + 2] * m3; + } + const float k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2]; + float transform; + kernel_transform(kernel, terms, gaussians, k_sq, transform, log_slope); + influence = row_factor * transform * factor[n1 + n2 + i3]; } - const double k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2]; - double transform; - kernel_transform(kernel, coefficient, gaussians, k_sq, transform, - log_slope); - influence = prefactor * transform / mesh.volume * factor[i1] * - factor[n1 + i2] * factor[n1 + n2 + i3]; + // An interior plane of the half spectrum stands for its mirror as well. + const float multiplicity = m3 == 0 ? 1.f : 2.f; + const int index = row * half + i3; + float energy = 0.f; + for (int c = 0; c < channels; ++c) { + float2& value = spectrum[int64_t{c} * modes + index]; + energy += value.x * value.x + value.y * value.y; + value.x *= influence; + value.y *= influence; + } + energy *= 0.5f * multiplicity * influence; + const float twice = 2.f * log_slope * energy; + row_sum[0] += energy; + row_sum[1] += energy + twice * k[0] * k[0]; + row_sum[2] += energy + twice * k[1] * k[1]; + row_sum[3] += energy + twice * k[2] * k[2]; + row_sum[4] += twice * k[0] * k[1]; + row_sum[5] += twice * k[0] * k[2]; + row_sum[6] += twice * k[1] * k[2]; } - // An interior plane of the half spectrum stands for its mirror as well. - const double multiplicity = m3 == 0 ? 1.0 : 2.0; - double energy = 0.0; - for (int c = 0; c < channels; ++c) { - double2& value = spectrum[c * modes + index]; - energy += value.x * value.x + value.y * value.y; - value.x *= influence; - value.y *= influence; +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + sum[s] += row_sum[s]; } - energy *= 0.5 * multiplicity * influence; - const double twice = 2.0 * log_slope * energy; - sum[0] += energy; - sum[1] += energy + twice * k[0] * k[0]; - sum[2] += energy + twice * k[1] * k[1]; - sum[3] += energy + twice * k[2] * k[2]; - sum[4] += twice * k[0] * k[1]; - sum[5] += twice * k[0] * k[2]; - sum[6] += twice * k[1] * k[2]; } __shared__ double partial[kModeSums][kThreads]; for (int s = 0; s < kModeSums; ++s) { @@ -541,7 +669,7 @@ __device__ double contact_value(int kernel, // Adds ``value`` over the lanes of a warp in a fixed butterfly order; every // lane receives the total. -__device__ __forceinline__ double warp_sum(double value) { +__device__ __forceinline__ float warp_sum(float value) { #pragma unroll for (int stride = kWarp / 2; stride > 0; stride /= 2) { value += __shfl_xor_sync(0xffffffffu, value, stride); @@ -549,124 +677,134 @@ __device__ __forceinline__ double warp_sum(double value) { return value; } -// Interpolates the potential mesh at every owner node, one warp per node: -// its potential, less that of its charge at itself, its force at fixed -// charges and the share F (x) r of its virial. Owners that carry no charge get -// zeros. The lanes take the points of the stencil in turn and the warp adds -// their sums. +// Interpolates the potential meshes at every owner node, one warp per node, +// one pass per channel: the potential phi^rec_c of the mesh, which still holds +// that of the charge at itself, and the force at fixed charges. Owners that +// carry no charge get zeros. Every lane sums its points in float32, the +// precision of the meshes, and the warp adds the lane sums. template __global__ void __launch_bounds__(kStencilWarps* kWarp) - interpolate_owners(const double* __restrict__ coord, - const float* __restrict__ node_charge, - const bool* __restrict__ node_mask, + interpolate_owners(const bool* __restrict__ node_mask, const int64_t* __restrict__ frame_ptr, const FrameMesh* __restrict__ frames, const AxisStencil* __restrict__ stencils, + const float* __restrict__ node_charge, const float* __restrict__ potential_mesh, - const double* __restrict__ coefficient, int frame_count, int owner_begin, int owner_count, int channels, - int kernel, - int gaussians, - double prefactor, - double* __restrict__ potential, - double* __restrict__ force, - double* __restrict__ moment, - double* __restrict__ atom_energy) { + float* __restrict__ potential, + float* __restrict__ force) { + using Slots = LaneSlots; + __shared__ AxisStencil shared[kStencilWarps][3]; const int warp = threadIdx.x / kWarp; const int lane = threadIdx.x % kWarp; - const double self = prefactor * contact_value(kernel, coefficient, gaussians); + const AxisStencil* axes = shared[warp]; for (int owner = blockIdx.x * kStencilWarps + warp; owner < owner_count; owner += gridDim.x * kStencilWarps) { const int64_t node = int64_t{owner_begin} + owner; const int f = frame_of(frame_ptr, frame_count, node); - const bool charged = f >= 0 && node_mask[node]; - const float* q = node_charge + node * channels; - double phi[kMaxChannels] = {}; - double f_node[3] = {0.0, 0.0, 0.0}; - if (charged) { - const FrameMesh grid = frames[f]; - const AxisStencil* axes = stencils + 3 * node; - const int64_t points = mesh_points(grid); - const float* frame_mesh = potential_mesh + grid.offset; - // Every lane sums its points in float32, the warp the lane sums in - // float64. ``gradient`` is the derivative of sum_c q_c phi_c with - // respect to u_d. - float lane_phi[kMaxChannels] = {}; - float lane_gradient[3] = {0.f, 0.f, 0.f}; - for (int t = lane; t < kOrder * kOrder * kOrder; t += kWarp) { - const StencilPoint point(t, grid, axes); - const float v1 = axes[0].value[point.k1]; - const float v2 = axes[1].value[point.k2]; - const float v3 = axes[2].value[point.k3]; - float charge_potential = 0.f; -#pragma unroll - for (int c = 0; c < kMaxChannels; ++c) { - if (c < channels) { - const float value = frame_mesh[c * points + point.offset]; - lane_phi[c] += v1 * v2 * v3 * value; - charge_potential += q[c] * value; - } - } - lane_gradient[0] += - axes[0].slope[point.k1] * v2 * v3 * charge_potential; - lane_gradient[1] += - v1 * axes[1].slope[point.k2] * v3 * charge_potential; - lane_gradient[2] += - v1 * v2 * axes[2].slope[point.k3] * charge_potential; + if (f < 0 || !node_mask[node]) { + for (int c = lane; c < channels; c += kWarp) { + potential[int64_t{owner} * channels + c] = 0.f; } + if (lane < 3) { + force[int64_t{owner} * 3 + lane] = 0.f; + } + continue; + } + load_stencil(stencils, node, lane, shared[warp]); + const FrameMesh& grid = frames[f]; + // The derivative of sum_c q_c phi_c with respect to u_d. + float gradient[3] = {0.f, 0.f, 0.f}; + for (int c = 0; c < channels; ++c) { + const float* channel_mesh = + potential_mesh + grid.offset + c * mesh_points(grid); + float phi = 0.f; + float slope[3] = {0.f, 0.f, 0.f}; +#pragma unroll + for (int round = 0; round < Slots::kRounds; ++round) { + const int slot = lane + kWarp * round; + if (slot < Slots::kSlots) { + const int group = slot / Slots::kPairs; + const int pair = slot - group * Slots::kPairs; + const int k2 = pair / kOrder; + const int k3 = pair - k2 * kOrder; + const float v2 = axes[1].value[k2]; + const float v3 = axes[2].value[k3]; + const float v23 = v2 * v3; + const float s2v3 = axes[1].slope[k2] * v3; + const float v2s3 = v2 * axes[2].slope[k3]; + const int column = axes[1].offset[k2] + axes[2].offset[k3]; #pragma unroll - for (int c = 0; c < kMaxChannels; ++c) { - if (c < channels) { - phi[c] = warp_sum(lane_phi[c]) - self * q[c]; + for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) { + const float value = channel_mesh[axes[0].offset[k1] + column]; + const float v1_value = axes[0].value[k1] * value; + phi += v1_value * v23; + slope[0] += axes[0].slope[k1] * value * v23; + slope[1] += v1_value * s2v3; + slope[2] += v1_value * v2s3; + } } } - double gradient[3]; + phi = warp_sum(phi); + const float q = node_charge[node * channels + c]; #pragma unroll for (int d = 0; d < 3; ++d) { - gradient[d] = warp_sum(lane_gradient[d]) * grid.n[d]; + gradient[d] += q * warp_sum(slope[d]); } - // d E / d r_j = sum_d d E / d u_d n_d (B^-1)_jd -#pragma unroll - for (int j = 0; j < 3; ++j) { - f_node[j] = -(gradient[0] * grid.inverse[3 * j] + - gradient[1] * grid.inverse[3 * j + 1] + - gradient[2] * grid.inverse[3 * j + 2]); + if (lane == 0) { + potential[int64_t{owner} * channels + c] = phi; } } - if (lane == 0) { - double energy = 0.0; - for (int c = 0; c < channels; ++c) { - potential[int64_t{owner} * channels + c] = phi[c]; - energy += q[c] * phi[c]; - } - atom_energy[owner] = 0.5 * energy; - const double* r = coord + 3 * node; - for (int a = 0; a < 3; ++a) { - force[int64_t{owner} * 3 + a] = f_node[a]; - for (int b = 0; b < 3; ++b) { - moment[int64_t{owner} * 9 + 3 * a + b] = f_node[a] * r[b]; - } - } + if (lane < 3) { + // F_j = -sum_d dE/du_d n_d (B^-1)_jd + force[int64_t{owner} * 3 + lane] = + gradient[0] * grid.gradient_to_force[3 * lane] + + gradient[1] * grid.gradient_to_force[3 * lane + 1] + + gradient[2] * grid.gradient_to_force[3 * lane + 2]; } + // The stencil of the next node of the warp overwrites this one. + __syncwarp(); } } +// The potential prefactor * K(0) of a unit charge at itself, which the +// potential mesh holds and the potential of a node excludes, for every thread +// of the block. +__device__ double self_potential(int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor) { + __shared__ double self; + if (threadIdx.x == 0) { + self = prefactor * contact_value(kernel, coefficient, gaussians); + } + __syncthreads(); + return self; +} + // Per-frame sums over one slice of the owners: the potential of the charged // owners in the constrained channel, the number of charged owners and the -// nine components of ``moment``. +// nine components of sum_i F_i (x) r_i. constexpr int kOwnerSums = 11; __global__ void __launch_bounds__(kFrameThreads) - owner_slices(const double* __restrict__ potential, - const double* __restrict__ moment, + owner_slices(const float* __restrict__ potential, + const float* __restrict__ force, + const double* __restrict__ owner_coord, + const float* __restrict__ owner_charge, const bool* __restrict__ owner_mask, const int64_t* __restrict__ owner_frame_ptr, int channels, int channel, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, double* __restrict__ slice_sum) { + const double self = self_potential(kernel, coefficient, gaussians, prefactor); const NodeRange slice = slice_of_frame(owner_frame_ptr); double value[kOwnerSums] = {}; for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end; @@ -675,11 +813,15 @@ __global__ void __launch_bounds__(kFrameThreads) continue; } if (channel >= 0) { - value[0] += potential[owner * channels + channel]; + value[0] += double{potential[owner * channels + channel]} - + self * owner_charge[owner * channels + channel]; } value[1] += 1.0; - for (int k = 0; k < 9; ++k) { - value[2 + k] += moment[owner * 9 + k]; + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + value[2 + 3 * a + b] += + double{force[owner * 3 + a]} * owner_coord[owner * 3 + b]; + } } } write_slice_sum(value, slice_sum); @@ -702,27 +844,52 @@ __global__ void owner_totals(const double* __restrict__ slice_sum, total[blockIdx.x * kOwnerSums + threadIdx.x] = value; } -// Writes the outputs of every owner: the charge gradient w (phi - [c = ch] -// mean), the force, and the virial F (x) r plus the share of the rest of the -// virial of the frame over its charged nodes, then the owner sums of that -// virial per frame. -__global__ void finish_owners(const double* __restrict__ potential, - const double* __restrict__ force, - const double* __restrict__ moment, +// Writes the outputs of every owner: the atomic energy q . phi / 2 with phi = +// phi^rec - prefactor K(0) q, the charge gradient w (phi - [c = ch] mean), and +// the virial F (x) r plus the share of the rest of the virial W of the frame +// over its charged nodes. The first threads write the share of the virial of +// every frame held by the owners, sum_i F_i (x) r_i + n (W - sum_j F_j (x) +// r_j) / N over the n charged owners of the N charged nodes of the frame. +// ``owner_local`` holds the owner sums of this domain, ``owner_total`` those +// of every domain. +__global__ void finish_owners(const float* __restrict__ potential, + const float* __restrict__ force, + const double* __restrict__ owner_coord, + const float* __restrict__ owner_charge, const bool* __restrict__ owner_mask, const int64_t* __restrict__ owner_frame_ptr, const double* __restrict__ frame_virial, + const double* __restrict__ owner_local, const double* __restrict__ owner_total, const double* __restrict__ frame_stat, int frame_count, int owner_count, int channels, int channel, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + double* __restrict__ atom_energy, float* __restrict__ charge_gradient, - float* __restrict__ force_out, - float* __restrict__ atom_virial) { - for (int owner = blockIdx.x * blockDim.x + threadIdx.x; owner < owner_count; - owner += gridDim.x * blockDim.x) { + float* __restrict__ atom_virial, + float* __restrict__ virial) { + const double self = self_potential(kernel, coefficient, gaussians, prefactor); + const int start = blockIdx.x * blockDim.x + threadIdx.x; + const int threads = gridDim.x * blockDim.x; + for (int index = start; index < 9 * frame_count; index += threads) { + const int f = index / 9; + const int k = index - 9 * f; + const double count = frame_stat[f * kStatWidth + kStatCount]; + const double* local = owner_local + f * kOwnerSums; + const double rest = + count > 0.0 + ? (frame_virial[index] - owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + virial[index] = static_cast(local[2 + k] + local[1] * rest); + } + for (int owner = start; owner < owner_count; owner += threads) { const int f = frame_of(owner_frame_ptr, frame_count, owner); const bool charged = f >= 0 && owner_mask[owner]; // The charged nodes of the frame, of every domain. @@ -730,61 +897,33 @@ __global__ void finish_owners(const double* __restrict__ potential, charged ? frame_stat[f * kStatWidth + kStatCount] : 1.0; const double mean = charged && channel >= 0 ? owner_total[f * kOwnerSums] / count : 0.0; + double energy = 0.0; for (int c = 0; c < channels; ++c) { - const double phi = - charged ? potential[int64_t{owner} * channels + c] : 0.0; - charge_gradient[int64_t{owner} * channels + c] = + const int64_t entry = int64_t{owner} * channels + c; + const double q = owner_charge[entry]; + const double phi = charged ? potential[entry] - self * q : 0.0; + energy += q * phi; + charge_gradient[entry] = static_cast(c == channel ? phi - mean : phi); } + atom_energy[owner] = 0.5 * energy; for (int a = 0; a < 3; ++a) { - force_out[int64_t{owner} * 3 + a] = - static_cast(force[int64_t{owner} * 3 + a]); - } - for (int k = 0; k < 9; ++k) { - const double rest = charged ? (frame_virial[f * 9 + k] - - owner_total[f * kOwnerSums + 2 + k]) / - count - : 0.0; - atom_virial[int64_t{owner} * 9 + k] = - static_cast(moment[int64_t{owner} * 9 + k] + rest); - } - } -} - -// Sums, over one slice of the owners of one frame, their virial. -__global__ void __launch_bounds__(kFrameThreads) - virial_slices(const float* __restrict__ atom_virial, - const int64_t* __restrict__ owner_frame_ptr, - double* __restrict__ slice_sum) { - const NodeRange slice = slice_of_frame(owner_frame_ptr); - double value[9] = {}; - for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end; - owner += kFrameThreads) { - for (int k = 0; k < 9; ++k) { - value[k] += atom_virial[owner * 9 + k]; + for (int b = 0; b < 3; ++b) { + const int k = 3 * a + b; + const double moment = double{force[int64_t{owner} * 3 + a]} * + owner_coord[int64_t{owner} * 3 + b]; + const double rest = charged ? (frame_virial[f * 9 + k] - + owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + atom_virial[int64_t{owner} * 9 + k] = static_cast(moment + rest); + } } } - write_slice_sum<9>(value, slice_sum); -} - -__global__ void virial_total(const double* __restrict__ slice_sum, - int slice_count, - float* __restrict__ virial) { - if (threadIdx.x >= 9) { - return; - } - double value = 0.0; - for (int slice = 0; slice < slice_count; ++slice) { - value += - slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * 9 + - threadIdx.x]; - } - virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value); } -// The weights M_p(1), ..., M_p(p - 1) of the assignment factor. -std::vector integer_bspline_weights(int order) { - std::vector weight(order - 1); +IntegerWeights integer_bspline_weights(int order) { + IntegerWeights weight{}; for (int l = 0; l < order - 1; ++l) { const double x = l + 1.0; double value = 0.0; @@ -800,11 +939,101 @@ std::vector integer_bspline_weights(int order) { for (int k = 2; k < order; ++k) { factorial *= k; } - weight[l] = value / factorial; + weight.value[l] = value / factorial; } return weight; } +void check_cufft(cufftResult result, const char* operation) { + TORCH_CHECK(result == CUFFT_SUCCESS, "latent_charge_sum: cuFFT ", operation, + " failed with error ", static_cast(result)); +} + +// The real-to-complex and complex-to-real transforms of the meshes of one +// shape, batched over the channels. The plans allocate no work area: every +// transform runs on a work area of the workspace of its call, set with the +// stream under ``mutex`` just before the transform is enqueued, so that +// concurrent calls never share one. +struct MeshPlans { + cufftHandle forward = 0; + cufftHandle inverse = 0; + size_t work_bytes = 0; + c10::DeviceIndex device = 0; + std::mutex mutex; + + MeshPlans() = default; + MeshPlans(const MeshPlans&) = delete; + MeshPlans& operator=(const MeshPlans&) = delete; + // A plan leaves the cache only past kPlanCacheSize shapes; its transforms + // may still be in flight, so the device finishes them first. The + // destructor reports no error. + ~MeshPlans() { + int current = 0; + cudaGetDevice(¤t); + cudaSetDevice(device); + cudaDeviceSynchronize(); + if (forward != 0) { + cufftDestroy(forward); + } + if (inverse != 0) { + cufftDestroy(inverse); + } + cudaSetDevice(current); + } +}; + +// The plans of the meshes of ``mesh`` with ``channels`` channels on the +// current device, kept across calls; beyond kPlanCacheSize shapes, the least +// recently used are released. The cache itself lives until the process ends +// and is never destroyed, since its plans would otherwise be released after +// the CUDA runtime shuts down. +constexpr size_t kPlanCacheSize = 16; + +std::shared_ptr mesh_plans(const FrameMesh& mesh, int channels) { + using Key = std::array; + using Cache = std::list>>; + static std::mutex cache_mutex; + static Cache& cache = *new Cache(); + const Key key = {c10::cuda::current_device(), mesh.n[0], mesh.n[1], mesh.n[2], + channels}; + // A plan that leaves the cache is released after the lock, since its + // release waits for the device. + std::shared_ptr evicted; + const std::lock_guard lock(cache_mutex); + for (auto entry = cache.begin(); entry != cache.end(); ++entry) { + if (entry->first == key) { + cache.splice(cache.begin(), cache, entry); + return entry->second; + } + } + auto plans = std::make_shared(); + plans->device = c10::cuda::current_device(); + long long shape[3] = {mesh.n[0], mesh.n[1], mesh.n[2]}; + const long long points = shape[0] * shape[1] * shape[2]; + const long long modes = shape[0] * shape[1] * (shape[2] / 2 + 1); + size_t forward_bytes = 0; + size_t inverse_bytes = 0; + check_cufft(cufftCreate(&plans->forward), "plan creation"); + check_cufft(cufftSetAutoAllocation(plans->forward, 0), "plan setup"); + check_cufft( + cufftMakePlanMany64(plans->forward, 3, shape, nullptr, 1, points, nullptr, + 1, modes, CUFFT_R2C, channels, &forward_bytes), + "plan creation"); + check_cufft(cufftCreate(&plans->inverse), "plan creation"); + check_cufft(cufftSetAutoAllocation(plans->inverse, 0), "plan setup"); + check_cufft( + cufftMakePlanMany64(plans->inverse, 3, shape, nullptr, 1, modes, nullptr, + 1, points, CUFFT_C2R, channels, &inverse_bytes), + "plan creation"); + plans->work_bytes = std::max(forward_bytes, inverse_bytes); + cache.emplace_front(key, plans); + if (cache.size() > kPlanCacheSize) { + evicted = std::move(cache.back().second); + cache.pop_back(); + } + return plans; +} + int blocks_for(int64_t work) { return static_cast(std::max( 1, std::min((work + kThreads - 1) / kThreads, 65535))); @@ -850,6 +1079,15 @@ void dispatch_order(int64_t order, const Launch& launch) { } } +// The value of ``evaluate`` at the B-spline order as a compile-time constant. +template +auto dispatch_order_value(int64_t order, const Evaluate& evaluate) { + decltype(evaluate(std::integral_constant{})) value{}; + dispatch_order( + order, [&](auto order_constant) { value = evaluate(order_constant); }); + return value; +} + } // namespace namespace deepmd { @@ -900,12 +1138,23 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, // === Step 1. The meshes of the frames, sized on the host === std::vector frames(frame_count); + std::vector> plans(frame_count); const double* cells = box.data_ptr(); int64_t mesh_total = 0; + int64_t spectrum_size = 0; + int64_t factor_size = 0; + size_t work_bytes = 0; for (int f = 0; f < frame_count; ++f) { frames[f] = frame_mesh(cells + 9 * f, spacing); frames[f].offset = mesh_total; mesh_total += channels * mesh_points(frames[f]); + const int half = frames[f].n[2] / 2 + 1; + spectrum_size = std::max(spectrum_size, int64_t{channels} * frames[f].n[0] * + frames[f].n[1] * half); + factor_size = + std::max(factor_size, int64_t{frames[f].n[0]} + frames[f].n[1] + half); + plans[f] = mesh_plans(frames[f], channels); + work_bytes = std::max(work_bytes, plans[f]->work_bytes); } const auto byte_options = torch::TensorOptions().dtype(torch::kUInt8); auto host_frames = @@ -913,23 +1162,15 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, byte_options.pinned_memory(true)); std::memcpy(host_frames.data_ptr(), frames.data(), sizeof(FrameMesh) * frame_count); - const auto device_frames = - host_frames.to(coord.device(), /*non_blocking=*/true); - const auto* frames_ptr = - reinterpret_cast(device_frames.data_ptr()); - const std::vector weight_host = integer_bspline_weights(order); - auto host_weight = torch::from_blob(const_cast(weight_host.data()), - {order - 1}, torch::kFloat64) - .pin_memory(); - const auto integer_weight = - host_weight.to(coord.device(), /*non_blocking=*/true); + const IntegerWeights integer_weight = + integer_bspline_weights(static_cast(order)); + // A few blocks per multiprocessor fill the device; every thread of the + // influence function then walks several modes, which keeps the block sums + // few. + const int mode_blocks = + 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; - // === Step 2. The constraint and the charged nodes of every frame === - const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); - auto frame_ptr = torch::empty({frame_count + 1}, counts.options()); - frame_offsets<<<1, kOffsetThreads, 0, stream>>>( - counts.data_ptr(), frame_count, frame_ptr.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + // === Step 2. The buffers of the call === const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; const auto slices_for = [](int64_t nodes) { return static_cast(std::min( @@ -937,176 +1178,197 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, std::max(1, (nodes + kSliceNodes - 1) / kSliceNodes))); }; const int slice_count = slices_for(average_nodes); - auto slice_charge = torch::empty( - {static_cast(frame_count) * slice_count * 2}, double_options); - auto frame_stat = torch::empty({frame_count * kStatWidth}, double_options); + const int owner_slice_count = slices_for( + (int64_t{owner_count} + frame_count - 1) / std::max(frame_count, 1)); + const int64_t stencil_bytes = dispatch_order_value(order, [](auto constant) { + return static_cast(sizeof(AxisStencil)); + }); + Workspace workspace; + const int64_t frames_at = + workspace.reserve(sizeof(FrameMesh) * frame_count); + const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); + const int64_t slice_charge_at = + workspace.reserve(int64_t{frame_count} * slice_count * 2); + const int64_t frame_stat_at = + workspace.reserve(int64_t{frame_count} * kStatWidth); + const int64_t stencils_at = + workspace.reserve(3 * node_count * stencil_bytes); + const int64_t node_charge_at = + workspace.reserve(node_count * channels); + const int64_t fixed_mesh_at = workspace.reserve(mesh_total); + // The potential mesh of a frame overwrites its charge mesh, which its + // forward transform has consumed. + const int64_t mesh_at = workspace.reserve(mesh_total); + const int64_t spectrum_at = workspace.reserve(spectrum_size); + const int64_t work_at = workspace.reserve(work_bytes); + const int64_t factor_at = workspace.reserve(factor_size); + const int64_t block_sum_at = + workspace.reserve(int64_t{mode_blocks} * kModeSums); + const int64_t frame_virial_at = + workspace.reserve(int64_t{frame_count} * 9); + const int64_t potential_at = + workspace.reserve(int64_t{owner_count} * channels); + const int64_t owner_slice_at = workspace.reserve( + int64_t{frame_count} * owner_slice_count * kOwnerSums); + const int64_t owner_local_at = + workspace.reserve(int64_t{frame_count} * kOwnerSums); + workspace.allocate(coord.device()); + auto* frames_ptr = workspace.at(frames_at); + auto* frame_ptr = workspace.at(frame_ptr_at); + auto* slice_charge = workspace.at(slice_charge_at); + auto* frame_stat = workspace.at(frame_stat_at); + auto* node_charge = workspace.at(node_charge_at); + auto* fixed_mesh = workspace.at(fixed_mesh_at); + auto* charge_mesh = workspace.at(mesh_at); + float* potential_mesh = charge_mesh; + auto* spectrum = workspace.at(spectrum_at); + void* work = workspace.at(work_at); + auto* factor = workspace.at(factor_at); + auto* block_sum = workspace.at(block_sum_at); + auto* frame_virial = workspace.at(frame_virial_at); + auto* potential = workspace.at(potential_at); + auto* owner_slice = workspace.at(owner_slice_at); + auto* owner_local = workspace.at(owner_local_at); + // The copy from the pinned staging buffer records its completion, before + // which the allocator does not hand the buffer out again. + torch::from_blob(frames_ptr, {host_frames.numel()}, + byte_options.device(coord.device())) + .copy_(host_frames, /*non_blocking=*/true); + C10_CUDA_CHECK( + cudaMemsetAsync(fixed_mesh, 0, sizeof(int64_t) * mesh_total, stream)); + + // === Step 3. The constraint, and the stencils and charges of the nodes === + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(), + frame_count, frame_ptr); + C10_CUDA_KERNEL_LAUNCH_CHECK(); charge_slices<<>>( - charge.data_ptr(), node_mask.data_ptr(), - frame_ptr.data_ptr(), channels, - static_cast(std::max(channel, 0)), - slice_charge.data_ptr()); + charge.data_ptr(), node_mask.data_ptr(), frame_ptr, channels, + static_cast(std::max(channel, 0)), slice_charge); C10_CUDA_KERNEL_LAUNCH_CHECK(); charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, kFrameThreads, 0, stream>>>( - slice_charge.data_ptr(), slice_count, frame_count, - total_charge.data_ptr(), frame_stat.data_ptr()); + slice_charge, slice_count, frame_count, total_charge.data_ptr(), + frame_stat); C10_CUDA_KERNEL_LAUNCH_CHECK(); - - // === Step 3. The stencils of the nodes and the charge meshes === - auto mesh = torch::zeros({mesh_total}, double_options.dtype(torch::kInt64)); - auto node_charge = torch::empty({node_count, channels}, float_options); - torch::Tensor stencils; dispatch_order(order, [&](auto order_constant) { constexpr int kOrder = decltype(order_constant)::value; - stencils = torch::empty( - {3 * node_count * static_cast(sizeof(AxisStencil))}, - byte_options.device(coord.device())); - auto* node_stencil = - reinterpret_cast*>(stencils.data_ptr()); + auto* node_stencil = workspace.at>(stencils_at); node_stencils<<>>( - coord.data_ptr(), node_mask.data_ptr(), - frame_ptr.data_ptr(), frames_ptr, frame_count, node_count, - node_stencil); + coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frames_ptr, frame_stat, + frame_count, node_count, channels, static_cast(channel), + node_stencil, node_charge); C10_CUDA_KERNEL_LAUNCH_CHECK(); + // === Step 4. The charge meshes === spread_charges <<>>( - charge.data_ptr(), node_mask.data_ptr(), - frame_ptr.data_ptr(), frames_ptr, node_stencil, - frame_stat.data_ptr(), frame_count, node_count, channels, - static_cast(channel), node_charge.data_ptr(), - reinterpret_cast(mesh.data_ptr())); + node_mask.data_ptr(), frame_ptr, frames_ptr, node_stencil, + node_charge, frame_count, node_count, channels, fixed_mesh); C10_CUDA_KERNEL_LAUNCH_CHECK(); }); - fixed_to_double<<>>( - mesh_total, - reinterpret_cast(mesh.data_ptr())); + fixed_to_float<<>>( + mesh_total, fixed_mesh, charge_mesh); C10_CUDA_KERNEL_LAUNCH_CHECK(); - const auto charge_mesh = mesh.view(torch::kFloat64); - // === Step 4. Influence function and potential meshes, frame by frame === - auto potential_mesh = torch::empty({mesh_total}, double_options); - auto frame_virial = torch::zeros({frame_count, 9}, double_options); + // === Step 5. Influence function and potential meshes, frame by frame === for (int f = 0; f < frame_count; ++f) { const FrameMesh& grid = frames[f]; - const int64_t points = mesh_points(grid); - const std::vector shape = {channels, grid.n[0], grid.n[1], - grid.n[2]}; - auto spectrum = at::fft_rfftn( - charge_mesh.narrow(0, grid.offset, channels * points).view(shape), - c10::nullopt, std::vector{1, 2, 3}); + MeshPlans& plan = *plans[f]; + { + const std::lock_guard lock(plan.mutex); + check_cufft(cufftSetStream(plan.forward, stream), "stream setup"); + check_cufft(cufftSetWorkArea(plan.forward, work), "work area setup"); + check_cufft(cufftExecR2C(plan.forward, charge_mesh + grid.offset, + reinterpret_cast(spectrum)), + "transform"); + } const int half = grid.n[2] / 2 + 1; - auto factor = torch::empty({grid.n[0] + grid.n[1] + half}, double_options); - assignment_factor<<>>( - grid, static_cast(order), integer_weight.data_ptr(), - factor.data_ptr()); + const int factor_count = grid.n[0] + grid.n[1] + half; + assignment_factor<<>>( + grid, static_cast(order), integer_weight, factor); C10_CUDA_KERNEL_LAUNCH_CHECK(); - const int64_t modes = int64_t{grid.n[0]} * grid.n[1] * half; - // A few blocks per multiprocessor fill the device; every thread then - // walks several modes, which keeps the block sums few. - const int mode_blocks = std::min( - blocks_for(modes), - 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount); - auto block_sum = - torch::empty({int64_t{mode_blocks} * kModeSums}, double_options); apply_influence<<>>( grid, channels, static_cast(kernel), - coefficient.data_ptr(), gaussians, prefactor, - factor.data_ptr(), - reinterpret_cast(spectrum.data_ptr>()), - block_sum.data_ptr()); + coefficient.data_ptr(), gaussians, prefactor, factor, spectrum, + block_sum); C10_CUDA_KERNEL_LAUNCH_CHECK(); - mode_virial<<<1, kThreads, 0, stream>>>(block_sum.data_ptr(), - mode_blocks, - frame_virial[f].data_ptr()); + mode_virial<<<1, kThreads, 0, stream>>>(block_sum, mode_blocks, + frame_virial + 9 * f); C10_CUDA_KERNEL_LAUNCH_CHECK(); - auto frame_potential = - potential_mesh.narrow(0, grid.offset, channels * points).view(shape); - at::fft_irfftn_out(frame_potential, spectrum, - std::vector{grid.n[0], grid.n[1], grid.n[2]}, - std::vector{1, 2, 3}, "forward"); + { + const std::lock_guard lock(plan.mutex); + check_cufft(cufftSetStream(plan.inverse, stream), "stream setup"); + check_cufft(cufftSetWorkArea(plan.inverse, work), "work area setup"); + check_cufft( + cufftExecC2R(plan.inverse, reinterpret_cast(spectrum), + potential_mesh + grid.offset), + "transform"); + } } - // === Step 5. Potential, force and virial of the owners === - // The stencils read the potential meshes in float32, the precision of the - // outputs. - const auto potential_mesh_float = potential_mesh.to(torch::kFloat32); - auto potential = torch::empty({owner_count, channels}, double_options); - auto force = torch::empty({owner_count, 3}, double_options); - auto moment = torch::empty({owner_count, 9}, double_options); - auto atom_energy = torch::empty({owner_count, 1}, double_options); + // === Step 6. Potential, force and virial of the owners === + auto force = torch::empty({owner_count, 3}, float_options); if (owner_count > 0) { dispatch_order(order, [&](auto order_constant) { constexpr int kOrder = decltype(order_constant)::value; interpolate_owners<<>>( - coord.data_ptr(), node_charge.data_ptr(), - node_mask.data_ptr(), frame_ptr.data_ptr(), frames_ptr, - reinterpret_cast*>(stencils.data_ptr()), - potential_mesh_float.data_ptr(), - coefficient.data_ptr(), frame_count, owner_begin, owner_count, - channels, static_cast(kernel), gaussians, prefactor, - potential.data_ptr(), force.data_ptr(), - moment.data_ptr(), atom_energy.data_ptr()); + node_mask.data_ptr(), frame_ptr, frames_ptr, + workspace.at>(stencils_at), node_charge, + potential_mesh, frame_count, owner_begin, owner_count, channels, + potential, force.data_ptr()); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } // The frames of the owners, on the axis of the owners: the frames // themselves when every node is an owner, else the one frame they share. - torch::Tensor owner_frame_ptr = frame_ptr; + torch::Tensor owner_frame_tensor; + const int64_t* owner_frame_ptr = frame_ptr; if (!every_node) { - owner_frame_ptr = - torch::full({2}, int64_t{owner_count}, frame_ptr.options()); - owner_frame_ptr.select(0, 0).zero_(); + owner_frame_tensor = + torch::full({2}, int64_t{owner_count}, counts.options()); + owner_frame_tensor.select(0, 0).zero_(); + owner_frame_ptr = owner_frame_tensor.data_ptr(); } const bool* owner_mask = node_mask.data_ptr() + owner_begin; - const int owner_slice_count = slices_for( - (int64_t{owner_count} + frame_count - 1) / std::max(frame_count, 1)); - auto owner_slice = torch::empty( - {int64_t{frame_count} * owner_slice_count * kOwnerSums}, double_options); + const double* owner_coord = coord.data_ptr() + 3 * owner_begin; + const float* owner_charge = node_charge + int64_t{owner_begin} * channels; owner_slices<<>>( - potential.data_ptr(), moment.data_ptr(), owner_mask, - owner_frame_ptr.data_ptr(), channels, static_cast(channel), - owner_slice.data_ptr()); + stream>>>(potential, force.data_ptr(), owner_coord, + owner_charge, owner_mask, owner_frame_ptr, channels, + static_cast(channel), static_cast(kernel), + coefficient.data_ptr(), gaussians, prefactor, + owner_slice); C10_CUDA_KERNEL_LAUNCH_CHECK(); - auto owner_total = - torch::empty({int64_t{frame_count} * kOwnerSums}, double_options); - owner_totals<<>>(owner_slice.data_ptr(), - owner_slice_count, - owner_total.data_ptr()); + owner_totals<<>>(owner_slice, owner_slice_count, + owner_local); C10_CUDA_KERNEL_LAUNCH_CHECK(); + // The mean potential of the constraint and the virial share run over the + // charged nodes of every domain. + torch::Tensor owner_sum; + const double* owner_total = owner_local; if (communicator != nullptr) { - // The mean potential of the constraint and the virial share run over the - // charged nodes of every domain. - owner_total.copy_(sum_over_domains(*communicator, owner_total)); + owner_sum = sum_over_domains( + *communicator, torch::from_blob(owner_local, {frame_count * kOwnerSums}, + double_options)); + owner_total = owner_sum.data_ptr(); } + auto atom_energy = torch::empty({owner_count, 1}, double_options); auto charge_gradient = torch::empty({owner_count, channels}, float_options); - auto force_out = torch::empty({owner_count, 3}, float_options); auto atom_virial = torch::empty({owner_count, 9}, float_options); - if (owner_count > 0) { - finish_owners<<>>( - potential.data_ptr(), force.data_ptr(), - moment.data_ptr(), owner_mask, - owner_frame_ptr.data_ptr(), frame_virial.data_ptr(), - owner_total.data_ptr(), frame_stat.data_ptr(), - frame_count, owner_count, channels, static_cast(channel), - charge_gradient.data_ptr(), force_out.data_ptr(), - atom_virial.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - } - auto virial_slice = torch::empty( - {int64_t{frame_count} * owner_slice_count * 9}, double_options); - virial_slices<<>>(atom_virial.data_ptr(), - owner_frame_ptr.data_ptr(), - virial_slice.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); auto virial = torch::empty({frame_count, 9}, float_options); - virial_total<<>>(virial_slice.data_ptr(), - owner_slice_count, - virial.data_ptr()); + finish_owners<<(owner_count, 9 * frame_count)), + kThreads, 0, stream>>>( + potential, force.data_ptr(), owner_coord, owner_charge, owner_mask, + owner_frame_ptr, frame_virial, owner_local, owner_total, frame_stat, + frame_count, owner_count, channels, static_cast(channel), + static_cast(kernel), coefficient.data_ptr(), gaussians, + prefactor, atom_energy.data_ptr(), + charge_gradient.data_ptr(), atom_virial.data_ptr(), + virial.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {atom_energy, charge_gradient, force_out, atom_virial, virial}; + return {atom_energy, charge_gradient, force, atom_virial, virial}; } } // namespace long_range diff --git a/source/tests/pt_expt/kernels/test_long_range_sum.py b/source/tests/pt_expt/kernels/test_long_range_sum.py index 18552e553f..b61c7390c9 100644 --- a/source/tests/pt_expt/kernels/test_long_range_sum.py +++ b/source/tests/pt_expt/kernels/test_long_range_sum.py @@ -11,8 +11,9 @@ cancel, and the energy, force, charge gradient and virial follow from one another. -For periodic frames the CUDA kernels evaluate the mesh sum in float64, and are -held to the mesh evaluation of the training path, differentiated by autograd. +For periodic frames the CUDA kernels evaluate the mesh sum in float32 on +B-spline weights formed in float64, and are held to the float64 mesh +evaluation of the training path, differentiated by autograd. """ from __future__ import ( @@ -30,6 +31,9 @@ from deepmd.dpmodel.long_range import ( COULOMB_CONSTANT, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, +) from deepmd.pt_expt.kernels.long_range import ( _reference, kernel_operator_arguments, @@ -322,17 +326,17 @@ def test_mesh_virial_is_the_strain_derivative(kernel_name: str) -> None: Straining the positions and the cell together by ``I + eps`` changes the energy by ``-W : eps``; the atomic virial of the charged nodes sums to W. - The mesh interpolates in float32, whose rounding the central difference - divides by its step, so that the step stays large against it. + The central difference takes the float64 mesh energy, whose rounding stays + far below that of the float32 outputs of the operator. """ system = _system([60], 2, 0, 13, 0.0) coord, charge, node_mask, n_node, total_charge = system box = _cells(1, 10.0) kernel = KERNELS[kernel_name]() - atom_energy, _, _, atom_virial, virial = latent_charge_sum( + _, _, _, atom_virial, virial = latent_charge_sum( *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.5 ) - h = 1e-3 + h = 1e-5 finite = np.zeros((3, 3)) for a in range(3): for b in range(3): @@ -348,28 +352,50 @@ def test_mesh_virial_is_the_strain_derivative(kernel_name: str) -> None: total_charge, ) cell = box.reshape(3, 3) @ strain.T + # Both strained cells keep the mesh of the cell. + assert mesh_dims_of_cells(cell.reshape(1, 9).numpy(), 0.5) == ( + mesh_dims_of_cells(box.numpy(), 0.5) + ) energies.append( - latent_charge_sum( - *strained, - kernel, - 0, - COULOMB_CONSTANT, - box=cell.reshape(1, 9), - spacing=0.5, - )[0].sum() + _reference_of(strained, kernel, 0, cell.reshape(1, 9), 8, 0.5)[ + 0 + ].sum() ) finite[a, b] = -(energies[0] - energies[1]).item() / (2 * h) scale = np.abs(finite).max() np.testing.assert_allclose( - virial[0].double().cpu().numpy().reshape(3, 3), finite, atol=2e-4 * scale + virial[0].double().cpu().numpy().reshape(3, 3), finite, atol=2e-6 * scale ) np.testing.assert_allclose( atom_virial.double().sum(0).cpu().numpy().reshape(3, 3), finite, - atol=2e-4 * scale, + atol=2e-6 * scale, ) +@pytest.mark.parametrize("order", [2, 3, 5, 7, 9, 10, 11, 12]) +def test_mesh_orders(order: int) -> None: + """Every B-spline order agrees with the mesh evaluation of training. + + The atomic energies of these small cells are a fraction of the self + energy, to a few 1e-7 of which the float32 mesh resolves them. + """ + system = _system([40, 23], 2, 3, 5, -5.0) + box = _cells(2, 9.0) + actual, expected = _evaluate(system, KERNELS["sog"](), 1, box, order, 0.7) + _assert_close(actual, expected, 1e-5) + + +def test_mesh_stencil_wider_than_the_axis() -> None: + """A stencil wider than an axis of the mesh wraps around it more than once.""" + system = _system([12], 1, 0, 9, 0.0) + box = _cells(1, 3.0) + # Four points per axis, against eight of the stencil. + assert mesh_dims_of_cells(box.numpy(), 1.0) == (4, 4, 4) + actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 1.0) + _assert_close(actual, expected, 1e-6) + + @pytest.mark.parametrize("periodic", [False, True]) def test_frame_without_charged_nodes(periodic: bool) -> None: """A frame whose nodes carry no charge contributes zeros, the others are unaffected.""" From c75fde9b5c681c6352c275294ccf909d3259688a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 05:05:33 +0800 Subject: [PATCH 28/30] perf(long-range): spread the kernel-parameter gradient of the mesh energy The gradient of a parameter of the kernel transform sums over every mode of the mesh, which the compiled graph reduced in a single block: two such reductions took 70% of the device time of a training step. The transform now reads its parameters through one copy per row of the modes, so that the reduction runs along the rows first. A compiled training step of DPA4C with a periodic LES module in the pme mode on the water example takes 19.6 ms instead of 32.5 ms (10.6 ms without long range). --- deepmd/dpmodel/long_range/kernel.py | 32 ++++++++++++++++--------- deepmd/dpmodel/long_range/reciprocal.py | 5 +++- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py index 67065a6eec..911dfba5dd 100644 --- a/deepmd/dpmodel/long_range/kernel.py +++ b/deepmd/dpmodel/long_range/kernel.py @@ -115,13 +115,13 @@ def kspace(self, k_sq: Array) -> Array: Parameters ---------- k_sq - Squared wave numbers with shape ``(nk,)``, in Å⁻²; every entry - must be positive. + Squared wave numbers with shape ``(nk,)`` or ``(rows, nk)``, in + Å⁻²; every entry must be positive. Returns ------- Array - The transformed kernel with shape ``(nk,)``, in ų. + The transformed kernel with the shape of ``k_sq``, in ų. """ raise NotImplementedError @@ -162,6 +162,18 @@ def _parameter(self, name: str, like: Array) -> Array: device=array_api_compat.device(like), ) + @staticmethod + def _row_values(value: Array, k_sq: Array) -> Array: + """A scalar broadcast to shape ``(..., 1)`` over the leading axes of ``k_sq``. + + The gradient of a parameter of the transform sums over every wave + number. Broadcast over the rows, it sums along every row first, a + reduction a compiled graph spreads over the rows, and then over the + rows, rather than over every wave number in one block. + """ + xp = array_api_compat.array_namespace(k_sq) + return xp.broadcast_to(value, (*k_sq.shape[:-1], 1)) + def serialize(self) -> dict: return { "@class": "InteractionKernel", @@ -270,7 +282,7 @@ def kspace(self, k_sq: Array) -> Array: # erf(alpha r)/r is the interaction of two Gaussian charge clouds of # width 1/(sqrt(2) alpha); its transform is the damped Coulomb # 4 pi exp(-k^2 / 4 alpha^2) / k^2. - alpha = xp.exp(self._parameter("log_alpha", k_sq)) + alpha = self._row_values(xp.exp(self._parameter("log_alpha", k_sq)), k_sq) return 4.0 * np.pi * xp.exp(-k_sq / (4.0 * alpha * alpha)) / k_sq def contact_value(self, like: Array) -> Array: @@ -382,13 +394,11 @@ def kspace(self, k_sq: Array) -> Array: # whose ln-sigma trapezoid sum approaches 4 pi / k^2. The mode loop is # unrolled: a (nk, M) intermediate breaks inductor's double backward of # the k-space virial (NaN), while per-mode (nk,) terms compile fine. - terms = [ - amplitude[m] - * (2.0 * np.pi) ** 1.5 - * width[m] ** 3 - * xp.exp(-0.5 * width[m] ** 2 * k_sq) - for m in range(width.shape[0]) - ] + terms = [] + for m in range(width.shape[0]): + weight = self._row_values(amplitude[m] * width[m] ** 3, k_sq) + rate = self._row_values(0.5 * width[m] ** 2, k_sq) + terms.append((2.0 * np.pi) ** 1.5 * weight * xp.exp(-rate * k_sq)) return sum(terms[1:], terms[0]) def contact_value(self, like: Array) -> Array: diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py index 8643ed3473..0fd3439a43 100644 --- a/deepmd/dpmodel/long_range/reciprocal.py +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -352,7 +352,10 @@ def mesh_energy( keep = keep[None, ...] & mesh.mask[:, None, None, None] safe_k_sq = xp.where(keep, k_sq, xp.ones_like(k_sq)) shape = safe_k_sq.shape - transform = xp.reshape(kernel.kspace(xp.reshape(safe_k_sq, (-1,))), shape) + # One row per frame and first mode index, along which the gradients of + # the kernel parameters sum first. + rows = xp.reshape(safe_k_sq, (nf * n1, n2 * (n3 // 2 + 1))) + transform = xp.reshape(kernel.kspace(rows), shape) f1, f2, f3 = _assignment_factor(m_half, (n1, n2, n3), order) factor = f1[:, None, None] * f2[None, :, None] * f3[None, None, :] influence = (COULOMB_CONSTANT * transform * factor[None, ...]) / volume[ From 2b729e4fa103cb628c44fa681fd54afea4b593ee Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 15:00:50 +0800 Subject: [PATCH 29/30] feat(long-range): per-frame blocked k grid, mesh evaluation by default The explicit wave numbers of a periodic batch were one ragged list over the frames, each paired with the nodes of its frame through gathers and segment sums whose cost and kernel count grew with the batch. The k grid now holds the wave numbers of every frame in blocks of K_BLOCK points, padded with masked guard points, beside the nodes of every frame padded with -1, so that the structure factors and the potentials are batched products of one frame's wave numbers with its own nodes, and every sum over the wave numbers runs over the points of a block first. The ragged and fixed grids share the layout; the fixed grid masks the points beyond the cutoff of a frame. A periodic module evaluates the reciprocal sum on the mesh unless `k_grid_mode` says otherwise, as frozen models already do; a periodic record that names no mode is one of the explicit wave numbers. The pair topology of a periodic batch carries four guard pairs: with two, a compiled graph identified the pair axis with the two rows of the pair index and split at the reductions over the pairs. For the same reason the trace-time samples of the k grid give its slots, blocks and points lengths of their own, of at least four. Compiled training of a small DPA4C on the water example (batch of 4, RTX PRO 6000) takes 37.5 ms per step with the ragged grid instead of 62 ms, 46.5 ms with the fixed grid instead of 127 ms, and 19.0 ms on the mesh, against 9.6 ms without long range. --- deepmd/dpmodel/long_range/latent_charge.py | 28 ++-- deepmd/dpmodel/long_range/reciprocal.py | 84 ++++++---- .../dpmodel/utils/neighbor_graph/__init__.py | 2 + .../utils/neighbor_graph/pair_graph.py | 14 +- .../utils/neighbor_graph/reciprocal_graph.py | 150 +++++++++++------- .../dpmodel/utils/neighbor_graph/topology.py | 7 +- deepmd/pt_expt/utils/serialization.py | 38 +++-- deepmd/utils/argcheck.py | 27 ++-- doc/model/long-range.md | 18 +-- source/api_cc/src/longRangePTExpt.h | 3 +- .../tests/common/dpmodel/test_long_range.py | 130 +++++++++++---- source/tests/pt_expt/model/test_long_range.py | 30 ++++ 12 files changed, 356 insertions(+), 175 deletions(-) diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py index 0480c92f06..4a2d16b2f4 100644 --- a/deepmd/dpmodel/long_range/latent_charge.py +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -134,16 +134,17 @@ class LatentChargeLongRange(NativeOP, BaseLongRange): wave-number cutoff :math:`k_{max}` of a periodic batch. k_grid_mode How training evaluates the reciprocal sum of a periodic batch: + ``"pme"``, the default, by the smooth particle-mesh Ewald method on a + mesh of spacing ``pme_spacing`` sized by the cells of the batch; ``"ragged"`` over the wave numbers of every frame within :math:`k_{max}`, enumerated eagerly, so that the k axis is a dynamic - dimension; ``"fixed"`` over one integer grid shared by the frames, + dimension; or ``"fixed"`` over one integer grid shared by the frames, sized by the largest cell of the data statistics or by - ``k_grid_bounds``, with the points a frame does not reach masked; or - ``"pme"`` by the smooth particle-mesh Ewald method on a mesh of - spacing ``pme_spacing`` sized by the cells of the batch. A frame - whose cell needs indices beyond the fixed grid is truncated to it, - with a warning. A frozen model evaluates every periodic frame on the - mesh, whatever the mode. + ``k_grid_bounds``, with the points a frame does not reach masked. A + frame whose cell needs indices beyond the fixed grid is truncated to + it, with a warning. A frozen model evaluates every periodic frame on + the mesh, whatever the mode. A module that serves no periodic system + evaluates no reciprocal sum and takes no mode. k_grid_bounds The per-axis integer bounds :math:`n_d^{max}` of the fixed grid, in place of the data statistics; only meaningful with @@ -171,7 +172,7 @@ def __init__( charge_constraint: dict | None = None, periodic: bool = False, kspace_tolerance: float = 1e-8, - k_grid_mode: str = "ragged", + k_grid_mode: str | None = None, k_grid_bounds: list[int] | None = None, decouple_k_grid: bool = False, pme_spacing: float | None = None, @@ -199,6 +200,8 @@ def __init__( # The frame charge/spin condition a compressed model is specialized # to; see :meth:`fold_charge_state`. self.charge_state = None + if k_grid_mode is None: + k_grid_mode = "pme" if self.periodic else "ragged" self.k_grid_mode = str(k_grid_mode) if self.k_grid_mode not in ("ragged", "fixed", "pme"): raise ValueError( @@ -344,7 +347,7 @@ def from_config( config.get("charge_constraint"), periodic=config.get("periodic", False), kspace_tolerance=config.get("kspace_tolerance", 1e-8), - k_grid_mode=config.get("k_grid_mode", "ragged"), + k_grid_mode=config.get("k_grid_mode"), k_grid_bounds=config.get("k_grid_bounds"), decouple_k_grid=config.get("decouple_k_grid", False), pme_spacing=config.get("pme_spacing"), @@ -504,9 +507,7 @@ def call_graph( self.pme_order, ) else: - reciprocal = kspace_energy( - charge_64, topologies[K_GRID], frame_id, self.kernel - ) + reciprocal = kspace_energy(charge_64, topologies[K_GRID], self.kernel) energy = energy + xp.astype(reciprocal, energy.dtype) return xp.reshape(energy, (n, 1)) @@ -686,7 +687,6 @@ def serialize(self) -> dict: if self.periodic: data["periodic"] = True data["kspace_tolerance"] = self.kspace_tolerance - if self.k_grid_mode != "ragged": data["k_grid_mode"] = self.k_grid_mode if self.k_grid_mode == "fixed" and self._k_grid_sized(): data["k_grid_bounds"] = to_numpy_array(self.k_grid_bounds).tolist() @@ -711,6 +711,8 @@ def deserialize(cls, data: dict) -> "LatentChargeLongRange": data["charge_constraint"], periodic=data.get("periodic", False), kspace_tolerance=data.get("kspace_tolerance", 1e-8), + # A periodic record without a mode is one of the explicit wave + # numbers. k_grid_mode=data.get("k_grid_mode", "ragged"), k_grid_bounds=data.get("k_grid_bounds"), decouple_k_grid=data.get("decouple_k_grid", False), diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py index 0fd3439a43..9965958b33 100644 --- a/deepmd/dpmodel/long_range/reciprocal.py +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -43,7 +43,6 @@ from deepmd.dpmodel.utils.neighbor_graph import ( KGrid, MeshGrid, - segment_sum, ) from .kernel import ( @@ -68,7 +67,6 @@ def _self_potential( def kspace_energy( charge: Array, k_grid: KGrid, - frame_id: Array, kernel: BaseInteractionKernel, ) -> Array: r"""Energy of the periodic frames over explicit wave numbers, per atom. @@ -78,16 +76,17 @@ def kspace_energy( :math:`\phi_i = \frac{k_e}{V} \sum_{\mathbf{k}} \hat K(\mathbf{k}) [\cos(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Re}\,\hat\rho + \sin(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Im}\,\hat\rho] - k_e K(0) q_i`. - The geometry is strained by :attr:`KGrid.strain` (see :class:`KGrid`). + Every frame pairs the wave numbers of its row of :attr:`KGrid.vec` with + the nodes of its row of :attr:`KGrid.slot`, in products batched over the + frames and the blocks of their wave numbers. The geometry is strained by + :attr:`KGrid.strain` (see :class:`KGrid`). Parameters ---------- charge Node charges with shape ``(N, C)``, zero on the nodes that carry none. k_grid - The wave numbers of the periodic frames. - frame_id - The frame of every node with shape ``(N,)``. + The wave numbers and the nodes of the periodic frames. kernel The pair kernel. @@ -98,38 +97,55 @@ def kspace_energy( without wave numbers. """ xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) dtype = charge.dtype - nf = k_grid.volume.shape[0] - coord = xp.astype(k_grid.coord, dtype) - k_vec = xp.astype(k_grid.vec, dtype) - volume = xp.astype(k_grid.volume, dtype) - # === Step 1. Strained geometry: r' = S r, k' = S^-T k, V' = V det S === - eye = xp.eye(3, dtype=dtype, device=array_api_compat.device(charge)) + n_node, n_channel = charge.shape + nf, n_slot = k_grid.slot.shape + # === Step 1. The nodes of every frame, (nf, ns, 3) and (nf, ns, C) === + real = k_grid.slot >= 0 + index = xp.reshape(xp.where(real, k_grid.slot, xp.zeros_like(k_grid.slot)), (-1,)) + coord = xp.take(xp.astype(k_grid.coord, dtype), index, axis=0) + coord = xp.reshape(coord, (nf, n_slot, 3)) + charge_slot = xp.reshape(xp.take(charge, index, axis=0), (nf, n_slot, n_channel)) + charge_slot = charge_slot * xp.astype(real, dtype)[:, :, None] + # === Step 2. Strained wave numbers and volumes: k' = S^-T k, V' = V det S === + # Every product over the wave numbers of a frame broadcasts the per-frame + # factor over the blocks, (nf, nb, kb, .), so that its sum runs over the + # points of a block. + eye = xp.eye(3, dtype=dtype, device=device) strain = xp.astype(k_grid.strain, dtype) + eye - coord = xp.matmul(xp.take(strain, frame_id, axis=0), coord[:, :, None])[:, :, 0] - strain_inv_t = xp.matrix_transpose(xp.linalg.inv(strain)) - k_vec = xp.matmul(xp.take(strain_inv_t, k_grid.frame, axis=0), k_vec[:, :, None]) - k_vec = k_vec[:, :, 0] - volume = volume * xp.linalg.det(strain) - # === Step 2. Structure factors, (nk, C) === - phase = xp.matmul(k_vec, xp.matrix_transpose(coord)) - same_frame = xp.astype(frame_id[None, :] == k_grid.frame[:, None], dtype) - cos_phase = xp.cos(phase) * same_frame - sin_phase = xp.sin(phase) * same_frame - rho_real = xp.matmul(cos_phase, charge) - rho_imag = xp.matmul(sin_phase, charge) - # === Step 3. Potential at the nodes, (N, C) === - weight = COULOMB_CONSTANT * kernel.kspace(xp.sum(k_vec * k_vec, axis=-1)) - weight = weight / xp.take(volume, k_grid.frame, axis=0) - weight = xp.where(k_grid.mask, weight, xp.zeros_like(weight))[:, None] - potential = xp.matmul(xp.matrix_transpose(cos_phase), weight * rho_real) - potential = potential + xp.matmul(xp.matrix_transpose(sin_phase), weight * rho_imag) + k_vec = xp.astype(k_grid.vec, dtype) + # The rows of k_vec S^-1 are the strained wave numbers. + strained_k = xp.matmul(k_vec, xp.linalg.inv(strain)[:, None, :, :]) + volume = xp.astype(k_grid.volume, dtype) * xp.linalg.det(strain) + # === Step 3. Structure factors, (nf, nb, kb, C) === + # The phase k' . r' = k . r does not change under the strain. + phase = xp.matmul(k_vec, xp.matrix_transpose(coord)[:, None, :, :]) + cos_phase = xp.cos(phase) + sin_phase = xp.sin(phase) + rho_real = xp.matmul(cos_phase, charge_slot[:, None, :, :]) + rho_imag = xp.matmul(sin_phase, charge_slot[:, None, :, :]) + # === Step 4. Potential at the nodes of every frame, (nf, ns, C) === + weight = COULOMB_CONSTANT * kernel.kspace(xp.sum(strained_k * strained_k, axis=-1)) + weight = weight / volume[:, None, None] + weight = xp.where(k_grid.mask, weight, xp.zeros_like(weight))[..., None] + block_potential = xp.matmul(xp.matrix_transpose(cos_phase), weight * rho_real) + block_potential = block_potential + xp.matmul( + xp.matrix_transpose(sin_phase), weight * rho_imag + ) + potential = xp.sum(block_potential, axis=1) # The self term belongs to the lattice sum of a frame with wave numbers; # a non-periodic batch evaluated by a periodic module carries none. - has_points = segment_sum(xp.astype(k_grid.mask, dtype), k_grid.frame, nf) > 0 - periodic_node = xp.take(has_points, frame_id, axis=0) - potential = potential - _self_potential(kernel, charge, periodic_node) - return 0.5 * xp.sum(charge * potential, axis=-1) + has_points = xp.astype(xp.any(k_grid.mask, axis=(1, 2)), dtype)[:, None, None] + potential = ( + potential + - COULOMB_CONSTANT * kernel.contact_value(charge) * charge_slot * has_points + ) + slot_energy = 0.5 * xp.sum(charge_slot * potential, axis=-1) + # === Step 5. Energy on the node axis === + # A padding slot carries no charge and adds a zero to the first node. + energy = xp.zeros((n_node,), dtype=dtype, device=device) + return xp_add_at(energy, index, xp.reshape(slot_energy, (-1,))) def bspline_integer_weights(order: int) -> np.ndarray: diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index f5667a4510..f13d80cb46 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -69,6 +69,7 @@ center_edge_pairs, ) from .reciprocal_graph import ( + K_BLOCK, K_GRID, PME_MESH, KGrid, @@ -98,6 +99,7 @@ __all__ = [ "ALL_PAIRS", + "K_BLOCK", "K_GRID", "PME_MESH", "TOPOLOGY_TYPES", diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index 8e68625d78..e474abd4ed 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -182,13 +182,23 @@ def build_pair_graph( return PairGraph(index=index, vec=vec, mask=mask) +# Guard pairs of the pair topology of a batch without non-periodic frames. +_GUARD_PAIRS = 4 + + def empty_pair_graph(like: Array) -> PairGraph: - """The pair topology of a batch without non-periodic frames: guard pairs only.""" + """The pair topology of a batch without non-periodic frames: guard pairs only. + + The guard pairs number :data:`_GUARD_PAIRS`, unlike the two rows of the + index and the three components of the vectors: a compiled graph that met + a pair axis of the length of one of those would identify the two, and + split at the reductions over the pairs. + """ xp = array_api_compat.array_namespace(like) device = array_api_compat.device(like) index = xp.zeros((2, 0), dtype=xp.int64, device=device) vec = xp.zeros((0, 3), dtype=like.dtype, device=device) - index, vec, mask = pad_and_guard_edges(index, vec, None) + index, vec, mask = pad_and_guard_edges(index, vec, None, min_edges=_GUARD_PAIRS) return PairGraph(index=index, vec=vec, mask=mask) diff --git a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py index fdcd2c57d0..766dfd45bc 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py @@ -54,27 +54,33 @@ @dataclass class KGrid: - r"""Reciprocal lattice points of the periodic frames of a batch. - - The wave numbers :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` of every frame, - for integer :math:`\mathbf{n}`, are built eagerly outside any compiled - region and concatenated over the frames. The k axis ends in masked guard - points, which keep its length a dynamic dimension of its own however few - points a batch holds. + r"""Reciprocal lattice points of the periodic frames of a batch, frame by frame. + + Row ``f`` of :attr:`vec` holds the wave numbers :math:`\mathbf{k} = 2\pi + B_f^{-1}\mathbf{n}` of frame ``f``, for integer :math:`\mathbf{n}`, built + eagerly outside any compiled region, in blocks of :data:`K_BLOCK` points + padded with masked guard points to the blocks of the longest row; row + ``f`` of :attr:`slot` holds the nodes of frame ``f``, padded with ``-1``. + A frame thus pairs its own wave numbers with its own nodes alone, at a + cost of the frame count times the padded wave-number count and the + padded node count, and every sum over the wave numbers of a frame runs + over the points of a block first. The block, point and slot axes are dynamic dimensions of + their own, of at least two entries. The energy is evaluated on the strained geometry of :attr:`strain`: positions by ``S = I + strain``, wave numbers by :math:`S^{-T}` and volumes by :math:`\det S`, so that the gradient with respect to the strain is the - negative virial. The graph lower makes :attr:`coord` and :attr:`strain` - autograd leaves at its inputs. + negative virial; the phases :math:`\mathbf{k} \cdot \mathbf{r}` do not + change under the strain. The graph lower makes :attr:`coord` and + :attr:`strain` autograd leaves at its inputs. """ - frame: Array - """(nk,) int64 the frame of every wave number.""" vec: Array - """(nk, 3) float the wave numbers, in Å⁻¹.""" + """(nf, nb, kb, 3) float the wave numbers of every frame in blocks, in Å⁻¹.""" mask: Array - """(nk,) bool real (1) vs guard (0) points.""" + """(nf, nb, kb) bool real (1) vs guard (0) points.""" + slot: Array + """(nf, ns) int64 the nodes of every frame, ``-1`` past its last node.""" volume: Array """(nf,) float the cell volumes, in ų.""" coord: Array @@ -87,14 +93,14 @@ class KGrid: LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") """Fields an autograd backend makes leaves of the energy.""" AXES: ClassVar[dict[str, dict[int, str]]] = { - "frame": {0: "entry"}, - "vec": {0: "entry"}, - "mask": {0: "entry"}, + "vec": {0: "frame", 1: "entry", 2: "k_point"}, + "mask": {0: "frame", 1: "entry", 2: "k_point"}, + "slot": {0: "frame", 1: "slot"}, "volume": {0: "frame"}, "coord": {0: "node"}, "strain": {0: "frame"}, } - """Dynamic axes of every field: the k axis, the frame axis, the node axis.""" + """Dynamic axes of every field: the frame, block, point, slot and node axes.""" @dataclass @@ -269,22 +275,60 @@ def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: return n_grid[np.any(n_grid != 0, axis=-1)] +K_BLOCK = 128 +"""Wave numbers per block of a row of a :class:`KGrid`.""" +# A guard point carries a nonzero wave number, whose kernel transform stays +# finite, and no weight. +_GUARD_K = (1e-6, 0.0, 0.0) +# The fewest entries of the block and slot axes, which keeps each a dynamic +# dimension of its own however few points or nodes a batch holds. +_MIN_ENTRIES = 2 + + +def _k_blocks(rows: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Lay the wave numbers of every frame out in blocks, padded with guard points. + + Returns the wave numbers with shape ``(nf, nb, K_BLOCK, 3)`` and their + mask with shape ``(nf, nb, K_BLOCK)``. + """ + longest = max([row.shape[0] for row in rows], default=0) + blocks = max(-(-longest // K_BLOCK), _MIN_ENTRIES) + vec = np.tile(np.asarray(_GUARD_K), (len(rows), blocks * K_BLOCK, 1)) + mask = np.zeros((len(rows), blocks * K_BLOCK), dtype=bool) + for f, row in enumerate(rows): + vec[f, : row.shape[0]] = row + mask[f, : row.shape[0]] = True + return ( + vec.reshape(len(rows), blocks, K_BLOCK, 3), + mask.reshape(len(rows), blocks, K_BLOCK), + ) + + +def _frame_slots(n_node: Array) -> np.ndarray: + """The nodes of every frame, one row each, padded with ``-1`` to the longest frame.""" + counts = np.asarray(to_numpy_array(n_node), dtype=np.int64).reshape(-1) + width = max(int(counts.max(initial=0)), _MIN_ENTRIES) + first = np.cumsum(counts) - counts + column = np.arange(width, dtype=np.int64)[None, :] + return np.where(column < counts[:, None], first[:, None] + column, -1) + + def _k_grid( - box: Array, frame: np.ndarray, vec: np.ndarray, mask: np.ndarray, coord: Array + box: Array, vec: np.ndarray, mask: np.ndarray, slot: np.ndarray, coord: Array ) -> KGrid: xp = array_api_compat.array_namespace(box) device = array_api_compat.device(box) return KGrid( - frame=xp.asarray(frame, device=device), vec=xp.astype(xp.asarray(vec, device=device), box.dtype), mask=xp.asarray(mask, device=device), + slot=xp.asarray(slot, device=device), volume=xp.astype(_volume(box), box.dtype), coord=coord, strain=_zero_strain(box), ) -def build_k_grid(box: Array, k_max: float, coord: Array) -> KGrid: +def build_k_grid(box: Array, k_max: float, coord: Array, n_node: Array) -> KGrid: r"""Enumerate the reciprocal lattice points of the periodic frames. With the lattice vectors as the rows of the cell :math:`B`, the wave @@ -292,8 +336,9 @@ def build_k_grid(box: Array, k_max: float, coord: Array) -> KGrid: integer :math:`\mathbf{n}` with :math:`0 < |\mathbf{k}| \le k_{max}`, so that :math:`\mathbf{a}_d \cdot \mathbf{k} = 2\pi n_d`. Their count is data dependent, so the enumeration runs eagerly. The points of every frame - follow the lexicographic order of their integer indices, and the batch - ends in two masked guard points. + follow the lexicographic order of their integer indices, and the rows of + the frames end in masked guard points up to the blocks of the longest + row. Parameters ---------- @@ -304,40 +349,38 @@ def build_k_grid(box: Array, k_max: float, coord: Array) -> KGrid: The wave-number cutoff, in Å⁻¹. coord Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. Returns ------- KGrid - The reciprocal lattice points of every frame, concatenated. + The reciprocal lattice points of every frame, one row each. """ box = _cells(box) cells = np.asarray(to_numpy_array(box), dtype=np.float64) bounds = k_grid_bounds_of_cells(cells, k_max) - vec_list = [] - frame_list = [] + rows = [] for f, cell in enumerate(cells): n_grid = enumerate_k_grid_bounds(bounds[f]).astype(np.float64) k = n_grid @ (2.0 * math.pi * np.linalg.inv(cell)).T - k = k[np.sum(k * k, axis=-1) <= k_max * k_max] - vec_list.append(k) - frame_list.append(np.full(k.shape[0], f, dtype=np.int64)) - # Two guard points with a nonzero wave number, whose kernel transform - # stays finite. - vec_list.append(np.array([[1e-6, 0.0, 0.0], [1e-6, 0.0, 0.0]])) - frame_list.append(np.zeros(2, dtype=np.int64)) - vec = np.concatenate(vec_list, axis=0) - mask = np.arange(vec.shape[0], dtype=np.int64) < vec.shape[0] - 2 - return _k_grid(box, np.concatenate(frame_list), vec, mask, coord) - - -def build_fixed_k_grid(box: Array, bounds: Array, k_max: float, coord: Array) -> KGrid: + rows.append(k[np.sum(k * k, axis=-1) <= k_max * k_max]) + vec, mask = _k_blocks(rows) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def build_fixed_k_grid( + box: Array, bounds: Array, k_max: float, coord: Array, n_node: Array +) -> KGrid: r"""Evaluate a fixed integer k grid on the cells of the periodic frames. Every frame shares the integer indices within ``bounds`` (see - :func:`enumerate_k_grid_bounds`), so the k axis has the static length - ``nf * nk`` without guard points. The wave numbers of a frame are - :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}`, and the points beyond the - cutoff :math:`|\mathbf{k}| \le k_{max}` are masked out. + :func:`enumerate_k_grid_bounds`), so the blocks of a row have the static + count of the grid. The wave numbers of a frame are :math:`\mathbf{k} = + 2\pi B^{-1}\mathbf{n}`, and the points beyond the cutoff + :math:`|\mathbf{k}| \le k_{max}` and the guard points that complete the + last block are masked out. A frame whose cell needs indices beyond the grid, a larger cell than the grid was sized for or a kernel whose cutoff has grown since, is evaluated @@ -357,18 +400,19 @@ def build_fixed_k_grid(box: Array, bounds: Array, k_max: float, coord: Array) -> The wave-number cutoff, in Å⁻¹. coord Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. Returns ------- KGrid - The wave numbers of every frame, frame-major over the shared grid. + The wave numbers of every frame over the shared grid, one row each. """ box = _cells(box) cells = np.asarray(to_numpy_array(box), dtype=np.float64) n_grid = enumerate_k_grid_bounds(bounds) bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(3) - nf = cells.shape[0] - nk = n_grid.shape[0] required = k_grid_bounds_of_cells(cells, k_max) beyond = np.any(required > bounds[None, :], axis=1) if np.any(beyond): @@ -381,23 +425,19 @@ def build_fixed_k_grid(box: Array, bounds: Array, k_max: float, coord: Array) -> ) rec = 2.0 * math.pi * np.linalg.inv(cells) k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) - k_sq = np.sum(k * k, axis=-1) - mask = (k_sq > 0.0) & (k_sq <= k_max * k_max) - frame = np.repeat(np.arange(nf, dtype=np.int64), nk) - return _k_grid(box, frame, k.reshape(nf * nk, 3), mask.reshape(-1), coord) + vec, real = _k_blocks(list(k)) + mask = real & (np.sum(vec * vec, axis=-1) <= k_max * k_max) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) def empty_k_grid(nf: int, coord: Array) -> KGrid: - """The k grid of a batch without periodic frames: two guard points.""" + """The k grid of a batch without periodic frames: guard points and no nodes.""" xp = array_api_compat.array_namespace(coord) device = array_api_compat.device(coord) eye = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + vec, mask = _k_blocks([np.zeros((0, 3), dtype=np.float64)] * nf) return _k_grid( - eye, - np.zeros(2, dtype=np.int64), - np.array([[1e-6, 0.0, 0.0], [1e-6, 0.0, 0.0]]), - np.zeros(2, dtype=bool), - coord, + eye, vec, mask, np.full((nf, _MIN_ENTRIES), -1, dtype=np.int64), coord ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/topology.py b/deepmd/dpmodel/utils/neighbor_graph/topology.py index 0fe4127123..9a504fc7f2 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/topology.py +++ b/deepmd/dpmodel/utils/neighbor_graph/topology.py @@ -75,6 +75,7 @@ def _build_reciprocal( name: str, box: Array, coord: Array, + n_node: Array, params: dict[str, Any], ) -> KGrid | MeshGrid: """Build a reciprocal-space topology of periodic frames from its parameters.""" @@ -87,8 +88,8 @@ def _build_reciprocal( if name == PME_MESH: return build_mesh_grid(box, params["spacing"], coord) if "bounds" in params: - return build_fixed_k_grid(box, params["bounds"], params["k_max"], coord) - return build_k_grid(box, params["k_max"], coord) + return build_fixed_k_grid(box, params["bounds"], params["k_max"], coord, n_node) + return build_k_grid(box, params["k_max"], coord, n_node) def build_topologies( @@ -179,7 +180,7 @@ def build_topologies( name: ( empty_pair_graph(coord) if name == ALL_PAIRS - else _build_reciprocal(name, box, coord, params.get(name, {})) + else _build_reciprocal(name, box, coord, n_node, params.get(name, {})) ) for name in names } diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 13b15d63b7..6d511315a0 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -705,11 +705,14 @@ def build_synthetic_topologies( """Build trace-time samples of the topologies a model declares. Single source of the trace-time topology inputs, shared by graph export - and compiled training. Tracing reads only their shapes. Every topology - lists ``n_entry`` entries on its own axis, the pairs or the wave numbers, - a length the caller keeps distinct from the other dynamic axes; a mesh - takes the three smallest even dimensions of at least four that no other - axis and no entry of ``forbidden`` has. + and compiled training. Tracing reads only their shapes. The pair + topology lists ``n_entry`` pairs, a length the caller keeps distinct from + the other dynamic axes; a mesh takes the three smallest even dimensions, + and a k grid the three smallest lengths for its slots, its blocks and the + points of a block, of at least four that no other axis and no entry of + ``forbidden`` has. Every length thus differs from the two rows of a pair + index and the three components of a vector, and from every other dynamic + axis, which a compiled graph would otherwise identify with it. Parameters ---------- @@ -718,7 +721,7 @@ def build_synthetic_topologies( n_node_total : int Length of the node axis. n_entry : int - Length of the entry axis of every sample. + Length of the pair axis. dtype : torch.dtype Precision of the geometry. device : torch.device @@ -726,7 +729,8 @@ def build_synthetic_topologies( nframes : int Length of the frame axis. forbidden : set[int] - Sizes of the static dimensions of the model, which a mesh avoids. + Sizes of the static dimensions of the model, which the mesh and k-grid + axes avoid. Returns ------- @@ -748,6 +752,8 @@ def build_synthetic_topologies( for size in range(4, 4 + 2 * len(taken) + 6, 2): if size not in taken and len(mesh_dims) < 3: mesh_dims.append(size) + free = sorted(set(range(4, len(taken) + 10)) - taken - set(mesh_dims)) + n_slot, n_block, n_point = free[:3] coord = torch.ones(n_node_total, 3, dtype=dtype, device=device) # Contiguous like the inputs at run time: a compiled graph reads the # strides of its samples. @@ -755,11 +761,13 @@ def build_synthetic_topologies( topologies = {} for name in model.atomic_model.required_topologies(): if name == K_GRID: + slot = torch.arange(nframes * n_slot, dtype=torch.int64, device=device) + blocks = (nframes, n_block, n_point) topologies[name] = KGrid( - frame=torch.zeros(n_entry, dtype=torch.int64, device=device), # Nonzero wave numbers: the kernel transform carries 1/k^2. - vec=torch.ones(n_entry, 3, dtype=dtype, device=device), - mask=torch.ones(n_entry, dtype=torch.bool, device=device), + vec=torch.ones(*blocks, 3, dtype=dtype, device=device), + mask=torch.ones(*blocks, dtype=torch.bool, device=device), + slot=(slot % n_node_total).reshape(nframes, n_slot), volume=torch.ones(nframes, dtype=dtype, device=device), coord=coord, strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), @@ -964,8 +972,8 @@ def _topology_dynamic_shapes( Every field follows the ``AXES`` its dataclass declares. The frame and node axes are those of the other inputs in ``shared``; the entry axis of every topology is a dimension of its own, of at least the two guard - entries; a mesh axis is twice a dimension, since the mesh dimensions are - even. + entries, and so are the slot axis and the point axis of a k grid; a mesh + axis is twice a dimension, since the mesh dimensions are even. Parameters ---------- @@ -988,6 +996,8 @@ def _topology_dynamic_shapes( cls = TOPOLOGY_TYPES[name] dims = dict(shared) dims["entry"] = torch.export.Dim(f"{name}_entries", min=2) + dims["slot"] = torch.export.Dim(f"{name}_slots", min=2) + dims["k_point"] = torch.export.Dim(f"{name}_block_points", min=2) for axis in ("mesh_1", "mesh_2", "mesh_3"): dims[axis] = 2 * torch.export.Dim(f"{name}_{axis}_half", min=1) specs += [ @@ -2356,8 +2366,8 @@ def _trace_and_export_impl( device=torch.device("cpu"), want_spin=is_native_spin, ) - # The pair axis of every topology is dynamic; its sample length - # is kept distinct from the other axes and from the CSR offsets. + # The pair axis is dynamic; its sample length is kept distinct + # from the other axes and from the CSR offsets. topologies = build_synthetic_topologies( model, n_sample, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a9633e425f..0e3ac1873a 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3709,17 +3709,18 @@ def long_range_latent_charge() -> list[Argument]: ) doc_k_grid_mode = ( "How training evaluates the reciprocal sum of a periodic batch: " - "'ragged' over the wave numbers of every frame within k_max, " - "enumerated per batch, so that the k axis is a dynamic dimension; " - "'fixed' over one integer grid shared by the frames, sized by the " - "largest cell of the data statistics or by `k_grid_bounds`, with the " - "points a frame does not reach masked, so that the k axis has a static " - "length (a frame whose cell needs points beyond the grid is truncated " - "to it, with a warning); 'pme' by the smooth particle-mesh Ewald " - "method on a mesh sized by the cells of the batch at `pme_spacing`, " - "whose cost grows nearly linearly with the system. A frozen model " + "'pme', the default of a periodic module, by the smooth particle-mesh " + "Ewald method on a mesh sized by the cells of the batch at " + "`pme_spacing`, whose cost grows nearly linearly with the system; " + "'ragged' over the wave numbers " + "of every frame within k_max, enumerated per batch, so that the k axis " + "is a dynamic dimension; or 'fixed' over one integer grid shared by " + "the frames, sized by the largest cell of the data statistics or by " + "`k_grid_bounds`, with the points a frame does not reach masked, so " + "that the k axis has a static length (a frame whose cell needs points " + "beyond the grid is truncated to it, with a warning). A frozen model " "evaluates every periodic frame on the mesh of its cell, whatever the " - "mode. Requires `periodic: true` for 'fixed' and 'pme'." + "mode. 'fixed' and 'pme' require `periodic: true`." ) doc_k_grid_bounds = ( "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " @@ -3783,11 +3784,11 @@ def long_range_latent_charge() -> list[Argument]: ), Argument( "k_grid_mode", - str, + [str, None], optional=True, - default="ragged", + default=None, doc=doc_k_grid_mode, - extra_check=lambda x: x in ("ragged", "fixed", "pme"), + extra_check=lambda x: x in (None, "ragged", "fixed", "pme"), extra_check_errmsg="must be 'ragged', 'fixed' or 'pme'", ), Argument( diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 1f190dbb9a..5da7d4353e 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -147,15 +147,8 @@ keeps the cutoff of the initial kernel. {ref}`k_grid_mode ` selects how training evaluates the sum: -- `ragged`: over the wave numbers of every frame within $k_{max}$, listed for - every batch; their number grows with the volume of the cell. -- `fixed`: over one integer grid shared by the frames of a batch, sized by the - largest cell of the data statistics or by - {ref}`k_grid_bounds `; - a frame whose cell needs a wave number beyond the grid is truncated to it, - with a warning. -- `pme`: by the smooth particle-mesh Ewald method. The charges are spread - with B-splines of order +- `pme`, the default: by the smooth particle-mesh Ewald method. The charges + are spread with B-splines of order {ref}`pme_order ` to a mesh that resolves every cell of the batch at {ref}`pme_spacing `, @@ -163,6 +156,13 @@ selects how training evaluates the sum: number of atoms. With the default order 8 and spacing $0.75\pi/k_{max}$, the energy and the forces stay within a few parts in a million of the sum over the wave numbers. +- `ragged`: over the wave numbers of every frame within $k_{max}$, listed for + every batch; their number grows with the volume of the cell. +- `fixed`: over one integer grid shared by the frames of a batch, sized by the + largest cell of the data statistics or by + {ref}`k_grid_bounds `; + a frame whose cell needs a wave number beyond the grid is truncated to it, + with a warning. A frozen model evaluates every periodic frame on the mesh of its cell, whatever its training mode, so a model trained in another mode can differ diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h index 9e5e05faa7..352f85a643 100644 --- a/source/api_cc/src/longRangePTExpt.h +++ b/source/api_cc/src/longRangePTExpt.h @@ -62,7 +62,8 @@ inline std::vector buildAllPairsTopology( * @brief Build the all-pairs topology of a periodic frame, which holds no * pair: the reciprocal-space topology describes the interactions. * - * Twin of Python ``empty_pair_graph``: two masked guard pairs. + * Counterpart of Python ``empty_pair_graph`` for the exported lower: two + * masked guard pairs, the lower bound of its pair axis. */ inline std::vector buildGuardPairTopology( const torch::Device& device, const torch::ScalarType vec_dtype) { diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 2f826a1e3b..79320de4a3 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -23,6 +23,7 @@ PairExcludeMask, ) from deepmd.dpmodel.utils.neighbor_graph import ( + K_BLOCK, build_k_grid, build_pair_graph, build_topologies, @@ -266,15 +267,20 @@ def test_kernel_contact_value_matches_real_space_limit(kernel: str) -> None: def test_build_k_grid_orthogonality_and_guards() -> None: - """The k points satisfy k . a_d = 2 pi n_d, are ordered and carry guards.""" + """The k points satisfy k . a_d = 2 pi n_d, every frame fills a row of its own.""" box = np.array([[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]])[None] - grid = build_k_grid(box, 2.0, np.zeros((4, 3))) - # guard points are two masked (1e-6, 0, 0) entries - assert grid.mask[-2:].tolist() == [False, False] - assert grid.mask[:-2].all() - n_integer = grid.vec[:-2] @ box[0].T / (2 * np.pi) + grid = build_k_grid(box, 2.0, np.zeros((4, 3)), np.array([4])) + # the points fill whole blocks, the last one completed by guard points + assert grid.mask.shape == grid.vec.shape[:3] + assert grid.vec.shape[2] == K_BLOCK + flat_mask = grid.mask.reshape(-1) + count = int(flat_mask.sum()) + assert flat_mask[:count].all() and not flat_mask[count:].any() + np.testing.assert_array_equal(grid.slot, [[0, 1, 2, 3]]) + vec = grid.vec.reshape(-1, 3)[:count] + n_integer = vec @ box[0].T / (2 * np.pi) np.testing.assert_allclose(n_integer, np.round(n_integer), atol=1e-10) - k_sq = np.sum(grid.vec[:-2] ** 2, axis=-1) + k_sq = np.sum(vec**2, axis=-1) assert (k_sq > 0).all() and (k_sq <= 4.0 + 1e-12).all() np.testing.assert_allclose(grid.volume, np.abs(np.linalg.det(box[0]))) # every integer point inside the sphere is present exactly once @@ -308,7 +314,7 @@ def test_build_topologies_periodic_branching() -> None: with_box = build_topologies( names, coord, n_node, box, atype=atype, topology_params=params ) - assert with_box[ALL_PAIRS].index.shape[1] == 2 # guards only + assert with_box[ALL_PAIRS].index.shape[1] == 4 # guards only assert not with_box[ALL_PAIRS].mask.any() assert with_box[K_GRID].mask.sum() > 0 without_box = build_topologies( @@ -345,6 +351,58 @@ def test_build_topologies_periodic_branching() -> None: ) +def test_build_k_grid_pads_the_rows_of_the_frames() -> None: + """Every frame keeps its own points and nodes; shorter rows end in masked guards.""" + box = np.stack([np.diag([8.0, 7.5, 9.0]), np.diag([5.0, 5.5, 6.0])]) + n_node = np.array([3, 5]) + grid = build_k_grid(box, 2.0, np.zeros((8, 3)), n_node) + alone = [ + build_k_grid(box[f : f + 1], 2.0, np.zeros((n, 3)), np.array([n])) + for f, n in enumerate(n_node) + ] + counts = [int(single.mask.sum()) for single in alone] + blocks = max(-(-max(counts) // K_BLOCK), 2) + assert grid.vec.shape == (2, blocks, K_BLOCK, 3) + for f, single in enumerate(alone): + mask = grid.mask[f].reshape(-1) + vec = grid.vec[f].reshape(-1, 3) + np.testing.assert_array_equal(mask, np.arange(blocks * K_BLOCK) < counts[f]) + np.testing.assert_allclose( + vec[: counts[f]], single.vec[0].reshape(-1, 3)[: counts[f]] + ) + assert np.all(vec[counts[f] :] == [1e-6, 0.0, 0.0]) + np.testing.assert_array_equal(grid.slot, [[0, 1, 2, -1, -1], [3, 4, 5, 6, 7]]) + + +@pytest.mark.parametrize("kernel", [LESKernel(alpha=0.6), SOGKernel(n_gaussian=3)]) +def test_kspace_energy_of_padded_frames_matches_frames_alone( + kernel: BaseInteractionKernel, +) -> None: + """Frames of unequal node counts, an empty one among them, sum as they do alone.""" + from deepmd.dpmodel.long_range.reciprocal import ( + kspace_energy, + ) + + box = np.stack( + [np.diag([8.0, 7.5, 9.0]), np.diag([6.0, 6.5, 7.0]), np.diag([5.0, 5.5, 6.0])] + ) + n_node = np.array([4, 0, 6]) + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 5.0, (10, 3)) + charge = rng.normal(0.0, 0.5, (10, 2)) + energy = kspace_energy(charge, build_k_grid(box, 2.5, coord, n_node), kernel) + start = 0 + for f, n in enumerate(n_node): + if n == 0: + continue + nodes = slice(start, start + n) + alone = build_k_grid(box[f : f + 1], 2.5, coord[nodes], np.array([n])) + np.testing.assert_allclose( + energy[nodes], kspace_energy(charge[nodes], alone, kernel), rtol=1e-12 + ) + start += n + + # --------------------------------------------------------------------------- # Fixed k-grid mode # --------------------------------------------------------------------------- @@ -369,7 +427,9 @@ def _module_topologies( def _periodic_long_range(**kwargs: object) -> LatentChargeLongRange: + """A periodic module, over the explicit wave numbers unless ``k_grid_mode`` says otherwise.""" charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + kwargs.setdefault("k_grid_mode", "ragged") return LatentChargeLongRange( charge_net, LESKernel(alpha=0.3), None, periodic=True, **kwargs ) @@ -410,23 +470,19 @@ def test_fixed_k_grid_matches_ragged_enumeration() -> None: k_max = 2.0 bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) nk = enumerate_k_grid_bounds(bounds).shape[0] - fixed = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3))) - # a static, guard-free layout: nf * nk points, frame-major - assert fixed.vec.shape == (2 * nk, 3) - np.testing.assert_array_equal(fixed.frame, np.repeat([0, 1], nk)) - ragged = build_k_grid(box, k_max, np.zeros((9, 3))) + n_node = np.array([4, 5]) + fixed = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node) + # a static layout: the blocks of the grid in every row + assert fixed.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + ragged = build_k_grid(box, k_max, np.zeros((9, 3)), n_node) for frame in range(2): expected = { - tuple(row) - for row in np.round(ragged.vec[(ragged.frame == frame) & ragged.mask], 8) - } - got = { - tuple(row) - for row in np.round(fixed.vec[(fixed.frame == frame) & fixed.mask], 8) + tuple(row) for row in np.round(ragged.vec[frame][ragged.mask[frame]], 8) } + got = {tuple(row) for row in np.round(fixed.vec[frame][fixed.mask[frame]], 8)} assert got == expected # the smaller cell masks more points of the shared grid - assert fixed.mask[:nk].sum() < fixed.mask[nk:].sum() + assert fixed.mask[0].sum() < fixed.mask[1].sum() np.testing.assert_allclose(fixed.volume, np.abs(np.linalg.det(box)), rtol=1e-12) @@ -442,21 +498,22 @@ def test_fixed_k_grid_truncates_oversized_cells() -> None: k_max = 2.0 bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) nk = enumerate_k_grid_bounds(bounds).shape[0] + n_node = np.array([4, 5]) with pytest.warns(UserWarning, match="truncated"): - bigger = build_fixed_k_grid(box * 2.0, bounds, k_max, np.zeros((9, 3))) - # the layout is untouched: static, guard-free, every grid point present - assert bigger.vec.shape == (2 * nk, 3) - np.testing.assert_array_equal(bigger.frame, np.repeat([0, 1], nk)) + bigger = build_fixed_k_grid(box * 2.0, bounds, k_max, np.zeros((9, 3)), n_node) + # the layout is untouched: static, every grid point present + assert bigger.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) # every grid point is still a valid wave number of the larger cell; the # mask keeps the whole k_max sphere the grid reaches k2 = np.sum(bigger.vec**2, axis=-1) - np.testing.assert_array_equal(bigger.mask, (k2 > 0) & (k2 <= k_max * k_max)) + real = (np.arange(bigger.mask[0].size) < nk).reshape(bigger.mask.shape[1:]) + np.testing.assert_array_equal(bigger.mask, real & (k2 <= k_max * k_max)) # the larger cell activates more of the shared grid than the original did - original = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3))) + original = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node) assert bigger.mask.sum() > original.mask.sum() # a larger wave-number cutoff (a drifting kernel) truncates the same way with pytest.warns(UserWarning, match="truncated"): - build_fixed_k_grid(box, bounds, k_max * 2.0, np.zeros((9, 3))) + build_fixed_k_grid(box, bounds, k_max * 2.0, np.zeros((9, 3)), n_node) def test_fixed_k_grid_mode_configuration() -> None: @@ -500,7 +557,8 @@ def test_fixed_k_grid_stats_size_the_grid() -> None: topologies = _module_topologies( long_range, coord, np.array([4, 5]), _cells(), atype ) - assert topologies["k_grid"].vec.shape == (2 * (np.prod(2 * expected + 1) - 1), 3) + nk = np.prod(2 * expected + 1) - 1 + assert topologies["k_grid"].vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) # statistics without a cell cannot size the grid unsized = _periodic_long_range(k_grid_mode="fixed") with pytest.raises(ValueError, match="no simulation cell"): @@ -516,10 +574,20 @@ def test_fixed_k_grid_round_trip() -> None: restored = LatentChargeLongRange.deserialize(record) assert restored.k_grid_mode == "fixed" np.testing.assert_array_equal(restored.k_grid_bounds, long_range.k_grid_bounds) - # the default mode serializes no grid fields + # the ragged mode serializes no grid bounds plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) assert plain.k_grid_mode == "ragged" assert plain.k_grid_bounds is None + # a periodic module evaluates on the mesh unless told otherwise, and a + # periodic record without a mode is one of the explicit wave numbers + charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + default = LatentChargeLongRange(charge_net, LESKernel(), None, periodic=True) + assert default.k_grid_mode == "pme" + assert default.required_topologies() == ("all_pairs", "pme_mesh") + record = default.serialize() + assert record["k_grid_mode"] == "pme" + record.pop("k_grid_mode") + assert LatentChargeLongRange.deserialize(record).k_grid_mode == "ragged" # --------------------------------------------------------------------------- @@ -563,7 +631,7 @@ def test_decoupled_k_grid_ignores_kernel_drift() -> None: moved = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] np.testing.assert_array_equal(drifted.vec, initial.vec) np.testing.assert_array_equal(drifted.mask, initial.mask) - assert moved.vec.shape[0] > initial.vec.shape[0] + assert np.sum(moved.mask) > np.sum(initial.mask) def test_decoupled_fixed_grid_keeps_its_mask() -> None: @@ -646,7 +714,7 @@ def test_pme_mesh_grid_build_and_flatten() -> None: np.testing.assert_array_equal( mesh.mode_1, np.where(expected > n1 // 2, expected - n1, expected) ) - assert tops["all_pairs"].index.shape[1] == 2 # guards only + assert tops["all_pairs"].index.shape[1] == 4 # guards only flat = flatten_topologies(tops) assert len(flat) == 3 + 8 back = unflatten_topologies(("all_pairs", PME_MESH), flat) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 62f0d5eb58..c48dc0ae4c 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -765,6 +765,7 @@ def test_periodic_checkpoint_inference_matches_direct_forward( model = _periodic_model(kernel) config = _config(kernel) config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" checkpoint = tmp_path / f"periodic_{kernel}_{grid_mode}.pt" wrapper = ModelWrapper(model, model_params=config) torch.save({"model": wrapper.state_dict()}, checkpoint) @@ -1137,8 +1138,10 @@ def test_pair_graph_skips_padding() -> None: def _periodic_model(kernel: str, **kwargs: object) -> torch.nn.Module: + """A periodic model over the explicit wave numbers of every frame.""" config = _config(kernel, **kwargs) config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" return get_model(config).to(env.DEVICE).eval() @@ -1282,6 +1285,32 @@ def energy(c: np.ndarray, b: np.ndarray) -> float: np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) +@pytest.mark.parametrize("grid_mode", ["ragged", "pme"]) +def test_trace_samples_give_every_axis_its_own_length(grid_mode: str) -> None: + """No two dynamic axes of the trace-time topologies share a length, nor take 2 or 3.""" + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_topologies, + ) + + model = ( + _periodic_model("les") if grid_mode == "ragged" else _periodic_pme_model("les") + ) + nframes, n_node, n_pair = 5, 7, 11 + topologies = build_synthetic_topologies( + model, n_node, n_pair, dtype=torch.float64, device=env.DEVICE, nframes=nframes + ) + lengths = {("", "frame"): nframes, ("", "node"): n_node} + for name, topology in topologies.items(): + for field, axes in type(topology).AXES.items(): + for axis, label in axes.items(): + key = ("", label) if label in ("frame", "node") else (name, label) + size = getattr(topology, field).shape[axis] + assert lengths.setdefault(key, size) == size, (name, field, label) + sizes = list(lengths.values()) + assert len(set(sizes)) == len(sizes), lengths + assert not {2, 3} & set(sizes), lengths + + @pytest.mark.parametrize("kernel", ["les", "sog"]) @REQUIRES_SUPPORTED_COMPILE def test_periodic_compiled_matches_eager(kernel: str) -> None: @@ -1760,6 +1789,7 @@ def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: """ config = _config(kernel) config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" config["long_range"]["decouple_k_grid"] = True model = get_model(config).to(env.DEVICE).eval() coord, atype, box = _periodic_case() From ceefa8f806fa6e8607d79d10d8316b4ab325ad11 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 15:01:14 +0800 Subject: [PATCH 30/30] perf(long-range): first-order strain in the reciprocal-space energy The strain of a topology is zero, and the energy is differentiated once with respect to it, for the virial. The value and the first derivative at zero strain are those of S^-1 = I - eps and det S = 1 + tr eps, which replace the inverse and the determinant of I + eps in the mesh and k-grid energies: their batched factorizations, error checks, and the singular value decompositions and LU solves of their derivatives ran outside the fused kernels of a compiled training step. The eager cell volumes are triple products. Compiled training of the Air preset with a periodic SOG long range (five frames of about 390 atoms per batch, RTX PRO 6000) takes 25.7 ms per step instead of 28.2 ms, against 21.1 ms without long range. Building the topologies of a batch takes 0.40 ms instead of 0.50 ms. --- deepmd/dpmodel/long_range/reciprocal.py | 51 +++++++++++++++---- .../utils/neighbor_graph/reciprocal_graph.py | 11 ++-- 2 files changed, 48 insertions(+), 14 deletions(-) diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py index 9965958b33..27b74e00e4 100644 --- a/deepmd/dpmodel/long_range/reciprocal.py +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -64,6 +64,38 @@ def _self_potential( ) +def _strained_rows(xp: Any, rows: Array, strain: Array) -> Array: + r"""Wave numbers of the strained cells, to first order in the strain. + + The strain :math:`\varepsilon` of a frame maps its cell :math:`B` to + :math:`B S^\mathsf{T}`, :math:`S = I + \varepsilon`, and a wave number, + a row :math:`\mathbf{k}`, to :math:`\mathbf{k} S^{-1}`. The strain of a + topology is zero, and the energy is differentiated once with respect to + it, for the virial; the value and the first derivative at zero strain + are those of :math:`S^{-1} = I - \varepsilon`, which maps the rows here + and keeps linear algebra out of the energy and its derivatives. + + Parameters + ---------- + rows + Wave numbers or reciprocal lattice vectors with shape ``(nf, ..., 3)``. + strain + The strain of every frame with shape ``(nf, 3, 3)``. + + Returns + ------- + Array + The strained wave numbers, of the shape of ``rows``. + """ + eps = xp.reshape(strain, (strain.shape[0],) + (1,) * (rows.ndim - 2) + (3, 3)) + return rows - xp.sum(rows[..., :, None] * eps, axis=-2) + + +def _strained_volume(volume: Array, strain: Array) -> Array: + r"""Volumes of the strained cells, :math:`V \det S = V (1 + \mathrm{tr}\,\varepsilon)` to first order.""" + return volume * (1.0 + strain[:, 0, 0] + strain[:, 1, 1] + strain[:, 2, 2]) + + def kspace_energy( charge: Array, k_grid: KGrid, @@ -112,12 +144,10 @@ def kspace_energy( # Every product over the wave numbers of a frame broadcasts the per-frame # factor over the blocks, (nf, nb, kb, .), so that its sum runs over the # points of a block. - eye = xp.eye(3, dtype=dtype, device=device) - strain = xp.astype(k_grid.strain, dtype) + eye + strain = xp.astype(k_grid.strain, dtype) k_vec = xp.astype(k_grid.vec, dtype) - # The rows of k_vec S^-1 are the strained wave numbers. - strained_k = xp.matmul(k_vec, xp.linalg.inv(strain)[:, None, :, :]) - volume = xp.astype(k_grid.volume, dtype) * xp.linalg.det(strain) + strained_k = _strained_rows(xp, k_vec, strain) + volume = _strained_volume(xp.astype(k_grid.volume, dtype), strain) # === Step 3. Structure factors, (nf, nb, kb, C) === # The phase k' . r' = k . r does not change under the strain. phase = xp.matmul(k_vec, xp.matrix_transpose(coord)[:, None, :, :]) @@ -336,11 +366,12 @@ def mesh_energy( # === Step 3. Influence function on the half spectrum, (nf, n1, n2, h3) === m3 = m_full[2][: n3 // 2 + 1] m_half = [m_full[0], m_full[1], m3] - eye = xp.eye(3, dtype=dtype, device=device) - strain = xp.astype(mesh.strain, dtype) + eye - a = xp.matmul(xp.matrix_transpose(xp.linalg.inv(strain)), box_inv) - metric = xp.matmul(xp.matrix_transpose(a), a) - volume = volume * xp.linalg.det(strain) + strain = xp.astype(mesh.strain, dtype) + # The rows of A^T = B^-T S^-1 are the reciprocal vectors of the strained + # cell, strained like wave numbers; W = A^T A. + a_t = _strained_rows(xp, xp.matrix_transpose(box_inv), strain) + metric = xp.sum(a_t[:, :, None, :] * a_t[:, None, :, :], axis=-1) + volume = _strained_volume(volume, strain) g1 = m_half[0][:, None, None] g2 = m_half[1][None, :, None] g3 = m_half[2][None, None, :] diff --git a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py index 766dfd45bc..c337cf40f4 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py @@ -71,8 +71,9 @@ class KGrid: positions by ``S = I + strain``, wave numbers by :math:`S^{-T}` and volumes by :math:`\det S`, so that the gradient with respect to the strain is the negative virial; the phases :math:`\mathbf{k} \cdot \mathbf{r}` do not - change under the strain. The graph lower makes :attr:`coord` and - :attr:`strain` autograd leaves at its inputs. + change under the strain. The strain is zero, and the evaluation takes it + to first order, which fixes the energy and that gradient. The graph lower + makes :attr:`coord` and :attr:`strain` autograd leaves at its inputs. """ vec: Array @@ -121,7 +122,7 @@ class MeshGrid: :math:`B`, so that the fractional coordinates are invariant and the gradient with respect to the strain is the negative virial. The volumes of the unstrained cells are computed eagerly, and the evaluation scales - them by :math:`\det S`. + them by :math:`\det S`, to first order in the strain. """ box: Array @@ -203,8 +204,10 @@ def _cells(box: Array) -> Array: def _volume(box: Array) -> Array: + """The cell volumes |a_1 . (a_2 x a_3)| with shape ``(nf,)``.""" xp = array_api_compat.array_namespace(box) - return xp.abs(xp.linalg.det(box)) + triple = box[:, 0, :] * xp.linalg.cross(box[:, 1, :], box[:, 2, :]) + return xp.abs(xp.sum(triple, axis=-1)) def _zero_strain(box: Array) -> Array: