Ensemble Riemannian data assimilation: towards large-scale dynamical systems

This paper presents the results of the ensemble Riemannian data assimilation for relatively high-dimensional nonlinear dynamical systems, focusing on the chaotic Lorenz-96 model and a two-layer quasi-geostrophic (QG) model of atmospheric circulation. The analysis state in this approach is inferred from a joint distribution that optimally couples the background probability distribution and the likelihood function, enabling formal treatment of systematic biases without any Gaussian assumptions. Despite the risk of the curse of dimensionality in the computation of the coupling distribution, comparisons with the classic implementation of the particle filter and the stochastic ensemble Kalman filter demonstrate that, with the same ensemble size, the presented methodology could improve the predictability of dynamical systems. In particular, under systematic errors, the root mean squared error of the analysis state can be reduced by 20 % (30 %) in the Lorenz-96 (QG) model.

1 Introduction

The science of data assimilation (DA) aims to optimally estimate the probability distribution of a state variable of interest in an Earth system model (ESM) given the information content of observations and previous time forecasts to improve their predictive abilities (Kalnay, 2003; Carrassi et al., 2018). Current DA methodologies, either variational (Lorenc, 1986; Zupanski, 1993; Courtier et al., 1994; Rabier et al., 2000; Poterjoy and Zhang, 2014) or filtering (Kalman, 1960; Bishop et al., 2001; Anderson, 2001; Tippett et al., 2003; Janjić et al., 2011; Carrassi and Vannitsem, 2011; Anderson and Lei, 2013; Lei et al., 2018), largely rely on penalization of second-order statistics of errors over the Euclidean space without explicitly accounting for systematic model biases. For example, in the three-dimensional variational (3D-Var) DA (Lorenc, 1986; Courtier et al., 1998; Lorenc et al., 2000; Li et al., 2013), a least-squares cost function comprising weighted Euclidean distances of the state from the previous model forecasts (background state) and the observations is formulated. Solution of this cost function leads to an analysis state which is a weighted average of the forecasts and observations across multiple dimensions of the problem with the weights dictated by prescribed background and observation error covariance matrices. The variants of the Kalman filtering DA methods (Evensen, 1994; Reichle et al., 2002; Evensen, 2003; Nerger et al., 2012b; Houtekamer and Zhang, 2016) also follow the same principle, but in these methods, the background covariance contains information from past observations and model evolution.

Apart from the Euclidean distance, other measures and distance metrics, including the quadratic mutual information (Kapur, 1994), Kullback–Leibler (KL) divergence (Kullback and Leibler, 1951), Hellinger distance (Hellinger, 1909), and Wasserstein distance (Villani, 2003) have also been utilized in DA frameworks. Among others, Tagade and Ravela (2014) introduced a nonlinear filter, where the analysis is obtained through maximization of the quadratic mutual information. Maclean et al. (2017) utilized the Hellinger distance to measure the difference between the predicted and observed spatial patterns in oceanic flows. Chianese et al. (2018) introduced a variational DA method in which minimization of the KL divergence led to an approximation of the bias terms and model parameters. Similarly, Li et al. (2019) employed the KL divergence in an optimization framework to incorporate inequality constraints into the ensemble Kalman filter (EnKF, Evensen, 1994). Recently, Pulido and van Leeuwen (2019) developed a mapping particle filter in which particles are pushed towards the posterior density by minimizing the KL divergence between the posterior and a series of intermediate probability densities.

In filtering class of DA methodologies, coupling techniques have been proposed as an alternative to the classic Bayesian inference (El Moselhy and Marzouk, 2012; Spantini et al., 2019). El Moselhy and Marzouk (2012) presented a new approach to find an optimal transport map that pushes forward the background to the posterior distribution. The approach was extended for generalization of the EnKF by deriving nonlinear coupling between the forecast and posterior distributions (Spantini et al., 2019). In recent years, the Wasserstein or Earth mover's distance, originating from the theory of optimal mass transport (OMT, Monge, 1781; Kantorovich, 1942; Villani, 2003; Kolouri et al., 2017; Chen et al., 2017; Y. Chen et al., 2018; B. Chen et al., 2019), has also been gaining attention in the DA community. Reich (2013) introduced a new resampling approach in particle filters, using the OMT, to maximize the correlation between the prior and posterior ensemble members. Ning et al. (2014) further utilized the Wasserstein distance to treat position errors arising from uncertain model parameters. Following on this work, Feyeux et al. (2018) proposed to replace the weighted Euclidean distance with the Wasserstein distance in variational DA frameworks to treat position error. Tamang et al. (2020) proposed to use the Wasserstein distance to regularize a variational DA framework for treating systematic errors arising from the model forecast in chaotic systems.

However, DA frameworks utilizing the Wasserstein distance are computationally expensive as they require a joint distribution to be obtained that couples two marginal distributions. Finding this joint distribution often relies on interior-point optimization methods (Altman and Gondzio, 1999) or Orlin's algorithm (Orlin, 1993) that have super-cubic run time – making the Wasserstein DA computationally challenging even for relatively low-dimensional problems. More recently, to reduce the computational cost, Tamang et al. (2021) used entropic regularization of the OMT formulation (Cuturi, 2013) through a new framework, called ensemble Riemannian data assimilation (EnRDA), to cope with systematic errors and tested it on a three-dimensional Lorenz-63 model (Lorenz, 1963).

Unlike Euclidean DA with a known connection with the family of Gaussian distributions through Bayes' theorem, the EnRDA does not rely on any parametric assumptions about the input probability distributions. Therefore, it does not guarantee an analysis state with a minimum mean squared error. However, it enables us to optimally (i) interpolate between the forecast distribution and the normalized likelihood function without any parametric assumptions about their shapes and (ii) formally penalize systematic translations between them arising due to potential geophysical biases.

However, the computational complexity of finding an optimal joint coupling between two m-dimensional probability distributions supported on d points in each dimension using the entropic regularization is O(d^2m). This might impose a significant limitation on the direct use of EnRDA for high-dimensional geophysical problems, where the problem dimension easily exceeds millions. As will be discussed later, the joint distribution in EnRDA is sampled at N² support points, with N number of ensembles, reducing the computational complexity to O(N²) at the expense of losing accuracy in optimal estimation of the joint distribution. Therefore, beyond implementation in a low-dimensional dynamical system, such as the three-dimensional Lorenz-63, the key questions that we aim to answer are as follows. Does the effectiveness of the EnRDA implementation still remain valid in high-dimensional DA problems where the ensemble size is smaller than the problem dimension? How does EnRDA perform, under systematic errors, in comparison to classic ensemble DA techniques with a comparable ensemble size? To answer the above questions, we implement EnRDA in the relatively high-dimensional chaotic Lorenz-96 system (Lorenz, 1995) and a two-layer quasi-geostrophic (QG) model of atmospheric circulation (Pedlosky, 1987). The results demonstrate that EnRDA can potentially enhance predictability of high-dimensional dynamical systems, when the state variables are not necessarily Gaussian and are corrupted with systematic errors. Nevertheless, extensive future research is necessary to test the applicability of the EnRDA for large-scale Earth system models in which the problem dimension is significantly larger than the examined dynamical systems.

The outline of the paper is as follows. Section 2 provides a brief background on optimal mass transport and Wasserstein distance. A brief review of the EnRDA methodology is presented in Sect. 3. Section 4 presents different test cases of implementation on the Lorenz-96 and QG models and documents the performance of the presented approach in comparison with the classic implementation of the standard particle filter with resampling and the stochastic ensemble Kalman filter (SEnKF). A summary and concluding remarks are presented in Sect. 5. The details of the entropic regularization for the EnRDA and covariance inflation and localization procedures for the SEnKF are provided in Appendix A.

2 Background on OMT and the Wasserstein barycenter

We provide a brief background on the theory of optimal mass transport (OMT) and Wasserstein barycenters. The OMT theory, first put forward by Monge (1781), aims to find the minimum cost of transporting distributed masses of materials from known source points to target points. The theory was later expanded as a new tool to compare probability distributions (Brenier, 1987; Villani, 2003) and since then has found its applications in the field of data assimilation (Ning et al., 2014; Feyeux et al., 2018; Li et al., 2018; Tamang et al., 2020), subsurface geophysical inverse problems (J. Chen et al., 2018; Yang et al., 2018; Yang and Engquist, 2018; Yong et al., 2019) and comparisons of climate model simulations (Vissio et al., 2020).

