MolecularDiffusion.modules.models.diffpharma.egnn¶

Three-graph EGNN backbone for DiffPharma.

Faithful port of others/DiffPharma/equivariant_diffusion/egnn_new.py with import rewrites and dead code (the GNN variant, EGNN.noh_skip, the commented-out debug scaffolding) removed. Parameter names and module attribute names are unchanged so the released checkpoint maps 1:1.

Aggregation uses plain Tensor.scatter_add_ – no torch_scatter here.

Classes¶

`EGNN`	Three parallel EGNN stacks with periodic ligand-node fusion.
`EquivariantBlock`
`EquivariantUpdate`
`GCL`
`SinusoidsEmbeddingNew`

Functions¶

`coord2cross`(x, edge_index, batch_mask[, norm_constant])
`coord2diff`(x, edge_index[, norm_constant])
`unsorted_segment_sum`(data, segment_ids, num_segments, ...)	TensorFlow-style `unsorted_segment_sum` ('sum' or 'mean').

Module Contents¶

class MolecularDiffusion.modules.models.diffpharma.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', reflection_equiv=True)¶

Bases: torch.nn.Module

Three parallel EGNN stacks with periodic ligand-node fusion.

Graph 1 is ligand+pocket, graph 2 ligand+H-bond particles, graph 3 ligand+hydrophobic particles. The ligand block of h/x (the first len(mask_atom) rows of each graph) is averaged across graphs every layer – pairwise onto graphs 2/3 normally, three-way at the midpoint and final layer.

forward(h, x, h2, x2, h3, x3, edge_index, edge_index_2, edge_index_3, mask_atom=None, mask_intersh=None, mask_intershp=None, node_mask=None, edge_mask=None, update_coords_mask=None, update_coords_mask_2=None, update_coords_mask_3=None, batch_mask=None, batch_mask_2=None, batch_mask_3=None, edge_attr=None)¶

aggregation_method = 'sum'¶

coords_range_layer¶

device = 'cpu'¶

embedding¶

embedding2¶

embedding2_out¶

embedding3¶

embedding3_out¶

embedding_out¶

hidden_nf¶

n_layers = 3¶

norm_diff = True¶

normalization_factor = 100¶

reflection_equiv = True¶

class MolecularDiffusion.modules.models.diffpharma.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', reflection_equiv=True)¶

Bases: torch.nn.Module

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

aggregation_method = 'sum'¶

coords_range_layer¶

device = 'cpu'¶

hidden_nf¶

n_layers = 2¶

norm_constant = 1¶

norm_diff = True¶

normalization_factor = 100¶

reflection_equiv = True¶

sin_embedding = None¶

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

Bases: torch.nn.Module

coord_model(h, coord, edge_index, coord_diff, coord_cross, edge_attr, edge_mask, update_coords_mask=None)¶

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

aggregation_method¶

coord_mlp¶

coords_range = 10.0¶

cross_product_mlp¶

normalization_factor¶

reflection_equiv = True¶

tanh = False¶

class MolecularDiffusion.modules.models.diffpharma.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.diffpharma.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.diffpharma.egnn.coord2cross(x, edge_index, batch_mask, norm_constant=1)¶

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

MolecularDiffusion.modules.models.diffpharma.egnn.unsorted_segment_sum(data, segment_ids, num_segments, normalization_factor, aggregation_method)¶: TensorFlow-style unsorted_segment_sum (‘sum’ or ‘mean’).