* Lanzieri  _et al._ [2022] D. Lanzieri, F. Lanusse, and J.-L. Starck, Hybrid Physical-Neural ODEs for Fast N-body Simulations, in  _Machine Learning for Astrophysics_ (2022) p. 60, arXiv:2207.05509 [astro-ph.CO] .
  * Kaushal  _et al._ [2022] N. Kaushal, F. Villaescusa-Navarro, E. Giusarma, Y. Li, C. Hawry, and M. Reyes, NECOLA: Toward a Universal Field-level Cosmological Emulator, The Astrophysical Journal 930, 115 (2022), arXiv:2111.02441 [astro-ph.CO] .
  * Jamieson  _et al._ [2023] D. Jamieson, Y. Li, R. A. de Oliveira, F. Villaescusa-Navarro, S. Ho, and D. N. Spergel, Field-level Neural Network Emulator for Cosmological N-body Simulations, Astrophys. J. 952, 145 (2023), arXiv:2206.04594 [astro-ph.CO] .
  * Virtanen  _et al._ [2020] P. Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, _et al._ , SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python, Nature Methods 17, 261 (2020).
  * Vlah  _et al._ [2016] Z. Vlah, U. Seljak, M. Yat Chu, and Y. Feng, Perturbation theory, effective field theory, and oscillations in the power spectrum, Journal of Cosmology and Astroparticle Physics 2016 (3), 057, arXiv:1509.02120 [astro-ph.CO] .
  * Spurio Mancini  _et al._ [2022] A. Spurio Mancini, D. Piras, J. Alsing, B. Joachimi, and M. P. Hobson, COSMOPOWER: emulating cosmological power spectra for accelerated Bayesian inference from next-generation surveys, Mon. Notices Royal Astron. Soc. 511, 1771 (2022), arXiv:2106.03846 [astro-ph.CO] .
  * Alsing  _et al._ [2020] J. Alsing, H. Peiris, J. Leja, C. Hahn, R. Tojeiro, _et al._ , SPECULATOR: Emulating Stellar Population Synthesis for Fast and Accurate Galaxy Spectra and Photometry, The Astrophysical Journal Supplement 249, 5 (2020), arXiv:1911.11778 [astro-ph.IM] .
  * Kingma and Ba [2014] D. P. Kingma and J. Ba, Adam: A Method for Stochastic Optimization, arXiv e-prints , arXiv:1412.6980 (2014), arXiv:1412.6980 [cs.LG] .
  * Blatman and Sudret [2010] G. Blatman and B. Sudret, An adaptive algorithm to build up sparse polynomial chaos expansions for stochastic finite element analysis, Probabilistic Engineering Mechanics 25, 183–197 (2010).
  * Blatman and Sudret [2011] G. Blatman and B. Sudret, Adaptive sparse polynomial chaos expansion based on least angle regression, Journal of Computational Physics 230, 2345–2367 (2011).
  * Lüthen  _et al._ [2020] N. Lüthen, S. Marelli, and B. Sudret, Automatic selection of basis-adaptive sparse polynomial chaos expansions for engineering applications, arXiv e-prints , arXiv:2009.04800 (2020), arXiv:2009.04800 [stat.CO] .
  * Marelli and Sudret [2018] S. Marelli and B. Sudret, An active-learning algorithm that combines sparse polynomial chaos expansions and bootstrap for structural reliability analysis, Structural Safety 75, 67–74 (2018).
  * Torre  _et al._ [2019] E. Torre, S. Marelli, P. Embrechts, and B. Sudret, Data-driven polynomial chaos expansion for machine learning regression, Journal of Computational Physics 388, 601–623 (2019).
  * Kaintura  _et al._ [2018] A. Kaintura, T. Dhaene, and D. Spina, Review of polynomial chaos-based methods for uncertainty quantification in modern integrated circuits, Electronics 7, 30 (2018).
  * Eifler  _et al._ [2015] T. Eifler, E. Krause, S. Dodelson, A. R. Zentner, A. P. Hearin, and N. Y. Gnedin, Accounting for baryonic effects in cosmic shear tomography: determining a minimal set of nuisance parameters using PCA, Mon. Notices Royal Astron. Soc. 454, 2451 (2015), arXiv:1405.7423 [astro-ph.CO] .
  * Huang  _et al._ [2019] H.-J. Huang, T. Eifler, R. Mandelbaum, and S. Dodelson, Modelling baryonic physics in future weak lensing surveys, Mon. Not. of the Royal Astron. Soc. 488, 1652 (2019), arXiv:1809.01146 [astro-ph.CO] .
  * Huang  _et al._ [2021] H.-J. Huang, T. Eifler, R. Mandelbaum, G. M. Bernstein, A. Chen, _et al._ , Dark energy survey year 1 results: Constraining baryonic physics in the Universe, Mon. Not. of the Royal Astron. Soc. 502, 6010 (2021), arXiv:2007.15026 [astro-ph.CO] .

