Abstract

Next-generation astronomical surveys like the Legacy Survey of Space and Time (LSST) will deliver billions of galaxy observations crucial for understanding dark matter and dark energy. However, extracting reliable galaxy properties like redshifts from these data requires scalable computational approaches that can handle these enormous datasets while maintaining scientific rigour and reproducibility. We present pop-cosmos, a forward-modelling framework for photometric galaxy survey data that constrains population-level galaxy properties up to redshift 6. Galaxies are modelled as draws from a population prior over physical parameters (redshift, stellar mass, star formation history, dust properties), mapped to observed colors and brightness using neural emulators of complex astrophysical models—achieving 10000x speedups. We use simulation-based inference to calibrate this population prior on deep multi-wavelength data (COSMOS2020), training a diffusion model to match the statistical properties of real survey data. The resulting model helps us understand and probe various astrophysical and cosmological phenomena. Central to our framework is flowfusion, a general-purpose library for density estimation and generative modelling that implements state-of-the-art machine learning methods including diffusion models and flow-matching. I will demonstrate how our open-source toolkit enables reproducible results from our scientific applications and discuss ongoing work with the Kilo-Degree Survey in preparation for LSST .