Solve

Stochastic block solver for ODE initial value problems.

The ODE-IVP to be solved is defined as

\[W X_t = f(X_t, t, \theta)\]

on the time interval \(t \in [a, b]\) with initial condition \(X_a = x_0\). In the paper, we use the notation \(a = t_{\mathrm{min}}\) and \(b = t_{\mathrm{max}}\).

The stochastic solver proceeds via Kalman filtering and smoothing of “interrogations” of the ODE model as described in Chkrebtii et al 2016, Schober et al 2019. In the context of the underlying Kalman filterer/smoother, the Gaussian state-space model is

\[ \begin{align}\begin{aligned}X_n = c_n + Q_n x_{n-1} + R_n^{1/2} \epsilon_n\\Z_n = W_n X_n - f(X_n, t, \theta) + V_n^{1/2} \eta_n.\end{aligned}\end{align} \]

We assume that \(c_n = c, Q_n = Q, R_n = R\), and \(W_n = W\) for all \(n\).

This module optimizes the calculations when \(Q\), \(R\), and \(W\), are block diagonal matrices of conformable and “stackable” sizes. That is, recall that the dimension of these matrices are n_state x n_state, n_state x n_state, and n_meas x n_state, respectively. Then suppose that \(Q\) and \(R\) consist of n_block blocks of size n_bstate x n_bstate, where n_bstate = n_state/n_block, and \(W\) consists of n_block blocks of size n_bmeas x n_bstate, where n_bmeas = n_meas/n_block. Then \(Q\), \(R\), \(W\) can be stored as 3D arrays of size n_block x n_bstate x n_bstate and n_block x n_bmeas x n_bstate. It is under this paradigm that the ode module operates.

rodeo.solve.solve_mv(key, ode_fun, ode_weight, ode_init, t_min, t_max, n_steps, interrogate, prior_pars, kalman_type='standard', **params)

Mean and variance of the stochastic ODE solver.

Parameters:

• key (PRNGKey) – PRNG key or None.

• ode_fun (Callable) – Higher order ODE function \(W X_t = F(X_t, t)\) taking arguments \(X\) and \(t\).

• ode_weight (ndarray(n_block, n_bmeas, n_bstate)) – Weight matrix defining the measure prior; \(W\).

• ode_init (ndarray(n_block, n_bstate)) – Initial value of the state variable \(X_t\) at time \(t = a\).

• t_min (float) – First time point of the time interval to be evaluated; \(a\).

• t_max (float) – Last time point of the time interval to be evaluated; \(b\).

• n_steps (int) – Number of discretization points (\(N\)) of the time interval that is evaluated, such that discretization timestep is \(dt = (b-a)/N\).

• interrogate (Callable) – Function defining the interrogation method.

• prior_pars (tuple) – A tuple containing the weight matrix and the variance matrix defining the solution prior; \(Q, R\).

• kalman_type (str) – Determine which type of Kalman (standard, square-root) to use.

• params (kwargs) – Optional model parameters.

Returns:

• mean_state_smooth (ndarray(n_steps+1, n_block, n_bstate)): Posterior mean of the solution process \(X_t\) at times \(t \in [a, b]\).

• var_state_smooth (ndarray(n_steps+1, n_block, n_bstate, n_bstate)): Posterior variance of the solution process at times \(t \in [a, b]\).

Return type:

(tuple)

rodeo.solve.solve_sim(key, ode_fun, ode_weight, ode_init, t_min, t_max, n_steps, interrogate, prior_pars, kalman_type='standard', **params)

Draw sample solution. Same arguments as solve_mv().

Returns:

Sample solution for \(X_t\) at times \(t \in [a, b]\).

Return type:

(ndarray(n_steps+1, n_blocks, n_bstate))