Efficient high-dimensional variational data assimilation with machine-learned reduced-order models

The first step in the surrogate construction is the reduction of the degrees of freedom of the original data set to enable rapid forecasts. Projection-based reduced order models (ROMs) can effectively compress a high-dimensional model while still preserving its spatio-temporal structure. The dimensionality reduction is performed through the projection step where the high-dimensional model is projected onto a set of optimally chosen basis functions with respect to the $L_{2}$ norm [3]. The mechanics of a POD-ROM, can be illustrated for a state variable $\mathbf{x}\in\mathbb{R}^{N}$, where $N$ represents the size of the computational grid. The POD-ROM then approximates $\mathbf{x}$ as the linear expansion on a finite number of $k$ orthonormal basis vectors $\boldsymbol{\phi}_{i}\in\mathbb{R}^{N}$, a subset of the POD basis:

where $r_{i}\in\mathbb{R}$ is the $i$th component of ${\mathbf{r}}\in\mathbb{R}^{k}$, which are the coefficients of the basis expansion. The $\{\boldsymbol{\phi}_{i}\},~{}i=1,\ldots,k$, $~{}\boldsymbol{\phi}_{i}\in\mathbb{R}^{N}$ are the POD modes. POD modes in Equation 1 can be shown to be the left singular vectors of the snapshot matrix (obtained by stacking $M$ snapshots of $\mathbf{x}$), $\mathbf{X}=[\mathbf{x}_{1},\ldots,\mathbf{x}_{M}]$, extracted by performing a singular value decomposition (SVD) on $\mathbf{X}$ [22, 10]. That is,

where $\mathbf{U}\in\mathbb{R}^{N\times M}$ and $\mathbf{\Phi}_{k}$ represent the first $k$ columns of $\mathbf{V}$ after truncating the last $M-k$ columns based on the relative magnitudes of the cumulative sum of their singular values (see e.g., [31]). The total $L_{2}$ error in approximating the snapshots via the truncated POD basis is then

where $\sigma_{i}$ is the singular value corresponding to the $i$th column of $\mathbf{V}$ and is also the $i$th diagonal element of $\boldsymbol{\Sigma}$. It is well known that POD bases are $L_{2}$-optimal and present a good choice for efficient compression of high-dimensional data [3]. A recent alternative for compressing the data consists of spectral POD [49, 32, 24].

For forecasting a dynamical system from examples of time-series data, we utilize an encoder-decoder framework coupled with long short-term memory neural networks (LSTMs) [21]. The encoder-decoder formulation is given by a first step where a latent representation is derived through information from historical observations (in our case of the compressed representation of the full state), i.e.,

where $h$ is the output of a function approximated by an LSTM. The LSTM neural architecture is devised to account for long and short-term correlations in time-series data through the specification of a hidden state that evolves over time and is affected by each observation of the dynamical system. The value of the hidden state at the end of the input sequence of length $T$ may then be employed as an encoded representation, $\mathbf{h}$, of the input window. After $\mathbf{h}$ is obtained, forecasts at different distances in the future may be obtained by functional predictions. In this study, the ‘decoding’ component of the architecture is also given by an LSTM cell that is provided the encoded state information for each timestep of the output, i.e.,

where $\tilde{h}$ is an LSTM cell as well and $T_{o}$ is the length of the output window. Once forecasts in the POD-compressed space are obtained using the encoded-decoder LSTM network, the full-order state variable can be reconstructed by using the precomputed basis functions. A schematic of the architecture is shown in Figure 1. In the past several dynamical systems forecasts have been performed solely with the use of POD-LSTM type learning [39, 33, 30, 29]. However, as will be demonstrated shortly, this approach can be enhanced significantly with the use of real-time DA from sparse observations.

Strong-constraint four-dimensional variational (4D-Var) DA processes simultaneously all observations at all times $t_{1},t_{2},\dots,t_{T}$ within the assimilation window (see 2 for a schematic representation). The control parameters are typically the initial conditions $\mathbf{x}_{0}$, which uniquely determine the state of the system at all future times under the assumption that the model (6) perfectly represents reality. The background state is the prior best estimate of the initial conditions $\mathbf{x}_{0}^{\rm b}$, and has an associated initial background error covariance matrix $\mathbf{B}_{0}$. The 4D-Var problem provides the estimate $\mathbf{x}_{0}^{\rm a}$ of the true initial conditions as the solution of the following optimization problem

The first term of the sum (8b) quantifies the departure of the solution $\mathbf{x}_{0}$ from the background state $\mathbf{x}_{0}^{\rm b}$ at the initial time $t_{0}$. The second term measures the mismatch between the forecast trajectory (model solutions $\mathbf{x}_{i}$) and observations $\mathbf{y}_{i}$ at all times $t_{i}$ in the assimilation window. The covariance matrices $\mathbf{B}_{0}$ and $\mathbf{R}_{i}$ need to be predefined, and their quality influences the accuracy of the resulting analysis.

