# The Impact of Baryons on Cosmological Inference from Weak Lensing Statistics

Tianhuan Lu Department of Astronomy, Columbia University, New York, NY 10027,
USA  Zoltán Haiman Department of Astronomy, Columbia University, New York, NY
10027, USA

###### Abstract

As weak lensing surveys are becoming deeper and cover larger areas,
information will be available on small angular scales down to the arcmin
level. To extract this extra information, accurate modelling of baryonic
effects is necessary. In this work, we adopt a baryonic correction model,
which includes gas both bound inside and ejected from dark matter (DM) haloes,
a central galaxy, and changes in the DM profile induced by baryons. We use
this model to incorporate baryons into a large suite of DM-only $N$-body
simulations, covering a grid of 75 cosmologies in the
$\Omega_{\mathrm{m}}-\sigma_{8}$ parameter space. We investigate how baryons
affect Gaussian and non-Gaussian weak lensing statistics and the cosmological
parameter inferences from these statistics. Our results show that
marginalizing over baryonic parameters degrades the constraints in
$\Omega_{\mathrm{m}}-\sigma_{8}$ space by a factor of $2-4$ compared to those
with baryonic parameters fixed. We investigate the contribution of each
baryonic component to this degradation, and find that the distance to which
gas is ejected (from AGN feedback) has the largest impact due to its
degeneracy with cosmological parameters. External constraints on this
parameter, either from other datasets or from a better theoretical
understanding of AGN feedback, can significantly mitigate the impact of
baryons in an HSC-like survey.

gravitational lensing: weak – cosmology: theory – cosmological parameters –
large-scale structure of Universe

##  1 Introduction

Weak gravitational lensing (hereafter weak lensing, or WL) describes the
bending of light rays coming from distant objects by inhomogeneities in the
foreground matter distribution. Since the matter distribution depends on the
underlying cosmological model, WL was proposed and has been proven to be a
powerful tool in inferring cosmological parameters, such as
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$ in a $\Lambda$CDM universe (see, e.g.
Bartelmann & Schneider, 2001; Refregier, 2003; Hoekstra & Jain, 2008;
Kilbinger, 2015, for reviews). Recent studies using data from WL surveys have
yielded especially strong constraints on a certain combination of cosmological
parameters (e.g. $S_{8}\equiv\sigma_{8}(\Omega_{\mathrm{m}}/0.3)^{0.5}$) and
are approaching percent-level precision (Joudaki et al., 2016; Köhlinger et
al., 2017; Hikage et al., 2019; Hamana et al., 2020).

One of the trends in the existing and planned WL survey programs is the
increase in depth—the number density of galaxies has increased from $\sim
10\,\mathrm{arcmin^{-2}}$ in the past to $\sim 20\,\mathrm{arcmin^{-2}}$ at
present, and up to $50\,\mathrm{arcmin^{-2}}$ is expected in the future
(Heymans et al., 2012; de Jong et al., 2013; Abbott et al., 2005; Aihara et
al., 2018; Laureijs et al., 2011; Spergel et al., 2015; Ivezić et al., 2019).
Although this will allow accessing more information theoretically on small
scales (down to a few arcmin), making use of this information is challenging
since many systematic errors are more prominent on small scales if not treated
accurately. These effects include uncertainties in baryonic physics (Zhan &
Knox, 2004; White, 2004; Jing et al., 2006) and intrinsic alignments of galaxy
shape orientations (Schneider & Bridle, 2010; Sifón et al., 2015).

In this paper, we focus on investigating how the baryonic physics affect WL
statistics — specifically the (convergence) power spectrum and WL peak counts
(as an example of a small-scale non-Gaussian statistic). Previous studies have
incorporated baryonic physics into WL observables using two broadly different
approaches. The first is to utilize hydrodynamical simulations that assume
certain treatments of baryonic physics and predict modifications to WL
statistics. For example, Hikage et al. (2019) take the fitting formula based
on the active galactic nucleus (AGN) feedback model in the OverWhelmingly
Large Simulations project (Schaye et al., 2010), considering AGN feedback
having the largest effect, and adjust the power spectrum. Osato et al. (2021)
have recently studied the baryonic effects using ray-traced convergence maps
based on the IllustrisTNG hydrodynamics simulations, which include several
baryonic effects (Nelson et al., 2019). The second approach is to modify the
shape of the dark matter (DM) haloes in a DM-only $N$-body simulations via so-
called “baryonic corrections”. An early attempt following this approach
studying baryonic effects on cosmological parameter estimates include Yang et
al. (2013). This work analysed how varying the concentration parameters of
Navarro-Frenk-White (NFW) haloes affects the shear power spectrum and the
lensing peak counts. More recently, Schneider et al. (2019) compared power
spectra and shear correlation functions adopting more sophisticated baryonic
corrections, and showed that they can accurately mimic the observables
inferred from hydrodynamical simulations.

In this paper, we take the latter approach, employing a state-of-the-art
baryonic correction model (hereafter BCM) proposed by Aricò et al. (2020, a
modification of an earlier version by ), to answer two questions:

  * •

How much does baryonic physics affect the lensing power spectrum and peak
counts?

  * •

If we infer baryonic model parameters simultaneously with cosmological
parameters, how well can they be constrained?

The BCM offers a parametrized way to characterize how baryonic physics alter
the shape of each halo on the basis of a DM-only $N$-body simulation, where
each parameter has a physically motivated interpretation. Compared to
hydrodynamical simulations, the BCM is much less computationally expensive,
while retaining flexibility to model various poorly constrained aspects of
baryonic physics. Aricò et al. (2020) have shown that the matter power
spectrum residuals between various hydrodynamical simulations and their best-
fitting BCMs are under 1%.

We organize this paper as follows. In § 2, we introduce our simulation suite,
how we derive statistics from simulations, and the BCM. In § 3, we show the
effects of baryonic physics and the posterior distributions of cosmological
and baryonic parameters from WL. In § 4, we then discuss how factors such as
different noise levels and restrictions of power spectrum measurements to
different angular scales affect parameter inferences, as well as the
implications of our findings. Finally, we summarize our main conclusions in §
5.

##  2 Methods

###  2.1 $N$-body simulations

Figure 1: The cosmologies of the 75 $N$-body simulations in the suite that we
use in this study. The red dot represents the fiducial cosmology
($\Omega_{\mathrm{m}}=0.311,\sigma_{8}=0.789$).

Our simulation suite consists of DM-only $N$-body simulations of $\Lambda$CDM
cosmologies with 75 different combinations of $\Omega_{\mathrm{m}}$ (the
matter density) and $\sigma_{8}$ (the magnitude of matter fluctuation). Our
parameter grid is illustrated in Figure 1.

The rest of the cosmological parameters are fixed as follows: Hubble constant
$H_{0}=72\,\mathrm{km\,s^{-1}\,{Mpc}^{-1}}$, baryon density
$\Omega_{\mathrm{b}}=0.046$, scalar spectral index $n_{\mathrm{s}}=0.96$,
effective number of relativistic degrees of freedom $n_{\mathrm{eff}}=3.04$,
and neutrino masses $m_{\nu}=0$.

The initial conditions are computed using CAMB (Lewis et al., 2000), and the
simulations are run with the $N$-body code Gadget-2 (Springel, 2005). The size
of each simulation run is $240\,h^{-1}\mathrm{Mpc}$ (comoving) with $512^{3}$
DM particles, each of which has a mass of $\approx
10^{10}\,\mathrm{M}_{\odot}$. In the fiducial cosmology
($\Omega_{\mathrm{m}}=0.311,\sigma_{8}=0.789$), the size of the simulation box
corresponds to a field of view of $6\times 6\,\mathrm{{deg}^{2}}$ at redshift
$z=1$.

###  2.2 Ray-tracing

We follow the procedures described by Petri (2016) to generate the WL
convergence ($\kappa$) maps from $N$-body simulations using the multiple lens
plane algorithm (Jain et al., 2000; Hilbert et al., 2009). We briefly
introduce the main steps in this procedure below and refer the reader to the
above papers for more details.

In this study, we assume that all lensed galaxies are at $z=1$. We take
snapshots of each $N$-body simulation at a series of redshifts between $z=0$
and $z=1$, such that the difference in the comoving distance between the
redshifts of adjacent snapshots is $80\,h^{-1}\mathrm{Mpc}$. For each snapshot
in the series, we cut a slab of width $80\,h^{-1}\mathrm{Mpc}$ along a random
axis ($x$, $y$ or $z$), and we generate a density plane by calculating the
column density of a slab on a $4096\times 4096$ grid followed by random shifts
and rotations (the latter in multiples of $90^{\circ}$). To obtain a potential
plane, we solve the two-dimensional Poisson equation on each density plane for
the gravitational potential field. Finally, we follow light rays from $z=0$ to
$z=1$ using the multiple lens plane algorithm on each series of potential
planes following Petri (2016) within a fixed field of view of $3.5\times
3.5\,\mathrm{{deg}^{2}}$ (consistent with previous works on this simulation
suite which ray-traced to $z=2$). The result of this procedure is a
$2048\times 2048$ pixelized converge map. Through cutting slabs and shifting
and rotating density planes randomly, we can repeat this process to generate
multiple realizations of the convergence map for each simulation, which can be
considered statistically independent (Petri et al., 2016).

To imitate the convergence maps calculated from real data, we add shape noise
to our simulated convergence maps by changing the value of each pixel by a
random number, drawn from a Gaussian distribution:

$\left\\{\mu=0,\sigma=\frac{\sigma_{\epsilon}}{\sqrt{2n_{\mathrm{gal}}A_{\mathrm{pix}}}}\right\\},$
(1)

where $\sigma_{\epsilon}=0.4$ is the typical mean intrinsic ellipticity of
galaxies, $n_{\mathrm{gal}}$ the surface density of lensed galaxies
controlling the level of noise, and $A_{\mathrm{pix}}$ the area of each pixel.
We also smooth the maps with a $1\,\mathrm{arcmin}(\sim 10\,\mathrm{pixels})$
Gaussian filter to remove artefacts caused by pixelation and to facilitate
peak counting statistics.

###  2.3 Modification by baryons

Parameter | Fiducial value | Prior bound
---|---|---
$M_{\mathrm{c}}$ | $3.3\times 10^{13}h^{-1}\mathrm{M}_{\odot}$ | $\left[5.9\times 10^{11},4.4\times 10^{15}\right]h^{-1}\mathrm{M}_{\odot}$
$M_{1,0}$ | $8.63\times 10^{11}h^{-1}\mathrm{M}_{\odot}$ | $\left[9.3\times 10^{10},1.1\times 10^{13}\right]h^{-1}\mathrm{M}_{\odot}$
$\eta$ | 0.54 | $\left[0.12,2.7\right]$
$\beta$ | 0.12 | $\left[0.026,3.8\right]$
Table 1: The fiducial values and prior bounds of the baryonic parameters

