MolecularDiffusion.modules.models.geoldm.egnn¶

E(n)-equivariant graph conv layers, ported from GeoLDM’s egnn/egnn_new.py (commit 03ae2031c712a1a6c1678e747bdcdc7a7560e00b).

Self-contained per repo precedent (modules/models/difflinker/egnn.py, modules/models/shepherd_arch’s own EGNN) – not shared with modules/models/en_diffusion.py’s own EnVariationalDiffusion machinery.

Classes¶

`EGNN`
`EquivariantBlock`
`EquivariantUpdate`
`GCL`
`GNN`
`SinusoidsEmbeddingNew`

Functions¶

`coord2diff`(x, edge_index[, norm_constant])
`unsorted_segment_sum`(data, segment_ids, num_segments, ...)	Custom PyTorch op to replicate TensorFlow's unsorted_segment_sum.

Module Contents¶

class MolecularDiffusion.modules.models.geoldm.egnn.EGNN(in_node_nf, in_edge_nf, hidden_nf, device='cpu', act_fn=nn.SiLU(), n_layers=3, attention=False, norm_diff=True, out_node_nf=None, tanh=False, coords_range=15, norm_constant=1, inv_sublayers=2, sin_embedding=False, normalization_factor=100, aggregation_method='sum')¶

Bases: torch.nn.Module

forward(h, x, edge_index, node_mask=None, edge_mask=None)¶

aggregation_method = 'sum'¶

coords_range_layer¶

device = 'cpu'¶

embedding¶

embedding_out¶

hidden_nf¶

n_layers = 3¶

norm_diff = True¶

normalization_factor = 100¶

class MolecularDiffusion.modules.models.geoldm.egnn.EquivariantBlock(hidden_nf, edge_feat_nf=2, device='cpu', act_fn=nn.SiLU(), n_layers=2, attention=True, norm_diff=True, tanh=False, coords_range=15, norm_constant=1, sin_embedding=None, normalization_factor=100, aggregation_method='sum')¶

Bases: torch.nn.Module

forward(h, x, edge_index, node_mask=None, edge_mask=None, edge_attr=None)¶

aggregation_method = 'sum'¶

coords_range_layer¶

device = 'cpu'¶

hidden_nf¶

n_layers = 2¶

norm_constant = 1¶

norm_diff = True¶

normalization_factor = 100¶

sin_embedding = None¶

class MolecularDiffusion.modules.models.geoldm.egnn.EquivariantUpdate(hidden_nf, normalization_factor, aggregation_method, edges_in_d=1, act_fn=nn.SiLU(), tanh=False, coords_range=10.0)¶

Bases: torch.nn.Module

coord_model(h, coord, edge_index, coord_diff, edge_attr, edge_mask)¶

forward(h, coord, edge_index, coord_diff, edge_attr=None, node_mask=None, edge_mask=None)¶

aggregation_method¶

coord_mlp¶

coords_range = 10.0¶

normalization_factor¶

tanh = False¶

class MolecularDiffusion.modules.models.geoldm.egnn.GCL(input_nf, output_nf, hidden_nf, normalization_factor, aggregation_method, edges_in_d=0, nodes_att_dim=0, act_fn=nn.SiLU(), attention=False)¶

Bases: torch.nn.Module

edge_model(source, target, edge_attr, edge_mask)¶

forward(h, edge_index, edge_attr=None, node_attr=None, node_mask=None, edge_mask=None)¶

node_model(x, edge_index, edge_attr, node_attr)¶

aggregation_method¶

attention = False¶

edge_mlp¶

node_mlp¶

normalization_factor¶

class MolecularDiffusion.modules.models.geoldm.egnn.GNN(in_node_nf, in_edge_nf, hidden_nf, aggregation_method='sum', device='cpu', act_fn=nn.SiLU(), n_layers=4, attention=False, normalization_factor=1, out_node_nf=None)¶

Bases: torch.nn.Module

forward(h, edges, edge_attr=None, node_mask=None, edge_mask=None)¶

device = 'cpu'¶

embedding¶

embedding_out¶

hidden_nf¶

n_layers = 4¶

class MolecularDiffusion.modules.models.geoldm.egnn.SinusoidsEmbeddingNew(max_res=15.0, min_res=15.0 / 2000.0, div_factor=4)¶

Bases: torch.nn.Module

forward(x)¶

dim¶

frequencies¶

n_frequencies¶

MolecularDiffusion.modules.models.geoldm.egnn.coord2diff(x, edge_index, norm_constant=1)¶

MolecularDiffusion.modules.models.geoldm.egnn.unsorted_segment_sum(data, segment_ids, num_segments, normalization_factor, aggregation_method: str)¶

Custom PyTorch op to replicate TensorFlow’s unsorted_segment_sum.

Normalization: ‘sum’ or ‘mean’.