xnn.common.models.les.EwaldSummation

class xnn.common.models.les.EwaldSummation(dl=2.0, sigma=1.0, exponent=1, remove_self_interaction=False, method='kspace', pme_order=8, pme_oversample=2.0)

Bases: Module, FastPathModule

Ewald energy of a (latent) per-atom variable q (paper eqs 3-5).

For periodic structures, the reciprocal-space sum over a k-grid limited by |k| <= 2*pi/dl (upstream convention: the paper’s k_c equals 2*pi/dl); for non-periodic structures (no cell), the equivalent erf-converged real-space direct sum. Both handle a multi-dimensional q, summing the energies of the channels.

Parameters:

• dl (float, optional) – Reciprocal grid resolution; the k-space cutoff is 2*pi/dl. By default 2.0 (k_c = pi, the paper’s bulk-water setting; its dimer/NaCl runs used dl = 3). It applies to periodic structures only: a structure without a cell takes the real-space branch, which is exact and reads sigma but never dl, so on an all-molecular dataset this knob is inert (and is therefore untested by the fit, whatever it is set to). The cutoff is recomputed from dl on every call, so it can be retuned on a loaded model.

• sigma (float, optional) – Gaussian smearing width in Angstrom, by default 1.0 (paper Methods: values between ~0.5 and 2 are reasonable; 1 was best for water).

• exponent (int, optional) – Interaction exponent p of 1/r^p: 1 (electrostatics, default) or 6 (London dispersion, paper eq 5).

• remove_self_interaction (bool, optional) – Subtract the Gaussian self energy sum q^2 / (sigma (2 pi)^{3/2}) from the reciprocal sum, by default False (the reference training scripts keep it; the term is short-ranged and can be absorbed by the short-range model either way). Note: for multi-channel q upstream subtracts the total sum q^2 once per channel (an n_channels-fold over-subtraction); xnn subtracts it once. The two agree for 1-dimensional q and always when the flag is off.

• method (str, optional) – "kspace" (default): the reference reciprocal sum over the wave-vector sphere (or its factorized fast path); "pme": the same sum through a particle mesh (reciprocal_pme()), N log N in the atoms, approximate to the spline aliasing error.

• pme_order (int, float, optional) – B-spline order (8) and mesh density relative to the sphere (2.0) of the mesh sum; at these defaults the energy agrees with the reference sum to about 1e-7 relative and the forces and stress to about 1e-6.

• pme_oversample (int, float, optional) – B-spline order (8) and mesh density relative to the sphere (2.0) of the mesh sum; at these defaults the energy agrees with the reference sum to about 1e-7 relative and the forces and stress to about 1e-6.

Notes

Charges are in scaled units (upstream norm_factor = 1): a physical charge Q in e corresponds to q = Q * LATENT_CHARGE_PER_E (sqrt(2 pi * 14.3996) = 9.5118) for energies in eV and distances in Angstrom; see LATENT_CHARGE_PER_E.

k_set(cell)

The wave vectors of the reciprocal sum of one periodic structure.

Parameters:

cell (Tensor) – Row-vector cell matrix, shape (3, 3) (triclinic allowed).

Returns:

• grid (Tensor) – Integer triples (M, 3) (long) of one half space within the spherical cutoff |k| <= 2 pi / dl.

• kpts (Tensor) – The wave vectors k = grid @ B, (M, 3).

• k2 (Tensor) – |k|^2, (M,).

Return type:

Tuple[Tensor, Tensor, Tensor]

reciprocal(pos, q, cell)

Reciprocal-space Ewald energy of one periodic structure.

Parameters:

• pos (Tensor) – Cartesian positions, shape (n, 3).

• q (Tensor) – Hidden variable, shape (n, n_channels).

• cell (Tensor) – Row-vector cell matrix, shape (3, 3) (triclinic allowed).

Returns:

Scalar long-range energy (summed over channels).

Return type:

Tensor

pme_mesh(cell)

Mesh dimensions of reciprocal_pme(): the smallest FFT-friendly sizes (factors 2, 3 and 5) holding the wave-vector sphere |k| <= 2 pi / dl pme_oversample times over.

Parameters:

cell (Tensor)

Return type:

List[int]

reciprocal_pme(pos, q, cell)

reciprocal() through a particle mesh (smooth PME).

