# Coupling techniques for nonlinear ensemble filtering

Alessio<EMAIL_ADDRESS>[ Ricardo
<EMAIL_ADDRESS>[ Youssef<EMAIL_ADDRESS>label=u1
[[ url]http://uqgroup.mit.edu  Massachusetts Institute of Technology
Massachusetts Institute of Technology
Cambridge, MA 02139 USA
, e2, e3

###### Abstract

We consider filtering in high-dimensional non-Gaussian state-space models with
intractable transition kernels, nonlinear and possibly chaotic dynamics, and
sparse observations in space and time. We propose a novel filtering
methodology that harnesses transportation of measures, convex optimization,
and ideas from probabilistic graphical models to yield robust ensemble
approximations of the filtering distribution in high dimensions. Our approach
can be understood as the natural generalization of the ensemble Kalman filter
(EnKF) to _nonlinear_ updates, using stochastic or deterministic couplings.
The use of nonlinear updates can reduce the intrinsic bias of the EnKF at a
marginal increase in computational cost. We avoid any form of importance
sampling and introduce non-Gaussian localization approaches for dimension
scalability. Our framework achieves state-of-the-art tracking performance on
challenging configurations of the Lorenz-96 model in the chaotic regime.

nonlinear filtering,

state-space models,

couplings,

transport maps,

ensemble Kalman filter,

graphical models,

localization,

approximate Bayesian computation,

###### keywords:

, , and

###### Contents

  1. 1 Introduction
    1. 1.1 Related work
    2. 1.2 Organization of the paper
    3. 1.3 Remarks on notation
  2. 2 State-space models
  3. 3 Non-Gaussian generalizations of the ensemble Kalman filter
  4. 4 Analysis step with couplings
  5. 5 Knothe–Rosenblatt rearrangement
  6. 6 Stochastic map filter
  7. 7 Estimation of the Knothe–Rosenblatt rearrangement from samples
  8. 8 Regularization in high dimensions: localization of nonlinear maps
  9. 9 Remarks on the stochastic map filter
  10. 10 Deterministic map filter
  11. 11 Numerical experiments
    1. 11.1 Common aspects of the problem configurations
    2. 11.2 Lorenz-63
    3. 11.3 Lorenz-96 with long inter-observation time
    4. 11.4 Lorenz-96 with heavy-tailed observational noise
  12. 12 Discussion
  13. A Parameterization and computation of monotone triangular maps
  14. B Deterministic map filter with local observations

##  1 Introduction

State-space models formalize the probabilistic description of a time-dependent
latent process—the state—observed indirectly at discrete times [22]. These
models can approximate a wide variety of stochastic processes, ranging from
the evolution of atmospheric variables in meteorology to the volatility of
financial assets. This paper is concerned with the problem of (discrete-time)
filtering, i.e., characterizing the sequence of conditional distributions of
the latent field at observation times, given all available measurements up to
that time. Filtering arises in virtually every discipline that seeks an online
integration of models with data, e.g., imaging, pharmacology, atmospheric
sciences, and oceanography. In geophysical applications, filtering and several
closely related inference problems (e.g., smoothing, sequential parameter
inference) fall under the broad label of _data assimilation_.

Despite its importance and ubiquity, filtering remains a challenging task,
particularly when (i) the state is high dimensional; (ii) the state transition
dynamic is nonlinear, expensive to simulate, and intractable, e.g., if it
involves the integration of a chaotic partial differential equation; and (iii)
observations of the state are sparse, both in space and time [42, 57]. These
constraints reflect typical challenges faced in numerical weather prediction
or geophysical data assimilation, and pose severe challenges for consistent
sequential Monte Carlo (SMC) algorithms, which invariably face particle
degeneracy or impoverishment [62, 20]. In such scenarios, state-of-the-art
results111At least in terms of time-averaged errors in _point estimates_ of
the state [36].  are typically obtained with the ensemble Kalman filter
(EnKF), which is the workhorse of modern ensemble-based data assimilation
[23]. The EnKF implements a two-step Monte Carlo approximation of the
classical Kalman recursions. In the forecast step, a particle approximation of
the filtering distribution is propagated through the transition dynamic to
yield a “forecast” ensemble at the next observation time. In the analysis
step, the forecast ensemble is updated via the action of a linear
transformation to yield an empirical approximation of the new filtering
marginal. The linear transformation is estimated under Gaussian assumptions.
Hence, the EnKF cannot yield consistent estimators of the filtering
distribution for non-Gaussian models [43]. In essence, the EnKF trades
consistency of the estimators for lower variance, and thus robustness in high
dimensions. Yet the intrinsic bias of the EnKF—due both to the linear
transformation and the way it is estimated—implies that increasing the
ensemble size beyond a certain threshold does not improve accuracy; increasing
computational effort does not yield better inference. We wish to address these
limitations.

In this work, we introduce non-Gaussian generalizations of the EnKF by
considering nonlinear transformations derived from couplings. Couplings
provide a link between distributions. That is, for a pair of distributions
$(\pi_{1},\pi_{2})$, a coupling is defined by a pair of random variables
$(\bm{X}_{1},\bm{X}_{2})$, which admit $\pi_{1}$ and $\pi_{2}$ as marginal
distributions [70]. We consider couplings that are induced by a continuous
transformation $T$ such that $\bm{X}_{2}=T(\bm{X}_{1})$. The transformation
can be either deterministic or stochastic. When $T$ is deterministic, we call
it a transport map. The transformation allows one to sample $\pi_{2}$ by
evaluating $T$ at samples from $\pi_{1}$.

We interpret the analysis step of the EnKF as a problem of coupling the
forecast distribution $\pi_{\bm{X}}$ with the filtering distribution
$\pi_{\bm{X}|\bm{Y}}$. That is, given samples $\bm{x}^{1},\ldots,\bm{x}^{M}$
from $\pi_{\bm{X}}$ and a likelihood function $\pi_{\bm{Y}|\bm{X}}$, we seek a
transformation $T$ that yields samples
$\smash{T(\bm{x}^{1}),\ldots,T(\bm{x}^{M})}$ from $\pi_{\bm{X}|\bm{Y}}$.
Couplings are not unique. We are particularly interested in transformations
that can be estimated efficiently without resorting to importance
sampling—perhaps using only convex optimization—and that are easy to
“localize” in high dimensions. We want to avoid the use of weights and thus
issues of particle degeneracy. Moreover, the computation of the transformation
should not become increasingly challenging as the variance of the observation
noise decreases. The latter is a typical concern of filtering algorithms that
rely on some attributes—e.g., moments—of the bootstrap particle filter
approximation. _This paper proposes two new algorithms_ , called the
stochastic and deterministic map filters. The former represents a non-Gaussian
generalization of the EnKF with “perturbed observations” [13], while the
latter is a nonlinear extension of the square-root EnKF [67].

The _stochastic map filter_ seeks a non-deterministic transformation from
forecast to filtering distribution. This transformation is given by a
transport map that pushes forward the joint distribution of state and data at
a given observation time ($\pi_{\bm{X},\bm{Y}}$) to the filtering
distribution. We target a specific transport map derived from the
Knothe–Rosenblatt (KR) rearrangement that pushes forward $\pi_{\bm{X},\bm{Y}}$
to a “reference” distribution with independent components, e.g., a standard
normal. The KR rearrangement is the unique monotone triangular transport map
that defines a deterministic coupling between two distributions [10], and can
be estimated within a finite-dimensional function space using only convex
optimization and samples from $\smash{\pi_{\bm{X},\bm{Y}}}$ [51]. These
samples are easy to obtain given a forecast ensemble $(\bm{x}^{i})_{i\leq M}$:
it suffices to simulate the likelihood $\pi_{\bm{Y}|\bm{X}}$ at each particle
$\smash{\bm{x}^{i}}$. If we restrict the estimator of the rearrangement to be
linear, then we recover the stochastic EnKF [23]. If we consider nonlinear
parameterizations of the estimator, however, we obtain a whole new class of
nonlinear filtering algorithms, which rely on fast and robust convex
optimization.

In the _deterministic map filter_ , we adopt a complementary strategy and seek
a deterministic transformation from the forecast to the filtering
distribution. We view the transformation as a composition of two functions.
The first function is a KR rearrangement that pushes forward the forecast
$\pi_{\bm{X}}$ to a reference distribution. This rearrangement can be
estimated via convex optimization given only forecast samples, and yields an
implicit approximation of the forecast density, $\widehat{\pi}_{\bm{X}}$. The
second function is a KR rearrangement that pushes forward the reference to an
approximation of the filtering distribution given by
$\widehat{\pi}_{\bm{X}|\bm{Y}}\propto\pi_{\bm{Y}|\bm{X}}\,\widehat{\pi}_{\bm{X}}$,
which can be evaluated in closed form up to a normalizing constant. The
estimation of the latter rearrangement leads to a variational problem, first
proposed by [49], that is in general non-convex. This is an important
difference with respect to the stochastic map filter, which relies entirely on
convex optimization, at the cost of estimating higher-dimensional
rearrangements. Yet, in the special case of local observations that are
conditionally independent, we show how to implement the deterministic map
filter using only convex optimization (Appendix B). The resulting scheme has a
flavor similar to the multivariate rank histogram filter [46]. For a linear
parameterization of the estimators and Gaussian likelihoods, the deterministic
map filter reduces to a (deterministic) square-root EnKF [67]; otherwise, it
is far more general.

In principle, both the stochastic and deterministic map filters can
approximate a nonlinear update of arbitrary complexity: it suffices to enrich
the function space for the estimator of the KR rearrangements. In practice,
however, there is a tradeoff. Depending on the number of forecast particles
$M$, overly complex parameterizations can lead to estimators with low bias but
unacceptable variance, and vice-versa. Skillful filtering requires a careful
balance between bias and variance. A key ingredient of the functional
framework proposed in this paper is the ability to depart gradually from the
linear ansatz. For example, in our numerical experiments we consider maps
whose components are additively separable, as sums of univariate nonlinear
functions. These parameterizations represent a natural extension of linear
functions in terms of complexity, and lead to filtering algorithms that can
outperform the EnKF at a marginal increase in computational cost and ensemble
requirements (Section 11). Of course, more general parameterizations are also
possible. In essence, we provide a natural tradeoff between computational cost
and statistical accuracy—a tradeoff which is very limited in the EnKF, due to
the restrictive linear ansatz.

In high-dimensional problems with limited ensemble sizes, we must further
regularize the estimation of the rearrangements. In this paper, we introduce a
notion of “localization” for nonlinear updates, by dropping some variable
dependence from each component of the transport map. We show how sparsity of
the KR rearrangement—a nonlinear function—is linked to properties of the
filtering distribution, such as decay of correlation and (approximate)
conditional independence [64]. Intuitively, these sparse transformations
approximate the projection of the filtering distribution onto a manifold of
non-Gaussian Markov random fields.

###  1.1 Related work

The idea of transporting measures via couplings has a long history [47, 70],
and has found applications in fluid dynamics, economics, statistics, machine
learning [21, 32, 49, 26], and many other fields. One of the first filtering
algorithms to rely on the explicit construction of (optimal) couplings is the
ensemble transform particle filter [56, 1], which uses an importance sampling
approximation of the filtering distribution—the same given by the bootstrap
particle filter [27]—to obtain a consistent particle approximation of Monge’s
optimal map [70] (cf. Section 4). Earlier work on particle flows investigated
the use of transport maps induced by flows of ODEs [17, 18]. Related ideas
include the feedback particle filter [74] and the “Gibbs flow” [29]. In
particular, [29] defines an approximation of the KR rearrangement between an
input and a target measure via the solution of an ODE whose drift term depends
only on the full conditionals of the target distribution. This approximation
can be used as a proposal density for sequential Monte Carlo (SMC) methods.
Also, the implicit sampling algorithm of [14, 48] defines a proposal density
using an approximate transport map, which can be evaluated by solving a one-
dimensional optimization problem: in this case, optimization is used to define
the action of a single map, not to search for the best rearrangement within a
class of candidate maps (see [45, Sec. 6] for more details). Among the flow
methods, the variational mapping particle filter implements a particle flow
given by a particular form of functional gradient descent on a reproducing
kernel Hilbert space [54, 39]. We clarify its connection with the
deterministic map filter in Section 10.

In our filters, we use a variational approximation of the KR rearrangement to
Gaussianize a collection of particles. The idea of Gaussianizing multivariate
data is well rooted in statistics [11], and is known as Gaussian anamorphosis
(GA) in the geostatistics literature [71]. Usually, these approaches
Gaussianize marginals of the data by estimating one-dimensional cumulative
distribution functions. A typical application of GA within the context of non-
Gaussian extensions of the EnKF is described in [2]: the idea is to (1)
estimate a nonlinear transformation that Gaussianizes joint samples from the
distribution of state and data at a given observation time, (2) derive an
equivalent nonlinear observation operator in the transformed space, (3) apply
EnKF formulas for nonlinear observation operators in the transformed space,
and (4) map the analysis ensemble back to the original space. As explained in
[2], this approach cannot recover the exact filtering distribution for non-
Gaussian models, even if we had the exact nonlinear transformation. The
problem is that we insist on using EnKF formulas in the transformed space. In
the stochastic map filter introduced in this paper, on the other hand, we show
how to bypass this issue by designing a nonlinear update that depends only on
the nonlinear transformation that Gaussianizes the data. If we had the exact
transformation, then the stochastic map filter would be exact.

Yet another measure transport approach to filtering is that of [64], which
uses the idea of decomposable couplings to yield a recursive approximation of
the full Bayesian solution to the sequential inference problem, using purely
variational techniques (i.e., no sampling). Outside the coupling literature,
an early attempt to devise non-Gaussian generalizations of the EnKF is the
Gaussian mixture filter of [5]. Here, the idea is to learn the forecast
density via a regularized mixture of Gaussians. Finally, we note that
complementary approaches to non-Gaussian high dimensional filtering include
hybrid filters [24] and local weights particle filters [52]. These algorithms
provide an alternative route to the regularization of SMC methods, which are
usually consistent but plagued by large estimation variance in high
dimensions. See [69] for a recent review.

###  1.2 Organization of the paper

The rest of the paper is organized as follows. In Section 2, we review the
notion of state-space models. In Section 2, we describe the EnKF and discuss
possible non-Gaussian generalizations. In Section 4, we interpret the analysis
step of the EnKF as a problem of coupling the forecast distribution with the
filtering distribution, and analyze the difference between stochastic and
deterministic couplings. In Section 5, we recall properties of the KR
rearrangement—a key coupling for our analysis. In Section 6, we introduce the
stochastic map filter. In Section 7, we address estimation of the KR
rearrangement via convex optimization, while in Section 8 we introduce
regularization ideas for high dimensions. Specific parameterizations of the KR
rearrangement are discussed in Appendix A. In Section 9, we collect a few
remarks on the stochastic map filter, including the connection with the
stochastic EnKF and additional localization ideas. In Section 10, we introduce
the deterministic map filter. We report on our numerical experiments with
chaotic dynamical systems in Section 11, evaluating both state estimation
error and fidelity to the underlying Bayesian solution, among other
performance metrics. We conclude in Section 12, where we discuss open issues
and outline directions for future work. Code and numerical examples are
available online.222https://github.com/map-filters

###  1.3 Remarks on notation

For a pair of functions $f$ and $g$, we denote their composition by $f\circ
g$. We denote by $\partial_{k}f$ the partial derivative of $f$ with respect to
its $k$th input variable. We use boldface capital letters, e.g., $\bm{X}$, to
denote random variables on $\mathbb{R}^{n}$, while we write scalar-valued
random variables as $X$. For a density $\pi$, $\bm{X}\sim\pi$ means that
$\bm{X}$ is distributed according to $\pi$. For all $n>0$, we let
$\mathbb{N}_{n}=\\{1,\ldots,n\\}$ denote the set of the first $n$ integers. If
$\bm{X}=(X_{1},\ldots,X_{p})$ is a collection of random variables and
$\mathcal{A}\subset\mathbb{N}_{p}$, then
$\bm{X}_{\mathcal{A}}=(X_{i},i\in\mathcal{A})$ denotes a subcollection of
$\bm{X}$. In the same way, for $j<k$,
$\bm{X}_{j:k}=(X_{j},X_{j+1},\ldots,X_{k})$. The same index notation applies
to dummy variables $\bm{x}\in\mathbb{R}^{n}$. If $\bm{X}=(X_{1},\ldots,X_{p})$
has joint density $\pi$ and $\mathcal{A}\subset\mathbb{N}_{p}$, we denote by
$\pi_{\bm{X}_{\mathcal{A}}}$ the marginal of $\pi$ along
$\bm{X}_{\mathcal{A}}$. If $\pi$ is the density of $\bm{Z}=(\bm{X},\bm{Y})$,
we denote by $\pi_{\bm{X}|\bm{Y}}$ the density of $\bm{X}$ given $\bm{Y}$. We
denote independence of a pair of random variables $\bm{X},\bm{Y}$ by
$\bm{X}\perp\\!\\!\\!\perp\bm{Y}$. In the same way,
$\bm{X}\perp\\!\\!\\!\perp\bm{Y}|\bm{R}$ means that $\bm{X}$ and $\bm{Y}$ are
independent given a third random variable $\bm{R}$.

##  2 State-space models

We consider a nonlinear, non-Gaussian state-space model given by a pair of
discrete-time stochastic processes, $(\bm{Z}_{k},\bm{Y}_{k})$ for $k\geq 1$,
where $\bm{Z}_{k}$ denotes the unobserved state of a Markov chain taking
values in $\smash{\mathbb{R}^{n}}$, while $(\bm{Y}_{k})$ refers to an observed
process in $\mathbb{R}^{d}$ that is conditionally independent given the states
at all times. In filtering, we want to leverage realizations
$(\bm{y}_{k})_{k\geq 1}$ of the observed process to infer the distribution of
the latent state conditioned on all available measurements, i.e., we want to
estimate the sequence of conditionals

$\pi_{\bm{Z}_{k}|\bm{y}_{1:k}}(\bm{z})\coloneqq\pi_{\bm{Z}_{k}|\bm{Y}_{1},\ldots,\bm{Y}_{k}}(\bm{z}|\bm{y}_{1},\ldots,\bm{y}_{k}),$
(2.1)

recursively in time, for all $k\geq 1$.

The joint law of the process $(\bm{Z}_{k},\bm{Y}_{k})$ is fully specified by
the sequence of transition kernels $(\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}})$ and
likelihood functions $(\pi_{\bm{Y}_{k}|\bm{Z}_{k}})$, together with a
distribution on the initial conditions of the latent process. We can think of
$\bm{Y}_{k}$ as an indirect and noisy observation of $\bm{Z}_{k}$, e.g.,

$\bm{Y}_{k}=h(\bm{Z}_{k},\bm{\mathcal{E}}_{k})$ (2.2)

for some mapping $h$ and “noise” random variable $\bm{\mathcal{E}}_{k}$. In
particular, the likelihood function $\pi_{\bm{Y}_{k}|\bm{Z}_{k}}$ need not be
Gaussian. Moreover, in many cases $(\bm{Z}_{k})_{k\geq 1}$ represents a
subcollection of random variables from a larger state process—e.g.,
$(\bm{Z}_{j/\Delta})_{j\in\mathbb{N}}$ for some integer $\Delta$—that is only
observed at certain times $k=j/\Delta\in\mathbb{N}$ (see Figure 1). We thus
consider settings with sparse measurements, both in space and time.

######  Example 1 (State-space model).

A possible dynamic for the state is given by $\bm{Z}_{0}\sim\mathcal{N}({\bf
0},{\bf I})$, and by a set of stochastic difference equations of the form

$\bm{Z}_{j+\frac{1}{\Delta}}=f\left(\bm{Z}_{j}\right)+\bm{\mathcal{W}}_{j},\qquad\forall\,j=0,\,1/\Delta,\,2/\Delta,\ldots,$
(2.3)

for some deterministic forward model
$f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ and i.i.d.
$(\bm{\mathcal{W}}_{j})\sim\mathcal{N}({\bf 0},{\bf I}/\Delta)$. Sampling from
$\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}}(\cdot|\bm{z})$ requires iterating over (2.3)
for $\Delta$ times with initial conditions $\bm{Z}_{0}=\bm{z}$. We can then
consider that the state is observed through $(\bm{Y}_{k})$ only at integer
times.

