Grad DFT: a software library for machine learning enhanced density functional theory

Abstract

Density functional theory (DFT) stands as a cornerstone method in computational quantum chemistry and materials science due to its remarkable versatility and scalability. Yet, it suffers from limitations in accuracy, particularly when dealing with strongly correlated systems. To address these shortcomings, recent work has begun to explore how machine learning can expand the capabilities of DFT; an endeavor with many open questions and technical challenges. In this work, we present Grad DFT: a fully differentiable JAX-based DFT library, enabling quick prototyping and experimentation with machine learning-enhanced exchange-correlation energy functionals. Grad DFT employs a pioneering parametrization of exchange-correlation functionals constructed using a weighted sum of energy densities, where the weights are determined using neural networks. Moreover, Grad DFT encompasses a comprehensive suite of auxiliary functions, notably featuring a just-in-time compilable and fully differentiable self-consistent iterative procedure. To support training and benchmarking efforts, we additionally compile a curated dataset of experimental dissociation energies of dimers, half of which contain transition metal atoms characterized by strong electronic correlations. The software library is tested against experimental results to study the generalization capabilities of a neural functional across potential energy surfaces and atomic species, as well as the effect of training data noise on the resulting model accuracy.

arXiv comment22 pages, 10 figures. The following article has been submitted to the Journal of Chemical Physics. After it is published, it will be found at https://publishing.aip.org/resources/librarians/products/journals/

References in corpus (32)

• Adam: A Method for Stochastic Optimization
• Nonlocal van der Waals density functional: The simpler the better
• Recent developments in the PySCF program package
• Power of data in quantum machine learning
• By-passing the Kohn-Sham equations with machine learning
• Finding Density Functionals with Machine Learning
• Libxc: a library of exchange and correlation functionals for density functional theory
• Self-interaction errors in density functional calculations of electronic transport
• Computational Complexity of interacting electrons and fundamental limitations of Density Functional Theory
• The discontinuous nature of the exchange-correlation functional -- critical for strongly correlated systems
• Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics
• Deep Learning and Density Functional Theory
• Machine learning and density functional theory
• Machine Learning the Physical Non-Local Exchange-Correlation Functional of Density-Functional Theory
• Learning the exchange-correlation functional from nature with fully differentiable density functional theory
• DeePKS: a comprehensive data-driven approach towards chemically accurate density functional theory
• Automatic differentiation in quantum chemistry with an application to fully variational Hartree-Fock
• Design and Analysis of Machine Learning Exchange-Correlation Functionals via Rotationally Invariant Convolutional Descriptors
• Predictive Power of the Exact Constraints and Appropriate Norms in Density Functional Theory
• Neural-network Kohn-Sham exchange-correlation potential and its out-of-training transferability
• Kinetic energy densities based on the fourth order gradient expansion: performance in different classes of materials and improvement via machine learning
• Accelerating Finite-temperature Kohn-Sham Density Functional Theory with Deep Neural Networks
• Complete basis set extrapolation of electronic correlation energies using the Riemann zeta function
• Evolving symbolic density functionals
• Differentiable quantum chemistry with PySCF for molecules and materials at the mean-field level and beyond
• Artificial neural networks for density-functional optimizations in fermionic systems
• Range-separated double-hybrid density-functional theory with coupled-cluster and random-phase approximations
• Kohn-Sham theory with paramagnetic currents: compatibility and functional differentiability
• Efficient Learning of a One-dimensional Density Functional Theory
• What the foundations of quantum computer science teach us about chemistry
• How Well Does Kohn-Sham Regularizer Work for Weakly Correlated Systems?
• D4FT: A Deep Learning Approach to Kohn-Sham Density Functional Theory

Cited by in corpus (7)

• Performant Automatic Differentiation of Local Coupled Cluster Theories: Response Properties and Ab Initio Molecular Dynamics
• MOLPIPx: an end-to-end differentiable package for permutationally invariant polynomials in Python and Rust
• Learning Density Functionals from Noisy Quantum Data
• Leveraging Normalizing Flows for Orbital-Free Density Functional Theory
• Algorithmic differentiation for plane-wave DFT: materials design, error control and learning model parameters
• Neural network distillation of orbital dependent density functional theory
• Quantum-Enhanced Neural Exchange-Correlation Functionals