MolecularDiffusion.modules.models.pmdm.epsnet¶

PMDM’s pocket-conditioned epsilon network.

Port of models/epsnet/MDM_pocket_coor_shared.py::MDM_full_pocket_coor_shared (the only network PMDM’s get_model can return), re-exported as PMDMEpsNet.

Shape of the model, unchanged from upstream:

the pocket is embedded once by a SchNet tower and is never noised;
the ligand is the diffused variable – coordinates score-matching style, atom features VP style;
a cross-attention block mixes ligand and pocket tokens, then two EGNN stacks (global @ g_cutoff, local @ cutoff) run over the concatenated ligand+pocket point set, with only ligand coordinates updated;
two MLP heads produce grad_{global,local}_node; the position score is the EGNN coordinate output itself.

Differences from upstream, all of them removals of code that this integration’s scope (see docs/model_integrations/pmdm/INTEGRATION_PLAN.md) puts out of reach:

the vae_context VAE latent branch and the context property list (both default-off upstream);
atom_num_emb (default-off), is_sidechain (always None), and the gaussian edge encoder (config is mlp);
self.model_global/model_local, which were ModuleList aliases of already-registered submodules and only duplicated state_dict keys.

The config object is replaced by explicit keyword arguments so the whole network is Hydra-instantiable from configs/tasks/diffusion_pmdm.yaml.

Classes¶

PMDMEpsNet

Dual-encoder, pocket-conditioned score network.

Functions¶

get_beta_schedule(→ numpy.ndarray)

Upstream's noise schedules (sigmoid is what the released config uses).

Module Contents¶

class MolecularDiffusion.modules.models.pmdm.epsnet.PMDMEpsNet(num_atom: int = 10, protein_feature_dim: int = 31, hidden_dim: int = 128, protein_hidden_dim: int = 128, num_convs: int = 3, num_convs_local: int = 3, protein_num_convs: int = 2, cutoff: float = 3.0, g_cutoff: float = 6.0, encoder_cutoff: float = 6.0, edge_order: int = 3, mlp_act: str = 'relu', edge_encoder: str = 'mlp', soft_edge: bool = True, norm_coors: bool = True, beta_schedule: str = 'sigmoid', beta_start: float = 1e-07, beta_end: float = 0.002, num_diffusion_timesteps: int = 1000)¶

Bases: torch.nn.Module

Dual-encoder, pocket-conditioned score network.

forward(batch, return_unreduced_loss: bool = False)¶

Score-matching loss for one batch of pocket-ligand complexes.

batch is any attribute-access object carrying the keys pmdm_collate emits (a SimpleNamespace in practice).

inpainting_sample(ligand_atom_type: torch.Tensor, ligand_pos_init: torch.Tensor, ligand_bond_index: torch.Tensor, ligand_bond_type: torch.Tensor | None, ligand_batch: torch.Tensor, frag_mask: torch.Tensor, protein_atom_feature_full: torch.Tensor, protein_pos: torch.Tensor, protein_batch: torch.Tensor, num_graphs: int, n_steps: int = 100, step_lr: float = 1e-06, clip: float = 1000.0, clip_local: float | None = None, clip_pos: float | None = None, global_start_sigma: float = float('inf'), local_start_sigma: float = float('inf'), w_global_pos: float = 1.0, w_global_node: float = 1.0, w_local_pos: float = 1.0, w_local_node: float = 1.0, sampling_type: str = 'generalized', eta: float = 1.0, keep_traj: bool = False) → Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor, List[torch.Tensor]]¶

RePaint-style constrained sampling (“lead optimisation” upstream): keep frag_mask atoms of a starting ligand fixed, regenerate the rest, inside the same fixed pocket.

frag_mask (bool, (n_ligand,)): True = keep exactly as given in ligand_pos_init/ligand_atom_type; False = regenerate. Every reverse step re-noises the kept fragment’s atom type to the current timestep before scoring, then force-restores both its position and type afterward – classic RePaint. Position is never re-noised (upstream’s matching line is dead/commented-out code, not ported). Port of upstream inpainting_sample, MDM_pocket_coor_shared.py:929-1183.

