2090741942justin/mp_20_pxrdnet

XRDnet: Ab Initio Nanostructure Solutions from PXRD via Score-Based Generative Modeling

This is the repository for XRDnet, the world's first end-to-end nanostructure solver from powder x-ray diffraction (PXRD) patterns. Associated article: Ab Initio Structure Solutions from Nanocrystalline Powder Diffraction Data via Deep Generative Modeling.

All code blocks assume you start from this directory.

Much thanks to CDVAE.

Visualization of Process

Li₂V₂F₁₂

Tm₆Sc₂

Requirements

Use Python 3.9.18 with Linux.

In our experience, depending on the system, you may have trouble with:

• PyTorch Scatter
• PyTorch Sparse

If so, just follow the instructions on their GitHub repos to install the versions that align with your CUDA version. Here are the suggested ways to do so:

pip install torch==2.0.0
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.html
pip install torch-sparse -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.html

After that, run the following command to install remaining requirements.

pip install -r requirements.txt

Setup environment variables by making a copy of the .env.template file and rename it to .env. Modify the following environment variables in .env.

• PROJECT_ROOT: path to the folder that contains this repo
• HYDRA_JOBS: path to a folder to store hydra outputs
• WABDB: path to a folder to store wabdb outputs

Finally, install this package with

pip install -e .

This process, if done correctly, should take less than an hour.

Pre-Trained Model Weights

You can find our trained model weights on Huggingface.

Data Procurement

Use Our Data

See data/mp_20 for MP-20-PXRD, our modification of MP-20 with PXRD patterns (no broadening added).

See data/experimental_data for instructions on obtaining the crystal structures with experimentally observed PXRD patterns (we do not own them; IUCr does).

Despite the filenames, they're not .csv files (they're pickle files), because the data contains PyTorch tensors.

No other action needs to be taken on data, as the scripts automatically load the datasets.

Generate It Yourself

