MESS: Modern Electronic Structure Simulations
This work addresses the problem of software incompatibility between traditional ESS tools and modern ML frameworks for researchers in chemistry, biology, and materials science, representing an incremental improvement by adapting existing practices to new contexts.
The authors tackled the challenge of integrating machine learning (ML) with electronic structure simulation (ESS) by developing MESS, a package implemented in JAX that ports ESS to the ML world, resulting in significant speedups on hardware accelerators and enabling easier combination of ESS with ML.
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.