Characterizing Structure Formation through Instance Segmentation

Dark matter haloes form from small perturbations to the almost homogeneous density field of the early universe. Although it is known how large these initial perturbations must be to form haloes, it is rather poorly understood how to predict which particles will end up belonging to which halo. However, it is this process that determines the Lagrangian shape of protohaloes and is therefore essential to understand their mass, spin and formation history. Here, we present a machine-learning framework to learn how the protohalo regions of different haloes emerge from the initial density field. This involves one neural network to distinguish semantically which particles become part of any halo and a second neural network that groups these particles by halo membership into different instances. This instance segmentation is done through the Weinberger method, in which the network maps particles into a pseudo-space representation where different instances can be distinguished easily through a simple clustering algorithm. Our model reliably predicts the masses and Lagrangian shapes of haloes object-by-object, as well as summary statistics like the halo-mass function. We find that our model extracts information close to optimal by comparing it to the degree of agreement between two N-body simulations with slight differences in their initial conditions. We publish our model open-source and suggest that it can be used to inform analytical methods of structure formation by studying the effect of systematic manipulations of the initial conditions.