MolecularDiffusion.modules.tasks.diffusion_kgdiff

KGDiff task: pocket-conditioned ligand diffusion with self-guidance.

Three objects, mirroring the platform’s usual layout:

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

• ModelTaskFactory – the _target_ of configs/tasks/diffusion_kgdiff.yaml.

• KGDiffPocketGenerator – the _target_ of configs/interference/gen_kgdiff_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 (PMDMPocketGenerator, DiffPharmaPocketGenerator) and needs no core change – cli/generate.py only does instantiate(cfg.interference, task=task) then .run().

The batch is NOT a PointCloud dict: KGDiff needs a diffused ligand cloud, a fixed pocket cloud, and one affinity label per complex, flat-concatenated with scatter indices. The collate in data/component/kgdiff_data.py already emits KGDiff’s own argument names, so the adapter here is a .to(device) map.

Two deviations from the generic contract, both deliberate and both shared with the other pocket models in-tree:

• sample() requires a pocket. There is no unconditional path – the network cannot run without protein_pos/protein_v.

• Sampled coordinates come back in the input pocket’s frame, not zero-CoM: center_pos_mode='protein' subtracts the pocket centroid and the sampler adds it back.

Out of scope this pass (see the integration plan): the valuenet* / target_diff / vina guide modes, PDBBind2020, and CFG / inpainting / trajectory export.

Attributes

DEFAULT_POCKET_EXTENT
INT_TYPE

Classes

KGDiffDiffusionTask	Task contract around ScorePosNet3D.
KGDiffPocketGenerator	Pocket-conditioned generation behind interference/gen_kgdiff_pocket.
ModelTaskFactory	Hydra entry point for configs/tasks/diffusion_kgdiff.yaml.
PocketSizePrior	Ligand-size prior conditioned on the pocket's spatial extent.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_kgdiff.KGDiffDiffusionTask(model: MolecularDiffusion.modules.models.kgdiff.ScorePosNet3D, atom_vocab: List[str] | None = None)

Bases: torch.nn.Module

Task contract around ScorePosNet3D.

evaluate(pred, target)

forward(batch)

predict_and_target(batch)

sample(batch_size=None, nodesxsample=None, num_steps=None, batch=None, guide_mode: str = 'joint', type_grad_weight: float = 100.0, pos_grad_weight: float = 25.0, progress: bool = True, **kwargs)

Sample ligands inside the pocket carried by batch.

The signature deliberately deviates from Section 2.1: it takes a pocket. batch must hold protein_pos / protein_v / protein_batch from a collated KGDiff batch (KGDiffPocketGenerator builds it).

Returns (one_hot, charges, coords, node_mask) padded to (B, N, .) in the ORIGINAL pocket frame. one_hot is over the 8-element atom_vocab, not the model’s 13 (element, aromatic) classes. charges is zeros: KGDiff has no charge channel.

atom_vocab

property device

model

property n_node_dist: Dict[int, float]

property node_dist_model: PocketSizePrior

prop_dist_model = None

split = 'train'

class MolecularDiffusion.modules.tasks.diffusion_kgdiff.KGDiffPocketGenerator(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, guide_mode: str = 'joint', type_grad_weight: float = 100.0, pos_grad_weight: float = 25.0, output_path: str = 'generated_kgdiff', seed: int = 42, device: str | None = None, **kwargs: Any)

Bases: MolecularDiffusion.modules.tasks.pocket_generator.PocketGenerator

Pocket-conditioned generation behind interference/gen_kgdiff_pocket.

The pocket comes from one row of a converted ASE db (docs/model_integrations/kgdiff/scripts/convert_dataset.py); the sampled ligand comes back in that pocket’s own frame.

guide_mode='joint' is KGDiff’s headline self-guided sampler (the model’s own affinity head is the classifier guide); 'wo' runs the same loop unguided, which is the paper’s ablation and a useful control.

The sampling loop itself lives in PocketGenerator.

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

guide_mode = 'joint'

pos_grad_weight = 25.0

tag = 'kgdiff'

type_grad_weight = 100.0

class MolecularDiffusion.modules.tasks.diffusion_kgdiff.ModelTaskFactory(task_type: str = 'diffusion_kgdiff', protein_atom_feature_dim: int = PROTEIN_FEATURE_DIM, ligand_atom_feature_dim: int = NUM_LIGAND_CLASSES, model_mean_type: str = 'C0', beta_schedule: str = 'sigmoid', beta_start: float = 1e-07, beta_end: float = 0.002, pos_beta_s: float = 0.01, v_beta_schedule: str = 'cosine', v_beta_s: float = 0.01, num_diffusion_timesteps: int = 1000, loss_v_weight: float = 100.0, loss_exp_weight: float = 1.0, sample_time_method: str = 'symmetric', use_classifier_guide: bool = True, time_emb_dim: int = 0, time_emb_mode: str = 'simple', center_pos_mode: str = 'protein', node_indicator: bool = True, model_type: str = 'uni_o2', num_blocks: int = 1, num_layers: int = 9, hidden_dim: int = 128, n_heads: int = 16, edge_feat_dim: int = 4, num_r_gaussian: int = 20, knn: int = 32, num_node_types: int = 8, act_fn: str = 'relu', norm: bool = True, cutoff_mode: str = 'knn', ew_net_type: str = 'global', num_x2h: int = 1, num_h2x: int = 1, r_max: float = 10.0, x2h_out_fc: bool = False, sync_twoup: bool = False, pred_exp_from_all: bool = False, atom_vocab: List[str] | None = None, **kwargs: Any)

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

No train_set parameter: KGDiff’s size prior is the static atom_num_config table, so nothing has to be measured at build time (see docs/adding_new_models.md Section 2.5 – the seam is opt-in).

build() → KGDiffDiffusionTask

atom_vocab

condition_names: List[str] = []

model_kwargs

task: KGDiffDiffusionTask | None = None

task_type = 'diffusion_kgdiff'

class MolecularDiffusion.modules.tasks.diffusion_kgdiff.PocketSizePrior(space_size: float = DEFAULT_POCKET_EXTENT)

Ligand-size prior conditioned on the pocket’s spatial extent.

Wraps KGDiff’s static atom_num_config table, so unlike the other in-tree node distributions it needs neither a train_set nor a checkpoint buffer – the table ships with the code and is identical at train and generate time.

Same interface as TabascoNodeDistribution etc.: .n_node_dist and .sample(n). space_size is set by whoever knows the pocket.

sample(n_samples: int) → torch.Tensor

n_node_dist

space_size

MolecularDiffusion.modules.tasks.diffusion_kgdiff.DEFAULT_POCKET_EXTENT = 30.0

MolecularDiffusion.modules.tasks.diffusion_kgdiff.INT_TYPE