MolecularDiffusion.data.component.pmdm_data

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

PMDM carries two node sets per complex – a diffused ligand point cloud and a fixed full-atom 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/pmdm_dataset.yaml. Direct precedent: data/component/diffpharma_data.py, which does the same for a four-node-set model 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/pmdm/scripts/convert_dataset.py:

Atoms -> ligand only (positions + raw element Z) key_value_pairs -> name, n_lig, n_pocket data -> lig_feat (N_lig, 8) rdkit atom-family flags,

pocket_coords, pocket_element, pocket_aa_type, pocket_is_backbone

Everything in the db is a raw integer, never one-hot. This module owns the one-hot expansion, so changing a vocabulary does not mean reconverting – and it keeps PMDM’s PDB/SDF parsing (and its np.bool/np.long calls, all removed in modern numpy) entirely out of the runtime path.

The collate deliberately emits PMDM’s own attribute names, so the task adapter is a one-line SimpleNamespace(**batch).

Attributes

AA_NAME_NUMBER
BACKBONE_NAMES
MAX_NUM_AA
PMDM_ATOM_VOCAB
PMDM_ELEMENTS
logger

Classes

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

Functions

fully_connected_edges(→ torch.Tensor)	Per-molecule fully-connected, self-loop-free edge index with offsets.
parse_pocket_pdb(→ Dict[str, numpy.ndarray])	Minimal PDB ATOM-line parser for a live pocket file, at generate time.
pmdm_collate(→ Dict[str, Any])	Concatenate complexes; emit PMDM's own attribute names.
write_pocket_pdb(→ None)	Write a pocket back to a PDB that parse_pocket_pdb() can read.

Module Contents

class MolecularDiffusion.data.component.pmdm_data.PMDMDataModule(root: str, train_file: str = 'pmdm_smoke.db', val_file: str | None = None, test_file: str | None = None, batch_size: int = 4, num_workers: int = 0, center: bool = True, limit: int | None = None, atom_vocab: List[str] | None = None, task_type: str = 'diffusion_pmdm', **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 31 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_pmdm'

test_file = 'pmdm_smoke.db'

test_set: PMDMDataset | None = None

train_file = 'pmdm_smoke.db'

train_set: PMDMDataset | None = None

val_file = 'pmdm_smoke.db'

valid_set: PMDMDataset | None = None

class MolecularDiffusion.data.component.pmdm_data.PMDMDataset(db_path: str, center: bool = True, limit: int | None = None)

Bases: torch.utils.data.Dataset

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

center=True subtracts the joint ligand+pocket centroid (upstream’s training-time behaviour, center_pos_pl re-centres anyway); generation uses center=False so the sampled ligand lands in the input pocket’s own frame.

# ponytail: rows are read eagerly into RAM at construction, like # DiffPharmaDataset. 31 complexes here, fine to ~100k. 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.pmdm_data.fully_connected_edges(sizes: List[int]) → torch.Tensor

Per-molecule fully-connected, self-loop-free edge index with offsets.

This is what PMDM’s GetAdj() transform builds (utils/transforms.py: get_adj_matrix(n), bond type constant 2) – the real bonds are thrown away before the model ever sees them.

MolecularDiffusion.data.component.pmdm_data.parse_pocket_pdb(path: str) → Dict[str, numpy.ndarray]

Minimal PDB ATOM-line parser for a live pocket file, at generate time.

Relocated from docs/model_integrations/pmdm/scripts/convert_dataset.py (that script now imports this instead of keeping its own copy) so PMDMConstrainedGenerator can read a plain pocket_file PDB directly, the same way PMDMDataset reads a converted db row. Returns raw arrays (element Z, amino-acid index, backbone flag) – one-hot expansion is the caller’s job, same convention as PMDMDataset.

MolecularDiffusion.data.component.pmdm_data.pmdm_collate(batch: List[Dict[str, Any]]) → Dict[str, Any]

Concatenate complexes; emit PMDM’s own attribute names.

MolecularDiffusion.data.component.pmdm_data.write_pocket_pdb(pocket: Dict[str, numpy.ndarray], path: str) → None

Write a pocket back to a PDB that parse_pocket_pdb() can read.

Only what the parser actually checks round-trips exactly: element, amino-acid identity, coordinates, and backbone-set membership (not the specific atom name – "CA"/"CB" stand in for backbone/sidechain). Used by docs/model_integrations/pmdm/scripts/extract_smoke_pair.py to reconstruct a smoke-test pocket from a converted db row (the original raw PDBs were deleted after conversion).

MolecularDiffusion.data.component.pmdm_data.AA_NAME_NUMBER

MolecularDiffusion.data.component.pmdm_data.BACKBONE_NAMES = ['CA', 'C', 'N', 'O']

MolecularDiffusion.data.component.pmdm_data.MAX_NUM_AA = 20

MolecularDiffusion.data.component.pmdm_data.PMDM_ATOM_VOCAB = ['H', 'C', 'N', 'O', 'F', 'P', 'S', 'Cl', 'Se', 'others']

MolecularDiffusion.data.component.pmdm_data.PMDM_ELEMENTS = [1, 6, 7, 8, 9, 15, 16, 17, 34, 119]

MolecularDiffusion.data.component.pmdm_data.logger