MolecularDiffusion.modules.tasks.diffusion_diffsbdd¶

DiffSBDD task: pocket-conditioned (and joint) 3D ligand diffusion.

Four objects, mirroring the platform’s usual layout:

DiffSBDDTask – the duck-typed Task (docs/adding_new_models.md Section 2.1) wrapping either ConditionalDDPM or EnVariationalDiffusion, selected by mode.
ModelTaskFactory – the _target_ of configs/tasks/diffusion_diffsbdd.yaml (and of its 4-line diffusion_diffsbdd_joint.yaml sibling, which only overrides mode).
DiffSBDDPocketGenerator – the _target_ of BOTH configs/interference/gen_diffsbdd_pocket.yaml and gen_diffsbdd_inpaint.yaml; they differ only in whether fixed_atom_indices is null. GenerativeFactory’s sample(batch_size, nodesxsample, ...) has no channel for “which pocket”, so pocket-conditioned models get their own generator behind their own _target_ – the established in-tree pattern (KGDiffPocketGenerator, PMDMPocketGenerator), needing no core change: cli/generate.py only does instantiate(cfg.interference, task=task) then .run().
LigandPocketSizePrior – the ligand-size marginal of the model’s 2D joint (n_lig, n_pocket) histogram, wrapped to satisfy the node_dist_model contract.

Deviations from the generic contract, both shared with the other pocket models in-tree:

sample() requires a pocket. ConditionalDDPM.sample() itself raises without one, so interference: gen_unconditional does not apply here.
Sampled coordinates come back in the input pocket’s frame – the generator restores the pocket CoM exactly as lightning_modules.py:845-852 does.

## The mode guard (read this before touching mode_id)

Both modes share task_type: diffusion_diffsbdd AND – verified against the released weights – identical state-dict keys and shapes (122 tensors each; update_pocket_coords is a flag, not a parameter). The integration plan expected a shape mismatch to catch a joint checkpoint loaded under a pocket_conditioning config; there is none, and cli/generate.py:322 loads with strict=False, so such a mix-up would run silently on the wrong sampler. The mode_id buffer below is the fix: it round-trips through the checkpoint, and DiffSBDDTask._check_mode() refuses to train or sample when the loaded value disagrees with the configured mode.

Out of scope this pass (see the integration plan): optimize.py / diversify(), SimpleConditionalDDPM, CA pocket representation, virtual nodes, the auxiliary Lennard-Jones loss, trajectory export, Vina docking, and EnVariationalDiffusion.sample() (the fully unconditional path that hallucinates a pocket).

Attributes¶

`FLOAT_TYPE`
`INT_TYPE`
`MODES`

Classes¶

`DiffSBDDPocketGenerator`	Pocket-conditioned generation and inpainting.
`DiffSBDDTask`	Task contract around `ConditionalDDPM` / `EnVariationalDiffusion`.
`LigandPocketSizePrior`	Ligand-size prior over the model's 2D joint histogram.
`ModelTaskFactory`	Hydra entry point for `configs/tasks/diffusion_diffsbdd*.yaml`.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_diffsbdd.DiffSBDDPocketGenerator(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, fixed_atom_indices: list | None = None, add_n_nodes: int | None = None, resamplings: int = 1, center: str = 'ligand', jump_length: int = 1, output_path: str = 'generated_diffsbdd', seed: int = 42, device: str | None = None, **kwargs: Any)¶

Bases: MolecularDiffusion.modules.tasks.pocket_generator.PocketGenerator

Pocket-conditioned generation and inpainting.

One class behind both interference configs – they differ only in whether fixed_atom_indices is null, and a DiffSBDDInpaintGenerator would be an interface with one implementation.

The pocket (and, for inpainting, the reference ligand that defined it) comes from one row of the converted ASE db, read with center=False so sampled ligands come back in that pocket’s own frame.

Fixed atoms are selected by index into the stored reference ligand (fixed_atom_indices: [0, 1, 2, 3]) rather than by PDB atom name (upstream’s --fix_atoms "C1 N6 C5"): the db has no atom names, and the same lig_fixed tensor reaches inpaint() either way.

The sampling loop itself lives in PocketGenerator.

add_n_nodes = None¶

center = 'ligand'¶

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

fixed_atom_indices¶

jump_length = 1¶

resamplings = 1¶

tag = 'diffsbdd'¶

class MolecularDiffusion.modules.tasks.diffusion_diffsbdd.DiffSBDDTask(model: torch.nn.Module, mode: str = 'pocket_conditioning', atom_vocab: List[str] | None = None)¶

Bases: torch.nn.Module

Task contract around ConditionalDDPM / EnVariationalDiffusion.

evaluate(pred, target)¶

forward(batch)¶

Loss assembly ported from lightning_modules.py:236-302.

