xnn.common.models.les.LatentEwald#

class xnn.common.models.les.LatentEwald(model, n_channels=4, hidden=None, q_bias=False, q_add_linear=True, dl=2.0, sigma=1.0, exponent=1, remove_self_interaction=False, constrain_charge=False, charge_channel=0, charge_weights='uniform', charge_solve=False, hardness='element', hardness_init=10.0, fragments=False, bond_factor=1.2, ion_charges=None, method='kspace', pme_order=8, pme_oversample=2.0, dipole_origin='com', dielectric=1.0, learn_dielectric=False)[source]#

Bases: InteratomicPotential

Wrap any xnn model with a Latent-Ewald long-range energy (CACE-LR).

The wrapped model must expose invariant per-atom features through the "node_features" key of its output dict and a node_feature_dim attribute – every built-in xnn model does. A bias-free MLP plus a parallel bias-free linear layer (the head used by the reference cace-lr-fit scripts) maps the features to the hidden variable q (paper eq 2), and EwaldSummation turns q into the long-range energy added to the model’s short-range prediction.

Enable from a config with model.extra["long_range"], e.g.:

extra: {..., long_range: {n_channels: 4, sigma: 1.0, dl: 2.0}}
Parameters:

  • model (InteratomicPotential) – The short-range model to wrap.

  • n_channels (int, optional) – Dimension of the hidden variable q, by default 4 (the paper’s bulk-water/NaCl setting; the reference charged-dimer script uses 1).

  • hidden (list of int, optional) – Hidden widths of the q MLP, by default [24, 12] (upstream).

  • q_bias (bool, optional) – Use a bias in the q MLP, by default False (the water script’s head). The reference charged-dimer script uses True – the bias lets the latent charge carry a per-structure offset, which matters for net-charged systems like ionic dimers.

  • q_add_linear (bool, optional) – Add a parallel bias-free linear layer to the q head, by default True (the water script). The charged-dimer script uses False (MLP only).

  • constrain_charge (bool, optional) – Pin each structure’s latent charges to its net charge after the q head: the channel charge_channel sums to total_charge * LATENT_CHARGE_PER_E (the batch’s total_charge, neutral when absent) and every other channel to zero, by a per-structure shift q_i <- q_i - w_i (sum q - Q') / sum w (the correction of PhysNet). The sum is then exact for any system size, where a bias in the q head (q_bias) only learns an offset per atom, and in the molecular path the monopole energy of charged structures follows. Off by default (the reference behavior). Charged training structures must carry their total_charge label, or they are forced neutral.

  • charge_channel (int, optional) – The channel that carries the physical net charge, by default 0.

  • charge_weights (str, optional) – "uniform" (default): the same shift for every atom; "learned": w_i = softplus(linear(features_i)), a learned per-atom capacity for the residual.

  • charge_solve (bool, optional) – Global charge solve for channel charge_channel (charge_solve): the head’s output of that channel is an electronegativity chi_i and the charges minimise chi.q + 1/2 sum J q^2 + E_lr(q) with sum q = total_charge (in latent units), so every charge responds to every other one through the Ewald kernel. Adds chi.q + 1/2 J q^2 to the energy (E_lr is the usual one). Off by default. With the coupling it contains, constrain_charge is its coupling-free limit.

  • hardness (str, optional) – Where the hardness J_i > 0 of the charge solve comes from: "element" (default), one learned value per element, or "features", softplus(linear(features_i)).

  • hardness_init (float, optional) – Initial hardness in eV per e^2 (converted to latent units), by default 10.0, about the atomic values of charge equilibration.

  • fragments (bool, optional) – One constraint row per fragment instead of one per structure, so no charge flows between separate molecules (plain charge equilibration is metallic at long range). Fragments are the covalently connected components (bonded_fragments()), with the ions of ion_charges always on their own; each fragment’s charge comes from the structure’s fragment_charges label (the charge of the fragment each atom belongs to) when present, else from ion_charges (neutral otherwise), and must add up to total_charge. Off by default.

  • bond_factor (float, optional) – Bond criterion: distance below bond_factor times the sum of the covalent radii, by default 1.2.

  • ion_charges (dict, optional) – Formal charges by element symbol or atomic number, by default ION_CHARGES (the alkali, alkaline-earth and halide ions).

  • dipole_origin (str, optional) – Origin of the "les_dipole" output: "com" (default) the center of mass, "centroid" the mean position, "origin" the coordinate origin. The choice only matters for charged structures.

  • dielectric (float, optional) – Screening of the charge channel’s long-range energy: its Ewald energy (and, with the charge solve, its Coulomb coupling) is divided by dielectric, by default 1 (vacuum). Set it to the solvent’s constant (78.4 for water) when the labels come from a continuum solvent model, so charges fitted to physical ones carry the screened interaction those labels contain. The other channels are unscreened.

  • learn_dielectric (bool, optional) – Make the dielectric a trainable parameter (its logarithm, so it stays positive), starting from dielectric. Off by default.

  • dl (float) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

  • sigma (float) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

  • exponent (int) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

  • remove_self_interaction (bool) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

  • method (str) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

  • pme_order (int) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

  • pme_oversample (float) – Passed to EwaldSummation (method: pme runs the periodic sum on a particle mesh). Note that dl acts on periodic structures only and exponent = 6 is periodic-only as well (the real-space branch that molecular structures take implements the 1/r kernel), so a dataset without cells trains neither.

Variables:

  • model (InteratomicPotential) – The wrapped short-range model.

  • q_net (torch.nn.Module) – The latent-charge MLP.

  • q_linear (torch.nn.Module or None) – The optional parallel linear layer; q = q_net(B) [+ q_linear(B)].

  • ewald (EwaldSummation) – The long-range energy module.

  • cutoff (float) – The wrapped model’s neighbor-list cutoff (proxied).

Notes

forward returns the combined "energy"; the long-range energy is spread uniformly over the atoms of each structure in "node_energy" so it still sums to the total. The additional keys "energy_sr", "energy_lr", and "latent_charges" expose the decomposition.

"les_charges" (N,) is the channel charge_channel in e (after the constraint or the charge solve) and "les_dipole" (B, 3) its dipole moment sum_i Q_i (r_i - o) in e Angstrom about the dipole_origin, for the molecular structures (zero for periodic ones, whose dipole is undefined). The loss terms les_charge_weight and les_dipole_weight fit them to charges and dipole labels; the keys are distinct from the "charges" / "dipole" outputs of models with charges of their own. With several channels only that channel is exported. Both report the charges before the dielectric screening, which acts on the energy only.

property dielectric: Tensor#

The current dielectric constant of the charge channel (a scalar tensor).

screened(q)[source]#

q with the charge channel divided by sqrt(dielectric), so the Ewald energy of that channel is divided by the dielectric.

Parameters:

q (Tensor)

Return type:

Tensor

set_use_fast(use_fast)[source]#

Choose the implementation of the reciprocal sum.

Parameters:

use_fast (bool or str) – "auto" (default): the factorized reciprocal sum (EwaldSummation.reciprocal_fast()) on a CUDA device for the periodic structures of at least AUTO_POLICY atoms (3000); True: wherever it runs; False: always the reference. The wrapped model’s own fast paths are set by set_use_fast(), which reaches both.

Returns:

self.

Return type:

LatentEwald

forward(data)[source]#

Short-range prediction plus the latent-Ewald long-range energy.

Parameters:

data (AtomicGraph) – The batched atomic graph.

Returns:

The wrapped model’s outputs with "energy" and "node_energy" including the long-range term, plus "energy_sr" (B,), "energy_lr" (B,), "latent_charges" (N, n_channels), "les_charges" (N,) and "les_dipole" (B, 3).

Return type:

dict of str to torch.Tensor

dipole(charges, pos, data, cell)[source]#

Dipole moment sum_i Q_i (r_i - o) of every structure, (B, 3).

Parameters:

  • charges (Tensor) – Per-atom charges (N,) in e.

  • pos (Tensor) – Positions (N, 3) in the charges’ dtype.

  • data (AtomicGraph) – The batch (structure index, atomic numbers, periodic flags).

  • cell (Tensor or None) – The cells (B, 3, 3), or None for a molecular batch.

Returns:

The dipole moments in e Angstrom about the dipole_origin, zero for periodic structures.

Return type:

Tensor

hardness_of(features, atomic_numbers)[source]#

The hardness J_i > 0 of every atom, (N,), in latent units.

Parameters:
Return type:

Tensor

solve_charges(chi, q, features, data, pos, cell)[source]#

Replace channel charge_channel of q by the solved charges.

One augmented solve per structure (solve_charges() with the Coulomb matrix of its own kernel: the direct sum for a cluster, the reciprocal sum for a cell), with one constraint row per structure or, with fragments, per fragment.

Returns:

The charges (N, n_channels), the per-atom energy chi_i q_i + 1/2 J_i q_i^2 (N,), the chemical potentials (one per constraint row), the hardness (N,) and the fragment index of every atom (N,) (None without fragments).

Return type:

(Tensor, Tensor, Tensor, Tensor, Tensor or None)

Parameters:
constrain(q, features, batch, num_graphs, total_charge)[source]#

Shift the latent charges so each structure sums to its target.

Parameters:

  • q (Tensor) – Latent charges, (N, n_channels).

  • features (Tensor) – Per-atom features, (N, F) (the learned weights read them).

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

  • num_graphs (int) – Number of structures.

  • total_charge (Tensor or None) – Net charge of every structure in e, (B,); None is neutral.

Returns:

q with channel charge_channel summing to total_charge * LATENT_CHARGE_PER_E and the others to zero.

Return type:

Tensor

classmethod from_config(cfg)[source]#

Not registered directly; built by build_model via model.extra['long_range'].

Return type:

LatentEwald