# Accurate Baryon Acoustic Oscillations reconstruction via semi-discrete
optimal transport

Sebastian von Hausegger1,2,3<EMAIL_ADDRESS>Bruno
Levy2 Roya Mohayaee1,3 1Rudolf Peierls Centre for Theoretical Physics,
University of Oxford, Parks Road, Oxford OX1 3PU, United Kingdom  2Université
de Lorraine, CNRS, Inria, LORIA, F-54000 Nancy, France  3Institut
d’Astrophysique de Paris, CNRS, Sorbonne Université, 98bis Bld Arago, 75014
Paris, France

###### Abstract

Optimal transport theory has recently reemerged as a vastly resourceful field
of mathematics with elegant applications across physics and computer science.
Harnessing methods from geometry processing, we report on the efficient
implementation for a specific problem in cosmology — the reconstruction of the
linear density field from low redshifts, in particular the recovery of the
Baryonic Acoustic Oscillation (BAO) scale. We demonstrate our algorithm’s
accuracy by retrieving the BAO scale in noise-less cosmological simulations
that are dedicated to cancel cosmic variance; we find uncertainties to be
reduced by a factor of 4.3 compared with performing no reconstruction, and a
factor of 3.1 compared with standard reconstruction.

††preprint: APS/123-QED

##  I Introduction

Linear perturbations in the primordial Universe propagate as sound waves
through the photon-baryon plasma until light and matter decouple. In a balance
of radiation pressure and gravity, the baryon-to-photon and the matter-to-
radiation ratios define a correlation length that is imprinted as a distinct
feature onto the density fields, detectable in both the Cosmic Microwave
Background (CMB) at early and the Large Scale Structure (LSS) at late times
Sakharov (1966); Sunyaev and Zeldovich (1970); Peebles and Yu (1970). However,
non-linear clustering of galaxies at low redshifts distort this imprint and
therewith impedes unbiased detection Eisenstein  _et al._ (2007a), calling for
methods to undo these non-linear effects, to ‘reconstruct’ the linear density
field, and thereby to enable accurate and precise measurement Eisenstein  _et
al._ (2007b). This so-called Baryon Acoustic Oscillation (BAO) scale serves as
a unique tool to map the expansion history of the Universe, and hence has
taken up a central role in cosmological analyses; most notably, its detection
Cole  _et al._ (2005); Eisenstein  _et al._ (2005) has provided constraints on
standard quantities such as the Hubble constant and dark energy density Zhao
_et al._ (2019), with promising capabilities to, with future observations,
even constrain more nuanced theories of, e.g., light degrees-of-freedom
Baumann  _et al._ (2017). Therefore, and especially in light of present and
upcoming large galaxy surveys 111To date, these include, but are not limited
to, e.g., DESI DESI Collaboration  _et al._ (2016), LSST LSST Dark Energy
Science Collaboration (2012), Euclid Euclid Collaboration  _et al._ (2011).,
the development of fast, accurate, and scalable reconstruction methods is
highly relevant to optimizing the yield of LSS studies. As a prime example of
the wide and successful applicability of modern Optimal transport theory, this
Letter promotes a recent implementation of one such method.

Optimal transport theory describes mappings between probability measures that
minimise a total cost function while satisfying a volume conservation
constraint Villani (2009, 2003). Recent advances in both the mathematical and
algorithmic aspects of the theory paved the way for breakthroughs in various
fields, such as artificial intelligence, economics, meteorology, biology, and
physics — due to the universality of processes minimizing an action. In the
context of this work, we exploit natural connections between optimal transport
and physics Frisch  _et al._ (2002); Levy  _et al._ (2021): the to-be-
minimised quantity and volume conservation are an action integral and the
continuity equation respectively. The Lagrange multiplier associated with the
constraint turns out to be the gravitational potential. Drawing from and
building on all these advances, we recently developed a deterministic
algorithm Levy  _et al._ (2021) that efficiently reconstructs the sought-for
linear density field. After having fully characterized its behavior in Ref.
Levy  _et al._ (2021), we here focus on the accuracy with which it retrieves
the BAO scale by applying it to noise-free cosmic variance cancelling
cosmological simulations generated with the FastPM Feng  _et al._ (2016)
algorithm.

This Letter recapitulates the mathematical basis of our method, illustrates
its particular application to cosmological density field reconstruction, and
finally showcases its excellent results in comparison with standard
reconstruction by the use of noise-free cosmic variance cancelling
cosmological simulations.

##  II Monge-Ampère-Kantorovich Reconstruction