Mode-independent: both models return the same keys, the conditional one substituting a constant 0 for the pocket terms. l2 with per-example dimension normalisation while training, the full VLB otherwise.

predict_and_target(batch)¶

sample(batch_size=None, nodesxsample=None, num_steps=None, batch=None, ligand=None, lig_fixed=None, resamplings: int = 1, center: str = 'ligand', jump_length: int = 1, **kwargs)¶

Sample ligands inside the pocket carried by batch.

The signature deliberately deviates from Section 2.1: it takes a pocket. batch must hold pocket_coords / pocket_one_hot / pocket_mask / num_pocket_nodes from a collated DiffSBDD batch (DiffSBDDPocketGenerator builds it).

The whole cost of joint mode + inpainting is the three-way branch below; the last two rows of the plan’s table are literally the same call, which is upstream’s own design (lightning_modules.py:814-834), not a shortcut:

mode	lig_fixed	call
pocket_conditioning	None	`sample_given_pocket`
pocket_conditioning	given	`ConditionalDDPM.inpaint`
joint	None->zeros	`EnVariationalDiffusion.inpaint`
joint	given	same call

NB the mode kwarg that GenerativeFactory passes ("ddpm" / "ddim") is swallowed by **kwargs and is NOT this task’s generative self.mode.

Returns (one_hot, charges, coords, node_mask) padded to (B, N, .) in the INPUT POCKET’s frame. charges is zeros: DiffSBDD has no charge channel.

atom_vocab¶

property device¶

mode = 'pocket_conditioning'¶

model¶

property n_node_dist: Dict[int, float]¶

property node_dist_model: LigandPocketSizePrior¶

prop_dist_model = None¶

split = 'train'¶

class MolecularDiffusion.modules.tasks.diffusion_diffsbdd.LigandPocketSizePrior(histogram: torch.Tensor)¶

Ligand-size prior over the model’s 2D joint histogram.

.sample(n) returns the ligand-axis marginal, which is what the Section 2.1 node_dist_model contract asks for. The pocket generator does not use it: it calls sample_conditional(n2=pocket_sizes) on the model’s own DistributionNodes, which is the native behaviour and the only one that respects pocket size.

sample(n_samples: int) → torch.Tensor¶

marginal¶

n_node_dist¶

class MolecularDiffusion.modules.tasks.diffusion_diffsbdd.ModelTaskFactory(task_type: str = 'diffusion_diffsbdd', mode: str = 'pocket_conditioning', atom_nf: int = NUM_ATOM_CLASSES, residue_nf: int = NUM_ATOM_CLASSES, n_dims: int = 3, joint_nf: int = 32, hidden_nf: int = 128, n_layers: int = 5, inv_sublayers: int = 1, attention: bool = True, tanh: bool = True, norm_constant: float = 1, sin_embedding: bool = False, normalization_factor: float = 100, aggregation_method: str = 'sum', reflection_equivariant: bool = False, edge_cutoff_ligand: float | None = None, edge_cutoff_pocket: float | None = 5.0, edge_cutoff_interaction: float | None = 5.0, edge_embedding_dim: int | None = None, diffusion_steps: int = 500, diffusion_noise_schedule: str = 'polynomial_2', diffusion_noise_precision: float = 0.0005, normalize_factors: Sequence[float] = (1.0, 4.0), max_n_lig: int = 107, max_n_pocket: int = 1671, atom_vocab: List[str] | None = None, train_set: Any = None, **kwargs: Any)¶

Hydra entry point for configs/tasks/diffusion_diffsbdd*.yaml.

train_set is declared so cli/train.py’s declarative seam (docs/adding_new_models.md Section 2.5) injects the training set: the 2D joint size histogram is a genuine dataset statistic and is needed at train time too, for the log_pN VLB term. It is stored as a buffer on the model so it survives the checkpoint and generation can rebuild with train_set=None. The buffer’s shape is fixed by max_n_lig / max_n_pocket, so checkpoint shapes never depend on which db was used; the defaults are the released CrossDocked histogram’s own shape.

build() → DiffSBDDTask¶

atom_vocab¶

condition_names: List[str] = []¶

ddpm_kwargs¶

dynamics_kwargs¶

max_n_lig = 107¶

max_n_pocket = 1671¶

mode = 'pocket_conditioning'¶

task: DiffSBDDTask | None = None¶

task_type = 'diffusion_diffsbdd'¶

train_set = None¶

MolecularDiffusion.modules.tasks.diffusion_diffsbdd.FLOAT_TYPE¶

MolecularDiffusion.modules.tasks.diffusion_diffsbdd.INT_TYPE¶

MolecularDiffusion.modules.tasks.diffusion_diffsbdd.MODES: Tuple[str, ...] = ('pocket_conditioning', 'joint')¶