dig.threedgraph.method

Method classes under dig.threedgraph.method.

ComENet	The ComENet from the "ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs" 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 Molecular Graphs" paper.
ProNet	The ProNet from the "Learning Protein Representations via Complete 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 Molecular Graphs" paper.
SphereNet	The spherical message passing neural network SphereNet from the "Spherical Message Passing for 3D Molecular Graphs" paper.
run	The base script for running different 3DGN methods.

class ComENet(cutoff=8.0, num_layers=4, hidden_channels=256, middle_channels=64, out_channels=1, num_radial=3, num_spherical=2, num_output_layers=3)

The ComENet from the “ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs” paper.

Parameters

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

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

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

• middle_channels (int, optional) – Middle embedding size for the two layer linear block. (default: 256)

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

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

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

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

forward(batch_data)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

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')

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 Molecular Graphs” 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)

forward(batch_data)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ProNet(level='aminoacid', num_blocks=4, hidden_channels=128, out_channels=1, mid_emb=64, num_radial=6, num_spherical=2, cutoff=10.0, max_num_neighbors=32, int_emb_layers=3, out_layers=2, num_pos_emb=16, dropout=0, data_augment_eachlayer=False, euler_noise=False)

The ProNet from the “Learning Protein Representations via Complete 3D Graph Networks” paper.

Parameters

• level – (str, optional): The level of protein representations. It could be aminoacid, obj:backbone, and allatom. (default: aminoacid)

• num_blocks (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)

• mid_emb (int, optional) – Embedding size used for geometric features. (default: 64)

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

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

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

• max_num_neighbors (int, optional) – Max number of neighbors during graph construction. (default: 32)

• int_emb_layers (int, optional) – Number of embedding layers in the interaction block. (default: 3)

• out_layers (int, optional) – Number of layers for features after interaction blocks. (default: 2)

• num_pos_emb (int, optional) – Number of positional embeddings. (default: 16)

• dropout (float, optional) – Dropout. (default: 0)

• data_augment_eachlayer (bool, optional) – Data augmentation tricks. If set to True, will add noise to the node features before each interaction block. (default: False)

• euler_noise (bool, optional) – Data augmentation tricks. If set to True, will add noise to Euler angles. (default: False)

forward(batch_data)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SchNet(energy_and_force=False, cutoff=10.0, num_layers=6, hidden_channels=128, out_channels=1, num_filters=128, num_gaussians=50)

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 Molecular Graphs” 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)

• out_channels (int, optional) – Output embedding size. (default: 1)

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

forward(batch_data)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

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', use_node_features=True)

The spherical message passing neural network SphereNet from the “Spherical Message Passing for 3D Molecular Graphs” 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)

forward(batch_data)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class run

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='')

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)

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)

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)