DiffCSP++

Official Implementation of Space Group Constrained Crystal Generation (DiffCSP++).

Setup

Complete .env with the following correctly filled out

PROJECT_ROOT="${THIS_FOLDER}"
HYDRA_JOBS="${THIS_FOLDER}/results"
WANDB_DIR="${THIS_FOLDER}"
USE_WANDB_LOGGING=1

One can specifiy USE_WANDB_LOGGING=0 to disable wandb.

micromamba env create -f environment.yml
micromamba activate diffcsp_pp

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. Support for the alex_mp_20 dataset can be added by following the extension instructions available in this repository. Pre-trained checkpoints are provided here.

Evaluation

Crystal Structure Prediction

python scripts/evaluate.py --model_path <model_path> --dataset <dataset> # with GT Wyckoff positions
python script/csp_from_template.py --model_path <csp model> --csv_path <path containing *_comp_fps.csv> --finder_model_path <CSPML model> # with CSPML templates
python scripts/compute_metrics.py --root_path <model_path> --tasks csp --gt_file data/<dataset>/test.csv 

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 pre-defined symmetries

One can sample structures given the symmetry information by the following codes:

python scripts/sample.py --model_path <model_path> --save_path <save_path> --spacegroup <spacegroup> --wyckoff_letters <wyckoff_letters> --atom_types <atom_types>

One example can be

python scripts/sample.py --model_path mp_csp --save_path MnLiO --spacegroup 58 --wyckoff_letters 2a,2d,4g --atom_types Mn,Li,O
# Or simplify WPs as
python scripts/sample.py --model_path mp_csp --save_path MnLiO --spacegroup 58 --wyckoff_letters adg --atom_types Mn,Li,O

If multiple structures are required to be generated, one can first collect the inputs into a json file like example/example.json:

[
    {
        "spacegroup_number": 58,
        "wyckoff_letters": ["2a","2d","4g"],
        "atom_types": ["Mn","Li","O"]
    },
    {
        "spacegroup_number": 194,
        "wyckoff_letters": "abff",
        "atom_types": ["Tm","Tm","Ni","As"]
    }
]

And parallelly generate the structures via

python scripts/sample.py --model_path <model_path> --save_path <save_path> --json_file <json_file>

Notably, the atom_types parameter is not required if a model for ab initio generation is used.