Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MESS: Modern Electronic Structure Simulations

Hatem Helal, Andrew Fitzgibbon

TL;DR

The paper addresses the barrier of integrating machine learning with electronic structure simulations, which are traditionally implemented in compiled languages, by introducing MESS. MESS is a JAX-based framework that enables end-to-end differentiable electronic structure simulations, hardware-accelerated DFT, and seamless ML-ESS integration. Key contributions include differentiable energy minimisation, a high-throughput batched DFT framework, and a modular open-source environment under the MIT license, with preliminary benchmarks showing about 16x speedups and enabling batch exploration of configurational space. This work provides a practical pathway to accelerate discovery in chemistry and materials by unifying ESS with modern ML tooling and hardware accelerators, while promoting further research into scalable and differentiable electronic structure representations.

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 https://github.com/graphcore-research/mess.

MESS: Modern Electronic Structure Simulations

TL;DR

The paper addresses the barrier of integrating machine learning with electronic structure simulations, which are traditionally implemented in compiled languages, by introducing MESS. MESS is a JAX-based framework that enables end-to-end differentiable electronic structure simulations, hardware-accelerated DFT, and seamless ML-ESS integration. Key contributions include differentiable energy minimisation, a high-throughput batched DFT framework, and a modular open-source environment under the MIT license, with preliminary benchmarks showing about 16x speedups and enabling batch exploration of configurational space. This work provides a practical pathway to accelerate discovery in chemistry and materials by unifying ESS with modern ML tooling and hardware accelerators, while promoting further research into scalable and differentiable electronic structure representations.

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 https://github.com/graphcore-research/mess.
Paper Structure (8 sections, 5 equations, 2 figures)

This paper contains 8 sections, 5 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Electronic Structure Theory in a Nutshell
  4. Software Architecture
  5. Motivating Examples
  6. Direct Minimisation of Density Functional
  7. Batching over Conformations
  8. Concluding Thoughts

Figures (2)

  • Figure 1: Total energy convergence for methane molecule with 6-31G basis set and PBE functional by direct optimisation using the Adam optimiser.
  • Figure 2: Batched evaluation of DFT energy for $H_2$ molecule within LDA approximation and the minimal STO-3G basis set.