Relation-Aware Pretraining and Reaction Center Modeling for Chemical Reaction Graph Representation Learning

Official PyTorch code snapshot for the final KAGT manuscript package.

Included Modules

Shared core

• train_pretrain.py
• data_process/
• dataset/
• models/
• utils/

Reaction classification

• reaction_class/train_rc.py
• reaction_class/datasets.py
• reaction_class/hash_data.py

Reaction condition prediction

• reaction_condition_prediction/process_uspto_condition_graph.py
• reaction_condition_prediction/dataset/graph_condition_dataset.py
• reaction_condition_prediction/models/graph_condition_model.py
• reaction_condition_prediction/training/train_from_processed_graph.py
• reaction_condition_prediction/configs/uspto_condition_graph_full_lmdb_mapped_center.yaml

Yield prediction

• yield_pred/process_yield_graph.py
• yield_pred/datasets.py
• yield_pred/train_yp.py
• yield_pred/train_buchwald_hartwig_kg.py
• yield_pred/configs/kg_bh__buchwald_hartwig.yaml

Environment

From this directory:

cd code
conda env create -f environment.yml
conda activate kg_mol

Data Root

All dataset paths are resolved from:

export KG_MOL_DATA_ROOT=/path/to/react_data

Pretraining

The pretraining entry uses full-feature masked atom modeling by default: all encoder-visible atom descriptors at selected atoms are masked or zeroed, while 3D coordinates are kept as noisy denoising inputs.

Expected ORDerly files:

$KG_MOL_DATA_ROOT/ORDerly/processed/ORDerly_processed_with_rare_train_filtered.pkl
$KG_MOL_DATA_ROOT/ORDerly/processed/ORDerly_processed_with_rare_test_filtered.pkl
$KG_MOL_DATA_ROOT/ORDerly/processed/orderly_mol_features.lmdb

Run:

python train_pretrain.py \
  --train_reaction_pkl_path "$KG_MOL_DATA_ROOT/ORDerly/processed/ORDerly_processed_with_rare_train_filtered.pkl" \
  --test_reaction_pkl_path "$KG_MOL_DATA_ROOT/ORDerly/processed/ORDerly_processed_with_rare_test_filtered.pkl" \
  --lmdb_path "$KG_MOL_DATA_ROOT/ORDerly/processed/orderly_mol_features.lmdb" \
  --mask_ratio 0.15 \
  --noise_scale 0.2 \
  --gpu 0

Reaction Classification

If needed, build the hashed Schneider dataset:

python reaction_class/hash_data.py \
  --input_dir "$KG_MOL_DATA_ROOT/schneider" \
  --output_dir "$KG_MOL_DATA_ROOT/schneider/processed_reaction_class"

Train:

python reaction_class/train_rc.py \
  --dataset schneider \
  --reaction_repr_mode center \
  --center_source auto \
  --pretrain_path ./results/pretrain/<your_pretrain_run>/pretrain_best.pt \
  --gpu 0

Reaction Condition Prediction

Build the mapped LMDB benchmark used in the final runs:

python reaction_condition_prediction/process_uspto_condition_graph.py \
  --csv_path "$KG_MOL_DATA_ROOT/RG/uspto_condition_full/USPTO_condition.csv" \
  --mapped_rxn_column auto \
  --output_dir "$KG_MOL_DATA_ROOT/RG/uspto_condition_full/processed_graph_lmdb_mapped" \
  --storage_format lmdb

Train:

python -m reaction_condition_prediction.training.train_from_processed_graph \
  --config reaction_condition_prediction/configs/uspto_condition_graph_full_lmdb_mapped_center.yaml

Yield Prediction

Build the Buchwald-Hartwig KG dataset:

python yield_pred/process_yield_graph.py \
  --source_mode kg_bh \
  --output_dir "$KG_MOL_DATA_ROOT/processed_yield"

Train with the final preset:

python yield_pred/train_buchwald_hartwig_kg.py

Equivalent direct command:

python yield_pred/train_yp.py \
  --config yield_pred/configs/kg_bh__buchwald_hartwig.yaml \
  --dataset_type kg_bh \
  --dataset_name buchwald_hartwig \
  --gpu 0