gt4sd.algorithms.generation.diffusion.geodiff.model.core module

Summary

Classes:

MoleculeGNN

Graph Neural Network Model for molecule conformation.

Reference

class MoleculeGNN(hidden_dim=128, num_convs=6, num_convs_local=4, cutoff=10.0, mlp_act='relu', edge_order=3, edge_encoder='mlp', smooth_conv=True)

Bases: ModelMixin, ConfigMixin

Graph Neural Network Model for molecule conformation.

Code adapted from colab:

https://colab.research.google.com/drive/1pLYYWQhdLuv1q-JtEHGZybxp2RBF8gPs#scrollTo=-3-P4w5sXkRU written by Nathan Lambert.

__init__(hidden_dim=128, num_convs=6, num_convs_local=4, cutoff=10.0, mlp_act='relu', edge_order=3, edge_encoder='mlp', smooth_conv=True)

Constructor for MoleculeGNN. :type hidden_dim: int :param hidden_dim: hidden dimension of the GNN. :type num_convs: int :param num_convs: number of GNN layers. :type num_convs_local: int :param num_convs_local: number of GNN layers for local edges. :type cutoff: float :param cutoff: nutoff radius for the GNN. :type mlp_act: str :param mlp_act: activation function for the MLP. :type edge_order: int :param edge_order: order of the edge features. :type edge_encoder: str :param edge_encoder: type of edge encoder. :type smooth_conv: bool :param smooth_conv: whether to use smooth convolution.

_forward(atom_type, pos, bond_index, bond_type, batch, edge_index=None, edge_type=None, edge_length=None, return_edges=False, extend_order=True, extend_radius=True, is_sidechain=None)

Forward pass for edges features.

Parameters

• atom_type (Tensor) – types of atoms, (N, ).

• pos (Tensor) – positions of atoms, (N, 3).

• bond_index (Tensor) – indices of bonds (not extended, not radius-graph), (2, E).

• bond_type (Tensor) – bond types, (E, ).

• batch (Tensor) – node index to graph index, (N, ).

• edge_index (Optional[Tensor, None]) – indices of edges (extended, radius-graph), (2, E’). If None, determined by extend_graph_order_radius.

• edge_type (Optional[Tensor, None]) – edge types, (E’, ). If None, determined by extend_graph_order_radius.

• edge_length (Optional[int, None]) – edge lengths, (E’, ). If None, determined by extend_graph_order_radius.

• return_edges (bool) – whether to return edge_index, edge_type, edge_length.

• extend_order (bool) – whether to extend the graph by bond order.

• extend_radius (bool) – whether to extend the graph by radius.

• is_sidechain (Optional[bool, None]) – whether the atom is a sidechain atom.

Return type

Tuple[Any, …]

Returns

tuple with local and global invariant edge features.

If return_edges is True, it also returns edge_index, edge_type, edge_length, local_edge_index.

forward(sample, timestep, return_dict=True, sigma=1.0, global_start_sigma=0.5, w_global=1.0, extend_order=False, extend_radius=True, clip_local=None, clip_global=1000.0)

Forward pass for the model.

Parameters

• sample (Data) – packed torch geometric object.

• timesteps. (timestep) –

• return_dict (bool) – whether or not to return a MoleculeGNNOutput instead of a plain tuple.

• sigma (float) – the noise variance scale.

• global_start_sigma (float) – the noise variance scale for global edges.

• w_global (float) – the weight for global edges.

• extend_order – whether to extend the graph by bond order.

• extend_radius – whether to extend the graph by radius.

• clip_local – the clip value for local edges.

• clip_global (float) – the clip value for global edges.

Return type

Union[MoleculeGNNOutput, Tuple]

Returns

dict with samples. If return_dict is false, return positions.

__annotations__ = {}

__doc__ = 'Graph Neural Network Model for molecule conformation.\n\n    Code adapted from colab:\n        https://colab.research.google.com/drive/1pLYYWQhdLuv1q-JtEHGZybxp2RBF8gPs#scrollTo=-3-P4w5sXkRU written by Nathan Lambert.\n    '

__module__ = 'gt4sd.algorithms.generation.diffusion.geodiff.model.core'