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)[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).

  • dl (float) – Passed to EwaldSummation. 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. 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. 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. 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.

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,) and "latent_charges" (N, n_channels).

Return type:

dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

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

Return type:

LatentEwald