##  Appendix A Emulation Implementation Details

###  A.1 Separating the BAO Signal from Linear Power Spectrum

In this Appendix, we explain in detail our algorithm to remove the BAO signal
from the linear power spectrum, a step required in emulating the quantity
$Q^{\rm S}(k,z)$ (Eq. 11). We follow the bacco emulator and employ the
methodology outlined in Appendix C of [118], with slight modifications. The
process can be summarized as follows:

  1. 1.

Standardize $\mathbf{k}$-range: Ensure $P_{\rm L}(k)$ covers the range of
$10^{-4}~{}h\mathrm{Mpc}^{-1}\leq k\leq 5~{}h\mathrm{Mpc}^{-1}$. If necessary,
linearly extrapolate the power spectrum.

  2. 2.

Linear Log-Log Interpolation: Interpolate the power spectrum $P_{\rm L}(k)$
using cubic spline at $2^{n}$ logarithmically-spaced points in $k$, ensuring
data is evenly spaced in log-log space for further analysis.

  3. 3.

Fast Sine Transform: Generate an array of $\log[kP_{\rm L}(k)]$ evaluated at
the $2^{n}$ $k$-bins and apply an orthonormalized type-II fast sine transform
to that array.

  4. 4.

Separating Even and Odd Entries: Denoting the index of the sine-transformed
array by $i$, divide the array into two separate arrays of size $2^{n-1}$,
corresponding to components with even $i$ and odd $i$. Save the even and odd
indices into another two arrays.

  5. 5.

Interpolation and Differentiation: Interpolate both even and odd sine-
transformed arrays using the respective set of indices. Then, calculate the
second derivative by either differentiating the cubic spline or computing the
finite differences to identify the locations of baryonic features.

  6. 6.

Noise Reduction: Average the second derivative points with their immediate
neighbors to reduce noise.

  7. 7.

Identifying Bumps: Locate the BAO bumps by finding the minimum and maximum
values of the second derivatives for both even and odd arrays, locating the
bump to be removed.

  8. 8.

Removing Bumps: Remove the region between the minimum and maximum second
derivatives from both even and odd sine-transformed arrays, filling the gaps
by interpolating the remaining data, scaled by $(i+1)^{2}$, with cubic
splines.

  9. 9.

Inverse Fast Sine Transform: Recombine the even and odd sine-transformed
arrays, dividing the $(i+1)^{2}$ scaling, and apply an inverse fast sine
transform to reconstruct $\log[kP_{\rm L}^{\mathrm{NBAO}}(k)]$.

The scipy101010https://scipy.org/ [137] library was used to perform the sine
transforms on the power spectra, GSL was used for interpolation. Our
implementation is given in this
repository111111https://github.com/bernardo7crf/Cython_Filter_Files. We remark
that the BAO signal separation is not unique but rather depends on the
specific methodology employed. Some examples, as discussed in [138], are the
Savitzky–Golay or Gaussian filters, fitting models such as B-Splines,
employing a power law model with the form $k^{n}$, or leveraging the output
from Boltzmann codes with no baryons. We found that the method of [118] was
sufficient to our needs, and we have not tested other methods.

###  A.2 Gaussian Process Regression

The first method tested to perform emulation is Gaussian Process regression.
GP regression is a non-parametric Bayesian method which assumes any set of
function values has a joint multivariate Gaussian distribution. We denote an
arbitrary principal component amplitude from Eq. 20 at an arbitrary redshift
as $\alpha$, which we treat as a function of a $D$-dimensional vector of
normalized cosmological parameters $\bm{\theta}$. Our training set provides a
vector $\bm{\alpha}_{\rm obs}$ of $N_{\rm train}$ observations of the function
at the training points $\bm{\theta}_{i}$, which generally contain noise:

$\alpha_{\mathrm{obs},i}=\alpha(\bm{\theta}_{i})+\epsilon_{i}.$ (30)

For simplicity, we assume the function values have been centered to have mean
0. We take the noise to be Gaussian, such that
$\epsilon_{i}\sim\mathcal{N}(0,\sigma_{\epsilon}^{2})$, where
$\sigma_{\epsilon}^{2}$ is called the noise variance.

