EDP Sciences logo
Subscriber Authentication Point
Search
Menu
Home All issues Volume 711 (July 2026) A&A, 711 (2026) A53 Full HTML
Open Access
Issue
A&A
Volume 711, July 2026
Article Number A53
Number of page(s) 11
Section Cosmology (including clusters of galaxies)
DOI https://doi.org/10.1051/0004-6361/202558143
Published online 02 July 2026

© The Authors 2026

Licence Creative CommonsOpen Access article, published by EDP Sciences, under the terms of the Creative Commons Attribution License (https://creativecommons.org/licenses/by/4.0), which permits unrestricted use, distribution, and reproduction in any medium, provided the original work is properly cited.

This article is published in open access under the Subscribe to Open model. This email address is being protected from spambots. You need JavaScript enabled to view it. to support open access publication.

1. Introduction

Within the standard model of cosmology, Lambda cold dark matter (ΛCDM) and its extensions, values of the cosmological parameters encode the fundamental properties of our Universe, ranging from the matter-energy content to the expansion history. Therefore, achieving sub-per cent level precision in measuring these parameters is essential for a detailed understanding of the origin and evolution of the Universe. To this end, the large-scale structure (LSS) of the Universe – the cosmic web traced by galaxies and halos – provides a rich and direct observable imprint of these parameters, making it a powerful probe for constraining cosmological parameters (see e.g. d’Amico et al. 2020; Colas et al. 2020; Uhlemann et al. 2020; Villaescusa-Navarro et al. 2020; Gualdi et al. 2021a; Valogiannis & Dvorkin 2022; Liu et al. 2022; Ajani et al. 2020, for recent studies). This is one of the primary science objectives for many new and upcoming surveys, for example CMB-S4 (Abazajian et al. 2022), the Roman Space Telescope (Lam et al. 2023), the Dark Energy Spectroscopic Instrument (DESI) (DESI Collaboration 2016), Euclid (Laureijs et al. 2011; Racca et al. 2019; Euclid Collaboration: Castro et al. 2023), the Large Synoptic Survey Telescope (LSST) (Ivezić et al. 2019), and J–PAS (Benitez et al. 2014).

To extract cosmological information from LSS data, traditional approaches have relied on summary statistics, usually the two-point correlation function or its Fourier space equivalent, the power spectrum P(k). While the two-point correlation function is computationally efficient and well-established, it captures, by definition, only the Gaussian aspect of clustering of matter or tracers. This makes the statistic suboptimal and inadequate on non-linear scales where the density field develops non-Gaussian features. In light of this, multiple approaches have been developed to go beyond the two-point function (see e.g. Marques et al. 2019; Villaescusa-Navarro et al. 2020; Gualdi et al. 2021b; Hahn et al. 2020; Friedrich et al. 2020; Giri & Smith 2022; Harnois-Déraps et al. 2021; Samushia et al. 2021; Naidoo et al. 2022; Bayer et al. 2021; Eickenberg et al. 2022). While higher-order statistics can, in principle, capture more of the information from the underlying fields, they often come at a significant computational cost. These considerations will play an important role in the deployment of these techniques to the vast datasets that will soon be available in cosmology.

In parallel, recent advances in machine learning (ML) have opened new avenues to analyse LSS data by leveraging neural networks to learn complex, high-dimensional features (Ravanbakhsh et al. 2017; Fluri et al. 2019; Makinen et al. 2022; Jeffrey et al. 2021; Gillet et al. 2019; Hortua 2021; Villanueva-Domingo & Villaescusa-Navarro 2022; Villanueva-Domingo et al. 2023; Hassan et al. 2022; Anagnostidis et al. 2022; Villaescusa-Navarro et al. 2022a; Cuesta-Lazaro & Mishra-Sharma 2024; Roncoli et al. 2023; Ho et al. 2024; Lee & Villaescusa-Navarro 2025). This represents an alternative route to the summary statistics in terms of capturing the total information in cosmological fields, thereby constraining the values of the parameters of interest. However, standard image-based ML approaches such as convolutional neural networks (CNNs) have been demonstrated to be sub-optimal for large but spatially sparse datasets. They suffer from information loss when mapping sparse galaxy or halo catalogues to dense grids. On the other hand, point cloud methods, including graph neural networks (GNNs) (de Santi et al. 2023; Cuesta-Lazaro & Mishra-Sharma 2024; Shao et al. 2023; Makinen et al. 2022; Lee & Villaescusa-Navarro 2025) and PointMLP variants (Anagnostidis et al. 2022; Chatterjee & Villaescusa-Navarro 2025), although directly applied to discrete data, face severe computational and memory constraints when applied to large galaxy catalogues with ∼105 objects.

For this paper we formulated a new field-level representation of discrete point datasets that converts them to continuous spatial maps by assigning to every point in space a value equal to the distance to the k-th NN in the original point dataset. This approach was inspired by the k-nearest neighbour cumulative distribution functions (kNN-CDFs), introduced in Banerjee & Abel (2021b), which have emerged as a promising cosmological summary statistic, offering an efficient and robust measure on discrete data that captures higher-order spatial correlations beyond two-point statistics. We fed the nearest neighbours (NN) distance maps created from halo fields of N-body simulations run at different cosmological parameters to train a neural network and study the possible constraints on cosmological parameters. We also studied the effects of combining the NN distance maps with the CDFs, using them as joint inputs to a hybrid neural network for cosmological parameter inference and to demonstrate tighter constraints than possible with each input individually. Our method represents a promising avenue for applying ML techniques effectively to large discrete datasets from large-scale cosmological surveys.

2. Nearest neighbour maps from simulation data

Banerjee & Abel (2021b) introduced the k-NN cumulative distribution function as a set of useful summary statistics for the clustering of discrete data points (e.g. a catalogue of halo or galaxy positions in a typical cosmological dataset). At any scale, r, the value of the kNN-CDF represents the fraction of all points (covering the entire space over which the clustering is to be measured) that contain at least k data points within a sphere of radius r. Using data structures like search trees, the kNN-CDFs can be computed quickly on 𝒪(N log N) time. Crucially, Banerjee & Abel (2021b) demonstrated that each kNN-CDF is related to different combinations of all N-point functions of the underlying clustering, making these statistics sensitive to all orders of beyond-Gaussian clustering. This sensitivity translates to tighter constraints on cosmological parameters than the two-point function while using the same datasets, while the quick compute time allows these statistics to be easily evaluated from large datasets.

Banerjee & Abel (2021a) extended the formalism to capture cross-correlations (at all orders) between discrete datasets, while Banerjee & Abel (2022) further extended the formalism to include cross-correlations between discrete point data and continuous maps. The application of similar ideas to measure auto-clustering and cross-clustering in sets of continuous maps, specifically weak lensing mass maps from galaxy surveys, was addressed in Anbajagane et al. (2023). These summary statistics have already been measured in the context of various datasets (Wang et al. 2022; Gupta & Banerjee 2024; Coulton et al. 2024; Zhou et al. 2025; Chand et al. 2025), while modelling these statistics as a function of cosmological parameters and different galaxy-halo connection models has also been explored (Banerjee et al. 2022; Yuan et al. 2023). The k-NN statistics performed well in the community-wide “Beyond-two-point challenge” and were able to recover the true cosmology with tight error bars Krause (2025).

Gangopadhyay et al. (2025) demonstrated that the kNN-CDFs and their derivatives have geometric interpretations. For example, the 1NN-CDF and its first three derivatives at some scale r encode the geometry of intersections of spheres of radius r centred on the data points. Specifically, the value of the CDF is proportional to the volume enclosed; the first derivative is proportional to the exposed area; the second derivative contains information about the angles of intersections of the spheres; and finally, the third derivative is related to the Euler characteristic of the resultant configuration. These connections also relate the kNN-CDFs geometrically to the germ-grain Minkowski Functionals (see e.g. Schmalzing et al. 1996). However, while the kNN-CDFs capture a great deal of the information about the clustering of a set of points, it is important to ask whether they can capture all the information, and if they do not, how the missing information can be accessed. A simple way to see that the kNN-CDFs may not contain all clustering information is to realise that the construction of the CDF, by definition, throws away spatial information. That is, the value of the CDF at some scale r retains no information about the spatial distribution of points contributing to that value. The question arises of whether they come from specific locations in the volume, or if they are close to being uniformly distributed over the entire volume.

This motivates the construction and analysis of the NN distance maps, where each point in space is assigned a value equal to the distance from that point to the k-th NN data point. This results in a smooth, continuous field-level representation of the original set of discrete data points. These maps, in their raw form, preserve spatial information about how fast or slowly the NN distances change in different directions, which, in turn, is directly related to the exact configuration of the original data points. In fact, for 1NN distance maps in three dimensions, the points where the field has a maximum in one direction lie on the face of the Voronoi tessellation defined by the data points. That is, they are equidistant from two data points. Points with maxima along two directions define the Voronoi edges and are equidistant from three data points. Finally, maxima in this field, where the gradient vanishes, correspond to Voronoi nodes, or points that are equidistant from four data points. This is demonstrated in two-dimensional slices in the left panel of Fig. 1. Similarly, one can construct maps for larger k; the right-hand panel of Fig. 1 shows the 2D maps for k = 4. These ideas, especially for the first NN distances, have been explored in the field of image processing, and are referred to as the distance transform (see e.g. Jones et al. 2006).

Thumbnail: Fig. 1. Refer to the following caption and surrounding text. Fig. 1.

2D slice of the first (left) and fourth (right) nearest neighbour distance maps for one of the simulations in the Quijote simulation suites used in this study. Each pixel is coloured by the distance from the pixel to the nearest data point (left panel) and by the distance to the fourth nearest neighbour data point (right panel). As can be seen, this converts the discrete dataset into a smooth, continuous map. The colour bar represents the distance (in Gpc/h) from the halos. We note that these maps are only for the purpose of visualisation. They were produced with 2562 random query points in a 256 × 256 2D grid, whereas the actual maps used in this study were produced with 104 random query points in a 100 × 100 2D grid, as mentioned in Sect. 2.1.

Neural networks are particularly well-suited to analysing the information in these maps. The traditional failing of CNNs and related techniques for sparse data can now be circumvented, as the mapping converts the data to a densely sampled continuous field. For this study, we applied kNN statistics (CDFs as well as maps) on the halo catalogues obtained from the Quijote simulation suite1. In particular, we focused on the 2000 Latin Hypercube realisations, each containing 5123 dark matter simulation particles at z = 0. The cosmological parameters of these simulations vary in the range

0.1 ≤ Ω m ≤ 0.5 , 0.03 ≤ Ω b ≤ 0.07 , 0.5 ≤ h ≤ 0.9 , 0.8 ≤ n s ≤ 1.2 , 0.6 ≤ σ 8 ≤ 1.0 . Mathematical equation: $$ \begin{aligned} 0.1&\le \Omega _{\rm m} \le 0.5,\\ 0.03&\le \Omega _{\rm b} \le 0.07,\\ 0.5&\le h \le 0.9,\\ 0.8&\le n_s \le 1.2,\\ 0.6&\le \sigma _8 \le 1.0. \end{aligned} $$

In this work, we concentrate on just two of these parameters: Ωm and σ8. Further, we select 105 most massive halos2 from the halo catalogues from each of these simulations and compute the NN distance maps and CDFs as mentioned below3. This number cut translates to a mass cut of around 5 × 1013 h−1 M⊙ (Banerjee & Abel 2021a), i.e. about five times more massive than the least massive halos. This choice ensures that our result should not depend on the resolution of the simulations.

2.1. Generation of NN maps

For creating the maps, we start with the 3D positions of the 105 most massive halos from each simulation:

  1. We first divide the 3D data into 100 2D slices along the z-axis with an inter-plane distance of 10 Mpc.

  2. For each of these slices, we generate 104 query points in a 100 × 100 regular grid4 and calculate distances of first, second, third, and fourth NNs from those query points5.

  3. Once obtained, these distances are used to create 2D maps of kNN distances for the query points. Figure 1 shows an example of the first and fourth NN distance maps, with the positions of halos in that slice overlaid as white points. We note that the maps shown in Fig. 1 are produced with 2562 random query points in a 256 × 256 grid and are made only for the purpose of visualisation.

  4. Finally, for each simulation, we randomly select ten maps corresponding to each NN distance to be used in the subsequent analysis. As we are taking ten random slices from each simulation, we have a total of 20 000 maps.

2.2. Generation of the NN CDFs

To obtain NN CDFs from the downsampled halo catalogues of the 105 most massive halos, we used the following method:

  1. We generate 106 random query points in 3D within the simulation box. By taking 106 random query points, we ensure that the number of query points is more than the discrete data points.

  2. Using scipy.spatial.KDTree, we compute the distances to the first, second, third, and fourth NNs for these query points.

  3. We then construct cumulative distribution functions (CDFs) of the NN distances, binned into 50 intervals of width 1 Mpc h−1. In the left panel of Fig. 2, we illustrate the shapes of these functions obtained from one of our simulations6. For better visualisation, especially when the value of the CDFs is ∼1, we depict the peaked CDF distribution for the first, second, third, and fourth NNs in the right-hand panel of the same figure. Following Banerjee & Abel (2021b), we define the Peaked CDF as

    PCDF ( r ) = { CDF ( r ) , if CDF ( r ) ≤ 0.5 , 1 − CDF ( r ) , if CDF ( r ) > 0.5 . Mathematical equation: $$ \begin{aligned} \mathrm{PCDF} (r) = {\left\{ \begin{array}{ll} \mathrm{CDF} (r),&\text{ if} \mathrm{CDF} (r) \le 0.5, \\ 1 - \mathrm{CDF} (r),&\text{ if} \mathrm{CDF} (r) > 0.5. \end{array}\right.} \end{aligned} $$(1)

Thumbnail: Fig. 2. Refer to the following caption and surrounding text. Fig. 2.

CDF (left panel) and peaked CDF (right panel) for 1NN (orange), 2NN (red), 3NN (magenta), and 4NN (blue) corresponding to one of the Quijote simulations in this study.

3. Machine learning model

We explored three ML models: (1) Map-only: a standard ResNet model that learns from the 2D slices of NN distance maps, (2) CDF-only: a fully connected NN that learns from the NN CDFs, and (3) Map+CDF: a hybrid model that takes as input the combination of the NN CDFs and NN distance maps. We describe them below in more detail.

3.1. Architecture

Map-only: We employed the RESNET-18 (He et al. 2016) architecture, a widely used convolutional neural network designed for image-based tasks, to infer cosmological parameters from 2D NN distance maps. The main novelty of the residual network (ResNet) is in its use of skip connections that effectively mitigate the vanishing gradient problem and enable the training of very deep architectures. In our case, each input map corresponds to a 2D grid of NN distances computed from dark matter or halo slices, with separate channels for different NN distance maps (e.g. 1NN, 2NN, 3NN, 4NN). The RESNET-18 model processes these maps through a series of convolutional, batch normalisation, ReLU, and residual blocks, ultimately reducing them to a feature vector that is passed through a fully connected layer to predict the cosmological parameters.

CDF-only: For each simulation, once we calculate the CDFs of the first, second, third, and fourth NN distances, each representing a 1D array with 48 bins (uniform bin width of 1 h−1 Mpc), they are then concatenated into a single 1D array of length 192 (4 × 48), which serves as the input for the fully connected neural network containing multiple layers of perceptrons (MLPs). We use ReLU as the non-linear activation function in this model. The number of layers, neurons, learning rate, weight decay, and dropout are hyperparameters optimised with OPTUNA (Akiba et al. 2019) in over 100 trials, as mentioned in Table A.1.

Map+CDF: The concept of hybrid or multimodal neural network, combination of different architecture trained on different types of data (e.g. images, texts) in order to maximise the performance of a neural network, has shown a lucrative performance in different fields of artificial intelligence (Chen et al. 2023; Azevedo et al. 2024; Zeng et al. 2022; Shi et al. 2019; Halbouni et al. 2022; Siraj & Ahad 2020; Demiss & Elsaigh 2024; Dattilo et al. 2019) and in cosmology (Dattilo et al. 2019; Ntampaka et al. 2020; Lucas Makinen et al. 2025; Mobina Hosseini & Soleimanpour Salmasi 2025; Bairagi & Wandelt 2025). While Ntampaka et al. (2020) used the hybrid network combining the 2D images and two-point power spectrum obtained from a simulated galaxy density field and showed that the hybrid network performs better than the individual neural network, Mobina Hosseini & Soleimanpour Salmasi (2025) employed a combination of a convolutional neural networks and recurrent neural networks on 21 cm brightness temperatures and reconstructed a 21 cm global signal with a prediction accuracy of 99.93%.

Motivated by the aforementioned works, we implemented a hybrid neural network architecture for this study. The model integrates a RESNET-18 backbone with fully connected layers that incorporate summary statistics, i.e. the NN cumulative distribution functions (CDFs). We initialise the RESNET-18 model after modifying the input convolutional layer to accept multiple input channels, corresponding to the different NN distances (e.g. 1NN to 4NN). The fully connected classification head of ResNet is removed and replaced with an identity layer to extract the learned feature representation from the input maps. The length of the learned features, i.e. the number of output channels from the ResNet (as shown in the figure) is kept to 512. These features are then concatenated with the 1D summary statistics vector (of length 192, corresponding to the concatenated CDFs from 1NN–4NN). The combined feature vector (of dimension 704, i.e. 512 from ResNet and 192 from CDFs) is passed through an inference block consisting of a multi-layer perceptron followed by a final output layer that predicts the cosmological parameters as shown in Fig. 3. The number of layers and neurons in the multi-layer perceptron of the network, as well as the learning rate and weight decay, are hyperparameters optimised with OPTUNA (Akiba et al. 2019) in over 100 trials, as mentioned in Table A.17. The architecture allows the model to jointly learn from spatial features in the NN distance maps and from global statistical summaries.

Thumbnail: Fig. 3. Refer to the following caption and surrounding text. Fig. 3.

Hybrid network in this study. The NN distance maps are used as input to the ResNet block. The output of the ResNet is then concatenated with the NN CDFs, and the merged input then passes through the inference blocks (containing several linear, ReLU, and dropout layers) to predict the mean and standard deviation of the inferred cosmological parameters. The values in brackets show the dimension of the tensor in different stages of the architecture. Here B denotes the batch dimension.

3.2. Loss function

All the above-mentioned models are trained to perform likelihood-free inference on the value of the cosmological parameters. For each parameter, the model predicts the marginal posterior mean (μ) and standard deviation (σ), defined as

μ ( P ) = ∫ θ p ( θ | P ) θ d θ , Mathematical equation: $$ \begin{aligned} \mu (\mathcal{P} )&= \int _{\theta } p(\theta | \mathcal{P} ) \theta \,\mathrm{d}\theta , \end{aligned} $$(2)

σ 2 ( P ) = ∫ θ p ( θ | P ) ( θ − μ ) 2 d θ , Mathematical equation: $$ \begin{aligned} \sigma ^2(\mathcal{P} )&= \int _{\theta } p(\theta | \mathcal{P} ) (\theta - \mu )^2\,\mathrm{d}\theta , \end{aligned} $$(3)

where 𝒫 represents the input, which could be CDFs for CDF-only, Maps for Map-only, or a combination of CDF and Maps for Map+CDF. To achieve this, we minimise the following loss function (Jeffrey & Wandelt 2020; Villaescusa-Navarro et al. 2022b)

L = ∑ i = 1 N θ log ( ∑ j ∈ batch ( θ i , j − μ i , j ) 2 ) + ∑ i = 1 N θ log ( ∑ j ∈ batch ( ( θ i , j − μ i , j ) 2 − σ i , j 2 ) 2 ) , Mathematical equation: $$ \begin{aligned} \mathcal{L}&= \sum _{i = 1}^{N_{\theta }}\log \left(\sum _{j\in \mathrm{batch}}\left(\theta _{i,j} - \mu _{i,j}\right)^2\right)\nonumber \\&\quad + \sum _{i = 1}^{N_{\theta }} \log \left(\sum _{j\in \mathrm{batch}}\left(\left(\theta _{i, j} - \mu _{i,j}\right)^2 - \sigma _{i,j}^2 \right)^2\right), \end{aligned} $$(4)

where θi, j, μi, j, and σi, j represent the true, inferred mean, and inferred standard deviation of the parameter i for the sample j. Further, Nθ is the total number of cosmological parameters; in this study there are two: Ωm and σ8. We employ this loss function to directly obtain the estimates of moments of the marginalised distribution of all parameters without calculating the posterior density.

3.3. Training procedure

We split the available input data into training (80%), validation (10%), and testing (10%). For all the ML architectures, we used the Adam optimiser (Kingma & Ba 2014)8. We used a batch size of 32 and trained for 300 epochs. Training was performed on a single NVIDIA A100 GPU and took approximately 30–120 min per training run, depending on the specific ML model used.

3.4. Validation metrics

For each cosmological parameter, we employed four statistics to quantify the accuracy and precision of our models on the test dataset, given by the following:

  • The mean relative error, ϵ, defined as

    ϵ = 1 N ∑ j = 1 N | θ j − μ j | θ j , Mathematical equation: $$ \begin{aligned} \epsilon = \frac{1}{N} \sum _{j = 1}^{N} \frac{\left|\theta _{j} - \mu _{j}\right|}{\theta _{j}}, \end{aligned} $$(5)

    where N denotes the size of the test dataset. A smaller value of ϵ indicates greater precision of the network;

  • The coefficient of determination, R2, defined as

    R 2 = 1 − ∑ j = 1 N ( θ j − μ j ) 2 ∑ j = 1 N ( θ j − θ ¯ ) 2 , Mathematical equation: $$ \begin{aligned} R^2 = 1 - \frac{\sum _{j = 1}^{N} (\theta _{j} - \mu _{j})^2}{\sum _{j = 1}^{N} (\theta _{j} - \overline{\theta })^2}, \end{aligned} $$(6)

    where θ ¯ = 1 N ∑ j = 1 N θ j Mathematical equation: $ \overline{\theta}=\frac{1}{N}\sum_{j= 1}^{N} \theta_{j} $. A network with higher accuracy leads to an R2 value closer to 1.

  • The mean squared error, MSE, defined as

    MSE = 1 N ∑ j = 1 N ( θ j − μ j ) 2 . Mathematical equation: $$ \begin{aligned} \mathrm{MSE} = \frac{1}{N}\sum _{j = 1}^{N} (\theta _{j} - \mu _{j})^2. \end{aligned} $$(7)

    An accurate network indicates a smaller mean squared error;

  • The χ2 value, defined as

    χ 2 = 1 N ∑ j = 1 N ( θ j − μ j ) 2 σ j 2 · Mathematical equation: $$ \begin{aligned} \chi ^2=\frac{1}{N}\sum _{j = 1}^{N} \frac{(\theta _{j} - \mu _{j})^2}{\sigma _{j}^2}\cdot \end{aligned} $$(8)

    A value of χ2 close to 1 indicates that network errors are calibrated correctly.

4. Results

We now present the results obtained from the ML models studied in this work. Table 1 summarises the results for the different ML models. Figure 4 shows the constraints we derive on Ωm and σ8 from different ML models. We now describe the main findings in the different models:

  • CDF-only: From Table 1, we can see that the CDF-only scenario produces a reasonable value of relative error and R2 for Ωm. The prediction for σ8 is poor with large error bars.

  • Map-only: When we train on NN distance maps using ResNet, we find that our models perform more poorly than the CDF-only scenario in all the validation metrics. In fact, the network fails in predicting any sensible constraints on σ8. Although we note that the Map-only model is not a fully 3D statistic (we generate 2D slices of the NN maps as mentioned in Sect. 2.1), and therefore, a comparison with CDF-only should be made with caution, it is still quite contrary to our expectations. In fact, we argue in Sect. 2.1 that NN distance maps would contain more information than the CDFs as they preserve the spatial information. Interestingly, we note that a very similar behaviour is seen in Chatterjee & Villaescusa-Navarro (2025) and Makinen et al. (2022) when used on the Qujote simulation, although it uses a completely different network architecture. One possible explanation could be the limited number of training datasets (maps corresponding to 1600 different parameter combinations). It is plausible that due to the limited number of training datasets, the network is unable to break the bias-σ8 degeneracy. When we use the Map-only method on the maps produced from DM particles (Appendix A) rather than halos, we find that this method outperforms the CDF-only methods, as expected from an unbiased tracer. Since halo bias modifies the amplitude of large-scale statistics in a way that is degenerate with the effect of σ8, a network trained on NN distance maps may struggle to disentangle the two for a limited amount of training data.

  • Map+CDF: The hybrid network model, where NN CDFs along with the NN distance maps are used, performs the best. It results in an R2 score of 0.80 and 0.93 for Ωm and σ8, respectively. Further, we achieve a relative error of ∼15% and ∼3% for Ωm and σ8, respectively.

We note that in all the scenarios, the value of the χ2 is ∼0.8. However, this is reasonable for a parameter inference study with ML.

Thumbnail: Fig. 4. Refer to the following caption and surrounding text. Fig. 4.

Performance of different models when trained to predict likelihood-free inference on the values of Ωm (left column) and σ8 (right column) in three scenarios: Top row: CDF-only; Middle row: Map-only; bottom row: Map+CDF. The values for different validation metrics are given in the legend. As can be seen, the Map-only scenario (middle panel) performs worse than the CDF-only (top panel) scenario. Further, the Map+CDF model performs the best across all the validation metrics.

Table 1.

Summary of validation metrics for different configurations.

4.1. Comparison with the two-point correlation function

As a benchmark, we have demonstrated how the constraining power of the two-point correlation function (2ptCF), ξ(r), compares with the NN CDFs in Fig. 5. We employed Pylians9 to compute the 2ptCF on scales smaller than 50 h−1 Mpc (to ensure the same range of scales as for the kNN statistics) on the same halo catalogues of the 105 most massive halos that were used to calculate the kNN statistics. We used these as the input to a fully connected neural network with LeakyReLU activation functions to obtain constraints on the Ωm and σ8. The number of layers, neurons of the fully connected neural network, along with the weight decay and learning rate of the Adam optimiser (Kingma & Ba 2014) are kept as hyperparameters and optimised using OPTUNA. The results obtained are shown in Fig. 5. As is evident, the performance of the CDFs is significantly better (by almost a factor of two in some of the validation metrics, for example R2) compared to the 2ptCF for both the parameters. This is also expected since Banerjee & Abel (2021a) found a similar conclusion, although using the Fisher analysis.

Thumbnail: Fig. 5. Refer to the following caption and surrounding text. Fig. 5.

Comparison between ξ(r)-only and CDF-only. As shown, CDF-only performs much better compared to ξ(r)-only, as expected from (Banerjee & Abel 2021b).

5. Related works and comparison

In the following, we compare the performance of our hybrid network with results obtained by the existing models that also utilise halo and/or galaxy positions. While several studies have demonstrated the use of diverse machine learning techniques to constrain cosmological parameters, we limit our comparison to those based on the Quijote dataset as others (e.g. Shao et al. 2023; Anagnostidis et al. 2022) use different datasets, and therefore do not permit a meaningful direct comparison.

In Ho et al. (2024), the positions of ∼10 000 halos were used as the input of a GNN to infer the cosmological parameters. In the absence of any values of the accuracy metrics in their paper, we visually estimated (from Fig. 7 in their paper) that their relative error for Ωm is ∼20%, which is poorer compared to our hybrid study. Their estimate of σ8 is very close to our estimate.

In Cuesta-Lazaro & Mishra-Sharma (2024) the authors used the position of the 5000 most massive halos as the input features for their generative modelling, and they obtained a mean relative error of ∼5% and ∼3% on Ωm and σ8, respectively. While their constraints on Ωm are better than our study, the constraints on σ8 from our study are comparable with their constraints. Further, caution must be used as the learned likelihood from their model is not well calibrated for Ωm, whereas our study is not affect by any such issues.

In Chatterjee & Villaescusa-Navarro (2025), the authors found that a point cloud-based network trained on the position of the 8192 most massive halos constrained Ωm with a relative error of 15.5%, which is very similar to the results of this study (15.3%). Interestingly, this point cloud-based study (when trained on the position of the halos) failed to infer σ8 completely, whereas the present study has been able to recover σ8 with excellent accuracy across all the validation metrics. We would like to emphasise that the training time and memory requirements in this study are significantly lower than those reported in Chatterjee & Villaescusa-Navarro (2025). For instance, training with 8192 halos and 32 neighbours at a batch size of 32 took approximately 2 days in their study, whereas our current approach completes training in only about 2 h, which means it is 24 times more efficient than the point cloud-based method. Moreover, while Chatterjee & Villaescusa-Navarro (2025) required 64 GB of GPU memory for a batch size of 32, our current method requires only a 16 GB GPU. The main reason for this is that we use 2D slices of the NN maps in this study, whereas the point cloud-based method works in three dimensions.

In Lee & Villaescusa-Navarro (2025), the authors used topological neural networks trained on the positions of the 5000 most massive halos and recovered Ωm and σ8 with relative errors of 15.33% and 3.69% with their best performing FullTNN network. Both of their constraints are similar to the constraints we find in this analysis.

In Huang et al. (2025), the authors used a GNN-based neural network and performed cosmological parameter inference analysis on the Big Sobol Sequence of Quijote simulation suites. When using positions as the only feature for the halos, they obtained an R2 value of 0.80 and 0.77 for Ωm and σ8, respectively. In our hybrid model, while the R2 value for σ8 is better than their reported R2 values, the R2 value for Ωm is the same as theirs. Moreover, our hybrid model takes 2 h of GPU time, whereas their GNN-based model takes 24 h on TPU.

6. Discussion and conclusions

In this study we have demonstrated the use of kNN statistics as inputs to neural network models for inferring cosmological parameters from halo catalogues generated by the Latin Hypercube dataset of the Quijote simulation suite. Specifically, we employed a hybrid neural network architecture that leverages both nearest neighbour distance maps and nearest neighbour cumulative distribution functions (CDFs), successfully recovering Ωm and σ8 with excellent accuracy.

We conducted a detailed comparison with all existing studies that use the Quijote simulations for cosmological parameter inference and showed that our approach achieves very high accuracy, while remaining highly computationally efficient.

The advantages of our method over previous approaches, such as those based on GNNs, topological neural networks, or point cloud networks, are twofold. First, it provides a far more efficient framework for cosmological parameter inference, substantially reducing computational cost relative to GNN- or PointMLP-based methods. Second, it uniquely enables field-level representations from extremely large halo samples with remarkable efficiency. These features make our model particularly well-suited for forthcoming galaxy surveys, which are expected to produce unprecedentedly large datasets.

One limitation of this study is that NN distance maps or CDFs cannot incorporate the information about the mass and velocity of the halos. In Chatterjee & Villaescusa-Navarro (2025) we find that mass and velocity are very crucial in putting tighter constraints on the cosmological parameters. To overcome this, in future work, we plan to include the peculiar velocities of the halos in one of the spatial directions (using redshift space distortion) in the simulations before making the images. It may even provide us with a stronger correlation between the velocity and the matter field. It is therefore possible for the neural network to pick up the subtle anisotropies being produced by these redshift space distortions, providing even tighter cosmological constraints.

Acknowledgments

The work of AC was supported by the European Union’s Horizon Europe research and innovation programme under the Marie Skłodowska-Curie Postdoctoral Fellowship HORIZON-MSCA-2023-PF-01, grant agreement No 101151693 (LUPCOS). AB’s work was partially supported by the Startup Research Grant (SRG/2023/000378) from the Science and Engineering Research Board (SERB), India. This work was also supported by U.S. Department of Energy grant DE-AC02-76SF00515 to SLAC National Accelerator Laboratory managed by Stanford University. The work of FVN is supported by the Simons Foundation. The authors acknowledge the PARAM Brahma Facility under the National Supercomputing Mission, Government of India, at the Indian Institute of Science Education and Research, Pune, for providing the computing resources for this work. The ML architecture developed in this work is implemented in PYTORCH (Paszke et al. 2019). The authors acknowledge the use of CHATGPT for refining the text at the final stage of the manuscript.

References

  1. Abazajian, K., Abdulghafour, A., Addison, G. E., et al. 2022, arXiv e-prints [arXiv:2203.08024] [Google Scholar]
  2. Ajani, V., Peel, A., Pettorino, V., et al. 2020, Phys. Rev. D, 102, 103531 [Google Scholar]
  3. Akiba, T., Sano, S., Yanase, T., Ohta, T., & Koyama, M. 2019, in The 25th ACM SIGKDD International Conference on Knowledge Discovery& Data Mining, 2623 [Google Scholar]
  4. Anagnostidis, S., Thomsen, A., Kacprzak, T., et al. 2022, ArXiv e-prints [arXiv:2211.12346] [Google Scholar]
  5. Anbajagane, D., Chang, C., Banerjee, A., et al. 2023, MNRAS, 526, 5530 [Google Scholar]
  6. Azevedo, B. F., Rocha, A. M. A., & Pereira, A. I. 2024, Mach. Learn., 113, 4055 [Google Scholar]
  7. Bairagi, A., & Wandelt, B. 2025, ArXiv e-prints [arXiv:2509.03165] [Google Scholar]
  8. Banerjee, A., & Abel, T. 2021a, MNRAS, 504, 2911 [NASA ADS] [CrossRef] [Google Scholar]
  9. Banerjee, A., & Abel, T. 2021b, MNRAS, 500, 5479 [Google Scholar]
  10. Banerjee, A., & Abel, T. 2022, MNRAS, 519, 4856 [Google Scholar]
  11. Banerjee, A., Kokron, N., & Abel, T. 2022, MNRAS, 511, 2765 [Google Scholar]
  12. Bayer, A. E., Villaescusa-Navarro, F., Massara, E., et al. 2021, ApJ, 919, 24 [NASA ADS] [CrossRef] [Google Scholar]
  13. Benitez, N., Dupke, R., Moles, M., et al. 2014, ArXiv e-prints [arXiv:1403.5237] [Google Scholar]
  14. Chand, E., Banerjee, A., Foreman, S., & Villaescusa-Navarro, F. 2025, MNRAS, 538, 2204 [Google Scholar]
  15. Chatterjee, A., & Villaescusa-Navarro, F. 2025, ApJ, 985, 132 [Google Scholar]
  16. Chen, X., Yang, Q., Wu, J., Li, H., & Tan, K. C. 2023, ArXiv e-prints [arXiv:2305.16594] [Google Scholar]
  17. Colas, T., d’Amico, G., Senatore, L., Zhang, P., & Beutler, F. 2020, J. Cosmol. Astropart. Phys., 2020, 001 [Google Scholar]
  18. Coulton, W. R., Abel, T., & Banerjee, A. 2024, MNRAS, 534, 1621 [Google Scholar]
  19. Cuesta-Lazaro, C., & Mishra-Sharma, S. 2024, Phys. Rev. D, 109, 123531 [Google Scholar]
  20. d’Amico, G., Gleyzes, J., Kokron, N., et al. 2020, J. Cosmol. Astropart. Phys., 2020, 005 [CrossRef] [Google Scholar]
  21. Dattilo, A., Vanderburg, A., Shallue, C. J., et al. 2019, AJ, 157, 169 [NASA ADS] [CrossRef] [Google Scholar]
  22. de Santi, N. S. M., Shao, H., Villaescusa-Navarro, F., et al. 2023, ApJ, 952, 69 [NASA ADS] [CrossRef] [Google Scholar]
  23. Demiss, B. A., & Elsaigh, W. A. 2024, Eng. Res. Express, 6, 032102 [Google Scholar]
  24. DESI Collaboration (Aghamousa, A., et al.) 2016, ArXiv e-prints [arXiv:1611.00036] [Google Scholar]
  25. Eickenberg, M., Allys, E., Moradinezhad Dizgah, A., et al. 2022, ArXiv e-prints [arXiv:2204.07646] [Google Scholar]
  26. Euclid Collaboration (Castro, T., et al.) 2023, A&A, 671, A100 [CrossRef] [EDP Sciences] [Google Scholar]
  27. Fluri, J., Kacprzak, T., Lucchi, A., et al. 2019, Phys. Rev. D, 100, 063514 [Google Scholar]
  28. Friedrich, O., Uhlemann, C., Villaescusa-Navarro, F., et al. 2020, MNRAS, 498, 464 [NASA ADS] [CrossRef] [Google Scholar]
  29. Gangopadhyay, K., Banerjee, A., & Abel, T. 2025, MNRAS, 543, 3409 [Google Scholar]
  30. Gillet, N., Mesinger, A., Greig, B., Liu, A., & Ucci, G. 2019, MNRAS, 484, 282 [NASA ADS] [Google Scholar]
  31. Giri, U., & Smith, K. M. 2022, J. Cosmol. Astropart. Phys., 2022, 028 [CrossRef] [Google Scholar]
  32. Gualdi, D., Gil-Marín, H., & Verde, L. 2021a, J. Cosmol. Astropart. Phys., 2021, 008 [Google Scholar]
  33. Gualdi, D., Novell, S., Gil-Marín, H., & Verde, L. 2021b, J. Cosmol. Astropart. Phys., 2021, 015 [Google Scholar]
  34. Gupta, K. R., & Banerjee, A. 2024, MNRAS, 531, 4619 [Google Scholar]
  35. Hahn, C., Villaescusa-Navarro, F., Castorina, E., & Scoccimarro, R. 2020, J. Cosmol. Astropart. Phys., 2020, 040 [Google Scholar]
  36. Halbouni, A., Gunawan, T. S., Habaebi, M. H., et al. 2022, IEEE Access, 10, 99837 [CrossRef] [Google Scholar]
  37. Harnois-Déraps, J., Martinet, N., Castro, T., et al. 2021, MNRAS, 506, 1623 [CrossRef] [Google Scholar]
  38. Hassan, S., Villaescusa-Navarro, F., Wandelt, B., et al. 2022, ApJ, 937, 83 [NASA ADS] [CrossRef] [Google Scholar]
  39. He, K., Zhang, X., Ren, S., & Sun, J. 2016, in 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR, 1) [Google Scholar]
  40. Ho, M., Bartlett, D. J., Chartier, N., et al. 2024, Open J. Astrophys., 7, 54 [NASA ADS] [CrossRef] [Google Scholar]
  41. Hortua, H. J. 2021, ArXiv e-prints [arXiv:2112.11865] [Google Scholar]
  42. Huang, N., Stiskalek, R., Lee, J.-Y., et al. 2025, ArXiv e-prints [arXiv:2507.03707] [Google Scholar]
  43. Ivezić, Ž., Kahn, S. M., Tyson, J. A., et al. 2019, ApJ, 873, 111 [Google Scholar]
  44. Jeffrey, N., & Wandelt, B. D. 2020, ArXiv e-prints [arXiv:2011.05991] [Google Scholar]
  45. Jeffrey, N., Alsing, J., & Lanusse, F. 2021, MNRAS, 501, 954 [Google Scholar]
  46. Jones, M., Baerentzen, J., & Sramek, M. 2006, IEEE Trans. Visualization Comput. Graphics, 12, 581 [Google Scholar]
  47. Kingma, D. P., & Ba, J. 2014, arXiv e-prints [arXiv:1412.6980] [Google Scholar]
  48. Krause, E., et al. 2025, ApJ, 990, 99 [Google Scholar]
  49. Lam, C. Y., Abrams, N., Andrews, J., et al. 2023, ArXiv e-prints [arXiv:2306.12514] [Google Scholar]
  50. Laureijs, R., Amiaux, J., Arduini, S., et al. 2011, ArXiv e-prints [arXiv:1110.3193] [Google Scholar]
  51. Lee, J.-Y., & Villaescusa-Navarro, F. 2025, ApJ, 989, 47 [Google Scholar]
  52. Liu, W., Jiang, A., & Fang, W. 2022, J. Cosmol. Astropart. Phys., 2022, 045 [Google Scholar]
  53. Lucas Makinen, T., Heavens, A., Porqueres, N., et al. 2025, J. Cosmol. Astropart. Phys., 2025, 095 [Google Scholar]
  54. Makinen, T. L., Charnock, T., Lemos, P., et al. 2022, Open J. Astrophys., 5, 18 [NASA ADS] [CrossRef] [Google Scholar]
  55. Marques, G. A., Liu, J., Zorrilla Matilla, J. M., et al. 2019, J. Cosmol. Astropart. Phys., 2019, 019 [CrossRef] [Google Scholar]
  56. Mobina Hosseini, S., & Soleimanpour Salmasi, B. 2025, ArXiv e-prints [arXiv:2508.05842] [Google Scholar]
  57. Naidoo, K., Massara, E., & Lahav, O. 2022, MNRAS, 513, 3596 [NASA ADS] [CrossRef] [Google Scholar]
  58. Ntampaka, M., Eisenstein, D. J., Yuan, S., & Garrison, L. H. 2020, ApJ, 889, 151 [NASA ADS] [CrossRef] [Google Scholar]
  59. Paszke, A., Gross, S., Massa, F., et al. 2019, arXiv e-prints [arXiv:1912.01703] [Google Scholar]
  60. Racca, G. D., Laureijs, R., Stagnaro, L., et al. 2016, in Space Telescopes and Instrumentation 2016: Optical, Infrared, and Millimeter Wave, eds. H. A. MacEwen, G. G. Fazio, M. Lystrup, et al., SPIE Conf. Ser., 9904, 99040O [NASA ADS] [Google Scholar]
  61. Ravanbakhsh, S., Oliva, J., Fromenteau, S., et al. 2017, ArXiv e-prints [arXiv:1711.02033] [Google Scholar]
  62. Roncoli, A., Ćiprijanović, A., Voetberg, M., Villaescusa-Navarro, F., & Nord, B. 2023, ArXiv e-prints [arXiv:2311.01588] [Google Scholar]
  63. Samushia, L., Slepian, Z., & Villaescusa-Navarro, F. 2021, MNRAS, 505, 628 [NASA ADS] [CrossRef] [Google Scholar]
  64. Schmalzing, J., Kerscher, M., & Buchert, T. 1996, Proc. Int. Sch. Phys. Fermi, 132, 281 [Google Scholar]
  65. Shao, H., Villaescusa-Navarro, F., Villanueva-Domingo, P., et al. 2023, ApJ, 944, 27 [Google Scholar]
  66. Shi, X., Wang, T., Wang, L., Liu, H., & Yan, N. 2019, in Asia-Pacific Signal and Information Processing Association Annual Summit and Conference (APSIPA ASC) (IEEE), 939 [Google Scholar]
  67. Siraj, M. S., & Ahad, M. 2020, in 2020 Joint 9th International Conference on Informatics, Electronics& Vision (ICIEV) and 2020 4th International Conference on Imaging, Vision& Pattern Recognition (icIVPR) (IEEE), 1 [Google Scholar]
  68. Uhlemann, C., Friedrich, O., Villaescusa-Navarro, F., Banerjee, A., & Codis, S. 2020, MNRAS, 495, 4006 [NASA ADS] [CrossRef] [Google Scholar]
  69. Valogiannis, G., & Dvorkin, C. 2022, Phys. Rev. D, 105, 103534 [Google Scholar]
  70. Villaescusa-Navarro, F., Hahn, C., Massara, E., et al. 2020, ApJS, 250, 2 [CrossRef] [Google Scholar]
  71. Villaescusa-Navarro, F., Ding, J., Genel, S., et al. 2022a, ApJ, 929, 132 [NASA ADS] [CrossRef] [Google Scholar]
  72. Villaescusa-Navarro, F., Genel, S., Anglés-Alcázar, D., et al. 2022b, ApJS, 259, 61 [NASA ADS] [CrossRef] [Google Scholar]
  73. Villanueva-Domingo, P., & Villaescusa-Navarro, F. 2022, ApJ, 937, 115 [NASA ADS] [CrossRef] [Google Scholar]
  74. Villanueva-Domingo, P., Villaescusa-Navarro, F., Genel, S., et al. 2023, Phys. Rev. D, 107, 103003 [Google Scholar]
  75. Wang, Y., Banerjee, A., & Abel, T. 2022, MNRAS, 514, 3828 [NASA ADS] [CrossRef] [Google Scholar]
  76. Yuan, S., Zamora, A., & Abel, T. 2023, MNRAS, 522, 3935 [NASA ADS] [CrossRef] [Google Scholar]
  77. Zeng, C., Ma, C., Wang, K., & Cui, Z. 2022, IEEE Access, 10, 47361 [Google Scholar]
  78. Zhou, Z., Cisewski-Kehe, J., Fang, K., & Banerjee, A. 2025, ApJ, 979, 194 [Google Scholar]
2

The halos are identified using an FoF algorithm from the simulation particles.

3

We have done a proof-of-concept calculation for maps and CDFs corresponding to the dark matter particles as well in the Appendix A.

4

We have also checked our result with a 256 × 256 grid with 2562 query points, but did not find any significant improvement for the Map-only scenario. We then used the 100 × 100 grid with 104 query points for NN maps throughout the analysis.

5

As shown in Banerjee & Abel (2021a), most of the information is contained in 1NN, 2NN, 3NN, and 4NN CDF statistics. Keeping this in mind, we started with these four NNs, but then also included 8th, 16th, and even 64th NN in the analysis. We did not find any significant change in the results, and therefore reported the constraints obtained using 1NN, 2NN, 3NN, and 4NN statistics.

6

For one of the simulations, the minimum distance for first (1NN), second (2NN), third (3NN), and fourth (4NN) nearest neighbour are 0.18,  0.98,  2.08,  2.94,  5.61 h−1 Mpc, respectively, where as the mean separation between two halos (after downsampling the data to 105 most massive halos) is ∼22 h−1 Mpc.

7

In the figure, the structure of the inference block presents the best architecture achieved after OPTUNA hyper-parameterisation.

8

The learning rate and weight decay of the optimiser are kept as hyperparameters and later optimised using OPTUNA (Akiba et al. 2019).

Appendix A: Constraints on cosmological parameters from dark matter particles

As a proof of concept, we compute NN maps and NN CDFs for dark matter (DM) particles to test their constraining power on cosmological parameters, both individually and in combination. In this analysis, we follow the same procedures outlined in Sects. 2.2 and 2.1, but instead of DM halos, we use the three-dimensional positions of 105 DM particles. The resulting constraints are shown in Fig. A.1. We find that, while the Map-only model outperforms the CDF-only model for Ωm, the situation is reversed for σ8. When the Maps and CDF are combined, i.e. Map+CDF, their joint constraining power exceeds that of either method alone.

Thumbnail: Fig. A.1. Refer to the following caption and surrounding text. Fig. A.1.

Same as Fig 4 but for DM particles rather than halo catalogeus

Table A.1.

Final values of the hyperparameters obtained using OPTUNA for different ML models.

All Tables

Table 1.

Summary of validation metrics for different configurations.

In the text
Table A.1.

Final values of the hyperparameters obtained using OPTUNA for different ML models.

In the text

All Figures

Thumbnail: Fig. 1. Refer to the following caption and surrounding text. Fig. 1.

2D slice of the first (left) and fourth (right) nearest neighbour distance maps for one of the simulations in the Quijote simulation suites used in this study. Each pixel is coloured by the distance from the pixel to the nearest data point (left panel) and by the distance to the fourth nearest neighbour data point (right panel). As can be seen, this converts the discrete dataset into a smooth, continuous map. The colour bar represents the distance (in Gpc/h) from the halos. We note that these maps are only for the purpose of visualisation. They were produced with 2562 random query points in a 256 × 256 2D grid, whereas the actual maps used in this study were produced with 104 random query points in a 100 × 100 2D grid, as mentioned in Sect. 2.1.
In the text
Thumbnail: Fig. 2. Refer to the following caption and surrounding text. Fig. 2.

CDF (left panel) and peaked CDF (right panel) for 1NN (orange), 2NN (red), 3NN (magenta), and 4NN (blue) corresponding to one of the Quijote simulations in this study.
In the text
Thumbnail: Fig. 3. Refer to the following caption and surrounding text. Fig. 3.

Hybrid network in this study. The NN distance maps are used as input to the ResNet block. The output of the ResNet is then concatenated with the NN CDFs, and the merged input then passes through the inference blocks (containing several linear, ReLU, and dropout layers) to predict the mean and standard deviation of the inferred cosmological parameters. The values in brackets show the dimension of the tensor in different stages of the architecture. Here B denotes the batch dimension.
In the text
Thumbnail: Fig. 4. Refer to the following caption and surrounding text. Fig. 4.

Performance of different models when trained to predict likelihood-free inference on the values of Ωm (left column) and σ8 (right column) in three scenarios: Top row: CDF-only; Middle row: Map-only; bottom row: Map+CDF. The values for different validation metrics are given in the legend. As can be seen, the Map-only scenario (middle panel) performs worse than the CDF-only (top panel) scenario. Further, the Map+CDF model performs the best across all the validation metrics.
In the text
Thumbnail: Fig. 5. Refer to the following caption and surrounding text. Fig. 5.

Comparison between ξ(r)-only and CDF-only. As shown, CDF-only performs much better compared to ξ(r)-only, as expected from (Banerjee & Abel 2021b).
In the text
Thumbnail: Fig. A.1. Refer to the following caption and surrounding text. Fig. A.1.

Same as Fig 4 but for DM particles rather than halo catalogeus
In the text

Current usage metrics show cumulative count of Article Views (full-text article views including HTML views, PDF and ePub downloads, according to the available data) and Abstracts Views on Vision4Press platform.

Data correspond to usage on the plateform after 2015. The current usage metrics is available 48-96 hours after online publication and is updated daily on week days.

Initial download of the metrics may take a while.