The BCM assumes that each halo can be divided into four components—DM, stars,
bound gas, and ejected gas. The radial profiles and relative weights of these
components are controlled by four baryonic parameters: $M_{\mathrm{c}}$, the
halo mass for retaining half of the total gas; $M_{1,0}$, the characteristic
halo mass for a galaxy mass fraction of 0.023 at $z=0$; $\eta$, the maximum
distance to which the gas content of a halo is ejected (in units of a
characteristic escape distance that is roughly 5 virial radii); and $\beta$,
the slope of the gas fraction against the halo mass. We choose the fiducial
values for the baryonic parameters following Aricò et al. (2020), and adopt
wide prior bounds so that they cover sufficient variations in the matter power
spectrum (see Table 1).

To incorporate baryonic effects into the $N$-body simulations, we replace DM
haloes with analytical profiles that include baryons. Specifically, we remove
all particles which are considered to belong to DM haloes, and paint the
pixelated image of column density of the spherically symmetric haloes (with
baryons) on the corresponding density planes.

Figure 2: The cumulative halo mass function at $z=0$ from the fiducial
simulation compared to that given by Jenkins et al. (2001).

First, we build halo catalogues for every snapshot of each $N$-body simulation
with rockstar (Behroozi et al., 2012), a halo finder based on adaptive
hierarchical refinement of friends-of-friends (FoF) groups, see Figure 2 for
one of the halo mass functions. We set the FoF linking length to the default
value of 0.28. rockstar assigns a list of DM particles to each halo or
subhalo, where each particle belongs to no more than one list, and we catalog
all haloes with total particle mass greater than
$10^{12}\,\mathrm{M}_{\odot}$. Second, for each halo, we remove its particles
from the slab if the distance between the particle and the halo centre is
smaller than the virial radius ($r_{200}$). We refer to the total mass of this
region inside $r_{200}$ as the halo mass $M$ (note that some of the particles
assigned by rockstar to a halo are outside $r_{200}$; for our purposes those
particles remain outside the halo). We then fit the radial distribution of the
removed particles to the NFW profile with mass $M$, and obtain the best-
fitting concentration parameter $c$ of each halo from a maximum likelihood
estimate. Third, given a specific choice of baryonic parameters, we calculate
the column density of each halo (analytical profile with BCM) on the density
plane grid at their locations, and they are added to the density plane.

In this work, we generate one realization of the convergence map for each
choice of baryonic parameters for each cosmology. The baryonic parameters are
chosen according to a Sobol sequence (Sobol’, 1967) of 160 elements in the
4-dimensional prior hypercube. Sobol sequences have a property of low-
discrepancy, i.e. evenly distributed across the space, so they perform better
in polynomial fitting and interpolation compared to random points. Compared to
Latin hypercube sampling, Sobol sequences have the advantage of being
extendable, in case we need more elements in the parameter space to improve
accuracy. In addition, we generate 128 realizations at the fiducial baryonic
parameters to estimate the covariance matrix for each cosmology.

We note that there are, on average, $\sim 30,000$ haloes in an $N$-body
snapshot, and $\gtrsim 10^{10}$ 2D halo images with BCM need to be calculated
given our setup (75 cosmologies $\times$ $(160+128)$ realizations $\times$ 30
snapshots $\times$ $30,000$ haloes). It is too computationally expensive to
calculate every halo following the method proposed by Aricò et al. (2020) and
Schneider & Teyssier (2015), since that technique involves displacing
individual N-body particles. To reduce the computational cost, we compute the
column densities of the four components in BCM separately:

$\displaystyle\Sigma_{\mathrm{halo}}(\theta_{\mathrm{b}},z,M,r_{\mathrm{200}},c)$
$\displaystyle=M\,w_{\mathrm{dm}}(\theta_{\mathrm{b}},z)\,\Sigma_{\mathrm{dm}}(r_{\mathrm{200}},c)$
$\displaystyle+M\,w_{\mathrm{star}}(\theta_{\mathrm{b}},z)\,\Sigma_{\mathrm{star}}(r_{\mathrm{200}},c)$
$\displaystyle+M\,w_{\mathrm{bg}}(\theta_{\mathrm{b}},z)\,\Sigma_{\mathrm{bg}}(r_{\mathrm{200}},c)$
$\displaystyle+M\,w_{\mathrm{eg}}(\theta_{\mathrm{b}},z)\,\Sigma_{\mathrm{eg}}(\eta,r_{\mathrm{200}}),$
(2)

where $\theta_{\mathrm{b}}$ denotes all baryonic parameters.
$\Sigma_{\mathrm{dm}}$ and $\Sigma_{\mathrm{bg}}$ are the profiles of the DM
and the bounded gas respectively, almost all the masses of which reside inside
$r_{\mathrm{200}}$. $\Sigma_{\mathrm{star}}$ is the profile of stellar mass
(Mohammed et al., 2014) concentrated at the centre of the halo; it is also
referred to as the central galaxy component. $\Sigma_{\mathrm{eg}}$ is the
profile of the ejected gas modelled by particles from AGN feedback following
the Maxwell–Boltzmann distribution, the radius of which relative to the escape
distance $\approx 5\,r_{\mathrm{200}}$ is determined by $\eta$. Here, the
profiles of the baryonic components contain their contributions to the
relaxation of DM $\Sigma_{b}^{\mathrm{(rdm)}}(b=\mathrm{\\{star,bg,eg\\}})$:

$\Sigma_{b}^{\mathrm{(rdm)}}=\frac{1}{\delta}\left(\Sigma_{\mathrm{rdm}}(\delta\Sigma_{b})-\Sigma_{\mathrm{dm}}\right),$
(3)

where $\Sigma_{\mathrm{rdm}}(\delta\Sigma_{b})$ is the profile of relaxed DM
with $\delta$ of the halo mass replaced by the baryonic profile $\Sigma_{b}$.
We note that this simplification only retains the linear dependency of relaxed
DM on the baryonic components.

We calculate three of the four surface density profiles
$\Sigma_{\\{\mathrm{dm,star,bg}\\}}$ on a grid of

$\displaystyle\log_{10}\left(\frac{r_{200}}{l_{\mathrm{pix}}}\right)$
$\displaystyle=0.12,0.16,\cdots,1.8,$ (4) $\displaystyle\log_{10}c$
$\displaystyle=0.1,0.2,\cdots,1.8,$ (5)
$\displaystyle\log_{10}\left(\frac{r_{200}}{l_{\mathrm{pix}}}\right)$
$\displaystyle=0.12,0.16,\cdots,1.8,$ (6) $\displaystyle l_{\mathrm{pix}}$
$\displaystyle=\frac{240\,h^{-1}\mathrm{Mpc}}{4096},$ (7)

with the total mass of each profile being unity, and additionally the
contribution on relaxed DM by ejected gas
$\Sigma_{\mathrm{eg}}^{\mathrm{(rdm)}}$ on

$\log_{10}\eta=-1.0,-0.9,\cdots,1.0.$ (8)

These profiles are then rasterized into images with their centers shifted by
0.0, 0.25, 0.5, and 0.75 pixels along both axis. Whenever we need the image of
an actual halo with BCM, we find the images of its components with the closest
$M$ and $c$ from the pre-calculated gallery, and weight them by
$w_{\\{\mathrm{dm,star,bg,eg}\\}}$ according to BCM.

###  2.4 Statistics and parameter inference

In this study, we employ two statistics on the convergence maps—power spectrum
and peak counts. The power spectrum of a convergence map is defined as the
Fourier transform of the two-point correlation function. We calculate the
power spectra by first performing Fourier transforms on the convergence maps,
and then obtain the values in 18 equally-spaced logarithmic bins of the
angular Fourier mode $\ell$ within $100<\ell<12,000$. Peak counts are defined
as the number of peaks per unit solid angle at difference convergence
($\kappa$) values, where a peak refers to a pixel that have a higher $\kappa$
than all of its eight neighboring pixels. We divide $-0.03<\kappa<0.15$ into
18 equally spaced bins and count the number of peaks with $\kappa$ values
within each of these bins.

We estimate the posterior distribution with Bayes’ theorem on six parameters
$\theta$ with log-uniform priors: $0.2<\Omega_{\mathrm{m}}<0.6$,
$0.4<\sigma_{8}<1.1$ and the four baryonic parameters shown in Table 1. The
likelihood function is given by

$\displaystyle
p(\mathbf{y}_{0}|\theta)\propto\frac{1}{\sqrt{\det\mathbf{C}}}\exp\left(-\frac{1}{2}\Delta\mathbf{y}^{\mathrm{T}}\widehat{\mathbf{C}^{-1}}\Delta\mathbf{y}\right),$
(9) $\displaystyle\Delta\mathbf{y}=\mathbf{y}_{0}-\mathbf{y}(\theta),$ (10)
$\displaystyle\widehat{\mathbf{C}^{-1}}=\frac{N-d-2}{N-1}\mathbf{C}^{-1},$
(11) $\displaystyle\mathbf{C}=\left(\frac{A_{\mathrm{survey}}}{3.5\times
3.5\,\mathrm{deg^{2}}}\right)^{-1}\mathbf{C}_{\mathrm{sim}}(\Omega_{\mathrm{m}},\sigma_{8}),$
(12)

where $\mathbf{y}(\theta)$ denotes the statistics at parameters $\theta$,
$\mathbf{y}_{0}$ the statistics at the fiducial parameters, $N$ the number of
realizations, $d$ the number of observables, $A_{\mathrm{survey}}$ the area of
the supposed survey (e.g. $A_{\mathrm{survey}}=1,500\,\mathrm{deg^{2}}$ for a
survey such as the one by Hyper Suprime-Cam or HSC; Aihara et al. 2018), and
$\mathbf{C}_{\mathrm{sim}}(\Omega_{\mathrm{m}},\sigma_{8})$ the covariance
matrix measured from the simulated convergence maps. In general, the
covariance matrix depends on the choice of both cosmological and baryonic
parameters, but we ignore its dependence on baryonic parameters here because:
1) the dependence of the covariances on baryonic parameters is small (see §
4.5 below) and 2) a single realization for each choice of baryonic parameters
will not yield meaningful estimations. The factor $(N-d-2)/(N-1)$ before the
covariances makes $\widehat{\mathbf{C}^{-1}}$ and unbiased estimation of the
precision matrix (Hartlap et al., 2007).

