dig.threedgraph.method

Method classes under dig.threedgraph.method. All the methods are re-implemented under the unifying 3DGN framework from the “Spherical Message Passing for 3D Graph Networks” paper.

DimeNetPP

The re-implementation for DimeNet++ from the “Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules” paper under the 3DGN gramework from “Spherical Message Passing for 3D Graph Networks” paper.

SchNet

The re-implementation for SchNet from the “SchNet: A Continuous-filter Convolutional Neural Network for Modeling Quantum Interactions” paper under the 3DGN gramework from “Spherical Message Passing for 3D Graph Networks” paper.

SphereNet

The spherical message passing neural network SphereNet from the “Spherical Message Passing for 3D Graph Networks” paper.

run

The base script for running different 3DGN methods.
class DimeNetPP(energy_and_force=False, cutoff=5.0, num_layers=4, hidden_channels=128, out_channels=1, int_emb_size=64, basis_emb_size=8, out_emb_channels=256, num_spherical=7, num_radial=6, envelope_exponent=5, num_before_skip=1, num_after_skip=2, num_output_layers=3, act=<function swish>, output_init='GlorotOrthogonal')[source]

The re-implementation for DimeNet++ from the “Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules” paper under the 3DGN gramework from “Spherical Message Passing for 3D Graph Networks” paper.

Parameters

  • energy_and_force (bool, optional) – If set to True, will predict energy and take the negative of the derivative of the energy with respect to the atomic positions as predicted forces. (default: False)

  • cutoff (float, optional) – Cutoff distance for interatomic interactions. (default: 5.0)

  • num_layers (int, optional) – Number of building blocks. (default: 4)

  • hidden_channels (int, optional) – Hidden embedding size. (default: 128)

  • out_channels (int, optional) – Size of each output sample. (default: 1)

  • int_emb_size (int, optional) – Embedding size used for interaction triplets. (default: 64)

  • basis_emb_size (int, optional) – Embedding size used in the basis transformation. (default: 8)

  • out_emb_channels (int, optional) – Embedding size used for atoms in the output block. (default: 256)

  • num_spherical (int, optional) – Number of spherical harmonics. (default: 7)

  • num_radial (int, optional) – Number of radial basis functions. (default: 6)

  • envelop_exponent (int, optional) – Shape of the smooth cutoff. (default: 5)

  • num_before_skip (int, optional) – Number of residual layers in the interaction blocks before the skip connection. (default: 1)

  • num_after_skip (int, optional) – Number of residual layers in the interaction blocks before the skip connection. (default: 2)

  • num_output_layers (int, optional) – Number of linear layers for the output blocks. (default: 3)

  • act – (function, optional): The activation funtion. (default: swish)

  • output_init – (str, optional): The initialization fot the output. It could be GlorotOrthogonal and zeros. (default: GlorotOrthogonal)

class SchNet(energy_and_force=False, cutoff=10.0, num_layers=6, hidden_channels=128, num_filters=128, num_gaussians=50)[source]

The re-implementation for SchNet from the “SchNet: A Continuous-filter Convolutional Neural Network for Modeling Quantum Interactions” paper under the 3DGN gramework from “Spherical Message Passing for 3D Graph Networks” paper.

Parameters

  • energy_and_force (bool, optional) – If set to True, will predict energy and take the negative of the derivative of the energy with respect to the atomic positions as predicted forces. (default: False)

  • num_layers (int, optional) – The number of layers. (default: 6)

  • hidden_channels (int, optional) – Hidden embedding size. (default: 128)

  • num_filters (int, optional) – The number of filters to use. (default: 128)

  • num_gaussians (int, optional) – The number of gaussians \(\mu\). (default: 50)

  • cutoff (float, optional) – Cutoff distance for interatomic interactions. (default: 10.0).

class SphereNet(energy_and_force=False, cutoff=5.0, num_layers=4, hidden_channels=128, out_channels=1, int_emb_size=64, basis_emb_size_dist=8, basis_emb_size_angle=8, basis_emb_size_torsion=8, out_emb_channels=256, num_spherical=7, num_radial=6, envelope_exponent=5, num_before_skip=1, num_after_skip=2, num_output_layers=3, act=<function swish>, output_init='GlorotOrthogonal')[source]

The spherical message passing neural network SphereNet from the “Spherical Message Passing for 3D Graph Networks” paper.

