# Entropic approximations of the semigeostrophic shallow water equations

Jean-David Benamou, Colin J Cotter, Jacob J M Francis, Hugo Malamut

<figure><img src="image_1.png" /></figure>

## To cite this version:

Jean-David Benamou, Colin J Cotter, Jacob J M Francis, Hugo Malamut. Entropic approximations of the semigeostrophic shallow water equations. 2025. hal-05161741

HAL Id: hal-05161741

https://inria.hal.science/hal-05161741v1

Preprint submitted on 17 Jul 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Entropic approximations of the semigeostrophic shallow water equations

Jean-David Benamou

INRIA PARIS, MOKAPLAN Team

48 Rue Barrault

75013 Paris, France

jean-david.benamou@inria.fr

Colin J. Cotter

IMPERIAL COLLEGE LONDON, Department of Mathematics

South Kensington Campus

London SW7 2AZ, United Kingdom

colin.cotter@imperial.ac.uk

Jacob J.M. Francis

IMPERIAL COLLEGE LONDON, Department of Mathematics

South Kensington Campus

London SW7 2AZ, United Kingdom

jacob.francis18@imperial.ac.uk

Hugo Malamut

INRIA PARIS, MOKAPLAN Team

48 Rue Barrault

75013 Paris, France hugo.malamut@inria.fr

July 17, 2025

## Abstract

We develop a discretisation of the semigeostrophic rotating shallow water equations, based upon their optimal transport formulation. This takes the form of a Moreau-Yoshida regularisation of the Wasserstein metric. Solutions of which provide the shallow water layer depth represented as a measure, which is itself the push forward of an evolving measure under the semigeostrophic coordinate transformation. First, we propose and study an en-
tropic regularisation of the rotating shallow water equations. Second, we discretise the regularised problem by replacing both measures with weighted sums of Dirac measures, and approximating the (squared) $L^2$ norm of the layer depth appearing as the potential energy. We propose an iterative method to solve the discrete optimisation problem relating the two measures, and analyse its convergence. The iterative method is demonstrated numerically and applied to the solution of the time-dependent shallow water problem in numerical examples.

# 1 Introduction

The semigeostrophic (SG) approximation describes the large scale evolution of fluid flows in the limit of the Rossby number (measuring relative size of the advective and Coriolis terms) going to zero in the distinguished limit where the Froude number (a ratio of transport to wave propagation timescales) is proportional to the square of the Rossby number. Originally proposed to explain the formation and subsequent evolution of atmospheric fronts (Hoskins, 1971), the SG equations have been proposed more recently as a tool for understanding weather models in the SG limit (Cullen, 2007, 2018). The SG equations are advantageous since they do not explicitly support fast wave motions, and they can be solved numerically using optimal transport techniques that do not require numerical dissipation for stability, even in the presence of fronts. These techniques were originally formulated as the "geometric algorithm" of Cullen and Roulstone (1993). At the time, these calculations were limited by available computer power. Further, there have recently been significant advances in computational methods for optimal transport. This is driving efforts to revisit the optimal transport approach for the SG equations. The geometric algorithm approximates the source density as a weighted sum of Dirac masses and computes the optimal transport to a Lebesgue measure in the target domain (referred to as a "semidiscrete" optimal transport problem). This corresponds to the construction of Laguerre cells (also known as power diagrams) which have subsequently been well studied, with Mérigot (2011); Kitagawa et al. (2019) providing a damped Newton method, for example. This was used to modernise the geometric algorithm in Bourne et al. (2022); Egan et al. (2022); Lavier (2024). Entropic regularisation is another popular approach to optimal transport. Here, the addition of a Kullback-Leibler regularisation term (scaled by a small parameter $\epsilon$) allows the fully discrete optimal transport problem (both the source and the target measures are weighted sums of Dirac masses) to be solved via the dual problem and the Sinkhorn iterative procedure that alternates on the optimality conditions for the two dual potentials in the problem. The Sinkhorn iteration is well suited to fast implementation on GPUs Cuturi (2013), and can be accelerated using scaling techniques (Schmitzer, 2019; Chizat et al., 2018). Further, the $O(\epsilon \log(\epsilon))$ error due to regularisation can be efficiently corrected ("debiased") to $O(\epsilon^2)$ (Feydy et al., 2019). Benamou et al. (2024) applied this approach to the
incompressible Boussinesq SG equations in the Eady vertical slice configuration.

In this paper we present a fully discrete optimal transport approach to the rotating shallow water SG equations, as a stepping stone towards the compressible Euler SG equations. The shallow water equations and compressible Euler equations share a feature which is that the divergence-free condition of incompressible flow is replaced by a density that is transported by the flow according to the continuity equation. (In the case of the shallow water model, the “density” is the volume of water per unit horizontal area which can change even though the fluid is assumed incompressible due to the vertical motion of the upper surface.) In this case, the associated optimal transport problem for the SG approximation of these models is generalised by removing the constraint on the target density and adding a potential energy term that penalises it instead (Shutts and Cullen, 1987). The resulting formulation for the shallow water SG model was rigorously analysed in Cullen and Gangbo (2001), and Cullen and Maroofi (2003) for the compressible Euler SG model. This paper is an adaptation to the entropic solution approach of the SG equation (Benamou et al., 2024) to the SG shallow water case. Entropy debiasing however raises new and interesting questions that are only tackled partially here. The optimal transport and entropic optimal transport tools were described in (Benamou et al., 2024) and we will refer the reader to this paper when needed.

The rest of this paper is structured as follows. Section 2 revisits the Shallow water SG model derivation using OT tools. Section 3 presents the Entropic regularization, debiasing, and solution algorithms. Finally, numerical results and comments are given in Section 4.

# 2 The OT formulation of the SWSG equations

## 2.1 Shallow water (SW) equations

The shallow water equations are a standard model of geophysical fluid dynamics which can serve as a simplified model of the ocean or a layer in the atmosphere. They are derived under the assumption of an inviscid, incompressible fluid in 3D with a free surface, in hydrostatic balance and assuming columnar motion. Here, the fluid is defined on 2D planar geometry (in Cartesian coordinates $x_1, x_2$) with Coriolis parameter $f$ and gravity parameter $g$ both assumed constant. Then the shallow water equations for unknown velocity $U_t = (u_1, u_2)$ and height $h$, are given by

$$ \frac{Du_1}{Dt} - fu_2 = -g \frac{\partial h}{\partial x_1}, \qquad (2.1) $$

$$ \frac{Du_2}{Dt} + fu_1 = -g \frac{\partial h}{\partial x_2}, \qquad (2.2) $$

$$ \frac{Dh}{Dt} = -h \left( \frac{\partial u_1}{\partial x_1} + \frac{\partial u_2}{\partial x_2} \right). \qquad (2.3) $$
The first two equations are Newton's balance of forces and the third is the mass conservation expressed in terms of the water column height, $h$. These equations are accompanied by initial conditions for $U_0$ and $h$. In this paper we consider the solution in a convex domain $\Omega \subset \mathbb{R}^2$ with boundary $\partial\Omega$ with either periodic or rigid boundary conditions $(u_1, u_2) \cdot n = 0$, where $n$ is the unit outward pointing normal to the boundary, or a mix of the two.

We recall the definition of the material derivative

$$ \frac{D}{Dt} = \frac{\partial}{\partial t} + u_1 \frac{\partial}{\partial x_1} + u_2 \frac{\partial}{\partial x_2}, \qquad (2.4) $$

so that (2.3) can also be written as a Eulerian conservative continuity equation,

$$ \frac{\partial}{\partial t}h + \frac{\partial}{\partial x_1}(h u_1) + \frac{\partial}{\partial x_2}(h u_2) = 0. \qquad (2.5) $$

The total volume of water is conserved but the height of water columns varies according to the non divergence free velocity. We will model $t \to h_t$ as a curve of densities over $\Omega$.

## 2.2 SG SW (SGSW) equations

The geostrophic approximation, which holds at large scales for slow moving flows, neglects $Du_1/Dt$ and $Du_2/Dt$ in (2.1-2.2), leading to the following formulae that define the (divergence free) geostrophic velocity $U_{t,g} = (u_{1,g}, u_{2,g})$ according to

$$ -f u_{2,g} = -g \frac{\partial h}{\partial x_1}, \quad f u_{1,g} = -g \frac{\partial h}{\partial x_2}. \qquad (2.6) $$

However, this is a purely diagnostic equation that does not predict dynamics. In the semigeostrophic approximation (see Cullen (2006) for a historical review) of (2.1-2.2) we neglect the acceleration of the "ageostrophic" part of the velocity $U_{ag} := U - U_g$, and hence we replace $Du_1/Dt$, $Du_2/Dt$ by $Du_{1,g}/Dt$, $Du_{2,g}/Dt$ (whilst retaining $u_1, u_2$ in 2.4), leading to

$$ \frac{Du_{1,g}}{Dt} - f u_2 = -g \frac{\partial h}{\partial x_1} = -f u_{2,g}, \qquad (2.7) $$

$$ \frac{Du_{2,g}}{Dt} + f u_1 = -g \frac{\partial h}{\partial x_2} = f u_{1,g}, \qquad (2.8) $$

$$ \frac{Dh}{Dt} = -h \left( \frac{\partial u_1}{\partial x_1} + \frac{\partial u_2}{\partial x_2} \right). \qquad (2.9) $$

To clarify how this equation might be solved, we use 2.6 to eliminate $u_{1,g}$, $u_{2,g}$ to obtain

$$ \frac{D}{Dt} \left( \frac{g}{f} \frac{\partial h}{\partial x_2} + f x_2 \right) = g \frac{\partial h}{\partial x_1}, \qquad (2.10) $$

$$ \frac{D}{Dt} \left( \frac{g}{f} \frac{\partial h}{\partial x_1} + f x_1 \right) = -g \frac{\partial h}{\partial x_2}, \qquad (2.11) $$
whilst the last equation (2.9) is unchanged.

