MolecularDiffusion.modules.tasks.diffusion_reactot¶

React-OT: transition states in ten deterministic network evaluations.

Duan, C.; Liu, G.-H.; Du, Y.; …; Kulik, H. J., Optimal transport for generating transition states in chemical reactions, Nature Machine Intelligence (2025), doi:10.1038/s42256-025-01010-0. Preprint, under a different title: arXiv:2404.13430, React-OT: Optimal Transport for Generating Transition State in Chemical Reactions. Ported from https://github.com/deepprinciple/react-ot at commit 6dfccd0.

Three pieces live here:

ReactOTTask: The platform task contract around EnSB. Its forward reproduces upstream’s SBModule.compute_loss (pl_trainer.py:993-1004) – which is a two-line unpacking, because the bridge returns a finished scalar. There is no ELBO to reduce, no loss_0, no kl_prior and no per-object scales weighting, which is why this is a fraction of the size of OAReactDiffTask’s.
ModelTaskFactory: Hydra entry point for configs/tasks/diffusion_reactot.yaml.
ReactOTTSGenerator: A thin TSGenerator subclass. No ``run()`` loop, and no ``PocketGenerator`` hook overridden. The loop is the shared pocket loop, the reaction plumbing is the shared TS layer, and what is left here is only React-OT’s own: which corpus and filters, which collate, and the solver knobs.

Relationship to OA-ReactDiff¶

Same authors, same backbone, same corpus, same reactant/TS/product triple, same batch layout – so this integration imports rather than duplicates: LEFTNet / EGNNDynamics / Normalizer / the graph helpers from modules/models/oareactdiff/, OAReactDiffTS1xDataset / oareactdiff_collate / OAReactDiffDataModule from data/component/oareactdiff_data.py, and _pad_ts() from its task module. None of those files is edited.

What differs is the process: OA-ReactDiff runs a stochastic reverse diffusion with RePaint re-noising and ranks five samples afterwards; React-OT integrates a Schrodinger bridge from the reactant/product midpoint to the transition state, deterministically, in nfe network evaluations (10 in the README’s published command). Run it twice and you get the same structure – which is why gen_reactot_ts.yaml defaults num_generate: 1.

The transition-state leak this class closes on purpose¶

EnSB.sample builds its node features from representations[1], whose position columns are the reference transition state when the batch comes from a corpus. Upstream gets away with it because net_out_fn overwrites those columns before every network call – i.e. it does not leak today, but only by accident, and it means upstream’s sampler cannot be pointed at an R/P pair with no known TS at all. ReactOTTask.sample() substitutes the midpoint into that slot explicitly, exactly as upstream’s own R/P-only deployment path does (run_model.py:122), so leak-freedom is structural.

Verified: on CPU, replacing representations[1]["pos"] with random noise before sampling changes the output by exactly 0.0 A, and two identical calls agree to exactly 0.0 A.

On CUDA both of those become ~5e-3 A, and that is not a leak. It is torch_scatter’s atomic-add nondeterminism (remove_mean_batch and LEFTNet’s aggregations) amplified through ten ODE steps: the repeat-run spread is the same size as the corrupted-slot spread, which is exactly what noise looks like and exactly what information would not. Run the leak check on CPU, where the answer is a clean zero. This is pre-existing platform behaviour shared with OA-ReactDiff, not something this integration introduced.

Scope of this integration¶

Transition state given a reactant and a product, and nothing else. Not in scope, each for a stated reason: the ei exponential-integrator solver (unpublished research variant needing a Monte-Carlo quadrature per coefficient); the other three mapping modes and five other mapping_initial values (no released weights, no data path); ts_guess conditioning (unreachable – the checkpoint records ts_guess=None); PBC / zeolite support (a different corpus, and the part of the backbone the shared vendored LEFTNet omits); validation-time sampling RMSD (runs a full sampler inside the validation loop); the size-aware DynamicBatchSampler; and the GFN2-xTB-pretrained variant the NMI paper reports (no second checkpoint ships, and its pretraining corpus is not in the Zenodo record).

Attributes¶

FRAG_FIXED

Classes¶

`ModelTaskFactory`	Hydra entry point for `configs/tasks/diffusion_reactot.yaml`.
`ReactOTTSGenerator`	Transition-state generation behind `interference/gen_reactot_ts`.
`ReactOTTask`	Task contract around `EnSB`.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_reactot.ModelTaskFactory(task_type: str = 'diffusion_reactot', 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), timesteps: int = 3000, beta_max: float = 0.3, power: float = 0.5, inv_power: float = 1, noise_schedule: str = 'cosine', 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, mapping: str = 'R+P->TS', mapping_initial: str = 'RP', sigma: float = 0.0, ts_guess: Any | None = None, idx: int = 1, nfe: int = 10, ot_ode: bool = True, solver: str = 'ode', method: str = 'midpoint', atol: float = 0.01, rtol: float = 0.01, clip_denoise: bool = True, atom_vocab: List[str] | None = None, **kwargs: Any)¶

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

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

