# A universal equation to predict $\Omega_{\rm m}$ from halo and galaxy
catalogues

Helen Shao Department of Astrophysical Sciences, Princeton University, 4 Ivy
Lane, Princeton, NJ 08544 USA  Center for Computational Astrophysics, Flatiron
Institute, 162 5th Avenue, New York, NY 10010, USA  Natalí S. M. de Santi
Center for Computational Astrophysics, Flatiron Institute, 162 5th Avenue, New
York, NY, 10010, USA  Instituto de Física, Universidade de São Paulo, R. do
Matão 1371, 05508-900, São Paulo, Brasil  Francisco Villaescusa-Navarro Center
for Computational Astrophysics, Flatiron Institute, 162 5th Avenue, New York,
NY, 10010, USA  Department of Astrophysical Sciences, Princeton University, 4
Ivy Lane, Princeton, NJ 08544 USA  Romain Teyssier Department of Astrophysical
Sciences, Princeton University, 4 Ivy Lane, Princeton, NJ 08544 USA  Yueying
Ni Harvard-Smithsonian Center for Astrophysics, 60 Garden Street, Cambridge,
MA 02138, US  McWilliams Center for Cosmology, Department of Physics, Carnegie
Mellon University, Pittsburgh, PA 15213, US  Daniel Anglés-Alcázar Department
of Physics, University of Connecticut, 196 Auditorium Road, U-3046, Storrs,
CT, 06269, USA  Center for Computational Astrophysics, Flatiron Institute, 162
5th Avenue, New York, NY 10010, USA  Shy Genel Center for Computational
Astrophysics, Flatiron Institute, 162 5th Avenue, New York, NY, 10010, USA
Columbia Astrophysics Laboratory, Columbia University, New York, NY, 10027,
USA  Ulrich P. Steinwandel Center for Computational Astrophysics, Flatiron
Institute, 162 5th Avenue, New York, NY 10010, USA  Elena Hernández-Martínez
Universitäts-Sternwarte, Fakultät für Physik, Ludwig-Maximilians-Universität
München, Scheinerstr. 1, 81679 München, Germany  Klaus Dolag  Universitäts-
Sternwarte, Fakultät für Physik, Ludwig-Maximilians-Universität München,
Scheinerstr. 1, 81679 München, Germany  Max-Planck-Institut für Astrophysik,
Karl-Schwarzschild-Straße 1, 85741 Garching, Germany  Christopher C. Lovell
Institute of Cosmology and Gravitation, University of Portsmouth, Burnaby
Road, Portsmouth, PO1 3FX, UK  Centre for Astrophysics Research, School of
Physics, Engineering & Computer Science, University of Hertfordshire, Hatfield
AL10 9AB, UK  Lehman H. Garrison Center for Computational Astrophysics,
Flatiron Institute, 162 5th Avenue, New York, NY 10010, USA  Eli Visbal
Department of Physics and Astronomy and Ritter Astrophysical Research Center,
University of Toledo, 2801 W Bancroft Street, Toledo, OH 43606, USA  Mihir
Kulkarni Department of Physics and Astronomy and Ritter Astrophysical Research
Center, University of Toledo, 2801 W Bancroft Street, Toledo, OH 43606, USA
Lars Hernquist  Harvard-Smithsonian Center for Astrophysics, 60 Garden Street,
Cambridge, MA 02138, USA  Tiago Castro INAF-Osservatorio Astronomico di
Trieste, Via G. B. Tiepolo 11, I-34143 Trieste, Italy  INFN, Sezione di
Trieste, Via Valerio 2, I-34127 Trieste TS, Italy  IFPU, Institute for
Fundamental Physics of the Universe, via Beirut 2, 34151 Trieste, Italy  Mark
Vogelsberger  Kavli Institute for Astrophysics and Space Research, Department
of Physics, MIT, Cambridge, MA 02139, USA  The NSF AI Institute for Artificial
Intelligence and Fundamental Interactions, Massachusetts Institute of
Technology, Cambridge MA 02139, USA

###### Abstract

We discover analytic equations that can infer the value of $\Omega_{\rm m}$
from the positions and velocity moduli of halo and galaxy catalogues. The
equations are derived by combining a tailored graph neural network (GNN)
architecture with symbolic regression. We first train the GNN on dark matter
halos from Gadget N-body simulations to perform field-level likelihood-free
inference, and show that our model can infer $\Omega_{\rm m}$ with $\sim 6\%$
accuracy from halo catalogues of thousands of N-body simulations run with six
different codes: Abacus, CUBEP3M, Gadget, Enzo, PKDGrav3, and Ramses. By
applying symbolic regression to the different parts comprising the GNN, we
derive equations that can predict $\Omega_{\rm m}$ from halo catalogues of
simulations run with all of the above codes with accuracies similar to those
of the GNN. We show that by tuning a single free parameter, our equations can
also infer the value of $\Omega_{\rm m}$ from galaxy catalogues of thousands
of state-of-the-art hydrodynamic simulations of the CAMELS project, each with
a different astrophysics model, run with five distinct codes that employ
different subgrid physics: IllustrisTNG, SIMBA, Astrid, Magneticum, SWIFT-
EAGLE. Furthermore, the equations also perform well when tested on galaxy
catalogues from simulations covering a vast region in parameter space that
samples variations in 5 cosmological and 23 astrophysical parameters. We
speculate that the equations may reflect the existence of a fundamental
physics relation between the phase-space distribution of generic tracers and
$\Omega_{\rm m}$, one that is not affected by galaxy formation physics down to
scales as small as $10~{}h^{-1}{\rm kpc}$.

$N$-body simulations – magnetohydrodynamics (MHD) – cosmology: cosmological
parameters – galaxies: statistics – methods: statistical

##  1 Introduction

$\Lambda$CDM is the current standard model in cosmology that describes the
evolution and expansion of the Universe, where CDM denotes cold dark matter
and $\Lambda$ represents the cosmological constant. This model explains how
primordial density perturbations in the early Universe were amplified by
gravity and eventually lead to the formation of the large-scale structures
that we observe today. To accomplish this, the model relies on several
cosmological parameters that characterize the composition and other
fundamental properties of our Universe. One of them is $\Omega_{\rm m}$, which
quantifies the fractional energy density of total matter, and obtaining an
accurate constraint for it is crucial for improving our understanding of the
foundational physics that governs the Universe.

Historically, the statistics used to analyze the density and velocity fields
of matter and galaxies have been useful probes for $\Omega_{\rm m}$ (Peebles,
1980; Davis et al., 1985; Angulo & Hahn, 2022). This includes the analysis of
redshift-space distortions of galaxy redshift surveys caused by virial and
peculiar velocities that deviate from cosmic expansion (Kaiser, 1987). Such
distortions strongly affect the statistical properties of galaxy clustering
because they break the symmetry in the line-of-sight direction. These
anisotropies directly probe the growth factor, that depends on $\Omega_{\rm
m}$ as described in Sargent & Turner (1977); Tonegawa et al. (2020). Another
useful statistic is the pairwise velocity metric defined for galaxies and
galaxy clusters as the peculiar velocity difference of pairs along their
radial separation vector. Its strong dependence on cosmology has allowed it to
effectively provide constraints on various cosmological parameters including
$\Omega_{\rm m}$ (Cen et al., 1994; Ma et al., 2015). These methods
demonstrate that valuable cosmological information is embedded on the small
scales ($\lesssim 5~{}h^{-1}$Mpc).

On large scales ($\gtrsim 10~{}h^{-1}$Mpc), methods that analyze cosmic flows
(Dekel, 1994) such as the skewness in the divergence of galaxy velocity fields
(Bernardeau et al., 1995) have led to constraints on $\Omega_{\rm m}$
independent of the biasing relation between the distribution of galaxies and
the underlying matter density field. A similar method is using the Zel’dovich
approximation to recover the initial density fluctuation field from observed
galaxy peculiar velocity and density fields. With this, one can then compute
the one-point probability distribution function (IPDF) which is sensitive to
$\Omega_{\rm m}$. Thus, one can tune the value of $\Omega_{\rm m}$ assumed for
the observed density fields to fit the IPDF of the observed velocity field
(Nusser & Dekel, 1992, 1993).

In recent years, there have been significant advances in building detailed
numerical simulations that accurately describe the distribution and dynamics
of galaxies and dark matter. These include both $N$-body and state-of-the-art
hydrodynamic simulations, and they have become powerful tools for constraining
cosmological parameters such as $\Omega_{\rm m}$. However, the optimal method
that can extract the maximum amount of information from this variety of data
is still unknown for non-Gaussian density fields. Fortunately, the advent of
revolutionary machine learning techniques provides an alternative way to
extract information from large amounts of data. By training neural networks to
learn cosmology directly from generic fields, one can achieve tight
constraints on the values of cosmological parameters without relying on
summary statistics (Ravanbakhsh et al., 2017; Schmelzle et al., 2017; Gupta et
al., 2018; Ntampaka et al., 2019; Ribli et al., 2019; Fluri et al., 2019;
Villaescusa-Navarro et al., 2020; Villaescusa-Navarro et al., 2021a;
Villanueva-Domingo & Villaescusa-Navarro, 2022).

In particular, Graph Neural Networks (GNNs), which are constructed to handle
graph representations of irregular data structures, are especially useful for
this purpose because of their unique ability to exploit relational knowledge
between nodes in the graphs down to arbitrarily small scales (Hamilton, 2020;
Battaglia et al., 2018; Bronstein et al., 2021). Specifically, in our previous
paper (Shao et al., 2022a) we showed that GNNs are able to infer $\Omega_{\rm
m}$ with a $6\%$ accuracy from halo catalogues of $N$-body simulations
containing information about the spatial distribution and velocity modulus of
the dark matter halos. More importantly, this network was shown to be robust
across various $N$-body simulations that are run with different numerical
codes, as well as various hydrodynamic simulations that each employ distinct
sub-grid physics models and astrophysical processes. This suggests that the
GNN is employing a fundamental relation between the halo properties and
$\Omega_{\rm m}$ that is not affected by numerical errors from the $N$-body
simulations or baryonic effects. Moreover, in our companion paper (de Santi et
al., 2023), we show that GNNs are able to perform robust inference of
$\Omega_{\rm m}$ from the $3D$ positions and $1D$ velocities galaxies of five
different hydrodynamic simulation codes while marginalizing over cosmologies,
astrophysical effects, sub-grid physics models, and sub-halo definitions.
These results demonstrate the abundance of robust information contained in the
phase space distribution of halos and galaxies.

However, the learned relation is hard to understand because the GNN encodes
information in high dimensional latent space representations that are not
associated with obvious physical interpretations. On the other hand, one can
use techniques in symbolic regression to reveal the physics underlying neural
networks via mathematical formulae. Symbolic regression algorithms can be
trained to approximate any learned network by fitting analytic expressions to
the input and output of neural network components. Such approximations may
also generalize better to data that exists outside the range of the data
distribution used for training because they possess stronger extrapolation
properties than neural networks, whose complex functional forms have the
tendency to overfit and learn uninformative priors used during training
(Villaescusa-Navarro et al., 2020). This method has been recently used to
rediscover physical laws in planetary motion, uncover new relations in matter
over-density fields, and more (Cranmer et al., 2019a, 2020; Wadekar et al.,
2020; Villaescusa-Navarro et al., 2021b; Lemos et al., 2022; Wadekar et al.,
2022; Shao et al., 2022b; Delgado et al., 2022; Bartlett et al., 2022).

Hence, in this paper, we attempt to understand the physical relations employed
by the GNNs presented in Shao et al. (2022a) and de Santi et al. (2023) by
providing an explicit mathematical formula that approximates the learned
networks. To achieve this, we follow a two-step method. First, we train a GNN
on halo positions and velocity moduli to show that a model with reduced latent
space dimensionality can recover the accuracy and robustness of the model
discussed in Shao et al. (2022a). Its compressed architecture will aid the use
of symbolic regression and decrease the complexity of the approximating
expressions. In the second step, we train a symbolic regressor to find
mathematical equations that approximate each component of the GNN model. We
show that the discovered analytic expressions are able to preserve the
accuracy and robustness of the relation found by the GNN by testing them on
halos from thousands of $N$-body and hydrodynamic simulations of varying
cosmological and astrophysical parameters. More surprisingly, we also
demonstrate that the equations are able to predict the value of $\Omega_{\rm
m}$ from galaxy catalogues of five different hydrodynamic simulations. This
suggests that the equations may be independent of the complex connection
between the spatial and velocity distributions of halos and galaxies. Finally,
we attempt to interpret the physical meaning of the equations. Since the
expressions reveal that the network is exploiting rotationally-symmetric
information encoded in the relative velocity modulus of the halo pairs on
small scales $\sim 1.35~{}h^{-1}$Mpc, we draw connections to traditional
techniques that rely on phase space distributions for galaxies and halos to
constrain $\Omega_{\rm m}$.

This paper is structured as follows. We first describe the data used for this
project in Section 2. In Section 3, we describe the architecture of our GNN
models, the symbolic regression algorithm, and the methods used to train,
validate, and test both models. In Section 4, we present the results of our
models and equations. We then provide a discussion of plausible physical
interpretations of the equations in Section 5. Finally, we summarize the main
findings in Section 6.

##  2 Data

We train our models using halo catalogues from high-resolution cosmological
simulations that contain two halo properties. First, the halo positions, r,
which are defined for the halo center using Cartesian coordinates in comoving-
space. Second, the halo velocity modulus, $V$, which is defined as the modulus
of the 3D peculiar velocity vector computed with respect to the velocity of
the simulation box. In this work we focus on halo and galaxy catalogues at
$z=0$. We describe the methods to generate the halo and galaxy catalogues we
use to train, validate, and test the model in Section 3.1.

###  2.1 Simulations

We follow the scheme used in Shao et al. (2022a) to test the accuracy and
robustness of our models. This strategy is composed of two parts. First, we
use cosmological $N$-body and hydrodynamic simulations that contain different
$\Omega_{\rm m}$ values organized in Latin hypercubes and varying initial
random seed conditions to quantify the percentage constraints and level of
precision achieved by the models. Specifically, $\Omega_{\rm m}$ varies in the
range

$0.1\leq\Omega_{\rm m}\leq 0.5$ (1)

for both the $N$-body and hydrodynamic simulations. Note that these
simulations also vary $\sigma_{\rm 8}$ in the range $0.6\leq\sigma_{\rm 8}\leq
1.0$. Furthermore, for the hydrodynamic simulations, we vary several
astrophysical parameters; most of them just alter four astrophysical
parameters controlling the efficiency of supernova and active galactic nucleus
(AGN) feedback but we also made use of a new set that varies 23 astrophysical
parameters controlling most of the free parameters in the considered
hydrodynamic code. The hydrodynamic simulations have been run with five
different codes that not only solve the hydrodynamic equations using different
methods, but they made use of different subgrid models. These simulations are
part of the CAMELS project and we refer the reader to Villaescusa-Navarro et
al. (2021b, 2022a) for further details.

Second, we use simulations that are generated with the same cosmologies and
initial seeds for a control set-up in which we can determine the robustness of
the models when evaluated on halos generated with different codes. For this,
we run 6 $N$-body simulations that have the same initial random seed and value
of $\Omega_{\rm m}=0.3175$ (all other cosmological parameters are shared among
codes), but each is run with a different code. Additionally, we run 4
hydrodynamic simulations that have the same value of $\Omega_{\rm m}=0.3$,
initial random seed, and employ their fiducial sub-grid physics model using 4
distinct codes.

