MolecularDiffusion.modules.models.diffsbdd.dynamics

EGNNDynamics, ported from DiffSBDD equivariant_diffusion/dynamics.py.

Two node sets – a diffused ligand and a protein pocket – are projected into a shared joint_nf space by separate encoder/decoder MLPs, concatenated into one flat node set, message-passed over a distance-cutoff graph, then split and decoded again.

The graph is rebuilt every forward pass from batch membership plus optional distance cutoffs (EGNNDynamics.get_edges()). It carries no chemistry: edge_embedding is a 3-way node provenance label (lig-lig / lig-pocket / pocket-pocket), not a bond type, and is None unless edge_embedding_dim is set – which no shipped DiffSBDD config does.

mode='gnn_dynamics' is not ported (unused by every shipped config), so self.mode is gone and the EGNN branch is unconditional.

Submodule names are state-dict keys of the released checkpoints; do not rename atom_encoder / atom_decoder / residue_encoder / residue_decoder / egnn.

Classes

EGNNDynamics

(xh_ligand, xh_pocket, t, mask_lig, mask_pocket) -> (out_lig, out_pocket).

Module Contents

class MolecularDiffusion.modules.models.diffsbdd.dynamics.EGNNDynamics(atom_nf: int, residue_nf: int, n_dims: int = 3, joint_nf: int = 16, hidden_nf: int = 64, act_fn: torch.nn.Module | None = None, n_layers: int = 4, attention: bool = False, condition_time: bool = True, tanh: bool = False, norm_constant: float = 0, inv_sublayers: int = 2, sin_embedding: bool = False, normalization_factor: float = 100, aggregation_method: str = 'sum', update_pocket_coords: bool = True, edge_cutoff_ligand: float | None = None, edge_cutoff_pocket: float | None = None, edge_cutoff_interaction: float | None = None, reflection_equivariant: bool = True, edge_embedding_dim: int | None = None)

Bases: torch.nn.Module

(xh_ligand, xh_pocket, t, mask_lig, mask_pocket) -> (out_lig, out_pocket).

forward(xh_atoms: torch.Tensor, xh_residues: torch.Tensor, t: torch.Tensor, mask_atoms: torch.Tensor, mask_residues: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

get_edges(batch_mask_ligand: torch.Tensor, batch_mask_pocket: torch.Tensor, x_ligand: torch.Tensor, x_pocket: torch.Tensor) → torch.Tensor

Complete graph within each complex, optionally distance-masked.

atom_decoder

atom_encoder

condition_time = True

edge_cutoff_i = None

edge_cutoff_l = None

edge_cutoff_p = None

edge_embedding

edge_nf = None

egnn

n_dims = 3

node_nf = 17

residue_decoder

residue_encoder

update_pocket_coords = True