MolecularDiffusion.modules.models.pmdm.encoders¶

Encoders used by the PMDM epsilon network.

Ported from PMDM’s models/encoders/{attention,edge,egnn,egnn_pytorch,schnet}.py, keeping only what MDM_full_pocket_coor_shared reaches:

CrossAttentionBlock – upstream BasicTransformerBlock, the ligand<->pocket token mixer.
SchNetProteinEncoder – upstream SchNetEncoder_protein, used for both the pocket and the ligand tower.
MLPEdgeEncoder – upstream MLPEdgeEncoder (edge_encoder: mlp).
EGNNSparseNetwork – upstream EGNN_Sparse_Network.

Two deliberate deviations from upstream:

No einops. Upstream’s attention uses rearrange for three reshapes; those are plain view/permute here, so PMDM needs no new dependency.
No ``MessagePassing``. Upstream’s EGNN_Sparse overrides propagate and reaches into PyG privates (self.inspector.distribute, _collect) that moved in PyG 2.5. The message/aggregate pair is written directly with scatter_add, which is what aggr="add" did anyway. Every graph fed to these layers is symmetric, so the source/target convention cannot flip a result.

Dropped as unreachable under this integration’s scope: the fourier-feature branch (fourier_features=0), global linear attention (global_linear_attn_every=0), the node/edge embedding-table branch (empty dims), and upstream’s unused ligandemb/proteinemb/inner atten_layer parameters in EGNN_Sparse_Network (the live code path is h = z).

linker_mask (used by PMDMEpsNet.linker_sample – see modules/tasks/diffusion_pmdm.py’s PMDMConstrainedGenerator, mode: linker) zeros the coordinate-update contribution for every atom it does not cover, biasing the EGNN toward moving only the region being regenerated (others/PMDM/models/encoders/egnn.py:268-274). Upstream additionally concatenates a zero mask for the pocket half before this point, but that pocket-half masking is a no-op on the coordinates this port returns (the pocket rows of coors_out never receive +mhat_i regardless of the mask), so only the ligand-sized mask is threaded through here.

Attributes¶

SiLU

Classes¶

`CrossAttentionBlock`	Upstream `BasicTransformerBlock`: self-attend both towers, then cross.
`EGNNSparseLayer`	One E(n)-equivariant message-passing layer.
`EGNNSparseNetwork`	Stack of `EGNNSparseLayer` over the joined ligand+pocket cloud.
`MLPEdgeEncoder`	Embed an edge length into `hidden_dim`.
`SchNetProteinEncoder`	Continuous-filter tower over a raw point cloud (its own radius graph).

Module Contents¶

class MolecularDiffusion.modules.models.pmdm.encoders.CrossAttentionBlock(dim: int, n_heads: int, d_head: int, dropout: float = 0.0, context_dim: int | None = None)¶

Bases: torch.nn.Module

Upstream BasicTransformerBlock: self-attend both towers, then cross.

forward(x: torch.Tensor, context: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]¶

attn1¶

attn2¶

attn2p¶

attn_p¶

ff¶

ffp¶

norm1¶

norm1p¶

norm2¶

norm2p¶

norm3¶

norm3p¶

class MolecularDiffusion.modules.models.pmdm.encoders.EGNNSparseLayer(feats_dim: int, pos_dim: int = 3, edge_attr_dim: int = 0, m_dim: int = 16, soft_edge: int = 0, norm_feats: bool = False, norm_coors: bool = False, norm_coors_scale_init: float = 0.01, dropout: float = 0.0)¶

Bases: torch.nn.Module

One E(n)-equivariant message-passing layer.

Only the ligand block of the coordinate update is applied: the pocket is a fixed conditioning cloud and must not move. Ligand nodes are the first n_ligand rows of every tensor (the caller concatenates ligand-then-pocket), which is exactly upstream’s convention.

forward(x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor | None, batch: torch.Tensor, n_ligand: int, linker_mask: torch.Tensor | None = None) → torch.Tensor¶

coors_mlp¶

coors_norm¶

edge_mlp¶

edge_weight¶

feats_dim¶

m_dim = 16¶

node_mlp¶

node_norm¶

pos_dim = 3¶

soft_edge = 0¶

class MolecularDiffusion.modules.models.pmdm.encoders.EGNNSparseNetwork(n_layers: int, feats_dim: int, pos_dim: int = 3, edge_attr_dim: int = 0, m_dim: int = 16, soft_edge: int = 0, norm_feats: bool = True, norm_coors: bool = False, norm_coors_scale_init: float = 0.01, dropout: float = 0.0)¶

Bases: torch.nn.Module

Stack of EGNNSparseLayer over the joined ligand+pocket cloud.

Returns the ligand slice only: (node_feats, coord_update), where the coordinate output is the displacement from the input positions (the equivariant score), as upstream does.

forward(z: torch.Tensor, pos: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor, batch: torch.Tensor, n_ligand: int, linker_mask: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]¶

feats_dim¶

mpnn_layers¶

pos_dim = 3¶

class MolecularDiffusion.modules.models.pmdm.encoders.MLPEdgeEncoder(hidden_dim: int = 100, activation: str = 'relu')¶

Bases: torch.nn.Module

Embed an edge length into hidden_dim.

Upstream also multiplies in a learned bond-type embedding, but both call sites in MDM_full_pocket_coor_shared.net pass edge_type=None (the joined ligand+pocket graphs are purely geometric), leaving that embedding permanently untrained. It is dropped here rather than shipped as dead weight that DDP would then flag as an unused parameter.

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

hidden_dim = 100¶

mlp¶

property out_channels: int¶

class MolecularDiffusion.modules.models.pmdm.encoders.SchNetProteinEncoder(hidden_channels: int = 128, num_filters: int = 128, num_interactions: int = 6, edge_channels: int = 64, cutoff: float = 10.0, input_dim: int = 31)¶

Bases: torch.nn.Module

Continuous-filter tower over a raw point cloud (its own radius graph).

forward(node_attr: torch.Tensor, pos: torch.Tensor, batch: torch.Tensor) → torch.Tensor¶

cutoff = 10.0¶

distance_expansion¶

emblin¶

hidden_channels = 128¶

interactions¶

property out_channels: int¶

MolecularDiffusion.modules.models.pmdm.encoders.SiLU¶