Nonlinear ensemble filtering with diffusion models:
Application to the surface quasi-geostrophic dynamics

In a diffusion model, there is a forward ${\mathbb{R}}^{N_{x}}$-valued stochastic differential equation (SDE) defined as

where $\boldsymbol{W}_{t}$ is a standard ${\mathbb{R}}^{N_{x}}$-valued Brownian motion (Wiener process) corresponding to an Itô type stochastic integral $\int\cdot d\boldsymbol{W}_{t}$, while $b$ and $\sigma$ are two explicitly given functions referred to as the drift and diffusion coefficients, respectively. The initial condition $\boldsymbol{Z}_{0}$ of the SDE in Eq. (1) follows some target distribution with its probability density function (pdf) denoted by $Q({\mathbf{z}}_{0})$. One can show that with properly chosen $b$ and $\sigma$, the diffusion process $\{\boldsymbol{Z}_{t}\}_{0\leq t\leq T}$ can transform any target pdf ${Q}({\mathbf{z}}_{0})$ to a standard Gaussian, i.e. $\boldsymbol{Z}_{T}\sim\mathcal{N}({\mathbf{0}},{\mathbf{I}})$, where $[0,T]$ is often referred to as the pseudo-time interval.

Generating samples $\{\boldsymbol{Z}_{0}^{i}\}_{i=1}^{N}$ of the target random variable $\boldsymbol{Z}_{0}$ pertains to simulating the following reverse-time SDE:

where $\int\cdot d\overleftarrow{\boldsymbol{W}}_{t}$ is a backward Itô stochastic integral (Kloeden and Platen 1992; Bao et al. 2016), and ${\mathbf{s}}(\cdot,t)$ is the so-called score function given by

where $Q({\mathbf{z}}_{t})$ denotes the pdf of $\boldsymbol{Z}_{t}$. Note that the reverse-time SDE in Eq. (2) is also a diffusion process except that the propagation direction is backwards in time from $T$ to $0$ with initial condition $\boldsymbol{Z}_{T}$ given at time $T$. An important result from the literature is that the solution $\boldsymbol{Z}_{0}$ of the reverse-time SDE follows the target distribution (Bao et al. 2023b). The practical implications are that we can generate samples from a standard Gaussian distribution (which can be done efficiently) and use the reverse-time SDE to transform them to samples of the target distribution. The score function ${\mathbf{s}}(\cdot,t)$ has an important role in this mapping process as it stores information about the distribution of the samples, which in turn helps guide their transformation over the pseudo-time interval as $T\rightarrow 0$. In particular, having the score function associated with the target pdf $Q({\mathbf{z}}_{0})$ and the predefined forward SDE allows us to generate an unlimited number of target samples by running the reverse-time SDE in Eq. (2).

In both the forward SDE and the reverse-time SDE, one needs to carefully choose the drift and diffusion coefficients $b$ and $\sigma$. While there are multiple options for $b$ and $\sigma$, in this work we let

with $\alpha_{t}=1-t$ and $\beta_{t}=\sqrt{t}$ for $t\in[0,1]$, which is consistent with the choice made in Song et al. (2021). Nevertheless, in future work we plan to explore the sensitivity of EnSF to different options for the drift and diffusion coefficients.

The traditional use of diffusion models in the machine learning (ML) literature is to generate highly realistic images and videos (Song et al. 2021). Later in our exposition, we will discuss how this powerful approach can be also used in the context of data assimilation (DA).

The formulation of every DA method requires two main ingredients – (stochastic) dynamic and observation models. Let $k=0,1,2,...,K$ be a time index. A general representation of the discretized system evolution is given by

Given the model and observation equations, a rather general formulation of the DA problem is to seek the filtering (posterior) pdf $P({\mathbf{x}}_{k}|{\mathbf{y}}_{1:k})$¹¹1The notation $1:k$ is shorthand for the set of integers from $1$ to $k$., which can be done recursively by iterating through a prediction and an update step:

The central idea in score-based filtering is to create score models which represent the filtering pdfs at different times and then generate analysis ensemble members according to the reverse-time SDE in Eq. (2). To this end, we introduce the score functions ${\mathbf{s}}_{k|k-1}$ and ${\mathbf{s}}_{k|k}$ corresponding to the prior pdf $P({\mathbf{x}}_{k}|{\mathbf{y}}_{1:k-1})$ and the posterior pdf $P({\mathbf{x}}_{k}|{\mathbf{y}}_{1:k})$:

• •

  Prior filtering score ${\mathbf{s}}_{k|k-1}$ such that $\boldsymbol{Z}_{0}\equiv\boldsymbol{X}_{k}|\boldsymbol{Y}_{1:k-1}$; i.e., $Q({\mathbf{z}}_{0})\equiv P({\mathbf{x}}_{k}|{\mathbf{y}}_{1:k-1})$.

• •

  Posterior filtering score ${\mathbf{s}}_{k|k}$ such that $\boldsymbol{Z}_{0}\equiv\boldsymbol{X}_{k}|\boldsymbol{Y}_{1:k}$; i.e., $Q({\mathbf{z}}_{0})\equiv P({\mathbf{x}}_{k}|{\mathbf{y}}_{1:k})$.

Following the standard workflow of diffusion models, the reverse-time SDE is initialized with samples drawn from a standard Gaussian distribution; that is, $\{\boldsymbol{Z}^{m}_{t_{N}}\}_{m=1}^{M}\sim\mathcal{N}({\mathbf{0}},{\mathbf{I}})$. Leveraging the Euler scheme in Eq. (10), these initial samples are mapped to the analysis ensemble $\{\boldsymbol{X}^{m}_{k-1|k-1}\}_{m=1}^{M}$ which follows the filtering pdf $P({\mathbf{x}}_{k-1}|{\mathbf{y}}_{1:k-1})$ by construction.

The generative nature of the diffusion process needed to obtain the analysis ensemble $\{\boldsymbol{X}^{m}_{k-1|k-1}\}_{m=1}^{M}$ is worth elaborating on. With an appropriately estimated score function $\hat{{\mathbf{s}}}_{k-1|k-1}$, we can generate unlimited samples from the posterior pdf $P({\mathbf{x}}_{k-1}|{\mathbf{y}}_{1:k-1})$, which can be seen as a non-Gaussian extension of the class of resampling EnKFs discussed by Anderson (2001). In addition, the sample size $M$ can be an arbitrary number depending on the specific application and computing resources.

The estimation of score functions is a central topic in the score-based filtering approach discussed so far. The standard technique to estimate scores is via deep learning (Song et al. 2021; Bao et al. 2023b). Here, we introduce an ensemble approximation referred to as the Ensemble Score Filter (EnSF) which does not require the training of neural networks.

Since the likelihood score appearing in Eq. (12) can be computed explicitly in the Gaussian case considered here, the major computational effort in EnSF lies in evaluating the prior score $\hat{{\mathbf{s}}}_{k|k-1}$. The latter can be achieved by setting $\boldsymbol{Z}_{0}$ to the forecast ensemble $\{\boldsymbol{X}^{m}_{k|k-1}\}_{m=1}^{M}$. From the definition of score functions and the choice of $b$ and $\sigma$ in Eq. (4), we have that the conditional density $Q({\mathbf{z}}_{t}|{\mathbf{z}}_{0})$ needed in the forward SDE is given by

Then, we can apply a Monte Carlo approximation for ${\mathbf{s}}_{k|k-1}$ based on Eq. (14) at a given ${\mathbf{z}}$ and $t\in[0,1]$:

