[1]StephanRasp

1]Technical University of Munich, Germany

Stephan Rasp<EMAIL_ADDRESS>
# Coupled online learning as a way to tackle instabilities and biases in
neural network parameterizations: general algorithms and Lorenz 96 case study
(v1.0)

###### Abstract

Over the last couple of years, machine learning parameterizations have emerged
as a potential way to improve the representation of sub-grid processes in
Earth System Models (ESMs). So far, all studies were based on the same three-
step approach: first a training dataset was created from a high-resolution
simulation, then a machine learning algorithm was fitted to this dataset,
before the trained algorithm was implemented in the ESM. The resulting online
simulations were frequently plagued by instabilities and biases. Here, coupled
online learning is proposed as a way to combat these issues. Coupled learning
can be seen as a second training stage in which the pretrained machine
learning parameterization, specifically a neural network, is run in parallel
with a high-resolution simulation. The high-resolution simulation is kept in
sync with the neural network-driven ESM through constant nudging. This enables
the neural network to learn from the tendencies that the high-resolution
simulation would produce if it experienced the states the neural network
creates. The concept is illustrated using the Lorenz 96 model, where coupled
learning is able to recover the "true" parameterizations. Further, detailed
algorithms for the implementation of coupled learning in 3D cloud-resolving
models and the super parameterization framework are presented. Finally,
outstanding challenges and issues not resolved by this approach are discussed.

The representation of subgrid processes, especially clouds, is the main cause
of uncertainty in climate projections and a large error source in weather
predictions (Schneider et al., 2017b). Models that explicitly resolve the most
difficult processes are now available but are too expensive for operational
forecasting. Machine learning (ML) has emerged as a potential shortcut which
would allow using short-term high-resolution simulations in order to improve
climate and weather models. However, two issues have plagued all approaches so
far: First, simulations with neural networks turned out to be unstable at
times. Second, even if stable, the resulting simulations had biases compared
to the reference model. In pre-ML climate model development, biases were
reduced by manual tuning of a handful of well-known parameters (Hourdin et
al., 2017). With thousands of non-physical parameters in a neural network,
this is no longer possible. In this paper, I propose coupled online learning
as a potential mechanism to tackle these two issues and illustrate the
principle using the two-level Lorenz 96 (L96) model, a common (but probably
too simple) model of multi-scale atmospheric flow (Lorenz,
1995)111Confusingly, even though the paper appears to have been published in
1995, most people refer to the model as the Lorenz 96 model..

##  1 Review of online machine learning parameterizations

Over the last couple of years, several attempts have been made at building ML
subgrid parameterizations, all of which followed a similar approach (Fig. 1).
The first step is to create a training dataset from a reference simulation. In
step two, this dataset is then used to train a ML algorithm. After training,
the predictions of the algorithm can then be compared offline against a
validation dataset. Promising offline results have been obtained for a number
of subgrid processes (Krasnopolsky et al., 2013; Bolton and Zanna, 2019). Step
three is to implement the ML algorithm in the climate model code where it
replaces the traditional subgrid schemes and is coupled to the dynamical core
and non-ML parameterizations. These hybrid models are then integrated forward
in what I will call online mode. The first study to implement an online ML
parameterization was done by Chevallier et al. (2000) who successfully
emulated the ECMWF radiation scheme. More recently, three studies have
implemented all three steps for moist convection in the context of an
atmospheric model (Brenowitz and Bretherton, 2018; O’Gorman and Dwyer, 2018;
Rasp et al., 2018). Note that all of these three studies used a simplified
aquaplanet world and the ML parameterizations only included the most important
variables in their input/output vectors. Cloud water and ice, for example,
were omitted for the sake of simplicity.

Figure 1: Schematic overview of ML parameterization workflow with and without
coupled online learning.

###  1.1 Rasp et al. (2018) – Super-parameterization with a neural network

The three attempts differ in training data and ML algorithms used. In Rasp et
al. (2018)(RPG18), we used a super-parameterized climate model, the Super-
Parameterized Community Atmosphere Model (SPCAM) as our training model
(Khairoutdinov and Randall, 2001). In super-parameterization (SP), a 2D cloud-
resolving model (CRM; $\Delta x$ = 4 km) is embedded in each global
circulation model (GCM; $\Delta x\approx$ 200 km, $\Delta t$ = 30 min) grid
column. The CRM handles convection, turbulence and microphysics, while
radiation222In some SPCAM versions radiation is computed on the CRM grid.,
surface processes and the dynamics are computed on the GCM grid as usual.
Compared to a global 3D CRM, SP is obviously less realistic but has several
conceptual and technical advantages. First, sub-grid and grid scale processes
are clearly separated, which makes it easy to define the parameterization task
for a ML algorithm. Second, because the CRM lives in isolation, it exactly
conserves certain quantities (e.g. energy and mass). A third, very practical
advantage is that SP simulations are significantly cheaper than global 3D
CRMs. In our study we trained a deep neural network to emulate the CRM
tendencies. The offline validation scores were very encouraging (Gentine et
al., 2018) even though the deterministic ML parameterization was unable to
reproduce the variability in the boundary layer. When we subsequently
implemented the ML parameterization in the climate model and ran it
prognostically (online), we managed to engineer a stable model that produced
results close to the original SP-GCM. However, small changes, either to the
training dataset or in the input and output vectors, quickly led to
unpredictable blow-ups. In these cases the network would output increasingly
unrealistic tendencies at individual grid columns. Further, some biases to the
reference model were evident (Fig. 1 in RPG18).

###  1.2 Brenowitz and Bretherton (2018) – Global 3D CRM with a neural network

Brenowitz and Bretherton (2018)(BB18) (extended in Brenowitz and Bretherton
(2019)) used a 3D CRM ($\Delta x$ = 4 km, $\Delta t$ = 10 s) to create their
reference simulation. This requires an additional spatial and temporal coarse-
graining step to generate the training data for a ML parameterization for a
coarser resolution model (in their case $\Delta x$ = 160 km, $\Delta t$ = 3
h). The challenge is to find the apparent subgrid tendencies. BB18 computed
the subgrid tendency $\left(\partial\bar{\phi}/\partial t\right)_{sg}$ of an
arbitrary variable $\phi$ (e.g. temperature or humidity) as the residual of
the total coarse-grained tendency and the coarse-grained advection term:

$\underbrace{\frac{\partial\bar{\phi}}{\partial t}}_{\text{Total coarse-
grained
tendency}}+\underbrace{\overline{\mathbf{v}}\cdot\overline{\nabla}\bar{\phi}}_{\text{Coarse-
grained advection}}=\left(\frac{\partial\bar{\phi}}{\partial t}\right)_{sg}$
(1)

where $\mathbf{v}$ is the 3D wind vector. This coarse-graining procedure
assumes that the coarse-grained advection term closely resembles the dynamics
of the coarse-grid GCM. This assumption might not be true in many cases.
Further, the residual "sub-grid" terms do not obey any conservation
constraints.

BB18 then fitted a neural network to the coarse-grained data, which produces
good results in offline mode. In online mode, however they also experienced
instabilities. Brenowitz and Bretherton (2019) identified unphysical
correlations learned by the network as the cause for the instabilities and
used two fixes to produce stable longer-term simulations. The first fix is to
cut off upper levels from the input vector. The second fix involves a multi-
time-step loss-function that integrates the network predictions forward in a
single-column model setup. This essentially penalizes unstable feedback loops.
Despite these improvements, the simulation drifts, potentially as a result of
the coarse-graining issues mentioned in the previous paragraph. For a further
exploration of the instabilities see Brenowitz et al. (2020).

###  1.3 O’Gorman and Dwyer (2018) – Traditional parameterization with a
random forest

The third online parameterization by O’Gorman and Dwyer (2018) uses a
traditional parameterization as reference. For cloud parameterizations, this
is mainly a proof-of-concept. For other, computationally expensive
parameterizations, such as line-by-line radiation parameterizations, pure ML
emulation is a promising target. As with our super-parameterization, this way
the parameterization task is clearly defined. The main difference of O’Gorman
and Dwyer (2018) to RPG18 and BB18 is the ML method: a random forest (Breiman,
2001). Rather than learning a regression, as neural networks do, random
forests essentially learn a multi-dimensional lookup table. Advantages of this
approach are: 1) The predictions of a random forest are limited by what it has
seen in the training dataset. This means it cannot produce "unphysical"
tendencies which could lead to model blow-ups. 2) Since the training data
obeys physical constraints, so will the random forest predictions by
default.333At least to a good degree of approximation. Predictions of decision
trees and therefore also random forests are averages over several training
targets. Each target will perfectly obey constraints. Since the conservation
constraints are likely non-linear, an average does not necessarily keep this
property but probably comes close. Comparing the results of O’Gorman and Dwyer
(2018) to RPG18 or BB18, it also seems like random forests are competitive
with neural networks for the parameterization problem. Note also a recently
published follow-up paper (Yuval and O’Gorman, 2020) in which the authors
trained a random forest from coarse-grained CRM data. Their coarse training
procedure differs from that in BB18 since it makes direct use of the
parameterized tendencies. Here, I will not further discuss random forests,
since they do not lend themselves to incremental online learning in their most
common implementations. Note, however, that there are online learning
algorithms for random forests (Saffari et al., 2009).

##  2 Coupled online learning – the general concept

Coupled online learning is essentially a second training step after the first
offline training on a reference dataset. The basic idea of coupled learning is
to run the low-resolution model with the machine learning parameterization
(ML-LR) model in parallel with the high-resolution (HR) model and train the
network every or every few time steps (3b. in Fig. 1)444A note on the
terminology: I will use the terms HR (high-resolution) and LR (low-resolution)
here when speaking about the general algorithm. When talking specifically
about atmospheric science applications, I will use the more common terms CRM
and GCM.. The HR model is continuously nudged towards the LR model state
keeping the two simulations close to each other. How close the two runs are
depends on the nudging time scale $\tau_{\text{nudging}}$. A small nudging
time scale forces the models closer together but the HR model might respond
unrealistically or eventually blow up if the nudging is too strong. Assuming
that the HR and LR model state are close together, this method allows the ML
parameterization to see what the HR model would do if it lived in the ML-LR
model world. This should help in reducing biases and preventing instabilities.
Take as an example a neural network parameterization that develops an unstable
feedback loop and starts producing highly unrealistic tendencies. With offline
learning only, the model will eventually blow up. In coupled learning, such
unrealistic predictions would result in large losses. In the next gradient
descent step the network will learn not to produce such tendencies any more.
The hope is that during this coupled learning phase the errors of the network
will become smaller and smaller, so that eventually the ML-LR model can be run
without supervision. Ideally, one could intermittently turn on the
"supervising" HR model for cases where the ML parameterization starts to
produce undesired tendencies. However, one has to consider that HR models
require a spin up phase, which prohibits immediate deployment from a cold
start. This problem might be less pronounced in the case of an embedded HR
model (as in SP) but nevertheless motivates the approach in this paper of
continuously running the two models in parallel.

The instability issues in previous studies can also be seen as a consequence
of overfitting to the reference simulation used for training. Once the ML
parameterization is coupled to the LR model it will create its own climate
which likely lies somewhat outside the training manifold. This can easily lead
to problems because neural networks struggle to extrapolate beyond what they
have seen during training. Coupled learning combats this problem by extending
the training with HR targets for each state that the ML-LR model produces.

Figure 2: Evolution of a tracer $\phi$ during one LR model time step. This
schematic applies the the L96 and 3D HR model case.

