MolecularDiffusion.modules.models.diffsbdd.egnn¶

EGNN backbone ported from DiffSBDD’s equivariant_diffusion/egnn_new.py.

Module and submodule names are load-bearing: they are the state-dict keys of the released CrossDocked checkpoints (Zenodo 8183747), so embedding, embedding_out, e_block_<i>, gcl_<i>, gcl_equiv, edge_mlp, node_mlp, att_mlp, coord_mlp and cross_product_mlp must not be renamed. docs/model_integrations/diffsbdd/scripts/convert_checkpoint.py asserts the whole key set, because cli/generate.py loads with strict=False and a bad remap would load nothing in silence.

Differences from upstream, all deliberate:

GNN / mode='gnn_dynamics' is not ported – no shipped DiffSBDD config selects it (configs/*.yml all leave mode at its egnn_dynamics default).
The flat-scatter batching, update_coords_mask and reflection_equiv cross-product term are why this is a separate EGCL from the platform’s existing one in modules/layers/ rather than a reuse of it.

Classes¶

`EGNN`	E(n)-equivariant GNN over the flat ligand+pocket node set.
`EquivariantBlock`	`inv_sublayers` GCLs followed by one coordinate update.
`EquivariantUpdate`	Coordinate update. `update_coords_mask` freezes the pocket nodes.
`GCL`	Invariant message-passing layer.
`SinusoidsEmbeddingNew`	Sinusoidal distance embedding (off in every shipped config).

Functions¶

`coord2cross`(→ torch.Tensor)
`coord2diff`(→ Tuple[torch.Tensor, torch.Tensor])
`unsorted_segment_sum`(→ torch.Tensor)	TensorFlow's `unsorted_segment_sum`, with 'sum' or 'mean' scaling.

Module Contents¶

class MolecularDiffusion.modules.models.diffsbdd.egnn.EGNN(in_node_nf: int, in_edge_nf: int, hidden_nf: int, act_fn: torch.nn.Module | None = None, n_layers: int = 3, attention: bool = False, out_node_nf: int | None = None, tanh: bool = False, coords_range: float = 15, norm_constant: float = 1, inv_sublayers: int = 2, sin_embedding: bool = False, normalization_factor: float = 100, aggregation_method: str = 'sum', reflection_equiv: bool = True)¶

Bases: torch.nn.Module

E(n)-equivariant GNN over the flat ligand+pocket node set.

forward(h: torch.Tensor, x: torch.Tensor, edge_index: torch.Tensor, update_coords_mask: torch.Tensor | None = None, batch_mask: torch.Tensor | None = None, edge_attr: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]¶

embedding¶

embedding_out¶

hidden_nf¶

n_layers = 3¶

class MolecularDiffusion.modules.models.diffsbdd.egnn.EquivariantBlock(hidden_nf: int, edge_feat_nf: int = 2, act_fn: torch.nn.Module | None = None, n_layers: int = 2, attention: bool = True, tanh: bool = False, coords_range: float = 15, norm_constant: float = 1, sin_embedding: torch.nn.Module | None = None, normalization_factor: float = 100, aggregation_method: str = 'sum', reflection_equiv: bool = True)¶

Bases: torch.nn.Module

inv_sublayers GCLs followed by one coordinate update.

forward(h: torch.Tensor, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor, update_coords_mask: torch.Tensor | None = None, batch_mask: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]¶

coords_range_layer¶

hidden_nf¶

n_layers = 2¶

norm_constant = 1¶

reflection_equiv = True¶

sin_embedding = None¶

class MolecularDiffusion.modules.models.diffsbdd.egnn.EquivariantUpdate(hidden_nf: int, normalization_factor: float, aggregation_method: str, edges_in_d: int = 1, act_fn: torch.nn.Module | None = None, tanh: bool = False, coords_range: float = 10.0, reflection_equiv: bool = True)¶

Bases: torch.nn.Module

Coordinate update. update_coords_mask freezes the pocket nodes.

forward(h: torch.Tensor, coord: torch.Tensor, edge_index: torch.Tensor, coord_diff: torch.Tensor, coord_cross: torch.Tensor | None, edge_attr: torch.Tensor, update_coords_mask: torch.Tensor | None = None) → torch.Tensor¶

aggregation_method¶

coord_mlp¶

coords_range = 10.0¶

cross_product_mlp = None¶

normalization_factor¶

reflection_equiv = True¶

tanh = False¶

class MolecularDiffusion.modules.models.diffsbdd.egnn.GCL(input_nf: int, output_nf: int, hidden_nf: int, normalization_factor: float, aggregation_method: str, edges_in_d: int = 0, act_fn: torch.nn.Module | None = None, attention: bool = False)¶

Bases: torch.nn.Module

Invariant message-passing layer.

forward(h: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor | None = None) → torch.Tensor¶

aggregation_method¶

attention = False¶

edge_mlp¶

node_mlp¶

normalization_factor¶

class MolecularDiffusion.modules.models.diffsbdd.egnn.SinusoidsEmbeddingNew(max_res: float = 15.0, min_res: float = 15.0 / 2000.0, div_factor: int = 4)¶

Bases: torch.nn.Module

Sinusoidal distance embedding (off in every shipped config).

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

dim¶

frequencies¶

n_frequencies¶

MolecularDiffusion.modules.models.diffsbdd.egnn.coord2cross(x: torch.Tensor, edge_index: torch.Tensor, batch_mask: torch.Tensor, norm_constant: float = 1) → torch.Tensor¶

MolecularDiffusion.modules.models.diffsbdd.egnn.coord2diff(x: torch.Tensor, edge_index: torch.Tensor, norm_constant: float = 1) → Tuple[torch.Tensor, torch.Tensor]¶

MolecularDiffusion.modules.models.diffsbdd.egnn.unsorted_segment_sum(data: torch.Tensor, segment_ids: torch.Tensor, num_segments: int, normalization_factor: float | None, aggregation_method: str) → torch.Tensor¶: TensorFlow’s unsorted_segment_sum, with ‘sum’ or ‘mean’ scaling.