We assume that we can sample the transition kernel
$\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}}$, but not evaluate its density, which might not
even exist when dealing with deterministic dynamics, e.g.,
$\bm{\mathcal{W}}_{j}=0$ in (2.3).

$\bm{Z}_{k}$$\bm{Z}_{k+\frac{1}{\Delta}}$$\bm{Z}_{k+\frac{\Delta-1}{\Delta}}$$\bm{Z}_{k+1}$$\bm{Y}_{k}$$\bm{Y}_{k+1}$$\bm{Z}_{(k-1)+\frac{\Delta-1}{\Delta}}$$\bm{Z}_{(k-1)+\frac{1}{\Delta}}$$\bm{Z}_{k-1}$$\bm{Y}_{k-1}$
(a)  Figure 1:  Markov structure for a state-space model consisting of a
latent field $(\bm{Z}_{\frac{j}{\Delta}})_{j\in\mathbb{N}}$ which is observed
every $\Delta$ time steps via a second stochastic process
$(\bm{Y}_{k})_{k\in\mathbb{N}}$.

##  3 Non-Gaussian generalizations of the ensemble Kalman filter

We work with a general class of ensemble filtering algorithms, wherein the
filtering distribution is approximated by a collection of particles that are
updated recursively over time via a two-step process consisting of a forecast
and an analysis step (Figure 2). The recursion takes a particle approximation
$\bm{z}^{1},\ldots,\bm{z}^{M}$ of the filtering distribution
$\pi_{\bm{Z}_{k}|\bm{y}_{1:k}}$ at time $k$, and turns it into a particle
approximation of the filtering distribution at the next observation time,
$\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k+1}}$, by considering the following identity,

$\pi_{\bm{Z}_{k+1}|\bm{Y}_{1:k+1}}\propto\pi_{\bm{Y}_{k+1}|\bm{Z}_{k+1}}\,\int\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}}\,\pi_{\bm{Z}_{k}|\bm{Y}_{1:k}}\,{\rm
d}{\scriptstyle\bm{Z}_{k}},$ (3.1)

which links filtering marginals at consecutive observation times. In the
forecast step, each particle $\bm{z}^{j}$ is updated independently with a
sample from $\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}}(\cdot|\bm{z}^{j})$, obtained by
simulating the state dynamic until the next observation time. The resulting
ensemble $\bm{z}^{1}_{f},\ldots,\bm{z}^{M}_{f}$ yields a particle
approximation of the forecast (or predictive) distribution
$\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k}}$—i.e., the distribution of the state at time
$k+1$ conditioned on all observations available up to time $k$. The forecast
step might involve the evaluation of an expensive forward model, like a
partial differential equation, and represents the computational bottleneck for
high dimensional and complex models, like those employed in weather
forecasting. Yet, the forecast step is of little statistical concern, since we
can usually sample exactly from $\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}}$. The analysis
step, however, is more challenging, since it involves the approximation of an
inference task. We can think of the forecast distribution
$\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k}}$ as the prior of a Bayesian inverse problem,
where the likelihood function and posterior distribution are given by
$\pi_{\bm{y}_{k+1}|\bm{Z}_{k+1}}$ and $\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k+1}}$,
respectively. In the analysis step, we seek a particle approximation of the
posterior $\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k+1}}$ in a setting where the prior
density $\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k}}$ is not known explicitly, and where
we only have finitely many samples $\bm{z}^{1}_{f},\ldots,\bm{z}^{M}_{f}$ from
the prior.

The analysis step is thus a problem of approximate Bayesian inference, and can
be implemented in various ways. For instance, the EnKF yields a particle
approximation of the filtering distribution by pushing forward the forecast
ensemble through a linear transformation estimated under Gaussian assumptions.
In this paper, we leverage ideas from measure transport to generalize the EnKF
by estimating nonlinear transformations in function space. Our goal is to
reduce the estimation bias in the EnKF, while retaining applicability and
robustness of the methodology in high dimensions.

Figure 2:  A typical ensemble filtering algorithm first propagates a particle
approximation of the filtering distribution at time $k$ until the next
observation time $k+1$ using the transition dynamic (forecast step). Then, it
updates the ensemble via some transformation that incorporates information
from the most recent observation $\bm{y}_{k+1}$ (analysis step). The analysis
step can be regarded as a problem of “static” Bayesian inference

##  4 Analysis step with couplings

The analysis step of the ensemble algorithms described in Section 2 can be
formulated as an abstract Bayesian inference problem, defined by a pair of
random variables $(\bm{X},\bm{Y})$ on $\mathbb{R}^{n}\times\mathbb{R}^{d}$,
which denote, respectively, the unobserved state variable and the data at the
current assimilation time. The joint law $\pi_{\bm{Y},\bm{X}}$ is specified by
the prior marginal $\pi_{\bm{X}}$, which corresponds to the forecast
distribution, and by the likelihood function $\pi_{\bm{Y}|\bm{X}}$. For
simplicity, we assume that $\pi_{\bm{Y},\bm{X}}$ is fully supported on
$\mathbb{R}^{d}\times\mathbb{R}^{n}$. In filtering, we are only given $M$
samples $\bm{x}^{1},\ldots,\bm{x}^{M}$ from the prior $\pi_{\bm{X}}$—the so-
called forecast ensemble—and we wish to sample the posterior (i.e., the
filtering distribution)

$\pi_{\bm{X}|\bm{y}^{*}}(\bm{x})=\pi_{\bm{X}|\bm{Y}}(\bm{x}|\bm{y}^{*})$ (4.1)

for some realization $\bm{y}^{*}\in\mathbb{R}^{d}$ of the data.

A possible approach to posterior sampling is to construct a coupling between
prior and posterior distributions using a transformation $T$ such that
$\bm{Z}\coloneqq T(\bm{X})$ follows the posterior law [70]. This
transformation and the associated coupling may or may not be deterministic.

In a deterministic coupling, $\bm{Z}=T(\bm{X})$ for some deterministic
transformation $T:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ called a transport
map. In this case, we say that $T$ pushes forward the prior to the posterior
distribution. For any $\bm{X}\sim\pi_{\bm{X}}$ and map $T$, the pushforward of
$\pi_{\bm{X}}$ by $T$ is given by the law of $T(\bm{X})$ and is denoted by
$T_{\sharp}\pi_{\bm{X}}$. The existence of a map $T$ such that
$T(\bm{X})\sim\pi_{\bm{X}|\bm{y}^{*}}$ is guaranteed under very mild
conditions. For instance, it suffices to have a prior distribution without
atoms [70]. In fact, there are infinitely many ways to push forward one
distribution to another, some of which minimize a notion of integrated cost.
Monge’s optimal map, for example, minimizes the Euclidean $L_{2}$ cost
$\mathbb{E}[\|T(\bm{X})-\bm{X}\|^{2}]$ over all admissible transformations
[47]. Transport maps are interesting because they enable direct posterior
sampling. That is, if $\bm{x}^{1},\ldots,\bm{x}^{M}$ are i.i.d. samples from
the prior, then $T(\bm{x}^{1}),\ldots,T(\bm{x}^{M})$ are also i.i.d. samples
from the posterior (cf. Figure 3). We will see in Section 10 that
deterministic couplings are associated with nonlinear generalizations of the
square-root EnKF [67].

######  Remark 1 (Ensemble transform particle filter).

In his pioneering approach to filtering [56], Reich proposes a nonparametric
estimator for Monge’s map, by solving the Kantorovich relaxation of a discrete
optimal transport problem between the empirical measure given by the prior
ensemble and the empirical approximation of the posterior distribution
obtained by reweighting prior particles according to their likelihood. In this
paper, we consider continuous rather than discrete transport, and look for
prior-to-posterior transformations—not necessarily optimal with respect to any
cost—that can be estimated without resorting to importance sampling, thus
avoiding the root cause of particle degeneracy.

A coupling need not be deterministic. For instance, we can have
$\bm{Z}=T(\bm{\Xi},\bm{X})\sim\pi_{\bm{X}|\bm{y}^{*}}$ for some deterministic
transformation $T$ and additional random variable $\bm{\Xi}$. As a function of
$\bm{X}$ alone, we can think of $T$ as a “stochastic” map. In this work, we
consider non-deterministic couplings induced by a map
$T:\mathbb{R}^{d}\times\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ such that
$\bm{Z}=T(\bm{Y},\bm{X})$, where $\bm{Y}$ represents the data random variable
of the Bayesian inference problem. These couplings can be used for posterior
sampling like their deterministic counterpart: if
$(\bm{y}^{1},\bm{x}^{1}),\ldots,(\bm{y}^{M},\bm{x}^{M})$ are i.i.d. samples
from the joint distribution $\pi_{\bm{Y},\bm{X}}$, then
$T(\bm{y}^{1},\bm{x}^{1}),\ldots,T(\bm{y}^{M},\bm{x}^{M})$ are also i.i.d.
samples from the posterior. In fact, the transformation $T$ is a transport map
that pushes forward $\pi_{\bm{Y},\bm{X}}$ to the posterior. Inference with
non-deterministic couplings is interesting because we can build estimators for
$T$ using only convex optimization. We will see in Section 6 that non-
deterministic couplings are linked to nonlinear generalizations of the
stochastic EnKF [13].

There is no unique way to couple prior and posterior distributions in a
Bayesian inference problem. In the context of filtering, we need couplings
that can be estimated efficiently in high dimensions from finitely many prior
samples. In Sections 6 and 10, we discuss specific choices for these
couplings, which lead to novel nonlinear filtering algorithms. Before that,
however, we need to discuss the properties of an important transport for our
analysis, the Knothe–Rosenblatt rearrangement [59].

Figure 3:  Illustration of a transport map $T$ that pushes forward the prior
$\pi_{\bm{X}}$ to the posterior $\pi_{\bm{X}|\bm{Y}=\bm{y}^{*}}$. By the
definition of $T$, if $(\bm{x}^{i})$ are i.i.d. samples from $\pi_{\bm{X}}$,
then $(T(\bm{x}^{i}))$ are also i.i.d. samples from
$\pi_{\bm{X}|\bm{Y}=\bm{y}^{*}}$.

##  5 Knothe–Rosenblatt rearrangement

Given any pair of positive densities $\pi,\eta$ on $\mathbb{R}^{n}$, there
exists a unique monotone triangular transport map $S$—called the
Knothe–Rosenblatt (KR) rearrangement—that pushes forward one of the
corresponding distributions to the other [10], i.e., if $\bm{Z}\sim\pi$, then
$S(\bm{Z})\sim\eta$. A (lower) triangular map
$S:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ is a multivariate function whose
$k$th component, $S^{k}$, depends only on the first $k$ input variables, i.e.,

$S(\bm{z})=\left[\begin{array}[]{l}S^{1}(z_{1})\\\\[2.0pt]
S^{2}(z_{1},z_{2})\\\ \vdots\\\ S^{n}(z_{1},z_{2},\dots
z_{n})\end{array}\right].$ (5.1)

Monotonicity of the multivariate transformation is intended with respect to
the lexicographic order on $\mathbb{R}^{n}$, which is equivalent to each slice

$\xi\mapsto S^{k}(z_{1},\ldots,z_{k-1},\xi)$ (5.2)

being an increasing function. For $\pi,\eta>0$, each function (5.2) is also
strictly increasing and the KR rearrangement becomes an invertible map. The
existence of the KR rearrangement stem from intuitive factorization properties
of any density, e.g., if $\bm{Z}=(Z_{1},\ldots,Z_{n})\sim\pi$, then
$\pi=\prod_{k}\pi_{Z_{k}|\bm{Z}_{1:k-1}}$. In fact, perhaps the most important
property for our analysis is that the rearrangement provides an implicit
characterization of the marginal conditionals of $\pi$, whenever
$\eta=\mathcal{N}({\bf 0},{\bf I}_{n})$. In this case, one can show that each
one dimensional mapping (5.2) pushes forward the marginal conditional
$\xi\mapsto\pi_{Z_{k}|\bm{Z}_{1:k-1}}(\xi|\bm{z}_{1:k-1})$ to a one
dimensional standard normal, for all $\bm{z}_{1:k-1}\in\mathbb{R}^{k-1}$ [61].
We leverage this property in the next section, where we discuss our first
filtering algorithm, while we defer questions pertaining to the estimation of
the rearrangement to Section 7.

##  6 Stochastic map filter

In this section, we outline an algorithm to carry out the analysis step at a
given observation time, based on the construction of non-deterministic
couplings between forecast and filtering distributions. We refer to the
notation of Section 4 for an abstract Bayesian inverse problem.

Our goal is to find a map
$T:\mathbb{R}^{d}\times\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ such that
$T(\bm{Y},\bm{X})\sim\pi_{\bm{X}|\bm{y}^{*}}$. If we had such a map, then we
could easily sample the posterior $\pi_{\bm{X}|\bm{y}^{*}}$ by pushing forward
samples from the joint distribution $\pi_{\bm{Y},\bm{X}}$ through $T$. See
Figure 4 for an illustration. Samples from $\pi_{\bm{Y},\bm{X}}$ are also easy
to obtain: given a collection of prior (or forecast) samples
$(\bm{x}^{i})\sim\pi_{\bm{X}}$, let $\bm{y}^{i}$ be a sample from
$\pi_{\bm{Y}|\bm{X}}(\cdot|\bm{x}^{i})$, for all $i=1,\ldots,M$. The resulting
pairs $(\bm{y}^{i},\bm{x}^{i})$ define samples from $\pi_{\bm{Y},\bm{X}}$. We
are then left with the problem of finding a suitable $T$. We call $T$ the
“analysis” map.

Let $S$ be the KR rearrangement that pushes forward $\pi_{\bm{Y},\bm{X}}$ to
$\mathcal{N}({\bf 0},{\bf I}_{d+n})$. Since the rearrangement is a triangular
map, we can always partition it as

$S(\bm{y},\bm{x})=\left[\begin{array}[]{l}S^{\bm{\mathcal{Y}}}(\bm{y})\\\\[2.0pt]
S^{\bm{\mathcal{X}}}(\bm{y},\bm{x})\end{array}\right]$ (6.1)

