Combining machine learning and data assimilation to forecast dynamical systems from noisy partial observations

Recent years have seen an increased interest in data-driven methods with the aim to develop cheap surrogate models to perform forecasting of dynamical systems. A main driver behind this endeavour is the potential saving of computational running time required for simulation. This is particularly important for stiff multi-scale systems for which the fastest time-scale puts strong restrictions on the time step which an be employed (Han et al., 2018, 2019, Raissi and Karniadakis, 2018, Raissi et al., 2019). Promising applications for such surrogate models come from atmospheric and ocean dynamics and from climate dynamics (Schneider et al., 2017, Dueben and Bauer, 2018, Rasp et al., 2018, Gagne II et al., 2020, Bolton and Zanna, 2019, Brajard et al., 2021, Nadiga, 2020, Cleary et al., 2021).

A particularly simple and efficient machine learning architecture are random feature maps (Rahimi and Recht, 2008, 2008, Bach, 2017a, b, Sun et al., 2019). Random feature maps provide a representation of the surrogate propagator for a dynamical system by a linear combination of randomly generated high-dimensional nonlinear functions of the input data. The training of random feature map networks only requires linear least-square regression and it was proven rigorously that random feature maps enjoy the so called universal approximation property which states that they can approximate any continuous function arbitrarily close (Park and Sandberg, 1991, Cybenko, 1989, Barron, 1993). The framework of random feature maps was extended to include internal dynamics in so called echo-state networks and reservoir computers with remarkable success in forecasting dynamical systems (Maass et al., 2002, Jaeger, 2002, Jaeger and Haas, 2004, Pathak et al., 2018a, Jüngling et al., 2019, Algar et al., 2019, Nadiga, 2020, Bollt, 2021, Gauthier et al., 2021, Platt et al., 2021).

These linear regression based training methods assume model errors rather than measurement errors in the data (Gottwald and Reich, 2021). However, the skill of random feature maps in providing a reliable surrogate forecast model can be severely impeded when the data used for training are contaminated by measurement noise. In previous work Gottwald and Reich (2021) showed that by combining random feature maps with ensemble-based sequential data assimilation, in a framework coined RAndom Feature map Data Assimilation (RAFDA), the noise can be effectively controlled leading to remarkable forecast skills for noisy observations. Moreover, RAFDA was shown to be able to provide model closure in multi-scale systems allowing for subgridscale parametrization as well as to provide reliable ensembles for probabilistic forecasting. In RAFDA the parameters of the random feature maps are learned sequentially with incoming observations in conjunction with the analysis of the state variables using an ensemble Kalman filter (EnKF) (Evensen, 2006). RAFDA thus combines a powerful approximation tool in the form of random feature maps with an advanced sequential and derivative-free data assimilation technique, which is particularly well suited for high-dimensional state and parameter estimation problems.

Several alternative combinations of data assimilation and machine learning techniques have recently been explored. Recently proposed methodologies include purely minimization-based approaches (Bocquet et al., 2019), approaches which combine ensemble-based sequential state estimation with minimization-based parameter estimation (Brajard et al., 2020, Bocquet et al., 2020), and purely ensemble-based sequential joined state and parameter estimation (Bocquet et al., 2021). Furthermore, an intimate formal equivalence between certain neural network architectures and data assimilation has been noted by Abarbanel et al. (2018). Finally, sequential data assimilation has also been extended to echo-state networks and reservoir computers (Tomizawa and Sawada, 2021, Wikner et al., 2021).

In this work, we extend RAFDA to the situation when only partial observations are available. In this case, the propagator map, that updates the system in the partially observed subspace only, is non-Markovian and knowledge of the past history is required. This problem was recently addressed within a machine learning context using appropriate model closures (Levine and Stuart, 2021). To account for the non-Markovian nature of the propagator map, we follow here a different path instead and employ phase-space reconstruction and Takens’ embedding theorem (Takens, 1981). Takens’ embedding theorem has been successfully used in data assimilation before where a forecast model was non-parametrically constructed using analogs (Hamilton et al., 2016). Recently reservoir computers, which may be viewed as providing an embedding of the dynamics in the reservoir space (Lu et al., 2017, Hart et al., 2020), were used to reconstruct dynamics in the reconstructed phase space using delay coordinates and linear regression for parameter estimation (Nakai and Saiki, 2021). Here we use delay coordinates in combination with sequential data assimilation to learn a surrogate forecast model directly in the reconstructed phase-space from noisy and partial observations.

While we follow a non-parametric or model-free approach in this paper, another promising direction is to combine machine learning and knowledge-based modelling components (Pathak et al., 2018b, Gagne II et al., 2020, Bolton and Zanna, 2019, Wikner et al., 2020, Brajard et al., 2021, Wikner et al., 2021). We note that RAFDA also allows for such extensions (Gottwald and Reich, 2021).