Record the released checkpoint’s own hyperparameters.

Every default here is the released reactot-pretrained.ckpt’s hyper_parameters block, not the repo’s train_rpsb_ts1x.py; where the two disagree the checkpoint wins. See the plan’s Hyperparameter Provenance table.

Parameters:

task_type – "diffusion_reactot".
model_config – LEFTNet’s own architecture block. Required.
node_nfs – per-object input widths, [9, 9, 9].
edge_nf – 0; no edge features anywhere in this model.
condition_nf – 1, and inert – fed constant zeros.
fragment_names – ["R", "TS", "P"]; the order binds each object to its encoder/decoder pair in the state dict.
pos_dim –
update_pocket_coords – False raises NotImplementedError.
condition_time – True.
edge_cutoff – None => fully connected within a reaction.
norm_values – (1,1,1): no normalisation.
norm_biases – (0,0,0).
timesteps – the bridge grid. 3000; do not change – nothing in the state dict pins it, so a wrong value loads silently and samples wrong, and it is what makes the ODE solver’s constant-beta assertion hold.
beta_max – 0.3; do not change, same assertion.
power – 0.5.
inv_power
noise_schedule – recorded by the checkpoint and marked # not used in upstream’s train script. Carried for fidelity; nothing reads it.
precision – as noise_schedule.
loss_type – asserted in {"vlb", "l2"}, then unread by EnSB.
pos_only – coordinates only; atom identities are supplied by the input reaction and copied through.
scales – per-object loss weights – inert in ``EnSB``. They reach only upstream’s dead DDPMModule path. Recorded.
fixed_idx – [0, 2] – inert; see FRAG_FIXED.
mapping – only "R+P->TS".
mapping_initial – only "RP".
sigma – 0.0; no endpoint jitter.
ts_guess – must stay None; out of scope.
idx – 1, the transition state.
nfe – default network evaluations at sampling time.
ot_ode – True – this is what makes it deterministic.
solver – "ode", the README’s published default.
method – "midpoint".
atol – dead for a fixed-grid method.
rtol – as atol.
clip_denoise – clamp predicted structures into +/-10 A.
atom_vocab – defaults to [H, C, N, O, F].
**kwargs – swallowed; cli/train.py injects node_feature*.

Raises:

ValueError – if model_config is missing.

build() → ReactOTTask¶

Assemble dynamics + bridge schedule + normaliser into the task.

Returns:: The built ReactOTTask.

atol¶

atom_vocab¶

beta_max¶

clip_denoise = True¶

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',)¶

idx = 1¶

inv_power¶

loss_type = 'l2'¶

mapping = 'R+P->TS'¶

mapping_initial = 'RP'¶

method = 'midpoint'¶

model_config¶

nfe = 10¶

node_nfs¶

noise_schedule = 'cosine'¶

norm_biases¶

norm_values¶

ot_ode = True¶

pos_dim = 3¶

pos_only = True¶

power¶

precision¶

rtol¶

scales¶

sigma¶

solver = 'ode'¶

task: ReactOTTask | None = None¶

task_type = 'diffusion_reactot'¶

timesteps = 3000¶

ts_guess = None¶

update_pocket_coords = True¶

class MolecularDiffusion.modules.tasks.diffusion_reactot.ReactOTTSGenerator(task: Any, reaction_index: int = 0, num_generate: int = 1, batch_size: int = 1, num_steps: int | None = 10, solver: str = 'ode', method: str = 'midpoint', atol: float = 0.01, rtol: float = 0.01, output_path: str = 'generated_reactot', seed: int = 42, device: str | None = None, **kwargs: Any)¶

Bases: MolecularDiffusion.modules.tasks.ts_generator.TSGenerator

Transition-state generation behind interference/gen_reactot_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 React-OT’s: which corpus, which filters, which collate, and the solver knobs. No ``PocketGenerator`` hook is overridden.

The sampler is deterministic (ot_ode=True, sigma=0.0), so num_generate: N yields N identical structures. That is honest but wasteful, which is why gen_reactot_ts.yaml defaults num_generate: 1 / batch_size: 1 and why _settings_note() says so in the header. It is a config default and a print statement, not a missing hook.

The corpus filters differ from OAReactDiffTSGenerator in exactly one value: single_frag_only=False. React-OT’s checkpoint was trained on multi-fragment reactions too – that is its capability claim – so its generator must be able to index them. swapping_react_prod stays off for the same reason as OA-ReactDiff’s: a generator should not silently hand back the reverse reaction.