for some functions
$S^{\bm{\mathcal{Y}}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ and
$S^{\bm{\mathcal{X}}}:\mathbb{R}^{d}\times\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$,
where $\bm{y}\in\mathbb{R}^{d}$ and $\bm{x}\in\mathbb{R}^{n}$. By the
definition of KR rearrangement, we have that the function
$S^{\bm{\mathcal{X}}}$ pushes forward $\pi_{\bm{Y},\bm{X}}$ to a standard
normal on $\mathbb{R}^{n}$, i.e.,

$S^{\bm{\mathcal{X}}}(\bm{Y},\bm{X})\sim\mathcal{N}({\bf 0},{\bf I}_{n}).$
(6.2)

Moreover, since the KR rearrangement implicitly characterizes the conditionals
of $\pi_{\bm{X},\bm{Y}}$ (cf. Section 5), we have that the map
$\bm{\xi}\mapsto S^{\bm{\mathcal{X}}}(\bm{y}^{*},\bm{\xi})$—obtained by fixing
the first input of $S^{\bm{\mathcal{X}}}$ to the observation
$\bm{y}^{*}\in\mathbb{R}^{d}$—pushes forward the posterior
$\pi_{\bm{X}|\bm{y}^{*}}$ to $\mathcal{N}({\bf 0},{\bf I}_{n})$. By combining
these properties, we arrive at a definition for the analysis map $T$ as

$T\coloneqq S^{\bm{\mathcal{X}}}(\bm{y}^{*},\cdot)^{-1}\circ
S^{\bm{\mathcal{X}}},$ (6.3)

where $S^{\bm{\mathcal{X}}}(\bm{y}^{*},\cdot)^{-1}$ denotes the inverse
function of the mapping $\bm{\xi}\mapsto
S^{\bm{\mathcal{X}}}(\bm{y}^{*},\bm{\xi})$. Inverting a triangular map is a
computationally trivial task, since it reduces to a sequence of one-
dimensional root findings. See Algorithm 1 for details. It is immediate to
verify that the map $T$ defined in (6.3) pushes forward $\pi_{\bm{Y},\bm{X}}$
to $\pi_{\bm{X}|\bm{y}^{*}}$, and thus defines a non-deterministic coupling
between prior and posterior distributions.

In practice, we need to build an estimator, $\widehat{T}$, for $T$. We propose
to use

$\widehat{T}\coloneqq\widehat{S}^{\bm{\mathcal{X}}}(\bm{y}^{*},\cdot)^{-1}\circ\widehat{S}^{\bm{\mathcal{X}}},$
(6.4)

where we replace $S$ in (6.3) by a constrained maximum likelihood estimator
$\widehat{S}$. We describe this estimator in Section 7. It turns out that the
components of $\widehat{S}$ can be computed quickly—in parallel, and via
convex optimization—from a collection of samples
$(\bm{y}^{i},\bm{x}^{i})\sim\pi_{\bm{Y},\bm{X}}$.

We summarize the resulting analysis step with stochastic maps in Algorithm 2.
The term stochastic maps refers to the use of non-deterministic couplings
between forecast and filtering distributions.

Figure 4:  Illustration of the analysis step in the stochastic map filter. For
each prior particle $\bm{x}^{i}$, we generate a sample $\bm{y}^{i}$ from the
likelihood function $\pi_{\bm{Y}|\bm{X}}(\cdot|\bm{x}^{i})$. The pairs
$(\bm{y}^{i},\bm{x}^{i})$ represent i.i.d. samples from the joint distribution
$\pi_{\bm{Y},\bm{X}}$ and are used to estimate a transport map $T$ that pushes
forward $\pi_{\bm{Y},\bm{X}}$ to the posterior
$\pi_{\bm{X}|\bm{Y}=\bm{y}^{*}}$. By the definition of $T$,
$(T(\bm{y}^{i},\bm{x}^{i}))$ gives the desired samples from the posterior
distribution.  Algorithm 1 (Invert a triangular map)
Given a map $S$ of the form (5.1) and an $\bm{x}\in\mathbb{R}^{n}$, solve the
nonlinear system $S(\bm{z})=\bm{x}$ for $\bm{z}$.

1:procedure InvertMap( S, $\bm{x}$ )

2: for $k\leftarrow 1$ to $n$ do

3: $z_{k}\leftarrow$ Solve $S^{k}(z_{1},\ldots,z_{k-1},\xi)=x_{k}$ for
$\xi\in\mathbb{R}$ $\triangleright$ one-dimensional root finding [53]

4: return ${\bm{z}=(z_{1},\ldots,z_{n})}$

Algorithm 2 (Analysis step with stochastic maps)
Given $M$ samples $\bm{x}^{1},\ldots,\bm{x}^{M}$ from the forecast
distribution, a likelihood function $\pi_{\bm{Y}|\bm{X}}$, and an observation
$\bm{y}^{*}$, generate a particle approximation $\bm{z}^{1},\ldots,\bm{z}^{M}$
of the filtering distribution $\pi_{\bm{X}|\bm{y}^{*}}$.

1:procedure AnalysisStochastic( $\bm{y}^{*}$, $\pi_{\bm{Y}|\bm{X}}$,
$\bm{x}^{1},\ldots,\bm{x}^{M}$ )

2: for $i\leftarrow 1$ to $M$ do $\triangleright$ Generate samples from
$\pi_{\bm{Y},\bm{X}}$

3: $\bm{y}^{i}\leftarrow$ sample from $\pi_{\bm{Y}|\bm{X}}(\cdot|\bm{x}^{i})$

4: $\widehat{S}^{\bm{\mathcal{X}}}\leftarrow$ estimator of
$S^{\bm{\mathcal{X}}}$ from $(\bm{y}^{i},\bm{x}^{i})_{i=1}^{M}$ given by
(7.2). $S^{\bm{\mathcal{X}}}$ is defined in (6.1)

5: for $i\leftarrow 1$ to $M$ do $\triangleright$ Action of $\widehat{T}$ as
defined in (6.4)

6: $\bm{z}^{i}\leftarrow\widehat{S}^{\bm{\mathcal{X}}}(\bm{y}^{i},\bm{x}^{i})$

7: $\bm{z}^{i}\leftarrow$ InvertMap(
$\widehat{S}^{\bm{\mathcal{X}}}({\bm{y}^{*}},\cdot)$, $\bm{z}^{i}$ )

8: return $\bm{z}^{1},\ldots,\bm{z}^{M}$

##  7 Estimation of the Knothe–Rosenblatt rearrangement from samples

In this section, we address the estimation of a KR rearrangement $S$ that
pushes forward an arbitrary density $\pi$ on $\mathbb{R}^{n}$ to a standard
normal $\eta=\mathcal{N}({\bf 0},{\bf I}_{n})$, given only $M$ samples
$\bm{z}^{1},\ldots,\bm{z}^{M}$ from $\pi$. This problem is at the heart of the
filtering methodologies proposed in this paper. For instance, in the
stochastic map algorithm of Section 6, we need to estimate a KR rearrangement
that pushes forward the joint distribution of state and data,
$\pi_{\bm{X},\bm{Y}}$, to a standard normal. In that case,
$\bm{z}^{1},\ldots,\bm{z}^{M}$ represent samples from $\pi_{\bm{X},\bm{Y}}$.

We recall a construction first proposed by [51]. We can build a constrained
maximum likelihood estimator for $S$ as follows: let $\bm{\mathcal{H}}$ be a
finite dimensional approximation space for the KR rearrangement, i.e., let
$\bm{\mathcal{H}}$ be a cone of monotone triangular maps
$U:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ that can be described with
finitely many parameters. For each $U$ in $\bm{\mathcal{H}}$ consider the
density given by pulling back $\eta$ by $U$, defined as
$U^{\sharp}\,\eta\coloneqq(U^{-1})_{\sharp}\,\eta$ and given explicitly by

$U^{\sharp}\,\eta(\bm{z})=\eta(U(\bm{z}))\,\det\nabla U(\bm{z}),$ (7.1)

where $\det\nabla U$ denotes the determinant of the gradient of the map. Note
that $\pi$ can be written as $S^{\sharp}\,\eta$ by the definition of the KR
rearrangement $S$. Thus, the family of densities
$(U^{\sharp}\,\eta)_{U\in\bm{\mathcal{H}}}$ defines a parametric model for
$\pi$—the law of the samples $\bm{z}^{1},\ldots,\bm{z}^{M}$. A maximum
likelihood estimator, $\widehat{S}$, for $S$ can then be written as

$\widehat{S}\in\arg\max_{U\in\bm{\mathcal{H}}}\,\frac{1}{M}\sum_{i=1}^{M}\,\log
U^{\sharp}\,\eta(\bm{z}^{i}).$ (7.2)

Indeed, $\widehat{S}^{\sharp}\,\eta$ minimizes an empirical estimate of the
Kullback–Leibler divergence between $\pi$ and the class of densities given by
(7.1) for $U\in\bm{\mathcal{H}}$.

Since $\eta=\mathcal{N}({\bf 0},{\bf I}_{n})$, a simple calculation shows that
each component, $\widehat{S}^{k}$, of $\widehat{S}$ can be computed
independently as [51]

$\widehat{S}^{k}\in\arg\min_{U^{k}\in\bm{\mathcal{H}}_{k}}\,\frac{1}{M}\sum_{i=1}^{M}\,\left(\frac{1}{2}\,U^{k}(\bm{z}^{i})^{2}-\log\partial_{k}U^{k}(\bm{z}^{i})\right),$
(7.3)

where $\bm{\mathcal{H}}_{k}$ is a function space for the $k$th component,
$U^{k}$, of $U\in\bm{\mathcal{H}}$, and where $\partial_{k}U^{k}$ denotes the
derivative of $U^{k}$ with respect to its $k$th input variable.
$\bm{\mathcal{H}}_{k}$ contains functions
$U^{k}:\mathbb{R}^{k}\rightarrow\mathbb{R}$ that are monotone with respect to
the last input variable, i.e., where each $\xi\mapsto
U^{k}(z_{1},\ldots,z_{k-1},\xi)$ is an increasing map.

In practice, we need to pick a parameterization for $U^{k}$ in (7.3), i.e., we
need to specify a choice for $\bm{\mathcal{H}}_{k}$. Different choices lead to
different filtering algorithms (cf. Algorithm 2). For instance, if we restrict
$\bm{\mathcal{H}}_{k}$ to contain only affine maps—i.e.,
$U^{k}(\bm{z})=u_{0}+\sum_{i\leq k}u_{i}\,z_{i}$ for some unknown coefficients
$(u_{i})$—then we obtain an EnKF-type algorithm. (We clarify this connection
in Section 9.) If we depart from the linear ansatz, however, we obtain a whole
new class of nonlinear filtering algorithms. For example, in this paper we
experiment with separable parameterizations of the form
$U^{k}(\bm{z})=\sum_{i\leq k}\bm{\mathfrak{u}}_{i}(z_{i})$ for some collection
of nonlinear functions $(\bm{\mathfrak{u}}_{i})$,

$\bm{\mathfrak{u}}_{i}(z)=\sum_{j}u_{ij}\,\psi_{j}(z),$ (7.4)

which can be expressed in terms of finitely many basis functions
$(\psi_{j})$—e.g., Hermite polynomials or radial basis functions—and unknown
coefficients $(u_{ij})$. These parameterizations are perhaps the simplest way
to increase complexity with respect to a linear function, and lead to
filtering algorithms that can outperform the EnKF in several challenging
scenarios, with only a marginal increase in computational cost (see Section
11). We describe separable parameterizations in Appendix A, together with a
general way to represent monotone triangular maps.

For the parameterizations that we consider in this paper, (7.3) is a finite-
dimensional convex program that can be solved efficiently—even in high
dimensions—via standard convex optimization methods (e.g., Newton’s method).
If we make certain structural assumptions on the estimator, however, (7.3)
admits a closed form solution without any need for numerical optimization. For
instance, if we let
$U^{k}(\bm{z})=\bm{\mathfrak{u}}(z_{1},\ldots,z_{k-1})+\bm{\mathfrak{u}}_{k}(z_{k})$
for some nonlinear function $\bm{\mathfrak{u}}$ and affine
$\bm{\mathfrak{u}}_{k}$, then (7.3) reduces to a standard linear regression
problem, provided that we consider parameterizations of $\bm{\mathfrak{u}}$
that are linear in the unknown coefficients. We refer to Appendix A for
additional details on this topic.

##  8 Regularization in high dimensions: localization of nonlinear maps

In this section, we discuss how to regularize the estimation of the KR
rearrangement $S$ in high dimensions, by imposing certain functional
constraints—e.g., sparsity—on the structure of the maximum likelihood
estimator $\smash{\widehat{S}}$ defined in (7.2). Regularization is the
cornerstone of filtering algorithms that aim to work in high dimensions with a
limited ensemble size: examples include covariance tapering and inflation in
the EnKF, or weight localization in particle filters [23, 52]. In this work,
we do not regularize the estimation of a covariance matrix, but rather the
estimation of a more general nonlinear function: the KR rearrangement $S$.

Like every estimator, $\widehat{S}$ has a bias—which quantifies how well we
can approximate the true KR rearrangement in the limit of infinite data—and a
variance—which measures the variability of the estimate with respect to
different realizations of the samples $\bm{z}^{1},\ldots,\bm{z}^{M}$ from
$\pi$. Both bias and variance of $\widehat{S}$—which is a random function—can
be controlled by changing the complexity of the approximation space for the
estimator. Very rich parameterizations can yield estimators with low bias, but
potentially unacceptable variance, if the number of samples
$\bm{z}^{1},\ldots,\bm{z}^{M}$ is insufficient. There is no point in trying to
learn arbitrarily nonlinear transformations. In fact, that would be harmful.
On the other hand, overly simple parameterizations that can be fitted with low
variance might have an unnecessarily large estimation bias, leading to a waste
of computational resources. Successful estimation requires addressing this
bias-variance tradeoff, especially when dealing with high dimensions and low
ensemble regimes.

There are many ways to regularize the estimation of the KR rearrangement. For
instance, one could augment the average log-likelihood in (7.2) with a convex
penalty on the coefficients of $\widehat{S}$—e.g., a ridge or lasso functional
[28]—yielding a penalized maximum likelihood estimator. In addition, one could
constrain the function space for the estimator by using only low-order
nonlinear basis functions (see Appendix A). In this section, we focus on the
complementary idea of regularizing the estimation by imposing sparsity
constraints on $\widehat{S}$, i.e., by dropping some variable dependence from
each component of the map, such that

$\widehat{S}^{k}(z_{1},\ldots,z_{k})=\widehat{S}^{k}(z_{j_{1}},\ldots,z_{j_{p}},z_{k})$
(8.1)

for some indices $j_{1},\ldots,j_{p}$ and integer $p<k-1$ [64]—effectively
introducing a notion of “localization” for nonlinear transformations. We now
discuss model assumptions that justify a sparsity pattern like (8.1).

The sparsity of a KR rearrangement $S$ that pushes forward a density $\pi$ on
$\mathbb{R}^{n}$ to $\mathcal{N}({\bf 0},{\bf I}_{n})$ depends on marginal
conditional independence properties of $\pi$. As explained in [64], if
$\bm{Z}=(Z_{1},\ldots,Z_{n})\sim\pi$, then the $k$th component of $S$ is not a
function of the $j$th input variable ($j<k$) if $Z_{j}$ and $Z_{k}$ are
conditionally independent given the remaining variables with index less than
$k$, i.e., if

$Z_{j}\perp\\!\\!\\!\perp Z_{k}|\bm{Z}_{\\{1:k-1\\}\setminus\,j}.$ (8.2)

Whether (8.2) holds, perhaps only approximately, depends on properties of
$\pi$ such as decay of dependence and conditional independence. These
properties are fairly common when filtering spatially distributed processes
since they encode an intuitive notion of local probabilistic interaction
between the random variables [57, 68].

#### Decay of dependence and banded estimators.

Let $d(\cdot,\cdot)$ be a distance function on the index set
$\\{1,\ldots,n\\}$, which defines a natural ordering of the variables
$\bm{Z}$, so that larger values of $d(i,j)$ are associated with a weaker
statistical dependence between $Z_{i}$ and $Z_{j}$. For instance, in data
assimilation $d(\cdot,\cdot)$ could be derived from a metric in physical
space. The first observation is that the $k$th component of the KR
rearrangement should depend weakly on variables that are far away from
$Z_{k}$, according to $\bm{d}(\cdot,\cdot)$. To exploit this structure, let
$\mathcal{A}_{k,r}$ denote the set of indices of state variables that are
within a distance $r$ from $Z_{k}$, i.e.,

$\mathcal{A}_{k,r}\coloneqq\\{i\in(1,\ldots,k)\,\,{\rm s.t.}\,\,d(i,k)\leq
r\\}.$ (8.3)

We propose to build an estimator $\widehat{S}$ of $S$ for which each component
$\widehat{S}^{k}$ depends only on the variables in $\mathcal{A}_{k,r}$,

$\widehat{S}(\bm{z})=\left[\begin{array}[]{l}\widehat{S}^{1}(\bm{z}_{\mathcal{A}_{1,r}})\\\\[2.0pt]
\vdots\\\ \widehat{S}^{n}(\bm{z}_{\mathcal{A}_{n,r}})\\\ \end{array}\right],$
(8.4)

for an “optimal” value of $r\geq 0$. This $r$ is obtained by minimizing some
prediction metric (e.g., RMSE) in an offline calibration phase, as if it were
the optimal localization radius for the EnKF [23].

In the case of linear maps, (8.4) corresponds to a “banded” estimator for the
Cholesky factor of the precision matrix of the ensemble
$\bm{z}^{1},\ldots,\bm{z}^{M}$ [73, 7], and has been used for the purpose of
regularized covariance estimation in the context of data assimilation [50].
Indeed, if $\pi=\mathcal{N}({\bf 0},\bm{\Sigma})$ and
$\bm{\Sigma}=\bm{L}\bm{L}^{\top}$ is a Cholesky decomposition of
$\bm{\Sigma}\succ 0$, then $\bm{S}\coloneqq\bm{L}^{-1}$ represents a linear KR
rearrangement that pushes forward $\pi$ to a standard normal. In this paper,
we extend these banded estimators to the nonlinear/non-Gaussian case, using
the KR rearrangement to generalize the notion of a Cholesky factor.

#### Markov properties and graphical models.

The sparsity pattern of the estimator $\widehat{S}$ can also be justified on
the grounds of conditional independence properties of $\pi$—the so-called
Markov properties. These properties can be collected in a simple undirected
graph $\bm{\mathcal{G}}=(\mathcal{V},\mathcal{E})$, where each node
$i\in\mathcal{V}$ is associated with one of the random variables in the
collection $\bm{Z}=(Z_{1},\ldots,Z_{n})$, and where each edge
$(i,j)\in\mathcal{E}$ encodes a notion of probabilistic interaction between
the pair $Z_{i},Z_{j}$. We say that $\bm{Z}$ (or $\pi$) is a Markov random
field (MRF) with respect to $\bm{\mathcal{G}}$ if for any partition
$\mathcal{A},\mathcal{S},\mathcal{B}$ of the vertex set—for which
$\mathcal{S}$ is a separator set333$\mathcal{S}$ is a separator set for
$\mathcal{A}$ and $\mathcal{B}$ if every path from a node in $\mathcal{A}$ to
a node in $\mathcal{B}$ goes through $\mathcal{S}$.  for $\mathcal{A}$ and
$\mathcal{B}$—we have that $\bm{Z}_{\mathcal{A}}$ and $\bm{Z}_{\mathcal{B}}$
are conditionally independent given $\bm{Z}_{\mathcal{S}}$, i.e.,
$\bm{Z}_{\mathcal{A}}\perp\\!\\!\\!\perp\bm{Z}_{\mathcal{B}}|\bm{Z}_{\mathcal{S}}$
[35]. See Figure 6 (left) for an example. For any strictly positive $\pi$, we
can find the sparsest $\bm{\mathcal{G}}$ that is compatible with the
conditional independence structure of $\pi$ as follows: let
$(i,j)\notin\mathcal{E}$ if and only if
$\bm{Z}_{i}\perp\\!\\!\\!\perp\bm{Z}_{j}|\bm{Z}_{\mathcal{V}\setminus\\{i,j\\}}$
[34].

A key result in [64] links the sparsity of the KR rearrangement $S$ to the
sparsity of $\bm{\mathcal{G}}$. For a given $\bm{\mathcal{G}}$, define the
sequence of marginal graphs $\bm{\mathcal{G}}^{n},\ldots,\bm{\mathcal{G}}^{1}$
using a simple recursion: (1) $\bm{\mathcal{G}}^{n}\coloneqq\bm{\mathcal{G}}$,
and (2) $\bm{\mathcal{G}}^{k-1}$ is obtained from $\bm{\mathcal{G}}^{k}$ by
removing node $k$ and by turning its neighborhood ${\rm
Nb}(k,\bm{\mathcal{G}}^{k})$ into a fully connected subset,444For any graph
$\bm{\mathcal{G}}$, we denote by ${\rm Nb}(k,\bm{\mathcal{G}})$ the
neighborhood of a node $k$ in $\bm{\mathcal{G}}$.  i.e., a clique. By [64,
Theorem 5.1], each component $S^{k}$ of the rearrangement can only depend on
variables in ${\rm Nb}(k,\bm{\mathcal{G}}^{k})$, for all $k=1,\ldots,n$. See
Figure 5 for an illustration (adapted from [64]). Hence, if we know that $\pi$
is Markov with respect to a sparse graph, then we can translate this
information into sparsity of the estimator $\widehat{S}$, possibly increasing
the sparsity pattern in (8.4).

If $\pi=\mathcal{N}({\bf 0},\bm{\Sigma})$, then describing the sparsity of a
linear KR rearrangement due to conditional independence reduces to the problem
of characterizing the fill-in for the Cholesky factor of the precision matrix
$\bm{\Sigma}^{-1}$—a well-known result in the field of Gaussian graphical
models [34], which has also been exploited in the context of data assimilation
to devise regularized estimators of the ensemble covariance [68]. Here, we
consider the non-Gaussian generalization provided by [64], where the Cholesky
factor is replaced by a nonlinear KR rearrangement. In particular, we note
that the sparsity of $S$ depends on the ordering of the input variables in the
same way as the fill-in of the Cholesky factor does. Different orderings can
yield different sparsity patterns. See Figure 6 (right) for an example.
Whenever possible, we then look for orderings that promote sparsity (cf. [64]
for additional details).

In many cases, $\pi$ is only approximately Markov with respect to a sparse
graph $\bm{\mathcal{G}}$, e.g., there might exist a second density
$\pi^{\prime}$ which is Markov with respect to $\bm{\mathcal{G}}$ and such
that $\mathcal{D}_{\rm KL}(\pi^{\prime}||\pi)<\varepsilon$, for some
$\varepsilon>0$. In particular, we know that the filtering marginals will not
be “exactly” Markov with respect to a sparse $\bm{\mathcal{G}}$, since
conditional independence properties do not persist over time for general
models [34]. Yet, when filtering spatially distributed processes (e.g., in
data assimilation) it is reasonable to expect that the filtering (or forecast)
marginals will be well-approximated by sparse MRFs, since the latter encode
probabilistic models with local or neighboring interactions. For example, in
Section 11 we consider the Lorenz-96 model, which arises from the
discretization of a one-dimensional PDE on a periodic domain. In that case, we
can expect that the filtering marginal is approximately Markov with respect to
a cycle graph with a low-degree of connectivity between the nodes. When
dealing with two-dimensional PDEs, it is natural to consider grid graphs [68,
38], or more general lattices on the plane, and so on for higher dimensional
domains. In practice, we can then sparsify the estimator $\widehat{S}$
according to a graph $\bm{\mathcal{G}}$, which implies Markov properties for
$\pi$ that are only approximately satisfied, and treat the exact geometry of
the graph—e.g., its neighborhood structure—as a calibration parameter of the
algorithm, as if it were a localization radius. This is equivalent to
regularizing the estimation of the KR rearrangement by projecting $\pi$ onto a
manifold of sparse MRFs [34].

$3$$2$$5$$1$$4$ (a) $\bm{\mathcal{G}}^{5}\coloneqq\bm{\mathcal{G}}$

$3$$2$$5$$1$$4$ (b) $\bm{\mathcal{G}}^{4}$

$3$$2$$5$$1$$4$ (c) $\bm{\mathcal{G}}^{3}$

$5$$2$$4$$1$$3$ (d) $\bm{\mathcal{G}}^{2}$

(e) $\partial_{j}S^{k}$

Figure 5: Here, we assume that $\pi$ is Markov with respect to the leftmost
graph $\bm{\mathcal{G}}$, and show the sequence of marginals graphs
$(\bm{\mathcal{G}}^{k})$ defined in Section 8. We then represent the sparsity
of the corresponding KR rearrangement $S$ using symbolic matrix notation: the
$(j,k)$-th entry of the matrix is not colored if the $k$th component of the
map does not depend on the $j$th input variable.

$\mathcal{A}$$\mathcal{B}$$\mathcal{S}$ (a)

$\mathcal{A}$$\mathcal{B}$$\mathcal{S}$$5$$2$$3$$1$$4$ (b)
$\bm{\mathcal{G}}^{\prime}$

$\mathcal{A}$$\mathcal{B}$$\mathcal{S}$ (c) $\partial_{j}S^{k}$

Figure 6: (left) Example of a graphical model for $\bm{Z}$, together with a
partition $\mathcal{A},\mathcal{S},\mathcal{B}$ of the vertex set, where
$\mathcal{S}$ is a separator for $\mathcal{A}$ and $\mathcal{B}$. By
definition,
$\bm{Z}_{\mathcal{A}}\perp\\!\\!\\!\perp\bm{Z}_{\mathcal{B}}|\bm{Z}_{\mathcal{S}}$.
(center) Re-ordering of the input variables for the example of Figure 5,
leading to a rearrangement $S$ with no sparsity (right).

##  9 Remarks on the stochastic map filter

The stochastic map algorithm of Section 6 is a general tool for Bayesian
inference: given $M$ samples $(\bm{x}^{i})$ from the prior, (1) generate $M$
samples $\bm{y}^{i}\sim\pi_{\bm{Y}|\bm{X}}(\cdot|\bm{x}^{i})$ from the
likelihood, (2) use the pairs $(\bm{y}^{i},\bm{x}^{i})$ to build an estimator
$\widehat{T}$ of the analysis map $T$ as

$T\coloneqq S^{\bm{\mathcal{X}}}(\bm{y}^{*},\cdot)^{-1}\circ
S^{\bm{\mathcal{X}}},\qquad\widehat{T}\coloneqq\widehat{S}^{\bm{\mathcal{X}}}(\bm{y}^{*},\cdot)^{-1}\circ\widehat{S}^{\bm{\mathcal{X}}},$
(9.1)

and (3) obtain a particle approximation of the posterior as
$\widehat{T}(\bm{y}^{1},\bm{x}^{1}),\ldots,\widehat{T}(\bm{y}^{M},\bm{x}^{M})$.

A few remarks are in order. First, we only need to sample the likelihood
function and do not need explicit access to the density $\pi_{\bm{Y}|\bm{X}}$,
which might be intractable or expensive to evaluate. We are thus in the
setting of approximate Bayesian computation (ABC) methods [44] or, more
broadly, inference in _generative models_. Second, we can compute
$\smash{\widehat{T}}$ using convex optimization: we only need to estimate a
few components, $\smash{S^{\bm{\mathcal{X}}}}$, of the KR rearrangement $S$
that pushes forward $\pi_{\bm{Y},\bm{X}}$ to a standard normal, using the
constrained maximum likelihood estimator defined in (7.2). In particular, the
block $\smash{S^{\bm{\mathcal{Y}}}}$ defined in (6.1) is not needed, while the
components of $\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ can be computed in
parallel using (7.3). Third, the estimation of $\smash{S^{\bm{\mathcal{X}}}}$
in high dimensions can be regularized using the localization techniques of
Section 8. (For the purpose of posterior sampling, we may only care about the
bias-variance tradeoff in the estimation of $T$, and not necessarily that of
$\smash{S^{\bm{\mathcal{X}}}}$.) Fourth, if we constrain the estimator
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ to be linear, then we recover an EnKF
with “perturbed observations,” i.e., the stochastic EnKF [13].

######  Remark 2 (Connection with the EnKF).

Assume that the samples $(\bm{y}^{i},\bm{x}^{i})$ have been centered, so that
they have zero mean. Then, the best linear estimator of
$\smash{S^{\bm{\mathcal{X}}}}$—according to (7.2)—is given by

$\widehat{S}^{\bm{\mathcal{X}}}(\bm{y},\bm{x})=\bm{A}\left(\bm{x}-\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}\widehat{\bm{\Sigma}}_{\bm{Y}}^{-1}\bm{y}\right),$
(9.2)