Both statistics (as well as their combination) and their covariances are
interpolated across the parameter space (on logarithmic scales) by a fifth-
order polynomial to cover arbitrary points over the full range of models. We
find that to fit 12,000 points (75 cosmologies $\times$ 160 baryon models), a
fifth-order 6-variate polynomial (462 parameters) captures the variation of
the statistics on all parameters with good accuracy without overfitting. Then,
we sample the posterior distribution with a Monte-Carlo Markov chain using
differential evolution for $10^{6}$ steps, which is sufficient for the chains
to converge.

##  3 Results

###  3.1 Effects of halo replacement

Figure 3: The column density of the removed DM particles, the NFW profile
images, and their differences in a $(25\,h^{-1}\mathrm{Mpc})^{2}$ region of a
slab at $z=0$. The maps are smoothed at a $100\,\mathrm{kpc}$ scale for better
visibility. Figure 4: The average column density power spectra from the full
$N$-body simulations (the original in black, and with halo quadrupoles removed
in orange) compared to the power spectrum $P_{\rm NFW}$ after halo replacement
in the fiducial cosmology at $z=0$. Figure 5: The residuals caused by halo
replacement (solid lines) compared to the uncertainties of the observables
(dashed lines) in the two statistics, a) power spectrum and b) peak counts.

The process of halo replacement speeds up the calculations of baryon models,
but it will introduce systematics since the haloes in the $N$-body simulations
does not match NFW profile exactly. In this section, we will study by how much
this mismatch affects the statistics.

Figure 3 shows an example of how replacing haloes with NFW profiles changes
the column density in a simulation slab, and Figure 4 shows the change in
column density power spectrum for the fiducial cosmology at $z=0$. We find
that halo replacement lowers the column density power spectrum by $\sim 7\%$
mostly at very small scales ($k\gtrsim 3\,h\,\mathrm{Mpc^{-1}}$). It can be
seen from the residual map shown in right panel of Figure 3 that the most
noticeable difference for each halo between the $N$-body particles and the NFW
profile is the quadrupole component. This means that projecting haloes to a
2-dimensional map leaves their quadrupole components stronger than the other
multipole moments and they are apparent when comparing to isotropic NFW
profiles. To investigate the importance of this quadrupole residual, we
average the original column density map and its duplicate but with the halo
particles rotated around the projected halo centre by $90^{\circ}$, and we
repeat the comparison of the column density power spectrum (see Figure 4). As
the discrepancy shrinks to $\sim 2\%$, we conclude that the quadrupole
residuals are the primary causes of the difference in the small-scale power
spectrum introduced by the adoption of spherical NFW haloes.

Although quadrupole residuals introduce systematics in general, we did not
attempt to correct them in this work because they have a limited impact on the
lensing power spectrum and peak counts as all maps are affected in the same
way (but we note that they may be a concern for other statistics or when our
method is applied on real data). In Figure 5, we compare the residuals caused
by halo replacement to the uncertainty of the observables from 64 realizations
of the fiducial cosmology with a noise level of
$n_{\mathrm{gal}}=20\,\mathrm{arcmin^{-2}}$. We find that for both statistics,
the residuals are mostly less than 10% of the uncertainties of the observables
in a single convergence map, but they are projected to be comparable to the
uncertainties in a $1,500\,\mathrm{deg^{2}}$ survey. As a result, a correction
for these quadrupolar asymmetries will need to be incorporated into the
spherical halo-based baryon corrections, when real data is fit in a large
($\gtrsim 1500~{}{\rm deg^{2}}$) future survey.

###  3.2 Impact of baryons on statistics

Figure 6: The column density power spectra with each of the four baryon
components and with all components at $z=0$, taking fiducial values. They are
compared to the power spectrum from replacing haloes with NFW profiles without
baryons. Figure 7: The responses of the power spectrum (top panels) and peak
counts (bottom panels) to variations in the cosmological (left) and baryonic
(right) parameters at their fiducial values in an HSC-like survey area (1,500
deg2) and galaxy density (20 amin-2.

In Figure 6, we show the column density power spectra with each of the four
baryon components included. The trends in this figure are very similar to
Figure 3 in Aricò et al. (2020) except that they were showing 3-dimensional
matter power spectra. A key difference in implementing BCM between our work
and Aricò et al. (2020) is that we replace halo particles by the images of
analytical profiles with baryons, while Aricò et al. (2020) moves individual
DM particles to match the effects caused by baryonic physics. We note that our
method has the advantage of being more accurate in representing baryonic
effects when there are only a small number of particles in the haloes, but it
has the disadvantage of not retaining the original shapes of the haloes and
replaces them with spherically symmetric profiles (see § 3.1 above), which can
induce biases in statistics. Regardless of these differences in
baryonification, our method produces matter power spectra of baryon components
that are very similar to Aricò et al. (2020). Figure 6 shows that among these
components, the ejected gas (AGN feedback) has by far the largest affect on
the power spectrum, suppressing it by $7\%$ at $k=1h\,\mathrm{Mpc^{-1}}$.

Figure 7 shows how much the statistics will change if we tweak each
cosmological and baryonic parameter, where we choose an HSC-like scenario:
survey area $A_{\mathrm{survey}}=1,500\,\mathrm{deg^{2}}$ and galaxy density
$n_{\mathrm{gal}}=20\,\mathrm{arcmin^{-2}}$. We find that the sensitivity of
both statistics to baryonic parameters are much smaller than to cosmological
parameters by a factor of $\sim 10-100$, which means the capability of power
spectrum and peak counts in constraining baryonic parameters is
correspondingly weaker. We also find that the responses show degeneracies
between certain cosmological and baryonic parameters. Most prominently, the
response of both the power spectrum and the peak counts to $\eta$ is
qualitatively similar to its response to variations in $\Omega_{\mathrm{m}}$,
in the opposite direction.

###  3.3 Parameter inference with baryons

Figure 8: The posteriors of cosmological and baryonic parameters in an HSC-like survey area and galaxy density. The thick and thin contours show the $1\sigma$ (68%) and $2\sigma$ (95%) credible region respectively, and the black crosses and lines show the fiducial values of the parameters. Methods | $n_{\mathrm{gal}}=20\,\mathrm{arcmin^{-2}}$ | $n_{\mathrm{gal}}=50\,\mathrm{arcmin^{-2}}$
---|---|---
$S_{\mathrm{full}}/10^{-5}$ | $S_{\mathrm{fid}}/10^{-5}$ | $S_{\mathrm{full}}/S_{\mathrm{fid}}$ | $S_{\mathrm{full}}/10^{-5}$ | $S_{\mathrm{fid}}/10^{-5}$ | $S_{\mathrm{full}}/S_{\mathrm{fid}}$
Power spectrum | 12.5 | 4.5 | 2.8 | 8.3 | 2.8 | 3.0
Peak counts | 14.0 | 4.2 | 3.4 | 13.5 | 2.8 | 4.7
Combined | 7.3 | 2.4 | 3.0 | 4.2 | 1.5 | 2.9
PS ($\ell<2,000$) | 13.9 | 5.9 | 2.4 | 11.0 | 4.3 | 2.5
Table 2: The area of the 68% posterior distributions in
$\Omega_{\mathrm{m}}-\sigma_{8}$ space in a survey with an area of
$1500\,\mathrm{deg^{2}}$ and two different galaxy densities, marginalized over
baryonic parameters ($S_{\mathrm{full}}$) or conditioned on fiducial baryonic
parameters ($S_{\mathrm{fid}}$).  Figure 9: The area of the posterior
distributions using the same setting as Table 2 but including higher galaxy
densities.

Figure 8 shows the posteriors of cosmological and baryonic parameters in an
HSC-like survey calculated as described in § 2. To quantify the ability in
constraining cosmological parameters, we marginalize over baryonic parameters
and calculate the area of the posterior distribution $S_{\mathrm{full}}$ in
$\Omega_{\mathrm{m}}-\sigma_{8}$ space enclosed by the $1\sigma$ (68%)
credible contour. The areas of the three statistics (see Table 2) are
$1.2\times 10^{-4}$ (power spectrum), $1.4\times 10^{-4}$ (peak counts), and
$0.7\times 10^{-4}$ (combined).

Among all 2-dimensional marginal distributions of a cosmological and a
baryonic parameter, $\eta$ shows the largest correlation with cosmological
parameters, which is consistent with the analysis on sensitivities in § 3.2.
We define another area of the posterior in $\Omega_{\mathrm{m}}-\sigma_{8}$
space $S_{\mathrm{fid}}$ by fixing the baryonic parameters to their fiducial
values, so that $S_{\mathrm{full}}/S_{\mathrm{fid}}$ represents the factor by
which the joint constraints on cosmological parameters would degrade if a
method were to constrain baryonic parameters at the same time. We find that in
an HSC-like survey, the factor of this degradation is 2.8 for the power
spectrum, 3.4 for the peak counts, and 3.0 for these two combined (see Figure
9). This suggests that the two statistics have similar degeneracies between
cosmological and baryonic parameters, and combining them does not
significantly reduce the degradation. This is again consistent with the
findings in Section 3.2, which show similar degeneracies with $\eta$ for both
the power spectrum and the peak counts.

##  4 Discussion

###  4.1 Parameter inference at different noise levels

In Figure 9, we show the $1\sigma$ area of the posterior in
$\Omega_{\mathrm{m}}-\sigma_{8}$ space at different noise levels marginalized
over baryonic parameters or conditioned on fiducial baryonic parameters. We
assume a survey with HSC-like area ($1500\,\mathrm{deg^{2}}$) and galaxy
densities from $10\,\mathrm{arcmin^{-2}}$ to $100\,\mathrm{arcmin^{-2}}$. We
find that at higher galaxy densities (lower noise levels), the constraints on
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$ generally becomes tighter, meanwhile
the degradation due to marginalizing over the baryonic parameters is larger,
especially for peak counts. We conclude that the statistics can extract more
information from small scales/high peaks but only significantly so when the
baryonic physics is specified.

###  4.2 Comparing power spectrum in full range with $\ell<2,000$

When the lensing power spectrum is used to constrain cosmological parameters,
the upper limit of $\ell$ is typically restricted to $\lesssim 2,000$ (e.g.
Köhlinger et al., 2016, 2017; Hikage et al., 2019). The motivation is that at
even higher $\ell$ (smaller scales), the matter and lensing power spectra
might be affected by baryonic physics to a degree that is hard to model
accurately. In this section, we explore the improvements in constraining
cosmological parameters when the power spectrum at small scales is included.
Figure 9 compares the results with the power spectrum over the full range
($100<\ell<12,000$) with that restricted to large scales only
($100<\ell<2,000$), where we only take the first 11 bins and the covariance
matrices are adjusted accordingly.

When the baryonic parameters are free, we find that the small-scale power
spectrum provides an 11% reduction in the $1\sigma$ area in an HSC-like
survey; when the baryonic parameters are fixed, this reduction increases to
24%. If the galaxy density is increased from $20\,\mathrm{arcmin^{-2}}$ to
$50\,\mathrm{arcmin^{-2}}$ matching the noise level of the Nancy Grace Roman
Space Telescope, the reductions will be further increased to 25% (baryons
free) and 36% (baryons fixed). As comparison, Fang & Haiman (2007) found
($n_{\mathrm{gal}}\approx 30\,\mathrm{arcmin^{-2}}$, no baryons) that by
increasing the upper limit of $\ell$ from 1,000 to 3,000, the uncertainty in
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$ are improved by 30% and 50%
respectively.

Nominally, a comparison by using the fixed-baryons constraints restricted to
large angular scales gives better constraints ($S_{\rm fid}\approx 6$ for at
$20\,\mathrm{arcmin^{-2}}$; blue curve in Figure 9) than including smaller
scales but at the cost of marginalising over the baryons ($S_{\rm full}\approx
12$; red curve). However, we find this conclusion to be unjustified. Contrary
to the notion that baryons only affect small scales, we find that restricting
the power spectrum to $\ell<2,000$ does not reduce the degradation caused by
uncertainties in the baryonic parameters. Table 2 shows that
$S_{\mathrm{full}}/S_{\mathrm{fid}}$ for the large-scale power spectrum equals
$2.4$, which is only modestly lower than that of the full-range power spectrum
(by $0.4$). This means that taking baryonic physics into consideration is
necessary even for the large-scale power spectrum, at least in the context of
the BCM models adopted here.

###  4.3 Priors on baryonic parameters

Figure 10: $1\sigma$ confidence area in $\Omega_{\mathrm{m}}-\sigma_{8}$ space
for the power spectrum and peak counts combined. The solid lines show one case
where the prior of one baryonic parameter is varied and the others are free.
The dashed lines show another case, where the prior of one baryonic parameters
is varied and the other parameters are fixed at their fiducial values. Figure
11: The same as Figure 10, except that only the large-scale power spectrum
($\ell<2,000$) is used.

The results we show in § 3.3 assume wide priors on the baryonic parameters
constrained only by WL, but we note that external observations or theoretical
models can give stricter priors. Ideally, the accuracy of these parameters can
be high enough so that the uncertainties of the observables caused by baryonic
physics is negligible relative to those of cosmological parameters. In this
section, we explore which of the baryonic parameters are most important in
degrading our ability to constrain $\Omega_{\mathrm{m}}$ and $\sigma_{8}$.

We refer to the full width of the prior on each baryonic parameter as 100%,
and we then gradually shrink the width of each prior linearly on a logarithmic
scale until both the upper and lower bounds reach the fiducial value (we refer
to the width of which as 0%). We can manipulate the priors on the baryonic
parameters in two simple ways: 1) adjust the width of prior on one parameter
between 100% and 0% while keeping the other three priors at 100%, or 2) adjust
the width of prior on one parameter and fix the other three parameters at
their fiducial values (i.e. assume 0% priors on the latter). The resulting
$1\sigma$ credible areas are shown in Figure 10 as a function of the prior
width. They indicate the following:

  * •

