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.

Module Attributes

`FAST_BLOCK_ENTRIES`	complex entries of one column block of the factorized reciprocal sum (fast path)
`AUTO_POLICY`	minimum atoms of a periodic structure (LES's cost does not depend on the neighbor graph, so this policy counts atoms, per structure).
`LATENT_CHARGE_PER_E`	two latent charges interact as `q1 q2 / (2 pi r)`, so `q = Q * sqrt(2 pi k_e)` with the Coulomb constant `k_e = HARTREE * BOHR` in eV Angstrom; 9.5118 per e

Functions

`intra_structure_pairs`(batch, num_graphs)	Every ordered pair `(i, j)`, `i != j`, of atoms in the same structure.
`periodic_structures`(cell, pbc, num_graphs, ...)	Which structures of a batch are periodic, `(B,)` bool.
`smooth_size`(n)	The smallest integer of at least `n` (and 8) with no prime factor above 5.

Classes

`EwaldSummation`([dl, sigma, exponent, ...])	Ewald energy of a (latent) per-atom variable `q` (paper eqs 3-5).
`LatentEwald`(model[, n_channels, hidden, ...])	Wrap any xnn model with a Latent-Ewald long-range energy (CACE-LR).