The paper is organised as follows. In Section 2 we develop our RAFDA methodology for partial observations. In Section 3 we show how RAFDA performs on the Lorenz 63 system and that RAFDA exhibits significantly increased forecast skill when compared to standard random feature maps trained on delay coordinates using linear regression. We close in Section 4 with a discussion and an outlook.

Consider a $D$-dimensional dynamical system $\dot{{\bm{u}}}=\mathcal{\bm{F}}({\bm{u}})$ which is accessed at equidistant times $t_{n}=n\Delta t$ of interval length $\Delta t>0$, $n\geq 0$, by partial noisy observations

with ${\bm{u}}_{n}={\bm{u}}(t_{n})$, observation operator $\bm{G}:{\mathbb{R}}^{D}\to{\mathbb{R}}^{d}$, measurement error covariance matrix $\bm{\mathrm{\Gamma}}\in{\mathbb{R}}^{d\times d}$ and $d$-dimensional independent and normally distributed noise $\bm{\eta}_{n}$, that is, $\bm{\eta}_{n}\sim{\mathcal{N}}({\bf 0},{\bf I})$.

The aim we set out to pursue is the following: using noisy observations ${\bm{y}}^{\rm{o}}_{j}$ for $0\leq j\leq n$ find an approximation to the propagator map which maps previous observations (or a suitable subset of them) to the unobserved variable ${\bm{y}}_{n+1}$ at future time $t_{n+1}$. In the Markovian case $d=D$ with $\bm{G}=\bm{I}$ the propagator map is only a function of the observation at the current time $t_{n}$. This case was studied in our previous work (Gottwald and Reich, 2021). Here the focus is on partial observations with $d<D$ for which the propagator map is non-Markovian. Our method judiciously incorporates three separate methods which we discuss in the following three subsections: 1.) Takens embedding and delay coordinates to deal with the aspect of having only access to partial observations and the resulting non-Markovian propagator map, 2.) random feature maps which form our approximation for the propagator map and 3.) data assimilation and EnKFs to control the observation noise and to determine the parameters of the random feature maps leading to the desired surrogate model.

We consider the Lorenz-63 system (Lorenz, 1963)

with ${\bm{u}}=(x,y,z)^{\rm T}\in\mathbb{R}^{3}$. Observations are only available for $x$, i.e. ${\bm{G}}=(1\;0\;0)$ and $d=1$ in (1), and are taken every $\Delta t=0.02$ time units. We employ observational noise with measurement error covariance $\bm{\mathrm{\Gamma}}=\eta{\bf I}$ and use $\eta=0.2$ unless stated otherwise (compare (1)). We ensure that the dynamics evolves on the attractor by discarding an initial transient of $40$ model time units. The optimal embedding dimension and delay time for these parameters are estimated as $m=3$ and $\tau=10$. In the following times are measured in units of the Lyapunov time $t\lambda_{\rm{max}}$ with the maximal Lyapunov exponent $\lambda_{\rm{max}}=0.91$.

We use a reservoir of size $D_{r}=300$ with internal parameters $w=0.005$ and $b=4$, and an ensemble size of $M=300$ for the ensemble Kalman filter. The regularization parameter is set to $\beta=2\times 10^{-5}$ and the inflation parameter to $\alpha=1+0.01\Delta t=1.0002$. We employ a training set of length $N=4,000$.

To test the propensity of RAFDA to learn a surrogate model (5) we generate a validation data set $\bm{\zeta}_{\rm valid}(t_{n})$, $n\geq 0$ sampled with the same rate as the training data set with $\Delta t=0.02$ from the $x$-component of the Lorenz-63 system (26). To quantify the forecast skill we measure the forecast time $\tau_{f}$, defined as the largest time such that the relative forecast error $\mathcal{E}(t_{n})=||\bm{\zeta}_{\rm valid}(t_{n})-\bm{\zeta}_{n}||^{2}/||\bm{\zeta}_{\rm valid}(t_{n})||^{2}\leq\theta$, where the $\bm{\zeta}_{n}$ are generated by the learned surrogate model. We choose here $\theta=40$.

The forecast skill depends crucially on the choice of the randomly chosen internal parameters $({\bm{W}_{\rm{in}}},{\bm{b}_{\rm{in}}})$ as well as on the training and validation data set. We therefore report on the mean behaviour over $500$ realisations, differing in the training and validation data set, in the random draws of the internal parameters $({\bm{W}_{\rm{in}}},{\bm{b}_{\rm{in}}})$ and in the initial ensembles for RAFDA. Each training and validation data set is generated from randomly drawn initial conditions which are evolved independently over $40$ model time units to ensure that the dynamics has settled on the attractor.

