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

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 (8)

• 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 \ref{['tab:downstream_compare']} 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 \ref{['tab:downstream_compare']} 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.