2.1. Data assimilation as Bayesian inference

DA comprises a set of techniques in NWP that aim to combine earth observations with assumptions about the state of the weather to produce an updated estimate of the weather. In this work, we simplify the full DA problem to spatial inference of a single variable. Let $ \boldsymbol{f}=\left({f}_1,\dots, {f}_n\right)\in {\mathrm{\mathbb{R}}}^n $ be a weather variable, such as surface temperature, that is discretized on a large 2D spatial grid consisting of points $ {\boldsymbol{x}}_1,\dots, {\boldsymbol{x}}_n $ . We also have $ m $ observations, $ \boldsymbol{y}\in {\mathrm{\mathbb{R}}}^m $ , which may either be derived from remote sensing products or direct observations of atmospheric variables. From a probabilistic perspective, DA can then be understood as a Bayesian inference problem: our assumptions about the weather can be encoded as a prior $ p\left(\boldsymbol{f}\right) $ , the observations as the likelihood $ p\left(\boldsymbol{y}|\boldsymbol{f}\right) $ , and our updated estimate as the posterior computed via Bayes’ theorem as $ p\left(\boldsymbol{f}|\boldsymbol{y}\right)\propto p\left(\boldsymbol{y}|\boldsymbol{f}\right)p\left(\boldsymbol{f}\right) $ . $ n $ is typically in the billions and $ m $ in the tens of millions (MetOffice 2024), making direct inference using Bayes’ theorem impossible. Thus, the choice of prior and likelihood is heavily influenced by the need to make inference efficient, and we discuss several options in the next section.

In operational DA systems, the problem is more complex, consisting of a 3D spatiotemporal grid and multiple weather variables at each grid point. However, inference in the simplified setting above is still challenging for large $ n $ and $ m $ . In Section 5, we discuss extensions to the full DA problem.

2.2. Existing methods for large-scale inference

2.2.3. 3D-Var

Another alternative that reduces the cost of GP regression is to only compute the MAP estimate, rather than the full posterior, that is find $ {\boldsymbol{f}}_{\mathrm{MAP}}=\arg {\max}_{\boldsymbol{f}}p\left(\boldsymbol{f}|\mathbf{y}\right) $ . This is the approach taken by 3D-Var, which is known as a “variational” method. The maximization can also be expressed as minimizing the cost function $ J\left[\boldsymbol{f}\right]=-\log p\left(\boldsymbol{f}|\mathbf{y}\right)=-\log p\left(\mathbf{y}|\boldsymbol{f}\right)-\log p\left(\boldsymbol{f}\right)+C $ , where $ C $ is a constant that does not depend on $ \boldsymbol{f} $ . Substituting in the GP prior and Gaussian likelihood, the cost function becomes

(1) $$ J\left[\boldsymbol{f}\right]=\frac{1}{2}{\left(\mathbf{y}-\boldsymbol{H}\left(\boldsymbol{f}\right)\right)}^T{\boldsymbol{R}}^{-1}\left(\mathbf{y}-\boldsymbol{H}\left(\boldsymbol{f}\right)\right)+\frac{1}{2}{\left(\boldsymbol{f}-{\boldsymbol{f}}_b\right)}^T{\mathtt{\varSigma}}^{-1}\left(\boldsymbol{f}-{\boldsymbol{f}}_b\right). $$

In practice, this is minimized using an optimizer, for example, L-BFGS (Liu and Nocedal Reference Liu and Nocedal1989) (as seen in D’Amore et al. (Reference D’Amore, Laccetti, Romano, Scotti and Murli2015)) or a Krylov solver (as seen in Bauer et al. (Reference Bauer, Quintino, Wedi, Bonanni, Chrust, Deconinck, Diamantakis, Düben, English, Flemming, Gillies, Hadade, Hawkes, Hawkins, Iffrig, Kühnlein, Lange, Lean, Marsden, Müller, Saarinen, Sarmany, Sleigh, Smart, Smolarkiewicz, Thiemert, Tumolo, Weihrauch, Zanna and Maciel2020)). 3D-Var is more flexible than optimal interpolation, as it can handle nonlinear observation models $ \boldsymbol{H} $ . However, it has the downside of not being able to provide uncertainty estimates, as it is a MAP estimator. In weather forecasting applications, 3D-Var is usually extended to 4D-Var, which includes a time component.

2.3. Factor graphs and inference with message passing

