MolecularDiffusion.modules.tasks.diffusion_diffdec

DiffDec: pocket-aware scaffold decoration with an end-to-end diffusion model.

DiffDec (https://github.com/biomed-AI/DiffDec, Xie et al., J. Chem. Inf. Model. 64(7) 2554-2564, 2024, doi:10.1021/acs.jcim.3c01466) grows an R-group off a named anchor atom of a fixed 3D scaffold, inside a fixed protein pocket. Scaffold and pocket atoms are never noised; only the R-group rows are diffused. See docs/model_integrations/diffdec/INTEGRATION_PLAN.md for the full integration plan (scope, data adapters, task-contract mapping).

No new model package. DiffDec is a fork of DiffLinker, and its src/edm_single.py, src/noise.py and src/egnn.py are the ported MolecularDiffusion.modules.models.difflinker modules modulo formatting and a fragment/linker -> scaffold/rgroup rename (verified by normalized diff – INTEGRATION_PLAN.md “Repo Inspection”). So this task reuses modules.models.difflinker rather than copying a fourth near-identical EDM/EGNN into the tree. That package is frozen ported upstream code now shared by two tasks (diffusion_difflinker.py and this one) – changing it changes both.

The mask-name mapping across that seam is:

DiffDec                     difflinker.EDM / Dynamics
---------------------------------------------------------
scaffold_mask (scaf+pocket) fragment_mask   (never noised)
rgroup_mask                 linker_mask     (diffused)
context[..., -2] scaffold_only_mask         fragment_only_mask
context[..., -1] pocket_only_mask           pocket_only_mask

DynamicsWithPockets is used with graph_type="4A", i.e. a 4 A radius graph rebuilt each forward pass, which is exactly DiffDec’s get_dist_edges (egnn.py l. 531-539).

Generation goes through DiffDecScaffoldGenerator, not through GenerativeFactory: DiffDec has no unconditional mode, and sample(batch_size, nodesxsample) has no channel for “which scaffold, which pocket, which anchor”. That is the established pattern for every pocket-conditioned model here (DiffSBDDPocketGenerator, gen_kgdiff_pocket.yaml, gen_pmdm_pocket.yaml).

Classes

DiffDecScaffoldGenerator	Scaffold decoration inside a fixed pocket.
DiffDecTask	Plain nn.Module task wrapper (TABASCO-style -- see
DiffDecTaskFactory	Factory matching cli/train.py's task_module.build() /

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_diffdec.DiffDecScaffoldGenerator(task: DiffDecTask, data_file: str | None = None, complex_index: int = 0, num_generate: int = 20, batch_size: int = 4, num_steps: int | None = None, output_path: str = 'generated_diffdec', save_reference: bool = True, seed: int = 42, device: str | None = None, **kwargs: Any)

Scaffold decoration inside a fixed pocket.

cli/generate.py:673 does instantiate(cfg.interference, task=task) then .run(), so every key of the interference config lands straight in this __init__ – nothing needs registering. Precedent: DiffSBDDPocketGenerator.

The scaffold, its anchor atom and the pocket all come from ONE row of upstream’s preprocessed .pt (data_file + complex_index). num_generate R-groups are then sampled for that one input, which is upstream’s own protocol (sample_single.py: --n_samples decorations per complex).

run() → None

batch_size = 4

complex_index = 0

data_file = None

device = 'cuda'

num_generate = 20

num_steps = None

output_path = 'generated_diffdec'

save_reference = True

seed = 42

task

class MolecularDiffusion.modules.tasks.diffusion_diffdec.DiffDecTask(in_node_nf: int, n_dims: int, hidden_nf: int, activation: str, tanh: bool, n_layers: int, attention: bool, norm_constant: float, inv_sublayers: int, sin_embedding: bool, normalization_factor: float, aggregation_method: str, model: str, normalization: str | None, condition_time: bool, anchors_context: bool, diffusion_steps: int, diffusion_noise_schedule: str, diffusion_noise_precision: float, diffusion_loss_type: str, normalize_factors: tuple, center_of_mass: str, atom_vocab: list | None = None)

Bases: torch.nn.Module

Plain nn.Module task wrapper (TABASCO-style – see docs/adding_new_models.md §2.6) implementing the §2.1 contract around DiffLinker’s EDM driven by DiffDec’s mask/context convention.

self.edm is deliberately named to match upstream’s own attribute path (DDPM.edm), so the released diffdec_single.ckpt’s edm.* keys land without renaming – see docs/model_integrations/diffdec/scripts/convert_checkpoint.py.

decorate(batch: Dict[str, Any], rgroup_sizes: torch.Tensor | None = None, keep_frames: int = 1)

Grow R-groups onto the scaffolds in batch.

Port of model_single.py::DDPM.sample_chain (l. 289-340). Returns (one_hot, positions, ligand_mask) where ligand_mask is atom_mask - pocket_mask – scaffold + R-group, pocket dropped – matching upstream sample_single.py l. 148.

Coordinates come back in the centred frame (partial COM removed); the caller adds the offset back, as upstream does at sample_single.py l. 146.

rgroup_sizes defaults to rgroup_mask.sum(1), i.e. the fixed 10-slot budget – upstream’s sample_fn = None path. Unused slots come back as the fake '#' atom.

evaluate(pred: torch.Tensor, target: torch.Tensor) → dict

forward(batch: Dict[str, Any])

get_extra_state() → dict

predict_and_target(batch: Dict[str, Any])

sample(*args: Any, **kwargs: Any)

Not reachable through GenerativeFactory – see the module docstring and INTEGRATION_PLAN.md’s Task-contract mapping.

The §2.1 signature is kept so the contract check passes, but DiffDec cannot generate without a scaffold, a pocket and an anchor atom, and sample(batch_size, nodesxsample, ...) carries none of those. Use DiffDecScaffoldGenerator (configs/interference/ gen_diffdec_scaffold.yaml), which calls decorate().

set_extra_state(state: dict) → None

anchors_context

property atom_count_histogram: dict | None

atom_vocab

center_of_mass

property device: torch.device

edm

in_node_nf

loss_type

property model

n_dims

property n_node_dist: dict | None

ndim_extra = 0

property node_dist_model: MolecularDiffusion.modules.models.difflinker.linker_size.DistributionNodes | None

norm_values

prop_dist_model = None

class MolecularDiffusion.modules.tasks.diffusion_diffdec.DiffDecTaskFactory(task_type: str = 'diffusion_diffdec', in_node_nf: int = 10, n_dims: int = 3, hidden_nf: int = 128, activation: str = 'silu', tanh: bool = False, n_layers: int = 6, attention: bool = False, norm_constant: float = 1e-06, inv_sublayers: int = 2, sin_embedding: bool = False, normalization_factor: float = 100, aggregation_method: str = 'sum', model: str = 'egnn_dynamics', normalization: str | None = 'batch_norm', condition_time: bool = True, anchors_context: bool = True, diffusion_steps: int = 500, diffusion_noise_schedule: str = 'polynomial_2', diffusion_noise_precision: float = 1e-05, diffusion_loss_type: str = 'l2', normalize_factors: tuple = (1, 4, 10), center_of_mass: str = 'anchors', atom_vocab: list | None = None, **kwargs: Any)

Factory matching cli/train.py’s task_module.build() / task_module.task instantiation pattern (precedent: diffusion_difflinker.py::DiffLinkerTaskFactory).

Defaults follow DiffDec configs/single.yml + train_single.py’s argparse defaults, cross-checked against the released diffdec_single.ckpt’s own hyper_parameters.

build() → DiffDecTask

activation = 'silu'

aggregation_method = 'sum'

anchors_context = True

atom_vocab

attention = False

center_of_mass = 'anchors'

condition_time = True

diffusion_loss_type = 'l2'

diffusion_noise_precision = 1e-05

diffusion_noise_schedule = 'polynomial_2'

diffusion_steps = 500

hidden_nf = 128

in_node_nf = 10

inv_sublayers = 2

kwargs

model = 'egnn_dynamics'

n_dims = 3

n_layers = 6

norm_constant = 1e-06

normalization = 'batch_norm'

normalization_factor = 100

normalize_factors = (1, 4, 10)

sin_embedding = False

tanh = False

task: DiffDecTask | None = None

task_type = 'diffusion_diffdec'