MolecularDiffusion.data.component.oareactdiff_data¶

Dataset / collate / DataModule for OA-ReactDiff on Transition1x.

Ported from oa_reactdiff/dataset/base_dataset.py and oa_reactdiff/dataset/transition1x.py (commit 543aaa8, MIT), restructured so the item is the shared Reaction container and the batch is upstream’s own (representations, conditions) tuple.

Why that split. Upstream’s dataset materialises a flat dict of index-suffixed keys – pos_0, one_hot_1, charge_2 – where the suffix silently encodes “reactant / transition state / product”. That is unreadable and unshareable: nothing in the layout says which index is the target. Reaction names the three roles; the positional layout is rebuilt in oareactdiff_collate(), immediately before the tensors reach EnVariationalDiffusion.forward, which is the only code that wants it. Nothing under data/ was modified to make this work – the module is selected purely by _target_ from configs/data/oareactdiff_ts1x_dataset.yaml, exactly as kgdiff_data.py and pmdm_data.py are.

The split is the published one, and it is not two files¶

train_addprop.pkl and valid_addprop.pkl contain the same 10,073 reactions. They differ only in their use_ind list – 9,000 entries in one, 1,073 in the other. So use_by_ind: true is what actually splits the data; pointing both at the same file would train and validate on the same reactions. Intersected with single_fragment, the effective sizes under the released settings are 6,733 train / 783 valid reactions (doubled by swapping_react_prod).

The five-wide atom one-hot is derived, not stored¶

Upstream builds it in the dataset via ATOM_MAPPING; here it is built in the collate from the reaction’s z. Keeping it derived means the shared container never carries a model-specific vocabulary, and changing the vocabulary never means reconverting anything. There is nothing to convert in the first place: the data module reads the upstream pickles directly.

`charge` is the atomic number¶

Upstream’s charges column holds Z, not a formal charge (base_dataset.py:170-176), and the ninth column of the model’s xh tensor is that same Z. There is no formal-charge channel anywhere in this model, which is why Reaction has no fc field.

Attributes¶

`ATOM_MAPPING`
`FRAGMENT_ORDER`
`N_ELEMENT`
`OAREACTDIFF_ATOM_VOCAB`
`TS_INDEX`
`logger`

Classes¶

`OAReactDiffDataModule`	DataModule contract: `load()` + `train_set`/`valid_set`/`test_set`.
`OAReactDiffTS1xDataset`	Transition1x reactions, one `Reaction` per item.

Functions¶

oareactdiff_collate(→ Tuple[List[Dict[str, ...)

Reactions -> upstream's (representations, conditions) tuple.

Module Contents¶

class MolecularDiffusion.data.component.oareactdiff_data.OAReactDiffDataModule(train_pkl: str, valid_pkl: str | None = None, center: bool = True, zero_charge: bool = False, single_frag_only: bool = True, swapping_react_prod: bool = True, use_by_ind: bool = True, batch_size: int = 4, num_workers: int = 0, limit: int | None = None, atom_vocab: List[str] | None = None, task_type: str = 'diffusion_oareactdiff', **kwargs: Any)¶

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

Two dataset instances over two pickles that hold the same reactions; the use_ind list inside each is what makes them a split. See the module docstring. test_set reuses the validation pickle because upstream ships no test.pkl.

load() → None¶: Build the train / valid / test datasets.

atom_vocab¶

batch_size = 4¶

center = True¶

collate_fn¶

kwargs¶

limit = None¶

num_workers = 0¶

single_frag_only = True¶

swapping_react_prod = True¶

task_type = 'diffusion_oareactdiff'¶

test_set: OAReactDiffTS1xDataset | None = None¶

train_pkl¶

train_set: OAReactDiffTS1xDataset | None = None¶

use_by_ind = True¶

valid_pkl¶

valid_set: OAReactDiffTS1xDataset | None = None¶

zero_charge = False¶

class MolecularDiffusion.data.component.oareactdiff_data.OAReactDiffTS1xDataset(pkl_path: str, center: bool = True, zero_charge: bool = False, single_frag_only: bool = True, swapping_react_prod: bool = True, use_by_ind: bool = True, limit: int | None = None)¶

Bases: torch.utils.data.Dataset

Transition1x reactions, one Reaction per item.

Parameters:

pkl_path – an upstream *_addprop.pkl / train.pkl.
center – subtract each object’s own centre of mass. On by default and effectively mandatory – the diffusion process operates on the zero-CoM subspace per object, and inpaint re-centres anyway.
zero_charge – replace the Z channel with zeros. The released checkpoint was trained with real Z, so leave this off unless you are training from scratch and know why.
single_frag_only – keep only reactions whose reactant and product are each a single connected fragment (the released setting).
swapping_react_prod – also emit every reaction backwards, with reactant and product exchanged and the same transition state. Doubles the set; the released checkpoint was trained with it on.
use_by_ind – honour the pickle’s published use_ind split. This is what separates train from validation – see the module docstring.
limit – cap the number of source reactions (before swapping). For smoke tests; None uses everything.

center = True¶

pkl_path¶

raw¶

swapping_react_prod = True¶

zero_charge = False¶

MolecularDiffusion.data.component.oareactdiff_data.oareactdiff_collate(batch: Sequence[MolecularDiffusion.data.component.reaction_data.Reaction], zero_charge: bool = False) → Tuple[List[Dict[str, torch.Tensor]], torch.Tensor]¶

Reactions -> upstream’s (representations, conditions) tuple.

Reproduces BaseDataset.collate_fn (base_dataset.py:52-88) but reads named fields instead of parsing index suffixes out of key strings. The output is exactly what EnVariationalDiffusion.forward takes, so no adapter runs on the training path.

Each of the three representations entries carries:

size (B,) atoms per sample in this object pos (sum n, 3) float32 coordinates one_hot (sum n, 5) int64, derived from z via ATOM_MAPPING charge (sum n, 1) int64 atomic number (or zeros) mask (sum n,) int64 scatter index -> which sample a row is

Dtypes are upstream’s: the one-hot and charge stay integral and are promoted to float by torch.cat when xh is assembled.

conditions is (B, 1) of zeros. The released checkpoint has condition_nf=1 and was fed constant zeros throughout (transition1x.py:141-147), so the channel exists in the weights but carries no information – do not mistake it for a place to put a property.

Parameters:

batch – reactions to collate. Every one needs ts_pos.
zero_charge – zero the atomic-number channel.

Returns:

(representations, conditions).

Raises:

ValueError – if any reaction is missing its transition state. There is no meaningful placeholder – the TS slot supplies the atom identities the pos_only sampler copies through.

MolecularDiffusion.data.component.oareactdiff_data.ATOM_MAPPING: Dict[int, int]¶

MolecularDiffusion.data.component.oareactdiff_data.FRAGMENT_ORDER = ('reactant', 'transition_state', 'product')¶

MolecularDiffusion.data.component.oareactdiff_data.N_ELEMENT = 5¶

MolecularDiffusion.data.component.oareactdiff_data.OAREACTDIFF_ATOM_VOCAB: List[str] = ['H', 'C', 'N', 'O', 'F']¶

MolecularDiffusion.data.component.oareactdiff_data.TS_INDEX = 1¶

MolecularDiffusion.data.component.oareactdiff_data.logger¶