If we have little knowledge about baryons and have to resort to wide priors,
$\eta$ is the most important parameter, with the largest potential to improve
the constraint on $\Omega_{\mathrm{m}}$ and $\sigma_{8}$ (solid purple curve
in Figure 10).

  * •

If we can determine baryonic parameters accurately, $M_{\mathrm{c}}$ becomes
the most important limiting parameter, i.e. a biased $M_{\mathrm{c}}$ is the
least tolerable (dashed orange curve in Figure 10).

  * •

$M_{1,0}$ and $\beta$ affect the constraint the least, regardless of our prior
knowledge on the other parameters.

We also show the corresponding results in Figure 11 restricted to the large-
scale only power spectrum ($\ell<2,000$). We find that in both ways of
changing priors, $\eta$ is then the most important parameter, followed by
$M_{1,0}$. This is consistent with the response curves in Figure 7, where
$\eta$ is dominant, but only dominant at large scales.

A few observational results can be used to find the best-fitting values of the
baryonic parameters. Here, we list some of these with uncertainties and their
width relative to those of the priors we employ. As mentioned by Schneider &
Teyssier (2015), Ade et al. (2013) suggests $\eta=0.3-0.7$ ($30\%$ of the
width of our prior) from the radial profile of the cluster gas fraction; X-ray
data (Sun et al., 2009; Vikhlinin et al., 2009; Gonzalez et al., 2013)
suggests $\beta=0.3-0.8$ ($20\%$ width) and $M_{\mathrm{c}}=0.2-2.0\times
10^{14}h^{-1}\mathrm{M}_{\odot}$ ($26\%$ width) from the bounded gas
fraction–$M_{500}$ relation. Stellar mass–halo mass relations from various
sources listed by Behroozi et al. (2013) suggests $M_{1,0}=(2-10)\times
10^{11}h^{-1}\mathrm{M}_{\odot}$ ($34\%$ width). Additionally, the matter
spectrum using BCM has been fit to those obtained in hydrodynamical
simulations to infer baryonic parameters by Aricò et al. (2020).
Interestingly, the range of simulations examined in that study generally agree
on the value of $\eta$ to be $0.1-1.0$ (70% width) but not on the other
baryonic parameters.

###  4.4 Born approximation

Figure 12: Same as Figure 8, but the model $\mathbf{y}(\theta)$ used to fit
the (fully ray-traced) mock data $\mathbf{y}_{0}$ employs the Born
approximation. The colored crosses in the $\Omega_{\mathrm{m}}-\sigma_{8}$
panel shows the maxima of the posterior distributions marginalized over the
baryonic parameters.  Figure 13: The posterior distributions with
$\mathbf{y}(\theta)$ using the Born approximation and fixing baryonic
parameters to their fiducial values. The supposed survey has a galaxy density
of $n_{\mathrm{gal}}=20\,\mathrm{arcmin^{-2}}$ and an area of $1,500$,
$5,000$, or $15,000\,\mathrm{deg^{2}}$. The thick and thin contours show the
$1\sigma$ and $2\sigma$ credible regions, respectively. The black crosses mark
the fiducial cosmology, and the colored crosses mark the maximum likelihood
cosmology inferred from the different statistics.

The large majority of previous works investigating baryonic effects obtained
lensing power spectra via a direct projection of the three-dimensional matter
power spectrum (Rudd et al., 2008; Zentner et al., 2008; Mead et al., 2015),
or used the Born approximation in place of full ray-tracing (Semboloni et al.,
2011; Fong et al., 2019; Huang et al., 2019; Weiss et al., 2019, but see Osato
et al. 2021 for a recent exception).

In the Born approximation, the convergence at position $\boldsymbol{\theta}$
from source comoving distance $\chi_{\mathrm{s}}$ is approximated by taking
the lowest order term of the ray-tracing calculations:

$\kappa_{\mathrm{born}}(\boldsymbol{\theta})=\frac{3H_{0}^{2}\Omega_{\mathrm{m}}}{2c^{2}}\int_{0}^{\chi_{\mathrm{s}}}\frac{\chi\,\mathrm{d}\chi_{\mathrm{s}}}{a(\chi)}W(\chi,\chi_{\mathrm{s}})\delta(\boldsymbol{\theta},\chi),$
(13)

where $a$ denotes the scale factor,
$W(\chi,\chi_{\mathrm{s}})=1-\chi/\chi_{\mathrm{s}}$ the lensing kernel, and
$\delta$ the density contrast. Note that the calculation of
$\kappa_{\mathrm{born}}$ only depends on the density contrast values along the
line of sight, which itself is approximated as a straight line. It makes the
Born approximation much cheaper than ray-tracing in terms of computational
cost. But assuming that the observed lensing effects can only be resembled by
ray-tracing, we need to assess the biases introduced by the Born approximation
whenever applied to WL surveys.

We build a set of convergence maps using the Born approximation in parallel
with the ray-tracing maps, and we apply them to two scenarios of parameter
inference. First, we replace the model $\mathbf{y}(\theta)$ in Equation 9 with
the model derived from the Born approximation convergence maps (but not
replacing $\mathbf{y}_{0}$), and infer cosmological and baryonic parameters,
the sampled posteriors of which is shown in Figure 12. We find that the Born
approximation can cause a large bias in this scenario for an HSC-like survey:
the power spectrum method is only slightly biased ($<0.5\sigma$ for
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$), but peak counts have a $1\sigma$ bias
in the two cosmological parameters and $\sim 2\sigma$ biases on $M_{1,0}$ and
$\beta$. Second, we fix the baryonic parameters to their fiducial values and
infer $\Omega_{\mathrm{m}}$ and $\sigma_{8}$. As is shown in Figure 13,
cosmological parameters are biased by a similar absolute amount, but since the
constraints are tighter with larger survey area, the biases are more
statistically significant. For the power spectrum, the significance reaches
$1\sigma$ when $A_{\mathrm{survey}}=15,000\,\mathrm{deg^{2}}$, and for peak
counts, the significance reaches $1\sigma$ when
$A_{\mathrm{survey}}=5,000\,\mathrm{deg^{2}}$. Similar results have also been
noted by Petri et al. (2017) for the power spectrum and moments of the
convergence field (in the absence of baryons).

In summary, using the Born approximation instead of ray-tracing has a greater
impact on the peak counts than on the power spectrum, possibly because peak
counts rely more heavily on small-scale information where there are larger the
deviations between the Born approximation and ray-tracing. In an HSC-like
survey, ray-tracing is necessary for inferring baryonic parameters or to
correctly marginalize over baryons but only marginally necessary for inferring
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$ for fixed baryonic parameters.

###  4.5 The dependence of covariances on baryonic parameters

Figure 14: The correlation matrix $\rho$ of the observables in the fiducial
model (upper panels), the difference of $\rho$ between the $\eta=0.54$ model
(all other parameters remain unchanged) and the fiducial model (middle
panels), and the difference between the
($\Omega_{\mathrm{m}}=0.315,\sigma_{8}=0.717$) model and the fiducial model.
The determinants of the covariance matrices in all three models are shown at
the bottom left corner of each panel.

In Figure 14, we show how covariances change with $\eta$, the most important
baryonic parameter. As an indication for the importance of the
$\eta$-dependence, we compare the determinants of the covariance matrices and
the correlations between the observables with two $\eta$ values—0.54 (fiducial
value) and 1.5. We find that there are only slight differences between them:
the determinant changes by only 7% and 20% for the two statistics. By
comparison, the determinant varies by 10% and a factor of $\approx 2$,
respectively, due to a 10% change in $\sigma_{8}$ ($0.789\to 0.717$), which is
a typical level of uncertainty from previous weak lensing surveys (see e.g.
Hikage et al., 2019).

