MolecularDiffusion.modules.models.shepherd_arch.egnn¶

E(n) Equivariant Graph Neural Network (EGNN) for ShEPhERD.

Classes¶

`EGNN`
`GaussianSmearing`
`MultiLayerPerceptron`

Module Contents¶

class MolecularDiffusion.modules.models.shepherd_arch.egnn.EGNN(node_embedding_dim, node_output_embedding_dim, edge_attr_dim, distance_expansion_dim=32, normalize_distance_vectors=True, num_MLP_hidden_layers=2, MLP_hidden_dim=32)¶

Bases: torch.nn.Module

forward(x, pos, edge_index, batch, edge_attr=None, pos_update_mask=None, residual_pos_update=False, residual_x_update=False)¶

coordinate_mlp¶

distance_expansion¶

message_mlp¶

message_mlp_embedding_dim¶

node_mlp¶

node_mlp_embedding_dim¶

node_output_embedding¶

node_output_embedding_dim¶

normalize_distance_vectors = True¶

class MolecularDiffusion.modules.models.shepherd_arch.egnn.GaussianSmearing(start: float = 0.0, stop: float = 5.0, num_gaussians: int = 50)¶

Bases: torch.nn.Module

forward(dist: torch.Tensor) → torch.Tensor¶

coeff¶

class MolecularDiffusion.modules.models.shepherd_arch.egnn.MultiLayerPerceptron(input_dim, hidden_dim, output_dim, num_hidden_layers, activation=torch.nn.LeakyReLU(0.2), include_final_activation=True)¶

Bases: torch.nn.Module

forward(x)¶

mlp¶