Predicting charge density distribution of materials using a local-environment-based graph convolutional network

The electron charge density distribution of materials is one of the key quantities in computational materials science as theoretically it determines the ground state energy and practically it is used in many materials analyses. However, the scaling of density functional theory calculations with number of atoms limits the usage of charge-density-based calculations and analyses. Here we introduce a machine-learning scheme with local-environment-based graphs and graph convolutional neural networks to predict charge density on grid points from the crystal structure. We show the accuracy of this scheme through a comparison of predicted charge densities as well as properties derived from the charge density, and that the scaling is $O$($N$). More importantly, the transferability is shown to be high with respect to different compositions and structures, which results from the explicit encoding of geometry.

Figure 1

(a) Crystal structure of crystalline ethylene. The blue plus symbol in the center denotes a grid point we are interested in. (b) Crystalline ethylene with the imaginary atom. Highlighted atoms are those within the cutoff radius. (c) Local environment around the imaginary atom. (d) Sketch of local-environment-based graph and CGCNN architecture. Color coding: Green, carbon; gray, hydrogen; blue, imaginary atom; yellow, highlighted atoms within the cutoff radius.

Figure 2

(a) Mean average error (MAE, in $e/{\AA }^3$) of the ML-predicted charge density of the test sets (grid points) from the training structures versus training set size for polymer and zeolite materials. (b) CPU time (in seconds) for DFT calculations and GPU time (in seconds) for ML prediction versus number of atoms in the cell for crystalline p-xylylene. DFT calculations are performed by 24 Intel Xeon CPUs with RAM of 128 GB, while ML calculations are carried out on a single NVIDIA GeForce GTX 1070 GPU with RAM of 2 GB.

Figure 3

(a) and (b) Appearance frequency of coordinated atoms of carbon atoms in the training set for the case of crystalline polymers versus the test set as a whole and nomex and nomex_defect, respectively. Here the “$X$”s in “C2X1” and “C2X2” denote rare elements in our case (Cl, F, S, Si, Hg). (c) Appearance frequency of oxygen coordinated atoms in the training set for the case of zeolites versus the structure of NPO_defect.

Figure 4

Visualization of electron charge density ($\rho$, in $e/{\AA }^3$). Panels (a), (b), (c) and (d), (e), (f) show crystal structure, ML-predicted $\rho$, and difference between ML-predicted $\rho$ and DFT-calculated $\rho$ on the C six-ring plane of pristine nomex and nomex with a carbon and a hydrogen vacancy, respectively. Panels (g), (h), (i) and (j), (k), (l) show crystal structure, ML-predicted $\rho$, and difference between ML-predicted $\rho$ and DFT-calculated $\rho$ on the Si-O six-ring plane of pristine NPO and NPO with a Si vacancy, respectively. Atom color coding: Green, carbon; gray, hydrogen; red, oxygen; blue, nitrogen; yellow, silicon.

Figure 5

Panels (a), (b), (c), and (d) show ML-predicted charge density ($\rho$, in $e/{\AA }^3$) versus DFT calculated $\rho$ for pristine nomex, nomex_defect, pristine NPO, and NPO_defect, respectively.

Figure 6

(a) Sketch of two different local environments with similar sum of atom contributions. (b) Geometries of central carbon atoms with coordinated C1H3, C2H2, C3H1, and C4 atoms. (c) Shape of C-C-H and C-O-H and their charge density distributions ($\rho$, in $e/{\AA }^3$). Atom color coding: Green, carbon; gray, hydrogen; red, oxygen.

Figure 7

(a) Illustration of the first toy-model experiment. The top and bottom MAEs (in $e/{\AA }^3$) are from the predictions to the orthogonal C-C-C molecule by one of the two training molecules (linear C-C-C and orthogonal C-O-C), while the middle one is from the prediction trained on both of the training molecules. (b) Illustration of the second toy-model experiment. The top and bottom MAEs (in $e/{\AA }^3$) are from the predictions to the orthogonal C-C-C molecule by one of the two training molecules (linear C-C-C and linear C-O-C), while the middle one is from the prediction trained on both of the training molecules (linear C-C-C and linear C-O-C). Atom color coding: Green, carbon; red, oxygen.