MolecularDiffusion.modules.tasks.diffusion_diffpharma

DiffPharma task: pocket- and pharmacophore-conditioned ligand diffusion.

Three objects live here, mirroring the platform’s usual layout:

  • DiffPharmaTask – the duck-typed Task (docs/adding_new_models.md Section 2.1) wrapping ConditionalDDPM.

  • DiffPharmaTaskFactory – the _target_ of configs/tasks/diffusion_diffpharma.yaml.

  • DiffPharmaPocketGenerator – the _target_ of configs/interference/gen_diffpharma_pocket.yaml. GenerativeFactory cannot express “sample inside this pocket”; a generator behind its own _target_ is the in-tree pattern for that (cf. DiffSMolShapeGenerator) and needs no core change – cli/generate.py only does instantiate(cfg.interference, task=task) then run().

The batch is NOT a PointCloud dict: DiffPharma needs four node sets (ligand, pocket, H-bond particles, hydrophobic particles), flat-concatenated with scatter masks. See data/component/diffpharma_data.py.

Attributes

DIFFPHARMA_ATOM_VOCAB

FLOAT_TYPE

INT_TYPE

NODE_SETS

Classes

DiffPharmaPocketGenerator

Pocket-conditioned generation behind interference/gen_diffpharma_pocket.

DiffPharmaTask

Task contract around ConditionalDDPM.

DiffPharmaTaskFactory

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

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_diffpharma.DiffPharmaPocketGenerator(task, pocket_db: str | None = None, pocket_index: int = 0, pocket_pdb: str | None = None, ref_sdf: str | None = None, num_generate: int = 20, batch_size: int = 4, num_steps: int | None = None, mol_size: list | None = None, output_path: str = 'generated_diffpharma', seed: int = 42, device: str | None = None, **kwargs: Any)

Bases: MolecularDiffusion.modules.tasks.pocket_generator.PocketGenerator

Pocket-conditioned generation behind interference/gen_diffpharma_pocket.

Two mutually exclusive pocket sources:

  • pocket_db + pocket_index – a row of a converted ASE db (scripts/convert_dataset.py), read with center=False.

  • pocket_pdb + ref_sdf – a novel pocket, built on the fly by data.component.diffpharma_prep.complex_from_files. The SDF is required, not optional: it is both the 8 A pocket-selection reference and the ligand ODDT detects the interactions against.

Either way the downstream dict is identical.

The sampling loop itself lives in PocketGenerator.

pocket_pdb = None
ref_sdf = None
seed_numpy = False
tag = 'diffpharma'
class MolecularDiffusion.modules.tasks.diffusion_diffpharma.DiffPharmaTask(model: MolecularDiffusion.modules.models.diffpharma.ConditionalDDPM, atom_vocab: List[str] | None = None)

Bases: torch.nn.Module

Task contract around ConditionalDDPM.

evaluate(pred, target)
forward(batch)
predict_and_target(batch)
sample(batch_size=None, nodesxsample=None, num_steps=None, batch=None, **kwargs)

Sample ligands inside the pocket carried by batch.

batch must hold the pocket_*/interh_*/interhp_* keys of a collated DiffPharma batch (the generator builds it). There is no unconditional mode – mirroring upstream’s own hard refusal.

Returns (one_hot, charges, coords, node_mask) padded to (B, N, .), in the ORIGINAL pocket frame (the reverse process re-centres on the ligand CoM; the shift is undone here).

atom_vocab
property device
model
property n_node_dist: Dict[int, float]

Ligand-size marginal. Not on the generation path (see module doc).

node_dist_model
prop_dist_model = None
split = 'train'
class MolecularDiffusion.modules.tasks.diffusion_diffpharma.DiffPharmaTaskFactory(task_type: str = 'diffusion_diffpharma', size_distribution_path: str | None = None, atom_nf: int = 11, residue_nf: int = 11, interh_nf: int = 3, interhp_nf: int = 6, n_dims: int = 3, joint_nf: int = 128, hidden_nf: int = 256, n_layers: int = 8, attention: bool = True, tanh: bool = True, norm_constant: float = 1, inv_sublayers: int = 1, sin_embedding: bool = False, aggregation_method: str = 'sum', normalization_factor: float = 100, edge_cutoff_ligand: float | None = None, edge_cutoff_pocket: float | None = 5.0, edge_cutoff_interaction: float | None = 5.0, reflection_equivariant: bool = False, diffusion_steps: int = 500, diffusion_noise_schedule: str = 'polynomial_2', diffusion_noise_precision: float = 0.0005, diffusion_loss_type: str = 'l2', normalize_factors=(1, 1), atom_vocab: List[str] | None = None, **kwargs: Any)

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

build() → DiffPharmaTask
atom_vocab
condition_names: List[str] = []
ddpm_kwargs
egnn_kwargs
size_distribution_path = None
task: DiffPharmaTask | None = None
task_type = 'diffusion_diffpharma'
MolecularDiffusion.modules.tasks.diffusion_diffpharma.DIFFPHARMA_ATOM_VOCAB = ['C', 'N', 'O', 'S', 'B', 'Br', 'Cl', 'P', 'I', 'F', 'others']
MolecularDiffusion.modules.tasks.diffusion_diffpharma.FLOAT_TYPE
MolecularDiffusion.modules.tasks.diffusion_diffpharma.INT_TYPE
MolecularDiffusion.modules.tasks.diffusion_diffpharma.NODE_SETS = ('lig', 'pocket', 'interh', 'interhp')