Efficient implementation of atom-density representations

TL;DR

This paper introduces librascal, a modular and optimized implementation for atom-density representations like SOAP, enabling faster computations and easier development of rotationally equivariant atomistic ML models.

Contribution

The paper presents librascal, a flexible, efficient software framework for atom-density representations, with optimized density expansion and data reduction techniques for improved ML model performance.

Findings

Librascal enables faster computation of atom-density features.

Data reduction in feature space reduces computational cost by up to 5x.

Optimized density expansion improves efficiency without loss of accuracy.

Abstract

Physically-motivated and mathematically robust atom-centred representations of molecular structures are key to the success of modern atomistic machine learning (ML) methods. They lie at the foundation of a wide range of methods to predict the properties of both materials and molecules as well as to explore and visualize the chemical compound and configuration space. Recently, it has become clear that many of the most effective representations share a fundamental formal connection: that they can all be expressed as a discretization of N-body correlation functions of the local atom density, suggesting the opportunity of standardizing and, more importantly, optimizing the calculation of such representations. We present an implementation, named librascal, whose modular design lends itself both to developing refinements to the density-based formalism and to rapid prototyping for new developments of rotationally equivariant atomistic representations. As an example, we discuss SOAP features, perhaps the most widely used member of this family of representations, to show how the expansion of the local density can be optimized for any choice of radial basis set. We discuss the representation in the context of a kernel ridge regression model, commonly used with SOAP features, and analyze how the computational effort scales for each of the individual steps of the calculation. By applying data reduction techniques in feature space, we show how to further reduce the total computational cost by at up to a factor of 4 or 5 without affecting the model's symmetry properties and without significantly impacting its accuracy.