MolecularDiffusion.data.component.apo2mol_data¶

Dataset / collate / DataModule for Apo2Mol apo-holo-ligand complexes.

Apo2Mol carries three things per complex: a diffused ligand point cloud, an apo pocket, and the holo pocket the model is trained to recover. The pocket is not just coordinates – it also carries a per-residue rigid transform (quaternion + translation) and five side-chain chi angles, because the pocket conformation is generated, not conditioned on. None of that fits the platform’s padded single-node-set PointCloud batch, so this module brings its own Dataset and collate_fn and is selected via _target_ from configs/data/apo2mol_dataset.yaml. Direct precedent: data/component/kgdiff_data.py, whose schema this extends, and pmdm_data.py / diffpharma_data.py, which go through the same seams (cli/train.py -> .load() / .train_set / .collate_fn, then data/lightning_data_module.py’s collate_fn or graph_collate).

Storage is one ASE db row per complex, written offline by docs/model_integrations/apo2mol/scripts/convert_dataset.py:

Atoms            -> ligand only (positions + raw element Z)
key_value_pairs  -> name, n_lig, n_pocket, n_res
data             -> lig_aromatic,
                    pocket_coords (apo), pocket_coords_holo,
                    pocket_element, pocket_aa_type, pocket_is_backbone,
                    pocket_atom_name, pocket_aa_name,
                    pocket_atom_to_aa_group,
                    res_rotations, res_translations,
                    chi_apo, chi_holo, chi_mask

Everything numeric is a raw integer (or float), never one-hot: the vocabularies live here, so changing one does not mean reconverting. The 13-class ligand index and the 27-dim protein feature are imported from kgdiff_data – Apo2Mol’s are byte-for-byte the same tables.

## Two invariants the collate must not break

protein_atom_name and protein_atom_to_aa_name stay nested per complex (list[list[str]]). Flattening them makes apply_transforms_tensor_batch index the wrong atoms and produce wrong coordinates without raising.
protein_atom_to_aa_group is not offset across complexes – it restarts at 0 in every one. Both the backbone (uni_transformer._global_residue_index) and the residue-transform code re-derive global ids themselves, so offsetting here would double-count.

Attributes¶

`APO2MOL_ATOM_VOCAB`
`MAX_CHI`
`logger`

Classes¶

`Apo2MolDataModule`	DataModule contract: `load()` + `train_set`/`valid_set`/`test_set`.
`Apo2MolDataset`	One converted ASE db -> one apo/holo/ligand complex per index.

Functions¶

apo2mol_collate(→ Dict[str, Any])

Concatenate complexes; emit Apo2Mol's own argument names.

Module Contents¶

class MolecularDiffusion.data.component.apo2mol_data.Apo2MolDataModule(root: str, train_file: str = 'apo2mol_smoke.db', val_file: str | None = None, test_file: str | None = None, batch_size: int = 2, num_workers: int = 0, limit: int | None = None, atom_vocab: List[str] | None = None, task_type: str = 'diffusion_apo2mol', **kwargs: Any)¶

DataModule contract: load() + train_set/valid_set/test_set.

# ponytail: one db, no split – val_file/test_file default to the # training db. Point them at real split dbs for a full run.

load() → None¶

atom_vocab¶

batch_size = 2¶

collate_fn¶

kwargs¶

limit = None¶

num_workers = 0¶

root¶

task_type = 'diffusion_apo2mol'¶

test_file = 'apo2mol_smoke.db'¶

test_set: Apo2MolDataset | None = None¶

train_file = 'apo2mol_smoke.db'¶

train_set: Apo2MolDataset | None = None¶

val_file = 'apo2mol_smoke.db'¶

valid_set: Apo2MolDataset | None = None¶

class MolecularDiffusion.data.component.apo2mol_data.Apo2MolDataset(db_path: str, limit: int | None = None)¶

Bases: torch.utils.data.Dataset

One converted ASE db -> one apo/holo/ligand complex per index.

Coordinates are never re-centred at load time: ScorePosNet3D subtracts the apo pocket centroid inside both get_diffusion_loss and sample_diffusion (center_pos_mode='protein') and the sampler adds the offset back, so samples land in the input pocket’s own frame.

# ponytail: rows are read eagerly into RAM at construction, like # KGDiffDataset. Fine for the smoke set and up to ~100k complexes; stream # by row id in __getitem__ if the full 24k-complex set with per-residue # arrays ever gets tight.

db_path¶

entries: List[Dict[str, Any]] = []¶

MolecularDiffusion.data.component.apo2mol_data.apo2mol_collate(batch: List[Dict[str, Any]]) → Dict[str, Any]¶

Concatenate complexes; emit Apo2Mol’s own argument names.

Three scatter indices come out, one per node set: ligand_batch (atoms), protein_batch / protein_element_batch (pocket atoms; the same tensor under both names, because the model reads it under each), and protein_translations_batch (residues).

MolecularDiffusion.data.component.apo2mol_data.APO2MOL_ATOM_VOCAB = ['H', 'C', 'N', 'O', 'F', 'P', 'S', 'Cl']¶

MolecularDiffusion.data.component.apo2mol_data.MAX_CHI = 5¶

MolecularDiffusion.data.component.apo2mol_data.logger¶