MolecularDiffusion.data.component.kgdiff_data¶

Dataset / collate / DataModule for KGDiff pocket-ligand complexes.

KGDiff carries two node sets per complex – a diffused ligand point cloud and a fixed pocket point cloud – flat-concatenated with torch_scatter-style batch indices, plus one scalar affinity label per complex. That is not 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/kgdiff_dataset.yaml. Direct precedent: data/component/pmdm_data.py and data/component/diffpharma_data.py, which do the same 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/kgdiff/scripts/convert_dataset.py:

Atoms -> ligand only (positions + raw element Z) key_value_pairs -> name, n_lig, n_pocket, affinity data -> lig_aromatic, pocket_coords, pocket_element,

pocket_aa_type, pocket_is_backbone

Everything in the db is a raw integer (or a bool), never one-hot. This module owns the vocabulary expansion – the 13-class (element, is_aromatic) ligand index and the 27-dim protein feature – so changing a vocabulary does not mean reconverting.

The collate emits KGDiff’s own argument names, so the task-side adapter is a one-line .to(device) map rather than a container conversion.

Attributes¶

`KGDIFF_ATOM_VOCAB`
`KGDIFF_ELEMENTS`
`LIGAND_INDEX_TO_Z`
`MAP_ATOM_TYPE_AROMATIC_TO_INDEX`
`MAX_NUM_AA`
`NUM_LIGAND_CLASSES`
`PROTEIN_ELEMENTS`
`PROTEIN_FEATURE_DIM`
`logger`

Classes¶

`KGDiffDataModule`	DataModule contract: `load()` + `train_set`/`valid_set`/`test_set`.
`KGDiffDataset`	One converted ASE db -> one pocket-ligand complex per index.

Functions¶

`kgdiff_collate`(→ Dict[str, Any])	Concatenate complexes; emit KGDiff's own argument names.
`ligand_class_index`(→ torch.Tensor)	`(element Z, is_aromatic)` -> the 13-class index, as `int64`.
`protein_features`(→ torch.Tensor)	`(M,) x3` raw integers -> the `(M, 27)` float feature matrix.

Module Contents¶

class MolecularDiffusion.data.component.kgdiff_data.KGDiffDataModule(root: str, train_file: str = 'kgdiff_smoke.db', val_file: str | None = None, test_file: str | None = None, batch_size: int = 4, num_workers: int = 0, limit: int | None = None, atom_vocab: List[str] | None = None, task_type: str = 'diffusion_kgdiff', **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. The smoke set is 100 complexes; add a real split (and LMDB # streaming) when a full CrossDocked run needs one.

load() → None¶

atom_vocab¶

batch_size = 4¶

collate_fn¶

kwargs¶

limit = None¶

num_workers = 0¶

root¶

task_type = 'diffusion_kgdiff'¶

test_file = 'kgdiff_smoke.db'¶

test_set: KGDiffDataset | None = None¶

train_file = 'kgdiff_smoke.db'¶

train_set: KGDiffDataset | None = None¶

val_file = 'kgdiff_smoke.db'¶

valid_set: KGDiffDataset | None = None¶

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

Bases: torch.utils.data.Dataset

One converted ASE db -> one pocket-ligand complex per index.

Unlike the other pocket models here, coordinates are never re-centred at load time: ScorePosNet3D centres on the pocket centroid inside both get_diffusion_loss and sample_diffusion (center_pos_mode='protein'), and sample_diffusion adds the offset back, so samples land in the input pocket’s own frame either way.

# ponytail: rows are read eagerly into RAM at construction, like # PMDMDataset. 100 complexes here, fine to ~100k. Stream by row id in # __getitem__ if a full CrossDocked split ever needs it.

db_path¶

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

MolecularDiffusion.data.component.kgdiff_data.kgdiff_collate(batch: List[Dict[str, Any]]) → Dict[str, Any]¶: Concatenate complexes; emit KGDiff’s own argument names.

MolecularDiffusion.data.component.kgdiff_data.ligand_class_index(z: numpy.ndarray, aromatic: numpy.ndarray) → torch.Tensor¶

(element Z, is_aromatic) -> the 13-class index, as int64.

Out-of-vocabulary pairs fall back to class 0 (H), which is exactly what upstream’s get_index does (utils/transforms.py:109-113) – it prints the offender and returns (1, False).

MolecularDiffusion.data.component.kgdiff_data.protein_features(element: numpy.ndarray, aa_type: numpy.ndarray, is_backbone: numpy.ndarray) → torch.Tensor¶

(M,) x3 raw integers -> the (M, 27) float feature matrix.

[element one-hot (6) | amino acid one-hot (20) | is_backbone (1)], matching FeaturizeProteinAtom.__call__ exactly. An element outside the vocabulary yields an all-zero element block, as upstream’s boolean == comparison does.

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

MolecularDiffusion.data.component.kgdiff_data.KGDIFF_ELEMENTS = [1, 6, 7, 8, 9, 15, 16, 17]¶

MolecularDiffusion.data.component.kgdiff_data.LIGAND_INDEX_TO_Z¶

MolecularDiffusion.data.component.kgdiff_data.MAP_ATOM_TYPE_AROMATIC_TO_INDEX¶

MolecularDiffusion.data.component.kgdiff_data.MAX_NUM_AA = 20¶

MolecularDiffusion.data.component.kgdiff_data.NUM_LIGAND_CLASSES = 13¶

MolecularDiffusion.data.component.kgdiff_data.PROTEIN_ELEMENTS = [1, 6, 7, 8, 16, 34]¶

MolecularDiffusion.data.component.kgdiff_data.PROTEIN_FEATURE_DIM = 27¶

MolecularDiffusion.data.component.kgdiff_data.logger¶