The same wave vectors (the sphere |k| <= 2 pi / dl) and weights, with the structure factors from B-spline-spread charges on a mesh (pme_weighted_sum()), so the cost grows as N log N instead of N times the number of vectors. It approximates the reference sum to the spline aliasing error, set by pme_order and pme_oversample.

Parameters:

• pos (Tensor)

• q (Tensor)

• cell (Tensor)

Return type:

Tensor

fast_supported(device, dtype)

The factorized reciprocal sum runs on CUDA devices.

Parameters:

• device (device)

• dtype (dtype)

Return type:

bool

reciprocal_fast(pos, q, cell)

reciprocal() through factorized phases (the fast path).

The same wave vectors and weights; the structure factors of each column of fixed (m_1, m_2) come from one complex matrix product over tabulated phase factors (reciprocal) instead of (n, M) sines and cosines, in column blocks that are recomputed in the backward pass. Differentiable to every order that reciprocal() is.

Parameters:

• pos (Tensor)

• q (Tensor)

• cell (Tensor)

Return type:

Tensor

coulomb_matrix(pos, cell=None)

gamma with E_lr(q) = 1/2 q^T gamma q for one structure.

Differentiable in pos and cell. For a cluster it is the erf(r / sqrt(2) sigma) / r kernel over 2 pi (twice the pair term of realspace(), which sums every ordered pair over 4 pi) plus the Gaussian self term on the diagonal; for a cell it is the Hessian of reciprocal(), 4/V sum_k f(k) cos k.(r_i - r_j) over the half grid, with k = 0 omitted. exponent must be 1.

Parameters:

• pos (Tensor) – Positions (n, 3) in Angstrom.

• cell (Tensor or None, optional) – Row-vector cell (3, 3) for a periodic structure.

Returns:

(n, n), symmetric positive definite.

Return type:

Tensor

realspace(pos, q)

Direct-sum equivalent for a non-periodic structure.

The pair interaction is erf(r / (sqrt(2) sigma)) / r – the potential of the Gaussian-smeared charges – normalized by 1/(4 pi) exactly as upstream, so periodic and molecular structures share the same energy scale. Only exponent = 1 is supported (as upstream).

Parameters:

• pos (Tensor)

• q (Tensor)

Return type:

Tensor

realspace_batch(pos, q, batch, num_graphs)

realspace() of every structure of a molecular batch at once.

The same sum over the pairs of each structure, taken over all of them in one pass instead of a structure at a time: a training batch of molecules then costs a few kernels instead of a loop with a host synchronization per structure. The i == j terms, which vanish identically in realspace() (erf(0) = 0), are left out.

Parameters:

• pos (Tensor) – Positions (N, 3).

• q (Tensor) – Hidden variable (N, n_channels).

• batch (Tensor) – Structure index of each atom (N,), in any order.

• num_graphs (int) – Number of structures B (those without atoms get zero).

Returns:

Energies (B,).

Return type:

Tensor

realspace_single(pos, q)

realspace() of one structure, with finite second derivatives.

The same pair sum as realspace(), contracted as q^T K q per channel instead of through an (N, N, n_channels) coupling tensor. The i == j distances, whose terms vanish in realspace(), are replaced before the square root, whose second derivative is not finite at zero. forward() takes it for a batch of one molecule (molecular dynamics, optimization), where the pair lists of realspace_batch() would cost more memory.

Parameters:

• pos (Tensor) – Positions (N, 3).

• q (Tensor) – Hidden variable (N, n_channels).

Returns:

The energy, a scalar.

Return type:

Tensor

forward(q, pos, batch, num_graphs, cell, pbc=None)

Long-range energy per structure for a batched graph.

Parameters:

• q (Tensor) – Hidden variable, shape (N,) or (N, n_channels).

• pos (Tensor) – Cartesian positions, shape (N, 3).

• batch (Tensor) – Structure index of each atom, shape (N,).

• num_graphs (int) – Number of structures B in the batch.

• cell (Tensor or None) – Cells of shape (B, 3, 3), or None for molecular batches.

• pbc (Tensor or None, optional) – Per-structure periodic flags (B, 3); a structure is treated as periodic when its cell is nonzero and any flag is set.

Returns:

Long-range energies, shape (B,).

Return type:

Tensor