JAX-IDEM - Hands-on Examples

1 Hands-on with JAX-IDEM

JAX-IDEM, (GitLab link for the repository), is a python package built for the JAX library (a numerical computation library built for GPU/TPU acceleration, just-in-time compilation and auto-differentiation; for more, see this quick start quide), to simulate from, and perform inference for IDE models easily. It is an additional contribution alongside the paper ‘Efficient Bayesian Filtering and Inference for Integro-Difference Equation Models’; this hands-on also serves to suppliment the paper, and should be enough to recreate the experiments therein, covering the methodology in more detail.

It should be noted that, in much of the code, the model is abbreviated as IDEM (Integro-Difference Equation Models) as opposed to the simpler IDE; this is to avoid confusion with Integrated Development Environments.

All code assumes that these packages are imported like so;

import jax
import jax.numpy as jnp
import jax.random as jr

import jaxidem
import jaxidem.idem as idem
import jaxidem.utils as utils

import examples_plots as plotter

jax.config.update("jax_enable_x64", True) # Some of the square root filters can survive 32 bit precision, but generally 64 is much safer

It is expected that a reader have basic knowledge of Python and JAX, esecially jit-compilation, autodifferentiation, and vectorisation of functions. This will also not cover IDE models themselves; see this page for a brief description.

1.1 Basis Expansions

The class jaxidem.utils.Basis stores basis functions of varying types for use in the rest of JAX-IDEM.

JAX-IDEM has two built-in basis types, bisquare and cosine, which can be created by utils.place_basis and utils.place_cosine_basis respectively;

# bisquare implementation is based on the R package FRK.
bisquare_basis = utils.place_basis(
    bounds = jnp.array([[0, 1], [0, 1]]), # unit square
    nres = 2,                             # number of resolutions for the basis
    min_knot_num=3                        # the number of basis functions placed on
                                          # the smallest axis of the coursest resolution
)

# we use an expanded unit square for cosine to improve performance at the 
# boundaries of a unit square
cosine_basis = utils.place_cosine_basis(
    bounds = jnp.array([[-0.1, 1.1], [-0.1, 1.1]]), # expanded unit square
    N = 10                                # the number of frequencies per direction
) 

print(bisquare_basis.nbasis) 
print(cosine_basis.nbasis)

90
100

utils.Basis objects provide a few methods for using these functions; namely vfun and mfun, evaluating the basis functions on a single input or an array of inputs, respectively;

s = jnp.array([0.5, 0.5])
key = jr.PRNGKey(seed=1)
coords = jr.uniform(key, shape=(20, 2), minval=0, maxval=1) # 20 random points on the unit square

phi = cosine_basis.vfun(s)
PHI = cosine_basis.mfun(coords)

print(phi.shape)
print(PHI.shape)

(100,)
(20, 100)

So, we can create an approximation of Franke’s function using either basis like so;

# Franke's function (Franke 1979, NPS-53-79-003): Two Gaussian peaks of
# different widths, a broad diagonal ridge, and a narrow negative dip.
def f(x):
    x1, x2 = 9 * x[0], 9 * x[1]
    return (0.75 * jnp.exp(-((x1 - 2)**2 + (x2 - 2)**2) / 4)
            + 0.75 * jnp.exp(-((x1 + 1)**2 / 49 + (x2 + 1) / 10))
            + 0.5 * jnp.exp(-((x1 - 7)**2 + (x2 - 3)**2) / 4)
            - 0.2 * jnp.exp(-((x1 - 4)**2 + (x2 - 7)**2)))

# evaluate it on a reasonably dense grid
grid = utils.create_grid(bounds = jnp.array([[0, 1], [0, 1]]),
                         ngrids = jnp.array([100,100]))
f_on_grid = jax.vmap(f)(grid.coords)

PHI_cos = cosine_basis.mfun(grid.coords)
PHI_bis = bisquare_basis.mfun(grid.coords)

alpha_cos = jnp.linalg.pinv(PHI_cos) @ f_on_grid
alpha_bis = jnp.linalg.pinv(PHI_bis) @ f_on_grid

# the approximations themselves, for the figures below
def f_cos(x): return cosine_basis.vfun(x) @ alpha_cos
def f_bis(x): return bisquare_basis.vfun(x) @ alpha_bis

Figure 1: Interactive plot of Franke’s function, along with cosine and bisquare approximations.

For this simple target, both approximations are near the truth.

1.2 Transition Kernels

Transition kernels are stored as the class utils.Kernel. This requires a basis for each kernel parameter, and values for each parameter. The built-in kernel for JAX-IDEM is the exponential kernel, as discussed in the main paper.

To create the kernel basis, create a dictionary with the following shape;

# Support for other link functions is being worked on as of version 0.3.5

K_basis = {
        'logit decay': utils.unit_basis, 
        'log bandwidth': utils.unit_basis,
        'x offset': utils.unit_basis,
        'y offset': utils.unit_basis,
    }