The SQG model belongs to a special class of quasi-geostrophic models in which a fluid of constant potential vorticity (PV) is bounded between two flat, rigid surfaces (Tulloch and Smith 2009b). Despite its idealized nature, the system is capable of simulating turbulent motions similar to those occurring in real geophysical flows (Smith et al. 2023). It is also suitable for DA studies because the SQG flow is inherently chaotic and sensitive to initial condition errors (Rotunno and Snyder 2008; Durran and Gingrich 2014).

In this study, we adopt the SQG formulation proposed by Tulloch and Smith (2009a) in which the dynamics reduce to the nonlinear Eady model with an f-plane approximation as well as uniform stratification and shear. For this case, the governing equations simplify to the advection of potential temperature on the bounding surfaces. Those equations are solved numerically by first applying a fast Fourier transform (FFT) to map model variables to spectral space. They are integrated forward with a 4${}^{\text{th}}$-order Runge Kutta scheme that uses a $2/3$ dealiasing rule and implicit treatment of hyperdiffusion. Further numerical details and open-source code can be found on the GitHub page shared at the end of the paper.

The EnSF performance is assessed with $20$ ensemble members using a standard observing system simulation experiment (OSSE) framework where synthetic observations are generated by corrupting the true (nature) run with random noise. The construction of the nature run mostly follows the details of Wang et al. (2021). Two notable exceptions are that we perform additional experiments with a higher resolution version of the model (96 $\times$ 96 points) and further carry out imperfect model experiments where the nature run is contaminated with unpredictable model errors.

More significant differences appear in the generation of synthetic observations. First, we utilize a 12-hour assimilation window. Compared to the 3-hour window of Wang et al. (2021), this constitutes a more challenging scenario as it creates larger differences with the true state and likely causes more significant departures from Gaussianity in the forecast ensemble. Moreover, we test a wider range of observing networks, starting with the simpler case of a fully observed state with linear observations and finishing with a partially and nonlinearly observed state (50$\%$ coverage) that is additionally subject to unexpected model errors. More specific details about the different observing networks are deferred to Section 4.

While there is a vast array of available EnKF methods we could compare EnSF’s performance against, here we opt for the Local Ensemble Transform Kalman Filter (LETKF; Hunt et al. 2007). LETKF is a popular EnKF variant that is currently implemented in several major operational centers (e.g., Schraff et al. 2016; Frolov et al. 2024). One of the most appealing features of this method is its efficiency – the model state is decomposed into overlapping subdomains which are updated separately (and in parallel) using the analysis equations of the Ensemble Transform Kalman Filter (ETKF; Bishop et al. 2001). Our numerical implementation considers local regions surrounding single grid points defined by a cutoff radius. The cut-off radius is determined by standard Gaspari-Cohn (GC) taper functions (Gaspari and Cohn 1999), and we also apply the original $R$-localization strategy of Hunt et al. (2007)²²2With R-localization, the local observation error variances are multiplied by the inverse GC function. to smoothly decrease the impact of observation away from each model grid point. Analogous to Wang et al. (2021), the vertical and horizontal localization scales are coupled dynamically through the Rossby radius of deformation. We also apply the relaxation to prior spread (RTPS) inflation discussed in Whitaker and Hamill (2012) in order to prevent underdispersive analysis ensembles as a result of the finite ensemble size. In most experiments, the inflation factor is optimally tuned for LETKF while EnSF currently sets RTPS to 1.

While the simultaneous use of localization and inflation can greatly improve LETKF’s performance, achieving optimal analysis results necessitates careful parameter tuning, which is computationally demanding in operational contexts. Moreover, sensitivity tests are required each time one makes changes to the NWP system (model resolution, type of assimilated observations, etc). Our initial results in Section 4 emphasize how LETKF’s skill is highly sensitive to suboptimal choices of localization and inflation even in the conceivably simpler case of a fully and linearly observed state. This will be contrasted with EnSF’s performance which leads to stable performance across all experiments and without any further ensemble regularization strategies (although additional experiments not presented in the paper suggest that optimal tuning of RTPS can further improve EnSF’s performance).

