\section*{\abstractname}
    \else
      \begin{center}
        {\bfseries \Large\abstractname\vspace{\z@}}
      \end{center}
      \quotation
    \fi}
    {\if@twocolumn\else\endquotation\fi}
\makeatother
\titlespacing\section{0pt}{5pt plus 4pt minus 2pt}{3pt plus 2pt minus 2pt}
\titleformat*{\section}{\large\bfseries\rmfamily}
\usepackage{adjustbox}
\linespread{1} 
\DeclareMathOperator*{\argmin}{arg\,min}
\usepackage{wrapfig}
\setcounter{tocdepth}{4}
\setcounter{secnumdepth}{4}
\title{ \Large \textbf{Advancing the Predictability of the Earth System Processes:\\ \large{Insights from Optimal Mass Transport Theory}}}



\begin{document}
\begin{center}
{\bf \Large Ensemble Riemannian Data Assimilation over the Wasserstein Space}
\end{center}

\begin{center}
    {\normalsize Sagar K. Tamang$^1$, Ardeshir Ebtehaj$^1$, Peter J. van Leeuwen$^2$, Dongmian Zou$^3$, Gilad Lerman$^4$}
\end{center}

\begin{center}
    $^1$Department of Civil, Environmental and Geo-Engineering and Saint Anthony Falls Laboratory, University of Minnesota-Twin Cities, Twin Cities, Minnesota\\
    $^2$Department of Atmospheric Science, Colorado State University, Fort Collins, Colorado, USA \\
    $^3$Duke Kunshan University, Kunshan, China\\
    $^4$School of Mathematics, University of Minnesota-Twin Cities, Twin Cities, Minnesota, USA
\end{center}

\begin{abstract}
In this paper, we present an ensemble data assimilation paradigm over a Riemannian manifold equipped with the Wasserstein metric. Unlike the Eulerian penalization of error in the Euclidean space, the Wasserstein metric can capture translation and difference between the shapes of square-integrable probability distributions of the background state and observations -- enabling to formally penalize geophysical biases in state-space with non-Gaussian distributions. The new approach is applied to dissipative and chaotic evolutionary dynamics and its potential advantages and limitations are highlighted compared to the classic variational and filtering data assimilation approaches under systematic and random errors.
\end{abstract}

\section{Introduction}

Extending the forecast skill of Earth System Models (ESM) relies on advancing the science of Data Assimilation (DA) \citep{tsuyuki2007recent,carrassi2018data}. A large body of current DA methodologies, either filtering \citep{kalman1960new,evensen1994sequential,reichle2002hydrologic,evensen2003ensemble} or variational approaches \citep{lorenc1986analysis,le1986variational,talagrand1987variational,park2003four,trevisan2010four,carrassi2010accounting,Ebtehaj2013}, are derived from basic principles of Bayesian inference under the assumption that the state-space is unbiased and can be represented well with Gaussian distributions, which are not often consistent with reality \hbox{\citep{bocquet2010beyond,pires2010diagnosis}}. It is well documented that this drawback often limits forecast skills of DA systems \citep{walker2001one,dee2005bias,ebtehaj2014variational,chen2019minimum} especially under the presence of systematic errors \citep{dee2003detection}. 

Apart from particle filters \citep{spiller2008modified,van2010nonlinear}, which are intrinsically designed for state-space with non-Gaussian distribution, numerous modifications to the variational DA (VDA) and ensemble-based filtering methods have been made to tackle non-Gaussianity of geophysical processes \citep{pires1996extending,han2008evaluation,mandel2009ensemble,anderson2010non}. As a few examples, in four-dimensional VDA, a quasi-static VDA is proposed to ensure convergence by gradually increasing the assimilation intervals \citep{pires1996extending}. \citet{kim2003ensemble} proposed modifications to the ensemble Kalman filter \cite[EnKF;][]{evensen1994sequential,li2009simultaneous} using approximate implementation of Bayes' theorem in lieu of linear interpolation via Kalman gain to deal with multimodal systems. For ensemble-based filters, \citet{anderson2010non} proposed a new approach to account for non-Gaussian priors and posteriors by utilizing rank histograms \citep{anderson1996method,hamill2001interpretation}. A hybrid ensemble approach was also suggested to combine advantages of both EnKF and particle filter \citep{mandel2009ensemble}. 

Even though particle filters can handle non-Gaussian likelihood functions, when observations lie away from the support set of the particles, the ensemble variance tends to zero over time and can render the filter degenerate \citep{poterjoy2016efficient}. In recent years, significant progress has been made to treat systematic errors through numerous ad hoc methods such as the field alignment technique \citep{ravela2007data} and morphing EnKF \citep{beezley2008morphing} that can tackle position errors between observations and forecast. Dual state-parameter EnKF \citep{moradkhani2005dual} was also developed to resolve systematic errors originating from parameter uncertainties. Additionally, bias aware variants of the Kalman filter were designed \citep{drecourt2006bias,de2007state,de2007correcting,kollat2008addressing} to simultaneously update the state-space and an {\it a priori} estimate of the additive biases. In parallel, the cumulative distribution function matching \citep{reichle2004bias} has garnered widespread attention in land DA.

From a geometrical perspective, Gaussian statistical inference methods exhibit a flat geometry \citep{Amari1985}. In particular, it is proved that linear auto-regressive and moving average Markov stochastic models, which are driven by Gaussian noise, form dually flat manifolds \citep{amari2012differential}. The notion of distance over such a geometrically flat space is defined over a straight line, which can be quantified by the Euclidean distance.  Consequently, the Euclidean space has served as a major tool in explaining statistical inference techniques using linear Gaussian models and has been used as a cornerstone of DA techniques. It is important to note that the Euclidean distance is ``Eulerian'' \citep{ning2014coping} and thus remains insensitive to the magnitude of translation between probability distributions with disjoint support sets -- when used to interpolate between them. 

Non-Gaussian statistical models often form geometrical manifolds. In the case of nonlinear regression, it is demonstrated that the statistical manifold exhibits a Riemannian geometry \citep{lauritzen1987statistical} over which the notion of distance between probability distributions is geodesic. Such a distance metric shall be Lagrangian to not only capture translation but also the difference between the entire shape of probability distributions \citep{Pennec2006}. How can we equip DA with a Riemannian geometry? To answer this question, inspired by the theories of optimal mass transport \citep{villani2003topics}, this paper presents the Ensemble Riemannian Data Assimilation (EnRDA) framework using the Wasserstein distance metric. 

In recent years, a few attempts have been made to utilize the Wasserstein metric in geophysical DA. \citet{reich2013nonparametric} introduced an ensemble transform particle filter, where the optimal transport framework was utilized to guide the resampling phase of the filter. \citet{ning2014coping} used the Wasserstein distance to reduce forecast uncertainty due to parameter estimation errors in dissipative evolutionary equations. \citet{feyeux2018optimal} suggested a novel approach employing the Wasserstein distance in lieu of the Euclidean distance to penalize the position error between state and observations. More recently, \citet{tamang2020regularized} introduced a Wasserstein regularization in a variational setting to correct for geophysical biases under chaotic dynamics. 

The EnRDA extends the previous work through the following main contributions: (a) EnRDA defines DA as a discrete barycenter problem over the Wasserstein space for assimilation in probability domain without any parametric or Gaussian assumption. The framework provides a continuum of non-parametric analysis probability histograms that naturally span between the distributions of the background state and observations through optimal transport of probability masses. (b) EnRDA operates in an ensemble setting using the entropic regularization by utilizing the Sinkhorn algorithm \citep{cuturi2013sinkhorn} for improving computational efficiency. (c) The paper studies advantages and limitations of DA over the Wasserstein space for dissipative advection-diffusion dynamics and nonlinear chaotic Lorenz-63 model in comparison with 3D Variational (3D-Var) DA as well as filtering techniques. 