Let us consider a discrete source probability distribution $p\left(\mathit{x}\right)={\sum }_{i=\mathrm{1}}^M{p}_{{\mathit{x}}_i}{\mathit{\delta }}_{{\mathit{x}}_i}$ and a target distribution $p\left(\mathit{y}\right)={\sum }_{j=\mathrm{1}}^N{p}_{{\mathit{y}}_j}{\mathit{\delta }}_{{\mathit{y}}_j}$ with their respective probability masses $\mathit{\{}{\mathbf{p}}_x\in {\mathbb{R}}_{+}^M:{\sum }_i{p}_{{\mathit{x}}_i}=\mathrm{1}\mathit{\}}$ and $\mathit{\{}{\mathbf{p}}_y\in {\mathbb{R}}_{+}^N:{\sum }_j{p}_{{\mathit{y}}_j}=\mathrm{1}\mathit{\}}$ supported on m- and n-element column vectors x_i∈ℝ^m and y_j∈ℝⁿ, respectively. The notation ${\mathbf{p}}_x\in {\mathbb{R}}_{+}^M$ represents probability masses p_x containing non-negative real numbers supported on M points, whereas δ_x is the Dirac function at x. In the Monge formulation, the goal is to seek an optimal surjective transportation map $T_{\mathit{\#}}^{a}p\left(\mathit{x}\right)=p\left(\mathit{y}\right)$ that “pushes forward” the source distribution p(x) towards the target distribution p(y), with a minimum transportation cost as follows:

$$\begin{array}{}\text{(1)}& T^a=\underset{T}{\mathrm{argmin}}{\sum }_{i=\mathrm{1}}^{M}c({\mathit{x}}_i,T({\mathit{x}}_i\left)\right)\mathrm{s}.\mathrm{t}.T_{\mathit{\#}}^{a}p\left(\mathit{x}\right)=p\left(\mathit{y}\right),\end{array}$$

where $c(\cdot ,\cdot )\in {\mathbb{R}}_{+}$ represents the cost of transporting a unit mass from one support point in x to another one in y.

The problem formulation by Monge as expressed in Eq. (1), however, is non-convex, and the existence of an optimal transportation map is not guaranteed (Y. Chen et al., 2019), especially when the number of support points for the target distribution exceeds that of the source distribution (N>M) (Peyré and Cuturi, 2019). This limitation was overcome by Kantorovich (1942), who introduced a probabilistic formulation of OMT – allowing splitting of probability mass from a single source point to multiple target points. The Kantorovich formalism recasts the OMT problem in a linear programming framework that finds an optimal joint distribution or coupling ${\mathbf{U}}^a\in {\mathbb{R}}_{+}^{M\times N}$ that couples the marginal source and target distributions with the following optimality criterion:

$$\begin{array}{}\text{(2)}& {\mathbf{U}}^a=\underset{\mathbf{U}}{\mathrm{argmin}}\mathrm{tr}\left({\mathbf{C}}^{\mathrm{T}}\mathbf{U}\right)\mathrm{s}.\mathrm{t}.\left\{\begin{array}{l}\mathbf{U}\in {\mathbb{R}}_{+}^{M\times N}\\ \mathbf{U}{\text{𝟙}}_N={\mathbf{p}}_x,\\ {\mathbf{U}}^{\mathrm{T}}{\text{𝟙}}_M={\mathbf{p}}_y\end{array}\right.\end{array}$$

where tr(⋅) is the trace of a matrix, (⋅)^T is the transposition operator and 𝟙_M represents an M-element column vector of ones. In the above formulation, the known $\left\{\mathbf{C}\in {\mathbb{R}}_{+}^{M\times N}:c_{ij}=\Vert {\mathit{x}}_i-{\mathit{y}}_j{\Vert }_{\mathrm{2}}^{\mathrm{2}}\right\}$ denotes the so-called transportation cost matrix which is defined based on the ℓ₂ norm $\Vert \cdot {\Vert }_{\mathrm{2}}$ or the Euclidean distance between the support points of the source and target distributions. Here, the (i,j)th element $u_{ij}^a$ of optimal joint distribution U^a represents the respective amount of mass transported from support point x_i to y_j. Then, the 2-Wasserstein distance or metric between the marginal probability distributions is defined as the square root of the optimal transportation cost $d_{\mathcal{W}}({\mathbf{p}}_x,{\mathbf{p}}_y)=\left(\mathrm{tr}\right({\mathbf{C}}^{\mathrm{T}}{\mathbf{U}}^a){)}^{\frac{\mathrm{1}}{\mathrm{2}}}$ (Dobrushin, 1970; Villani, 2008). It should be noted that due to the linear equality and non-negativity constraints in Eq. (2), the family of joint distributions that satisfy these constraints forms a bounded convex polytope (Cuturi and Peyré, 2018) and, consequently, the optimal joint distribution U^a is located on one of the extreme points of such a polytope (Peyré and Cuturi, 2019).

Recall that, over the Euclidean space, the barycenter of a group of points is equivalent to their (weighted) mean value. The Wasserstein metric offers a Riemannian generalization of this problem and allows us to define the barycenter of a family of probability distributions (Rabin et al., 2011; Bigot and Klein, 2018; Srivastava et al., 2018). In particular, for a group of K probability mass functions ${\mathbf{p}}_{\mathrm{1}},\mathrm{\dots },{\mathbf{p}}_K$, a Wasserstein barycenter p_η is defined as their Fréchet mean (Fréchet, 1948) as follows (Agueh and Carlier, 2011):

$$\begin{array}{}\text{(3)}& {\mathbf{p}}_{\mathit{\eta }}=\underset{\mathbf{p}}{\mathrm{argmin}}\sum _{k=\mathrm{1}}^K{\mathit{\eta }}_k{d}_{\mathcal{W}}^{\mathrm{2}}(\mathbf{p},{\mathbf{p}}_k),\end{array}$$

where $\left\{({\mathit{\eta }}_{\mathrm{1}},\mathrm{\dots },{\mathit{\eta }}_K{)}^{\mathrm{T}}\in {\mathbb{R}}_{+}^K:{\sum }_k{\mathit{\eta }}_k=\mathrm{1}\right\}$ represent the weights associated with the respective distributions. In special cases where the group of K distributions is Gaussian $\left\{\mathcal{N}({\mathit{\mu }}_{\mathrm{1}},{\mathbf{\Sigma }}_{\mathrm{1}}),\mathrm{\dots },\mathcal{N}({\mathit{\mu }}_K,{\mathbf{\Sigma }}_K)\right\}$ with mean ${\mathit{\mu }}_{\mathrm{1}},\mathrm{\dots },{\mathit{\mu }}_K$ and positive definite covariance ${\mathbf{\Sigma }}_{\mathrm{1}},\mathrm{\dots },{\mathbf{\Sigma }}_K$, the Wasserstein barycenter is also a Gaussian density 𝒩(μ_η,Σ_η) with ${\mathit{\mu }}_{\mathit{\eta }}=\sum _k{\mathit{\eta }}_k{\mathit{\mu }}_k$ and Σ_η is the unique positive definite root of the matrix equation $\mathbf{\Sigma }=\sum _k{\mathit{\eta }}_k({\mathbf{\Sigma }}^{\frac{\mathrm{1}}{\mathrm{2}}}{\mathbf{\Sigma }}_k{\mathbf{\Sigma }}^{\frac{\mathrm{1}}{\mathrm{2}}}{)}^{\frac{\mathrm{1}}{\mathrm{2}}}$ (Agueh and Carlier, 2011).

3 EnRDA

In this section, to be self content, we provide a brief summary of the EnRDA methodology, while more details can be found in Tamang et al. (2021). Let us assume that the evolution of the ith ensemble member x_i∈ℝ^m of ESM simulations can be presented as the following stochastic dynamical system:

$$\begin{array}{}\text{(4)}& {\mathit{x}}_i^t=\mathcal{M}\left({\mathit{x}}_i^{t-\mathrm{1}}\right)+{\mathit{\omega }}_i^{t}i=\mathrm{1},\mathrm{\dots },M,\end{array}$$

where $\mathcal{M}:{\mathbb{R}}^m\to {\mathbb{R}}^m$ is the deterministic nonlinear model operator, evolving the model state in time with a stochastic error term ${\mathit{\omega }}_i^t\in {\mathbb{R}}^m$. This dynamical system is observed at time t through an observation equation ${\mathit{y}}^t=\mathcal{H}\left({\mathit{x}}^t\right)+{\mathit{\upsilon }}^t$, where $\mathcal{H}:{\mathbb{R}}^m\to {\mathbb{R}}^n$ maps the state to the observation space and υ^t∈ℝⁿ represents an additive observation error. Note that the error terms are not necessarily drawn from Gaussian distributions but need to have finite second-order moments.

Hereafter, we drop the time superscript for brevity and represent the model (or background) probability distribution as $p\left(\mathit{x}\right)={\sum }_{i=\mathrm{1}}^M{p}_{{\mathit{x}}_i}{\mathit{\delta }}_{{\mathit{x}}_i}$ with its probability mass vector $\mathit{\{}{\mathbf{p}}_x\in {\mathbb{R}}_{+}^M:{\sum }_i{p}_{{\mathit{x}}_i}=\mathrm{1}\mathit{\}}$. Furthermore, the normalized likelihood function is represented as $\stackrel{\mathrm{̃}}{p}\left(\mathit{y}\right|\mathit{x})$ centered at the given observation y with its probability mass vector $\mathit{\{}{\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}\in {\mathbb{R}}_{+}^N:{\sum }_j{\stackrel{\mathrm{̃}}{p}}_{\mathit{y}|{\mathit{x}}_j}=\mathrm{1}\mathit{\}}$. The probability distribution of the analysis state, p(x_a), is then defined as the Wasserstein barycenter between forecast distribution and the normalized likelihood function:

$$\begin{array}{}\text{(5)}& \begin{array}{rl}p\left({\mathit{x}}_a\right)=\underset{p\left(\mathit{z}\right)}{\mathrm{argmin}}\mathit{\{}& \mathit{\eta }{d}_{\mathcal{W}}^{\mathrm{2}}\left[p\left(\mathit{x}\right),p\left(\mathit{z}\right)\right]+(\mathrm{1}-\mathit{\eta })d_{\mathcal{W}}^{\mathrm{2}}\\ & \left[\stackrel{\mathrm{̃}}{p}\left(\mathit{y}\right|\mathit{x}),p(\mathit{z})\right]\mathit{\}},\end{array}\end{array}$$

where $\mathit{\eta }\in [\mathrm{0},\mathrm{1}]$ is a displacement parameter that controls the relative weight of the background and observation. The displacement parameter η is a hyperparameter that captures the relative weights of the histogram of the background state and likelihood function in characterization of the analysis state distribution as a Wasserstein barycenter. The optimal value of η needs to be determined offline using reference data through cross-validation studies. It is important to note that the above formalism requires all dimensions to be observable, and thus those dimensions with no observations cannot be updated, which is a limitation of the current formalism compared to the Euclidean DA. This limitation is further discussed later on in Sect. 5.

To solve the above DA problem, we need to characterize the background distribution and the normalized likelihood function. Similar to the approach used in the particle filter (Gordon et al., 1993; van Leeuwen, 2010), we suggest approximating them through ensemble realizations. To construct the histogram of the normalized likelihood function, we can draw N samples at each assimilation cycle by perturbing the available observation y with the observation error 𝒩(0,R).

To obtain the Wasserstein barycenter p(x_a) in Eq. (5), we use the McCann formalism (McCann, 1997; Peyré and Cuturi, 2019):

$$\begin{array}{}\text{(6)}& p\left({\mathit{x}}_a\right)=\sum _{i=\mathrm{1}}^M\sum _{j=\mathrm{1}}^N{u}_{ij}^a{\mathit{\delta }}_{{\mathit{z}}_{ij}},\end{array}$$

where ${\mathit{z}}_{ij}=\mathit{\eta }{\mathit{x}}_i+(\mathrm{1}-\mathit{\eta }){\mathit{y}}_j$ represent the support points of the analysis distribution and $u_{ij}^a$ are the elements of the joint distribution $\left\{{\mathbf{U}}^a\in {\mathbb{R}}_{+}^{M\times N}:{\sum }_i{\sum }_j{u}_{ij}=\mathrm{1}\right\}$. It is important to note that the analysis state histogram, at each assimilation cycle, is supported on at most $M+N-\mathrm{1}$ points, which is the maximum number of non-zero entries in the optimal joint coupling (Peyré and Cuturi, 2019). To keep the number of ensemble members constant throughout, M ensemble members are resampled from p(x_a) using the multinomial resampling scheme (Li et al., 2015).

Computation of the joint distribution in Eq. (2) is computationally expensive as explained previously and can be prohibitive for high-dimensional geophysical problems. As suggested by Cuturi (2013), to reduce the computational cost, we regularize the cost function in the optimal transportation plan formulation of EnRDA by a Gibbs–Boltzmann entropy function:

$$\begin{array}{}\text{(7)}& \begin{array}{rl}{\mathbf{U}}^a=& \underset{\mathbf{U}}{\mathrm{argmin}}\mathrm{tr}\left({\mathbf{C}}^{\mathrm{T}}\mathbf{U}\right)-\mathit{\gamma }\mathrm{tr}\left({\mathbf{U}}^{\mathrm{T}}\left[\log (\mathbf{U}-{\text{𝟙}}_M{\text{𝟙}}_N^{\mathrm{T}})\right]\right)\\ & \mathrm{s}.\mathrm{t}.\left\{\begin{array}{l}\mathbf{U}\in {\mathbb{R}}_{+}^{M\times N}\\ \mathbf{U}{\text{𝟙}}_N={\mathbf{p}}_x,\\ {\mathbf{U}}^{\mathrm{T}}{\text{𝟙}}_M={\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}\end{array}\right.\end{array}\end{array}$$

where $\mathit{\gamma }\in {\mathbb{R}}_{+}$ is a regularization parameter. The entropic regularization transforms the original OMT formulation to a strictly convex problem, which can then be efficiently solved using Sinkhorn's algorithm (Sinkhorn, 1967). The details of Sinkhorn's algorithm for solving regularized optimal transportation problems are presented in Appendix A1. The regularization parameter γ balances the solution between the optimal joint distribution and the one that maximizes the relative entropy. It is evident from Eq. (7) that at the limit γ→0, the solution of Eq. (7) converges to the analysis joint distribution with a minimum morphing cost. However, as the value of γ increases, the convexity of the problem also increases, enabling the deployment of more efficient optimization algorithms than classic solvers of linear programming problems (Dantzig et al., 1955; Orlin, 1993). At the same time, the number of non-zero entries of the joint coupling increases from $M+N-\mathrm{1}$ to MN points as γ→∞, which results in a maximum entropy solution that converges to ${\mathbf{U}}^a\to {\mathbf{p}}_x{\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}^{\mathrm{T}}$. For a more comprehensive explanation of EnRDA, one can refer to Tamang et al. (2021). It should be noted that EnRDA formulation allows the number of ensemble members to be different from the number of perturbed observations, i.e., M≠N. The values of M and N need to be chosen to adjust the trade-off between accuracy and computational cost.

As an example, we examine here the solution of Eq. (5) between a banana-shaped distribution denoted by $\mathcal{F}({\mathit{\xi }}_{\mathrm{1}},{\mathit{\xi }}_{\mathrm{2}},{\mathit{\xi }}_{\mathrm{3}},b)\propto \exp (-{\mathit{\xi }}_{\mathrm{1}}(\mathrm{4}-b{x}_{\mathrm{1}}-x_{\mathrm{2}}^{\mathrm{2}})-{\mathit{\xi }}_{\mathrm{2}}(x_{\mathrm{1}}^{\mathrm{2}}-{\mathit{\xi }}_{\mathrm{3}}{x}_{\mathrm{2}}^{\mathrm{2}}\left)\right)$ and a bivariate Gaussian distribution as a function of the displacement parameter $\mathit{\eta }\in [\mathrm{0},\mathrm{1}]$ – resembling the background distribution p(x) and the normalized likelihood function $\stackrel{\mathrm{̃}}{p}\left(\mathit{y}\right|\mathit{x})$, respectively, with regularization parameter γ=1000. As seen from Fig. 1, for lower values of η, the analysis state distribution is closer to the observation, and its shape resembles the Gaussian distribution. However, as the value of η increases, the analysis state distribution moves closer to the background distribution and starts morphing into a banana-shaped distribution. Therefore, the analysis state distribution is defined as the one that is sufficiently close to the background distribution and the normalized likelihood function not only based on their shape, but also their central location – depending on the displacement parameter. Thus, unlike the Euclidean barycenter, this approach does not guarantee that the mean or mode of the analysis state probability distribution is a minimum mean-squared error estimate of the initial condition.

It is important to note that in the original OMT formulation, the number of support points required for the optimal joint coupling U scales with the problem dimension (d^m), making it potentially restrictive for the high-dimensional problems, where d represents the number of support points in each dimension and m is the number of dimensions. The presented EnRDA setting bypasses this problem by sampling the joint distribution using only N ensemble members. However, such approximation might lead to errors in optimal estimation of the joint coupling that are translated into the analysis state. In the next section, we present results from systems of well-known dynamics to investigate whether EnRDA can lead to a proper approximation of the analysis state under systematic errors in relatively high-dimensional nonlinear problems when compared to classic ensemble DA approaches with comparable ensemble size.

Figure 1 The analysis distribution obtained as a Wasserstein barycenter for different values of the displacement parameter $\mathit{\eta }\in [\mathrm{0},\mathrm{1}]$ between a background distribution represented by a banana-shaped distribution $p\left(\mathit{x}\right):\mathcal{F}({\mathit{\xi }}_{\mathrm{1}},{\mathit{\xi }}_{\mathrm{2}},{\mathit{\xi }}_{\mathrm{3}},b)$ with ξ₁=0.02, ξ₂=0.06, ξ₃=1.6, and b=8 and the normalized likelihood function represented by a bivariate Gaussian $\stackrel{\mathrm{̃}}{p}\left(\mathit{y}\right|\mathit{x}):\mathcal{N}({\mathit{\mu }}_{\mathrm{1}},{\mathbf{\Sigma }}_{\mathrm{1}})$, where ${\mathit{\mu }}_{\mathrm{1}}=\left[\begin{array}{c}-\mathrm{35}\\ \mathrm{0}\end{array}\right]$ and ${\mathbf{\Sigma }}_{\mathrm{1}}=\left[\begin{array}{cc}\mathrm{3}& \mathrm{0}\\ \mathrm{0}& \mathrm{2}\end{array}\right]$.

4 Numerical experiments and results

4.1 Lorenz-96

The Lorenz model (Lorenz-96, Lorenz, 1995), which is widely adopted as a test bed for numerous DA experiments (Trevisan and Palatella, 2011; Tang et al., 2014; Shen and Tang, 2015; Lguensat et al., 2017; Tian et al., 2018), offers a simplified representation of the extra-tropical dynamics in the Earth's atmosphere. The model coordinates $\left\{\mathit{x}=(x_{\mathrm{1}},\mathrm{\dots },x_K{)}^{\mathrm{T}}\in {\mathbb{R}}^K\right\}$ at K dimensions represent the state of an arbitrary atmospheric quantity measured along the Earth's latitudes at K equally spaced longitudinal slices. The model is designed to mimic the continuous-time variation in atmospheric quantities due to interactions between three major components, namely, advection, internal dissipation, and external forcing. The model dynamics is represented as follows:

$$\begin{array}{}\text{(8)}& {\displaystyle \frac{\mathrm{d}{x}_k}{\mathrm{d}t}}=\left(x_{k+\mathrm{1}}-x_{k-\mathrm{2}}\right)x_{k-\mathrm{1}}-x_k+F,k=\mathrm{1},\mathrm{\dots },K,\end{array}$$

where $F\in {\mathbb{R}}_{+}$ is a constant external forcing independent of the model state. The Lorenz-96 model has cyclic boundaries with $x_{-\mathrm{1}}=x_{K-\mathrm{1}}$, x₀=x_K, and $x_{K+\mathrm{1}}=x_{\mathrm{1}}$. It is known that for small values of $F<\mathrm{8}/\mathrm{9}$, the system approaches a steady-state condition with each coordinate value converging to the external forcing x_k→F, ∀k, whereas for $F>\mathrm{8}/\mathrm{9}$, chaos develops (Lorenz and Emanuel, 1998). For a standard model setup with F=8, the system is known to exhibit highly chaotic behavior with the largest Lyapunov exponent of 1.67 (Brajard et al., 2020).

Experimental setup, results, and discussion

We focus on the 40-dimensional Lorenz-96 system (i.e., K=40) and compare EnRDA results with the classic implementation of the PF (Gordon et al., 1993; Van Leeuwen, 2009; van Leeuwen, 2010; Poterjoy and Anderson, 2016) and the SEnKF (Evensen, 1994; Houtekamer and Mitchell, 1998; Burgers et al., 1998; Janjić et al., 2011; Anderson, 2016; Van Leeuwen, 2020). Similarly to the experimental setting suggested in Lorenz and Emanuel (1998) and Nerger et al. (2012a), we initialize the model by choosing x₂₀=8.008 and x_k=8 for all other model coordinates. In order to avoid any initial transient effect, the model in Eq. (8) is integrated for 1000 time steps using the fourth-order Runge–Kutta approximation (Runge, 1895; Kutta, 1901) with a non-dimensional time step of Δt=0.01, and the endpoint of the run is utilized as the initial condition for DA experimentation.

Similar to the suggested experimental setting in van Leeuwen (2010), we obtain the ground truth by integrating Eq. (8) with a time step of Δt over a time period of T= 0–20 in the absence of any model error. The observations are assumed to be available at each assimilation time interval of 10Δt and deviated from the ground truth by a Gaussian error ${\mathit{\upsilon }}_t\sim \mathcal{N}(\mathrm{0},{\mathit{\sigma }}_{\mathrm{obs}}^{\mathrm{2}}{\mathbf{\Sigma }}_{\mathit{\rho }})$, with ${\mathit{\sigma }}_{\mathrm{obs}}^{\mathrm{2}}=\mathrm{1}$ and the correlation matrix ${\mathbf{\Sigma }}_{\mathit{\rho }}\in {\mathbb{R}}_{+}^{\mathrm{40}\times \mathrm{40}}$ with 1 on the diagonals, 0.5 on the first sub- and super-diagonals, and 0 everywhere else. The observation time step of 10Δt is equivalent to 12 h in global ESMs (Lorenz, 1995).

To characterize the distribution of the background state for each DA methodology, 50 (5000) ensemble members (particles) for the SEnKF and EnRDA (PF) are generated using model errors ${\mathit{\omega }}_t\sim \mathcal{N}(\mathrm{0},{\mathit{\sigma }}_t^{\mathrm{2}}{\mathbf{I}}_{\mathrm{40}})$ with ${\mathit{\sigma }}_t^{\mathrm{2}}=\mathrm{0.25}$ for t>0 and ${\mathit{\sigma }}_{\mathrm{0}}^{\mathrm{2}}=\mathrm{4}$, where throughout I_m represents an m×m identity matrix. To alleviate the known degeneracy problem in the PF, a higher number of particles was used. Furthermore, to introduce additional systematic background error, we utilize an erroneous external forcing of F_m=6 instead of the “true” forcing value F=8. To have a robust inference, the average values of the error metrics are reported for 50 experiments using different random realizations. As will be elaborated later on, we set the EnRDA displacement parameter η=0.44, determined through a cross-validation study based on a minimum mean-squared error criterion. This tuning is similar to tuning inflation and localization parameters in a typical EnKF or tuning length scales in 3D- or 4D-Var. Note that we already introduced some systematic error because the truth has zero model error, while the prior does have model errors. In a fully unbiased setup the truth and the prior are drawn from the same distribution.

The results of EnRDA are shown in Fig. 2. In Fig. 2a, the temporal evolution of the ground truth and EnRDA analysis state is shown over all dimensions of Lorenz-96, while a snapshot at time 10 (t) is presented in Fig. 2b. The analysis state obtained from EnRDA follows the ground truth reasonably well during all time steps, with a root mean squared error (RMSE) of 0.85. The comparison of EnRDA with the classic implementations of the SEnKF and PF are shown in Fig. 3. It can be seen that the RMSE of the PF increases sharply over time, suggesting that the problem of filter degeneracy still exists despite the higher number of particles. This problem is exacerbated due to the presence of bias causing a rapid collapse of the ensemble variance over time as more particles fall outside of the support set of the likelihood function. The root mean squared error of both the SEnKF and EnRDA is stabilized over time and is smaller by ∼20 % (80 %) in EnRDA compared to the SEnKF (PF). It is important to note that the presence of systematic bias due to erroneous choice of the external forcing inherently favors EnRDA over SEnKF since the latter is a minimum variance unbiased estimator at the limit M→∞, where M represents the number of ensemble members.

Figure 2 (a) Temporal evolution of the ground truth x_tr and analysis state x_a by ensemble Riemannian data assimilation (EnRDA) for K=40 dimensions of Lorenz-96 over T=0–20 (t) and (b) their snapshots at T=10 (t) together with the available observations y.

Figure 3 Temporal evolution of the root mean squared error (RMSE) for the (a) particle filter (PF) with 5000 particles, (b) stochastic ensemble Kalman filter (SEnKF), and (c) EnRDA, each with 50 ensemble members in a 40-dimensional Lorenz-96 system. The results report the mean values of 50 independent simulations.

As previously noted, the displacement parameter η plays an important role in EnRDA as it controls the shape and position of the analysis state distribution relative to the background distribution and the normalized likelihood function. Currently, there exists no known closed-form solution for optimal approximation of this parameter. Therefore, in this paper, we focus on determining its optimal value through heuristic cross-validation by an offline bias-variance trade-off analysis. Specifically, we quantify the RMSE of the EnRDA analysis state for different values of η for 50 independent simulations.

The bias and RMSE, together with their respective 5th–95th percentile bounds, as functions of the displacement parameter η are shown in Fig. 4a. As explained earlier, when η increases, the analysis distribution moves towards the background distribution. Since the background state is systematically biased due to the erroneous external forcing, the analysis bias increases monotonically with η, while the RMSE shows a minimum point. Therefore, there exists a form of bias-variance trade-off in the analysis error which leads to an approximation of an optimal value of η based on a minimum RMSE criterion. It is important to note that the background uncertainty and thus the optimal value of η vary in response to the ensemble size as shown in Fig. 4b. The reason is that a larger number of ensemble members reduces the uncertainty in the characterization of the background, but the bias is not affected. To compensate, a larger optimal value for η is needed. This optimal value approaches an asymptotic value as the ensemble sample size increases and will achieve the highest value at the limit M→∞, when the sample moments converge to the biased forecast moments.

Figure 4 (a) Bias and RMSE for a range of displacement parameter $\mathit{\eta }\in [\mathrm{0.1},\mathrm{0.6}]$ in EnRDA with 50 ensemble members, obtained across 40 dimensions of the Lorenz-96 system. The shaded regions indicate the 5th–95th percentile bound for the respective error metrics obtained from 50 independent simulations. (b) Variation of RMSE as a function of the number of ensemble members and η.

One may argue that such a tuning favors EnRDA since it explicitly accounts for the effects of bias, either in background or observations, while there is no bias-correction mechanism in the implementation of the SEnKF and the PF. To make a fairer comparison, we investigate an alternative approach to approximate the displacement parameter solely based on the known error covariance matrices at each assimilation cycle. Recall that in classic DA, the analysis state is essentially the Euclidean barycenter, where the relative weights of the background state and observations are optimally characterized based on the error covariances under zero bias assumptions. However, over the Wasserstein space, the displacement parameter determines the weight between the entire distribution of the background and the normalized likelihood function. Theoretically, knowing the Wasserstein distances from ground truth to both likelihood function and forecast probability distribution enables us to obtain an optimal value for η. Even though such distances are not known in reality, the total Wasserstein distance between the normalized likelihood function and the forecast distribution is known at each assimilation cycle. Therefore, an estimate of the distance between the ground truth and the normalized likelihood function or the forecast distribution leads to an approximation of η.

It is known that the square of the Wasserstein distance between two equal-mean Gaussian distributions 𝒩(μ,Σ₁) and 𝒩(μ,Σ₂) is $d_{\mathcal{W}}^{\mathrm{2}}=\mathrm{tr}({\mathbf{\Sigma }}_{\mathrm{1}}+{\mathbf{\Sigma }}_{\mathrm{2}}-\mathrm{2}({\mathbf{\Sigma }}_{\mathrm{1}}^{\frac{\mathrm{1}}{\mathrm{2}}}{\mathbf{\Sigma }}_{\mathrm{2}}{\mathbf{\Sigma }}_{\mathrm{1}}^{\frac{\mathrm{1}}{\mathrm{2}}}{)}^{\frac{\mathrm{1}}{\mathrm{2}}})$ (Y. Chen et al., 2019). Therefore, under the assumption that only the background state is biased, the square of the Wasserstein distance between the true state x_tr, as a Dirac delta function, and the normalized likelihood function reduces to tr(R). At the same time, the square of the Wasserstein distance between the normalized likelihood function and forecast distribution is tr(C^TU^a). Therefore, we can approximate the interpolation parameter as $\stackrel{\mathrm{‾}}{{\mathit{\eta }}_a}=\mathrm{tr}\left(\mathbf{R}\right){\left(\mathrm{tr}\left({\mathbf{C}}^{\mathrm{T}}{\mathbf{U}}^a\right)+\mathrm{tr}\left(\mathbf{R}\right)\right)}^{-\mathrm{1}}$ without any explicit a priori knowledge of bias.

Comparisons of the RMSE values for the studied DA methodologies as a function of ensemble size are shown in Fig. 5. For EnRDA, the displacement parameter is obtained from the bias-aware cross-validation (η=0.44, EnRDA-I) and from the known error covariances as explained above (EnRDA-II). The SEnKF and EnRDA result in smaller error metrics with a much smaller ensemble size than the PF. As seen, EnRDA can perform well even for smaller ensemble sizes as low as 20. Its results quickly stabilize with more than 40 ensemble members and exhibit a marginal improvement over the SEnKF (12 %–24 %) in the presence of bias. The RMSE of the SEnKF also stabilizes quickly but remains above the standard deviation of the observation error, indicating that in the presence of bias, the lowest possible variance, known as the Cramer–Rao lower bound (Cramér, 1999; Rao et al., 1973), cannot be met.

Figure 5 The RMSE for the different number of ensemble members/particles in the PF, SEnKF, and EnRDA when the displacement parameter is obtained from bias-aware cross-validation (ENRDA-I) and a dynamic approach without a priori knowledge of bias (EnRDA-II) for the Lorenz-96 system. The dashed line is the standard deviation of the observation error.

It is also important to note that the higher RMSE of the PF compared to the SEnKF and EnRDA is due to the problem of filter degeneracy, which is further exacerbated by the presence of systematic errors in model forecasts (Poterjoy and Anderson, 2016). To alleviate this problem, one may investigate the use of methodologies suggested in recent years, including the auxiliary particle filter where the weights of the particles at each assimilation cycle are defined based on the likelihood function from the next cycle using a pre-model run (Pitt and Shephard, 1999), the backtracking particle filter in which the analysis state is backtracked to identify the time step when the filter became degenerate (Spiller et al., 2008) as well as sampling from a transition density to pull back particles towards observations (van Leeuwen, 2010).

To further test the efficiency of EnRDA, another configuration of the Lorenz-96 is implemented using a Laplace-distributed observation error at each assimilation interval of 10Δt with variance ${\mathit{\sigma }}_{\mathrm{obs}}^{\mathrm{2}}=\mathrm{2}$ (Lei and Bickel, 2011; Spantini et al., 2019). Similar to the setting of earlier implementations using Gaussian observation error, 50 (5000) ensemble members (particles) for the SEnKF and EnRDA (PF) are generated. On average, the EnRDA reduces the RMSE by 26 % (47 %) compared to the SEnKF (PF) using 50 random realizations.

4.2 Quasi-geostrophic model

The multilayered QG (Pedlosky, 1987) model is known as one of the simplest circulation models capable of providing a reasonable representation of the mesoscale variability in geophysical flows. In its simplified form, the QG model describes the conservation of potential vorticity $\mathit{\{}{\mathit{\zeta }}_k{\mathit{\}}}_{k=\mathrm{1}}^K$ in K vertically mixed vertical layers:

$$\begin{array}{}\text{(9)}& ({\displaystyle \frac{\partial }{\partial t}}+u_k{\displaystyle \frac{\partial }{\partial \mathit{\lambda }}}+v_k{\displaystyle \frac{\partial }{\partial \mathit{\varphi }}}){\mathit{\zeta }}_k=\mathrm{0},k=\mathrm{1},\mathrm{\dots },K,\end{array}$$

where $u_k=-\frac{\partial {\mathrm{\Psi }}_k}{\partial \mathit{\varphi }}$ and $v_k=\frac{\partial {\mathrm{\Psi }}_k}{\partial \mathit{\lambda }}$ represent the zonal and meridional components of the velocity field, obtained from the geostrophic approximation, $\mathit{\{}{\mathrm{\Psi }}_k{\mathit{\}}}_{k=\mathrm{1}}^K$ is the streamfunction in K layers, and λ and ϕ are the zonal and meridional coordinates, respectively.

For a two-layer QG model (K=2), the potential vorticity at any time step is the sum of the relative vorticity, the planetary vorticity and the stretching term, given by

$$\begin{array}{}\text{(10)}& \begin{array}{rl}& {\mathit{\zeta }}_k={\mathrm{\nabla }}^{\mathrm{2}}{\mathrm{\Psi }}_k+f+\left(\mathrm{1}-\mathrm{2}{\mathit{\delta }}_{\mathrm{2}k}\right){\displaystyle \frac{f_{\mathrm{0}}^{\mathrm{2}}}{g^{\prime }{h}_k}}\left({\mathrm{\Psi }}_{\mathrm{2}}-{\mathrm{\Psi }}_{\mathrm{1}}\right),\\ & k=\mathrm{1},\mathrm{\dots },\mathrm{2},\end{array}\end{array}$$

where ${\mathrm{\nabla }}^{\mathrm{2}}(\cdot )=\frac{{\partial }^{\mathrm{2}}(\cdot )}{\partial {\mathit{\lambda }}^{\mathrm{2}}}+\frac{{\partial }^{\mathrm{2}}(\cdot )}{\partial {\mathit{\varphi }}^{\mathrm{2}}}$ is the Laplace operator, $f=f_{\mathrm{0}}+\mathit{\beta }(\mathit{\varphi }-{\mathit{\varphi }}_{\mathrm{0}})$ is the Coriolis parameter linearly varying with the meridional coordinate ϕ (β-plane approximation), f₀ is the Coriolis parameter at mid-basin where ϕ=ϕ₀, $g^{\prime }=\frac{g({\mathit{\rho }}_{\mathrm{2}}-{\mathit{\rho }}_{\mathrm{1}})}{{\mathit{\rho }}_{\mathrm{2}}}$ is the reduced value of the gravitational acceleration g, and ρ_k and h_k are the density and thickness of the kth layer, respectively. The QG model has been the subject of numerous experiments to test the performance of DA techniques (Evensen, 1994; Evensen and Van Leeuwen, 1996; Fisher and Gürol, 2017; Penny et al., 2019; Cotter et al., 2020).

Experimental setup, results and discussion

Due to the high dimensionality of the QG model and the well-known problem of filter degeneracy in the PF, we chose to omit its application to the QG model. Similarly to the study conducted in Evensen (1992, 1994), the streamfunction is chosen as the state variable for the DA experiments. The streamfunction field, at each vertical layer, is discretized over a uniform grid of dimension m_λ×m_ϕ with a spacing of $\mathrm{\Delta }\mathit{\lambda }=\mathrm{\Delta }\mathit{\varphi }=\mathrm{100}$ km, where m_λ=65 and m_ϕ=33. The model domain is assumed to have periodic boundaries along the zonal direction and free-slip conditions, that is, $v_k=\mathrm{0},\forall k$ holds on the northern and southern boundaries. The standard model parameter values of $f_{\mathrm{0}}=\mathrm{7.28}\times {\mathrm{10}}^{-\mathrm{5}}$ s⁻¹, $\mathit{\beta }=\mathrm{2}\times {\mathrm{10}}^{-\mathrm{11}}$ m⁻¹ s⁻¹, and g=9.81 m s⁻² are used. The total depth of the atmospheric column is set to 10 km with depths and densities of the top and bottom layers as $h_{\mathrm{1}}=h_{\mathrm{2}}=\mathrm{5}$ km and ρ₁=1 and ρ₂=1.05 kg m⁻³, respectively. We first initialize the streamfunction in the two layers as a function of the zonal and meridional coordinates by setting ${\mathrm{\Psi }}_{\mathrm{1}}(\mathit{\lambda },\mathit{\varphi })=-\mathrm{12.5}\times {\mathrm{10}}^{\mathrm{6}}{\tan }^{-\mathrm{1}}\left(\mathrm{20}(\mathit{\varphi }/\mathrm{\Delta }\mathit{\varphi }-m_{\mathit{\varphi }}/\mathrm{2})m_{\mathit{\varphi }}^{-\mathrm{1}}\right)-\mathrm{1.25}\times {\mathrm{10}}^{\mathrm{6}}\sin \left(\mathrm{2}\mathit{\pi }(\mathit{\lambda }/\mathrm{\Delta }\mathit{\lambda }-\mathrm{1})m_{\mathit{\lambda }}^{-\mathrm{1}}\right){\sin }^{\mathrm{2}}\left(\mathrm{2}\mathit{\pi }\right(\mathit{\varphi }/\mathrm{\Delta }\mathit{\varphi }-\mathrm{1}\left)\right(m_{\mathit{\varphi }}-\mathrm{1}{)}^{-\mathrm{1}})$ m² s⁻¹ and ${\mathrm{\Psi }}_{\mathrm{2}}(\mathit{\lambda },\mathit{\varphi })=\mathrm{0.3}{\mathrm{\Psi }}_{\mathrm{1}}(\mathit{\lambda },\mathit{\varphi })$.

From the initial value of the streamfunction field in each layer, potential vorticity is obtained using a nine-point second-order finite difference scheme to compute the Laplacian in Eq. (10). The model in Eq. (9) is then integrated with a time step of Δt=0.5 h using the fourth-order Runge–Kutta approximation to advect and obtain potential vorticity at internal grid points for the next time step. The streamfunction at the next time step is then calculated from this potential vorticity by solving the set of Helmholtz equations (Eq. 10). To avoid any form of initial transient behavior and to create vortex structures in the streamfunction, the QG model is integrated first for 720 time steps, and then the endpoint of the run is used as the initial condition for subsequent DA experimentation.

The ground truth of the streamfunction is obtained by integrating the QG model with a time step of Δt over a time period of T= 0–15 d in the absence of any model error. Observations are assumed to be available at an assimilation time interval of 24Δt or 12 h. To construct observations, representative, random and systematic errors are applied to the ground truth. The representative error is applied by lowering the resolution of the ground truth through box averaging over a window of size n_λ×n_ϕ, where n_λ=5 and n_ϕ=3. Then a heteroscedastic Gaussian observation noise with bias 0.6×10⁶ m² s⁻¹ and a standard deviation of 10 % of the mean magnitude of the ground truth is applied.

Figure 6 (a) The true state x_tr, (b) background state x_b, and (c) observations y for the bottom layer field of the streamfunction in the quasi-geostrophic model at the first assimilation cycle T=12 h. The black plus (grey cross) signs show the location of the global extrema for the true state (background and observation).

To characterize the distribution of the background state, 50 ensemble members for both SEnKF and EnRDA are generated using model errors ${\mathit{\omega }}_t\sim \mathcal{N}(\mathrm{0},\mathit{\alpha }{\mathit{\sigma }}_t^{\mathrm{2}}{\mathbf{I}}_{m_{\mathit{\lambda }}\times m_{\mathit{\varphi }}})$ for each layer with ${\mathit{\sigma }}_{\mathrm{0}}^{\mathrm{2}}={\mathrm{10}}^{\mathrm{8}}$ m⁴ s⁻² and ${\mathit{\sigma }}_t^{\mathrm{2}}=\mathrm{5}\times {\mathrm{10}}^{\mathrm{6}}$ m⁴ s⁻² for t>0, where the factor $\mathit{\alpha }\in [\mathrm{0},\mathrm{1}]$ grows linearly from 0 at the northern and southern boundaries to 1 at mid-basin. To introduce systematic errors in the forecast, we utilize a multiplicative error of 0.015 % in the QG model by multiplying the potential vorticity obtained from Eq. (10) at every Δt by a factor of 1.00015. At each assimilation cycle, N=500 samples of the observations are obtained by perturbing the observations with the heteroscedastic Gaussian observation noise with standard deviation 10 % of the mean magnitude of the ground truth.

In the SEnKF, to alleviate the well-known problem of undersampling (Anderson, 2012) and improve its performance, we utilize covariance inflation (Anderson and Anderson, 1999) and localization (Houtekamer and Mitchell, 2001; Hamill, 2001) as discussed in Appendix A2. For EnRDA, similar to the Lorenz-96 setup (Sect. 4.1), the displacement parameter is set to η=0.4 through a cross-validation study based on a minimum RMSE criterion as shown in Table 1. To increase the robustness of the inference about the results, the quality metrics are averaged using 10 simulations with different random realizations.

Figure 7 The streamfunction analysis state x_a by (a) SEnKF and (d) EnRDA as well as (b, e) their respective absolute error fields and (c, f) zonal mean of the error for the bottom layer of the quasi-geostrophic model at the first assimilation cycle T=12 h. The RMSE values (×10⁶ m² s⁻¹) for the entire fields are also reported in (a) and (d).

The true state, background state, and observations of the bottom layer streamfunction at the first assimilation cycle T=12 h are shown in Fig. 6. It can be seen that both the background state and the observations show possible systematic biases as the position and the values of their global extrema are significantly different from the ground truth.

The results of the DA experiments using the SEnKF and EnRDA at the first assimilation cycle for the bottom layer are also shown in Fig. 7. It can be seen that, in the SEnKF, the streamfunction values are slightly overestimated, signaling the persistence of bias in the analysis state (Fig. 7a). This is further evident as the analysis error field is coherent and structured (Fig. 7b). On the other hand, it appears that EnRDA (Fig. 7d) results in a more incoherent error field with a reduced bias (Fig. 7e). The RMSE for the EnRDA (0.28×10⁶ m² s⁻¹) is lower than the one by the SEnKF (0.46×10⁶ m² s⁻¹). However, the difference between the two methods shrinks over T= 0–15 d, and the mean analysis RMSE over both layers by the EnRDA (SEnKF) reaches 0.21×10⁶ (0.25×10⁶) m² s⁻¹. Furthermore, in the SEnKF, due to the presence of systematic error, the zonal mean of the absolute error is consistently higher than that of the EnRDA (see Fig. 7c and f). To test the effects of non-Gaussian errors on the performance of EnRDA, analogous to the experiments conducted for Lorenz-96, we examined a Laplace noise with the same variance used for the Gaussian errors. The results show that EnRDA performance does not change appreciably and that the improvement remains on the same order of magnitude as reported for the Gaussian observation errors.

Table 1 Average root mean squared error (RMSE) values as a function of the displacement parameter $\mathit{\eta }\in [\mathrm{0.25},\mathrm{0.6}]$ for ensemble Riemannian data assimilation (EnRDA) from 10 independent simulations of the two-layer quasi-geostrophic model.

Figure 8 The average RMSE values as a function of assimilation intervals 6, 12 and 18 h in the SEnKF and EnRDA for the two-layer quasi-geostrophic model.

We further examined the performance of the EnRDA and the SEnKF on the QG model with a ±50 % change in the assimilation interval of 12 h as shown in Fig. 8. To make the comparison fair between different assimilation intervals which have a different number of assimilation cycles and to eliminate the impact of transient behavior, we only report the statistics for the last 15 assimilation steps. With the increase in assimilation interval, the systematic error grows in the forecast largely due to the multiplicative error being added to the forecast at every time step. Therefore, as is expected, with the increase in assimilation interval, the RMSE grows monotonically and the performance of the DA methodologies degrades. However, the EnRDA demonstrates consistent improvement over a bias-blind implementation of the SEnKF (20 %–33 %) across the range of assimilation intervals. On average, using a cluster with 24 cores and a clock rate of 2.5 GHz, it took around 3.5 h to complete one independent simulation on the QG model for EnRDA compared to 2.5 h for EnKF, each with 50 ensemble members.

5 Summary and concluding remarks

In this study, we demonstrated that data assimilation (DA) over the Wasserstein space through the EnRDA (Tamang et al., 2021) can be properly scaled and result in improved predictability of non-Gaussian geophysical dynamics at relatively high dimensions, under systematic errors. In particular, we applied the EnRDA to the 40-dimensional chaotic Lorenz-96 system and a two-layer quasi-geostrophic representation of atmospheric circulation and compared its results with the stochastic ensemble Kalman filter and the particle filter with comparable ensemble size. Under the made assumptions and experimental settings, EnRDA improved the root mean squared error by almost 20 % (30 %) for the Lorenz-96 (QG) model when compared to the classic Euclidean DA techniques. We need to emphasize that in the absence of systematic errors, Euclidean DA methodologies definitely demonstrate improved performance over EnRDA in terms of the root mean squared error. Despite the reported improvements, further comprehensive comparisons with the Euclidean DA methodologies equipped with bias-correction methodologies, such as the cumulative distribution function (CDF) matching (Reichle et al., 2004), are required to fully characterize the pros and cons of DA over the Wasserstein space.

One of the major weaknesses of the presented methodology in its current form is that all dimensions of the problem are assumed to be observable. This is an important issue when it comes to the assimilation of sparse data. Future research is needed to address partial observability in DA over the Wasserstein space. A possible direction is through multi-marginal optimal mass transport (Pass, 2015), which could enable us to couple different dimensions of the problem and propagate the information content of sparse observations to unobserved dimensions. Moreover, currently, the displacement parameter is constant across multiple dimensions of the problem. Future research is needed to understand how the displacement parameter can be estimated differently depending on the error structure across different dimensions of the state space. Another option is to perform the EnRDA only in that part of the state space that is directly observed and use the ensemble covariance to update the unobserved part of state space, similar to a SEnKF. We anticipate that expanding the application of the presented methodology for assimilating satellite data into land–atmosphere models could be another promising future direction of research given the fact that these models are often markedly biased (Dee and Da Silva, 1998; Chepurin et al., 2005; De Lannoy et al., 2007; Lin et al., 2017).

It should be noted that the experimental settings presented here only deal with the univariate state variable. The use of a scalar regularization parameter in the EnRDA penalizes the transportation cost matrix elements uniformly even when the physical variables of interest are different by orders of magnitude. A possible future solution to this problem can be obtained by rather utilizing Mahalanobis or a weighted Euclidean distance (Olver et al., 2006) in lieu of the Euclidean distance to obtain a modified ground transportation cost matrix.

Although EnRDA demonstrated a reasonable performance on the presented dynamical systems without significant computational burden, the computational complexity might be a limiting factor for its large-scale implementation. In Earth system models, where the dimension easily exceeds hundred of millions, dimensionality reduction might be necessary. One might hypothesize that the optimal transportation plan remains unaltered for the change in the basis. Thus, future research can be devoted to examining the optimal transportation plan for the principal components (Olver et al., 2006) of the geophysical state variables of interest to significantly lower its computational cost.

Appendix A

A1 Sinkhorn's algorithm for optimal mass transport

To solve the regularized optimal mass transport problem in Eq. (7), we utilize Sinkhorn's algorithm (Sinkhorn, 1967). To that end, first, the Lagrangian form of Eq. (7) using two Lagrange multipliers a∈ℝ^M and b∈ℝ^N is obtained as follows:

$$\begin{array}{}\text{(A1)}& \begin{array}{rl}\mathcal{L}\left(\mathbf{U},\mathit{a},\mathit{b}\right)=& \mathrm{tr}\left({\mathbf{C}}^{\mathrm{T}}\mathbf{U}\right)-\mathit{\gamma }\mathrm{tr}\left({\mathbf{U}}^{\mathrm{T}}\right[\log (\mathbf{U}-{\text{𝟙}}_M{\text{𝟙}}_N^{\mathrm{T}})\left]\right)\\ & -{\mathit{a}}^{\mathrm{T}}\left(\mathbf{U}{\text{𝟙}}_N-{\mathbf{p}}_x\right)-{\mathit{b}}^{\mathrm{T}}\left({\mathbf{U}}^{\mathrm{T}}{\text{𝟙}}_M-{\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}\right).\end{array}\end{array}$$

Now, we set the first-order derivative of the Lagrangian form in Eq. (A1) with respect to the (i,j)th element of the joint distribution (u_ij) to zero:

$$\begin{array}{}\text{(A2)}& {\displaystyle \frac{\partial \mathcal{L}(\mathbf{U},\mathit{a},\mathit{b})}{\partial u_{ij}}}=c_{ij}+\mathit{\gamma }\log \left(u_{ij}\right)-a_i-b_j=\mathrm{0}\forall i,j,\end{array}$$

which ultimately leads to $u_{ij}=\exp \left(\frac{a_i}{\mathit{\gamma }}\right)$ $\exp (-\frac{c_{ij}}{\mathit{\gamma }})$ $\exp \left(\frac{b_j}{\mathit{\gamma }}\right)$. This can be rewritten in matrix form as U^a=diag(s)Vdiag(t), where $\mathit{\{}\mathbf{V}\in {\mathbb{R}}_{+}^{M\times N}:v_{ij}=\exp (-\frac{c_{ij}}{\mathit{\gamma }}\left)\mathit{\right\}}$ is Gibb's kernel of the cost matrix C and s∈ℝ^M and t∈ℝ^N are the unknown scaling vectors. The notation $\mathrm{diag}\left(\mathit{x}\right)\in {\mathbb{R}}^{M\times M}$ represents a diagonal matrix with its diagonal entries provided by x∈ℝ^M.

By setting the derivatives of the Lagrangian with respect to the Lagrange multipliers to zero, we recover the two conditions, which we can write as p_x=diag(s)Vdiag(t)𝟙_N and ${\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}=\mathrm{diag}\left(\mathit{t}\right){\mathbf{V}}^{\mathrm{T}}\mathrm{diag}\left(\mathit{s}\right){\text{𝟙}}_M$, leading to

$$\begin{array}{}\text{(A3)}& \mathit{s}={\mathbf{p}}_x\oslash \left(\mathbf{V}\mathit{t}\right)\mathrm{and}\mathit{t}={\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}\oslash \left({\mathbf{V}}^{\mathrm{T}}\mathit{s}\right),\end{array}$$

where the notation x⊘y represents a Hadamard element-wise division of equal-length vectors. The form presented in Eq. (A3) is known as the matrix scaling problem (Borobia and Cantó, 1998) and can be efficiently solved iteratively:

$$\begin{array}{}\text{(A4)}& {\mathit{s}}^{\left(i\right)}={\mathbf{p}}_x\oslash \left(\mathbf{V}{\mathit{t}}^{(i-\mathrm{1})}\right)\mathrm{and}{\mathit{t}}^{\left(i\right)}={\stackrel{\mathrm{̃}}{\mathbf{p}}}_{y|x}\oslash \left({\mathbf{V}}^{\mathrm{T}}{\mathit{s}}^{\left(i\right)}\right),\end{array}$$

where i is the iteration count and the algorithm is initialized with a positive vector t⁽⁰⁾=𝟙_N. In our implementation, we set the iteration termination criterion as $\frac{\Vert {\mathit{s}}^{\left(i\right)}-{\mathit{s}}^{(i-\mathrm{1})}{\Vert }_{\mathrm{2}}}{\Vert {\mathit{s}}^{(i-\mathrm{1})}{\Vert }_{\mathrm{2}}}\le {\mathrm{10}}^{-\mathrm{4}}$ or i>300. After the convergence of the solution for s and t, the optimal joint distribution can be obtained as U^a=diag(s)Vdiag(t).

A2 Covariance inflation and localization in the ensemble Kalman filter

The ensemble size in the SEnKF, if much smaller than the state dimension, such as in the presented case of the quasi-geostrophic model, leads to underestimation of the forecast error covariance matrix and subsequently filter divergence problems. To alleviate this problem, a covariance inflation procedure can be implemented by multiplying the forecast error covariance matrix by an inflation factor τ>1 (Anderson and Anderson, 1999) where its optimal value depends on the ensemble size (Hamill et al., 2001) and other characteristics of the problem at hand.

The covariance localization procedure in the SEnKF further attempts to improve its performance by ignoring the spurious long-range dependence in the ensemble background covariance by applying a prespecified cutoff threshold to the correlation structure of the field. An SEnKF equipped with a tuned localization procedure can be efficiently used in high-dimensional atmospheric and ocean models even with fewer than 100 ensemble members (Anderson, 2012). The covariance localization in an SEnKF is accomplished by modifying the Kalman gain matrix $\mathbf{K}\in {\mathbb{R}}^{m\times m}$ through implementation of a Hadamard element-wise product of the forecast error covariance matrix $\mathbf{B}\in {\mathbb{R}}^{m\times m}$ with a distance-based correlation matrix $\mathit{\rho }\in {\mathbb{R}}^{m\times m}$:

$$\begin{array}{}\text{(A5)}& \mathbf{K}=(\mathit{\rho }\odot \mathbf{B}){\mathbf{H}}^{\mathrm{T}}\left(\mathbf{H}\right(\mathit{\rho }\odot \mathbf{B}){\mathbf{H}}^{\mathrm{T}}+\mathbf{R}{)}^{-\mathrm{1}},\end{array}$$

where X⊙Y represent the Hadamard element-wise product between equal size matrices X and Y.

Following the work of Gaspari and Cohn (1999), we utilized the fifth-order piece-wise rational function that depends on a single length scale parameter d and a Euclidean distance matrix $\left\{\mathbf{L}\in {\mathbb{R}}^{m\times m}:l_{ij}=\Vert x_i-x_j{\Vert }_{\mathrm{2}}\right\}$ for obtaining the (i,j)th element of the localizing correlation matrix ρ:

$$\begin{array}{}\text{(A6)}& \begin{array}{rl}& {\mathit{\rho }}_{ij}=\\ & \left\{\begin{array}{ll}-\frac{\mathrm{1}}{\mathrm{4}}{r}^{\mathrm{5}}+\frac{\mathrm{1}}{\mathrm{2}}{r}^{\mathrm{4}}+\frac{\mathrm{5}}{\mathrm{8}}{r}^{\mathrm{3}}-\frac{\mathrm{5}}{\mathrm{3}}{r}^{\mathrm{2}}+\mathrm{1},& \mathrm{0}\le r\le \mathrm{1},\\ {\displaystyle \frac{\mathrm{1}}{\mathrm{12}}}{r}^{\mathrm{5}}-{\displaystyle \frac{\mathrm{1}}{\mathrm{2}}}{r}^{\mathrm{4}}+{\displaystyle \frac{\mathrm{5}}{\mathrm{8}}}{r}^{\mathrm{3}}+{\displaystyle \frac{\mathrm{5}}{\mathrm{3}}}{r}^{\mathrm{2}}-\mathrm{5}r+\mathrm{4}-{\displaystyle \frac{\mathrm{2}}{\mathrm{3}}}{r}^{-\mathrm{1}},& \mathrm{1}<r\le \mathrm{2},\\ \mathrm{0},& r>\mathrm{2},\end{array}\right.\end{array}\end{array}$$

where $r=\frac{l_{ij}}{d}$, and d is the length scale.

In our implementation of the SEnKF in the QG model, the inflation factor and length scale were chosen between τ= 1.01–1.08 and d= 400–1800 (km), respectively, depending on the experimental setup through trial and error analysis to minimize the root mean squared error.