Joint cosmological parameter inference and initial condition reconstruction with Stochastic Interpolants

In this work, we present a unified approach to cosmological parameter inference and initial condition reconstruction using Stochastic Interpolants. We apply this method to jointly reconstruct simulations of non-linear dark matter fields and infer simulator parameters, demonstrating its accuracy and scalability with dataset size. We show how the amortized learned distribution reproduces the posterior obtained with Hamiltonian Monte Carlo without the need for a differentiable forward model or explicit likelihood. Additionally, we introduce a flexible framework for controllable simulators that impose partial constraints, showcasing its potential in generating tailored simulations. This work provides a scalable and accurate approach for reconstructing initial conditions in cosmological simulations, with broad implications for upcoming galaxy surveys like DESI.