ChargeFlow: Flow-Matching Refinement of Charge-Conditioned Electron Densities

Tri Minh Nguyen; Sherif Abdulkader Tawfik; Truyen Tran; Svetha Venkatesh

cond-mat.mtrl-sciMar 25, 2026

ChargeFlow: Flow-Matching Refinement of Charge-Conditioned Electron Densities

Authors: Tri Minh Nguyen, Sherif Abdulkader Tawfik, Truyen Tran, Svetha Venkatesh

Organizations: Deakin University, Geelong, VIC 3216, Australia

Abstract

Accurate charge densities are central to electronic-structure theory, but computing charge-state-dependent densities with density functional theory remains too expensive for large-scale screening and defect workflows. We present ChargeFlow, a flow-matching refinement model that transforms a charge-conditioned superposition of atomic densities into the corresponding DFT electron density on the native periodic real-space grid using a 3D U-Net velocity field. Trained on 9,502 charged Materials Project-derived calculations and evaluated on an external 1,671-structure benchmark spanning perovskites, charged defects, diamond defects, metal-organic frameworks, and organic crystals, ChargeFlow is not uniformly best on every in-distribution class but is strongest on problems dominated by nonlocal charge redistribution and charge-state extrapolation, improving deformation-density error from 3.62% to 3.21% and charge- response cosine similarity from 0.571 to 0.655 relative to a ResNet baseline. The predicted densities remain chemically useful under downstream analysis, yielding successful Bader partitioning on all 1,671 benchmark structures and high-fidelity electrostatic potentials, which positions flow matching as a practical density-refinement strategy for charged materials.

Figures & tables

Split / class	Samples	Unique parents	Charge states
MP charged-density train	9502	4065	$\{-3,-2,-1,0,1,2,3\}$
MP charged-density dev hold-out	2376	1952	$\{-3,-2,-1,0,1,2,3\}$
Perovskites	159	40	$\{-2,-1,1,2\}$
Charged defects	1149	588	$\{-3,-2,-1,1,2,3\}$
Multisite diamond defects	54	18	$\{-2,-1,1,2\}$
Special diamond defects	42	14	$\{-2,-1,1,2\}$

Table 1: Dataset breakdown for the Materials Project charged-density corpus and the external periodic benchmark. “Unique parents” are counted after removing the _charge_Q suffix from the structure identifier.

System Type	ChargeFlow	ChargeFlowL2	ResNet	U-Net	CDeepDFT
Perovskites	5.41	7.12	3.74	3.58	3.03
Charged Defects	6.46	9.00	4.68	5.53	5.26
Multisite Defects	6.15	6.43	8.60	8.59	9.67
Special Defects	6.78	7.73	8.50	7.34	9.54
MOFs	7.56	11.29	7.86	8.78	10.80
Extreme MOFs	7.84	10.80	8.50	7.96	nan

Table 2: Mean Absolute Error

\varepsilon_{\text{MAE}}

(%) on the external periodic test classes. Sample counts are given in parentheses. Bold indicates the best result per row. ChargeFlowL2 is a variant trained without the auxiliary MAE loss (

\alpha=0

). Extreme MOFs contain

Q\in\{-20,-10,+10,+20\}

; extreme organic crystals contain

Q\in\{-10,-5,+5,+10\}

. “nan” indicates that the model failed to produce valid predictions for that test set.

Metric	ChargeFlow	ResNet
Bader successful materials	1671 / 1671	1590 / 1671
Bader atom-level $R^{2}$	0.9901	0.9833
Bader atom-level MAE (e)	0.2369	0.2853
Potential mean per-material MAE (eV)	1.3326	1.1700
Potential mean per-material $R^{2}$	0.9954	0.9928

Table 3: Head-to-head downstream comparison on the 1,671-material periodic benchmark. Bader metrics are computed directly on predicted densities. ResNet Bader statistics are reported on the 1,590 structures for which the Bader partition completed successfully; potential metrics are reported on all 1,671 materials.

Figure 1: Detailed Bader charge analysis for ChargeFlow across 1,671 periodic materials (67,274 atoms). (A) Atom-level parity plot of predicted versus ground-truth Bader charges, showing $R^{2}=0.9901$ and MAE $=0.237$ e. (B) Distribution of per-material Bader-charge MAE by material class. (C) Correlation between electron density error ( $\varepsilon_{\text{MAE}}$ ) and per-material Bader-charge MAE, demonstrating that density accuracy translates to accurate charge partitioning. Table 3 reports the ChargeFlow–ResNet comparison.

Figure 2: Detailed electrostatic (Hartree) potential analysis for ChargeFlow across 1,671 periodic materials. (A) Mean per-material potential $R^{2}$ as a function of system net charge $|Q|$ , showing stable accuracy even for extreme charge states (overall mean $R^{2}=0.9954$ ). (B) Distribution of per-material potential MAE by material class. Table 3 reports the ChargeFlow–ResNet comparison.

Figure 3: Per-class breakdown of electrostatic potential MAE as a function of system net charge $|Q|$ . Each panel corresponds to a different material class, with the number of test structures ( $n$ ), mean MAE, standard deviation, and median indicated. The diamond defect classes exhibit the tightest distributions, while MOFs show the largest variability owing to their complex pore geometries and extreme charge states.

Figure 4: Charge-state extrapolation for metal-organic frameworks.

\varepsilon_{\text{MAE}}

is plotted as a function of charge magnitude

|Q|

for ChargeFlow and ResNet. The dashed vertical line marks the training boundary (

|Q|=3

). ChargeFlow maintains lower error across all charge magnitudes, with the gap widening in the extrapolative regime (

|Q|=10,20

). Error bars indicate the standard error of the mean.

