ParticleDA.jl v.1.0: a distributed particle-filtering data assimilation package

Digital twins of physical and human systems informed by real-time data are becoming ubiquitous across weather forecasting, disaster preparedness, and urban planning, but researchers lack the tools to run these models effectively and efficiently, limiting progress. One of the current challenges is to assimilate observations in highly non-linear dynamical systems, as the practical need is often to detect abrupt changes. We have developed a software platform to improve the use of real-time data in non-linear system representations where non-Gaussianity limits the applicability of data assimilation algorithms such as the ensemble Kalman filter and variational methods. Particle-filter-based data assimilation algorithms have been implemented within a user-friendly open-source software platform in Julia – ParticleDA.jl. To ensure the applicability of the developed platform in realistic scenarios, emphasis has been placed on numerical efficiency and scalability on high-performance computing systems. Furthermore, the platform has been developed to be forward-model agnostic, ensuring that it is applicable to a wide range of modelling settings, for instance unstructured and non-uniform meshes in the spatial domain or even state spaces that are not spatially organized. Applications to tsunami and numerical weather prediction demonstrate the computational benefits and ease of using the high-level Julia interface with the package to perform filtering in a variety of complex models.

1 Introduction

Data assimilation (DA) focuses on optimally combining observations with a dynamical model of a physical system to estimate how the system state evolves over time. The field of research has its origins within the numerical weather prediction (NWP) community, where DA techniques are applied iteratively to update current best estimates of the state of the atmosphere. Recently the methods and practices developed have been employed in diverse areas of geosciences, with Carrassi et al. (2018) and Vetra-Carvalho et al. (2018) providing recent overviews. Further DA has seen a huge expansion into other scientific disciplines with applications in, for example, robotics (Berquin and Zell, 2022), economic modelling (Nadler et al., 2019), and plasma physics (Sanpei et al., 2021). In the era of digital twinning, which involves combining high-fidelity representations of reality with the optimal use of observations, real-time data have become vital and DA frameworks have naturally been incorporated. The area of data learning has also emerged where DA approaches are integrated with machine learning techniques (Buizza et al., 2022).

There are various popular DA techniques, with variational methods (3D-Var and 4D-Var) (Thépart et al., 1993) and ensemble Kalman filters (EnKFs) (Evensen, 1994; Burgers et al., 1998) being extensively used in operational and research settings. Bannister (2017) provides a good overview of operational methods. However, these methods have difficulties with handling non-linear problems and with representing uncertainties accurately (Lei et al., 2010; Bocquet et al., 2010). For instance Miyoshi (2005) and Kondo and Miyoshi (2019) updated an ensemble of 10 240 particles using the EnKF to demonstrate the bimodality of some distributions due to inherent non-linearities. Furthermore, the continuing growth in compute hardware performance has allowed running increasingly complex and high-resolution models which are able to resolve non-linear processes happening at a fine spatial scale (Vetra-Carvalho et al., 2018), creating an increasing demand for DA methods which are able to accurately quantify uncertainty in such settings. Particle filters (PFs) (Gordon et al., 1993) are an alternative approach which offer the promise of consistent DA for problems with non-linear dynamics and non-Gaussian noise distributions. Traditionally the main difficulty with particle filtering techniques has been the “curse of dimensionality” (Bengtsson et al., 2008; Bickel et al., 2008; Snyder, 2011), where in high-dimensional settings, filtering leads to degeneracy of the importance weights associated with each particle and loss of diversity within an ensemble unless the ensemble size scales exponentially with the observation dimension. To improve the applicability of PFs, there have been many recent developments involving localization techniques (e.g. the reviews in Farchi and Bocquet, 2018; Graham and Thiery, 2019), incorporation of tempering and/or mutation steps (e.g. Cotter et al., 2020; Ruzayqat et al., 2022), hybrid approaches, improved computational implementations, and combination of the above with improved proposal distributions. The above ongoing efforts have extended the applicability of PF methods within geoscientific domains. Leeuwen et al. (2019) provide an overview on the integration of particle filters in high-dimensional geoscience applications.

The DA paradigm of optimally combining observations and model has a wide range of applications. However, an existing hurdle impeding the incorporation of data into pre-existing dynamical models is the lack of readily available software packages capable of bridging the two sources of information: model and observations. This is the motivation behind ParticleDA.jl: to provide a generic and user-friendly framework to enable the incorporation of particle filtering techniques with pre-existing numerical models. ParticleDA.jl is an open-source package in Julia which provides efficient implementations of several particle filter algorithms and also importantly offers an extensible framework to allow the simple addition of new filter implementations. It has been developed to be agnostic to the forward model to ensure applicability in a wide range of settings, and emphasis has been placed on computational efficiency and scalability on high-performance computing systems. Initial efforts have been focused on integration with spatially dependent numerical models; however the implementation is applicable to a much more general class of state space models (see Sect. 2), allowing incorporation in a broad range of applications.

For a specific class of state space models with additive Gaussian state and observation noise, as well as linear observation operators, ParticleDA.jl allows particle filtering with the so-called “locally optimal” proposal distribution. Though the latter is not amongst the latest contributions in the PF literature (e.g. Doucet et al., 2000), it has not been extensively explored in the systems driven by high-dimensional partial differential equations (PDEs), which we use as case studies in this work. As noted in Snyder (2011), improved proposals used within PFs can in practice significantly improve the performance of the algorithm but by themselves do not overcome the curse of dimensionality. Our numerical experiments illustrate that ParticleDA.jl can already be usefully applied in practice to models with moderately high dimensions; however, an important line of future work will be extending the framework with additional filter implementation incorporating approaches such as localization and tempering to allow scaling to very high-dimensional settings.

(Barth et al., 2016)(Le Provost, 2016)(Schauer et al., 2018)(Schoenbrod, 2018)(Carlson et al., 2018)(Sunberg et al., 2017)(Lee and Piibeleht, 2017)

Table 1 Summary of algorithms implemented and parallelism support in existing Julia data assimilation packages.

Within the Julia ecosystem, there are several existing packages which implement data assimilation algorithms. DataAssim.jl (Barth et al., 2016) provides implementations of a range of EnKF and extended Kalman filter (KF) methods and an incremental variant of 4DVar. EnKF.jl (Le Provost, 2016) implements stochastic and deterministic (square-root) variants of the EnKF which can be combined with various approaches (e.g. covariance inflation) to avoid ensemble collapse in models with deterministic dynamics. Kalman.jl (Schauer et al., 2018) and KalmanFilters.jl (Schoenbrod, 2018) both provide implementations of the exact KF algorithm for linear Gaussian models, with KalmanFilters.jl additionally implementing unscented variants of the KF for use in models with non-linear dynamics or observation operators. ParticleFilters.jl (Sunberg et al., 2017) and LowLevelParticleFilters.jl (Carlson et al., 2018) both provide PF implementations, with LowLevelParticleFilters.jl additionally providing KF implementations. SequentialMonteCarlo.jl (Lee and Piibeleht, 2017) provides an interface for implementing (and example implementations of) the wider class of sequential Monte Carlo (SMC) methods, of which PFs can be considered a special case, with the ability to run particle ensembles in parallel on multiple threads. EnsembleKalmanProcesses.jl (Dunbar et al., 2022) implements several derivative-free optimization algorithms based on the EnKF, mainly targeted at Bayesian inverse problem settings. Another package to note in the Julia ecosystem is DataAssimilationBenchmarks.jl (Grudzien and Bocquet, 2022; Grudzien et al., 2022), which offers a framework to empirically validate and develop novel DA techniques.

Table 1 summarizes the algorithm and parallelism support of existing Julia data assimilation packages along with our package ParticleDA.jl. The existing packages LowLevelParticleFilters.jl and SequentialMonteCarlo.jl, which support parallelization of operations across ensemble members, both use shared-memory parallelism, with tasks run simultaneously across multiple threads on the same device. In contrast, as described in Sect. 3.3, our package ParticleDA.jl supports both shared- and distributed-memory parallelism, which enables efficient deployment on high-performance computing (HPC) systems.