For these two above steps, we employ thousands of $N$-body and hydrodynamic
simulations that have volumes of (25 $h^{-1}$Mpc)3 and have been run with 11
different codes. We briefly describe these codes below, but for more detailed
information, we refer the reader to Shao et al. (2022a) and the listed
paper(s) for each code. Note that at the end of the descriptions for each
code, we include the number of simulations generated to contain the same
cosmology and initial random seed as the other codes, and the number of
simulations that contain varying cosmologies, initial seeds, and/or
astrophysical parameters arranged in a Latin-hypercube (or Sobol sequence),
respectively.

####  2.1.1 $N$-body codes

The different $N$-body codes follow the evolution of dark matter particles
(that represent the cold dark matter plus baryonic fluid) under the effect of
self-gravity in a given expanding cosmological background using different
numerical techniques and approximations. The six codes we use to run the
$N$-body simulations are described briefly below.

  1. 1.

Abacus. This code computes the long-range gravitational potential by
decomposing the near-field and far-field forces in which the near-field forces
are reduced to a $r^{-2}$ summation (or an appropriately softened form) and
the far-field forces to a discrete convolution over multipoles (Garrison et
al., 2021). We run 51 simulations with Abacus: 1 simulation with a shared
cosmology and initial random seed among codes and 50 simulations in a Latin-
hypercube with varying values of $\Omega_{\rm m}$ and $\sigma_{\rm 8}$.

  2. 2.

CUBEP3M. This code employs a particle-particle particle-mesh (P3M) scheme,
described in Harnois-Déraps et al. (2013), where long-range gravitational
forces are computed via a two-level particle mesh calculation. We ran 51
CUBEP3M simulations: 1 simulation with shared cosmology and initial random
seed among codes and 50 simulations in a Latin-hypercube. For the simulation
sharing the cosmology and initial random seed, we used the exact same initial
particles as in the other codes, whereas the CUBEP3M initial conditions,
generated using the Zeldovich approximation, were used for the 50 simulations
in the Latin-hypercube.

  3. 3.

Enzo. This is an Adaptative Mesh Refinement (AMR) code, as described in Bryan
et al. (2014), that solves the Poisson equation via a fast Fourier technique
(Hockney & Eastwood, 1988) on the root grid and a multigrid solver on the
individual sub-mesh. We only have one Enzo simulation which shares the same
cosmology and initial random seed with the other codes.

  4. 4.

Gadget. This code utilizes a TreePM algorithm to compute short-range forces
and Fourier techniques to calculate long-distance forces, as described in
Springel (2005). We use the halo catalogues from these simulations to train
the models. We run 1,001 of the Gadget simulations: 1 simulation with shared
cosmology and initial random seed among codes and 1,000 simulations that have
different values of $\Omega_{\rm m}$, $\sigma_{\rm 8}$, and initial random
seed. We use the halo catalogues from these simulations to train the models.

  5. 5.

PKDGrav3. This code computes forces using Fast Multipole Method (FMM,
Greengard & Rokhlin 1987) as described in Potter et al. (2017). We run 1,001
$N$-body simulations with this code: 1 simulation with shared cosmology and
initial random seed among codes and 1,000 simulations with different values of
$\Omega_{\rm m}$, $\sigma_{\rm 8}$, and initial random seed that are organized
in a Latin-hypercube.

  6. 6.

Ramses. This code uses the Adaptive Particle Mesh technique described in
Teyssier (2002). It solves Poisson’s equation level by level using Dirichlet
boundary conditions and a Multigrid relaxation solver. We have run 1,001
Ramses simulations: 1 simulation with shared cosmology and initial random seed
among codes, and 1,000 simulations with different values of $\Omega_{\rm m}$,
$\sigma_{\rm 8}$, and initial random seed that are organized in a Latin-
hypercube.

####  2.1.2 Hydrodynamic codes

The hydrodynamic simulations have been run using codes that solve the
hydrodynamic equations with different numerical methods and employ distinct
models to describe astrophysical processes such as star formation and feedback
from supernova and AGN. The hydrodynamic simulations have been run with the
codes MP-Gadget, Arepo, OpenGadget, Gizmo, and SWIFT-EAGLE. In these
simulations, which are part of the CAMELS project (Villaescusa-Navarro et al.,
2021b), we vary the values of $\Omega_{\rm m}$, $\sigma_{\rm 8}$, the initial
random seed, and several astrophysical parameters that we describe below.
Instead of referring to these simulations by the name of the code used to run
them, we will call them by name of the flagship simulations associated with
them and their sub-grid model; i.e. ASTRID, IllustrisTNG, Magneticum, SIMBA,
and SWIFT-EAGLE respectively. We note that the SB28 simulations have been run
with the Arepo code and employ the IllustrisTNG sub-grid model, but since they
vary 28 parameters we use a special name for them. Below, we briefly describe
the simulations from the different codes:

  1. 7.

ASTRID. These simulations employ the MP-Gadget code to solve the gravity (with
TreePM), hydrodynamics (with the pressure-entropy formulation of SPH), and
astrophysical processes (Ni et al., 2022; Bird et al., 2022). We have run
1,001 simulations with this code: 1 simulation with shared cosmology and
initial random seed among codes, and 1,000 simulations with different values
of $\Omega_{\rm m}$, $\sigma_{\rm 8}$, four astrophysical parameters that
control the efficiency of supernova and AGN feedback, and initial random seed
that are organized in a Latin-hypercube.

  2. 8.

IllustrisTNG. These simulations have been run with the Arepo code Springel
(2010); Weinberger et al. (2020), making use of a TreePM plus moving-mesh
finite volume (MMFV) method (Weinberger et al., 2017; Pillepich et al.,
2018a). We have run 1,029 simulations with this code: 1 simulation with shared
cosmology and initial random seed among codes, and 1,000 simulations with
different values of $\Omega_{\rm m}$, $\sigma_{\rm 8}$, four astrophysical
parameters that control the efficiency of supernova and AGN feedback, and
initial random seed that are organized in a Latin-hypercube. We also run 27
simulations using this code that only differs in the value of their initial
random seed to study the effect of cosmic variance, which we refer to as the
CV set. Finally, we have 1 simulation of this code containing a periodic
comoving volume of $(205~{}h^{-1}{\rm Mpc})^{3}$. This simulation is part of
the IllustrisTNG-300 set (Marinacci et al., 2018; Nelson et al., 2018; Naiman
et al., 2017; Nelson et al., 2019; Springel et al., 2018; Pillepich et al.,
2018b) and we use it to quantify how our analytic expressions behave in the
presence of super-sample covariance effects.

  3. 9.

Magneticum. This simulation is run with the code OpenGadget3 and implements
the SPH-scheme following Beck et al. (2016). For more details, see Dolag et
al. (2004); Jubelgas et al. (2004); Hirschmann et al. (2014); Groth et al.
(2023). We have run 51 Magneticum simulations: 1 simulation with shared
cosmology and initial random seed among codes, and 50 simulations with
different values of $\Omega_{\rm m}$, $\sigma_{\rm 8}$, four astrophysical
parameters that control the efficiency of supernova and AGN feedback, and
initial random seed that are organized in a Latin-hypercube.

  4. 10.

SIMBA. These simulations have been run with the GIZMO code (Hopkins, 2015)
with a TreePM plus Meshes Finite Mass method (MFM), see Davé et al. (2019). We
have run 1,001 SIMBA simulations: 1 simulation with shared cosmology and
initial random seed among codes, and 1,000 simulations with different values
of $\Omega_{\rm m}$, $\sigma_{\rm 8}$, four astrophysical parameters that
control the efficiency of supernova and AGN feedback, and initial random seed
that are organized in a Latin-hypercube.

  5. 11.

SB28. These simulations have been run with Arepo and employ the IllustrisTNG
model. They contain 1,024 simulations and we place them in a different
category as they vary the value of 5 cosmological parameters ($\Omega_{\rm
m}$, $\Omega_{\rm b}$, $h$, $n_{s}$, $\sigma_{\rm 8}$), and 23 astrophysical
parameters controlling most of the code free parameters. The values of the 28
parameters are organized in a Sobol sequence Sobol’ (1967).

  6. 12.

SWIFT-EAGLE. These simulations have been run with the SWIFT-EAGLE code
(Schaller et al., 2016, 2018) and employ a sub-grid physics model that aims at
mimicking the original Gadget-EAGLE model (Schaye et al., 2015; Crain et al.,
2015), with some parameter and implementation differences (Borrow et al.,
2022). The full model will be described in Borrow & et. al. (2023). The suite
contains $64$ simulations varying eight subgrid parameters that control
stellar and AGN feedback on a Latin-hypercube (parameter ranges are given in
square brackets):

     * •

$f_{\mathrm{E,min}}$, the minimal stellar feedback fraction, [0.18, 0.6]

     * •

$f_{\mathrm{E,max}}$, the maximal stellar feedback fraction, [5, 10]

     * •

$N_{\mathrm{H,0}}$, pivot point in density that the feedback energy fraction
plane rotates around, [$10^{-0.6}$, $10^{-0.15}$]

     * •

$\sigma_{\mathrm{n}}$ and $\sigma{\mathrm{Z}}$, energy fraction sigmoid width,
controlling the density and metallicity dependence, [0.1, 0.65]

     * •

$\varepsilon_{\mathrm{f}}$, coupling coefficient of radiative efficiency of
AGN feedback, [$10^{-2}$, $10^{-1}$]

     * •

$\Delta T_{\mathrm{AGN}}$, AGN heating temperature, [$10^{8.3}$, $10^{9.0}$]

     * •

$\alpha$, black hole accretion suppression / enhancement factor, [0.2, 1.1].

##  3 Methods

Figure 1: This is a schematic of our methodology which is explained in Section
3. We begin by constructing a graph from a halo catalogue using halo positions
and velocity moduli. We then feed the graphs to a GNN and train it to perform
parameter inference for $\Omega_{\rm m}$. After training the model, we use
symbolic regression to extract the equations from each component of the GNN
architecture. Finally, we assemble the equations into one expression and use
it to predict $\Omega_{\rm m}$ from halos and galaxies of various $N$-body and
hydrodynamic simulations.

In Shao et al. (2022a) we found that GNNs are not only able to infer
$\Omega_{\rm m}$ with a $5.6\%$ precision but are also robust across different
$N$-body and hydrodynamic codes, suggesting that the learned relation might be
physically fundamental. In this work, we build upon this previous study to
understand the found relation and search for an analytic formula that can
approximate the mapping from the halo positions and velocities, r and $V$, to
the cosmological parameter, $\Omega_{\rm m}$. To accomplish this, we make use
of both GNNs and symbolic regression algorithms. We refer the reader to
Cranmer et al. (2019a, b) for similar methodologies devised to extract
symbolic relations from trained neural networks.

We begin by training a GNN with the goal of obtaining a low-dimensional latent
space network to learn a relation between the input halo properties and
$\Omega_{\rm m}$ that can approximate the previously found model. We can do
this by fixing certain hyperparameters of the GNN so that it has a reduced
architecture depth and width. This step is key to aiding the search for
analytic expressions when we use symbolic regression to approximate the GNN,
as we later explain. We then evaluate this GNN model on halo catalogues from
the $N$-body and hydrodynamic codes described in the previous section to
ensure that the sparse architecture is able to achieve comparable precision
and accuracy to the model obtained in Shao et al. (2022a). Finally, we use
symbolic regression to fit mathematical formulae to each component of the
architecture in the trained GNN model to obtain approximate analytic
equations. To improve their interpretability, we also make modifications
motivated by physical principles, such as preserving the symmetries present in
the data and model and simplifying the found expressions. We refer the reader
to Fig. 1 which depicts this methodology schematically.

In the following sections, we describe in detail the ingredients we use to
perform this procedure: 1) the method for constructing the halo (training,
validating, and testing) and galaxy catalogues (testing), 2) the graph data
used to train the GNN, 3) the GNN architecture and training procedure, 4) the
data and procedure used to train the symbolic regressor, and finally, 5) the
metrics used to evaluate the accuracy and precision of the models.

###  3.1 Halo and Galaxy Catalogues

Here, we describe the procedures for constructing the halo and galaxy
catalogues that we use to train, validate, and test the GNN and symbolic
expressions.

  * •

Halo Catalogues for Training and Validating. For training and validation, we
use halo catalogues from the Gadget simulations. For each simulation, we
generate 10 halo catalogues by taking all halos with masses larger than
$M_{\rm X}$, where $M_{\rm X}$ is a randomly chosen number between
$100\,m_{\rm p}$ and $500\,m_{\rm p}$. Here, $m_{p}$ is the mass of a single
dark matter particle. As explained in Shao et al. (2022a), using different
dark matter particle thresholds is key to achieving a model that is robust to
different simulations. These halo catalogues are generated by running Rockstar
(Behroozi et al., 2013) on snapshots from the numerical simulations described
above.

  * •

Halo Catalogues for Testing. We use all $N$-body simulations described in the
previous section and two hydrodynamic simulations: IllustrisTNG and SIMBA. For
each simulation, we generate 5 halo catalogues for the five different dark
matter particle thresholds: $\\{100,200,300,400,500\\}$. Note that for
hydrodynamic simulations, the mass of a dark matter halo contains
contributions from various mass sources. Hence, instead of considering only
the amount of dark matter mass to make our mass cuts, we define $m_{p}$ as the
effective particle mass: $m_{p}=\frac{1}{N_{c}}\Omega_{\rm m}V\rho_{c}$, where
$V$ is the volume of the simulation, $\rho_{c}$ is the Universe’s critical
density today, and $N_{c}=256^{3}$ is the effective number of particles. These
halo catalogues are generated by running Rockstar (Behroozi et al., 2013) on
snapshots from the numerical simulations described above. However, for one
test where we gauge the robustness of the train models to different halo
definitions, we run Subfind (Dolag et al., 2009) to generate halo catalogues
from the Gadget $N$-body and Illustris-TNG simulations.

  * •

Galaxy Catalogues for Testing. We use galaxy catalogues from all the
hydrodynamic simulations described in the previous section. We define a galaxy
as a subhalo (can either be a central or satellite) that contains a stellar
mass of at least $N\times m_{*}$ where $N\in{3,4,5,6}$ and $m_{*}=1.3\times
10^{7}~{}h^{-1}M_{\odot}$. For each simulation, we construct 4 catalogues,
each using a different $N$. We limit the range of the stellar mass thresholds
to be no larger than $6\times m_{*}$ because we find that using larger cuts
result in catalogues with galaxy number densities that are smaller than the
number densities (from the halo catalogues) used to train the network and
equations. We find that using catalogues with number densities that are
outside the training range can lead to inaccurate predictions. These galaxy
catalogues are generated by running Rockstar (Behroozi et al., 2013) on
snapshots from the six hydrodynamic simulations described above, with the
exception of the catalogues from the SWIFT-EAGLE simulations which were
generated using the halo finder VELOCIraptor (Elahi et al., 2019; Cañas et
al., 2019).

###  3.2 GNNs

The methods described in this section closely follow those presented in Shao
et al. (2022a) to infer $\Omega_{\rm m}$. We emphasize the key changes that we
implement in this work are: 1) using only the summation operator as the
aggregation function and 2) reducing the depth and width of the GNN
architecture with constrained hyperparameter optimization. These steps
decrease the complexity of the model and allow for easier interpretation of
the learned relations.