##  5 Conclusions

In this work, we adopted a start-of-the-art baryonic correction model (BCM) to
investigate the impact of baryons on the inference of cosmological parameters
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$ from Gaussian and non-Gaussian WL
statistics (the power spectrum and peak counts). In order to be able to
incorporate the baryonic effects into a large suite of ray-tracing
simulations, we employed a low computational-cost methodology where we
“painted” baryons onto $N$-body simulations by replacing the particles in DM
haloes by images of analytical profiles that contain baryons.

We summarize our main findings as follows:

  1. 1.

The difference in column densities between DM haloes from the $N$-body
simulations and their NFW-halo replacements is mainly the quadrupole
component, which explains most of the discrepancy in the column density power
spectrum at $k\gtrsim 3\,h\,\mathrm{Mpc^{-1}}$.

  2. 2.

If the baryonic parameters are inferred simultaneously with $\Omega_{m}$ and
$\sigma_{8}$ and being marginalized over, the area of the $1\sigma$ region in
$\Omega_{\mathrm{m}}-\sigma_{8}$ space will be $2-4$ times larger than that
with baryonic parameters fixed. This factor is larger for peak counts and for
higher galaxy densities (less noisy data).

  3. 3.

The quality of constraints on $\Omega_{\mathrm{m}}$ and $\sigma_{8}$ in terms
of the $1\sigma$ area improves the most when $\eta$ and $M_{\mathrm{c}}$ have
tighter priors, while $M_{1,0}$ and $\beta$ barely affects the $1\sigma$ area.

  4. 4.

In an HSC-like survey, expanding the range of lensing power spectrum from
$\ell<2,000$ to $\ell<12,000$ can improve the constraints on
$\Omega_{\mathrm{m}}$ and $\sigma_{8}$ by 11% when marginalized over baryonic
parameters and by 24% when baryonic parameters are fixed. These improvements
will be larger at higher galaxy densities.

  5. 5.

Baryonic physics, with the BCM models we adopted, and in the absence of strong
priors, affects the power spectrum on both small and large scales; a limit of
$\ell<2,000$ still requires the modelling of baryons.

  6. 6.

Using the Born approximation to generate convergence maps has a greater impact
on peak counts than power spectrum, but is necessary to correctly model
baryonic effects and marginalize over them, for both statistics in an HSC-like
survey.

We note two caveats in this work that do not affect our finding above but can
be a problem if we apply the same methods to observational data in the future.
We need to carry the shape information of the DM haloes from the $N$-body
simulation to the analytical profiles to remove the halo replacement
residuals. We also need to generate more realizations of the convergence map
for each combination of parameters so that when comparing to the real data,
the statistics are not biased due to insufficient randomization. More
generally, before they can be applied to real data, a better calibration of
the specific parameterized BCM will be necessary against physical
hydrodynamical models, including a range of baryon physics and their impact on
peak counts.

## Acknowledgements

We thank Colin Hill and Lam Hui for useful discussions and acknowledge support
by NASA ATP grant 80NSSC18K1093. We acknowledge support from the NSF XSEDE
facility for data analysis in this study.

## References

  * Abbott et al. (2005) Abbott, T., Aldering, G., Annis, J., et al. 2005, arXiv preprint astro-ph/0510346
  * Ade et al. (2013) Ade, P. A., Aghanim, N., Arnaud, M., et al. 2013, Astronomy & Astrophysics, 550, A131
  * Aihara et al. (2018) Aihara, H., Arimoto, N., Armstrong, R., et al. 2018, Publications of the Astronomical Society of Japan, 70, S4
  * Aricò et al. (2020) Aricò, G., Angulo, R. E., Hernández-Monteagudo, C., et al. 2020, Monthly Notices of the Royal Astronomical Society, 495, 4800
  * Bartelmann & Schneider (2001) Bartelmann, M., & Schneider, P. 2001, Physics Reports, 340, 291
  * Behroozi et al. (2013) Behroozi, P. S., Wechsler, R. H., & Conroy, C. 2013, The Astrophysical Journal, 770, 57
  * Behroozi et al. (2012) Behroozi, P. S., Wechsler, R. H., & Wu, H.-Y. 2012, The Astrophysical Journal, 762, 109
  * de Jong et al. (2013) de Jong, J. T., Kleijn, G. A. V., Kuijken, K. H., Valentijn, E. A., et al. 2013, Experimental Astronomy, 35, 25
  * Fang & Haiman (2007) Fang, W., & Haiman, Z. 2007, Physical Review D, 75, 043010
  * Fong et al. (2019) Fong, M., Choi, M., Catlett, V., et al. 2019, Monthly Notices of the Royal Astronomical Society, 488, 3340
  * Gonzalez et al. (2013) Gonzalez, A. H., Sivanandam, S., Zabludoff, A. I., & Zaritsky, D. 2013, The Astrophysical Journal, 778, 14
  * Hamana et al. (2020) Hamana, T., Shirasaki, M., Miyazaki, S., et al. 2020, Publications of the Astronomical Society of Japan, 72, 16
  * Hartlap et al. (2007) Hartlap, J., Simon, P., & Schneider, P. 2007, Astronomy & Astrophysics, 464, 399
  * Heymans et al. (2012) Heymans, C., Van Waerbeke, L., Miller, L., et al. 2012, Monthly Notices of the Royal Astronomical Society, 427, 146
  * Hikage et al. (2019) Hikage, C., Oguri, M., Hamana, T., et al. 2019, Publications of the Astronomical Society of Japan, 71, 43
  * Hilbert et al. (2009) Hilbert, S., Hartlap, J., White, S., & Schneider, P. 2009, Astronomy & Astrophysics, 499, 31
  * Hoekstra & Jain (2008) Hoekstra, H., & Jain, B. 2008, Annual Review of Nuclear and Particle Science, 58, 99
  * Huang et al. (2019) Huang, H.-J., Eifler, T., Mandelbaum, R., & Dodelson, S. 2019, Monthly Notices of the Royal Astronomical Society, 488, 1652
  * Ivezić et al. (2019) Ivezić, Ž., Kahn, S. M., Tyson, J. A., et al. 2019, The Astrophysical Journal, 873, 111
  * Jain et al. (2000) Jain, B., Seljak, U., & White, S. 2000, The Astrophysical Journal, 530, 547
  * Jenkins et al. (2001) Jenkins, A., Frenk, C., White, S. D., et al. 2001, Monthly Notices of the Royal Astronomical Society, 321, 372
  * Jing et al. (2006) Jing, Y., Zhang, P., Lin, W., Gao, L., & Springel, V. 2006, The Astrophysical Journal Letters, 640, L119
  * Joudaki et al. (2016) Joudaki, S., Blake, C., Heymans, C., et al. 2016, Monthly Notices of the Royal Astronomical Society, stw2665
  * Kilbinger (2015) Kilbinger, M. 2015, Reports on Progress in Physics, 78, 086901, doi: 10.1088/0034-4885/78/8/086901
  * Köhlinger et al. (2016) Köhlinger, F., Viola, M., Valkenburg, W., et al. 2016, Monthly Notices of the Royal Astronomical Society, 456, 1508
  * Köhlinger et al. (2017) Köhlinger, F., Viola, M., Joachimi, B., et al. 2017, Monthly Notices of the Royal Astronomical Society, 471, 4412
  * Laureijs et al. (2011) Laureijs, R., Amiaux, J., Arduini, S., et al. 2011, arXiv preprint arXiv:1110.3193
  * Lewis et al. (2000) Lewis, A., Challinor, A., & Lasenby, A. 2000, The Astrophysical Journal, 538, 473
  * Mead et al. (2015) Mead, A., Peacock, J., Heymans, C., Joudaki, S., & Heavens, A. 2015, Monthly Notices of the Royal Astronomical Society, 454, 1958
  * Mohammed et al. (2014) Mohammed, I., Martizzi, D., Teyssier, R., & Amara, A. 2014, arXiv preprint arXiv:1410.6826
  * Nelson et al. (2019) Nelson, D., Springel, V., Pillepich, A., et al. 2019, Computational Astrophysics and Cosmology, 6, 1
  * Osato et al. (2021) Osato, K., Liu, J., & Haiman, Z. 2021, Monthly Notices of the Royal Astronomical Society, 502, 5593
  * Petri (2016) Petri, A. 2016, Astronomy and Computing, 17, 73
  * Petri et al. (2016) Petri, A., Haiman, Z., & May, M. 2016, Phys. Rev. D, 93, 063524, doi: 10.1103/PhysRevD.93.063524
  * Petri et al. (2017) Petri, A., Haiman, Z., & May, M. 2017, Physical Review D, 95, 123503
  * Refregier (2003) Refregier, A. 2003, Annual Review of Astronomy and Astrophysics, 41, 645
  * Rudd et al. (2008) Rudd, D. H., Zentner, A. R., & Kravtsov, A. V. 2008, The Astrophysical Journal, 672, 19
  * Schaye et al. (2010) Schaye, J., Vecchia, C. D., Booth, C., et al. 2010, Monthly Notices of the Royal Astronomical Society, 402, 1536
  * Schneider & Teyssier (2015) Schneider, A., & Teyssier, R. 2015, Journal of Cosmology and Astroparticle Physics, 2015, 049
  * Schneider et al. (2019) Schneider, A., Teyssier, R., Stadel, J., et al. 2019, Journal of Cosmology and Astroparticle Physics, 2019, 020
  * Schneider & Bridle (2010) Schneider, M. D., & Bridle, S. 2010, Monthly Notices of the Royal Astronomical Society, 402, 2127
  * Semboloni et al. (2011) Semboloni, E., Hoekstra, H., Schaye, J., van Daalen, M. P., & McCarthy, I. G. 2011, Monthly Notices of the Royal Astronomical Society, 417, 2020
  * Sifón et al. (2015) Sifón, C., Hoekstra, H., Cacciato, M., et al. 2015, Astronomy & Astrophysics, 575, A48
  * Sobol’ (1967) Sobol’, I. M. 1967, Zhurnal Vychislitel’noi Matematiki i Matematicheskoi Fiziki, 7, 784
  * Spergel et al. (2015) Spergel, D., Gehrels, N., Baltay, C., et al. 2015, arXiv preprint arXiv:1503.03757
  * Springel (2005) Springel, V. 2005, Monthly notices of the royal astronomical society, 364, 1105
  * Sun et al. (2009) Sun, M., Voit, G., Donahue, M., et al. 2009, The Astrophysical Journal, 693, 1142
  * Vikhlinin et al. (2009) Vikhlinin, A., Burenin, R., Ebeling, H., et al. 2009, The Astrophysical Journal, 692, 1033
  * Weiss et al. (2019) Weiss, A. J., Schneider, A., Sgier, R., et al. 2019, Journal of Cosmology and Astroparticle Physics, 2019, 011
  * White (2004) White, M. 2004, Astroparticle Physics, 22, 211
  * Yang et al. (2013) Yang, X., Kratochvil, J. M., Huffenberger, K., Haiman, Z., & May, M. 2013, Physical Review D, 87, 023511
  * Zentner et al. (2008) Zentner, A. R., Rudd, D. H., & Hu, W. 2008, Physical Review D, 77, 043507
  * Zhan & Knox (2004) Zhan, H., & Knox, L. 2004, The Astrophysical Journal Letters, 616, L75