The organization of the paper is as follows: Section \ref{sec:2} provides a brief background on Bayesian DA formulations and optimal mass transport. The mathematical formalism of the EnRDA is described in Section \ref{sec:3}. Section \ref{sec:4} presents the results and compares them with their Euclidean counterparts. Section 5 discusses the findings and ideas for future research.

\section{Background} \label{sec:2}
\subsection{Notations}
Throughout, small bold letters represent $m$-element column vectors $\mathbf{x} = (x_1,\ldots,x_m)^\text{T}\in\mathbb{R}^m$, where $(\cdot)^\text{T}$ is the transposition operator. The $m$-by-$n$ matrices $\mathbf{X}\in\mathbb{R}^{m\times n}$ are denoted by capital bold letters, whereas $\mathbb{R}^m_+ \,(\mathbb{R}_+^{m\times n})$ denotes those vectors (matrices) only containing non-negative real numbers. The $\mathbbm{1}_m$ refers to an $m$-element vector of ones and ${\mathbf{I}}_m$ is an $m\times m$ identity matrix. A diagonal matrix with entries given by $\mathbf{x}\in\mathbb{R}^m$ is represented by $\diag(\mathbf{x}) \in \mathbb{R}^{m \times m}$. Notation $\mathbf{x}\sim\mathcal{N}(\boldsymbol{\mu},\,\boldsymbol{\Sigma})$ denotes that the random vector $\mathbf{x}$ is drawn from a Gaussian distribution with mean $\boldsymbol{\mu}$ and covariance $\boldsymbol{\Sigma}$ and $\mathbb{E}_X(\mathbf{x})$ is the expectation of $\mathbf{x}$. The $\ell_q$-norm of $\mathbf{x}$ is defined as $\Vert \mathbf{x}\Vert_q = \big(\sum_{i=1}^m |x_i|^q \big)^{1/q}$ with $q>0$ and the square of the weighted $\ell_2$-norm of $\mathbf{x}$ is represented as $ \norm{\mathbf{x}}_{\mathbf{B}^{-1}}^2=\mathbf{x}^\text{T}\mathbf{B}^{-1}\mathbf{x}$, where $\mathbf{B}$ is a positive definite matrix. Notations of $\mathbf{x}\odot \mathbf{y}$ and $\mathbf{x}\oslash\mathbf{y}$ represent the element-wise Hadamard product and division between equal length vectors $\mathbf{x}$ and $\mathbf{y}$, respectively. Notation $\langle \mathbf{A},\mathbf{B} \rangle = \text{tr}(\mathbf{A}^\text{T}\mathbf{B})$ denotes the Frobenius inner product between matrices $\mathbf{A}$ and $\mathbf{B}$ and $\text{tr}(\cdot)$ and $\text{det}[\cdot]$ represent trace and determinant of a square matrix, respectively. Here, $p(\mathbf{x})=\sum_{i=1}^M p_{\mathbf{x}_{i}}\delta_{\mathbf{x}_{i}}$ represents a discrete probability distribution with respective histogram $\{\mathbf{p}_x \in \mathbb{R}^{M}_+\,:\,\sum_i p_{\mathbf{x}_i} =1\}$ supported on $\mathbf{x}_i$, where $\delta_{\mathbf{x}_i}$ represents a Kronecker delta function at $\mathbf{x}_i$. Throughout, the dimension of the state or observations is denoted by little letters such as $\mathbf{x}\in\mathbb{R}^m$ while the number of ensembles or support points of their respective probability distribution is shown by capital letters such as $\mathbf{p}_x \in \mathbb{R}^{M}_+$. 

\subsection{Data Assimilation on Euclidean Space}
In this section, we provide a brief review of the derivation of classic variational DA and particle filters based on the Bayes' theorem to set the stage for the presented Ensemble Riemannian DA formalism.
\subsubsection{Variational Formulation}
Let us consider a discrete-time Markovian dynamics and its observations as follows:
\begin{equation}
    \begin{split}
        \mathbf{x}^{t} &= \mathcal{M}(\mathbf{x}^{t-1}) + \boldsymbol{\omega}^t\,, \qquad \quad \boldsymbol{\omega}^t \sim \mathcal{N}(\mathbf{0},\mathbf{B})\\
        \mathbf{y}^t &= \mathcal{H}(\mathbf{x}^t) + \boldsymbol{v}^t\,, \qquad \qquad \,\,\,
        \boldsymbol{v}^t \sim \mathcal{N}(\mathbf{0},\mathbf{R})\,,
    \end{split}
    \label{eq:1}
\end{equation}
where $\mathbf{x}^t \in \mathbb{R}^m$ and $\mathbf{y}^t \in \mathbb{R}^n$ represent the state variables and the observations at time $t$, $\mathcal{M}:\mathbb{R}^m\rightarrow{\mathbb{R}^m}$ and $\mathcal{H}:\mathbb{R}^m\rightarrow{\mathbb{R}^n}$ are the deterministic forward model and observation operator, and $\boldsymbol{\omega}^t\in \mathbb{R}^m$ and $\boldsymbol{v}^t\in \mathbb{R}^n$ are the independent and identically distributed model and observation errors, respectively. 

Recalling the Bayes' theorem, dropping the time superscript, without loss of generality, the posterior probability density function (pdf) of the state given the observation can be obtained as $p(\mathbf{x}|\mathbf{y}) \propto p(\mathbf{y}|\mathbf{x})\, p(\mathbf{x})/p(\mathbf{y})$, where $p(\mathbf{y}|\mathbf{x})$ is proportional to the likelihood function, $p(\mathbf{x})$ is the prior density and $p(\mathbf{y})$ denotes the distribution of observations.
Letting $\mathbf{x}_b = \mathbb{E}_X(\mathbf{x}) \in \mathbb{R}^m$ represents the background state, ignoring the constant term $\log\:p(\mathbf{y})$ and assuming Gaussian distributions for the observation error and the prior, logarithm of the posterior density leads to the known three-dimensional variational (3D-Var) cost function \citep{lorenc1986analysis}:
\begin{equation}
    \begin{aligned}
    -\log \:p(\mathbf{x}|\mathbf{y})  &\propto \frac{1}{2}(\mathbf{x}-\mathbf{x}_b)^\text{T} \mathbf{B}^{-1}(\mathbf{x}-\mathbf{x}_b)+\frac{1}{2}(\mathbf{y}-\mathcal{H}(\mathbf{x}))^\text{T} \mathbf{R}^{-1}(\mathbf{y}-\mathcal{H}(\mathbf{x}))\\
    &\propto {\norm{\mathbf{x}-\mathbf{x}_b}_{\mathbf{B}^{-1}}^2+ \norm{\mathbf{y}-\mathcal{H}(\mathbf{x})}_{\mathbf{R}^{-1}}^2}\,.
    \end{aligned}
    \label{eq:2}
\end{equation}
As a result, the analysis state obtained by minimization of the 3D-Var cost function in Eq.~(\ref{eq:2}) is the mode of the posterior distribution that coincides with the posterior mean when errors are drawn from unbiased Gaussian densities and $\mathcal{H}$ is a linear operator. Using the Woodbury matrix inversion lemma \citep{woodbury1950inverting}, it can be easily demonstrated that for a linear observation operator, the analysis states in the 3D-Var and Kalman filter are equivalent \citep{Tarantola1987}. As is evident, zero-mean Gaussian assumptions lead to penalization of the error through the weighted Euclidean norm.


