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

TL;DR

ChargeFlow is a novel flow-matching model that refines charge-conditioned atomic densities into accurate DFT electron densities, enabling efficient large-scale screening of charged materials with improved accuracy over baseline methods.

Contribution

This work introduces ChargeFlow, the first flow-matching refinement model for charge densities, trained on extensive DFT data, and demonstrates its effectiveness across diverse charged material systems.

Findings

Improves deformation-density error from 3.62% to 3.21%.

Enhances charge-response cosine similarity from 0.571 to 0.655.

Successfully performs Bader partitioning and electrostatic potential calculations.

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.