In this work, we perform inference using message passing (Kschischang et al. Reference Kschischang, Frey and Loeliger2001), which computes the marginals of any joint probability model that can be expressed as a factor graph. We introduce message passing for a general continuous distribution $ g\left(\boldsymbol{f}\right)=g\left({f}_1,\dots, {f}_n\right) $ , and specialize it to our posterior in Section 3.2. $ g\left(\boldsymbol{f}\right) $ can be expressed as a factor graph if it has a known decomposition

(2) $$ g\left(\boldsymbol{f}\right)\propto \prod \limits_{i=1}^n{\phi}_i\left({f}_i\right)\prod \limits_{j=1}^n{\phi}_{ij}\left({f}_i,{f}_j\right), $$

where $ {\phi}_i $ and $ {\phi}_{i,j} $ , referred to as the nodewise and pairwise factors, respectively, are functions from a variable, or a pair of variables, to $ \mathrm{\mathbb{R}} $ . These do not have to be probability distributions. The factorization in (2) induces a sparse graph on the variables $ {\left\{{f}_i\right\}}_{i=1}^n $ , where two nodes $ {f}_i $ and $ {f}_j $ are connected if and only if $ {\phi}_{i,j} $ is nonconstant. Given such a graph, the algorithm associates a pair of messages on each edge at each iteration $ t $ : a message $ {m}_{ij}^t=\left({a}_{ij}^t,{b}_{ij}^t\right)\in \mathrm{\mathbb{R}}\times \mathrm{\mathbb{R}} $ from $ {f}_i $ to $ {f}_j $ , and a similar message $ {m}_{ji}^t $ from $ {f}_j $ to $ {f}_i $ . We use the message-passing variant introduced by Ruozzi and Tatikonda (Reference Ruozzi and Tatikonda2013) (in turn a generalization of Wiegerinck and Heskes (Reference Wiegerinck and Heskes2002)). This is summarized in Algorithm 1 and illustrated for our application in Figure 2, and Appendix A describes it in more detail.

Figure 2. Illustration of node A sending a message to node B. Black arrows indicate incoming messages that are combined to compute the outgoing message in blue.

We use the notation $ \left\{{f}_k\sim {f}_i\right\} $ to denote the set of all variables $ {f}_k $ that are connected with $ {f}_i $ , $ \mathrm{C}\mathrm{O}\mathrm{M}\mathrm{PUTE}\ \mathrm{O}\mathrm{UTGOING}\ \mathrm{M}\mathrm{ESSAGE}\left(\right) $ and $ \mathrm{C}\mathrm{OMPUTE}\ \mathrm{M}\mathrm{ARGINAL}\left(\right) $ are defined formally in Appendix B, $ T $ is the total number of iterations, and $ c\in \mathrm{\mathbb{R}} $ is a hyperparameter to re-weight contributions of the messages, necessary to aid convergence. Note that each iteration of the inner for loops is independent, and each node only writes and reads messages with the nodes directly connected to it. The algorithm is therefore very amenable to distributed computation.

While this instance of message passing can theoretically compute both the mean and variance of the marginals, in practice, the variance estimates are biased and do not provide useful estimates of the uncertainty (Weiss and Freeman Reference Weiss and Freeman1999). Thus, we only use message passing to compute the posterior mean, the MAP estimate of the posterior. This makes message passing an alternative to 3D/4D-Var.

3.1. Derivation of the factor graph

The Matérn GP prior can be characterized as the solution to a stochastic partial differential equation (SPDE) of the form

(3) $$ {\left({\kappa}^2-\Delta \right)}^{\alpha /2}f=\mathbf{W}, $$

where $ f $ is the process, $ \Delta $ is the Laplacian operator, $ \kappa $ and $ \alpha $ are the positive hyperparameters, and $ \mathbf{\mathcal{W}} $ is the spatial white-noise process with spectral density $ {\sigma}^2q $ , for hyperparameters $ \sigma, q\in \mathrm{\mathbb{R}} $ . Following Lindgren et al. (Reference Lindgren, Rue and Lindström2011), we first derive a GMRF representation of the Matérn GP by discretizing this SPDE using finite differences (finite elements would also be possible). On a uniform 2D grid with step sizes $ \Delta x $ and $ \Delta y $ in the $ x $ and $ y $ directions, respectively, this yields a random matrix–vector system

