Struct2ML: End-to-End Prediction of Electronic Charge Densities Using 3D U-Net

Abstract

We present Struct2ML, a scalable, low-latency volumetric inference pipeline that predicts three-dimensional electronic charge densities directly from atomic configurations. By leveraging a 3D U-Net architecture, our method provides a surrogate route to approximate charge-density fields without performing a full self-consistent Kohn–Sham calculation for each queried structure, mapping multichannel structural descriptors with explicit periodic-boundary handling to volumetric charge fields. Unlike graph-based approaches that often require complex basis functions to reconstruct continuous fields, our grid-based representation maintains a computational complexity ofO(N), offering a natural interface for large-scale periodic and molecular systems optimized for real-time analysis. Key contributions include (i) scalable infrastructure, a high-throughput disk-to-GPU streaming dataloader optimized for terabyte-scale HDF5 archives, ensuring efficient training on massive data sets; (ii) robust training protocol, a reproducible workflow featuring automated voxel normalization, physics-informed multiobjective loss functions, and rigorous checkpointing; and (iii) multi-domain validation, a comprehensive evaluation suite providing both voxel-wise quantitative metrics and qualitative spatial analysis tailored for charge-density comparison. By substantially reducing the inference time to the millisecond scale (1–10 ms), this framework enables real-time charge-density exploration that reproduces key spatial and topological features of reference density functional theory (DFT) densities with data set-dependent accuracy, serving as a high-throughput utility complementary to equivariant models that primarily target static, maximal accuracy.