\subsubsection{Particle Filters}
Particle filters \citep{gordon1993novel,doucet2009tutorial,van2019particle} in DA were introduced to address the issue of non-Gaussian distribution of the state by representing the prior and posterior distributions through a weighted ensemble of model outputs referred to as ``particles''. In its standard discrete setting, using Monte Carlo simulations, the prior distribution $p(\mathbf{x})$ is represented by a sum of equal-weight Kronecker delta functions as $\displaystyle p(\mathbf{x}) = \frac{1}{M} \sum_{i=1}^{M} \delta_{\mathbf{x}_i}$, where $\mathbf{x}_i \in \mathbb{R}^m$ is the state variable represented by the $i^\textrm{th}$ particle.

Each of these $M$ particles are then evolved through the nonlinear model in Eq.~(\ref{eq:1}). Assuming that the conditional distribution $\displaystyle p(\mathbf{y}|\mathbf{x}_i) = \frac{1}{(2\pi)^{n/2}|\mathbf{R}|^{1/2}} \exp\bigg\{-\frac{1}{2}[\mathbf{y}-\mathcal{H}(\mathbf{x}_i)]^\text{T}\mathbf{R}^{-1}[\mathbf{y}-\mathcal{H}(\mathbf{x}_i)]\bigg\}$ is Gaussian, using the Bayes' theorem, it can be shown that the posterior distribution $p(\mathbf{x}|\mathbf{y})$ can be approximated using a set of weighted particles as $p(\mathbf{x}|\mathbf{y}) = \sum_{i=1}^M\, w_i\, \delta_{\mathbf{x}_i}$, where $\displaystyle w_i = \frac{p(\mathbf{y}|\mathbf{x}_i)}{\sum_{j=1}^M p(\mathbf{y}|\mathbf{x}_j)}$. The particles are then resampled from the posterior distribution $p(\mathbf{x}|\mathbf{y})$ based on their relative weights and propagated forward in time according to the model dynamics.

As is evident, in particle filters, weights of each particle are updated using the Gaussian likelihood function under a zero-mean error assumption. However, in the presence of systematic biases, when the support sets of particles and the observations are disjoint, only the weights of a few particles become significantly large and weights of other particles tend to zero. As the underlying dynamical system progresses in time, only those few particles, with relatively larger weights, are resampled and the filter can become degenerate gradually in time \citep{poterjoy2016efficient}. 

\subsection{Optimal Mass Transport}

The theory of optimal mass transport (OMT), coined by Gaspard Monge \citep{monge1781memoire} and later extended by Kantorovich \citep{kantorovich1942translocation}, was developed to minimize transportation cost in resource allocation problems with purely practical motivations. Recent developments in mathematics discovered that OMT provides a rich ground to compare and morph probability distributions and uncovered new connections to partial differential equations \citep{Jordan1998,Otto2001} and functional analysis \citep{brenier1987decomposition,benamou2000computational,villani2003topics}. 

In a discrete setting, let us define two discrete probability distributions $p(\mathbf{x})=\sum_{i=1}^M p_{\mathbf{x}_{i}}\delta_{\mathbf{x}_{i}}$ and $p(\mathbf{y})=\sum_{j=1}^N p_{\mathbf{y}_j}\delta_{\mathbf{y}_j}$ with their respective histograms  $\{\mathbf{p}_{x} \in \mathbb{R}^{M}_+\,:\,\sum_i p_{\mathbf{x}_{i}} =1\}$ and $\{\mathbf{p}_{y}\in \mathbb{R}^{N}_+\,:\,\sum_j p_{\mathbf{y}_j} =1\}$ supported on ${\mathbf{x}_{i}}$ and ${\mathbf{y}_j}$. A ``ground'' transportation cost matrix $\mathbf{C}\in \mathbb{R}_+^{M\times N}$ is defined such that its elements $c_{ij} = \norm{{\mathbf{x}_{i}}-{\mathbf{y}_j}}^q_q$ represent the cost of transporting unit probability masses from location $\mathbf{x}_{i}$ to $\mathbf{y}_j$. The Kantorovich OMT problem determines an optimal ``transportation plan'' $\mathbf{U}^a\in\mathbb{R}^{M\times N}_+$ that can linearly map two probability measures onto each other with minimum amount of total transportation cost as follows:
\begin{equation}
            \mathbf{U}^a = \underset{\mathbf{U}}{\text{argmin}} \:\: \langle\mathbf{C},\mathbf{U}\rangle
            \qquad \text{s.t.} \qquad \mathbf{U}\ge 0, \quad\mathbf{U}\mathbbm{1}_{N}=\mathbf{p}_{{x}},\quad \mathbf{U}^\text{T}\mathbbm{1}_{M}=\mathbf{p}_{y}\,.
    \label{eq:3}
\end{equation}

The transportation plan can be interpreted as a ``joint distribution'' that couples the marginals histograms $\mathbf{p}_{{x}}$ and $\mathbf{p}_{y}$. For the transportation cost with $q=2$, the OMT problem in Eq.~(\ref{eq:3}) is convex and defines the square of the 2-Wasserstein distance between the distributions as $d^2_\mathcal{W}\left(\mathbf{p}_{{x}},\mathbf{p}_{{y}}\right) = \langle\mathbf{C},\mathbf{U^a}\rangle$. 

What is the advantage of the Wasserstein distance for interpolating between probability distributions compared to other measures of proximity -- such as the Hellinger distance \citep{hellinger1909neue} or the Kullback–Leibler (KL) divergence \citep{kullback1951information}? To elaborate on this question, we confine our consideration to the Gaussian densities over which the Wasserstein distance can be represented in a closed form. In particular, interpolating over the 2-Wasserstein space using parameter $\eta$, between $\mathcal{N}(\boldsymbol{\mu}_x,\boldsymbol{\Sigma}_x)$ and $\mathcal{N}(\boldsymbol{\mu}_y,\boldsymbol{\Sigma}_y)$, results in a Gaussian distribution $\mathcal{N}(\boldsymbol{\mu}_\eta,\boldsymbol{\Sigma}_\eta)$, where $\boldsymbol{\mu}_\eta = \eta\,\boldsymbol{\mu}_x+(1-\eta)\,\boldsymbol{\mu}_y$ and $\boldsymbol{\Sigma}_\eta = \boldsymbol{\Sigma}_x^{-1/2}\big(\eta\,\boldsymbol{\Sigma}_x+(1-\eta)\,(\boldsymbol{\Sigma}_x^{1/2}\boldsymbol{\Sigma}_y\boldsymbol{\Sigma}_x^{1/2})^{1/2}\big)^2\,\boldsymbol{\Sigma}_x^{-1/2}$ \citep{chen2019optimal}.  


Fig.~\ref{fig:1} shows the spectrum of interpolated distributions between two Gaussian pdfs for a range of the interpolation parameter $\eta \in [0,1]$. As shown, the interpolated densities using the Hellinger distance, which is Euclidean in the space of probability measure, are bimodal. Although the Gaussian shape of the interpolated densities using the KL divergence is preserved, the variance of the interpolants is not necessarily bounded by the variances of the input Gaussian densities. Unlike these metrics, as shown, the Wasserstein distance moves the mean and preserves the shape of the interpolants through a natural morphing process. 

\begin{figure}
    \centering
    \includegraphics[width=1\linewidth]{Fig1.png}
    \caption{Interpolation between two Gaussian distributions $\mathcal{N}(\mu_1, \sigma_1^2)$ and $\mathcal{N}(\mu_2, \sigma_2^2)$ where, $\mu_1=-1.1, \mu_2 = 1.4, \sigma_1^2=0.4$, and $\sigma_2^2=0.01$ as a function of displacement parameter ${\eta} \in [0,1]$ for the (a) Hellinger distance, (b) Kullback-Leibler divergence, and (c) 2-Wasserstein distance \citep{peyre2019}.}
    \label{fig:1}