The semi-geostrophic shallow water system is (2.9-2.11) with unknowns $(u_1, u_2, h)$ which are functions of $(t, x_1, x_2)$ (time and space). Note here that given $h$, $(u_1, u_2)$ can be obtained from (2.10-2.11), and hence initial conditions are required for $h$ only. This reflects the SG approximation as a next order correction to the geostrophic balance condition that determines $u$ from $h$.

## 2.3 Hoskins' transform and Cullen stability principle

The optimal transport formulation is based on the change of physical coordinates $X = (x_1, x_2) \in \Omega$ into geostrophic coordinates $Y = (y_1, y_2)^1$, also known as Hoskins' transformation (Hoskins, 1975), defined by

$$y_1 = x_1 + \frac{g}{f^2} \frac{\partial h}{\partial x_1}, \quad y_2 = x_2 + \frac{g}{f^2} \frac{\partial h}{\partial x_2}. \qquad (2.12)$$

The geostrophic domain is the image, deforming in time, of the physical domain under this map. We notice that it is a gradient

$$Y = \nabla P_t(X),$$

where

$$P_t(X) = \frac{1}{2} \|X\|^2 + \frac{g}{f^2} h_t(X). \qquad (2.13)$$

Since $h$ depends on time, so does the domain in geostrophic coordinates $P_t$.

Under this change of variables, (2.10-2.11) becomes

$$\frac{DY}{Dt} = f \underbrace{\begin{pmatrix} 0 & 1 \\ -1 & 0 \end{pmatrix}}_{=J} (Y - X) = U_{t,g}(X), \qquad (2.14)$$

where the advective derivative (2.4) is still governed by the physical velocity $U_t = (u_1, u_2)$.

It is convenient from there to switch to a Lagrangian description of the equations coupling particles in the physical and geostrophic domains. The Lagrangian description of fluid dynamics is formulated in terms of a time dependent flow map $\mathcal{X}_t: \Omega \to \Omega$ such that $\mathcal{X}_t(\mathcal{X}_0)$ describes the time evolution of a moving fluid particle for each fixed $\mathcal{X}_0 \in \Omega$. Let us consider particles $\mathcal{X}_t$ moving with velocity $U_t = (u_1, u_2)$ and their images $\mathcal{Y}_t(\mathcal{Y}_0) = \nabla P_t(\mathcal{X}_t(\mathcal{X}_0))$ in geostrophic space. From

<footnote n="1">¹Note that the geostrophic coordinates were denoted G in (Benamou et al., 2024), we use Y to avoid confusion with the gravity constant g.</footnote>
(2.4-2.14) we obtain the following system of ODEs,

$$ \frac{\partial}{\partial t} \chi_t = U_t(\chi_t), \qquad (2.15) $$

$$ \frac{\partial}{\partial t} \mathcal{Y}_t = f J \cdot (\mathcal{Y}_t - \mathcal{X}_t) = U_{t,g}(\mathcal{X}_t). \qquad (2.16) $$

This is complemented by the Lagrangian form of (2.5), which we reformulate now as follows. Let $\mu = \mathcal{L} \perp \Omega$ where $\mathcal{L}$ is the Lebesgue measure on $\Omega$ and define, for all time $t$, $(h_t \mu)$ as the pushforward measure of $(h_0 \mu)$ by the flow map $\mathcal{X}_0 \mapsto \mathcal{X}_t$ satisfying

$$ \int_B h_t d\mu = \int_{\mathcal{X}_t^{-1}(B)} h_0 d\mu, \text{ for all } \mu \text{ measurable sets } B. $$

We will use the shorthand notation for the pushforward,

$$ (h_t \mu) = (\mathcal{X}_t)_\# (h_0 \mu). \qquad (2.17) $$

The Lagrangian system (2.15-2.17) is closed by (2.13) applied to $(\mathcal{X}_t, \mathcal{Y}_t)$,

$$ \mathcal{Y}_t = \mathcal{X}_t + \frac{g}{f^2} \nabla h_t(\mathcal{X}_t), \qquad (2.18) $$

as $U_t$ is also an unknown.

A dynamic solution in geostrophic coordinates alone, can be constructed, based on an additional assumption now known as the Cullen Stability Principle.

**Assumption 1.** (Cullen Stability Principle) The solution of (2.7-2.9) is such that $P_t$ (depending on $h$ as in (2.13)) remains a strictly convex potential for all times in $[0, T]$. In other words, $h$ is $(-f/g)$-convex.

Then the map $X \to \nabla P_t(X)$ is injective and bijective onto $\nabla P_t(\Omega)$. Its inverse is given by $Y \to \nabla Q_t(Y)$ where $Q_t = P_t^*$ is the Legendre-Fenchel transform of $P_t$.

If we are given $Q_t$, we can decouple the dynamics in geostrophic coordinates, according to

$$ \frac{\partial}{\partial t} \mathcal{Y}_t = f J \cdot (\mathcal{Y}_t - \nabla Q_t(\mathcal{Y}_t)) = U_{t,g}(\nabla Q_t(\mathcal{Y}_t)), \qquad (2.19) $$

the last equality being a consequence of (2.18). The computation of $Q_t$ independently of $\mathcal{X}_t$ is explained in the next section.

**Remark 1.** We can also return to a Eulerian description of the flow in geostrophic coordinates. For an initial geostrophic distribution $\sigma_0$, the curve in time of measures $\sigma_t = (\mathcal{Y}_t)_\#\sigma_0$ is a distributional solution of

$$ \frac{\partial}{\partial t} \sigma_t(Y) + \nabla \cdot (\sigma_t(Y) U_{t,g}(\nabla Q_t(Y))) = 0 \qquad (2.20) $$
## 2.4 Optimal Transport Formulation

The following result links the flows in physical and geostrophic coordinates, showing that the dynamics in physical coordinates can be recovered from the geostrophic one assuming the maps {$\nabla Q_t$} are given.

