Abstract

Data assimilation is concerned with sequentially estimating a temporally evolving state. This task, which arises in a wide range of scientific and engineering applications, is particularly challenging when the state is high-dimensional and the state-space dynamics are unknown. This paper introduces a machine learning framework for learning dynamical systems in data assimilation. Our auto-differentiable ensemble Kalman filters (AD-EnKFs) blend ensemble Kalman filters for state recovery with machine learning tools for learning the dynamics. In doing so, AD-EnKFs leverage the ability of ensemble Kalman filters to scale to high-dimensional states and the power of automatic differentiation to train high-dimensional surrogate models for the dynamics. Numerical results using the Lorenz-96 model show that AD-EnKFs outperform existing methods that use expectation-maximization or particle filters to merge data assimilation and machine learning. In addition, AD-EnKFs are easy to implement and require minimal tuning.

Keywords

  1. ensemble Kalman filters
  2. autodifferentiation
  3. data assimilation
  4. machine learning

MSC codes

  1. 62M05
  2. 68T09
  3. 68T07
  4. 86-08

Get full access to this article

View all available purchase options and get full access to this article.

Supplementary Material


PLEASE NOTE: These supplementary files have not been peer-reviewed.


Index of Supplementary Materials

Title of paper: Auto-differentiable Ensemble Kalman Filters

Authors: Yuming Chen, Daniel Sanz-Alonso, and Rebecca Willett

File: supplement.pdf

Type: PDF

Contents: Additional proofs, additional figures and implementation details.

References