In the following, we compare RAFDA with standard random feature maps using linear regression (LR). Figure 1 shows the empirical histogram for both RAFDA and LR when the surrogate model (5) is trained on data contaminated with noise of strength $\eta=0.2$. The mean forecast time obtained by RAFDA is $\tau_{f}=2.12$ and is roughly three times larger than the forecast time obtained using standard random feature maps which yields $\tau_{f}=0.77$. Furthermore, it is seen that the distribution of forecast times for RAFDA is heavily skewed towards larger forecast times. As expected, these forecast times are smaller than for the fully observed case with ${\bm{G}}={\bf I}$, as considered in Gottwald and Reich (2021), where extreme values of $\tau_{f}=9.28$ were reported. Figure 2 shows a typical example of $x(t)$ for a forecast time of $\tau_{f}=2.6$, which is close to the mean forecast time. Figure 3 shows that the surrogate model (5) obtained using RAFDA produces a model which very well approximates the long-time statistics of the full Lorenz system (26) in the sense that its trajectories reproduce the attractor in delay-coordinate space. Contrary, LR does not lead to surrogate models which are consistent with the dynamics of the Lorenz system (26). Reproducing dynamically consistent trajectories is, of course, paramount for the purpose of using such surrogate models for long-time integration when the interest is rather on the overall statistical behaviour rather than on accurate short-term forecasts.

The ability of RAFDA to control noise present in a training data set depends on the strength $\eta$ of the noise. We show in Figure 4 results of the mean forecast time for a range of $\eta\in[6.25\times 10^{-6},4.3\times 10^{5}]$. RAFDA outperforms standard random feature maps with linear regression for noise levels $\log\eta<5$. Moreover, there is a robust plateau with mean forecast times of around $\tau_{f}\approx 2$ for a large range of noise strengths with $\eta\lessapprox 1$. Notice that for $\eta\to 0$, maybe surprisingly, the mean forecast time of RAFDA does not converge to the one obtained by standard random feature maps but remains threefold larger. This was discussed in Gottwald and Reich (2021): although LR achieves a smaller value of the cost function (2.2) for $\eta\ll 1$ it does not generalise as well to unseen data. The cost function is not zero as generically the output data do not lie in the span of the random feature maps. It is important to note that LR with $\beta>0$ assumes model error (which is not present in our application of simulating the Lorenz-63 system (26) rather than observational error. We further note that, for very large noise levels $\log\eta>5$, the forecast times $\tau_{f}$ become zero for RAFDA. This is due to the data assimilation component of our method and is an instance of the well-documented problem of filter divergence (Ehrendorfer, 2007, Nadiga et al., 2013). More precisely, in finite-size ensembles, most ensemble members may align with the most unstable direction (Ng et al., 2011) implying a small forecast error covariance. In combination with large observation noise this leads to the filter trusting its own forecast and the analysis is not corrected by incoming new observations. Filter divergence can be avoided by increasing the ensemble size and/or employing covariance inflation. We remark that we abstained here from fine-tuning all hyperparameters.

We proposed a new data-driven method to estimate surrogate one-step forecast maps from noisy partial observations. Our method determines the surrogate map in the reconstructed phase space for delay vector coordinates, thereby dealing with the problem of having only access to partial observations. We employed the RAFDA framework in the reconstructed phase space, in which the surrogate map is constructed as a linear combination of random feature maps the coefficients of which are learned sequentially using a stochastic EnKF. We showed that learning the coefficients of the linear combination of random feature maps sequentially with incoming new observations rather than by learning them using the method of least-squares on the whole data set greatly increases the forecast capability of the surrogate model and significantly controls the measurement noise.

Our numerical results showed that the quality of the surrogate model exhibits some degree of variance, depending on the draw of the arbitrarily fixed internal parameters of the random feature map model ${\bm{W}_{\rm{in}}}$ and ${\bm{b}_{\rm{in}}}$. Being able to choose good candidates for those parameters or learning them in conjunction with the rest of the surrogate model would greatly improve the applicability of the method. In our previous work we provided some guidance on what constitutes good hyperparameters enabling good learning (Gottwald and Reich, 2021). How to choose actual optimal values for these hyperparameters requires, for example, an optimization procedure wrapped around the RAFDA method described here, and is planned for future research. Apart from ${\bm{W}_{\rm{in}}}$ and ${\bm{b}_{\rm{in}}}$ there are numerous hyperparameters which require tuning, such as the regularization parameter $\beta$, the inflation parameter $\alpha$, the size and the variance of the initial ensemble, as well as the assumed observational error covariance $\Gamma$. To obtain optimal performance of RAFD these would need to be tuned which can be done by additional optimization procedures (Nadiga et al., 2019).

The random feature map architecture is closely related to reservoir computers (Jaeger, 2002, Jaeger and Haas, 2004, Pathak et al., 2018a, Jüngling et al., 2019). In reservoir computers the incoming signal is used in conjunction with an internal reservoir dynamics to produce the output. It is argued that this helps to take into account non-trivial memory of the underlying dynamical system. It will be interesting to see if random feature maps in the space of delay vectors performs as well as reservoir computers. For systems with a fast decay of correlation such as the Lorenz 63 system the additional reservoir dynamics did not lead to an improvement of the surrogate map (Gottwald and Reich, 2021), but incorporating memory for systems with slow decay of correlation taking into account information from past observations may be important.

SR is supported by Deutsche Forschungsgemeinschaft (DFG) - Project-ID 318763901 - SFB1294