EMBER-2: emulating baryons from dark matter across cosmic time with deep modulation networks

Galaxy formation is a complex problem that connects large-scale cosmology with small-scale astrophysics over cosmic time-scales. Hydrodynamical simulations are the most principled approach to model galaxy formation, but have large computational costs. Recently, emulation techniques based on convolutional neural networks (CNNs) have been proposed to predict baryonic properties directly from dark matter simulations. The advantage of these emulators is their ability to capture relevant correlations, but at a fraction of the computational cost compared to simulations. However, training basic CNNs over large redshift ranges is challenging, due to the increasing non-linear interplay between dark matter and baryons paired with the memory inefficiency of CNNs. This work introduces EMBER-2, an improved version of the EMBER (EMulating Baryonic EnRichment) framework, to simultaneously emulate multiple baryon channels including gas density, velocity, temperature, and H i density over a large redshift range, from z=6 to z=0. EMBER-2 incorporates a context-based styling network paired with Modulated Convolutions for fast, accurate, and memory efficient emulation capable of interpolating the entire redshift range with a single CNN. Although EMBER-2 uses fewer than 1/6 the number of trainable parameters than the previous version, the model improves in every tested summary metric including gas mass conservation and cross-correlation coefficients. The EMBER-2 framework builds the foundation to produce mock catalogues of field level data and derived summary statistics that can directly be incorporated in future analysis pipelines. We release the source code at the official website https://maurbe.github.io/ember2/.