We use a Radial Basis Function (RBF) kernel to model the covariance between
values of the function:

$k(\bm{\theta},\bm{\theta}^{\prime})\equiv\sigma^{2}\exp\Bigg{[}-\frac{1}{2}\sum^{D}_{j=1}\frac{(\theta_{j}-\theta^{\prime}_{j})^{2}}{l_{j}^{2}}\Bigg{]}.$
(31)

where we have the hyperparameters $\sigma^{2}$, called the signal variance,
and the $\\{l_{j}\\}$, called the correlation lengths. We can form a
covariance matrix, $\bm{K}$, from the training points defined by $K_{pq}\equiv
k(\bm{\theta}_{p},\bm{\theta}_{q})$. However, we modify this matrix to account
for noise $\hat{\bm{K}}\equiv\bm{K}+\sigma_{\epsilon}^{2}\bm{I}$, where
$\bm{I}$ is the identity matrix. The hyperparameters $\sigma^{2}$,
$\\{l_{j}\\}$, and $\sigma_{\epsilon}^{2}$, are determined from the training
data by maximizing the marginal likelihood:

$p(\bm{\alpha}_{\rm obs}|K)=\frac{1}{\sqrt{(2\pi)^{N_{\rm
train}}\left|\hat{\bm{K}}\right|}}\exp\Big{[}-\frac{1}{2}\bm{\alpha}_{\rm
obs}^{T}\hat{\bm{K}}^{-1}\bm{\alpha}_{\rm obs}\Big{]}.$ (32)

Aside from the training points, we consider a test point $\bm{\theta}^{*}$ at
which we wish to infer the value of the function $\alpha(\bm{\theta}^{*})$.
The GP assumption states that the value of the function at the test point is
normally distributed along with the observed values at the training points:

$\left[\begin{array}[]{c}\bm{\alpha}_{\rm obs}\\\
\alpha(\bm{\theta}^{*})\end{array}\right]\sim\mathcal{N}\left(\left[\begin{array}[]{c}\bm{0}\\\
0\end{array}\right],\left[\begin{array}[]{cc}\hat{\bm{K}}&\bm{K}^{*}\\\
\bm{K}^{*T}&K^{**}\end{array}\right]\right)$ (33)

The vector $\bm{K}^{*}$ contains the covariances between the function values
at the test point and training points, $\bm{K}^{*}_{p}\equiv
k(\bm{\theta}^{*},\bm{\theta}_{p})$, and we define $K^{**}\equiv
k(\bm{\theta}^{*},\bm{\theta}^{*})$. In order to predict the value of the
function at $\bm{\theta}^{*}$, and the prediction uncertainty, we take the
mean and variance of the conditional distribution
$p(\alpha(\bm{\theta}^{*})|\\{\bm{\theta}_{i}\\},\bm{\alpha}{\rm obs})$
respectively:

$\displaystyle\alpha_{\rm pred}(\bm{\theta}^{*})$
$\displaystyle=\bm{K}^{*T}\hat{\bm{K}}^{-1}\bm{\alpha}_{\rm obs}$ (34)
$\displaystyle\sigma^{2}_{\rm pred}$
$\displaystyle=K^{**}-\bm{K}^{*T}\hat{\bm{K}}^{-1}\bm{K}^{*}.$ (35)

Here $\alpha_{\rm pred}(\bm{\theta}^{*})$ serves as our (centered) prediction
of the principal component amplitude at the test point, with $\sigma^{2}_{\rm
pred}$ providing the uncertainty in the prediction. We perform a different
regression for every principal component, for every redshift, using the
gpy121212https://gpy.readthedocs.io/en/deploy/ library.

###  A.3 Neural Network

Another emulation strategy we have employed is through the use of neural
networks. We use a fully-connected, feed-forward network consisting of the
input layer of normalized cosmological parameters, three hidden layers with
$512$ neurons each, and the output layer of principal components $\alpha_{\ell
j}$ according to Eq. 20. We train one NN model for each redshift, and thus the
index $j$ will be omitted in the following. The value of the neurons in each
subsequent layer are calculated via linear transformations (i.e. the
application of weights), the addition of biases to the previous layer, and the
application of an activation function used by [139, 140], which has the form :

$y^{m+1}_{n}=\left[\gamma^{m}_{n}+(1-\gamma^{m}_{n})\frac{1}{1+e^{-\beta^{m}_{n}y^{m}_{n}}}\right]\tilde{y}^{m}_{n},$
(36)

