CNN+FoF: application of deep learning to the identification of dark matter haloes

ABSTRACT We present a deep-learning-based approach for identifying dark matter haloes in cosmological N-body simulations. Our framework consists of a volumetric convolutional neural network to classify individual simulation particles as either halo or non-halo members, followed by a highly optimized and parallelized friends-of-friends clustering algorithm that groups the classified halo members into distinct haloes. The training data comprise simulations generated using GADGET-4, with labels obtained with the ROCKSTAR halo finder. Our models incorporate two main halo mass definitions, $M_{200\mathrm{b}}$ and $M_{\text{vir}}$, with similar performance. For haloes defined by the ROCKSTAR$M_{200\mathrm{b}}$ criterion, the classification network demonstrated stable performance across multiple simulation resolutions. For the highest resolution, it achieved over 98 per cent across all primary performance metrics when identifying halo particles. Furthermore, the FoF algorithm yielded halo catalogues with a purity generally exceeding 99 per cent and a stable completeness of 94 per cent for masses above $5\times 10^{11} \, M_\odot$. Our pipeline recovered the centre-of-mass positions, velocities and halo masses with high fidelity, yielding a halo mass function consistent to within 5 per cent of the reference while faithfully reconstructing the internal density profiles. The primary objective of this study is to offer a faster and scalable alternative to conventional halo finders, achieving a speed-up of approximately one order of magnitude relative to ROCKSTAR. While traditional N-body simulations remain dominated by the time required to solve gravitational dynamics, this post-processing acceleration becomes essential when coupled with fast field-level emulators, offering a promising pathway towards efficient, end-to-end forward modelling in modern simulation-based inference.