MolecularDiffusion.modules.models.diffpharma.dynamics

DiffPharma denoiser: four node-set encoders around a three-graph EGNN.

Port of others/DiffPharma/equivariant_diffusion/dynamics.py. Only the egnn_dynamics mode is kept (the gnn_dynamics branch is unreachable in every shipped config and its GNN backbone is not part of the released checkpoint).

Classes

EGNNDynamics

forward returns (lig_out, pocket_out), both flat + scatter-masked.

Module Contents

class MolecularDiffusion.modules.models.diffpharma.dynamics.EGNNDynamics(atom_nf, residue_nf, interh_nf, interhp_nf, n_dims, joint_nf=16, hidden_nf=64, device='cpu', act_fn=torch.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', update_pocket_coords=False, edge_cutoff_ligand=None, edge_cutoff_pocket=None, edge_cutoff_interaction=None, reflection_equivariant=True, edge_embedding_dim=None)

Bases: torch.nn.Module

forward returns (lig_out, pocket_out), both flat + scatter-masked.

lig_out is (n_lig_total, 3 + atom_nf): the predicted epsilon for the ligand. pocket_out is returned for signature compatibility with the joint (unconditional) mode and is unused when the pocket is context.

forward(xh_atoms, xh_residues, xh_intersh, xh_intershp, t, mask_atoms, mask_residues, mask_intersh, mask_intershp)

get_edges(batch_mask_ligand, batch_mask_pocket, x_ligand, x_pocket)

atom_decoder

atom_encoder

condition_time = True

device = 'cpu'

edge_cutoff_i = None

edge_cutoff_l = None

edge_cutoff_p = None

edge_embedding

edge_nf = None

egnn

interh_decoder

interh_encoder

interhp_decoder

interhp_encoder

n_dims

node_nf = 17

residue_decoder

residue_encoder

update_pocket_coords = False