The algorithmic details of coupled learning differ depending on the exact
model setup. The main contribution of this paper will be to describe coupled
learning algorithms for the simple L96 model as well as global 3D HR models
and SP models. To understand how coupled learning actually works it is helpful
to draw diagrams for the evolution of a tracer $\phi$ (e.g. temperature) at
one grid point during one LR model time step. I will start with the case of
the 3D HR setup, which also applies to the L96 model (see Fig. 2 for notation;
for more technical details refer to algorithms 1 and 2). At the beginning of
the time step $\phi$ will generally have different values in the LR and HR
model (the HR values are coarse-grained to the LR model grid). The difference
$\Delta\phi=$ LR - HR is used to compute the nudging tendency555Note that
”tendencies” are defined per unit of time, while ”increments” are tendencies
multiplied by a time step. $\Delta\phi/\tau_{\text{nudging}}$ which is
constant during the HR model integration. The total increment from nudging
then is $\Delta\phi_{\text{nudging}}=\Delta\phi/\tau_{\text{nudging}}*\Delta
t_{\text{LR}}$. In addition to the nudging increment, the HR model also
evolves on its own. Assuming that during this short time interval the nudging
and the HR-internal evolution (i.e. the HR increment that would be in the
absence of nudging) are independent, the state of the HR model at the end of
the LR model time step (HR′) is a linear superposition666I will call this the
linear superposition assumption in the rest of the paper.. The "assumed" HR-
internal increment can be computed as $\Delta\phi_{\text{HR-
internal}}=\text{HR}^{\prime}-\text{HR}-\Delta\phi_{\text{nudging}}$. In the
meantime, the LR model will first execute its dynamical core and any other
parameterizations that are not represented by a ML algorithm777Typically, in a
LR model time step the physics is run before the dynamics. But where the time
step starts and ends is arbitrary, so the two can be switched without
problems.. The resulting state is LR′. Then the ML parameterization will be
called and the resulting tendencies will be added to give LR′′ = LR′ \+
$\mathcal{P}$(LR). One open question is whether the input to the ML
parameterization should be LR or LR′. In this study LR is used but the
differences are small. If the ML-LR model was a perfect emulation of the HR
model, the total LR increment $\text{LR}^{\prime\prime}-\text{LR}$ should be
equal to the HR increment $\Delta\phi_{\text{HR-internal}}$. Therefore, the
target for the parameterization is $y=\Delta\phi_{\text{HR-
internal}}-(\text{LR}^{\prime}-\text{LR})$. and the mean squared error loss is
$\mathcal{L}=\frac{1}{N}\sum_{i}(y-\mathcal{P}(\text{LR}))^{2}$. The ML
parameterization is then optimized every few time steps.

##  3 Parameterization experiments using the Lorenz 96 model

###  3.1 The L96 model

The L96 model (Lorenz, 1995) is an idealized model of atmospheric circulation
that, in its two-level variant, has been extensively used for parameterization
research (Wilks, 2005; Crommelin and Vanden-Eijnden, 2008). Here, I use the
model as described in Schneider et al. (2017a). Briefly, the model consists of
a slow variable $X_{k}$ ($k=1,\dots,K$) and a coupled fast variable $Y_{j,k}$
($j=1,\dots,J$):

$\frac{dX_{k}}{dt}=\underbrace{-X_{k-1}(X_{k-2}-X_{k+1})}_{\text{Advection}}\underbrace{-X_{k}}_{\text{Diffusion}}\underbrace{+F}_{\text{Forcing}}\underbrace{-hc\bar{Y}_{k}}_{\text{Coupling}}$
(2)

$\frac{1}{c}\frac{dY_{j,k}}{dt}=\underbrace{-bY_{j+1,k}(Y_{j+2,k}-Y_{j-1,k})}_{\text{Advection}}\underbrace{-Y_{j,k}}_{\text{Diffusion}}\underbrace{+\frac{h}{J}X_{k}}_{\text{Coupling}}$
(3)

Here, the overbar denotes an average over all $J$ fast variables for a given
$K$. Both, $X$ and $Y$ are periodic. $K=36$, $J=10$, $h=1$ and $F=c=b=10$.
These parameters indicate that the fast variable evolves 10 times faster than
the slow variable and has one tenth of the amplitude. A Runge-Kutta 4th order
scheme with a time step of 0.001 is used to integrate these equations. The
one-level model consists only of equation 2 without the coupling term on the
right hand side888For animations of the L96 system, see
https://raspstephan.github.io/blog/lorenz-96-is-too-easy/.

For parameterization research, $X$ represents the large-scale, resolved
variables, whereas $Y$ represents the small-scale, unresolved variables. The
job of a parameterization $\mathcal{P}$ is to approximate the coupling term in
the $X$ equation:

$-hc\bar{Y}_{k}:=B_{k}\approx\mathcal{P}(X_{k})$ (4)

The parameterization task is shown in Fig. 3. Here, I only consider
deterministic parameterizations that are local in space and time. Studies
(Wilks, 2005; Dueben and Bauer, 2018; Pathak et al., 2018; Bocquet et al.,
2019) suggest that non-local and stochastic parameterizations achieve better
results. However, the focus here is on developing a learning algorithm rather
than achieving building the best parameterization which is why I opted for the
simplest setup.

Figure 3: Blue dots are data points from a reference simulation with the real
L96 parameters. The solid orange and green lines are the linear regression and
neural network parameterization fitted to this data. The red dots are data
points from the L96 simulations with "wrong" parameter values used for
pretraining. The dashed lines are the parameterization fits for these "wrong"
values, which serve as a starting point for the coupled learning experiments.

###  3.2 Machine learning parameterizations

Two parameterizations will be considered: a linear regression and a neural
network. The linear regression case is easily interpretable and helps to
illustrate the learning procedure, while the neural network is a more
realistic case.

The linear regression parameterization looks as follows:

$B_{k}=aX_{k}+b$ (5)

When fitted to the points shown in Fig. 3, $a=-0.31$ and $b=-0.20$.

Algorithm 1 Online learning algorithm for the L96 model. Bold-face
$\mathbf{X}$($\mathbf{Y}$) indicate vectors with all $K$($J$) elements

0: Pretrained LR-parameterization $\mathcal{P}_{\theta}$ with parameters
$\theta$

0: Initial conditions $\mathbf{X}_{0}$ and $\mathbf{Y}_{0}$

0: Two-level HR model with time step $\Delta t_{\text{HR}}$

0: One-level LR model with time step $\Delta t_{\text{LR}}=N\Delta
t_{\text{HR}}$

0: Nudging time scale $\tau_{\text{nudging}}$

0: Feature memory $\mathcal{F}$ and target memory $\mathcal{T}$

0: Training frequency $M$; learning rate $\alpha$ and batch size $m$

Initialize HR with $\mathbf{X}_{0}$ and $\mathbf{Y}_{0}$; initialize "LR" with
$\mathbf{X}_{0}$

for $t$ = 1,… do

Store difference at beginning of time step
$\Delta\mathbf{X}=\mathbf{X}_{\text{LR}}-\mathbf{X}_{\text{HR}}$

Store LR state at beginning of time step
$\hat{\mathbf{X}}_{\text{LR}}=\mathbf{X}_{\text{LR}}$

Store HR state at beginning of time step
$\hat{\mathbf{X}}_{\text{HR}}=\mathbf{X}_{\text{HR}}$

for $n=1,N$ do

Integrate HR model with forcing term $\Delta\mathbf{X}/\tau_{\text{nudging}}$
added to the RHS of Eq. 2 to get new $\mathbf{X}_{\text{HR}}$

end for

Compute "assumed" HR increment $\Delta\mathbf{X}_{\text{HR-
internal}}=\mathbf{X}_{\text{HR}}-\hat{\mathbf{X}}_{\text{HR}}-\Delta\mathbf{X}/\tau_{\text{nudging}}*\Delta
t_{\text{LR}}$

Store $\hat{X}_{\text{LR},k}$ for $k\in{1,\dots,K}$ in $\mathcal{F}$

Integrate LR model without parameterization term to get new
$\mathbf{X}_{\text{LR}}$

Store $\Delta X_{\text{HR-
internal},k}-(X_{\text{LR},k}-\hat{X}_{\text{LR},k})$ for $k\in{1,\dots,K}$ in
$\mathcal{T}$

Update LR state
$\mathbf{X}_{\text{LR}}\leftarrow\mathbf{X}_{\text{LR}}+\mathcal{P}_{\theta}(\hat{\mathbf{X}}_{\text{LR}})$

if $t$ mod $M$ = 0 then

Optimize loss for samples in $\mathcal{F}$ and $\mathcal{T}$:
$\mathcal{L}_{\theta}=(\mathcal{P}_{\theta}(\mathcal{F})-\mathcal{T})^{2}$
using stochastic gradient descent with learning rate $\alpha$ and batch size
$m$

Empty $\mathcal{F}$ and $\mathcal{T}$

end if

end for

Figure 4: (a) Evolution of linear regression parameters $a$ (slope) and $b$
(intercept). An iteration on the x-axis corresponds to one gradient descent
update which in this case is equal to ten ML-LR model time steps. (b)
Evolution of the neural network parameterization starting with the "wrongly"
pretrained fit. See the Jupyter notebook for an animated version of this.

Neural networks consist of one or multiple layers of linearly connected nodes,
modified by non-linear activation functions.999For a great introduction to
neural networks, see Nielsen (2015) Here, I use a neural network with 2 hidden
layers of 32 nodes in-between the input and output layer, which both have size
1. The total number of parameters is 1,153. The hidden layers are passed
through an exponential linear unit (ELU) activation function. A neural network
fit to real data is also shown in Fig. 3.

###  3.3 Coupled online learning101010All experiments were done in a Jupyter
notebook that can be launched via Binder from the GitHub repository:
https://github.com/raspstephan/Lorenz-Online. There, an interactive instance
of the notebook can be run in the cloud.

To mimic the situation in a real climate model where the parameterization
would first be pretrained offline on a traditional parameterization, super-
parameterization or coarse-grained dataset, a training dataset using the full
L96 equations but with different parameters was created: $F=7$, $h=2$,
$c=b=5$. The resulting, "wrong" data points along with the linear regression
and neural network parameterizations are also shown in Fig. 3.

Algorithm 1 outlines the workflow for coupled learning in the L96 framework.
There are several hyper-parameters. First, the time steps $\Delta
t_{\text{HR}}$ and $\Delta t_{\text{LR}}$. In the easiest case, they are the
same. However, more realistically, the HR model has a finer time step than the
ML-LR model. For the experiments here, I used $N=10$, i.e. $\Delta
t_{\text{ML}}=0.01$.

The experiments indicate that coupled learning works well in both cases (Fig.
4). One slight difference is that the learned linear regression intercept
parameter $b$ is slightly different from the reference in the case where the
HR time step is smaller. This is likely an indication that the linear
superposition assumption during the HR integration is not perfect. However,
the differences are very small.

Another hyper-parameter is the update frequency of the neural network $M$. The
experiments show that updating every time step causes the parameters to change
a lot every update step. This is likely because the batch, which has size $m$,
is only a small sample of the parameter space that is also potentially
correlated. To combat this, we can gather the features and targets over
several ML-LR model time steps before doing an update step. Here, I used
$M=10$. This results in significantly smoother parameter convergence. Another
potential advantage of updating only every few time steps is that the ML model
can evolve more freely, thereby covering a larger fraction of the state space.
At the extreme end, setting $M$ to an infinitely large value amounts to not
learning during the simulation and simply collecting the training data in
$\mathcal{F}$ and $\mathcal{T}$. Theoretically one could do this and then
train the ML parameterization offline after the simulation. However, with a
nudging time scale greater than the model time step (see discussion in the
next paragraph), the collected training data does not represent the true fit
(see experiments in Jupyter notebook).

Finally, the nudging time scale $\tau_{\text{nudging}}$ has an impact on the
fit. When the time scale is equal to the time step of the LR model, the HR
model will be fully pulled towards the LR state. This works well for the L96
model but for complex ESMs this nudging will likely throw the HR model too
much off its attractor because the LR+ML model will exhibit different
dynamics, particularly at the start of learning. It is also common practice to
use larger nudging time scales for GCMs (e.g. 24 h in Bretherton et al.,
2019). Weaker nudging, however, means that the HR and LR states at the
beginning of the time step are further apart, which introduces an error. If
the nudging time scale is too large, the ML targets will eventually become
meaningless (see sensitivity experiments in the accompanying Jupyter
notebook). For the experiments shown in Fig. 4,
$\tau_{\text{nudging}}=0.1=10\Delta t_{\text{LR}}$.

The same algorithm can be used to train much more complicated
parameterizations such as a neural network (Fig. 4b). The $X$–$B$ curve
gradually approaches the one learned offline using data generated with the
correct L96 parameters. One final note on the L96 experiments: I did not
exhaustively search for the best combination of hyper-parameters because the
L96 experiments mainly serve as proof-of-concept. For coupled learning in a
real modeling setup, the parameters are likely very different.

