MolecularDiffusion.modules.tasks.diffusion_pmdm¶

PMDM task: pocket-conditioned 3D ligand diffusion.

Three objects, mirroring the platform’s usual layout:

PMDMDiffusionTask – the duck-typed Task (docs/adding_new_models.md Section 2.1) wrapping PMDMEpsNet.
ModelTaskFactory – the _target_ of configs/tasks/diffusion_pmdm.yaml.
PMDMPocketGenerator – the _target_ of configs/interference/gen_pmdm_pocket.yaml. GenerativeFactory’s sample(batch_size, nodesxsample, ...) has no channel for “which pocket”, so pocket-conditioned models get their own generator behind their own _target_; that is the established in-tree pattern (DiffPharmaPocketGenerator, the DiffSMol shape generator) and needs no core change – cli/generate.py only does instantiate(cfg.interference, task=task) then .run().

The batch is NOT a PointCloud dict: PMDM needs a diffused ligand cloud plus a fixed pocket cloud, flat-concatenated with scatter indices. The collate in data/component/pmdm_data.py already emits PMDM’s own attribute names, so the adapter here is a one-line SimpleNamespace(**batch).

PMDMConstrainedGenerator adds the two constrained-sampling modes upstream calls “lead optimisation” and “linker design” (mode: lead_opt / mode: linker) – keep part of a starting ligand fixed, regenerate the rest, inside the same fixed pocket. Both modes read the pocket and the starting ligand from plain files (pocket_file/mol_file), not a converted db – see docs/model_integrations/pmdm/INTEGRATION_PLAN.md’s revision-3 Q1/Q1b. One class for both modes, not two – the same pattern DiffSBDD already uses for its own de novo/inpaint split.

Out of scope this pass (see the integration plan): the VAE-latent and property-context branches, and CFG / gradient guidance.

Attributes¶

`ATOMIC_NUMBERS`
`BOND_TOLERANCE`
`INT_TYPE`
`MAX_LIGAND_SIZE`
`MIN_ATOM_SEPARATION`

Classes¶

`LigandSizeDistribution`	Ligand-atom-count prior, sampled at generation time.
`ModelTaskFactory`	Hydra entry point for `configs/tasks/diffusion_pmdm.yaml`.
`PMDMConstrainedGenerator`	Constrained generation (lead optimisation OR linker design), behind
`PMDMDiffusionTask`	Task contract around `PMDMEpsNet`.
`PMDMPocketGenerator`	Pocket-conditioned generation behind `interference/gen_pmdm_pocket`.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_pmdm.LigandSizeDistribution(histogram: Dict[int, float])¶

Ligand-atom-count prior, sampled at generation time.

Same shape of interface as the other in-tree node distributions (TabascoNodeDistribution, DiffSMolShapeNodeDistribution): .histogram / .n_node_dist / .sample(n). Written here rather than imported from diffusion_tabasco because that module hard-imports tensordict and the whole TABASCO flow stack at module scope.

sample(n_samples: int) → torch.Tensor¶

histogram¶

n_node_dist¶

class MolecularDiffusion.modules.tasks.diffusion_pmdm.ModelTaskFactory(task_type: str = 'diffusion_pmdm', 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, atom_vocab: List[str] | None = None, train_set: Any = None, **kwargs: Any)¶

Hydra entry point for configs/tasks/diffusion_pmdm.yaml.

build() → PMDMDiffusionTask¶

atom_vocab¶

condition_names: List[str] = []¶

model_kwargs¶

task: PMDMDiffusionTask | None = None¶

task_type = 'diffusion_pmdm'¶

train_set = None¶

class MolecularDiffusion.modules.tasks.diffusion_pmdm.PMDMConstrainedGenerator(task, mode: str = 'lead_opt', pocket_db: str | None = None, pocket_index: int = 0, mol_size: list | None = None, pocket_file: str | None = None, mol_file: str | None = None, fixed_atoms: List[int] | None = None, atoms_to_replace: List[int] | None = None, atoms_to_add: int | None = None, num_generate: int = 20, batch_size: int = 4, num_steps: int | None = None, sampling_type: str = 'generalized', step_lr: float = 1e-06, clip: float = 1000.0, clip_pos: float | None = None, global_start_sigma: float | None = None, w_global_pos: float = 1.0, w_global_node: float = 1.0, w_local_pos: float = 1.0, w_local_node: float = 1.0, output_path: str = 'generated_pmdm_constrained', seed: int = 42, device: str | None = None, validity_filter: bool = True, max_retries: int = 10, **kwargs: Any)¶

Bases: PMDMPocketGenerator

Constrained generation (lead optimisation OR linker design), behind interference/gen_pocket (mode: lead_opt / mode: linker) – see pocket_generator.py’s _TASK_TO_GENERATOR for the dispatch.