####  3.2.1 Model input: Halo Graphs

The input of the GNN is a graph defined as
$\mathcal{G}=(\mathcal{V},\mathcal{E})$, where $\mathcal{V}$ is the set of
nodes and $\mathcal{E}$ is the set of edges. The nodes represent the halos (or
galaxies) and an edge is created between two nodes if their distance is
smaller than the linking radius, $r_{\rm link}$. This property is considered a
hyperparameter that we optimize during training, as we explain later. Thus,
two nodes $i,j\in\mathcal{V}$ are referred to as neighbors if they are
connected via an edge, $(i,j)\in\mathcal{E}$. As in Shao et al. (2022a), we do
not consider self-loops and account for periodic boundary conditions when
computing distances and angles between nodes.

The nodes and the edges can have different properties associated with them,
that we denote as $\textbf{v}_{i}^{(n)}$ and $\textbf{e}_{ij}^{(n)}$,
respectively. Since the architecture of the GNN models may consist of multiple
layers that take a graph as the input and outputs an updated graph. For this
reason, we denote the node and edge features at the $n^{th}$ layer with the
superscript $n$.

The initial node feature, represented by $\textbf{v}_{i}^{(0)}$, that we use
is the halo velocity modulus, $V$. Since the velocities are defined with
respect to the simulation box, the node features preserve Galilean invariance.
The edge features between nodes $i$ and $j$ at the $n^{th}$ layer are
represented by $\textbf{e}_{ij}^{(n)}$ and they contain information about the
spatial distribution of halos. To ensure that the model preserves the
rotational and translational invariance of the data, we use the following
vector for the edge features:

$\textbf{e}^{(0)}=[\alpha_{ij},\beta_{ij},\gamma_{ij}]$ (2)

where

$\displaystyle\alpha_{ij}$
$\displaystyle=\frac{\textbf{r}_{i}-\textbf{c}}{|\textbf{r}_{i}-\textbf{c}|}\cdot\frac{\textbf{r}_{j}-\textbf{c}}{|\textbf{r}_{j}-\textbf{c}|}$
(3) $\displaystyle\beta_{ij}$
$\displaystyle=\frac{\textbf{r}_{i}-\textbf{c}}{|\textbf{r}_{i}-\textbf{c}|}\cdot\frac{\textbf{d}_{ij}}{|\textbf{d}_{ij}|}$
(4) $\displaystyle\gamma_{ij}$ $\displaystyle=\frac{|\textbf{d}_{ij}|}{r_{\rm
link}},$ (5)

with $\textbf{d}_{ij}=\textbf{r}_{i}-\textbf{r}_{j}$ being the relative
distance between the nodes $i$ and $j$ and c is the centroid of the
halo/galaxy distribution. $\alpha_{ij}$ defines the angle between the
positions of node $i$ and its neighbor node $j$, while $\beta_{ij}$ describes
the angle between positions of node $i$ and the separation between nodes $i$
and $j$. Note that we have normalized the distance, $\textbf{d}_{ij}$, by
dividing it with the linking radius, $r_{\rm link}$, to have dimensionless
edge features. We refer the reader to Villanueva-Domingo & Villaescusa-Navarro
(2022) for more details on this construction.

####  3.2.2 Architecture

The architecture of our GNN model closely follows CosmoGraphNet111
https://github.com/PabloVD/CosmoGraphNet (Villanueva-Domingo, 2022), presented
in Villanueva-Domingo & Villaescusa-Navarro (2022) and used in Shao et al.
(2022a). However, our model only includes one message-passing layer and a
final aggregation layer. We arrived at this architecture by experimenting with
different numbers of hidden layers to optimize the simplicity of the model
while maintaining the precision and accuracy of its predictions. We explain
this in more details in Section 3.2.3.

In the message-passing layer, information from the input node and edge
features are encoded with multilayer perceptrons (MLP) and recursively
exchanged and aggregated between each node’s neighbors and edges. Afterwards,
the node and edge features are updated. This creates hidden feature vectors
that are ultimately used to predict the target parameter. For this reason, we
denote the edge and node features that are input to the message-passing layer
(the initial halo properties) with the superscript $(0)$ and output (hidden)
features by the message-passing layer with the superscript $(1)$.

For our compressed GNN, we restricted to two hidden features for each node and
edge because the number of hidden features scales in proportion to the number
of analytic expressions needed to approximate the network, as we explain
later.

For the message-passing layer, the input of the edge model are the initial
features of the node $i$, the neighboring node $j$, and their shared edge. In
this case, the initial node features are $V$ as defined in Section 2 and the
initial edge features are ${e}^{(0)}$ as defined in Equation 2. This
information is passed through an MLP, denoted by $\phi^{e}$ and the output are
the updated hidden edge features:

$\textbf{e}_{ij}^{(1)}=\phi^{e}\left(\left[\textbf{v}_{i}^{(0)},\textbf{v}_{j}^{(0)},\textbf{e}^{(0)}\right]\right).$
(6)

This hidden edge feature, along with the initial node feature of node $i$, is
then passed to the node model, where another MLP, denoted by $\phi^{v}$,
outputs the hidden node features:

$\textbf{v}_{i}^{(1)}=\phi^{v}\left(\left[\textbf{v}_{i}^{(0)},\sum_{j\in\mathcal{N}_{i}}\textbf{e}^{(1)}\right]\right)~{}.$
(7)

Here, we use a permutationally invariant aggregation function - the summation
- to aggregate the node features of the neighbor nodes $j\in\mathcal{N}_{i}$
that are connected to node $i$. In Shao et al. (2022a), the aggregation
function used was a concatenation of the maximum, summation, and mean
operators. In this work, we reduce this function to just the summation to
decrease the complexity of the learned relations. This choice is motivated by
the fact that the summation can serve as a proxy for the other two operators.
Using only one aggregation operator as opposed to three decreases the number
of hidden channels by a factor of three and thus reduces the number of
equations we for our model.

The final layer in the architecture aggregates the hidden node features output
by the message passing layer to make the prediction y:

$\textbf{y}=\phi^{u}\left(\left[\sum_{i\in\mathcal{G}}\textbf{v}_{i}^{(1)}\right]\right),$
(8)

where $\sum_{i\in\mathcal{G}}$ operates over all nodes in the graph and
$\phi^{u}$ is another MLP that extracts the target information.

####  3.2.3 Training procedure

We train and test the models using graphs constructed from halo catalogues of
the Gadget simulations. For each simulation, we construct 10 catalogues using
the procedure described in Section 3.1 to marginalize over the halo number
density. Once trained, the model is tested using catalogues from all
simulations. For Gadget, we split the simulations into training (80%),
validation (10%), and testing (10%) data sets before creating halo catalogues
for each simulation. For the other codes, we use the entirety of the dataset
for testing.

We standardize the values of input node features as

$\tilde{x}=\frac{x-{\mu}}{{\delta}},$ (9)

where $\mu$ and $\delta$ denote the mean and standard deviation of the feature
$x$. However, we explain in later sections that the value of $\delta$ must be
tuned for when evaluating the symbolic equations. We also normalize the values
of the target cosmological parameter, $\Omega_{\rm m}$:

$\bar{\Omega}_{\rm m}=\frac{\Omega_{\rm m}-{\rm min}(\Omega_{\rm m})}{{\rm
max}(\Omega_{\rm m})-{\rm min}(\Omega_{\rm m})},$ (10)

where the minimum and maximum values of the ranges of $\Omega_{\rm m}$ are
listed in Eq. 1.

As we did in Shao et al. (2022a), we train the GNN to perform likelihood-free
inference so the output of the model is $\mathbf{y}=[\mu_{i},\sigma_{i}]$,
where $\mu_{i}$ is the posterior mean and $\sigma_{i}$ is the posterior
standard deviation of $\Omega_{\rm m}$. To achieve this, we employ the
following loss function:

$\begin{split}\mathcal{L}&=\log{\bigg{(}\sum_{j\in{\rm
batch}}(\theta_{i,j}-\mu_{i,j})\bigg{)}^{2}}+\\\ &\log{\bigg{(}\sum_{j\in{\rm
batch}}\big{(}(\theta_{i,j}-\mu_{i,j})^{2}-\sigma_{i,j}^{2}\big{)}\bigg{)}^{2}}\end{split}$
(11)

where the sums are performed over the halo catalogues in the batch. Further
details on this can be found in Jeffrey & Wandelt (2020) and Villaescusa-
Navarro et al. (2022b).

Our model is implemented in PyTorch (Paszke et al., 2019) and PyTorch
Geometric (Fey & Lenssen, 2019). We use the AdamW optimizer (Loshchilov &
Hutter, 2017) with beta values equal to 0.9 and 0.999. We train the network
using a batch size of 8 for 500 epochs. The hyperparameters for our model are:
1) the learning rate, 2) the weight decay, and 3) the linking radius. We use
the optuna code (Akiba et al., 2019) to perform Bayesian optimization and find
the best value of these hyper-parameters for each model. As mentioned earlier,
we aim to reduce the depth and width of our GNN architecture to obtain a
compressed network so we restrict to only one layer and two hidden neurons.
For each model, we run 100 trials, where each trial consists of training the
model using selected values of the hyper-parameters. We perform the
optimization of the hyper-parameters required to achieve the lowest validation
loss possible and use early stopping to save only the model with a minimum
validation error.

###  3.3 Symbolic Regression

While neural networks can provide precise and accurate approximations of
complex relations in the data, interpreting them is often challenging because
they employ a large number of parameters to make predictions. Therefore, it is
desirable to extract mathematical expressions that characterize, or
approximate, the relation learned by the neural network because it is easier
to understand the physics of the found relationships in such forms. Moreover,
analytic equations have been found to generalize better, than neural networks,
to data with characteristics not presented in the training set, which can give
us more robust predictions and possibly illuminate fundamental properties of
the model (Shao et al., 2022b).

For this purpose, we first train a symbolic regression algorithm designed to
approximate functions with analytic formulae. We then modify the expressions
using reasoning based on physical principles - such as that the model should
preserve rotational and translational symmetries of the data - to improve the
interpretability of the equations and reduce their complexity. In this
section, we describe the symbolic regression algorithm we use and the
procedure for fitting functions to components of the learned GNN.

We use the package pysr (Cranmer, 2020) to train a symbolic regression
algorithm with the ability to fit mathematical formulas to the learn GNN
relations. This package implements genetic programming which searches for the
optimal analytic expression creating combinations between the sets of given
operators and input variables. The found expressions of each so-called
generation are evaluated, and the most accurate ones survive to the next
generation. Throughout this iterative process, mutations and crossovers take
place to explore the entire equation space and find an accurate expression.

However, a key limitation of symbolic regression is that its tractability and
accuracy are restricted to low-dimensional spaces of input data. To circumvent
this, we limit the size of the latent space produced by the GNN, as described
in Section 3.2.2. Using the learned parameters and relations from the low-
dimensional GNN architecture, we search for equations that characterize the
model by approximating the individual MLPs used in the node model, edge model,
and final layer described in Equations 6, 7, and 8, respectively. We emphasize
that since there is only one message-passing layer, we only need to
approximate one node model MLP and one edge model MLP. Moreover, for each of
the node and edge models, we search for two equations because there are two
hidden features. The data and procedure used to obtain these equations are
described below.

  * •

Approximating Edge Model: To approximate the edge model, we train a symbolic
regressor to map from the input variables, $\textbf{x}^{e}$, to the target
variables, $\textbf{y}^{e}$, defined as:

$\displaystyle\textbf{x}^{e}$
$\displaystyle=\left(v_{i}^{(0)},v_{j}^{(0)},\alpha_{ij},\beta_{ij},\gamma_{ij}\right)$
(12) $\displaystyle\textbf{y}^{e}$
$\displaystyle=\left(e_{1}^{(1)},e_{2}^{(1)}\right).$ (13)

The input variables are the initial features of the nodes and their neighbors,
as well as the initial edge features as described in Section 3.2.1. The
corresponding target variables are the edge features of the MLP in the edge
model defined in Equation 6. Since the GNN employs only two hidden features
for each message-passing layer, we denote the first component of the edge
feature as $e_{1}^{(1)}$ and the second component as $e_{2}^{(1)}$. To obtain
this data, we randomly select 10 $(\textbf{x}^{e},\textbf{y}^{e})$ pairs from
each graph in the training set. This selection is done to ensure that we have
a representative sample of the training set without using every node pair of
all graphs which would result in too large of a dataset.

  * •

Approximating Node Model: Similarly, to approximate the node model, the input
variables, $\textbf{x}^{n}$, and the target variables, $\textbf{y}^{n}$, of
the symbolic regressor, are:

$\displaystyle\textbf{x}^{n}$
$\displaystyle=\left(v_{i}^{(0)},\sum_{j\in\mathcal{N}_{i}}{e}_{1}^{(1)},\sum_{j\in\mathcal{N}_{i}}{e}_{2}^{(1)}\right)$
(14) $\displaystyle\textbf{y}^{n}$
$\displaystyle=\left(v_{1}^{(1)},v_{1}^{(1)}+v_{2}^{(1)}\right).$ (15)

As seen above, the inputs are the initial node feature and the neighborhood-
wise sums of the hidden edge features because the output of the edge model is
aggregated using the summation operator before being passed onto the node
model. The corresponding target variables are the hidden node features of the
MLP in the node model defined in Equation 7. We denote the first and second
hidden node features as $v_{1}^{(1)}$ and $v_{2}^{(1)}$, respectively.
However, instead of directly finding an equation for the second node feature,
$v_{2}^{(1)}$, we instead search for a formula for the sum
$v_{1}^{(1)}+v_{2}^{(1)}$. This is because we find that the change of
variables allows us to obtain more accurate approximations than with the
original target variable. Ultimately, to obtain the expression of
$v_{2}^{(1)}$, we subtract from it $v_{1}^{(1)}$. To obtain this data, we
randomly sample 10 $(\textbf{x}^{n},\textbf{y}^{n})$ pairs from each graph in
the training set as we did with the edge model data.

  * •

Approximating Final MLP: Lastly, to approximate the MLP in the final
aggregation layer, the input and target variables are:

$\displaystyle\textbf{x}^{u}$
$\displaystyle=\left(\sum_{i\in\mathcal{G}}v_{1}^{(1)},\sum_{i\in\mathcal{G}}v_{2}^{(1)}\right)$
(16) $\displaystyle\textbf{y}^{u}$ $\displaystyle=\mu_{i}.$ (17)

Here, the inputs are the graph-wise sums of the hidden node features because
the output of the node model is aggregated using the summation operator before
being passed onto the final MLP. The corresponding target is the mean
posterior. We do not attempt to find an expression for the posterior standard
deviation as it is solely a component of the parameter inference methodology
and does not contribute additional physical understanding. We obtain this data
from each graph in the training set. Note that this time there is no need to
select a sub-sample of nodes from each graph because $x^{u}$ and $y^{u}$ are
global properties of the graph so we can use every graph in the training set.

In each of the above approximation steps, the symbolic regression algorithm
searches for analytic expressions that can map from the given input variables
to the desired target. For the training, the regressor is allowed to employ
the following binary operators: ”add”, ”sub”, ”mult”, ”div”, ”pow”222The
listed operators perform addition, subtraction, multiplication, and division.
”pow” takes the power of $X$ to the input variable, where $X$ is any number.
and the following unary operators: ”1/x” (the inverse of a variable), ”abs”,
”log”, ”log10”, ”sqrt”. We employ a standard Mean Squared Error (MSE) loss
function to optimize the fitting defined as,

