DOCSEARCHbyAlgolia
Submit Docs

DiffCSP

repository·main·Indexed 41 days ago

https://github.com/jiaor17/diffcsp

An implementation of Crystal Structure Prediction by Joint Equivariant Diffusion (NeurIPS 2023). DiffCSP enables crystal structure prediction (CSP), ab initio generation, and property optimization using equivariant diffusion models.

Tokens
1.2K
Snippets
6
Records
6
Agent score
2%

What's inside DiffCSP

  1. Install DiffCSP dependencies and configure environment

    main

    To set up the DiffCSP environment, ensure you have the following Python packages installed:

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

    After installation, configure the environment by renaming .env.template to .env and setting the following absolute paths:

    • PROJECT_ROOT: The absolute path of the repository.
    • HYDRA_JOBS: The absolute path where Hydra outputs should be saved.
    • WABDB_DIR: The absolute path where WABDB outputs should be saved.
    python==3.8.13
torch==1.9.0
torch-geometric==1.7.2
pytorch_lightning==1.3.8
pymatgen==2023.8.10

  2. Perform Property Optimization

    main

    Property optimization involves three steps: training an energy model, running the optimization, and evaluating the results.

    1. Train Energy Model: Train a time-dependent energy prediction model using diffcsp/run.py with model=energy.
    2. Optimization: Run scripts/optimization.py using the trained energy model (--model_path) and the unconditional diffusion model (--uncond_path).
    3. Evaluation: Run scripts/compute_metrics.py with --tasks opt.
    # 1. Train energy model
python diffcsp/run.py data=<dataset> model=energy expname=<expname> data.datamodule.batch_size.test=100

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

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

  3. Train DiffCSP models for CSP or Ab Initio Generation

    main

    You can train models using diffcsp/run.py.

    Available Datasets (data):

    • perov_5
    • mp_20
    • mpts_52
    • carbon_24

    Tasks:

    1. Crystal Structure Prediction (CSP): Use the standard training command.
    2. Ab Initio Generation: Use the model=diffusion_w_type flag.

    Note: <expname> is an arbitrary identifier for your experiment.

    # For CSP task
python diffcsp/run.py data=<dataset> expname=<expname>

# For Ab Initio Generation task
python diffcsp/run.py data=<dataset> model=diffusion_w_type expname=<expname>

  4. Evaluate Stable Structure Prediction

    main

    To evaluate stable structure prediction, use scripts/evaluate.py followed by scripts/compute_metrics.py.

    Single Sample Evaluation:

    1. Run scripts/evaluate.py with the model path and dataset.
    2. Run scripts/compute_metrics.py specifying --tasks csp and providing the ground truth file (--gt_file).

    Multiple Samples Evaluation:

    1. Run scripts/evaluate.py with the --num_evals flag.
    2. Run scripts/compute_metrics.py with the --multi_eval flag.
    # Single 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

  5. Evaluate Ab Initio Generation

    main

    To evaluate Ab Initio Generation, use scripts/generation.py and then compute metrics.

    1. Generate samples using scripts/generation.py.
    2. Compute metrics using scripts/compute_metrics.py with --tasks gen.
    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

  6. Sample from an arbitrary chemical formula

    main

    Use scripts/sample.py to generate crystal structures for a specific chemical composition.

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