MolecularDiffusion.modules.tasks.diffusion_oareactdiff

OA-ReactDiff: transition-state generation given a reactant and a product.

Duan, Du, Jia & Kulik, Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model, Nature Computational Science 3, 1045-1055 (2023). Ported from https://github.com/chenruduan/OAReactDiff at commit 543aaa8 (MIT).

Three pieces live here:

OAReactDiffTask

The platform task contract around EnVariationalDiffusion. Its forward reproduces upstream’s DDPMModule.compute_loss (trainer/pl_trainer.py:208-282) – the diffusion module returns raw loss terms, and the pos_only denominators plus the per-object scales weighting are applied here, where upstream applies them.

ModelTaskFactory

Hydra entry point for configs/tasks/diffusion_oareactdiff.yaml.

OAReactDiffTSGenerator

A thin TSGenerator subclass. It writes no ``run()`` loop and no reaction plumbing. “Load one fixed structural context, tile it to a batch, sample the unknown part into it, write .xyz” is the shared pocket loop (pocket_generator.py); reading one Reaction, sizing the transition state from it and writing the reference structures beside the samples is the shared TS layer on top of it (ts_generator.py). What is left here is OA-ReactDiff’s own: which corpus and filters, which collate, and the RePaint knobs.

Scope of this integration

Transition state given a reactant and a product, and nothing else. Not in scope, each for a stated reason: unconditional whole-reaction generation (pos_only=True means the network cannot invent atom identities, and no generation seam carries a chemical formula); the confidence / recommender model that ranks 5 samples per reaction (no checkpoint for it ships); React-OT (a separate repository and paper); energy-difference conditioning (condition_nf=1 was fed constant zeros, so no trained conditional variant exists); validation-time inpainting RMSD (deep-copies the model and runs a 150-step inpaint per validation epoch); and trajectory frames (the frame-saving lines are commented out inside upstream’s own inpaint).

Two facts from the released checkpoint that constrain everything

• It is pos_only=True: coordinates only. Atom identities are supplied by the input reaction and copied straight through, never generated.

• ddpm.schedule.gamma_module.gamma is a non-learned nn.Parameter of shape (5001,) living in the state dict, so training timesteps is pinned at 5000 on load. Sampling does not use it: OAReactDiffTask. sample() builds a fresh polynomial_2 schedule at num_steps (250 by default) and swaps it in for the duration of the call, which is exactly what upstream’s evaluate/utils.py:14-32 set_new_schedule does.

Attributes

FRAG_FIXED

Classes

ModelTaskFactory	Hydra entry point for configs/tasks/diffusion_oareactdiff.yaml.
OAReactDiffTSGenerator	Transition-state generation behind interference/gen_oareactdiff_ts.
OAReactDiffTask	Task contract around EnVariationalDiffusion.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_oareactdiff.ModelTaskFactory(task_type: str = 'diffusion_oareactdiff', model_config: Dict[str, Any] | None = None, node_nfs: Sequence[int] = (9, 9, 9), edge_nf: int = 0, condition_nf: int = 1, fragment_names: Sequence[str] = ('R', 'TS', 'P'), pos_dim: int = 3, update_pocket_coords: bool = True, condition_time: bool = True, edge_cutoff: float | None = None, norm_values: Sequence[float] = (1.0, 1.0, 1.0), norm_biases: Sequence[float] = (0.0, 0.0, 0.0), noise_schedule: str = 'cosine', timesteps: int = 5000, precision: float = 1e-05, loss_type: str = 'l2', pos_only: bool = True, scales: Sequence[float] = (1.0, 2.0, 1.0), fixed_idx: List[int] | None = None, atom_vocab: List[str] | None = None, **kwargs: Any)

Hydra entry point for configs/tasks/diffusion_oareactdiff.yaml.

No train_set parameter: the only construction-time statistic upstream has a slot for is size_histogram, which is passed None and never read (en_diffusion.py:31,42; pl_trainer.py:117). The Section 2.5 seam is therefore not used and cli/train.py is untouched.

build() → OAReactDiffTask

Assemble dynamics + schedule + normaliser into the task.

atom_vocab

condition_names: List[str] = []

condition_nf = 1

condition_time = True

edge_cutoff = None

edge_nf = 0

fixed_idx

fragment_names

generation_time_keys = ('reaction_pkl',)

loss_type = 'l2'

model_config

node_nfs

noise_schedule = 'cosine'

norm_biases

norm_values

pos_dim = 3

pos_only = True

precision

scales

task: OAReactDiffTask | None = None

task_type = 'diffusion_oareactdiff'

timesteps = 5000

update_pocket_coords = True

class MolecularDiffusion.modules.tasks.diffusion_oareactdiff.OAReactDiffTSGenerator(task: Any, reaction_index: int = 0, num_generate: int = 20, batch_size: int = 4, num_steps: int | None = 250, resamplings: int = 5, jump_length: int = 5, noise_schedule: str = 'polynomial_2', output_path: str = 'generated_oareactdiff', seed: int = 42, device: str | None = None, **kwargs: Any)

Bases: MolecularDiffusion.modules.tasks.ts_generator.TSGenerator

Transition-state generation behind interference/gen_oareactdiff_ts.

