xnn.cnn.models.schnet.SchNet#

class xnn.cnn.models.schnet.SchNet(n_features=64, n_interactions=3, n_rbf=301, cutoff=30.0, gamma=10.0, cutoff_fn=None, energy_shift=0.0, energy_scale=1.0, species=None, atomic_energies=None)[source]#

Bases: InteratomicPotential

SchNet continuous-filter convolutional interatomic potential.

Faithful to the NIPS 2017 manuscript (see the module docstring for the equation-by-equation walk-through): embeds atoms by nuclear charge, refines their features through T residual interaction blocks built around the continuous-filter convolution, and reads out a per-atom energy through a two-layer atom-wise network with shifted-softplus nonlinearities. Per-atom energies are standardized with the training-set statistics (E_i = energy_scale * E^hat_i + energy_shift, the DTNN convention) plus a per-element reference atom_ref, then sum-pooled into the total energy. Works unchanged for molecules and periodic solids, since periodicity enters only through the edge vectors.

The defaults reproduce the paper’s architecture: F = 64 feature maps, T = 3 interaction blocks, and Gaussian RBFs on a 0.1-Angstrom grid from 0 to 30 Angstrom with gamma = 10 (301 centers – the grid the paper states as “centers 0 <= mu_k <= 30 every 0.1 Angstrom”).

Parameters:

  • n_features (int, optional) – Dimension F of the per-atom feature vectors, by default 64 (the paper’s value; kept constant through the interaction blocks).

  • n_interactions (int, optional) – Number of stacked interaction blocks T (no weight sharing), by default 3.

  • n_rbf (int, optional) – Number of Gaussian radial basis functions expanding the distances, by default 301 (0.1-Angstrom spacing on [0, 30]).

  • cutoff (float, optional) – Neighbor-list radius and upper end of the RBF center grid (Angstrom), by default 30.0. The paper uses no explicit cutoff; 30 Angstrom covers all pairs of its molecular datasets. For condensed phases use a finite cutoff (e.g. 5.0) with n_rbf ~ cutoff / 0.1 and cutoff_fn="cosine".

  • gamma (float or None, optional) – Width parameter of the Gaussian RBFs, by default 10.0 (the paper’s value, per Angstrom^2). None ties the width to the center spacing instead (see GaussianRBF).

  • cutoff_fn (str or None, optional) – "cosine" multiplies the generated filters by a smooth CosineCutoff envelope (recommended with finite cutoffs); None (default) is the paper’s unmodulated filter.

  • energy_shift (float, optional) – Additive per-atom energy standardization E_mu (the training-set mean energy per atom, DTNN Methods step 4), by default 0.0. Stored as a non-trainable buffer; see set_energy_scale_shift().

  • energy_scale (float, optional) – Multiplicative per-atom energy standardization E_sigma (the training-set standard deviation of the energy per atom), by default 1.0.

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

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

Variables:

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

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

  • rbf (GaussianRBF) – Gaussian radial basis expansion of interatomic distances.

  • interactions (torch.nn.ModuleList) – Stack of _Interaction blocks.

  • readout (torch.nn.Sequential) – Atom-wise MLP (F -> F/2 -> 1) mapping final features to the unstandardized per-atom energy E^hat_i; its last layer is zero-initialized (the DTNN convention) so initial predictions equal energy_shift + atom_ref.

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

set_energy_scale_shift(scale, shift)[source]#

Set the DTNN per-atom energy standardization from training stats.

Parameters:

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

  • shift (float) – E_mu, the 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

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

TorchScript-compatible core: tensors in, features + energy out.

The single implementation reused by node_energy() (the deploy entry point) and forward(), so it must avoid the AtomicGraph dataclass and any Python-only constructs.

Parameters:

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

  • edge_index (Tensor) – Edge index of shape (2, E); row 0 is the source (neighbor) and row 1 the destination (center) node of each edge.

  • edge_vec (Tensor) – Edge displacement vectors, shape (E, 3) (already accounting for any periodic cell shifts).

Returns:

The invariant node features after the last interaction block (N, node_feature_dim) and the per-atom energy (N,) (standardized and including the per-element reference shift).

Return type:

tuple of Tensor

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

Per-atom energy, shape (N,) (thin wrapper over node_features_energy(); the deploy wrappers call this).

Parameters:
Return type:

Tensor

forward(data)[source]#

Compute per-atom and total energies for a batch of structures.

Parameters:

data (AtomicGraph) – Batched atomic graph providing atomic numbers, edge index and edge vectors.

Returns:

Dictionary with "node_energy" (per-atom energy, shape (N,)), "energy" (per-structure total energy, the sum pooling of Fig. 2) and "node_features" (invariant per-atom features, shape (N, node_feature_dim)).

Return type:

dict[str, Tensor]

classmethod from_config(cfg)[source]#

Build a SchNet from a configuration object.

Core fields: cfg.n_features -> F, cfg.n_interactions -> T, cfg.n_rbf and cfg.cutoff -> the RBF grid / neighbor-list radius. Everything else is read from cfg.extra (gamma, cutoff_fn, energy_shift, energy_scale, species, atomic_energies); schnetpack 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:

SchNet