${\rm MSE}=\frac{1}{N}\sum_{i=1}^{N}\left(y_{\rm true}-y_{\rm
pred}\right)^{2},$ (18)

where $\textbf{y}_{\rm pred}$ denotes the predicted value of the target
variable and $\textbf{y}_{\rm true}$ is corresponding true value. The model
was trained for 100,000 trials with a batch size of 64.

During training, the algorithm outputs a list of equations found by the
regressor. For each equation, PYSR provides three values to quantify the fit
of the equation: its complexity, MSE, and score. The complexity of the
equation takes into account the number of operators, constants, and variables
used. The MSE and the complexity are combined into an overall metric that
gives the equation’s score, akin to Occam’s Razor (Cranmer, 2020).
Specifically, the algorithm sorts the found equations from the least to the
most complex, and for each equation, it computes the fractional decrease in
MSE relative to the next (more complex) equation. The score is maximized if
this fractional decrease is large. We evaluate several candidate equations on
a test set for each hidden feature before selecting one that optimizes the
tradeoff between complexity and accuracy with these metrics in mind.

###  3.4 Performance Metrics

For the graph $i$, with the true value of the considered parameter $y_{\rm
truth,i}$, our models output the posterior mean, $y_{\rm infer,i}$, and
standard deviation $\sigma_{i}$. To evaluate the accuracy and precision of our
models, we follow Villanueva-Domingo & Villaescusa-Navarro (2022) and Shao et
al. (2022a), and employ four different metrics:

  1. 1.

Mean relative error, $\epsilon$, defined as

$\epsilon=\frac{1}{N}\sum_{i}^{N}\frac{|y_{{\rm truth},i}-y_{{\rm
infer},i}|}{y_{{\rm truth},i}},$ (19)

where $N$ is the number of halo catalogues in the test set.

  2. 2.

Coefficient of determination, $R^{2}$, defined as

$R^{2}=1-\frac{\sum_{i}^{N}(y_{{\rm truth},i}-y_{{\rm
infer},i})^{2}}{\sum_{i}^{N}(y_{{\rm truth},i}-\overline{y}_{{\rm
truth}})^{2}},$ (20)

  3. 3.

Root mean squared error, RMSE, defined as:

${\rm RMSE}=\sqrt{\frac{1}{N}\sum_{i=1}^{N}\left(y_{{\rm truth},i}-y_{\rm
infer}\right)^{2}}$ (21)

  4. 4.

Chi squared, $\chi^{2}$, defined as:

$\chi^{2}=\frac{1}{N}\sum_{i=1}^{N}\frac{(y_{\rm truth,i}-y_{\rm
infer,i})^{2}}{\sigma_{i}^{2}}~{}.$ (22)

Note that a value of $\chi^{2}$ that is close to one suggests that the
standard deviations are accurately predicted. On the other hand, a larger or
lower value indicates that the uncertainties are under - or overestimated,
respectively.

Note that the sums in all expressions above run over the graphs in the test
set.

##  4 Results

In this section, we present the results we obtain from training the GNN model.
We then show the analytic approximations that were found using symbolic
regression.

###  4.1 GNN Results

Figure 2: This figure compares the accuracies of the GNN model with the
analytic formulae obtained from symbolic regression and the modified equations
using physical principles. Left: We first train a GNN with a compressed latent
space representation to perform likelihood-free inference for the cosmological
parameter $\Omega_{\rm m}$ with halo catalogues containing the positions and
velocity moduli of the halos. Evidently, the model is able to achieve very
high accuracy with a mean relative error of only $\sim 6.4\%$ when evaluated
on the test set of Gadget simulations. Despite its reduced dimensionality,
this accuracy is comparable to the model found in Shao et al. (2022a). Right:
We then use symbolic regression to extract analytic expressions for each MLP
in the message-passing and final aggregation layers of the GNN. After
modifying them to reduce their complexity and to preserve the symmetries of
the model, we evaluate the expressions on the Gadget test set. As shown, the
expressions are able to maintain the accuracy of the GNN, with an error of
only $\sim 6.7\%$, indicating that the equations are close approximations for
the learned GNN relations.

We first train a GNN with a single message-passing layer and fix the number of
hidden features to two. Using Bayesian optimization of the hyperparameters, we
find that the optimal linking radius is $\sim 1.35~{}h^{-1}{\rm Mpc}$ which
describes the characteristic length scale of the model. When we evaluate the
trained model on a test set of Gadget simulations, we find that it is able to
attain very accurate predictions of $\Omega_{\rm m}$ with a mean relative
error of $6\%$ and a $\chi^{2}$ of 1.37. This indicates that both the
posterior mean and standard deviations are accurately inferred. These results
are depicted in the left panel of Fig. 2. Hence, we see that the accuracy of
the model is not significantly compromised by the reduction in the dimensions
of its latent space with respect to the model used in Shao et al. (2022a),
which was $\sim 5.6\%$. In the following two sections, we present the results
for testing the equations on halos from the six different $N$-body simulations
and four hydrodynamic simulations. We then present the predictions for the
model tested on galaxies from six different hydrodynamic simulation suites.

####  4.1.1 Halos

We first find that the model is robust to different $N$-body codes despite
being trained on halo catalogues from only the Gadget simulations, agreeing
with the results discussed in Shao et al. (2022a). The simulations we use for
this test are Abacus, Ramses, PKDGrav3, Enzo, and CUBEP3M which share the same
cosmology and initial conditions but employ different numerical methods, as
described in Section 2. As shown in the top panel of Fig. 3, the model obtains
similar constraints on $\Omega_{\rm m}$ for these simulations. We present more
detailed results of this test in Appendix A, where the figures depict the
accuracies of the model when tested on 50 catalogues of different cosmologies
from each simulation code.

Figure 3: Top: We train a GNN model with a compressed latent space to perform
likelihood-free inference for the cosmological parameter $\Omega_{\rm m}$. The
input to the model are halo catalogues from Gadget that only carry information
about halo positions and peculiar velocity moduli. Once trained, we test the
model on halo catalogues from different $N$-body and hydrodynamic simulations
as indicated in the legend. We note that simulations of the same type, either
$N$-body or hydrodynamic, are run with the same initial conditions, cosmology
(and fiducial astrophysics for the hydrodynamic simulations). For each
simulation, we generate 5 catalogues. Each halo catalogue contains all halos
with masses above $Nm_{p}$, where $m_{p}$ is the particle mass and $N$ can be
100, 200, 300, 400, or 500 (see legend). The y-axis represents the difference
between the truth and the inference. As can be seen, this model exhibits
surprising extrapolation properties and is robust to all simulation codes
despite only containing one message-passing layer and two latent features.
Bottom: Same as above but for the analytic equations obtained using symbolic
regression and modified to preserve rotational and translational symmetries in
the data, as described in Section 3.3. As can be seen, the formulae maintain
the robustness of the GNN model and achieve a very similar accuracy compared
to the GNN.

Moreover, the model is robust to different hydrodynamic codes. When tested on
halo catalogues from the IllustrisTNG, SIMBA, Astrid, and Magneticum
simulations, the GNN is able to achieve similar precision and accuracy
compared to the predictions for the $N$-body codes, as seen in the top panel
of Fig. 3. This demonstrates that the model is robust even to hydrodynamics,
varying astrophysical parameters, and different subgrid physics models. This
agrees with the results from Shao et al. (2022a) and shows that even with a
reduced latent space dimensionality, the model could possibly still be
learning a fundamental relation between the halo properties and $\Omega_{\rm
m}$. However, note that by reducing the size of the latent space the precision
of the predictions decreases slightly, which is expected.

Another test that we performed to gauge the extent of the robustness of the
GNN is evaluating our model on halos generated using a halo finder that is
different (Subfind) from the one used during training (Rockstar). We find that
the model is able to extrapolate to these halos and we present the details of
this test in Appendix B.

####  4.1.2 Galaxies

We also asked if the network would extrapolate to galaxy distributions after
being trained on only the positions and velocities of $N$-body halos. Hence,
we test the GNN on galaxy catalogues from the following hydrodynamic
simulations: Astrid, IlustrisTNG, Magneticum, SB28, SIMBA, and SWIFT-EAGLE. As
per the halo catalogues employed in the previous sections, the galaxy
catalogues used to perform the following tests contain the galaxy positions
and velocity moduli.

We find that the GNN is unable to accurately predict $\Omega_{\rm m}$ for all
galaxy catalogues of each simulation. We include the results in Fig. 9 of
Appendix C. This is not surprising given that the GNN was trained on $N$-body
simulations and hence was not given any information regarding the intricate
astrophysical and baryonic processes in galaxy distributions. Moreover, the
halo-galaxy connection is known to be a complex and challenging relation
(Moster et al., 2018; Behroozi et al., 2019).

###  4.2 Analytic Approximations

Here we present the equations extracted from the trained GNN model using the
symbolic regression method explained in Section 3.3. The formulae for each of
the hidden edge and node features, as well as for the predicted posterior mean
from the final MLP, are listed in Table 1. The listed RMSE values are computed
by individually replacing the corresponding component in the GNN architecture
with each expression while keeping all other components of the GNN unchanged
and evaluating them on halo catalogues of the Gadget test set. The computed
RMSE values are used to gauge the error that each approximate equation
introduces.

It is important to note that the variables $v_{i}$ and $v_{j}$ in the
equations represent the initial edge features or velocity moduli. As explained
in Section 3.2.3, these variables were normalized by the mean and standard
deviation of the velocity modulus for the halos from the training set to
ensure that all terms in the equations are dimensionless. Hence, the velocity
modulus terms in the equations are $v_{i}=\frac{v_{i}-\mu}{\delta}$ and
$v_{j}=\frac{v_{j}-\mu}{\delta}$ where $\mu=189$ km$\rm s^{\rm-1}$ is a fixed
value that was the computed mean velocity modulus for all halos in the
training set and $\delta$ is treated as a free parameter. For testing on halo
catalogues, we set $\delta=129$ km$\rm s^{\rm-1}$ which is equal to the value
used during training and was the standard deviation computed for all halos in
the training set. On the other hand, for testing on galaxy catalogues, we tune
$\delta$ to fit to each hydrodynamic simulation set as listed in Section 2
because we find that using the value $\delta=129$ km$\rm s^{\rm-1}$ leads to
inaccurate predictions. This is not surprising given that this value was
computed for $N$-body halos which would not be expected to extrapolate to
galaxies. Hence, it is possible that tuning it for different simulations can
account for the halo-galaxy bias. We discuss this in more detail in Section
4.2.2.

We also note that the presented edge equations were modified to include terms
that depend only on the relative velocity moduli of the halos and their
neighbors. This was done to simplify the equations and improve their
interpretability. Moreover, including only the relative velocity modulus as
opposed to arbitrary linear combinations of $v_{i}$ and $v_{j}$ (see equations
in Table 3) enforces the symmetry between the information from the velocity of
a halo and its neighbor. Furthermore, as described in Section 2, the halo
velocity moduli that appear in all the equations are defined with respect to
the simulation box, implying that the equations also preserve Galilean
invariance. We note that this modification improves the accuracy of the
equations compared to the original expressions found by the symbolic
regression algorithm. We include more details on this result, as well as the
original equations found by the symbolic regression algorithm, in Appendix D.
In the following discussions, we only refer to the modified equations.

Table 1: This table lists the analytic formulae obtained using symbolic
regression for each component of the learned GNN model: the edge model, node
model, and the MLP in the final aggregation layer. The last column lists the
RMSE values of the analytic expressions when they are individually substituted
into the GNN architecture. This evaluation is done by replacing the
corresponding MLP in the edge model, node model, or final aggregation layer
with the symbolic approximation while keeping all other components of the GNN
unchanged. When these approximations replace all components of the GNN
architecture, the RMSE of the predictions is 0.026, as shown in Fig. 2. We
note that the edge equations have been modified based on physical motivations
to preserve the symmetries of the data. Specifically, we modified the edge
equations to depend only on relative velocity moduli $v_{i}-v_{j}$, rather
than individual halo velocity modulus terms. This is done to enforce the
parity between the information from the velocity of a halo and its neighbor.
Compared to the predictions shown in Fig. 10, we see that using these modified
equations improves the overall accuracy of the predictions.
The way to use these equations is as follows. First, given a halo/galaxy catalogue, a mathematical graph is constructed by considering the halos/galaxies as nodes and linking nodes by edges if their distance is smaller than $r_{\rm link}=1.35~{}h^{-1}{\rm Mpc}$ (see Sec. 3.2.1 for details). Second, the feature of node $i$ is defined as $v_{i}=(|\vec{v}_{i}|-\mu)/\delta$, where $|\vec{v}_{i}|$ is the velocity modulus of halo/galaxy $i$, $\mu=189~{}{\rm km~{}s^{-1}}$, and $\delta$ is a free parameter with units of km $\rm s^{-1}$ that needs to be adjusted for galaxy catalogues (see Section 4.2.2 and table LABEL:tab:param_values for more details). Third, the edge features $\beta_{ij}$ and $\gamma_{ij}$ between nodes $i$ and $j$ are computed using Eqs. 4 and 5, respectively. Fourth, the updated edge features of the graph are computed using the below first two equations. Fifth, the updated node features are computed using the below third and fourth equations. Finally, from the updated graph we can estimate $\Omega_{\rm m}$ by using the below fifth equation. GNN Component | Formula | RMSE
---|---|---
Edge Model: $e^{(1)}_{1}$ | $1.32|v_{i}-v_{j}+0.21|+0.12(v_{i}-v_{j})-0.12(\gamma_{ij}+\beta_{ij}-1.73)$ | 0.03
Edge Model: $e^{(1)}_{2}$ | $|1.62(v_{i}-v_{j})+0.45|+1.98(v_{i}-v_{j})+0.55$ | 0.04
Node Model: $v^{(1)}_{1}$ | $1.21^{v_{i}}(0.77^{3.29\sum_{j\in\mathcal{N}_{j}}e^{(1)}_{1}+\sum_{j\in\mathcal{N}_{j}}e^{(1)}_{2}})+0.12$ | 0.02
Node Model: $v^{(1)}_{1}+v^{(1)}_{2}$ | $0.78-\sqrt{\log(0.16^{\sum_{j\in\mathcal{N}_{j}}e_{2}+\sum_{j\in\mathcal{N}_{j}}e_{1}-0.41v_{i}-1.05})}+1.45$ | 0.03
Final MLP: $\mu_{\Omega_{\rm m}}$ | $4\times 10^{-4}\cdot(-5.5\sum_{i\in\mathcal{G}}v^{(1)}_{2}+2.21\sum_{i\in\mathcal{G}}v^{(1)}_{1}+|0.96\sum_{i\in\mathcal{G}}v^{(1)}_{2}+0.82\sum_{i\in\mathcal{G}}v^{(1)}_{1}|)-0.103\ $ | 0.03

The accuracy of the equations when evaluated on the halo catalogues of the
Gadget simulations is shown in the right panel of Fig. 2. As can be seen,
these analytic approximations achieve similar mean relative error ($6.7\%$)
and RMSE ($2.6\times 10^{-2}$) as the GNN, suggesting that they are accurate
representations of the trained network. We emphasize that our analytic formula
predicts the posterior mean while the error bars (posterior standard
deviation) are obtained from the GNN discussed in Section 4.1.

