Neural ensemble Kalman filter: Data assimilation for compressible flows with shocks
Authors: Xu-Hui Zhou, Lorenzo Beronilla, Michael K. Sleeman, Hangchuan Hu, Matthias Morzfeld, Andrew M. Stuart, Tamer A. Zaki
Organizations: Scripps Institution of Oceanography, University of California San Diego, La Jolla, 92093, CA, USA · Computing + Mathematical Sciences, California Institute of Technology, Pasadena, 91125, CA, USA · Department of Mechanical Engineering, Johns Hopkins University, Baltimore, 21218, MD, USA
Data assimilation (DA) for compressible flows with shocks is challenging because many classical DA methods generate spurious oscillations and nonphysical features near uncertain shocks. We focus here on the ensemble Kalman filter (EnKF). We show that the poor performance of the EnKF may be attributed to the bimodal forecast distribution that can arise in the vicinity of an uncertain shock location; this violates the assumptions underpinning the EnKF, which assume a forecast which is close to Gaussian. To address this issue we introduce the new neural EnKF. The basic idea is to systematically embed neural function approximations within ensemble DA by mapping the forecast ensemble of shocked flows to the parameter space (weights and biases) of a deep neural network (NN) and to subsequently perform DA in that space. The nonlinear mapping encodes sharp and smooth flow features in an ensemble of NN parameters. Neural EnKF updates are therefore well-behaved only if the NN parameters vary smoothly within the neural representation of the forecast ensemble. We show that such a smooth variation of network parameters can be enforced via physics-informed transfer learning, and demonstrate that in so-doing the neural EnKF avoids the spurious oscillations and nonphysical features that plague the EnKF. The applicability of the neural EnKF is demonstrated through a series of systematic numerical experiments with the inviscid Burgers' equation, the Sod shock tube, and a two-dimensional blast wave.
Figures & tables
Figure 1: Behavior of the EnKF for a shock-like transition, illustrated using a hyperbolic-tangent example: (a) forecast ensemble with 50 members; (b) analysis ensemble obtained using the EnKF, showing spurious oscillations near the shock; and (c) marginal probability density functions (PDFs) of the forecast ensemble at two representative locations indicated in (a), showing an approximately Gaussian distribution away from the shock ( x=−0.5 ) and a bimodal distribution near the shock location ( x=0 ). For clarity, the PDFs in (c) are estimated using 10,000 forecast ensemble members. Here, the error bars represent ±2 standard deviations of the observational noise; the same convention is adopted in all subsequent figures containing error bars.
Figure 2: Schematic of the neural EnKF framework. At each DA cycle, forward simulations generate a forecast ensemble {zif(x)}i=1ne in physical space. Each forecast ensemble member is then parameterized by a deep neural network FNN(θif;x) , which maps the spatial coordinates to the corresponding physical state (e.g., density, velocity, and pressure in compressible flows). The network parameters θif are then updated to θia through assimilation of the physical-space observations d (e.g., pressure measurements). Finally, the updated networks reconstruct the analysis ensemble {zia(x)}i=1ne back in physical space.
Figure 3: Schematic illustration of neural-network training across ensemble members on a non-convex loss landscape. (a) Independent training: each ensemble member is initialized from a different random starting point θi(0) and optimized independently, leading to convergence toward distinct local minima θi∗ and a poorly aligned ensemble in neural space (e.g., ensemble members 1 and 3 converge to nearby minima, whereas member 2 converges to a separated minimum). (b) Transfer learning: ensemble members are initialized from the trained parameters of a nearby ensemble member (e.g., θ3∗ ), guiding successive optimization processes toward neighboring regions of the loss landscape and promoting alignment of ensemble parameters in neural space.
Figure 4: Schematic illustration of the nearest-neighbor chain construction: (a) ensemble members in the chosen metric space; (b) initialization of the chain by selecting the most “central” ensemble member, defined as the member with the smallest average distance to all others; and (c)–(h) iterative construction of the chain, where at each step the ensemble member closest to the currently selected set is appended. Unselected members are shown as grey circles, previously selected members in light blue, and the member added at the current step is highlighted in dark blue. Solid/black lines indicate the resulting nearest-neighbor chain, with numbers denoting the order of insertion. Dashed/red arrows indicate the parent–child relationship used for training initialization, pointing from the nearest previously selected member (shown in light blue with a dark outline) to the newly added member.
Figure 5: Comparison of neural EnKF updates under two training strategies for the hyperbolic-tangent surrogate shock example. Panels (a1–c1) correspond to independent training with random initialization , while panels (a2–c2) correspond to the proposed nearest-neighbor chain training . (a1) Forecast ensemble together with the corresponding neural-network fits for each ensemble member, indicating negligible fitting errors. (b1) Analysis ensemble obtained using independent training, exhibiting irregular and spatially incoherent updates across ensemble members. (c1) Variation of neural-network parameters across ensemble members under independent training, showing highly irregular variability and poor alignment of parameters in neural space. (a2–b2) Forecast and analysis ensembles obtained using nearest-neighbor chain training, in which the analysis ensemble preserves the shock-like structure and better reflects the observation data. (c2) Smooth neural-network parameter variations under nearest-neighbor chain training, showing well-aligned parameter distributions in neural space. In panels (c1) and (c2), the horizontal axis indexes the ensemble members (50 samples), while the vertical axis indexes the flattened neural-network parameters (169 trainable parameters).
Figure 6: DA results for the inviscid Burgers’ equation with a structurally heterogeneous initial ensemble under neural EnKF updates. (a) Initial ensemble, where several members (red) lack the left trapezoidal segment present in the true profile (black), while the remaining members (grey) retain the complete structure. (b) Forecast and analysis ensembles at the first DA cycle, illustrating rapid recovery of the overall profile. (c) and (d) Analysis ensembles at the 20th and 50th DA cycles, respectively, showing progressive structural convergence toward the true solution during intermediate and late stages of assimilation.
Figure 7: Evolution of RMSE and ensemble spread over DA cycles for the inviscid Burgers’ equation with an initial ensemble exhibiting structural heterogeneity. Both metrics decrease rapidly during the initial assimilation stage and subsequently stabilize at comparable low levels. Markers are shown every five DA cycles for clarity.
Figure 8: Schematic of the reference shock tube configuration. The diaphragm at xd=0.5 separates the high-pressure left state (ρL,uL,pL) and the low-pressure right state (ρR,uR,pR) .
Figure 9: Analysis ensemble for the Sod shock tube problem at representative DA steps, showing increasing agreement with the analytical solution. Columns show density ρ , velocity u , and pressure p . Rows correspond to the initial ensemble ( t=0 ), the forecast ensemble prior to the first DA step ( t=0.025 ), and the analysis ensembles at the 1st, 2nd, 4th, and 8th DA steps ( t=0.025 , 0.05 , 0.1 , and 0.2 ). The density field exhibits three flow features, from left to right: a rarefaction wave, a contact interface, and a shock, whereas the velocity and pressure fields contain a rarefaction wave and a shock but do not include a contact discontinuity.
Figure 10: Comparison of the space–time evolution under DA for the Sod shock tube problem, with neural EnKF updates applied every Δtobs=0.025 . Columns show density ρ , velocity u , and pressure p . Rows correspond to the analytical solution, ensemble mean, a “farthest” ensemble member, and the absolute error between the analytical solution and the farthest ensemble member. The farthest ensemble member, defined once as the last member in the nearest-neighbor chain at the first DA step, corresponds to the member with the largest initial structural mismatch. The analytical solution and error fields are plotted on a fine grid with 3200 cells to enhance visualization of sharp discontinuities, while the ensemble mean and farthest ensemble member are shown on the numerical grid with 400 cells.
Figure 11: RMSE and ensemble spread over successive DA cycles for the Sod shock tube problem under the neural EnKF, showing rapid reduction of both metrics across all three variables during the first few DA cycles and sustained comparably low levels thereafter. The value at cycle 0 is computed from the forecast ensemble prior to the first DA update ( t=0.025 ), while cycles 1–8 correspond to the analysis ensembles after successive DA steps.
Figure 12: Schematic illustration of the 2D blast-wave initial configuration and ensemble sampling. (a) Computational domain Ω=Ω1∪Ω2 , where Ω1 denotes the circular high-pressure blast region and Ω2 the surrounding low-pressure region. (b) Ensemble of sampled initial discontinuity locations obtained by Gaussian sampling of the blast center and radius. For clarity, uncertainty in the density and pressure within Ω1 is not shown.
Figure 13: Pressure fields for the 2D blast wave over the full assimilation window, demonstrating rapid correction after the first neural EnKF update and sustained agreement with the reference solution thereafter. Rows show the reference solution, the ensemble mean (statistical state estimate), and the farthest ensemble member, defined as the last member in the nearest-neighbor chain constructed at the first DA step. Columns correspond to the initial condition ( t=0 ), the forecast states prior to the first DA update ( t=4.0×10−3 ), and the analysis states after five successive DA steps up to the final time ( t=2.0×10−2 ).
Figure 14: Velocity fields ( x -component) for the 2D blast wave over the full assimilation window, demonstrating effective correction after the first neural EnKF update and sustained agreement with the reference solution thereafter. Rows show the reference solution, the ensemble mean (statistical state estimate), and the farthest ensemble member. Columns correspond to the initial static condition ( t=0 ), the forecast states prior to the first DA update ( t=4.0×10−3 ), and the analysis states after five successive DA steps up to the final time ( t=2.0×10−2 ).
Figure 15: Density fields for the 2D blast wave over the full assimilation window, demonstrating effective correction after the first two neural EnKF updates and sustained agreement with the reference solution thereafter. Rows show the reference solution, the ensemble mean (statistical state estimate), and the farthest ensemble member. Columns correspond to the initial condition ( t=0 ), the forecast state prior to the first DA update ( t=4.0×10−3 ), and the analysis states after five successive DA steps up to the final time ( t=2.0×10−2 ).
Figure 16: Evolution of RMSE and ensemble spread over successive DA cycles for the 2D blast wave. Pressure and velocity exhibit a rapid reduction in both metrics after the first DA cycle, followed by sustained low levels. Density shows a comparable trend in ensemble spread, while its RMSE decreases more gradually and remains relatively low throughout the assimilation window. Cycle 0 corresponds to the forecast ensemble prior to the first DA update ( t=4×10−3 ), and cycles 1–5 denote analysis ensembles after successive DA steps.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Burgers’ equation
Sod shock tube
Blast wave
Input variables
x
x
x,y
Output variables
u
ρ,u,p
ρ,u,v,p
# input neurons
1
1
2
# hidden layers
4
4
6
# neurons per hidden layer
50
64
100
# output neurons
1
3
4
Appendix
Table 1: Neural network architectures used in the numerical experiments.
Figure 17: Training convergence histories and wall-clock times for the Sod shock tube case: (a,b) training loss histories for selected ensemble members during the first and final DA cycles, respectively; and (c) training wall-clock times for each DA cycle, comparing the first ensemble member in the nearest-neighbor chain with the average over the remaining 39 ensemble members and demonstrating the significant reduction in training cost enabled by warm-start initialization in the chain-based training strategy. Here, the first ensemble member in the chain is trained using a learning rate of 10−2 , while the remaining ensemble members use a learning rate of 10−3 . All reported wall-clock times correspond to single-core CPU execution.
Figure 18: EnKF performance for the shock tube problem at the first DA step ( t=0.025 ), showing strong oscillations and the emergence of nonphysical states. The shaded regions indicate the nonphysical range ( ρ<0 and p<0 ).
Figure 19: Forward propagation of uncertainty in the left-state density ρL for the shock tube problem, illustrating its weak influence on the pressure field. The ensemble is evolved from t=0.0 to 0.2 without DA. Columns show density ρ , velocity u , and pressure p ; rows correspond to the initial and final times.
Obs. noise
# Obs. loc.
Δxobs
# DA cycles
Δtobs
Baseline
5%
10
0.1
8
0.025
Exp. A1
1%
10
0.1
8
0.025
Exp. A2
10%
10
0.1
8
0.025
Exp. B1
5%
5
0.2
8
0.025
Exp. B2
5%
20
0.05
8
0.025
Exp. C1
5%
10
0.1
4
0.05
Appendix
Table 2: DA configurations used in the parametric study of the shock tube problem. The baseline corresponds to the setup used in the main text. Experiments A1–A2 vary the observation noise level; B1–B2 vary the number of observation locations, thereby changing the observation spacing Δxobs ; and C1–C2 vary the observation interval Δtobs , leading to different numbers of DA cycles. Parameters that differ from the baseline are highlighted in bold .
Figure 20: Final-time ( t=0.2 ) analysis ensembles for the shock tube problem under the DA configurations summarized in Table 2 . Columns show density ρ , velocity u , and pressure p . Rows correspond to the baseline case used in the main text; Experiments A1–A2 (low and high observation noise); Experiments B1–B2 (reduced and increased numbers of observation locations); and Experiments C1–C2 (lower and higher observation frequency).
The ensemble Kalman filter (EnKF) is widely adopted for sequential data assimilation, but fails for solutions with discontinuities, such as shocks in compressible flows. For ensembles with uncertainty in discontinuity location, the EnKF analysis can lose physical consistency, which can be viewed through three lenses: (i) bimodal pointwise statistics in the forecast ensemble, (ii) mixed-sign analysis weights in the affine superposition of forecast members, and (iii) spatially distributed Kalman-gain profiles that modify the state around a discontinuity rather than translating it coherently. In order to preserve flow features in the analysis state and, as such, maintain physical consistency, a feature-preserving latent-EnKF is introduced. The proposed method performs the ensemble update in a low-dimensional latent space learned by a coordinate-conditioned auto-decoder. The latent representation retains near-Gaussian ensemble statistics and provides coordinates in which flow features vary coherently across the forecast ensemble. The EnKF updates performed on the latent codes are decoded into coherent changes in the represented flow features, translating shocks through coordinate morphing rather than superposing them at displaced locations. The corresponding Kalman-gain profiles generated by the latent-EnKF exhibit an approximately step-like variation across the discontinuities, allowing the fronts to translate while preserving sharp features and monotonicity. Sod shock-tube, Mach~2 shock-cylinder, and immersed-body shape-inference tests with sparse, noisy observations demonstrate accurate recovery of shocks and contact discontinuities without spurious oscillations.
Data assimilation estimates a dynamical state from partial and noisy observations. Classical ensemble filters are efficient but restrict analysis updates through finite sample covariance and affine Gaussian distribution. We introduce the Flow Ensemble Filter (FlowEF), which uses conditional flow matching to transport the forecast ensemble from a classical baseline filter to an analysis ensemble. FlowEF uses a localized Gaussian source during training, transports forecast ensemble members from a baseline filter at deployment, and conditions its velocity field on ensembles from that baseline filter and the observation. The proposed model therefore learns a nonlinear update while mapping each baseline ensemble independently. For sparsely observed dynamical systems, FlowEF improves both deterministic and probabilistic metrics over all four classical ensemble filters. It also achieves the best performance among the state-of-the-art generative data assimilation models.
Haoyuan Chen, Alexandre Thiéry
Department of Statistics and Data Science National University of Singapore 6 Science Drive 2, Singapore 117546
Data assimilation estimates the state of a dynamical system from model forecasts and incoming observations. Many observation mechanisms, however, are many-to-one, implicit, non-smooth, or accessible only through simulation, and need not provide the residual structures or likelihood guidance required by existing ensemble filters. We introduce implicit data assimilation, in which the analysis law is defined as an energy tilt of the forecast distribution. We then propose the Ensemble Controlled-flow Filter (EnCF), which realizes this update through a stochastic controlled flow and learns the observation-dependent control by adjoint matching from terminal energy gradients. For simulator-defined observations, EnCF-LF learns a surrogate conditional energy from samples and applies the same controlled-flow solver. We prove ideal exactness, derive a one-step error decomposition, and establish non-accumulation of local errors under filter stability. Numerical results show that Kalman-type filters remain preferable for smooth additive-Gaussian observations, while the proposed methods are better suited to non-Gaussian, many-to-one, multimodal, and implicit observation models.
Zhuoyuan Li, Yue Zhao, Ming Li
Institute for Functional Intelligent Materials, National University of Singapore, Singapore 117544, Singapore · School of Mathematical Sciences, Fudan University, Shanghai 200433, China