CheMeleon-Rxn

Descriptor-regression pretraining for reaction property prediction.

CheMeleon-Rxn is the reaction analog of CheMeleon (Burns, Zalte et al., 2026). It pretrains a Chemprop CGR (Condensed Graph of Reaction) D-MPNN by regressing a reaction's graph embedding onto a dense reaction-level descriptor built from classical per-molecule descriptors, then fine-tunes the pretrained encoder on small barrier-height and reactivity benchmarks. No labels are used during pretraining — the target is a self-supervised descriptor vector — so the corpus can be scaled to millions of unlabeled reactions.

This repository accompanies the manuscript "Descriptor-Based Pre-Training Improves Reaction Property Prediction" (Zalte, Fang, Metwally, Cheng, Glisman, Pang) and is released for reproducibility. A preprint is coming soon.

How it works

1. Corpus. Pool atom-mapped reactant→product reaction SMILES from several public sources into one deduplicated, leakage-guarded set (pipeline/corpus/).
2. Descriptor target. For each reaction, compute classical per-molecule descriptors (Mordred or RDKit-2D), sum-pool each side, and combine as [Σ_reactants ‖ (Σ_products − Σ_reactants)] — a dense, permutation- and count-invariant reaction descriptor (~3226 dims), then winsorize + z-score (pipeline/descriptors/).
3. Pretrain. Train a CGR D-MPNN to regress its reaction embedding onto that descriptor, with random target masking (models/gnn/pretrain.py).
4. Fine-tune. Load the pretrained encoder and fine-tune on each downstream task; the from-scratch baseline is the same code path with no checkpoint (models/gnn/evaluate.py).

Reactants and products only — no reagents, solvents, or catalysts enter the target or the model input.

Installation

The code is imported via sys.path (top-level directories), so there is nothing to pip install -e. Always run commands from the repository root. Create the environment with conda:

conda env create -f environment.yml
conda activate chemeleon_rxn

or with pip into an existing Python 3.12 environment:

pip install -r requirements.txt

This single chemeleon_rxn environment runs everything except the RXNFP baseline: corpus building, descriptor baking, pretraining, fine-tuning, and all other baselines. scripts/setup_env.sh reproduces it from scratch.

cuik-molmaker (pretraining only). Pretraining featurizes the full ~2M-reaction corpus on the fly, so it uses the accelerated cuik-molmaker backend in chemprop==2.3.1 (the first public release carrying it). This requires cuik-molmaker-pin and rdkit pinned to the same release (2026.3.4); a mismatch breaks its C++ ABI. Fine-tuning and every baseline use chemprop's standard CGR featurizer and do not need cuik-molmaker — so if you only fine-tune from a published checkpoint, the pin is irrelevant. For CUDA training, install a CUDA-enabled PyTorch build for your platform first (see pytorch.org); otherwise chemprop pulls in a CPU-only torch wheel.

RXNFP baseline (separate environment). The RXNFP baseline pins an older transformers/tokenizers stack that conflicts with this environment, so it runs in its own rxnfp conda env (conda env create -f models/rxnfp/environment.yml). It is a two-stage flow: precompute fingerprints in the rxnfp env, then train the MLP head in chemeleon_rxn; see models/rxnfp/README.md.

Data and checkpoints

The downstream benchmarks fetch reproducibly through the registry in pipeline/corpus/sources.py — each source exposes a fetch(raw_dir) that downloads and unpacks into its own folder under data/raw/. The barrier benchmarks come from the "Benchmark Data for Chemprop" archive (Zenodo 10.5281/zenodo.8174268, CC-BY-4.0). data/ is gitignored.

The pretraining corpus is scaled up with reactions from Pistachio, which is licensed and cannot be redistributed here — so the full pretraining corpus is not publicly reproducible. To let others use the pretrained model without rebuilding the corpus, the pretrained checkpoints will be published on Hugging Face (link added on release). The corpus-building and pretraining code is included so the pipeline is fully auditable, and the fine-tuning benchmarks above are entirely public.

Downstream tasks

A dataset is a reaction set with one or more targets; a task is a (dataset, target) pair keyed <dataset>_<target>. Splits are per-dataset and task-agnostic: 5 cross-validation splits with a fixed test set and reshuffled train/val, shared by every target of a dataset. The eight regression tasks, in ascending training size:

Task	Dataset	Target
e2_ea	E2 barriers	activation energy
sn2_ea	SN2 barriers	activation energy
cycloadd_gact	cycloaddition	activation free energy
cycloadd_grxn	cycloaddition	reaction free energy
rdb7_ea	RDB7 (CCSD(T)-F12)	activation energy
rdb7_dhrxn	RDB7 (CCSD(T)-F12)	reaction enthalpy
ha_cn_yield_lcap	Ha et al. C–N coupling HTE	yield (LCAP)
lograte_logk	H-abstraction rates	rate coefficient (log k)