We implemented ParticleDA.jl in the Julia programming language (Bezanson et al., 2017) because of its combination of performance and productivity, which enables rapid prototyping and development of high-performance numerical applications (Churavy et al., 2022; Giordano et al., 2022), with the possibility of using both shared- and distributed-memory parallelism strategies. In particular, Julia makes use of the multiple-dispatch programming paradigm, which is particularly well-suited for designing a program which combines different models with different filtering algorithms, keeping the two concerns separated. This allows domain experts and software engineers to collaborate on the code using the same high-level language.

The rest of this paper is organized as follows. The mathematical set-up of the particle-filtering algorithm is defined in Sect. 2 along with the various filtering proposal distributions implemented. Section 3 outlines the code structure and parallelization schemes. Sections 4 and 5 illustrate applications of the framework to simple low-dimensional state space models, namely a stochastically driven damped simple harmonic oscillator model and a stochastic variant of the Lorenz 1963 chaotic attractor model (Lorenz, 1963). Section 6 introduces an application to a spatially extended state space model, specifically a tsunami modelling test case formulated as a linear Gaussian state space model, with validation of the filtering approaches and parallel performance scaling results. In Sect. 7 the incorporation with a more complex non-linear atmospheric dynamical model is investigated along with some results. Finally in Sect. 8 concluding remarks and future work are outlined.

2 Particle filtering

Let ${\mathit{x}}_t\in {\mathbb{R}}^{d_x}$ represent the state of the model at an integer time index t and ${\mathit{y}}_t\in {\mathbb{R}}^{d_y}$ the vector of observations of the system at this time index. We assume a state space model formulation, with the states following a Markov process and the observations depending only on the state at the corresponding time index; that is

$$\begin{array}{}\text{(1)}& {\mathit{x}}_{\mathrm{0}}\sim p_{\mathrm{0}}(\cdot ),{\mathit{x}}_t\sim p_t(\cdot \mathrm{|}{\mathit{x}}_{t-\mathrm{1}}),{\mathit{y}}_t\sim g_t(\cdot \mathrm{|}{\mathit{x}}_t),t\ge \mathrm{1},\end{array}$$

where $p_{\mathrm{0}}:{\mathbb{R}}^{d_x}\to {\mathbb{R}}_{\ge \mathrm{0}}$ is the density of the initial state distribution; $p_t:{\mathbb{R}}^{d_x}\times {\mathbb{R}}^{d_x}\to {\mathbb{R}}_{\ge \mathrm{0}},t\ge \mathrm{1}$ are the densities of the state transition distributions; and $g_t:{\mathbb{R}}^{d_y}\times {\mathbb{R}}^{d_x}\to {\mathbb{R}}_{\ge \mathrm{0}},t\ge \mathrm{1}$ are the densities of the conditional distribution of the observations given the current states.

A key assumption of the state space model formulation is that the state of the system evolves stochastically in time. In geophysical applications, commonly the models of interest are specified as the solution to time-dependent PDEs or systems of ordinary differential equations (ODEs), for which the state dynamics are inherently deterministic. Without a stochastic element to the dynamics, the evolution of the state over time is entirely determined by the initial state. To perform particle filtering in such models, we must therefore augment the deterministic dynamics of the model with stochastic updates. These stochastic updates can be considered random forcings of the model representing physical processes not modelled in the deterministic model as well as the discretization errors introduced when simulating ODE and PDE models (Leeuwen et al., 2019). Importantly the stochastic updates should maintain any constraints or relationships between the state variables in the underlying physical phenomena being modelled – for example for state vectors corresponding to spatial discretizations of a continuous field, the stochastic updates should maintain any assumed smoothness properties of the field.

For the most part, we will concentrate on a specialization of this general state space model class, whereby the state and observations are both subject to additive Gaussian noise and the observations depend linearly on the state, which covers a wide range of modelling scenarios in practice. Concretely we consider a state update of the form

$$\begin{array}{}\text{(2)}& {\mathit{x}}_t=F_t\left({\mathit{x}}_{t-\mathrm{1}}\right)+{\mathit{u}}_t,{\mathit{u}}_t\sim \mathcal{N}(\mathrm{0},\mathbf{Q}),t\ge \mathrm{1},\end{array}$$

where $F_t:{\mathbb{R}}^{d_x}\to {\mathbb{R}}^{d_x}$ is the forward operator at time index t, representing the deterministic dynamics of the system, and ${\mathit{u}}_t\in {\mathbb{R}}^{d_x}$ is the additive Gaussian state noise at time index t, representing stochastic aspects of the system dynamics. The observations are modelled as being generated according to

$$\begin{array}{}\text{(3)}& {\mathit{y}}_t=\mathbf{H}\left({\mathit{x}}_t\right)+{\mathit{v}}_t,{\mathit{v}}_t\sim \mathcal{N}(\mathrm{0},\mathbf{R}),t\ge \mathrm{1},\end{array}$$

where $\mathbf{H}\in {\mathbb{R}}^{d_y\times d_x}$ is a linear observation operator and ${\mathit{v}}_t\in {\mathbb{R}}^{d_y}$ is the additive Gaussian observation noise. The distributions of the state and observation noise are parameterized by positive definite covariance matrices $\mathbf{Q}\in {\mathbb{R}}^{d_x\times d_x}$ and $\mathbf{R}\in {\mathbb{R}}^{d_y\times d_y}$ respectively.

The objective of the particle filter is to estimate the filtering distribution for each time index t, which is the conditional probability distribution of the state x_t given observations ${\mathit{y}}_{\mathrm{1}},\mathrm{\dots },{\mathit{y}}_t$ up to time index t, with the density of the filtering distribution at time index t denoted ${\mathit{\pi }}_t\left({\mathit{x}}_t\right|{\mathit{y}}_{\mathrm{1}:t})$.

The particle-filtering algorithm builds on sequential importance sampling by introducing additional resampling steps. See Doucet et al. (2000) for an in-depth introduction, but the key features are introduced in Algorithm 1. An ensemble of particles $\mathit{\{}{\mathit{x}}_t^{\left(i\right)}{\mathit{\}}}_{i=\mathrm{1}}^N$ represents an approximation of the filtering distribution at each time index t≥1 as ${\mathit{\pi }}_t\left(\mathrm{d}{\mathit{x}}_t\right|{\mathit{y}}_{\mathrm{1}:t})\approx \frac{\mathrm{1}}{N}{\sum }_{i=\mathrm{1}}^N{\mathit{\delta }}_{{\mathit{x}}_t^{\left(i\right)}}(\mathrm{d}{\mathit{x}}_t)$. In each filtering step, new values for the particles are sampled from a proposal distribution (more details about the proposals implemented in ParticleDA.jl are given in Sect. 2.1) and importance weights computed for each proposed particle value. At the end of the filtering step, the weighted proposed particle ensemble is resampled to produce a new uniformly weighted ensemble to use as the input to the next filtering step.

Algorithm 1 Particle filter.

1:

Initialize particles $\mathit{\{}{\mathit{x}}_{\mathrm{0}}^{\left(i\right)}{\mathit{\}}}_{i=\mathrm{1}}^N$, with ${\mathit{x}}_{\mathrm{0}}^{\left(i\right)}\sim p_{\mathrm{0}}(\cdot )$, for $\mathrm{1}\le i\le N$.

2:

for time index t=1 to T do

3:

for particle index i=1 to N do

4:

Sample proposed particle ${\stackrel{\mathrm{̃}}{\mathit{x}}}_t^{\left(i\right)}\sim q_t(\cdot |{\mathit{x}}_{t-\mathrm{1}}^{\left(i\right)},{\mathit{y}}_t)$.

5:

Compute (unnormalized) importance weight $w_t^{\left(i\right)}=W_t({\stackrel{\mathrm{̃}}{\mathit{x}}}_t^{\left(i\right)},{\mathit{x}}_{t-\mathrm{1}}^{\left(i\right)},{\mathit{y}}_t)=\frac{p_t\left({\stackrel{\mathrm{̃}}{\mathit{x}}}_t^{\left(i\right)}\right|{\mathit{x}}_{t-\mathrm{1}}^{\left(i\right)}\left)g_t\right({\mathit{y}}_t\left|{\stackrel{\mathrm{̃}}{\mathit{x}}}_t^{\left(i\right)}\right)}{q_t\left({\stackrel{\mathrm{̃}}{\mathit{x}}}_t^{\left(i\right)}\right|{\mathit{x}}_{t-\mathrm{1}}^{\left(i\right)},{\mathit{y}}_t)}$.