It is recommended that logit decay and log bandwidth are kept using unit (constant) bases, as with this the construction of the transition matrix \(M\) becomes much easier and more stable. If they are not, then later methods for constructing \(M\) rely on Gaussian Quadrature methods that can be very unreliable when the bandwidth is low.

The offset parameters, however, can be anything; here, we can use cosine basis functions; two \(3\times 3\) basis grids, one for each direction:

K_basis_var = {
    'logit decay': utils.unit_basis,
    'log bandwidth': utils.unit_basis,
    'x offset': utils.place_cosine_basis(N=3),
    'y offset': utils.place_cosine_basis(N=3)
}

The exponential kernel,

\[ \kappa(\boldsymbol s,\boldsymbol r\mid\boldsymbol\theta_{\kappa}(\boldsymbol s)) = \frac{\gamma}{2\pi \delta^2} \exp{\left(-\frac{1}{2\delta^2} \Vert\boldsymbol r-\boldsymbol s-(\tau^{(1)}(\boldsymbol s), \tau^{(2)}(\boldsymbol s))^\intercal\Vert^2 \right)}. \]

can then be created using idem.exp_kernel. For a spatially invariant kernel (the direction and strength is the same across space);

s = 0.001
decay = 0.99
k = {
    'logit decay': jnp.log(jnp.array([decay / (1 - decay)])),
    'log bandwidth': jnp.log(jnp.array([s])),
    'x offset': jnp.array([-0.1]),
    'y offset': jnp.array([0.1]),
}
kernel_inv = idem.exp_kernel(K_basis, k)

Figure 2: Interactive plot showing kernel strength of the invariant example. Click anywhere within the field to to see how the motion an diffusion is distributed; the mass is the kernel that gets convoluted at the previous time point; effectively, this is the area from the previous time point that affects the most recent time point at the blue cross.

Figure 3: Vector field showing the direction of the invariant kernel across 100 points; this represents the direction of the flow across the space.

1

Or we can allow the offsets to vary across space; we will approximate a spiral vector field with the cosine basis;

# At any particular point, the vector field is a scaled normal vector to the centre [0.5,0.5]
def f(x):
    return 0.4 * jnp.array([-(x[1] - 0.5), (x[0] - 0.5)])

f_on_grid = jax.vmap(f)(grid.coords)

PHI_x = K_basis_var['x offset'].mfun(grid.coords)
PHI_y = K_basis_var['y offset'].mfun(grid.coords)

# Use least-squares to get the cosine-basis approximation of the normal vector field
x_offset = jnp.linalg.pinv(PHI_x) @ f_on_grid[:,0]
y_offset = jnp.linalg.pinv(PHI_y) @ f_on_grid[:,1]

k_var = {
    'logit decay': jnp.log(jnp.array([decay / (1 - decay)])),
    'log bandwidth': jnp.log(jnp.array([s])),
    'x offset': x_offset,
    'y offset': y_offset,
}
kernel_var = idem.exp_kernel(K_basis_var, k_var)

We can visualise this kernel by making a plot for its strength from a specific point and the vector field of its offset.

Figure 4: Interactive plot showing kernel strength of the variant example. Click anywhere within the field to to see how the motion an ddiffusion is distributed; the mass is the kernel that gets convoluted at the previous time point; effectively, this is the area from the previous time point that affects the most recent time point at the blue cross.

Figure 5: Vector field showing the direction of the variant kernel across 100 points; this represents the direction of the flow across the space.

1

From here on, we will only use the spatially invariant kernel for the simulated data and a variant kernel for the real data example, but it is mostly the same for both.

1.3 Model Setup

In this subsection we present code to create simple models and simulate data from the models.

_keys = jr.split(jr.PRNGKey(10), 3)  # matches seed_invariant=10 in config_supplement / config.toml
key = _keys[0]

process_basis = utils.place_cosine_basis(N=10)

model_inv = idem.gen_example_idem(key,
                                  process_basis = process_basis,
                                  beta = jnp.array([100.0, -80.0]),
                                  covariate_labels = ["x", "y"],
                                  kernel = kernel_inv,
                                  S2_eta = 2.**2,
                                  S2_eps = 1.**2)

model_var = idem.gen_example_idem(key,
                                  process_basis = process_basis,
                                  beta = jnp.array([100.0, -80.0]),
                                  covariate_labels = ["x", "y"],
                                  kernel = kernel_var,
                                  S2_eta = 2.**2,
                                  S2_eps = 1.**2)

print(idem.format_params(model_inv.params))
print(idem.format_params(model_var.params))

Kernel Parameters: 
     {'log bandwidth': '-6.90776', 'logit decay': '4.59512', 'x offset': '-0.10000', 'y offset': '0.10000'}
