xnn.hybrid.models.spookynet.SpookyNet#

class xnn.hybrid.models.spookynet.SpookyNet(n_features=128, n_interactions=6, n_rbf=16, cutoff=5.291772109044924, lr_cutoff=None, n_residual=1, charge_embedding=True, spin_embedding=True, nonlocal_interactions=True, attention='favor', n_random_features=None, zbl_repulsion=True, electrostatics=True, d4_dispersion=True, exp_weighting=False, ewald=False, ewald_accuracy=1e-06, seed=None, species=None, atomic_energies=None)[source]#

Bases: InteratomicPotential

SpookyNet potential with charge/spin embeddings, nonlocal attention and physical terms.

The defaults are the paper’s architecture: F = 128 features, T = 6 interaction modules, K = 16 exponential Bernstein radial functions, a cutoff of 10 bohr (5.29 Angstrom), FAVOR+ attention with F random features, and the repulsion, electrostatics and D4 terms. See the module docstring for the equations.

Parameters:

  • n_features (int, optional) – Feature width F, by default 128.

  • n_interactions (int, optional) – Interaction modules T, by default 6.

  • n_rbf (int, optional) – Radial functions K, by default 16.

  • cutoff (float, optional) – Local cutoff r_c (Angstrom), by default 10 bohr.

  • lr_cutoff (float or None, optional) – Cutoff of the electrostatics and dispersion; None (default) sums over all pairs of every (molecular) structure. Required for periodic structures. The neighbor-list radius cutoff attribute is the larger of the two.

  • n_residual (int, optional) – Residual blocks in every residual stack and residual MLP, by default 1.

  • charge_embedding (bool, optional) – Include e_Q and e_S (eq 1), by default True.

  • spin_embedding (bool, optional) – Include e_Q and e_S (eq 1), by default True.

  • nonlocal_interactions (bool, optional) – Include the attention blocks (eq 18), by default True.

  • attention (str, optional) – "favor" (default; FAVOR+, eq 20) or "exact" (eq 19).

  • n_random_features (int or None, optional) – FAVOR+ features, by default n_features.

  • zbl_repulsion (bool, optional) – Include the physical terms of eq 4, by default True.

  • electrostatics (bool, optional) – Include the physical terms of eq 4, by default True.

  • d4_dispersion (bool, optional) – Include the physical terms of eq 4, by default True.

  • exp_weighting (bool, optional) – Weight the radial functions by exp(-gamma r) (reference-code option), by default False.

  • ewald (bool, optional) – Ewald-sum the electrostatics of periodic structures (needs lr_cutoff, the real-space cutoff), by default False.

  • ewald_accuracy (float, optional) – Relative accuracy of the Ewald sum, by default 1e-6.

  • seed (int or None, optional) – Seed of the FAVOR+ random features (None: the global generator).

  • species (list of int or None, optional) – Only used to interpret atomic_energies.

  • atomic_energies (array-like or None, optional) – Per-species energies loaded into atom_ref (the element energy biases of eq 21).

Variables:

  • cutoff (float) – Neighbor-list radius.

  • local_cutoff (float) – The cutoff of the local interactions and the repulsion.

  • charge_ref (atom_ref,) – Element energy and charge biases (eqs 21, 24), starting at zero.

  • output (torch.nn.Linear) – Bias-free map of f to the atomic energy and charge.

  • use_nonlocal (use_p_orbitals, use_d_orbitals,) – Evaluation switches (all True): setting one to False drops the p- or d-like terms of eq 12 or the nonlocal features n of eq 11, the decompositions of the paper’s Fig. 2.

set_atomic_energies(species, values)[source]#

Set the element energy biases atom_ref of species to values.

Raises:

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

Return type:

None

node_features_energy_charges(atomic_numbers, edge_index, edge_vec, pos, cell, pbc, total_charge, spin_multiplicity)[source]#

TorchScript core of one structure: descriptors, atomic energies and charges.

Parameters:

  • atomic_numbers (Tensor) – (N,).

  • edge_index (Tensor) – (2, E) within self.cutoff.

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

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

  • cell (Tensor) – (3, 3) (ignored unless pbc has a true entry).

  • pbc (Tensor) – (3,) bool.

  • total_charge (float) – Net charge of the structure.

  • spin_multiplicity (float) – Its spin multiplicity 2S + 1.

Returns:

node_features (N, F), node_energy (N,) (all terms) and charges (N,).

Return type:

tuple of Tensor

forward(data)[source]#

Energies, charges and the dipole moment of a (batched) graph.

Parameters:

data (AtomicGraph) – Batched graph built with self.cutoff. total_charge (None: neutral) and spin_multiplicity (None: singlet) set the electronic state.

Returns:

"node_energy" (N,), "energy" (B,), "node_features" (the descriptors f, (N, F)), "charges" (N,), "dipole" (B, 3) (sum q_i r_i, e Angstrom) and the per-structure terms "energy_repulsion", "energy_electrostatics" and "energy_dispersion".

Return type:

dict of str to Tensor

classmethod from_config(cfg)[source]#

Build a SpookyNet from a configuration object.

Core fields: cfg.n_features, cfg.n_interactions, cfg.n_rbf, cfg.cutoff. The other constructor arguments are read from cfg.extra; the reference code’s spellings (num_features, num_modules, use_zbl_repulsion, …) are translated by xnn.common.config.translate.

Parameters:

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

Returns:

Instantiated model.

Return type:

SpookyNet

classmethod from_reference_checkpoint(source)[source]#

Load a model saved by the reference code (SpookyNet.save).

Reads the hyperparameters and the state dict of the reference format (the current layout: linear electronic embeddings, irreducible d functions, exponential Bernstein radial basis, swish activations), renames the parameters and reverses the order of the radial functions. The model is returned in eval mode, in float32 or float64 as stored.

Parameters:

source (str, pathlib.Path or dict) – Path of the .pth file, or its loaded content.

Returns:

The model.

Return type:

SpookyNet

Raises:

ValueError – For a checkpoint layout this implementation does not cover.