\end{figure}

As is previously noted, this metric is not limited to any Gaussian assumption. Fig.~\ref{fig:2} shows the 2-Wasserstein interpolation between a gamma and a Gaussian distribution. The results show the Lagrangian nature of the Wasserstein metric that penalizes the translation and mismatch between the shapes of the pdfs. It can be shown that $d_\mathcal{W}^2(\mathbf{p}_{{x}},\mathbf{p}_{y}) = d_\mathcal{W}^2(\overline{\mathbf{p}}_{{x}},\overline{\mathbf{p}}_{y}) +\norm{\boldsymbol{\mu}_{{x}}-\boldsymbol{\mu}_{y}}_2^2$, where $\overline{\mathbf{p}}_{{x}}$ and $\overline{\mathbf{p}}_{y}$ are the centered zero-mean probability masses and $\boldsymbol{\mu}_{{x}}$ and $\boldsymbol{\mu}_{y}$ are the respective mean values \citep{peyre2019}. 

\begin{figure}
    \centering
    \includegraphics[width=0.9\linewidth]{Fig2.png}
    \caption{(a) Optimal transportation plan or the joint distribution $\mathbf{U}^a$ between a gamma $\Gamma(2,2)$ and a Gaussian marginal distribution $\mathcal{N}(6.5,1)$ as well as (b) the 2-Wasserstein interpolation between them for different values of the displacement parameter $\eta\in[0,1]$.}
    \label{fig:2}
\end{figure}

\section{Ensemble Riemannian Data Assimilation} \label{sec:3}
\subsection{Problem Formulation}
First, let us recall that the weighted mean of a cloud of points $\left\{\mathbf{x}_i\right\}_{i=1}^M\in \mathbb{R}^m$ in the Euclidean space is $\boldsymbol{\mu}_{x}=\sum_{i=1}^M \eta_i\,\mathbf{x}_i/M$ for a given family of non-negative weights $\sum_i \eta_i=1$. This expected value is equivalent to solving the following variational problem:
\begin{equation}
 \boldsymbol{\mu}_{x}=\underset{\mathbf{x}}{\text{argmin}} \sum_{i=1}^M \eta_i \norm{\mathbf{x}_i-\mathbf{x}}^2_2\,.
\end{equation}
Thus, the 3D-Var problem in Eq.~(\ref{eq:2}), after Cholesky decomposition \citep{nash1990compact} of the error covariance matrices and rearrangement of the terms, can be interpreted as a ``barycenter problem'' in the Euclidean space, where the analysis state is the weighted mean of the background state and observation. 

By changing the distance metric from Euclidean to the Wasserstein \citep{Agueh2011}, a Riemannian barycenter can be defined as the Fr\'echet mean \citep{frechet1948elements} of $N_p$ probability histograms with finite second-order moments as follows:
\begin{equation}
    {\mathbf p}_\eta = \argmin_{\mathbf{p}} \sum_{k=1}^{N_p} \eta_k\: d^2_\mathcal{W}(\mathbf{p},\mathbf{p}_k)\,.
    \label{eq:5}
\end{equation}

Inspired by \citep{feyeux2018optimal}, the EnRDA defines the probability distribution of the analysis state $p({\mathbf{x}_a})\in \mathbb{R}^{M}$ as the Fréchet barycenter over the Wasserstein space as follows:
\begin{equation}
        p(\mathbf{x}_a) =\argmin_{\mathbf{p}_{x}} \left\{ \eta \:d^2_\mathcal{W}(\mathbf {p}_{x},\,\mathbf{p}_{{x}_b}) + (1-\eta) \:d^2_\mathcal{W}(\mathbf{p}_{x},\,|\text{det}\left[\mathcal{H}^{\prime}(\mathbf{x})\right]|\,\mathbf{p}_{y})\right\}\,,\label{eq:6}
\end{equation}
where the displacement parameter $\eta>0$ assigns the relative weights to the observation and background term to capture their respective geodesic distances from the true state. Here $\mathcal{H}^{\prime}(\cdot)$ is the Jacobian of the observation operator assuming that $\mathcal{H}:\mathbf{x}\xrightarrow{}\mathbf{y}$ is a smooth and a square (i.e., $m=n$) bijective map. The $\eta$ is a hyperparameter and its optimal value should be determined empirically using some reference data through cross-validation studies. It is also important to note that due to the bijective assumption for the observation operator, the above formalism currently lacks the ability to propagate the information content of observed dimensions to unobserved ones. This limitation is further discussed later on in the section \ref{sec:5}.

The solution of the above DA formalism involves finding the optimal analysis transportation plan or the joint distribution $\mathbf{U}^a\in\mathbb{R}^{M\times N}$, using Eq.~(\ref{eq:3}), which couples the background and observation marginal histograms. From the joint histogram $\mathbf{U}^a$, we use the McCann's method \citep{mccann1997convexity,peyre2019computational} to obtain the analysis probability distribution: 
\begin{equation}
    p(\mathbf{x}_a) = \sum_{i=1}^M\sum_{j=1}^N u^a_{ij}\, \delta_{\mathbf{z}_{{ij}}}\,,
    \label{eq:7}
\end{equation}
where the analysis support points are $\mathbf{z}_{ij}= \eta\,\mathbf{x}_{i}+(1-\eta)\,\mathbf{y}_{j}$. The widely used interior-point methods \citep{altman1999regularized} and the Orlin's \citep{orlin1993faster} algorithm which are used to solve Eq.~(\ref{eq:3}), have super-cubic run time with a computational complexity of $O(M^3 \: \log \,M)$, where $M=N$. This is a limitation in high-dimensional geophysical DA problems that will be addressed in the next subsection. 


To solve Eq.~(\ref{eq:6}) in an ensemble setting, let us assume that in the absence of any a priori information, initially the background probability distribution is represented by $i=1\ldots M$ ensemble members of the state variable $\mathbf{x}_{i} \in \mathbb{R}^m$ as $p({\mathbf{x}_b}) = \frac{1}{M} \sum_{i=1}^{M} \delta_{\mathbf{x}_{i}}$. An a priori assumption is needed to reconstruct the observation distribution $p(\mathbf{y})=\sum_{i=1}^{N} \mathbf{p}_{y_j} \delta_{\mathbf{y}_{j}}$ at $j=1\ldots N$ supporting points. To that end, one may choose a parametric or a non-parametric model based on the past climatological information. Here, for simplicity, we assume a zero-mean Gaussian representation with covariance $\mathbf{R}\in\mathbb{R}^{n\times n}$ similar to the suggested approach in \citep{Burgers1998} that can be used to perturb the given observation at each assimilation cycle. After each assimilation cycle, the probability histogram of the analysis state $p(\mathbf{x}_a)$ is recovered from Eq.~(\ref{eq:8}) over $\mathbf{z}_{ij}$ at $M\times N$ support points. Then $p(\mathbf{x}_a)$ is resampled at $M$ points using the multinomial sampling scheme \citep{li2015resampling} to initialize the next time step forecasts. 


\subsection{Entropic Regularization of EnRDA}
In order to speed up the computation of coupling between $\mathbf{p}_{x_b}$ and $\mathbf{p}_y$, the problem in Eq.~(\ref{eq:3}) can be regularized \citep{cuturi2013sinkhorn} as follows:
\begin{equation}
    \mathbf{U}^a =  \underset{\mathbf{U}}{\text{argmin}} \:\: \langle\mathbf{C},\mathbf{U}\rangle - \gamma\,H(\mathbf{U})
            \qquad \text{s.t.} \:\: \: \mathbf{U}\ge 0, \:\:\:\mathbf{U}\mathbbm{1}_N=\mathbf{p}_{x_b},\:
                \: \mathbf{U}^\text{T}\mathbbm{1}_M=\mathbf{p}_y\,,
    \label{eq:8}