**Proposition 1.** Let $P_0$ and $h_0$ be an initial convex potential and height defined over $\Omega$ and define $\sigma_0 = (\nabla P_0)_{\#}(h_0 \mu)$. Additionally, consider the family of maps {$Y_0 \to Y_t$}$_t$ (a solution of (2.16), for instance) and a family of maps {$\nabla Q_t$} defined on the support of $\sigma_t$.

Defining

$$ \sigma_t = (\mathcal{Y}_t)_{\#} \sigma_0 $$

$$ \mathcal{X}_0 \to \mathcal{X}_t(\mathcal{X}_0) := \nabla Q_t \circ \mathcal{Y}_t \circ \nabla P_0(\mathcal{X}_0), $$

and assuming that

$$ (h_t \mu) = (\nabla Q_t)_{\#} \sigma_t, \qquad (2.21) $$

implies that

$$ (h_t \mu) = (\mathcal{X}_t)_{\#}(h_0 \mu). $$

*Proof.* The result follows from the decomposition of pushforward for composition of maps, as follows.

$$ \begin{array}{ccc} \mathcal{Y}_t(\mathcal{Y}_0) & \xrightarrow{\nabla Q_t} & \mathcal{X}_t(\mathcal{X}_0) \\ \updownarrow t & & \updownarrow t \\ \mathcal{Y}_0 & \leftarrow \nabla P_0 & \mathcal{X}_0 \end{array} \qquad \begin{array}{ccc} \sigma_t & \xrightarrow{(\nabla Q_t)_{\#}} & h_t \mu \\ (\mathcal{Y}_t)_{\#} & \updownarrow & \updownarrow (\mathcal{X}_t)_{\#} \\ \sigma_0 & \leftarrow \leftarrow \nabla P_0 & h_0 \mu \end{array} $$

<figure><img src="image_3.png" /></figure>

We now need to define the family of maps $(\nabla Q_t)_t$ and heights $(h_t)_t$ that are compatible with assumptions (2.13-2.21) and (CSP). This was obtained based on energy minimisation considerations in Cullen and Purser (1989); Cullen and Gangbo (2001); Cullen (2006). It is also contained in the more recent "unbalanced" optimal Transport problem (Chizat, 2017) or the concept of Moreau Yoshida[JF: hyphen or no] envelope in Wasserstein space (Sarrazin, 2022). This is summarised in the following proposition.

**Proposition 2.** Given a measure $\sigma$ and $\mu$ the volume measure on $\Omega$ defined as above, we consider the minimisation problem

$$ \inf_{h \in \mathcal{P}(\Omega)} \mathcal{E}_{\sigma}(h), \qquad (2.22) $$

where

$$ \mathcal{E}_{\sigma} = \mathcal{K}_{\sigma} + \mathcal{P}, $$

$$ \mathcal{K}_{\sigma}(h) = f^2 \text{ OT}(h \mu, \sigma) := f^2 \inf_{G, \sigma=G \# h \mu} \frac{1}{2} \int_{\Omega} \|G(X) - X\|^2 h \mu(\mathrm{d}X), \qquad (2.23) $$

$$ \mathcal{P}(h) = \frac{g}{2} \int_{\Omega} \|h(X)\|^2 d\mu(X), \text{ if } h \in \mathcal{L}_{\mu}^2(\Omega) \text{ and } +\infty \text{ otherwise,} $$
are the total, kinetic and potential energies respectively, and OT is the usual 2-Wasserstein distance squared (see Proposition 2 in (Benamou et al., 2024)). The following properties hold.

(i) (2.22) has a unique solution and also admits the dual formulation,

$$ \inf_{h \in \mathcal{P}(\Omega)} \mathcal{E}_{\sigma}(h) = \sup_{\substack{\varphi \in \mathcal{C}(\Omega), \psi \in \mathcal{C}(\mathbb{R}^d) \text{ s.t.} \\ \varphi(X) + \psi(Y) \le \frac{1}{2} \|Y - X\|^2, \forall (X, Y)}} \mathcal{D}_{\sigma}(\varphi, \psi), \quad (2.24) $$

where

$$ \mathcal{D}_{\sigma}(\varphi, \psi) := \int \psi \sigma(\mathrm{d} Y) - \frac{f^2}{2g} \int_{\Omega} \|\varphi(X)\|^2 d\mu(X). $$

check with JD

(ii) Strong Fenchel Rockafellar duality holds and there exist unique minimisers/maximisers satisfying

$$ h = -\partial_{\varphi} \mathcal{D}_{\sigma}(\varphi, \psi) = -\frac{f^2}{g} \varphi, \quad \psi = \partial_{\sigma} \mathcal{E}_{\sigma}(h). \quad (2.25) $$

(iii) The optimal $(\varphi, \psi)$ are the Kantorovich potentials for the optimal transport problem between $(h\mu)$ and $\sigma$. The associated “Brenier” potentials,

$$ P(X) = \frac{\|X\|^2}{2} - \varphi(X) \quad \text{and} \quad Q(Y) = \frac{\|Y\|^2}{2} - \psi(Y), $$

satisfy

$$ P \text{ is convex, } Q = P^*, (h\mu) = (\nabla Q)_{\#} \sigma \text{ and } \sigma = (\nabla P)_{\#} (h\mu). $$

We are now ready to formulate the SWSG equation in geostrophic space.

**Theorem 3 (OT SWSG).** Assume $(h_0, \sigma_0)$ and $\mu$ are given as per Proposition 1. Further, assume that we are given $\{\mathcal{Y}_t, h_t\}$ solutions for all time $t \in [0, T]$ of the following system,

$$ \left\{ \begin{array}{l} \frac{\partial}{\partial t} \mathcal{Y}_t = f J \cdot (\mathcal{Y}_t - \nabla Q_t(\mathcal{Y}_t)) \quad \mathcal{Y}_0 = \mathrm{Id}, \quad \sigma_t = (\mathcal{Y}_t)_{\#} \sigma_0, \\ (\varphi_t, \psi_t) = \arg \sup_{\varphi \in \mathcal{C}(\Omega), \psi \in \mathcal{C}(\mathbb{R}^d) \text{ s.t.}} \mathcal{D}_{\sigma_t}(\varphi, \psi), \\ \varphi(X) + \psi(Y) \le \frac{1}{2} \|Y - X\|^2, \forall (X, Y) \\ h_t := -\frac{f^2}{g} \varphi_t, \quad Q_t(Y) = \frac{\|Y\|^2}{2} - \psi_t(Y). \end{array} \right. \quad (2.26) $$

Then, for $\mathcal{X}_t := \nabla Q_t(\mathcal{Y}_t)$, the following properties hold.

i) $\{\mathcal{X}_t, \mathcal{Y}_t, h_t\}$ are solutions of the full Lagrangian SWSG system (2.15-2.18).
ii) $U_{g,t}(\mathcal{X}_t) = f J \cdot (\nabla P_t(\mathcal{Y}_t) - \mathcal{X}_t)$ where $P_t = Q_t$.

iii) The geostrophic flow equation can be written in compact form using the Wasserstein $\sigma$ gradient, according to

$$ \frac{\partial}{\partial t} \mathcal{Y}_t = f J \cdot \nabla \partial_{\sigma} \mathcal{E}_{\sigma_t}(h) = f J \cdot \nabla \psi_t(\mathcal{Y}_t). \qquad (2.27) $$

*Proof.* We apply Propositions 1 and 2 in sequence. In particular, (2.18) is obtained by taking the gradient of $P$ (which is convex hence a.e-differentiable) defined in (iii) of Proposition 2, and using the characterisation of $h$ in (ii) of Proposition 2. Point iii) is likewise established by taking the gradient of $Q$ defined in (iii) of Proposition 2. $\square$

**Remark 2.** The full velocity $U_t(\mathcal{X}_t)$ (rather than the geostrophic velocity $U_g$) is not needed and not computed to solve (2.27) but can be recovered from the solution as $\frac{\partial}{\partial t}\mathcal{X}_t$.

**Remark 3.** Combining Proposition 2 (iii) with Proposition 1 (ii), (2.20) can be written as

$$ \frac{\partial}{\partial t} \sigma_t + \nabla \cdot (\sigma_t J \cdot \nabla \partial_{\sigma} \mathcal{E}_{\sigma_t}) = 0 \qquad (2.28) $$

which satisfies the Wasserstein Hamiltonian system definition of Ambrosio and Gangbo (2008). Consequently, the total energy is conserved. Further, the vector field $J \cdot \nabla \partial_{\sigma_t} \mathcal{E}_{\sigma_t}$ is divergence free. Then, provided that $Y_t$ remains smooth and invertible, $\frac{\partial}{\partial t} \sigma_t(Y_t) = 0$.

# 3 Entropic SWSG approximation and Sinkhorn algorithm

In this section, we adapt the entropic solution approach to the incompressible SG equation of (Benamou et al., 2024) to the shallow water SG equation. This consists of one simple change: we add an entropic regularisation term to the optimal transport problem (2.24). It convexifies the problem and enforces the inequality constraint on the potentials in a soft way. A further approximation is required to make a computational method, namely the grid approximation of the potential energy, but we leave this aspect until later. The convergence analysis of the time discretised and entropy regularised problem (as in Carlier and Malamut (2024)) is also left for further studies. In this section, we focus on the solution of (3.3) (and its Sinkhorn Divergence debiased version) using the Sinkhorn Algorithm.

## 3.1 Entropic regularisation

Starting from the primal formulation (2.22), we replace $\text{OT}(.,.)$ in the kinetic energy with $\text{OT}_{\varepsilon}(.,.)$ (as defined in section 3.2 in Benamou et al. (2024)), for $0 < \varepsilon <<$
1. Passing to the dual formulation, (2.24) is replaced by the entropy regularised dual problem,

$$ \sup_{\varphi \in \mathcal{C}(\Omega), \psi \in \mathcal{C}(\mathbb{R}^d)} \mathcal{D}_{\sigma, \varepsilon}(\varphi, \psi) := \int_{\mathbb{R}^d} \psi(Y) d\sigma(Y) - \frac{f^2}{2g} \int_{\Omega} \|\varphi(X)\|^2 d\mu(X) - \\ \varepsilon \int_{\Omega \times \mathbb{R}^d} \left( e^{(\varphi(X)+\psi(Y))/\varepsilon} K_{\varepsilon}(X,Y) - 1 \right) d\mu(X) d\sigma(Y), \quad (3.1) $$

where we define the heat kernel

$$ K_{\varepsilon}(X, Y) = e^{-\|X-Y\|^2/\varepsilon}. \quad (3.2) $$

The extra third term is strictly convex, ensuring existence and uniqueness. It also enforces the positivity constraint in (2.24) in a soft way. The solution of the regularised problem (3.1) is characterised as follows.

**Proposition 4.** (3.1) has a unique solution $(\varphi_{\varepsilon}, \psi_{\varepsilon})$ with the following properties.

(i) The optimality conditions for the optimiser are

$$ \left\{ \begin{array}{l} \psi_{\varepsilon}(Y) = -\varepsilon \log \left( \int_{\Omega \times \mathbb{R}^d} (e^{\varphi_{\varepsilon}(X)/\varepsilon} K_{\varepsilon}(X, Y) d\mu(X)) \right), \quad \forall Y, \\ \\ \varphi_{\varepsilon}(X) = -\varepsilon W_0 \left( \frac{g}{\varepsilon f^2} \int_{\Omega \times \mathbb{R}^d} e^{\psi_{\varepsilon}(Y)/\varepsilon} K_{\varepsilon}(X, Y) d\sigma(Y) \right), \quad \forall X, \end{array} \right. \quad (3.3) $$

where $W_0$ is the 0$^{th}$ branch of the Lambert function.

(ii) We define the "entropic" water height $h_{\varepsilon} := -\frac{f^2}{g} \varphi_{\varepsilon}$. Then $(\varphi_{\varepsilon}, \psi_{\varepsilon})$ solve

$$ \mathrm{OT}_{\varepsilon}(\sigma, h_{\varepsilon}\mu) := \sup_{\varphi_{\varepsilon} \in \mathcal{C}(\Omega), \psi_{\varepsilon} \in \mathcal{C}(\mathbb{R}^d)} \int \psi_{\varepsilon}(Y) d\sigma(Y) + \int_{\Omega} \varphi_{\varepsilon}(X) h_{\varepsilon}(X) d\mu(X) - \\ \varepsilon \int_{\Omega \times \mathbb{R}^d} (e^{(\varphi_{\varepsilon}(X)+\psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X, Y) - 1) d\mu(X) d\sigma(Y), \quad (3.4) $$

and $h_{\varepsilon}$ is the solution of the primal problem

$$ \min_{h} \mathcal{E}_{\sigma, \varepsilon}(h) := \mathrm{OT}_{\varepsilon}(\sigma, h) + \frac{g}{2f^2} \int \|h\|^2 d\mu. \quad (3.5) $$

(iii) The gradient of $\psi_{\varepsilon}$ is also the barycentric map, given by

$$ \nabla \psi_{\varepsilon}(Y) = \mathcal{Y} - \frac{1}{\sigma(Y)} \int_{\Omega \times \mathbb{R}^d} X e^{(\varphi_{\varepsilon}(X)+\psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X, Y) d\mu(X). \quad (3.6) $$
Proof. Property (i) follows directly from the optimality conditions of the strictly concave maximisation (3.1). For (ii), it is sufficient to plug the Legendre-Fenchel dual,

$$ \frac{1}{2} \|\varphi_{\varepsilon}\|^2 = \sup_{h_{\varepsilon}} h_{\varepsilon} \varphi_{\varepsilon} - \frac{1}{2} \|h_{\varepsilon}\|^2, \qquad (3.7) $$

into 3.1. We take the gradient in the first equation of (3.3) to get (iii). □

We now gather a few results on the convergence as $\varepsilon \to 0$ of these quantities.

**Proposition 5.** As $\varepsilon \to 0$, we have the following properties of convergence to the solution of the unregularised problem (which also apply when the measures $\mu$ and $\sigma$ are discrete).

(i) Convergence of the value: Let $(\varphi_{\varepsilon}, \psi_{\varepsilon})$ solve the regularised problem with parameter $\varepsilon$, and let $(\varphi, \psi)$ solve the unregularised problem. Then $0 \le \mathcal{D}_{\sigma}(\varphi_{\varepsilon}, \psi_{\varepsilon}) - \mathcal{D}_{\sigma}(\varphi, \psi) \le C\varepsilon \log(1/\varepsilon)$, for some constant $C$ depending on $\sigma$.

(ii) $\|h_{\varepsilon} - h\|_{L^2(\mu)} \le C' \sqrt{\varepsilon \log(1/\varepsilon)}$ for some constant $C'$ depending on $\sigma$.

(iii) If the Brenier map from $\sigma$ to $h$ is globally Lipschitz, then $\nabla \psi_{\varepsilon} \to \nabla \psi$ in $L^2(\sigma)$.

**Proof.** (i) By duality, we have that
$$ \mathcal{D}_{\sigma}(\varphi, \psi) = \min_{h} \mathcal{E}_{\sigma}(h) \quad \text{and} \quad \mathcal{D}_{\sigma}(\varphi_{\varepsilon}, \psi_{\varepsilon}) = \min_{h} \mathcal{E}_{\sigma, \varepsilon}(h). \qquad (3.8) $$

Since the entropic optimal transport cost $\text{OT}_{\varepsilon}$ is increasing in $\varepsilon$, it is clear that $\mathcal{E}_{\sigma, \varepsilon}$ in (3.5) is also increasing in $\varepsilon$. Thus, the minima in $h$ described in (3.8) are also ordered according to

$$ \mathcal{D}_{\sigma}(\varphi, \psi) \le \mathcal{D}_{\sigma}(\varphi_{\varepsilon}, \psi_{\varepsilon}). $$

For the second inequality, choose $h$ as a candidate in the minimisation of $\mathcal{E}_{\sigma, \varepsilon}$. Then, equation (3.8) gives $\mathcal{D}_{\sigma}(\varphi_{\varepsilon}, \psi_{\varepsilon}) \le \mathcal{E}_{\sigma, \varepsilon}(h)$. So,

$$ \mathcal{D}_{\sigma}(\varphi_{\varepsilon}, \psi_{\varepsilon}) - \mathcal{D}_{\sigma}(\varphi, \psi) \le \mathcal{E}_{\sigma, \varepsilon}(h) - \mathcal{E}_{\sigma}(h) = \text{OT}_{\varepsilon}(\sigma, h) - \text{OT}_{0}(\sigma, h), $$

and it is known that this suboptimality converges in $\varepsilon \log(1/\varepsilon)$ (see Carlier et al. (2017)).

(ii) The function $\mathcal{E}_{\sigma}$ is $\frac{g}{f^2}$-strongly convex, with minimum at $h$, so

$$ \|h - h_{\varepsilon}\|_{L^2(\mu)}^2 \le \frac{2f^2}{g} (\mathcal{E}_{\sigma}(h_{\varepsilon}) - \mathcal{E}_{\sigma}(h)). $$

This left hand side is equal to $\mathcal{D}_{\sigma}(\varphi_{\varepsilon}, \psi_{\varepsilon}) - \mathcal{D}_{\sigma}(\varphi, \psi)$, which is dominated by $C\varepsilon \log(1/\varepsilon)$ as seen above. So

$$ \|h - h_{\varepsilon}\|_{L^2(\mu)} \le C' \sqrt{\varepsilon \log(1/\varepsilon)}. $$
(iii) First, note that, as a result of the relation between the Kantorovich potential $\psi$ and the Brenier potential $Q$, $\nabla\psi(Y) = Y - \nabla Q(Y)$. Combining this fact with relation (3.6) yields

$$ \nabla\psi_{\varepsilon}(Y) - \nabla\psi(Y) = \nabla Q(Y) - \frac{1}{\sigma(Y)} \int_{\Omega \times \mathbb{R}^d} X e^{(\varphi_{\varepsilon}(X) + \psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X, Y) d\mu(X). $$

By Jensen's inequality,

$$ \|\nabla\psi_{\varepsilon} - \nabla\psi\|_{L^2(\sigma)}^2 \le \int |X - \nabla Q(Y)|^2 d\gamma_{\varepsilon}(X, Y), $$

where $d\gamma_{\varepsilon}(X, Y) := e^{(\varphi_{\varepsilon}(X) + \psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X, Y) d\mu(X)d\sigma(Y)$ is the entropic optimal transport plan between $h_{\varepsilon}$ and $\sigma$ (see Peyré et al. (2019) for more information). The measure $\gamma_{\varepsilon}$ has support on the product space $\Omega \times \mathbb{R}^d$ and marginals $h_{\varepsilon}, \sigma$.

Arguing as in Berman (2020a) and Li and Nochetto (2020), who build upon an earlier argument of Gigli (2011), we can use the following inequality,

$$ \frac{|X - \nabla Q(Y)|^2}{2L} \le \left( \frac{|X - Y|^2}{2} - \phi(X) - \psi(Y) \right), $$

where $L$ is the Lipschitz constant of the Brenier map, so that

$$ \frac{\|\nabla\psi_{\varepsilon} - \nabla\psi\|_{L^2(\sigma)}^2}{2L} \le \int \frac{|X - Y|^2}{2} d\gamma_{\varepsilon}(X, Y) - \underbrace{\int \phi(X) d\gamma_{\varepsilon}(X, Y)}_{\int \phi(X) h_{\varepsilon}(X) d\mu(X)} - \underbrace{\int \psi(Y) d\gamma_{\varepsilon}(X, Y)}_{\int \psi(Y) d\sigma(Y)}. \quad (3.9) $$

The first term converges to $OT_0(\sigma, h)$. Since $\phi = -\frac{2g}{f^2}h$ is bounded in $L^2(\mu)$ and $h_{\varepsilon}$ converges to $h$ in $L^2(\mu)$, then

$$ \int \phi h_{\varepsilon} d\mu \to \int \phi h d\mu, $$

by the Cauchy-Schwarz inequality. By Kantorovich duality,

$$ \begin{aligned} \int (\phi h_{\varepsilon})(X) d\mu(X) + \int \psi(Y) d\sigma(Y) & \to \int (\phi h)(X) d\mu(X) + \int \psi(Y) d\sigma(Y) \\ &= OT_0(\sigma, h). \end{aligned} \quad (3.10) $$

Hence, the right-hand side of equation (3.9) tends to 0, and so does the left-hand side.

<figure><img src="image_4.png" /></figure>
The computation of $(\varphi_{\varepsilon}, \psi_{\varepsilon})$ relies on an iterative relaxation of (3.3), in the manner of Sinkhorn. In our setting this amounts to the following.

$$ \left\{ \begin{array}{ll} \psi_{\varepsilon}^{k+1}(Y) = -\varepsilon \log \left( \int_{\Omega \times \mathbb{R}^d} e^{\varphi_{\varepsilon}^k(X)/\varepsilon} K_{\varepsilon}(X, Y) d\mu(X) \right), & \forall Y, \\ \varphi_{\varepsilon}^{k+1}(X) = -\varepsilon W_0 \left( \frac{g}{\varepsilon f^2} \int e^{\psi_{\varepsilon}^{k+1}(Y)/\varepsilon} K_{\varepsilon}(X, Y) d\sigma(Y) \right), & \forall X. \end{array} \right. \quad (3.11) $$

The convergence proof of this variant of Sinkhorn is available in a more general setting in Chizat et al. (2018). The proof in DiMarino and Gerolin (2019) can also be adapted. Replacing the log by the Lambert function in the second equation actually improves the contraction rate of the method.

**Remark 4.** *From a computational perspective, it is important to notice that the formula (3.6), giving both the transport map and flow speed, does not involve derivatives, providing a natural extension of the gradient when $\mu$ is a discrete probability measure.*

## 3.2 Debiasing OT$_{\varepsilon}$ with Sinkhorn divergence

The Entropic regularisation introduces an $\varepsilon$ dependent bias in the kinetic energy $\mathcal{E}_{K,\sigma} = f^2$ OT and associated Wasserstein gradient $\nabla\partial_{\sigma}\mathcal{E}_{K,\sigma}$ (this is (3.6)). A debiased version, referred to as “Sinkhorn divergence”, has been proposed (Genevay et al., 2018; Feydy et al., 2019; Chizat et al., 2020; Pooladian et al., 2022), giving the following correction of OT$_{\varepsilon}$,

$$ S_{\varepsilon}(h\mu, \sigma) = \mathrm{OT}_{\varepsilon}(h\mu, \sigma) - \frac{1}{2} (\mathrm{OT}_{\varepsilon}(h\mu, h\mu) + \mathrm{OT}_{\varepsilon}(\sigma, \sigma)). \quad (3.12) $$

Further, Theorem 1 of Feydy et al. (2019) shows that $S_{\varepsilon}$ is positive and convex in its two variables and continuous w.r.t the Wasserstein topology. It is also built to recover the natural identity $S_{\varepsilon}(v, v) = 0$, which holds for OT but does not hold for $\mathrm{OT}_{\varepsilon}$. Based on the small $\varepsilon$ asymptotic expansion of $\mathrm{OT}_{\varepsilon}$ (Carlier et al., 2017; Conforti and Tamanini, 2021; Pal, 2019), $S_{\varepsilon}$ is an approximation of OT of order $O(\varepsilon^2)$ compared with $O(\varepsilon \log \varepsilon)$ for $\mathrm{OT}_{\varepsilon}$ (Feydy et al., 2019; Chizat et al., 2020). This has been established rigorously in the continuous setting under technical assumptions on the marginal measures, such as compact support. Given both marginal measures, $S_{\varepsilon}$ makes a better proxy of the unregularised Wasserstein distance and is still easy to compute, with just three independent $\mathrm{OT}_{\varepsilon}$ problems to solve and combine.

Going further, the debiased entropic map, given by $\sigma$ Wasserstein gradient $\nabla\partial_{\sigma}\{\mathrm{OT}_{\varepsilon}(h\mu, \sigma) - \mathrm{OT}_{\varepsilon}(\sigma, \sigma)/2\}$, is computable at the same cost. In our problem this means replacing $\nabla\psi$ by $\nabla(\psi_{\varepsilon} - \psi_{\varepsilon}^S)$ where $\psi_{\varepsilon}^S$ is the Kantorovich potential solution of the dual formulation of the symmetric $\mathrm{OT}_{\varepsilon}(\sigma, \sigma)$ problem. As $\sigma$ is given, this is easily and independently computed by the Sinkhorn iteration,

$$ \psi_{\varepsilon}^{S,k+1}(Y) = -\varepsilon \log \left( \int e^{(\psi_{\varepsilon}^{S,k}(Y'))/\varepsilon} K_{\varepsilon}(Y', Y) d\sigma(Y') \right), \quad (3.13) $$
followed by applying the barycentric map formula (3.6) to the resulting potential $\psi_{\varepsilon}^{S,*}$ at convergence $k \to \infty$,

$$ \nabla \psi_{\varepsilon}^{S,*}(Y) = Y - \frac{1}{\sigma(Y)} \int_{\Omega \times \mathbb{R}^d} X e^{(\psi_{\varepsilon}^{S,*}(X) + \psi_{\varepsilon}^{S,*}(Y))/\varepsilon} K_{\varepsilon}(X, Y) d\mu(X). \quad (3.14) $$

Based again on the asymptotic characterisation in the continuous setting of Conforti and Tamanini (2021), Pooladian et al. (2022) show that the correction does not degrade the accuracy of the approximation for smooth continuous transport maps. The bias remains of order $O(\varepsilon)$ in $L^2$ norm. They also show that discretisation by a weighted sum of Dirac masses may degrade the correction. Nevertheless, our experiments (see later sections) with our discretisation of the SWSG shows that debiasing improves the solution, at least for the water height. For samplings of compactly supported marginal measure in particular, the Sinkhorn divergence map corrects at least the entropic regularisation induced contraction at the boundary (Feydy et al., 2019), which we otherwise observed to be significant in our experiments.

## 3.3 Debiasing the Entropic SWSG problem

The Entropic approximation of SWSG consists of replacing OT by $OT_{\varepsilon}$ in (2.22-2.23) and then $\nabla\psi_t$ by $\nabla\psi_{t,\varepsilon}$ (pointwise in $t$) in 2.27. Debiasing the Entropic SWSG problem consists of replacing $OT_{\varepsilon}$ by $S_{\varepsilon}$ and then $\nabla\psi_{t,\varepsilon}$ by $\nabla(\psi_{t,\varepsilon} - \psi_{t,\varepsilon}^S)$. Since we minimise over $h_{\varepsilon}$ (2.22), the debiasing symmetric part $OT_{\varepsilon}(h_{\varepsilon}\mu, h_{\varepsilon}\mu)$ is coupled to the entropic kinetic energy $OT_{\varepsilon}(\sigma, h_{\varepsilon}\mu)$ and the potential energy $P(h_{\varepsilon})$. Thus, unlike the other symmetric part $OT_{\varepsilon}(\sigma, \sigma)$, it cannot be computed independently. Below, we develop iterative methods for this new problem.

First, we recall that while the convexity of $OT_{\varepsilon}(h_{\varepsilon}\mu, \sigma)$ in $h_{\varepsilon}$ is well known, the less obvious convexity of the negative symmetric terms has been established using the change of variable $u_{\varepsilon} = e^{\varphi_{\varepsilon}^S/\varepsilon}$ in the dual formulation ($\varphi_{\varepsilon}^S$ being the classical Kantorovich potential),

$$ -\mathrm{OT}_{\varepsilon}(h_{\varepsilon}\mu, h_{\varepsilon}\mu) = \varepsilon \min_{u_{\varepsilon} \in C^{+}(\Omega)} - \int h_{\varepsilon}(X) \log(u_{\varepsilon}(X)) d\mu(X) - \frac{1}{2} + \frac{1}{2} \int u_{\varepsilon}(X) u_{\varepsilon}(X') K_{\varepsilon}(X, X') d\mu(X) d\mu(X'). \quad (3.15) $$

Proposition 3 and 4 of Feydy et al. (2019) carefully check the positivity of the optimal $u_{\varepsilon}$. Therefore, we can use (3.15) inside the Sinkhorn divergence (3.12) and remove the constant symmetric $\sigma$ part, before plugging the result in place of OT in (2.23). We arrive at the convex/concave saddle point problem,

$$ \inf_{h_{\varepsilon} \in \mathcal{P}(\Omega), u_{\varepsilon} \in C^{+}(\Omega)} \sup_{\varphi_{\varepsilon}, \psi_{\varepsilon}} F(h_{\varepsilon}, u_{\varepsilon}, \varphi_{\varepsilon}, \psi_{\varepsilon}), \quad (3.16) $$
where

$$
\begin{align*}
F(h_{\varepsilon}, u_{\varepsilon}, \varphi_{\varepsilon}, \psi_{\varepsilon}) &:= \int \psi_{\varepsilon}(Y) d\sigma(Y) + \int \varphi_{\varepsilon}(X) h(X) d\mu(X) + \\
&\quad - \varepsilon \int h_{\varepsilon}(X) \log(u_{\varepsilon}(X)) d\mu(X) - \\
&\quad \varepsilon \int \left( e^{(\varphi_{\varepsilon}(X)+\psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X,Y) - 1 \right) d\mu(X) d\sigma(Y) + \\
&\quad \frac{\varepsilon}{2} \int u_{\varepsilon}(X) u_{\varepsilon}(X') K_{\varepsilon}(X,X') d\mu(X) d\mu(X') + \\
&\quad \frac{g}{2} \int \|h_{\varepsilon}(X)\|^2 d\mu(X),
\end{align*}
$$

where the Coriolis parameter is set to $f = 1$ and we have added/removed constants to ease the computations.

Existence and uniqueness of a saddle point solution to the system

$$ \nabla_{h_{\varepsilon}, u_{\varepsilon}, \varphi_{\varepsilon}, \psi_{\varepsilon}} F(h_{\varepsilon}, u_{\varepsilon}, \varphi_{\varepsilon}, \psi_{\varepsilon}) = 0 \quad (3.17) $$

follows from the strict convexity/concavity of $F$ (see Ekeland and Temam (1976) chap.6 for example) as well as the convergence of steepest ascent-descent gradient methods of the form

$$ (h_{\varepsilon}^{k+1}, u_{\varepsilon}^{k+1}) = (h_{\varepsilon}^{k}, u_{\varepsilon}^{k}) - t \nabla_{h_{\varepsilon}, u_{\varepsilon}} F(h_{\varepsilon}^{k}, u_{\varepsilon}^{k}, \varphi_{\varepsilon}^{k}, \psi_{\varepsilon}^{k}), \quad (3.18) $$

$$ (\varphi_{\varepsilon}^{k+1}, \psi_{\varepsilon}^{k+1}) = (\varphi_{\varepsilon}^{k}, \psi_{\varepsilon}^{k}) + t \nabla_{\varphi_{\varepsilon}, \psi_{\varepsilon}} F(h_{\varepsilon}^{k}, u_{\varepsilon}^{k}, \varphi_{\varepsilon}^{k}, \psi_{\varepsilon}^{k}), \quad (3.19) $$

with relaxation parameter $t$ chosen such that $0 < t < \frac{\lambda^2}{2L}$ (in finite dimension, $L$ the Lipschitz constant of $F$ and $\lambda$ bounding below the modulus of convexity/concavity). There exist many refinements of this method in the literature (See Mokhtari et al. (2020) and references therein).

## 3.4 A Sinkhorn approach for the solution of (3.17)

We rewrite (3.17) explicitly as the system

$$ 1 = \int e^{(\varphi_{\varepsilon}(X)+\psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X,Y) d\mu(X), \quad (3.20) $$

$$ h(X) = \int \left( e^{(\varphi_{\varepsilon}(X)+\psi_{\varepsilon}(Y))/\varepsilon} K_{\varepsilon}(X,Y) \right) d\sigma(Y) \quad (3.21) $$

$$ u(X) = e^{\frac{\varphi_{\varepsilon}(X)+g h(X)}{\varepsilon}}, \quad (3.22) $$

$$ h(X) = u(X) \int u(X') K_{\varepsilon}(X, X') d\mu(X'). \quad (3.23) $$
As (3.22-3.23) arise from a strictly convex problem, it has a unique solution $(h_\epsilon[\varphi_\epsilon], u_\epsilon[\varphi_\epsilon])$ depending on $\varphi_\epsilon$, but cannot be solved explicitly. Formally, eliminating $(h_\epsilon, u_\epsilon)$ yields

$$ \sup_{\varphi_\epsilon, \psi_\epsilon} \int \psi_\epsilon(Y) d\sigma(Y) + \mathcal{J}_h(\varphi_\epsilon) - \epsilon \int \left( e^{(\varphi_\epsilon(X)+\psi_\epsilon(Y))/\epsilon} K_\epsilon(X,Y) - 1 \right) d\mu(X) d\sigma(Y), \quad (3.24) $$

and

$$ \mathcal{J}_h(\varphi_\epsilon) = -\frac{g}{2} \int \|h[\varphi_\epsilon](X)\|^2 d\mu(X) - \frac{\epsilon}{2} \int h[\varphi_\epsilon](X) d\mu(X). \quad (3.25) $$

We do not have proof of concavity in $\varphi_\epsilon$ of this new problem. The convergence of a coordinatewise ascent in $(\varphi_\epsilon, \psi_\epsilon)$ algorithm similar to Sinkhorn is unclear and left for further research.

Instead, we develop a heuristic iterative relaxation method “à la Sinkhorn” to solve the optimality system (3.20-3.23), and provide experimental evidence of convergence for it. We notice that $-\varphi_\epsilon$ and $g h_\epsilon$ are comparable (and equal when $\epsilon = 0$), $u_\epsilon$ depends on the symmetric potential. It is therefore tempting to eliminate $h_\epsilon$ using (3.22) in (3.20-3.23). We find (taking the log when necessary),

$$ \psi_{\epsilon}(Y) = -\epsilon \log \left( \int e^{\varphi_{\epsilon}/\epsilon} K_{\epsilon}(X, Y) d\mu(X) \right), \quad (3.26) $$

$$ \varphi_{\epsilon}(X) = -\epsilon W_{0} \left( \frac{g}{\epsilon} u(X) \int e^{\psi_{\epsilon}(Y)}/\epsilon K_{\epsilon}(X, Y) d\sigma(Y) \right), \quad (3.27) $$

$$ g u(X) = \frac{\epsilon \log(u(X)) - \varphi_{\epsilon}(X)}{\int u(X') K_{\epsilon}(X, X') d\mu(X')}. \quad (3.28) $$

If $u=1$ then (3.26,3.27) is exactly (3.3), so we suggest to try and adapt Sinkhorn. This leads us to the following iterative algorithm,

$$ \psi_{\epsilon}^{k+1}(Y) = -\epsilon \log \left( \int e^{\varphi_{\epsilon}^{k}(X)/\epsilon} K_{\epsilon}(X, Y) d\mu(X) \right), \quad (3.29) $$

$$ \varphi_{\epsilon}^{k+1}(X) = -\epsilon W_0 \left( \frac{g}{\epsilon} u^k(X) \int e^{\psi_{\epsilon}^{k+1}(Y)}/\epsilon K_{\epsilon}(X, Y) d\sigma(Y) \right), \quad (3.30) $$

$$ g u^{k+1}(X) = \frac{\epsilon \log(u^k(X)) - \varphi_{\epsilon}^{k+1}(X)}{\int u^k(X') K_{\epsilon}(X, X') d\mu(X')}. \quad (3.31) $$

Numerically, we observe a better convergence compared with (3.18-3.19). We also note that Sinkhorn does not require a choice of relaxation parameter $t$.
# 4 Numerical study

## 4.1 Discretisation and Algorithm

The physical domain $\Omega$ is discretised using a time independent (i.e. Eulerian) regular Cartesian grid $(X_i)_{i \in [1, N]}$. We make the approximation $\mu = \frac{1}{N} \sum_i \delta_{X_i}$ and

$$ (h_t \mu) = \frac{1}{N} \sum_i h_{t,i} \delta_{X_i}, \qquad (4.1) $$

where $(h_{t,i})_i$ is the height X-grid function. Given a continuous initial height measure $h_0$, the initial geostrophic measure is $\sigma_0 = (Id + \frac{g}{f^2} \nabla h_0)_{\#} h_0 \mu$. This is approximated as

$$ \sigma_0 = \frac{1}{N} \sum_i h_{0,i} \delta_{Y_{0,i}}, \text{ where } Y_{0,i} = \{Id + \frac{g}{f^2} \nabla h_0\} (X_i). \qquad (4.2) $$

The family $(Y_{0,i})_i$ specifies the initial position of the Lagrangian particles $(Y_{t,i})_i$ following the flow (2.27). This flow is the Lagrangian discretisation for $\sigma_t = (Y_t)_{\#} \sigma_0$ according to

$$ \sigma_t = \frac{1}{N} \sum_i h_{0,i} \delta_{Y_{t,i}}. \qquad (4.3) $$

The dynamics of the particles is given by the debiased velocity (2.27),

$$ \frac{\partial}{\partial t} Y_{t,i} = f J \cdot \{ \nabla \psi_{t,\varepsilon}(Y_{t,i}) - \nabla \psi_{t,\varepsilon}^S(Y_{t,i}) \}, \qquad (4.4) $$

for all $i$. Here, for all $t$, $(\nabla \psi_{t,\varepsilon}, \nabla \psi_{t,\varepsilon}^S)$ is given statically by (3.6) and (3.14). The discretised potentials in the expression of the maps are given by the Sinkhorn solutions of (3.16) and (3.13) with marginals given and discretised as (4.1-4.3).

Finally the system of ODEs (4.4) is discretised and integrated in time using a time stepping schemes, e.g. Euler, Heun (RK2) or RK4. Figure (1) shows the number of Sinkhorn iterations to obtain converged potentials at successive time steps. When using the previous time step solutions as a warm initialisation, we see improved convergence.

## 4.2 Test case descriptions

In this section we describe the test problems that we used in our numerical study, focussed on jet solutions in a periodic domain. The 2D (horizontal) domain $\Omega = [0, 1] \times [0, 1]$ is periodic in the $x_1$ direction, and we use a suitable periodic (squared) distance,

$$ c(x, y) = \min \left\{ |x_1 - y_1|^2, |x_1 - y_1 + 1|^2, |x_1 - y_1 - 1|^2 \right\} + |x_2 - y_2|^2. \qquad (4.5) $$
<figure><img src="image_5.png" /><figcaption>Figure 1: Illustration of the convergence of the Sinkhorn iterates for the perturbed jet case (see below) with $dt = 0.1, \varepsilon = 0.01, \alpha = 0.001$ and the reinitialisation strategy we employ, which recycles the potentials from the previous iterations (when including dynamics). The index $t$ represents the time steps in our Sinkhorn iterates, which have converged sufficiently below the set tolerance (1e-11 here). The residual is measured as $\max\{|\phi^{k+1} - \phi^k|_\infty, |\psi^{k+1} - \psi^k|_\infty\}$ at iteration $k$. Time steps at a half correspond to the second order intermediate time steps in the two stage Runge-Kutta method used.</figcaption></figure>

Before each optimal transport calculation, particles are remapped into the primary domain $0 \le x_2 \le 1$. Rigid boundary conditions on the bottom and top $x_2 = 0, 1$ boundaries are implicitly enforced by the transport target domain.

In all cases we select units with Coriolis and gravity constants $\hat{f} = 1, \hat{g} = 0.1$. These are appropriately chosen to enforce the SG rotationally dominated regime, shallow water scales and correspond to Burger number and Rossby number of 0.1, which is required to be small for SWSG.

Our first testcase uses a $x_1$-independent $h_0$ initialisation that satisfies (CSP) Cullen stability principle (assumption 1). It generates a stationary (in time) jet height profile that remains in geostrophic balance. Indeed the geostrophic speed (2.6) is the rotation of the $h$ gradient and therefore $u_{2,g}$ vanishes and $u_{1,g}(x_1, x_2) = C(x_2)$ is constant in $x_1$. The density of $\sigma_0$ is therefore also independent of $x_1$ as well as the two transport maps $\nabla P_0$ and $\nabla Q_0$. Speeds and densities are constant in $x_1$, so the $x_1$ independence is preserved in time by the geostrophic dynamics. Geostrophic particles retain their initial $x_1$ independent speed, there is no acceleration and $U = U_g$, see (2.7-2.8). Our stationary jet is obtained from the initial height,

$$h_0(x_1, x_2) = a \tanh(b(x_2 - c)) + d, \quad (4.6)$$

with damping term, $a$, slope scale term, $b$, shift away from zero, $c$, average height,
d. Cullen Stability Principle holds provided, $\frac{3\sqrt{3}f}{4g} > ab^2$. We consider two cases, a shallow jet and a steeper jet. For both cases, $a = 0.1, c = 0.5, d = 1.0$, but $b = 10$ for the steeper jet and $b = 5$ for the shallow jet. This height field yields a pressure gradient and a subsequent stationary jet.

Our second test case corresponds to a perturbation of the initial height with a Gaussian bump,

$$h_0(x_1, x_2) = a \tanh(b(x_2 - c)) + d + \frac{\alpha}{2\pi\sigma_0^2} e^{-\frac{1}{2\sigma_0^2}|x_1 - \mu_1|^2 + |x_2 - \mu_2|^2} \quad (4.7)$$

Throughout we fix $\mu_1 = 0.5, \mu_2 = 0.3, \sigma_0 = 0.1, \alpha = 0.001$. Now necessarily $u_{2,g} \neq 0$, and instabilities will form leading to waves and eventually front formation. Cullen stability principle holds provided $\sigma_0, \alpha$ are chosen to keep $P$ convex. For this, it is sufficient to choose $\alpha$ small; here we take $\sigma = 0.1$, and $\alpha = 0.001$. The profile begins as a stationary jet while the perturbation grows, leading to the formation of waves. Eventually, these waves break, forming fronts.

## 4.3 Numerical results

First we investigate the accuracy of the reconstructed height field corresponding to the steady initial jet profiles. At time 0, the data $\sigma_0$ is constructed from the analytically provided initial height $h_0$, see (4.2). Provided it satisfies CSP, the solution of (2.22) $(\varphi, \psi)$ reconstructs $h_0(X)$ as $-f^2 \varphi(X)/g$ and the geostrophic velocity in geostrophic coordinates $U_{g,t}(\nabla Q_0(X_t)) = -g/f J \nabla h_0(Y := X + \nabla h_0(X))$ as $J \nabla \psi(Y)$.

Figure 6 shows the height error, for different initialisations and as a function of $\varepsilon$ ($N$ given as above), when using $OT_\varepsilon$ as an approximation of $W_2^2$,

$$E_{OT_\varepsilon}^h = \sqrt{S_{0.01}\left(\frac{1}{N}\sum_{i=1}^N\left(-\varphi_{\varepsilon,i}\delta_{X_i}, \frac{1}{M}\sum_{k=1}^M h_{0,k}\delta_{X_k}\right)\right)}$$

and

$$E_{S_\varepsilon}^h = \sqrt{S_{0.01}\left(\frac{1}{N}\sum_{i=1}^N\left((\varphi_{\varepsilon,i}^S - \varphi_{\varepsilon,i})\delta_{X_i}, \frac{1}{M}\sum_{k=1}^M h_{0,k}\delta_{X_k}\right)\right)}$$

when using the debiased $S_\varepsilon$. We measure the errors using a fixed $\varepsilon = 0.01$ Sinkhorn divergence $S_{0.01}$ as a loss between these discrete measures. In both cases, the exact $h_0$ is approximated on a fine grid ($M = 640^2 \gg N$). Debiasing yields lower errors but does not seems to improve the order of convergence, which is roughly $O(\varepsilon^{1.5})$. The steeper (harder) jet profile produces larger errors as expected.
In Figure 7, we investigate the approximation error in $\epsilon$ at time 0 for the velocity of the geostrophic flow $\frac{\partial}{\partial t}\mathcal{Y}_t = U_{g,t}(\nabla Q_0^{-1}(\mathcal{Y}_t))$. Since we use a Lagrangian discretisation, we use the transport loss $S_{0.01}$ as a proxy for $W_2^2$ to measure errors, but this time in the (4D) phase space $(\mathcal{Y}_0, U_{g,0}(\mathcal{Y}_0))$. The plotted errors, for different initialisations are

$$E_{OT_{\epsilon}}^{U} = \sqrt{S_{0.01} \left( \frac{1}{N} \sum_{i=1}^{N} (h_{0,i} \delta_{(\mathcal{Y}_{0,i}, -\frac{g}{f} J \cdot \nabla \psi_{0,\epsilon}(\mathcal{Y}_{0,i}))'} \frac{1}{M} \sum_{k=1}^{M} h_{0,k} \delta_{(\nabla Q_0(\mathcal{X}_{0,k}), U_{g,0}(\mathcal{X}_{0,k}))} \right)}$$

when using $OT_{\epsilon}$ as an approximation of $W_2^2$, and

$$E_{S_{\epsilon}}^{U} = \sqrt{S_{0.01}(\mu_{N}, \nu_{M})},$$

$$\mu_N = \frac{1}{N} \sum_{i=1}^{N} (h_{0,i} \delta_{(\mathcal{Y}_{0,i}, -\frac{g}{f} J \cdot \nabla (\psi_{0,\epsilon}(\mathcal{Y}_{0,i}) - \psi_{0,\epsilon}^{S}(\mathcal{Y}_{0,i})))'}$$

$$\nu_N = \frac{1}{M} \sum_{k=1}^{M} h_{0,k} \delta_{(\nabla Q_0(\mathcal{X}_{0,k}), U_{g,0}(\mathcal{X}_{0,k}))'}$$

when using the debiased $S_{\epsilon}$. The fine $\mathcal{Y}_{0,k}$ discretisation is constructed via Hoskins' transformation of a fine $X_k$ grid. The debiased version achieves a lower error but convergence rate is now $O(\epsilon^{0.75})$ except for the shallow jet case with a weaker gradient.

Figure 2 demonstrates the stationary jet with geostrophic velocities (shown by the colour and quivers). In the central column the barycentric projection into the physical domain shows how physical trajectories are formed, and then the 3D plot illustrates the reconstructed height profile of the tanh curve. Clearly the stability of the jet is maintained with only small oscillations around the base state observed, due to asymmetries from the Lagrangian discretisation.

Figure 3 demonstrates the nonstationary behaviour of the perturbed jet, leading to front-like structures. This is evident in the final two frames, where the large velocities demonstrate the front being supported by strong geostrophic winds. This behaviour is also reflected in the height profile, where a bulge moves along the underlying perturbed jet.

To begin studying the correctness of these numerical solutions, we consider if the total energy is conserved, and explore its breakdown into potential and kinetic energy. The stable jet approximately conserves energy with variations on the scale of 1e-7 (not shown). Figure 4 illustrates shows the energy exchange and conservation for the perturbed jet, when using the debiased Sinkhorn divergence methodology, in terms of the normalised energy

$$\frac{\mathcal{E}_{\sigma}(h_t) - \mathcal{E}_{\sigma}(h_0)}{\mathcal{E}_{\sigma}(h_0) - \mathcal{E}_{\sigma}} \qquad (4.8)$$
where $\mathcal{E}_U$ is the minimum background energy corresponding to a uniform height profile with no velocity. The spatial semidiscretisation conserves energy exactly, so any energy errors will arise from time discretisation or by truncation of the iterative scheme to find the optimal transport solution. Hence, we vary the size of time steps and the stepping scheme (2nd order Heun or 4th order Rk4). Crucially, the system appears to converse its energy with variations in the default energy conversation on the order of 1e-5. There is marginal improvement in conservation from Heun to RK4, but a greater contrast in stability of energy is through smaller time stepping. Overall all approaches keeps the energy error small, around 1e-4.

We remark that the method generates two approximations of the height. The (Eulerian, prognostic) grid function $(h_{t,i})_i$ on the right in Figure 3. A (diagnostic) weighted point cloud Lagrangian sampling,

$$ \frac{1}{N} \sum_{i} h_{0,i} \delta_{\mathcal{X}_{t,i}} = \mathcal{Y}_{t,i} + \nabla (\psi_{t,\epsilon}(\mathcal{Y}_{t,i}) - \psi_{t,\epsilon}^{S}(\mathcal{Y}_{t,i}))' \quad (4.9) $$

in the middle panel of Figure 3 corresponding to the push-forward of geostrophic density sampling (4.3) by the barycentric map approximating (2.21). The weights are fixed at initialisation and cannot capture large density variations. This is seen in particular in the last two snapshots where a “hole” appears in the domains in the last two snapshots.

The SG approximation of the SW equations neglects the acceleration of the ageostrophic component of the velocity $U_{t,ag} := U_t - U_{t,\hat{g}}$ (see Section 2.2). In Figure 5 we measure numerically the time evolution of $\|U_{t,ag}\|_2/\|U_{t,\hat{g}}\|_2$, the relative size of the ageostrophic versus geostrophic velocity. The full velocity needed to compute $U_{t,ag}$ is approximated via second order central finite difference in time $U_t(\mathcal{X}_t) = (\mathcal{X}_{t+1} - \mathcal{X}_{t-1})/(2dt)$ (see remark 2). For the stationary jet (on the left) the velocity is purely geostrophic (see the discussion on the initialisation above) and for the perturbed jet (on the right) the ratio remains around 0.1.

At a later positive time $t$ we do not have a reference analytical solution for the non stationary perturbed jet. Instead, we examine the “pseudoconvergence” towards a fine grid solution. Proposition 5 provides convergence results in the continuous case but little information on the convergence rates. We follow Berman (2020b) (as discussed in Benamou et al. (2024) in section 4.2), parameterising the space discretisation $N$ with $\epsilon$ according to $N = 1/\epsilon^2$. We stop the Sinkhorn iterations when $\mathcal{L}^\infty$ difference of $1e-11$ in the potential increments is reached.

The fine grid solution corresponds to $N_0 = 2^{16}$. Figure 8 shows convergence for the discrete geostrophic density (4.3),

$$ E_{\text{OT}_\epsilon, S_\epsilon}^{\sigma}(\sigma_t^N) = \sqrt{S_{0.01}(\sigma_t^N, \sigma_t^{N_0})}. $$

The added $N$ upper script on $\sigma_t^N$ clarifies the dependence on the discretisation, again using the Sinkhorn divergence loss as a proxy for $W_2$. Both the entropic $(\text{OT}_\epsilon)$ and debiased entropic $(S_\epsilon)$ are tested.
Figure 9 shows the pseudoconvergence in the reconstructed height on the X grid. Either using $h^N = -f/g \varphi_\varepsilon^N$ for the entropic $OT_\varepsilon$ or $h^N = -f/g (\varphi_\varepsilon^N - \varphi_\varepsilon^{S,N})$ for the $S_\varepsilon$ debiased approach. The fine reference solution $h^{N_0}$ is computed with the debiased approach. As the X grid is structured we a $L^p$ loss combined with interpolation in addition of the $S_{0.01}$ transport loss. We observe a pseudoconvergence of order $O(\varepsilon^{3/2})$ suggesting the discretisation error dominates in the $t=0$ convergence against the exact solution (Figure 7). We also remark that debiasing improves the accuracy but not the order of convergence. This was already observed for the SG Eady slice test case Benamou et al. (2024). Heuristically, it can be explained by noticing that the transport correction of the symmetric part in 3.12 is asymptotically irrotational as $\varepsilon \to 0$. See remark 6 in Benamou et al. (2024).

## Acknowledgements

JD Benamou gratefully acknowledges the support of the CNRS Imperial College Abraham de Moivre International Laboratory Fellowship. J.J.M. Francis gratefully acknowledges the support of this work by the Natural Environment Research Council [grant number NE/S007415/1] and their CASE Studentship "Transport methods for the verification of Numerical Weather Prediction (NWP) forecasts" at the Met Office.

## References

Ambrosio, L., Gangbo, W., 2008. Hamiltonian ODEs in the Wasserstein space of probability measures. Communications on Pure and Applied Mathematics 61, 18–53.

Benamou, J.D., Cotter, C., Malamut, H., 2024. Entropic optimal transport solutions of the semigeostrophic equations. Journal of Computational Physics 500, 112745. URL: https://www.sciencedirect.com/science/article/pii/S0021999123008410, doi:https://doi.org/10.1016/j.jcp.2023.112745.

Berman, R.J., 2020a. Convergence rates for discretized Monge–Ampère equations and quantitative stability of optimal transport. Foundations of Computational Mathematics 21, 1099–1140. URL: https://doi.org/10.1007/s10208-020-09480-x, doi:10.1007/s10208-020-09480-x.

Berman, R.J., 2020b. The Sinkhorn algorithm, parabolic optimal transport and geometric Monge–Ampère equations. Numerische Mathematics 145, 771–836. URL: https://doi.org/10.1007/s00211-020-01127-x, doi:10.1007/s00211-020-01127-x.
Bourne, D., Egan, C., Pelloni, B., Wilkinson, M., 2022. Semi-discrete optimal transport methods for the semi-geostrophic equations. Calculus of Variations and Partial Differential Equations 61. URL: http://cvgmt.sns.it/paper/4811/, doi:10.1007/s00526-021-02133-z. cvgmt preprint.

Carlier, G., Duval, V., Peyré, G., Schmitzer, B., 2017. Convergence of entropic schemes for optimal transport and gradient flows. SIAM Journal on Mathematical Analysis 49, 1385–1418.

Carlier, G., Malamut, H., 2024. Well-posedness and convergence of entropic approximation of semi-geostrophic equations. URL: https://arxiv.org/abs/2404.17387, arXiv:2404.17387.

Chizat, L., 2017. Unbalanced optimal transport: Models, numerical methods, applications. Ph.D. thesis. Université Paris sciences et lettres.

Chizat, L., Peyré, G., Schmitzer, B., Vialard, F.X., 2018. Scaling algorithms for unbalanced optimal transport problems. Mathematics of computation 87, 2563–2609.

Chizat, L., Roussillon, P., Léger, F., Vialard, F.X., Peyré, G., 2020. Faster Wasserstein distance estimation with the Sinkhorn divergence. Advances in Neural Information Processing Systems 33, 2257–2269.

Conforti, G., Tamanini, L., 2021. A formula for the time derivative of the entropic cost and applications. URL: https://www.sciencedirect.com/science/article/pii/S002212362100046X, doi:https://doi.org/10.1016/j.jfa.2021.108964.

Cullen, M., 2018. The use of semigeostrophic theory to diagnose the behaviour of an atmospheric GCM. Fluids 3, 72.

Cullen, M., Gangbo, W., 2001. A variational approach for the 2-dimensional semi-geostrophic shallow water equations. Archive for rational mechanics and analysis 156, 241–273.

Cullen, M., Maroofi, H., 2003. The fully compressible semi-geostrophic system from meteorology. Archive for rational mechanics and analysis 167, 309–336.

Cullen, M., Purser, R., 1989. Properties of the Lagrangian semigeostrophic equations. Journal of Atmospheric Sciences 46, 2684–2697.

Cullen, M., Roulstone, I., 1993. A geometric model of the nonlinear equilibration of two-dimensional Eady waves. Journal of Atmospheric Sciences 50, 328–332.

Cullen, M.J.P., 2006. A mathematical theory of large-scale atmosphere/ocean flow. World Scientific.
Cullen, M.J.P., 2007. Modelling atmospheric flows. Acta Numerica 16, 67–154.

Cuturi, M., 2013. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems 26.

DiMarino, S., Gerolin, A., 2019. An optimal transport approach for the Schrödinger bridge problem and convergence of Sinkhorn algorithm. Journal of Scientific Computing 85. URL: https://api.semanticscholar.org/CorpusID: 208139296.

Egan, C.P., Bourne, D.P., Cotter, C.J., Cullen, M.J., Pelloni, B., Roper, S.M., Wilkinson, M., 2022. A new implementation of the geometric method for solving the Eady slice equations. Journal of Computational Physics 469, 111542.

Ekeland, I., Temam, R., 1976. Convex Analysis and Variational Problems. volume 1 of *Studies in Mathematics and its Applications*. North-Holland, Amsterdam.

Feydy, J., Séjourné, T., Vialard, F.X., Amari, S.i., Trouvé, A., Peyré, G., 2019. Interpolating between optimal transport and MMD using Sinkhorn divergences, in: Chaudhuri, K., Sugiyama, M. (Eds.), Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS 2019), PMLR. pp. 2681–2690.

Genevay, A., Peyre, G., Cuturi, M., 2018. Learning generative models with Sinkhorn divergences, in: Storkey, A., Perez-Cruz, F. (Eds.), Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, PMLR. pp. 1608–1617. URL: https://proceedings.mlr.press/v84/genevay18a.html.

Gigli, N., 2011. On Hölder continuity-in-time of the optimal transport map towards measures along a curve. Proceedings of the Edinburgh Mathematical Society 54, 401–409. URL: https://doi.org/10.1017/s001309150800117x, doi:10.1017/s001309150800117x.

Hoskins, B.J., 1971. Atmospheric frontogenesis models: Some solutions. Quarterly Journal of the Royal Meteorological Society 97, 139–153.

Hoskins, B.J., 1975. The geostrophic momentum approximation and the semi-geostrophic equations. Journal of Atmospheric Sciences 32, 233–242.

Kitagawa, J., Mérigot, Q., Thibert, B., 2019. Convergence of a Newton algorithm for semi-discrete optimal transport. Journal of the European Mathematical Society 21, 2603–2651.

Lavier, T., 2024. A semi-discrete optimal transport scheme for the 3D incompressible semi-geostrophic equations. arXiv preprint arXiv:2411.00575.
Li, W., Nochetto, R.H., 2020. Quantitative stability and error estimates for optimal transport plans. IMA Journal of Numerical Analysis 41, 1941–1965. URL: https://doi.org/10.1093/imanum/draa045, doi:10.1093/imanum/draa045.

Mérigot, Q., 2011. A multiscale approach to optimal transport. Computer Graphics Forum 30, 1584–1592. URL: https://hal.archives-ouvertes.fr/hal-00604684, doi:10.1111/j.1467-8659.2011.02032.x.

Mokhtari, A., Ozdaglar, A., Pattathil, S., 2020. A unified analysis of extra-gradient and optimistic gradient methods for saddle point problems: Proximal point approach, in: Chiappa, S., Calandra, R. (Eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, PMLR. pp. 1497–1507. URL: https://proceedings.mlr.press/v108/mokhtari20a.html.

Pal, S., 2019. On the difference between entropic cost and the optimal transport cost. arXiv: Probability URL: https://api.semanticscholar.org/CorpusID: 168170192.

Peyré, G., Cuturi, M., et al., 2019. Computational optimal transport: with applications to data science. Foundations and Trends® in Machine Learning 11, 355–607.

Pooladian, A.A., Cuturi, M., Niles-Weed, J., 2022. Debiaser beware: pitfalls of centering regularized transport maps. URL: https://arxiv.org/abs/2202.08919, arXiv:2202.08919.

Sarrazin, C., 2022. Lagrangian discretization of variational problems in Wasserstein spaces. Ph.D. thesis. Université Paris-Saclay.

Schmitzer, B., 2019. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM J. Sci. Comput. 41.

Shutts, G., Cullen, M., 1987. Parcel stability and its relation to semigeostrophic theory. Journal of Atmospheric Sciences 44, 1318–1330.
<figure><img src="image_6.png" /><figcaption>Figure 2: Stationary Jet profile, with parameters (a,b,c,d) = (0.1, 10, 0.5, 1.0). The Figure illustrates the integration (using Heun's stepping) at $t = 0, 30, 60$. The left column shows the geoverse points, coloured by their initial $y_1$ position illustrating mixing over time. The middle column displays the corresponding barycentric projection and points are coloured by the approximate universe velocity from a 2nd order finite difference method, with illustrative quivers. The right column shows the reconstructed height fields.</figcaption></figure>
Perturbed Jet Profile, $\alpha=0.001$ | $f,g,\epsilon,N, dt = 1, 0.1, 0.01, 100^2, 0.05$

<figure><img src="image_7.png" /><figcaption>Figure 3: Perturbed jet profile, with parameters (a,b,c,d,) = (0.1, 10, 0.5, 1.0) and ($\mu_{x_1}, \mu_{x_2}, \sigma, \alpha$) = (0.5,0.3,0.1,0.001). The Figure illustrates the integration (using Heun's stepping) at $t = 0, 30, 45, 60$. The left column shows the geoverse points, coloured by their initial $y_1$ position illustrating mixing over time. The middle column displays the corresponding barycentric projection and points are coloured by the approximate universe velocity from a 2nd order finite difference method, with illustrative quivers. The right column shows the reconstructed height fields for the perturbed jet, where fronts start to form.</figcaption></figure>
<figure><img src="image_8.png" /><figcaption>Figure 4: Energy conservation test for the perturbed jet, using the debiased Sinkhorn divergence. The rows compare two time-stepping schemes: Heun (second-order Runge-Kutta) and classical Runge-Kutta 4. The left column shows how potential and kinetic energy vary from their initial values over time. All four time steps closely overlap. The right column plots the absolute, normalised default energy conservation (See equation 4.8) on a semi-logarithmic scale. The entropy contribution, not shown here, stays below the set tolerance of $1e-11$. The energy variation is defined as the current energy (kinetic or potential) minus the initial energy.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>Figure 5: Time evolution of the relative size of ageostrophic/geostrophic velocities
the jet and perturbed jet (see section 4.3).</figcaption></figure>
<figure><img src="image_10.png" /><figcaption>Figure 6: Convergence plots $\epsilon \to E_{OT_\epsilon}^h$ and $\epsilon \to E_{S_\epsilon}^h$ (see section 4.3) at time step 0 for different initialisations. Note that the shallow jet is also stationary.</figcaption></figure>

<figure><img src="image_11.png" /><figcaption>Figure 7: Convergence plots $\epsilon \to E_{OT_\epsilon}^U$ and $\epsilon \to E_{S_\epsilon}^U$ (section 4.3) at time step 0 for different initialisations. Note that the shallow jet is also stationary.</figcaption></figure>

<figure><img src="image_12.png" /><figcaption>Figure 8: Pseudo convergence against a fine solution for the perturbed jet, this is $E_{OT_\epsilon, S_\epsilon}^\sigma(\sigma_t^N)$ for $N \nearrow 2^{16}$ and different times (section 4.3).</figcaption></figure>
<figure><img src="image_13.png" /><figcaption>Figure 9: Same as figure 8 but for the reconstructed height and different losses. (section 4.3).</figcaption></figure>