Abstract
The calculation of electron density distribution using density functional theory (DFT) in materials and molecules is central to the study of their quantum and macroscale properties, yet accurate and efficient calculation remains a longstanding challenge. We introduce ChargE3Net, an E(3)equivariant graph neural network for predicting electron density in atomic systems. ChargE3Net enables the learning of higherorder equivariant features to achieve high predictive accuracy and model expressivity. We show that ChargE3Net exceeds the performance of prior work on diverse sets of molecules and materials. When trained on the massive dataset of over 100K materials in the Materials Project database, our model is able to capture the complexity and variability in the data, leading to a significant 26.7% reduction in selfconsistent iterations when used to initialize DFT calculations on unseen materials. Furthermore, we show that nonselfconsistent DFT calculations using our predicted charge densities yield nearDFT performance on electronic and thermodynamic property prediction at a fraction of the computational cost. Further analysis attributes the greater predictive accuracy to improved modeling of systems with high angular variations. These results illuminate a pathway towards a machine learningaccelerated ab initio calculations for materials discovery.
Similar content being viewed by others
Introduction
Quantum mechanical interactions of electrons in atoms provide a foundation for understanding and predicting properties of materials^{1}. Modern density functional theory (DFT), particularly the KohnSham (KS) formulation of DFT^{2}, is by far the most widely used method for performing such electronic structure calculations due to its balanced tradeoff between accuracy and computational cost. Because KSDFT obtains the electronic structure by iteratively solving KohnSham equations and recalculating the electron charge density until convergence to a selfconsistent field (SCF), extensive use of DFT is hindered by the unfavorable O(N^{3}) scaling with the system size, making it infeasible for largescale calculations beyond a few hundred atoms.
While prior works on linearscaling DFT^{3,4,5}, local coupled cluster methods^{6,7}, and orbitalfree formulation of DFT^{8} have nominally reduced the scaling limit, largescale calculation is still unobtainable for general problems. In recent years, many efforts have focused on the potential of machine learning (ML) approaches to overcome the quantum scaling limit of DFT. ML models have been developed to map chemical and material systems to properties to enable highthroughput screening^{9,10,11,12,13,14}. Once trained, ML models are computationally fast, but they are restricted to properties that already exist in a database (typically at the ground state), limiting their applicability to general problems. ML models have also been proven useful for predicting atomic energies and forces for largescale molecular dynamic simulations^{15,16,17,18,19}. The process generally requires a dedicated computational campaign with DFT to generate highquality data for a subset of materials with a predefined potential or temperature range^{18,19,20,21,22,23,24}. More recent work has trained a foundation model across the periodic table for predicting interatomic potentials^{25}, but it still requires additional DFT data and model finetuning for applications beyond structure relaxation. An alternative approach is to directly predict converged DFT quantities (e.g., electron charge density) to reduce the number of SCF iterations needed to solve KS equations. Since the number of SCF iterations is dependent on the accuracy of the initialized charge density, a fast and accurate charge density prediction may be able to bypass the need for selfconsistent iterations, enabling accurate computation of downstream properties at a fraction of the computational cost.
Several works have used ML for charge density prediction. One approach is to expand the charge density with a local atomic basis set and use ML to predict the expansion coefficients^{26,27,28,29,30,31,32}. Local atomic orbital basis sets are often parametrized for specific elements. While these models are, in principle, more efficient, scaling on the order of the number of basis functions, they are restricted to the basis set they are trained on, thus impacting generalizability. Furthermore, any model that attempts to find the parameters of atomic basis sets that best fit gridbased data is fundamentally limited by the expressivity of the basis set^{32}. A parallel line of research focuses on the prediction of the Hamiltonian operator^{33,34,35,36}, however these are also restricted by the choice of basis set, and often not obtainable through DFT codes.
An alternative approach is to learn the electron charge density directly from a discretized grid of density points in 3D Cartesian space. The gridbased approach is agnostic to the basis set (planewave or atomic basis) used for the quantum chemistry calculation and are a natural input format for DFT codes^{37,38,39,40}. Existing gridbased methods^{41,42,43,44,45} have shown promising results on small molecules and small subsets of material classes. Pioneering works^{41,42} have focused on learning the charge density using graph neural networks invariant to translation, rotation and permutation^{9,46}. Jørgensen and Bhowmik^{42} later demonstrate an improvement in accuracy using an equivariant graph neural network, PaiNN^{43}, through \({{\mathbb{R}}}^{3}\) vector representations and rotationally equivariant operations^{47}. While these models have been shown to achieve high accuracy on small, specialized molecular or material datasets, no work has demonstrated applicability on larger, diverse datasets across the periodic table. In addition, these models use scalar and/or vector representations to construct rotationally equivariant networks, which could be further improved by incorporating higherorder equivariant features, as has been shown in other tasks, such as the prediction of atomic forces^{16}. Finally, limited work has demonstrated the utility of MLpredicted charge densities to initialize DFT calculations, with existing work showing only marginal reduction of SCF steps^{44,48}.
In this work, we develop a gridbased model, ChargE3Net, that uniquely combines graph neural networks, SO(3) equivariant features, and higherorder tensor operations, to accurately predict the electron charge density. Our model only uses atomic species and positions as input without the need for additional features. It enables the use of higherorder equivariant features to achieve more accurate performance and higher model expressivity. In a headtohead comparison with prior gridbased and basis functionbased methods, we show that our method not only achieves greater accuracy on existing benchmark data, but also is expressive enough to capture the complexity and variability present in Materials Project (MP) data^{49}, consisting of more than 100K inorganic materials across the periodic table and for all materials classes. We demonstrate the application of MPtrained ChargE3Net to initialize DFT calculations, resulting in a median of 26.7% reduction in SCF steps on unseen MP data and 28.6% on novel GNoME materials^{50}. Both have surpassed the reductions demonstrated by prior work^{44,48}. When evaluating the utility of ChargE3Net in computing a wide range of electronic and thermodynamic properties nonselfconsistently, we show a great portion of materials achieving nearDFT accuracy. Further analysis of materials with less accurate predictions provides insights into a potential cause of the observed inconsistency between selfconsistent and nonselfconsistent calculations to inform future research. By enabling higherorder equivariant models, we can also investigate the impact of equivariant representations in model performance. We show that higherorder representations yield more accurate charge density predictions for materials with high angular variation. Lastly, We demonstrate the linear time complexity of our model with respect to system size, showing the capability of predicting density on systems with > 10^{4} atoms, surpassing what is feasible with ab initio calculations.
Results
Development of higherorder equivariant neural networks for charge density prediction
The charge density, along with other properties such as forces, is equivariant under rotation and translation of the atomic system. That is, a rotation or translation of the atomic system in Euclidean space will result in an equivalent rotation or translation to the charge density or forces. Formally, this is known as equivariance with respect to E(3), which includes rotations, translations, and reflections in 3D space. This can be achieved in principle by using only invariant scalar features^{46}, such as interatomic distances; however, this prevents the model from using angular information, limiting the accuracy that can be attained^{51}. More recent work^{47} has introduced equivariance methods using vector \({{\mathbb{R}}}^{3}\) features, such as relative atomic positions, to incorporate angular information, and has shown to improve the performance of charge density prediction^{43}. We hypothesize that higherorder representations could help improve model performance and accelerate DFT calculations.
ChargE3Net achieves equivariance through means outlined by tensor field networks^{52}. Translation equivariance is achieved by using relative atomic positions; rotation equivariance is achieved by restricting features to be irreducible representations, irreps, of SO(3), which are operated upon by equivariant functions. These features take the form \({V}_{cm}^{(\ell ,p)}\), a dictionary of tensors with keys representing rotation order ℓ ∈ {0, 1, 2, . . . } and parity p ∈ { − 1, 1}. Each tensor has a channel index c ∈ [0, N_{channels}) and an index m ∈ [ − ℓ, ℓ]. In this way, the representation at a given rotation order ℓ and parity p would have a size of \({{\mathbb{R}}}^{{N}_{{{{\rm{channels}}}}}\times (2\ell +1)}\). These representations are combined with the equivariant tensor product operation ⊗ , using ClebschGordan coefficients C^{52} as implemented in e3nn^{53}:
where ℓ_{o} and p_{o} are given by ∣ℓ_{1} − ℓ_{2}∣≤ℓ_{o}≤∣ℓ_{1} + ℓ_{2}∣ and p_{o} = p_{1}p_{2}. We maintain only those representations with ℓ_{o}≤L where L is a maximum allowed rotation order. For ease of notation, we will omit the keys ℓ, p and indices c, m for these tensor representations for the rest of the paper. Unless otherwise indicated, we use an L = 4 maximum rotation order for ChargE3Net.
Figure 1 shows the model architecture of ChargE3Net. To predict charge density of a given atomic system, a set of M probe points \(\{{\vec{{{{\bf{g}}}}}}_{1},...,{\vec{{{{\bf{g}}}}}}_{M}\}\in {{\mathbb{R}}}^{3}\) are introduced to locations where charge densities are to be predicted. We represent predicted charge density \(\hat{\rho }(\vec{{{{\bf{g}}}}})\) as an equivariant graph neural network with inputs of atomic numbers \(\{{z}_{1},...,{z}_{N}\}\in {\mathbb{N}}\) with respective locations of the atoms \(\{{\vec{{{{\bf{r}}}}}}_{1},...,{\vec{{{{\bf{r}}}}}}_{N}\}\in {{\mathbb{R}}}^{3}\) where N denotes the total number of atoms, as well as probe locations \(\{{\vec{{{{\bf{g}}}}}}_{1},...,{\vec{{{{\bf{g}}}}}}_{M}\}\in {{\mathbb{R}}}^{3}\) and an optional periodic boundary cell \({{{\bf{B}}}}\in {{\mathbb{R}}}^{3\times 3}\) for periodic systems.
As illustrated in Fig. 1b, c, a graph is first constructed with atoms and probe points as vertices, with edges formed via proximity with a cutoff of 4 Å. Edges connecting atoms are unidirectional and edges connecting atoms and probe points are directed pointing from atoms to probe points. Vertex features for atoms \({\{{{{{\bf{A}}}}}_{i}\}}_{i = 0}^{N}\) are initialized as a onehot encoding of the atomic number, represented as an ℓ = 0, p = 1 tensor with N_{channels} equal to the number of atomic species. Features for probe points \({\{{{{{\bf{P}}}}}_{k}\}}_{k = 0}^{M}\) are initialized as a single scalar zero (i.e., ℓ = 0, p = 1 and N_{channels} = 1).
ChargE3Net uses a graph neural network (GNN) to update these tensor features (Fig. 1d). Within each layer of GNN, the network performs message passing between neighboring atoms and probes, where probe points can only receive messages from neighboring atoms (Fig. 1e). The core of the message passing in ChargE3Net is a convolution function that updates the vertex tensor feature with neighboring vertices by the convolution operation (Eq. (4) in “Methods” section). Since the proposed convolution operation replaces the linear message passing with equivariant operation through the tensor product defined in Eq. (1), the operation is rotation and translation equivariant; see “Methods” section for more detail.
ChargE3Net conducts a series of equivariant convolution and gate operations to obtain representations for each atom, with an addition series of equivariant convolution and gate operations to obtain representations for each probe point. The learned tensor feature associated to each probe point is passed to a postprocessing layer where a simple regression is used to predict the charge density at the corresponding probe location \(\hat{\rho }({\vec{{{{\bf{g}}}}}}_{k})\).
Performance evaluation
We train and validate ChargE3Net on a diverse set of DFTcomputed charge density data, including organic molecules (QM9)^{54,55,56}, nickel manganese cobalt battery cathode materials (NMC)^{57}, and inorganic materials collected from Materials Project (MP)^{49,58}. The first two datasets have been used in the past for benchmarking charge density models for small molecules and a subset of crystal systems with boundary conditions^{30,43}. To develop a charge density model for the entire periodic table, we investigate one of the largest open databases of DFT, the Materials Project (MP) database. It includes charge densities for over 100K bulk materials with diverse compositions and structures (Supplementary Figs. 1–3 and “Methods” section).
We assess the performance our models using mean absolute error normalized by the total number of electrons within a volume of the atomic system^{27,28,30,43}, ε_{mae}, which is computed via numerical integration on the charge density grid points G:
Unless otherwise noted, the probe points G hereby reference the full set of discretized unitcell grid points for which we have DFTcomputed charge density values.
Table 1 compares the performance of ChargE3Net to models introduced in OrbNetEqui^{30} and DeepDFT^{43}. For QM9 and NMC datasets, we use identical training, validation, and test splits as Jørgensen and Bhowmik^{43}. All reported ε_{mae} are computed using publicly released models and verified to match the numbers reported in the original work. For the MP dataset, we train the DeepDFT models using the same experimental settings from the original work (see Supplementary Note 1 for more detail). We observe that all equivariant models perform better than invariant models, demonstrating the importance of respecting rotation equivariance. Moreover, ChargE3Net model significantly outperforms prior equivariant DeepDFT models^{43} on the Materials Project (MP) and QM9 datasets, while achieving similar performance on the NMC dataset. This suggests that the inclusion of higherorder representations can improve model performance. More remarkably, our model achieves a lower ε_{mae} on the QM9 dataset than OrbNetEqui^{30} model, which leverages additional features based on quantum mechanical calculations while ChargE3Net only uses atomic species and position information. This shows how the model is capable of learning highly informative representations of quantum interactions in a purely datadriven approach.
ChargE3Net’s performance on the MP data demonstrates that it is capable of successfully modeling charge density across the periodic table and a wide variety of material classes that are not present in other datasets. When we compare the prediction performance with respect to the number of unique species in a material, the MPtrained ChargE3Net is capable of predicting charge density in systems with many unique species, without a significant change in error distribution as number of species increases (Supplementary Fig. 4). This suggests ChargE3Net can reproduce a broad range of chemical interactions at the electronic level without incurring the cost of an ab initio calculation.
Using predicted charge densities to accelerate DFT
DFT calculations require an initial charge density and wavefunction prior to iteratively solving the KohnSham equations until selfconsistency^{59}. Charge density initializations that closely match the groundstate charge density often save compute time by reducing the number of selfconsistent field (SCF) steps needed to find a selfconsistent solution. Many schemes for initializing the charge density exist in the literature, but the most common one is the superposition of atomic densities (SAD), where the initial charge density is computed by summing over the atomic charge density of each atom in the system^{60}.
We demonstrate that charge densities predicted by ChargE3Net reduce the number of SCF steps, thereby accelerating DFT calculations, on unseen nonmagnetic materials (Fig. 2a, b) compared to a SAD initialization. The median number of SCF steps needed to converge a singlepoint calculation selfconsistently for a nonmagnetic material starting from the SAD initialization is 15, while the same starting from a ChargE3Net initialization is 11, representing a 26.7% decrease in SCF steps. As an upperbound baseline, we also include the distribution of the number of SCF steps needed to converge a calculation starting from a converged (i.e. groundtruth) charge density (SC), showing that our model still has room for improvement. About 96% of calculations done with a predicted charge density ran with fewer SCF steps than initializing with SAD, and the compute time improvement is robust to a wide range of ε_{mae} (Supplementary Fig. 5).
To evaluate how well ChargE3Net’s predicted charge densities reproduce physical properties, we compute the energy, forces, band energies, and bandgaps of the test set nonselfconsistently (Fig. 2c, d). 40% of materials in the test set have peratom energy errors less than 1 meV atom^{−1}. 70% of materials have forces <0.03 eV Å^{−1}, a common threshold for structural relaxations; 15% of the test set have force errors of 0.0 (Supplementary Fig. 6). All of the structures with perfect agreement have selfconsistent forces of 0.0 eV Å^{−1} for each atom and Cartesian component, possibly because of cancellation of charge density errors due to symmetry. 45% and 76% of materials recover band energies at each kpoint and the direct band gap within chemical accuracy (43 meV), respectively. 45% of the band gaps obtained nonselfconsistently exactly matched those of selfconsistent calculations (Supplementary Fig. 6). Similar to the force data described above, calculations with perfect agreement all have a selfconsistent band gap of 0 eV, indicating the model is able to readily identify metallic materials. We note that our use of chemical accuracy for band gaps is a more stringent criterion than experimental measurement errors for wellstudied semiconductors^{61}. In addition, nonselfconsistent calculations using the Materials Project parameters as done with this test set have an additional source of error (Supplementary Fig. 7 and “Methods” section) as a result of an incorrect calculation parameter set by the Materials Project. This error could explain why not all NSC calculations computed on a selfconsistent charge density achieve chemical accuracy (see Fig. 2d and Supplementary Fig. 8).
To examine the generalization capabilities of ChargE3Net on outofdistribution data, we analyze a random subset of 1924 materials published by GNoME^{50} with 5 or more unique species to differentiate it from the MP distribution. The average ε_{mae} found for this set is 0.632%, which is higher than that of the MP test set for ChargE3Net but lower than either DeepDFT variant trained on the Materials Project (Table 1). We conduct a similar experiment to the MP test set above with nonselfconsistent DFT calculations on 1510 nonmagnetic materials. The median energy error is 4.3 meV atom^{−1} for the GNoME set, whereas the same is 1.8 meV atom^{−1} for the MP test set. Only 8.4% of materials in GNoME computed nonselfconsistently from the ChargE3Net charge density obtained the total energy within chemical accuracy (<1 meV atom^{−1}), though nearly all selfconsistent calculations meet this criteria (Supplementary Fig. 9). The median number of SCF steps needed for selfconsistency falls from 14 to 10 steps when initializing with ChargE3Net, a 28.6% drop (Supplementary Fig. 10) and a similar acceleration found for the Materials Project test set. Moreover, 91% of calculations saw a decreased number of SCF steps to reach convergence. While further work is needed to improve the generalizability of ChargE3Net for nonselfconsistent property predictions, selfconsistent calculations on outofdistribution materials are still accelerated when initialized by the predicted charge density.
In addition to the model’s robustness to materials outside its training set, we show that ChargE3Net can predict the charge densities of materials after small atomic perturbations with minor increase to the charge density error (Supplementary Fig. 11). This allows us to accurately obtain phonon properties, which may rely on finite differences as the main method to obtain approximations to the local potential energy surface, from nonselfconsistent calculations. We compute the constantvolume heat capacity (C_{v}) under the harmonic approximation for two zeolites, which are not included in the Materials Project dataset, and show that C_{v}(T) obtained closely match those obtained by selfconsistent calculations (Supplementary Note 2). As another example of transferability, we compute the electronic bandstructure of the Si(111)2 × 1 surface nonselfconsistently from the ChargE3Net prediction and show that it almost exactly matches that of a selfconsistent calculation (Supplementary Note 3).
Higherorder representations lead to improved model accuracy
To demonstrate the impact of introducing higherorder equivariant tensor representations in our model, we examine the effect of training models while varying the maximum rotation order L, with ℓ ∈ {0, . . . , L} up to L = 4 used in the final model. We determine the number of channels for each rotation order by \({N}_{{{{\rm{channels}}}}}=\lfloor \frac{500/(L+1)}{\ell * 2+1}\rfloor\). For example, the L = 0 model has representations consisting of 500 even scalars, and 500 odd scalars, while the L = 1 model has representations consisting of 250 even scalars, 250 odd scalars, 83 even vectors, and 83 odd vectors. In this way, each model is constructed such that the total representation size and the proportion of the total representation used by each rotation order are approximately equal. We train each model on 1000 and 10,000 material subsets and the full MP dataset, using the same validation and test sets. We observe a consistent increase in performance on each training subset as the maximum rotation order increases, and a similar scaling behavior with respect to the training set size (Supplementary Fig. 12). This trend suggests that higherorder representations are necessary for accurate charge density modeling, and can match the performance of lower order models with less data.
To gain intuition behind why and when higherorder representations yield higher performance, we investigate two factors contributing to the variance of charge density within a material: (1) radial dependence, or a dependence on the distance from neighboring atoms, and (2) angular dependence, or dependence on angular orientation with respect to those atoms and the larger atomic system. While most materials exhibit strong radial dependence, some also exhibit strong angular dependence. For example, materials with covalent bonds share electrons between nuclei, leading to higher charge density along bonding axes and more angular dependence. In Fig. 3, we illustrate this concept with two materials. Cs(H_{2}PO_{4}) exhibits high angular dependence, where charge density is dependent on angle as well as radial distance from the nearest atom. The complex nature of the charge density shown here arises from covalent bonding among the H, P, and O atoms. Conversely, Rb_{2}Sn_{6} does not exhibit this, as its density appears to be almost entirely a function of the atomic distance, suggesting that a lowerorder equivariant or invariant architecture could model its electron density well. This may be indicative of ionic interactions rather than the covalent interactions found in Cs(H_{2}PO_{4}). We find this intuition to be consistent with model performance, as Fig. 3 shows similar performance for L = 0 and L = 4 models for Rb_{2}Sn_{6}, while the L = 4 model exceeds the performance of the L = 0 model for Cs(H_{2}PO_{4}).
At the atomic level, we find more evidence of the benefits of higherorder equivariant representations in modeling charge density. We examine the perelement errors by computing ε_{mae} on grid points G_{local,i} ⊂ G, a set of points that are within 2Å of an atom of species i. Fig. 4 shows this perelement error for both L = 0 (lowerleft triangles) and L = 4 (upperright triangles), averaged across the test set. While the L = 0 model struggles with reactive nonmetals and metalloids, which often form covalent bonds, L = 4 shows strong improvement among these elements. Nearly all elements show improvement, but the median change in ε_{mae} values for those materials containing a nonmetal or metalloid is −44.6% (n = 1628) compared to −23.0% for materials with only metals (n = 346; Supplementary Fig. 13).
To further quantify angular dependence within the system, we develop a metric ζ to determine to what extent a particular system exhibits more angular variation in its charge density with respect to atom locations. This is achieved by measuring the dot product of the unit vector from a probe point to its nearest neighboring atom and the gradient of the density at that probe point:
where G is a set of probe points and \({\hat{{{{\bf{r}}}}}}_{ki}\) is a unit vector from probe point at location \({\vec{{{{\bf{g}}}}}}_{k}\) to the closest atom at location \({\vec{{{{\bf{r}}}}}}_{i}\). Intuitively, a material with charge density gradients pointing directly towards or away from the nearest atom will have dot products larger in magnitude and ζ → 0, whereas a material with density that has high angular variance with respect to the nearest atom will have dot products smaller in magnitude, and ζ → 1; see Supplementary Fig. 14 for the test set average ζ over G_{local,i} for every element. We observe that the differences in performance between the lower and higher rotation order networks correlates strongly with ζ (Supplementary Fig. 15). With covalent bonding and associated high angular variance frequently present in the Materials Project data, this further justifies the need for introducing higher rotation order representations into charge density prediction networks. In addition, we find the higherorder model maintains higher accuracy under perturbations to atomic positions (Supplementary Fig. 11), an imperative trait when using models for phonon calculations.
Runtime and scalability
To demonstrate the scalability of our model, we analyze the runtime duration of our model on systems with increasing number of atoms, and compare to the duration of DFT calculations. We run each method on a single material, BaVS_{3} (mp3451 in MP), creating supercells from 1 × 1 × 1 to 10 × 10 × 10 and recording the runtime to generate charge density of the system. Supplementary Fig. 16 shows an approximate linear, O(N), scaling of our model with respect to number of atoms, while DFT exhibits approximately O(N^{2.3}) before exceeding our computational budget. This is to be expected, as the graph size increases linearly with cell volume if the resolution remains the same, while DFT has shown to scale at a cubic rate with respect to system size^{43}. Furthermore, our method can be fully parallelized up to each point in the density grid with no communication overhead.
Discussion
This work introduces ChargE3Net, an architecture for predicting charge density through equivariance with higherorder representations. We demonstrate that introducing higherorder tensor representations over scalars and vectors used in prior work achieves greater predictive accuracy, and show how this is achieved through the improved modeling of systems with high angular variance. ChargE3Net enables accurate charge density calculations on atomic systems larger than what is computationally feasible with DFT, but can also be used in combination with selfconsistent DFT to reduce the number of SCF steps needed for convergence. Our findings hold across the large and diverse Materials Project dataset, paving the way towards a universal charge density model for materials. Future work to expand the model to charge density calculations of unrelaxed structures such as those in MPtrj^{24} would support general application of ML models to accelerate ab initio calculations and should be pursued.
The work of Hohenberg and Kohn proves that all material properties could, in principle, be derived from the electron density^{1}. Because this quantity is so fundamental to predicting material properties, we view models that compute electron densities as a potential intermediate towards other material properties that may improve the generalizability of machine learning models in material science. As evidence, our nonselfconsistent calculations, which can be seen as a groundtruth calculation of the electronic strucuture (and derived properties) using the crystal structure and charge density as inputs, demonstrate that one could indeed obtain electronic and thermal properties of materials using highly accurate chargedensity estimates. On a broader dataset, we show that nonselfconsistent calculations are able to capture the total energy, atomic forces, electronic bandstructure, and the band gap within chemical accuracy for a majority of structures in the test set, though there is room for improvement. This leads us to believe that large models trained on a charge density dataset that spans the known material space could be a promising direction towards a foundation model for DFT calculations that can be finetuned for downstream property prediction.
We have also explored the connection between errors in charge density predictions and the property error. We show in Supplementary Figs. 17–22 that errors in the charge density establish a lower bound to the property error, and that small charge density errors can give rise to substantial property error. Examining charge density errors more closely, specifically analyzing ε_{mae} localization in low and higherror nonselfconsistent calculations, we find difficulty identifying localized regions where charge density prediction errors lead to property errors (Supplementary Fig. 23). While the errors in nonselfconsistent property predictions are certainly affected by errors in charge density predictions from ChargE3Net, the errors may in part be explained by an incorrect DFT setting (the VASP tag LMAXMIX. See the Methods section for more information). The errors introduced by this calculation parameter may mean achieving chemical accuracy is improbable for some materials using nonselfconsistent calculations even when initializing with a selfconsistent (i.e. ground truth) charge density.
We demonstrate that higherorder E(3) equivariant tensor operations yield significant improvement in charge density predictions and are more data efficient (Figs. 3 and 4). However, the E(3) convolution used in this work are limited by the O(L^{6}) time complexity of tensor products^{62}. Recent work^{62} has showed that SO(3) convolutions can be reduced to SO(2), with a time complexity of O(L^{3}). In the future, our approach can be combined with the methodology presented in^{62} to improve computational efficiency, and enable the practical usage of higher rotation orders (L > > 4) in atomistic modeling.
We also anticipate extending our model to include the spin density, which will allow for acceleration of DFT calculations of magnetic materials. Training ChargE3Net with energies and forces would allow for rapid identification of the lowest energy magnetic configuration, a normally laborious and timeintensive process^{63,64}, and the ability to obtain magnetic properties, such as the Curie temperature.
Methods
ChargE3Net architecture
As described in Section 2.1, the inputs of the model consist of probe coordinates \(\vec{{{{\bf{g}}}}}\) and atom coordinates \(\vec{{{{\bf{r}}}}}\) with species z. Each atom and probe point are represented by a dictionary of tensors, where A^{n} denotes an atom representation and P^{n} denotes a probe representation, each at layer n. A^{n=0} is initialized as a onehot encoded z, represented as an ℓ = 0, p = 1 tensor with N_{channels} equal to the number of atomic species. P^{n=0} is initialized as single scalar zero, with ℓ = 0, p = 1, and N_{channels} = 1. Each representation is updated through a series of alternating convolution Conv(⋅) and nonlinearity Gate(⋅) layers, a total of N_{interactions} = 3 times. With each convolution layer, all possible outputs of the tensor product at different ℓ and p values with ℓ_{o}≤L are concatenated, increasing the representation size past the scalar initializations of A^{n} and P^{n}. Atom representations are updated with \({{{{\bf{A}}}}}_{i}^{n+1}=\,{{\mbox{Gate}}}({{{\mbox{Conv}}}\,}_{{{{\rm{atom}}}}}^{n}({\vec{{{{\bf{r}}}}}}_{i},{{{{\bf{A}}}}}_{i}^{n}))\). Conv_{atom} is defined as:
where W_{1}, W_{2}, W_{3} are learned weights applied as a linear mix or selfinteraction^{52}. The set N(i) includes all atoms within the cutoff distance, including those outside potential periodic boundaries. r_{ij} is the distance from \({\vec{r}}_{i}\) and \({\vec{r}}_{j}\), with unit vector \({\hat{r}}_{ij}\). \(Y({\hat{{{{\bf{r}}}}}}_{ij})\) are spherical harmonics, and \(R({r}_{ij})\in {{\mathbb{R}}}^{{N}_{{{{\rm{basis}}}}}}\) is a learned radial basis function defined as:
where MLP(⋅) is a twolayer multilayer perceptron with SiLU nonlinearity^{65}, ϕ(⋅) is a Gaussian basis function \(\phi {({r}_{ij})}_{k}\propto \exp ({({r}_{ij}{\mu }_{k})}^{2})\) with μ_{k} uniformly sampled between zero and the cutoff, then normalized to have a zero mean and unit variance.
The convolution updates the representation of each atom to be the sum of tensor products between neighboring atom representations and the corresponding spherical harmonics describing their relative angles, weighted by a learned radial representation. This sum is followed by an additional selfconnection, then a residual selfconnection. The output of the convolution is then passed though an equivariant nonlinearity Gate( ⋅ ) operation as described in ref. ^{66}. We use SiLU and \(\tanh\) activation functions for even and odd parity scalars respectively, as is done in^{16}.
Probe representations are updated similarly as the atoms for each layer, except their representations depend solely on the representations of neighboring atoms, with no probeprobe interaction. Each probe representation is updated with \({{{{\bf{P}}}}}_{k}^{n+1}=\,{{\mbox{Gate}}}({{{\mbox{Conv}}}\,}_{{{{\rm{probe}}}}}^{n}({\vec{{{{\bf{g}}}}}}_{k},{{{{\bf{P}}}}}_{k}^{n}))\), where Conv_{probe} is defined as:
Note that weights W are not shared with those for \({\,{{\mbox{Conv}}}\,}_{{{{\rm{atom}}}}}^{n}(\cdot )\). Since atom representations are computed independently of probe positions, they can be computed once per atomic configuration, even if multiple batches are required to obtain predictions for all probe points. Finally, predicted charge density \(\hat{\rho }({\vec{{{{\bf{g}}}}}}_{k})\) is computed as a linear combination of the scalar elements of the final representation \({{{{\bf{P}}}}}_{k}^{n = {N}_{{{{\rm{layers}}}}}}\).
Training ChargE3Net models
We develop an algorithm to achieve more efficient atomprobe graph construction, enabling faster training and inference times. Prior work constructs graphs by adding probe points directly to the atomic structure, computing edges with atomic simulation software such as ase or asap^{67}, then removing edges between probes. While this is performant for small systems with few probes, we find it does not scale well due to the redundant computation of probeprobe edges, which grow quadratically with respect to the number of probe points. To address this, our algorithm leverages two kd trees^{68}; one to partition atoms, and one to partition probe points. This enables an efficient radius search for atomprobe edges exclusively. We implement this through modification of the kd tree implementation from scipy^{69} to handle periodic boundary conditions. Constructing graphs with 1000 probe points, we observe an average graph construction time ( ± standard error) of 0.00758 ± 0.00087, 0.0103 ± 0.0005, and 0.00462 ± 0.00056 seconds for MP, QM9, and NMC respectively. This is at least an order of magnitude faster than methods from ref. ^{43}, with which we measured graph construction times of 12.6 ± 0.9, 0.114 ± 0.006, and 0.141 ± 0.010 s.
To train the ChargE3Net model, Supplementary Table 1 outlines the experimental setup on each of the datasets (e.g., learning rate, batch size, number of probe points and number of training steps). Specifically, at each gradient step, a random batch of materials is selected, from which a random subset of N_{points} probe points are selected. We use the Adam optimizer^{70}, and decay the learning rate by 0.96^{s/β} at step s. We find that optimizing for L1 error improves performance and training stability over mean squared error (Supplementary Table 2). Supplementary Fig. 24 shows the evolution of training loss and validation error as the model is trained on each dataset. All models were implemented in PyTorch and trained on NVIDIA Volta V100 GPUs using a distributed compute architecture.
Datasets
The QM9 and NMC datasets, along with train, validation and test splits, are provided by Jørgensen and Bhowmik^{43}. The QM9 dataset^{56} contains VASP calculations of the 133,845 small organic molecules introduced by Ramakrishnan et al.^{55}, with training, validation, and test splits of size 123,835, 50, and 10,000 respectively. The NMC dataset^{57} consists of VASP calculations of 2000 randomly sampled nickel manganese cobalt oxides containing varying levels of lithium content, with training, validation, and test splits of size 1450, 50, and 500 respectively. For the MP data, we collected 122,689 structures and associated charge densities from api.materialsproject.org^{49}. Structures in the dataset that shared composition and space group were identified as duplicates and only the highest material_id structure was included, leaving 108,683 materials. The data was split randomly into training, validation, and test splits with sizes 106,171, 512, and 2000 respectively. Later, analyzing the test set materials for thermodynamic stability and structural abnormalities (Supplementary Fig. 25), 26 materials were found to have unrealistic volume per atom, in excess of 100 Å^{3} atom^{−1} and mostly from space group Immm. These materials were removed from the test set and excluded from all reported results and figures.
We note to the reader a known issue with the Materials Project dataset. Original calculation parameters may have an incorrectly set VASP parameter called LMAXMIX, which sets the maximum angular momentum number for the onecenter PAW charge densities saved in the charge density output of VASP and passed into the charge density mixer. This affects elements with d and felectrons when no Hubbard U parameter was set by the Materials Project. In Supplementary Fig. 7, we correct this error for the randomly sampled test set, then compare the distribution of properties between nonselfconsistent calculations of the original and corrected input parameters. Note that nonselfconsistent calculations with a potentially incorrectly set LMAXMIX, as is set in the current Materials Project charge density dataset, makes a peratom energy error about two orders of magnitude larger than that of selfconsistent calculations. This represents an additional source of error to all nonselfconsistent calculations based on Materials Project data.
In addition to a heldout test set randomly sampled from the Materials Project, we obtain a random set of 2000 published structures from GNoME^{50} as a second test set. Specifically, we choose a random set of materials from GNoME with 5 or 6 unique species as these are less wellrepresented in the Materials Project training data. DFT calculations were set up using pymatgen^{71} and otherwise using the same methodology as detailed in Section 4.4. 76 calculations failed, leaving 1924 materials for analysis, 1510 of which were nonmagnetic. Nonmagnetic structures were identified via DFT calculations initialized by SAD and magnetic moment guesses provided by pymatgen with an overall absolute magnetic moment <0.1 μB and absolute atomic magnetic moments <0.1 μB for every atom in the structure.
Supplementary Figure 1 shows a visualization of SOAP descriptors of materials in the MP train and test set and GNoME subset, projected into 2d using tSNE. We do not find a significant difference in distributions across the MP train and test sets, which is to be expected as the data is divided randomly. Supplementary Figure 2 shows the distribution of crystal systems and space groups within the MP training set. As seen with the SOAP descriptors, in Supplementary Fig. 3 we note a similar distribution of crystal systems across the MP train and test sets, with the GNoME subset showing significantly more monoclinic systems and significantly fewer tetragonal and cubic systems.
Density functional theory calculations
All density functional theory (DFT) calculations were done with VASP version 5.4.4^{37,38,39,72} using projectoraugmented wave (PAW) pseudopotentials^{73}. For all calculations done on the test set of Materials Project data, we use the input files provided by the Materials Project to ensure these are computed in a manner consistent with how the charge density grids are originally generated. For the GNoME dataset, we use a kpoint density of 100 points Å^{3}, the same density that the original Materials Project charge density dataset was computed with, and generated the input files with pymatgen^{71}. We set LMAXMIX to 4 for delements and 6 for felements only for DFT+U calculations as set by pymatgen, which, while incorrect (see Section 4.3), allows us to compare our performance on this set of materials against that of the Materials Project test set. For the experiments with perturbed structures, including those where electronic and thermal properties were computed, we use all the same settings as the Materials Project but with different kgrids, all generated with the MonkhorstPack scheme unless otherwise specified. We use a 520 eV energy cutoff and converge each calculation until the change in energy between SCF steps is less than 10^{−}5 eV. For the Si(111)(2 × 1) calculation, we use a kgrid of 4 × 8 × 1, while all zeolite calculations were done at the gammapoint only. Electronic bandstructures were plotted using pymatgen^{71}. We use the same Hubbard U values as the Materials Project, which were obtained from Wang et al.^{74} and Jain et al.^{75}. PAW augmentation occupancies were not predicted by ChargE3Net and were taken from the Materials Project dataset or computed selfconsistently for the GNoME dataset.
Thermal properties were computed using phonopy^{76,77} for the heat capacity. Zeolite and bulk Si structures were relaxed until the maximum norm of the HellmannFeynman forces fell below 0.001 eV Å^{−1} prior to computing thermal properties. Structures were perturbed by 0.03Å to find the Hessian of the potential energy surface via finite differences, and the heat capacity at constant volume can then be found via the harmonic approximation with the equation
where θ_{k} = hν_{k}/k_{B}, h is Planck’s constant, ν_{k} is the vibrational frequency of vibrational mode k, k_{B} is Boltzmann’s constant, R is the gas constant, and T is the temperature in Kelvin.
Data availability
Materials Project data was collected on 7 May 2023. Associated task identifiers will be included in our provided repository, along with the necessary download script, and our train, validation, and test splits. We also provide the identifiers for the subset of GNoME data used in our work. VASP computations on the QM9 and NMC datasets were downloaded from refs. ^{56} and ^{57} respectively.
Code availability
Our model implementation, along with pretrained model weights is available at https://github.com/AIforGreatGood/charge3net.
References
Hohenberg, P. & Kohn, W. Inhomogeneous electron gas. Phys. Rev. 136, B864–B871 (1964).
Kohn, W. & Sham, L. J. Selfconsistent equations including exchange and correlation effects. Phys. Rev. 140, A1133–A1138 (1965).
Prentice, J. C. A. et al. The ONETEP linearscaling density functional theory program. J. Chem. Phys. 152, 174111 (2020).
Soler, J. M. et al. The SIESTA method for ab initio orderN materials simulation. J. Condens. Matter Phys. 14, 2745 (2002).
Mohr, S. et al. Accurate and efficient linear scaling DFT calculations with universal applicability. Phys. Chem. Chem. Phys. 17, 31360–31370 (2015).
Riplinger, C. & Neese, F. An efficient and near linear scaling pair natural orbital based local coupled cluster method. J. Chem. Phys. 138, 034106 (2013).
Riplinger, C., Sandhoefer, B., Hansen, A. & Neese, F. Natural triple excitations in local coupled cluster calculations with pair natural orbitals. J. Chem. Phys. 139, 134101 (2013).
Witt, W. C., del Rio, B. G., Dieterich, J. M. & Carter, E. A. Orbitalfree density functional theory for materials research. J. Mater. Res. 33, 777–795 (2018).
Xie, T. & Grossman, J. C. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Phys. Rev. Lett. 120, 145301 (2018).
Chen, C., Ye, W., Zuo, Y., Zheng, C. & Ong, S. P. Graph networks as a universal machine learning framework for molecules and crystals. Chem. Mater. 31, 3564–3572 (2019).
Louis, S.Y. et al. Graph convolutional neural networks with global attention for improved materials property prediction. Phys. Chem. Chem. Phys. 22, 18141–18148 (2020).
Dunn, A., Wang, Q., Ganose, A., Dopp, D. & Jain, A. Benchmarking materials property prediction methods: the Matbench test set and Automatminer reference algorithm. Npj Comput. Mater. 6, 1–10 (2020).
Choudhary, K. & DeCost, B. Atomistic line graph neural network for improved materials property predictions. Npj Comput. Mater. 7, 185 (2021).
Koker, T., Quigley, K., Spaeth, W., Frey, N. C. & Li, L. Graph contrastive learning for materials. Preprint at https://arxiv.org/abs/2211.13408 (2022).
Drautz, R. Atomic cluster expansion for accurate and transferable interatomic potentials. Phys. Rev. B 99, 014104 (2019).
Batzner, S. et al. E (3)Equivariant graph neural networks for dataefficient and accurate interatomic potentials. Nat. Commun. 13, 2453 (2022).
Musaelian, A. et al. Learning local equivariant representations for largescale atomistic dynamics. Nat. Commun. 14, 579 (2023).
Chmiela, S. et al. Machine learning of accurate energyconserving molecular force fields. Sci. Adv. 3, e1603015 (2017).
Du, X. et al. Machinelearningaccelerated simulations to enable automatic surface reconstruction. Nat. Comput. Sci. 3, 1034–1044 (2023).
Chanussot, L. et al. Open Catalyst 2020 (OC20) dataset and community challenges. ACS Catal. 11, 6059–6072 (2021).
Tran, R. et al. The Open Catalyst 2022 (OC22) dataset and challenges for oxide electrocatalysts. ACS Catal. 13, 3066–3084 (2023).
Sriram, A.et al. The Open DAC 2023 dataset and challenges for sorbent discovery in direct air capture. arXiv https://doi.org/10.48550/arXiv.2311.00341 (2023).
Pengmei, Z., Liu, J. & Shu, Y. Beyond MD17: the reactive xxMD dataset. Sci. Data 11, 222 (2024).
Deng, B. et al. CHGNet as a pretrained universal neural network potential for chargeinformed atomistic modelling. Nat. Mach. Intell. 5, 1031–1041 (2023).
Batatia, I. et al. A foundation model for atomistic materials chemistry. Preprint at https://arxiv.org/abs/2401.00096 (2023).
Brockherde, F. et al. Bypassing the KohnSham equations with machine learning. Nat. Commun. 8, 872 (2017).
Grisafi, A. et al. Transferable machinelearning model of the electron density. ACS Cent. Sci. 5, 57–64 (2018).
Fabrizio, A., Grisafi, A., Meyer, B., Ceriotti, M. & Corminboeuf, C. Electron density learning of noncovalent systems. Chem. Sci. 10, 9424–9432 (2019).
Lewis, A. M., Grisafi, A., Ceriotti, M. & Rossi, M. Learning electron densities in the condensed phase. J. Chem. Theory Comput. 17, 7203–7214 (2021).
Qiao, Z. et al. Informing geometric deep learning with electronic interactions to accelerate quantum chemistry. Proc. Natl. Acad. Sci. USA 119, e2205221119 (2022).
Rackers, J. A., Tecot, L., Geiger, M. & Smidt, T. E. A recipe for cracking the quantum scaling limit with machine learned electron densities. Mach. Learn.: Sci. Technol. 4, 015027 (2023).
del Rio, B. G., Phan, B. & Ramprasad, R. A deep learning framework to emulate density functional theory. Npj Comput. Mater. 9, 158 (2023).
Li, H. et al. Deeplearning density functional theory Hamiltonian for efficient ab initio electronicstructure calculation. Nat. Comput. Sci. 2, 367–377 (2022).
Unke, O. et al. SE (3)equivariant prediction of molecular wavefunctions and electronic densities. Adv. Neural Inf. Process. Syst. 34, 14434–14447 (2021).
Nigam, J., Willatt, M. J. & Ceriotti, M. Equivariant representations for molecular Hamiltonians and Ncenter atomicscale properties. J. Chem. Phys. 156, 14115 (2022).
Zhong, Y., Yang, J., Xiang, H. & Gong, X. Universal machine learning kohnsham hamiltonian for materials. Chin. Phys. Lett. 41, 077103 (2024).
Kresse, G. & Hafner, J. Ab initio molecular dynamics for liquid metals. Phys. Rev. B 47, 558–561 (1993).
Kresse, G. & Furthmüller, J. Efficiency of abinitio total energy calculations for metals and semiconductors using a planewave basis set. Comput. Mater. Sci. 6, 15–50 (1996).
Kresse, G. & Furthmüller, J. Efficient iterative schemes for ab initio totalenergy calculations using a planewave basis set. Phys. Rev. B 54, 11169–11186 (1996).
Giannozzi, P. et al. QUANTUM ESPRESSO: a modular and opensource software project for quantum simulations of materials. J. Condens. Matter Phys. 21, 395502 (2009).
Gong, S. et al. Predicting charge density distribution of materials using a localenvironmentbased graph convolutional network. Phys. Rev. B 100, 184103 (2019).
Jørgensen, P. B. & Bhowmik, A. DeepDFT: Neural message passing network for accurate charge density prediction. Preprint at https://arxiv.org/abs/2011.03346 (2020).
Jørgensen, P. B. & Bhowmik, A. Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids. Npj Comput. Mater. 8, 183 (2022).
Sunshine, E. M., Shuaibi, M., Ulissi, Z. W. & Kitchin, J. R. Chemical properties from graph neural networkpredicted electron densities. J. Phys. Chem. C 127, 23459–23466 (2023).
Focassio, B., Domina, M., Patil, U., Fazzio, A. & Sanvito, S. Linear jacobilegendre expansion of the charge density for machine learningaccelerated electronic structure calculations. Npj Comput. Mater. 9, 87 (2023).
Schütt, K. et al. Schnet: a continuousfilter convolutional neural network for modeling quantum interactions. Adv. Neural Inf. Process. Syst. 30, 991–1001 (2017).
Schütt, K., Unke, O. & Gastegger, M. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, 9377–9388 (PMLR, 2021).
Pope, P. & Jacobs, D. Towards combinatorial generalization for catalysts: a kohnsham chargedensity approach. Adv. Neural Inf. Process. Syst. 36, 60585–60598 (2024).
Shen, J.X. et al. A representationindependent electronic charge density database for crystalline materials. Sci. Data 9, 661 (2022).
Merchant, A. et al. Scaling deep learning for materials discovery. Nature 624, 80–85 (2023).
Pozdnyakov, S. N. & Ceriotti, M. Incompleteness of graph neural networks for points clouds in three dimensions. Mach. Learn.: Sci. Technol. 3, 045020 (2022).
Thomas, N. et al. Tensor field networks: rotationand translationequivariant neural networks for 3d point clouds. Preprint at https://arxiv.org/abs/1802.08219 (2018).
Geiger, M. & Smidt, T. e3nn: Euclidean neural networks. Preprint at https://arxiv.org/abs/2207.09453 (2022).
Ruddigkeit, L., Van Deursen, R., Blum, L. C. & Reymond, J.L. Enumeration of 166 billion organic small molecules in the chemical universe database GDB17. J. Chem. Inf. Model 52, 2864–2875 (2012).
Ramakrishnan, R., Dral, P. O., Rupp, M. & Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Sci. Data 1, 1–7 (2014).
Jørgensen, P. B. & Bhowmik, A. QM9 Charge Densities And Energies Calculated With VASP. https://doi.org/10.11583/DTU.16794500.v1 (2022).
Jørgensen, P. B. & Bhowmik, A. NMC liion Battery Cathode Energies And Charge Densities. https://doi.org/10.11583/DTU.16837721.v1 (2022).
Jain, A.et al. Commentary: The Materials Project: a materials genome approach to accelerating materials innovation. APL Mater.1, 11002 (2013).
Sholl, D. S. & Steckel, J. A. Density Functional Theory (John Wiley & Sons, Ltd, 2009).
Lehtola, S. Assessment of initial guesses for selfconsistent field calculations. superposition of atomic potentials: simple yet efficient. J. Chem. Theory Comput. 15, 1593–1604 (2019).
Wing, D. et al. Band gaps of crystalline solids from Wannierlocalization–Based optimal tuning of a screened rangeseparated hybrid functional. Proc. Natl. Acad. Sci. USA 118, e2104556118 (2021).
Passaro, S. & Zitnick, C. L. Reducing SO(3) convolutions to SO(2) for efficient equivariant gnns. In International Conference on Machine Learning, 27420–27438 (PMLR, 2023).
Horton, M. K., Montoya, J. H., Liu, M. & Persson, K. A. Highthroughput prediction of the groundstate collinear magnetic order of inorganic materials using Density Functional Theory. Npj Comput. Mater. 5, 1–11 (2019).
Lee, K., Youn, Y. & Han, S. Identification of groundstate spin ordering in antiferromagnetic transition metal oxides using the Ising model and a genetic algorithm. Sci. Technol. Adv. Mater. 18, 246–252 (2017).
Hendrycks, D. & Gimpel, K. Gaussian error linear units (GELUs). Preprint at https://arxiv.org/abs/1606.08415 (2016).
Weiler, M., Geiger, M., Welling, M., Boomsma, W. & Cohen, T. S. 3d steerable cnns: learning rotationally equivariant features in volumetric data. Adv. Neural Inf. Process. Syst. 31, 10381–10392 (2018).
Larsen, A. H. et al. The atomic simulation Environment—a Python library for working with atoms. J. Condens. Matter Phys. 29, 273002 (2017).
Bentley, J. L. Multidimensional binary search trees used for associative searching. Commun. ACM 18, 509–517 (1975).
Virtanen, P. et al. SciPy 1.0: fundamental algorithms for scientific computing in python. Nat. Methods 17, 261–272 (2020).
Kingma, D. P. & Ba, J. Adam: a method for stochastic optimization. Int. Conf. Learn. Represent. 3, 1–15 (2015).
Ong, S. P. et al. Python Materials Genomics (pymatgen): a robust, opensource python library for materials analysis. Comput. Mater. Sci. 68, 314–319 (2013).
Hafner, J. Abinitio simulations of materials using VASP: densityfunctional theory and beyond. J. Comput. Chem. 29, 2044–2078 (2008).
Kresse, G. & Joubert, D. From ultrasoft pseudopotentials to the projector augmentedwave method. Phys. Rev. B 59, 1758–1775 (1999).
Wang, L., Maxisch, T. & Ceder, G. Oxidation energies of transition metal oxides within the GGA+U framework. Phys. Rev. B 73, 195107 (2006).
Jain, A. et al. Formation enthalpies by mixing GGA and GGA+U calculations. Phys. Rev. B 84, 045115 (2011).
Togo, A., Chaput, L., Tadano, T. & Tanaka, I. Implementation strategies in phonopy and phono3py. J. Condens. Matter Phys. 35, 353001 (2023).
Togo, A. Firstprinciples phonon calculations with phonopy and phono3py. J. Phys. Soc. Jpn 92, 012001 (2023).
Acknowledgements
We thank Tess Smidt and Mark Polking for helpful discussions. DISTRIBUTION STATEMENT A. Approved for public release. Distribution is unlimited. This material is based upon work supported by the Under Secretary of Defense for Research and Engineering under Air Force Contract No. FA870215D0001. Any opinions, findings, conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the Under Secretary of Defense for Research and Engineering. © 2023 Massachusetts Institute of Technology. Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part 252.2277013 or 7014 (Feb 2014). Notwithstanding any copyright notice, U.S. Government rights in this work are defined by DFARS 252.2277013 or DFARS 252.2277014 as detailed above. Use of this work other than as specifically authorized by the U.S. Government may violate any copyrights that exist in this work.
Author information
Authors and Affiliations
Contributions
T.K. designed, implemented, and trained the ChargE3Net models. K.Q. collected and prepared the MP, QM9, and NMC datasets, and contributed to the software implementation and training of baseline models. K.Q. devised the angular variation metric and performed analysis on the effects of rotation order together with T.K. E.T. performed experimentation and analysis for initialization of DFT calculations and nonselfconsistent calculation of material properties. K.T. provided guidance on DFT calculations. L.L. conceived the project, supervised the research and advised experimentation and data analysis. All authors contributed to the writing and editing of the manuscript.
Corresponding authors
Ethics declarations
Competing interests
The authors declare no competing interests.
Additional information
Publisher’s note Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.
Supplementary information
Rights and permissions
Open Access This article is licensed under a Creative Commons Attribution 4.0 International License, which permits use, sharing, adaptation, distribution and reproduction in any medium or format, as long as you give appropriate credit to the original author(s) and the source, provide a link to the Creative Commons licence, and indicate if changes were made. The images or other third party material in this article are included in the article’s Creative Commons licence, unless indicated otherwise in a credit line to the material. If material is not included in the article’s Creative Commons licence and your intended use is not permitted by statutory regulation or exceeds the permitted use, you will need to obtain permission directly from the copyright holder. To view a copy of this licence, visit http://creativecommons.org/licenses/by/4.0/.
About this article
Cite this article
Koker, T., Quigley, K., Taw, E. et al. Higherorder equivariant neural networks for charge density prediction in materials. npj Comput Mater 10, 161 (2024). https://doi.org/10.1038/s41524024013431
Received:
Accepted:
Published:
DOI: https://doi.org/10.1038/s41524024013431
This article is cited by

Electronic structure prediction of multimillion atom systems through uncertainty quantification enabled transfer learning
npj Computational Materials (2024)