In EXP_L1, we compare the performance of EnSF and LETKF in the case where the state of the SQG model is fully and linearly observed. Measurements are corrupted by additive Gaussian errors such that

It was already discussed in Section 33.3 that one of the major obstacles with EnKFs is the need for a careful parameter tuning. The impact of suboptimal choices for these parameters is illustrated in Figure  1 where LETKF uses the same localization and inflation parameters as Fig. 11a of Wang et al. (2021) [horizontal localization scale of $L_{h}=2500$ km and RTPS $=0.6$]. Despite the physical realism of these parameter choices, we see that the LETKF experiment quickly diverges due to changes in the model resolution ($96\times 96$ vs. $64\times 64$ grid points in the horizontal), assimilation period (12h vs. 3h) and number of observations (fully vs. partially observed state). By contrast, EnSF demonstrates stable performance throughout all 300 assimilation cycles without requiring any additional parameter tuning.

To achieve a fair comparison between the two DA methods in this linear observation regime, our next step is to optimize LETKF’s localization and inflation parameters. We revert back to the coarser model resolution of $64\times 64$ horizontal grid points and see from Figure 2a that the best LETKF performance is achieved with $L_{h}=2000$ km and RTPS = $0.3$. Evidently, LETKF performs slightly better in this case. However, it is worth noting that EnSF continues to maintain stable performance despite the change in model resolution and lack of additional parameter tuning. EnSF’s sensitivity to different parameter settings will be the subject of future work and is expected to further reduce the gap between the two filtering techniques in this linear regime (EXP_L1).

In the next experimental setting EXP_L2, we address the more challenging but realistic scenario involving an imperfect model. While still assuming direct measurements of all state variables, we add an a-priori unknown stochastic noise $\boldsymbol{E}^{u}$ to the SQG state evolution. Assuming the noise is Gaussian, the modified stochastic-dynamic model can be written as

The covariance matrix ${\mathbf{\Sigma}}_{k}$ is taken to be diagonal as we want to impose spatially uncorrelated model noise. We also assume that the unknown model errors are white in time (i.e., no temporal correlations). More specifically, $\boldsymbol{E}^{u}_{k}$ in EXP_L2 is defined as the composition of four distinct state-dependent error processes: (1) $20\%$ chance of occurrence (in time) with $20\%$ model noise, (2) $15\%$ chance of occurrence with $30\%$ model noise, (3) $10\%$ chance of occurrence with $40\%$ model noise, and (4) $5\%$ chance of occurrence with $50\%$ model noise. The purpose of introducing this variety of model errors is to simulate the effects of flow-dependent model uncertainties. For example, the rare occurrence of high-amplitude model errors in the 4${}^{\text{th}}$ process might be associated with the rapid development of dynamical instabilities in the $\theta$ field. The study of Held et al. (1995) shows examples of such instabilities and discusses how they tend to form along temperature filaments (see their Fig. 2).

It is crucial to highlight that although we specified a particular structure for the time-dependent error process $\boldsymbol{E}^{u}_{k}$, obtaining an explicit expression for $\boldsymbol{E}^{u}_{k}$ is impractical in reality. Consequently, adjusting the LETKF inflation and localization parameters to account for the unknown model errors $\boldsymbol{E}^{u}_{k}$ is not feasible. In Figure 3, we visually illustrate the impact of adding model errors of various amplitude to the true state at the final time.