The shape of the request is identical to every other TS run – load one reaction, tile it to a batch, sample the transition state into it, write .xyz – so the loop is run(), the reaction plumbing is TSGenerator, and what is left here is only what is OA-ReactDiff’s: which corpus, which filters, which collate, and the RePaint knobs.

The reaction comes from one row of an upstream Transition1x pickle, read with the released training filters (single_frag_only, use_by_ind) but without swapping_react_prod: a generator should not silently hand back the reverse reaction. With valid_addprop.pkl that makes reaction_index an index into the 783 held-out single-fragment reactions.

The output folder is self-describing: reactant.xyz, product.xyz and reference_ts.xyz are written once alongside the generated molecule_NNN.xyz files, so a sample can be judged against its own inputs and against the DFT transition state without going back to the pickle.

The corpus key is reaction_pkl (declared as source_key, not as a constructor parameter, so there is exactly one of it); the shared interference/gen_ts.yaml reaches the same slot as reaction_source. Either works, both together raise.

Configure the run.

Parameters:

• task – the loaded TS task; its sample() is what gets called.

• reaction_source – where reactions are read from, forwarded to the base as pocket_db.

• reaction_index – which reaction in that source to generate for.

• num_generate – how many transition states to sample for it.

• batch_size – how many of those to sample at once.

• num_steps – reverse-process steps; None => the model’s default.

• output_path – directory for the .xyz files.

• seed – torch/random/numpy seed.

• device – None => cuda if available.

• **kwargs – rejected by the base, on purpose – an unknown interference key is a typo, not a no-op.

Note

There is no mol_size. A transition state has exactly the atoms of its reaction, so a size prior would be a lie; _sizes() reads the count off the reaction instead.

db_required_msg = 'interference.reaction_pkl (or reaction_source, on the shared gen_ts.yaml) is required:...

jump_length = 5

max_retries: int | None = None

noise_schedule = 'polynomial_2'

resamplings = 5

seed_numpy = True

source_key = 'reaction_pkl'

tag = 'oareactdiff'

class MolecularDiffusion.modules.tasks.diffusion_oareactdiff.OAReactDiffTask(ddpm: MolecularDiffusion.modules.models.oareactdiff.EnVariationalDiffusion, scales: Sequence[float] = (1.0, 2.0, 1.0), loss_type: str = 'l2', pos_only: bool = True, precision: float = 1e-05, atom_vocab: List[str] | None = None)

Bases: torch.nn.Module

Task contract around EnVariationalDiffusion.

evaluate(pred: torch.Tensor, target: torch.Tensor) → Dict[str, torch.Tensor]

Validation metric.

Upstream’s eval_inplaint_batch RMSD is deliberately not here: it deep-copies the whole model and runs a 150-step inpaint every eval_epochs, which is far too heavy for a training loop.

forward(batch: Tuple[List[Dict[str, torch.Tensor]], torch.Tensor]) → Tuple[torch.Tensor, Dict[str, Any]]

One training/validation step.

Reproduces DDPMModule.compute_loss (pl_trainer.py:208-282). The l2 branch is taken only while training; validation always takes the VLB branch, which is upstream’s behaviour and is why the reported validation loss is on a different scale from the training loss.

Parameters:

batch – (representations, conditions) from oareactdiff_collate().

Returns:

(loss, stats); stats carries upstream’s logged keys.

predict_and_target(batch: Tuple[List[Dict[str, torch.Tensor]], torch.Tensor]) → Tuple[torch.Tensor, torch.Tensor]

Pure-generative stub: the loss is the prediction, target is zero.

sample(batch_size: int | None = None, nodesxsample: torch.Tensor | None = None, num_steps: int | None = None, batch: Dict[str, Any] | None = None, resamplings: int = 5, jump_length: int = 5, noise_schedule: str = 'polynomial_2', **kwargs: Any) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Generate transition states for the reactions in batch.

The signature deviates from Section 2.1 the same way every pocket-conditioned task’s does: the conditioning structure arrives as batch, because there is no channel for “which reaction” in sample(batch_size, nodesxsample, ...).

Parameters:

• batch_size – ignored; taken from batch.

• nodesxsample – accepted and checked, not used to choose. The transition state has exactly as many atoms as the reaction, so a value disagreeing with it is a caller bug worth surfacing rather than silently overriding.

• num_steps – reverse-process steps. None falls back to self.T (5000) – both shipped generate configs set 250, which is upstream’s evaluation setting.

• batch – {"representations": [...], "conditions": tensor}.

• resamplings – RePaint re-noise rounds per jump.

• jump_length – steps to jump back on each resampling.

• noise_schedule – the sampling schedule to build fresh.

Returns:

(one_hot, charges, coords, node_mask) for the transition state only, padded to (B, N, .).

Raises:

ValueError – if batch is missing, or nodesxsample disagrees with the reaction’s atom count.

T

atom_vocab

ddpm

property device: torch.device

Device of the first parameter.

loss_type = 'l2'

property model: OAReactDiffTask

There is no separate inner module the generation code needs.

n_fragments

node_dist_model: Any = None

pos_only = True

precision = 1e-05

prop_dist_model: Any = None

scales

split = 'train'

MolecularDiffusion.modules.tasks.diffusion_oareactdiff.FRAG_FIXED = [0, 2]