In drug discovery, aqueous solubility is an important pharmacokinetic property which affects absorption and assay availability of drug.
Thus, in silico prediction of solubility has been studied for its utility in virtual screening and lead optimization.
Recently, machine learning (ML) methods using experimental data has been popular because physics-based methods like quantum mechanics and molecular dynamics are not suitable for high-throughput tasks due to its computational costs.
However, ML method can exhibit over-fitting problem in a data-deficient condition, and this is the case for most chemical property datasets.
In addition, ML methods are regarded as a black box function in that it is difficult to interpret contribution of hidden features to outputs, hindering analysis and modification of structure-activity relationship.
To deal with mentioned issues, we developed Bayesian graph neural networks (GNNs) with the self-attention readout layer.
Unlike most GNNs using self-attention in node updates, self-attention applied at readout layer enabled a model to improve prediction performance as well as to identify atom-wise importance, which can help lead optimization as exemplified for three FDA-approved drugs.
Also, Bayesian inference enables us to separate more or less accurate results according to uncertainty in solubility prediction task
We expect that our accurate, reliable and interpretable model can be used for more careful decision-making and various applications in the development of drugs.

§ INTRODUCTION

Aqueous solubility (hereinafter referred to as "solubility") is one of the most necessary pharmacokinetic properties in drug discovery, since it affects absorption, bioavailability and toxicity of drug candidates.
Predicting solubility is also essential for early-stage drug discovery, as insoluble compounds may not be available for biochemical assays.
For these two reasons, in silico prediction of solubility has been heavily investigated.
An alternative approach is to use freely available cheminformatics software such as RDKit[Landrum, 2013] to obtain water-octanol partition coefficient (cLogP).
However, as shown in Figure <ref>, cLogP is not a good alternative to experimentally-determined solubility.

Scatter plot between cLogP values obtained by RDKit and solubility values provided by AqSolDB[Sorkun et al., 2019].

More sophisticated approaches, like other chemical property prediction problems, can be categorized into three: quantum mechanics (QM), molecular dynamics (MD), and data-driven methods.

COSMO-RSol calculates solubility based on the free energy change of solid to liquid state using a quantitative structure-property relationship (QSPR) model to estimate free energy change of fusion and QM calculation to estimate the difference between chemical potential of pure and solution states.[Klamt et al., 2002]
MD predicts solubility through a two-phase crystal-liquid system simulation.
The effect of transition can be calculated by integrating simulations performed on a crystal or a liquid system independently or by metadynamics of a coexisting system directly.[Bjelobrk et al., 2021]
However, due to heavy computational cost, both QM and MD approaches may be insuitable for virtual screening.
On the other hand, data-driven methods, which are mostly based on machine learning (ML), learn the relationship between molecular features and target properties given training data samples.
However, as these models require training data samples, inferior size and quality of available data have hindered their implementation to virtual screening.[Sorkun et al., 2019]

Lessons learnt from the previous solubility prediction challenges, recent works has begun to focus on data curation and feature engineering.[Sorkun et al., 2019, Sorkun et al., 2021, Lovrić et al., 2021]
In the first challenge sponsored by Pfizer and hosted by Goodman group,[Llinàs et al., 2008] most submitted models showed lower than 60% accuracy,[Hopfinger et al., 2009] which was criticized to be due to improper algorithms or molecular features rather than the noisy data (i.e., participants were provided with only 100 training data samples).[Palmer and Mitchell, 2014]
However, recent research demonstrates that simple algorithms such as support vector machine and partial least squares with well-curated data showed comparable performance.[Lovrić et al., 2021]
Results from the second solubility challenge also back up the importance of data quality and informative descriptors.[Llinas et al., 2020]
There was no significant difference between models trained with thousands of samples and only few hundreds samples because both datasets were not able to sufficiently account for chemical space of the drug-like molecules.

In addition, state-of-the-art ML algorithms, such as graph neural networks and XGBoost, exhibited similar or even inferior performance to multiple linear regression.

Previous studies have mostly focused on improving prediction accuracy, but data-driven models show deteriorated performance on predicting unseen molecules as exemplified in the previous challenges.
Therefore, we would like to argue that it is necessary to assess models in terms of not only accuracy but also predictive uncertainty and interpretability for real task.
For example, if a given query is an out-of-distribution sample, a model should exhibit high predictive uncertainty, which can help users to agree or disagree the prediction results.[Ryu et al., 2019, Schwaller et al., 2019, Soleimany et al., 2021, Graff et al., 2021]
In addition, interpretability, such as contribution of molecular substructures to predictive solubility, enables users to infer how improve solubility by modifying substructures.[Harren et al., 2022]

In this work, we present a solubility prediction model with accurate, reliable and interpretable results.
In contrast to QSPR models based on molecular descriptors and old-fashioned ML models (e.g. support vector machine and random forest), we utilized graph neural networks in order to learn features from graph representation of molecular structures.
Also, we adopt self-attention[Vaswani et al., 2017] in the readout function, which aggregates atom features and produces a graph feature per single molecule.
Our designed readout not only improves prediction performance but also enables us to interpret the attention weight as importance value.
Lastly, we confirm that using Bayesian inference results in improved prediction performance than non-Bayesian approach and uncertainty quantification, which can be utilized to prioritize more accurate results based on predictive uncertainty.

§ METHODS AND IMPLEMENTATIONS

§.§ Dataset and data splitting

We used the AqSolDB dataset [Sorkun et al., 2019] provided by Therapeutic Data Commons [Huang et al., 2021] to develop our neural networks.
This dataset consists of 9,982 molecules, where the structures are given in simplified molecular-input line-entry system (SMILES) format and the labels are annotated with continuous solubility indicator value.
As described in [Huang et al., 2021][Huang et al., 2021], molecules can be categorized by the indicator value ($y_{true}$) as follows:

* $y_{true} > 0.0$ : highly soluble

* $ -2.0 < y_{true} \leq 0.0 $ : soluble

* $ -4.0 < y_{true} \leq -2.0 $ : partially soluble

* $y_{true} < -4.0$ : insoluble

We developed our models with those SMILES and label pairs with the four-fold scaffold splitting[Wu et al., 2018] of entire set into 70%/10%/20% for train/validation/test sets.
Both classification and regression models were obtained with and without converting the indicator values to the categorical labels, respectively.

§.§ Input featurization and Graph neural networks

Our baseline model framework is graph neural network (GNN) which utilizes the graph representation of molecular structure as their inputs.
To generate inputs for our model, we featurize molecular graphs $G(V,E)$, where $V$ is the set of node features $\{ h_i^{(0)} \}$, $E$ is the set of edge features ${\{ e^{(0)}_{ij}} \}$ and $i$ and $j$ are node indices, with RDKit[Landrum, 2013].
We assign initial features to each nodes by using the RDKit functions as follows:

* One-hot encoding of atom types: atom.GetSymbol()

* One-hot encoding of atom degree: atom.GetDegree()

* One-hot encoding of number of hydrogens: atom.GetTotalNumHs()

* One-hot encoding of implicit valence number: atom.ImplicitValence()

* Indicator value for whether given atom is aromatic: atom.GetIsAromatic()

Also, initial edge features are assigned as follows:

* One-hot encoding of atom bond types: see our code released in github.

* Indicator value for whether bond is conjugated: bond.GetIsConjugated()

* Indicator value for whether bond is in ring: bond.IsInRing()

[Hwang et al., 2020] shows that descriptors, which can appropriately describe atom and bond characters, significantly affects prediction performance of GNNs,
and thus we follow the input featurization scheme of the previous work.[Hwang et al., 2020]

Using the featurized input molecular graphs, the $l$-th message passing layer updates the $i$-th node feature $h_{i}^{(l-1)}$ to the output $h_{i}^{(l)}$ by aggregating the adjacent node and edge features:

\begin{equation}
h_i^{(l)} = \text{NodeUpdate} (h^{(l-1)}_{i}, \{h^{(l-1)}_{j}\}_{j \in \mathcal{N}_{i}}, \{e^{(l-1)}_{ij}\}_{j \in \mathcal{N}_{i}}),
\end{equation}

where $\mathcal{N}_{i}$ stands for the indices of the neighbor nodes of the $i$-th node.

After passing the $L$-stack of message passing layers, the node features are aggregated and summarized to a single vector $z_{G}$, so called the graph feature, by the readout layer:

\begin{equation}
z_{G} = \text{Readout}(\{h_i^{(L)}\}).
\end{equation}

Then, the final prediction layer converts the graph feature to a predictive value:

\begin{equation}
y_{pred} = \text{Linear}(z_{G}).
\end{equation}

§.§ Node updates

We implemented three different GNN node updates and then chose the best model based on the results from the first validation experiment, as shown in Figure <ref>.
Graph convolutional network (GCN)[Duvenaud et al., 2015, Kipf and Welling, 2016] is the simplest node update among the three implementations.
The node update equation of GCN can be described as follows:

\begin{equation}
h_i^{(l)} = \text{LN}(h_i^{(l-1)} + \sigma(W^{(l)}\sum_{j \in \mathcal{N}_{i}} h^{(l-1)}_{j})),
\end{equation}

where $\text{LN}$ stands for layer normalization[Ba et al., 2016] and $\sigma(\cdot)$ is non-linear activation.

Graph isomorphism network (GIN)[Xu et al., 2018] can be thought as more sophisticated version of GCN, which replaces the linear layer of GCN to the multi-layer perceptron (MLP):

\begin{equation}
h_i^{(l)} = \text{LN}(h_i^{(l-1)} + \text{MLP}^{(l)} (\sum_{j \in \mathcal{N}_{i}} h^{(l-1)}_{j})).
\end{equation}

where $\text{MLP}$ stands for

\begin{equation}
\text{MLP}^{(l)}(x) = W_2^{(l)}\sigma(W_1^{(l)}x + b_1^{(l)}) + b_2^{(l)}
\end{equation}

Furthermore, to utilize bond information, we include edge features to the GIN layer's aggregation of node features:

\begin{equation}
h_i^{(l)} = \text{LN}(h_i^{(l-1)} + \text{MLP}^{(l)} (\sum_{j \in \mathcal{N}_{i}} h^{(l-1)}_{j} + e_{ij}^{(l-1)})).
\end{equation}