In the following two sections, we present the results for testing the
equations on halos from the six different $N$-body simulations and four
distinct hydrodynamic simulation codes, as well as galaxies from six different
hydrodynamic simulation sets.

####  4.2.1 Halos

We first test the robustness of the analytic equations by evaluating them on
halos of the different $N$-body simulations, as we did with the GNN. We find
the analytic formulae to be accurate across all simulations, with predictions
of comparable mean relative errors as depicted in the lower panel of Fig. 3.
We note that in some cases, the analytic expressions are able to extrapolate
better than the GNN due to their known improved generalization abilities (e.g
see Shao et al. (2022b)). For instance, certain numerical artifacts that
appear in the predictions made by the GNN for boundary cases such as halo
catalogues generated with 100 or 500 minimum particle thresholds, are not
present in the predictions made by the analytic expressions. We elaborate on
this in Appendix A. Again, this suggests that the found formulae might
represent fundamental relations between the halo properties and the
cosmological parameter, $\Omega_{\rm m}$, as they are not affected by the
additional astrophysical processes such as gas cooling and AGN feedback.
Similar to the GNN, we perform a second robustness test using halo catalogues
generated with SUBFIND and find that the equations reach comparable
accuracies. See Appendix B for more details and plots. For all these tests,
the depicted errorbars are represent the inferred posterior standard deviation
values obtained by the GNN model trained on the halo catalogues since we do
not find an expression for this value, as discussed in Section 3.3.

####  4.2.2 Galaxies

We also test the equations on galaxy catalogues from the six hydrodynamic
simulation suites: Astrid, IlustrisTNG, Magneticum, SB28, SIMBA, and SWIFT-
EAGLE. We emphasize that this is not a trivial task as the GNN and the
corresponding equations were trained using dark matter halos from $N$-body
simulations that do not contain any information about the intergalactic
dynamics or baryonic processes present in hydrodynamic simulations. There is
also a complex galaxy-halo connection which can, for instance, be reflected in
the relative abundances of halos and galaxies where larger halos can contain
multiple galaxies while smaller halos may not contain any. These biases can
possibly leave a significant imprint in the relations between the relative
position and velocity terms of the equations found for halos. For these tests,
we follow the definitions of galaxies and stellar mass thresholds discussed in
Section 3.1 in constructing the galaxy catalogues where we include both
central and satellite galaxies.

Figure 4: We test the analytic equations that were trained for halo catalogues
of $N$-body simulations on thousands of galaxies from 6 different hydrodynamic
simulation sets: Astrid, IllustrisTNG, Magneticum, SB28, Simba, and SWIFT-
EAGLE, to predict the value of $\Omega_{\rm m}$ and plot the predicted against
truth for each simulation. To conserve space, we only present results for the
tests performed on catalogues constructed with a stellar mass threshold of
$4\times m_{*}$ where $m_{*}$ is a fixed mass for an individual stellar
particle as described in Section 3.1, but we reach similar accuracies for
catalogues constructed with other mass cuts. We also include only 50 randomly
selected catalogues for each simulation set for the clarity of the figures but
the reported metrics were computed for all simulations in the suites. Note
that for the bottom right panel, which depicts the predictions for the SWIFT-
EAGLE simulation set, we use simulations that are generated with the same
value of $\Omega_{\rm m}=0.3$. Thus, we plot the difference between the truth
and the prediction on the y-axis for these catalogues. As depicted in Fig. 5,
a large fraction of the catalogues, particularly for the ASTRID, IllustrisTNG,
and SB28 simulations, contain galaxy number densities that are outside the
range of the number densities exhibited by the halo catalogues used during
training of the network and equations. Hence, in this plot we remove these
outliers and find that the mean relative errors of the predictions
significantly decrease (see Fig. 5 for comparison). These results exhibit a
relatively high accuracy with mean errors that average around $\epsilon\sim
9.4\%$, comparable to the accuracies obtained by our companion paper (de Santi
et al., 2023) with model trained on galaxy properties. This further
demonstrates the robustness of the equations as well as their ability to use
halo properties to extrapolate to galaxy distributions. This is a surprising
result given the various astrophysical processes exhibited by the hydrodynamic
simulations and the complex mapping between galaxies and halos. Figure 5: This
figure follows the format of Fig. 4. Here, each scatterpoint (representing one
galaxy catalogue) is colored according to the number of galaxies the catalogue
contains. The colorbar depicts the range of galaxy number density present in
the catalogues for the corresponding stellar mass threshold of each column. As
it can be seen, a significant portion of the galaxy catalogues from
simulations such as Astrid, Illustris-TNG, and SB28 contain much smaller or
larger galaxy number densities than the number densities seen during training,
which were within the range of $(1000,6000)$. These catalogues account for the
the relatively larger errors in these predictions is expected because the halo
number density acts as an uninformative prior during the training of the GNN
and equations. When we omit these outlier catalogues, we obtain smaller
scatter in the results, as shown in Fig. 4.

We present the results for evaluating the equations on galaxy catalogues from
the different hydrodynamic simulations in Fig. 4. Each panel is labeled with
the corresponding simulation suite. For simplicity, we present the predictions
for only the galaxy catalogues generated with the stellar mass threshold of
$4\times m_{*}$ for a fixed $m_{*}$ denoting the mass of an individual stellar
particle as described in Section 3.1. However, we find that the equations are
able to perform with similar accuracies for catalogues constructed with
different mass cuts, which we discuss further in Appendix E. Moreover, since
the simulations from the SWIFT-EAGLE suite are run with the same value of
$\Omega_{\rm m}$, we plot the difference between the true ($\Omega_{\rm
m}=0.3$) and the predicted values on the y-axis for these catalogues. We note
that the presented errorbars for all simulations are the inferred posterior
standard deviation values obtained by the model trained and tested on galaxy
catalogues discussed in de Santi et al. (2023), since the equations predict
only the first moment of the posterior for $\Omega_{\rm m}$ (see Section 3.3).

There are are several important features to note for evaluating the equations
on galaxy catalogues from the different hydrodynamic simulations. First, for
each simulation we tune the parameter $\delta$ to improve the accuracy of the
predictions. As discussed in Section 4.2, this parameter appears in the
equation as a normalization of the velocity modulus terms $v_{i}$ and $v_{j}$,
and its value varies for different hydrodynamic simulations when testing on
galaxies. We tune this normalization because we noticed that using the
original value $\delta=129$ km$\rm s^{\rm-1}$, the standard deviation of the
velocity moduli for all halos in the training set, resulted in predictions
that deviated from the truth in terms of a slope and bias, which varies for
each simulation. Thus, in Table LABEL:tab:param_values we list the values of
$\delta$ that we optimize for each simulation using non-linear least squares
with scipy-optimize333https://docs.scipy.org/doc/scipy/reference/optimize.html
for the catalogues constructed using the $4\times m_{*}$ stellar mass
threshold. We also compare these found values with the $\delta$ used to
evaluate on halo catalogues in the table and in later discussions.

Second, after tuning this parameter, we find that the equations are able to
predict $\Omega_{\rm m}$ with mean relative errors of $15.35\%$ for ASTRID,
$12.85\%$ for Illustris-TNG, $6.89\%$ for Magneticum, $16.17\%$ for SB28,
$8.50\%$ for SIMBA, and $4.08\%$ for SWIFT-EAGLE, across the four stellar mass
thresholds. Evidently, the predictions for the galaxy catalogues from ASTRID,
SB28, and Illustris-TNG exhibit significantly larger error than for the halo
catalogues. This can be explained by two reasons. One, there are additional
astrophysical processes and dynamics present in the thousands of hydrodynamic
simulations that can interfere with the equations’ extrapolation ability.
Given that the equations can only encode information regarding the
gravitational interactions between halos from $N$-body simulations, the
effects of these various astrophysical parameters may impede on the accuracy
of the predictions. Moreover, there are likely to be significantly more
outliers for simulations such as SB28, where we vary 28 cosmological and
astrophysical parameters at a time. This is also true for the ASTRID
simulations which encompass a wider range of galaxy properties and are able to
encapsulate the variations found in the other simulation suites. A more
detailed discussion of the wide range of characteristics in the ASTRID
simulations can be found in our companion papers, de Santi et al. (2023) and
in Ni et al. (2023) (in preparation).

Two, there is a large fraction of the galaxy catalogues that contain galaxy
number densities outside the scope of the halo number densities seen by the
GNN and equations during training. For instance, the number of halos in
catalogues from the Gadget simulations used for training ranges from
$\sim$1,000 to 6,000. However, there are galaxy catalogues that contain fewer
than 500 galaxies at this stellar mass threshold. These outliers are
particularly dominant in the IllustrisTNG, Astrid, and SB28 simulations, which
leads to under-predicted values of $\Omega_{\rm m}$. This effect can be seen
in Fig. 5 which contains the same plots as Fig. 4 but with each scatter-point
colored according to the galaxy number density that the catalogue contains.
The colorbars accompanying each plot indicate the range of the galaxy number
densities present in the catalogues. As it can be seen, in catalogues with
significantly lower (higher) galaxy number densities compared to those seen in
training, the value of $\Omega_{\rm m}$ is often under- (over-) predicted that
contributes to the large scatter. On the other hand, if one removes these
outliers, the mean relative errors significantly decrease. Hence, Fig. 4
depicts the results for only the catalogues with galaxy number densities that
fall within the range of $(1000,6000)$. Restricting to these catalogues
decreases the mean relative errors to: $9.76\%$ for Astrid, $10.34\%$ for
IllustrisTNG, $7.02\%$ for Magneticum, $12.24\%$ for SB28, $8.29\%$ for SIMBA,
and $4.08\%$ for SWIFT-EAGLE. Thus, we conclude that the equations are able to
extrapolate to galaxies with with accuracies that are comparable to those
attained for the halo catalogues from hydrodynamic simulations. These results
are also comparable to those obtained by our companion paper (de Santi et al.,
2023) where we trained a model directly on galaxy properties. We note that the
effect of the number density being an uninformative prior during the learning
process can be diminished by broadening the range of halo number densities
used to train the network and equations, but we leave this for future work.

Table 2: In this table, we list the optimized values of the free parameter, $\delta$, which is the normalization of the velocity modulus terms used in the analytic expressions (see 4.2). We list the values for the six different hydrodynamic sets, ASTRID, Illustris-TNG, Magneticum, SB28, SIMBA, and SWIFT-EAGLE. These values were obtained using linear least squares optimization with scipy-optimize as described in Section 4.2.2 to achieve robustness across various simulation codes. We also include the $\delta$ used for testing on $N$-body halos for comparison. Simulation | $\delta$ | Simulation | $\delta$
---|---|---|---
$N$-body codes | $129.2$ | SB28 | $100.0$
ASTRID | $126.5$ | SIMBA | $122.5$
Illustris-TNG | $99.6$ | SWIFT-EAGLE | $114.5$
Magneticum | $147.2$ |  |

After accounting for the aforementioned details, we conclude that the
equations are able to accurately predict $\Omega_{\rm m}$ for galaxy
catalogues. We emphasize that the ability of the equations to achieve a
reasonable inference of $\Omega_{\rm m}$, being trained on halo catalogues
from $N-$body codes, is a surprising result because it is expected that
baryonic effects will affect the abundance and clustering of galaxies in a
complex and unknown manner. This is particularly astounding for simulations
such as those from the SB28 suite that covers a vast volume in parameter space
with many regions not covered by the training set (e.g. cosmological
parameters like $h$, $n_{s}$, and $\Omega_{\rm b}$). Furthermore, the
equations work really well for SWIFT-EAGLE catalogues that were created
running a different halo/subhalo finder than the one used for training.
Furthermore, the equations are robust to the nontrivial galaxy-halo connection
as they can map the information learned about the halo position and velocity
fields to those for galaxies. The ability of the equations to remain robust to
these variations provide strong indication that they may be relying on
fundamental relations in the galaxy and halo phase-space distribution that
encodes effective information on $\Omega_{\rm m}$. Another possibility is that
the equations are extracting information on scales unaffected by astrophysical
dynamics. In the next section, we explore possible interpretations of these
equations in more detail.

##  5 Discussion

Here, we discuss some speculative interpretations of the found equations. We
attempt to only explain the formulae for the edge models because their
functional forms are simpler than those for the node models. The edge model
also solely employs physical information about the halo positions and velocity
moduli so they are responsible for directly leveraging the clustering and
distribution of the halos. This aligns with the analysis from Cranmer et al.
(2019a), where it was argued that the relations used in the edge models of
GNNs are analogous to describing the force laws between pairs of particles in
physical systems. We will elaborate on how the edge equations found in this
work may also reflect physical relations pertaining to the halo and galaxy
populations. The node model, on the other hand, exhibits a more complex form
because it introduces non-linearities to the formulae and makes use of
information pertaining to the aggregate features from all neighboring halos.
However, this should not suggest that the equations for the node model contain
information that is less important than those for the edge model.

###  5.1 Relative Peculiar Velocity Modulus

In both edge model equations, $e^{(1)}_{1}$ and $e^{(1)}_{2}$, the information
regarding the velocities of the halos appear in terms in the form of
$(v_{i}-v_{j})$, which indicates that the model is taking advantage of the
relative velocity moduli of the halos and their neighbors. This dependence
also preserves the parity between the information content of a halo and that
of its neighbor. The ability for the edge model in the GNN to employ
relational information between pairs of bodies of a system has been a
recognized advantage (Cranmer et al., 2019a, 2020) towards understanding the
physical principles underlying the model predictions. We believe that in this
case, using the relative velocities allows the models to gauge the local
gravitational forces where the relative velocity moduli between two halos can
serve as a proxy for the depth of the potential wells in the bound system.
This is reasonable since larger relative speeds of interacting bodies can
result from the presence of stronger attractive forces between them. From
this, the model may be learning a representation of the masses of the halos.
An analogous discussion in Cen et al. (1994) reached similar conclusions
pertaining to the pairwise peculiar velocities and speeds which were found to
have strong dependence on $\Omega_{\rm m}$ at the same small scale as that
used by the models in this work ($\lesssim 5~{}h^{-1}$Mpc).

We also speculate that the presence of these terms reflect the strong
dependence of $\Omega_{\rm m}$ on the information available in the cosmic
velocity fields (Bernardeau et al., 1995; Dekel, 1994). For instance,
Bernardeau et al. (1995) discusses a derived relation between the moments of
the scalar field of the peculiar velocity divergence and $\Omega_{\rm m}$ that
is independent of the biasing between the distribution of galaxies and the
underlying dark matter density field. It is possible that the found
expressions in this work reflect a similar relationship because our models
have been trained using the scalar halo velocity modulus and demonstrate an
accuracy that is not significantly affected by the presence of astrophysical
and baryonic effects. We speculate that the network and equations may be
correcting for the non-linearities of the galaxy velocity fields on smaller
scales by considering the galaxy distribution and number densities.
Specifically, the equations may be obtaining stochastic velocities from the
relative positions of galaxies using the baryonic physics present in the
hydrodynamic simulations. This information, coupled with the input pairwise
velocity moduli, may then be used to compute the contribution of the galaxy
velocities from the bulk flows that trace the large-scale structure of the
Universe. Since the bulk flows are a consequence of the mass continuity
equation which relates the large-scale density and growth rate, the equations
are able to extract cosmological information on $\Omega_{\rm m}$. A similar
argument was made in the formulation of the cosmic virial theorem Peebles
(1976, 1980) which constructs a relation between the mean square relative
peculiar velocity computed for galaxy pairs and the galaxy correlation
functions. Hence, we emphasize the importance of leveraging both the positions
and velocities of the halos/galaxies in the analytic expressions. This aligns
with previous findings that using only the positions or only the velocities
fails to achieve accurate inference (Villanueva-Domingo et al., 2022). Our
companion paper, de Santi et al. (2023), also reaches similar conclusions
about the amount of information contained in the galaxy phase-space. Moreover,
we have found that introducing additional halo properties such as the halo
mass and maximum circular velocity eliminates the generalization of the
expressions to various simulation codes (Shao et al., 2022a), which further
indicates the robustness of the information contained in peculiar velocities
for inferring $\Omega_{\rm m}$.

