Jrystal

A JAX-based Differentiable Density Functional Theory Framework for Materials

Reference docs

Core Features

Differentiable: Leveraging JAX's automatic differentiation for efficient gradient computation of quantum properties, enabling straightforward optimization workflows.

GPU-Accelerated: Optimized for modern GPU architectures, delivering high-performance quantum calculations with automatic hardware acceleration.

Solid-State Calculation: Full-featured framework for periodic systems using plane wave basis sets, supporting precise electronic structure calculations of crystalline materials.

Direct Optimization: A direct minimization approach that avoids SCF iterations, enabling smooth convergence and natural integration of machine learning methods, and advanced quantum chemistry methods into density functional theory calculations.

Installation

You can install Jrystal directly from GitHub:

pip install [email protected]:sail-sg/jrystal.git

Development Installation

To install in development mode:

git clone [email protected]:sail-sg/jrystal.git
cd jrystal
pip install -e .

This allows you to modify the source code and have the changes take effect without rebuilding and reinstalling.

To build and serve the documentation locally:

make doc-dev

Then visit localhost:8000.

Usage

Run calculations using the command line interface:

jrystal -m energy -c config.yaml

The following modes are supported:

• energy: compute the ground state energy of a system
• band: compute the band structure of a system

You can customize the calculation by modifying the config.yaml file.

Band Structure Benchmarks

Jrystal has been benchmarked against Quantum ESPRESSO for various materials. Below are the comparison results for both all-electron and norm-conserving pseudopotential calculations.

All-electron Calculations

Silicon (Si)

Aluminum (Al)

Graphene

Sodium (Na)

Norm-conserving Pseudopotential Calculations

Silicon (Si)

Aluminum (Al)

Development Status

Note: The pseudopotential model is currently experimental.

The Team

This project is jointly developed by SEA AI LAB (SAIL) and NUS I-FIM.

Citation

If you find this project useful, please cite:

@inproceedings{ml4ps2024,
    title = {Jrystal: A JAX-based Differentiable Density Functional Theory Framework for Materials},
    author = {Li, Tianbo and Shi, Zekun and Dale, Stephen Gregory and Vignale, Giovanni and Lin, Min},
    booktitle = {Machine Learning and the Physical Sciences Workshop at NeurIPS 2024},
    year = {2024},
}

@article{li2024diagonalization,
    title={Diagonalization without Diagonalization: A Direct Optimization Approach for Solid-State Density Functional Theory},
    author={Li, Tianbo and Lin, Min and Dale, Stephen and Shi, Zekun and Neto, AH Castro and Novoselov, Kostya S and Vignale, Giovanni},
    journal={arXiv preprint arXiv:2411.05033},
    year={2024}
}

License

This project is licensed under the Apache License 2.0.

Documentation

For detailed documentation, tutorials, and API reference, please visit our documentation page.