\end{equation}
where $\gamma>0$ is the regularization parameter and $H(\mathbf{U}): = 
\langle \mathbf{U},\log\:\mathbf{U}-\mathbbm{1}_M \mathbbm{1}_N^\text{T}\rangle$ represents the Gibbs-Boltzmann relative entropy function. Note that the relative entropy is a concave function and thus its negative value is convex.

The Lagrangian function ($\mathcal{L}$) of Eq.~(\ref{eq:8}) can be obtained by adding two dual variables or Lagrangian multipliers $\mathbf{q}\in \mathbb{R}^M$ and $\mathbf{r}\in \mathbb{R}^N$ as follows:
\begin{equation}
    \mathcal{L}(\mathbf{U},\mathbf{q},\mathbf{r}) = \langle\mathbf{C},\mathbf{U}\rangle - \gamma\,H(\mathbf{U})-\langle\mathbf{q},\mathbf{U}\mathbbm{1}_N-\mathbf{p}_{x_b}\rangle-\langle \mathbf{r},\mathbf{U}^\text{T}\mathbbm{1}_M-\mathbf{p}_y\rangle\,.
    \label{eq:9}
\end{equation}

\noindent Setting the derivative of the Lagrangian function to zero, we have

\begin{equation}
    \frac{\partial \mathcal{L}(\mathbf{U},\mathbf{q},\mathbf{r})}{\partial u_{ij}} = c_{ij}+\gamma \log (u_{ij})-q_i-r_j =0\,, \qquad \forall{i,j}\,.
    \label{eq:10}
\end{equation}

The convexity of the entropic regularization keeps the problem in Eq.~(\ref{eq:8}) strongly convex and it can be shown \citep{peyre2019computational} that Eq.~(\ref{eq:10}) leads to a unique optimal joint density with the following form:
\begin{equation}
    \mathbf{U}^a = \diag(\mathbf{v})\,\mathbf{K}\, \diag(\mathbf{w})\,,
    \label{eq:11}
\end{equation}
where $\mathbf{v} = \exp(\mathbf{q})\oslash (\gamma \mathbbm{1}_M) \in \mathbb{R}^M$ and $\mathbf{w} = \exp(\mathbf{r})\oslash (\gamma \mathbbm{1}_N) \in \mathbb{R}^{N}$ are the unknown scaling variables and $\mathbf{K} \in \mathbb{R}^{M \times N}_+$ is the Gibbs kernel, associated with cost matrix $\mathbf{C}$ with element $\displaystyle k_{ij} = \exp{(-\frac{c_{ij}}{\gamma})}$.

From the mass conservation constraints in Eq.~(\ref{eq:8}) and scaling form of the optimal joint density in Eq.~(\ref{eq:11}), we can derive
\begin{equation}
    \diag(\mathbf{v})\,\mathbf{K}\, \diag(\mathbf{w})\mathbbm{1}_N=\mathbf{p}_{x_b} \qquad \text{and}
                \qquad \diag(\mathbf{w})\,\mathbf{K}^\text{T}\, \diag(\mathbf{v})\mathbbm{1}_M=\mathbf{p}_y\,.
\label{eq:13}
\end{equation}

The two unknown scaling variables $\mathbf{v}$ and $\mathbf{w}$ in Eq.~(\ref{eq:11}) can be iteratively solved using the Sinkhorn's algorithm \citep{cuturi2013sinkhorn} as follows:

\begin{equation}
    \mathbf{v}^{l+1} = \mathbf{p}_{x_b}\oslash(\mathbf{K}\mathbf{w}^l) \qquad \text{and} \qquad \mathbf{w}^{l+1} = \mathbf{p}_y\, \oslash(\mathbf{K}^\text{T}\mathbf{v}^l)\,.
    \label{eq:13}
\end{equation}
A summary of the EnRDA implementation is demonstrated in Algorithm~\ref{alg:EnRDA}.

\begin{algorithm}[!ht]
	{\caption{Ensemble Riemannian Data Assimilation \label{alg:EnRDA}}}
	\begin{algorithmic}[1]
	    \State {\bf Inputs:} Ensemble size $M$, number of perturbed observations $N$ from a chosen observation pdf, displacement parameter $\eta$, entropic regularization parameter $\gamma$ and total number of time steps $T$.
	    \State {\bf Initialize:} $\mathbf{x}^0_{i} \sim p(\mathbf{x}^0), \:\: i = 1,\ldots, M$.
	    \For{$t=1,\ldots,T$}
	        \State  $\mathbf{x}^t_{i}= \mathcal{M}(\mathbf{x}^{t-1}_i)+\boldsymbol{\omega}^t_i, \:\: i = 1,\ldots,M.$
        \State Generating ensemble of observations $\mathbf{y}^t_{j}, \:\: j = 1,\ldots,N$.
	      \State At initial time obtain probability histogram of the background state and observations: \begin{equation*}
	      p(\mathbf{x}_b) = \frac{1}{M} \sum_{i=1}^{M} \delta_{\mathbf{x}^t_i}\:, \quad p(\mathbf{y}) = \sum_{j=1}^{N} \mathbf{p}_{y_j}\delta_{\mathbf{y}^t_{j}}. 
	      \end{equation*}
	   \State{Compute the joint histogram as follows:
			\begin{equation*}
			\mathbf{U}^a = \underset{\mathbf{U}}{\text{argmin}} \: \sum_{i=1}^M \sum_{j=1}^N u_{ij}\,c_{ij} -\gamma \langle\mathbf{U},\log{\mathbf{U}}-\mathbbm{1}_M\mathbbm{1}_N^\text{T}\rangle \quad \text{s.t.} \:\: \mathbf{U}\ge 0,\: \mathbf{U}\mathbbm{1}_N = \mathbf{p}_{{x}_b},\: \mathbf{U}^\text{T}\mathbbm{1}_{M}=\mathbf{p}_{y}\,,
			\end{equation*}
			\begin{equation*}
			    \hspace{10mm} \text{where}\:\: c_{ij} = \norm{{\mathbf{x}_{i}^t}-{\mathbf{y}_j^t}}_2^2\,.
			\end{equation*}
			}
	   \State Obtain analysis probability distribution $ p({\mathbf{x}_a}) = \sum_{i=1}^M\sum_{j=1}^N u^a_{ij}\, \delta_{\mathbf{z}_{ij}}$ where $\mathbf{z}_{ij}={\eta\,\mathbf{x}^t_{i}+(1-\eta)\,\mathbf{y}^t_{j}}$.
	   \State Obtain $M$ analysis ensemble members $\mathbf{x}_{ai} \in \mathbb{R}^m$ by multinomial sampling from $p({\mathbf{x}_a})$.
	   \State Set $\mathbf{x}^t_i:=\mathbf{x}_{ai}$.
	   \EndFor
	\end{algorithmic}
\end{algorithm}

The entropic regularization parameter $\gamma$ plays an important role in characterization of the joint density; however, there exists no closed-form solution for its optimal selection. Generally speaking, increasing the value of $\gamma$ will increase convexity of the cost function and thus computational efficiency; however, at the expense of reduced coupling between the marginal histograms, consistent with the second law of thermodynamics. 


