Jrystal#

A JAX-based Differentiable Density Functional Theory Framework for Materials.

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.

Call Graph#

This project is a JAX-based package for differantiable density functional theory computation of solids.

Installation#

You may install directly from GitHub, using the following command. This can be used to obtain the most recent version of Optax:

pip install git+https://github.com/sail-sg/jrystal.git

If you want to install the package in development mode, you can use the following command. This allows you to modify the source code and have the changes take effect without you having to rebuild and reinstall.

git clone git@github.com:sail-sg/jrystal.git
cd jrystal
pip install -e .

For more installation instructions, please refer to the installation page.

Run with Command Line#

Once jrystal is installed, you can run it with the following command.

jrystal -m energy -c config.yaml

The -m or --mode option is used to specify the mode to run. Currently, we support two modes: energy and band.

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

The -c or --config option is used to specify the configuration file. If it is not provided, the program will look for the file named config.yaml in the current directory. You may modify the config.yaml file to customize the calculation.

Band Structure Benchmark#

We provide a benchmark for the band structure calculation of bulk silicon, graphene, aluminum, and sodium with Quantum ESPRESSO for both all-electron and norm-conserving pseudopotentials.

All-electron calculation#

Silicon (Si)#	Aluminum (Al)#
Graphene#	Sodium (Na)#

Norm-conserving Psuedopotential calculation#

The Team#

This project is developed by SEA AI LAB (SAIL). We are also grateful to researchers from NUS I-FIM for contributing ideas and theoretical support.

Citation#

If you find this project useful, please cite our preprints:

@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},
}

and

@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.

Jrystal

Contents