Given a low-dimensional differentiable surrogate model, for instance the proposed POD-LSTM, that approximates the dynamics of the full-order numerical model, the variational DA process can be accelerated dramatically. If such a surrogate model has predicted $\boldsymbol{\hat{x}}_{1}$ to $\boldsymbol{\hat{x}}_{T}$ for a range of timesteps in the reduced-space (i.e., $\boldsymbol{\hat{x}}\approx P\mathbf{x}$, $\boldsymbol{\hat{x}}\in\mathbb{R}^{K}$, with $K\ll N$) and observations from the real system, $\mathbf{y}_{1}$ to $\mathbf{y}_{T}$ are available, the objective function of the DA problem may be obtained by reconstructing the model state from the compressed representation. Here, $P:\mathbb{R}^{N}\rightarrow\mathbb{R}^{K}$ is the function that maps from physical space to reduced space. In case of POD based compression $P$ is a linear operation that projects from the physical space to the basis spanned by the truncated set of the left singular vectors of an SVD performed on snapshots from full-state model evaluation. If inversion from the truncated subspace is represented as $P^{\dagger}$, this gives us

Note that $\hat{\mathbf{x}}_{i}$, which is an approximation to the state at time $t_{i}$, is evaluated using the surrogate model—that is, by propagating $\boldsymbol{\hat{x}_{0}}$ using the surrogate model (LSTM in this case) and applying $P^{\dagger}$ to the result. This can be written mathematically as

Given a fixed $P^{\dagger}$ with an effective compression ratio (obtained from our dimensionality reduction technique), Equation 8b becomes a cost function (9) expressed in $K$ dimensions which is amenable to rapid update of the initial conditions. Moreover, gradients of this objective function with respect to the initial conditions are trivially computable because of the use of automatic differentiation of our LSTM neural network. The overall approach is as follows:

• •

  For each forecast window, collect random observations from the true state.

• •

  Perform an optimization of Equation 9 by perturbing the inputs to the LSTM neural network. This input is the projected version of the initial condition state.

• •

  If optimization has converged, perform one forecast with the optimized initial condition. This is the DA improved forecast.

The optimization methodology used in this article is the sequential least-squares programming approach [35] implemented in SciPy. The neural architecture was deployed using TensorFlow 2.4. Our overall implementation is in Python.

The data that we use in this study are a subset of 20 years of output from the regional climate model WRF version 3.3.1, prepared with methods and configurations described by Wang and Kotamarthi [52]. WRF, a fully compressible, nonhydrostatic, regional numerical prediction system with proven suitability for a broad range of applications, is driven by reanalysis data NCEP-Reanalysis2 (NCEP-R2) for the period 1984–2003. The NCEP-R2 dataset used in  [52] includes many of the observational datasets available to build dynamically consistent gridded fields that are typically used for forecast model initialization and boundary condition setting.

We use daily averaged geopotential height at 500 hPa (Z500) from WRF to demonstrate the approach developed in this study. We choose Z500 because it is a standard field for analysing the quality of weather forecasts, it does not have very strong local gradients (in contrast to fields such as humidity or precipitation) and does not depend on local conditions such as topography, yet most of the important global flow patterns—such as midlatitude jets and a gradient between poles and Equator—are visible in Z500.

We retrieve the Z500 data from the WRF output at a grid spacing of 12km and at 3 hour intervals with 515$\times$599 grid cells in the south-north and west-east directions, respectively. We then calulate per-grid-point daily averages to obtain one data value per grid point per day. Spatially, we coarsen the data by five strides which still maintains the spatial patterns of Z500 but reduces the data size significantly. Each daily Z500 snapshot therefore comprises 102$\times$119 60km$\times$60km grid points. We have 3287 such daily snapshots for the period 1 January 1984 to 31 December 1991. Despite the coarsening of the original 12km data, the typical eastward propagating waves are clearly visible in the northern hemisphere. Such waves are also observed in the real atmosphere and are one of the main features of midlatitude weather variability on timescales of several days.

We use the first six years of the daily averaged Z500 WRF data (1984–1989) for surrogate training and optimization. Of those data, we select 70% at random as training data, for use in training the supervised ML formulation, and keep the remaining 30% as validation data, for use in tuning the neural architecture hyperparameters and to control overfitting (a situation in which the trained network predicts well on the training data but not on the test data). We also construct a test data set consisting of one year of records (1991) for prediction and evaluation. We skip 1990 so as to ensure that there is no overlap in input windows. We note that we choose identity as our observation error covariance matrix and we use a scaling to improve convergence. In our future studies, in addition to using a realistic observation error covariance matrix, we will also pursue using a flow-dependent background error covariance matrix as described in [6].