Figure 1: Semi-discrete Monge-Ampére-Kantorovich Reconstruction. Illustration
of the reconstruction’s flow beginning with an input of point masses at
$\tau_{F}$ (here $z=0$) that are mapped to their corresponding regions at
$\tau_{I}$ (here $z\rightarrow\infty$) from where they draw their mass through
gravitational collapse. The regions are the polyhedral cells defined by the
Laguerre diagram as defined in the main body. In a subsequent step the cells
evolve until a time $\tau_{F}>\tau>\tau_{I}$ (here $z=9$) and according to a
LPT description – in the present paper the Zel’dovich approximation. In the
final step a mass density is expressed on a Eulerian mesh. The figure shows a
tiny subset of a reconstruction from AbacusCosmos, that comprises 9.7 million
Laguerre cells.

Monge-Ampère-Kantorovich (MAK) reconstruction uniquely determines the
Lagrangian trajectories of a given Eulerian distribution of particle positions
by solving an optimal transport (OT) problem Brenier  _et al._ (2003); Levy
_et al._ (2021). To this effect, consider self-gravitating matter in an
Einstein-de Sitter universe, where particle trajectories are the solution to
extremizing the action $I$ subject to the Poisson equation, mass conservation
and appropriate boundary conditions,

$I=\frac{1}{2}\int_{\tau_{I}}^{\tau_{F}}\int_{V}\,\left(\rho|{\bf
v}|^{2}+\frac{3}{2}|\nabla_{x}\phi|^{2}\right)\tau^{3/2}\;d^{3}x\,d\tau.$ (1)

Here, ${\bf x}={\bf x}({\bf q},\tau)$ are the trajectories in Lagrangian
coordinates of particles initially at ${\bf q}$, ${\bf v}$ is the Eulerian,
co-moving velocity field, $\rho$ is the Eulerian density field, and $\phi$ is
the gravitational potential. Serving as a time variable, $\tau$ denotes the
amplitude of the growing linear mode, normalized such that initially
$\tau_{I}=0$ and finally $\tau_{F}=1$. Correspondingly, the final density
$\rho(\tau_{F})$ is considered to have evolved from an initially uniform
state, $\rho(\tau_{I})=1$. At final time $\tau_{F}$, we suppose that the
density field is clustered into a set of points $({\bf x}_{i})_{i=1}^{N}$ with
masses $(m_{i})_{i=1}^{N}$. Ultimately, we aim to recover the (linear) density
field at high redshift, or $\tau=\tau_{I}+\epsilon$, $\epsilon\ll 1$, from a
low-redshift, or present-time distribution of clustered matter. In practice,
this will require finding the optimal map that assigns initial positions ${\bf
q}$ to the input particle positions ${\bf x}({\bf q},\tau_{F})$.

Linearizing around the stationary points of the action, Eq. (1), one finds
only the kinetic term to remain, resulting in uniform rectilinear motion
Landau and Lifshitz (1975). Interestingly, this case is proven Benamou and
Brenier (2000) to be equivalent to the $L_{2}$ Monge-Kantorovich optimal
transport problem Monge (1784),

$\inf_{{\bf x}_{F}}\int_{V}\,\rho({\bf q})|{\bf x}_{F}({\bf q})-{\bf
q}|^{2}\;d^{3}q,$ (2)

in which the integrated squared distance is minimized, subject to mass
conservation, $\rho_{F}({\bf x})/\rho_{I}({\bf q})=\textrm{det}\left[d^{3}{\bf
x}/d^{3}{\bf q}\right]$. In the dual formulation of the problem, that
exchanges the unknown application ${\bf x}_{F}$ with the mass conservation
constraint, one solves the _Monge-Ampère equation_ for a scalar function
$\Phi$, associated with the constraint, and called the _Kantorovich potential_
:

$\frac{\rho_{F}({\bf x})}{\rho_{I}({\bf
q})}=\textrm{det}\left[\frac{\partial^{2}\Phi}{\partial{\bf q}_{i}\partial{\bf
q}_{j}}\right],$ (3)

The optimal assignment map ${\bf q}\mapsto{\bf x}_{F}({\bf q})$ turns out to
be uniquely determined as the gradient of $\Phi$ Brenier (1991). From the
second-order optimality condition, one can prove that $\Phi$ is a _convex_
function, which is interesting because it implies the absence of shell-
crossing during structure formation: intuitively, one may think of the
velocity field as the normals to the graph of $\Phi$, and the normal vectors
to a convex body do not intersect. More formally, the proof is obtained by
contradiction: the existence of a collision would imply an inequality that
would in turn contradict the convexity of $\Phi$ (details in Levy  _et al._
(2021)).

It is interesting to express the optimization problem as a function of
physical quantities instead of $\Phi$, using the relation between $\Phi$, the
initial gravitational potential $\phi_{I}$ and the final gravitational
potential $\phi_{F}$:

$\displaystyle\Phi({\bf q})$ $\displaystyle=$ $\displaystyle 1/2\,{\bf
q}^{2}-\phi_{I}({\bf q})$ (4) $\displaystyle\phi_{I}({\bf q})$
$\displaystyle=$ $\displaystyle\inf_{{\bf x}}[1/2|{\bf x}-{\bf
q}|^{2}-\phi_{F}({\bf x})].$ (5)