Variance Parameters: 1.00000, 4.00000
Coefficient Parameters: 100.00000,-80.00000
Kernel Parameters: 
     {'log bandwidth': '-6.90776', 'logit decay': '4.59512', 'x offset': '0.00000,-0.00000,-0.00000,0.16288,-0.00000,0.00000,0.00000,0.00000,0.00000', 'y offset': '-0.00000,-0.16288,-0.00000,0.00000,-0.00000,-0.00000,0.00000,0.00000,-0.00000'}
Variance Parameters: 1.00000, 4.00000
Coefficient Parameters: 100.00000,-80.00000

This creates an idem.model object that provides many useful methods for simulation and inference. To simulate from the model, we specify the location and time of every observation. The code below simulates \(T=4\) time steps at \(200\) randomly-placed observation locations, initialising the process as a mass concentrated at the top-left of the domain, and writes the latent process (evaluated on the model’s internal \(100\times100\) grid) and the noisy observations to CSV files for further analysis. The initial coefficients \(\boldsymbol\alpha_0\) are found by least-squares on a fine \(50\times50\) grid, fitting a Gaussian bump centred at \((0.3, 0.7)\) with width \(\sigma=0.1\); a single-point target would instead produce Gibbs oscillations because the cosine basis is globally supported. The centre is placed at least \(3\sigma\) from every boundary so that the even-periodic extension of the Gaussian (implicit in the cosine basis) remains negligible inside the domain.

import pandas as pd
from pathlib import Path

T, nobs = 6, 1000

# random points in space where observations are made
coords_sim = jr.uniform(_keys[1], shape=(nobs, 2))
x = jnp.tile(coords_sim[:, 0], T)
y = jnp.tile(coords_sim[:, 1], T)
times = jnp.repeat(jnp.arange(1, T + 1), nobs)

# Create an initial 'ball' mass in the top left corner; using this we can see
# how each kernel effects this mass across time
_bump_height, _bump_centre, _bump_sd = 300.0, jnp.array([0.3, 0.7]), 0.15

def _bump(s): return _bump_height * jnp.exp(-0.5 * jnp.sum((s - _bump_centre)**2) / _bump_sd**2)
alpha_0 = jnp.linalg.pinv(process_basis.mfun(grid.coords)) @ jax.vmap(_bump)(grid.coords)

process_data_inv, obs_data_inv, alphas_inv = model_inv.simulate(_keys[2], x, y, times,
                                                                covariates=jnp.column_stack([x, y]),
                                                                alpha_0=alpha_0)

process_data_var, obs_data_var, alphas_var = model_var.simulate(_keys[2], x, y, times,
                                                                covariates=jnp.column_stack([x, y]),
                                                                alpha_0=alpha_0)

Path("data").mkdir(exist_ok=True)
pd.DataFrame({"x": process_data_inv.x, "y": process_data_inv.y,
              "z": process_data_inv.z, "time": process_data_inv.times}
             ).to_csv("data/supplement_process_inv.csv", index=False)
pd.DataFrame({"x": obs_data_inv.x, "y": obs_data_inv.y,
              "z": obs_data_inv.z, "time": obs_data_inv.times}
             ).to_csv("data/supplement_obs_inv.csv", index=False)


Path("data").mkdir(exist_ok=True)
pd.DataFrame({"x": process_data_var.x, "y": process_data_var.y,
              "z": process_data_var.z, "time": process_data_var.times}
             ).to_csv("data/supplement_process_var.csv", index=False)
pd.DataFrame({"x": obs_data_var.x, "y": obs_data_var.y,
              "z": obs_data_var.z, "time": obs_data_var.times}
             ).to_csv("data/supplement_obs_var.csv", index=False)

# The filtering, inference and sampling sections below all work with a single
# synthetic dataset, and none of them depends on which kernel produced it. They
# take the invariant simulation, which is the one the simulation study
# corresponds to; point this at obs_data_var to run them on the spiral instead.
obs_data = obs_data_inv

1.4 Simulated Data

