Predicting Halo Mass Function Proxies from Merger Tree Distributions using a Hybrid GNN and Gaussian Mixture Model

Abstract: The dark matter halo mass function (HMF) is a fundamental cosmological probe, reflecting the number density of dark matter halos as a function of mass, and is intimately linked to the hierarchical assembly histories encoded in merger trees. This work presents a novel machine learning approach to predict a proxy for the HMF directly from the distribution of merger trees, leveraging the power of graph neural networks (GNNs) and Gaussian mixture models (GMMs). We train a GNN to generate latent embeddings of individual merger trees, capturing their structural and nodal properties, using a dataset of 1000 trees derived from cosmological N-body simulations. Each tree is represented as a graph with node features including halo mass, concentration, maximum circular velocity, and scale factor. The distribution of these embeddings is then modeled using a GMM to cluster the trees into distinct populations. Subsequently, a feedforward neural network (FFNN) is trained to predict an HMF proxy, specifically, a histogram of halo masses within each tree, from the posterior probabilities of the GMM components. Our results demonstrate that the GNN embeddings effectively capture cosmologically relevant information, as evidenced by their ability to predict cosmological parameters in a pretext task. Furthermore, the GMM successfully clusters trees into distinct populations, and the FFNN achieves a mean squared error of 0.000522 on the test set when predicting the HMF proxy. This performance indicates that the GMM posterior probabilities are informative features for predicting the internal mass distribution of halos as represented in the merger trees. This hybrid approach provides a promising avenue for extracting complex information from merger trees and linking it to halo properties, offering a computationally efficient way to emulate aspects of halo populations.

Subjects:	astro-ph.CO; cs.LG
Cite as:	PX:2508.00002