6:

end for

7:

Generate new (equally weighted) particles $\mathit{\{}{\mathit{x}}_t^{\left(i\right)}{\mathit{\}}}_{i=\mathrm{1}}^N$ by resampling from the weighted empirical distribution

$${\mathit{x}}_t^{\left(i\right)}\sim {\displaystyle \frac{{\sum }_{i=\mathrm{1}}^N{w}_t^{\left(i\right)}{\mathit{\delta }}_{{\stackrel{\mathrm{̃}}{\mathit{x}}}_t^{\left(i\right)}}(\cdot )}{{\sum }_{i=\mathrm{1}}^N{w}_t^{\left(i\right)}}}.$$

8:

end for

2.1 Proposal distributions

Two forms of proposal distributions are implemented in ParticleDA.jl: the “naive” bootstrap proposal, applicable to general state space models described by Eq. (1), and the locally optimal proposal, which can be tractably computed only for a restricted class of state space models, including importantly those described by Eqs. (2) and (3).

The bootstrap proposal ignores the observations with the particle proposals sampled from the state transition distributions,

$$\begin{array}{}\text{(4)}& q_t\left({\mathit{x}}_t\right|{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)=p_t({\mathit{x}}_t\left|{\mathit{x}}_{t-\mathrm{1}}\right),\end{array}$$

with the unnormalized importance weights at time index t≥1 then simplifying to

$$\begin{array}{}\text{(5)}& W_t({\mathit{x}}_t,{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)=g_t\left({\mathit{y}}_t\right|{\mathit{x}}_t\left)\right(=W_t({\mathit{x}}_t,{\mathit{y}}_t)).\end{array}$$

While appealingly simple and applicable to a wide class of models, the bootstrap particle filter performs poorly when observations are informative about the state due to the observations being ignored in the proposal. In such cases the proposed particles will typically be far away from the mass of the true filtering distribution, with the importance weights in this setting tending to have high variance, leading to weight degeneracy whereby all the normalized importance weights but one are close to zero.

To alleviate the tendency to weight degeneracy, we can use alternative proposal distributions which decrease the variance of the importance weights. For proposal distributions $q_t\left({\mathit{x}}_t\right|{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)$ which condition only on the previous state x_t−1 and current observation y_t, the optimal proposal, in the sense of minimizing the variance of the importance weights, can be shown (Doucet et al., 2000) to be

$$\begin{array}{}\text{(6)}& q_t\left({\mathit{x}}_t\right|{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)={\displaystyle \frac{p_t\left({\mathit{x}}_t\right|{\mathit{x}}_{t-\mathrm{1}}\left)g_t\right({\mathit{y}}_t\left|{\mathit{x}}_t\right)}{\int p_t\left({\stackrel{\mathrm{̃}}{\mathit{x}}}_t\right|{\mathit{x}}_{t-\mathrm{1}}\left)g_t\right({\mathit{y}}_t\left|{\stackrel{\mathrm{̃}}{\mathit{x}}}_t\right)\mathrm{d}{\stackrel{\mathrm{̃}}{\mathit{x}}}_t}},\end{array}$$

with corresponding unnormalized importance weights

$$\begin{array}{}\text{(7)}& \begin{array}{rl}{W}_t({\mathit{x}}_t,{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)& =\int p_t\left({\stackrel{\mathrm{̃}}{\mathit{x}}}_t\right|{\mathit{x}}_{t-\mathrm{1}}\left)g_t\right({\mathit{y}}_t\left|{\stackrel{\mathrm{̃}}{\mathit{x}}}_t\right)\mathrm{d}{\stackrel{\mathrm{̃}}{\mathit{x}}}_t\\ & (=W_t({\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t\left)\right).\end{array}\end{array}$$

Note that in this case the importance weights are independent of the sampled values of the particle proposals.

For general state space models, sampling from this locally optimal proposal and computing the importance weights can be infeasible due to the integral in Eqs. (6) and (7) not having a closed-form solution. However, for the specific case of a state space model of the form described by Eqs. (2) and (3), the proposal distribution has the tractable form

$$\begin{array}{}\text{(8)}& \begin{array}{rl}{q}_t\left({\mathit{x}}_t\right|{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)& =\mathcal{N}\left({\mathit{x}}_t\left|F_t\right({\mathit{x}}_{t-\mathrm{1}})+{\mathbf{QH}}^{\mathsf{T}}\right.\\ & \cdot ({\mathbf{HQH}}^{\mathsf{T}}+\mathbf{R}{)}^{-\mathrm{1}}({\mathit{y}}_t-\mathbf{H}{F}_t\left({\mathit{x}}_{t-\mathrm{1}}\right)),\\ & \left.\mathbf{Q}-{\mathbf{QH}}^{\mathsf{T}}({\mathbf{HQH}}^{\mathsf{T}}+\mathbf{R}{)}^{-\mathrm{1}}\mathbf{HQ}\right),\end{array}\end{array}$$

with corresponding importance weights

$$\begin{array}{}\text{(9)}& W_t({\mathit{x}}_t,{\mathit{x}}_{t-\mathrm{1}},{\mathit{y}}_t)=\mathcal{N}\left({\mathit{y}}_t\left|\mathbf{H}{F}_t\right({\mathit{x}}_{t-\mathrm{1}}),{\mathbf{HQH}}^{\mathsf{T}}+\mathbf{R}\right).\end{array}$$

To generate samples from the locally optimal proposal distribution, we exploit the fact that for ${\stackrel{\mathrm{̃}}{\mathit{x}}}_t\sim \mathcal{N}\left(F_t\right({\mathit{x}}_{t-\mathrm{1}}),\mathbf{Q})$ and ${\stackrel{\mathrm{̃}}{\mathit{y}}}_t\sim \mathcal{N}(\mathbf{H}{\stackrel{\mathrm{̃}}{\mathit{x}}}_t,\mathbf{R})$ – that is $({\stackrel{\mathrm{̃}}{\mathit{x}}}_t,{\stackrel{\mathrm{̃}}{\mathit{y}}}_t)$ sampled from the joint distribution on the state and observation given the previous state x_t−1 under the state space model –

$$\begin{array}{}\text{(10)}& {\mathit{x}}_t={\stackrel{\mathrm{̃}}{\mathit{x}}}_t+{\mathbf{QH}}^{\mathsf{T}}({\mathbf{HQH}}^{\mathsf{T}}+\mathbf{R}{)}^{-\mathrm{1}}({\mathit{y}}_t-{\stackrel{\mathrm{̃}}{\mathit{y}}}_t)\end{array}$$

is distributed according to the locally optimal proposal distribution in Eq. (8). Importantly this means that to use the locally optimal proposal distribution when filtering, we only need to implement functions for sampling from the state transition and observation models and functions for evaluating the matrix terms in Eq. (10) – that is QH^𝖳 and HQH^𝖳+R. Note that unlike a direct implementation of sampling from Eq. (8) by performing a Cholesky factorization of the proposal covariance matrix, we do not need to explicitly evaluate or store a d_x×d_x covariance matrix, and we only need to perform an $\mathcal{O}\left(d_y^{\mathrm{3}}\right)$ linear solver and 𝒪(d_xd_y) matrix–vector multiplication rather than an $\mathcal{O}\left(d_x^{\mathrm{3}}\right)$ Cholesky decomposition. As well as reducing the time and memory complexity of the linear algebra operations, this approach reduces the implementation burden on a user wishing to apply the locally optimal proposal, by reusing functions required for simulating the forward model, and ensures any algorithmic efficiencies used in the implementation of simulating the forward model are also leveraged in sampling from the locally optimal proposal distribution.

In both the bootstrap and the locally optimal proposals, the stochastic nature of the state transitions is essential to maintaining diversity in the ensemble, ensuring any particles duplicated in the previous resampling step give rise to distinct proposals. For state space models with the state update and observation model described by Eqs. (2) and (3) specifically, we can see that the bootstrap and locally optimal proposals converge to the same degenerate distribution ${\mathit{\delta }}_{F_t\left({\mathit{x}}_{t-\mathrm{1}}\right)}$ as the state noise vanishes. This emphasizes the importance of using stochastic state dynamics for the PF algorithms used here to remain valid.

2.2 Resampling

A vital aspect of all PF algorithms is the resampling step shown in Algorithm 1. Resampling multiplies particles found at good positions in space that agree with observations and removes unwanted particles, concentrating computational effort on the more plausible ensemble members. It is key for establishing analytical results showing that Monte Carlo errors in estimates of expectations under the filtering distribution are controlled uniformly in time; see, for example, the standard reference Del Moral (2004). Such a result provides a critical justification for the powerful performance of PF-based algorithms in many applications. ParticleDA.jl implements a systematic resampling scheme (Douc and Cappé, 2005), which uses a single uniform random variate to resample all the particle indices.

A useful metric for capturing the variability in the weights before resampling and indicating whether weight degeneracy has occurred, is the estimated effective sample size (ESS), which is defined as

$$\begin{array}{}\text{(11)}& {\mathrm{ESS}}_t={\displaystyle \frac{\mathit{\{}{\sum }_{i=\mathrm{1}}^N{w}_t^{\left(i\right)}{\mathit{\}}}^{\mathrm{2}}}{{\sum }_{i=\mathrm{1}}^N\mathit{\{}{w}_t^{\left(i\right)}{\mathit{\}}}^{\mathrm{2}}}},\end{array}$$

where $\mathit{\{}{w}_t^{\left(i\right)}{\mathit{\}}}_{i=\mathrm{1}}^N$ denotes the unnormalized particle importance weights. The estimated ESS approximates the number of independent samples that would produce estimates of similar variance as the ones obtained by the available (correlated) particles.

3 Code structure

As previously stated, ParticleDA.jl is designed to be forward-model agnostic (i.e. capable of running with arbitrary state space models). To enable this, the model and filter portions of ParticleDA.jl are carefully delineated. The high-level structure of the main run_particle_filter function used to perform filtering with a state space model given a sequence of observations ${\mathit{y}}_{\mathrm{1}},\mathrm{\dots },{\mathit{y}}_T$ is summarized in Algorithm 2, where T is the total number of observation times. Operations which use thread-based parallelism are labelled with (∥). When run across multiple processes, operations involving communication across ranks are labelled with (↔) and those which only run on the coordinating rank are labelled with (∘). Further details on the filter and model interface and parallelization implementation are given in the following sections.

Algorithm 2 Structure of ParticleDA.jl
run_particle_filter function.

1:

Initialize model

2:

Initialize states (∥)

3:

Initialize filter

4:

Initialize summary statistics (∥)

5:

for time index t=1 to T do

6:

Sample proposal and compute (unnormalized) particle weights for current observation y_t (∥)

7:

Gather particle weights (↔)

8:

Normalize particle weights (∘)

9:

Resample particle indices according to weights (∘)

10:

Broadcast new particle indices (↔)

11:

Copy particle states according to new indices (↔)

12:

Update summary statistics ($\leftrightarrow ,\parallel$)

13:

Write outputs (∘)

14:

end for

3.1 Filter interface

Implementing a filtering algorithm in ParticleDA.jl requires providing implementations of two functions, init_filter and sample_proposal_and_compute_log_weights!, with the corresponding methods dispatched on a filter type argument which is a concrete subtype of the ParticleFilter abstract type. Implementations are currently provided for particle filters with bootstrap proposals (with corresponding type BootstrapFilter) and locally optimal proposals (with corresponding type OptimalFilter).

The init_filter method deals with initializing any filter-specific data structures, including allocating arrays to hold the particle weights and resampling indices and ensemble summary statistics. A set of shared filter parameters are passed to the init_filter method as an instance of a dedicated FilterParameters type, with functionality provided for reading these parameters from a YAML file. Key parameters include the number of particles in the ensemble, the number of tasks used when scheduling parallelizable operations in multi-threaded code segments (see Sect. 3.3), the seed for the pseudo-random number generator used to generate random variates during filtering, and the file path to write filtering outputs to, as well as options for controlling the verbosity of the filter output.

The sample_proposal_and_compute_log_ weights! method provides an implementation of generating new values for an ensemble of particles from the proposal distribution associated with the filter type and computing the corresponding unnormalized particle weights (in particular their logarithms to maintain numerical stability), corresponding to lines 4 and 5 in Algorithm 1 respectively. As hinted at by the presence of an ! suffix in the function name, a convention in Julia for indicating functions which mutate one or more of their arguments, the sample_proposal_and_compute_log_weights! function computes the updates to the arrays representing the particle states and weights in place. The proposal generation and weight computation stages are combined into a single function rather than two separate functions to allow the filter implementation to avoid redundant computations of quantities required in both sampling the proposals and computing the weights. For example, both the locally optimal proposal distribution in Eq. (8) and corresponding weights in Eq. (9) require the values of the particles after the deterministic update F_t but before addition of the state noise.

Table 2 Summary of the main functions defining the model interface. Functions in rows marked ⋆ are only required when using the locally optimal proposal.

3.2 Model interface

To support filtering in general state space models while still allowing filters to exploit additional structure in the model when present, we define an extensible model interface, with a core set of functions requiring implementation for all state space models, with model classes with additional structure able to extend this core interface. In particular we exploit this approach for conditionally Gaussian state space models having a state update and observation models of the form described by Eqs. (2) and (3) respectively to allow filtering using the locally optimal proposal distribution in Eq. (8). Table 2 summarizes the key functions requiring implementation within the core interface both for general state space models and for the restricted class of models for which the locally optimal proposal can be applied, along with a brief description of what operations they perform in the notation of Sect. 2.

A key pair of functions in Table 2 are update_state_deterministic! and update_state_stochastic!, which for general state space models when applied in sequence correspond to sampling from the state transition distribution ${\mathit{x}}_t\sim p_t(\cdot \mathrm{|}{\mathit{x}}_{t-\mathrm{1}})$, while for models with state updates of the specific form in Eq. (2), they correspond to updating the state by applying the deterministic forward operator F_t and incrementing by a state noise vector ${\mathit{u}}_t\sim \mathcal{N}(\mathrm{0},\mathbf{Q})$ respectively. While the state transitions for general state space models may not factor into a composition of deterministic and stochastic updates, full generality is still maintained as update_state_deterministic! can leave the state vector unchanged (corresponding to an identity operation) with update_state_stochastic! then solely responsible for sampling from the state transition distribution.

The functions get_covariance_state_observation _given_previous_state and get_covariance_observation_observation _given_previous_state are required for computing the locally optimal proposal update in Eq. (10). In both cases the models' implementations of these functions are required to evaluate a scalar covariance value for a pair of integer state or observation indices in 1:d_x and 1:d_y respectively. Providing the state noise covariance Q_ij can be evaluated in 𝒪(1) (with respect to d_x and d_y) operations for pairs of state indices $i,j\in \mathrm{1}:d_x$, which will typically be the case if the state noise corresponds to the discretization of a Gaussian random field with an explicit covariance function and the observation operator H is sparse (for example each observation depends on the state at only one or a few indices); then both functions can be evaluated in 𝒪(1) cost. The d_y×d_x and d_y×d_y matrices HQ and HQH^𝖳+R are then evaluated by calling the functions for grids of state and observation indices, with 𝒪(d_yd_x) and $\mathcal{O}\left(d_y^{\mathrm{2}}\right)$ costs respectively, without requiring explicit construction of the d_x×d_x matrix Q, which is particularly important when d_x is large.

As a further optimization, models can implement a get_state_indices_correlated_to_ observations function, which returns a subset of the state indices 1:d_x that excludes indices for which the corresponding state variable is uncorrelated to all observation variables (that is $i\in \mathrm{1}:d_x$ such that (HQ)_ij=0 for all values of $j\in \mathrm{1}:d_y$). For example, if the state corresponds to the discretization of a set of spatial fields with the observations corresponding to only one of these fields, then only the state indices corresponding to the observed field would need to be returned. This is then used to avoid needing to compute the known zero-covariance terms.

The functions get_covariance_observation_ noise and get_covariance_state_noise are not required directly for the locally optimal proposal update. However, typically they will be used in the definitions of other model methods, and models are required to provide implementations to allow testing the model methods for internal consistency and to support a Kalman filter implementation for linear Gaussian models.

As well as providing implementation of the functions in Table 2, models are required to implement an initialization function which is passed to the top-level run_particle_filter function and used to initialize an instance of the model data structure type that the model interface functions are dispatched on. This initialization function is passed a dictionary of model-specific parameter values read from a YAML file and the number of tasks that may be simultaneously scheduled when running model functions in parallel (see Sect. 3.3), allowing the assignment of per-task buffers for use in computing intermediate results while remaining thread-safe.

3.3 Parallelization scheme

As stated, both shared- and distributed-memory parallelization approaches can be leveraged within ParticleDA.jl to exploit both multiple processing elements sharing memory on a single node (for example central processing unit (CPU) cores) and multiple nodes (potentially each with multiple processing elements) in a cluster. Particle and weight updates are parallelized across multiple threads on shared-memory systems using the native task-based multi-threading support in Julia. In distributed-memory environments, ParticleDA.jl allows parallelizing across processes (ranks) with communication between processes performed using the Julia package MPI.jl (Byrne et al., 2021), which acts as a wrapper around a message passing interface (MPI) implementation installed on the system. MPI.jl has been found to be able to achieve little to no overhead in applications with thousands of MPI ranks (Giordano et al., 2022). ParticleDA.jl uses HDF5 files for file-based input (of the observed data used for filtering) and output (of statistics computed during filtering), using the HDF5.jl Julia package; when running in the distributed setting, file input–output is performed only on a single coordinating rank.

Figure 1 Visualization of the hierarchical parallelization model in ParticleDA.jl for an example case where updates to 24 particles are distributed across three MPI ranks. Each rank splits the particles across four tasks, with these tasks scheduled to run in parallel across two or more threads on each rank.

Figure 2 An overview of how the key stages in the filtering loop are distributed across R ranks in ParticleDA.jl. Rank 0 is the coordinating rank which mediates communication across ranks and performs file input and output. Shaded nodes indicate stages which are run in parallel across multiple threads on each rank. Edges between nodes indicate stages which involve communication across ranks.

An illustration of how per-particle operations are distributed in the two-level parallelization scheme is given in Fig. 1. Each MPI rank is assigned an equal proportion of the total number of particles in the ensemble. Within each MPI rank, operations which can be parallelized across particles are scheduled across multiple tasks, each associated with a subset of the particles assigned to the rank. The tasks are run simultaneously across multiple threads, with the flexibility in number of tasks per rank allowing a trade-off between improved load balancing across processing elements on a rank by having multiple tasks scheduled per parallel thread and the increased overhead involved in scheduling more tasks.

A sketch of the key operations in the main filtering loop and how they are distributed across multiple ranks is shown in Fig. 2. A key principle is to reduce as much as possible the requirement to communicate the full particle state vectors between ranks. Particles remain local to specific ranks for all operations other than when copying states as part of the resampling step, with this step potentially requiring particles with large weight which are duplicated after resampling to be copied point-to-point to other ranks. Communication between ranks is also required when gathering the unnormalized particle weights to the coordinating rank to allow normalization and when broadcasting the resampled particle indices from the coordinating rank to other ranks; however these operations only require communicating a single scalar per particle.

Communication between ranks is also required when computing any summary statistics of the estimated filtering distributions at each time index. ParticleDA.jl currently supports estimating the mean and, optionally, the variance of the filtering distributions for each state dimension, with a summary statistic type argument passed to the top-level run_particle_filter function allowing specification of which summary statistics to compute. Sufficient statistics of the local particles for the relevant summary statistics are computed on each rank before these local sufficient statistics are accumulated on the coordinating rank using an MPI reduce operation and are used to compute the statistics of interest. For CPU architectures for which MPI.jl supports using custom reductions (https://github.com/JuliaParallel/MPI.jl/issues/404, last access: 27 February 2023), a more numerically stable “pooled” algorithm (Chan et al., 1982) is used for computing the mean and variance (adapting the example code given in Byrne et al., 2021); implementations of the less numerically stable naive algorithms which directly accumulate the sum and sum of squares, which can be performed using standard MPI sum reductions, are also provided as a fallback for running on other CPU architectures.

Figure 3 RMSE in particle filter estimates of filtering distribution means and (log) variances against number of particles for the damped simple harmonic oscillator model. The RMSE values are calculated against ground truth values computed using a Kalman filter and are computed for the mean of the squared errors across all state components and time steps.

4 Stochastically driven damped simple harmonic oscillator

As a tractable first test case, we consider a two-dimensional state space model corresponding to the time discretization of a stochastic differential equation

$$\mathrm{d}\mathit{x}\left(\mathit{\tau }\right)=\left(\begin{array}{cc}\mathrm{0}& \mathrm{1}\\ -{\mathit{\omega }}_{\mathrm{0}}^{\mathrm{2}}& -{\mathit{\omega }}_{\mathrm{0}}/Q\end{array}\right)\mathit{x}\left(\mathit{\tau }\right)\mathrm{d}\mathit{\tau }+\left(\begin{array}{c}\mathrm{0}\\ \mathrm{1}\end{array}\right)\mathrm{d}W\left(\mathit{\tau }\right),$$

representing a damped simple harmonic oscillator driven by a Wiener noise process W(τ), with τ the (continuous) time coordinate, ω₀ the frequency of the undamped oscillator, and Q a quality factor for the oscillator. This process has been proposed as a model for astronomical time series data (Foreman-Mackey et al., 2017), with details of its formulation as a state space model given in Jordán et al. (2021). Importantly, the state space model is linear Gaussian, so we can use a Kalman filter to exactly compute the true Gaussian filtering distributions.

We use an instance of the model with parameters ω₀=1 and Q=2. We assume an observation model ${\mathit{y}}_t\sim \mathcal{N}({\mathit{x}}_t,{\mathrm{0.5}}^{\mathrm{2}}\mathbf{I})$, with x_t=x(0.2t) (that is a fixed time step 0.2 between observation times), and simulate observations from the model for T=200 time steps with initial state distribution ${\mathit{x}}_{\mathrm{0}}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$. Figure 3 shows the root mean squared errors (RMSEs) in particle filter estimates of the means and log variances of the Gaussian filtering distributions (compared to ground truth values computed using a Kalman filter), as a function of the number of particles used in the ensemble, for filters using both the bootstrap and the locally optimal proposal. We see that the locally optimal proposal gives a small but consistent improvement in RMSE for a given ensemble size, reflecting the lower variance in the empirical estimates compared to the filtering distributions. As expected, the errors in the filter estimates appear to be asymptotically tending to zero at a polynomial rate in the ensemble size, providing some assurance of the correctness of the ParticleDA.jl filter implementations.

5 Lorenz system

The Lorenz 63 system was introduced by Lorenz (1963) and is a non-linear dynamical model capturing a simplified representation of thermal convection. The model is defined by the ODE system

$$\begin{array}{}\text{(12)}& \begin{array}{rl}& {\displaystyle \frac{\mathrm{d}{x}_{\mathrm{1}}\left(\mathit{\tau }\right)}{\mathrm{d}\mathit{\tau }}}=\mathit{\sigma }\left(x_{\mathrm{2}}\right(\mathit{\tau })-x_{\mathrm{1}}(\mathit{\tau }\left)\right),\\ & {\displaystyle \frac{\mathrm{d}{x}_{\mathrm{2}}\left(\mathit{\tau }\right)}{\mathrm{d}\mathit{\tau }}}=\mathit{\rho }{x}_{\mathrm{1}}\left(\mathit{\tau }\right)-x_{\mathrm{2}}\left(\mathit{\tau }\right)-x_{\mathrm{1}}\left(\mathit{\tau }\right)x_{\mathrm{3}}\left(\mathit{\tau }\right),\\ & {\displaystyle \frac{\mathrm{d}{x}_{\mathrm{3}}\left(\mathit{\tau }\right)}{\mathrm{d}\mathit{\tau }}}=x_{\mathrm{1}}\left(\mathit{\tau }\right)x_{\mathrm{2}}\left(\mathit{\tau }\right)-\mathit{\beta }{x}_{\mathrm{3}}\left(\mathit{\tau }\right),\end{array}\end{array}$$

where τ is the time coordinate; x₁(τ),x₂(τ), and x₃(τ) are the prognostic variables of the model; and σ, ρ, and β are free parameters. As outlined by Lorenz (1963), we have set the free parameters to σ=10, ρ=28, and $\mathit{\beta }=\frac{\mathrm{8}}{\mathrm{3}}$ as this set-up will lead to chaotic behaviour. To formulate this as a state space model with state transitions of the form described by Eq. (2), we set F_t to the flow map corresponding to numerically solving the initial value problem for the ODE system in Eq. (12) over a fixed inter-observation time interval of 0.1 time units such that ${\mathit{x}}_t=\left(x_{\mathrm{1}}\right(\mathrm{0.1}t),x_{\mathrm{2}}(\mathrm{0.1}t),x_{\mathrm{3}}(\mathrm{0.1}t\left)\right)$ and use additive isotropic state noise with covariance Q=0.5²I. The Tsit5 solver with adaptive time stepping from the Julia package DifferentialEquations.jl (Rackauckas and Nie, 2017) is used to solve the ODE system. The initial state distribution is taken to be ${\mathit{x}}_{\mathrm{0}}\sim \mathcal{N}(\mathbf{0},{\mathrm{0.5}}^{\mathrm{2}}\mathbf{I})$. We assume an observation model ${\mathit{y}}_t\sim \mathcal{N}({\mathit{x}}_t,{\mathrm{1.0}}^{\mathrm{2}}\mathbf{I})$ and simulate observations for T=500 times from the model to use for filtering.

Figure 4 illustrates the performance of a filtering run with N=100 particles applied to the simulated observations using the locally optimal proposal. The left subplot shows the (noisy) observations and estimated mean of the filtering distributions at each of the observation times; note the appearance of the Lorenz attractor. The right subplot shows the evolution of the RMSE calculated for the estimated mean against the observations for different number of particles (N) with the locally optimal proposal.

Figure 4 (a) Mean of the particles and the observations at each time step in state space for N=100. (b) The RMSE calculated through time for different numbers of particles (N).

5.1 Non-linear observation operator

The effect of a non-linear observation operator on the performance of the bootstrap filter proposal is explored; note that the locally optimal proposal cannot be readily applied in this set-up, and therefore only the bootstrap filter is used. Two observation operators are introduced: the linear H(x_t)=x_t and the non-linear $H\left({\mathit{x}}_t\right)=\log \left|{\mathit{x}}_t\right|$ case.

Simulations of the Lorenz system (Eq. 12) are carried out with a similar set-up to in Sect. 5: an initial state distribution taken to be ${\mathit{x}}_{\mathrm{0}}\sim \mathcal{N}(\mathbf{0},{\mathrm{0.5}}^{\mathrm{2}}\mathbf{I})$, an observation model ${\mathit{y}}_t\sim \mathcal{N}\left(H\right({\mathit{x}}_t),{\mathrm{1.0}}^{\mathrm{2}}\mathbf{I})$, and additive isotropic state noise with covariance Q=0.5²I. The bootstrap filter is used, and the observations are assimilated every 0.1 time units. For the non-linear case, the observation operator acts upon the observation values, which are then perturbed by draws from the independent observation error.

The system is simulated for T=5000 time steps, and a time-averaged RMSE is calculated for the last 4500 steps for the linear observation, the non-linear observation, and a no-assimilation case. The time-averaged RMSE results are plotted in Fig. 5.

Figure 5 Time-averaged RMSE for three different runs with various numbers of particles: a naive ensemble (no assimilation), the linear observation case H(x_t)=x_t, and the non-linear observation case $H\left({\mathit{x}}_t\right)=\log \left|{\mathit{x}}_t\right|$.

6 Tsunami model

As a more complex test case, we now consider a tsunami modelling example. Tsunamis are rare events which have the capacity to cause severe loss of life and damage. At present, tsunami warning centres rely on crude decision matrices, pre-computed databases of high-resolution simulations, or “on-the-fly” real-time simulations to rapidly deduce the hazard associated with an event (Gailler et al., 2013). These existing approaches have been developed with seismically generated tsunamis in mind, and the alternative tsunamigenic sources (landslide and volcanic eruptions) are less well constrained. The ongoing efforts of incorporating data assimilation techniques within tsunami modelling could augment warning centres' capability in this regard (Maeda et al., 2015; Gusman et al., 2016). It should be noted that the tsunami model built into ParticleDA.jl is a drastic simplification of tsunami models used in operational practice but provides a useful test case for users and showcases the potential of particle filters within tsunami modelling efforts. Further, it allows for direct validation with a Kalman filter as the resulting state space model is linear Gaussian.

As a first-order approximation, two-dimensional linear long-wave equations (Goto, 1984), corresponding to a linearization of the shallow-water equations, are used to capture the tsunami dynamics. Assuming a state space model with state transitions of the form described in Eq. (2), the (linear) deterministic forward operator F_t(x_t) for the test case is defined by numerically solving the PDEs

$$\begin{array}{}\text{(13)}& \begin{array}{rl}{\displaystyle \frac{\partial \mathit{\eta }(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial \mathit{\tau }}}& =-{\displaystyle \frac{\partial u(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial s_{\mathrm{1}}}}-{\displaystyle \frac{\partial v(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial s_{\mathrm{2}}}},\\ {\displaystyle \frac{\partial u(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial \mathit{\tau }}}& =-gh(s_{\mathrm{1}},s_{\mathrm{2}}){\displaystyle \frac{\partial \mathit{\eta }(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial s_{\mathrm{1}}}},\\ {\displaystyle \frac{\partial v(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial \mathit{\tau }}}& =-gh(s_{\mathrm{1}},s_{\mathrm{2}}){\displaystyle \frac{\partial \mathit{\eta }(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})}{\partial s_{\mathrm{2}}}},\end{array}\end{array}$$

where τ is the time coordinate, (s₁,s₂) denotes the spatial coordinates, $\mathit{\eta }(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})$ is the free-surface elevation (wave height), h(s₁,s₂) is the (static) water depth, g is the acceleration due to gravity, and $u(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})$ and $v(\mathit{\tau },s_{\mathrm{1}},s_{\mathrm{2}})$ are the components of the depth-averaged horizontal velocities. The system of linear PDEs in Eq. (13) is solved using a first-order finite-difference scheme with absorbing boundary conditions, using a Julia reimplementation of the tsunami data assimilation code (TDAC) accompanying Gusman et al. (2016).

The state vector x_t is defined as the concatenation of the flattened vectors formed by the spatial discretizations of the fields η, u, and v on a 51×51 uniform grid over a square spatial domain $[\mathrm{0},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}]\times [\mathrm{0},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}]$ (resulting in an overall state dimension $d_x=\mathrm{3}\times {\mathrm{51}}^{\mathrm{2}}=\mathrm{7803}$), with uniform interval of 2 time units between observation times; that is

$$\begin{array}{rl}{\mathit{x}}_t=(& \mathit{\eta }(\mathrm{2}t,\mathrm{0},\mathrm{0}),\mathit{\eta }(\mathrm{2}t,\mathrm{0},\mathrm{4}\times {\mathrm{10}}^{\mathrm{3}}),\mathrm{\dots },\mathit{\eta }(\mathrm{2}t,\mathrm{0},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}),\\ & \mathit{\eta }(\mathrm{2}t,\mathrm{4}\times {\mathrm{10}}^{\mathrm{3}},\mathrm{0}),\mathrm{\dots },\mathit{\eta }(\mathrm{2}t,\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}),\\ & u(\mathrm{2}t,\mathrm{0},\mathrm{0}),u(\mathrm{2}t,\mathrm{0},\mathrm{4}\times {\mathrm{10}}^{\mathrm{3}}),\mathrm{\dots },u(\mathrm{2}t,\mathrm{0},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}),\\ & u(\mathrm{2}t,\mathrm{4}\times {\mathrm{10}}^{\mathrm{3}},\mathrm{0}),\mathrm{\dots },u(\mathrm{2}t,\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}),\\ & v(\mathrm{2}t,\mathrm{0},\mathrm{0}),v(\mathrm{2}t,\mathrm{0},\mathrm{4}\times {\mathrm{10}}^{\mathrm{3}}),\mathrm{\dots },v(\mathrm{2}t,\mathrm{0},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}}),\\ & v(\mathrm{2}t,\mathrm{4}\times {\mathrm{10}}^{\mathrm{3}},\mathrm{0}),\mathrm{\dots },v(\mathrm{2}t,\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}},\mathrm{2}\times {\mathrm{10}}^{\mathrm{5}})).\end{array}$$

The additive state noise is chosen as the spatial discretizations of independent Gaussian random fields for each of the variables η, u, and v, with a Matérn covariance kernel with length-scale parameter λ=500, smoothness parameters μ=2.5, and marginal standard deviation parameter σ=0.01 used for all three fields. The spatially correlated nature of the state noise distribution ensures the perturbed spatial fields remain smooth. A circulant embedding method (Dietrich and Newsam, 1997) implemented in the Julia package GaussianRandomFields.jl (Robbe, 2017) is used to efficiently simulate Gaussian random fields on a uniform grid using fast Fourier transforms, resulting in an 𝒪(d_xlog d_x) operation cost complexity for each realization. For filtering, the initial state distribution is also chosen to correspond to a zero-mean Gaussian distribution corresponding to the spatial discretizations of independent Gaussian random fields for each of the variables η, u, and v, with a Matérn covariance kernel with the same parameters $(\mathit{\lambda },\mathit{\mu },\mathit{\sigma })$ as above.

We assumed noisy pointwise observations of the free-surface-elevation field η at 15 “station” locations $\mathit{\{}{s}_{\mathrm{1}}^{\left(m\right)},s_{\mathrm{2}}^{\left(m\right)}{\mathit{\}}}_{m=\mathrm{1}}^{\mathrm{15}}$, chosen as grid points randomly sampled from a uniform distribution over the spatial grid for simplicity, with independent observation noise with standard deviation 0.01; that is ${\mathit{y}}_t\sim \mathcal{N}\left(\left(\mathit{\eta }\right(\mathrm{2}t,s_{\mathrm{1}}^{\left(m\right)},s_{\mathrm{2}}^{\left(m\right)}){)}_{m=\mathrm{1}}^{\mathrm{15}},{\mathrm{0.01}}^{\mathrm{2}}\mathbf{I}\right)$. For the simulation of the observations, to produce an initial wave-producing perturbation, the mean of the initial state distribution for the free-surface-elevation components is altered to correspond to the function

$$\begin{array}{}\text{(14)}& {\overline{\mathit{\eta }}}_{\mathrm{0}}(s_{\mathrm{1}},s_{\mathrm{2}})=\left\{\begin{array}{ll}\left((\mathrm{1}+\cos (\mathit{\pi }(s_{\mathrm{1}}-a)/c\left)\right)\right.& \\ \left.(\mathrm{1}+\cos (\mathit{\pi }(s_{\mathrm{2}}-a)/c\left)\right)\right)d/\mathrm{4}& (s_{\mathrm{1}}-a{)}^{\mathrm{2}}+(s_{\mathrm{2}}-a{)}^{\mathrm{2}}\le c^{\mathrm{2}},\\ \mathrm{0}& \text{otherwise},\end{array}\right.\end{array}$$

evaluated at the grid points with a=10⁴, b=10⁴, $c=\mathrm{3}\times {\mathrm{10}}^{\mathrm{4}}$, and d=30, with the mean of the velocity components left as zero. As the initial state distribution assumed when filtering differs, we have a small degree of model mismatch. The observations are simulated for T=640 times, with the PDE system numerically integrated in time for four time steps of 0.5 time units between each pair of observation times.

Figure 6 Snapshots of the surface elevation used to generate the simulated observations (a) and the corresponding estimated filtering distribution means using N=50 particles (b).

6.1 Validation

We performed an initial filtering run on the simulated observations using an ensemble of N=50 particles using the locally optimal proposals. Snapshots of the simulated free-surface-elevation field used to generate the observations and the corresponding particle estimate of the mean of the filtering distribution on the free-surface-elevation field are shown in Fig. 6.

Figure 7 (a) RMSE in particle filter estimates of filtering distribution across observation times for the tsunami models with a fixed ensemble size of N=50 for filters using both bootstrap and locally optimal proposals. (b) RMSEs in particle filter estimates of the filtering distribution mean for both the bootstrap and the locally optimal proposal at τ=200 for varying ensemble size N. Note that the RMSE values are calculated against the true mean of the filtering distributions coming from a Kalman filter run.

As the tsunami state space model implemented here is linear Gaussian, a Kalman filter was used to compute ground truth values for the means (and covariances) of the filtering distributions, and these were then compared to the filtering estimates for various ensemble sizes N and proposal distributions. Figure 7a shows the RMSE in the estimate of the filtering distribution mean for each observation time for PF using both the locally optimal and the bootstrap proposals with the same ensemble size (N=50). The filter using the locally optimal proposal can be observed to give a consistent improvement in the accuracy of the filtering distribution estimates across time. Figure 7b instead shows the RMSE in the estimate of the filtering distribution mean at a single observation time τ=200, for filtering runs with varying ensemble sizes N for both bootstrap and locally optimal proposals; the results indicate a consistent gain in accuracy of the filtering estimates when using the locally optimal compared to bootstrap proposal, across a range of different ensemble sizes N.

6.2 Parallelization performance

As discussed in Sect. 3.3, ParticleDA.jl is capable of leveraging both shared- and distributed-memory parallelism. Scaling runs on ARCHER2, which is the UK's Tier-1 supercomputer, have been carried out to highlight the performance in practice. A weak-scaling study, using the same experimental set-up as described in Sect. 6.1, is run with the bootstrap proposal, keeping the number of particles per core constant while increasing the number of nodes. The compute nodes on ARCHER2 consist of 2× AMD EPYC 7742, 2.25 GHz, 64-core, with 8 non-uniform memory access (NUMA) regions per node (16 cores per NUMA region, 8 cores per core complex die (CCD) and 4 cores per core complex (CCX) (shared L3 cache)). The weak-scaling runs try to optimize for this hardware architecture with various runs targeting an MPI rank per NUMA/CCD/CCX region and an appropriate number of threads per MPI rank (Fig. 8). The weak-scaling efficiency is defined as $E\left(N\right)=\frac{T\left(\mathrm{2}\right)}{T\left(N\right)}$, where T(N) is the wall time for running on N MPI ranks. There are 2 particles per core, so at the maximum number of cores (2048) and ranks (128) tested here, there are 4096 particles.

As stated in Sect. 3.3, the main performance bottlenecks are the communication steps: the copying of states and the gathering of particle weights, which require point-to-point communications and a global communication step respectively. Another component which contributes to poor scaling for large node counts is updating the summary statistics, which requires a reduction in the mean and optionally the variance for each state dimension over all MPI ranks. For the results presented here, we have mostly remedied this loss of performance by collecting statistics at the final filtering iteration only. However, it should be noted that for cases which need frequent outputted statistics, this will contribute to a degradation of the parallel performance.

Figure 8 Weak-scaling parallel efficiency for the tsunami model test case with different set-ups of ranks per node and cores per rank on ARCHER2. A drop-off in performance can be seen when moving from single-node to multi-node runs.

7 Atmospheric general circulation model (AGCM)

An integration of ParticleDA.jl with an atmospheric dynamical model, Simplified Parameterizations, primitivE-Equation DYnamics (SPEEDY), showcases the efforts involved in coupling the software with pre-existing model implementations. SPEEDY is an AGCM which was developed by Molteni (2003) and consists of a spectral primitive-equation dynamic core along with a set of simplified physical parameterization schemes. The SPEEDY model retains the core characteristics of the current state-of-the-art AGCMs but requires drastically reduced (orders of magnitude) computational resources (Molteni, 2003). This computational efficiency allows one to utilize the model to carry out large ensemble and/or data assimilation experiments. According to Molteni (2003) the SPEEDY model accurately simulates the general structure of global atmospheric circulation and exhibits similar systematic errors to the state-of-the-art AGCM, albeit with larger error amplitudes. The model implementation used here (Hatfield, 2018) is written in Fortran and provides an interesting example of the integration steps required to interface with ParticleDA.jl. The coupling with ParticleDA.jl relies on the SPEEDY implementation being set up to output its data fields at set intervals.

As stated, SPEEDY is a simplified AGCM model. The prognostic variables consist of the zonal and meridional wind velocity components (u,v), temperature (T), specific humidity (q), and surface pressure (p_s). A T30 resolution of the model is used here, which corresponds to a horizontal grid size of 96×48 with eight vertical layers. The vertical layers are defined by sigma levels, where the pressure is normalized by the surface pressure ($p/p_{\mathrm{s}}$).

We extend the deterministic SPEEDY model to a state space model setting using a state transition update of the form described by Eq. (2), with numerical simulation of the SPEEDY model forward in time by 6 simulated hours corresponding to the deterministic forward operator F_t. The state vector x_t is defined as the concatenation of the flattened vectors corresponding to the spatial discretizations of the prognostic variable fields u, v, T, q, and p_s, with each of the first four variables being defined in three dimensions across a $\mathrm{96}\times \mathrm{48}\times \mathrm{8}$ spatial grid, while the final surface pressure variable p_s is defined in two dimensions on a 96×48 spatial grid, resulting in an overall state dimension of $d_x=\mathrm{4}\times \mathrm{96}\times \mathrm{48}\times \mathrm{8}+\mathrm{96}\times \mathrm{48}=\mathrm{152}\mathrm{064}$.

The additive Gaussian state noise is assumed to correspond to spatial discretizations of independent two-dimensional Gaussian random fields for the surface pressure p_s and for each vertical level for the prognostic variables u, v, T, and q. To reflect the underlying spherical geometry over which the spatial grid is defined, a non-stationary covariance function using a Matérn kernel applied to the geodesic (great-circle) distance between the points on the sphere the grid points correspond to is used, with the Matérn kernels using common values of λ=1 and μ=2.5 for the length scale and smoothness parameters respectively, while the marginal standard deviation parameter σ is set separately for each prognostic variable, with σ=1 m s⁻¹ for u and v, σ=1 K for T, σ=0.001 kg kg⁻¹ for q, and σ=100 Pa for p_s. The GaussianRandomFields.jl package is again used to generate realizations of the (spatially discretized) random fields, with the use of a non-stationary covariance function in this case necessitating an approach which uses an eigendecomposition of the full 4608×4608 covariance matrix for each discretized two-dimensional field to generate the samples. The geodesic-distance-based covariance function used is not guaranteed to be positive definite, which is heuristically dealt with by setting all negative eigenvalues to zero. We recognize that this approach of introducing state noise into the dynamics has some limitations but for the purposes here is sufficient.

The data assimilation experiments carried out here followed a similar set-up to that used in Miyoshi (2005). A linear Gaussian observation model of the form described by Eq. (3) is used, with observations assumed to be available only for the surface pressure field (in this regard differing from the set-up used by Miyoshi, 2005) at 50 spatial locations corresponding to randomly sampled grid points, with additive independent observation noise with standard deviation 1000 Pa. The initial state used to generate the simulated observations is generated by performing 1-simulated-year “spin-up” of the deterministic SPEEDY model from a resting atmosphere ($u=v=\mathrm{0}$) initial condition, with simulated observations then generated for 250 observation times at 6-hourly intervals using the state space model. Initial state values for a filtering run using N=256 particles and locally optimal proposals were generated by performing a long-term (10 simulated years) run of the deterministic SPEEDY model, with the state selected randomly from the simulated times in the final month of simulation and state noise of the same distribution used in state transitions added.

7.1 Results

In Fig. 9 snapshots of the true surface pressure (top left) and the ensemble estimate of the mean surface pressure (top right) after 250 assimilation cycles are shown. Minimal differences can be observed. The subplots in the bottom row showcase the time-averaged L₂ error for ensemble mean estimates with and without assimilation. The L₂ error is calculated against the true-surface-pressure fields used to simulate the observations at each grid cell over the 250 assimilation cycles. The time-averaged errors are dominated by mid-latitude patterns, but the ensemble run without assimilation exhibits larger errors. This error comparison validates the claim that the assimilation is giving improved estimates of the state of the system.

As stated in the Introduction, one of the key benefits of particle filters is to provide the promise of non-linear and non-Gaussian DA. To highlight this, sample distributions of the surface pressure at various observation locations at different time points are shown in Fig. 10. The distributions across the N=256 particles exhibit heavy tails towards the true surface pressure at the given locations. It should be noted that similar non-Gaussian distributions were showcased by Miyoshi et al. (2014) and Kondo and Miyoshi (2019) in a near-identical experimental set-up but with an ensemble Kalman filter. However, a key difference to be highlighted here is the relative size of the ensembles used, 256 particles here versus a 10 240 ensemble size used by Miyoshi et al. (2014) to generate the non-Gaussian distributions.

Figure 9 (a, b) Snapshot of the true surface pressure (a) and mean assimilated surface pressure (b) (N=256) after 250 assimilation cycles (12:00:00 UTC, 4 March 1982). (c, d) Time-averaged L₂ error for the mean of the assimilation run (c) and for the mean of an ensemble run without assimilation (d). The 50 observation locations are highlighted by the black stars.

Figure 10 Normalized histograms corresponding to the estimates of the marginal filtering distributions of the surface pressure at various observation locations and at different points in time from a filtering run with an ensemble of N=256 particles. The true surface pressures are highlighted by the vertical red line, and the mean surface pressures of the particles are highlighted by the vertical black line. The green line represents a fitted Gaussian distribution.

8 Conclusions

We have developed a flexible Julia package, ParticleDA.jl, for performing particle-filter-based data assimilation, with the potential of offering improved filtering accuracy when working with models exhibiting non-Gaussianity in the filtering distributions. The use of a high-level language Julia simplifies the process both for users wanting to apply the package to their own models and for developers wishing to extend the package with new filter implementations while still maintaining similar computational efficiency to lower-level compiled languages like Fortran and C++.

Particular attention has been paid to ensuring ParticleDA.jl is suitable for performing filtering on HPC systems, with a versatile two-level model used to support both shared- and distributed-memory parallelism. This is important in allowing efficient exploitation of the typically complex hierarchies of processing elements used in modern HPC systems (see for example the description of the hardware architecture of ARCHER2 in Sect. 6.2), both when running large ensembles of models where each particle can be simulated on a single processing element and for the perhaps more practically relevant setting of running smaller ensembles of more complex models which require multiple processing elements to simulate a single particle.

ParticleDA.jl currently provides implementations of particle filters using bootstrap and locally optimal proposals, with the former applicable to general state space models and the latter to a more restricted subset of state space models with Gaussian state transition distributions and linear Gaussian observation models. As illustrated in our numerical experiments, particle filters using the locally optimal proposal distributions can offer significantly improvements in the accuracy of filtering estimates for a given ensemble size where applicable. However, as noted in the Introduction, particle filters using the locally optimal proposal distribution are known to still suffer from a curse of dimensionality requiring the ensemble size to scale exponentially with the system dimension to avoid weight degeneracy (Snyder, 2011). An important future extension to ParticleDA.jl will therefore be in providing implementations of filtering algorithms exploiting approaches such as spatial localization (Farchi and Bocquet, 2018) to allow scaling to very high-dimensional geophysical applications. Implementations of filters exploiting spatial localization will be necessarily applicable to a restricted subclass of spatially extended state space models; similar to the approach used for implementing the locally optimal proposal filter, the extensible nature of the model interface in ParticleDA.jl should allow model-agnostic localized filter implementations to be added by simply defining additional functions required to be implemented by the model interface.

Another key element that a user should be aware of is the generation of state noise and the role that it plays (Evensen et al., 2022). For some geophysical applications, this can be a non-trivial task as the definition of the state noise should respect the smoothness of the state variable and any underlying physical constraints. For example, particular efforts have been made in the AGCM case (Sect. 7) to capture the underlying spherical geometry by implementing a non-stationary covariance function using a Matérn kernel based on the geodesic distance between points on the sphere. This approach of introducing state noise into the dynamics offers an improvement to a stationary covariance function but still has limitations and does not guarantee that physical constraints are conserved. This example highlights the important role of the state noise, and potential users should be aware of the efforts needed to accurately capture this in their systems of interest.

Overall, the aim of our platform is to enable easily accessible and accurate, fast DA for a wide range of users. We hope that various scientific communities will adopt ParticleDA.jl, possibly leading to fast step changes in some geoscientific investigations and beyond.