1.
M. Abadi, P. Barham, J. Chen, Z. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, G. Irving, M. Isard, et al., Tensorflow: A system for large-scale machine learning, in Proceedings of the 12th USENIX Symposium on Operating Systems Design and Implementation, 2016, pp. 265--283.
2.
S. Agapiou, O. Papaspiliopoulos, D. Sanz-Alonso, and A. M. Stuart, Importance sampling: Intrinsic dimension and computational cost, Statist. Sci., 32 (2017), pp. 405--431.
3.
J. L. Anderson, An ensemble adjustment Kalman filter for data assimilation, Monthly Weather Review, 129 (2001), pp. 2884--2903.
4.
J. L. Anderson and S. L. Anderson, A Monte Carlo implementation of the nonlinear filtering problem to produce ensemble assimilations and forecasts, Monthly Weather Review, 127 (1999), pp. 2741--2758.
5.
C. Andrieu, A. Doucet, and R. Holenstein, Particle Markov chain Monte Carlo methods, J. R. Stat. Soc. Ser. B Stat. Methodol., 72 (2010), pp. 269--342.
6.
I. Ayed, E. de Bézenac, A. Pajot, J. Brajard, and P. Gallinari, Learning Dynamical Systems from Partial Observations, preprint, arXiv:1902.11136, 2019.
7.
C. M. Bishop, Pattern Recognition and Machine Learning, Springer, New York, 2006.
8.
M. Bocquet, J. Brajard, A. Carrassi, and L. Bertino, Data assimilation as a learning tool to infer ordinary differential equation representations of dynamical models, Nonlinear Processes in Geophysics, 26 (2019), pp. 143--162.
9.
M. Bocquet, J. Brajard, A. Carrassi, and L. Bertino, Bayesian Inference of Chaotic Dynamics by Merging Data Assimilation, Machine Learning and Expectation-Maximization, preprint, arXiv:2001.06270, 2020.
10.
M. Bocquet, C. A. Pires, and L. Wu, Beyond Gaussian statistical modeling in geophysical data assimilation, Monthly Weather Review, 138 (2010), pp. 2997--3023.
11.
J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: Composable Transformations of Python+NumPy Programs, http://github.com/google/jax, 2018.
12.
J. Brajard, A. Carrassi, M. Bocquet, and L. Bertino, Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: A case study with the Lorenz 96 model, J. Comput. Sci., 44 (2020), 101171.
13.
S. L. Brunton and J. N. Kutz, Data-driven Science and Engineering: Machine Learning, Dynamical Systems, and Control, Cambridge University Press, Cambridge, UK, 2019.
14.
A. Carrassi, M. Bocquet, A. Hannart, and M. Ghil, Estimating model evidence using data assimilation, Quart. J. Royal Meteorological Society, 143 (2017), pp. 866--880.
15.
N. K. Chada, Y. Chen, and D. Sanz-Alonso, Iterative ensemble kalman methods: A unified perspective with some new variants, Foundations of Data Science, 3 (2021), pp. 331--369.
16.
R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud, Neural ordinary differential equations, in Advances in Neural Information Processing Systems, Vol. 31, S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, eds., Curran Associates, 2018, pp. 6571--6583, https://proceedings.neurips.cc/paper/2018/file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf.
17.
T. J. Cocucci, M. Pulido, M. Lucini, and P. Tandeo, Model error covariance estimation in particle and ensemble Kalman filters using an online expectation--maximization algorithm, Quart. J. Royal Meteorological Society, 147 (2021), pp. 526--543.
18.
A. Corenflos, J. Thornton, A. Doucet, and G. Deligiannidis, Differentiable Particle Filtering via Entropy-Regularized Optimal Transport, preprint, arXiv:2102.07850, 2021.
19.
D. Crisan, P. Del Moral, A. Jasra, and H. Ruzayqat, Log-Normalization Constant Estimation Using the Ensemble Kalman-Bucy Filter with Application to High-Dimensional Models, preprint, arXiv:2101.11460, 2021.
20.
E. De Brouwer, J. Simm, A. Arany, and Y. Moreau, GRU-ODE-Bayes: Continuous Modeling of Sporadically-Observed Time Series, preprint, arXiv:1905.12374, 2019.
21.
P. Del Moral, Feynman-Kac Formulae: Genealogical and Interacting Particle Systems with Applications, Probab. Appl., Springer-Verlag, New York, 2004.
22.
P. Del Moral, J. Tugaut, et al., On the stability and the uniform propagation of chaos properties of ensemble Kalman--Bucy filters, Ann. Appl. Probab., 28 (2018), pp. 790--850.
23.
T. DelSole and X. Yang, State and parameter estimation in stochastic dynamical models, Phys. D, 239 (2010), pp. 1781--1788.
24.
A. P. Dempster, N. M. Laird, and D. B. Rubin, Maximum likelihood from incomplete data via the EM algorithm, J. R. Stat. Soc. Ser. B Stat. Methodol., 39 (1977), pp. 1--22.
25.
A. Doucet, S. Godsill, and C. Andrieu, On sequential Monte Carlo sampling methods for Bayesian filtering, Statist. Comput., 10 (2000), pp. 197--208.
26.
A. Doucet and A. M. Johansen, A tutorial on particle filtering and smoothing: Fifteen years later, Handbook Nonlinear Filtering, 12 (2009), 3.
27.
D. Dreano, P. Tandeo, M. Pulido, B. Ait-El-Fquih, T. Chonavel, and I. Hoteit, Estimating model-error covariances in nonlinear state-space models using Kalman smoothing and the expectation--maximization algorithm, Quart. J. Royal Meteorological Society, 143 (2017), pp. 1877--1885.
28.
C. Drovandi, R. G. Everitt, A. Golightly, D. Prangle, et al., Ensemble MCMC: Accelerating pseudo-marginal MCMC for state space models using the ensemble Kalman filter, Bayesian Anal., to appear.
29.
G. Evensen, Sequential data assimilation with a nonlinear quasi-geostrophic model using Monte Carlo methods to forecast error statistics, J. Geophysical Research Oceans, 99 (1994), pp. 10143--10162.
30.
G. Evensen, Data Assimilation: The Ensemble Kalman Filter, Springer, New York, 2009.
31.
M. Fraccaro, S. Kamronn, U. Paquet, and O. Winther, A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning, preprint, arXiv:1710.05741, 2017.
32.
R. Furrer and T. Bengtsson, Estimation of high-dimensional prior and posterior covariance matrices in Kalman filter variants, J. Multivariate Anal., 98 (2007), pp. 227--255.
33.
G. Gaspari and S. E. Cohn, Construction of correlation functions in two and three dimensions, Quart. J. Royal Meteorological Society, 125 (1999), pp. 723--757.
34.
M. B. Giles, Collected matrix derivative results for forward and reverse mode algorithmic differentiation, in Advances in Automatic Differentiation, Springer, New York, 2008, pp. 35--44.
35.
N. J. Gordon, D. J. Salmond, and A. F. Smith, Novel approach to nonlinear/non-Gaussian Bayesian state estimation, IEE Proc. F Radar and Signal Processing, 140 (1993), pp. 107--113.
36.
G. A. Gottwald and S. Reich, Supervised learning from noisy observations: Combining machine-learning techniques with data assimilation, Phys. D, 423 (2021), 132911.
37.
T. M. Hamill, J. S. Whitaker, and C. Snyder, Distance-dependent filtering of background error covariance estimates in an ensemble Kalman filter, Monthly Weather Review, 129 (2001), pp. 2776--2790.
38.
A. Hannart, A. Carrassi, M. Bocquet, M. Ghil, P. Naveau, M. Pulido, J. Ruiz, and P. Tandeo, DADA: Data assimilation for the detection and attribution of weather and climate-related events, Climatic Change, 136 (2016), pp. 155--174.
39.
J. Harlim, S. W. Jiang, S. Liang, and H. Yang, Machine learning for prediction with missing dynamics, J. Comput. Phys., 428 (2021), 109922.
40.
P. L. Houtekamer and H. L. Mitchell, Data assimilation using an ensemble Kalman filter technique, Monthly Weather Review, 126 (1998), pp. 796--811.
41.
P. L. Houtekamer and H. L. Mitchell, A sequential ensemble Kalman filter for atmospheric data assimilation, Monthly Weather Review, 129 (2001), pp. 123--137.
42.
P. L. Houtekamer and F. Zhang, Review of the ensemble Kalman filter for atmospheric data assimilation, Monthly Weather Review, 144 (2016), pp. 4489--4532.
43.
T. Ishizone, T. Higuchi, and K. Nakamura, Ensemble Kalman Variational Objectives: Nonlinear Latent Trajectory Inference with a Hybrid of Variational Inference and Ensemble Kalman Filter, preprint, arXiv:2010.08729, 2020.
44.
A. H. Jazwinski, Stochastic Processes and Filtering Theory, Courier Corporation, 2007.
45.
R. E. Kalman, A new approach to linear filtering and prediction problems, Trans. ASME J. Basic Engineering, 82 (1960).
46.
M. Katzfuss, J. R. Stroud, and C. K. Wikle, Understanding the ensemble Kalman filter, Amer. Statist., 70 (2016), pp. 350--357.
47.
M. Katzfuss, J. R. Stroud, and C. K. Wikle, Ensemble Kalman methods for high-dimensional hierarchical dynamic space-time models, J. Amer. Statist. Assoc., 115 (2020), pp. 866--885.
48.
D. P. Kingma and M. Welling, Auto-Encoding Variational Bayes, in Proceedings of the 2nd International Conference on Learning Representations, 2014.
49.
A. Kloss, G. Martius, and J. Bohg, How to train your differentiable filter, Autonomous Robots, (2021), pp. 1--18.
50.
R. Krishnan, U. Shalit, and D. Sontag, Structured inference networks for nonlinear state space models, in Proceedings of the AAAI Conference on Artificial Intelligence, 2017.
51.
E. Kwiatkowski and J. Mandel, Convergence of the square root ensemble Kalman filter in the large ensemble limit, SIAM/ASA J. Uncertain. Quantif., 3 (2015), pp. 1--17.
52.
K. J. Law, D. Sanz-Alonso, A. Shukla, and A. M. Stuart, Filter accuracy for the Lorenz 96 model: Fixed versus adaptive observation operators, Phys. D, 325 (2016), pp. 1--13.
53.
K. J. Law and A. M. Stuart, Evaluating data assimilation algorithms, Monthly Weather Review, 140 (2012), pp. 3757--3782.
54.
K. J. Law, A. M. Stuart, and K. Zygalakis, Data Assimilation, Springer, New York, 2015.
55.
K. J. Law, H. Tembine, and R. Tempone, Deterministic mean-field ensemble Kalman filtering, SIAM J. Sci. Comput., 38 (2016), pp. A1251--A1279.
56.
T. A. Le, M. Igl, T. Rainforth, T. Jin, and F. Wood, Auto-Encoding sequential Monte Carlo, in Proceedings of the International Conference on Learning Representations, 2018, https://openreview.net/forum?id=BJ8c3f-0b.
57.
F. Le Gland and L. Mevel, Recursive identification in hidden markov models, in Proceedings of the 36th Conference on Decision and Control, Vol. 4, 1997, pp. 3468--3473.
58.
F. Le Gland, V. Monbet, and V.-D. Tran, Large Sample Asymptotics for the Ensemble Kalman Filter, Ph.D. thesis, INRIA, 2009.
59.
M. E. Levine and A. M. Stuart, A Framework for Machine Learning of Model Error in Dynamical Systems, preprint, arXiv:2107.06658, 2021.
60.
E. N. Lorenz, Predictability: A problem partly solved, in Proceedings of the Seminar on Predictability, 1996.
61.
C. J. Maddison, D. Lawson, G. Tucker, N. Heess, M. Norouzi, A. Mnih, A. Doucet, and Y. W. Teh, Filtering Variational Objectives, preprint, arXiv:1705.09279, 2017.
62.
A. J. Majda and J. Harlim, Filtering Complex Turbulent Systems, Cambridge University Press, Cambridge, UK, 2012.
63.
H. P. McKean, Propagation of chaos for a class of non-linear parabolic equations, in Stochastic Differential Equations, Lecture Series in Differential Equations, Session 7, 1967, pp. 41--57.
64.
S. Metref, A. Hannart, J. Ruiz, M. Bocquet, A. Carrassi, and M. Ghil, Estimating model evidence using ensemble-based data assimilation with localization---The model selection problem, Quart. J. Royal Meteorological Society, 145 (2019), pp. 1571--1588.
65.
C. Naesseth, S. Linderman, R. Ranganath, and D. Blei, Variational sequential Monte Carlo, in International Conference on Artificial Intelligence and Statistics, PMLR, 2018, pp. 968--977.
66.
D. Nguyen, S. Ouala, L. Drumetz, and R. Fablet, EM-Like Learning Chaotic Dynamics from Noisy and Partial Observations, preprint, arXiv:1903.10335, 2019.
67.
E. Ott, B. R. Hunt, I. Szunyogh, A. V. Zimin, E. J. Kostelich, M. Corazza, E. Kalnay, D. Patil, and J. A. Yorke, A local ensemble Kalman filter for atmospheric data assimilation, Tellus A Dynamic Meteorology and Oceanography, 56 (2004), pp. 415--428.
68.
O. Papaspiliopoulos and M. Ruggiero, Optimal filtering and the dual process, Bernoulli, 20 (2014), pp. 1999--2019.
69.
R. Pascanu, T. Mikolov, and Y. Bengio, On the difficulty of training recurrent neural networks, in International Conference on Machine Learning, PMLR, 2013, pp. 1310--1318.
70.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al., PyTorch: An imperative style, high-performance deep learning library, in Advances in Neural Information Processing Systems, 32, 2019.
71.
M. Pulido, P. Tandeo, M. Bocquet, A. Carrassi, and M. Lucini, Stochastic parameterization identification using ensemble Kalman filtering combined with maximum likelihood methods, Tellus A Dynamic Meteorology and Oceanography, 70 (2018), pp. 1--17.
72.
M. Raissi, P. Perdikaris, and G. E. Karniadakis, Multistep Neural Networks for Data-Driven Discovery of Nonlinear Dynamical Systems, preprint, arXiv:1801.01236, 2018.
73.
R. Ranganath, S. Gerrish, and D. Blei, Black box variational inference, in Artificial Intelligence and Statistics, PMLR, 2014, pp. 814--822.
74.
S. S. Rangapuram, M. Seeger, J. Gasthaus, L. Stella, Y. Wang, and T. Januschowski, Deep state space models for time series forecasting, in Proceedings of the 32nd International Conference on Neural Information Processing Systems, 2018, pp. 7796--7805.
75.
S. Reich and C. Cotter, Probabilistic Forecasting and Bayesian Data Assimilation, Cambridge University Press, Cambridge, UK, 2015.
76.
D. J. Rezende, S. Mohamed, and D. Wierstra, Stochastic backpropagation and approximate inference in deep generative models, in International Conference on Machine Learning, PMLR, 2014, pp. 1278--1286.
77.
M. Roth, G. Hendeby, C. Fritsche, and F. Gustafsson, The Ensemble Kalman filter: A signal processing perspective, EURASIP J. Advances in Signal Processing, 2017 (2017), pp. 1--16.
78.
Y. Rubanova, R. T. Chen, and D. K. Duvenaud, Latent ordinary differential equations for irregularly-sampled time series, Advances in Neural Information Processing Systems, 32, 2019.
79.
P. Sakov and L. Bertino, Relation between two common localisation methods for the EnKF, Comput. Geosci., 15 (2011), pp. 225--237.
80.
D. Sanz-Alonso, Importance sampling and necessary sample size: An information theory approach, SIAM/ASA J. Uncertain. Quantif., 6 (2018), pp. 867--879.
81.
D. Sanz-Alonso and A. M. Stuart, Long-time asymptotics of the filtering distribution for partially observed chaotic dynamical systems, SIAM/ASA J. Uncertain. Quantif., 3 (2015), pp. 1200--1220.
82.
D. Sanz-Alonso, A. M. Stuart, and A. Taeb, Inverse Problems and Data Assimilation, preprint, arXiv:1810.06191, 2019.
83.
D. Sanz-Alonso and Z. Wang, Bayesian update with importance sampling: Required sample size, Entropy, 23 (2021), https://doi.org/10.3390/e23010022.
84.
C. Schillings and A. M. Stuart, Analysis of the ensemble Kalman filter for inverse problems, SIAM J. Numer. Anal., 55 (2017), pp. 1264--1290.
85.
C. Snyder, T. Bengtsson, P. Bickel, and J. Anderson, Obstacles to high-dimensional particle filtering, Monthly Weather Review, 136 (2008), pp. 4629--4640.
86.
J. R. Stroud and T. Bengtsson, Sequential state and variance estimation within the ensemble Kalman filter, Monthly Weather Review, 135 (2007), pp. 3194--3208.
87.
J. R. Stroud, M. Katzfuss, and C. K. Wikle, A Bayesian adaptive ensemble Kalman filter for sequential state and parameter estimation, Monthly Weather Review, 146 (2018), pp. 373--386.
88.
J. R. Stroud, M. L. Stein, B. M. Lesht, D. J. Schwab, and D. Beletsky, An ensemble Kalman filter and smoother for satellite data assimilation, J. Amer. Statist. Assoc., 105 (2010), pp. 978--990.
89.
I. Sutskever, O. Vinyals, and Q. V. Le, Sequence to sequence learning with neural networks, in Advances in Neural Information Processing Systems, 2014, pp. 3104--3112.
90.
A.-S. Sznitman, Topics in propagation of chaos, in Ecole d'été de probabilités de Saint-Flour XIX---1989, Springer, New York, 1991, pp. 165--251.
91.
I. Szunyogh, E. J. Kostelich, G. Gyarmati, E. Kalnay, B. R. Hunt, E. Ott, E. Satterfield, and J. A. Yorke, A local ensemble transform Kalman filter data assimilation system for the NCEP global model, Tellus A Dynamic Meteorology and Oceanography, 60 (2008), pp. 113--130.
92.
P. Tandeo, M. Pulido, and F. Lott, Offline parameter estimation using EnKF and maximum likelihood error covariance estimates: Application to a subgrid-scale orography parametrization, Quart. J. Royal Meteorological Society, 141 (2015), pp. 383--395.
93.
G. Ueno and N. Nakamura, Iterative algorithm for maximum-likelihood estimation of the observation-error covariance matrix for ensemble-based filters, Quart. J. Royal Meteorological Society, 140 (2014), pp. 295--315.
94.
G. Ueno and N. Nakamura, Bayesian estimation of the observation-error covariance matrix in ensemble-based filters, Quart. J. Royal Meteorological Society, 142 (2016), pp. 2055--2080.
95.
G. C. Wei and M. A. Tanner, A Monte Carlo implementation of the EM algorithm and the poor man's data augmentation algorithms, J. Amer. Statist. Assoc., 85 (1990), pp. 699--704.
96.
J. S. Whitaker, T. M. Hamill, X. Wei, Y. Song, and Z. Toth, Ensemble data assimilation with the ncep global forecast system, Monthly Weather Review, 136 (2008), pp. 463--482.
97.
R. J. Williams and D. Zipser, Gradient-based learning algorithms for recurrent networks and their computational complexity, in Back-propagation: Theory, Architectures and Applications, Y. Chauvin and D. E. Rumelhart, eds., Erlbaum, Hillsdale, NJ, 1995, pp. 433--486.
98.
K. Xu and C. K. Wikle, Estimation of parameterized spatio-temporal dynamic models, J. Statist. Plann. Inference, 137 (2007), pp. 567--588.

