MolecularDiffusion.modules.tasks.diffusion_apo2mol¶

Apo2Mol task: apo-pocket-conditioned ligand diffusion with a co-generated pocket.

Three objects, the layout every pocket-conditioned model here uses:

Apo2MolDiffusionTask – the duck-typed Task (docs/adding_new_models.md Section 2.1) wrapping ScorePosNet3D and the frozen PMINet prior.
ModelTaskFactory – the _target_ of configs/tasks/diffusion_apo2mol.yaml.
Apo2MolPocketGenerator – the _target_ of configs/interference/gen_apo2mol_pocket.yaml. GenerativeFactory’s sample(batch_size, nodesxsample, ...) has no channel for “which pocket”, so pocket models get their own generator behind their own _target_; cli/generate.py only does instantiate(cfg.interference, task=task) then .run(), so no core change is needed.

What Apo2Mol adds over KGDiff / IPDiff, in one sentence. Those two treat the pocket as a fixed condition; Apo2Mol takes an apo (ligand-free) pocket and generates the ligand and a holo-like pocket conformation together, by diffusing per-residue rigid transforms and side-chain chi angles alongside the ligand point cloud.

Where the generated pocket goes. The Section 2.1 sample() return tuple has no protein channel, so it stays the platform-standard (one_hot, charges, coords, node_mask) and the co-generated pocket is handed back out-of-band via Apo2MolDiffusionTask.last_pocket, which the generator writes out as a .pdb sidecar per sample together with its RMSD / TM-score against the input pocket.

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

sample() requires a pocket; there is no unconditional path.
Sampled coordinates come back in the input pocket’s frame (center_pos_mode='protein').

Out of scope this pass (see the integration plan): guidance of every kind (Apo2Mol has none – prop_dist_model is None), unconditional generation, trajectory / n_frames export, the retrieval-prompt branch, pos_only sampling, and Vina docking.

Attributes¶

`DEFAULT_PMINET_CKPT`
`INT_TYPE`
`logger`

Classes¶

`Apo2MolDiffusionTask`	Task contract around `ScorePosNet3D` + the frozen PMINet prior.
`Apo2MolPocketGenerator`	Pocket-conditioned generation behind `interference/gen_apo2mol_pocket`.
`ModelTaskFactory`	Hydra entry point for `configs/tasks/diffusion_apo2mol.yaml`.

Functions¶

write_pocket_pdb(→ None)

Write a minimal, viewer-loadable PDB for one generated pocket.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_apo2mol.Apo2MolDiffusionTask(model: MolecularDiffusion.modules.models.apo2mol.ScorePosNet3D, net_cond: MolecularDiffusion.modules.models.apo2mol.BAPNet, atom_vocab: List[str] | None = None, pos_noise_std: float = 0.1)¶

Bases: torch.nn.Module

Task contract around ScorePosNet3D + the frozen PMINet prior.

evaluate(pred, target)¶

forward(batch)¶

predict_and_target(batch)¶

sample(batch_size=None, nodesxsample=None, num_steps=None, batch=None, progress: bool = True, **kwargs)¶

Sample ligands inside the apo pocket carried by batch.

The signature deliberately deviates from Section 2.1: it takes a pocket. batch must come from apo2mol_collate, or from Apo2MolPocketGenerator._repeat(), which builds the same keys.

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 (Apo2Mol has no charge channel).

The co-generated pocket lands in last_pocket as {"pos", "batch", "rmsd", "tmscore"}.

If batch carries no protein_pos_holo (pocket-only generation) the apo coordinates are used in its place, exactly as upstream’s sample_custom_pocket.py:690 does. The reported RMSD / TM-score is then displacement from the input apo structure, not accuracy against a true holo structure.

atom_vocab¶

property device¶

last_pocket: Dict[str, Any] | None = None¶

model¶

property n_node_dist: Dict[int, float]¶

net_cond¶

property node_dist_model: MolecularDiffusion.modules.tasks.diffusion_kgdiff.PocketSizePrior¶

pos_noise_std = 0.1¶

prop_dist_model = None¶

split = 'train'¶

class MolecularDiffusion.modules.tasks.diffusion_apo2mol.Apo2MolPocketGenerator(task, pocket_db: str | None = None, pocket_index: int = 0, num_generate: int = 20, batch_size: int = 2, num_steps: int | None = None, mol_size: list | None = None, use_holo_reference: bool = False, output_path: str = 'generated_apo2mol', seed: int = 42, device: str | None = None, **kwargs: Any)¶

Bases: MolecularDiffusion.modules.tasks.pocket_generator.PocketGenerator

Pocket-conditioned generation behind interference/gen_apo2mol_pocket.

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

Two things it writes that no other in-tree generator does:

pocket_<i>.pdb – the co-generated holo-like pocket for each sample, which is Apo2Mol’s headline contribution;
pocket_metrics.csv – per-sample RMSD and TM-score of that pocket against the reference the model was given.

use_holo_reference decides what that reference is. Default false reproduces upstream’s custom-inference behaviour (sample_custom_pocket.py:686-690): holo is set to apo, so the metrics measure displacement from the input apo structure, not accuracy. Set it to true on a db row converted from a real apo/holo pair to get a genuine accuracy number instead.

The sampling loop itself lives in PocketGenerator.

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

tag = 'apo2mol'¶

use_holo_reference = False¶

class MolecularDiffusion.modules.tasks.diffusion_apo2mol.ModelTaskFactory(task_type: str = 'diffusion_apo2mol', net_cond_ckpt: str | None = DEFAULT_PMINET_CKPT, cond_dim: int = 128, topk_prompt: int = 0, pos_noise_std: float = 0.1, 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, lambda_schedule: str = 'sigmoid', num_diffusion_timesteps: int = 1000, loss_v_weight: float = 100.0, loss_chi_weight: float = 5.0, sample_time_method: str = 'symmetric', 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 = 5, 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, num_protein_update_steps: int = 5, atom_vocab: List[str] | None = None, **kwargs: Any)¶

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

No train_set parameter: like KGDiff and IPDiff, the ligand-size prior is a static table conditioned on pocket extent, so nothing has to be measured at build time (docs/adding_new_models.md Section 2.5 – that seam is opt-in).

build() → Apo2MolDiffusionTask¶

atom_vocab¶

cond_dim = 128¶

condition_names: List[str] = []¶

model_kwargs¶

net_cond_ckpt = 'docs/model_integrations/apo2mol/checkpoints/pminet'¶

pos_noise_std = 0.1¶

task: Apo2MolDiffusionTask | None = None¶

task_type = 'diffusion_apo2mol'¶

MolecularDiffusion.modules.tasks.diffusion_apo2mol.write_pocket_pdb(path: str, coords: torch.Tensor, atom_names: List[str], aa_names: List[str], aa_group: torch.Tensor, elements: List[str] | None = None) → None¶

Write a minimal, viewer-loadable PDB for one generated pocket.

Residue numbering is the local 0-based aa_group + 1, so the sidecar is self-consistent but is NOT the original PDB numbering.

MolecularDiffusion.modules.tasks.diffusion_apo2mol.DEFAULT_PMINET_CKPT = 'docs/model_integrations/apo2mol/checkpoints/pminet'¶

MolecularDiffusion.modules.tasks.diffusion_apo2mol.INT_TYPE¶

MolecularDiffusion.modules.tasks.diffusion_apo2mol.logger¶