A comparison between the analysis RMSE errors of EnSF and LETKF in EXP_L2 is presented in Figure 4. Unlike our previous experiment, LETKF diverges from the true state as model errors accumulate in time, with analysis RMSEs now comparable to the free forecast run where no observations are assimilated. This finding offers additional evidence for the enhanced sensitivity of LETKF to model imperfections – without the implementation of more advanced regularization techniques (e.g., adaptive inflation), LETKF will be susceptible to analysis errors even when nearly optimal parameters have been used and the system is fully and linearly observed. This is to be contrasted with EnSF’s performance where we achieve stable performance throughout all analysis cycles despite the lack of any further tuning. A snapshot of these differences at the last analysis time is shown in Figure 5. The top 3 figures offer a visual confirmation that EnSF’s analysis mean (panel b) is much closer to the truth (panel a). While LETKF (panel c) does well at representing large-scale patterns, it struggles to capture some of the small-scale features and extreme values in the potential temperature field. The larger error magnitude in LETKF is also confirmed by the two error plots in the second row.

In the second round of experiments, the synthetically generated observations are nonlinear functions of the SQG state. Specifically, we consider an arctangent nonlinearity and additive Gaussian observation errors. The observation model relevant to EXP_NL1 from Table 1 can be written as

Notice that the observation error variance is scaled relative to the linear observation experiments in order to reflect the narrower range of the arctangent function. In this case, we choose $\tilde{{\mathbf{R}}}_{k}=0.01\times{\mathbf{R}}_{k}=0.01\times{\mathbf{I}}$. Moreover, EXP_NL1 tests are carried out in the absence of unpredictable model shocks.

Analogous to the linear observation setting, stabilizing LETKF’s performance requires extensive parameter tuning. Figure 6a indicates the existence of a very narrow range of optimal localization and inflation parameters for which LETKF does not diverge – a manifestation of the strong deviations from non-Gaussianity in the posterior distribution. The red curve in Figure 6b confirms that this best tuned LETKF maintains stability during the entire experiment. On the other hand, a naive substitution of the optimal LETKF parameters from the linear observation setting results in a rapidly diverging filter (magenta curve in Figure 6b). While this comparison admittedly constitutes an extreme modification of the observing system, it reaffirms our earlier point that any changes in operational NWP systems require costly calibrations of the underlying EnKF algorithm. By contrast, EnSF systematically outperforms the best LETKF configuration without further application of ensemble regularization strategies.

Our last experimental setting EXP_NL2 considers the most complex DA scenario – only half of the SQG state is observed nonlinearly such that

We also work under the imperfect model assumption described by Eq. (20), but choose a simpler definition for the unknown model errors $\boldsymbol{E}^{u}_{k}$. In particular, a single stochastic noise is added 10% of the time with a magnitude that equals 30% of the nature run’s values.

Under the challenging DA setting of EXP_NL2, we find that EnSF’s analysis RMSEs remain nearly intact (compare blue curves in Figures 6b and 8b), while LETKF (Figure 8a) exhibits a visible skill deterioration (compare red curves in Figures 6b and 8b) despite our efforts to fine tune its performance for this case.

Given the multiscale characteristics of the SQG model, we conclude this section by presenting a spectral comparison of the EnSF and LETKF results in Figure 9. The plotted diagnostics still refer to the EXP_NL2 setting (partially and nonlinearly observed state corrupted by unpredictable model errors) but are now averaged in time. Examination of the solid red and blue curves in panel (a) indicates that EnSF has consistently lower analysis errors for all wavenumbers, but the most significant improvements occur at large scales (also refer to panel b). Focusing our attention on the LETKF results (red cuferves), we notice a systematic underdispersion of the analysis ensemble, which is most pronounced at large scales (panel c). This is perhaps one of the important reasons why a large number of LETKF experiments diverge in EXP_NL2. A different situation emerges for EnSF (blue curves) which tends to be overdispersive on average (panel c). This behavior is connected to the specific inflation settings used in EnSF (RTPS $=1$) but is likely also influenced by our choice of SDE and score parameters such as $b(t),\sigma(t)$ and $h(t)$. We hypothesize that an optimal selection of these parameters will further improve EnSF’s performance and lead to a better spread-error consistency. Nevertheless, it is clear that even this “vanilla” implementation of EnSF offers significant performance benefits over LETKF.