###  5.2 Velocity normalization

Here we also discuss the implications of tuning the normalization of the
velocity modulus terms, $\delta$, for galaxies from each simulation set.
Previous findings in Juszkiewicz et al. (1999, 2000) indicate that the halo-
galaxy distribution bias can induce biases in pairwise velocity statistics
defined using the radial separation between galaxies. Thus, we speculate that
the normalization of the velocity modulus terms $v_{i}$ and $v_{j}$ in our
equations reflect a similar correction to account for the fact that the
spatial clustering of galaxies may not trace that of the matter field. In that
case, it would expected for the values of $\delta$ to differ for various
galaxy populations. Since the optimal value of $\delta$ varies across
different hydrodynamic codes, we hypothesize that this parameter relates the
kinematics of the galaxy velocities to their abundances. For instance, as seen
in Table LABEL:tab:param_values, the value of $\delta$ is largest for the
Magneticum simulations which have been found to contain significantly higher
galaxy number densities compared to the other codes (de Santi et al., 2023).
Consequently, the disparity in optimal $\delta$ values can possibly reflect
the variations in the abundances of satellites in simulations of difference
codes since the peculiar motion of satellites are more sensitive to small
scale dynamics and their presence would thus contribute to a larger spread in
the dispersion of the peculiar velocity. On the other hand, the mean galaxy
number densities are smallest for IllustrisTNG and SB28, which can explain why
$\delta$ is smallest for these two simulations (see Table
LABEL:tab:param_values). We leave for future work to further investigate the
role of $\delta$ in the context of galaxy abundances, populations, and
cosmological inference.

###  5.3 Spatial distribution and clustering

Next, we discuss the implications of halo clustering and spatial distribution
in the found edge equations. In the first edge equation, $e_{1}$, the presence
of the terms $\beta$ and $\gamma$ reflect the spatial distribution of the
halos in the catalogues. Specifically, the variable $\gamma\in(0,1]$ describes
the distance between two halos where its range is restricted due to its
normalization by the linking radius, $r_{\rm link}\sim 1.35~{}h^{-1}{\rm
Mpc}$, as described in 3.2.1. Thus, a smaller $\gamma$ would indicate a denser
distribution of halos. Meanwhile, the variable $\beta\in[-1,1]$ describes the
angular orientation of a halo with respect to its neighbor and can provide
information about the shape of the distribution, e.g. the filamentary
structure of the cosmic web. Both parameters are used by the model to learn
about the presence of large scale structures such as superclusters and
filaments.

##  6 Conclusions

In this work we have found an analytic expression that approximates the
relation employed by a GNN that was trained to infer $\Omega_{\rm m}$ from
dark mater halo catalogues. This was motivated by the results of Shao et al.
(2022a) which found that GNNs are able to perform accurate field-level
inference of $\Omega_{\rm m}$ using halo catalogues from various $N$-body and
hydrodynamic simulations. These results imply that the found relation could be
a fundamental one as it is not affected by varying numerical errors,
astrophysical processes, subgrid physics, or even halo definitions. This
motivates us to gain a better understanding of the learned relation by
approximating it with symbolic equations that are more physically
interpretable than a neural network.

To derive the analytic approximations, we followed a two-step approach. We
first simplified the model that was used in the previous work to obtain a GNN
with reduced latent space dimensionality. The intention for this step was to
maintain the accuracy and precision of the model discussed in Shao et al.
(2022a) while building a less complex, and hence more easily interpretable,
architecture. We train our compressed model on catalogues that only contain
the positions and peculiar velocity moduli of dark matter halos from $N$-body
simulations. Next, we trained a symbolic regressor to fit equations to each
component of the trained GNN (see Scheme 1).

We summarize the main results of this work below:

  * •

We train a compressed GNN architecture composed of only 1 message passing
layer and 2 hidden features on halo catalogues from the Gadget $N$-body
simulations. We find that it is able to achieve precise constraints on
$\Omega_{\rm m}$ with a mean relative error of $\epsilon\sim 6.5\%$, similar
to the GNN model with larger latent space dimensionality as discussed in Shao
et al. (2022a), which achieved a mean relative error of $\epsilon\sim 5.6\%$.

  * •

The compressed GNN model, trained on Gadget simulations, is also robust across
thousands of halo catalogues generated from five different $N$-body codes
–Abacus, CUBEP3M, Enzo, PKDGrav3, Ramses– and four different hydrodynamic
codes that employ different galaxy formation implementations: Astrid,
IllustrisTNG, Magneticum, SIMBA. This model reproduces the results of Shao et
al. (2022a) where the non-triviality of this robustness was discussed.

  * •

We use symbolic regression to find equations that approximate the different
MLPs that our GNN model is comprised of. These analytic equations can
approximate the learned relation between $\Omega_{\rm m}$ and the input halo
properties with a mean relative error of $\epsilon\sim 6.7\%$ when evaluated
on halos from Gadget $N$-body simulations. We then evaluate the equations on
thousands of $N$-body and hydrodynamic simulations run with the different
codes listed above. Thus, we demonstrate that the equations are able to
reproduce the preciseness and robustness of the GNN, concluding that they are
successful approximations of the learned network.

  * •

We further find that the equations are able to extrapolate better than the GNN
in certain cases. Specifically, we test on galaxy catalogues from six
different hydrodynamic simulation suites and find that while the equations are
able to predict the value of $\Omega_{\rm m}$ accurately while the GNN is
unable to. This is a surprising feat given that the equations were trained
only on halo properties from $N$-body simulations and were not given any
information regarding the complex baryonic effects and astrophysical feedback
processes present in galaxy interactions. This also demonstrates that the
equations may be exploiting a relation between positions, velocities, and
$\Omega_{\rm m}$ that is independent of the halo-galaxy connection.

  * •

To obtain good accuracies in the galaxy catalogues we need to tune one single
free-parameter, $\delta$, which is the normalization of the velocity modulus
terms used in the analytic expressions. The value of $\delta$ appears to be
sensitive to the characteristics of the considered galaxy population. We leave
for future work in studying its physical role as well as the best strategy to
constrain it - such as fitting it using a subset of data, marginalizing over
its values, or others.

  * •

As in our companion paper (de Santi et al., 2023), we find some robustness to
super-sample covariance effects, although further work is needed to properly
assest it taking into account the setup we used to train our models. Further
details are presented in Appendix F.

  * •

We attempt to provide physical interpretation of the equations for the edge
component of the GNN, which could reflect physical laws and forces between
interacting objects represented by the nodes of the graph. Specifically, the
equations demonstrate an explicit dependence on the pairwise velocity modulus
and relative positions of halos/galaxies at separation distances $\lesssim
1.35~{}h^{-1}$Mpc. These dependencies illustrate how the rotational and
translational symmetries present in the data are maintained and exploited by
the model. Moreover, the dual reliance on the spatial and velocity fields of
the halos indicate that there is robust information embedded in the phase-
space distribution of halos, perhaps reflecting some underlying physical law
like the continuity equation. We draw speculative connections to past works
that have analyzed similar information in observational fields at the same
scales, such as the pairwise velocity and speed statistics as analyzed in Cen
et al. (1994); Juszkiewicz et al. (1999, 2000), and the use of cosmic velocity
fields as seen in Bernardeau et al. (1995); Dekel (1994).

##  7 Acknowledgments

We thank Lucy Reading-Ikkanda for creating Fig. 1. We thank Ravi Sheth, Oren
Slone, David Spergel, Ben Wandelt, Michael Strauss, Oliver Philcox, Gigi
Guzzo, Marina Silvia Cagliari, and Miles Cranmer for the enlightening
discussions. NSMS acknowledges financial support from FAPESP, grants
2019/13108-0 and 2022/03589-4. The CAMELS project is supported by the NSF
grant AST 2108078. EV is supported by NSF grant AST-2009309 and NASA grant
80NSSC22K0629. EH acknowledge supported by the grant agreements
ANR-21-CE31-0019 / 490702358 from the French Agence Nationale de la Recherche
/ DFG for the LOCALIZATION project. KD acknowledges support by the COMPLEX
project from the European Research Council (ERC) under the European Union’s
Horizon 2020 research and innovation program grant agreement ERC-2019-AdG
882679 as well as by the Deutsche Forschungsgemeinschaft (DFG, German Research
Foundation) under Germany’s Excellence Strategy - EXC-2094 - 390783311. TC is
supported by the INFN INDARK PD51 grant and the FARE MIUR grant
‘ClustersXEuclid’ R165SBKTMA. The research in this paper made use of the
SWIFT-EAGLE open-source simulation code (http://www.swiftsim.com, Schaller et
al., 2018) version 1.2.0.