###  3.4 Purpose and limitations of L96 experiments

The L96 model, while commonly used to test parameterization and data
assimilation approaches, only represent a small fraction of the challenges
that algorithms are faced with in real GCMs. In particular, L96 does not
exhibit any of the issues that require a coupled learning approach in the
first place: an offline parameterization for the L96 model is stable and does
not show major biases. The purpose of demonstrating the method using the L96
model is mostly a sanity check. Having confirmed that coupled learning works
in this simple framework now gives us more confidence to try to apply it for
more complex systems.

##  4 Algorithms for online learning in the super-parameterization and 3D HR
frameworks

In this section, I will outline how coupled learning algorithms can be applied
to 3D CRMs and super-parameterized GCMs.

###  4.1 3D high-resolution models

The 3D HR case is similar to the L96 setup (algorithm 2). The key difference
is that the scale separation is not clearly defined as in L96 or SP but rather
downscaling (coarse-graining) and upscaling is required to get the HR state on
the LR model grid and, reversely, apply the forcing term, which is computed on
the LR model grid, in the HR model. Issues with this will be further discussed
in Section 5. The other difference between algorithms 1 and 2 is the way the
gradient update is computed. In the L96 case the features and targets are
stored in memory. This is unpractical for the HR setup since it requires
storing several 3D fields over several time steps. Rather, in algorithm 2 the
gradients are computed directly at each time step and collected in a single
gradient vector $\mathcal{G}$, which is then used to update the parameters
every $M$ steps. This also allows computing the gradients locally on each node
and then collecting them. The size of $\mathcal{G}$ is equal to the number of
network parameters and, therefore, manageable. There is also no explicit batch
size in this version of the algorithms. Rather, the batch size is implicitly
given by $M\times K$.

One major conceptual difference of the 3D HR case to SP (see below) lies in
what is actually learned by the neural network during coupled learning. In SP,
the CRM is purely responsible for clouds and turbulence while a 3D HR model
also evolves globally according to its own set of physics. What this means is
that the neural network essentially learns a sub-grid correction term that
compensates for everything(!) missing from the LR model dynamics and non-ML
physics in comparison to the HR model (HR${}^{\prime}\to$ LR′′ in Fig. 2). So
even if all parameterizations except for convection are present in the LR, the
network will not only learn convective tendencies. On the one hand, this is
exactly what is required to get the LR closer to the expensive HR simulation.
On the other hand, this makes the interpretation of what the network does a
little more complicated.

Algorithm 2 Online learning algorithm for 3D HR models. Bold-face vectors
indicate state vectors for all grid columns $k\in{1,\dots,K}$. Overbars denote
vectors on the coarse LR model grid.

0: Pretrained ML-parameterization $\mathcal{P}_{\theta}$ with parameters
$\theta$

0: Initial conditions on the HR grid $\mathbf{x}_{0}$

0: Downscaling and upscaling algorithms $\mathcal{D}$ and $\mathcal{U}$

0: HR model with time step $\Delta t_{\text{HR}}$

0: LR model with time step $\Delta t_{\text{LR}}=N\Delta t_{\text{HR}}$

0: Nudging time scale $\tau_{\text{nudging}}$

0: Gradient memory $\mathcal{G}=0$

0: Training frequency $M$; learning rate $\alpha$

Initialize HR with $\mathbf{x}_{0}$; initialize LR with
$\bar{\mathbf{x}}_{0}=\mathcal{D}(\mathbf{x}_{0})$

for $t$ = 1,… do

Store delta at beginning of time step
$\Delta\bar{\mathbf{x}}=\mathcal{D}(\mathbf{x}_{\text{HR}})-\bar{\mathbf{x}}_{\text{LR}}$

Store LR state at beginning of time step
$\hat{\mathbf{x}}_{\text{LR}}=\bar{\mathbf{x}}_{\text{LR}}$

Store HR state at beginning of time step
$\hat{\mathbf{x}}_{\text{HR}}=\mathbf{x}_{\text{HR}}$

for $n=1,N$ do

Integrate HR with nudging term
$-\mathcal{U}(\Delta\bar{\mathbf{x}})/\tau_{\text{nudging}}$ to get new
$\mathbf{x}_{\text{HR}}$

end for

Compute "assumed" HR increment $\Delta\mathbf{\bar{x}}_{\text{HR-
internal}}=\mathcal{D}(\mathbf{x}_{\text{HR}}-\hat{\mathbf{x}}_{\text{HR}}-\mathcal{U}(\Delta\mathbf{\bar{x}})/\tau_{\text{nudging}}*\Delta
t_{\text{LR}})$

Integrate LR model (only dynamics and non-ML physics) to get intermediate
$\bar{\mathbf{x}}_{\text{LR}}$

Compute target $\mathbf{\mathcal{T}}=\Delta\mathbf{\bar{x}}_{\text{HR-
internal}}-(\bar{\mathbf{x}}_{\text{LR}}-\hat{\mathbf{x}}_{\text{LR}})$

Compute loss for each column
$\mathcal{L}_{\theta,k}=(\mathcal{P}_{\theta}(\hat{x}_{\text{LR},k})-\mathcal{T}_{k})^{2}$

Store gradient:
$\mathcal{G}\leftarrow\mathcal{G}+\frac{1}{M}\frac{1}{K}\sum_{k}\nabla_{\theta}\mathcal{L}_{\theta}$

Add ML increment:
$\bar{\mathbf{x}}_{\text{LR}}\leftarrow\bar{\mathbf{x}}_{\text{LR}}+\mathcal{P}_{\theta}(\hat{\mathbf{x}}_{\text{LR}})$

if $t$ mod $M$ = 0 then

Update parameters $\theta$ using gradients $\mathcal{G}$ with rate $\alpha$

$\mathcal{G}=0$

end if

end for

Algorithm 3 Online learning algorithm for super-parameterized LR models. This
algorithm is specific to the SPCAM code structure. Note that the notation is
slightly different from algorithm 2: the LR model state $\mathbf{x}$ now does
not have an overbar and $\bar{x}_{\text{CRM}}$ denotes the averaged CRM state.

0: Pretrained ML-parameterization $\mathcal{P}_{\theta}$ with parameters
$\theta$

0: Initial conditions $\mathbf{x}_{0}$

0: Embedded SP-CRM with time step $\Delta t_{\text{CRM}}$

0: GCM with time step $\Delta t_{\text{GCM}}=N\Delta t_{\text{CRM}}$,
$N\in\mathbb{Z}^{+}$

0: Gradient memory $\mathcal{G}=0$

0: Training frequency $M$; learning rate $\alpha$

Initialize GCM with $\mathbf{x}_{0}$

Uniformly initialize each CRM grid column from $\mathbf{x}_{0}$

for $t$ = 1,… do

Call CRM but do not update GCM state; internally this computes and applies the
forcing term.

Loss
$\mathcal{L}_{\theta,k}=(x_{\text{GCM},k}+\mathcal{P}_{\theta}(x_{\text{GCM},k})-(\bar{x}_{\text{CRM}})_{k})^{2}$

Store gradient:
$\mathcal{G}\leftarrow\mathcal{G}+\frac{1}{M}\frac{1}{K}\sum_{k}\nabla_{\theta}\mathcal{L}_{\theta}$
(can be done within each MPI thread)

Add ML tendency:
$\mathbf{x}_{\text{GCM}}\leftarrow\mathbf{x}_{\text{GCM}}+\mathcal{P}_{\theta}(\mathbf{x}_{\text{GCM}})\Delta
t_{\text{GCM}}$

if $t$ mod $M$ = 0 then

Globally collect gradients $\mathcal{G}$

Update parameters $\theta$ using gradients $\mathcal{G}$ with rate $\alpha$

$\mathcal{G}=0$

end if

end for

Figure 5: Evolution of a tracer $\phi$ during a regular SP step on the left
and for coupled learning on the right.

###  4.2 Super-parameterization

Similar to L96, SP has the advantage of a clean scale-separation, which makes
the parameterization learning task easier. It also provides a good framework
for coupled learning since SP already has the LR model and the embedded CRMs
running in parallel. Because the embedded CRMs do not have any large-scale
dynamics on their own, the time step schematic in Fig. 5 looks different to
Fig. 2. In contrast to regular SP, the LR model state is not set to the CRM
state after the CRM integration. Instead, the LR model evolves on its own
according the the ML physics and the difference between CRM′ and LR′ is the
loss to minimize. Algorithm 3 describes coupled learning specifically for
super-parameterized models like SPCAM. The interactions between the LR model
and CRM are already contained in the CRM function call. This means that only
few changes to the code are required: the neural network forward and backward
passes have to be implemented, in addition to the optimizer and the
communication of the gradients between the threads.

##  5 Discussion

###  5.1 Which variables have to be forced/predicted by the neural network?

In the three original ML parameterization studies, of the prognostic
variables, only temperature and humidity were used in the input and output.
This was done to reduce the complexity of the problem to the fewest prognostic
variables necessary to produce a general circulation. In coupled learning, the
variables used by the ML parameterization also have to be forced in the HR
model. The HR model will typically have many more prognostic variables
compared to the LR model (e.g. hydrometeors) but it is alright for those to
evolve without forcing. In fact, this might be necessary since the HR and LR
models might have different prognostic variables. This is the case in SP where
only the LR model prognostic variables are forced during CRM integration. If
the variables predicted by the neural network differ, for example temperature
vs. moist static energy, an additional conversion step has to be added to the
up- and downscaling described below.

So theoretically coupled learning should work fine even if only temperature
and humidity are forced/predicted. However, there are reasons for going beyond
this. First, it is likely that the network skill suffers from not having
information about e.g. cloud water. We saw this in RPG18 where the network was
essentially unable to produce a shallow cloud heating signature in the sub-
tropics. Second, to implement physical constrains it is necessary to add more
variables in order to close the conservation budgets, which we will discuss
now.

###  5.2 Physical constraints

A major critique of machine learning and especially neural network
parameterizations is that they do not obey physical constraints. However,
Beucler et al. (2019) recently showed that it is possible to encode physical
constraints in neural networks if the conservation equations are known. There
are two ways of doing so: First, violation of constraints can be added to the
loss term during neural network training. This does not guarantee that the
constraints are exactly obeyed, particularly outside of the training regime,
but in practice might come close. The second method is to hard-code the
conservation constraints into the last layers of the neural network. This
ensures exact conservation and has been shown to only hurt skill of the
network slightly.

One downside of implementing physical constraints in Beucler et al. (2019) is
that it requires predicting all prognostic variables that occur in the
conservation budget equations. In effect, this increased the size of the
output vector from 65 in RPG18 to 218. This now also includes variables that
we might not actually care about like the snow storage term. Anecdotally, more
variables also means more potential for things to go wrong, e.g. instabilities
to develop. One possibility to reign in this complexity in offline and coupled
learning is to omit some of these terms from the output vector and simply set
them to zero in the budget equations. While this makes it impossible for the
network to exactly reproduce the target (where all terms of the budget
equation are used), this essentially forces the network to make the closest
prediction to the target that lies on its own manifold of physically
conserving solutions. If the omitted terms are small, this should still yield
good results.

When using coarse-grained HR output as training data as in BB18, the residuals
(Eq. 1) do not obey any conservation relations. In coupled learning, physical
constraints could still be encoded however. All one needs to know is the
budget equations valid on the LR model grid, i.e. the equations a traditional
parameterization would also obey. The network will then learn the best
physically conserving sub-grid correction term to bring the LR model closer to
the HR model.

###  5.3 Up- and downscaling