We expect that the GIN shows better results than the GCN as the former transforms both node and edge features of adjacent nodes.

Lastly, we implemented a modified version of graph attention network[Veličković et al., 2017], which adopts self-attention in node updates.
Our modified version, namely graph transformer (GT), incorporates edge features in both computation of attention weights and aggregation of adjacent node features:

\begin{equation}
h_i^{(l)} = \text{LN}(h_i^{(l-1)} + \text{MLP}^{(l)} (\hat{h}_{i}^{(l)})),
\end{equation}

where the intermediate node feature $\hat{h}_{i}^{(l)}$ is given by

\begin{equation}
\hat{h}_{i}^{(l)} = \text{LN}(W_{6}^{(l)}h^{(l-1)}_{i} + \sum_{j \in \mathcal{N}_{i}} w_{ij}^{(l)} (W_{5}^{l} h^{(l-1)}_{j} + W_{4}^{(l)} e_{ij}^{(l-1)}) ).
\end{equation}

The attention coefficient $w_{ij}^{(l)}$ between two adjacent nodes $i$ and $j$ is computed by softmax activation of scaled-dot products between the source node feature and the adjacent node and edge features.:

\begin{equation}
w_{ij}^{(l)} = \text{softmax}(\frac{1}{\sqrt{d}} (W_{3}^{(l)} h_{i}^{(l-1)}) (W_{2}^{(l)} h_{j}^{(l-1)} + W_{1}^{(l)} e_{ij}^{(l-1)})^{T} ).
\end{equation}

We note that our implementation of GIN and GT is inspired by that introduced in PyTorch Geometric[Fey and Lenssen, 2019] implementations.

§.§ Readout with self-attention

In most GNN implementations, the readout (or sometimes referred to as pooling) step adopts mean and sum operations to aggregate node features that are updated by a stack of message passing layers.
Both operations do not assign different importance to each node, but aggregates node features with equal weights.

On the other hand, we hypothesize that assigning different atom-wise importance to each node results in more accurate prediction:

\begin{equation}
z_{G} = \sum_{v} \alpha_v h_v^{(L)},
\end{equation}

where $\alpha_v$ is the attention weight given to the $v$-th node, since attention weights can be considered as the importance of objects – in our case node (atom) features.
We expect that adopting attention weights is not only improves prediction performance but also helps model interpretability, as is validated by computer vision and language understanding works[Dosovitskiy et al., 2020, Vaswani et al., 2017].

Therefore, we utilize pooling by multi-head attention (PMA) firstly introduced in Set Transformer[Lee et al., 2019],
where attention weight $\alpha_v$ is given by

\begin{equation}
\alpha_v = \text{softmax}(\frac{1}{\sqrt{d}} \textbf{1} (W_{\text{pma}} h_{v}^{(L)})^{T} ).
\end{equation}

Although PMA implemented in Set Trasanformer uses a  randomized seed vector as a query vector,
we found that a fixed vector 1 initialized with one shows better performance in our solubility prediction task.

In our experiments, we compare mean and PMA readouts as shown in Figure <ref>.
As described in [Xu et al., 2018], mean and sum readouts are appropriate to summarize node statistics in terms of fraction and exact number, respectively.[Xu et al., 2018]

In solubility prediction, using ratio between hydrophilic and hydrophobic groups is more appropriate than exact number of those groups to summarize node statistics.

In this work, we conjecture and experimentally confirm that sum readout leads to unstable training and mean readout is appropriate.

§.§ Bayesian learning for uncertainty quantification and reliable prediction

Our model quantifies uncertainty and calibrates predictive label with Bayesian learning and inference.
As shown in [Ryu et al., 2019], uncertainty quantification enables detecting erroneously labeled samples in a dataset based on the data-driven (aleatoric) uncertainty.
This offers that uncertainty can be utilized to prioritize predictions expected to be more accurate.
Following a benchmark study that uses Bayesian inference for reliable prediction of molecular properties[Hwang et al., 2020],
we used Monte-Carlo dropout (MC-DO)[Gal and Ghahramani, 2016] and stochastic weight averaging (SWA)[Maddox et al., 2019] to develop regression and classification models, respectively.

§.§ Code availability

We release our code in <https://github.com/SeongokRyu/gnn_solubility> to clarify implementations and reproducibility of experimental results.

§ RESULTS AND DISCUSSION

§.§ Attention-based readout and Bayesian learning leads to more accurate results

Prediction performance of regression models in terms of (a) root-mean-square-error (RMSE) and (b) coefficient of determination ($R^2$).

Firstly, we compare the performance of solubility regression models built on various model architectures and Bayesian learning algorithm in order to seek the best model specification.
We investigate three node feature updates (i.e., GCN, GIN and GT) and two graph feature updates (i.e., mean and PMA).
Also, we test the usefulness of approximate Bayesian learning with MC-DO.

In Figure <ref>, we show the performances in terms of root-mean-square-error (RMSE) and coefficient of determination ($R^2$).
Contrary to previous researches, which showed that using self-attention (i.e. GT) or edge features (i.e. GIN) improve prediction power,
the simplest node update (i.e. GCN and GT) shows the best predictive performance among three node updates.
On the other hand, regardless of node updating methods, the attention-based readout (i.e. PMA) consistently outperforms the mean readout.
The results imply that fine-tuning the readout layer of GNN architectures is more helpful than tuning the node update layers to improve the predictive power, at least in our solubility prediction task.

