MolecularDiffusion.modules.tasks.pocket_generator¶

The one run() loop behind every pocket-conditioned generator.

Seven models (DiffSBDD, DiffInt, KGDiff, IPDiff, PMDM, Apo2Mol, DiffPharma) each need “load one pocket, tile it, sample ligands into it, write .xyz”. That loop used to be copy-pasted seven times. It lives here once; what genuinely differs per model lives in a handful of hooks:

hook	what it decides
`tag`	log prefix
`seed_numpy`	whether `np.random.seed` is called (see below)
`max_retries`	bound on resampling rounds; `None` => unbounded
`db_required_msg`	the “pocket_db is required” error, or `None`
`_pocket()`	read one pocket item (which Dataset, which centering)
`_repeat()`	tile that item `n` times, with fresh scatter indices
`_sizes()` /
`_pick_sizes()`	how many atoms each sampled ligand gets
`_sample_kwargs()`	the model-specific kwargs of `task.sample()`
`_accept()`	which of a batch’s molecules are written out
`_after_batch()`	side outputs per batch (Apo2Mol’s pocket .pdb)
`_start()` /
`_summary()`	the header / footer the run prints

``seed_numpy`` is not a style knob. PMDM and DiffPharma have never seeded numpy, the other five always have. Seeding it everywhere would change what those two generate; that is a behaviour change, not a fix.

The loop is not pocket-only. A second family rides it through TSGenerator, which fills these hooks once for transition-state models – the fixed structural context there is a reaction rather than a protein pocket. Add a new TS model by subclassing that, not this.

Two ways to reach a generator:

interference/gen_<model>_pocket.yaml names the concrete subclass, and carries that model’s own knobs (guidance weights, sampler settings, inpainting);
interference/gen_pocket.yaml names PocketGenerator itself and carries only the nine keys all seven share – the concrete subclass is then picked from the loaded task (see _for_task()).

Because that second path bypasses per-model Hydra key checking, the base __init__ rejects any key that no generator in the resolved chain declares, instead of ignoring it: a typo’d guide_mdoe must not silently do nothing. Note this is not a claim that per-model keys are refused on the shared config – dispatch reaches the concrete subclass first, so a key that subclass declares is consumed and honoured there.

Attributes¶

INT_TYPE

Classes¶

PocketGenerator

Load one pocket, sample num_generate ligands into it, write .xyz.

Module Contents¶

class MolecularDiffusion.modules.tasks.pocket_generator.PocketGenerator(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, output_path: str = 'generated_pocket', seed: int = 42, device: str | None = None, **kwargs: Any)¶

Load one pocket, sample num_generate ligands into it, write .xyz.

Instantiating this class directly (what gen_pocket.yaml does) returns the concrete subclass for the task that was passed in.

run() → None¶

batch_size = 4¶

db_required_msg: str | None = None¶

device = 'cuda'¶

max_retries: int | None = None¶

mol_size¶

num_generate = 20¶

num_steps = None¶

output_path = 'generated_pocket'¶

pocket_db = None¶

pocket_index = 0¶

seed = 42¶

seed_numpy = True¶

tag = 'pocket'¶

task¶

MolecularDiffusion.modules.tasks.pocket_generator.INT_TYPE¶