Another issue is how to convert 3D fields from the LR model to the HR grid and
vice versa. I already mentioned downscaling or coarse-graining along with some
issues in the context of discussing BB18. For coupled learning in the 3D HR
setup (Algorithm 2) a downscaling algorithm $\mathcal{D}$ is required to
transform the HR state $\mathbf{x}_{\text{HR}}$ to the LR model grid to
compute the ML targets. Upscaling $\mathcal{U}$ is used to apply the forcing
term, which is computed on the LR model grid, in the HR model. The simplest
method for downscaling is to simply average the HR values onto the LR model
grid and interpolate if necessary. In signal processing this is the equivalent
of applying a rectangular filter which potentially leads to aliasing. It might
be worth investigating common filtering methods, such as using a Gaussian low-
pass filter. 121212See https://dsp.stackexchange.com/questions/6313/low-pass-
filter-parameters-for-image-downsampling for a related discussion. For
upscaling, simply taking the LR model grid value that corresponds
geographically to each HR grid point will results in sharp boundaries for the
HR forcing field. A different way would be to use a smoother interpolation
function, for example a spline. In practice, how problematic sharp boundaries
in the forcing would be is hard to say without trying it out. Note also that
up- and downscaling is done in operational data assimilation, for example
4DVAR, where the adjoint model is run on a lower resolution.

###  5.4 Technical challenges

Depending on the setup, there are some daunting technical challenges for the
implementation of coupled learning. SPCAM represents the easiest case because
it already has the embedded CRMs running in parallel with the LR model with
coupling. The key challenge here would be the implementation of the neural
network forward and backward pass. We have already implemented the forward
pass in RPG18 by hard-coding it in Fortran. This works but is error-prone,
hard to debug and cumbersome. Backpropagation along with a modern gradient
descent algorithm like Adam (Kingma and Ba, 2014) would add to the complexity.
Another option is to call Python from Fortran131313see Noah Brenowitz’s blog
post: https://www.noahbrenowitz.com/post/calling-fortran-from-python/ but this
is potentially slow. Further, since the network parameters are global, the
gradient descent step has to happen globally as well requiring communication
between the nodes. The Python-Fortran interface currently is a major obstacle
in ML parameterization research that begs for a simpler solution.141414CLIMA
might be just that eventually: https://github.com/climate-machine/CLIMA; or
alternatively the Sympl and CliMT frameworks (Merwin Monteiro et al., 2018).

For the 3D HR setup, in addition to the neural network implementation and the
up-/downscaling issues, coupled learning requires two models to be run in
parallel communicating every few time steps. This potentially requires quite a
lot of engineering. My guess is that a successful and relatively quick
implementation of coupled learning requires extensive working knowledge with
the atmospheric models used.

###  5.5 How efficient is the online learning algorithm?

Running a HR model is expensive. Therefore, it is essential that the coupled
learning algorithm is efficient enough to learn from a limited number of
coupled HR simulations. To judge this, L96 is a bad toy model because it is so
far removed from the actual problem. On the one hand, the parameterization
task is exceedingly easy (one input, one output). On the other hand, it has 32
"LR" grid points while a 2-degree global LR model has more than 8,000,
yielding a much larger sample for each gradient descent update. Further, there
are a large range of hyper-parameters to tune. For a dry run, one could use a
network trained offline on a reference dataset and then simulate coupled
training by using a different, non-shuffled dataset (e.g. the +4K run from
RPG18). This should provide guidance for choosing hyper-parameters and give a
rough estimate of how many iterations are required.

Coupled learning is a potential method to combat some of the main obstacles in
ML parameterization research: instabilities and tuning. In this paper my aim
was to present the algorithms and challenges as clearly as possible and
demonstrate the general feasibility in the L96 case. The next step will be to
test coupled learning in a more realistic framework. Some open questions are:
How much weight should be given to new samples, particularly if the tendencies
are substantially chaotic? Are the HR and ML-LR model guaranteed to converge?
Will the linear superposition assumption break down if the forcing becomes too
large? How should situations be handled where the model crashes after all?
Finally, coupled learning can only fix short term prediction errors, which
raises the question to which degree this would lead to a decrease in long-term
biases.

There are a number of problems with ML parameterizations that coupled learning
cannot address. First and foremost for climate modeling: generalization, i.e.
the ability of a neural network parameterization to perform well outside its
training regime. Neural networks are essentially non-linear regression
algorithms and should not be expected to learn anything beyond what they have
encountered during training (Scher and Messori, 2019). The research area of
learning physical laws with deep learning is still in its infancy. For this
reason Schneider et al. (2017a) advocate sticking to physically motivated
parameterizations and improve the tuning process. Note that coupled learning
can still be used to tune parameters in existing parameterizations if they are
coded up in differentiable fashion.

Another issue unsolved by coupled learning is stochasticity. Any deterministic
ML model that minimizes a mean error will be unable to represent random
fluctuations in the training dataset. This leads to smoothed out predictions.
The case for stochastic parameterizations has been growing steadily (Berner et
al., 2015; Palmer, 2019) raising the question how stochasticity can be
incorporated into ML parameterizations. Two possible approaches could be using
generative adversarial networks (GANs; Subramanian et al., 2018) or using a
parametric distribution.151515Parametric approaches have been commonly used
for post-processing of NWP forecasts (Rasp and Lerch, 2018), however mostly
for single output tasks. Realistic multi-variate predictions need to take into
account covariances, which might require further research. How to combine
coupled learning with GANs, however, is not readily apparent.

Finally, high-resolution models might be better than coarse models but they
still are not the truth. Our best knowledge of the true behavior of the
atmosphere comes from observations. The problem is that observations are
intermittent in space and time and, in the case of remote sensing, indirect.
So how to learn from such data? Schneider et al. (2017a) propose a parameter
estimation approach using an ensemble Kalman inversion,a gradient free method
for parameter optimization (Garbuno-Inigo et al., 2019). The second best guess
of the truth would be a reanalysis, such as the
ERA5161616https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5
dataset, which provides 3D fields every 3 hours. It could well be worth
spending some thoughts on exploring how re-analyses could be used for ML
parameterization training.

Clouds are incredibly complex. No wonder then that we humans have such trouble
shoving them into mathematical concepts. We need any assistance we can get.
Could ML provide us with such? The verdict is still out. First studies show
that ML models are, in general, capable of representing sub-grid tendencies
but the way towards actually improving weather and climate models poses
several obstacles. Coupled online learning could be one potential solution out
of many to overcome some of these obstacles.

All code (version 1.0) is available here:
https://github.com/raspstephan/Lorenz-Online. The L96 experiments are all
contained in a single Jupyter notebook which anyone can launch and interact
with here: https://mybinder.org/v2/gh/raspstephan/Lorenz-
Online/master?filepath=coupled-learning.ipynb

The author declares no competing interests.

###### Acknowledgements.

I thank Chris Bretherton, Noah Brenowitz, Tapio Schneider, Sebastian Scher,
David John Gagne, Tom Beucler, Mike Pritchard and Pierre Gentine for their
valuable input. I acknowledge funding from the German Research Foundation,
partly through the project SFB/TRR 165 “Waves to Weather.”