## References

  * Akiba et al. (2019) Akiba, T., Sano, S., Yanase, T., Ohta, T., & Koyama, M. 2019, in Proceedings of the 25rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining
  * Angulo & Hahn (2022) Angulo, R. E., & Hahn, O. 2022, Living Reviews in Computational Astrophysics, 8, doi: 10.1007/s41115-021-00013-z
  * Bartlett et al. (2022) Bartlett, D. J., Desmond, H., & Ferreira, P. G. 2022, arXiv e-prints, arXiv:2211.11461.  https://arxiv.org/abs/2211.11461
  * Battaglia et al. (2018) Battaglia, P. W., Hamrick, J. B., Bapst, V., et al. 2018, CoRR, abs/1806.01261
  * Beck et al. (2016) Beck, A. M., Murante, G., Arth, A., et al. 2016, MNRAS, 455, 2110, doi: 10.1093/mnras/stv2443
  * Behroozi et al. (2019) Behroozi, P., Wechsler, R. H., Hearin, A. P., & Conroy, C. 2019, MNRAS, 488, 3143, doi: 10.1093/mnras/stz1182
  * Behroozi et al. (2013) Behroozi, P. S., Wechsler, R. H., & Wu, H.-Y. 2013, ApJ, 762, 109, doi: 10.1088/0004-637X/762/2/109
  * Bernardeau et al. (1995) Bernardeau, F., Juszkiewicz, R., Dekel, A., & Bouchet, F. R. 1995, Monthly Notices of the Royal Astronomical Society, 274, 20, doi: 10.1093/mnras/274.1.20
  * Bird et al. (2022) Bird, S., Ni, Y., Di Matteo, T., et al. 2022, MNRAS, 512, 3703, doi: 10.1093/mnras/stac648
  * Borrow & et. al. (2023) Borrow, J., & et. al. 2023, In preparation.
  * Borrow et al. (2022) Borrow, J., Schaller, M., Bahe, Y. M., et al. 2022, The impact of stochastic modeling on the predictive power of galaxy formation simulations, Tech. rep.  https://ui.adsabs.harvard.edu/abs/2022arXiv221108442B
  * Bronstein et al. (2021) Bronstein, M. M., Bruna, J., Cohen, T., & Velickovic, P. 2021, CoRR, abs/2104.13478
  * Bryan et al. (2014) Bryan, G. L., Norman, M. L., O’Shea, B. W., et al. 2014, ApJS, 211, 19, doi: 10.1088/0067-0049/211/2/19
  * Cañas et al. (2019) Cañas, R., Elahi, P. J., Welker, C., et al. 2019, MNRAS, 482, 2039, doi: 10.1093/mnras/sty2725
  * Cen et al. (1994) Cen, R., Bahcall, N. A., & Gramann, M. 1994, ApJ, 437, L51, doi: 10.1086/187680
  * Crain et al. (2015) Crain, R. A., Schaye, J., Bower, R. G., et al. 2015, MNRAS, 450, 1937, doi: 10.1093/mnras/stv725
  * Cranmer (2020) Cranmer, M. 2020, PySR: Fast & Parallelized Symbolic Regression in Python/Julia, Zenodo, doi: 10.5281/zenodo.4041459
  * Cranmer et al. (2020) Cranmer, M. D., Sanchez-Gonzalez, A., Battaglia, P. W., et al. 2020, CoRR, abs/2006.11287
  * Cranmer et al. (2019a) Cranmer, M. D., Xu, R., Battaglia, P., & Ho, S. 2019a, Learning Symbolic Physics with Graph Networks, arXiv, doi: 10.48550/ARXIV.1909.05862
  * Cranmer et al. (2019b) —. 2019b, Learning Symbolic Physics with Graph Networks, arXiv, doi: 10.48550/ARXIV.1909.05862
  * Davé et al. (2019) Davé, R., Anglés-Alcázar, D., Narayanan, D., et al. 2019, MNRAS, 486, 2827, doi: 10.1093/mnras/stz937
  * Davis et al. (1985) Davis, M., Efstathiou, G., Frenk, C. S., & White, S. D. M. 1985, ApJ, 292, 371, doi: 10.1086/163168
  * de Putter et al. (2012) de Putter, R., Wagner, C., Mena, O., Verde, L., & Percival, W. J. 2012, Journal of Cosmology and Astroparticle Physics, 2012, 019, doi: 10.1088/1475-7516/2012/04/019
  * de Santi et al. (2023) de Santi et al., N. S. M. 2023, In preparation.
  * Dekel (1994) Dekel, A. 1994, ARA&A, 32, 371, doi: 10.1146/annurev.aa.32.090194.002103
  * Delgado et al. (2022) Delgado, A. M., Wadekar, D., Hadzhiyska, B., et al. 2022, Monthly Notices of the Royal Astronomical Society, 515, 2733, doi: 10.1093/mnras/stac1951
  * Dolag et al. (2009) Dolag, K., Borgani, S., Murante, G., & Springel, V. 2009, MNRAS, 399, 497, doi: 10.1111/j.1365-2966.2009.15034.x
  * Dolag et al. (2004) Dolag, K., Jubelgas, M., Springel, V., Borgani, S., & Rasia, E. 2004, ApJ, 606, L97, doi: 10.1086/420966
  * Elahi et al. (2019) Elahi, P. J., Cañas, R., Poulton, R. J. J., et al. 2019, PASA, 36, e021, doi: 10.1017/pasa.2019.12
  * Fey & Lenssen (2019) Fey, M., & Lenssen, J. E. 2019, Fast Graph Representation Learning with PyTorch Geometric, 2.0.2.  https://github.com/pyg-team/pytorch_geometric
  * Fluri et al. (2019) Fluri, J., Kacprzak, T., Lucchi, A., et al. 2019, Phys. Rev. D, 100, 063514, doi: 10.1103/PhysRevD.100.063514
  * Garrison et al. (2021) Garrison, L. H., Eisenstein, D. J., Ferrer, D., Maksimova, N. A., & Pinto, P. A. 2021, MNRAS, 508, 575, doi: 10.1093/mnras/stab2482
  * Greengard & Rokhlin (1987) Greengard, L., & Rokhlin, V. 1987, Journal of Computational Physics, 73, 325, doi: 10.1016/0021-9991(87)90140-9
  * Groth et al. (2023) Groth, F., Steinwandel, U. P., Valentini, M., & Dolag, K. 2023, arXiv e-prints, arXiv:2301.03612, doi: 10.48550/arXiv.2301.03612
  * Gupta et al. (2018) Gupta, A., Matilla, J. M. Z., Hsu, D., & Haiman, Z. 2018, Phys. Rev. D, 97, 103515, doi: 10.1103/PhysRevD.97.103515
  * Hamilton et al. (2006) Hamilton, A. J. S., Rimes, C. D., & Scoccimarro, R. 2006, Monthly Notices of the Royal Astronomical Society, 371, 1188, doi: 10.1111/j.1365-2966.2006.10709.x
  * Hamilton (2020) Hamilton, W. L. 2020, Synthesis Lectures on Artificial Intelligence and Machine Learning, 14, 1
  * Harnois-Déraps et al. (2013) Harnois-Déraps, J., Pen, U.-L., Iliev, I. T., et al. 2013, MNRAS, 436, 540, doi: 10.1093/mnras/stt1591
  * Hirschmann et al. (2014) Hirschmann, M., Dolag, K., Saro, A., et al. 2014, MNRAS, 442, 2304, doi: 10.1093/mnras/stu1023
  * Hockney & Eastwood (1988) Hockney, R. W., & Eastwood, J. W. 1988, Computer simulation using particles
  * Hopkins (2015) Hopkins, P. F. 2015, MNRAS, 450, 53, doi: 10.1093/mnras/stv195
  * Hu & Kravtsov (2003) Hu, W., & Kravtsov, A. V. 2003, The Astrophysical Journal, 584, 702, doi: 10.1086/345846
  * Jeffrey & Wandelt (2020) Jeffrey, N., & Wandelt, B. D. 2020, in 34th Conference on Neural Information Processing Systems, Online Conference, Canada.  https://hal.archives-ouvertes.fr/hal-03047530
  * Jubelgas et al. (2004) Jubelgas, M., Springel, V., & Dolag, K. 2004, MNRAS, 351, 423, doi: 10.1111/j.1365-2966.2004.07801.x
  * Juszkiewicz et al. (2000) Juszkiewicz, R., Ferreira, P. G., Feldman, H. A., Jaffe, A. H., & Davis, M. 2000, Science, 287, 109, doi: 10.1126/science.287.5450.109
  * Juszkiewicz et al. (1999) Juszkiewicz, R., Springel, V., & Durrer, R. 1999, The Astrophysical Journal, 518, L25, doi: 10.1086/312055
  * Kaiser (1987) Kaiser, N. 1987, MNRAS, 227, 1, doi: 10.1093/mnras/227.1.1
  * Lemos et al. (2022) Lemos, P., Jeffrey, N., Cranmer, M., Ho, S., & Battaglia, P. 2022, Rediscovering orbital mechanics with machine learning, arXiv, doi: 10.48550/ARXIV.2202.02306
  * Loshchilov & Hutter (2017) Loshchilov, I., & Hutter, F. 2017, arXiv e-prints, arXiv:1711.05101.  https://arxiv.org/abs/1711.05101
  * Ma et al. (2015) Ma, Y.-Z., Li, M., & He, P. 2015, A&A, 583, A52, doi: 10.1051/0004-6361/201526051
  * Marinacci et al. (2018) Marinacci, F., Vogelsberger, M., Pakmor, R., et al. 2018, Monthly Notices of the Royal Astronomical Society, 480, 5113, doi: 10.1093/mnras/sty2206
  * Moster et al. (2018) Moster, B. P., Naab, T., & White, S. D. M. 2018, MNRAS, 477, 1822, doi: 10.1093/mnras/sty655
  * Naiman et al. (2017) Naiman, J. P., Pillepich, A., Springel Enrico Ramirez-Ruiz, V., et al. 2017, ArXiv e-prints, 1707.03401.  https://arxiv.org/abs/1707.03401
  * Nelson et al. (2018) Nelson, D., Pillepich, A., Springel, V., et al. 2018, Monthly Notices of the Royal Astronomical Society, 477, 1206, doi: 10.1093/mnras/stx3040
  * Nelson et al. (2019) Nelson, D., Springel, V., Pillepich, A., et al. 2019, Computational Astrophysics and Cosmology, 6, 2, doi: 10.1186/s40668-019-0028-x
  * Ni et al. (2022) Ni, Y., Di Matteo, T., Bird, S., et al. 2022, MNRAS, 513, 670, doi: 10.1093/mnras/stac351
  * Ni et al. (2023) Ni et al., Y. 2023, In preparation.
  * Ntampaka et al. (2019) Ntampaka, M., Eisenstein, D. J., Yuan, S., & Garrison, L. H. 2019, arXiv e-prints, arXiv:1909.10527.  https://arxiv.org/abs/1909.10527
  * Nusser & Dekel (1992) Nusser, A., & Dekel, A. 1992, ApJ, 391, 443, doi: 10.1086/171360
  * Nusser & Dekel (1993) —. 1993, ApJ, 405, 437, doi: 10.1086/172376
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., et al. 2019, arXiv e-prints, arXiv:1912.01703.  https://arxiv.org/abs/1912.01703
  * Peebles (1976) Peebles, P. J. E. 1976, Ap&SS, 45, 3, doi: 10.1007/BF00642136
  * Peebles (1980) —. 1980, The large-scale structure of the universe (Princeton University Press)
  * Pillepich et al. (2018a) Pillepich, A., Springel, V., Nelson, D., et al. 2018a, MNRAS, 473, 4077, doi: 10.1093/mnras/stx2656
  * Pillepich et al. (2018b) Pillepich, A., Nelson, D., Hernquist, L., et al. 2018b, MNRAS, 475, 648, doi: 10.1093/mnras/stx3112
  * Potter et al. (2017) Potter, D., Stadel, J., & Teyssier, R. 2017, Computational Astrophysics and Cosmology, 4, 2, doi: 10.1186/s40668-017-0021-1
  * Ravanbakhsh et al. (2017) Ravanbakhsh, S., Oliva, J., Fromenteau, S., et al. 2017, Estimating Cosmological Parameters from the Dark Matter Distribution, arXiv, doi: 10.48550/ARXIV.1711.02033
  * Ribli et al. (2019) Ribli, D., Pataki, B. Á., Zorrilla Matilla, J. M., et al. 2019, MNRAS, 490, 1843, doi: 10.1093/mnras/stz2610
  * Sargent & Turner (1977) Sargent, W. L. W., & Turner, E. L. 1977, ApJ, 212, L3, doi: 10.1086/182362
  * Sato et al. (2009) Sato, M., Hamana, T., Takahashi, R., et al. 2009, The Astrophysical Journal, 701, 945, doi: 10.1088/0004-637x/701/2/945
  * Schaller et al. (2016) Schaller, M., Gonnet, P., Chalk, A. B. G., & Draper, P. W. 2016, SWIFT: Using Task-Based Parallelism, Fully Asynchronous Communication, and Graph Partition-Based Domain Decomposition for Strong Scaling on more than 100,000 Cores (eprint: arXiv:1606.02738), doi: 10.1145/2929908.2929916
  * Schaller et al. (2018) Schaller, M., et al. 2018, SWIFT: SPH With Inter-dependent Fine-grained Tasking, Astrophysics Source Code Library.  http://ascl.net/1805.020
  * Schaye et al. (2015) Schaye, J., Crain, R. A., Bower, R. G., et al. 2015, MNRAS, 446, 521, doi: 10.1093/mnras/stu2058
  * Schmelzle et al. (2017) Schmelzle, J., Lucchi, A., Kacprzak, T., et al. 2017, arXiv e-prints, arXiv:1707.05167.  https://arxiv.org/abs/1707.05167
  * Shao et al. (2022a) Shao, H., Villaescusa-Navarro, F., Villanueva-Domingo, P., et al. 2022a, Robust field-level inference with dark matter halos, arXiv, doi: 10.48550/ARXIV.2209.06843
  * Shao et al. (2022b) Shao, H., Villaescusa-Navarro, F., Genel, S., et al. 2022b, The Astrophysical Journal, 927, 85, doi: 10.3847/1538-4357/ac4d30
  * Sobol’ (1967) Sobol’, I. 1967, USSR Computational Mathematics and Mathematical Physics, 7, 86, doi: https://doi.org/10.1016/0041-5553(67)90144-9
  * Springel (2005) Springel, V. 2005, Mon. Not. Roy. Astron. Soc., 364, 1105, doi: 10.1111/j.1365-2966.2005.09655.x
  * Springel (2010) Springel, V. 2010, MNRAS, 401, 791, doi: 10.1111/j.1365-2966.2009.15715.x
  * Springel et al. (2018) Springel, V., Pakmor, R., Pillepich, A., et al. 2018, MNRAS, 475, 676, doi: 10.1093/mnras/stx3304
  * Takada & Bridle (2007) Takada, M., & Bridle, S. 2007, New Journal of Physics, 9, 446, doi: 10.1088/1367-2630/9/12/446
  * Takada & Hu (2013) Takada, M., & Hu, W. 2013, Physical Review D, 87, doi: 10.1103/physrevd.87.123504
  * Teyssier (2002) Teyssier, R. 2002, A&A, 385, 337, doi: 10.1051/0004-6361:20011817
  * Tonegawa et al. (2020) Tonegawa, M., Park, C., Zheng, Y., et al. 2020, The Astrophysical Journal, 897, 17, doi: 10.3847/1538-4357/ab95ff
  * Villaescusa-Navarro et al. (2020) Villaescusa-Navarro, F., Wandelt, B. D., Anglés-Alcázar, D., et al. 2020, arXiv e-prints, arXiv:2011.05992.  https://arxiv.org/abs/2011.05992
  * Villaescusa-Navarro et al. (2021a) Villaescusa-Navarro, F., Genel, S., Angles-Alcazar, D., et al. 2021a, arXiv e-prints, arXiv:2109.10360.  https://arxiv.org/abs/2109.10360
  * Villaescusa-Navarro et al. (2021b) Villaescusa-Navarro, F., Anglés-Alcázar, D., Genel, S., et al. 2021b, ApJ, 915, 71, doi: 10.3847/1538-4357/abf7ba
  * Villaescusa-Navarro et al. (2022a) Villaescusa-Navarro, F., Genel, S., Anglés-Alcázar, D., et al. 2022a, arXiv e-prints, arXiv:2201.01300.  https://arxiv.org/abs/2201.01300
  * Villaescusa-Navarro et al. (2022b) —. 2022b, ApJS, 259, 61, doi: 10.3847/1538-4365/ac5ab0
  * Villanueva-Domingo (2022) Villanueva-Domingo, P. 2022, PabloVD/CosmoGraphNet, v1.0, Zenodo, doi: 10.5281/zenodo.6485804
  * Villanueva-Domingo & Villaescusa-Navarro (2022) Villanueva-Domingo, P., & Villaescusa-Navarro, F. 2022, arXiv e-prints, arXiv:2204.13713.  https://arxiv.org/abs/2204.13713
  * Villanueva-Domingo et al. (2022) Villanueva-Domingo, P., Villaescusa-Navarro, F., Anglés-Alcázar, D., et al. 2022, ApJ, 935, 30, doi: 10.3847/1538-4357/ac7aa3
  * Wadekar et al. (2020) Wadekar, D., Villaescusa-Navarro, F., Ho, S., & Perreault-Levasseur, L. 2020, arXiv e-prints, arXiv:2012.00111.  https://arxiv.org/abs/2012.00111
  * Wadekar et al. (2022) Wadekar, D., Thiele, L., Hill, J. C., et al. 2022, The SZ flux-mass ($Y$-$M$) relation at low halo masses: improvements with symbolic regression and strong constraints on baryonic feedback, arXiv, doi: 10.48550/ARXIV.2209.02075
  * Weinberger et al. (2020) Weinberger, R., Springel, V., & Pakmor, R. 2020, Astrophys. J. Suppl., 248, 32, doi: 10.3847/1538-4365/ab908c
  * Weinberger et al. (2017) Weinberger, R., Springel, V., Hernquist, L., et al. 2017, MNRAS, 465, 3291, doi: 10.1093/mnras/stw2944
  * Yu et al. (2011) Yu, H.-R., Harnois-Dé raps, J., Zhang, T.-J., & Pen, U.-L. 2011, Monthly Notices of the Royal Astronomical Society, no, doi: 10.1111/j.1365-2966.2011.20360.x

##  Appendix A Additional $N$-body and Hydrodynamic Simulations

In Section 4, we presented the perfromance of the GNN and analytic expressions
on the different $N$-body and hydrodynamic simulations run with the same
cosmologies and initial conditions. Here, we present additional results
demonstrating the accuracy and robustness of both model predictions for both
$\Omega_{\rm m}$ and $\sigma_{\rm 8}$ for different minimum halo particle
thresholds. For these plots, we evaluate the models on 50 simulations
containing different cosmologies and initial conditions for four different
$N$-body codes: Abacus, CUBEP3M, PKDGrav, and Ramses (in Fig. 7). We also test
the models on 1,000 simulations from two hydrodynamic codes: IllustrisTNG and
SIMBA, but plot the results for 50 randomly selected simulations to conserve
space (in Fig. 8). As before, we perform these tests using halo catalogues
created with different minimum halo particle thresholds as indicated in the
plots. Each of these plots depict the predictions plotted against the truth
minus the inference.

As it can be seen, the GNN is able to infer $\Omega_{\rm m}$ accurately for
all $N$-body simulations with similar mean relative errors of $\sim 7\%$ and
the analytic expressions have comparable accuracies of $\sim 8\%$ (see Figs. 6
and 7).

For the hydrodynamic simulations IllustrisTNG and SIMBA, we obtain concurring
results where both the GNN and analytic expressions are able to attain mean
relative errors of $\sim 7\%$ (Fig. 8). An interesting note is that the GNN
predictions for halo catalogues constructed with 100 or 500 minimum particle
thresholds exhibit tail biases due to the effects of the prior distribution,
as seen in the right panels of 8. These numerical artifacts are not present in
the inferences made by the analytic expressions due to their known better
generalization capabilities.

Figure 6: We train a GNN with a low-dimensional latent space to infer
$\Omega_{\rm m}$ from catalogues of the Gadget $N$-body simulations using the
halo relative positions and velocity moduli. We then evaluate this model on
different $N$-body simulations: Abacus, CUBEP3M, PKDGrav3, and Ramses using
catalogues created with particle thresholds indicated next to the plots. As
can be seen, the model is able to extrapolate well to different $N$-body codes
and is able to predict with similar accuracy compared to that of the halo
catalogues from Gadget.

Figure 7: This follows the format as Fig. 6 but for the analytic equations
discussed in Section D. The equations were found using symbolic regression and
modified using physical principles to preserve the rotational and
translational symmetries of the data. As can be seen, the equations maintain
the accuracy and robustness exhibited by the GNN in Fig. 6, indicating that
the formulae offer good approximations to the model.

Figure 8: Similar to Figures 6 and 7, we test the GNN and the analytic
equations on halo catalogues generated from the SIMBA and IllustrisTNG
hydrodynamic simulations. For clarity, we plot the predictions for 50 randomly
selected catalogues in each panel. It can be seen that both models remain
robust to the additional astrophysical effects present in these simulations,
indicating that they are employing a possibly fundamental relation between the
relevant halo properties and $\Omega_{\rm m}$. Moreover, the analytic
equations are able to capture this as its accuracies for all hydrodynamic
simulations are similar to that of the GNN.