where $\bm{A}$ is the inverse of the Cholesky factor of the “residual”
covariance matrix,

$\widehat{\bm{\Sigma}}_{\bm{X}|\bm{Y}}\coloneqq\widehat{\bm{\Sigma}}_{\bm{X}}-\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}\,\widehat{\bm{\Sigma}}_{\bm{Y}}^{-1}\,\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}^{\top},$
(9.3)

where $\widehat{\bm{\Sigma}}_{\bm{X}}$, $\widehat{\bm{\Sigma}}_{\bm{Y}}$, and
$\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}$ are empirical (maximum likelihood)
estimators of ${\mathbb{V}\mathrm{ar}}(\bm{X})$,
${\mathbb{V}\mathrm{ar}}(\bm{Y})$, and
${\mathbb{C}\mathrm{ov}}(\bm{X},\bm{Y})$, respectively.555In the case of a
linear–Gaussian observation model, e.g., if $\bm{Y}=\bm{H}\bm{X}+\mathcal{E}$
for $\mathcal{E}\sim\mathcal{N}({\bf 0},\bm{R})$ independent of $\bm{X}$, and
with $\bm{H}$ known, we could regularize the estimation of
${\mathbb{C}\mathrm{ov}}(\bm{X},\bm{Y})$ and ${\mathbb{V}\mathrm{ar}}(\bm{Y})$
by letting
$\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}\coloneqq\widehat{\bm{\Sigma}}_{\bm{X}}\bm{H}^{\top}$
and
$\widehat{\bm{\Sigma}}_{\bm{Y}}=(\bm{H}\widehat{\bm{\Sigma}}_{\bm{X}}\bm{H}^{\top}+\bm{R})$.
For a given observation $\bm{y}^{*}$ of the data, we can then compute a closed
form expression for the inverse of the map
$\bm{\xi}\mapsto\smash{\widehat{S}^{\bm{\mathcal{X}}}}(\bm{y}^{*},\bm{\xi})$
and for the estimator of the analysis map,

$\widehat{T}(\bm{y},\bm{x})=\bm{x}-\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}\widehat{\bm{\Sigma}}_{\bm{Y}}^{-1}\left(\bm{y}-\bm{y}^{*}\right),$
(9.4)

which corresponds to the forecast-to-analysis update of the stochastic EnKF
[60]. The stochastic map filter offers a framework to generalize this linear
ansatz to more general nonlinear transformations.

At any assimilation time, we need to condition the forecast distribution on
$d$ scalar observations $\bm{y}^{*}=(y^{*}_{1},\ldots,y^{*}_{d})$ from
$\bm{Y}$. In Section 6, we presented an algorithm that can assimilate all
these observations simultaneously. A possible alternative if the observations
are conditionally independent given the state, i.e.,

$\pi_{\bm{Y}|\bm{X}}=\prod_{k=1}^{d}\pi_{Y_{k}|\bm{X}},$ (9.5)

is to process each scalar observation $y^{*}_{k}$ individually and
sequentially by iterating $d$ times over Algorithm 2: the posterior ensemble
associated with $y^{*}_{k}$ can be used as a prior ensemble when assimilating
$y^{*}_{k+1}$. Hence, the analysis step can be either performed by computing
$d$ analysis maps on $\smash{\mathbb{R}^{n+1}}$—each map associated with one
scalar observation—or by computing a single map on $\smash{\mathbb{R}^{n+d}}$
(cf. the EnKF with single versus multiple observations [30]).

######  Remark 3 (Localize the analysis map for a scalar observation with
local likelihood).

When assimilating a scalar observation (or a small batch thereof) it makes
sense to consider an additional form of localization for
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ with respect to those already
introduced in Section 8. The idea is that the analysis map $T$ should revert
to the identity function along components that correspond to unobserved
variables that are far from the observed ones, for instance with respect to
the metric $\bm{d}(\cdot,\cdot)$ defined in Section 8. We can enforce this
identity structure in the estimator $\smash{\widehat{T}}$ by requiring that
some components of $\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ be the identity
function. Specifically, assume that we are observing a single scalar
observation $Y$ which follows a local likelihood model, i.e.,
$\pi_{Y|\bm{X}}=\pi_{Y|X_{\ell}}$ for some scalar state variable $X_{\ell}$.
After a permutation, we can assume that $X_{\ell}=X_{1}$ and that the state
variables have been reordered in terms of their increasing distance from
$X_{1}$. We can then look for an estimator
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ with the following sparsity
pattern,666One can easily increase this sparsity pattern using the
localization ideas of Section 8.

$\widehat{S}^{\bm{\mathcal{X}}}(y,\bm{x})=\left[\begin{array}[]{l}s(y,\bm{x}_{1:j})\\\\[2.0pt]
x_{j+1}\\\ \vdots\\\ x_{n}\end{array}\right],\qquad
s:\mathbb{R}\times\mathbb{R}^{j}\rightarrow\mathbb{R}^{j},$ (9.6)

for some index $j$ that represents the analogue of a “localization” radius for
EnKF algorithms. The effect of (9.6) is to produce an estimator
$\smash{\widehat{T}}$ for the analysis map that reverts to the identity
function after $j$ components, thus producing a truly local coupling. Note
that the localization in (9.6) is only justified because we use
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ within the expression for
$\smash{\widehat{T}}$ in (9.1). We do not care if
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ is a poor estimator of the
rearrangement $\smash{S^{\bm{\mathcal{X}}}}$ as a whole, as long as the
resulting $\widehat{T}$ given by (9.1) is a good estimator of the analysis
map.

######  Remark 4 (Transform rather than sample).

