xnn.common.models.les#
Latent Ewald Summation (LES): long-range interactions for any xnn model.
Implements Cheng, npj Comput Mater 11, 80 (2025): short-range MLIPs miss
long-range physics (electrostatics, dispersion) beyond their receptive field.
LES fixes this generically – a small MLP maps each atom’s invariant
features to a low-dimensional hidden variable q (paper eq 2, analogous to
environment-dependent partial charges, but unconstrained), and an Ewald
summation over the structure factor of q (eqs 3-4) supplies the long-range
energy
E_lr = (1/V) sum_{0<k<k_c} exp(-sigma^2 k^2 / 2) / k^2 * abs(S(k))^2 .
Faithful to the reference implementation (cace.modules.EwaldPotential of
BingqingCheng/cace and the training scripts of
BingqingCheng/cace-lr-fit): EwaldSummation follows
the same algorithm (the triclinic-capable reciprocal-space sum with half-space
symmetry weights, the k = 0 and self-interaction conventions, the
1/r^6 dispersion variant of paper eq 5, and the real-space
erf-converged direct sum used for non-periodic structures), independently
implemented and verified against the reference to machine precision in
tests/test_les.py. One upstream wart is fixed rather than reproduced: the
reference always builds its k-vector grid in float32, which crashes float64
runs; here the grid follows the input dtype.
Because LatentEwald only needs invariant per-atom features, it wraps
any registered xnn model – every model exposes its features through the
"node_features" output key and a node_feature_dim attribute (CACE’s
symmetrized B features, the scalar channels of MACE / NequIP node features,
Allegro’s environment-aggregated edge latents, SchNet / PhysNet feature vectors,
HDNNP/ANI descriptors). Enable it from a config with
model.extra["long_range"] (see
build_model()) or wrap directly:
model = LatentEwald(build_model(cfg.model), n_channels=4, sigma=1.0)
out = ForceStressOutput(model)(graph) # forces/stress include E_lr
The wrapped energy cost is roughly twice the short-range cost.
Classes
|
Ewald energy of a (latent) per-atom variable |
|
Wrap any xnn model with a Latent-Ewald long-range energy (CACE-LR). |