MolecularDiffusion.modules.models.difflinker.egnn¶

EGNN backbone + DiffLinker’s Dynamics/DynamicsWithPockets wrapper, ported near-verbatim from DiffLinker’s src/egnn.py.

DynamicsWithPockets and the pocket-distance-graph helpers port as-is but are unused/untested this pass — pocket/protein conditioning is explicitly out of scope (see docs/model_integrations/difflinker/INTEGRATION_PLAN.md).

Classes¶

`Dynamics`
`DynamicsWithPockets`	Ports as-is; pocket/protein conditioning is out of scope this pass, so
`EGNN`
`EquivariantBlock`
`EquivariantUpdate`
`GCL`
`GNN`
`SinusoidsEmbeddingNew`

Functions¶

`coord2diff`(x, edge_index[, norm_constant])
`unsorted_segment_sum`(data, segment_ids, num_segments, ...)	Replicates TensorFlow's `unsorted_segment_sum`. `aggregation_method`

Module Contents¶

class MolecularDiffusion.modules.models.difflinker.egnn.Dynamics(n_dims, in_node_nf, context_node_nf, hidden_nf=64, device='cpu', activation=nn.SiLU(), n_layers=4, attention=False, condition_time=True, tanh=False, norm_constant=0, inv_sublayers=2, sin_embedding=False, normalization_factor=100, aggregation_method='sum', model='egnn_dynamics', normalization=None, centering=False, graph_type='FC')¶

Bases: torch.nn.Module

forward(t, xh, node_mask, linker_mask, edge_mask, context)¶

t: (B)
xh: (B, N, D), where D = 3 + nf
node_mask: (B, N, 1)
edge_mask: (B*N*N, 1)
context: (B, N, C)

get_edges(n_nodes, batch_size)¶

centering = False¶

condition_time = True¶

context_node_nf¶

edge_cache: dict¶

graph_type = 'FC'¶

model = 'egnn_dynamics'¶

n_dims¶

class MolecularDiffusion.modules.models.difflinker.egnn.DynamicsWithPockets(n_dims, in_node_nf, context_node_nf, hidden_nf=64, device='cpu', activation=nn.SiLU(), n_layers=4, attention=False, condition_time=True, tanh=False, norm_constant=0, inv_sublayers=2, sin_embedding=False, normalization_factor=100, aggregation_method='sum', model='egnn_dynamics', normalization=None, centering=False, graph_type='FC')¶

Bases: Dynamics

Ports as-is; pocket/protein conditioning is out of scope this pass, so this class is unused/untested — see INTEGRATION_PLAN.md.

forward(t, xh, node_mask, linker_mask, edge_mask, context)¶

t: (B)
xh: (B, N, D), where D = 3 + nf
node_mask: (B, N, 1)
edge_mask: (B*N*N, 1)
context: (B, N, C)

get_dist_edges(x, node_mask, batch_mask, linker_mask, fragment_only_mask, pocket_only_mask)¶

static get_dist_edges_4A(x, node_mask, batch_mask)¶

class MolecularDiffusion.modules.models.difflinker.egnn.EGNN(in_node_nf, in_edge_nf, hidden_nf, device='cpu', activation=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, linker_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.difflinker.egnn.EquivariantBlock(hidden_nf, edge_feat_nf=2, device='cpu', activation=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, linker_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.difflinker.egnn.EquivariantUpdate(hidden_nf, normalization_factor, aggregation_method, edges_in_d=1, activation=nn.SiLU(), tanh=False, coords_range=10.0)¶

Bases: torch.nn.Module

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

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

aggregation_method¶

coord_mlp¶

coords_range = 10.0¶

normalization_factor¶

tanh = False¶

class MolecularDiffusion.modules.models.difflinker.egnn.GCL(input_nf, output_nf, hidden_nf, normalization_factor, aggregation_method, activation, edges_in_d=0, nodes_att_dim=0, attention=False, normalization=None)¶

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¶

normalization_factor¶

class MolecularDiffusion.modules.models.difflinker.egnn.GNN(in_node_nf, in_edge_nf, hidden_nf, aggregation_method='sum', device='cpu', activation=nn.SiLU(), n_layers=4, attention=False, normalization_factor=1, out_node_nf=None, normalization=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.difflinker.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.difflinker.egnn.coord2diff(x, edge_index, norm_constant=1)¶

MolecularDiffusion.modules.models.difflinker.egnn.unsorted_segment_sum(data, segment_ids, num_segments, normalization_factor, aggregation_method: str)¶: Replicates TensorFlow’s unsorted_segment_sum. aggregation_method is ‘sum’ or ‘mean’.