Configure the run.

Parameters:

task – the loaded ReactOTTask.
reaction_index – which reaction in the corpus to generate for.
num_generate – how many transition states to sample. Defaults to 1 because the sampler is deterministic.
batch_size – how many to sample at once.
num_steps – this is upstream’s ``–nfe``: the number of network evaluations, 10 in the README’s published command.
solver – "ode" (published) or "ddpm" (upstream’s train script default, and what their FastAPI service runs).
method – ODE method; only "midpoint" is implemented.
atol – dead for a fixed-grid method; carried for fidelity.
rtol – as atol.
output_path – directory for the .xyz files.
seed – torch/random/numpy seed. Does not affect the structure.
device – None => cuda if available.
**kwargs – rejected by the base, on purpose.

atol = 0.01¶

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

max_retries: int | None = None¶

method = 'midpoint'¶

rtol = 0.01¶

seed_numpy = True¶

solver = 'ode'¶

source_key = 'reaction_pkl'¶

tag = 'reactot'¶

class MolecularDiffusion.modules.tasks.diffusion_reactot.ReactOTTask(ddpm: MolecularDiffusion.modules.models.reactot.EnSB, nfe: int = 10, ot_ode: bool = True, solver: str = 'ode', method: str = 'midpoint', atol: float = 0.01, rtol: float = 0.01, clip_denoise: bool = True, atom_vocab: List[str] | None = None)¶

Bases: torch.nn.Module

Task contract around EnSB.

Bind the bridge to the platform’s training/generation contract.

Parameters:

ddpm – the built bridge. Named ddpm for the same reason upstream does (SBModule.ddpm): it is what the released checkpoint’s tensor prefix says, and renaming it would mean remapping 246 keys for nothing. React-OT is not a diffusion model.
nfe – default network evaluations per sample. The interference config’s num_steps overrides it.
ot_ode – True => fully deterministic, the released setting.
solver – "ode" (published default) or "ddpm".
method – ODE method; only "midpoint" is implemented.
atol – dead for a fixed-grid method; carried for fidelity.
rtol – as atol.
clip_denoise – clamp each predicted x0 into +/-10 A.
atom_vocab – output vocabulary; defaults to [H, C, N, O, F].

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

Validation metric.

Upstream’s validation-time sampling RMSD is deliberately not here: eval_sample_batch runs a full sampler inside the validation loop every epoch, and eval_rmsd sweeps the whole no-swap loader every ten (pl_trainer.py:1087-1140). Far too heavy for a training loop.

Parameters:

pred – what predict_and_target() returned.
target – ignored.

Returns:

{"val_loss": ...}.

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

One training/validation step.

This is the whole of upstream’s SBModule.compute_loss (pl_trainer.py:993-1004): unpack, call, report. loss is already a 0-dim scalar (F.mse_loss reduces by default), so nothing is reduced again here.

Parameters:: batch – (representations, conditions) from oareactdiff_collate().
Returns:: (loss, stats). scaled_err is reported because it is what upstream’s checkpoint callback actually monitors (val_ep_scaled_err).

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.

Parameters:: batch – as forward().
Returns:: (pred, target), each (1,).

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 sample(batch_size, nodesxsample, ...) has no channel for “which reaction”.

``num_steps`` IS upstream’s ``–nfe`` – the number of network evaluations, not a fraction of the 3000-step bridge grid. There is deliberately no second key spelled nfe; the family key carries it and ReactOTTSGenerator._settings_note() prints it as nfe= so the run log speaks upstream’s vocabulary.

Parameters:

batch_size – ignored; taken from batch.
nodesxsample – accepted and checked, not used to choose. A transition state has exactly as many atoms as its reaction, so a disagreeing value is a caller bug worth surfacing.
num_steps – network evaluations. None => the task’s nfe.
batch – {"representations": [...], "conditions": tensor}.
solver – "ode" / "ddpm"; None => the task’s default.
method – ODE method; None => the task’s default.
atol – None => the task’s default. Dead for midpoint.
rtol – as atol.
**kwargs – swallowed, as every other task’s sample does.

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¶

atol¶

atom_vocab¶

clip_denoise = True¶

ddpm¶

property device: torch.device¶: Device of the first parameter.

method = 'midpoint'¶

property model: ReactOTTask¶: There is no separate inner module the generation code needs.

nfe = 10¶

node_dist_model: Any = None¶

ot_ode = True¶

prop_dist_model: Any = None¶

rtol¶

solver = 'ode'¶

split = 'train'¶

MolecularDiffusion.modules.tasks.diffusion_reactot.FRAG_FIXED = [0, 2]¶