## References

  * Berner et al. (2015) Berner, J., Fossell, K. R., Ha, S.-Y., Hacker, J. P., and Snyder, C.: Increasing the Skill of Probabilistic Forecasts: Understanding Performance Improvements from Model-Error Representations, Monthly Weather Review, 143, 1295–1320, 10.1175/MWR-D-14-00091.1, URL http://journals.ametsoc.org/doi/abs/10.1175/MWR-D-14-00091.1, 2015\.
  * Beucler et al. (2019) Beucler, T., Rasp, S., Pritchard, M., and Gentine, P.: Achieving Conservation of Energy in Neural Network Emulators for Climate Modeling, URL http://arxiv.org/abs/1906.06622, 2019.
  * Bocquet et al. (2019) Bocquet, M., Brajard, J., Carrassi, A., and Bertino, L.: Data assimilation as a learning tool to infer ordinary differential equation representations of dynamical models, Nonlinear Processes in Geophysics, 26, 143–162, 10.5194/npg-26-143-2019, URL https://www.nonlin-processes-geophys.net/26/143/2019/, 2019.
  * Bolton and Zanna (2019) Bolton, T. and Zanna, L.: Applications of Deep Learning to Ocean Data Inference and Subgrid Parameterization, Journal of Advances in Modeling Earth Systems, 11, 376–399, 10.1029/2018MS001472, URL http://doi.wiley.com/10.1029/2018MS001472, 2019.
  * Breiman (2001) Breiman, L.: Random Forests, Machine Learning, 45, 5–32, URL https://link.springer.com/content/pdf/10.1023/A:1010933404324.pdf, 2001\.
  * Brenowitz and Bretherton (2018) Brenowitz, N. D. and Bretherton, C. S.: Prognostic Validation of a Neural Network Unified Physics Parameterization, Geophysical Research Letters, 45, 6289–6298, 10.1029/2018GL078510, URL http://doi.wiley.com/10.1029/2018GL078510, 2018.
  * Brenowitz and Bretherton (2019) Brenowitz, N. D. and Bretherton, C. S.: Spatially Extended Tests of a Neural Network Parametrization Trained by Coarse-graining, Journal of Advances in Modeling Earth Systems, p. 2019MS001711, 10.1029/2019MS001711, URL https://onlinelibrary.wiley.com/doi/abs/10.1029/2019MS001711, 2019\.
  * Brenowitz et al. (2020) Brenowitz, N. D., Beucler, T., Pritchard, M., and Bretherton, C. S.: Interpreting and Stabilizing Machine-learning Parametrizations of Convection, URL http://arxiv.org/abs/2003.06549, 2020.
  * Bretherton et al. (2019) Bretherton, C. S., McCoy, I. L., Mohrmann, J., Wood, R., Ghate, V., Gettelman, A., Bardeen, C. G., Albrecht, B. A., and Zuidema, P.: Cloud, aerosol, and boundary layer structure across the northeast Pacific stratocumulus-cumulus transition as observed during CSET, Monthly Weather Review, 147, 2083–2103, 10.1175/MWR-D-18-0281.1, 2019.
  * Chevallier et al. (2000) Chevallier, F., Morcrette, J.-J., Chéruy, F., and Scott, N. A.: Use of a neural-network-based long-wave radiative-transfer scheme in the ECMWF atmospheric model, Quarterly Journal of the Royal Meteorological Society, 126, 761–776, 10.1002/qj.49712656318, URL http://doi.wiley.com/10.1002/qj.49712656318, 2000.
  * Crommelin and Vanden-Eijnden (2008) Crommelin, D. and Vanden-Eijnden, E.: Subgrid-Scale Parameterization with Conditional Markov Chains, Journal of the Atmospheric Sciences, 65, 2661–2675, 10.1175/2008JAS2566.1, URL http://journals.ametsoc.org/doi/abs/10.1175/2008JAS2566.1, 2008\.
  * Dueben and Bauer (2018) Dueben, P. D. and Bauer, P.: Challenges and design choices for global weather and climate models based on machine learning, Geosci. Model Dev., 10.5194/gmd-2018-148, URL https://www.geosci-model-dev-discuss.net/gmd-2018-148/gmd-2018-148.pdf, 2018\.
  * Garbuno-Inigo et al. (2019) Garbuno-Inigo, A., Hoffmann, F., Li, W., and Stuart, A. M.: Gradient Structure Of The Ensemble Kalman Flow With Noise 1. PROBLEM SETTING, URL https://arxiv.org/pdf/1903.08866.pdf, 2019.
  * Gentine et al. (2018) Gentine, P., Pritchard, M., Rasp, S., Reinaudi, G., and Yacalis, G.: Could Machine Learning Break the Convection Parameterization Deadlock?, Geophysical Research Letters, 45, 5742–5751, 10.1029/2018GL078202, URL http://doi.wiley.com/10.1029/2018GL078202, 2018.
  * Hourdin et al. (2017) Hourdin, F., Mauritsen, T., Gettelman, A., Golaz, J.-C., Balaji, V., Duan, Q., Folini, D., Ji, D., Klocke, D., Qian, Y., Rauser, F., Rio, C., Tomassini, L., Watanabe, M., and Williamson, D.: The Art and Science of Climate Model Tuning, Bulletin of the American Meteorological Society, 98, 589–602, 10.1175/BAMS-D-15-00135.1, URL http://journals.ametsoc.org/doi/10.1175/BAMS-D-15-00135.1, 2017\.
  * Khairoutdinov and Randall (2001) Khairoutdinov, M. F. and Randall, D. A.: A cloud resolving model as a cloud parameterization in the NCAR Community Climate System Model: Preliminary results, Geophysical Research Letters, 28, 3617–3620, 10.1029/2001GL013552, URL http://doi.wiley.com/10.1029/2001GL013552, 2001.
  * Kingma and Ba (2014) Kingma, D. P. and Ba, J.: Adam: A Method for Stochastic Optimization, arXiv, 1412.6980, URL http://arxiv.org/abs/1412.6980, 2014.
  * Krasnopolsky et al. (2013) Krasnopolsky, V. M., Fox-Rabinovitz, M. S., and Belochitski, A. A.: Using Ensemble of Neural Networks to Learn Stochastic Convection Parameterizations for Climate and Numerical Weather Prediction Models from Data Simulated by a Cloud Resolving Model, Advances in Artificial Neural Systems, 2013, 1–13, 10.1155/2013/485913, URL https://www.hindawi.com/archive/2013/485913/, 2013.
  * Lorenz (1995) Lorenz, E. N.: Predictability: a problem partly solved, in: Seminar on Predictability, 4-8 September 1995, vol. 1, pp. 1–18, ECMWF, ECMWF, Shinfield Park, Reading, URL https://www.ecmwf.int/node/10829, 1995\.
  * Merwin Monteiro et al. (2018) Merwin Monteiro, J., McGibbon, J., and Caballero, R.: Sympl (v. 0.4.0) and climt (v. 0.15.3) - Towards a flexible framework for building model hierarchies in Python, Geoscientific Model Development, 11, 3781–3794, 10.5194/gmd-11-3781-2018, URL https://doi.org/10.5194/gmd-11-3781-2018, 2018.
  * Nielsen (2015) Nielsen, M. A.: Neural Networks and Deep Learning, Determination Press, URL http://neuralnetworksanddeeplearning.com, 2015.
  * O’Gorman and Dwyer (2018) O’Gorman, P. A. and Dwyer, J. G.: Using Machine Learning to Parameterize Moist Convection: Potential for Modeling of Climate, Climate Change, and Extreme Events, Journal of Advances in Modeling Earth Systems, 10, 2548–2563, 10.1029/2018MS001351, URL http://doi.wiley.com/10.1029/2018MS001351, 2018.
  * Palmer (2019) Palmer, T. N.: Stochastic weather and climate models, Nature Reviews Physics, p. 1, 10.1038/s42254-019-0062-2, URL http://www.nature.com/articles/s42254-019-0062-2, 2019.
  * Pathak et al. (2018) Pathak, J., Hunt, B., Girvan, M., Lu, Z., and Ott, E.: Model-Free Prediction of Large Spatiotemporally Chaotic Systems from Data: A Reservoir Computing Approach, Physical Review Letters, 120, 24 102, 10.1103/PhysRevLett.120.024102, URL https://doi.org/10.1103/PhysRevLett.120.024102, 2018.
  * Rasp and Lerch (2018) Rasp, S. and Lerch, S.: Neural Networks for Postprocessing Ensemble Weather Forecasts, Monthly Weather Review, 146, 3885–3900, 10.1175/MWR-D-18-0187.1, URL http://journals.ametsoc.org/doi/10.1175/MWR-D-18-0187.1, 2018\.
  * Rasp et al. (2018) Rasp, S., Pritchard, M. S., and Gentine, P.: Deep learning to represent subgrid processes in climate models., Proceedings of the National Academy of Sciences of the United States of America, 115, 9684–9689, 10.1073/pnas.1810286115, URL http://www.ncbi.nlm.nih.gov/pubmed/30190437http://www.pubmedcentral.nih.gov/articlerender.fcgi?artid=PMC6166853, 2018.
  * Saffari et al. (2009) Saffari, A., Leistner, C., Santner, J., Godec, M., and Bischof, H.: On-line Random Forests, in: 2009 IEEE 12th International Conference on Computer Vision Workshops, ICCV Workshops, pp. 1393–1400, IEEE, 10.1109/ICCVW.2009.5457447, URL http://ieeexplore.ieee.org/document/5457447/, 2009.
  * Scher and Messori (2019) Scher, S. and Messori, G.: Generalization properties of neural networks trained on Lorenzsystems, Nonlinear Processes in Geophysics Discussions, pp. 1–19, 10.5194/npg-2019-23, URL https://www.nonlin-processes-geophys-discuss.net/npg-2019-23/, 2019\.
  * Schneider et al. (2017a) Schneider, T., Lan, S., Stuart, A., and Teixeira, J.: Earth System Modeling 2.0: A Blueprint for Models That Learn From Observations and Targeted High-Resolution Simulations, Geophysical Research Letters, 44, 396–12, 10.1002/2017GL076101, URL http://doi.wiley.com/10.1002/2017GL076101, 2017a.
  * Schneider et al. (2017b) Schneider, T., Teixeira, J., Bretherton, C. S., Brient, F., Pressel, K. G., Schär, C., and Siebesma, A. P.: Climate goals and computing the future of clouds, Nature Climate Change, 7, 3–5, 10.1038/nclimate3190, URL http://www.nature.com/articles/nclimate3190, 2017b.
  * Subramanian et al. (2018) Subramanian, A. C., Gagne, D. J., I., Christensen, H. M., and Monahan, A. H.: Evaluating Generative Adversarial Network Stochastic Parameterizations of the Lorenz ’96 Model at Climate and Weather Scales, American Geophysical Union, Fall Meeting 2018, URL http://adsabs.harvard.edu/abs/2018AGUFMIN14A..02S, 2018.
  * Wilks (2005) Wilks, D. S.: Effects of stochastic parametrizations in the Lorenz ’96 system, Quarterly Journal of the Royal Meteorological Society, 131, 389–407, 10.1256/qj.04.03, URL http://doi.wiley.com/10.1256/qj.04.03, 2005.
  * Yuval and O’Gorman (2020) Yuval, J. and O’Gorman, P. A.: Use of machine learning to improve simulations of climate, URL http://arxiv.org/abs/2001.03151, 2020.


# Entanglement enhancement in cavity magnomechanics by an optical parametric
amplifier

Bakht Hussain  Department of Physics and Applied Mathematics, Pakistan
Institute of Engineering and Applied Sciences (PIEAS), Nilore $45650$,
Islamabad, Pakistan.  Shahid Qamar  Department of Physics and Applied
Mathematics, Pakistan Institute of Engineering and Applied Sciences (PIEAS),
Nilore $45650$, Islamabad, Pakistan.  Center for Mathematical Sciences, PIEAS,
Nilore, Islamabad $45650$, Pakistan.  Muhammad Irfan  Department of Physics
and Applied Mathematics, Pakistan Institute of Engineering and Applied
Sciences (PIEAS), Nilore $45650$, Islamabad, Pakistan.  Center for
Mathematical Sciences, PIEAS, Nilore, Islamabad $45650$, Pakistan.

###### Abstract

We propose a method to enhance bipartite and tripartite entanglement in cavity
magnomechanics using an optical parametric amplifier (OPA). We analyze this
system and identified parametric regimes where different types of entanglement
are enhanced. We show that, a proper choice of phase of the parametric
amplifier leads to the enhancement of the bipartite entanglements. Moreover,
the tripartite entanglement is also significantly enhanced in the presence of
OPA. The OPA not only enhances the strength of entanglement but also increases
the domain of entanglement over a wider space of detunings as compared to the
system when no OPA is present. Similarly, the robustness of entanglement
against temperature is also enhanced. Another important consequence of OPA is
the fact that it relaxes the requirement of strong magnon-phonon coupling to
generate cavity-magnon entanglement which is necessary for the case when OPA
is not present. We believe that the presented scheme is a step forward to
realize robust quantum entanglement using current technology.

##  I Introduction

Cavity magnomechanics, where a ferrimagnetic crystal (e.g., yttrium iron
garnet (YIG)) is coupled with a microwave cavity has drawn considerable
attention during last few years. It provides a unique platform to study
interactions among cavity, magnon, and phonon modes. Magnetic materials
generally have high spin density and low dissipation rate. This results in
strong [1, 2, 3, 4, 5, 6] and ultrastrong [7, 8] couplings among the magnon or
Kittel mode [9] (collective excitation of the spins) and a microwave cavity
mode. This strong coupling enables the coherent information transfer between
different excitations which has a wide range of applications in quantum
information processing.

Macroscopic quantum entanglement plays an important role not only in quantum
information processing but also to understand fundamental theories. In this
regard, cavity optomechanics attracted a lot of attention during past decade
[10]. Vitali et al. studied the existence of quantum entanglement between the
cavity mirror and the cavity field [11]. Recent experiments showed promising
results where entanglement between cavity-phonon [12, 13] and two mechanical
oscillators [14, 15] is demonstrated. Such macroscopic entangled states are
proposed to help understand fundamental issues of modern physics like quantum-
gravity [16], quantum-to-classical transition [17], testing models of wave-
function collapse [18], etc. On applications side, such macroscopic quantum
systems allow precision measurement and sensing, better quantum memory, and
single photon detection [10, 19].

Recently, in a seminal work, Li et al., [20] showed that similar to
optomechanics, cavity magnomechanics also offers possibility to generate
macroscopic entanglement. This system supports both bipartite and tripartite
entanglement. Subsequent studies discussed entanglement between two magnon
modes [21, 22, 23], two vibrational modes of YIG spheres [24], steady Bell
state generation [25], among others. Later, it is shown that entanglement can
be further enhanced by including Kerr effect [26, 27]. In a very recent study,
it is shown that the presence of injected weak squeezed vacuum field [28]
enhances entanglement and one-way steering, especially in the red-detuned
driving regime [29]. Similarly, it is shown that the resonant interaction of
the antiferromagnet with light enhances the magnon-magnon entanglement [30].

In this study, we consider the effects of a degenerate optical parametric
amplifier on entanglement properties of the cavity magnomechanical system. We
show that the presence of the optical parametric amplifier (OPA) enhances the
bipartite as well as tripartite entanglement for properly chosen parameters.
Apart from enhancement, it also widens the parameter space where entanglement
can exist. Therefore, makes the system more feasible towards experimental
realization. Moreover, entanglement in the presence of OPA is more robust
against the thermal excitations. Li et al., showed that the bipartite
entanglement between cavity and magnon modes exist only for a non-zero value
of magnon-phonon coupling [20] and maximum entanglement is achieved when the
coupling is sufficiently strong. We found that remarkably, this condition is
no longer needed in the presence of OPA and strong cavity-magnon entanglement
exists even for zero magnon-phonon coupling.

Figure 1:  (a) Schematic diagram of a single-mode cavity with OPA and a YIG
sphere. YIG sphere is placed in cavity such that it is simultaneously near the
maximum magnetic field of the cavity mode and in a uniform bias field, which
is responsible for magnon-photon coupling. A microwave field (not shown) is
applied to enhance magnon-phonon coupling. At YIG sphere site, the magnetic
field (along $x$ axis) of the cavity mode, the drive magnetic field (in $y$
direction), and bias magnetic field ($z$ direction) are mutually
perpendicular. (b) System frequencies and linewidths. A microwave field at
frequency $\omega_{0}$ drives the magnon mode $m$ with frequency $\omega_{m}$
and bandwidth $\kappa_{m}$, and the photons scatter by mechanical motion at
frequency $\omega_{b}$ onto the two sidebands at $\omega_{0}\pm\omega_{b}$.

##  II system model and Hamiltonian

We consider a hybrid cavity-magnomechanical system [31] which consists of a
single-mode cavity of frequency $\omega_{c}$ with an OPA and a YIG sphere as
shown in Fig. 1. This system supports three different types of excitation
namely, photons, magnons, and phonons. The magnon mode is due to the
collective motion of a large number of spins in a ferrimagnetic system e.g., a
YIG sphere with a typical diameter of $250\,{\mu}m$ [31]. The coupling between
the cavity mode and magnon mode is mediated by the magnetic dipole. The
magnetostrictive interaction leads to the coupling of magnons and phonons. The
phonon mode arises due to the geometric deformation of YIG sphere. The
deformation is caused because of the magnon excitation inside the YIG sphere
which induces a varying magnetization [32].

The Hamiltonian of the system under rotating-wave approximation in a frame
rotating with the frequency of the drive field is given by:

