Draft:ODE filter

Visualization of the filtering procedure for the logistic ordinary differential equation. The first column shows the prior, which is a two-times-integrated Wiener process. The last row shows the residual, which is equal to the measurement operator (i.e., x(t) − f(x(t))). The second column shows the prior after being initialized by the initial condition. The third row shows the posterior distribution obtained by running an extended Kalman filter of order one for the ODE. All subplots display the mean (solid, thick line) and the 95% confidence interval (shaded area), as well as ten samples from the corresponding distribution (solid lines). The dashed black line shows the ground truth solution of the logistic ODE. The blue dots represent the observed measurements.

ODE filters are a class of probabilistic numerical methods for solving ordinary differential equations (ODEs) that frame the problem of finding a solution to an initial value problem as a Bayesian inference task. Solutions are modeled as probability distributions that also account for the discretization error introduced through the numerical approximation. This probabilistic treatment provides an computation-aware alternative to classical numerical ODE-solvers, which typically only provide point-wise error bounds. It further enables sampling of joint trajectories from the posterior, as well as quantifying uncertainty about the underlying ODE itself. ODE filters offer a flexible framework for incorporating additional information, such as measurements or conservation laws.

The general ODE filtering procedure consists of two steps: The first is the definition of a Prior distribution, for the ODE solution, in the form of a stochastic processes, more specifically, as Gauss–Markov processes. Discretization results in nonlinear Gaussian state space models. Second, use a general Bayesian filtering and smoothing algorithm, from the field of recursive Bayesian estimation, to find numerical solutions to the ODE. The naming convention of ODE filters typically reflects the underlying filtering algorithm: For instance, combining the particle filter with this framework yields the particle ODE filter. Since smoothers are extensions of filters, the literature often refers to both under the umbrella term "ODE filters".

Theory

Problem setup

Consider a first-order ordinary differential equation (ODE) initial value problem (IVP) of the form

$${\dot {y}}(t)=f(y(t),t),\quad t\in [0,T],\quad y(0)=y_{0}\in \mathbb {R} ^{d}.$$

The vector field ${\textstyle f:\mathbb {R} ^{d}\times [0,T]\rightarrow \mathbb {R} ^{d}}$ is assumed to be Lipschitz-continuous such that a unique global solution to this ODE-IVP exists according to the Pickard-Lindelöf theorem . Classical Numerical ODE Solvers, are algorithms that compute approximate solutions ${\textstyle {\hat {y}}(t)\approx y(t)}$ on a discrete mesh ${\textstyle \{t_{n}\}_{n=0}^{N},\quad 0=t_{0}<\cdots <t_{N}=T,\quad n\in \mathbb {N} }$. Probabilistic ODE Solvers additionally estimate the uncertainty introduced through the discretization. ODE Filters are a type of probabilistic ODE solvers, that adopt a Bayesian Inference framework to compute a posterior:

$$p\left(y(t)\mid y(0),\{{\dot {y}}(t_{n})=f(y(t_{n}),t_{n})\}_{n=0}^{N}\right).$$

This is done in two steps. First, prior data and likelihood is defined by modeling solutions to the ODE as a Gauss-Markov process. Second, the posterior is computed with Bayesian filtering and smoothing algorithms.

Step 1: Gauss-Markov process

Prior

The prior is specified by a linear time invariant (LTI) stochastic differential equation (SDE) of the form

$$X(0)\sim {\mathcal {N}}(\mu _{0},\Sigma _{0})$$

$$\mathrm {d} X(t)=FX(t)\mathrm {d} t+L\mathrm {d} B_{t}.$$

It describes a stochastic process ("the system"):

$$X(t)=\left[\left(X^{(0)}(t)\right)^{T},\cdots ,\left(X^{(q)}(t)\right)^{T}\right]^{T}\in \mathbb {R} ^{(q+1)d}.$$

The state of the system are realizations $x(t)\sim X(t)$, where $X^{(0)}(t)$ and $X^{(1)}(t)$ model ,$y(t)$ and ${\dot {y}}(t)$ respectively. Te remaining $q-1$ sub-vectors may be used to model higher-order derivatives. $F$ is a state transition matrix, $L$ a diffusion matrix, and $B(t)$ is a vector standard Wiener process. Note that $F$ and $L$ need to be chosen such that $\mathrm {d} X^{(0)}(t)=X^{(1)}(t)\mathrm {d} t$ holds. Solutions to LTI-SDEs are Gauss-Markov processes (GMP), of the form:

$$p\left(x(t)\mid X^{(0)}(0)=x_{0}=y_{0}\right)=GP\left(x(t)\mid \mu (t),k(t,t')\right)$$

$$\mu (t)=exp(Ft)x_{0}$$

$$k(t,t')=\int _{0}^{\mathrm {min(t,t')} }exp(F(t-\tau ))LL^{T}exp(F(t'-\tau ))^{T}\mathrm {d} \tau .$$

Where $\mu :[0,T]\rightarrow \mathbb {R} ^{(q+1)d}$ and $k:[0,T]\times [0,T]\rightarrow \mathbb {R} ^{(q+1)d\times (q+1)d}$. Gauss-Markov processes are computationally favorable over generic Gaussian process priors, as they allow for inference with linear time complexity ${\mathcal {O}}(N)$ compared to the general GP inference complexity of ${\mathcal {O}}(N^{3})$. Discretization of the GMP results in a linear Gaussian state space model (LGSSM):

$$x(t_{0})\sim {\mathcal {N}}(\mu _{0},\Sigma _{0})$$

$$x(t_{n+1})\mid x(t_{n})\sim {\mathcal {N}}(A(t_{n+1}-t_{n})x(t_{n}),Q(t_{n+1}-t_{n}))$$

$$A(h)=exp(Fh)$$

$$Q(h)=\int _{0}^{h}e^{F\tau }LL^{T}e^{F^{T}\tau }\mathrm {d} \tau .$$

Data

Information about the ODE is introduced through a measurement operator (i.e. state misalignment), known as an information operator in this context,

$$Z(t)=g(X(t))=X^{(1)}(t)-f(X^{(0)}(t),t).$$

If it were possible to condition $X(t)$ on the event $0=z(t)\sim Z(t)$for all $t\in [0,T]$, this would result in a point-mass on the true solution . This is computationally intractable (just like for all other forms of numerical simulation), which motivates the discretization of the problem and conditioning the state on discrete observations ${\textstyle \{z(t_{n})=0\}_{n=0}^{N}.}$ The information operator can be modified to solve higher-order ODEs. Multiple information operators can be added to incorporate additional information, such as conservation laws or measurement data.

Likelihood

The likelihood of the observations given the GMP prior is:

$$z_{n}\mid x(t_{n})\sim \delta (z_{n}-g(x(t_{n}))).$$

To improve the model's generality and account for numerical errors introduced by discretization, a measurement variance $R_{n}$ can be added to $z_{n}$ . The observations can be incorporated into the LGSSM as a nonlinear measurement model, yielding the final nonlinear Gaussian state space model (NLGSSM), that is the basis of all ODE filters:

$$x_{0}\sim {\mathcal {N}}(\mu _{0},\Sigma _{0})$$

$$x_{n+1}\mid x_{n}\sim {\mathcal {N}}(A(h_{n+1})x_{n},Q(h_{n+1}))$$

$$z_{n}\mid x_{n}\sim {\mathcal {N}}(g_{n}(x_{n}),R_{n}).$$

Here $h_{n+1}=t_{n+1}-t_{n}$, $x_{n}=x(t_{n})$, $z_{n}=z(t_{n})$, and $g_{n}(x_{n})=x_{n}^{(1)}-f(x_{n}^{(0)},t_{n})$. The desired posterior $p\left(x(t)\mid x_{0},\{z_{n}\}_{n=0}^{N}\right)$ can now be computed iteratively by any filtering and smoothing algorithms known from signal processing or recursive Bayesian estimation.

Step 2: Bayesian filter and smoother

The most general ODE filter and smoother algorithms to solve the NLGSSM are equivalent to the general recursive Beyesian estimation algorithms:

Algorithm: ODE Filter$(f,x_{0},p(x_{n+1}\mid x_{n}))$
 1: Initialize $p(x_{0})$
 2: for ${\textstyle n=0,1,\cdots ,N-1}$ do:
 3:     optional: adapt dynamic model ${\textstyle p(x_{n+1}\mid x_{n})}$
 4:     optional: choose step size $h_{n}>0$
 5:     predict ${\textstyle p(x_{n+1}\mid z_{1:n})=\int p(x_{n+1}\mid x_{n})p(x_{n}\mid z_{1:n})\mathrm {d} x_{n}}$
 6:     observe the ODE $z_{n+1}=0$
 7:     update $p(x_{n+1}\mid z_{1:n+1})=p(z_{n+1}\mid x_{n+1})p(x_{n+1}\mid z_{1:n})p(z_{n+1})^{-1}$
 8:     optional: backward transition $p(x_{n}|x_{n+1},z_{1:n})=p(x_{n+1}|x_{n})p(x_{n}|z_{1:n})p(x_{n+1}|z_{1:n})^{-1}$
 9: end for
10: return ${\textstyle \{p(x_{n}\mid z_{1:n}),p(x_{n}\mid z_{1:n-1})\}_{n=0,\dots ,N}}$

Algorithm: ODE Smoother$(f,x_{0},p(x_{n+1}\mid x_{n}))$
 1: ${\textstyle \{p(x_{n}\mid z_{1:n}),p(x_{n}\mid z_{1:n-1})\}_{n=0,\dots ,N}\;=\;{\text{ODE filter}}(f,x_{0},p(x_{n+1}\mid x_{n}))}$ 
 2: for $n=N-1,N-2,\cdots ,0$ do:
 3:     compute ${\textstyle p(x_{n}\mid z_{1:N})=p(x_{n}\mid z_{1:n})\int p(x_{n+1}\mid x_{n})p(x_{n+1}\mid z_{1:N})p(x_{n+1}\mid z_{1:n})^{-1}\mathrm {d} x_{n+1}}$
 4: end for
 5: return $\{p(x_{n}\mid z_{1:N})\}_{n=0,\dots ,N}$ 

For nonlinear observation models, the probability distributions become non-Gaussian, rendering the above equations intractable. Therefore, to define a specific ODE filter and smoother, the predict, update, and compute steps must be specified in a tractable form. The full posterior can be computed from

$$p(x_{0:N}\mid z_{1:N})=p(x_{N}\mid z_{1:N})\prod _{n=0}^{N-1}p(x_{n}\mid x_{n+1},z_{n}).$$

Having access to the full posterior enables sampling of joint trajectories that classical numerical methods do not allow.

Implementation Details

Choice of prior

The choice of a prior distribution is replaced by choosing the free parameters of the GMP. That is choosing $\mu _{0},\Sigma _{0},F,L,q$, however, under some constraints. One of the most popular and most widely used GMPs is the q-times integrated Wiener process (i.e. integrated Brownian motion) . For simplicity we can set w.l.o.g. $d=1$. The q-times Integrated Wiener process (IWP) is defined such that the sub-coordinates model the first q-derivatives of $X^{(0)}(t),$ i.e.$dX^{(n-1)}=X^{(n)}dt\quad \forall n\in [1,q]$. Further, $X^{(q)}$ is a standard Wiener Process. This is equivalent to the drift and dispersion matrices:

$${\breve {F}}_{\mathrm {IWP} }=\left[{\begin{array}{cccc}0&1&\cdots &0\\\vdots &\vdots &\ddots &\vdots \\0&0&\cdots &1\\0&0&\cdots &0\end{array}}\right],\quad {\breve {L}}_{\mathrm {IWP} }=\sigma \left[{\begin{array}{cccc}0\\\vdots \\0\\1\end{array}}\right].$$

Note that for $d>1$ the full drift and dispersion matrices follow from the one-dimensional versions with a Kronecker product $F_{\mathrm {IWP} }={\breve {F}}_{\mathrm {IWP} }\otimes I_{d}$, and $L_{\mathrm {IWP} }={\breve {L}}_{\mathrm {IWP} }\otimes I_{d}$. This yields the following known closed form solutions for the transition matrices:

$$[A(h)]_{ij}=[\exp(Fh)]_{ij}=\mathbb {I} (j\geq i){\frac {h^{j-i}}{(j-i)!}},$$

$$[Q(h)]_{ij}=\sigma ^{2}{\frac {h^{2q+3-i-j}}{(2q+3-i-j)(q+1-i)!(q+1-j)!}}.$$

The ideal initialization of the IWP that considers all available information from the ODE with initial value $x_{0}$ is to evaluate all derivatives at $x_{0}$:

$$\mu _{0}=[x_{0},f(x_{0}),f^{<2>}(x_{0}),\cdots ,f^{<q>}(x_{0})]^{T},\quad \Sigma _{0}={\boldsymbol {0}}.$$

Here $f^{<i>}(a)=[\nabla _{x}f^{<i-1>}\odot f](a)$ is recursively defined from $f^{<0>}(a)=a$ and $f^{<1>}(a)=f(a)$, and $\odot$ is the elementwise product. Computationally this can be efficiently computed via Taylor mode automatic differentiation .

Choice of filter and smoother

One possible approach to building efficient inference algorithms for nonlinear measurement operators $g(X(t))$ is via linearization. This results in a Gaussian inference framework, which is computationally desirable because it allows efficient closed-form inference. Taylor approximation provides one viable linearization technique, whereas the the degree and the locality/globality result in different subcategories of extended Kalman filters. Choosing a local, first order Taylor expansion corresponds to the extended Kalman filter of order 1 (EKF1). Other options are a Taylor approximation of order 0, leading to the EKF0 algorithm, or a global linearization, known as the iterated extended Kalman smoother (IEKS) .

Example: EKF1

The first-order Taylor expansion of $g$ around a point $\xi \in \mathbb {R} ^{(q+1)d}$ is:

$${\hat {g}}(X)=g(\xi )+J_{g}(\xi )(X-\xi ),$$

where $J_{g}(\xi )\in \mathbb {R} ^{d\times (q+1)d}$ is the Jacobian of $g$. This leads to an affine observation model

$$z_{n}\mid x_{n}\sim g(z_{n}\mid x_{n})={\mathcal {N}}(H_{n}x_{n}+c_{n},R_{n}),\quad H_{n}=J_{g}(\xi _{n}),\quad c_{n}=g(\xi _{n})-J_{g}(\xi _{n})\xi _{n}.$$

The linearization point is chosen as the predicted mean from the previous observation $\xi _{n}=m_{n}^{-}=\mathbb {E} [p(x_{n}\mid z_{1:n-1})]$. With this, the predict, update, and compute equations all take on Gaussian form, for which closed form inference is possible. See extended Kalman filter for more details.

$$\mathrm {\textbf {predict}} :p(x_{n+1}\mid z_{1:n})={\mathcal {N}}(x_{n+1}\mid m_{n+1}^{-},P_{n+1}^{-})$$

$$\mathrm {\textbf {update}} :p(x_{n+1}\mid z_{1:n+1})={\mathcal {N}}(x_{n+1}\mid m_{n+1},P_{n+1})$$

$$\mathrm {\textbf {compute}} :p(x_{n}\mid z_{1:N})={\mathcal {N}}(x_{n}\mid m_{n}^{s},P_{n}^{s}).$$

Examples

Visualisation of the ODE filtering and smoothing procedure for the logistic ordinary differential equation using a two-times integrated Wiener process prior and an extended Kalman filter of order one. First, iterative forward filtering (brown) with forward prediction (black), then iterative smoothing (rose), and finally sampling from the full posterior distribution. The thick line shows the mean, with the shaded area showing the 95% confidence interval. The blue dots represent the observed data and the dashed line represents the ground truth.

Logistic ODE

As a one dimensional example, consider the following logistic differential equation IVP:

$${\dot {y}}(t)=y(t)(1-y(t)),\quad t\in [0,10]\quad y(0)=y_{0}=0.01,$$

for which a analytical solution exists:

$$y(t)={\frac {1}{1+\left({\frac {1+y_{0}}{y_{0}}}\right)y_{0}}}.$$

To solve this ODE-IVP with a PDE filter, the prior, discretization, as well as the predict, update, and compute steps need to be specified.

Prior

2-times IWP:

$$A(h)=\left[{\begin{array}{ccc}1&h&h^{2}/2\\0&1&h\\0&0&1\\\end{array}}\right],\quad \sigma =1,\quad Q(h)=\left[{\begin{array}{ccc}h^{5}/20&h^{4}/8&h^{3}/6\\h^{4}/8&h^{3}/3&h^{2}/2\\h^{3}/6&h^{2}/2&h\\\end{array}}\right],$$

with Taylor mode initialization:

$$X(0)=\left[y_{0},y_{0}(1-y_{0}),y_{0}(1-y_{0})(1-2y_{0})\right]^{T}.$$

Data and lielihood

We discretized the time domain uniformly with $N=20\rightarrow h=0.5\rightarrow t_{n}=n*h$ and observe data $\{z(t_{n})=0\}_{n=1}^{N}$ without any noise, hence $R_{n}=0$.

Filter and smoother

As filtering algorithm we use the EKF1 algorithm with the Jacobian is $J_{g}(X)=[-X^{(0)}(1-2X^{(0)}),1,0]$.

Software

• probdiffeq JAX based ODE filter library.
• ProbNumDiffEq.jl ODE filter library in Julia.
• probnum Probabilistic Numerics in Python.

See also

• Probabilistic numerics
• Stochastic differential equations
• Gauss-Markov process
• Recursive Bayesian estimation
• Extended Kalman filter