xnn.common.models.les.EwaldSummation#

class xnn.common.models.les.EwaldSummation(dl=2.0, sigma=1.0, exponent=1, remove_self_interaction=False)[source]#

Bases: Module

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.

Notes

Charges are in scaled units (upstream norm_factor = 1): a physical charge Q in e corresponds to q = Q * sqrt(90.0474) for energies in eV and distances in Angstrom.

reciprocal(pos, q, cell)[source]#

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

realspace(pos, q)[source]#

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:
Return type:

Tensor

forward(q, pos, batch, num_graphs, cell, pbc=None)[source]#

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