Figure 5: Charge-state extrapolation for organic crystals. ChargeFlow exhibits a remarkably flat error profile (

\varepsilon_{\text{MAE}}\approx 7.1

–

7.4

%), maintaining a consistent

\sim

1 percentage point advantage over ResNet across both in-range and extrapolative charge states. Error bars indicate the standard error of the mean.

Figure 6: Spatial decomposition of prediction error for ChargeFlow and ResNet. Mean absolute error (e/Bohr 3 ) is plotted as a function of distance to the nearest atom, with shaded bands indicating $\pm 1$ standard deviation across samples. Green and purple backgrounds denote the near-atomic core ( $d<1$ Å) and interstitial ( $d>2.5$ Å) regions, respectively. (A) Within the training charge range ( $|Q|\leq 3$ ), ChargeFlow lowers the peak near-atomic error relative to ResNet. (B) For extrapolative charge states ( $|Q|>3$ ), the near-atomic advantage becomes larger, consistent with improved modeling of charge redistribution under stronger perturbations.

Figure 7: Deformation density analysis comparing ChargeFlow and ResNet. The deformation density

\Delta\rho=\rho_{\text{DFT}}-\rho_{\text{promol}}

isolates the bonding and charge-transfer contribution after subtracting a promolecular reference. (A) Scatter of per-structure

R^{2}

on total density versus deformation density; points below the diagonal indicate that

\Delta\rho

is a harder target than the full density. (B) Mean

R^{2}

\Delta\rho

by material class. (C) Normalized mean absolute error

\varepsilon_{\text{MAE}}

\Delta\rho

by material class. (D) Per-class deformation-density advantage, reported as

\Delta R^{2}=R^{2}_{\mathrm{ChargeFlow}}-R^{2}_{\mathrm{ResNet}}

, where positive values favor ChargeFlow. ChargeFlow outperforms ResNet on 7 of 8 classes, with the largest gains on defective diamonds, organic crystals, and extreme-charge systems.

Figure 8: Charge-response analysis comparing ChargeFlow and ResNet across 862 sampled charge-state pairs. (A) Distribution of cosine similarity between predicted and ground-truth differential density

\delta\rho=\rho(Q_{2})-\rho(Q_{1})

. (B) Cosine similarity as a function of charge-jump magnitude

|\Delta Q|

, with the dashed line marking the training boundary (

|\Delta Q|=3

). (C) MAE of

\delta\rho

versus

|\Delta Q|

, showing that ChargeFlow achieves lower error at large charge perturbations. (D) Mean cosine similarity by material class.

Computational Materials Science

Computational materials science uses computation and data-driven models to predict, explain, and design materials and their properties.

See all 141 papers

Equivariant Flow Matching for Electron Density Prediction
Oct 2, 2026Chenxing Liang, Chengdong Wang, Yuchao Lin +2Computational Quantum Chemistry Quantum Chemistry
Department of Computer Science and Engineering, Texas A&M University · Department of Materials Science and Engineering, Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University +2
Neural-Network Solutions to Real-Space Charge Density and Generalization
Sep 14, 2026Yuxuan Zeng, Tong Xie, Taoyuze Lv +1Neural Surrogate Modeling Equivariant Neural Networks
School of Artificial Intelligence & Data Science, University of Science and Technology of China · Suzhou Institute for Advanced Research, University of Science and Technology of China
All-electron machine-learned orbital-free density functional for periodic CHNO systems
Oct 13, 2026Marc K. Ickler, Andreas Hermann, Fred A. HamprechtDensity Functional Theory Computational Quantum Chemistry
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University, Heidelberg, 69120, Germany. · Centre for Science at Extreme Conditions and SUPA, School of Physics and Astronomy, University of Edinburgh, Edinburgh, EH9 3FD, United Kingdom.

Oct 2, 2026J/K move · Enter open · S save

Equivariant Flow Matching for Electron Density Prediction

Chenxing Liang, Chengdong Wang, Yuchao Lin, Xiaofeng Qian+1

Department of Computer Science and Engineering, Texas A&M University · Department of Materials Science and Engineering, Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University+2

Computational Quantum Chemistry Quantum Chemistry Density Functional Theory Flow-Based Generative Modeling Molecular Property Prediction Flow Matching

Open paper AskPDF arXiv

Machine learning surrogates for density functional theory (DFT) have been increasingly used to reduce the cost of first-principles calculations. In this arena, predicting real-space electron densities offers a scalable and transferable initialization for self-consistent field (SCF) procedures. However, current methods face a clear dilemma. That is, grid-based architectures incur a high computational cost, while basis-set methods fail to capture the structural correlations inherent in the coefficient space. Here, we develop OrbFlow, an

\mathrm{SE}(3)

-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients via flow matching. OrbFlow retains the efficiency of a compact atom-centered basis while replacing pointwise regression with a learned probability path over the full coefficient space. It is trained through a two-phase trajectory curriculum that mitigates discretization drift during numerical integration. OrbFlow achieves state-of-the-art accuracy on QM9, reducing density error by 13.6% relative to the previous best model, and reduces error by 51% to 63% on every molecule of the MD benchmark relative to the strongest prior method sharing its basis. The predicted density also cuts SCF iterations by up to 68% with zero-shot transfer to unseen exchange-correlation functionals and recovers dipole and quadrupole moments to within a few percent of DFT references without any SCF calculation.

Abstract

Computational Materials Science

Explore similar work

Equivariant Flow Matching for Electron Density Prediction

Neural-Network Solutions to Real-Space Charge Density and Generalization

All-electron machine-learned orbital-free density functional for periodic CHNO systems