$\displaystyle H/\hbar=$
$\displaystyle\Delta_{c}c^{{\dagger}}c+\Delta_{m}m^{{\dagger}}m+\frac{\omega_{b}}{2}(q^{2}+p^{2})+g_{mb}m^{{\dagger}}mq$
(1)
$\displaystyle+g_{mc}(c^{{\dagger}}m+cm^{{\dagger}})+iG(e^{i\theta}c^{{\dagger}2}-e^{-i\theta}c^{2})$
$\displaystyle+i\Omega(m^{{\dagger}}-m).$

The first two terms in Eq. (1) describe the energy of the bare cavity and
magnon modes, respectively with annihilation operators $c$ and $m$ and
creation operators $c^{\dagger}$ and $m^{\dagger}$. The detuning of the cavity
frequency $\omega_{c}$ from drive field is $\Delta_{c}=\omega_{c}-\omega_{0}$,
while of magnon mode frequency $\omega_{m}$ is
$\Delta_{m}=\omega_{m}-\omega_{0}$. The third term represents the mechanical
mode of frequency $\omega_{b}$ (phonon mode), where $q$ and $p$ are
dimensionless position and momentum quadratures. The interaction between the
magnon and phonon modes is characterized by the magnomechanical coupling rate
$g_{mb}$ which is typically small. However, applying a strong drive microwave
field enhances this rate. Magnon and microwave cavity coupling is described by
the fifth term with magnon-cavity coupling rate $g_{mc}$. The condition for
strong coupling is: $g_{mc}>\kappa_{c},\kappa_{m}$ [1, 2, 3, 4, 5, 6], where
$\kappa_{c}$ and $\kappa_{m}$ are dissipation rates of cavity and magnon
modes, respectively. The parameter $G$ represents the nonlinear gain of the
OPA with $\theta$ the phase of the driving field. A microwave field of
amplitude $B_{0}$ and frequency $\omega_{0}$ drives the magnon mode. The
strength of the drive is characterized by the Rabi frequency
$\Omega=\frac{\sqrt{5}}{4}\gamma\sqrt{N}B_{0}$ [20], where
$\gamma/2\pi=28~{}GHz/T$ is gyromagnetic ratio and $N=\rho{V}$ is the total
number of spins with YIG sphere spin density $\rho=4.22\times 10^{27}m^{-3}$
having volume $V$. The derivation of $\Omega$ assumes low-lying excitations
i.e., ${\langle}m^{{\dagger}}m\rangle{\ll}2Ns$, where ground state spin number
of $Fe^{3+}$ ion in YIG is $s=\frac{5}{2}$. Here, we neglect the effects of
the radiation pressure because the wavelength of the microwave field is
assumed to be much larger than the size of the sphere. Similarly, the
rotating-wave approximation is justified under the condition: $\omega_{c}$,
$\omega_{m}\gg g_{mc}$, $\kappa_{c}$, $\kappa_{m}$ [31]. The quantum Langevin
equations (QLEs) [33] describing the system are given by:

$\displaystyle\dot{c}=$
$\displaystyle-(i\Delta_{c}+\kappa_{c})c-ig_{mc}m+2Ge^{i\theta}c^{\dagger}+\sqrt{2\kappa_{c}}c^{in},$
$\displaystyle\dot{m}=$
$\displaystyle-(i\Delta_{m}+\kappa_{m})m-ig_{mc}c-ig_{mb}mq+\Omega+\sqrt{2\kappa_{m}}m^{in},$
$\displaystyle\dot{q}=$ $\displaystyle~{}\omega_{b}p,$ $\displaystyle\dot{p}=$
$\displaystyle-\omega_{b}q-\gamma_{b}p-g_{mb}m^{{\dagger}}m+\xi,$ (2)

with $\gamma_{b}$ the mechanical damping rate, whereas $c^{in}$, $m^{in}$ and
$\xi$ are input noise operators for cavity, magnon, and mechanical modes,
respectively with zero mean and characterized by correlation functions [34]:
${\langle}c^{in}(t)c^{in{\dagger}}(\acute{t})\rangle=[N_{c}(\omega_{c})+1]\delta(t-\acute{t})$,
${\langle}c^{in{\dagger}}(t)c^{in}(\acute{t})\rangle=N_{c}(\omega_{c})\delta(t-\acute{t})$,
and
${\langle}m^{in}(t)m^{in{\dagger}}(\acute{t})\rangle=[N_{m}(\omega_{m})+1]\delta(t-\acute{t})$,
${\langle}m^{in{\dagger}}(t)m^{in}(\acute{t})\rangle=N_{m}(\omega_{m})\delta(t-\acute{t})$,
and
${\langle}\xi(t)\xi(\acute{t})+\xi(\acute{t})\xi(t)\rangle/2\simeq\gamma_{b}[2N_{b}(\omega_{b})+1]\delta(t-\acute{t})$,
where a Markovian approximation has been made, which is valid for a large
mechanical quality factor $Q=\frac{\omega_{b}}{\gamma_{b}}{\gg}1$ [35, 33].
$N_{j}(\omega_{j})=[\exp{({\hbar\omega_{j}}/{k_{B}T})}-1]^{-1}\quad(j=c,m,b)$
are equilibrium mean number of thermal photon, magnon, and phonon,
respectively.

We can linearize the dynamics of the system if we strongly drive the magnon
mode such that $\left|{\langle}m\rangle\right|\gg 1$. Similarly, a strong
cavity-magnon interaction leads to large cavity field amplitude
$\left(\left|{\langle}c\rangle\right|\gg 1\right)$. Therefore, we can write an
operator in terms of its steady-state value and a fluctuation part upto first
order: $O={\langle}O\rangle+{\delta}O\quad(O=c,m,q,p)$, where
${\langle}O\rangle$ is steady-state mean value and ${\delta}O$ is a small
fluctuation operator with zero mean value. Inserting, the operator forms in
Eq. (II), we obtain a set of steady-state values which are nonlinear algebraic
equations and a set of equations for fluctuation operators. We write the
fluctuation operator in the form of quadrature fluctuations:

$\displaystyle
f(t)^{T}=[{\delta}x_{1}(t),{\delta}y_{1}(t),{\delta}x_{2}(t),{\delta}y_{2}(t),{\delta}q,{\delta}p],$
(3)

with ${\delta}x_{1}=({\delta}c+{\delta}c^{{\dagger}})/{\sqrt{2}}$,
${\delta}y_{1}=i({\delta}c^{{\dagger}}-{\delta}c)/{\sqrt{2}}$,
${\delta}x_{2}=({\delta}m+{\delta}m^{{\dagger}})/{\sqrt{2}}$, and
${\delta}y_{2}=i({\delta}m^{{\dagger}}-{\delta}m)/{\sqrt{2}}$. The resulting
set of equations for quadrature fluctuations are given by:

$\displaystyle\dot{f(t)}=Af(t)+\sigma(t),$ (4)

where

$\displaystyle\sigma(t)^{T}=[\sqrt{2\kappa_{c}}x_{1}^{in}(t),\sqrt{2\kappa_{c}}y_{1}^{in}(t),\sqrt{2\kappa_{m}}x_{2}^{in}(t),\sqrt{2\kappa_{c}}y_{2}^{in}(t),0,\xi(t)],$
(5)

is the noise sources column vector, and $A$ is the drift matrix which is given
by:

$\displaystyle
A=\begin{pmatrix}-\kappa_{c}+2G\cos{\theta}&\Delta_{c}+2G\sin{\theta}&0&g_{mc}&0&0\\\
-\Delta_{c}+2G\sin{\theta}&-\kappa_{c}-2G\cos{\theta}&-g_{mc}&0&0&0\\\
0&g_{mc}&-\kappa_{m}&\tilde{\Delta}_{m}&-G_{mb}&0\\\
-g_{mc}&0&-\tilde{\Delta}_{m}&-\kappa_{m}&0&0\\\ 0&0&0&0&0&\omega_{b}\\\
0&0&0&G_{mb}&-\omega_{b}&-\gamma_{b}\end{pmatrix}.$ (6)

The parameter $\tilde{\Delta}_{m}=\Delta_{m}+g_{mb}{\langle}q\rangle$ is
effective magnon-drive detuning including the frequency shift due to the
magnomechanical interaction, where
${\langle}q\rangle=-(g_{mb}/{\omega_{b}})\left|{\langle}m\rangle\right|^{2}$
is the steady-state position quadrature value. The effective magnomechanical
coupling rate is $G_{mb}=i\sqrt{2}g_{mb}{\langle}m\rangle$, where the steady-
state value of magnon amplitude ${\langle}m\rangle$ read:

$\displaystyle{\langle}m\rangle={\frac{\Omega[(i\Delta_{c}+\kappa_{c})-2Ge^{i\theta}]}{g_{mc}^{2}+(i\tilde{\Delta}_{m}+\kappa_{m})[(i\Delta_{c}+\kappa_{c})-2Ge^{i\theta}]}}.$
(7)

The drift matrix $A$ (Eq. (6)) is written under the conditions:
$\left|{\tilde{\Delta}_{m}}\right|,\,\left|\Delta_{c}\right|\gg\kappa_{c},\,\kappa_{m},G$.
Since $\tilde{\Delta}_{m}$ contains $\left|{\langle}m\rangle\right|^{2}$, Eq.
(7) is intrinsically nonlinear. To reach a unique steady-state solution, the
system should be stable. The stability of the system is ensured by enforcing
the Routh-Hurwitz criteria [36], which states that the system is stable if the
real parts of all the eigenvalues of the drift matrix are negative. We have
ensured this condition in our numerical analysis.

Since, the nature of the quantum noises is assumed to be Gaussian with zero
mean and we are dealing with the linearized dynamics of the quantum
fluctuations, the steady-state of the quantum fluctuations is a continuous
variable tripartite Gaussian state. This state is fully characterized by a
$6\times 6$ covariance matrix (CM) $\mathcal{C}$, whose entries are defined as
$\mathcal{C}_{ij}=\frac{1}{2}{\langle}f_{i}(t)f_{j}(\acute{t})+f_{j}(\acute{t})f_{i}(t)\rangle\,(i,j=1,2,...,6)$.
The solution of the following Lyapunov equation (for a stable system) then
gives the covariance matrix [11]:

$\displaystyle A\mathcal{C}+\mathcal{C}A^{T}=-D,$ (8)

with diffusion matrix
$D=diag[\kappa_{c}(2N_{c}+1),\kappa_{c}(2N_{c}+1),\kappa_{m}(2N_{m}+1),\kappa_{m}(2N_{m}+1),0,\gamma_{b}(2N_{b}+1)]$
defined through
${\langle}\sigma_{i}(t)\sigma_{j}(\acute{t})+\sigma_{j}(\acute{t})\sigma_{i}(t)\rangle/2=D_{ij}\delta(t-\acute{t})$.
Analytical solution of the Lyapunov equation for covariance matrix is
cumbersome, therefore, we solve it numerically. To quantify entanglement
between any two modes of the system, we use logarithmic negativity $E_{N}$
[20, 27, 37, 38]:

$\displaystyle E_{N}=\max{\left[0,-\ln{2\tilde{\eta}}\right]},$ (9)

with $\tilde{\eta}$ the minimum symplectic eigenvalue of the reduced
covariance matrix of the two sub-systems of interest [20]. For tripartite
entanglement, we use minimal residual contangle given by [20, 27, 39, 40]:

$\displaystyle
R_{\tau}^{min}=min[R_{\tau}^{c|mb},R_{\tau}^{m|cb},R_{\tau}^{b|cm}],$ (10)

with residual contangle given by:

$\displaystyle R_{\tau}^{i|jk}=C_{i|jk}-C_{i|j}-C_{i|k}\geq 0.$ (11)

The residual contangle is a continuous variable analogous of tangle used for
discrete-variable tripartite entanglement [41]. In Eq. (11), the notation
$C_{u|v}$ represents the contangle of sub-systems $u$ and $v$ ($v$ contains
one or two modes). For the presence of genuine tripartite entanglement in the
system, there must be nonzero minimum residual contangle $R_{\tau}^{min}>0$.
Further detail on calculation of $E_{N}$ and $R_{\tau}$ is given in the
supplementary material of Ref. [20].