The estimation of $\smash{S^{\bm{\mathcal{X}}}}$ by
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ yields an implicit approximation of
the posterior distribution: by definition, the mapping
$\bm{\xi}\mapsto\smash{S^{\bm{\mathcal{X}}}}(\bm{y}^{*},\bm{\xi})$ pushes
forward $\pi_{\bm{X}|\bm{y}^{*}}$ to a standard normal. In principle, we could
then generate approximate posterior samples by pushing forward samples from a
standard normal $\mathcal{N}({\bf 0},{\bf I}_{n})$ through the inverse of the
map
$\bm{\xi}\mapsto\smash{\widehat{S}^{\bm{\mathcal{X}}}}(\bm{y}^{*},\bm{\xi})$.
In practice, however, we do not do this; rather, we use
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ to build an estimator
$\smash{\widehat{T}}$ of the prior-to-posterior update. This is in analogy
with the EnKF, where the posterior ensemble is given by a (linear)
transformation applied to the prior ensemble, rather than by samples from a
Gaussian approximation of the posterior.

In numerical experiments—where transformations can only be approximate—we
observe that pushing forward samples from $\pi_{\bm{Y},\bm{X}}$ through
$\smash{\widehat{T}}$ often yields far more accurate posterior approximations
than pushing standard normal samples through
$\widehat{S}^{\bm{\mathcal{X}}}(\bm{y}^{*},\cdot)^{-1}$. This improvement can
be attributed to the cancellation of errors in the composition of
$\smash{\widehat{S}^{\bm{\mathcal{X}}}}$ with its partial inverse. As a
limiting but nonetheless illustrative example: when the data are uninformative
and $S^{\bm{\mathcal{X}}}$ does not depend on $\bm{y}$, the Gaussian samples
mapped through the inverse of the estimated map will be affected by bias and
variance in $\smash{\widehat{S}^{\bm{\mathcal{X}}}}$, while $\widehat{T}$ will
be the identity function and generate exact posterior samples.

We conclude this section with a final remark. We presented the stochastic map
filter using the KR rearrangement as our building-block coupling, because the
latter can be easily approximated in high dimensions using nonlinear
transformations and convex optimization. Yet the algorithm would still make
sense if we were to replace the KR rearrangement with any other transport map
$S$ that has the “block” structure in (6.1), and that pushes forward
$\pi_{\bm{Y},\bm{X}}$ to $\mathcal{N}({\bf 0},{\bf I}_{d+n})$.

##  10 Deterministic map filter

In this section, we explore the use of deterministic couplings in the analysis
step of Section 4. In particular, we seek a transport map
$T:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ that pushes forward prior to
posterior, i.e., such that $T(\bm{X})\sim\pi_{\bm{X}|\bm{y}^{*}}$. See Figure
3 for an illustration. (Note that in the stochastic map filter of Section 6,
the map $T$ was defined on $\mathbb{R}^{n+d}$ rather than on
$\mathbb{R}^{n}$.)

For all monotone triangular maps $U$ on $\mathbb{R}^{n}$, define $\pi_{U}$ to
be the density that is proportional to the function
$\bm{\xi}\mapsto\pi_{\bm{Y}|\bm{X}}(\bm{y}^{*}|\bm{\xi})\,U^{\sharp}\,\eta(\bm{\xi})$,
which is the product of the likelihood and the pullback density
$U^{\sharp}\,\eta$. Now, let $S$ be the KR rearrangement that pushes forward
the prior $\pi_{\bm{X}}$ to $\eta\coloneqq\mathcal{N}({\bf 0},{\bf I}_{n})$.
Hence, the pullback density of $\eta$ by $S$—see the definition in (7.1)—is
the prior, i.e., $S^{\sharp}\,\eta=\pi_{\bm{X}}$. It follows that $\pi_{S}$ is
the posterior density $\pi_{\bm{X}|\bm{y}^{*}}$.

We define a candidate analysis map $T$ as follows:

$T\coloneqq\mathcal{T}\circ S,$ (10.1)

where $\mathcal{T}$ is a KR rearrangement that pushes forward $\eta$ to
$\pi_{S}$. It is immediate to verify that (10.1) pushes forward $\pi_{\bm{X}}$
to $\pi_{\bm{X}|\bm{y}^{*}}$, and thus defines a valid deterministic coupling
between the prior and posterior distributions.

We propose to estimate $T$ by

$\smash{\widehat{T}}\coloneqq\widehat{\mathcal{T}}\circ\widehat{S},$ (10.2)

where $\widehat{S}$ and $\smash{\widehat{\mathcal{T}}}$ are (constrained)
maximum likelihood estimators of two different KR rearrangements: the former
pushes forward $\pi_{\bm{X}}$ to $\eta$, while the latter pushes forward
$\eta$ to $\pi_{\widehat{S}}$. Intuitively, $\pi_{\widehat{S}}$ is an
approximation of the posterior distribution, which depends on the estimator
$\smash{\widehat{S}}$. We can compute $\smash{\widehat{S}}$ efficiently via
convex optimization from a collection of prior samples
$(\bm{x}^{i})\sim\pi_{\bm{X}}$ using the construction of Section 7 and the
regularization techniques of Section 8. The computation of
${\widehat{\mathcal{T}}}$, however, is less straightforward, as we explain in
the next remark. In particular, note that we do not have samples from
$\pi_{\widehat{S}}$; instead, given $\widehat{S}$, we can only evaluate this
density up to a normalizing constant, as long as the likelihood
$\pi_{\bm{Y}|\bm{X}}(\bm{y}^{*}|\cdot)$ can itself be evaluated explicitly.

######  Remark 5 (Estimation of the Knothe–Rosenblatt rearrangement from
densities).

${\widehat{\mathcal{T}}}$ is a maximum likelihood estimator of a KR
rearrangement that pushes forward a standard normal $\mathcal{N}({\bf 0},{\bf
I}_{n})$ to a density $\pi\coloneqq\pi_{\widehat{S}}$ on $\mathbb{R}^{n}$,
i.e., ${\widehat{\mathcal{T}}}$ is defined analogously to (7.2) as

$\widehat{\mathcal{T}}\in\arg\max_{U\in\bm{\mathcal{H}}}\,\frac{1}{N}\sum_{i=1}^{N}\,\log
U^{\sharp}\,\pi(\bm{z}^{i}),$ (10.3)

for a collection of $N$ samples $(\bm{z}^{i})_{i=1}^{N}$ from
$\mathcal{N}({\bf 0},{\bf I}_{n})$ and for some approximation space
$\bm{\mathcal{H}}$. There is an important difference between (7.2) and (10.3):
the density $\eta$ in (7.2) is always log-concave, while $\pi$ in (10.3) need
not be. As a result, (10.3) is in general not a convex problem [33]. Moreover,
we cannot compute the components of ${\widehat{\mathcal{T}}}$ independently as
in (7.3), since $\pi$ need not factorize as a product of its marginals. Hence,
the computation of ${\widehat{\mathcal{T}}}$ is inherently harder than that of
${\widehat{S}}$, but still feasible. The numerical solution of (10.3) in the
context of Bayesian inference was pioneered by [49], which uses gradient-based
optimization (e.g., BFGS or Newton-CG [72]) to minimize an equivalent
reformulation of (10.3),

$\widehat{\mathcal{T}}\in\arg\min_{U\in\bm{\mathcal{H}}}\,-\frac{1}{N}\sum_{i=1}^{N}\,\left(\log\,\pi(U(\bm{z}^{i}))+\sum_{k=1}^{n}\log\partial_{k}U^{k}(\bm{z}^{i})\right),$
(10.4)

over a finite dimensional space of polynomial maps (see Appendix A and [45,
8]). Three aspects of (10.4) are particularly important. First, we only need
to evaluate $\pi_{\widehat{S}}$ up to a normalizing constant. Second, the
number $N$ of samples from $\mathcal{N}({\bf 0},{\bf I}_{n})$ need not be
limited by the cardinality of the forecast ensemble, in contrast to (7.2).
Third, we must evaluate the likelihood function $\pi_{\bm{Y}|\bm{X}}$ for each
sample in (10.4). The latter is a crucial contrast with the stochastic map
filter of Section 6, which only requires sampling from (i.e., simulating) the
likelihood.

In the special case of a Gaussian likelihood $\pi_{\bm{Y}|\bm{X}}$ and linear
$\widehat{S}$, the estimator $\pi_{\widehat{S}}$ of the posterior distribution
becomes Gaussian. The corresponding KR rearrangement that pushes forward
$\mathcal{N}({\bf 0},{\bf I}_{n})$ to $\pi_{\widehat{S}}$ is linear, meaning
that without loss of generality we can solve (10.4) over the space of linear
maps. In this case, (10.4) is a convex (quadratic) problem with a closed form
expression for $\smash{\widehat{\mathcal{T}}}$, and the deterministic map
filter reduces to a particular ensemble square-root filter [67].

If either $\pi_{\bm{Y}|\bm{X}}$ is non-Gaussian or $\smash{\widehat{S}}$ is
nonlinear—e.g., if we seek a non-Gaussian approximation of the prior
distribution—then the linear ansatz for $\smash{\widehat{\mathcal{T}}}$ is no
longer optimal. In such cases, we might want to capture non-Gaussian structure
by considering parameterizations of $\smash{\widehat{\mathcal{T}}}$ that
depart gradually and incrementally from that of a linear function, seeking a
balance between accuracy and computational cost (see Section 7). In
particular, we can regularize the estimation of the rearrangement in high
dimensions using localization ideas similar777Most of the discussion of
Section 8 is specific to a KR rearrangement whose target distribution is a
product measure on $\mathbb{R}^{n}$—e.g., $\mathcal{N}({\bf 0},{\bf
I}_{n})$—and thus does not apply to the estimation of $\mathcal{T}$. See [64]
for additional details.  to those of Section 8. We propose to use a “banded”
estimator analogous to (8.4) to leverage decay of dependence in $\pi$, i.e.,

$\widehat{\mathcal{T}}(\bm{z})=\left[\begin{array}[]{l}\widehat{\mathcal{T}}^{1}(\bm{z}_{\mathcal{A}_{1,r}})\\\\[2.0pt]
\vdots\\\ \widehat{\mathcal{T}}^{n}(\bm{z}_{\mathcal{A}_{n,r}})\\\
\end{array}\right],$ (10.5)

for some value of the localization radius $r\geq 0$. For a Gaussian $\pi$,
(10.5) corresponds to a banded estimator for the Cholesky factor of the
covariance matrix of $\pi$, rather than its precision (cf. Section 8).

We summarize the analysis step with deterministic maps in Algorithm 3. Note
that there are many possible variations for the definition of the prior-to-
posterior update (10.1). Most notably, $\smash{\widehat{\mathcal{T}}}$ need
not approximate a KR rearrangement: in fact any map that pushes forward $\eta$
to $\pi_{S}$ would work, e.g., maps induced by the flows of ODEs [17, 3, 29]
or by the composition of many simple functions [39, 19], including deep neural
networks [58].

######  Remark 6 (Connection with the variational mapping particle filter).

The variational mapping particle filter [54] uses the nonparametric
variational inference algorithm proposed by [39] to generate approximate
samples from a non-Gaussian approximation of the filtering distribution. The
idea is to minimize precisely (10.3)—a non-convex functional—by implementing a
form of functional gradient descent on a reproducing kernel Hilbert space
(RKHS) of transformations (with no restrictions to triangular maps). The non-
Gaussian approximation of the unnormalized filtering distribution is defined
by the product of the likelihood function with an approximation of the
forecast distribution at time $k$ given by888For ease of exposition, we assume
that the dynamic of the latent field is only given at the observation times.

$\pi_{\bm{Z}_{k+1}|\bm{y}_{1:k}}\approx\frac{1}{M}\,\sum_{i=1}^{M}\pi_{\bm{Z}_{k+1}|\bm{Z}_{k}}(\cdot|\bm{z}^{i}),$
(10.6)

where $(\bm{z}^{i})$ represent (approximate) samples from the previous
filtering marginal $\pi_{\bm{Z}_{k}|\bm{y}_{1:k}}$. This approximation of the
forecast density contrasts with the regularized maximum likelihood estimation
of the KR rearrangement $S$ in (10.1) from the forecast ensemble, using (7.2).
Indeed, it would be interesting to approximate the rearrangement
$\smash{\widehat{T}}$ in the deterministic map algorithm using the technique
of [39] (or its second-order extension [19]), after “localizing” the RKHS of
transformations by imposing a sparsity pattern like the one in (10.5).
Regularization is critical in high dimensions.

The deterministic map filter presented in this section features an inherently
non-convex step. This is an important distinction from the stochastic map
filter of Section 6, which relies solely on convex optimization (at the price
of computing a higher dimensional transformation). In Appendix B, we discuss a
common structural assumption on the state-space model that allows us to easily
bypass non-convex optimization in the deterministic map algorithm, leading to
a scheme with close ties to the multivariate rank histogram filter [46].

Algorithm 3 (Analysis step with deterministic maps)
Given $M$ samples $\bm{x}^{1},\ldots,\bm{x}^{M}$ from the forecast
distribution, a likelihood function $\pi_{\bm{Y}|\bm{X}}$, and an observation
$\bm{y}^{*}$, generate a particle approximation $\bm{z}^{1},\ldots,\bm{z}^{M}$
of the filtering distribution $\pi_{\bm{X}|\bm{y}^{*}}$. Let
$\eta\coloneqq\mathcal{N}({\bf 0},{\bf I}_{n})$.

1:procedure AnalysisDeterministic( $\bm{y}^{*}$, $\pi_{\bm{Y}|\bm{X}}$,
$\bm{x}^{1},\ldots,\bm{x}^{M}$ )

2: $\widehat{S}\leftarrow$ estimator of $S$ from $(\bm{x}^{i})_{i=1}^{M}$
given by (7.2) $\triangleright$ $S_{\sharp}\,\pi_{\bm{X}}=\eta$

3: Define $\pi_{\widehat{S}}$ as
$\pi_{\widehat{S}}(\bm{\xi})\propto\pi_{\bm{Y}|\bm{X}}(\bm{y}^{*}|\bm{\xi})\,\widehat{S}^{\sharp}\,\eta(\bm{\xi})$

4: $\widehat{\mathcal{T}}\leftarrow$ estimator of $\mathcal{T}$ from
$\pi_{\widehat{S}}$ given by (10.4) $\triangleright$
$\mathcal{T}_{\sharp}\,\eta=\pi_{\widehat{S}}$

5: for $i\leftarrow 1$ to $M$ do $\triangleright$ Action of $\widehat{T}$ as
defined in (10.2)

6: $\bm{z}^{i}\leftarrow\widehat{S}(\bm{x}^{i})$

7: $\bm{z}^{i}\leftarrow\widehat{\mathcal{T}}(\bm{z}^{i})$

8: return $\bm{z}^{1},\ldots,\bm{z}^{M}$

##  11 Numerical experiments

In this section, we numerically investigate the performance of the stochastic
map filter over a range of examples. We show that nonlinear configurations of
the new filter can have tracking performance consistently superior to that of
the EnKF. We also show that the stochastic map filter provides better
approximations of the true (Bayesian) filtering distribution. We focus on the
stochastic map filter because it can be implemented using only convex
optimization, regardless of any nonlinear/non-Gaussian structure in the state-
space model. Numerical experiments with earlier versions of the deterministic
map filter can be found in related work by the authors [64, Ch. 6], for the
special case of conditionally independent and local observations (see Appendix
B). Many aspects of our problem setup and our performance assessment follow
[37], as detailed below.

###  11.1 Common aspects of the problem configurations

First we discuss aspects of the numerical experiments that are shared among
the test problems of Sections 11.2–11.4.

#### Prior models.

We consider two dynamical models of increasing dimension: the Lorenz-63 and
the Lorenz-96 models, which are widely used testbeds for filtering in chaotic
dynamical systems. The Lorenz-63 model is a three-dimensional system
introduced in [40] that describes the natural convection of a heated fluid.
The state at time $t$ is a three-dimensional vector
$\bm{Z}(t)=(Z_{1}(t),Z_{2}(t),Z_{3}(t))$ whose dynamics are given by the ODE
system,

$\frac{dZ_{1}}{dt}=-\sigma Z_{1}+\sigma
Z_{2},\;\;\;\frac{dZ_{2}}{dt}=-Z_{1}Z_{3}+\rho
Z_{1}-Z_{2},\;\;\;\frac{dZ_{3}}{dt}=Z_{1}Z_{2}-\beta Z_{3},$ (11.1)

where $(\beta,\rho,\sigma)$ are fixed parameters. In our experiments, we set
$\beta=8/3$, $\rho=28$, and $\sigma=10$, which produces chaotic solutions
tending to the well-known Lorenz attractor.

The Lorenz-96 model is a popular testbed for numerical weather prediction
algorithms, reproducing coarse features of the mid-latitude atmosphere [41,
42]. The state at time $t$ is a 40-dimensional vector,
$\bm{Z}(t)=(Z_{1}(t),\ldots,Z_{40}(t))$. The state dynamics are defined by a
set of nonlinear ODEs that represent the spatial discretization of a time-
dependent PDE, i.e.,

$\frac{{\rm d}Z_{j}}{{\rm d}t}=(Z_{j+1}-Z_{j-2})Z_{j-1}-Z_{j}+F,\qquad
j=1,\ldots,40,$ (11.2)

with periodic boundary conditions and constant forcing parameter $F$ [57]. In
our experiments we use $F=8$, which leads to a fully chaotic dynamic [42].