Alongside this document is a collection of methods in `plotter.py’ which provides functions to visualise various object in JAX-IDEM as plotly plots. Figure 6 and Figure 8 show the latent process and noisy observations at four selected time steps.

This exact simulation setup is not exactly what the simulation study for the invariant case actually does; this figure is chosen to make the motion of the IDE clear.

Figure 6: Simulated process data across the 4 data points using the spatially invariant kernel. You can clearly see the motion and diffusion of the gaussian mass.

Figure 7: Simulated process data across the 4 data points using the spatially variant kernel. This time advection clearly follows a circular flow.

Figure 8: Sampled bubble-plot of observed values at each 200 points across the 4 time points for the spatially invariant simulation.

Figure 9: Sampled bubble-plot of observed values at each 200 points across the 4 time points for the spatially variant simulation.

1.5 Pre-processing Real Data

For real data, the method utils.pd_to_st can be used to convert pandas dataframes into the st_data format. For example, for the sydney radar data;

radar_df = pd.read_csv('data/radar_df.csv')
# No intercept: with a cosine basis the constant covariate is collinear with the
# k = (0,0) basis function, so the paper's Sydney fits drop it.
radar_data = utils.pd_to_st(radar_df, 's1', 's2', 'time', 'z', add_intercept=False)

Time inputted are strings. Attempting to coerce into datetime objects, then into floats...

This can then be used in the same way as data simulated synthetically from model.simulate in the following sections.

Figure 10: The Sydney Radar Reflectivity data set.

To make things more interesting, we can censor this data; removing the one time point from the middle and end

T = radar_data.T
mid = T // 2
keep = jnp.array([i for i in range(T) if i not in (mid, T - 1)])
print(f"Censoring time indices {mid}, {T - 1}  ({T} → {len(keep)} obs times)")
radar_data_cen = radar_data[keep]

Censoring time indices 6, 11  (12 → 10 obs times)

Figure 11

This way, we can judge our predictions for forecasting, backcasting and ‘intercasting’ (prediciton between two time points; ‘interpolation’ is a more common, but has spatial connotations).

1.6 Inference Model Details

Given some data, either synthetically produced as above with obs_data or a real dataset like sydney, we can initialise a new model using idem.init_model;

inf_model_synth = idem.init_model(obs_data,
                            k_spat_inv = True) 

# Matches the paper's Sydney example (run 251): r = 15^2 cosine process basis,
# a 3x3 cosine basis for each kernel offset, and the domain padded 10% beyond
# the data extent.
_sydn_args = dict(k_spat_inv = False,
                  basis_type = "cosine", basis_args = [15],
                  k_basis_type = "cosine", k_basis_args = [3, 3],
                  expand_bounds = True)

inf_model_sydn  = idem.init_model(radar_data_cen, **_sydn_args)

# Sigma_eta proportional to Psi^-1, i.e. white noise in space rather than on the
# basis coefficients. The shape is read off the Gram matrix, so the model is
# built once to get Psi and then rebuilt carrying the diagonal.
_psi_inv = 1.0 / jnp.diag(inf_model_sydn.GRAM)
inf_model_sydn  = idem.init_model(radar_data_cen,
                                  S_eta_diag = _psi_inv / jnp.max(_psi_inv),
                                  **_sydn_args)

print(idem.format_params(inf_model_synth.params))
print(idem.format_params(inf_model_sydn.params))

Kernel Parameters: 
     {'log bandwidth': '-6.21156', 'logit decay': '2.00000', 'x offset': '0.00000', 'y offset': '0.00000'}
Variance Parameters: 49.89625, 49.89625
Coefficient Parameters: 0.00000,0.00000
Kernel Parameters: 
     {'log bandwidth': '-10.42674', 'logit decay': '2.00000', 'x offset': '0.00000,0.00000,0.00000,0.00000,0.00000,0.00000,0.00000,0.00000,0.00000', 'y offset': '0.00000,0.00000,0.00000,0.00000,0.00000,0.00000,0.00000,0.00000,0.00000'}
Variance Parameters: 0.99135, 0.99135
Coefficient Parameters: 

We will use the synthetic data as our spatially invariant kernel case and the Sydney radar model as our spatially variant kernel case.

The idem.model class gives various methods for filtering and smoothing of the models in a jit-compatible way. Firstly, model.get_log_likelihood is the method that creates a jit-compatible function to use the model parameters to run a filter to get the marginal likelihood of the model with those parameters;

nll_synth = inf_model_synth.get_log_like(
                obs_data,
                method="inf", # information is generally the best; see the paper and later filtering sections
                likelihood='partial', # whether to include the normalising constant
                negative=True, # defaults to positive but we generally want negative
                prior_on_a0 = False, # put the diffuse prior on alpha_1, not alpha_0
            )

nll_sydn  = inf_model_sydn.get_log_like(
                radar_data,
                method="inf",
                likelihood='partial',
                negative=True,
                prior_on_a0 = False,
            )

print(nll_synth(inf_model_synth.params))
print(nll_synth(model_inv.params))

15076.933973975405
46754.98206571441

This is effectively all we need for a basic maximum likelihood estimator; using OPTAX to do a few optimisation steps on the synthetic data;

import optax

optimizer = optax.adam(1e-1)
nll_grad = jax.grad(nll_synth)

params = inf_model_synth.params
opt_state = optimizer.init(params)

for _ in range(5):
    grads = nll_grad(params)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

print(idem.format_params(params))
print(nll_synth(params))

Kernel Parameters: 
     {'log bandwidth': '-5.78844', 'logit decay': '2.22947', 'x offset': '-0.13202', 'y offset': '0.07594'}
Variance Parameters: 30.35721, 51.80979
Coefficient Parameters: 0.48398,-0.47127
14130.182561925143

While the offsets have effectively already locked in to the true simulation values, direct likelihood maximisation struggles to identify the other parameters, which is where more sophisticated methods become important.

1.7 Prior Specification

For prior specification;

# S2_eta prior parameters
mu_eta_0 = 0.01
s2_eta_0 = 10.
alp_eta_0 = 2 + mu_eta_0**2/s2_eta_0
lam_eta_0 = mu_eta_0 * (1 + mu_eta_0**2/s2_eta_0)

# S2_eps prior parameters
mu_eps_0 = 0.1
s2_eps_0 = 10.
alp_eps_0 = 2 + mu_eps_0**2/s2_eps_0
lam_eps_0 = mu_eps_0 * (1 + mu_eps_0**2/s2_eps_0)

# beta prior
p = model_inv.beta.size
mu_beta_0 = 0
s2_beta_0 = 1000.

mu_ker_0 = {
    "logit decay": jnp.array([4]),
    "log bandwidth": jnp.array([-10]),
    "x offset": jnp.array([0.]),
    "y offset": jnp.array([0.]),
}
s2_ker_0 = {
    "logit decay": jnp.array([10.]),
    "log bandwidth": jnp.array([5]),
    "x offset": jnp.array([1.]),
    "y offset": jnp.array([1.]),
}

Using these, we can define the prior and posterior density functions;

from jax.scipy.stats.multivariate_normal import logpdf as lpmvn
from jax.scipy.stats.norm import logpdf as lpnorm

@jax.jit
def inv_gamma_logpdf(x, alpha, beta):
    return (-alpha - 1) * jnp.log(x) - beta / x

@jax.jit
def kern_prior(ker_params):
    return (jnp.sum(lpnorm(ker_params['logit decay'], mu_ker_0['logit decay'], s2_ker_0['logit decay']))
            + jnp.sum(lpnorm(ker_params['log bandwidth'], mu_ker_0['log bandwidth'], s2_ker_0['log bandwidth']))
            + jnp.sum(lpnorm(ker_params['x offset'], mu_ker_0['x offset'], s2_ker_0['x offset']))
            + jnp.sum(lpnorm(ker_params['y offset'], mu_ker_0['y offset'], s2_ker_0['y offset'])))

@jax.jit
def neg_log_post(params):
    log_beta_prior = lpmvn(params['beta'], mu_beta_0 * jnp.ones(p), s2_beta_0 * jnp.eye(p))
    log_eta_prior = inv_gamma_logpdf(jnp.exp(params['log_S2_eta']), alp_eta_0, lam_eta_0) + params['log_S2_eta']
    log_eps_prior = inv_gamma_logpdf(jnp.exp(params['log_S2_eps']), alp_eps_0, lam_eps_0) + params['log_S2_eps']
    log_kern_prior = kern_prior(params['kernel_params'])
    return nll_synth(params) - log_beta_prior - log_eta_prior - log_eps_prior - log_kern_prior

For the sampling methods;

from jax.random import normal as snorm

model = inf_model_synth
r = model.nbasis
T_obs = obs_data.T
PHI_obs_tree = jax.tree.map(model.process_basis.mfun, obs_data.coords_tree)
X_obs_tree = obs_data.X_obs_tree
tobs = sum(len(z) for z in obs_data.zs_tree)


def sample_invgamma(key, alpha, beta):
    return 1 / (jr.gamma(key, alpha) / beta)

def sample_log_s2_eta(key, M, alphas):
    A = alphas[1:] - alphas[:-1] @ M.T
    alp = alp_eta_0 + 0.5 * r * T_obs
    lam = lam_eta_0 + 0.5 * jnp.sum(jax.lax.map(lambda a: a.T @ a, A))
    return jnp.log(sample_invgamma(key, alp, lam))

def sample_log_s2_eps(key, params, alphas):
    ztildes = obs_data.tildify(params['beta'])
    alphas_tree = list(alphas[1:])
    B = 0.5 * jnp.sum(jnp.array(jax.tree.map(lambda z, phi, a: (z - phi @ a).T @ (z - phi @ a), ztildes, PHI_obs_tree, alphas_tree)))
    return jnp.log(sample_invgamma(key, alp_eps_0 + 0.5 * tobs, lam_eps_0 + B))

def sample_beta(key, S2_eps, alphas):
    alphas_tree = list(alphas[1:])
    B = jnp.eye(p) / s2_beta_0 + jnp.sum(jnp.array(jax.tree.map(lambda x: x.T @ x, X_obs_tree))) / S2_eps
    diffs = jax.tree.map(lambda z, x, phi, a: x.T @ (z - phi @ a), obs_data.zs_tree, X_obs_tree, PHI_obs_tree, alphas_tree)
    L = jnp.linalg.cholesky(B)
    mu = jax.scipy.linalg.cho_solve((L, True), jnp.sum(jnp.array(diffs), axis=0)) / S2_eps
    return mu + jax.scipy.linalg.solve_triangular(L.T, snorm(key, shape=(p,)), lower=False)

idem.model also provides the method to get the JIT-compatible simulation smoother (FFBS) method;

sample_alphas = model.get_sim_smoother(obs_data, method='inf', prior_on_a0=False)

1.8 Posterior sampling

With these we have all the necessary tools to construct either a marginal sampler on the marginal likelihood or a Gibbs sampler as described in the paper. For example, an HMC kernel can be constructed using BLACKJAX as follows

import blackjax

bj_hmc_kernel = blackjax.hmc.build_kernel()

@jax.jit
def hmc_kernel(key, fparams, imm, step_size):
    state = blackjax.hmc.init(
        position = params,
        logdensity_fn = neg_log_post
        )
    new_state, hmc_info = bj_hmc_kernel(
        key,
        state,
        logdensity_fn = neg_log_post,
        step_size = step_size,
        inverse_mass_matrix=imm,
        num_integration_steps=20,
    )
    return new_state.position, hmc_info.acceptance_rate

Similarily for a HMC-Within-Gibbs;

hmc_kernel = blackjax.hmc.build_kernel()

@jax.jit
def mwg_kernel(key, state, imm, step_size):

    params, alphas = state

    theta_k = params['kernel_params']
    log_S2_eta = params['log_S2_eta']
    log_S2_eps = params['log_S2_eps']
    beta = params['beta']

    key_theta, key_eta, key_eps, key_beta, key_alphas = jax.random.split(key, num=5)

    # Sample from theta using a HMC step
    def logdensity_theta(theta): return kern_log_prob(theta_k = theta,
                                                      alphas=alphas,
                                                      log_S2_eta=log_S2_eta)
    
    theta_state = blackjax.hmc.init(
        position=theta_k,
        logdensity_fn=logdensity_theta
        )

    new_state, hmc_info = hmc_kernel(
        key_theta,
        theta_state,
        logdensity_theta,
        step_size = step_size,
        inverse_mass_matrix=imm,
        num_integration_steps=20,
    )
    
    theta_k = new_state.position
    acceptance_rate = hmc_info.acceptance_rate
    
    params['kernel_params'] = theta_k
    _, M = kern_post_with_M(theta_k, alphas, log_S2_eta)

    # Sample from the variances with the conditional posteriors

    # process variance
    if "log_S2_eta" not in fixed_keys:
        log_S2_eta = sample_log_s2_eta(key_eta, M, alphas)
        params['log_S2_eta'] = log_S2_eta
    
    if "log_S2_eps" not in fixed_keys:
        log_S2_eps = sample_log_s2_eps(key_eps, params, alphas)
        params['log_S2_eps'] = log_S2_eps        

    # Linear coefficients — collapsed Gibbs: marginalise over alpha via Kalman filter.
    # beta = sample_beta(key_beta, jnp.exp(log_S2_eps), alphas)  # old conditional step
    if "beta" not in fixed_keys:
        beta_proposal = sample_beta_marginal(key_beta, params)
        beta = jnp.where(jnp.any(jnp.isnan(beta_proposal)), beta, beta_proposal)
        params['beta'] = beta

    
    # Sample from alphas using FFBS
    alphas = sample_alphas(key_alphas, params, M)

    new_state = (params, alphas)
    
    return new_state, acceptance_rate

Both of these samplers will likely mix poorly by default; the simulation study uses a three-step procedure to tune the samplers, as follows: - An initial step using dual averaging to set step size such that acceptance rate is appropriate (around 0.63 for HMC) - A ‘full adaption’ phase using both dual averaging and variance adaption (imm, the inverse mass matrix, is roughly equivalent to the proposal variance; optimaly, it should be a scalar mulitple of the true variance, so it can be estimated either through a warm-up like this or by other methods such as Laplace Approximation) - A finalisation of the step-size using dual-averaging after the inverse mass matrix is fixed

This is done both for marginal and full methods, where for the Gibbs method, the gradient-based method within the Gibbs sampler is tuned.

For NUTS, poorly chosen initial tuning causes the algorithm to take a long time to find a U-turn; this leads quickly to an excesive number of computations of the likelihood and its gradient. For this reason, the adaption phase to estimate Variance is still done with HMC and a fixed number of integration steps; once variance is fixed, the final step-size adaption is done targetting the optimal acceptance rate for NUTS (around 0.8).

It should also be noted that BLACKJAX does provide it’s own mechanisms for tuning; see their tutorial here. Their mechanisms do not work for more customised samplers, such as our block-Gibbs sampler, so set up the tuning steps manually and use the same logic for the marginal method for parity.

1.9 Note on Time-Invariance and IID Errors

For most of the filtering algorithms, there are various optimisations possible depending on the complexity of the system which can be made. For the purposes of this project, we have generally considered only the case where the variance of the observation error is independant and identically distributed across the spatial locations observed and each time point, with some single variance \(\sigma_\epsilon^2\). The filters in JAX-IDEM do allow for different shapes, and two allow for optimisation two variants of the information filters have two versions; one at idem.filters and one at idem.filters.invariant (to be read as ‘temporally invariant’). The invariant variants assume IID errors in both process and observation, and also assume that the spatial locations remain the same across all time points. These assumptions are true for the simulation study in question, although it should be noted that they do have different performance profiles; in particular, the invariant versions are \(\mathcal O(n)\), whereas inversions of \(\Sigma_\epsilon^{(t)}\) may be necessary for the temporally variant/non-IID versions.

1.10 Filtering and Smoothing

Using a chain drawn by NUTS for the Sydney \(r=225\) model, we will now see how to perform some basic inference on the marginal posterior process.

Filters with complete outputs (not just log-likelihood) are also available with the method get_filter, but also the full filter-smoother output can be made with get_filter_smoother

# only the Covariance-based Kalman-RTS smoother is implemented.
filtsmooth = inf_model_sydn.get_filter_smoother(radar_data_cen,
                                   nforecast = 3, # Forecast as computed by the filter
                                   prior_on_a0 = False,
                                   )

Using a 5,000-long chain from the posterior using NUTS, we computed this posterior mean for the model

import pickle

with open("data/sydn_r225_var/post_mean_params.pkl", 'rb') as file:
    post_mean_params = pickle.load(file)

filtsmooth_results = filtsmooth(post_mean_params)

ms = filtsmooth_results['ms']
ms_smooth = filtsmooth_results['m_smooth']

This is only the smoothed results using the posterior mean fo the parameters, which is obviously not ideal and not taking full advantage of the posterior sample.

Instead, we can run the entire filter-smoother on the chain and take the mean of the process means that come out to create a Roa-Blackwellised estimate of the process; for example, using a NUTS chain for the Sydney variant \(r=225\) case;

filtsmooth = inf_model_sydn.get_filter_smoother(radar_data_cen, prior_on_a0=False, nforecast=3)

# Load the chain from the run
nuts_chain = jnp.array(
    pd.read_csv("data/sydn_r225_var/nuts_chain.csv").drop(columns=['iter', 'accepted'])
)

# thin the chain
n = 10
nuts_chain_thin = nuts_chain[:: len(nuts_chain) // n][:n]

# The chains contain 'flattened' parameter values which need to be unflattened
# into the dictionary format accepted by filtsmooth.
# The simulations used `utils.flatten_and_unflatten` to flatten so using the
# same here gets us the unflattening function
flatten, unflatten = utils.flatten_and_unflatten(inf_model_sydn.params)

# The curried version that takes flat inputs like of the chain
filtsmooth_flat = lambda theta: filtsmooth(unflatten(theta))

r = inf_model_sydn.nbasis
T = radar_data_cen.T

acc_ms, acc_m2s, acc_vars = jnp.zeros((T+3,r)), jnp.zeros((T+3,r,r)), jnp.zeros((T+3,r,r))

# outer product vectorised across the first axis
vec_outer = jax.vmap(jnp.outer)

for i in range(n): # thinning to only n samples
    res = filtsmooth_flat(nuts_chain_thin[i])

    # The forecast states come back separately from the smoothed ones, so stack
    # them to get a single series over all T + nforecast time points.
    ms = jnp.concatenate([res['m_smooth'], res['m_forecast']])
    Ps = jnp.concatenate([res['P_smooth'], res['P_forecast']])

    
    acc_ms = acc_ms + ms
    acc_m2s = acc_m2s + vec_outer(ms, ms)
    acc_vars = acc_vars + Ps

marginal_mean = acc_ms / n
# using the law of total variance
marginal_var = acc_vars / n + \
    acc_m2s / n - vec_outer(marginal_mean, marginal_mean)

We can cast this onto an st_data object, evaluating the mean basis on the same grid that the original radar_data grid is on;

# We need locations at which to evaluate the basis. We can grab coordinates for
# each of the 12 time points in the radar data like so

coords_tree = radar_data.coords_tree

# This is a list containing jax arrays of the coordinates at each time point.
# We also want to add the two extra forecast points; repeating the grid two
# extra times;
coords_tree = coords_tree + [coords_tree[0]] * 2

RB_mean_data = idem.basis_params_to_st_data(marginal_mean, inf_model_sydn.process_basis, coords_tree)

For the 12 time points actually observed we can also compute the standardised residuals;

# Stack the data into a JAX array; possible only because there is no missing 
# data here, otherwise we would have to work with the PyTrees themselves
zs_true = jnp.stack(radar_data.zs_tree)
zs_mean = jnp.stack(RB_mean_data.zs_tree[:12])

resids = zs_true - zs_mean

# Again, coords are the same at every time point so we only need one
PHI = inf_model_sydn.process_basis.mfun(coords_tree[0])

# Compute the variances at each location
cast_var = jax.vmap(lambda P: jnp.diag(PHI @ P @ PHI.T))

s2_eps_post_mean = jnp.mean(jnp.exp(nuts_chain[:,-2]))

# Using law of total variance to include observation error (at the posterior mean)
# to compute the total posterior variance at each observation location
obs_var = cast_var(marginal_var[:12]) + s2_eps_post_mean

std_resids = resids / jnp.sqrt(obs_var)

Figure 12: The uncensored Sydney Radar Data (left), the posterior predictive means at the same points (middle), and the standardised residuals (right).

We can clearly see that the standardised residuals are not very homoogeneous in space or time; the model clearly does not fit perfectly.

We can also update a model with new parameters with the `update’ method