Figure 2: Density plot of bipartite entanglement between cavity and magnon
modes $(E_{cm})$ versus dimensionless detunings $\Delta_{c}/{\omega_{b}}$ and
$\tilde{\Delta}_{m}/{\omega_{b}}$ (a) $G=0$, (b) - (d) $G=3.31\times
10^{6}~{}s^{-1}.$ The phase $\theta$ associated with OPA drive field is
assumed to be (b) $0$, (c) $\pi/2$, and (d) $\pi$, respectively.

##  III Results and Discussion

We now present results of our numerical simulations and show how the presence
of OPA not only enhances entanglement but also supports existence of
entanglement in a comparatively broader range of parameter space. This same
system without OPA was recently considered by Jie Li et al. [20]. Therefore,
we mostly use the same experimental parameters which they considered and also
reproduce some of their results for a comparison. We adopt the following
parameters [31]: $\omega_{c}/{2\pi}=\omega_{m}/{2\pi}=10\,$GHz,
$\omega_{b}/{2\pi}=10\,$MHz, $\gamma_{b}/{2\pi}=10^{2}\,$Hz,
$\kappa_{c}/{2\pi}=\kappa_{m}/{2\pi}=1\,$MHz,
$g_{mc}/{2\pi}=G_{mb}/{2\pi}=3.2\,$MHz, and temperature $T=10\,$mK. Once, we
choose these parameters, we found that the maximum value of $G$ which
satisfies the stability criteria is $G=3.31\times 10^{6}~{}s^{-1}$. We noticed
that the value of phase $\theta$ associated with OPA drive field also plays an
important role on the dynamics of the system, therefore, we choose three
different values of $\theta$ in our numerical simulations.

Figure 3: Density plot of bipartite entanglement between cavity and phonon
modes $(E_{cb})$ versus dimensionless detunings $\Delta_{c}/{\omega_{b}}$ and
$\tilde{\Delta}_{m}/{\omega_{b}}$ (a) $G=0$, (b)-(d) $G=3.31\times
10^{6}~{}s^{-1}.$ The phase $\theta$ associated with OPA drive field is
assumed to be (b) $0$, (c) $\pi/2$, and (d) $\pi$, respectively.

In Fig. (2), we present the bipartite entanglement $E_{cm}$ between the cavity
and magnon mode as a function of dimensionless detunings
$\Delta_{c}/{\omega_{b}}$ and $\tilde{\Delta}_{m}/{\omega_{b}}$ in the
presence as well as absence of OPA. We reproduce the entanglement result when
OPA is not present [20] in Fig. 2(a). In the presence of OPA, $E_{cm}$
enhances significantly as shown in Figs. 2(b) - (d). We note that the phase of
the drive field plays a crucial role as shown in Figs. 2(b) - (d), where the
maximum enhancement is obtained for $\theta=\pi$. While entanglement is
enhanced roughly two times in the presence of OPA (See Fig. 2(d)), the
parameter space where entanglement survives is also significantly increased.
This is an important feature which makes the experimental generation and
detection of entanglement more feasible. The enhancement in entanglement is
due to the fact that the presence of OPA enhances the cavity occupation number
and changes the photon statistics of the cavity mode thereby improving cavity-
magnon coupling. As a result, the steady-state value of the magnon mode (See
Eq. (7)) also depends on parametric gain $G$ and associated phase $\theta$.
Therefore, the phase associated with the drive field also plays an important
role. A proper choice of phase may lead to a maximum noise suppression thereby
resulting maximum entanglement. It is important to note that the effects of
OPA are previously considered in cavity optomechanics showing enhancement in
entanglement [42, 43, 44] and cooling [45].

Figure 4: Density plot of bipartite entanglement between phonon and magnon
modes $(E_{mb})$ versus dimensionless detunings $\Delta_{c}/{\omega_{b}}$ and
$\tilde{\Delta}_{m}/{\omega_{b}}$ (a) $G=0$, (b)-(d) $G=3.31\times
10^{6}~{}s^{-1}.$ The phase $\theta$ associated with OPA drive field is
assumed to be (b) $0$, (c) $\pi/2$, and (d) $\pi$, respectively.

Next, we show the effects of OPA on the bipartite entanglement between the
photons and phonons $E_{cb}$. Again, we plot $E_{cb}$ in the absence of OPA in
Fig. 3(a). The presence of the OPA enhances entanglement significantly. The
maximum enhancement occurs when the phase of the drive field is zero. In Fig.
4, we present the effects of OPA on the bipartite entanglement between the
magnon and phonon modes $E_{mb}$ of the YIG sphere. Because the magnon mode
and phonon mode are directly coupled within the YIG sphere, we do not see any
appreciable effect on $E_{mb}$ due to the presence of OPA.

Figure 5: Density plot of bipartite entanglement between cavity and magnon
modes $(E_{cm})$ versus dimensionless detuning $\Delta_{c}/{\omega_{b}}$ and
dimensionless magnon-phonon coupling $G_{mb}/g_{mc}$. (a) $G=0$, (b)-(d)
$G=3.31\times 10^{6}s^{-1}.$ The phase $\theta$ associated with OPA drive
field is (b) $0$, (c) $\pi/2$, and (d) $\pi$. In all cases, we chose
$\tilde{\Delta}_{m}=0.9\omega_{b}$.

In Fig. 5, we show the density plots of bipartite cavity-magnon entanglement
$E_{cm}$ versus $\Delta_{c}$ and $G_{mb}$ at fixed value of
$\tilde{\Delta}_{m}=0.9\omega_{b}$. In the absence of an OPA (i.e, $G=0$) and
magnon-phonon coupling, cavity and magnon modes interact via linear cavity-
magnon beam splitter interaction. It is well know that such interaction does
not lead to entanglement [20]. Therefore, cavity-magnon entanglement exists
only if there is a finite coupling between the magnons and phonons (a non-
linear magnetostrictive interaction) [See Fig. 5(a)]. The maximum entanglement
exists only if this coupling is comparable to the magnon-cavity coupling i.e,
$G_{mb}\approx g_{mc}$. Remarkably, this requirement is no longer valid if OPA
is present in the cavity. Instead we found that due to the presence of OPA,
even stronger entanglement is present when $G_{mb}=0$ [See Figs. 5(b)-(d)] as
compared to the $G_{mb}\approx g_{mc}$ case when there was no OPA. Figs.
5(b)-(d) show that entanglement enhances for all the three choices of OPA
phase with a maximum enhancement at $\theta=\pi$. Moreover, entanglement is
now present for a significantly wider range of $\Delta_{c}$, which is an
important improvement from experimental point of view. The presence of OPA
changes the photon statistics by generating squeezed states of light, which
then translate to magnon mode. The presence of squeezed states and non-linear
gain of OPA generate the cavity-magnon entanglement even when magnon-phonon
coupling is absent, which is quite interesting.

Figure 6: (a) Cavity-magnon and (b) Cavity-phonon entanglement as a function
of temperature in the presence ($G=3.31\times 10^{6}s^{-1}$) and absence
($G=0$) of OPA. The values of $\Delta_{c}/{\omega_{b}}$ and
$\tilde{\Delta}_{m}/{\omega_{b}}$ for each curve is optimized from Fig. 2 and
Fig. 3.

Interestingly, the presence of OPA also enhances the robustness of
entanglement against temperature as shown in Fig. 6. In Fig. 6(a), we plot
cavity-magnon entanglement as a function of temperature in the presence and
absence of OPA. We chose the optimum values of detunings $\Delta_{c}$ and
$\tilde{\Delta}_{m}$ from Fig. 2. Therefore, each curve has different optimum
values corresponding to the respective panel in Fig. 2. In the presence of
OPA, the robustness of entanglement against temperature enhances
significantly. The solid curve in Fig. 6(a) shows that entanglement vanishes
($E_{cm}=0$) at a temperature of 0.17 K in the absence of OPA. At this same
temperature, the dotted curve which is plotted for $\theta=\pi$ in the
presence of OPA, shows an entanglement value of $E_{cm}\approx 0.11$ which is
comparable to the non-OPA entanglement value of $E_{cm}\approx 0.13$ at
$T\approx 0$. Therefore, the presence of OPA facilitates strong entanglement
at a much higher temperature. This illustrates the importance of OPA for a
robust cavity-magnon entanglement. The dashed-dotted curve for $\theta=\pi/2$
also shows a similar behaviour where $E_{cm}\approx 0.08$ at 0.17 K. The
dashed curve for $\theta=0$ also shows robustness but it is relatively less
pronounced. For cavity-phonon entanglement (See Fig. 6(b)), we found that for
$\theta=0$ (dashed curve) and $\theta=\pi/2$ (dashed-dotted curve),
entanglement is slightly more robust as compared to the case when OPA is not
present (solid curve). The robustness decreases when the phase $\theta=\pi$.
This is also evident from the fact that entanglement at $T\approx 0$ for this
case is already smaller than the entanglement when no OPA is present.

One of the most important feature of cavity magnomechanics is the possibility
of generation of genuine tripartite entanglement in this system. Li et al.,
[20] showed that genuine tripartite magnon-phonon-photon entanglement exists
in the system, if the magnon mode is in resonance with anti-Stokes (blue
sideband) and the cavity mode is in resonance with Stokes (red sideband).
Next, we consider the effects of OPA on tripartite entanglement. In order to
study this, we plot the minimum of the residual contangle in Fig. 7 as a
function of detuning $\Delta_{c}$. Here, we use $G_{mb}/2\pi=4.8~{}MHz$ [20].
The solid curve represents the tripartite entanglement in the absence of OPA
[20]. For a non-zero value of gain parameter $G=3.61\times 10^{6}s^{-1}$, we
plot $R_{\tau}^{min}$ for $\theta=0$, $\pi/2$, and $\pi$ in dashed, dashed-
dotted, and dotted curve, respectively. For $\theta=\pi/2$, we obtained
maximum enhancement in tripartite entanglement. The enhancement is also
accompanied with an increase in domain of entanglement over a relatively wider
space of detuning values. For other two choices, there is no significant
enhancement in the maximum value, however, the overall area in which
tripartite entanglement exists increases which is quite interesting.

Figure 7: Tripartite entanglement between cavity-magnon-phonon modes as a
function of $\Delta_{c}/{\omega_{b}}$. The solid curve is for $G=0$, while
rest of the curves are for $G=3.61\times 10^{6}s^{-1}$ and three different
choices of phase $\theta$. The value of $G_{mb}/2\pi=4.80MHz$. The rest of the
parameters are the same as in Fig. (5).

It is important to note that we have neglected Kerr effect in our study. The
presence of the Kerr effect has very important consequences as discussed in
Ref. [27, 46]. In fact, Kerr effect strongly depends on the volume of the
sphere and the strength of drive field. For a YIG sphere of diameter $d=250\mu
m$ and using the other relevant parameters considered here, we show that such
nonlinear effects are negligible and can be ignored in our linearized model.
Since
$\left|{\tilde{\Delta}_{m}}\right|,\,\left|\Delta_{c}\right|\gg\kappa_{c},\,\kappa_{m},G$
and
$g_{mc}^{2}\ll\left|{\tilde{\Delta}_{m}\Delta_{c}}\right|\approx\omega_{b}^{2}$
for the parameters used in our simulations, the effective magnon-phonon
coupling is given by $G_{mb}=\sqrt{2}g_{mb}\Omega/\omega_{b}$ [20]. For
$G_{mb}/2\pi=3.2MHz$ and $g_{mb}/2\pi=0.2Hz$, the strength of the magnetic
field is approximated $B_{0}\approx 3.9\times 10^{-5}T$. The number of spins
in the YIG sphere considered here are $N\approx 3.5\times 10^{16}$ and
therefore, the corresponding drive strength $\Omega\approx 7.1\times
10^{14}Hz$. In an earlier study Li et. al., [20] estimated that the magnitude
of Kerr coefficient for a $250\mu m$ sphere is $K\approx
6.4nHz\ll\omega_{b}^{3}/\Omega^{2}$, which is negligible and our linearized
model is valid for the considered sphere size. Nevertheless, including both
Kerr effect and OPA may lead to interesting new results and may be considered
in a future study.