Next, by comparing the point-estimation of model parameters (labeled as `Vanila') and Bayesian model (labeled as `MC-DO'),
we confirm that adopting approximate Bayesian learning significantly improves the predictive performance.
Although using MC-DO gives us higher RMSE and lower $R^2$ on the train set, it gives us lower RMSE and $R^2$ on the validation and test sets,
and thereby reduces the generalization gap between the train set and validation/test sets.

Classification accuracy of the regression and classification models. We note that the prediction results from the regression model were converted to the categorical values and evaluated in terms of classification accuracy.

Since the solubility indicator value given by the AqSolDB dataset can be categorized into four classes, as described in the method section, we trained the classification model and evaluated the classification performance of the best regression model by converting true labels and predictive labels to categorical values.

Figure <ref> demonstrates the accuracy of both the regression and the classification models.
Though the regression model underperformed in the train set, opposite results were obtained in the validation and test sets.
We conjecture that the training setting of regression tasks might be more information-rich than that of classification tasks.
For example, indicator values -1.9 and -0.1 are both categorized into the `soluble' group, but the former may show smaller solubility than the latter.
On the other hand, categorical values are identical for both values, provides the same information to the model during the training stage.
Therefore, continuous indicator values provide more information than categorical values, resulting in better generalization results.

Based on aforementioned results, we chose the regression model built on the GCN node update, the PMA readout, and the MC-DO Bayesian inference as the baseline architecture for further investigation.

§.§ More reliable predictions can be prioritized with predictive uncertainty

Prediction performance of the best regression models for (a) train, (b) validation and (c) test sets. We report the prediction performance in terms of root-mean-square-error (RMSE), coefficient of determination ($R^2$) and classification accuracy. The right-most panel shows the number of samples in the low and high uncertainty groups and the total number of samples.

One of the biggest advantages of Bayesian learning is utilizing uncertainty to distinguish reliable and unreliable results predictions.
For example, [Ryu et al., 2019][Ryu et al., 2019] demonstrated that mis-labeled noisy data can be detected by high aleatoric (data-driven) uncertainty and elucidated the source of those errors.
Inspired from such applications, we investigated the reliability of our solubility prediction model.

As shown in Figure <ref>, we categorize the prediction results into high and low uncertainty groups and report the predictive performances in terms of RMSE, $R^2$ and classification accuracy.

The x-axis display the uncertainty thresholds, with which we divide the uncertainty groups
As expected, the lower the uncertainty groups, the better their predictive performances in terms of lower RMSE, higher $R^2$ and higher classification accuracy.
One exceptional was higher classification accuracy for higher uncertainty group in the validation set, which possibly due to extremely small number of samples used for validation, as is shown in the right-most panel of Figure <ref>(b).

§.§ Visualization and interpretation of prediction results with attention weights

Understanding how black-box machine learning models perform a given task is crucial point for reliability.
Therefore, we interpret solubility prediction results on FDA-approved drug compounds by investigating the changes in predictive solubility according to structural modification and visualizing of the attention weights computed by the PMA readout.

FDA approved drugs – imatinib, osimertinib, and afatinib – and their predictive solubility. We show the neutral form of the drugs (middle), their salt form (right) and the detached form of the solubility increasing group (left).

In Figure <ref>, we show three FDA-approved drugs (middle), i.e. imatinib, osimertinib and afatinib, their salt-form (right), and the compounds before structural modification (left).
Predictive solubility of each compound is noted below.
We highlight the attached substructures, which are commonly used to enhance solubility and bio-availability, to the left-most compounds in red circles (middle) and anionic salts in blue circles (right).
Consistent with medicinal chemistry knowledge, such structural modifications increase predictive solubility.
For example, N-methyl piperazine group is well-known for increasing bioavailibity and have been used to optimize lead compounds;
our model predicted that it significantly increases solubility through structural modification of the imatinib.
Similar results were observed in modifications of osimertinib and afatinib with hydrophilic motifs highlighted in red circles, respectively.
Not so surprisingly, input graphs represented with salt forms also showed increased predictive solubility compared to the neutral forms.
Albeit we cannot ensure that experimental values of those compounds, we propose that our model can be used for one part of lead optimization steps – improve solubility by structure modification.

The attention weights of the atoms of three FDA approved drugs – imatinib, osimertinib, and afatinib. We show the neutral form of the drugs (middle), their salt form (right) and the detached form of the solubility increasing group (left).

To interpret which atoms are considered by our prediction model to be important, in Figure <ref>,
we visualize the attention weights $\{\bar{\alpha}_v\}_{v \in V}$ computed by the PMA readout.
The attention weight of the $v$-th node is computed as follows:

\begin{equation}
\bar{\alpha}_v = \frac{N_{at}}{T \cdot H} \sum_{t=1}^{T} \sum_{h=1}^{H} \alpha_{vht},
\end{equation}

where $\alpha_{vht}$ is the attention weight for the $v$-th node from the $h$-th attention head at the $t$-th MC-sampling step, $T$ is the number of MC-samplings for MC-DO, $H$ is the number heads for multi-head attention, and $N_{at}$ is the number of atoms for a given molecule.

Since the softmax activation normalizes attention weights to one, we scale attention weights by multiplying the number of atoms to prevent attention weights to be small for large molecules.

For all three examples, we observe higher attention weights for hetero-aromatic rings than benzene rings, which is consistent to the fact that hetero-aromatic rings are more polar and hydrophilic than benzene rings.
Also, consistent to common medicinal chemistry knowledge and solubility prediction results, our model gives high attention weights to newly attached motifs, which are highlighted with red circles in the middle panel of Figure <ref>.
WE observe similar results for anionic salts, which are highlighted in blue circles in the right-most panel of Figure <ref>.

§ CONCLUSION

Here, we present a solubility prediction model based on graph neural networks, self-attention and Bayesian inference.
While previous studies focused on advancing node update parts, we find that using self-attention in the readout part is more effective in improving prediction accuracy.
This may be because, our attention weights are interpreted as importance of atoms for a given prediction task.
The experimental results demonstrate that such interpretation is valid, solubility-increasing substructures showed high attention weights.
Also, Bayesian inference via approximate MC-dropout enables us to quantify predictive uncertainty and differentiate between more and less accurate results.
Overall, our prediction results help us choose virtually screened candidates for experimental assays by analyzing uncertainty and attention weights.

We would like to assert that our prediction framework can be extended to any ligand-based (single-instance) prediction task.
Recently, Therapeutic Data Commons, the consortium to benchmark molecular machine learning,[Huang et al., 2021] has provided a variety of well-curated datasets including prediction tasks, with both single-instance and multi-instances, and generative tasks.
Among the single-instance tasks, prediction of ADME and toxicity are as important as solubility prediction to reduce failed trials in drug developments.
Since provided datasets are also data-deficient, prediction models may be vulnerable to over-fitting issue.
For those tasks, we highlight that assessment of reliability and interpretability is also essential for robust decision making and communication with experimentalists.[Schwaller et al., 2019, Soleimany et al., 2021]

In addition, we believe that our solubility prediction model can be extended to other machine learning applications, such as active learning and generative reinforcement learning.
Predictive uncertainty indicates which data points require additional labeling in supervised learning and exploration in (generative-) reinforcement learning.[Gal et al., 2017]
Such application have gained increasing attention to maximize performance with limited computational or experimental costs. [Graff et al., 2021, Soleimany et al., 2021, Gentile et al., 2020, Yang et al., 2021]

Studying how increases in training/explored data points change predictive uncertainty and attention weights enables us to draw better understanding of molecular machine learning with deep neural networks.
We leave aforementioned discussion as possible future research direction.

§ ACKNOWLEDGEMENTS

We would like to appreciate Juwon Hong and Chaok Seok for proofread of manuscripts.
This work was supported by Galux Inc. (No. Galux-20210001).

§ AUTHOR CONTRIBUTIONS

Seongok Ryu conceived the idea and performed implementation and experiments.
Seongok Ryu and Sumin Lee analyzed the results and wrote the manuscript together.

§ CONFLICTS OF INTEREST

The authors declare no competing financial interests.

[Landrum, 2013]
Landrum, G. RDKit: A software suite for cheminformatics, computational
chemistry, and predictive modeling. 2013
[Sorkun et al., 2019]
Sorkun, M. C.; Khetan, A.; Er, S. AqSolDB, a curated reference set of aqueous
solubility and 2D descriptors for a diverse set of compounds.
Scientific data 2019, 6, 1–8
[Klamt et al., 2002]
Klamt, A.; Eckert, F.; Hornig, M.; Beck, M. E.; Bürger, T. Prediction of
aqueous solubility of drugs and pesticides with COSMO-RS. Journal of
computational chemistry 2002, 23, 275–281
[Bjelobrk et al., 2021]
Bjelobrk, Z.; Mendels, D.; Karmakar, T.; Parrinello, M.; Mazzotti, M.
Solubility prediction of organic molecules with molecular dynamics
simulations. Crystal Growth & Design 2021, 21,
[Sorkun et al., 2021]
Sorkun, M. C.; Koelman, J. V. A.; Er, S. Pushing the limits of solubility
prediction via quality-oriented data selection. Iscience
2021, 24, 101961
[Lovrić et al., 2021]
Lovrić, M.; Pavlović, K.; Žuvela, P.; Spataru, A.;
Lučić, B.; Kern, R.; Wong, M. W. Machine learning in prediction of
intrinsic aqueous solubility of drug-like compounds: Generalization,
complexity, or predictive ability? Journal of Chemometrics
2021, 35, e3349
[Llinàs et al., 2008]
Llinàs, A.; Glen, R. C.; Goodman, J. M. Solubility challenge: can you
predict solubilities of 32 molecules using a database of 100 reliable
measurements? Journal of chemical information and modeling
2008, 48, 1289–1303
[Hopfinger et al., 2009]
Hopfinger, A. J.; Esposito, E. X.; Llinas, A.; Glen, R. C.; Goodman, J. M.
Findings of the challenge to predict aqueous solubility. Journal of
chemical information and modeling 2009, 49, 1–5
[Palmer and Mitchell, 2014]
Palmer, D. S.; Mitchell, J. B. Is experimental data quality the limiting factor
in predicting the aqueous solubility of druglike molecules? Molecular
Pharmaceutics 2014, 11, 2962–2972
[Llinas et al., 2020]
Llinas, A.; Oprisiu, I.; Avdeef, A. Findings of the second challenge to predict
aqueous solubility. Journal of chemical information and modeling
2020, 60, 4791–4803
[Ryu et al., 2019]
Ryu, S.; Kwon, Y.; Kim, W. Y. A Bayesian graph convolutional network for
reliable prediction of molecular properties with uncertainty quantification.
Chemical science 2019, 10, 8438–8446
[Schwaller et al., 2019]
Schwaller, P.; Laino, T.; Gaudin, T.; Bolgar, P.; Hunter, C. A.; Bekas, C.;
Lee, A. A. Molecular transformer: a model for uncertainty-calibrated chemical
reaction prediction. ACS central science 2019, 5,
[Soleimany et al., 2021]
Soleimany, A. P.; Amini, A.; Goldman, S.; Rus, D.; Bhatia, S. N.; Coley, C. W.
Evidential deep learning for guided molecular property prediction and
discovery. ACS central science 2021, 7,
[Graff et al., 2021]
Graff, D. E.; Shakhnovich, E. I.; Coley, C. W. Accelerating high-throughput
virtual screening through molecular pool-based active learning.
Chemical science 2021, 12, 7866–7881
[Harren et al., 2022]
Harren, T.; Matter, H.; Hessler, G.; Rarey, M.; Grebner, C. Interpretation of
Structure–Activity Relationships in Real-World Drug Design Data Sets Using
Explainable Artificial Intelligence. Journal of Chemical Information
and Modeling 2022,
[Vaswani et al., 2017]
Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.;
Kaiser, Ł.; Polosukhin, I. Attention is all you need. Advances in neural
information processing systems. 2017; pp 5998–6008
[Huang et al., 2021]
Huang, K.; Fu, T.; Gao, W.; Zhao, Y.; Roohani, Y.; Leskovec, J.; Coley, C. W.;
Xiao, C.; Sun, J.; Zitnik, M. Therapeutics data Commons: machine learning
datasets and tasks for therapeutics. arXiv preprint arXiv:2102.09548
[Wu et al., 2018]
Wu, Z.; Ramsundar, B.; Feinberg, E. N.; Gomes, J.; Geniesse, C.; Pappu, A. S.;
Leswing, K.; Pande, V. MoleculeNet: a benchmark for molecular machine
learning. Chemical science 2018, 9, 513–530
[Hwang et al., 2020]
Hwang, D.; Yang, S.; Kwon, Y.; Lee, K. H.; Lee, G.; Jo, H.; Yoon, S.; Ryu, S.
Comprehensive Study on Molecular Supervised Learning with Graph Neural
Networks. Journal of Chemical Information and Modeling 2020,
60, 5936–5945
[Duvenaud et al., 2015]
Duvenaud, D. K.; Maclaurin, D.; Iparraguirre, J.; Bombarell, R.; Hirzel, T.;
Aspuru-Guzik, A.; Adams, R. P. Convolutional networks on graphs for learning
molecular fingerprints. Advances in neural information processing
systems 2015, 28
[Kipf and Welling, 2016]
Kipf, T. N.; Welling, M. Semi-supervised classification with graph
convolutional networks. arXiv preprint arXiv:1609.02907 2016,

[Ba et al., 2016]
Ba, J. L.; Kiros, J. R.; Hinton, G. E. Layer normalization. arXiv
preprint arXiv:1607.06450 2016,
[Xu et al., 2018]
Xu, K.; Hu, W.; Leskovec, J.; Jegelka, S. How powerful are graph neural
networks? arXiv preprint arXiv:1810.00826 2018,
[Veličković et al., 2017]
Veličković, P.; Cucurull, G.; Casanova, A.; Romero, A.; Lio, P.;
Bengio, Y. Graph attention networks. arXiv preprint arXiv:1710.10903
[Fey and Lenssen, 2019]
Fey, M.; Lenssen, J. E. Fast graph representation learning with PyTorch
Geometric. arXiv preprint arXiv:1903.02428 2019,
[Dosovitskiy et al., 2020]
Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.;
Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.,
et al.  An image is worth 16x16 words: Transformers for image
recognition at scale. arXiv preprint arXiv:2010.11929 2020,

[Lee et al., 2019]
Lee, J.; Lee, Y.; Kim, J.; Kosiorek, A.; Choi, S.; Teh, Y. W. Set transformer:
A framework for attention-based permutation-invariant neural networks.
International Conference on Machine Learning. 2019; pp 3744–3753
[Gal and Ghahramani, 2016]
Gal, Y.; Ghahramani, Z. Dropout as a bayesian approximation: Representing model
uncertainty in deep learning. international conference on machine learning.
2016; pp 1050–1059
[Maddox et al., 2019]
Maddox, W. J.; Izmailov, P.; Garipov, T.; Vetrov, D. P.; Wilson, A. G. A simple
baseline for bayesian uncertainty in deep learning. Advances in Neural
Information Processing Systems 2019, 32
[Gal et al., 2017]
Gal, Y.; Islam, R.; Ghahramani, Z. Deep bayesian active learning with image
data. International Conference on Machine Learning. 2017; pp 1183–1192
[Gentile et al., 2020]
Gentile, F.; Agrawal, V.; Hsing, M.; Ton, A.-T.; Ban, F.; Norinder, U.;
Gleave, M. E.; Cherkasov, A. Deep docking: a deep learning platform for
augmentation of structure based drug discovery. ACS central science
2020, 6, 939–949
[Yang et al., 2021]
Yang, S.; Hwang, D.; Lee, S.; Ryu, S.; Hwang, S. J. Hit and Lead Discovery with
Explorative RL and Fragment-based Molecule Generation. Advances in
Neural Information Processing Systems 2021, 34