where $y^{m+1}_{n}$ is the value of the $n$-th neuron from the $(m+1)$-th
layer, $\tilde{y}^{m}_{n}$ is the is the value of the $n$-th neuron from the
$(m+1)$-th layer after the application of weights and biases, and
$\gamma^{m}_{n}$ and $\beta^{m}_{n}$ are parameters of the activation function
that can be backpropagated during the training. We have tested several
activation functions, and their emulation errors at $z=0$ and $k\lesssim
1h/\mathrm{Mpc}$ were within approximately $0.8\%$ (Rectified Linear Unit),
$0.5\%$ (sigmoid) and $0.2\%$ (Eq. 36). We have also tested different numbers
of layers, neurons and residual layers, finding no improvement on the
emulation performance.

In order to train the neural network emulators, we adjust the weights and
biases to minimize the $L_{1}$ loss function, or the mean absolute error:

$L_{1}=\sum_{\mathrm{training}}\sum_{\ell=1}^{N_{\mathrm{PC}}}|\alpha_{\mathrm{truth},\ell}-\alpha_{\mathrm{pred},\ell}|,$
(37)

where $\alpha_{\mathrm{truth},\ell}$ are the principal component weights from
the training sample, $\alpha_{\mathrm{pred},\ell}$ are the neural network
predictions for the training sample, and the outermost sum goes through all
samples in the training set.

In order to prevent overfitting, we perform a $L_{1}-L_{2}$ regularization,
also known as Elastic Net regularization, of the weights and biases. This
encourages the training process to keep weights and biases to a low absolute
value, limiting the complexity of the model. We keep $10\%$ of the training
cosmologies as a validation set for early stopping: a technique that prevents
overfitting by ensuring that the validation loss does not grow during the
training. The training process is halted after 2500 epochs, with an initial
learning rate of $10^{-3}$ which, starting from epoch $1500$, halves every
$200$ epochs. Our neural network is implemented using the
keras131313https://keras.io library, and the loss function is minimized using
the adam optimizer [141].

###  A.4 Polynomial Chaos Expansion

Polynomial Chaos Expansion is another surrogate modelling technique, which
expands the model’s output in orthogonal polynomials [142, 143, 144, 145, 146,
147]. These polynomials capture the functional relationship between inputs and
outputs. The choice of polynomial basis in PCE corresponds to the input
variables’ probability distribution, and if the inputs do not follow a
specific distribution then PCE reduces to a simple polynomial expansion.

To perform the PCE, we start with the principal components $\alpha_{\ell j}$
of Eq. 20, calculated for the training set. At fixed redshift, this represents
a $N_{\mathrm{train}}\times N_{\mathrm{PC}}$ matrix. We decompose this matrix
as a product of two matrices $X\times W$, where $W$ will contain the
polynomial coefficients to be determined, and $X$ is the polynomial basis
matrix representing the orthogonal polynomial expansions. To construct the
matrix $X$, we calculate its elements $X_{ij}$ using:

$X_{ij}=\prod^{N_{\mathrm{PC}}}_{k=1}L_{\beta_{jk}}(x_{ik}).$

Here, $L_{\beta_{jk}}$ are the orthogonal polynomial basis functions, $x_{ik}$
are powers of the input variables and $\bm{\beta}$ is a multi-index matrix.
The multi-index $\bm{\beta}_{j}=(\beta_{j1},\beta_{j2},\ldots,\beta_{jD})$
specifies the polynomial orders for each input dimension in the $j$th column
of $X$. Each $\beta_{jk}$ is a non-negative integer representing the degree of
the $k$th input parameter in the $j$th polynomial term. For example, consider
a 2-dimensional input $(x_{1},x_{2})$ expanded up to the 2nd order. The
$\bm{\beta}$ matrix would be:

$\bm{\beta}=\begin{pmatrix}0&0\\\ 1&0\\\ 0&1\\\ 2&0\\\ 1&1\\\
0&2\end{pmatrix}$

Each row in $\bm{\beta}$ represents a different combination of polynomial
terms, with values indicating the powers of $x_{1}$ and $x_{2}$ for that
combination. We define the candidate set $\mathcal{A}_{\text{cand}}$ of multi-
indices $\bm{\beta}$ to include in the expansion, based on constraints like
the maximum total polynomial degree $p$, maximum interaction order $r$, and an
optional $q$-norm penalty to induce sparsity:

$\mathcal{A}_{\text{cand}}=\left\\{\bm{\beta}_{i}:\left(\sum_{k=1}^{N_{\mathrm{PC}}}\beta_{ik}^{q}\right)^{\frac{1}{q}}\leq
p,\sum_{\begin{subarray}{c}k=1\\\ \beta_{ik}\neq
0\end{subarray}}^{N_{\mathrm{PC}}}1\leq r\right\\}$

We then use Elastic-Net regression to determine the optimal PCE coefficients
$W$ that balance the mean-squared error between the PCE model and training
data with $L_{1}$ (lasso) and $L_{2}$ (ridge) regularization terms to help
prevent overfitting.

After solving the regression problem to obtain the optimal PCE coefficient
matrix $W$, we can evaluate the principal component amplitudes $\alpha$ for a
new cosmology $\bm{\theta}^{*}$ as:

$\alpha_{\rm pred}(\bm{\theta}^{*})=X_{\bm{\theta}^{*}}\times W$

where $X_{\bm{\theta}^{*}}$ is the row of the polynomial basis matrix
corresponding to $\bm{\theta}^{*}$, evaluated using the selected multi-indices
in $\mathcal{A}$.

###  A.5 Neural PCE

During the emulation of the _enhanced-precision_ simulations, we observed
performance issues characterized by increased prediction errors. To address
this problem, we developed a hybrid model, combining PCE with neural networks,
which we have termed the Neural PCE (NPCE). In NPCE, PCE is used for initial
prediction, and the NN performs a correction phase, similar to the iterative
refinement process in ODE solvers.

A standalone PCE emulator was trained as described in A.4, and its outputs
served as inputs to a neural network model. The neural network was designed
with input and output layers sized to match the PCE output, incorporating 2 to
7 hidden layers based on problem complexity, with each layer containing 32 to
128 neurons. The activation function specified in 36 was used alongside a
residual connection between the input and output layers. The network’s weights
were initialized randomly and trained using the PCE predictions as inputs
against the data, aiming to minimize mean squared error via the L-BFGS
algorithm.

For predictions at new cosmologies, the NPCE model refined initial PCE
emulator predictions through the neural network, yielding improved accuracy.
The residual connection in the neural network allowed it to learn corrections
to the PCE predictions. It is important to note that we trained the PCE and
neural network components separately, rather than as an integrated model. By
adding the neural network, we increased the model parameters by approximately
$10\%$ compared to using PCE alone, without the need to modify the PCE
training process itself.

The effectiveness of this model is demonstrated through a comparison of the
prediction errors between the existing emulators and NPCE in Figure 3. The
results show a reduction in error rates using the new model, with the error
being under $\sim 0.25\%$ for most cosmologies on all scales, indicating a
successful improvement in performance.

Figure 12: 68% and 95% confidence contours of several cosmological parameters
from MCMCs run with three emulation techniques: Gaussian Process regression
(pink filled contours), neural network (blue dashed lines) and Polynomial
Chaos Expansion (maroon dotted-dashed lines). In both panels, the COLA
emulators are trained with _default-precision_ simulations calibrated to the
ee2 reference cosmology, which also serves as the fiducial cosmology. We use a
scale cut even more aggressive than Cutoff 3 for this test, with
$\theta_{\mathrm{min}}=2.7^{\prime}$ for $\xi_{+}$ and
$\theta_{\mathrm{min}}=8.7^{\prime}$ for $\xi_{-}$.

###  A.6 Testing Emulator Equivalence

We now test whether the three emulator techniques, namely Gaussian Process
regression, neural networks, and Polynomial Chaos Expansion, are equivalent at
the level of parameter inference. For this test, we ran MCMCs using the
LSST-Y1 simulated cosmic shear data. We utilized the three emulated methods,
where each emulator was trained with _default-precision_ simulations and
$N_{\rm refs}=1$. We use the ee2 reference cosmology as the fiducial, and
employ a scale cut even more aggressive than Cutoff 3, with
$\theta_{\mathrm{min}}=2.7^{\prime}$ for $\xi_{+}$ and
$\theta_{\mathrm{min}}=8.7^{\prime}$ for $\xi_{-}$. Fig. 12 shows the
confidence contours and marginalized posteriors of several cosmological
parameters obtained with the three emulators. We see strong agreement in the
1D distributions, and we conclude that any of the three emulation techniques
can be used interchangeably.

##  Appendix B $N_{\rm refs}^{\Lambda\rm CDM}=500$ and $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$, Comparisons at the Boost Level

