L EARNING D YNAMICAL C HARACTERISTICS WITH N EURAL O PERATORS FOR D ATA A SSIMILATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Data assimilation refers to a group of algorithms that combines numerical models with observations to obtain an optimal estimation of the system’s states. In areas like earth science, numerical models are usually formulated by differential equations, also known as the prior dynamics.
It is a great challenge for neural networks to properly exploit the dynamical characteristics for data assimilation, because ﬁrst, it is difﬁcult to represent complicated dynamical characteristics in neural networks, and second, the dynamics are likely to be biased.
The state-of-the-art neural networks borrow from the traditional method to introduce dynamical characteristics by optimizing the 4D-Var objective function in which the dynamics are inherently quantiﬁed, but the iterative optimization process leads to high computational cost.
In this paper, we develop a novel deep learning framework with neural operators for data assimilation.
The key novelty of our proposed approach is that we design a so-called ﬂow operator through self-supervised learning to explicitly learn dynamical characteristics for reconstructed states.
Numerical experiments on the Lorenz-63 and Lorenz-96 systems, which are the standard benchmarks for data assimilation performance evaluation, show that the proposed method is at least three times faster than state-of-the-art neural networks, and reduces the dynamic loss by two orders of magnitude.
It is also demonstrated that our method is well-adapted to biases in the prior dynamics.
I NTRODUCTION
Data assimilation can be essentially deﬁned as a statistical technique to combine the prior dynamics with a sequence of noisy and irregularly-sampled observations, which plays an important role in utilizing observational data, especially for numerical weather forecasting systems.
For instance, the
European Centre for Medium-Range Weather Forecasts (ECMWF) employs the variational assimilation algorithms in its operational systems to take full advantage of both in situ and satellite-derived data, as well as state-of-the-art numerical models (Rabier et al., 2000).
Classical data assimilation methods range from early empirical analysis (Bergthorsson et al., 1955) and optimal interpolation (Gandin, 1963) to later variationalbased assimilation algorithms (Sasaki, 1970) and the ﬁltering methods based on statistical estimation theory, such as Kalman ﬁlter (Welch et al., 1995), ensemble Kalman ﬁlter
(Evensen, 2003), and particle ﬁlter (Carpenter et al., 1999). Although classical methods have been essential tools for improving the capability of major numerical weather prediction centers worldwide (Kalnay, 2003), they generally do not account for the propagation of observation information along time and therefore are unable to utilize observations at different times effectively (Song et al., 2017). 4D-Var Assimilation Method
Four-dimensional variational assimilation algorithm is a cuttingedge method for data assimilation. Unlike previous assimilation algorithms, it explicitly takes into consideration the revolution of observation information along time, in other words, the timedependent dynamical characteristics (Song et al., 2017).
Experimental results prove the advantage of 4D-Var assimilation methods in utilizing observational information (Lorenc & Rawlins, 2005). However, due to the complicated modeling and solving process of the 4DVar algorithm, it is considered computationally expensive, especially for high resolution cases.
(Fisher, 1998). Machine Learning for data assimilation
Numerous studies developing machine learning for data assimilation have built their theories on classical assimilation methods to reduce computational cost, especially those prior to 2020 (Hsieh & Tang, 1998; Vijaykumar et al., 2002; de Campos Velho et al., 2002; H¨arter & de Campos Velho, 2008; Cintra & de Campos Velho, 2018; Mack et al., 2020). As mentioned above, these works are subject to the limitations of the classical methods that do not consider time-dependent dynamical characteristics.
Since 2020, the importance of 4DVar assimilation methods is gradually recognized, and much more novel and bolder research has been done in the area of machine learning for 4D-Var assimilation. These works can be generally divided into three categories: learning observation operator, learning model biases, and learning the optimization algorithm. Learning Observation Operator The aim of learning observation operator is to use neural network to construct a map from observations to reconstructed states.
A representative work is Frerix et al. (2021), in which the observational data is exploited by the neural operator to provide better initial states for the optimization of 4D-Var objective function. This work has at least two ﬂaws. The ﬁrst is that the neural network is used only to exploit observational information, while the integration of observation and dynamical characteristics is still implemented by 4D-Var, leading to a limited improvement in efﬁciency.
The second is that the strict constraints of time-dependent dynamical characteristics are the prerequisites for this work, which prevent it from generalizing to the cases in which the prior dynamics are biased.
Learning Model Biases
As mentioned above, the prior dynamics can be biased, and another group of works uses neural network to learn the biases of prior dynamic models for data assimilation Brajard et al.
; Arcucci et al. ; Farchi et al. ; Bonavita & Laloyaux (2020).
These works generally achieve a balance between accuracy and efﬁciency, but the improvement in computing efﬁciency is insigniﬁcant, because the assimilation of dynamical characteristics and observations is still done through iteratively optimizing 4D-Var objective function. Learning Optimization Algorithm Since the main computational overhead of the 4D-Var algorithm lies in the process of iteratively optimizing 4D-Var objective function, some studies have focused on replacing traditional optimization algorithms with neural networks. Among them, the representative one is Fablet et al.
(2021b;a). This work simultaneously learns both the dynamical characteristics of the model and the optimization algorithm of the 4D-Var objective function through neural networks.
It not only takes into consideration the weak constraints of time-dependent dynamical characteristics, but also successfully reduces the number of iterations compared with traditional 4D-Var algorithms, subsequently improving computational efﬁciency. We consider this work to be the state-of-the-art method for comparison. Despite the achievements of the work, there is still room for improvement in terms of both accuracy and efﬁciency, and moreover, quantitative exploration on how the prior bias affects the performance of neural network frameworks remains to be done. Our Contributions
We develop a new approach of combing data-driven strategies with modeldriven methods for time-dependent data assimilation. This framework consists of three operators implemented by neural networks, namely, observation operator, perturbator, and ﬂow operator. Among them, the cooperation of the perturbator and ﬂow operator decouples the two goals of learning observations and learning dynamical characteristics, and they succeed in continuously reﬁning reconstructed states.
• To our knowledge, this is the ﬁrst work that directly uses neural network to blend the dynamical prior with a sequence of observations without explicit formulations of 4D Var objective function for assimilation. • We test the proposed framework on the Lorenz-63 and Lorenz-96 systems. Experimental results support the effectiveness of our framework.
The proposed architecture is at least three times faster than the state-of-the-art neural network. It reduces the dynamic loss by two orders of magnitude and is comparable to the state-of-the-art neural networks in terms of reconstruction error. • We design experiments in which the prior dynamics are different from ground truth to simulate real-world scenarios.
Experimental results prove that our method is well-adapted to deviations of the prior dynamic model.
P ROBLEM S TATEMENT
We refer to Fablet et al. (2021b) to introduce the strict deﬁnition and assumptions under considered problem settings.
Denote x for the time-dependent state in some space X satisfying x ( t ) ∈ X , M gt [Figure] for the ground truth dynamic model, and then the dynamic system can be formulated as
[Equation]
In real-world scenarios, the true dynamic model of earth systems is unknown and meteorologists develop physical models to estimate it, based on the laws of kinetics and thermodynamics.
Denote M pr for the prior model.
In the ideal case, the bias of prior model should be bounded; that is, there exits a constant C , such that,
[Equation]
Observations { y ( t i ) } are assumed to be discrete and partial, with white Gaussian noise.
Suppose that all the observations are obtained at time t i =
t
+ i ∆ t , i = 0 , 1 , ..., N and at time t i , only components in domain Ω t i are observed.
The formulation can be written as
[Equation]
where (cid:15) ( t ) represents the white Gaussian noise.
The problem setting can be further explained within a deep learning paradigm.
Our main goal is to construct an inverse map, implemented in the form of a neural network, from a time sequence of noisy and partial observations y to its corresponding reconstructed states x .
The a priori knowledge is that the reconstructed states approximately satisfy certain dynamic equations with a bounded error.
The major challenge is that it is difﬁcult to incorporate the prior knowledge into neural network structures.
Most of the previous work has relied on the minimization of the 4D-Var target function to tackle it.
In our work, a pure neural network solution is given.
L EARNING F RAMEWORK WITH T HREE O PERATORS
To deal with the problem, we propose an end-to-end learning framework with three operators, including observation operator, ﬂow operator and perturbator, as shown in Fig 1.
3. O BSERVATION O PERATOR
The observation operator aims to learn lower-dimensional observations and convert them into physical states (Frerix et al., 2021).
We design the structure of the observation operator based on the
Residual U-Net structure (Ronneberger et al., 2015; Siddique et al., 2021).
The hierarchical structure of U-Net allows it to capture both the large-scale and small-scale information of observation c h h 2h 2h h+h h h c 2h 2h+2h 4h
KxT
KxT h
KxT o
Obs Time 4h+h
KxT o
Conv 3x3 + SiLU
Conv 3x
MaxPool
TransConv
Concat
Modiﬁed U-Net Structure
In addition, this paper also modiﬁes the network according to the speciﬁcity of the data assimilation issue, as illustrated below. Unlike images, of which the three dimensions are length, height, and channels, the inputs of data assimilation are sequences of observations (states), whose dimensions are time, location, and channels. The length and height of images are symmetric, while time and location are generally not.
To tackle with this difference, we carry out dimensionality reductions in U-Net for data assimilation in the temporal dimension.
What is more, our proposed observation operator adds a skip layer from the input to the bottom hidden variables. It is implemented by extracting valid observations in the observed moments from the input and concatenating them directly with the bottom hidden variables. The advantage of this operation is that the information of observations can be more effectively transferred through the network. The number of layers for U-Net can be set to three (Figure 2), four, or even more, depending on the sparsity of observations.
The reconstruction loss L rec is applied for the training of the observation operator.
If we denote the network by Ψ and the output by ˆ x = Ψ( y ) , the loss function is
[Equation]
With the help of the observation operator, we are already able to construct physical states, but with relatively low accuracy. Instead of seeking help from traditional 4D-Var algorithms to further improve accuracy Frerix et al. (2021), we construct the so-called perturbator and ﬂow operator for such purposes.
3. F LOW O PERATOR Dynamic Loss
Before proceeding to the introduction of the ﬂow operator, let us ﬁrst illustrate an important concept, dynamic loss. The dynamic loss measures how well the given sequence of states ﬁts the considered dynamics, which is formulated by
[Equation]
[Equation]
where Φ is deﬁned by
[Equation]
Φ is a ﬂow-dependent operator, which describes the hypothetical result of ﬁrst moving the sequence one step backward and then integrating one step forward. The smaller the dynamic loss, the better the sequence of states ﬁts the dynamics.
If the sequence of states x perfectly matches the dynamic system, then Φ( x ) = x holds, and the dynamic loss is zero.
Function of Flow Operator
The major function of the ﬂow operator is to restore the reconstructed sequence of states with a relatively high dynamic loss to the corresponding sequence that better conforms to the prior dynamics.
Suppose that x ( t ) ∈ X is a sequence of reconstructed system’s states, we require that the dynamic loss after the operator is smaller than a given value L 0 , which is usually orders of magnitude smaller than the original dynamic loss. Denote the ﬂow operator by F , and the requirement can be formulated as [Equation]
Network Structure
Instead of using the neural network to directly compute the improved output, we use the U-Net structure to compute the increment compared to the input state; that is,
[Equation]
This subtle modiﬁcation is to better ensure that the amount of modiﬁcation is small and make the training process converge faster.
Network Training and Stopping Criterion
We train the ﬂow operator with the dynamic loss, and the training process is self-supervised because no additional label is required.
For the case in which the prior dynamics are unbiased, we simply ﬁx the number of epochs in advance and do the training.
In contrast, if the prior dynamic model is inaccurate, we stipulate a stopping criterion for the training procedure:
[Equation]
where E x gt [ L dyn − pr ( x gt )] is the expectation of the prior dynamic loss of ground truth states, which measures the difference between the prior dynamics and ground truth.
Because the error of prior dynamics is bounded, E x gt [ L dyn − pr ( x gt )] is also a bounded value smaller than inﬁnity.
As long as inequality 9 is not satisﬁed, we end the current training phase and move on to the next one.
Inequality 9 guarantees that the dynamic loss of the reconstructed states is no smaller than that of the true states.
It is a necessary but not sufﬁcient condition to avoid overﬁtting in terms of dynamics, and any reconstructed states with a smaller dynamic loss than true states indicate that more information than expected has been learned from the prior dynamics.
3. P ERTURBATOR Function of Perturbator
The role of the ﬂow operator is to guarantee that the reconstructed ﬁeld conforms to the dynamics, but this is insufﬁcient for data assimilation, because there are inﬁnitely many sequences of states that conform to the dynamics, and we need to ﬁnd the one that also ﬁts the observations well. On the other hand, thanks to the ﬂow operator’s ability to guarantee that the output is consistent with the prior dynamics, we can add a new neural network, called the perturbator, in front, which has more free space to learn from observations and labels. It can use the learned knowledge to revise the reconstruction states without considering the problem of conforming to the prior dynamics.
In this way, we decouple the two goals of conforming to observations and conforming to prior dynamics, and both the ﬂow operator and perturbator can perform their own duties. Speciﬁcally, The perturbator uses the observations and labels it has learned to perturb the reconstructed states to make it deviate from the original ﬂow, and then the ﬂow operator corrects the deviated ﬂow to make it a new ﬂow that conforms to the prior dynamics. Cooperation with Flow Operator
The perturbator and ﬂow operator always work in pairs. Given a sequence of states with relatively low dynamic loss L dyn and relatively high reconstruction loss L rec , assuming that we have trained a ﬂow operator F , the function of the perturbator P can be Observation
Operator
Input
Output
[Figure]
Operator
[Figure]
Flow
Operator
Output
Pretrain
Pretrain
Observation
Operator
Input
Flow
[Figure]
Flow
[Figure]
Finetune
Finetune
Observation Operator Input Flow
[Figure]
Flow
Operator
×n ×n
We adopt this seemingly complicated training approach because these operators must be individually trained with parameters of other operators ﬁxed to perform as we expect.
formulated as
[Equation] The ﬁrst formulation requires that the reconstruction loss is reduced, and the second requires that the dynamic loss should remain in the same order of magnitude. Ideally, as more pairs of perturbators and ﬂow operators are applied to the input, the reconstruction loss will continuously decrease while maintaining the dynamic loss. We emphasize that the output of the perturbator P ( x ) is likely to have both high reconstruction loss and high dynamic loss, but it does not matter because we only require that the losses after the ﬂow operator F ( P ( x )) are small.
A typical U-Net structure is adopted to build the perturbator, and in this case we use a hybrid loss to learn from both the dynamics and labels. [Equation]
The training process of the proposed framework can be divided into two parts: pretraining and ﬁnetuning. During pretraining, network parameters of the three operators in the framework are roughly determined, while the purpose of ﬁnetuning is to further adjust and modify the parameters to ﬁt the recurrent structure. We let readers refer to Figure 3 for details.
N UMERICAL E XPERIMENTS
In this section, numerical experiments are reported for Lorenz-63 and Lorenz-96 systems.
Owing to their chaotic properties and similarities with meteorological systems, Lorenz systems are widely used for evaluation of various data assimilation algorithms, especially in data-driven and machine learning case-studies (Lguensat et al., 2017; Raissi, 2018; Fablet et al., 2021b;
Frerix et al., 2021).
[Table] 4. D ESIGN OF E XPERIMENTS
The Lorenz-63 and Lorenz-96 dynamic systems are governed by the ordinary differential equations (ODE) expressed as Equation 12 and Equation 13, respectively (Lorenz, 1963; 1996).
Both systems are one-dimensional. We further design two experiments to test the performance of our framework in different scenarios.
We refer to the previous work for the selection of parameterization (Fablet et al., 2021b), where for Lorenz-63, σ = 10 , ρ = 28 , β = 8 / 3 and for Lorenz-96, N = 40 , F = 8 .
We simulate time series with a 0.01 time step for Lorenz-63 and a 0.05 time step for Lorenz-96, using a RK45 integration scheme (Dormand & Prince, 1980).
The assimilation window length is 200 time steps.
For the Lorenz-63 system, the observation data is sampled every eight steps for only the ﬁrst component, and for the Lorenz-96 system, it is sampled every four steps for random half of the components.
In both cases, a Gaussian noise with a variance of 2.0 is involved.
[Equation]
For the Lorenz-63 system, We adopt a 4-layer U-Net architecture for the construction of the observation operator, perturbator, and ﬂow operator, where c = 3 , K = 1 , T = 200 , and h = 32 .
For the
Lorenz-96 system, we adopt a 3-layer architecture, where c = 1 , K = 40 , T = 200 , and h = 32 . 4. U NBIASED P RIOR D YNAMICS
First, we consider the case in which the prior dynamics are identical to the ground truth dynamic, that is M gt =
M pr .
We use a total of eight and ten blocks to train the deep learning framework for the Lorenz-63 and Lorenz-96 systems, respectively, including an observation operator and a ﬂow operator in the front, followed by pairs of perturbators and ﬂow operators, as shown in Figure 1 with
N = 3 and N = 4 .
In Table 1, we compare our ﬁnetuned neural network structure with the state-of-the-art networks (Fablet et al., 2021b) in terms of R-score, LossDyn, and time-to-solution on test set. Our proposed [Figure] [Figure]
method achieves better reconstruction performance than the state-of-the-art method on the metric of R-score.
We reduce the reconstruction error by 44% and 3% for Lorenz-63 and Lorenz-96 systems, respectively, and reduce the dynamic loss by two orders of magnitude, which are signiﬁcant achievements.
Moreover, the dynamic loss is at least one order of magnitude smaller than the stateof-the-art, regardless of the number of blocks.
In terms of time-to-solution, our proposed framework is four times and three times faster than the state-of-the-art method for the two dynamic systems.
We select an example from the test set in Figure 4 to further illustrate how our architecture works.
Focus on the blue oval box area in the upper right corner, where we ﬁnd how the cooperation of perturbator and ﬂow operator succeeds in correcting the initial states.
In the beginning, the states generated by two blocks are clearly inconsistent with the ground truth in the upper right region.
After a perturbator is added (block 3), disturbances with high frequencies are involved in the misaligned area to distinguish them from well-constructed areas.
If a ﬂow operator is further added behind the perturbator (block 4), the high-frequency noise disappears, and the inconsistency is corrected to some extent.
Such a perturbation-reﬁnement step is repeated (block 5 & block 6) and the misalignment becomes smaller with repetition.
The lower-left panel also shows a continuous decrease in the reconstruction error, from two to four to six blocks (orange line). 4. B IASED P RIOR D YNAMICS
In the following experiment, biases are introduced to the prior dynamics by perturbing the ODE parameters.
In the case of Lorenz-63 systems, σ in Equation 12 is selected to be the perturbed parameter, and in the case of Lorenz-96 system, F in Equation 13 is selected.
Speciﬁcally, we select ∆ σ from the set { 0.25, 0.50, 0.75, 1.00, 1.50, 2.00, 2.50, 3.00, 3.50, 4.00, 4.50, 5.00 } , select ∆ F from the set { 0.50, 1.00, 1.50, 2.00, 2.50, 3.00, 3.50, 4.00, 4.50, 5.00 } , and conduct numerical experiments accordingly.
The neural network structures and training losses are almost identical to those in the experiment with no bias, except that we use the prior dynamics rather than ground truth dynamics to construct the dynamic loss.
We report our results in Figure 5.
Take the Lorenz-63 system as an example. As illustrated in the upper-right ﬁgure, when ∆ σ is smaller than 1.0, the increase in R-score is smaller than 20%, which shows that our framework is robust to small biases.
When ∆ σ is smaller than 2.00, the Rscore increases as ∆ σ increases, whereas if ∆ σ is larger than 2.00, the R-score ﬂuctuates without monotonic changes.
Such a phenomenon can be explained by looking into how our network structure works: When ∆ σ is small, the difference between the prior dynamics and ground truth is small, and the information of the prior dynamics can be successfully exploited by the ﬂow operator.
Therefore, as ∆ σ gets larger, useful information provided by the prior dynamics decreases and the performance of our framework worsens.
However, if ∆ σ is too large, it is no longer necessary for the ﬂow operator to learn from the bad prior dynamics; it is cheaper to learn directly from ground truth labels.
Hence, in this case, the value of ∆ σ does not affect the performance of the neural network.
The results shown in the upper-left ﬁgure can further conﬁrm our explanation.
When ∆ σ < 2 , the difference between L dyn − pr (ˆ x ) and L dyn − pr ( x gt ) is close to zero, which is consistent with our stopping criterion (Inequality 9), so the network works as expected.
However, if ∆ σ > 2 , the bias of prior dynamics is too large for the ﬂow operator to learn, and solely learning from the ground truth is sufﬁcient for the neural network to generate states whose dynamic loss is smaller than L dyn − pr ( x gt ) .
The ﬁndings above also holds for the Lorenz-96 system, as shown in the lower ﬁgures.
The difference is that the threshold changes from 2.0 to around 2.5.
They provide further evidence that our network learns features from prior dynamics rather than simply remembering ground-truth labels, and that our framework works properly despite biases in prior dynamics.
C ONCLUSIONS
We have introduced a deep learning framework for data assimilation.
The key novelty of our research is that we develop the ﬂow operator to learn from the prior dynamics and continuously reﬁne the reconstructed states.
The results of numerical experiments on Lorenz-63 and Lorenz-96 systems support the relevance of our framework.
In terms of efﬁciency
, our framework has been found to be four times and three times faster than state-of-the-art networks on Lorenz-63 and Lorenzsystems, respectively. In terms of assimilation accuracy, our framework has reduced the dynamic loss by two orders of magnitude, and at least maintained the reconstruction accuracy, compared with the state-of-the-art method.
It has also been found to be well-adapted to biases in prior dynamics.
The results indicate that the proposed method can serve as a new method for integrating ODE information into neural networks.
They suggest that ODE information can be applied to modify coarse reconstruction states instead of being explicitly embedded into neural network structures.
D ATASET C ONSTRUCTION
We do a long integration over time and divide the time sequence into three parts. The ﬁrst part is for building the training set, the second is for the validation set, and the last is for the test set, as shown in Figure 6. For each part, we randomly sample time series of the same length as the assimilation window multiple times, and then use the sampled series as the dataset. [Figure]
Speciﬁcally, for the Lorenz-63 system, the training data comprises 10000 sequences sampled from time series with 12000 steps; the validation data comprises 2000 sequences from time series with 5000 steps and the test data comprises with 3000 sequences from 10000 time steps.
For the Lorenz96 system, the training data comprises 2000 sequences sampled from time series with 7500 steps; the validation data comprises 300 sequences from time series with 2500 steps and the test data comprises with 200 sequences from 10000 time steps. [Figure] [Figure] Reconstructed examples and associated absolute errors
[Figure] A. L OSS C HANGE D URING T RAINING The variations of reconstruction and dynamic loss against the number of epochs for Lorenz-63 and Lorenz-96 systems are reported in Figure 7 and Figure 8, respectively.
During pretraining, the observation operator learns a map from the observational space to the physical states at ﬁrst.
Since a shift between two different spaces takes place in this period, we see a steep decrease of loss within the ﬁrst 20 epochs.
After the pretraining of observation operator is ﬁnished, we start the pretraining of ﬂow operator to correct the ﬂow dynamics of reconstructed states.
The pretraining of the ﬂow operator is relatively slow because a self-supervised loss is employed in this procedure.
We claim that the ﬂow operator is indispensable in further decreasing the reconstruction loss and dynamic loss, for the reason that without the ﬂow operator, the observation operator would quickly lead to overﬁtting, in other words, the difference between the training loss and validation loss will diverge, as shown in the red and blue dashed lines in the ﬁgure.
During the ﬁnetuning procedure, the reconstruction loss drops steadily while the dynamic loss ﬂuctuates.
This is not surprising considering that the ﬁnetuning procedure alternates between the perturbator and ﬂow operator.
If the training object is the perturbator, the dynamic loss is relatively high because the optimized target is the hybrid loss; if the training object is ﬂow operator, the dynamic loss is relatively low because the optimized target is just the dynamic loss.
Such alternating ﬁnetuning process contributes to the continuous decrease of the R-score.
A. S AMPLE S TUDY OF THE L ORENZ -96 S YSTEM
Figure 9 demonstrates the physical states reconstructed by different neural networks and their associated absolute error maps.
Visually, the dark regions in the error map becomes smaller and lighter as the number of block grows, which conﬁrms the effects of the collaboration of the perturbator and ﬂow operator in this example.