post_model = inf_model_sydn.update(post_mean_params)

2 Stability and Timing

2.1 Stability

get_log_like can optionally take filter type (defaulting to sqinf, the Square-root Information Filter). More filters are included in the package than discussed in the paper, but hte four used for the stability result in Section 3.2. are skalman (Stabilised Kalman), sqrt (Square-Root), inf (Information), and sqinf.

Most of the numerical instability occurs when \(P_0\) increases, so we can test each filter in 32-bit like so;

from jax import config

# Drop to 32-bit for this block only. x64 is a global flag, so it has to be
# toggled rather than scoped, and restored below -- every chunk after this one
# would otherwise run in single precision too.
config.update("jax_enable_x64", False)

from itertools import product
 
filters = ['skalman', 'sqrt', 'inf', 'sqinf']
ps = [10**2, 10**4, 10**6, 10**8]

results = {}
for filt, p in product(filters, ps):
    nll = inf_model_synth.get_log_like(
        obs_data,
        method=filt,
        likelihood='partial',
        P_0 = p * jnp.eye(inf_model_synth.nbasis),
        negative=True,
        prior_on_a0 = False,
    )

    results[(filt, p)] = float(nll(inf_model_synth.params))

filter           P_0 = 1e2     P_0 = 1e4     P_0 = 1e6     P_0 = 1e8
--------------------------------------------------------------------
skalman       14448.900391  14617.850586           nan           nan
sqrt          14448.890625  14617.321289  14846.371094  15201.710938
inf           14448.466797  14616.658203  14846.126953  15076.429688
sqinf         14448.831055  14617.386719  14846.521484  15076.765625

