physics.flu-dynFeb 26, 2026
SaveNeural ensemble Kalman filter: Data assimilation for compressible flows with shocks
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
Abstract
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 ( ) and a bimodal distribution near the shock location ( ). For clarity, the PDFs in (c) are estimated using 10,000 forecast ensemble members. Here, the error bars represent 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 in physical space. Each forecast ensemble member is then parameterized by a deep neural network , which maps the spatial coordinates to the corresponding physical state (e.g., density, velocity, and pressure in compressible flows). The network parameters are then updated to through assimilation of the physical-space observations (e.g., pressure measurements). Finally, the updated networks reconstruct the analysis ensemble 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 and optimized independently, leading to convergence toward distinct local minima 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., ), 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 separates the high-pressure left state and the low-pressure right state .
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 , and pressure . Rows correspond to the initial ensemble ( ), the forecast ensemble prior to the first DA step ( ), and the analysis ensembles at the 1st, 2nd, 4th, and 8th DA steps ( , , , and ). 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 . Columns show density , velocity , and pressure . 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 ( ), 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 , where denotes the circular high-pressure blast region and 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 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 ( ), the forecast states prior to the first DA update ( ), and the analysis states after five successive DA steps up to the final time ( ).
Figure 14: Velocity fields ( -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 ( ), the forecast states prior to the first DA update ( ), and the analysis states after five successive DA steps up to the final time ( ).
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 ( ), the forecast state prior to the first DA update ( ), and the analysis states after five successive DA steps up to the final time ( ).
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 ( ), 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 | |||
| Output variables | |||
| # input neurons | 1 | 1 | 2 |
| # hidden layers | 4 | 4 | 6 |
| # neurons per hidden layer | 50 | 64 | 100 |
| # output neurons | 1 | 3 | 4 |
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 , while the remaining ensemble members use a learning rate of . 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 ( ), showing strong oscillations and the emergence of nonphysical states. The shaded regions indicate the nonphysical range ( and ).
Figure 19: Forward propagation of uncertainty in the left-state density for the shock tube problem, illustrating its weak influence on the pressure field. The ensemble is evolved from to without DA. Columns show density , velocity , and pressure ; rows correspond to the initial and final times.
| Obs. noise | # Obs. loc. | # DA cycles | |||
| 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 |
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 ; and C1–C2 vary the observation interval , leading to different numbers of DA cycles. Parameters that differ from the baseline are highlighted in bold .
Figure 20: Final-time ( ) analysis ensembles for the shock tube problem under the DA configurations summarized in Table 2 . Columns show density , velocity , and pressure . 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).
More in this topic
Data Assimilation
Data assimilation combines observations with dynamical models to estimate hidden system states and their uncertainty over time.
See all 58 papers