langevin_dynamics_sample(ligand_atom_type: torch.Tensor, ligand_pos_init: torch.Tensor, ligand_bond_index: torch.Tensor, ligand_bond_type: torch.Tensor | None, ligand_batch: torch.Tensor, protein_atom_feature_full: torch.Tensor, protein_pos: torch.Tensor, protein_batch: torch.Tensor, num_graphs: int, n_steps: int = 100, step_lr: float = 1e-06, clip: float = 1000.0, clip_local: float | None = None, clip_pos: float | None = None, global_start_sigma: float = float('inf'), local_start_sigma: float = float('inf'), w_global_pos: float = 1.0, w_global_node: float = 1.0, w_local_pos: float = 1.0, w_local_node: float = 1.0, sampling_type: str = 'generalized', eta: float = 1.0, keep_traj: bool = False) → Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor, List[torch.Tensor]]¶

Reverse process inside a fixed pocket.

Returns (ligand_pos, pos_traj, ligand_atom_type, atom_traj), flat and concatenated (not padded). Final coordinates are translated back into the input pocket’s frame.

linker_sample(ligand_atom_type: torch.Tensor, ligand_pos_init: torch.Tensor, ligand_bond_index: torch.Tensor, ligand_bond_type: torch.Tensor | None, ligand_batch: torch.Tensor, frag_mask: torch.Tensor, protein_atom_feature_full: torch.Tensor, protein_pos: torch.Tensor, protein_batch: torch.Tensor, num_graphs: int, n_steps: int = 100, step_lr: float = 1e-06, clip: float = 1000.0, clip_local: float | None = None, clip_pos: float | None = None, global_start_sigma: float = float('inf'), local_start_sigma: float = float('inf'), w_global_pos: float = 1.0, w_global_node: float = 1.0, w_local_pos: float = 1.0, w_local_node: float = 1.0, sampling_type: str = 'generalized', eta: float = 1.0, keep_traj: bool = False) → Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor, List[torch.Tensor]]¶

RePaint-style constrained sampling for linker design: keep two (or more) disjoint fragments of a starting ligand fixed (frag_mask), regenerate the region between them, inside the same fixed pocket.

Structurally identical to inpainting_sample() (same restore-after-every-step loop), with two differences: the kept fragment is fed to the network exactly as given, with no per-step forward-noise on its atom type (upstream computes a matching noise tensor here but never uses it – dead code, not ported); and the network additionally receives linker_mask (the region being regenerated), which biases its coordinate update toward that region (see encoders.py). Port of upstream linker_sample, MDM_pocket_coor_shared.py:1186-1425.

net(ligand_atom_type: torch.Tensor, ligand_pos: torch.Tensor, ligand_bond_index: torch.Tensor, ligand_bond_type: torch.Tensor | None, ligand_batch: torch.Tensor, protein_embeddings: torch.Tensor, protein_pos: torch.Tensor, protein_batch: torch.Tensor, time_step: torch.Tensor, linker_mask: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]¶

One denoiser pass. Returns the global/local position and node scores plus the ligand edge bookkeeping the loss needs.

linker_mask (bool, (n_ligand,), True = regenerate): used only by linker_sample(), to bias the EGNN coordinate update toward the region being regenerated (see encoders.py’s module docstring). None (default) reproduces training/default-sampling behaviour exactly.

T¶

alphas¶

atten_layer¶

betas¶

cutoff = 3.0¶

edge_encoder_global¶

edge_encoder_local¶

edge_order = 3¶

encoder_global¶

encoder_local¶

g_cutoff = 6.0¶

grad_global_node_mlp¶

grad_local_node_mlp¶

hidden_dim = 128¶

ligand_encoder¶

num_atom = 10¶

num_timesteps¶

protein_encoder¶

temb¶

temb_proj¶

MolecularDiffusion.modules.models.pmdm.epsnet.get_beta_schedule(beta_schedule: str, beta_start: float, beta_end: float, num_diffusion_timesteps: int) → numpy.ndarray¶: Upstream’s noise schedules (sigmoid is what the released config uses).