As an example, the effects of $\gamma$ on the coupling between two Gaussian mixture models $\mathbf{p}_{x_b}$ and $\mathbf{p}_y$ are demonstrated in Fig.~\ref{fig:3}. It can be seen that at smaller values of $\gamma=0.001$, the probability masses of the joint distribution are sparse and lie compactly along the main diagonal  -- capturing a strong coupling between the background state and observations. However, as the value of $\gamma$ increases, the probability masses of the joint distribution spread out -- reflecting less degree of dependencies between the marginals. It is important to note that in limiting cases, as $\gamma \rightarrow{0}$, the solution of Eq.~(\ref{eq:8}) converges to the true optimal joint histogram, while as $\gamma \rightarrow{\infty}$ the entropy of the analysis state increases and tends to $\mathbf{p}_{x_b} \mathbf{p}_y^\text{T}$. Throughout, we empirically choose a minimum value for $\gamma$ that leads to a stable solution by the Sinkhorn algorithm, assuring sufficient fidelity to the optimal transportation of probability masses according to Eq.~(\ref{eq:8}).

\begin{figure}
    \centering
    \includegraphics[width=6.5in]{Fig3.png}
    \caption{The effect of the entropic regularization parameter ${\gamma}$ on the optimal joint histogram coupling two Gaussian mixture models $\mathbf{p}_{x_b}: 0.5\, \mathcal{N}(-12,0.4) + 0.5\,\mathcal{N}(-8,0.8)$ and $\mathbf{p}_y:0.55\, \mathcal{N}(5,4) + 0.45\,\mathcal{N}(9.5,4)$.}
    \label{fig:3}
\end{figure}


\section{Numerical Experiments and Results} \label{sec:4}

In order to demonstrate the performance of the EnRDA and quantify its effectiveness, we focus on the linear advection-diffusion equation and the chaotic Lorenz-63 model \citep{lorenz1963deterministic}. The advection-diffusion model explains a wide range of heat, mass, and momentum transport across the land, vegetation, and atmospheric continuum, and has been utilized to evaluate the performance of DA methodologies \citep{zhang1997data,hurkmans2006numerical,ning2014coping,ebtehaj2014variational,berardi2016new}. Similarly, the Lorenz-63 model, as a chaotic model of atmospheric convection, has been widely used in testing the performance of DA methodologies \citep{miller1994advanced,Nakano2007,van2010nonlinear,goodliff2015comparing,tandeo2015combining,tamang2020regularized}. Throughout, under controlled experimental settings with foreknown model and observation errors, we run the forward models under systematic errors and compare the results of the EnRDA with 3D-Var for advection-diffusion dynamics and with the particle filter and EnKF for the Lorenz-63 system.   

\subsection{Advection-Diffusion Equation}

\subsubsection{State-space Characterization}
The advection-diffusion is a special case of the Navier-Stokes partial differential equation. In its linear form, with constant diffusivity in an incompressible fluid flow, it is expressed for a mass conserved physical quantity $\mathbf{x}(\mathbf{s},t)$ as follows: 
\begin{equation}
    \frac{\partial \mathbf{x}(\mathbf{s},t)}{\partial t} + \mathbf{a} \odot \nabla{\mathbf{x}(\mathbf{s},t)} =\mathbf{D}\, \nabla^2{\mathbf{x}(\mathbf{s},t)}\,,
    \label{eq:14}
\end{equation}
where $\mathbf{s}\in \mathbb{R}^n$ represents a $n-$dimensional spatial domain at time $t$. In the above expression, $ \mathbf{a}= (a_1,\ldots,a_n)^T \allowbreak \in \mathbb{R}^n$ is the advection velocity vector and $\mathbf{D} =\diag(D_1,\ldots,D_n)\in \mathbb{R}^{n\times n}$ represents the diffusivity matrix. Given initial condition $\mathbf{x}(\mathbf{s},t=0)$, owing to its linearity, the solution at time $t$ can be obtained by convolving the initial condition with a Kronecker delta function $\delta(\mathbf{s}-\mathbf{a}\,t)$ followed by a convolution with the fundamental Gaussian kernel $\displaystyle\mathcal{G}(\mathbf{s},t)=\frac{1}{(2\pi)^{n/2}\,|\boldsymbol{\Sigma}|^{1/2}} \exp{\left(-\frac{1}{2}\,\mathbf{s}^\textrm{T}\boldsymbol{\Sigma}^{-1}\mathbf{s}\right)}$, where $\boldsymbol{\Sigma} = 2\,\mathbf{D}\,t$.


\subsubsection{Experimental Setup and Results}
In this subsection, we present the results of DA experiments on 1-D and 2-D advection-diffusion equations. For the 1-D case, the state-space is characterized over a spatial domain $s\in(0,60]$ with a discretization of $\Delta s = 0.1$. The model parameters are chosen to be $a = 0.8$ [L/T] and $D=0.25$ [L$^2$/T]. The initial state resembles a bimodal mixture of Gaussian distributions obtained by superposition of two Kronecker delta functions $x(s,t=0) = 300\,\delta(s)$ -- evolved for time 15 and 25 [t], respectively. The ground truth of the trajectory is then obtained by evolving the initial state at a time step of $\Delta t=0.5$ over a time period of $T=0$--30 [t] in the absence of any model error.

The observations are obtained at assimilation intervals $10\Delta t$, assuming an identity observation operator, through corrupting the ground truth by a heteroscedastic Gaussian noise with a variance $\epsilon_y=5$\% of the squared values of the ground truth state. We introduce both systematic and random errors in model simulations. For the systematic error, model velocity and diffusivity coefficient are set to $a' = 0.12$ [L/T] and $D' = 0.4$ [L$^2$/T] respectively. To impose the random error, a heteroscedastic Gaussian noise with variance $\epsilon_b=2$\% is added at every $\Delta t$ to model simulations. One hundred ensembles are used in EnRDA and the regularization and displacement parameters are set to $\gamma=3$ and $\eta=0.2$ by trial and error. To obtain a robust conclusion about the comparison of the proposed EnRDA methodology with 3D-Var, experiments are repeated for 50 independent simulation scenarios.

The evolution of the initial state over a time period $T=0-30$ [t] and the results comparing the EnRDA with 3D-Var at 5, 15, and 25 [t] are shown in Fig.~\ref{fig:4}. As demonstrated, during all time steps, EnRDA reduces the analysis uncertainty, in terms of both bias and unbiased root mean squared error (ubrmse). The shape of the entire state-space is properly preserved and remains closer to the ground truth. As shown, in the 3D-Var, although the analysis state follows the true state reasonably well for initial time steps, as the system propagates over time under systematic errors, the analysis state deviates further away from the ground truth. It is important to note that the displacement parameter in EnRDA is largely determined by the bias while the relative weights in 3D-Var are solely based on the background and observation errors. In particular, under the 3D-Var experiment, the average value of the relative weight we assign to the background state $\alpha = \text{tr}(\mathbf{R})/\text{tr}(\mathbf{R}+\mathbf{B})$ is around 0.4 while in EnRDA this weight is $\eta=0.2$. Thus in EnRDA, we favored the unbiased observations more and thus the observed improvements in comparison with the 3D-Var might not be completely fair. Because the displacement parameter $\eta$ can be tuned for example based on the mean squared error that encompasses the effect of bias but there is no such a mechanism available in 3D-Var. Can EnRDA improve the analysis uncertainty even when $\alpha$ and $\eta$ are comparable?


\begin{figure}
    \centering
    \includegraphics[width=6.5in]{Fig4.png}
    \caption{(a) Temporal evolution of a bimodal initial state under a linear advection-diffusion equation and (b--d) the true state $\mathbf{x}_{tr}$, observations $\mathbf{y}$ and analysis states $\mathbf{x}_{a}$ by 3D-Var and EnRDA for three time snapshots at 5, 15 and 25 [t]. The bias and ubrmse of the analysis state by EnRDA (3D-Var) are reported in the legends.}
    \label{fig:4}
