xnn.hybrid.models.spookynet#

SpookyNet (Unke et al., Nat. Commun. 2021): electronic states and nonlocal effects.

An implementation of the SpookyNet potential of

  • O. T. Unke, S. Chmiela, M. Gastegger, K. T. Schuett, H. E. Sauceda, K.-R. Mueller, “SpookyNet: Learning force fields with electronic degrees of freedom and nonlocal effects”, Nat. Commun. 12, 7273 (2021),

written from the paper’s Methods section on the xnn abstractions (scatter_sum(), MollifierCutoff, ExponentialBernsteinRBF, the D4 tables of d4, the Ewald sums of electrostatics) and consistent with the behavior of the authors’ reference code (github.com/OUnke/SpookyNet), without copying it.

It belongs to the hybrid family: graph message passing combined with a transformer (self-attention) update and a physics-based split of the energy, like BAMBOO.

Inputs. Besides atomic numbers and positions the model reads the total charge Q (AtomicGraph.total_charge, neutral when absent) and the number of unpaired electrons S = spin_multiplicity - 1 (AtomicGraph.spin_multiplicity, singlet when absent).

Representation (eqs 1-3, 9-11). Every atom starts from x0 = e_Z + e_Q + e_S. The nuclear embedding e_Z = M d_Z + e~_Z maps a 20-entry descriptor of the ground-state electron configuration (electron_configurations()) through a learned matrix and adds a free element vector. The electronic embeddings spread Q and S over the atoms by an attention-like weighting (eq 10):

q_i = linear(e_Z,i),  a_i = softplus(q_i . k / sqrt(F)) / sum_j softplus(q_j . k / sqrt(F)),
e_Psi,i = resmlp(a_i |Psi| v)

with separate keys and values k, v for positive and negative charges. T interaction modules then refine the features (eq 11):

x~ = residual(x),  x' = residual(x~ + l + n),  y = resmlp(x')

and the atomic descriptors are f = sum_t y_t. Here residual is the pre-activation residual block x + W2 silu(W1 silu(x)) (eq 7), resmlp a residual block followed by an activation and a linear layer (eq 8) and silu the generalized SiLU alpha x sigmoid(beta x) with learned per-feature alpha and beta (eq 6, Swish).

Local interaction (eqs 12-17). Within the cutoff r_c every pair contributes through basis functions that look like s, p and d orbitals, g_lm(r) = rho_k(r) Y_lm(r / |r|) with the exponential Bernstein radial functions rho_k (times the cutoff exp(-r^2 / (r_c^2 - r^2))) and the real spherical harmonics without their normalization constant:

s_i = sum_j resmlp_s(x~_j) o G_s rho(r_ij)
p_i = sum_j resmlp_p(x~_j) o G_p rho(r_ij) Y_1(r_ij)          (3 x F)
d_i = sum_j resmlp_d(x~_j) o G_d rho(r_ij) Y_2(r_ij)          (5 x F)
l_i = resmlp_l(resmlp_c(x~_i) + s_i + <P1 p_i, P2 p_i> + <D1 d_i, D2 d_i>)

so angular information enters at linear cost in the number of neighbors.

Nonlocal interaction (eqs 18-20). Self-attention over all atoms of a structure, n = attention(resmlp_q(x~), resmlp_k(x~), resmlp_v(x~)), evaluated with the linear-scaling FAVOR+ approximation (positive orthogonal random features drawn once at initialization, f = F of them) or exactly (attention="exact", eq 19).

Energy (eqs 4, 21-27). Atomic energies and partial charges come from one linear layer on f plus element biases; the charges are shifted so that they sum to Q (eq 24). Three physical terms are added: a ZBL-like nuclear repulsion with learnable, positive parameters (eq 22), point-charge electrostatics that switch from 1 / sqrt(r^2 + 1) to 1 / r between r_c / 4 and 3 r_c / 4 (eq 23, 25) and the two-body D4 dispersion with the model’s charges, learnable s8, a1, a2 and a learnable scale s_q of the reference charges (eq 27). The dipole moment is sum_i q_i r_i (eq 26).

Conventions worth knowing:

  • The electronic embedding follows the reference code: the value is |Psi| v_+- (the paper writes Psi v_+-, the same model with v_- negated), the keys use the sign pattern of Psi (e / max(e, 1)) and the weights are normalized with +1e-8; k and v start orthogonal, since zero values would leave the branch without a gradient.

  • The radial functions are ordered by the paper’s k (the power of exp(-gamma r)); the reference code stores them in reverse order. Its extra weighting exp(-gamma r) is the exp_weighting option.

  • FAVOR+ subtracts, for stability, the largest key projection of the whole structure; the reference code’s batched path takes the largest over a zero-padded matrix instead, which differs (by the 1e-4 stabilizer of the features) only when every projection of a structure is negative. Here a structure’s energy never depends on the other structures of a batch.

  • The D4 reference polarizabilities are recomputed from s_q in every evaluation (differentiable), with the dftd4 reference data shipped with d4; physical constants are CODATA 2018 (the attributes k_e, bohr and hartree of the term modules).

  • Without lr_cutoff the electrostatics and dispersion run over all pairs of every structure (molecules). With lr_cutoff they use the neighbor list with the reference code’s shifted kernels, which also serves periodic structures; ewald=True replaces the shifted Coulomb kernel by the Ewald sum of the periodic structures (the bare 1 / r lattice sum plus the short-range damping of eq 23).

Module Attributes

MAX_Z

the electron-configuration descriptor and the reference models cover the elements up to Rn.

Functions

electron_configurations([max_z])

Element descriptors d_Z of the nuclear embedding (paper eq 9).

orthogonal_random_features(n_features, dim)

Projection matrix of positive orthogonal random features (FAVOR+).

reference_state_dict(sd)

Translate a reference-code state dict to SpookyNet parameter names.

smooth_switch(r, r_on, r_off)

Switch from 1 (r <= r_on) to 0 (r >= r_off), smooth to every order (eq 25).

Classes

SpookyNet([n_features, n_interactions, ...])

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

Swish(n_features[, alpha, beta])

Generalized SiLU alpha x sigmoid(beta x) with learned per-feature alpha, beta (eq 6).
Show Source