MESS: Modern Electronic Structure Simulations

Hatem Helal,Andrew Fitzgibbon

2024-06-05

Abstract:Electronic structure simulation (ESS) has been used for decades to provide quantitative scientific insights on an atomistic scale, enabling advances in chemistry, biology, and materials science, among other disciplines. Following standard practice in scientific computing, the software packages driving these studies have been implemented in compiled languages such as FORTRAN and C. However, the recent introduction of machine learning (ML) into these domains has meant that ML models must be coded in these languages, or that complex software bridges have to be built between ML models in Python and these large compiled software systems. This is in contrast with recent progress in modern ML frameworks which aim to optimise both ease of use and high performance by harnessing hardware acceleration of tensor programs defined in Python. We introduce MESS: a modern electronic structure simulation package implemented in JAX; porting the ESS code to the ML world. We outline the costs and benefits of following the software development practices used in ML for this important scientific workload. MESS shows significant speedups n widely available hardware accelerators and simultaneously opens a clear pathway towards combining ESS with ML. MESS is available at <a class="link-external link-https" href="https://github.com/graphcore-research/mess" rel="external noopener nofollow">this https URL</a>.

Machine Learning,Materials Science,Computational Physics

What problem does this paper attempt to address?

This paper introduces the MESS (Modern Electronic Structure Simulation) project, which is a modern electronic structure simulation package based on the JAX framework. It aims to solve the problem of integrating traditional electronic structure simulation software with machine learning models. Traditional electronic structure simulation software is usually written in compiled languages such as FORTRAN and C, while machine learning models are primarily implemented in Python, which requires establishing complex software bridges. MESS achieves seamless integration with machine learning frameworks by porting the simulation code to JAX and utilizing hardware acceleration and batch processing to improve efficiency.

The main contributions mentioned in the paper include:

1. Implementation of end-to-end differentiable electronic structure simulation, supporting algorithmic gradient evaluation, for routine tasks such as electronic energy minimization and atomic force computation.
2. Development of a high-throughput DFT (Density Functional Theory) simulation framework, effectively exploring the configuration space using hardware acceleration and batch processing.
3. Creation of a scalable and modular simulation environment, facilitating the integration of machine learning models into electronic structure calculations.

The paper also discusses related work, such as the application of DFT in solving electronic structure problems, and the role of machine learning in electronic structure simulation, such as the use of Machine Learning Interatomic Potentials (MLIPs). Additionally, the paper introduces the basic concepts of electronic structure theory and provides an overview of MESS's software architecture, emphasizing the importance of hardware acceleration, automatic differentiation, and high-level interpretive languages like Python.

Through MESS, researchers can more easily combine electronic structure methods with machine learning models, accelerating interdisciplinary research and lowering the learning curve, enabling more scientists to perform such simulations. The paper also provides several examples demonstrating how to utilize MESS for direct density functional optimization and batch processing electronic structure optimization to improve computational efficiency.