Combined Optimization of Dynamics and Assimilation with End-to-End Learning on Sparse Observations

TL;DR

CODA is presented, an end-to-end optimization scheme for jointly learning dynamics and DA directly from sparse and noisy observations that provides greater robustness to model misspecification than classical DA approaches.

Abstract

Fitting nonlinear dynamical models to sparse and noisy observations is fundamentally challenging. Identifying dynamics requires data assimilation (DA) to estimate system states, but DA requires an accurate dynamical model. To break this deadlock we present CODA, an end-to-end optimization scheme for jointly learning dynamics and DA directly from sparse and noisy observations. A neural network is trained to carry out data accurate, efficient and parallel-in-time DA, while free parameters of the dynamical system are simultaneously optimized. We carry out end-to-end learning directly on observation data, introducing a novel learning objective that combines unrolled auto-regressive dynamics with the data- and self-consistency terms of weak-constraint 4Dvar DA. By taking into account interactions between new and existing simulation components over multiple time steps, CODA can recover initial conditions, fit unknown dynamical parameters and learn neural network-based PDE terms to match both available observations and self-consistency constraints. In addition to facilitating end-to-end learning of dynamics and providing fast, amortized, non-sequential DA, CODA provides greater robustness to model misspecification than classical DA approaches.

Figures (7)

• Figure 1: Learning strategy through weak constrained neural network. a. example of observations used to compute the loss $\mathcal{L}_\text{coda}$ for a single element of a training batch; b. autoregressive rollout through the forward operator $\mathcal{M}_{\psi}$. The learned parametrization $\mathcal{B}_{\psi}$ is an additive correction defined by a neural network; c. training step for one batch element of observations.
• Figure 2: 1.5D Unet Architecture for 1D data assimilation applications.
• Figure 3: Comparison of CODA with baseline DA approaches. a. ground truth simulation generated using one level L96 simulator with $F=8$ and $dt=0.01$; b. observations generated from ground truth simulation applying Gaussian noise with standard deviation of $1.0$; c.-e analysis field using Optimal Interpolation (OI), Ensemble Kalman Smoother (EnKS) and CODA.; f. cumulative distribution of errors with respect to ground truth simulation (blue - OI, green - EnKS, red - CODA); g.-k. Forecast root-mean-square error started with estimated initial condition by OI, EnKS and CODA. Forecast RMSE mean (line) and standard deviation (shading) are computed for 1000 independent sets of observations generated from different initial conditions.
• Figure 4: Root mean square error of CODA analysis, a. different training setups including strong constraint ($\alpha=0$); b. generalization of best DA network trained on data with $75\%$ missing data and noise standard deviation $1.0$;
• Figure 5: Convergence of CODA-based parameter tuning task along training data assimilation network. a. 1 random initialization per task with various f; b.10 random initializations of parameter tuning task for $f=8$; c. root mean squared error of estimated initial state and loss terms: data mismatch, model error and total loss.