The initial and final potentials $\phi_{I}$ and $\phi_{F}$ are mutually
related through the _Legendre-Fenchel transform_ (Eq. 5) (the same one that
converts between Lagrangian and Hamiltonian mechanics). They turn out to be
the solution of the following optimization problem:

$\max_{\phi_{I},\phi_{F}}K=\int_{V}\phi_{I}({\bf q})\rho_{I}({\bf q})d{\bf
q}+\int_{V}\phi_{F}({\bf x})\rho_{F}({\bf x})d{\bf x},$ (6)

where the functional $K$ is called the _Kantorovich dual_.

While numerical methods for approximately solving the Monge-Ampère (equation
3) have been devised Shi  _et al._ (2018); Birkin  _et al._ (2019); Liu  _et
al._ (2021), elegant algorithmic developments in computer science allow for
efficiently constructing exact solutions (see Peyré and Cuturi (2019) for a
review).

Previous reconstruction methods that exactly solve the Monge-Ampère equation
Brenier  _et al._ (2003); Frisch  _et al._ (2002) considered a _discrete_
version of the Monge-Kantorovich problem (Eq. 2), i.e. to find the permutation
$j(i)$ between a finite set of $N$ homogeneously distributed particle
positions $({\bf q}_{j})_{j=1}^{N}$ at $\tau_{I}$ and their corresponding
positions $({\bf x}_{i})_{i=1}^{N}$ at $\tau_{F}$ that minimizes
$\sum_{i}|{\bf x}_{i(j)}-{\bf q}_{j}|^{2}$. Efficient combinatorial methods
Bertsekas (1992) avoid exploring the full set of $N!$ possible permutations;
however, they still scale as $\sim\\!N^{2}\log(N)$, rendering such algorithms
increasingly unfeasible for larger data sets. Moreover, using these
combinatorial methods, it is difficult to allocate a different mass $m_{i}$ to
each point.

However, exploiting the variational nature of the Monge-Kantorovich problem,
the present _semi-discrete_ approach replaces this exhaustive combinatorial
search by the optimization of a well-behaved objective function,
simultaneously making use of the geometric structure of the setting: Instead
of representing the initial condition as a discrete set of points $({\bf
q}_{j})_{j=1}^{N}$, we consider it as a continuum, while the density at
$\tau_{F}$ is concentrated on a set of points ${\bf x}_{i}$ with the
associated masses $m_{i}$. Replacing $\phi_{I}$ by its expression (Legendre
transform of $\phi_{F}$, Eq. 5), the Kantorovich dual $K$ in Equation 6
becomes:

$K(\psi)=\sum_{i}\int_{V^{\psi}_{i}}\left[\frac{1}{2}|{\bf x}_{i}-{\bf
q}|^{2}-\psi_{i}\right]\;d^{3}q+\sum_{i}m_{i}\psi_{i},$ (7)

that depends on the vector $\psi=(\phi_{F}({\bf x}_{i}))_{1}^{N}$, and where
the subsets $V_{i}^{\psi}$, also stemming from the Legendre-Fenchel transform
of $\phi_{F}$, are defined by:

$V^{\psi}_{i}=\left\\{{\bf q}\ \left|\ \frac{1}{2}|{\bf x}_{i}-{\bf
q}|^{2}-\psi_{i}<\frac{1}{2}|{\bf x}_{j}-{\bf q}|^{2}-\psi_{j},\ \forall j\neq
i\right.\right\\}.$ (8)

The regions $V_{i}^{\psi}$ are called _Laguerre cells_ , and they form a
_Laguerre diagram_ (more on this in Aurenhammer  _et al._ (1992); Mérigot
(2011); Lévy (2015); Mérigot and Thibert (2020)). Each Laguerre cell
$V_{i}^{\psi}$ corresponds to the region of space mapped to ${\bf x}_{i}$
through the reconstructed motion.

The vector $\psi$ that determines the Laguerre diagram can be obtained by
maximizing the Kantorovich dual $K$ (Eq. 7) subject to the constraint that no
cell is empty ($V^{\psi}_{i}\neq\emptyset,\;\forall i$). Finally, one
retrieves the gravitational potential at $\tau_{I}$ through the Legendre
transform of $\psi$:

$\displaystyle\phi(\tau_{I},{\bf q})$ $\displaystyle=\inf_{{\bf
x}_{F}}\left[\frac{1}{2}\left|{\bf q}-{\bf x}_{i}\right|^{2}-\psi_{i}\right]$
(9) $\displaystyle=\frac{1}{2}\left|{\bf q}-{\bf x}_{i({\bf
q})}\right|^{2}-\psi_{i({\bf q})},$

where $i({\bf q})$ is the index of the Laguerre cell $V_{i}$ that contains
${\bf q}$. In other words, maximizing $K(\psi)$ is equivalent to solving for
the Kantorovich potential $\Phi$ in the Monge-Ampère equation.

