xnn.gnn.models.dimenet.DimeNet#

class xnn.gnn.models.dimenet.DimeNet(n_features=128, n_interactions=6, n_rbf=6, n_spherical=7, cutoff=5.0, n_bilinear=8, n_output_features=None, p=6, n_before_skip=1, n_after_skip=2, n_output_layers=3, activation='silu', output_init='glorot_orthogonal', trainable_rbf=True, radial_gate='kj', reference_basis=False, species=None, atomic_energies=None)[source]#

Bases: InteratomicPotential

DimeNet: directional message passing with the bilinear interaction (ICLR 2020).

See the module docstring for the equation-by-equation walk-through. The defaults are the paper’s QM9/MD17 architecture (Appendix B): F = 128, six interaction blocks, N_SHBF = 7, N_SRBF = N_RBF = 6, N_bilinear = 8, cutoff 5 Angstrom, envelope exponent p = 6, one residual block before and two after the skip connection, three dense layers in the output blocks. Works for molecules and periodic systems (periodicity enters through the edge vectors).

Parameters:

  • n_features (int, optional) – Embedding and message width F, by default 128.

  • n_interactions (int, optional) – Number of interaction blocks T, by default 6.

  • n_rbf (int, optional) – Number of radial Bessel functions N_RBF (also the zeros per degree N_SRBF of the 2D basis, as in the paper), by default 6.

  • n_spherical (int, optional) – Number of spherical-harmonic degrees N_SHBF of the 2D basis, by default 7.

  • cutoff (float, optional) – Cutoff c in Angstrom, by default 5.0.

  • n_bilinear (int, optional) – Inner dimension of the bilinear layer, by default 8.

  • n_output_features (int or None, optional) – Width of the atom features in the output blocks; None (default, the paper) keeps F without an up-projection. DimeNet++ sets it.

  • p (int, optional) – Exponent of the envelope u(d) of eq 8, by default 6 (the reference code’s envelope_exponent + 1).

  • n_before_skip (int, optional) – Residual blocks before and after the skip connection of each interaction block, by default 1 and 2.

  • n_after_skip (int, optional) – Residual blocks before and after the skip connection of each interaction block, by default 1 and 2.

  • n_output_layers (int, optional) – Dense layers of each output block, by default 3.

  • activation (str or torch.nn.Module, optional) – Activation, by default "silu" (the paper’s swish); see make_activation().

  • output_init (str, optional) – Initialization of the final layer of the output blocks: "glorot_orthogonal" (default, the reference choice for energies) or "zeros" (the fresh model then predicts exactly atom_ref).

  • trainable_rbf (bool, optional) – Fine-tune the Bessel frequencies n pi by backpropagation (section 5), by default True.

  • radial_gate (str, optional) – Which distance gates the incoming message m_kj in the interaction block: "kj" (default) is the reference implementation, where the gate is W e_RBF(d_kj); "ji" is eq 4 and Fig. 4 as printed, W e_RBF(d_ji). The published models were trained with "kj".

  • reference_basis (bool, optional) – Reproduce the 2D basis of the reference implementation, whose envelope carries an extra factor c / d_kj relative to eq 6 (its radial and 2D bases share one envelope that includes the 1 / d of eq 7). False (default) is the paper’s u(d) a_SBF.

  • 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 (aligned with species).

Variables:

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

  • rbf (BesselRBF) – Radial basis of eq 7 (trainable frequencies).

  • sbf (SphericalBesselBasis) – 2D spherical Fourier-Bessel basis of eq 6.

  • envelope (PolynomialCutoff) – The envelope u(d) of eq 8.

  • embedding (_Embedding) – The embedding block (holds the atom-type embeddings).

  • interactions (torch.nn.ModuleList) – The T interaction blocks.

  • outputs (torch.nn.ModuleList) – The T + 1 output blocks (one after the embedding block).

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

  • node_feature_dim (int) – Width of the invariant per-atom features returned as "node_features" (the summed output-block features before their final layer; what e.g. the LES charges are read from).

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

Parameters:

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

  • edge_index (Tensor) – Edge index of shape (2, E); row 0 is the source j and row 1 the destination i of each edge (the message m_ji).

  • edge_vec (Tensor) – Edge vectors x_i - x_j, shape (E, 3) (cell shifts included).

Returns:

The invariant per-atom features (N, node_feature_dim) (the output-block features before their final layer, summed over the blocks) and the per-atom energy (N,), sum_l t_i^(l) + atom_ref[Z_i].

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:

"node_energy" (per-atom energy, shape (N,)), "energy" (per-structure total energy, the sum over atoms) and "node_features" (invariant per-atom features, shape (N, node_feature_dim)).

Return type:

dict[str, Tensor]

classmethod from_config(cfg)[source]#

Build a DimeNet from a configuration object.

Core fields: cfg.n_features -> F, cfg.n_interactions -> T, cfg.n_rbf -> N_RBF = N_SRBF, cfg.cutoff -> c. cfg.extra may hold n_spherical, n_bilinear, n_output_features, p (or the reference code’s envelope_exponent = p - 1), n_before_skip, n_after_skip, n_output_layers, activation, output_init, trainable_rbf, reference_basis, species and atomic_energies. The reference code’s key spellings (emb_size, num_blocks, …) are translated by xnn.common.config.translate.

Parameters:

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

Returns:

Instantiated model.

Return type:

DimeNet