This repository is the official implementation of the paper:
Tsypin, A., Ugadiarov, L. A., Khrabrov, K., Telepov, A., Rumiantsev, E., Skrynnik, A., ... & Kadurin, A. (2023, October).
Gradual Optimization Learning for Conformational Energy Minimization.
In The Twelfth International Conference on Learning Representations.
Experiments and results on the SPICE dataset can be found in the "GOLF-SPICE" branch.
Model |
|
|||||
---|---|---|---|---|---|---|
RDKit | ||||||
Torsional Diffusion | ||||||
ConfOpt | ||||||
Uni-Mol+ | ||||||
|
||||||
- Set up environment on the GPU machine.
# On the GPU machine ./scripts/setup_gpu_env.sh conda activate GOLF_schnetpack pip install -r requirements.txt
- Download training dataset
$\mathcal{D}_0$ and evaluation dataset$\mathcal{D}_{\text{test}}$ mkdir data && cd data wget https://sc.link/FpEvS -O D-0.db wget https://sc.link/W6RUA -O D-test.db cd ../
- Train baseline PaiNN model
Running this script will create a folder in the specified
cd scripts/train ./run_training_baseline.sh <cuda_device_number>
log_dir
directory (we use "./results" in our configs and scripts). The name of the folder is specified by theexp_name
hyperparameter. The folder will contain checkpoints, a metrics file and a config file with hyperparameters.
- Set up environment on the GPU machine like in the first section
- Download optimization trajectories datasets.
cd data wget https://sc.link/ZQRiV -O D-traj-10k.db wget https://sc.link/Z0ebo -O D-traj-100k.db wget https://sc.link/hj1JX -O D-traj-500k.db cd ../
- Train PaiNN.
cd scripts/train ./run_training_trajectories-10k.sh <cuda_device_number> ./run_training_trajectories-100k.sh <cuda_device_number> ./run_training_trajectories-500k.sh <cuda_device_number>
To speed up the training, we parallelize DFT computations using several CPU-rich machines. The training of the NNP takes place on the parent machine with a GPU.
-
Set up environment on the GPU machine like in the first section
-
Log in to CPU-rich machines. They must be accessible via
ssh
. -
Set up environments on CPU-rich machines.
# On CPU rich machines git clone https://github.com/AIRI-Institute/GOLF cd GOLF/scripts ./setup_dft_workers.sh <num_threads> <num_workers> <start_port>
Here,
n_ports
denotes number of workers on a CPU-rich machine, andports_range_begin
denotes the starting port numbers for workers. Workers calculate energies and forces usingpsi4
for newly generated conformations. For example,./setup_host.sh 24 20000
will launch a total of 48 workers listening to ports20000, ... , 20023
. You can change theports_range_begin
inenv/dft.py
.By default we assume that each worker uses 4 CPU-cores (can be changed in
env/dft_worker.py
, line 22) which means thatn_ports
must be less or equal tototal_cpu_number / 4
. -
Add ip addresses of CPU rich machines to a text file. We use
env/host_names.txt
.
Train PaiNN with GOLF.
cd scripts/train
./run_training_GOLF-10k.sh <cuda_device_number>
The evaluation can be done with or without psi4
energy estimation for NNP-optimization trajectories. The argument 'eval_early_stop_steps' controls for which conformations in the optimization trajectory to evaluate energy/forces with psi4
. For example, setting eval_early_stop_steps
to an empty list will result in no additional psi4
energy estimations, and setting it to [1 2 3 5 8 13 21 30 50 75 100]
will result in 13 additional energy evaluations for each conformation in evaluation dataset. Note that in order to compute the psi4.optimize
with spherical representation of the molecule was used (approximately 30 steps until convergence).
In this repo, we provide NNPs pre-trained on different datasets and with GOLF in the checkpoints
directory:
-
$f^{\text{baseline}}$ (checkpoints/baseline-NNP/NNP_checkpoint
) -
$f^{\text{traj-10k}}$ (checkpoints/traj-10k/NNP_checkpoint
) -
$f^{\text{traj-100k}}$ (checkpoints/traj-100k/NNP_checkpoint
) -
$f^{\text{traj-500k}}$ (checkpoints/trak-500k/NNP_checkpoint
) -
$f^{\text{GOLF-1k}}$ (checkpoints/GOLF-1k/NNP_checkpoint
) -
$f^{\text{GOLF-10k}}$ (checkpoints/GOLF-10k/NNP_checkpoint
)
For example, to evaluate GOLF-10k and additionally calculate psi4
energies/forces along the optimization trajectory, run:
python evaluate_batch_dft.py --checkpoint_path checkpoints/GOLF-10k --agent_path NNP_checkpoint_actor --n_parallel 240 --n_threads 24 --conf_number -1 --host_file_path env/host_names.txt --eval_db_path data/GOLF_test.db --timelimit 100 --terminate_on_negative_reward False --reward dft --minimize_on_every_step False --eval_early_stop_steps 1 2 3 5 8 13 21 30 50 75 100
Make sure that n_threads
is equal to the number of workers on each CPU-rich machine. Setting n_threads
to a larger number will result in optimization failures. If you wish to only evaluate the last state in each optimization trajectory, set timelimit
and eval_early_stop_steps
to the same number: --timelimit T --eval_early_stop_steps T
.
After the evaluation is finished, an evaluation_metrics.json
file with per-step metrics will be created. Each record in evaluation_metrics.json
describes optimization statistics for a single conformation and contains such metrics as: forces/energies MSE, percentage of optimized energy, predicted and ground-truth energies, etc. The final NNP-optimized conformations are stored in results.db
database.
To cite this work, please use:
@inproceedings{tsypin2023gradual,
title={Gradual Optimization Learning for Conformational Energy Minimization},
author={Tsypin, Artem and Ugadiarov, Leonid Anatolievich and Khrabrov, Kuzma and Telepov, Alexander and Rumiantsev, Egor and Skrynnik, Alexey and Panov, Aleksandr and Vetrov, Dmitry P and Tutubalina, Elena and Kadurin, Artur},
booktitle={The Twelfth International Conference on Learning Representations},
year={2023}
}