MESS: Modern Electronic Structure Simulations

TL;DR

MESS is a modern electronic structure simulation package built in JAX that leverages hardware acceleration, offering significant speedups and facilitating integration of electronic structure simulations with machine learning techniques.

Contribution

The paper introduces MESS, a new electronic structure simulation software implemented in JAX, bridging traditional ESS with modern ML frameworks for improved performance and usability.

Findings

MESS achieves significant speedups on hardware accelerators.

The software enables easier integration of ESS with machine learning.

MESS demonstrates the benefits of adopting ML development practices in scientific computing.

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.