rel. diff.       P_0 = 1e2     P_0 = 1e4     P_0 = 1e6     P_0 = 1e8
--------------------------------------------------------------------
skalman           4.80e-06      3.17e-05           nan           nan
sqrt              4.12e-06      4.48e-06      1.01e-05      8.29e-03
inf               2.52e-05      4.98e-05      2.66e-05      2.23e-05
sqinf             0.00e+00      0.00e+00      0.00e+00      0.00e+00

Note that these numbers are not the same numbers as in the paper’s experiments, but should be illustrative of the same point.

We can also confirm that in 64 bit the filters all agree to a reasonable degree;

# Return to 64-bit (will remain for the rest of the document)
config.update("jax_enable_x64", True)

results = {}
for filt, p in product(filters, ps):
    nll = inf_model_synth.get_log_like(
        obs_data,
        method=filt,
        likelihood='partial',
        P_0 = p * jnp.eye(inf_model_synth.nbasis),
        negative=True,
        prior_on_a0 = False,
    )

    results[(filt, p)] = float(nll(inf_model_synth.params))

filter           P_0 = 1e2     P_0 = 1e4     P_0 = 1e6     P_0 = 1e8
--------------------------------------------------------------------
skalman       14448.890847  14617.314631  14846.684393  15076.933975
sqrt          14448.890847  14617.314631  14846.684393  15076.933974
inf           14448.890847  14617.314631  14846.684393  15076.933974
sqinf         14448.890847  14617.314631  14846.684393  15076.933974

