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 (DFT) 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 9502 charged Materials Project-derived calculations and evaluated on an external 1671-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 1671 benchmark structures and high-fidelity electrostatic potentials, which position flow matching as a practical density-refinement strategy for charged materials.