Materials Learning Algorithms (MALA): Scalable Machine Learning for Electronic Structure Calculations in Large-Scale Atomistic Simulations

TL;DR

MALA presents a scalable, LDOS-centered machine learning framework to accelerate electronic structure calculations at scales well beyond conventional DFT. By learning the local density of states from grid-based descriptors (e.g., bispectrum and ACE) and integrating data sampling, training, and inference within a unified library, MALA delivers LDOS, DOS, density, and total energy with near-linear scaling in system size. The approach is demonstrated across boron, aluminum, and beryllium systems, including large Be slabs with stacking faults and solid–liquid phase transitions, with extensive benchmarking for accuracy, transferability, and scalability on multi-GPU hardware. The framework’s modularity and compatibility with Quantum ESPRESSO and LAMMPS position it as a versatile tool for large-scale materials modeling, offering pathways toward LDOS-enabled MD, band-structure calculations, and STM-like observables in complex materials.

Abstract

We present the Materials Learning Algorithms (MALA) package, a scalable machine learning framework designed to accelerate density functional theory (DFT) calculations suitable for large-scale atomistic simulations. Using local descriptors of the atomic environment, MALA models efficiently predict key electronic observables, including local density of states, electronic density, density of states, and total energy. The package integrates data sampling, model training and scalable inference into a unified library, while ensuring compatibility with standard DFT and molecular dynamics codes. We demonstrate MALA's capabilities with examples including boron clusters, aluminum across its solid-liquid phase boundary, and predicting the electronic structure of a stacking fault in a large beryllium slab. Scaling analyses reveal MALA's computational efficiency and identify bottlenecks for future optimization. With its ability to model electronic structures at scales far beyond standard DFT, MALA is well suited for modeling complex material systems, making it a versatile tool for advanced materials research.

Figures (25)

• Figure 1: Overview of the general workflow of a feed-forward neural network used in the MALA package. A network consists of $N_l$ layers each comprising $w_l$ neurons. Neurons are interconnected by weights $\bm{w}$. In the network, input data $\bm{x}$ is propagated to individual neurons, where it is multiplied by weights, summed, and added to a bias $b$. This value is then passed through a non-linear activation function $\sigma$, producing the neuron’s output. The output is subsequently passed to the next layer via different weights, and this process repeats until the final layer is reached, yielding the neural network’s output.
• Figure 2: Overview of MALA workflow. Both model training procedure and inference are shown. For the latter, relevant routines are bundled in a separate class within the MALA package to facilitate access. Gray boxes denote pure Python-based routines, while blue, red and orange boxes denote that the majority of the computational workload is offloaded to the external libraries LAMMPSLAMMPSLAMMPS2, PyTorchpytorch and Quantum ESPRESSOgiannozzi_quantum_2009giannozzi_advanced_2017giannozzi_q_2020, respectively. Data generation is typically performed with the VASPkresse_ab_1993kresse_efficiency_1996kresse_efficient_1996 and Quantum ESPRESSO simulation codes, in conjunction with a MALA specific DFT-MD trajectory analysis. MALA internal routines are implemented with the aid of NumPy Numpy and SciPy Scipy.
• Figure 3: Calculation of bispectrum descriptors, representing the local ionic configuration within the neighborhood surrounding a grid point in real space. Figure adapted from Ref. Lenz_Dissertation.
• Figure 4: Dependence of band energy error on Gaussian width (expressed in units of energy grid spacing $\delta\epsilon$) for aluminum (256 atoms) and beryllium (128 atoms) at room temperature. Band energy errors are averaged over three ionic configurations per element. A range of optimal width values is evident, beyond which errors increase noticeably. Previous publications malapaperhyperparameterpapertemperaturepapersizetransferpaperLenz_Dissertation have identified $w_\mathrm{d}=2\delta\epsilon$ as a balanced choice for both aluminum and beryllium.
• Figure 5: Comparison of different hyperparameter optimization methods for tuning MALA models. The MAE of total free energy prediction error across the test set is shown relative to the computational cost of determining model hyperparameters. An aluminum system at room temperature, containing 10 configurations with 256 atoms each, was modeled. For each set of hyperparameters, five neural networks were trained. Error bars denote the standard deviation across these five network initializations, while solid symbols indicate the average MAE. Results are presented for direct search, the tree-structured Parzen estimator (TPE), the orthogonal array tuning method (OATM), and neural architecture search without training (NASWOT). Also shown are two combinations of TPE and NASWOT: using NASWOT to prune unpromising trials during TPE optimization (NASWOT-based pruner) and using TPE to optimize training hyperparameters for which NASWOT provides no insights. Details on hyperparameter optimization methods are available in Ref. hyperparameterpaper.