Parameters

  • energy_and_force (bool, optional) – If set to True, will predict energy and take the negative of the derivative of the energy with respect to the atomic positions as predicted forces. (default: False)

  • cutoff (float, optional) – Cutoff distance for interatomic interactions. (default: 5.0)

  • num_layers (int, optional) – Number of building blocks. (default: 4)

  • hidden_channels (int, optional) – Hidden embedding size. (default: 128)

  • out_channels (int, optional) – Size of each output sample. (default: 1)

  • int_emb_size (int, optional) – Embedding size used for interaction triplets. (default: 64)

  • basis_emb_size_dist (int, optional) – Embedding size used in the basis transformation of distance. (default: 8)

  • basis_emb_size_angle (int, optional) – Embedding size used in the basis transformation of angle. (default: 8)

  • basis_emb_size_torsion (int, optional) – Embedding size used in the basis transformation of torsion. (default: 8)

  • out_emb_channels (int, optional) – Embedding size used for atoms in the output block. (default: 256)

  • num_spherical (int, optional) – Number of spherical harmonics. (default: 7)

  • num_radial (int, optional) – Number of radial basis functions. (default: 6)

  • envelop_exponent (int, optional) – Shape of the smooth cutoff. (default: 5)

  • num_before_skip (int, optional) – Number of residual layers in the interaction blocks before the skip connection. (default: 1)

  • num_after_skip (int, optional) – Number of residual layers in the interaction blocks before the skip connection. (default: 2)

  • num_output_layers (int, optional) – Number of linear layers for the output blocks. (default: 3)

  • act – (function, optional): The activation funtion. (default: swish)

  • output_init – (str, optional): The initialization fot the output. It could be GlorotOrthogonal and zeros. (default: GlorotOrthogonal)

class run[source]

The base script for running different 3DGN methods.

run(device, train_dataset, valid_dataset, test_dataset, model, loss_func, evaluation, epochs=500, batch_size=32, vt_batch_size=32, lr=0.0005, lr_decay_factor=0.5, lr_decay_step_size=50, weight_decay=0, energy_and_force=False, p=100, save_dir='', log_dir='')[source]

The run script for training and validation.

Parameters

  • device (torch.device) – Device for computation.

  • train_dataset – Training data.

  • valid_dataset – Validation data.

  • test_dataset – Test data.

  • model – Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.

  • loss_func (function) – The used loss funtion for training.

  • evaluation (function) – The evaluation function.

  • epochs (int, optinal) – Number of total training epochs. (default: 500)

  • batch_size (int, optinal) – Number of samples in each minibatch in training. (default: 32)

  • vt_batch_size (int, optinal) – Number of samples in each minibatch in validation/testing. (default: 32)

  • lr (float, optinal) – Initial learning rate. (default: 0.0005)

  • lr_decay_factor (float, optinal) – Learning rate decay factor. (default: 0.5)

  • lr_decay_step_size (int, optinal) – epochs at which lr_initial <- lr_initial * lr_decay_factor. (default: 50)

  • weight_decay (float, optinal) – weight decay factor at the regularization term. (default: 0)

  • energy_and_force (bool, optional) – If set to True, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: False)

  • p (int, optinal) – The forces’ weight for a joint loss of forces and conserved energy during training. (default: 100)

  • save_dir (str, optinal) – The path to save trained models. If set to '', will not save the model. (default: '')

  • log_dir (str, optinal) – The path to save log files. If set to '', will not save the log files. (default: '')

train(model, optimizer, train_loader, energy_and_force, p, loss_func, device)[source]

The script for training.

Parameters

  • model – Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in training.

  • train_loader (Dataloader) – Dataloader for training.

  • energy_and_force (bool, optional) – If set to True, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: False)

  • p (int, optinal) – The forces’ weight for a joint loss of forces and conserved energy during training. (default: 100)

  • loss_func (function) – The used loss funtion for training.

  • device (torch.device) – The device where the model is deployed.

Return type

Traning loss. ( mae)

val(model, data_loader, energy_and_force, p, evaluation, device)[source]

The script for validation/test.

Parameters

  • model – Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.

  • data_loader (Dataloader) – Dataloader for validation or test.

  • energy_and_force (bool, optional) – If set to True, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: False)

  • p (int, optinal) – The forces’ weight for a joint loss of forces and conserved energy. (default: 100)

  • evaluation (function) – The used funtion for evaluation.

  • device (torch.device, optional) – The device where the model is deployed.

Return type

Evaluation result. ( mae)