JAX-XC: Exchange Correlation Functionals Library in Jax

Kunhao Zheng, Min Lin

Keywords: Software Library, Density Functional Theory, Exchange-Correlation Functional

TL;DR: We present JAX-XC, an open-source library that provides exchange-correlation functionals in Jax.

Abstract: We present JAX-XC, an open-source library that provides exchange-correlation functionals in Jax. JAX-XC is built from LIBXC, its correctness has been verified numerically against LIBXC. Thanks to Jax, JAX-XC is end-to-end differentiable, computationally more efficient thanks to the vectorization provided by XLA, and also portable on various accelerators. More importantly, as more research is focusing on machine learning for density functional theory, we hope that JAX-XC could serve as a deep learning-friendly tool and a stepping-stone for researchers working in the intersection of deep learning and density functional theory.