\end{figure}


\begin{figure}[t]
    \centering
    \includegraphics[width=0.6\textwidth]{Fig5.png}
    \caption{The true state $\mathbf{x}_{tr}$ versus the background states $\mathbf{x}_b$ and observations $\mathbf{y}$ (a--c) with systematic errors under a 2-D advection-diffusion dynamics as well as the analysis state $\mathbf{x}_a$ by 3D-Var (d--f) and EnRDA (g--i) for different displacement parameters in the Euclidean $\alpha$ and Riemannian space $\eta$, where the entropic regularization parameter is set to $\gamma =0.003$. The black plus signs show the location of the modes for the true state.}
    \label{fig:5}
\end{figure}

Fig.~\ref{fig:5} shows the results of a 2-D assimilation into the advection-diffusion equation where the underlying state is bimodal. In this example problem, the state-space is characterized over a spatial domain $s_1=(0,10]$ and $s_2=(0,10]$ with a discretization of $\Delta s_1=\Delta s_2 =0.1$. The advection-diffusion is considered to be an isotropic process with the true model parameter values set as $a_1 = a_2= 0.08$ [L/T], and $D_1=D_2 = 0.02$ [L$^2$/T]. The shown state variable is obtained after evolving two Kronecker delta functions $\mathbf{x}(\mathbf{s},t)=1000\, \delta(s_1,s_2)$ and $\mathbf{x}(\mathbf{s},t)=4000\, \delta(s_1,s_2)$ for 25 and 35 [t], respectively.

To resemble a model with systematic errors, background state is obtained by increasing the advective velocity to 0.12 [L/T] while diffusivity is reduced to 0.01 [L$^2$/T] (Fig.~\ref{fig:5}b). Observations are not considered to have position biases; however, a systematic representative error is imposed assuming that the sensing system has a lower resolution than the model. To that end, we evolve two Kronecker delta functions,  $\mathbf{x}(\mathbf{s},t)=800\, \delta(s_1,s_2)$ and $\mathbf{x}(\mathbf{s},t)=2400\, \delta(s_1,s_2)$, with less mass than the true state for same time period of 25 and 35 [t] and then up-scaled the field by a factor of two through box averaging. 


As shown in Fig.~\ref{fig:5}, the EnRDA preserves the shape of the state variable well and gradually moves the mass towards the background state as the value of $\eta$ increases, while the bias remains almost constant and the ubrmse increases from 0.12 to 0.95. The error quality metrics are constantly below the 3D-Var counterpart. The shape of the analysis state for small values of $\alpha$ is not well recovered in 3D-Var due to the position bias. As $\alpha$ increases from 0.25 to 0.75, 3D-Var nudges the analysis state towards the background state and begins to recover the shape. The bias is reduced by more than 30\%, from 0.15 to 0.05; however, this occurs at the expense of almost three folds increase in ubrmse, from 0.3 to 1.1. The reason for reduction of the bias is that the positive differences between the analysis state and true state are compensated by their negative differences. However, ubrmse is quadratic and thus measures the average magnitude of the error irrespective of its signs. We should emphasize that the presented results do not imply that EnRDA is always superior to 3D-Var. Indeed, 3D-Var is a minimum mean squared estimator and cannot be outperformed by EnRDA in the absence of bias in a state space with Gaussian distribution. 


\subsection{Lorenz-63}

\subsubsection{State-space Characterization}
The Lorenz system \cite[Lorenz-63,][]{lorenz1963deterministic} is derived through truncation of the Fourier series of the Rayleigh-B\'{e}nard convection model. This model can be interpreted as a simplistic local weather system only involving the effect of local shear stress and buoyancy forces. The system is expressed using coupled ordinary differential equations that describe the temporal evolution of three coordinates  $x$, $y$, and $z$ representing the rate of convective overturn, horizontal, and vertical temperature variations as:
\begin{equation}
    \begin{split}
        \frac{dx}{dt} &= -\sigma (x-y)\\
        \frac{dy}{dt} &= \rho x-y-xz\\
        \frac{dz}{dt} &= xy-\beta z\,,\\
    \end{split}
    \label{eq:15}
\end{equation}
where $\sigma$ represents the Prandtl number, $\rho$ is a normalized Rayleigh number proportional to the difference in temperature gradient through the depth of the fluid and $\beta$ denotes a horizontal wave number of the convective motion. It is well established that for parameter values of $\sigma =10$, $\rho =28$ and $\beta = 8/3$, the system exhibits chaotic behavior with the phase space revolving around two unstable stationary points located at ($\sqrt{\beta(\rho -1)},\sqrt{\beta(\rho -1)},\rho -1$) and ($-\sqrt{\beta(\rho -1)},-\sqrt{\beta(\rho -1)},\rho -1$).

\subsubsection{Experimental Setup and Results}
In this subsection, we demonstrate the results of DA in the Lorenz system under systematic error using EnRDA, particle filter and EnKF. Throughout, we use the classic multinomial resampling for implementation of EnRDA and particle filter. Apart from the systematic error component, we utilize the standard experimental setting used in numerous DA studies \citep{miller1994advanced,furtado2008data,van2010nonlinear,amezcua2014ensemble}. In order to obtain the ground truth of the model trajectory, the system is initialized at $\mathbf{x}_0=(1.508870$, $-1.531271$, $25.46091$) and integrated with a time step of $\Delta t=0.01$ over a time period of $T=0$--20 [t] using the fourth-order Runge-Kutta approximation \citep{runge1895numerische,kutta1901beitrag}. The observations are obtained at every assimilation interval $40 \Delta t$ by assuming identity observation operator and perturbing the ground truth with Gaussian noise $\boldsymbol{v}_t \sim \mathcal{N}(0,\sigma_{obs}^2\,\boldsymbol{\Sigma}_\rho)$, where $\sigma_{obs}^2=2$ and the correlation matrix $\boldsymbol{\Sigma}_\rho\in \mathbb{R}^{3\times 3}$ is populated with 1 on the diagonal entries, 0.5 on the first sub and super diagonals, and 0.25 on the second sub and super diagonals. 

In order to characterize the distribution of the background state, 100 particles (ensemble members) of particle filter, EnKF, and EnRDA are generated by corrupting the ground truth at the initial time with a zero-mean Gaussian noise $\boldsymbol{\omega}_0 \sim \mathcal{N}(0,\sigma_{0}^2\: \mathbf{I}_3)$, where $\sigma_{0}^2=2$. For introducing systematic errors, the model parameters are set to $\sigma' = 10.5$, $\rho'=27$, and $\beta' =10/3$. The random errors are also introduced as the system evolves in time by adding a Gaussian noise
$\boldsymbol{\omega}_t \sim \mathcal{N}(0,\sigma_{b}^2\: \mathbf{I}_3)$ at every $\Delta t$, with $\sigma_{b}^2= 0.02$. Throughout, to draw a robust statistical conclusion about the error statistics, the DA experiments are repeated for 50 independent simulations. As described previously, to properly account for the effects of both bias and ubrmse, the optimal value of the displacement parameter $\eta$ in EnRDA can be selected based on an offline analysis of the minimum mean squared analysis or forecast error. However, to provide a fair comparison between the EnRDA and other filtering methods, at each assimilation cycle, we set $\eta = \text{tr}(\mathbf{R})/\text{tr}(\mathbf{R}+\mathbf{B})$ assuming that the observation operator is an identity matrix. Note that while the observation error covariance remains constant in time, the background error covariance is obtained from simulated ensembles by EnRDA and changes in time dynamically. This selection assures that the relative weights assigned to the background state and observations remain at the same order of magnitude among different methods.