Information & Authors

Information

Published In

cover image SIAM Journal on Mathematics of Data Science
SIAM Journal on Mathematics of Data Science
Pages: 801 - 833
ISSN (online): 2577-0187

History

Submitted: 19 July 2021
Accepted: 26 March 2022
Published online: 23 June 2022

Keywords

  1. ensemble Kalman filters
  2. autodifferentiation
  3. data assimilation
  4. machine learning

MSC codes

  1. 62M05
  2. 68T09
  3. 68T07
  4. 86-08

Authors

Affiliations

Funding Information

FBBVA
Air Force Office of Scientific Research https://doi.org/10.13039/100000181 : FA9550-18-1-0166
National Science Foundation https://doi.org/10.13039/100000001 : DMS-2027056, OAC-1934637, DMS-1925101, DMS-1930049, DMS-2023109
U.S. Department of Defense https://doi.org/10.13039/100000005 : FA9550-18-1-0166
U.S. Department of Energy https://doi.org/10.13039/100000015 : DE-SC0022232, DE-AC02-06CH11357

Metrics & Citations

Metrics

Citations

If you have the appropriate software installed, you can download article citation data to the citation manager of your choice. Simply select your manager software from the list below and click Download.

View Options

View options

PDF

View PDF

Media

Figures

Other

Tables

Share

Share

Copy the content Link

Share with email

Email a colleague

Share on social media

The SIAM Publications Library now uses SIAM Single Sign-On for individuals. If you do not have existing SIAM credentials, create your SIAM account https://my.siam.org.