##  Appendix B Robustness to different halo finder: SUBFIND

In Section 4, we discussed the accuracy and robustness of the GNN and analytic
expressions when evaluated on various simulation codes. Here, we present
another test for the robustness of these models where we evaluate the GNN and
the analytic approximations on halo catalogues generated using a different
halo finder (Subfind) than the one used for training (Rockstar). Subfind
identifies halos by determining local peaks in the three-dimensional density
field and separating them using saddle points. The overdense regions and their
surroundings are then examined for subhalos, which are gravitationally self-
bound regimes. Those that are not bound are attached to their neighboring
overdensities with whom they share saddle points. Subfind operates on all
particle types in the simulations, dark matter and baryonic alike Dolag et al.
(2009) .

To perform these tests, we consider the total mass of the halo contained in a
sphere with a mean density that is 200 times the mean density of the Universe
at redshift $z=0$. Same as the previous tests, we construct halo catalogues
with varying minimum particle thresholds in the range of $[100,500]$, as
explained in Section 2.

First, we perform this test for the 1,000 halo catalogues from the $N$-body
Gadget simulations. As it can be seen in the top plots of Fig. 11, both the
GNN and the analytic expressions provide accurate predictions of $\Omega_{\rm
m}$ overall, with mean relative errors of $\sim 8.8\%$ and $\sim 9.2\%$,
respectively, across the different halo catalogues. However, there are some
interesting features to emphasize. First, while the GNN predictions exhibit an
offset and a significant lower-tail bias for the halo catalogue generated with
a minimum particle threshold of 100, this is a boundary case considering the
interval of the minimum particle thresholds used to construct the catalogues.
Moreover, the identification of lower mass halos can vary across different
halo finders and this can influence the predictions more strongly than the
presence of more massive halos which are more likely to be commonly identified
in both halo finders. Second, it can be seen that the analytic approximation
demonstrate higher accuracy for this boundary case which is another indication
of the better extrapolation capabilities of analytic equations over neural
networks.

Likewise, we performed the same test with halo catalogues from the
IllustrisTNG hydrodynamic simulations. The results for this are shown in the
bottom plots of Fig. 11. As it can be seen, the mean relative errors are
similar between the GNN and the analytic expressions, averaging to be $\sim
9.5\%$ across the different halo catalogues. While these metrics indicate
slightly decreased precision of the predictions, this can be attributed to
additional baryonic effects present. Nevertheless, the overall accuracy
further demonstrates the generalization ability of the trained network as it
is able to extrapolate to both additional hydrodynamic simulations and varying
halo definitions. This agrees with the results discussed in Sections 4.2.2 and
C, where it was also found that both the network and the symbolic
approximations are able to obtain robust predictions for catalogues generated
from the SWIFT-EAGLE simulations which employ a different halo/subhalo finder
(VELOCIraptor).

##  Appendix C Additional Plots: Testing GNN on Galaxies

As discussed in Section 4.1.2, we trained a GNN on halo catalogues and tested
the learned network on galaxies from six different hydrodynamic simulation
suites: Astrid, IllustrisTNG, Magneticum, SB28, Simba, and SWIFT-EAGLE. Here,
we present the results for these predictions. As it can be seen in Fig. 9, the
GNN is unable to accurately predict the values of $\Omega_{\rm m}$ as all the
predictions exhibit a bias deviating from the true values. This is common
across all simulations, which is expected given that there is a nontrivial
connection between halo and galaxy distributions. On the other hand, as
explained in Section 4.2.2, the analytic equations that approximate the GNN
can be tuned to avoid this error.

We believe that these biases are due to the effects of the halo-galaxy
connection in addition to the differences in the abundance of galaxies found
in the catalogues used for testing and that of halos found in the training
dataset. As discussed in Section 4.2.2, the network is unable to extrapolate
to number densities outside of the training range. In the case of galaxy
catalogues, as shown in Fig. 5, there are many catalogues with galaxy number
densities that fall below the range of the halo number densities seen during
training, $(1000,6000)$. However, the under-predicted values of $\Omega_{\rm
m}$ cannot be solely attributed to the abundance of galaxies. As discussed in
our companion paper, de Santi et al. (2023), the full range of galaxy number
densities is exhibited for all values of $\Omega_{\rm m}$. Hence, there is no
strong correlation between $\Omega_{\rm m}$ and the number of galaxies in each
catalogue. This agrees further demonstrates that the biases present in the
network predictions are attributed to the intrinsic characteristics of the
galaxy population.

Figure 9: This plot shows the predictions of the GNN trained on halo
catalogues from Gadget $N$-body simulations being tested on galaxies from six
different hydrodynamic simulations as listed for each row. To construct the
galaxy catalogues, we follow the procedure discussed in Section 3.1 and use
four different stellar mass thresholds which are labeled for each column. For
clarity, we plot the predictions for 50 randomly selected catalogues in each
panel. As can be seen, the GNN is unable to accurately predict the values of
$\Omega_{\rm m}$ as all the predictions exhibit a bias deviating from the true
values. This is common across all simulations, which is expected given that
there is a nontrivial connection between halo and galaxy distributions.
However, as explained in Section 4.2.2, the analytic equations that
approximate the GNN can be easily tuned to avoid this error.

##  Appendix D Original Symbolic Regression Equations

Figure 10: This plot shows the predictions of the original equations found by
the symbolic regression algorithm evaluated on the halo catalogues of the
Gadget test set. As can be seen, while it achieves similar accuracy to the GNN
model, with a mean relative error of $7.1\%$, it is not as accurate as the
modified expressions which had an error of $6.6\%$. See Fig. 2.

In Section D, we presented the equations obtained by the symbolic regression
algorithm that were then modified based on motivations of physical principles.
Here, in Table 3, we present the equations originally found by the symbolic
regression algorithm that we trained following the procedure described in
Section 3.3. Note that the edge equations found by the algorithm contained
dependencies on the individual velocity modulus of halos in the form of linear
combinations of $v_{i}$ and $v_{j}$. These terms break the parity between a
halo and its neighbor. Moreover, we adapted terms that explicitly reflect
differences between the velocity moduli due to the known statistics between
pairwise velocities and $\Omega_{\rm m}$,

We also show the accuracy of these equations when evaluated on halo catalogues
from the Gadget test set simulations in Fig. 10. As it can be seen, these
formulas are able to achieve a mean relative error of $\sim 7.1\%$ which is
slightly higher than the error of the modified equations, possibly indicating
that the imposed symmetries offer an important constraint on the predictions
and play a significant role in achieving accurate inferences (see Fig. 2).

Table 3: This table follows the format of Table 1 and lists the original analytic formulas found by the symbolic regression algorithm. The key difference is that the edge model equations originally obtained by the algorithm depend on terms $v_{i}$ and $v_{j}$, which are the individual halo velocity moduli. As in the equations discussed in Section 4.2. the velocities used in the equations here have also been normalized to aid the model training and to ensure that they are dimensionless: $v_{i}=\frac{v_{i}-\mu}{\delta}$, $v_{j}=\frac{v_{j}-\mu}{\delta}$. For testing these equations on halo catalogues, we use the fixed values $\mu=189$ km $\rm s^{-1}$ and $\delta=129$ km $\rm s^{-1}$ computed from the mean and standard deviation of the velocity moduli for all halos in the training set. The accuracy of these equations are shown in Fig. 10 GNN Component | Formula | RMSE
---|---|---
Edge Model: $e^{(1)}_{1}$ | $1.32|1.05v_{i}-v_{j}+0.21|-0.12v_{j}-0.12(\gamma_{ij}+\beta_{ij}-1.73)$ | 0.028
Edge Model: $e^{(1)}_{2}$ | $|1.53(v_{i}-1.06v_{j})+0.45|+1.93(v_{i}-1.02v_{j})+0.55$ | 0.035
Node Model: $v^{(1)}_{1}$ | $1.21^{v_{i}}(0.77^{3.29\sum_{j\in\mathcal{N}_{j}}e^{(1)}_{1}+\sum_{j\in\mathcal{N}_{j}}e^{(1)}_{2}})+0.12$ | 0.02
Node Model: $v^{(1)}_{1}+v^{(1)}_{2}$ | $0.78-\sqrt{\log(0.16^{\sum_{j\in\mathcal{N}_{j}}e_{2}+\sum_{j\in\mathcal{N}_{j}}e_{1}-0.41v_{i}-1.05})}+1.45$ | 0.03
Final MLP: $\mu$ | $4\times 10^{-4}\times(-5.5\sum_{i\in\mathcal{G}}v^{(1)}_{2}+2.21\sum_{i\in\mathcal{G}}v^{(1)}_{1}+|0.96\sum_{i\in\mathcal{G}}v^{(1)}_{2}+0.82\sum_{i\in\mathcal{G}}v^{(1)}_{1}|)-0.103\ $ | 0.03

Figure 11: We trained a GNN using halo catalogues generated with the halo
finder Rockstar to infer $\Omega_{\rm m}$, and approximated the learned model
with analytic equations using symbolic regression. The top plots show the
accuracy of the model and the analytic approximations when evaluated on halo
catalogues from the $N$-body Gadget simulations using a different halo finder
- Subfind \- constructed with the varying minimum particle thresholds as
described earlier. It is overall able to accurately extrapolate to the
different halo finder with $\sim 8\%$ mean relative error across the different
catalogues. On the other hand, while the analytic expressions have a slightly
larger error of $\sim 9\%$, they do not exhibit the noticeable biases present
in the predictions from the GNN, demonstrating the known improved
extrapolation properties of analytic expressions over neural networks. The
bottom plots depict the same test as above but for halo catalogues from the
IllustrisTNG hydrodynamic simulations.

##  Appendix E Varying stellar mass thresholds

In this section, we discuss the results for testing the analytic equations
discussed in Section 4.2 on galaxy catalogues constructed with different
minimum stellar mass thresholds: $N\times m_{*}$ for $N\in\\{3,4,5,6\\}$ where
$m_{*}$ is a fixed mass for a single stellar particle. As explained in Section
3.1, the use of different mass cuts during training of the model and equations
enables the models to marginalize over the halo/galaxy number densities found
in each simulation due to different halo/galaxy mass functions. Here, we test
whether the equations are robust to this using the simulations from ASTRID and
SWIFT-EAGLE. To do this, we use the same $\delta$ values optimized for
catalogues of the single mass threshold $4\times m_{*}$ as discussed in
Section 4.2.2.

First, we present the results for ASTIRD, which are shown in Fig. 12. As it
can be seen, the accuracy of the equations are not largely affected by the
different mass thresholds used, as expected. Second, we perform these tests
for galaxy catalogues from SWIFT-EAGLE, as shown in Fig. 13. Here, we explain
the apparent trend of increasing scatter in the predicted $\Omega_{\rm m}$
values as the stellar mass threshold increases with the fact that all the
SWIFT-EAGLE simulations were run with the same random seed. Hence, the
predictions should be considered as highly correlated, which causes the small
bias for the catalogue with a larger stellar mass threshold for a fixed
$\delta$ value.

Figure 12: We evaluate the analytic equations discussed in Section 4.2 on
galaxy catalogues from the Astrid simulation set constructed using four
different minimum stellar mass thresholds: $N\times m_{*}$ for
$N\in\\{3,4,5,6\\}$ where $m_{*}$ is a fixed mass for a single stellar
particle. Each column is labeled with the corresponding mass cut. As it can be
seen, the accuracies of the equations are preserved for the different mass
thresholds demonstrating that the model has marginalized over the number
density of galaxies. Figure 13: We evaluate the analytic equations discussed
in Section 4.2 on galaxy catalogues from the SWIFT-EAGLE simulation set
constructed using four different minimum stellar mass thresholds: $N\times
m_{*}$ for $N\in\\{3,4,5,6\\}$ where $m_{*}$ is a fixed mass for a single
stellar particle. Each column is labeled with the corresponding mass cut. We
note that since the SWIFT-EAGLE simulations were generated using the same
initial random seed, there is a high correlation between the galaxy catalogues
of the different stellar mass thresholds for this simulation set that is
responsible for the trend of decreasing accuracy as the stellar mass threshold
increases.

##  Appendix F Testing with Super-Sample Covariance

Here we demonstrate that the analytic equations discussed in Section 4.2 are
robust to the effects of super-sample covariance. Quantifying how the analytic
equations behave in response to super-sample covariance is a critical step
towards being able to apply them to observational data from surveys that are
sampled with finite volume. This is because in galaxy surveys, the short-
wavelength modes that contain information on the non-linear dynamics are
coupled to long-wavelength, or super-sample, modes that extend beyond the
survey volume (Hamilton et al., 2006; Hu & Kravtsov, 2003; Takada & Hu, 2013).
This results in sample variances that dominate the non-linear regime (Sato et
al., 2009; Takada & Bridle, 2007; Yu et al., 2011; de Putter et al., 2012). In
the analysis that we have performed so far, we have not taken into
consideration of this effect because we have used simulations with periodic
boundary conditions which are not influenced by background modes that extend
outside the simulation box.

To test for these effects, we evaluate the analytic equations on galaxy
catalogues constructed from $(25~{}h^{-1}Mpc)^{3}$ sub-volumes randomly
selected from the IllustrisTNG-300 simulation to match the size of the
simulation boxes used for training. As described in Section 2.1.2, this
simulation has a total volume of $(205~{}h^{-1}Mpc)^{3}$ and was run with the
cosmology $\Omega_{\rm m}=0.3089$. It is important to note that unlike the
simulations used for training, we do not impose periodic boundary conditions
on the sub-volumes used in this test in order to account for the super-sample
modes.

We present the results of this test in the top panel of Fig. 14. Each plot in
the figure depicts the differences between the truth and predicted
$\Omega_{\rm m}$ made by the analytic equations for 100 randomly selected sub-
volumes. Following the same procedure used to perform the previous tests on
galaxies, we construct four catalogues for each sub-volume using the stellar
mass thresholds discussed in Section 3.1. Each column is thus labeled with the
corresponding stellar mass cut used. We note that the predictions across all
catalogues exhibit a common offset from the truth which we account for by
introducing to the final MLP equation an additive constant of $b=-0.19$ found
using $\chi^{2}$ minimization. This common bias is explained by the fact that
the equations are being evaluated on sub-volumes that do not contain periodic
boundary conditions but were trained only on simulations that contain periodic
boundary conditions. After correcting for this, the analytic expressions are
able to achieve mean relative errors of $\sim 11.1\%$.

Figure 14: Top: We quantify the behavior of the analytic expressions,
discussed in Section 4.2, in the presence of super-sample covariance. We test
the analytic expressions on 100 $(25~{}h^{-1}Mpc)^{3}$ sub-volumes randomly
selected from the IllustrisTNG-300 simulation without imposing periodic
boundary conditions. The simulation contains a total volume of
$(205~{}h^{-1}Mpc)^{3}$ and was run with a cosmology of $\Omega_{\rm
m}=0.308$. The plots depict the difference between the true $\Omega$ value and
the predicted for the galaxy catalogues constructed using each of the four
stellar mass thresholds as indicated at the top of each column. For all
catalogues, the predictions are corrected for their negative offset from the
truth by introducing a bias in the analytic expression for the final MLP,
$b=-0.19$, which shifts all predictions upwards by a constant. After adjusting
for this common offset, the predictions exhibit mean relative errors of $\sim