The Kantorovich dual $K$ has two properties that are important from the point
of view of numerical optimization. Firstly, $K$ corresponds to the lower
envelope of a family of affine functions, hence it is _concave_ Aurenhammer
_et al._ (1992); Mérigot (2011); Lévy (2015); Mérigot and Thibert (2020),
which ensures the existence and uniqueness of $\psi$. Secondly, by studying
the Taylor expansion of $K$ in two configurations differing from a
combinatorial change in the Laguerre diagram, one can observe that the terms
up to $2^{nd}$ order match Liu  _et al._ (2009), hence $K$ is $C^{2}$_-smooth_
, which allows for an efficient and convergent optimization Kitagawa  _et al._
(2016). It is the combination of the two analytic properties of $K$ (convexity
and smoothness) that allow replacing the expensive combinatorial computation
of previous methods with a Newton method that exploits the first- and second-
order derivatives of $K$ to efficiently minimize it.

Finally, each Laguerre cell $V_{i}^{\psi}$ at $\tau_{I}$ is mapped to its
corresponding point ${\bf x}_{i}$ at $\tau_{F}$. At an intermediate time
$\tau$ the trajectories of the mass elements can be deduced via a Lagrangian
perturbation theory (LPT) description, in our case the Zel’dovich
approximation Zel’dovich (1970). This, in turn, can be converted as a mass
density on a Eulerian grid, as depicted in Fig. 1 (see also Levy  _et al._
(2021)).

Figure 2: Fractional BAO signal in simulations and reconstructions. Top panel:
Ratio of power spectra from simulations with and without BAO signal, averaged
over ten FastPM simulations and their reconstructions, $\langle P_{\rm
wig}(k)-P_{\rm nw}(k)\rangle/\langle P_{\rm wig}(k)\rangle$. Bottom panel:
Deviations from the initial condition’s fractional BAO signal, $\Delta^{\rm
IC}(k)=\langle P_{\rm wig}(k)\rangle/\langle P_{\rm nw}(k)\rangle-\langle
P_{\rm wig}^{\rm IC}(k)\rangle/\langle P_{\rm nw}^{\rm IC}(k)\rangle$. The
shaded band indicates the standard deviation $\sigma^{\rm IC}(k)$ as estimated
from the simulations and defined in the text.

##  III BAO reconstruction

Baryonic Acoustic Oscillations imprint a signature in the matter power
spectrum by periodically modulating the large-scale power at a frequency
corresponding to the sound horizon at decoupling Sakharov (1966). Non-linear
gravitational evolution blurs and shifts this feature Eisenstein  _et al._
(2007a), complicating reliable detection and interpretation, wherefore
reconstruction algorithms were devised that correct for such influences. In
addition to physical effects, finite survey volumes, both in simulations and
in practice, induce sample variance (so-called cosmic variance) into the
measured quantities due to a reduced number of modes on these very scales.
Since these modulations can be seen as only small perturbations to the bulk of
the power spectrum, the variance is dominated by the un-modulated power
spectrum, suites of dedicated $N$-body simulations can effectively cancel
cosmic variance for analyses that only target the BAO signal Prada  _et al._
(2016); Schmittfull  _et al._ (2017); Wang  _et al._ (2017), e.g., to test
reconstruction methods. In this view, we apply our reconstruction algorithm
Levy  _et al._ (2021) to the FastPM Feng  _et al._ (2016) simulations of Ref.
Ding  _et al._ (2018).

This suite of FastPM simulations is comprised of pairs of simulations
initiated with the same random phases at redshift $z=9$, yet with power
spectra that differ by the presence of the BAO feature, in the following
referred to as ‘wiggle’ and ‘no-wiggle’ power spectra, or $P_{\rm wig}(k)$ and
$P_{\rm nw}(k)$. Each of the simulations traces $2048^{3}$ particles in a cube
of $1380\,h^{-1}$Mpc side length from $z=9$ to $z=0$, and saved at redshift
$z=0$ as a particle sample. In order to keep computing time low, we sub-sample
only $\sim\\!1\%$ of all particles in each of the ten simulation pairs we
consider, resulting in about $85\times 10^{6}$ particles per simulation. All
simulations follow a $\Lambda$CDM cosmology with parameters from Ref. Ade  _et
al._ (2016), which gives an expected BAO scale of $r_{\rm BAO}^{\rm
th}=147.5\,\rm{Mpc}$.