We integrate both ODE systems, (11.1) and (11.2), using a fourth-order
explicit Runge-Kutta method, with constant stepsizes $\Delta t=0.05$ and
$\Delta t=0.01$, respectively. For the Lorenz-63 model only, we also wish to
construct a “reference” Bayesian filtering solution via a consistent
formulation of the particle filter. In this case, we add a small amount of
Gaussian noise to the state at each integration step,
$\bm{\varepsilon}_{\Delta t}\sim\mathcal{N}(0,10^{-4}\mathbf{I}_{3})$, for the
purpose of delaying degeneracy of the particle filter. We note that the EnKF
and the stochastic map filter remain stable for the Lorenz-63 configurations
below _without_ this additive noise. We do not add any noise to the dynamics
of the Lorenz-96 model. Therefore, the Markov transition kernel for the
Lorenz-96 model is intractable (cf. Section 2).

#### Likelihood model.

The state of each system is observed indirectly every $\Delta t_{\rm obs}$
time units. Larger values of $\Delta t_{\rm obs}$ are associated with
increasingly non-Gaussian forecast statistics. We index the state at discrete
observation times by $(\bm{Z}_{k})_{k\geq 1}$ where $\bm{Z}_{k}=\bm{Z}(k\Delta
t_{\text{obs}})$ (e.g., see Figure 1). For any observation step $k$, the
likelihood model is specified by

$\bm{Y}_{k}=\bm{\mathrm{H}}\bm{Z}_{k}+\bm{\mathcal{E}}_{k},$ (11.3)

where $\bm{\mathrm{H}}\in\mathbb{R}^{d\times n}$ is a linear operator that
selects $d\leq n$ components of the state at uniform intervals (e.g., every
other component, or every fourth component), while $\bm{\mathcal{E}}_{k}$ is
an additive noise that is independent of $\bm{Z}_{k}$. In Sections 11.2 and
11.3, we consider configurations using Gaussian observational noise with zero
mean and covariance $\theta^{2}\mathbf{I}_{d}$. In Section 11.4, we employ
_heavy-tailed_ Laplace observational noise with zero mean and covariance
$2\theta^{2}\mathbf{I}_{d}$ (i.e., scale parameter $\theta$ for each
dimension).

#### Simulation setup.

Given a random initial condition $\bm{Z}_{0}$ drawn from
$\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$, we generate a sequence of “true”
hidden states $(\bm{z}^{*}_{k})_{k}$ and a sequence of synthetic observations
$(\bm{y}_{k}^{*})_{k}$ with the same model used for filtering. We are thus in
the usual “identical twin experiment” setting [57], wherein the filtering
algorithms must estimate the true state given only these synthetic
observations.

The initial ensemble for the filtering algorithms is generated through a
_spin-up_ phase. First, we draw $M$ i.i.d. random samples of the initial
condition, $\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$. Then, we run the
stochastic EnKF—without localization—for $2000$ assimilation steps. Using the
resulting ensemble at $k=2000$, we then apply the filter to be tested (e.g.,
an EnKF with optimally tuned localization, or a variant of the stochastic map
filter) for an additional $4000$ assimilation steps. At each step, we use the
ensemble mean
$\bar{\bm{z}}_{k}\coloneqq\frac{1}{M}{\sum_{i=1}^{M}}\bm{z}_{k}^{i}$ as a
point estimate of the true state $\bm{z}^{*}_{k}$. Lastly, to eliminate the
transient effect of switching between filtering algorithms, we analyze the
assimilation quality only using the last $2000$ assimilation cycles. We also
performed studies with up to $10^{4}$ assimilation cycles in this last phase,
but found no significant difference in the results.

#### Performance metrics.

We evaluate the tracking performance of the filter via the root-mean-squared
error (RMSE), defined at any assimilation step $k$ as
$\text{RMSE}_{k}\coloneqq\|\bar{\bm{z}}_{k}-\bm{z}^{*}_{k}\|_{2}/\sqrt{n}$. To
measure the concentration of the ensemble members, we compute the ensemble
spread, defined at any assimilation step $k$ as
$[\text{tr}(\widehat{\bm{\Sigma}}_{k})/n]^{1/2}$, where
$\widehat{\bm{\Sigma}}_{k}$ denotes the ensemble covariance matrix at step
$k$. In Sections 11.2–11.4, we report the _time average_ of each of these
metrics (e.g., over the 2000 test assimilation steps) as a function of the
ensemble size $M$. In addition, we compute the coverage probability of the
intervals given by the empirical $2.5\%$ and $97.5\%$ quantiles of each
marginal of the ensemble, i.e., the frequency with which the $i$th component
of $\bm{z}^{*}_{k}$ is contained in the $i$th marginal interval. For each
component, this frequency is of the form $k/T$ where $k\in\\{0,\dots,T\\}$
denotes the number of intervals that contain the corresponding component of
the true state, and $T$ represents the total number of assimilation times. We
then average these coverage probabilities over the $n$ components of the
state. Lastly, we also evaluate the quality of the ensemble by computing the
continuous ranked probability score (CRPS) [12, 25]. The CRPS for the $i$th
marginal component of the state at step $k$ compares the ensemble’s empirical
CDF $\widehat{F}_{k,i}(z)$ with a Heaviside function centered at the true
state, $H(z-(\bm{z}^{*}_{k})_{i})$, as given by
$\text{CRPS}_{k,i}=\int(\widehat{F}_{k,i}(z)-\mathbbm{1}((\bm{z}^{*}_{k})_{i}\leq
z))^{2}dz$. We average $\text{CRPS}_{k,i}$ over state components $i$ and
assimilation steps $k$.

#### Algorithm parameters.

The results below investigate the performance of the stochastic map filter
(Algorithm 2) as nonlinearities are gradually introduced in the prior-to-
posterior transformation $\widehat{T}$.

Each component function $U^{k}$ of the map $S^{\bm{\mathcal{X}}}$ is allowed
to depend on the data $y$ and on a subset of the state variables
$z_{1},\dots,z_{k}$ given by $\bm{z}_{\mathcal{A}_{k,r}}$. In particular, our
numerical experiments consider a separable representation
$U^{k}(y,\bm{z}_{\mathcal{A}_{k,r}})=\bm{\mathfrak{u}}^{k}_{0}(y)+\sum_{i\in\mathcal{A}_{k,r}}\bm{\mathfrak{u}}^{k}_{i}(z_{i})$
for $k=1,\dots,n$. For the linear version of the stochastic map filter,
labeled as “linear” in the figures, $\bm{\mathfrak{u}}^{k}_{0}$ and
$\bm{\mathfrak{u}}^{k}_{i}$ are restricted to be linear functions. To increase
the complexity of the map, we add radial basis functions (RBFs) to the
approximation spaces for $\bm{\mathfrak{u}}^{1}_{0}$,
$\bm{\mathfrak{u}}^{k}_{0}$, and $\bm{\mathfrak{u}}^{k}_{i}$ for $i<k$. For
$\bm{\mathfrak{u}}^{1}_{1}$, we use a monotone parameterization based on
integrals of radial basis functions (sigmoids), and we keep
$\bm{\mathfrak{u}}_{k}^{k}$ linear for $k>1$. See Appendix A for more details
on this representation. While many other nonlinear representations of $U^{k}$
are possible, this particular separable form was chosen because it gradually
adds degrees of freedom starting from a linear representation and allows (7.3)
to be solved in closed form for $k>1$. We label results for maps with linear
terms plus $p$ RBFs/sigmoids in each component’s functions as “linear + $p$
RBFs.” The figures below thus compare the filtering performance of four
ensemble methods: the stochastic EnKF, linear maps, and linear maps with
$p\in\\{1,2\\}$ RBFs.

For a fair comparison, all algorithms process the observations at a given
assimilation time sequentially (see Section 9) and do not evaluate the
likelihood function, which we assume is intractable. The stochastic EnKF
computes the Kalman gain and the linear map in (9.4) by estimating the cross-
covariance $\smash{\widehat{\bm{\Sigma}}_{\bm{X},\bm{Y}}}$ and observation
variance $\smash{\widehat{\bm{\Sigma}}_{\bm{Y}}}$ using data simulated from
the likelihood model for each forecast sample. Since we do not require direct
access to the observation operator $\bm{\mathrm{H}}$ or to the covariance of
$\bm{\mathcal{E}}_{k}$, these algorithms are generalizable to non-additive and
non-Gaussian likelihood models.

As described in Section 9, the linear map algorithm is related to the
stochastic EnKF. The only difference between these two algorithms lies in the
way they exploit structure for the purpose of regularization. While stochastic
EnKFs exploit decay of correlation between the variables with increasing
spatial distance, regularization in the stochastic map filter enforces
sparsity in the map $S^{\bm{\mathcal{X}}}$, based on conditional independence,
decay of correlation, and local likelihood structure (see Sections 8 and 9).
In the stochastic map filter, the sparsity structure is parameterized by the
distance $r$ that defines the set of input state variables for each component
in (8.3), and by the number of components $j\leq n$ where
$S^{\bm{\mathcal{X}}}$ departs from the identity map (see Remark 3). In the
stochastic EnKF, covariance localization is applied by multiplying the
forecast covariance elementwise with a Gaspari–Cohn tapering function that has
an optimally tuned localization radius [30]. While the linear maps and the
stochastic EnKF behave similarly with larger ensemble sizes—where less
regularization is needed—the results may differ for small $M$.

All algorithms considered in this section are implemented with multiplicative
prior inflation and localization. For the inflation, we increased the spread
of the forecast ensemble that is used to estimate $S^{\bm{\mathcal{X}}}$. We
inflated these samples by multiplying their deviations from the sample mean by
a scalar $\zeta\geq 1$. The parameters (i.e., inflation parameter $\zeta$ and
localization parameters $r$ and $j$) are tuned for each ensemble size and
problem configuration to minimize the time-averaged RMSE. Only the tuned
results are presented below. Lastly, for the $p\in\\{1,2\\}$ stochastic map
filter we also investigated the sensitivity of the results with respect to the
scaling parameter $\gamma$ of the RBFs (described in Appendix A) and found no
significant dependence. Thus, the results are presented below for a fixed
value of $\gamma=2.0$ across all of the experiments.

###  11.2 Lorenz-63

We first consider the Lorenz-63 system without any form of localization
applied to $S^{\bm{\mathcal{X}}}$. The idea is to isolate the effect of adding
nonlinear terms to the prior-to-posterior update. In this case, the states of
the system are fully observed (i.e., $n=d$) with Gaussian observational noise
and variance $\theta^{2}=4$. The time between observations is set to $\Delta
t_{\rm obs}=0.1$. The resulting average RMSE, for a range of $M$, is plotted
on the left of Figure 7.

For very small $M$, linear maps have the most robust performance. However, for
sufficiently large ensemble sizes, we observe a consistent improvement in RMSE
by adding nonlinearities to the approximation space of the map. Specifically,
the stochastic map filter with $p=1$ RBFs yields better tracking performance
for $M\geq 40$, while the filter with $p=2$ RBFs provides improved results for
$M\geq 200$. In both cases, the improvement in average RMSE is more than
$20\%$. The average RMSE values obtained with increasing map nonlinearity (for
large enough $M$) also tend towards the RMSE of our “reference”
solution—obtained with a consistent SIR particle filter using $M=10^{6}$
samples (with an average relative effective sample size of 0.89), which we
plot in Figure 7 as a dotted line.

While this experiment did not involve map localization or covariance tapering
in the EnKF, the difference between the stochastic EnKF and the linear map
results for small $M$ can be attributed to the local likelihood structure that
we exploit. When assimilating scalar observations $Y=y^{*}$ with a local
likelihood model (i.e., $Y$ only depends on one component of the state, say
$Z_{1}$ after a permutation), the states $Z_{k}$ for $k=2,\dots,n$ are
conditionally independent of $Y$ given $Z_{1}$. This Markov property results
in sparsity of the KR rearrangement from $\pi_{\bm{X},Y}$ to $\mathcal{N}({\bf
0},{\bf I}_{n+1})$, and hence the map $S^{\bm{\mathcal{X}}}$ (see Section 8).
In particular, the components $U^{k}$ do not depend on $y$ for $k\geq 2$. We
note that this sparsity is exact, and thus different from any approximate
conditional independence that could further be exploited to regularize the
estimation of $S^{\bm{\mathcal{X}}}$.

In a second experiment, we investigate the performance of the stochastic map
filter as a function of the time $\Delta t_{\rm obs}$ between observations,
for a fixed ensemble size. The idea is to test the filter for a sequence of
increasingly non-Gaussian forecast distributions. The average RMSE of the
various algorithms for $\Delta t_{\rm obs}\in[0.1,0.5]$ and for $M=1000$ is
shown in the right panel of Figure 7. The stochastic map filters demonstrate a
consistently lower RMSE over the full range of inter-observation times. This
shows that map filters of increasing complexity remain stable (and continue to
provide more accurate tracking) even as the state evolution becomes more
nonlinear.

Figure 7: Average RMSE of the Lorenz-63 model for $\Delta t_{\rm obs}=0.1$ as
a function of $M$ (left), and for $M=1000$ as a function of the inter-
observation time $\Delta t_{\rm obs}$ (right).

For this low-dimensional problem, we also investigate how the stochastic map
filter captures the true filtering distribution (i.e., the sequence of
posteriors $\pi_{\bm{X}_{k}|\bm{y}^{*}_{1:k}}$) over time. We first use an SIR
particle filter with $10^{6}$ samples to estimate the posterior mean
$\bm{z}_{k}^{\text{PF}}$ and posterior covariance matrix
$\bm{\Sigma}_{k}^{\text{PF}}$ at each step $k$, with $\Delta t_{\rm obs}=0.1$.
In Figure 8, we plot, for each algorithm, the normalized $L_{2}$ error in the
posterior mean, defined as
$\|\bar{\bm{z}}_{k}-\bm{z}_{k}^{\text{PF}}\|_{2}/\sqrt{n}$, and the normalized
error in the posterior covariance, defined as
$\|\widehat{\bm{\Sigma}}_{k}-\bm{\Sigma}_{k}^{\text{PF}}\|_{F}/n$, averaged
over 2000 assimilation times $k$. We observe a significant improvement in the
approximation of these posterior statistics, just by introducing a few
nonlinearities in the prior-to-posterior transformation. For a sufficiently
large $M$, the improvement can be more than $50\%$.

To characterize the variability of the particle filter estimates used above,
we repeat the assimilation exercise with the reference $M=10^{6}$ SIR filter
for 100 times. We compute the sample standard deviations of these replicated
filtering results (e.g., of the $L_{2}$ norm of the posterior mean and the
Frobenius norm of the posterior covariance matrix at each assimilation time)
and plot their maxima over time using dotted lines in Figure 8. For
comparison, the standard errors of the same norms of posterior moments
estimated with the stochastic map filter, for $M\in[20,600]$, are 4 to 10
times smaller than the standard errors of the $M=10^{6}$ SIR filter.

Finally, to assess the predictive performance of our ensembles, we compute the
CRPS and report the average, over all three state components and $2000$
assimilation cycles, for increasing $M$ in Figure 11 (left). For $M\geq 60$
samples, the nonlinear stochastic map filter yields lower values of CRPS,
which provides a simultaneous diagnostic of calibration and sharpness as
defined in [25]. With increasing nonlinearity in the prior-to-posterior
transformation, the average CRPS for the Lorenz-63 model approaches that of a
particle filter using $10^{6}$ samples.

Figure 8: Normalized $L_{2}$ error of the posterior mean (left) and Frobenius
error of the posterior covariance matrix (right) for the Lorenz-63 model.

###  11.3 Lorenz-96 with long inter-observation time

Here, we follow the “hard case” setup of [37, 5] for the Lorenz-96 model, with
sparse observations in space and time, i.e., $d=20$ (observing every other
component of the state) and $\Delta t_{\rm obs}=0.4$. Each observation has
independent additive Gaussian noise with variance $\theta^{2}=0.5$. For
reference, we note that a $\Delta t_{\rm obs}=0.05$ would correspond to
roughly $6$ hours in a global weather model [41, 42]. The large time interval
between observations makes the forecast distribution highly non-Gaussian. As a
result, this configuration typically requires ensemble sizes greater than the
state dimension $n$ to track $\bm{z}_{k}^{*}$. Furthermore, with a high-
dimensional state vector it is necessary to tune the map’s localization
parameters to achieve stable filtering performance over time. As described in
Section 11.1, we tune the localization parameters $r$ and $j$ for the
stochastic map filter and the localization radius for the EnKF, and the
inflation parameter $\zeta$ for both algorithms. We find that the filters
require essentially no forecast inflation for these identical twin
experiments.

Figure 9 shows the average and median values of $\text{RMSE}_{k}$ (over 2000
terminal assimilation cycles $k$) as a function of the ensemble size $M$. By
increasing the nonlinearity of the map $\widehat{S}^{\bm{\mathcal{X}}}$ and
hence $\widehat{T}$, the stochastic map filter can reduce the bias of the EnKF
given sufficient samples. For example, the plateaus of RMSE in Figure 9 for
the EnKF reflect the limitation of the EnKF’s affine prior-to-posterior map:
the RMSE initially decreases, then stagnates with increasing $M$. Nonlinear
maps reduce the values of these plateaus (by roughly 25% in this example) and
provide smaller RMSE even at intermediate ensemble sizes. By introducing
relatively few nonlinear basis functions (in $\widehat{S}^{\bm{\mathcal{X}}}$)
whose parameters can be _reliably_ learned, the stochastic map filter can
better capture the non-Gaussian forecast statistics and improve the tracking
of $\bm{z}_{k}^{*}$.

