xnn.gnn.models.dimenet.DimeNetPP#

class xnn.gnn.models.dimenet.DimeNetPP(n_features=128, n_interactions=4, n_rbf=6, n_spherical=7, cutoff=5.0, n_triplet_features=64, n_basis_features=8, n_output_features=256, 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: DimeNet

DimeNet++: the fast directional message passing of Gasteiger et al. (NeurIPS-W 2020).

Same trunk as DimeNet with the Hadamard interaction, the two-layer basis maps, the triplet down/up-projections and the output up-projection described in the module docstring. The defaults are the paper’s (Table 1, last row): F = 128, four blocks, N_triplet = 64, N_basis = 8, N_out = 256, N_SHBF = 7, N_SRBF = N_RBF = 6, cutoff 5 Angstrom.

Parameters:

n_features (int, optional) – Message width F, by default 128.
n_interactions (int, optional) – Number of interaction blocks, by default 4.
n_rbf (int, optional) – Number of radial Bessel functions (and zeros per degree of the 2D basis), by default 6.
n_spherical (int, optional) – Number of spherical-harmonic degrees of the 2D basis, by default 7.
cutoff (float, optional) – Cutoff in Angstrom, by default 5.0.
n_triplet_features (int, optional) – Width N_triplet of the down-projected triplet features, by default 64.
n_basis_features (int, optional) – Inner width N_basis of the two-layer basis maps, by default 8.
n_output_features (int, optional) – Width N_out of the up-projected atom features in the output blocks, by default 256.
**kwargs – The remaining options of DimeNet (p, n_before_skip, n_after_skip, n_output_layers, activation, output_init, trainable_rbf, reference_basis, species, atomic_energies).
p (int)
n_before_skip (int)
n_after_skip (int)
n_output_layers (int)
output_init (str)
trainable_rbf (bool)
radial_gate (str)
reference_basis (bool)

classmethod from_config(cfg)[source]#

Build a DimeNetPP from a configuration object.

As DimeNet.from_config(), with the extras n_triplet_features (default 64), n_basis_features (8) and n_output_features (256) in place of n_bilinear; the reference code’s int_emb_size / basis_emb_size / out_emb_size spellings are translated.

Parameters:: cfg (xnn.common.config.schema.ModelConfig) – The core model config.
Returns:: Instantiated model.
Return type:: DimeNetPP