xnn.gnn.models.painn.PaiNN#

class xnn.gnn.models.painn.PaiNN(n_features=128, n_interactions=3, n_rbf=20, cutoff=5.0, radial_basis='bessel', shared_filters=False, epsilon=1e-08, dipole=False, polarizability=False, n_output_blocks=2, correct_charges=True, atomic_dipoles=True, scalar_product=True, vector_propagation=True, vector_features=True, energy_shift=0.0, energy_scale=1.0, species=None, atomic_energies=None)[source]#

Bases: InteratomicPotential

PaiNN interatomic potential with optional dipole and polarizability heads.

Faithful to the ICML 2021 manuscript (see the module docstring for the equation-by-equation walk-through): embeds atoms by nuclear charge, refines scalar and vector features through T message/update block pairs and reads out per-atom energies through a two-layer network with a SiLU nonlinearity. Per-atom energies are standardized (energy_scale, energy_shift) and shifted by the per-element reference atom_ref, then sum-pooled into the total energy.

The defaults reproduce the paper’s architecture: F = 128 features, T = 3 blocks, 20 Bessel radial functions with a cosine cutoff at 5 Angstrom.

Parameters:

  • n_features (int, optional) – Width F of the scalar and vector features, by default 128.

  • n_interactions (int, optional) – Number of message/update block pairs T, by default 3.

  • n_rbf (int, optional) – Number of radial basis functions, by default 20.

  • cutoff (float, optional) – Neighbor-list radius and cosine cutoff (Angstrom), by default 5.0.

  • radial_basis (str, optional) – "bessel" (default; the paper’s sin(n pi r / r_cut) / r) or "gaussian" (the reference code’s default).

  • shared_filters (bool, optional) – Use one filter-generating layer for all blocks instead of one per block. Default False.

  • epsilon (float, optional) – Stabilizer of the vector norm in the update block, by default 1e-8.

  • dipole (bool, optional) – Add the dipole head of eq 13 (outputs "dipole" (B, 3) and the latent "charges" (N,)). Default False.

  • polarizability (bool, optional) – Add the polarizability head of eq 14 (output "polarizability" (B, 3, 3)). Default False.

  • n_output_blocks (int, optional) – Gated equivariant blocks per tensorial head, by default 2.

  • correct_charges (bool, optional) – Shift the latent charges so they sum to the structure’s net charge (total_charge, zero when absent). Default True.

  • atomic_dipoles (bool, optional) – Include the atomic dipoles mu_i(v_i) of eq 13 in the dipole moment. False keeps the latent charges alone (eq 12), the comparison of the paper’s spectra section. Default True.

  • scalar_product (bool, optional) – Keep the <U v, V v> term of eq 9 (Table IV ablation when off).

  • vector_propagation (bool, optional) – Keep the v_j o phi_vv term of eq 8 (Table IV ablation when off).

  • vector_features (bool, optional) – Keep the vector features at all; False is the invariant model of the last Table IV row (scalar messages of eq 7 and ds = a_ss(s) updates; the tensorial heads then see zero vectors). Ablated terms carry no parameters.

  • energy_shift (float, optional) – Additive per-atom energy standardization, by default 0.0.

  • energy_scale (float, optional) – Multiplicative per-atom energy standardization, by default 1.0.

  • species (list of int or None, optional) – Only used to interpret atomic_energies; the model handles all elements up to Z = 99.

  • atomic_energies (array-like or None, optional) – Per-species reference energies loaded into atom_ref.

Variables:

  • cutoff (float) – Neighbor-list cutoff radius.

  • embedding (torch.nn.Embedding) – Nuclear-charge-to-feature embedding a_Z.

  • rbf (torch.nn.Module) – The radial basis (BesselRBF or GaussianRBF).

  • cutoff_fn (CosineCutoff) – The cosine cutoff multiplying every filter.

  • interactions (torch.nn.ModuleList) – The T message/update block pairs (message and update submodules).

  • readout (torch.nn.Sequential) – Atom-wise energy network F -> F/2 -> 1.

  • atom_ref (torch.nn.Embedding) – Learnable per-element energy reference, initialized to zero.

  • polarizability_head (dipole_head,) – The tensorial heads, when requested.

set_energy_scale_shift(scale, shift)[source]#

Set the per-atom energy standardization from training statistics.

Parameters:

  • scale (float) – Standard deviation of the training-set energy per atom.

  • shift (float) – Mean training-set energy per atom.

Return type:

None

set_atomic_energies(species, values)[source]#

Initialize the per-element reference energies atom_ref.

Parameters:

  • species (list of int) – Atomic numbers the values refer to, in order.

  • values (array-like) – One reference energy per entry of species.

Raises:

ValueError – If the number of values does not match the number of species.

Return type:

None

representation(atomic_numbers, edge_index, edge_vec)[source]#

TorchScript-compatible trunk: the scalar and vector features of every atom.

Parameters:

  • atomic_numbers (Tensor) – Per-atom atomic numbers (N,).

  • edge_index (Tensor) – Edge index (2, E); row 0 the neighbor j, row 1 the center i.

  • edge_vec (Tensor) – Edge vectors pos[i] - pos[j] of shape (E, 3) (with any periodic image shift applied).

Returns:

Scalars s of shape (N, F) and vectors v of shape (N, 3, F) after the last update block.

Return type:

tuple of Tensor

node_features_energy(atomic_numbers, edge_index, edge_vec)[source]#

TorchScript-compatible core: scalar features and per-atom energies.

Parameters:

  • atomic_numbers (Tensor) – Per-atom atomic numbers (N,).

  • edge_index (Tensor) – Edge index (2, E).

  • edge_vec (Tensor) – Edge vectors (E, 3).

Returns:

The scalar features (N, F) and the per-atom energy (N,) (standardized and including the per-element reference).

Return type:

tuple of Tensor

node_energy(atomic_numbers, edge_index, edge_vec)[source]#

Per-atom energy (N,) (the LAMMPS / TorchScript entry point).

Parameters:
Return type:

Tensor

forward(data)[source]#

Energies and, when the heads are present, dipoles and polarizabilities.

Parameters:

data (AtomicGraph) – Batched atomic graph.

Returns:

"node_energy" (N,), "energy" (B,), "node_features" (the scalars, (N, F)), "node_vectors" ((N, 3, F)) and, with the heads, "dipole" (B, 3) with "charges" (N,) and "polarizability" (B, 3, 3).

Return type:

dict of str to Tensor

classmethod from_config(cfg)[source]#

Build a PaiNN from a configuration object.

Core fields: cfg.n_features, cfg.n_interactions, cfg.n_rbf, cfg.cutoff. Everything else is read from cfg.extra (radial_basis, shared_filters, epsilon, dipole, polarizability, n_output_blocks, correct_charges, atomic_dipoles, scalar_product, vector_propagation, vector_features, energy_shift, energy_scale, species, atomic_energies); the reference code’s key spellings are translated by xnn.common.config.translate.

Parameters:

cfg (xnn.common.config.schema.ModelConfig) – The core model config.

Returns:

Instantiated model.

Return type:

PaiNN