Beginning with these particles’ positions, ${\bf x}$, at $z=0$, we compute the
density field Levy  _et al._ (2021) at $z=9$ as described in the previous
section and via the choice of an appropriate linear growth factor 222The
present analysis is not very sensitive to the choice of growth factor Lukic
_et al._ (2007) as it is only concerned with the ratio of ‘wiggle’ and ‘no-
wiggle’ simulations. Nevertheless, when dealing with real data, a slight
mismatch between the true growth factor and a perhaps wrongly approximated one
will in effect show up as a bias factor, $b$, to the reconstructed power, and
can therefore easily be taken into account in a subsequent step, e.g. in a
parameter fit one can marginalize over this parameter.. Figure 2 demonstrates
the effectiveness of our reconstruction, where the relative differences,
$d(k)=P_{\rm wig}(k)/\langle P_{\rm nw}(k)\rangle-1$, of initial ($z=9$) and
final ($z=0$) power spectra are compared with those of our reconstructed
density fields, averaged over all ten simulations. Our reconstruction’s
excellent performance is further highlighted by contrasting it against the
result obtained from so-called standard reconstruction Eisenstein  _et al._
(2007b)333We computed the standard reconstructed density field using the
routine FFTRecon of the python suite nbodykit Hand  _et al._ (2018)., visible
through the agreement of initial condition and reconstruction out to wave
numbers $k\gg 0.15\,h^{-1}\textrm{Mpc}$, where discrepancies between initial
condition and standard reconstruction first appear.

Table 1: Bias and uncertainties of BAO scales recovered in the fractional power spectra of simulations and reconstructions with and without cancelling cosmic variance. The columns lists the mean values and standard deviations of $r_{\rm BAO}$ (left columns) $r_{\rm BAO}-r_{\rm BAO}^{\rm IC}$ (right columns) obtained from fitting $\alpha$ to the power spectra as described in the main body. |  $\qquad\langle\hat{r}_{\rm BAO}\rangle$ $[\textrm{Mpc}]$ |  $\qquad\sigma_{\hat{r}_{\rm BAO}}$ $[\textrm{Mpc}]$
---|---|---
|  vs. theory |  vs. IC |  vs. theory |  vs. IC
Initial cond. |  $+0.02$ [$+0.01\%$] |  $\pm 0.00$ [$\pm 0.00\%$] |  $0.34$ [$0.23\%$] |  $0.00$ [$0.00\%$]
MAK rec. |  $-0.11$ [$-0.08\%$] |  $-0.13$ [$-0.09\%$] |  $0.44$ [$0.30\%$] |  $0.29$ [$0.19\%$]
Standard rec. |  $-0.03$ [$-0.02\%$] |  $-0.05$ [$-0.03\%$] |  $0.69$ [$0.47\%$] |  $0.85$ [$0.58\%$]
Final cond. |  $+0.48$ [$+0.33\%$] |  $+0.46$ [$+0.31\%$] |  $1.06$ [$0.72\%$] |  $1.19$ [$0.81\%$]

In line with Ref. Schmittfull  _et al._ (2017) we $\chi^{2}$-fit templates,
$m(k)$, to the power spectrum ratios to obtain estimates of the BAO scale in
each of the ten simulations. We define $m(k)=e^{k^{2}\Sigma^{2}/2}\left[P^{\rm
IC}_{\rm wig}(k/\alpha)/P^{\rm IC}_{\rm nw}(k/\alpha)-1\right]$ as the
relative difference of power spectra of the initial, linear density field,
allowing for a shift $\alpha=\hat{r}_{\rm BAO}/r_{\rm BAO}^{\rm th}$ of the
BAO scale, and a Gaussian damping $\Sigma$ 444We define $P_{\rm wig}^{\rm
IC}(k/\alpha)/P_{\rm nw}^{\rm IC}(k/\alpha)$ as the average over all ten
simulations, shown in Fig. 2, and interpolated to accommodate values of
$\alpha\neq 1$., and up-weigh small scales via choice of the standard error
$\sigma^{2}(k)=2\,\textrm{Var}\left[d(k)\right]/N_{\rm modes}(k)$, where
$N_{\rm modes}(k)$ is the number of Fourier modes that contributes to the
computed power in each $k$-bin. We perform the fits over the full $k$-range
shown in Figure 2, and the best-fit values of $\alpha$ define the best-fit BAO
scales $\hat{r}_{\rm BAO}$ in each fractional power spectrum $d(k)$. Table 1
presents biases and uncertainties in retrieving $r^{\rm th}_{\rm BAO}$ in each
of the simulations and reconstructions. While comparison with the theory value
$r_{\rm BAO}^{\rm th}$ (left columns) confirms and restates more precisely the
results of Ref. Levy  _et al._ (2021), the right columns optimally and for the
first time showcase our algorithm’s accuracy and precision in reconstructing
the BAO scale from noiseless cosmological simulations, by subtracting from
each simulation the inherent BAO scale, $r_{\rm BAO}^{\rm IC}$, before
determining mean and spread, thereby cancelling cosmic variance.

