ECTS

Description

This is the official code of "ECTS: An ultra-fast diffusion model for exploring chemical reactions with equivariant consistency".

Unveiling reaction mechanisms through the exploration of reaction paths, including identification of transition states (TS), prediction of reaction energy barriers (Ea), and mapping of reaction pathways, is crucial for the study of chemical reactions. However, this process usually requires extensive and computationally demanding quantum chemistry calculations. Here, we propose an equivariant consistency generative model ECTS, an ultra-fast diffusion method that unifies TS generation, energy prediction, and pathway search within one framework. Our results highlight that the efficiency of ECTS is at least two orders of magnitude higher than conventional diffusion models. TS structures generated by ECTS exhibit an error margin of just 0.12 Å root mean square deviation compared to the ground truth. Additionally, by continuously refining the energy barrier predictions in the denoising process, ECTS achieves a median error of merely 2.4 kcal/mol without any post-DFT calculations. Moreover, as a novel feature, ECTS can also generate reaction paths which are in general agreement with the true reaction paths, indicating ECTS could potentially be useful for exploring reaction mechanisms.

Installations

There are two choice for installation of ECTS, the first is install ECTS from source code as following. Environment install:

conda env create -f environment.yaml 

EcTs installation:

cd ECTS

pip install -e .

Second, ECTS has been integrated in MLatom 3.18 (dralgroup/mlatom: AI-enhanced computational chemistry). user can install MLatom as following.

python3 -m pip install -U MLatom

The turtorial of ECTS in MLatom is available from https://xacs.xmu.edu.cn/docs/mlatom/tutorial_re_explore.html

Users can sign up a account of Aitomistic Platform (https://www.aitomistic.xyz) for online computations with ECTS.

Download pretrained models

The pretrained models of ECTS is avaliable from https://figshare.com/account/projects/255383/articles/29487914

Processed datasets

The RXNID of reactions in different datasets are provided in ECTS/datasets. The processed datasets is avaliable from https://figshare.com/account/projects/255383/articles/29487893

Training

Train the TS predictor $f_{Ts}$ and Energy predictor $f_E$ together:

cd scripts/train_ts

CUDA_VISIBLE_DEVICES=0,1 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 2 --rdzv_id 1 train_ts_and_e.py -i ctrl_ts_e.json

Train the TS predictor $f_{Ts}$ only:

cd scripts/train_ts

CUDA_VISIBLE_DEVICES=0,1 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 2 --rdzv_id 1 train_ts.py -i ctrl_ts.json

Train the energy predictor $f_E$ only:

cd scripts/train_ts

CUDA_VISIBLE_DEVICES=0,1 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 2 --rdzv_id 1 train_e.py -i ctrl_e.json

Train the path structure generator $f_{Path}$ only:

cd scripts/train_path

CUDA_VISIBLE_DEVICES=0,1 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 2 --rdzv_id 1 train_path.py -i ctrl_path.json

cd EcTs_Model/model 

mv online_model_perepoch.cpk path_online_model_perepoch.cpk

User should train the $f_{Path}$ in another directory to avoid the file overwriting.

Functions

Examples for TS generation

cd scripts/sample_ts

CUDA_VISIBLE_DEVICES=0 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 1 --rdzv_id 1 sample.py -i ctrl_sample.json -r ./rxn3086/r.xyz -p ./rxn3086/p.xyz -n rxn3086 --steps 1

Results were saved in samples/rxn3086 folder.

Examples for Path interpolation

cd scripts/sample_path 	

CUDA_VISIBLE_DEVICES=0 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 1 --rdzv_id 1 sample_path.py -i ctrl_sample.json -r ./rxn3086/r.xyz -p ./rxn3086/p.xyz -n rxn3086 --steps 1

Results were saved in sample_pathes/rxn3086 folder.

Evaluation

To evaluate TS structure and $E_a$ predicted by ECTS on the test set

cd scripts/eval_ts

CUDA_VISIBLE_DEVICES=0 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 1 --rdzv_id 1 eval.py -i ctrl_sample.json  --steps 1

To evaluate ECTS-Path on the test set. Users should download the refence pathes from https://figshare.com/account/projects/255383/articles/29503472 for evaluation.

cd scripts/eval_path

CUDA_VISIBLE_DEVICES=3 torchrun --rdzv_backend c10d --rdzv_endpoint localhost:0 --nnodes 1 --nproc_per_node 1 --rdzv_id 1 eval_path.py -i ctrl_sample.json  --steps 1

Link to papers

https://chemrxiv.org/engage/chemrxiv/article-details/67cf041bfa469535b9bc28d4

Citations

1. Xu M, Li B, Dong Z, Dral P, Zhu T, Chen H. ECTS: An ultra-fast diffusion model for exploring chemical reactions with equivariant consistency. ChemRxiv. 2025; doi:10.26434/chemrxiv-2025-f9vdp