Figure 13: Bounds on the best $50\%$, $90\%$, and $100\%$ of emulation errors
at $z=0$ for our NPCE $w$CDM emulator of the uncalibrated $B_{\rm COLA}(k,z)$
used for the $N_{\rm refs}^{\Lambda\rm CDM}=500$ and $N_{\rm refs}^{\Lambda\rm
CDM}=\infty$ approaches from Sec. III.5. The emulation errors in the
$\Lambda$CDM validation set quantify the error introduced in emulating the raw
COLA prediction $B_{\rm COLA}(k,z)$ at the $\Lambda$CDM reference cosmologies,
which is involved in both approaches. The errors in the $w$CDM validation set
do similarly for the $w$CDM prediction of $B_{\rm COLA}(k,z)$ in the $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$ approach. Figure 14: Following Fig. 1, the
comparisons of the boost from COLA to ee2, for the training set and validation
sets, using the $N_{\rm refs}^{\Lambda\rm CDM}=500$ and $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$ methodology. For $N_{\rm refs}^{\Lambda\rm
CDM}=500$ references we compute the boost using $B(k,z)=\tilde{B}(k,z)$ as in
Eq. 5, while for $N_{\rm refs}^{\Lambda\rm CDM}=\infty$ references we use
$B(k,z)=\bar{B}(k,z)$ as in Eq. 29. For these comparisons, we evaluate the
respective equations for the boost by computing $B_{\rm COLA}(k,z)$ directly
from simulations for the $w$CDM predictions, while using the NPCE emulator for
the $\Lambda$CDM-projected predictions as we lack the requisite $\Lambda$CDM
simulations.

In order to implement the $N_{\rm refs}^{\Lambda\rm CDM}=500$ and $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$ approaches, the same emulator for the
uncalibrated boost $B_{\rm COLA}(k,z)$ was used, trained on the $N_{\rm
train}^{w\rm CDM}=500$ $w$CDM cosmologies. In both cases emulation now enters
the pipeline in two places. In the $N_{\rm refs}^{\Lambda\rm CDM}=500$
approach, the $B_{\rm COLA}(k,z)$ emulator serves to compute the COLA
prediction at the reference cosmologies in Eq. 4, before the $\tilde{B}(k,z)$
emulator is then subsequently trained. Meanwhile, in the $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$ approach the $B_{\rm COLA}(k,z)$ emulator is
used in Eq. 29 both to evaluate the COLA prediction at the $w$CDM point and
again at the $\Lambda$CDM-projected point. Therefore, we use the more
effective NPCE method, and a larger $N_{\rm PC}=25$, in order to prevent a
compounding of emulation errors in our analysis. With the higher number of
principal components, there was little improvement observed in implementing
the BAO-smearing procedure, therefore we use $Q^{\rm NS}(k,z)$ as the
emulation variable for the $B_{\rm COLA}(k,z)$ emulator. This is advantageous
in the $N_{\rm refs}^{\Lambda\rm CDM}=\infty$ approach as one would otherwise
need to compute $P_{\rm L}(k,z|\bm{\theta}_{\Lambda\rm CDM})$ in addition to
$P_{\rm L}(k,z|\bm{\theta})$ merely to recover the boosts from the emulation
variable.

Figure 15:  Confidence contours (68% and 95%) for $\Omega_{m}$, $S_{8}$ and
$w$, in three cases: ee2 (pink filled contours), the $N_{\rm refs}^{\Lambda\rm
CDM}=101$ _default-precision_ COLA emulator (blue dashed lines), and the same
COLA emulator when marginalizing over 1 PC is employed (maroon dotted-dashed
lines). The top panel shows results using the fiducial cosmology
$(\Omega_{m}^{\uparrow},A_{s}^{\uparrow},w^{\downarrow})$. We observe that the
2D contours from the COLA chain employing marginalization, completely enclose
the ee2 contours and spread the 1D distributions in the direction of the ee2
distributions, relative to the COLA chains without PC marginalization.
However, in the bottom panel where the fiducial cosmology
$(\Omega_{m}^{\uparrow},A_{s}^{\downarrow},w^{\downarrow})$ is used, the 1D
distributions spread in the direction opposite of the ee2 constraints.  Figure
16: 1D differences between COLA emulators and ee2 for fiducial cosmologies
outside of the $\Lambda$CDM region, for the two most aggressive scale cuts. We
test the effect of marginalizing over a PC measuring the difference between
the COLA and ee2 data vector, using our _default-precision_ emulator with
$N_{\rm refs}^{\Lambda\rm CDM}=101$.