Due to the arising of shift terms in the non-linear power spectra Eisenstein
_et al._ (2007a); Sherwin and Zaldarriaga (2012); Schmittfull  _et al._
(2017), the BAO scale at $z=0$ appears biased by $\sim\\!0.3\%$, in accordance
with previous findings Seo  _et al._ (2010); Mehta  _et al._ (2011);
Schmittfull  _et al._ (2017). This is accompanied by a $\sim\\!0.8\%$
uncertainty that reflects the blurring of the BAO peak that as well is caused
by non-linear gravitational growth. Reconstruction reduces this bias Sherwin
and Zaldarriaga (2012) as we too see in both MAK and standard reconstruction,
and further sharpens the BAO peak increasing the precision with which
$\hat{r}_{\rm BAO}$ is determined; compared with the inherent uncertainty the
simulations carry at the final condition – including cosmic variance –
standard reconstruction improves the precision by a factor of $1.5$ while MAK
reconstruction gives a factor of $2.4$ of enhancement. In an idealised
scenario, without cosmic variance, the factor $1.4$ improvement of standard
reconstruction is surpassed by MAK reconstruction by as much as $4.3$. In all
cases we find a significant reduction of the bias as well.

As elaborated in Ref. Levy  _et al._ (2021), subsample variance has
significant impact on the overall error budget. Both shot noise and subsample
variance are virtually removed by the use of the present simulation suite, and
cosmic variance is further cancelled by direct simulation-to-simulation
comparison as we display in the right columns, Table 1. We therewith optimally
test our reconstructions’ accuracy.

##  IV Conclusions

This Letter demonstrates the application of Optimal Transport theory to a
specific problem in cosmology, the reconstruction of the BAO peak in the
matter power spectrum from low-$z$ observations. BAO analyses play a crucial
role in inferring cosmological parameters, and reconstruction methods have
long aided the accuracy with which this signal is extracted. Outperforming
many of the most promising algorithms, our method scales well ($\propto N\log
N$) with increased survey size, securing bright prospects in light of upcoming
large-scale galaxy surveys.

In specific, we found that our reconstruction improves on detecting the BAO
signal in the power spectrum by a factor of 4.3 compared with attempting to
extract the BAO scale without having performed any reconstruction. Even in the
case of having applied the so-called standard reconstruction technique, our
method reduces the uncertainties by more than a factor of 3. This is highly
promising especially given that in moving forward in time, we considered no
more than the Zel’dovich approximation. We therefore highly anticipate further
improvement of reported accuracy by amending the second step in Fig. 1 with
corrections from higher-order Lagrangian perturbation theory Bernardeau  _et
al._ (2002); Schmittfull  _et al._ (2017).

The next steps for optimal incorporation of our reconstruction method into
analyses of survey data include its adaptation to account for the surveys’
selection functions, halo masses, redshift-space distortions, and
characterization and computation of the reconstruction covariance matrices.
Our algorithm’s flexibility easily accommodates such modifications without
losing its efficiency.

In summary, our method makes direct use of the variational nature of
gravitational evolution and thereby its reconstruction. It finds a quick path
to the solution by leveraging first and second order information of the
problem (i.e. it being both smooth ($C^{2}$) and convex), while existing MAK
methods need to exhaustively explore a huge $(N^{2}\log(N))$ combinatorial
space. This is made possible by a fortuitous yet elegant convergence between
the physical, mathematical and computational aspects of the problem: the
specific cosmological setting that we considered (continuous mass transported
to a point set) has nice mathematical properties (semi-discrete Monge-Ampère
equation translated into a smooth and concave optimization problem), with an
underlying geometric structure (Laguerre diagram) that can be _exactly_
computed by our algorithm.

###### Acknowledgements.

The authors thank Zhejie Ding for sharing their simulations. SvH thanks Yu
Feng for assistance with FastPM. BL and SvH acknowledge support from an Inria
internal grant (Action exploratoire AeX EXPLORAGRAM). SvH is supported at
Oxford by the Carlsberg Foundation, and wishes to thank Linacre College for
the award of a Junior Research Fellowship. RM thanks the Rudolf Peierls Centre
for hospitality.

