Skip to content

[NeurIPS 2023] The implementation for the paper "Crystal Structure Prediction by Joint Equivariant Diffusion"

License

Notifications You must be signed in to change notification settings

jiaor17/DiffCSP

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

21 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Crystal Structure Prediction by Joint Equivariant Diffusion (NeurIPS 2023)

Implementation codes for Crystal Structure Prediction by Joint Equivariant Diffusion (DiffCSP).

License: MIT [Paper]

Overview

Demo

Dependencies and Setup

python==3.8.13
torch==1.9.0
torch-geometric==1.7.2
pytorch_lightning==1.3.8
pymatgen==2023.8.10

Rename the .env.template file into .env and specify the following variables.

PROJECT_ROOT: the absolute path of this repo
HYDRA_JOBS: the absolute path to save hydra outputs
WABDB_DIR: the absolute path to save wabdb outputs

Training

For the CSP task

python diffcsp/run.py data=<dataset> expname=<expname>

For the Ab Initio Generation task

python diffcsp/run.py data=<dataset> model=diffusion_w_type expname=<expname>

The <dataset> tag can be selected from perov_5, mp_20, mpts_52 and carbon_24, and the <expname> tag can be an arbitrary name to identify each experiment. Pre-trained checkpoints are provided here.

Evaluation

Stable structure prediction

One sample

python scripts/evaluate.py --model_path <model_path> --dataset <dataset>
python scripts/compute_metrics.py --root_path <model_path> --tasks csp --gt_file data/<dataset>/test.csv 

Multiple samples

python scripts/evaluate.py --model_path <model_path> --dataset <dataset> --num_evals 20
python scripts/compute_metrics.py --root_path <model_path> --tasks csp --gt_file data/<dataset>/test.csv --multi_eval

Ab initio generation

python scripts/generation.py --model_path <model_path> --dataset <dataset>
python scripts/compute_metrics.py --root_path <model_path> --tasks gen --gt_file data/<dataset>/test.csv

Sample from arbitrary composition

python scripts/sample.py --model_path <model_path> --save_path <save_path> --formula <formula> --num_evals <num_evals>

Property Optimization

# train a time-dependent energy prediction model 
python diffcsp/run.py data=<dataset> model=energy expname=<expname> data.datamodule.batch_size.test=100

# Optimization
python scripts/optimization.py --model_path <energy_model_path> --uncond_path <model_path>

# Evaluation
python scripts/compute_metrics.py --root_path <energy_model_path> --tasks opt

Acknowledgments

The main framework of this codebase is build upon CDVAE. For the datasets, Perov-5, Carbon-24 and MP-20 are from CDVAE, and MPTS-52 is collected from its original codebase.

Citation

Please consider citing our work if you find it helpful:

@article{jiao2023crystal,
  title={Crystal structure prediction by joint equivariant diffusion},
  author={Jiao, Rui and Huang, Wenbing and Lin, Peijia and Han, Jiaqi and Chen, Pin and Lu, Yutong and Liu, Yang},
  journal={arXiv preprint arXiv:2309.04475},
  year={2023}
}

Contact

If you have any questions, feel free to reach us at:

Rui Jiao: [email protected]

About

[NeurIPS 2023] The implementation for the paper "Crystal Structure Prediction by Joint Equivariant Diffusion"

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages