Jrystal: A JAX-based Differentiable Density Functional Theory Framework for Materials

Density functional theory (DFT) is crucial for studying materials at the atomic level, known for its accurate predictions and computational efficiency. However, integrating AI techniques into current DFT packages is challenging due to the iterative nature of the self-consistent field methods. To address these challenges, we developed Jrystal, a JAX-based differentiable DFT framework that integrates modern deep learning frameworks with automatic differentiation and adopts a direct optimization approach using plane wave bases. This makes Jrystal a powerful tool for designing new DFT algorithms with machine learning techniques.