Figure 9 also exposes the bias-variance tradeoff in learning the prior-to-
posterior map ${T}$, which corresponds to a link between $M$ and the
complexity of the map (i.e., the number of the basis functions in each
component of $S^{\bm{\mathcal{X}}}$). For smaller ensemble sizes, a richer
class of nonlinear maps yields estimators $\widehat{T}$ with less bias but
higher variance. The EnKF and linear maps, on the other hand, yield more
biased but lower-variance estimators that offer more stable tracking
performance for small $M$. In the present experiments, as $M$ increases, the
map class yielding the minimum RMSE evolves quickly from linear to $p=1$ and
then to $p=2$. Overall, the stochastic map filter provides a flexible
framework for adjusting the complexity of the map to the given ensemble size
in order to extract a good estimator of $T$.

Figure 9: Average (left) and median (right) RMSE (over $2000$ assimilation
cycles) for the “hard” Lorenz-96 configuration of Section 11.3, with $\Delta
t_{obs}=0.4$, observing every other component of the state ($d=20$) with
additive Gaussian observational noise. Dashed line is the standard deviation
of the observational noise.

Figure 10: Average ensemble spread (left) and average coverage probability of
the $[2.5\%,97.5\%]$ empirical intervals of each state marginal (right) for
the “hard” Lorenz-96 configuration. (See precise definitions of these
quantities in Section 11.1.)

In addition to improving the RMSE, the nonlinear stochastic map filter reduces
the ensemble spread while simultaneously increasing the probability of
covering the true state. We present these results for increasing $M$ in Figure
10. To further assess the predictive performance of the ensemble, we compute
the average CRPS (over the $n=40$ state components and $2000$ assimilation
cycles) for this Lorenz-96 configuration and report the results in Figure 11
(right). With $M\geq 100$ samples, the nonlinear stochastic map filter yields
a lower CRPS than its linear counterparts. Adding more nonlinearity to the
transformation allows the score to improve even further with increasing $M$.

Figure 11: Average continuous rank probability score (CRPS) of the analysis
ensemble for our Lorenz-63 experiments (left) and the “hard” Lorenz-96
configuration (right).

###  11.4 Lorenz-96 with heavy-tailed observational noise

A second challenging configuration of the Lorenz-96 model is obtained with
sparse and heavy-tailed observations. In this setting, we consider $\Delta
t_{\rm obs}=0.1$ and $d=10$ (observing every fourth component of the state),
and additive independent Laplace-distributed (i.e., double-exponential)
observational noise. The noise has variance $2\theta^{2}$, where $\theta=1$ is
the scale parameter of the Laplace distribution. The large variance of the
noise and limited number of observations make filtering difficult; for
instance, large RMSE values are obtained when using the EnKF with ensemble
sizes smaller than $M=100$.

Similar to the results in Section 11.3, we observe improved tracking
performance with increasing nonlinearity in the map $S^{\bm{\mathcal{X}}}$,
provided $M$ is not too small. Figure 12 shows the average and median RMSE for
the EnKF and different stochastic map parameterizations. As $M$ increases, the
bias-variance tradeoff shifts in favor of more complex maps. Furthermore, as
illustrated in Figure 13, higher order maps also increase the average coverage
probability of the ensemble across $2000$ assimilation steps without
increasing the ensemble spread. These results suggest that, by increasing the
complexity of the prior-to-posterior transformation, ensemble members from the
stochastic map filter better represent the true state $\bm{z}_{k}^{*}$ and its
uncertainty for problems with non-Gaussian forecast and analysis
distributions.

Figure 12: Average (left) and median (right) RMSE (over $2000$ assimilation
cycles) for the heavy-tailed (Laplace observational noise) Lorenz-96
configuration of Section 11.4, with $\Delta t_{obs}=0.1$ and $d=10$
observations. Dashed line is the standard deviation of the observational
noise.

Figure 13: Average ensemble spread (left) and average coverage probability of
the $[2.5\%,97.5\%]$ empirical intervals of each state marginal (right) for
the heavy-tailed Lorenz-96 configuration.

##  12 Discussion

We have introduced a class of nonlinear filtering algorithms based on the
construction of stochastic or deterministic couplings between continuous
distributions. These algorithms generalize the EnKF to nonlinear updates of
tunable and arbitrary complexity. The key computational task in applying these
algorithms is the numerical estimation of a KR rearrangement via convex
optimization. Our goal in using nonlinear updates is to reduce the intrinsic
bias of the EnKF, while retaining its robustness in high dimensions. We run
our numerical experiments on various configurations of the Lorenz-96 model,
which is a common testbed for numerical weather prediction systems; we also
include a comparison with the true (Bayesian) filtering distributions for the
lower-dimensional Lorenz-63 model. Through these experiments, we show how
filtering with nonlinear updates can dramatically outperform the EnKF in a
range of highly non-Gaussian settings, for a marginal increase in ensemble
requirements and computational effort. In what follows, we briefly discuss
limitations and possible extensions of this approach to filtering.

#### Nonlinear parameterizations.

Highly nonlinear updates do not automatically guarantee better filtering
performance. Trying to estimate complex transformations from very few samples
could yield unacceptable variance. The bias and variance of a transport map
estimator must be balanced, possibly by introducing additional information
through regularization. In this paper, we consider several regularization
techniques. For instance, we define a hierarchy of parameterizations for the
transport maps with an increasing level of nonlinearity and select the “best”
parameterization according to a prediction error criterion. We notice several
intuitive trends. Increasingly nonlinear transformations typically demand
larger ensembles, while for a fixed parameterization of the maps, the tracking
error plateaus beyond a certain ensemble size. Similarly, for any given
problem, we expect there to be a sufficiently small ensemble size for which
the linear ansatz—as in EnKF techniques—performs best. Moreover, for nonlinear
filtering problems that are well approximated by Gaussian distributions, we do
not expect major gains from using nonlinear updates, even for large ensemble
sizes. Nonlinear updates are mostly needed to capture non-Gaussian structure.
The complexity of the transformations should then be adapted to the number of
ensemble members and to the level of non-Gaussianity of the problem. In future
work, we wish to automate this adaptation rather than running an offline
calibration phase, i.e., we plan to develop nonparametric or semi-parametric
extensions of the proposed filtering algorithms.

#### Localization and other form of regularization.

A related regularization idea is that of exploiting decay of correlation and
approximate conditional independence in the filtering distribution. This
structure is common when filtering spatiotemporal processes. In this paper, we
introduced simple localization ideas for nonlinear updates by imposing
specific sparsity patterns in the transport maps composing the update. In
future work, one could explore a much wider range of regularization
techniques. For instance, it is natural to consider LASSO-type estimators in
order to “discover”—rather than to enforce—sparsity of the nonlinear
transformations. Or one might exploit low-rank structure in the change from
forecast to filtering distribution in order to reduce the effective dimension
of the nonlinear update, following the line of work in [65, 16]. Initial
accounts of low-rank structure in transport maps can be found in [63, Ch. 5]
and [9]. In general, we believe that practical high-dimensional filtering
relies crucially on the development of structure-exploiting regularization
techniques.

#### Consistency.

From a theoretical point of view, it would be interesting to study the
behavior of the proposed algorithms in a simultaneous limit of increasing
ensemble size and increasing complexity of the rearrangements. In particular,
a key goal should be to understand conditions under which it is possible to
establish some form of _consistency_ for the filters, much in the same spirit
as the convergence theory for SMC methods [15]. This question is meaningless
for the EnKF, which is inherently inconsistent due to the fixed linear
structure of the update.

It would also be interesting to understand the stochastic map filter in a
continuous-time filtering setting, and thus to explore its connections with
the feedback particle filter [66]. This link might offer further insight into
conditions that ensure consistency.

#### Further extensions.

The ideas presented in this paper can readily be extended to smoothing, i.e.,
characterizing the conditional distribution of past states, given all
available measurements. In this context, one can first estimate a (non-
singular) Markov transition kernel between observation times via an upper-
triangular KR rearrangement, and then use this approximation to propagate the
filtering ensemble backwards in time. The resulting scheme can be interpreted
as a Monte Carlo approximation of the transport map backward-smoothing
algorithm introduced in [64, Section 7]. Relatedly, it would be useful to
extend the present framework to problems of sequential _joint parameter and
state_ estimation.

Finally, we note that the analysis step of the stochastic map filter is an
instance of _approximate Bayesian computation_ , as it requires only sampling
from the likelihood (and prior) in order to construct a prior-to-posterior
transformation. Realizing and expanding this approach for generic problems of
inference in generative models represents an exciting avenue for future work.

## Acknowledgments

AS, RSB, and YMM gratefully acknowledge support from the AFOSR Computational
Mathematics program and AFOSR MURI award FA9550-15-1-0038. AS and YM are also
grateful for support from the US Department of Energy, Office of Advanced
Scientific Computing Research, AEOLUS project. RSB also acknowledges support
from an NSERC PGSD-D Fellowship. AS acknowledges support from the Deutsche
Forschungsgemeinschaft (DFG), through CRC 1294 “Data Assimilation” Project
(A02), for a visit to the Institute for Mathematics of the University of
Potsdam. The authors also thank Daniele Bigoni, Jana de Wiljes, Matthias
Morzfeld, Sebastian Reich, and Xin Tong for many insightful discussions and
suggestions.

##  Appendix A Parameterization and computation of monotone triangular maps

In this section, we first discuss possible parameterizations for a monotone
triangular map $U$ that is meant to approximate the KR rearrangement between a
pair of densities on $\mathbb{R}^{n}$ (Section 7). Then, we link specific
structures in the parameterization of $U$ to computational simplifications for
an estimator of the rearrangement.

The KR rearrangement between a pair of positive densities is a triangular
bijection on $\mathbb{R}^{n}$ such that each slice (5.2) is strictly
increasing and absolutely continuous [10]. Hence, it makes sense to consider
candidate rearrangements $U$ whose $k$th components, $U^{k}$, can be written
as

$U^{k}(z_{1},\ldots,z_{k})=a_{k}(z_{1},\ldots,z_{k-1})+\int_{0}^{z_{k}}\exp\left(b_{k}(z_{1},\ldots,z_{k-1},t)\right)\,{\rm
d}t,$ (A.1)

for some arbitrary functions $(a_{k},b_{k})$, and for all $k=1,\ldots,n$ [55,
45]. For any choice of $(a_{k},b_{k})$, the rearrangement $U$ is a strictly
monotone (increasing) triangular map, since
$\partial_{k}\,U^{k}=\exp(b_{k})>0$ for all $k$. In fact, (A.1) provides a
rather general way to represent monotone triangular functions.

The estimator in (7.3) relies on the definition of an approximation space,
$\bm{\mathcal{H}}_{k}$, for every component of the rearrangement. A viable
strategy is to use the monotone representation in (A.1) and to constrain each
function $(a_{k},b_{k})$ to lie in a finite dimensional space, i.e.,

$a_{k}(\bm{z})=\sum_{i}a_{i}\,\psi_{i}^{a}(\bm{z}),\qquad
b_{k}(\bm{z})=\sum_{j}b_{j}\,\psi_{j}^{b}(\bm{z}),$ (A.2)

for a collection of finitely many basis functions
$\psi_{i}^{a},\psi_{j}^{b}$—e.g., multivariate Hermite polynomials, Hermite
functions, splines, or radial basis functions—and unknown real valued
coefficients $(a_{i},b_{j})$.999We may substitute the exponential in (A.1)
with a square function or even a sum of squares [31]. This offers the
computational advantage of precomputing any integral in (A.1) once and for
all, thanks to parameterizations of $b_{k}$ that is linear in the unknown
coefficients (see [8]).  The minimization in (7.3) can then be cast in terms
of the unknown coefficients, which fully parameterize the maps in
$\bm{\mathcal{H}}_{k}$.

As explained in Section 8, the “complexity” of each $\bm{\mathcal{H}}_{k}$ is
a key ingredient for regularizing the estimation of the rearrangement; there
is no point in trying to learn arbitrarily nonlinear maps from a given finite
ensemble. For instance, the EnKF restricts the search to linear maps. In
general, it is possible to depart from the linear ansatz in a gradual way, by
imposing structural assumptions of increasing complexity on the rearrangement.
For example, we can consider maps whose components are separable as sums of
_univariate_ nonlinear functions $(\bm{\mathfrak{u}}_{i})$ that each admit a
linear parameterization like (A.2) or (A.1) (if monotone), i.e.,

$U^{k}(z_{1},\ldots,z_{k})=\bm{\mathfrak{u}}_{1}(z_{1})+\ldots+\bm{\mathfrak{u}}_{k}(z_{k}).$
(A.3)

Next, we can consider separable parameterizations in terms of bivariate
functions, i.e., $U^{k}(z_{1},\ldots,z_{k})=\sum_{i\leq
k}\bm{\mathfrak{u}}_{i}(z_{i},z_{k})$, and so on. In essence, the complexity
of the parameterization should be adapted (and calibrated) to the cardinality
of the ensemble and to the amount of prior information that is injected in the
estimation (e.g., see the localization ideas of Section 7).

In the numerical examples of this paper, we explore the separable
parameterization in (A.3), since we wish to work with ensemble sizes
comparable to those needed by the EnKF. We propose a specific strategy to
parameterize each $\bm{\mathfrak{u}}_{i}$ in (A.3), although there are plenty
of potential alternatives. For $i<k$, $\bm{\mathfrak{u}}_{i}$ is simply a
univariate function—not necessarily monotone. In this case, we parameterize
$\bm{\mathfrak{u}}_{i}$ as a linear combination of $p$ radial basis functions
(RBFs) plus a global linear term, i.e.,

$\bm{\mathfrak{u}}_{i}(z)=u_{i0}\,z+\sum_{j=1}^{p}u_{ij}\,\mathcal{N}(z;\xi_{j},\sigma_{j}^{2}),$
(A.4)

for a collection of centers $\bm{\xi}\coloneqq(\xi_{1},\ldots,\xi_{p})$, scale
parameters $\bm{\sigma}\coloneqq(\sigma_{1},\ldots,\sigma_{p})$, and unknown
coefficients $(u_{ij})$. The number $p$ of RBFs defines the degree of
nonlinearity of $\bm{\mathfrak{u}}_{i}$. For $p=0$, $\bm{\mathfrak{u}}_{i}$
reverts to a linear function. The centers $\bm{\xi}$ are chosen from $p$
(uniform) empirical quantiles of the marginal distribution of the ensemble
along the $i$th coordinate direction, i.e.,
$\xi_{j}\coloneqq\hat{q}_{j/(p+1)}$ for $j=1,\ldots,p$, where for any
$\alpha\in(0,1)$, $\hat{q}_{\alpha}$ denotes the $\alpha$-th empirical
quantile. The scale parameters $\bm{\sigma}$ are defined to be proportional to
the inter-distance between the centers $\bm{\xi}$, i.e.,
$\sigma_{j}\coloneqq\gamma\,(\xi_{j+1}-\xi_{j-1})/2$ for $j=1,\ldots,p$, where
$\xi_{0}\coloneqq\xi_{1}$, $\xi_{p+1}\coloneqq\xi_{p}$, and where $\gamma>0$
is a tuning parameter. In our experiments, we consider $\gamma\in[1,3]$. The
parameterization of $\bm{\mathfrak{u}}_{k}$ in (A.3) is a bit different, since
the monotonicity constraint on $U^{k}$ requires $\bm{\mathfrak{u}}_{k}$ to be
an increasing map. One option is to use the monotone representation in (A.1)
just for $\bm{\mathfrak{u}}_{k}$. In practice, however,
$\bm{\mathfrak{u}}_{k}$ is a univariate function and thus we can seek
alternative ways of imposing monotonicity, e.g., by expressing
$\bm{\mathfrak{u}}_{k}$ as a nonnegative combination of monotone basis
functions. For a monotone $\bm{\mathfrak{u}}_{k}$, its derivative
$\partial\bm{\mathfrak{u}}_{k}$ is always nonnegative. For $p>0$, we propose
to parameterize $\partial\bm{\mathfrak{u}}_{k}$ with $p+2$ positive basis
functions, i.e.,
$\partial\bm{\mathfrak{u}}_{k}=\smash{\sum_{j=0}^{p+1}}\,u_{kj}\,\psi_{j}^{\prime}$,
where $\psi_{j}^{\prime}(z)\coloneqq\mathcal{N}(z;\xi_{j},\sigma_{j}^{2})$ for
$j=1,\ldots,p$, while

$\psi_{0}^{\prime}(z)\coloneqq\frac{1}{2}-\frac{1}{2}\operatorname{erf}\left(\frac{z-\xi_{0}}{\sqrt{2}\,\sigma_{0}}\right),\quad{\rm
and}\quad\psi_{p+1}^{\prime}(z)\coloneqq\frac{1}{2}+\frac{1}{2}\operatorname{erf}\left(\frac{z-\xi_{p+1}}{\sqrt{2}\,\sigma_{p+1}}\right),$
(A.5)

