Learning Diffusion Model for Atomic Structure Generation

This project is for me to learn diffusion models for atomic structure generation, serving as a starting point and playground for testing different models.

Currently, almost all codes are adapted from the graphite repository.

Environment Setup

# Create conda environment with PyTorch 2.8 and CUDA 12.9
conda create -n genai python=3.11
conda activate genai

# Install PyTorch
pip install torch==2.8 --index-url https://download.pytorch.org/whl/cu129

# Install other packages
pip install lightning torch-geometric  torch-cluster
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.8.0+cu129.html
pip install ase dpdata numpy scipy matplotlib tqdm wandb jupyterlab

Data

The amorphous structures for Li-Si-P-S were salvaged from the DP-Gen process during the MLFF development for my previous work (ACS Appl. Mater. Interfaces, 2024) on Li-Si-P-S SSE.

The structures are from early stages of the DP-Gen process, so they might not be the most accurate. They serve for testing purposes only.

6 digits naming convention: 404020 means 40 Li2S - 40 SiS2 - 20 P2S5

Project Structure

loweffort_LiSiPS_amorph_gen/
├── data/
│   ├── structures/
│   │   └── train/
│   │       └── *.poscar          # POSCAR files for training (~2828 structures)
│   ├── SmolSSE/                  # Raw data salvaged from DP-Gen
│   └── preprocess.py             # Script to convert .data → .poscar
├── graphite/                     # Copied file from graphite  repo
│   ├── __init__.py
│   ├── diffusion.py              # VarianceExplodingDiffuser
│   └── graph.py                  # periodic_radius_graph, pbc utilities
├── data_module.py                # PyTorch Lightning DataModule
├── models.py                     # ScoreModel, LitScoreNet
├── gen.ipynb          # Main notebook for training & generation
├── lit_logs/                     # Training logs and checkpoints
│   ├── amorph-gen/               # WandB run checkpoints
│   └── wandb/                    # WandB logs
└── tmp/                          # Generated structures output
    ├── LiSiPS_gen_*.vasp         # Generated POSCAR files
    └── denoise_traj.extxyz       # Denoising trajectory

Usage

open main.ipynb

Results

Testing generating amorphous Li3PS4 (Li:30, Si:0, P:10, S:40) with 1024 denoiser steps:

The generated structure looks off, but not bad as a starting point.

References

• graphite LLNL/graphite
• DP-Gen