We conduct a number of experiments to evaluate the performance of our emulator-assisted DA approach. In the first, we perform a grid-based search in which we vary both the number of POD modes (5, 10, 15) and the size of the input window (7, 14, 28, or 42 days of lead-time information) to identify the optimal combination for forecasting. Other hyperparameters such as learning rate, learning rate scheduler, number of LSTM cells, and number of neurons are determined by a combination of experience in previous modeling tasks, considerations of computational efficiency, and limited manual tuning. We set the output forecast window to 20 days to represent a difficult forecast task for the emulator.

We give the complete set of hyperparameters for our problem in Table 1. We found that the lowest validation errors were achived when just five modes were retained, a result that matches earlier studies [24], where increasing the number of modes is seen to cause difficulties in long-term forecasting (for example, for more than two weeks). The training and validation data are used in this phase of our experimentation to identify the best model for performing accelerated DA. We train the neural network architecture by penalizing the loss function on the training data, while using the validation data to enforce an early stopping criterion (i.e., prevent overfitting). Once the different models have been trained, the best model is determined by studying the validation losses of all the different architectures. This model is then used for testing on unseen data and for DA. For a further sensitivity analysis of this model, we trained four other models with varying input window sizes (7, 28, 35, 42) and different random seeds, with other hyperparameters fixed, and utilized them for obtaining ensemble test results. We trained and validated architectures using multiple hyperparameters in parallel by using Ray [34], a TCP/IP-based parallelization protocol, on Nvidia A100 GPUs of the Argonne Leadership Computing Facility’s Theta supercomputer.

Our first set of assessments test the POD-LSTM emulator without any DA. We show in Figure 3 how emulator predictions for day 15 (for all forecasts in the test region) compare to climatology. Climatology, here, refers to the average forecast on a specific forecast day based on averaged observations from the training data. For instance, if geopotential height is to be forecast on December 7, 1991, the climatology prediction would be the average of December 7 geopotential height values for 1984-1989. The metrics we use for comparison first include the Cosine similarity improvement as shown in Figure 3(a). Here, the cosine similarity (which captures the alignment between prediction and truth) obtained from climatology is subtracted from that obtained from the emulator forecast. Thus, in this case, negative (blue) regions are where climatology captures the temporal trend of the forecast better than our emulator. In Figure 3(b), we subtract the MAE of the emulator predictions from that of climatology. Here, the blue regions are where climatology is more accurate on average than the emulator. Figure 3(c) merely shows the MAE for the emulator. From this analysis, it is clear that for large regions in the data domain (particularly in the North), the emulator performs quite poorly.

We next assess results from applying DA through random observations at 5000 locations of the domain, to see if results obtained with the emulator can be improved. Here, sparse observations at each timestep of the output window are used within the optimization statement introduced in Equation 9 to update the initial conditions (i.e., the input window). We emphasize that the observations are obtained randomly from the true state of the system during forecasts, which corresponds to the test data introduced previously.

The results, in Figure 4, show that the use of sensor information aids the forecast immensely, improving performance in all metrics significantly by virtue of DA. In particular, the proposed augmentation (i.e., the use of DA during forecasts) reduces forecast errors considerably in regions where the sole use of the emulator was not competitive. We provide MAE assessments for our 20 day output averaged over the testing data range in Figure 5, where we compare the raw emulator outputs, climatology, persistence, and the DA-corrected outputs. We also provide confidence intervals for the regular (i.e., without DA) and DA-corrected emulators where the five different emulators are trained with different random seeds and input window durations to assess the sensitivity to the initialization of the training as well as the automatic differentiation based optimization. The results indicate that the DA-corrected emulation is the most accurate and consistent forecast technique.

We also provide in Figure 6 a comparison with another benchmark, persistence, which uses the state of the last day in the input window as the forecast for each day of the output window. Persistence is seen to be more accurate for short lead times but is outperformed by both climatology and data-driven methods for longer forecast durations.

Our choice of 20 days for the forecast window was motivated by the results of Rasp et al. [44, 45], who performed analyses for five-day forecasts. Our study increased the forecast window to four times the original assessments for geopotential height to demonstrate the possibility for improved forecasting of any emulator using automatic differentiation-enabled variational DA. In many subregions of the computational domain, the use of DA with the emulator is also seen to improve on climatology predictions—even at large lead times (20 days out), where classical data-driven and numerical forecast models are typically less accurate than climatology.

Our method also delivers large computational gains. A one year forecast with the emulator-assisted DA performed each day takes approximately one hour on a single processor core without any accelerator hardware. In contrast, the original PDE-based simulation required 21,600 core hours for a one-year forecast with a 515$\times$599 grid. (A coarse-grained forecast for one year on the 103$\times$120 grid used here takes 172 core hours, but the resulting flow field does not adequately reproduce the fidelity of the 515$\times$599 grid.) One can observe significant speed up ($\sim$10⁴ times) for the emulation of the geopotential height, even without factoring the cost of an additional variational DA step. Furthermore, the solution to the 4D-Var problem, which yields an improved initial condition, requires on the order of 100 model runs, where each model run can be several orders of magnitude more expensive than our emulator.