The NPCE $B_{\rm COLA}(k,z)$ emulator utilizes a polynomial expansion up to
the 12th order in the cosmological parameters. Additionally, the neural
network responsible for the correction phase consists of 3 layers and 128
neurons, employing the activation function presented in Eq. 36. Fig. 13 shows
the emulation errors in $B_{\rm COLA}(k,z)$ at $z=0$. As the reference COLA
predictions in both $N_{\rm refs}^{\Lambda\rm CDM}=500$ and $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$ approaches are confined to $\Lambda$CDM, we
compute the emulation errors in the $\Lambda$CDM validation set, in addition
to the $w$CDM validation set. We see that the emulator of the uncalibrated
$B_{\rm COLA}(k,z)$ achieves an error of less than 0.3% on all scales used
$k\leq\pi h$Mpc-1 for either validation set. In the $N_{\rm refs}^{\Lambda\rm
CDM}=500$ approach, we opted to use the NPCE method for the $\tilde{B}(k,z)$
emulator as well following the choices detailed in Sec. II.2.2, using $N_{\rm
PC}=25$. We set the parameter $\sigma_{d}=0.2$ in Eq. 6, which yields
emulation errors similar to the PCE errors displayed in the middle panel of
Fig. 3. To improve predictive performance we increased the number of layers in
the neural network step to 7.

In Fig. 14 we present the disagreement between the COLA boosts and ee2 at
$z=0$, when using $\tilde{B}(k,z)$ with $N_{\rm refs}^{\Lambda\rm CDM}=500$
and $\bar{B}(k,z)$ for the $N_{\rm refs}^{\Lambda\rm CDM}=\infty$ method.
Here, we use the simulation output directly to compute $B_{\rm COLA}(k,z)$ for
$w$CDM cosmologies in the calculation of the respective boosts, and use the
NPCE $B_{\rm COLA}(k,z)$ emulator to obtain the COLA predictions at
$\Lambda$CDM references. We perform this comparison for both the 279 $w$CDM
training cosmologies inside the ee2 parameter space as was done for the right
panels in Fig. 1, as well as for the $w$CDM validation set. This is because
both approaches are similar in their treatment of the training set,
calibrating the training simulations, either exclusively or most heavily, to
their $\Lambda$CDM-projected cosmologies. However, for predictions of the
validation points, the $N_{\rm refs}^{\Lambda\rm CDM}=\infty$ strategy is able
to calibrate these predictions to their respective $\Lambda$CDM-projected
cosmologies, unlike the $N_{\rm refs}^{\Lambda\rm CDM}=500$ case.

For the training cosmologies, the $N_{\rm refs}^{\Lambda\rm CDM}=500$ approach
produced errors at $k=1h$Mpc-1 above $1\%$, $1.5\%$, and $2\%$ for only
$26.5\%$, $10.4\%$, and $2.5\%$ of cosmologies respectively. Meanwhile the
$N_{\rm refs}^{\Lambda\rm CDM}=\infty$ approach had $20.8\%$, $9.0\%$, and
$4.3\%$ of these cosmologies outside the $1\%$, $1.5\%$, and $2\%$ thresholds.
However, in the validation set, the $N_{\rm refs}^{\Lambda\rm CDM}=500$
references yielded $31.0\%$, $12.0\%$, and $1.0\%$ of cosmologies outside of
those respective criteria, while the $N_{\rm refs}^{\Lambda\rm CDM}=\infty$
method only yielded $17.0\%$, $3.0\%$, and $1.0\%$. Therefore, excluding the
$\mathcal{O}(0.1\%)$ emulation errors shown in Fig. 13, the $N_{\rm
refs}^{\Lambda\rm CDM}=\infty$ methodology manages to reproduce the ee2 boost
predictions at $k=1h$Mpc-1 at sub-percent error for $\sim 80\%$ of cosmologies
throughout the $w$CDM parameter space.

##  Appendix C Mitigating Biases Using Principal Component Analysis

As a final technique to mitigate the biases in the COLA analyses, we try a
procedure inspired by baryonic effect mitigation using Principal Component
Analysis [148, 149, 150]. As a high number of ee2 $\Lambda$CDM references
proved an effective approach to reducing the bias of COLA, we consider whether
differences in the $\Lambda$CDM data vectors between COLA and ee2 can
effectively be marginalized over in extended models. We start by generating an
LHS of 30 $\Lambda$CDM cosmologies in a space shrunken symmetrically by $30\%$
total from the ee2 limits in each parameter, in order to lessen the
variability of the data vectors.

