MolecularDiffusion.data.component.diffsbdd_data¶

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

DiffSBDD carries two node sets per complex – a diffused ligand point cloud and a protein pocket point cloud – flat-concatenated with torch_scatter style batch indices. 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/diffsbdd_dataset.yaml. Direct precedent: data/component/kgdiff_data.py and data/component/pmdm_data.py, 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 the kgdiff ASE db, reused unchanged (see the integration plan): one row per complex, written 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 (affinity unused here) data -> lig_aromatic, pocket_coords, pocket_element,

pocket_aa_type, pocket_is_backbone

Everything there is a raw integer, so this module owns DiffSBDD’s own vocabulary expansion and nothing has to be reconverted to change it.

Vocabulary: one 10-class element encoder shared by ligand and pocket, and hydrogens stripped from both. This is not a guess – it is read off the released full-atom CrossDocked checkpoints (Zenodo 8183747), whose hyper_parameters say dataset='crossdock' + pocket_representation='full-atom', which in lightning_modules.py:90-97 resolves both encoders to dataset_params['crossdock']['atom_encoder'] (10 classes, no others bucket), and whose atom_encoder / residue_encoder weights are both (20, 10). The 11-class crossdock_full table in constants.py:154 is unused by those weights. Upstream’s process_crossdock.py is run with --no_H, so hydrogens are absent from both node sets – hence the strip here, which is also what keeps the pocket inside a vocabulary that has no fallback class.

The collate emits DiffSBDD’s own key names (dataset.py:53-70), so the task-side adapter is one dict comprehension.

Attributes¶

`DIFFSBDD_ATOM_VOCAB`
`DIFFSBDD_ELEMENTS`
`NUM_ATOM_CLASSES`
`logger`

Classes¶

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

Functions¶

`diffsbdd_collate`(→ Dict[str, Any])	Concatenate complexes; emit DiffSBDD's own argument names.
`one_hot_from_z`(→ torch.Tensor)	Raw element Z -> `(N, 10)` float one-hot.

Module Contents¶

class MolecularDiffusion.data.component.diffsbdd_data.DiffSBDDDataModule(root: str, train_file: str = 'diffsbdd_smoke.db', val_file: str | None = None, test_file: str | None = None, batch_size: int = 4, num_workers: int = 0, limit: int | None = None, center: bool = True, atom_vocab: List[str] | None = None, task_type: str = 'diffusion_diffsbdd', **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¶

center = True¶

collate_fn¶

kwargs¶

limit = None¶

num_workers = 0¶

root¶

task_type = 'diffusion_diffsbdd'¶

test_file = 'diffsbdd_smoke.db'¶

test_set: DiffSBDDDataset | None = None¶

train_file = 'diffsbdd_smoke.db'¶

train_set: DiffSBDDDataset | None = None¶

val_file = 'diffsbdd_smoke.db'¶

valid_set: DiffSBDDDataset | None = None¶

class MolecularDiffusion.data.component.diffsbdd_data.DiffSBDDDataset(db_path: str, limit: int | None = None, center: bool = True)¶

Bases: torch.utils.data.Dataset

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

Coordinates are joint-centred on the combined ligand+pocket centroid at load time, exactly as upstream’s ProcessedLigandPocketDataset (dataset.py:35-41, center=True) does for both modes. The generator restores the original pocket frame afterwards.

# ponytail: rows are read eagerly into RAM, like KGDiffDataset. 100 # complexes here; stream by row id in __getitem__ if a full CrossDocked # split ever needs it.

center = True¶

db_path¶

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

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

MolecularDiffusion.data.component.diffsbdd_data.one_hot_from_z(z: numpy.ndarray) → torch.Tensor¶

Raw element Z -> (N, 10) float one-hot.

Raises on an out-of-vocabulary element rather than silently emitting an all-zero row: this vocabulary has no others bucket, so a zero row is a feature vector the model never saw in training.

MolecularDiffusion.data.component.diffsbdd_data.DIFFSBDD_ATOM_VOCAB = ['C', 'N', 'O', 'S', 'B', 'Br', 'Cl', 'P', 'I', 'F']¶

MolecularDiffusion.data.component.diffsbdd_data.DIFFSBDD_ELEMENTS = [6, 7, 8, 16, 5, 35, 17, 15, 53, 9]¶

MolecularDiffusion.data.component.diffsbdd_data.NUM_ATOM_CLASSES = 10¶

MolecularDiffusion.data.component.diffsbdd_data.logger¶