for a collection of $p+2$ centers $(\xi_{i})$ and scale parameters
$(\sigma_{i})$. The centers and scale parameters are chosen similarly to those
in (A.4), i.e., from uniform empirical quantiles of the marginal distribution
of the ensemble along the $k$th direction. Intuitively,
$\partial\bm{\mathfrak{u}}_{k}$ is the linear combination of $p$ RBFs
$(\psi^{\prime}_{1},\ldots,\psi^{\prime}_{p})$ plus a pair of positive
functions $(\psi^{\prime}_{0},\psi^{\prime}_{p+1})$ that revert to constants
in the tails. The idea is to promote robustness by constraining
$\bm{\mathfrak{u}}_{k}$ to be linear in the tails. Non-negativity of
$\partial\bm{\mathfrak{u}}_{k}$ is imposed by the set of linear inequalities
$u_{kj}\geq 0$, for $j=0,\ldots,p+1$. The resulting monotone parameterization
of $\bm{\mathfrak{u}}_{k}$ is obtained by integrating
$\partial\bm{\mathfrak{u}}_{k}$, i.e.,
$\bm{\mathfrak{u}}_{k}(z)\coloneqq\int_{0}^{z}\partial\bm{\mathfrak{u}}_{k}(x)\,{\rm
d}x=c+\smash{\sum_{j=0}^{p+1}}\,u_{kj}\,\psi_{j}(z)$, for an unknown constant
$c$ and antiderivatives $\psi_{j}\coloneqq\int\psi_{j}^{\prime}$ given by

$\left\\{\begin{array}[]{lr}\psi_{0}(z)=\frac{1}{2}\left(\left(z-\xi_{0}\right)\left(1-\operatorname{erf}\left(\Delta_{0}\right)\right)-\sigma_{0}\,\sqrt{2/\pi}\,\exp(-\Delta_{0}^{2})\right)\\\\[10.0pt]
\psi_{j}(z)=\frac{1}{2}\left(1+\operatorname{erf}\left(\Delta_{j}\right)\right)&\text{for
}1\leq j\leq p\\\\[10.0pt]
\psi_{p+1}=\frac{1}{2}\left(\left(z-\xi_{p+1}\right)\left(1+\operatorname{erf}\left(\Delta_{p+1}\right)\right)+\sigma_{p+1}\,\sqrt{2/\pi}\,\exp(-\Delta_{p+1}^{2})\right)\par\end{array}\right.\,,$
(A.6)

where $\Delta_{j}\coloneqq(z-\xi_{j})/(\sqrt{2}\,\sigma_{j})$ for all $j$. For
$p=0$, we set $\bm{\mathfrak{u}}_{k}$ to be an affine function. For the
parameterization of $U^{k}$ given by (A.4) and (A.6), the minimization in
(7.3) is a linearly constrained convex program that can be solved efficiently
via numerical optimization, e.g., using the projected Newton’s method [6].
Moreover, if $\bm{\mathfrak{u}}_{k}$ is parameterized by an affine function,
then (7.3) reduces to a standard linear regression problem with a closed form
solution. We analyze this fact in the more general case of parameterizations
that are separable only in the last (monotone) variable, i.e.,

$U^{k}(z_{1},\ldots,z_{k})=\bm{\mathfrak{u}}(z_{1},\ldots,z_{k-1})+\bm{\mathfrak{u}}_{k}(z_{k}),$
(A.7)

where $\bm{\mathfrak{u}}(\bm{z})=\sum_{j}u_{j}\,\psi_{j}(\bm{z})$ for a
collection of arbitrary basis functions $(\psi_{j})$ and where
$\bm{\mathfrak{u}}_{k}(z)=c+\alpha z$. In this case, we can rewrite (7.3) as

$\min_{(\tilde{u}_{j}),\tilde{c},\alpha}\,\frac{\alpha^{2}}{2}\,\left\\{\frac{1}{M}\sum_{i=1}^{M}\,\left(\sum_{j}\tilde{u}_{j}\,\psi_{j}(z^{i}_{1},\ldots,z^{i}_{k-1})+\tilde{c}+z^{i}_{k}\right)^{2}\right\\}-\log\alpha,$
(A.8)

where $\bm{z}^{i}=(z^{i}_{1},\ldots,z^{i}_{n})$, $\tilde{c}\coloneqq
c/\alpha$, and $\tilde{u}_{j}\coloneqq u_{j}/\alpha$ for all $j$. The
minimization of (A.8) with respect to the new variables
$(\tilde{c},\tilde{u}_{j})$ can be done independently of $\alpha$, and
corresponds to a standard regression problem, i.e.,

$\kappa^{*}\coloneqq\min_{(\tilde{u}_{j}),\tilde{c}}\,\frac{1}{M}\sum_{i=1}^{M}\,\left(\sum_{j}\tilde{u}_{j}\,\psi_{j}(z^{i}_{1},\ldots,z^{i}_{k-1})+\tilde{c}+z^{i}_{k}\right)^{2},$
(A.9)

while the optimal $\alpha$ is given by $1/\sqrt{\kappa^{*}}$.

##  Appendix B Deterministic map filter with local observations

In many problems of interest, we have conditionally independent and local
observations, i.e.,

$\pi_{\bm{Y}|\bm{X}}=\prod_{k=1}^{d}\pi_{Y_{k}|X_{\ell_{k}}},$ (B.1)

for a subset of $d$ state variables $(X_{\ell_{k}})$—possibly up to a
whitening transformation of the state [30, 46]. In Section 9, we saw that
(B.1) implies that we can process each of the $d$ scalar observations
individually and sequentially. For the deterministic map filter, there is a
strong incentive in processing scalar observations one at a time, because we
can avoid the explicit solution of a non-convex problem like (10.4).

Without loss of generality, let us assimilate a scalar observation, $y^{*}$,
of the first state variable, possibly up to a permutation. The corresponding
likelihood function is given by $\pi_{Y|X_{1}}$. We want to define a
structure-exploiting analysis map $T$, similar to (10.1). Let $S$ be a KR that
pushes forward $\pi_{\bm{X}}$ to a standard normal $\eta=\mathcal{N}({\bf
0},{\bf I}_{n})$—as in (10.1)—and let $\eta_{1}=\mathcal{N}(0,1)$ denote the
marginal of $\eta$ along the first variable. By the definition of KR
rearrangement, the pullback of $\eta_{1}$ by the first component, $S^{1}$, of
$S$ is equivalent to the prior marginal along the observed component, i.e.,
$(S^{1})^{\sharp}\,\eta_{1}=\pi_{X_{1}}$. Moreover, for all univariate
monotone maps $\mathfrak{u}$, define $\pi_{\mathfrak{u}}$ to be the density
that is proportional to the function
$\xi\mapsto\pi_{Y|X_{1}}(y^{*}|\xi)\,\mathfrak{u}^{\sharp}\,\eta_{1}(\xi)$,
which is the product of likelihood and pullback
$\smash{\mathfrak{u}^{\sharp}\,\eta_{1}}$. Note that $\pi_{S^{1}}$ is
equivalent to the posterior marginal $\pi_{X_{1}|y^{*}}$. Let $\mathfrak{m}$
be the increasing rearrangement on $\mathbb{R}$—really a one dimensional KR
rearrangement—that pushes forward $\eta_{1}$ to $\pi_{S^{1}}$, and notice that
the posterior distribution factorizes as

$\pi_{\bm{X}|y^{*}}=\pi_{X_{1}|y^{*}}\,\pi_{\bm{X}_{2:n}|X_{1}},$ (B.2)

because of the local structure of the likelihood function. The map
$\mathfrak{m}$ characterizes the posterior marginal $\pi_{X_{1}|y^{*}}$, while
$S$ gives access to the prior conditional $\pi_{\bm{X}_{2:n}|X_{1}}$. To see
this, notice that for all $z\in\mathbb{R}$, the map

$z_{2},\ldots,z_{n}\mapsto\left[\begin{array}[]{l}S^{2}(z,z_{2})\\\ \vdots\\\
S^{n}(z,z_{2},\dots z_{n})\end{array}\right]$ (B.3)

pushes forward the conditional $\pi_{\bm{X}_{2:n}|X_{1}}(\cdot|z)$ to a
standard normal on $\mathbb{R}^{n-1}$ (cf. Section 5). Hence, we can define
the analysis map $T$ as follows,

$T\coloneqq S_{\bf Id}^{-1}\circ\mathcal{M}\circ S,$ (B.4)

where $S_{\bf Id}$ coincides with $S$ except for the first component, which is
set to the identity function,

$S_{\bf Id}(\bm{z})\coloneqq\left[\begin{array}[]{l}z_{1}\\\\[2.0pt]
S^{2}(z_{1},z_{2})\\\ \vdots\\\ S^{n}(z_{1},z_{2},\dots
z_{n}),\end{array}\right],\quad{\rm
and}\quad\mathcal{M}(\bm{z})\coloneqq\left[\begin{array}[]{l}\mathfrak{m}(z_{1})\\\\[2.0pt]
z_{2}\\\ \vdots\\\ z_{n}\end{array}\right].$ (B.5)

It is immediate to verify that $T$ pushes forward $\pi_{\bm{X}}$ to
$\pi_{\bm{X}|y^{*}}$. In particular, the map $\mathcal{M}$ is responsible for
incorporating information from the measurement $y^{*}$ only in the observed
component of the state—this is a local assimilation step—while the map $S_{\bf
Id}^{-1}$ is responsible for propagating information to the remaining
components. In practice, we just showed that, when assimilating a local scalar
observation, the map $\mathcal{T}$ in (10.1) can be rewritten as
$\mathcal{T}=\smash{S_{\bf Id}^{-1}}\circ\mathcal{M}$.

We propose to estimate $T$ by $\widehat{T}$,

$\widehat{T}\coloneqq\widehat{S}_{\bf
Id}^{-1}\circ\widehat{\mathcal{M}}\circ\widehat{S},$ (B.6)

where $\widehat{S}$ is the usual maximum likelihood estimator of the KR
rearrangement $S$ (Section 7), $\smash{\widehat{S}_{\bf Id}}$ is an estimator
of $\smash{S_{\bf Id}}$ obtained (for free) by selecting the corresponding
components of $\widehat{S}$, while $\smash{\widehat{\mathcal{M}}}$ is defined
by an estimator, $\widehat{\mathfrak{m}}$, of the increasing rearrangement
$\mathfrak{m}$ that pushes forward a standard normal to
$\pi_{\widehat{S}^{1}}$. The latter is an approximation of the posterior
marginal, $\pi_{X_{1}|y^{*}}$, which depends only on the first component of
$\smash{\widehat{S}}$.

In principle, we could define $\widehat{\mathfrak{m}}$ to be the solution of a
non-convex minimization problem like (10.4), for some parameterization of the
rearrangement, which is now just a one dimensional function, and thus easier
to represent than a general rearrangement on $\mathbb{R}^{n}$. In practice,
however, we can simply target a different characterization of the increasing
rearrangement in terms of CDFs. If $F:\mathbb{R}\rightarrow[0,1]$ is the CDF
of $\mathcal{N}(0,1)$ and
$F_{\pi_{\widehat{S}^{1}}}:\mathbb{R}\rightarrow[0,1]$ is the CDF of
$\pi_{\widehat{S}^{1}}$, then we can express $\mathfrak{m}$ as

$\mathfrak{m}=F_{\pi_{\widehat{S}^{1}}}^{-1}\circ F,$ (B.7)

and thus $\widehat{\mathfrak{m}}$ can be any estimator of (B.7). There are
many options here, from parametric to fully nonparametric estimators. We
propose the following strategy. For a collection of $N$ grid points
$(x_{1},\ldots,x_{N})$ on $\mathbb{R}$, let $f_{i}$ denote an approximation of
$F_{\pi_{\widehat{S}^{1}}}(x_{i})$ obtained via numerical integration (recall
that we can evaluate the density $\pi_{\widehat{S}^{1}}$ everywhere up to a
normalizing constant). Notice that
$\smash{\mathfrak{m}^{-1}}=\smash{F^{-1}}\circ F_{\pi_{\widehat{S}^{1}}}$ and
let $\xi_{i}\coloneqq\smash{F^{-1}}(f_{i})$ for all $i=1,\ldots,N$. We can
think of $\xi_{i}$ as an approximation of $\mathfrak{m}^{-1}(x_{i})$. Hence,
the pairs $(\xi_{i},x_{i})_{i=1}^{N}$ yield approximate input-output
evaluations of the rearrangement $\mathfrak{m}$, and can be used to define a
convex regression problem for the estimator $\widehat{\mathfrak{m}}$,

$\widehat{\mathfrak{m}}\in\arg\min_{\mathfrak{u}\in\bm{\mathcal{H}}}\,\frac{1}{N}\sum_{i=1}^{N}\left(\,\mathfrak{u}(\xi_{i})-x_{i}\right)^{2},$
(B.8)

for some space $\bm{\mathcal{H}}$ of univariate monotone maps (Appendix A).

We can further use localization ideas from Section 7 and Remark 3 to
regularize the estimation of $T$ in high dimensions.

######  Remark 7 (Connection with the multivariate rank histogram filter).

Intuitively, the multivariate rank histogram filter implements a nonparametric
estimator of a deterministic update similar to (B.4), exploiting the
factorization (B.2) of the posterior density. The idea is to (1) apply Bayes’
rule locally—along the observed component—combining the histogram
approximation of $\pi_{X_{1}}$ suggested by [4] with the likelihood function
$\pi_{Y|X_{1}}$ to yield a piecewise constant approximation of the posterior
marginal $\pi_{X_{1}|y^{*}}$, and (2) propagate information to the unobserved
components by computing a sequence of (nonparametric) particle approximations
of the conditionals $\pi_{X_{k}|\bm{X}_{1:k-1}}$, for $k\geq 2$. Recall that
the conditionals $\pi_{X_{k}|\bm{X}_{1:k-1}}$ fully characterize the KR
rearrangement $S$ in (B.4). In high dimensions, it is increasingly challenging
to deploy particle approximations of the conditionals
$\pi_{X_{k}|\bm{X}_{1:k-1}}$, and thus [46] resorts to a mean-field
approximation of the posterior density,

$\pi_{\bm{X}|y^{*}}\approx\pi_{X_{1}|y^{*}}\,\prod_{k=2}^{n}\pi_{X_{k}|X_{1}},$
(B.9)

which relies only on the characterization of bivariate conditionals.

## References

  * [1] [author] Acevedo, WalterW., de Wiljes, JanaJ. and Reich, SebastianS. (2017). Second-order accurate ensemble transform particle filters. SIAM Journal on Scientific Computing 39 A1834–A1850.
  * [2] [author] Amezcua, JavierJ. and Van Leeuwen, Peter JanP. J. (2014). Gaussian anamorphosis in the analysis step of the EnKF: a joint state-variable/observation approach. Tellus A: Dynamic Meteorology and Oceanography 66 23493\.
  * [3] [author] Anderes, E.E. and Coram, M.M. (2012). A general spline representation for nonparametric and semiparametric density estimates using diffeomorphisms. arXiv:1205.5314.
  * [4] [author] Anderson, Jeffrey LJ. L. (2010). A non-Gaussian ensemble filter update for data assimilation. Monthly Weather Review 138 4186–4198.
  * [5] [author] Bengtsson, ThomasT., Snyder, ChrisC. and Nychka, DougD. (2003). Toward a nonlinear ensemble filter for high-dimensional systems. Journal of Geophysical Research: Atmospheres 108\.
  * [6] [author] Bertsekas, Dimitri PD. P. (1982). Projected Newton methods for optimization problems with simple constraints. SIAM Journal on control and Optimization 20 221–246.
  * [7] [author] Bickel, P. J.P. J. and Levina, E.E. (2008). Regularized estimation of large covariance matrices. The Annals of Statistics 36 199–227.
  * [8] [author] Bigoni, D.D., Spantini, A.A. and Marzouk, Y.Y. (2019). On the computation of monotone transports. In preparation.
  * [9] [author] Bigoni, DanieleD., Zahm, OlivierO., Spantini, AlessioA. and Marzouk, YoussefY. (2019). Greedy inference with layers of lazy maps. arXiv:1906.00031.
  * [10] [author] Bogachev, V. I.V. I., Kolesnikov, A. V.A. V. and Medvedev, K. V.K. V. (2005). Triangular transformations of measures. Sbornik: Mathematics 196 309\.
  * [11] [author] Box, George EPG. E. and Cox, David RD. R. (1964). An analysis of transformations. Journal of the Royal Statistical Society. Series B (Methodological) 211–252.
  * [12] [author] Bröcker, JochenJ. (2012). Evaluating raw ensembles with the continuous ranked probability score. Quarterly Journal of the Royal Meteorological Society 138 1611–1617.
  * [13] [author] Burgers, GerritG., Jan van Leeuwen, PeterP. and Evensen, GeirG. (1998). Analysis scheme in the ensemble Kalman filter. Monthly weather review 126 1719–1724.
  * [14] [author] Chorin, A. J.A. J. and Tu, X.X. (2009). Implicit sampling for particle filters. Proceedings of the National Academy of Sciences 106 17249–17254.
  * [15] [author] Crisan, D.D. and Doucet, A.A. (2002). A survey of convergence results on particle filtering methods for practitioners. IEEE Transactions on Signal Processing 50 736–746.
  * [16] [author] Cui, T.T., Martin, J.J., Marzouk, Y.Y., Solonen, A.A. and Spantini, A.A. (2014). Likelihood-informed dimension reduction for nonlinear inverse problems. Inverse Problems 30 114015\.
  * [17] [author] Daum, F.F. and Huang, J.J. (2008). Particle flow for nonlinear filters with log-homotopy. In SPIE Defense and Security Symposium 696918–696918. International Society for Optics and Photonics.