##  IV Conclusion

In conclusion, we studied a scheme to enhance bipartite and tripartite
entanglement in cavity magnomechanical system. We found that the presence of
OPA in cavity magnomechanical system significantly enhances the bipartite
cavity-magnon and cavity-phonon entanglement for a proper choice of OPA phase.
Interestingly, the presence of OPA has no significant effect on magnon-phonon
entanglement. This is due to the fact that the magnon and phonon modes are
directly coupled within YIG sphere and OPA only changes the photon statistics.
Another important finding is that the presence of OPA relaxes the condition of
strong magnon-phonon coupling to achieve cavity-magnon entanglement. We found
strong cavity-magnon entanglement even in the absence of magnon-phonon
coupling. The robustness of both cavity-magnon and cavity phonon entanglement
against temperature also enhances in the presence of OPA. Most importantly,
the OPA enhances the tripartite entanglement significantly in the limit of
experimentally feasible parameters. Apart from the above mentioned
enhancements, OPA also increases the domain of entanglement in the parameter
regime which is an important feature related to the possible experimental
implementation.

## References

  * Huebl  _et al._ [2013] H. Huebl, C. W. Zollitsch, J. Lotze, F. Hocke, M. Greifenstein, A. Marx, R. Gross, and S. T. Goennenwein, High cooperativity in coupled microwave resonator ferrimagnetic insulator hybrids, Phys. Rev. Lett. 111, 127003 (2013).
  * Tabuchi  _et al._ [2014] Y. Tabuchi, S. Ishino, T. Ishikawa, R. Yamazaki, K. Usami, and Y. Nakamura, Hybridizing ferromagnetic magnons and microwave photons in the quantum limit, Phys. Rev. Lett. 113, 083603 (2014).
  * Zhang  _et al._ [2014] X. Zhang, C.-L. Zou, L. Jiang, and H. X. Tang, Strongly coupled magnons and cavity microwave photons, Phys. Rev. Lett. 113, 156401 (2014).
  * Goryachev  _et al._ [2014] M. Goryachev, W. G. Farr, D. L. Creedon, Y. Fan, M. Kostylev, and M. E. Tobar, High-cooperativity cavity qed with magnons at microwave frequencies, Phys. Rev. Applied 2, 054002 (2014).
  * Bai  _et al._ [2015] L. Bai, M. Harder, Y. Chen, X. Fan, J. Xiao, and C.-M. Hu, Spin pumping in electrodynamically coupled magnon-photon systems, Phys. Rev. Lett. 114, 227201 (2015).
  * Zhang  _et al._ [2015] D. Zhang, X.-M. Wang, T.-F. Li, X.-Q. Luo, W. Wu, F. Nori, and J. You, Cavity quantum electrodynamics with ferromagnetic magnons in a small yttrium-iron-garnet sphere, npj Quantum Information 1, 1 (2015).
  * Bourhill  _et al._ [2016] J. Bourhill, N. Kostylev, M. Goryachev, D. Creedon, and M. Tobar, Ultrahigh cooperativity interactions between magnons and resonant photons in a YIG sphere, Phys. Rev. B 93, 144420 (2016).
  * Kostylev  _et al._ [2016] N. Kostylev, M. Goryachev, and M. E. Tobar, Superstrong coupling of a microwave cavity to yttrium iron garnet magnons, Appl. Phys. Lett. 108, 062402 (2016).
  * Kittel [1948] C. Kittel, On the theory of ferromagnetic resonance absorption, Phys. Rev. 73, 155 (1948).
  * Aspelmeyer  _et al._ [2014] M. Aspelmeyer, T. J. Kippenberg, and F. Marquardt, Cavity optomechanics, Rev. Mod. Phys. 86, 1391 (2014).
  * Vitali  _et al._ [2007] D. Vitali, S. Gigan, A. Ferreira, H. Böhm, P. Tombesi, A. Guerreiro, V. Vedral, A. Zeilinger, and M. Aspelmeyer, Optomechanical entanglement between a movable mirror and a cavity field, Phys. Rev. Lett. 98, 030405 (2007).
  * Palomaki  _et al._ [2013] T. A. Palomaki, J. D. Teufel, R. W. Simmonds, and K. W. Lehnert, Entangling Mechanical Motion with Microwave Fields, Science 342, 710 (2013).
  * Riedinger  _et al._ [2016] R. Riedinger, S. Hong, R. A. Norte, J. A. Slater, J. Shang, A. G. Krause, V. Anant, M. Aspelmeyer, and S. Gröblacher, Non-classical correlations between single photons and phonons from a mechanical oscillator, Nature 530, 313 (2016).
  * Ockeloen-Korppi  _et al._ [2018] C. F. Ockeloen-Korppi, E. Damskägg, J.-M. Pirkkalainen, M. Asjad, A. A. Clerk, F. Massel, M. J. Woolley, and M. A. Sillanpää, Stabilized entanglement of massive mechanical oscillators, Nature 556, 478 (2018).
  * Riedinger  _et al._ [2018] R. Riedinger, A. Wallucks, I. Marinković, C. Löschnauer, M. Aspelmeyer, S. Hong, and S. Gröblacher, Remote quantum entanglement between two micromechanical oscillators, Nature 556, 473 (2018).
  * Marletto and Vedral [2017] C. Marletto and V. Vedral, Gravitationally Induced Entanglement between Two Massive Particles is Sufficient Evidence of Quantum Effects in Gravity, Phys. Rev. Lett. 119, 240402 (2017).
  * Fröwis  _et al._ [2018] F. Fröwis, P. Sekatski, W. Dür, N. Gisin, and N. Sangouard, Macroscopic quantum states: Measures, fragility, and implementations, Rev. Mod. Phys. 90, 025004 (2018).
  * Bassi  _et al._ [2013] A. Bassi, K. Lochan, S. Satin, T. P. Singh, and H. Ulbricht, Models of wave-function collapse, underlying theories, and experimental tests, Rev. Mod. Phys. 85, 471 (2013).
  * Barzanjeh  _et al._ [2021] S. Barzanjeh, A. Xuereb, S. Gröblacher, M. Paternostro, C. A. Regal, and E. M. Weig, Optomechanics for quantum technologies, Nat. Phys. , 1 (2021).
  * Li  _et al._ [2018] J. Li, S.-Y. Zhu, and G. Agarwal, Magnon-photon-phonon entanglement in cavity magnomechanics, Phys. Rev. Lett. 121, 203601 (2018).
  * Li and Zhu [2019] J. Li and S.-Y. Zhu, Entangling two magnon modes via magnetostrictive interaction, New J. Phys. 21, 085001 (2019).
  * Yu  _et al._ [2020] M. Yu, S.-Y. Zhu, and J. Li, Macroscopic entanglement of two magnon modes via quantum correlated microwave fields, J. Phys. B: At. Mol. Opt. Phys 53, 065402 (2020).
  * Nair and Agarwal [2020] J. M. Nair and G. Agarwal, Deterministic quantum entanglement between macroscopic ferrite samples, Appl. Phys. Lett. 117, 084001 (2020).
  * Li and Gröblacher [2021] J. Li and S. Gröblacher, Entangling the vibrational modes of two massive ferromagnetic spheres using cavity magnomechanics, Quantum Sci. Technol. 6, 024005 (2021).
  * Yuan  _et al._ [2020a] H. Y. Yuan, P. Yan, S. Zheng, Q. Y. He, K. Xia, and M.-H. Yung, Steady Bell State Generation via Magnon-Photon Coupling, Phys. Rev. Lett. 124, 053602 (2020a).
  * Zhang  _et al._ [2019] Z. Zhang, M. O. Scully, and G. S. Agarwal, Quantum entanglement between two magnon modes via Kerr nonlinearity driven far from equilibrium, Phys. Rev. Research 1, 023021 (2019).
  * Yang  _et al._ [2020] Z.-B. Yang, J.-S. Liu, H. Jin, Q.-H. Zhu, A.-D. Zhu, H.-Y. Liu, Y. Ming, and R.-C. Yang, Entanglement enhanced by Kerr nonlinearity in a cavity-optomagnonics system, Opt. Express 28, 31862 (2020).
  * Li  _et al._ [2019] J. Li, S.-Y. Zhu, and G. S. Agarwal, Squeezed states of magnons and phonons in cavity magnomechanics, Phys. Rev. A 99, 021801 (2019).
  * Zhang  _et al._ [2022] W. Zhang, T. Wang, X. Han, S. Zhang, and H.-F. Wang, Quantum entanglement and one-way steering in a cavity magnomechanical system via a squeezed vacuum field, Opt. Express 30, 10969 (2022).
  * Yuan  _et al._ [2020b] H. Y. Yuan, S. Zheng, Z. Ficek, Q. Y. He, and M.-H. Yung, Enhancement of magnon-magnon entanglement inside a cavity, Phys. Rev. B 101, 014419 (2020b).
  * Zhang  _et al._ [2016] X. Zhang, C.-L. Zou, L. Jiang, and H. X. Tang, Cavity magnomechanics, Science advances 2, e1501286 (2016).
  * Kittel [1958] C. Kittel, Interaction of spin waves and ultrasonic waves in ferromagnetic crystals, Phys. Rev. 110, 836 (1958).
  * Benguria and Kac [1981] R. Benguria and M. Kac, Quantum langevin equation, Phys. Rev. Lett. 46, 1 (1981).
  * Gardiner and Zoller [2004] C. Gardiner and P. Zoller,  _Quantum Noise: A Handbook of Markovian and Non-Markovian Quantum Stochastic Methods with Applications to Quantum Optics_ , 3rd ed. (Springer, Berlin ; New York, 2004).
  * Giovannetti and Vitali [2001] V. Giovannetti and D. Vitali, Phase-noise measurement in a cavity with a movable mirror undergoing quantum brownian motion, Phys. Rev. A 63, 023812 (2001).
  * DeJesus and Kaufman [1987] E. X. DeJesus and C. Kaufman, Routh-hurwitz criterion in the examination of eigenvalues of a system of nonlinear ordinary differential equations, Phys. Rev. A 35, 5288 (1987).
  * Vidal and Werner [2002] G. Vidal and R. F. Werner, Computable measure of entanglement, Phys. Rev. A 65, 032314 (2002).
  * Plenio [2005] M. B. Plenio, Logarithmic negativity: A full entanglement monotone that is not convex, Phys. Rev. Lett. 95, 090503 (2005).
  * Adesso and Illuminati [2006] G. Adesso and F. Illuminati, Continuous variable tangle, monogamy inequality, and entanglement sharing in gaussian states of continuous variable systems, New J. Phys. 8, 15 (2006).
  * Adesso and Illuminati [2007] G. Adesso and F. Illuminati, Entanglement in continuous-variable systems: recent advances and current perspectives, J. Phys. A: Math. Theor. 40, 7821 (2007).
  * Coffman  _et al._ [2000] V. Coffman, J. Kundu, and W. K. Wootters, Distributed entanglement, Phys. Rev. A 61, 052306 (2000).
  * Shahidani  _et al._ [2014] S. Shahidani, M. H. Naderi, M. Soltanolkotabi, and S. Barzanjeh, Steady-state entanglement, cooling, and tristability in a nonlinear optomechanical cavity, J. Opt. Soc. Am. B 31, 1087 (2014).
  * Li  _et al._ [2015] J. Li, B. Hou, Y. Zhao, and L. Wei, Enhanced entanglement between two movable mirrors in an optomechanical system with nonlinear media, EPL (Europhysics Letters) 110, 64004 (2015).
  * Ahmed and Qamar [2017] R. Ahmed and S. Qamar, Optomechanical entanglement via non-degenerate parametric interactions, Physica Scripta 92, 105101 (2017).
  * Huang and Agarwal [2009] S. Huang and G. S. Agarwal, Enhancement of cavity cooling of a micromechanical mirror using parametric interactions, Phys. Rev. A 79, 013821 (2009).
  * Hyde  _et al._ [2018] P. Hyde, B. M. Yao, Y. S. Gui, G.-Q. Zhang, J. Q. You, and C.-M. Hu, Direct measurement of foldover in cavity magnon-polariton systems, Phys. Rev. B 98, 174423 (2018).