(4) $$ \boldsymbol{Lf}=\boldsymbol{w},\hskip1em \mathrm{where}\hskip1em \boldsymbol{w}=\sqrt{\frac{\sigma^2q}{\Delta x\Delta y}}\mathbf{z},\hskip1em \mathbf{z}\in {\mathrm{\mathbb{R}}}^n\sim \mathbf{\mathcal{N}}\left(\mathbf{0},{\mathbf{I}}_n\right). $$

Here, $ \boldsymbol{L}\in {\mathrm{\mathbb{R}}}^{n\times n} $ is a matrix representing the operator $ \mathrm{\mathcal{L}}:= {\left({\kappa}^2-\Delta \right)}^{\alpha /2} $ under discretization, which is guaranteed to be sparse if the exponent $ \alpha /2 $ is an integer (Lindgren et al. Reference Lindgren, Rue and Lindström2011), and $ \boldsymbol{f},\boldsymbol{w} $ are finite-dimensional vector representations of the random fields $ f $ and $ \mathcal{W} $ (see Appendix C for details on the discretization). Now, (4) implies that $ \boldsymbol{f} $ is a Gaussian random variable of the form

(5) $$ \boldsymbol{f}\sim \mathbf{\mathcal{N}}\left(\mathbf{0},{\left(\gamma {\boldsymbol{L}}^T\boldsymbol{L}\right)}^{-1}\right),\hskip1em \mathrm{where}\hskip1em \gamma \hskip0.5em := \frac{\Delta x\Delta y}{\sigma^2q}, $$

which is a GMRF, since its precision matrix $ \boldsymbol{P}\hskip0.5em := \hskip0.5em \gamma {\boldsymbol{L}}^T\boldsymbol{L} $ is sparse (here, $ \boldsymbol{L} $ is sparse and banded). We can build a graph from this GMRF by the following simple rule: Take all of $ {f}_1,\dots, {f}_n $ as nodes in the graph and connect two nodes $ {f}_i $ and $ {f}_j $ by an edge if $ {\left[\boldsymbol{P}\right]}_{ij}\ne 0 $ . Then, we have

(6) $$ p\left(\boldsymbol{f}\right)\propto \exp \left(-\frac{1}{2}{\boldsymbol{f}}^T\boldsymbol{Pf}\right)=\exp \left(-\sum \limits_{i=1}^n\frac{1}{2}{\left[\boldsymbol{P}\right]}_{ii}{f}_i^2-\sum \limits_{j\sim i}{\left[\boldsymbol{P}\right]}_{ij}\hskip0.5em {f}_i\hskip0.5em {f}_j\right), $$

where $ {\Sigma}_{j\sim i}\left(\cdot \right) $ denotes the sum over all indices $ j $ that are adjacent to $ i $ in the graph. Setting

(7) $$ {\phi}_i\left({f}_i\right) := \exp \left(-\frac{1}{2}{\left[\boldsymbol{P}\right]}_{ii}{f}_i^2\right)\hskip1em \mathrm{and}\hskip1em {\phi}_{ij}\left({f}_i,{f}_j\right) := \exp \left(-{\left[\boldsymbol{P}\right]}_{ij}{f}_i{f}_j\right), $$

we have that the prior $ p\left(\boldsymbol{f}\right)\propto {\prod}_{i=1}^N{\phi}_i\left({f}_i\right){\prod}_{j\sim i}{\phi}_{ij}\left({f}_i,{f}_j\right) $ , thus we have a factor graph representation.

3.2. Computing the posterior mean with message passing

Having calculated the factor graph corresponding to the prior, we can now apply message passing (Algorithm 1) to combine this with the observations and compute the posterior marginals. We make several modifications to tailor the algorithm to DA, which we summaries below and detail in Appendix B.

3.2.2. Update damping

To improve the stability of the algorithm, we follow Pretti (Reference Pretti2005) and dampen the updates of the messages, replacing line 6 of Algorithm 1 with

$$ {m}_{ij}^{t+1}=\left(1-\eta \right){m}^t+\eta\;\mathrm{C}\mathrm{O}\mathrm{M}\mathrm{PUTE}\ \mathrm{O}\mathrm{UTGOING}\ \mathrm{M}\mathrm{ESSAGE}\left(\cdot \right), $$

where $ \eta \in \left(0,1\right) $ is a hyperparameter, which we refer to as the learning rate.

3.2.4. Multigrid