One class for both – following the same pattern DiffSBDD already uses for its own de novo/inpaint split (gen_diffsbdd_pocket.yaml vs gen_diffsbdd_inpaint.yaml, same _target_): the two upstream sampling methods (inpainting_sample/linker_sample) share their entire loop, differing only in which atoms are held fixed and which internal EGNN masking mode fires (PMDMDiffusionTask.sample’s mode dispatch). A second class here would be exactly the “interface with one implementation” gen_diffsbdd_inpaint.yaml’s own header warns against.

mode picks the task and which of the two mutually-exclusive atom lists is read:

"lead_opt" – keep fixed_atoms of a starting ligand (mol_file) exactly as given, grow atoms_to_add brand-new atoms onto it. Mirrors upstream’s sample_frag.py --keep_index.
"linker" – delete atoms_to_replace (typically the gap between two fragments) and regrow atoms_to_add atoms in their place. Mirrors upstream’s sample_linker.py --mask.

Both read the pocket from pocket_file – a plain PDB, NOT a converted db, unlike PMDMPocketGenerator’s de novo path (which stays a separate class: its input genuinely is a different source, not just a different knob – a converted ASE db row, not a live file).

db_required_msg = None¶

keep_atoms¶

mode = 'lead_opt'¶

mol_file = None¶

n_new_atoms = 0¶

pocket_file = None¶

replace_atoms¶

tag = 'pmdm-leadopt'¶

class MolecularDiffusion.modules.tasks.diffusion_pmdm.PMDMDiffusionTask(model: MolecularDiffusion.modules.models.pmdm.PMDMEpsNet, atom_vocab: List[str] | None = None, size_histogram: torch.Tensor | None = None)¶

Bases: torch.nn.Module

Task contract around PMDMEpsNet.

evaluate(pred, target)¶

forward(batch)¶

predict_and_target(batch)¶

sample(batch_size=None, nodesxsample=None, num_steps=None, batch=None, mode: str | None = None, n_new_atoms: int | None = None, **kwargs)¶

Sample ligands inside the pocket carried by batch.

The signature deliberately deviates from Section 2.1: it takes a pocket. batch must hold the protein_* keys of a collated PMDM batch (PMDMPocketGenerator builds it). There is no unconditional mode – the network has no path that runs without a pocket.

mode (None | "lead_opt" | "linker"): None (default) is the existing de novo path, unchanged. Either other value is constrained sampling – batch must ALSO carry frag_ligand_pos/frag_ligand_atom_feature/frag_element_batch (the kept fragment, tiled once per sample – PMDMConstrainedGenerator builds these), and n_new_atoms is required (how many brand-new atoms to grow, mirroring upstream’s single --num_atom). Dispatches to PMDMEpsNet.inpainting_sample ("lead_opt") or linker_sample ("linker") instead of langevin_dynamics_sample.

Returns (one_hot, charges, coords, node_mask) padded to (B, N, .), in the ORIGINAL pocket frame. charges is zeros: PMDM’s pocket configuration has no charge channel.

atom_vocab¶

property device¶

model¶

property n_node_dist: Dict[int, float]¶

property node_dist_model: LigandSizeDistribution¶

prop_dist_model = None¶

split = 'train'¶

class MolecularDiffusion.modules.tasks.diffusion_pmdm.PMDMPocketGenerator(task, pocket_db: str | None = None, pocket_index: int = 0, num_generate: int = 20, batch_size: int = 4, num_steps: int | None = None, mol_size: list | None = None, sampling_type: str = 'generalized', step_lr: float = 1e-06, clip: float = 1000.0, clip_pos: float | None = None, global_start_sigma: float | None = None, w_global_pos: float = 1.0, w_global_node: float = 1.0, w_local_pos: float = 1.0, w_local_node: float = 1.0, output_path: str = 'generated_pmdm', seed: int = 42, device: str | None = None, validity_filter: bool = True, max_retries: int = 10, **kwargs: Any)¶

Bases: MolecularDiffusion.modules.tasks.pocket_generator.PocketGenerator

Pocket-conditioned generation behind interference/gen_pmdm_pocket.

The pocket comes from one row of a converted ASE db (docs/model_integrations/pmdm/scripts/convert_dataset.py), read with center=False so the sampled ligand comes back in that pocket’s own frame.

The only generator that may write FEWER than num_generate: unusable samples are filtered out and resampled, bounded by max_retries.

The sampling loop itself lives in PocketGenerator.

db_required_msg = 'interference.pocket_db is required: PMDM has no unconditional mode. Point it at a converted ASE...¶

max_retries = 10¶

sampler_kwargs¶

seed_numpy = False¶

tag = 'pmdm'¶

validity_filter = True¶

MolecularDiffusion.modules.tasks.diffusion_pmdm.ATOMIC_NUMBERS¶

MolecularDiffusion.modules.tasks.diffusion_pmdm.BOND_TOLERANCE = 0.4¶

MolecularDiffusion.modules.tasks.diffusion_pmdm.INT_TYPE¶

MolecularDiffusion.modules.tasks.diffusion_pmdm.MAX_LIGAND_SIZE = 128¶

MolecularDiffusion.modules.tasks.diffusion_pmdm.MIN_ATOM_SEPARATION = 0.9¶