Closed-form conditional diffusion models for data assimilation

Abstract

We propose closed-form conditional diffusion models for data assimilation. Diffusion models use data to learn the score function (defined as the gradient of the log-probability density of a data distribution), allowing them to generate new samples from the data distribution by reversing a noise injection process. While it is common to train neural networks to approximate the score function, we leverage the analytical tractability of the score function to assimilate the states of a system with measurements. To enable the efficient evaluation of the score function, we use kernel density estimation to model the joint distribution of the states and their corresponding measurements. The proposed approach also inherits the capability of conditional diffusion models of operating in black-box settings, i.e., the proposed data assimilation approach can accommodate systems and measurement processes without their explicit knowledge. The ability to accommodate black-box systems combined with the superior capabilities of diffusion models in approximating complex, non-Gaussian probability distributions means that the proposed approach offers advantages over many widely used filtering methods. We evaluate the proposed method on nonlinear data assimilation problems based on the Lorenz-63 and Lorenz-96 systems of moderate dimensionality and nonlinear measurement models. Results show the proposed approach outperforms the widely used ensemble Kalman and particle filters when small to moderate ensemble sizes are used.

Figures (6)

• Figure 1: Evolution of the filtering distribution of $x_1$ (vertical axis) for the Lorenz-63 system over $K=100$ assimilation steps (horizontal axis). Each column corresponds to a different ensemble size and simulation. Within each column, the rows show (top to bottom): the reference distribution obtained using SIR with $N_{\rm true}=100{,}000$ particles (Truth), the filtering distribution obtained with ensemble size $N$ using three different filters. The color intensity corresponds to the estimated probability density, with darker regions indicating higher density.
• Figure 2: Predicted and updated states at assimilation step $k=60$ and for $N=500$ with three different filters. Red particles represent the states after the prediction step, blue particles represent the states after the update step, and the black line indicates the observation.
• Figure 3: Visualization the performance of different filters with ensemble size $N=100$ for one test trajectory of the 10-dimensional Lorenz-96 system over data assimilation steps 300-400. In each plot, the vertical axis corresponds to the states, and the horizontal axis corresponds to the data assimilation step. Rows 1 through 5 correspond to the observations, true states, estimated mean of the assimilated states, absolute error of mean with respect to the ground truth, and estimated ensemble spread (one standard deviation about the estimated mean), respectively. Columns 1 through 3 correspond to the conditional diffusion model-based, EnKF, and SIR filters, respectively.
• Figure 4: Performance comparison of three filtering methods for a single degree of freedom, $x_1$, of the 10-dimensional Lorenz-96 system with ensemble size $N=100$ over data assimilation steps 300-400, using the same test trajectory as in \ref{['fig:L96-10D']}. The top, middle, and bottom plots correspond to the conditional diffusion model-based filter, EnKF, and SIR filter, respectively. In each plot, the vertical axis corresponds to the value of $x_1$, and the horizontal axis corresponds to the data assimilation steps. The solid blue line represents the ensemble mean of the assimilated state, the dashed red line denotes the true state trajectory, the black crosses indicate the observations, and the shaded gray region shows the ensemble spread (one standard deviation about the estimated mean).
• Figure 5: Visualization the performance of different filters with ensemble size $N=100$ for one test trajectory of the 20-dimensional Lorenz-96 system over data assimilation steps 300-400. In each plot, the vertical axis corresponds to the states, and the horizontal axis corresponds to the data assimilation step. Rows 1 through 5 correspond to the observations, true states, estimated mean of the assimilated states, absolute error of mean with respect to the ground truth, and estimated ensemble spread (one standard deviation about the estimated mean), respectively. Columns 1 through 3 correspond to the conditional diffusion model-based, EnKF, and SIR filters, respectively.