dig.threedgraph.method¶
Method classes under dig.threedgraph.method
.
The ComENet from the "ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs" paper. 

The reimplementation for DimeNet++ from the "Fast and UncertaintyAware Directional Message Passing for NonEquilibrium Molecules" paper under the 3DGN gramework from "Spherical Message Passing for 3D Molecular Graphs" paper. 

The reimplementation for SchNet from the "SchNet: A Continuousfilter Convolutional Neural Network for Modeling Quantum Interactions" paper under the 3DGN gramework from "Spherical Message Passing for 3D Molecular Graphs" paper. 

The spherical message passing neural network SphereNet from the "Spherical Message Passing for 3D Molecular Graphs" paper. 

The base script for running different 3DGN methods. 
 class ComENet(cutoff=8.0, num_layers=4, hidden_channels=256, out_channels=1, num_radial=3, num_spherical=2, num_output_layers=3)[source]¶
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
)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)[source]¶
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')[source]¶
The reimplementation for DimeNet++ from the “Fast and UncertaintyAware Directional Message Passing for NonEquilibrium 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
andzeros
. (default:GlorotOrthogonal
)
 forward(batch_data)[source]¶
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)[source]¶
The reimplementation for SchNet from the “SchNet: A Continuousfilter 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)[source]¶
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)[source]¶
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
andzeros
. (default:GlorotOrthogonal
)
 forward(batch_data)[source]¶
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[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
)