## References

  * Sakharov (1966) A. D. Sakharov, Soviet Journal of Experimental and Theoretical Physics 22, 241 (1966).
  * Sunyaev and Zeldovich (1970) R. A. Sunyaev and Y. B. Zeldovich, Astrophysics and Space Science 7, 3 (1970).
  * Peebles and Yu (1970) P. J. E. Peebles and J. T. Yu, Astrophys. J. 162, 815 (1970).
  * Eisenstein  _et al._ (2007a) D. J. Eisenstein, H.-J. Seo, and M. J. White, Astrophys. J. 664, 660 (2007a), arXiv:astro-ph/0604361 .
  * Eisenstein  _et al._ (2007b) D. J. Eisenstein, H.-J. Seo, E. Sirko, and D. Spergel, Astrophys. J. 664, 675 (2007b), arXiv:astro-ph/0604362 [astro-ph] .
  * Cole  _et al._ (2005) S. Cole, W. J. Percival, J. A. Peacock, P. Norberg, C. M. Baugh, C. S. Frenk, I. Baldry, J. Bland-Hawthorn, T. Bridges, R. Cannon, _et al._ , MNRAS 362, 505 (2005), arXiv:astro-ph/0501174 [astro-ph] .
  * Eisenstein  _et al._ (2005) D. J. Eisenstein, I. Zehavi, D. W. Hogg, R. Scoccimarro, M. R. Blanton, R. C. Nichol, R. Scranton, H.-J. Seo, M. Tegmark, Z. Zheng, _et al._ , Astrophys. J. 633, 560 (2005), arXiv:astro-ph/0501171 [astro-ph] .
  * Zhao  _et al._ (2019) G.-B. Zhao, Y. Wang, S. Saito, H. Gil-Marín, W. J. Percival, D. Wang, C.-H. Chuang, R. Ruggeri, E.-M. Mueller, F. Zhu, _et al._ , MNRAS 482, 3497 (2019), arXiv:1801.03043 [astro-ph.CO] .
  * Baumann  _et al._ (2017) D. Baumann, D. Green, and M. Zaldarriaga, JCAP 2017, 007 (2017), arXiv:1703.00894 [astro-ph.CO] .
  * Note (1) To date, these include, but are not limited to, e.g., DESI DESI Collaboration  _et al._ (2016), LSST LSST Dark Energy Science Collaboration (2012), Euclid Euclid Collaboration  _et al._ (2011).
  * Villani (2009) C. Villani,  _Optimal transport : old and new_, Grundlehren der mathematischen Wissenschaften (Springer, Berlin, 2009).
  * Villani (2003) C. Villani,  _Topics in optimal transportation_, Graduate studies in mathematics (American Mathematical Society, Providence (R.I.), 2003).
  * Frisch  _et al._ (2002) U. Frisch, S. Matarrese, R. Mohayaee, and A. Sobolevskii, Nature 417 (2002).
  * Levy  _et al._ (2021) B. Levy, R. Mohayaee, and S. von Hausegger, MNRAS 506, 1165 (2021), arXiv:2012.09074 [astro-ph.CO] .
  * Feng  _et al._ (2016) Y. Feng, M.-Y. Chu, U. Seljak, and P. McDonald, MNRAS 463, 2273 (2016), arXiv:1603.00476 [astro-ph.CO] .
  * Brenier  _et al._ (2003) Y. Brenier, U. Frisch, M. Henon, G. Loeper, S. Matarrese, R. Mohayaee, and A. Sobolevskii, MNRAS 346, 501 (2003), arXiv:astro-ph/0304214 .
  * Landau and Lifshitz (1975) L. Landau and E. Lifshitz,  _Course of Theoretical Physics, Volume 1: mechanics_ (Butterworth-Heinemann, 1975).
  * Benamou and Brenier (2000) J. Benamou and Y. Brenier, Numerische Mathematik 84, 375 (2000).
  * Monge (1784) G. Monge, Histoire de l’Académie Royale des Sciences (1781) , 666 (1784).
  * Brenier (1991) Y. Brenier, Communications on Pure and Applied Mathematics 44, 375 (1991).
  * Shi  _et al._ (2018) Y. Shi, M. Cautun, and B. Li, Phys. Rev. D 97, 023505 (2018), arXiv:1709.06350 [astro-ph.CO] .
  * Birkin  _et al._ (2019) J. Birkin, B. Li, M. Cautun, and Y. Shi, MNRAS 483, 5267 (2019), arXiv:1809.08135 [astro-ph.CO] .
  * Liu  _et al._ (2021) Y. Liu, Y. Yu, and B. Li, Astrophys. J. Suppl. 254, 4 (2021), arXiv:2012.11251 [astro-ph.CO] .
  * Peyré and Cuturi (2019) G. Peyré and M. Cuturi, Found. Trends Mach. Learn. 11, 355 (2019).
  * Bertsekas (1992) D. P. Bertsekas, Computational Optimization and Applications 1, 7 (1992).
  * Aurenhammer  _et al._ (1992) F. Aurenhammer, F. Hoffmann, and B. Aronov, in  _Symposium on Computational Geometry_ (1992) pp. 350–357.
  * Mérigot (2011) Q. Mérigot, Comput. Graph. Forum 30, 1583 (2011).
  * Lévy (2015) B. Lévy, ESAIM M2AN (Mathematical Modeling and Analysis) (2015).
  * Mérigot and Thibert (2020) Q. Mérigot and B. Thibert, “Optimal transport: discretization and algorithms,” (2020), arXiv:math/2003.00855 [math.OC] .
  * Liu  _et al._ (2009) Y. Liu, W. Wang, B. Lévy, F. Sun, D. Yan, L. Lu, and C. Yang, ACM Trans. Graph. 28, 101:1 (2009).
  * Kitagawa  _et al._ (2016) J. Kitagawa, Q. Mérigot, and B. Thibert, CoRR (2016), arXiv:1603.05579 .
  * Zel’dovich (1970) Y.-B. Zel’dovich, Astronomy and astrophysics 5, 84 (1970).
  * Prada  _et al._ (2016) F. Prada, C. G. Scóccola, C.-H. Chuang, G. Yepes, A. A. Klypin, F.-S. Kitaura, S. Gottlöber, and C. Zhao, MNRAS 458, 613 (2016), arXiv:1410.4684 [astro-ph.CO] .
  * Schmittfull  _et al._ (2017) M. Schmittfull, T. Baldauf, and M. Zaldarriaga, Phys. Rev. D 96, 023505 (2017), arXiv:1704.06634 [astro-ph.CO] .
  * Wang  _et al._ (2017) X. Wang, H.-R. Yu, H.-M. Zhu, Y. Yu, Q. Pan, and U.-L. Pen, Astrophys. J. Lett. 841, L29 (2017), arXiv:1703.09742 [astro-ph.CO] .
  * Ding  _et al._ (2018) Z. Ding, H.-J. Seo, Z. Vlah, Y. Feng, M. Schmittfull, and F. Beutler, MNRAS 479, 1021 (2018), arXiv:1708.01297 [astro-ph.CO] .
  * Ade  _et al._ (2016) P. A. R. Ade _et al._ (Planck), Astron. Astrophys. 594, A13 (2016), arXiv:1502.01589 [astro-ph.CO] .
  * Note (2) The present analysis is not very sensitive to the choice of growth factor Lukic  _et al._ (2007) as it is only concerned with the ratio of ‘wiggle’ and ‘no-wiggle’ simulations. Nevertheless, when dealing with real data, a slight mismatch between the true growth factor and a perhaps wrongly approximated one will in effect show up as a bias factor, $b$, to the reconstructed power, and can therefore easily be taken into account in a subsequent step, e.g. in a parameter fit one can marginalize over this parameter.
  * Note (3) We computed the standard reconstructed density field using the routine FFTRecon of the python suite nbodykit Hand  _et al._ (2018).
  * Note (4) We define $P_{\rm wig}^{\rm IC}(k/\alpha)/P_{\rm nw}^{\rm IC}(k/\alpha)$ as the average over all ten simulations, shown in Fig. 2, and interpolated to accommodate values of $\alpha\neq 1$.
  * Sherwin and Zaldarriaga (2012) B. D. Sherwin and M. Zaldarriaga, Phys. Rev. D 85, 103523 (2012), arXiv:1202.3998 [astro-ph.CO] .
  * Seo  _et al._ (2010) H.-J. Seo, J. Eckel, D. J. Eisenstein, K. Mehta, M. Metchnik, N. Padmanabhan, P. Pinto, R. Takahashi, M. White, and X. Xu, Astrophys. J. 720, 1650 (2010), arXiv:0910.5005 [astro-ph.CO] .
  * Mehta  _et al._ (2011) K. T. Mehta, H.-J. Seo, J. Eckel, D. J. Eisenstein, M. Metchnik, P. Pinto, and X. Xu, Astrophys. J. 734, 94 (2011), arXiv:1104.1178 [astro-ph.CO] .
  * Bernardeau  _et al._ (2002) F. Bernardeau, S. Colombi, E. Gaztanaga, and R. Scoccimarro, Phys. Rept. 367, 1 (2002), arXiv:astro-ph/0112551 .
  * DESI Collaboration  _et al._ (2016) DESI Collaboration, A. Aghamousa, J. Aguilar, S. Ahlen, S. Alam, L. E. Allen, C. Allende Prieto, J. Annis, S. Bailey, _et al._ , arXiv e-prints , arXiv:1611.00036 (2016), arXiv:1611.00036 [astro-ph.IM] .
  * LSST Dark Energy Science Collaboration (2012) LSST Dark Energy Science Collaboration, arXiv e-prints , arXiv:1211.0310 (2012), arXiv:1211.0310 [astro-ph.CO] .
  * Euclid Collaboration  _et al._ (2011) Euclid Collaboration, R. Laureijs, J. Amiaux, S. Arduini, J. L. Auguères, J. Brinchmann, R. Cole, M. Cropper, C. Dabin, L. Duvet, _et al._ , arXiv e-prints , arXiv:1110.3193 (2011), arXiv:1110.3193 [astro-ph.CO] .
  * Lukic  _et al._ (2007) Z. Lukic, K. Heitmann, S. Habib, S. Bashinsky, and P. M. Ricker, Astrophys. J. 671, 1160 (2007), arXiv:astro-ph/0702360 .
  * Hand  _et al._ (2018) N. Hand, Y. Feng, F. Beutler, Y. Li, C. Modi, U. Seljak, and Z. Slepian, Astron. J. 156, 160 (2018), arXiv:1712.05834 [astro-ph.IM] .