\begin{figure}
    \centering
    \includegraphics[width=6.5in]{Fig7_v3.png}
    \caption{Temporal evolution of the true state $\mathbf{x}_{tr}$ of the Lorenz-63, observations $\mathbf{y}$ as well as the analysis state $\mathbf{x}_a$ for the particle filter $(\textrm{PF})$ (first column), EnKF $(\textrm{EnKF})$ (second column) and EnRDA $(\textrm{EnRDA})$ (third column) with 100 particles (ensemble members) respectively. The temporal evolution of the particles and ensemble members are shown with solid gray lines. Also shown within dashed rectangles are the windows of time over which support sets of observations and particle spread are disjoint in particle filter as well as EnKF and EnRDA deviate from the ground truth.}
    \label{fig:6}
\end{figure}

Fig.~\ref{fig:6} shows the temporal evolution of the ground truth and the analysis state by the particle filter (first column), EnKF (second column), and EnRDA (third column) over a time period of 0 to 15 [t] for one simulation. As is evident, the particle filter is well capable of capturing the ground truth when the observations lie within the particle spread. However, when the observations lie far apart from the support set of particles (Fig.~\ref{fig:6}, dashed box) and the pdfs of the background state and observations become disjoint, the filter becomes degenerate and the analysis state (particle mean) deviates away from the ground truth. It is to note that due to the systematic error, the particles in the $z$-coordinate lie away from the observations and the trajectory fluctuates around the mean of the ground truth (Fig.~\ref{fig:6} g, dashed box). As a result, the bias of the particle filter along the $z$-dimension is markedly lower than that of the EnKF and the EnRDA while ubrmse is significantly higher. Whereas, both EnKF and EnRDA are capable of capturing the true state well even when ensemble spread and observations are far apart from each other. Although EnKF does not suffer from the same problem of filter degeneracy as the particle filter, in earlier time steps from 2.5 to 7.5 [t], it struggles to adequately nudge the analysis state towards the ground truth when ensemble members are far from the observations due to the imposed systematic bias. EnRDA seems to be robust to the propagation of systematic biases in this region and follows the true trajectory well.

\begin{figure}
    \centering
    \includegraphics[width=6.5in]{Fig8_v3.png}
    \caption{Temporal evolution of bias and ubrmse along three dimensions of the Lorenz-63 for (a, d) particle filter, (b, e) EnKF, and (c, f) EnRDA, each with 100 particles (ensemble members). The mean values are computed over 50 independent simulations.}
    \label{fig:8}
\end{figure}

The time evolution of the mean values of the bias and ubrmse for 50 independent simulations, with the same error structure, is color coded over the phase space in Fig.~\ref{fig:8}. As is evident, these forecast quality metrics are relatively lower for the EnRDA than the EnKF and particle filter throughout the simulation period. Nevertheless, we can see that the improvement compared to the EnKF is modest. In particular, across all dimensions of the problem, the mean bias and ubrmse are decreased in EnRDA by 68 (13)\% and 53 (27)\% compared to the particle filter (EnKF). More detailed information about the expected values of the bias and ubrmse are reported in Table \ref{tab:1}. We emphasize that the presented results shall be interpreted in light of the presence of systematic biases. In fact, EnRDA cannot reduce the analysis error variance beyond a minimum mean squared estimator such as EnKF in the absence of biases.

\begingroup
\setlength{\tabcolsep}{5pt}
\begin{table}
    \centering
        \caption{Expected values of the bias and ubrmse for the particle filter, EnKF and EnRDA from 50 independent simulations of Lorenz-63 across all problem dimensions.}
        \vspace{2mm}
        \begin{tabular}{ccccccccc}
        \toprule
        Methods & \multicolumn{4}{c}{bias} & \multicolumn{4}{c}{ubrmse}\\
        & $x$ & $y$ & $z$ & $x-z$ & $x$ & $y$ & $z$ & $x-z$\\
        \midrule
        Particle Filter & 2.24 & 2.41 & 0.59 & 1.75 & 6.25 & 7.95 & 7.88 & 7.36\\
        EnKF & 0.33 & 0.35 & 1.23 & 0.64 & 3.80 & 5.41 & 5.02 & 4.74 \\
        EnRDA & 0.17  & 0.24 & 1.25  & 0.56 & 2.63 &  4.0  & 3.78 & 3.47 \\
        \bottomrule
    \end{tabular}
    \label{tab:1}
\end{table}
\endgroup

\section{Discussion and Concluding Remarks} \label{sec:5}

In this study, we introduced an ensemble data assimilation (DA) methodology over a Riemannian manifold, namely Ensemble Riemannian DA (EnRDA), and illustrated its performance in comparison with a few  Euclidean  DA techniques for dissipative and chaotic dynamics.. We demonstrated that the presented methodology is capable of assimilating information in probability domain -- characterized by the families of distributions with finite second-order moments. The key message is that when the probability distribution of the forecast and observations exhibit non-Gaussian structure and their support sets are disjoint, due to the presence of systematic errors; the Wasserstein metric can be leveraged to potentially extend geophysical forecast skills. Even though, future research for a comprehensive comparison with existing filtering and bias correction methodologies is needed to completely characterize relative pros and cons of the proposed approach -- especially when it comes to the ensemble size and optimal selection of the displacement parameter $\eta$. 

We explained the role of regularization and displacement parameter in EnRDA and empirically examined their effects on the optimal joint histogram, coupling the background state and observations, and consequently on the analysis state. Nevertheless, future studies are required to characterize closed-form or heuristic expressions to expand our understating of their impacts on the forecast uncertainty. As was explained earlier, unlike the Euclidean DA methodologies that assimilate available information using different relative weights across multiple dimensions through the error covariance matrices; a scalar displacement parameter is utilized in the the EnRDA that interpolates uniformly between all dimensions of the problem. Future research can be devoted to developing a framework that utilizes a vector representation of the displacement parameters to effectively tackle possible heterogeneity of uncertainty across multiple dimensions.

In it's current form, the EnRDA requires the observation operator to be smooth and bijectve. This is a limitation when observations of all problem dimensions are not available and propagation of observations to non-observed dimensions is desired. Extending the EnRDA methodology to include partially observed systems seems to be an important future research area. This could include performing a rough inversion for unobserved components of the system offline or extending the methodology in the direction of particle flows \citep{van2019particle}. 

Lastly, we should mention that the EnRDA is computationally expensive as it involves estimation of the coupling through the Wasserstein distance. On a desktop machine with a 3.4 GHz CPU clock rate, it took around 1600\,s to complete 50 independent simulations on Lorenz-63 for the EnRDA compared to 651 (590)\,s for the particle filter (EnKF) with 100 particles (ensemble members). Since the computational cost is nonlinearly related to the problem dimension, it is expected that it grows significantly for large-scale geophysical DA and becomes a limiting factor. Although the entropic regularization works well for the presented low dimensional problems, future research is needed to test its efficiency in high-dimensional problems. Constraining the solution of the coupling on a submanifold of probability distributions with a Gaussian mixture structure \citep{chen2019optimal} can be a future research direction for lowering the computational cost.


\section*{Acknowledgements}
The first and second author acknowledge the grant from the National Aeronautics and Space Administration (NASA) Terrestrial Hydrology Program (THP, 80NSSC18K1528) and the New (Early Career) Investigator Program (NIP, 80NSSC18K0742). The third author acknowledges support from the European Research Council for funding via the Horizon2020 CUNDA project under number 694509. The fifth author also acknowledges support from National Science Foundation (NSF, DMS1830418).

\newpage