We apply a multigrid technique to speed up the convergence of the algorithm. These have been used extensively when solving partial differential equations numerically, and provide an efficient way of accelerating the convergence of iterative approaches if multiscale phenomena are being modeled, with grids at different resolutions capturing different spatial scales. In the case of message passing, we can intuitively view information propagating from the observation locations across the graph. If the density of the observations is low, this can take many iterations. To solve this, the multigrid approach starts with a low-resolution grid, and iterates message passing until convergence—this is fast on a low-resolution grid. It then doubles the size of the grid, initializing the messages to the converged messages from the previous grid. This process is repeated until we reach the target grid size. Observations are taken on the target resolution and introduced at each multigrid level when the observation location exactly coincides with the grid level coordinates. Figure 3 illustrates the procedure.

Figure 3. Illustration of the multigrid implementation, showing the marginal means computed at two levels ( $ 128\times 128 $ to $ 256\times 256 $ ) of resolution on the simulated data.

4.1. Effect of domain size and observation density

Our first set of experiments are performed on simulated data. We sample a square ground truth field from a GMRF prior, randomly select a given fraction of the grid points as observations, and perform inference against these observations and the same prior. Table 1 shows the RMSE between the posterior mean and the ground truth and the time taken, as we vary the grid size and the observation density. In the message-passing runs, the multigrid approach is used, with a base grid size of $ 32\times 32 $ in all cases.

Table 1. Comparison on simulated data. We give the mean over three ground truths; we do not observe significant variance (therefore omitted). Bold indicates where either 3D-Var or message passing performed better. R-INLA is included to show the minimal achievable error given the prior, as it computes an exact posterior

The results show that 3D-Var and message passing achieve similar RMSEs in all cases, although both have more error than R-INLA when only 1% of the grid is observed. When 5% and 10% of the grid is observed, 3D-Var and message passing take similar amounts of time; however, message passing is significantly slower for a larger grid when only 1% is observed. While the iteration time of message passing is independent of the number of observations, as the observation density falls it takes an increasing number of iterations for the information to propagate from the observed points across the grid, thus early stopping happens later. R-INLA is much slower than the other methods, because it is a sequential method that cannot take advantage of the GPU.

4.2. Large-scale example

For a more realistic use case, we consider the global surface temperature field. We take the ground truth data from a run of the Met Office’s Unified Model (Walters et al. Reference Walters, Baran, Boutle, Brooks, Earnshaw, Edwards, Furtado, Hill, Lock, Manners, Morcrette, Mulcahy, Sanchez, Smith, Stratton, Tennant, Tomassini, Van Weverberg, Vosper, Willett, Browse, Bushell, Carslaw, Dalvi, Essery, Gedney, Hardiman, Johnson, Johnson, Jones, Jones, Mann, Milton, Rumbold, Sellar, Ujiie, Whitall, Williams and Zerroukat2019) at N1280 resolution, where the data are valid for 06UTC 2020-01-01. To avoid issues with boundary conditions, we consider a clipped domain with dimensions $ 2500\times 1500=3.75M $ grid points. We use spherical polar coordinates. The observation locations are generated from the geographical positions (latitude, longitude) of weather-focused satellites calculated over a 3-hour window, which corresponds to $ \approx 8\% $ of the grid being observed. As the prior mean, we select a climatology mean of the global surface temperature calculated from ERA5 (Hersbach et al. Reference Hersbach, Bell, Berrisford, Hirahara, Horányi, Muñoz-Sabater, Nicolas, Peubey, Radu, Schepers, Simmons, Soci, Abdalla, Abellan, Balsamo, Bechtold, Biavati, Bidlot, Bonavita, De Chiara, Dahlgren, Dee, Diamantakis, Dragani, Flemming, Forbes, Fuentes, Geer, Haimberger, Healy, Hogan, Hólm, Janisková, Keeley, Laloyaux, Lopez, Lupu, Radnoti, Rosnay, Rozum, Vamborg, Villaume and Thépaut2020).

Figure 1 highlights the resultant mean estimates from the message-passing approach, while error plots are shown in Figure 7 in Appendix D. 3D-Var achieved an area-weighted RMSE of $ 2.33 $ K and took $ 16 $ seconds, while message passing achieved an area-weighted RMSE of $ 1.23 $ K and took $ 115 $ seconds. For comparison, the area-weighted RMSE calculated for the prior mean (ERA5) against the high-resolution temperature field is $ 2.78 $ K. R-INLA did not complete after $ >1 $ hour of processing time; thus, we do not include its results.