MolecularDiffusion.data.component.diffdec_data

Dataset / collate / DataModule for DiffDec scaffold-pocket-R-group complexes.

DiffDec carries three node sets per example – a fixed 3D scaffold, a protein pocket, and the diffused R-group – flat-concatenated into ONE padded node set with five per-atom masks (scaffold_only_mask, pocket_mask, scaffold_mask = scaffold+pocket, rgroup_mask, anchors). That is not the platform’s single-node-set PointCloud batch, so this module brings its own Dataset and collate_fn and is selected via _target_ from configs/data/diffdec_dataset.yaml. Direct precedent: data/component/diffsbdd_data.py, through the same seams (cli/train.py:597 -> .load() / .train_set / .collate_fn, then data/lightning_data_module.py’s collate_fn or graph_collate).

Storage is upstream’s own preprocessed ``.pt``, read unchanged. DiffDec publishes crossdocksingle_{train,test}_full.pt (Zenodo record 10527451) – a plain list[dict] of exactly the tensors its model consumes. Converting that into an ASE db and back would be a lossy round-trip for no gain: the five masks and the anchor flag have no home in the PointCloud db’s node-feature columns. So the DataModule reads the list directly and diffdec_collate() below is a port of upstream src/datasets.py collate (l. 126-166), pocket branch.

Each row carries:

uuid, name              -> list attrs, passed through untouched
positions   (N, 3)      -> float32 coordinates
one_hot     (N, 10)     -> element one-hot over DIFFDEC_ATOM_VOCAB
charges     (N,)        -> atomic number (unused by the model, h = one_hot)
anchors     (N,)        -> 1.0 on the single scaffold attachment atom
scaffold_only_mask (N,) -> scaffold atoms, pocket excluded
pocket_mask (N,)        -> pocket atoms
scaffold_mask (N,)      -> scaffold_only | pocket  (everything NOT noised)
rgroup_mask (N,)        -> the diffused R-group slots
num_atoms               -> int, list attr

The R-group is padded to a fixed 10 slots with a fake '#' atom (upstream parse_rgroup, datasets.py l. 84-124), so R-group size is implicit: the model emits '#' into unused slots and the generator strips those rows. '#' is therefore a real class of the vocabulary here, never a padding artefact to be filtered at load time.

Attributes

DATA_ATTRS_TO_ADD_LAST_DIM

DATA_ATTRS_TO_PAD

DATA_LIST_ATTRS

DIFFDEC_ATOM_VOCAB

DIFFDEC_CHARGES

FAKE_ATOM_INDEX

logger

Classes

DiffDecDataModule

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

DiffDecDataset

One upstream .pt -> one scaffold/pocket/R-group complex per index.

Functions

create_template(→ torch.Tensor)

Port of upstream datasets.py::create_template (l. 216-220).

create_templates_for_rgroup_generation_single(...)

Port of upstream datasets.py l. 222-244.

diffdec_collate(→ Dict[str, Any])

Port of upstream src/datasets.py::collate (l. 126-166).

Module Contents

class MolecularDiffusion.data.component.diffdec_data.DiffDecDataModule(root: str, train_file: str = 'crossdocksingle_test_full.pt', 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_diffdec', **kwargs: Any)

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

# ponytail: no split logic – val_file/test_file fall back to the # training file. The published test split is 43 complexes and is what the # smoke test uses; point each at its own .pt for a real run.

load() → None
atom_vocab
batch_size = 4
collate_fn
kwargs
limit = None
num_workers = 0
root
task_type = 'diffusion_diffdec'
test_file = 'crossdocksingle_test_full.pt'
test_set: DiffDecDataset | None = None
train_file = 'crossdocksingle_test_full.pt'
train_set: DiffDecDataset | None = None
val_file = 'crossdocksingle_test_full.pt'
valid_set: DiffDecDataset | None = None
class MolecularDiffusion.data.component.diffdec_data.DiffDecDataset(pt_path: str, limit: int | None = None)

Bases: torch.utils.data.Dataset

One upstream .pt -> one scaffold/pocket/R-group complex per index.

No re-centring at load time: DiffDecTask.forward removes the partial centre of mass itself (w.r.t. center_of_mass, anchors by default), exactly as upstream model_single.py:161-169 does, and the generator adds that offset back so samples land in the original pocket’s frame.

# ponytail: the whole list is read into RAM, like DiffSBDDDataset. The # test split is 43 complexes / 939 KB; the full train split is 1.6 GB, so # switch to a memory-mapped or sharded read if you train on it.

data: List[Dict[str, Any]]
pt_path
MolecularDiffusion.data.component.diffdec_data.create_template(tensor: torch.Tensor, scaffold_size: int, rgroup_size: int, fill: float = 0.0) → torch.Tensor

Port of upstream datasets.py::create_template (l. 216-220).

Keeps the first scaffold_size rows (scaffold + pocket, which the collate ordering puts first) and replaces the R-group rows with a constant block – zeros for coordinates/features, ones for rgroup_mask.

MolecularDiffusion.data.component.diffdec_data.create_templates_for_rgroup_generation_single(data: Dict[str, Any], rgroup_sizes: torch.Tensor) → Dict[str, Any]

Port of upstream datasets.py l. 222-244.

Blanks the ground-truth R-group out of a real batch, leaving the scaffold and pocket intact, so sampling starts from noise in the R-group slots only.

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

Port of upstream src/datasets.py::collate (l. 126-166).

Pocket branch only (pocket_mask is always present in the single-R-group CrossDocked data): edge_mask is emitted as a flat batch-index vector, not an adjacency mask. The actual graph is a 4 A radius graph rebuilt every forward pass by DynamicsWithPockets.get_dist_edges_4A, which uses this vector only to forbid edges between different samples.

MolecularDiffusion.data.component.diffdec_data.DATA_ATTRS_TO_ADD_LAST_DIM
MolecularDiffusion.data.component.diffdec_data.DATA_ATTRS_TO_PAD
MolecularDiffusion.data.component.diffdec_data.DATA_LIST_ATTRS
MolecularDiffusion.data.component.diffdec_data.DIFFDEC_ATOM_VOCAB = ['C', 'O', 'N', 'F', 'S', 'Cl', 'Br', 'I', 'P', '#']
MolecularDiffusion.data.component.diffdec_data.DIFFDEC_CHARGES = [6, 8, 7, 9, 16, 17, 35, 53, 15, 0]
MolecularDiffusion.data.component.diffdec_data.FAKE_ATOM_INDEX
MolecularDiffusion.data.component.diffdec_data.logger