Reproducing the results

Quick path — one script, all models. Download the data/ bundle into the repo root (pretrained checkpoints + benchmark splits; see "Data and checkpoints"), activate the chemeleon_rxn env on a machine with a CUDA GPU, then:

bash scripts/run_all_experiments.sh

This fine-tunes CheMeleon-Rxn and every ablation from the released checkpoints and runs all baselines, writing results/<name>/_overall.md (+ metrics.csv) for each. It does not re-run pretraining (that needs the licensed corpus — pass PRETRAIN=1 if you have it) or the RXNFP baseline (separate env — pass RXNFP=1). Then regenerate the paper's figures and tables with analysis/paper_figures.ipynb.

To run stages individually instead, every step is a plain script runnable as bash <script> from the repository root with the environment active — no scheduler assumptions. On a cluster (PBS, SLURM, …) wrap the same command in your submission script; on a workstation with a GPU, run it directly. The CPU-heavy steps parallelize over all available cores; pretraining and fine-tuning use a CUDA GPU.

# 0. Materialize benchmark splits ONCE (prerequisite for the corpus leakage guard
#    and for fine-tuning). CPU.
bash pipeline/splits/make_splits_all.sh

# 1. Build the pretraining corpus, then bake all four descriptor sets
#    (mordred/rdkit2d × sum/mean). CPU, many cores, several hours (Mordred is the pole).
#    NOTE: the full corpus draws on licensed Pistachio reactions (see "Data and
#    checkpoints"), so this stage is only fully reproducible with a Pistachio license.
#    To skip it, use the published pretrained checkpoints instead and jump to step 3.
bash pipeline/build_all_hpc.sh
#    (or run the stages individually: pipeline/corpus/prepare_hpc.sh then
#     pipeline/descriptors/{molecules,reactions,scale}_hpc.sh)

# 2. Pretrain the hero CGR D-MPNN on the scaled descriptor targets. GPU.
CONFIG=configs/pretrain/chemeleon_rxn.yaml bash models/gnn/pretrain.sh

# 3. Fine-tune the pretrained encoder on all 8 tasks (from step 2's checkpoint or a
#    published one). GPU.
CONFIG=configs/pretrain/chemeleon_rxn.yaml bash models/gnn/finetune.sh

# 4. Baselines (scratch D-MPNN, hybrid GNN+descriptor, descriptor MLP, MorganDiff,
#    GIN/GINE) — fully public, no corpus needed. GPU.
bash models/run_all_baselines.sh

Ensembling and scoring. For each (task, split), 3 models are trained with distinct seeds and their test predictions averaged, giving one ensemble score per split; the 5 per-split scores are the Tukey-HSD replicates. Each run writes results/<tag>/_overall.md, which the analysis notebooks in analysis/ glob to produce the significance tables and figures.

Featurization note. KEEP_H=True keeps explicit atom-mapped hydrogens: the barrier sets encode the reaction center as H transfer, so chemprop's default RemoveHs would delete the reaction-defining atoms. This is applied identically in pretraining and fine-tuning so the encoder sees the same CGR graph.

Repository layout

chemeleon_rxn/   shared core: config, reaction/CGR handling, descriptor backends
                 (registry), dataset loaders, split IO, zarr dataset, scaling
pipeline/        data prep: corpus building, descriptor Stages A/B/C, split generation
models/          model families: gnn/ (pretrain + evaluate + encoders), gnn_descriptors/,
                 descriptor_mlp/, morgan_diff/, rxnfp/
configs/         per-experiment YAML: pretrain/ (CheMeleon-Rxn + ablations),
                 finetune/ (freeze-schedule study), baselines/
scripts/         run_all_experiments.sh (one-command reproduction) + env setup, TMAP
analysis/        notebooks: Tukey-HSD significance, figures, TMAP visualization
results/         per-run outputs; only _overall.md + metrics.csv are tracked
reports/         generated EDA figures (gitignored; reproducible from analysis/)

The reaction D-MPNN and its from-scratch baseline are one code path in models/gnn/, toggled by an optional pretrained-checkpoint argument. Descriptor backends (mordred/rdkit2d) and reaction operators (difference/diff_concat/sum_concat) are registries, so the pretraining target is itself an experimental axis.

Citation

If you use this code, please cite the manuscript (preprint coming soon):

Zalte, A. S.; Fang, C.; Metwally, E.; Cheng, A.; Glisman, A.; Pang, H.-W. Descriptor-Based Pre-Training Improves Reaction Property Prediction.

and the original CheMeleon work it builds on:

Burns, J. W.; Zalte, A. S.; Abreu, C. R. A.; Sieg, J.; Feldmann, C.; Mathea, M.; Green, W. H. Deep Learning Foundation Models for Low-Data Regimes from Classical Molecular Descriptors. J. Chem. Inf. Model. 2026. 10.1021/acs.jcim.6c01546

License

Released under the MIT License.