2.2 Timing

By varying the number of observation points per time point, we can also measure how the computation time compares for each filter.

ns = [100, 200, 300, 400]

key = jr.PRNGKey(seed=1) # new PRNG key for randomly jittering the parameters

timings = {}
for filt, n in product(filters, ns):
    data_slice = obs_data[:, :n] # cuts off all observations beyond the nth at each time point

    nll = inf_model_synth.get_log_like(
        data_slice,
        method=filt,
        likelihood='partial',
        P_0 = 10**8 * jnp.eye(inf_model_synth.nbasis),
        negative=True,
        prior_on_a0 = False,
    )

    # time_jit compiles the function, then runs it 5 more times on jittered
    # inputs, returning the compile time and the average of the timed runs.
    res = utils.time_jit(key, nll, inf_model_synth.params, 5, desc=f"{filt} n={n}")

    timings[(filt, n)] = res.average_time

filter           n = 100     n = 200     n = 300     n = 400
------------------------------------------------------------
skalman           231.18      263.58      552.93     3047.30
sqrt             2767.39     3459.83      883.29     6693.30
inf               211.61      228.41       87.17       70.04
sqinf             134.87      865.48      272.38      650.22

It should be noted that the output of this document is unreliable; if this document is exectuted live, then the result depends on the CPU being used (the notebook is not set to run on a GPU, although it can be). Background processes are likely to colour the results as well.