This is unnecessary for MP-20, because we already have created and uploaded the datasets for you. However, if you wish to try something new on the MP-20 dataset (let's say, different wavelengths or data splits), you can follow these instructions.

Before running, replace /home/gabeguo/ in create_data.sh with your home directory. This takes less than an hour.

cd scripts
bash create_data.sh

Data should be saved in data/mp_20.

PXRDGen Raw 0-90/0.1 To 512 Experiment

This workspace adds a minimal-change path for training XRDnet on PXRDGen-style ideal pymatgen XRD patterns while keeping the original 512-dimensional XRDnet model interface.

What Was Changed

• cdvae/pl_data/dataset.py

  • Added xrd_filter: raw.
  • Raw mode skips sinc broadening, Gaussian broadening, and vertical noise.
  • Generalized max-pooling downsampling with np.array_split, so non-divisible lengths such as 900 -> 512 work correctly.
  • Removed hard-coded 512 shape assumptions from dataset item construction.

• cdvae/pl_modules/xrd.py and cdvae/pl_modules/model.py

  • Kept the default XRD head at 512 dimensions.
  • Parameterized the XRD head with data.n_postsubsample, so the raw experiment still uses 512 but future dimensions do not hit stale asserts.
  • Made XRD-head min-max normalization numerically stable with an epsilon and clamp, instead of stopping training on floating-point/debug assertions.

• visualization/visualize_materials.py, scripts/gen_xrd.py, and scripts/conditional_generation.py

  • Candidate structures generated during evaluation now use the configured XRD range and vector length.
  • For this experiment that means pymatgen ideal XRD over 2theta=0-90 with 900 source bins, then raw max-pooling to 512 for candidate ranking.
  • sample1 and sample20 are represented as best-of-1 and best-of-20 conditional generations using StructureMatcher for accuracy.

• Added conf/data/mp_20_pxrdgen_xrd90_0p1_raw512.yaml

  • Uses xrd_filter: raw.
  • Stores/loads 900-bin source XRD from 2theta=0-90, bin=0.1.
  • Trains/evaluates on 512 bins after dataset max-pooling.

• Added scripts/build_pxrdgen_raw512_dataset.py

  • Converts PXRDGen .pt splits to XRDnet pickle files.
  • Recomputes XRD from CIF with pymatgen.analysis.diffraction.xrd.XRDCalculator.
  • Uses ideal peak intensities, max intensity per 2theta bin, and max-intensity normalization.

• cdvae/run.py

  • Uses the current in-memory model for trainer.test(..., ckpt_path=None), which avoids Lightning 2.2 trying to load a nonexistent best checkpoint in short smoke runs.

• Added scripts/run_pxrdgen_raw512_end_to_end.sh

  • Runs data conversion, training, sample1 evaluation, and sample20 evaluation.
  • Logs each stage elapsed time in seconds and hours.
  • The final stage_summary.tsv line reports total training/evaluation time excluding data creation plus maximum GPU memory used.

• Added scripts/smoke_pxrdgen_raw512.sh

  • Runs the same flow on a tiny limited split with one training batch and one tested material.

Environment

Use the conda environment configured for this workspace:

source /workspace/miniforge/bin/activate pxrdnet
cd /workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval
export PROJECT_ROOT=/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval
export HYDRA_JOBS=/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/hydra
export WABDB_DIR=/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/wabdb
export PYTHONPATH=/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval:${PYTHONPATH}
export WANDB_MODE=disabled

Smoke Test

cd /workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval
bash scripts/smoke_pxrdgen_raw512.sh

Smoke outputs are written under:

/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/data/mp_20_pxrdgen_xrd90_0p1_raw512_smoke/
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/hydra/singlerun/<date>/<run_tag>/
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/paper_results_pxrdgen_raw512/<run_tag>/

Full End-To-End Run

cd /workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval
bash scripts/run_pxrdgen_raw512_end_to_end.sh

Useful overrides:

RUN_TAG=pxrdgen_raw512_run01 CUDA_VISIBLE_DEVICES=0 bash scripts/run_pxrdgen_raw512_end_to_end.sh
NUM_TESTED_MATERIALS=200 NUM_GRADIENT_STEPS=5000 N_STEP_EACH=100 bash scripts/run_pxrdgen_raw512_end_to_end.sh
FORCE_REBUILD_DATA=1 bash scripts/run_pxrdgen_raw512_end_to_end.sh

By default, the data stage reuses an existing data directory when train.csv, val.csv, test.csv, and manifest.json are already present. Changing RUN_TAG only changes logs, Hydra outputs, checkpoints, and result folders; it does not force data regeneration. Set FORCE_REBUILD_DATA=1 only when the source splits or XRD generation settings changed.

The main logs are:

/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/end_to_end.log
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/stage_summary.tsv
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/summary.txt
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/data.log
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/train.log
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/sample1.log
/workspace/cdvae_xrd_pxrdgen_raw512_pxrdgen_match_eval/logs/<run_tag>/sample20.log

end_to_end.log records every stage start, command, stage log path, elapsed seconds, elapsed hours, and max GPU memory seen so far. Each stage also streams stdout/stderr through tee, so tqdm/progress output is visible in the terminal while still being saved to the per-stage log file.

The sample result folders save ground-truth CIFs, predicted candidate CIFs, generated XRD tensors, per-material metrics JSON, and aggregate metrics JSON. These files are the intended starting point for later Rwp calculation.

Training

This trains the $10 &#8491$ and $100 &#8491$ models. They will be saved under ../hydra/singlerun/[today's date].

On a single GeForce RTX 3090 (24 GB), each model should take about one day to train.

cd scripts
CUDA_VISIBLE_DEVICES=x bash train_mp20_model_sinc10.sh
CUDA_VISIBLE_DEVICES=x bash train_mp20_model_sinc100.sh

Evaluation

You will have to change --model_path inside each script to have the appropriate home directory (rather than /home/gabeguo/) and date (rather than 2024-04-07).

On a single GeForce RTX 3090 (24 GB), each evaluation (per model) should take about one day to conduct.

MP-20

XRDnet

cd scripts
CUDA_VISIBLE_DEVICES=x bash conditional_generation_sinc10.sh
CUDA_VISIBLE_DEVICES=x bash conditional_generation_sinc100.sh

Semi-Random Baseline

cd scripts
CUDA_VISIBLE_DEVICES=x bash conditional_generation_random_baseline_sinc10.sh
CUDA_VISIBLE_DEVICES=x bash conditional_generation_random_baseline_sinc100.sh

Latent Space Search Baseline

cd scripts
CUDA_VISIBLE_DEVICES=x bash conditional_generation_sinc10_baseline_noOpt.sh
CUDA_VISIBLE_DEVICES=x bash conditional_generation_sinc100_baseline_noOpt.sh

Experimentally Collected Data

Getting Correct Configs

Your file directory should look something like this:

cdvae_xrd/
  ... [some stuff here] ...
hydra/singlerun/
  [whatever date you trained model on]/
    mp_20_sinc10/
      .hydra/
        config.yaml
        hydra.yaml
        overrides.yaml
      hparams.yaml
      ... [other stuff here] ...
    mp_20_sinc100/
      ... [same stuff here] ...

Run the following code (assuming you are in cdvae_xrd) to create the proper evaluation setup for experimental data:

cd ../hydra/singlerun/[whatever date you trained model on]
cp mp_20_sinc10 mp_20_sinc10_EXPERIMENTAL_TEST

Now, go into mp_20_sinc10_EXPERIMENTAL_TEST/.hydra/config.yaml and change line 7 to be:

root_path: ${oc.env:PROJECT_ROOT}/data/experimental_xrd

from root_path: ${oc.env:PROJECT_ROOT}/data/mp_20

Do exactly the same change for mp_20_sinc10_experimental/hparams.yaml.

Running the Evals

Again, remember to change --model_path inside each script to have the appropriate home directory (rather than /home/gabeguo/) and date (rather than 2024-04-07).

This should only take a few hours at most, due to there being fewer experimental PXRD patterns.

cd scripts
CUDA_VISIBLE_DEVICES=x bash conditional_generation_experimental.sh
CUDA_VISIBLE_DEVICES=x bash conditional_generation_baseline_noOpt.sh
CUDA_VISIBLE_DEVICES=x bash conditional_generation_random_baseline_experimental.sh

Getting R-Factors

As before, in the __main__ part, change the home directory from /home/gabeguo/ to whatever your home directory is.

This should take less than an hour.

cd scripts
python calculate_xrd_patterns_post_hoc.py
python calculate_r_factor_post_hoc.py
bash calc_r_value_distribution.sh

Calculating Results by Crystal System

Reiterating (as you've already guessed), in extract_results_by_crystal_system.sh, change the home directory from /home/gabeguo/ to whatever your home directory is.

This should take less than an hour.

cd scripts
bash extract_results_by_crystal_system.sh

Citation

If you find the code in this repository helpful, please cite the following:

@article{guo2024diffusion,
  title={Ab Initio Structure Solutions from Nanocrystalline Powder Diffraction Data},
  author={Guo, Gabe and Saidi, Tristan and Terban, Maxwell and Valsecchi, Michele and Billinge, Simon JL and Lipson, Hod},
  journal={arXiv preprint arXiv:2406.10796},
  year={2024}
}