For each cosmology, we compute two cosmic shear data vectors $\mathbf{d}$, one
using a COLA emulator for nonlinear corrections and another using ee2. We then
generate a difference matrix $\Delta$ such that each column is the difference
of the COLA and ee2 data vectors for one cosmology. The shape of $\Delta$ is
$N_{\mathrm{DV}}\times 30$, where $N_{\mathrm{DV}}$ is the number of elements
in the data vector. Schematically, this reads:

$\Delta=\left[\begin{array}[]{ccc}\vline&&\vline\\\
(\mathbf{d}_{1}^{\textsc{ee2}}-\mathbf{d}_{1}^{\mathrm{COLA}})&...&(\mathbf{d}_{30}^{\textsc{ee2}}-\mathbf{d}_{30}^{\mathrm{COLA}})\\\
\vline&&\vline\\\ \end{array}\right]_{N_{\mathrm{DV}}\times 30}.$ (38)

Performing a Cholesky decomposition on the survey covariance matrix
$C=LL^{T}$, we then weight the difference matrix by
$\Delta_{L^{-1}}=L^{-1}\Delta$. We then perform a Singular Value Decomposition
on $\Delta_{L^{-1}}$, such that $\Delta_{L^{-1}}=U\Sigma V$, where $U$ and $V$
are square orthogonal matrices of shapes $N_{\mathrm{DV}}\times
N_{\mathrm{DV}}$ and $30\times 30$, respectively. $\Sigma$ is a matrix of
shape $N_{\mathrm{DV}}\times 30$ whose first $30$ rows form a diagonal matrix
with the singular values, and the other elements are zero. The first $30$
columns of $U$ are eigenvectors of $\Delta_{L^{-1}}$, denoted by
$\mathbf{e}_{i}$:

$U=\left[\begin{array}[]{cccc}\vline&&\vline&...\\\
\mathbf{e}_{1}&...&\mathbf{e}_{30}&...\\\ \vline&&\vline&...\\\
\end{array}\right]_{N_{\mathrm{DV}}\times N_{\mathrm{DV}}}.$ (39)

By reweighting these vectors with $L$, they can be used to describe the
differences between COLA and ee2 data vectors in the $\Lambda$CDM model. The
Cholesky and Singular Value decompositions were performed using scipy [137].
To incorporate the principal components in the analysis, at each step of the
MCMC, we correct the final data vector as:

$\mathbf{d}=\mathbf{d}^{\text{No
PC}}+L\sum_{i=1}^{n_{\mathrm{PC}}}Q_{i}\mathbf{e}_{i},$ (40)

where $Q_{i}$ is a parameter sampled in the MCMC with large flat priors
$[-100,100]$, and $n_{\mathrm{PC}}$ is the number of principal components
chosen to correct the data vector. By marginalizing the cosmological parameter
constraints over $Q_{i}$, we expect to mitigate the differences between COLA
and ee2 results at the cost of increasing the error bars. We remark that this
method only relies on $\Lambda$CDM predictions of the matter power spectrum,
and we want to test whether this procedure can mitigate biases in extended
models, in our case wCDM.

In Fig. 15 we show 2D contours when marginalizing over PCs is employed and
when it is not. The $N_{\rm refs}^{\Lambda\rm CDM}=101$ _default-precision_
emulator was used for these tests. We include two fiducial cosmologies using
Cutoff 2, showing mixed results between the two cases. While marginalizing
over 1 PC broadened the posterior, which can alleviate tension, the posterior
did not always move in the correct direction so as to lessen the difference in
the means between the COLA and ee2 constraints. Fig. 16 shows a more complete
assessment, using the 4 fiducial cosmologies with $w=w^{\downarrow}$, and the
two more aggressive scale cuts. Here we plot the difference in the means
$\Delta\theta=\theta_{\rm COLA}-\theta_{\textsc{ee2}}$ and the error bars
separately
$\sigma_{\theta}=\sqrt{\sigma_{\theta,\textsc{ee2}}^{2}+\sigma_{\theta,\mathrm{COLA}}^{2}}$,
to show the effect of PC marginalization on $\sigma_{\theta}$.

The results in Fig. 16 are mixed, showing a broadening of error bars but
failing to consistently reduce the difference in means between the COLA
emulators and ee2. Hence, we did not extensively test marginalizing over a
larger number of PCs. This approach can likely be improved by adjusting the
principal components for each cosmology sampled in the chain.