Dimension Mismatch When Using ODEs With az.plot_hdi

Questions v5
,
1

Hello,

I am attempting to fit a simple model using ODEs and pymc. The model appears to run successfully, and I can plot the mean predicted posterior against the observed values. However, I can’t plot a HDI range from the sampled posterior. The error message is:

Dimension mismatch for x: (20,) and hdi: (4000, 2). Check the dimensions of y andhdi_kwargs to make sure they are compatible

I am probably doing something silly, and would appreciate any help please to understand and fix what I am doing wrong.

My code is:

%matplotlib inline
import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pymc.sampling.jax as pmjax
import xarray as xr
from pymc.ode import DifferentialEquation
from scipy.integrate import odeint
 
np.random.seed(20394)
 
def freefall(y, t, p):
    return 2.0 * p[1] - p[0] * y[0]
 
times = np.arange(0, 10, 0.5)
gamma, g, y0, sigma = 0.4, 9.8, -2, 2
y = odeint(freefall, t=times, y0=y0, args=tuple([[gamma, g]]))
yobs = np.random.normal(y, 2)

ode_model = DifferentialEquation(func=freefall, times=times, n_states=1, n_theta=2, t0=0)

with pm.Model() as model:
    sigma = pm.HalfCauchy("sigma", 1)
    gamma = pm.Lognormal("gamma", 0, 1)
    ode_solution = ode_model(y0=[0], theta=[gamma, 9.8])

    Y = pm.Normal("Y", mu=ode_solution, sigma=sigma, observed=yobs)

    idata = pm.sample_prior_predictive(draws=100)
    idata.extend(pm.sample(1000, chains=4))
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)

_, ax = plt.subplots()
ax.scatter(times, idata.observed_data["Y"], color='C1', alpha=0.6)
ax.plot(times, y, label='True', color='k', alpha=0.6)
ax.plot(times, idata.posterior_predictive["Y"].mean(("chain", "draw")), label='Mean', color='C0', alpha=0.6)
az.plot_hdi(x=xr.DataArray(times), y=idata.posterior_predictive["Y"])
2

Hi all,

I’ve solved this myself after spotting a related topic “Dimension error with using az.plot_lm() when running tutorial notebook?”

I had tried the squeeze function to drop the unwanted dimension Y_dim_3 previously, but including a drop=True was needed for it to work.

Updated code is:

_, ax = plt.subplots()
ax.scatter(times, idata.observed_data["Y"], color='C1', alpha=0.6)
ax.plot(times, y, label='True', color='k', alpha=0.6)
ax.plot(times, idata.posterior_predictive["Y"].mean(("chain", "draw")), label='Mean', color='C0', alpha=0.6)
az.plot_hdi(x=xr.DataArray(times), y=idata.posterior_predictive["Y"].squeeze(dim=['Y_dim_3'], drop=True))