Bayesian Vector Autoregressive Models#

import os

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import statsmodels.api as sm

from pymc.sampling_jax import sample_blackjax_nuts
/Users/nathanielforde/opt/miniconda3/envs/pymc_ar_ex/lib/python3.9/site-packages/pymc/sampling_jax.py:37: UserWarning: This module is experimental.
  warnings.warn("This module is experimental.")
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")
%config InlineBackend.figure_format = 'retina'

V(ector)A(uto)R(egression) Models#

In this notebook we will outline an application of the Bayesian Vector Autoregressive Modelling. We will draw on the work in the PYMC Labs blogpost (see Vieira [n.d.]). This will be a three part series. In the first we want to show how to fit Bayesian VAR models in PYMC. In the second we will show how to extract extra insight from the fitted model with Impulse Response analysis and make forecasts from the fitted VAR model. In the third and final post we will show in some more detail the benefits of using hierarchical priors with Bayesian VAR models. Specifically, we’ll outline how and why there are actually a range of carefully formulated industry standard priors which work with Bayesian VAR modelling.

In this post we will (i) demonstrate the basic pattern on a simple VAR model on fake data and show how the model recovers the true data generating parameters and (ii) we will show an example applied to macro-economic data and compare the results to those achieved on the same data with statsmodels MLE fits and (iii) show an example of estimating a hierarchical bayesian VAR model over a number of countries.

Autoregressive Models in General#

The idea of a simple autoregressive model is to capture the manner in which past observations of the timeseries are predictive of the current observation. So in traditional fashion, if we model this as a linear phenomena we get simple autoregressive models where the current value is predicted by a weighted linear combination of the past values and an error term.

\[ y_t = \alpha + \beta_{y0} \cdot y_{t-1} + \beta_{y1} \cdot y_{t-2} ... + \epsilon \]

for however many lags are deemed appropriate to the predict the current observation.

A VAR model is kind of generalisation of this framework in that it retains the linear combination approach but allows us to model multiple timeseries at once. So concretely this mean that \(\mathbf{y}_{t}\) as a vector where:

\[ \mathbf{y}_{T} = \nu + A_{1}\mathbf{y}_{T-1} + A_{2}\mathbf{y}_{T-2} ... A_{p}\mathbf{y}_{T-p} + \mathbf{e}_{t} \]

where the As are coefficient matrices to be combined with the past values of each individual timeseries. For example consider an economic example where we aim to model the relationship and mutual influence of each variable on themselves and one another.

\[\begin{split} \begin{bmatrix} gdp \\ inv \\ con \end{bmatrix}_{T} = \nu + A_{1}\begin{bmatrix} gdp \\ inv \\ con \end{bmatrix}_{T-1} + A_{2}\begin{bmatrix} gdp \\ inv \\ con \end{bmatrix}_{T-2} ... A_{p}\begin{bmatrix} gdp \\ inv \\ con \end{bmatrix}_{T-p} + \mathbf{e}_{t} \end{split}\]

This structure is compact representation using matrix notation. The thing we are trying to estimate when we fit a VAR model is the A matrices that determine the nature of the linear combination that best fits our timeseries data. Such timeseries models can have an auto-regressive or a moving average representation, and the details matter for some of the implication of a VAR model fit.

We’ll see in the next notebook of the series how the moving-average representation of a VAR lends itself to the interpretation of the covariance structure in our model as representing a kind of impulse-response relationship between the component timeseries.

A Concrete Specification with Two lagged Terms#

The matrix notation is convenient to suggest the broad patterns of the model, but it is useful to see the algebra is a simple case. Consider the case of Ireland’s GDP and consumption described as:

\[ gdp_{t} = \beta_{gdp1} \cdot gdp_{t-1} + \beta_{gdp2} \cdot gdp_{t-2} + \beta_{cons1} \cdot cons_{t-1} + \beta_{cons2} \cdot cons_{t-2} + \epsilon_{gdp}\]
\[ cons_{t} = \beta_{cons1} \cdot cons_{t-1} + \beta_{cons2} \cdot cons_{t-2} + \beta_{gdp1} \cdot gdp_{t-1} + \beta_{gdp2} \cdot gdp_{t-2} + \epsilon_{cons}\]

In this way we can see that if we can estimate the \(\beta\) terms we have an estimate for the bi-directional effects of each variable on the other. This is a useful feature of the modelling. In what follows i should stress that i’m not an economist and I’m aiming to show only the functionality of these models not give you a decisive opinion about the economic relationships determining Irish GDP figures.

Creating some Fake Data#

def simulate_var(
    intercepts, coefs_yy, coefs_xy, coefs_xx, coefs_yx, noises=(1, 1), *, warmup=100, steps=200
):
    draws_y = np.zeros(warmup + steps)
    draws_x = np.zeros(warmup + steps)
    draws_y[:2] = intercepts[0]
    draws_x[:2] = intercepts[1]
    for step in range(2, warmup + steps):
        draws_y[step] = (
            intercepts[0]
            + coefs_yy[0] * draws_y[step - 1]
            + coefs_yy[1] * draws_y[step - 2]
            + coefs_xy[0] * draws_x[step - 1]
            + coefs_xy[1] * draws_x[step - 2]
            + rng.normal(0, noises[0])
        )
        draws_x[step] = (
            intercepts[1]
            + coefs_xx[0] * draws_x[step - 1]
            + coefs_xx[1] * draws_x[step - 2]
            + coefs_yx[0] * draws_y[step - 1]
            + coefs_yx[1] * draws_y[step - 2]
            + rng.normal(0, noises[1])
        )
    return draws_y[warmup:], draws_x[warmup:]

First we generate some fake data with known parameters.

var_y, var_x = simulate_var(
    intercepts=(18, 8),
    coefs_yy=(-0.8, 0),
    coefs_xy=(0.9, 0),
    coefs_xx=(1.3, -0.7),
    coefs_yx=(-0.1, 0.3),
)

df = pd.DataFrame({"x": var_x, "y": var_y})
df.head()
x y
0 34.606613 30.117581
1 34.773803 23.996700
2 35.455237 29.738941
3 33.886706 27.193417
4 31.837465 26.704728
fig, axs = plt.subplots(2, 1, figsize=(10, 3))
axs[0].plot(df["x"], label="x")
axs[0].set_title("Series X")
axs[1].plot(df["y"], label="y")
axs[1].set_title("Series Y");
../_images/15561392e1e4bb0f00ed45dbc30709227e906498ea8825c46c238adfe8037049.png

Handling Multiple Lags and Different Dimensions#

When Modelling multiple timeseries and accounting for potentially any number lags to incorporate in our model we need to abstract some of the model definition to helper functions. An example will make this a bit clearer.

### Define a helper function that will construct our autoregressive step for the marginal contribution of each lagged
### term in each of the respective time series equations
def calc_ar_step(lag_coefs, n_eqs, n_lags, df):
    ars = []
    for j in range(n_eqs):
        ar = pm.math.sum(
            [
                pm.math.sum(lag_coefs[j, i] * df.values[n_lags - (i + 1) : -(i + 1)], axis=-1)
                for i in range(n_lags)
            ],
            axis=0,
        )
    ars.append(ar)
    beta = pm.math.stack(ars, axis=-1)

    return beta


### Make the model in such a way that it can handle different specifications of the likelihood term
### and can be run for simple prior predictive checks. This latter functionality is important for debugging of
### shape handling issues. Building a VAR model involves quite a few moving parts and it is handy to
### inspect the shape implied in the prior predictive checks.
def make_model(n_lags, n_eqs, df, priors, mv_norm=True, prior_checks=True):
    coords = {
        "lags": np.arange(n_lags) + 1,
        "equations": df.columns.tolist(),
        "cross_vars": df.columns.tolist(),
        "time": [x for x in df.index[n_lags:]],
    }

    with pm.Model(coords=coords) as model:
        lag_coefs = pm.Normal(
            "lag_coefs",
            mu=priors["lag_coefs"]["mu"],
            sigma=priors["lag_coefs"]["sigma"],
            dims=["equations", "lags", "cross_vars"],
        )
        alpha = pm.Normal(
            "alpha", mu=priors["alpha"]["mu"], sigma=priors["alpha"]["sigma"], dims=("equations",)
        )
        data_obs = pm.Data("data_obs", df.values[n_lags:], dims=["time", "equations"], mutable=True)

        betaX = calc_ar_step(lag_coefs, n_eqs, n_lags, df)
        betaX = pm.Deterministic(
            "betaX",
            betaX,
            dims=[
                "time",
            ],
        )
        mean = alpha + betaX

        if mv_norm:
            n = df.shape[1]
            ## Under the hood the LKJ prior will retain the correlation matrix too.
            noise_chol, _, _ = pm.LKJCholeskyCov(
                "noise_chol",
                eta=priors["noise_chol"]["eta"],
                n=n,
                sd_dist=pm.HalfNormal.dist(sigma=priors["noise_chol"]["sigma"]),
            )
            obs = pm.MvNormal(
                "obs", mu=mean, chol=noise_chol, observed=data_obs, dims=["time", "equations"]
            )
        else:
            ## This is an alternative likelihood that can recover sensible estimates of the coefficients
            ## But lacks the multivariate correlation between the timeseries.
            sigma = pm.HalfNormal("noise", sigma=priors["noise"]["sigma"], dims=["equations"])
            obs = pm.Normal(
                "obs", mu=mean, sigma=sigma, observed=data_obs, dims=["time", "equations"]
            )

        if prior_checks:
            idata = pm.sample_prior_predictive()
            return model, idata
        else:
            idata = pm.sample_prior_predictive()
            idata.extend(pm.sample(draws=2000, random_seed=130))
            pm.sample_posterior_predictive(idata, extend_inferencedata=True, random_seed=rng)
    return model, idata

The model has a deterministic component in the auto-regressive calculation which is required at each timestep, but the key point here is that we model the likelihood of the VAR as a multivariate normal distribution with a particular covariance relationship. The estimation of these covariance relationship gives the main insight in the manner in which our component timeseries relate to one another.

We will inspect the structure of a VAR with 2 lags and 2 equations

n_lags = 2
n_eqs = 2
priors = {
    "lag_coefs": {"mu": 0.3, "sigma": 1},
    "alpha": {"mu": 15, "sigma": 5},
    "noise_chol": {"eta": 1, "sigma": 1},
    "noise": {"sigma": 1},
}

model, idata = make_model(n_lags, n_eqs, df, priors)
pm.model_to_graphviz(model)
Sampling: [alpha, lag_coefs, noise_chol, obs]
INFO:pymc:Sampling: [alpha, lag_coefs, noise_chol, obs]
../_images/a31716b0d351402d70cedef2b3a907d3bef990fff0734c1efb4328cf0c2fdf84.svg

Another VAR with 3 lags and 2 equations.

n_lags = 3
n_eqs = 2
model, idata = make_model(n_lags, n_eqs, df, priors)
for rv, shape in model.eval_rv_shapes().items():
    print(f"{rv:>11}: shape={shape}")
pm.model_to_graphviz(model)
Sampling: [alpha, lag_coefs, noise_chol, obs]
INFO:pymc:Sampling: [alpha, lag_coefs, noise_chol, obs]
  lag_coefs: shape=(2, 3, 2)
      alpha: shape=(2,)
noise_chol_cholesky-cov-packed__: shape=(3,)
 noise_chol: shape=(3,)
../_images/5a36aae302712572deef73fdb4da967052fc3f9bfcc8f21f96e3f83f3e695b7c.svg

We can inspect the correlation matrix between our timeseries which is implied by the prior specification, to see that we have allowed a flat uniform prior over their correlation.

ax = az.plot_posterior(
    idata,
    var_names="noise_chol_corr",
    hdi_prob="hide",
    group="prior",
    point_estimate="mean",
    grid=(2, 2),
    kind="hist",
    ec="black",
    figsize=(10, 4),
)
../_images/b7887869e1b289f073b87ae70e24f06a326f2932abc5edd8c4eab10fef9030dc.png

Now we will fit the VAR with 2 lags and 2 equations

n_lags = 2
n_eqs = 2
model, idata_fake_data = make_model(n_lags, n_eqs, df, priors, prior_checks=False)
Sampling: [alpha, lag_coefs, noise_chol, obs]
INFO:pymc:Sampling: [alpha, lag_coefs, noise_chol, obs]
Auto-assigning NUTS sampler...
INFO:pymc:Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
INFO:pymc:Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
INFO:pymc:Multiprocess sampling (4 chains in 4 jobs)
NUTS: [lag_coefs, alpha, noise_chol]
INFO:pymc:NUTS: [lag_coefs, alpha, noise_chol]
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
100.00% [12000/12000 11:54<00:00 Sampling 4 chains, 0 divergences]
Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 748 seconds.
INFO:pymc:Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 748 seconds.
The acceptance probability does not match the target. It is 0.8896, but should be close to 0.8. Try to increase the number of tuning steps.
WARNING:pymc:The acceptance probability does not match the target. It is 0.8896, but should be close to 0.8. Try to increase the number of tuning steps.
Sampling: [obs]
INFO:pymc:Sampling: [obs]
100.00% [8000/8000 07:29<00:00]

We’ll now plot some of the results to see that the parameters are being broadly recovered. The alpha parameters match well, but the individual lag coefficients show differences.

az.summary(idata_fake_data, var_names=["alpha", "lag_coefs", "noise_chol_corr"])
/Users/nathanielforde/opt/miniconda3/envs/pymc_ar_ex/lib/python3.9/site-packages/arviz/stats/diagnostics.py:586: RuntimeWarning: invalid value encountered in double_scalars
  (between_chain_variance / within_chain_variance + num_samples - 1) / (num_samples)
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
alpha[x] 16.161 1.389 13.588 18.756 0.027 0.019 2620.0 3694.0 1.0
alpha[y] 9.630 1.389 7.022 12.209 0.027 0.019 2622.0 3813.0 1.0
lag_coefs[x, 1, x] 0.322 0.981 -1.482 2.224 0.011 0.011 7704.0 5796.0 1.0
lag_coefs[x, 1, y] 0.302 0.988 -1.539 2.147 0.011 0.011 8812.0 5915.0 1.0
lag_coefs[x, 2, x] 0.296 0.991 -1.496 2.242 0.011 0.011 7735.0 5277.0 1.0
lag_coefs[x, 2, y] 0.312 1.003 -1.571 2.182 0.012 0.011 7467.0 5164.0 1.0
lag_coefs[y, 1, x] 1.087 0.046 1.003 1.175 0.001 0.001 4173.0 5041.0 1.0
lag_coefs[y, 1, y] -0.462 0.055 -0.562 -0.352 0.001 0.001 2721.0 4089.0 1.0
lag_coefs[y, 2, x] -0.308 0.075 -0.445 -0.164 0.001 0.001 2818.0 3953.0 1.0
lag_coefs[y, 2, y] 0.130 0.055 0.028 0.233 0.001 0.001 2853.0 4215.0 1.0
noise_chol_corr[0, 0] 1.000 0.000 1.000 1.000 0.000 0.000 8000.0 8000.0 NaN
noise_chol_corr[0, 1] -0.221 0.068 -0.344 -0.088 0.001 0.001 6969.0 5044.0 1.0
noise_chol_corr[1, 0] -0.221 0.068 -0.344 -0.088 0.001 0.001 6969.0 5044.0 1.0
noise_chol_corr[1, 1] 1.000 0.000 1.000 1.000 0.000 0.000 8370.0 7576.0 1.0
az.plot_posterior(idata_fake_data, var_names=["alpha"], ref_val=[18, 8]);
../_images/83414014bcf3a716c2a5188311536ea464cb91e0b2e312728376ee0fe6350641.png

Next we’ll plot the posterior predictive distribution to check that the fitted model can capture the patterns in the observed data. This is the primary test of goodness of fit.

def shade_background(ppc, ax, idx, palette="cividis"):
    palette = palette
    cmap = plt.get_cmap(palette)
    percs = np.linspace(51, 99, 100)
    colors = (percs - np.min(percs)) / (np.max(percs) - np.min(percs))
    for i, p in enumerate(percs[::-1]):
        upper = np.percentile(
            ppc[:, idx, :],
            p,
            axis=1,
        )
        lower = np.percentile(
            ppc[:, idx, :],
            100 - p,
            axis=1,
        )
        color_val = colors[i]
        ax[idx].fill_between(
            x=np.arange(ppc.shape[0]),
            y1=upper.flatten(),
            y2=lower.flatten(),
            color=cmap(color_val),
            alpha=0.1,
        )


def plot_ppc(idata, df, group="posterior_predictive"):
    fig, axs = plt.subplots(2, 1, figsize=(25, 15))
    df = pd.DataFrame(idata_fake_data["observed_data"]["obs"].data, columns=["x", "y"])
    axs = axs.flatten()
    ppc = az.extract_dataset(idata, group=group, num_samples=100)["obs"]
    # Minus the lagged terms and the constant
    shade_background(ppc, axs, 0, "inferno")
    axs[0].plot(np.arange(ppc.shape[0]), ppc[:, 0, :].mean(axis=1), color="cyan", label="Mean")
    axs[0].plot(df["x"], "o", mfc="black", mec="white", mew=1, markersize=7, label="Observed")
    axs[0].set_title("VAR Series 1")
    axs[0].legend()
    shade_background(ppc, axs, 1, "inferno")
    axs[1].plot(df["y"], "o", mfc="black", mec="white", mew=1, markersize=7, label="Observed")
    axs[1].plot(np.arange(ppc.shape[0]), ppc[:, 1, :].mean(axis=1), color="cyan", label="Mean")
    axs[1].set_title("VAR Series 2")
    axs[1].legend()


plot_ppc(idata_fake_data, df)
../_images/87a0488d98fddf266cd66d945de66284e450a58cc13c5b753158af12e3dd4c6a.png

Again we can check the learned posterior distribution for the correlation parameter.

ax = az.plot_posterior(
    idata_fake_data,
    var_names="noise_chol_corr",
    hdi_prob="hide",
    point_estimate="mean",
    grid=(2, 2),
    kind="hist",
    ec="black",
    figsize=(10, 6),
)
../_images/a27585688c72de656440bd875452efd7faa678f65d78c886693368443d0a4c00.png

Applying the Theory: Macro Economic Timeseries#

The data is from the World Bank’s World Development Indicators. In particular, we’re pulling annual values of GDP, consumption, and gross fixed capital formation (investment) for all countries from 1970. Timeseries models in general work best when we have a stable mean throughout the series, so for the estimation procedure we have taken the first difference and the natural log of each of these series.

try:
    gdp_hierarchical = pd.read_csv(
        os.path.join("..", "data", "gdp_data_hierarchical_clean.csv"), index_col=0
    )
except FileNotFoundError:
    gdp_hierarchical = pd.read_csv(pm.get_data("gdp_data_hierarchical_clean.csv"), ...)

gdp_hierarchical
country iso2c iso3c year GDP CONS GFCF dl_gdp dl_cons dl_gfcf more_than_10 time
1 Australia AU AUS 1971 4.647670e+11 3.113170e+11 7.985100e+10 0.039217 0.040606 0.031705 True 1
2 Australia AU AUS 1972 4.829350e+11 3.229650e+11 8.209200e+10 0.038346 0.036732 0.027678 True 2
3 Australia AU AUS 1973 4.955840e+11 3.371070e+11 8.460300e+10 0.025855 0.042856 0.030129 True 3
4 Australia AU AUS 1974 5.159300e+11 3.556010e+11 8.821400e+10 0.040234 0.053409 0.041796 True 4
5 Australia AU AUS 1975 5.228210e+11 3.759000e+11 8.255900e+10 0.013268 0.055514 -0.066252 True 5
... ... ... ... ... ... ... ... ... ... ... ... ...
366 United States US USA 2016 1.850960e+13 1.522497e+13 3.802207e+12 0.016537 0.023425 0.021058 True 44
367 United States US USA 2017 1.892712e+13 1.553075e+13 3.947418e+12 0.022306 0.019885 0.037480 True 45
368 United States US USA 2018 1.947957e+13 1.593427e+13 4.119951e+12 0.028771 0.025650 0.042780 True 46
369 United States US USA 2019 1.992544e+13 1.627888e+13 4.248643e+12 0.022631 0.021396 0.030758 True 47
370 United States US USA 2020 1.924706e+13 1.582501e+13 4.182801e+12 -0.034639 -0.028277 -0.015619 True 48

370 rows × 12 columns

fig, axs = plt.subplots(3, 1, figsize=(20, 10))
for country in gdp_hierarchical["country"].unique():
    temp = gdp_hierarchical[gdp_hierarchical["country"] == country].reset_index()
    axs[0].plot(temp["dl_gdp"], label=f"{country}")
    axs[1].plot(temp["dl_cons"], label=f"{country}")
    axs[2].plot(temp["dl_gfcf"], label=f"{country}")
axs[0].set_title("Differenced and Logged GDP")
axs[1].set_title("Differenced and Logged Consumption")
axs[2].set_title("Differenced and Logged Investment")
axs[0].legend()
axs[1].legend()
axs[2].legend()
plt.suptitle("Macroeconomic Timeseries");
../_images/c6091e62189bcc6882ad0300120f267fb8115666b8d6c8e2e721476bc7bdbdca.png

Ireland’s Economic Situation#

Ireland is somewhat infamous for its GDP numbers that are largely the product of foreign direct investment and inflated beyond expectation in recent years by the investment and taxation deals offered to large multi-nationals. We’ll look here at just the relationship between GDP and consumption. We just want to show the mechanics of the VAR estimation, you shouldn’t read too much into the subsequent analysis.

ireland_df = gdp_hierarchical[gdp_hierarchical["country"] == "Ireland"]
ireland_df.reset_index(inplace=True, drop=True)
ireland_df.head()
country iso2c iso3c year GDP CONS GFCF dl_gdp dl_cons dl_gfcf more_than_10 time
0 Ireland IE IRL 1971 3.314234e+10 2.897699e+10 8.317518e+09 0.034110 0.043898 0.085452 True 1
1 Ireland IE IRL 1972 3.529322e+10 3.063538e+10 8.967782e+09 0.062879 0.055654 0.075274 True 2
2 Ireland IE IRL 1973 3.695956e+10 3.280221e+10 1.041728e+10 0.046134 0.068340 0.149828 True 3
3 Ireland IE IRL 1974 3.853412e+10 3.381524e+10 9.207243e+09 0.041720 0.030416 -0.123476 True 4
4 Ireland IE IRL 1975 4.071386e+10 3.477232e+10 8.874887e+09 0.055024 0.027910 -0.036765 True 5
n_lags = 2
n_eqs = 2
priors = {
    ## Set prior for expected positive relationship between the variables.
    "lag_coefs": {"mu": 0.3, "sigma": 1},
    "alpha": {"mu": 0, "sigma": 0.1},
    "noise_chol": {"eta": 1, "sigma": 1},
    "noise": {"sigma": 1},
}
model, idata_ireland = make_model(
    n_lags, n_eqs, ireland_df[["dl_gdp", "dl_cons"]], priors, prior_checks=False
)
idata_ireland
Sampling: [alpha, lag_coefs, noise_chol, obs]
INFO:pymc:Sampling: [alpha, lag_coefs, noise_chol, obs]
Auto-assigning NUTS sampler...
INFO:pymc:Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
INFO:pymc:Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
INFO:pymc:Multiprocess sampling (4 chains in 4 jobs)
NUTS: [lag_coefs, alpha, noise_chol]
INFO:pymc:NUTS: [lag_coefs, alpha, noise_chol]
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
100.00% [12000/12000 01:14<00:00 Sampling 4 chains, 0 divergences]
Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 114 seconds.
INFO:pymc:Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 114 seconds.
Sampling: [obs]
INFO:pymc:Sampling: [obs]
100.00% [8000/8000 01:46<00:00]
arviz.InferenceData
    • <xarray.Dataset>
      Dimensions:                (chain: 4, draw: 2000, equations: 2, lags: 2,
                                  cross_vars: 2, noise_chol_dim_0: 3, time: 49,
                                  betaX_dim_1: 1, noise_chol_corr_dim_0: 2,
                                  noise_chol_corr_dim_1: 2, noise_chol_stds_dim_0: 2)
      Coordinates:
        * chain                  (chain) int64 0 1 2 3
        * draw                   (draw) int64 0 1 2 3 4 5 ... 1995 1996 1997 1998 1999
        * equations              (equations) <U7 'dl_gdp' 'dl_cons'
        * lags                   (lags) int64 1 2
        * cross_vars             (cross_vars) <U7 'dl_gdp' 'dl_cons'
        * noise_chol_dim_0       (noise_chol_dim_0) int64 0 1 2
        * time                   (time) int64 2 3 4 5 6 7 8 9 ... 44 45 46 47 48 49 50
        * betaX_dim_1            (betaX_dim_1) int64 0
        * noise_chol_corr_dim_0  (noise_chol_corr_dim_0) int64 0 1
        * noise_chol_corr_dim_1  (noise_chol_corr_dim_1) int64 0 1
        * noise_chol_stds_dim_0  (noise_chol_stds_dim_0) int64 0 1
      Data variables:
          lag_coefs              (chain, draw, equations, lags, cross_vars) float64 ...
          alpha                  (chain, draw, equations) float64 0.02911 ... 0.003431
          noise_chol             (chain, draw, noise_chol_dim_0) float64 0.03778 .....
          betaX                  (chain, draw, time, betaX_dim_1) float64 0.03706 ....
          noise_chol_corr        (chain, draw, noise_chol_corr_dim_0, noise_chol_corr_dim_1) float64 ...
          noise_chol_stds        (chain, draw, noise_chol_stds_dim_0) float64 0.037...
      Attributes:
          created_at:                 2022-12-10T19:04:29.191143
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2
          sampling_time:              113.52139687538147
          tuning_steps:               1000

    • <xarray.Dataset>
      Dimensions:    (chain: 4, draw: 2000, time: 49, equations: 2)
      Coordinates:
        * chain      (chain) int64 0 1 2 3
        * draw       (draw) int64 0 1 2 3 4 5 6 ... 1993 1994 1995 1996 1997 1998 1999
        * time       (time) int64 2 3 4 5 6 7 8 9 10 11 ... 42 43 44 45 46 47 48 49 50
        * equations  (equations) <U7 'dl_gdp' 'dl_cons'
      Data variables:
          obs        (chain, draw, time, equations) float64 0.08365 0.0472 ... 0.02789
      Attributes:
          created_at:                 2022-12-10T19:06:48.368723
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:  (chain: 4, draw: 2000, time: 49)
      Coordinates:
        * chain    (chain) int64 0 1 2 3
        * draw     (draw) int64 0 1 2 3 4 5 6 7 ... 1993 1994 1995 1996 1997 1998 1999
        * time     (time) int64 2 3 4 5 6 7 8 9 10 11 ... 42 43 44 45 46 47 48 49 50
      Data variables:
          obs      (chain, draw, time) float64 4.255 4.911 5.119 ... -1.215 2.023
      Attributes:
          created_at:                 2022-12-10T19:04:31.278701
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:              (chain: 4, draw: 2000)
      Coordinates:
        * chain                (chain) int64 0 1 2 3
        * draw                 (draw) int64 0 1 2 3 4 5 ... 1995 1996 1997 1998 1999
      Data variables: (12/16)
          energy_error         (chain, draw) float64 0.1612 0.3398 ... 0.236 -0.05665
          energy               (chain, draw) float64 -185.0 -181.7 ... -181.0 -179.5
          lp                   (chain, draw) float64 189.9 189.6 186.5 ... 185.4 190.1
          acceptance_rate      (chain, draw) float64 0.548 0.279 ... 0.8959 0.7233
          step_size            (chain, draw) float64 0.3363 0.3363 ... 0.2929 0.2929
          max_energy_error     (chain, draw) float64 1.487 10.93 ... 0.24 0.7903
          ...                   ...
          tree_depth           (chain, draw) int64 4 4 3 4 4 4 4 3 ... 4 4 4 4 4 4 4 4
          diverging            (chain, draw) bool False False False ... False False
          index_in_trajectory  (chain, draw) int64 4 2 4 -4 -5 1 ... 7 -1 6 -15 6 4
          process_time_diff    (chain, draw) float64 0.04586 0.04647 ... 0.05515
          smallest_eigval      (chain, draw) float64 nan nan nan nan ... nan nan nan
          step_size_bar        (chain, draw) float64 0.3313 0.3313 ... 0.3453 0.3453
      Attributes:
          created_at:                 2022-12-10T19:04:29.214346
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2
          sampling_time:              113.52139687538147
          tuning_steps:               1000

    • <xarray.Dataset>
      Dimensions:                (chain: 1, draw: 500, noise_chol_stds_dim_0: 2,
                                  equations: 2, noise_chol_dim_0: 3, lags: 2,
                                  cross_vars: 2, time: 49, betaX_dim_1: 1,
                                  noise_chol_corr_dim_0: 2, noise_chol_corr_dim_1: 2)
      Coordinates:
        * chain                  (chain) int64 0
        * draw                   (draw) int64 0 1 2 3 4 5 ... 494 495 496 497 498 499
        * noise_chol_stds_dim_0  (noise_chol_stds_dim_0) int64 0 1
        * equations              (equations) <U7 'dl_gdp' 'dl_cons'
        * noise_chol_dim_0       (noise_chol_dim_0) int64 0 1 2
        * lags                   (lags) int64 1 2
        * cross_vars             (cross_vars) <U7 'dl_gdp' 'dl_cons'
        * time                   (time) int64 2 3 4 5 6 7 8 9 ... 44 45 46 47 48 49 50
        * betaX_dim_1            (betaX_dim_1) int64 0
        * noise_chol_corr_dim_0  (noise_chol_corr_dim_0) int64 0 1
        * noise_chol_corr_dim_1  (noise_chol_corr_dim_1) int64 0 1
      Data variables:
          noise_chol_stds        (chain, draw, noise_chol_stds_dim_0) float64 1.718...
          alpha                  (chain, draw, equations) float64 0.066 ... -0.1081
          noise_chol             (chain, draw, noise_chol_dim_0) float64 1.718 ... ...
          lag_coefs              (chain, draw, equations, lags, cross_vars) float64 ...
          betaX                  (chain, draw, time, betaX_dim_1) float64 -0.1237 ....
          noise_chol_corr        (chain, draw, noise_chol_corr_dim_0, noise_chol_corr_dim_1) float64 ...
      Attributes:
          created_at:                 2022-12-10T19:02:19.951274
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:    (chain: 1, draw: 500, time: 49, equations: 2)
      Coordinates:
        * chain      (chain) int64 0
        * draw       (draw) int64 0 1 2 3 4 5 6 7 ... 492 493 494 495 496 497 498 499
        * time       (time) int64 2 3 4 5 6 7 8 9 10 11 ... 42 43 44 45 46 47 48 49 50
        * equations  (equations) <U7 'dl_gdp' 'dl_cons'
      Data variables:
          obs        (chain, draw, time, equations) float64 2.329 1.173 ... 1.204
      Attributes:
          created_at:                 2022-12-10T19:02:19.958802
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:    (time: 49, equations: 2)
      Coordinates:
        * time       (time) int64 2 3 4 5 6 7 8 9 10 11 ... 42 43 44 45 46 47 48 49 50
        * equations  (equations) <U7 'dl_gdp' 'dl_cons'
      Data variables:
          obs        (time, equations) float64 0.04613 0.06834 ... 0.1264 0.05461
      Attributes:
          created_at:                 2022-12-10T19:02:19.961523
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:    (time: 49, equations: 2)
      Coordinates:
        * time       (time) int64 2 3 4 5 6 7 8 9 10 11 ... 42 43 44 45 46 47 48 49 50
        * equations  (equations) <U7 'dl_gdp' 'dl_cons'
      Data variables:
          data_obs   (time, equations) float64 0.04613 0.06834 ... 0.1264 0.05461
      Attributes:
          created_at:                 2022-12-10T19:02:19.964775
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

az.plot_trace(idata_ireland, var_names=["lag_coefs", "alpha", "betaX"], kind="rank_vlines");
../_images/b50a46359c387725f0e5b2d6aa697f6868258c7c15edf146637540e188f35ae2.png
def plot_ppc_macro(idata, df, group="posterior_predictive"):
    df = pd.DataFrame(idata["observed_data"]["obs"].data, columns=["dl_gdp", "dl_cons"])
    fig, axs = plt.subplots(2, 1, figsize=(20, 10))
    axs = axs.flatten()
    ppc = az.extract_dataset(idata, group=group, num_samples=100)["obs"]

    shade_background(ppc, axs, 0, "inferno")
    axs[0].plot(np.arange(ppc.shape[0]), ppc[:, 0, :].mean(axis=1), color="cyan", label="Mean")
    axs[0].plot(df["dl_gdp"], "o", mfc="black", mec="white", mew=1, markersize=7, label="Observed")
    axs[0].set_title("Differenced and Logged GDP")
    axs[0].legend()
    shade_background(ppc, axs, 1, "inferno")
    axs[1].plot(df["dl_cons"], "o", mfc="black", mec="white", mew=1, markersize=7, label="Observed")
    axs[1].plot(np.arange(ppc.shape[0]), ppc[:, 1, :].mean(axis=1), color="cyan", label="Mean")
    axs[1].set_title("Differenced and Logged Consumption")
    axs[1].legend()


plot_ppc_macro(idata_ireland, ireland_df)
../_images/82b1de5442ac4f8df6ae48045c0d547b4eea1b87dc968c6218d5f5bccba54e42.png
ax = az.plot_posterior(
    idata_ireland,
    var_names="noise_chol_corr",
    hdi_prob="hide",
    point_estimate="mean",
    grid=(2, 2),
    kind="hist",
    ec="black",
    figsize=(10, 6),
)
../_images/7fcae92bf3ea44db16431f4dd4ab16f1f4121db4cb50ed023d93ea470173d87f.png

Comparison with Statsmodels#

It’s worthwhile comparing these model fits to the one achieved by Statsmodels just to see if we can recover a similar story.

VAR_model = sm.tsa.VAR(ireland_df[["dl_gdp", "dl_cons"]])
results = VAR_model.fit(2, trend="c")
results.params
dl_gdp dl_cons
const 0.034145 0.006996
L1.dl_gdp 0.324904 0.330003
L1.dl_cons 0.076629 0.305824
L2.dl_gdp 0.137721 -0.053677
L2.dl_cons -0.278745 0.033728

The intercept parameters broadly agree with our Bayesian model with some differences in the implied relationships defined by the estimates for the lagged terms.

corr = pd.DataFrame(results.resid_corr, columns=["dl_gdp", "dl_cons"])
corr.index = ["dl_gdp", "dl_cons"]
corr
dl_gdp dl_cons
dl_gdp 1.000000 0.435807
dl_cons 0.435807 1.000000

The residual correlation estimates reported by statsmodels agree quite closely with the multivariate gaussian correlation between the variables in our Bayesian model.

az.summary(idata_ireland, var_names=["alpha", "lag_coefs", "noise_chol_corr"])
/Users/nathanielforde/opt/miniconda3/envs/pymc_ar_ex/lib/python3.9/site-packages/arviz/stats/diagnostics.py:586: RuntimeWarning: invalid value encountered in double_scalars
  (between_chain_variance / within_chain_variance + num_samples - 1) / (num_samples)
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
alpha[dl_gdp] 0.026 0.008 0.011 0.042 0.000 0.000 7455.0 5288.0 1.0
alpha[dl_cons] 0.008 0.007 -0.005 0.020 0.000 0.000 7816.0 5908.0 1.0
lag_coefs[dl_gdp, 1, dl_gdp] 0.300 1.007 -1.648 2.139 0.010 0.011 10791.0 6055.0 1.0
lag_coefs[dl_gdp, 1, dl_cons] 0.297 1.009 -1.610 2.161 0.010 0.012 10765.0 5523.0 1.0
lag_coefs[dl_gdp, 2, dl_gdp] 0.302 0.994 -1.522 2.194 0.010 0.011 10562.0 5908.0 1.0
lag_coefs[dl_gdp, 2, dl_cons] 0.300 0.988 -1.579 2.133 0.009 0.011 11157.0 5958.0 1.0
lag_coefs[dl_cons, 1, dl_gdp] 0.330 0.102 0.137 0.519 0.001 0.001 9264.0 6099.0 1.0
lag_coefs[dl_cons, 1, dl_cons] 0.272 0.169 -0.039 0.596 0.002 0.001 8862.0 5911.0 1.0
lag_coefs[dl_cons, 2, dl_gdp] -0.026 0.112 -0.228 0.190 0.001 0.001 8986.0 5885.0 1.0
lag_coefs[dl_cons, 2, dl_cons] -0.002 0.170 -0.314 0.327 0.002 0.002 9026.0 5129.0 1.0
noise_chol_corr[0, 0] 1.000 0.000 1.000 1.000 0.000 0.000 8000.0 8000.0 NaN
noise_chol_corr[0, 1] 0.416 0.124 0.184 0.648 0.001 0.001 7714.0 5323.0 1.0
noise_chol_corr[1, 0] 0.416 0.124 0.184 0.648 0.001 0.001 7714.0 5323.0 1.0
noise_chol_corr[1, 1] 1.000 0.000 1.000 1.000 0.000 0.000 7495.0 7594.0 1.0

We plot the alpha parameter estimates against the Statsmodels estimates

az.plot_posterior(idata_ireland, var_names=["alpha"], ref_val=[0.034145, 0.006996]);
../_images/b2b0f29db0e82ab20cbd636bdc6128ad7d23c54926e2a1c352f641d887140abf.png
az.plot_posterior(
    idata_ireland,
    var_names=["lag_coefs"],
    ref_val=[0.330003, -0.053677],
    coords={"equations": "dl_cons", "lags": [1, 2], "cross_vars": "dl_gdp"},
);
../_images/92be9dc49a6c2aa715d7e04e3ae5721961f6807fa15c1917b545ba77696abbce.png

We can see here again how the Bayesian VAR model recovers much of the same story. Similar magnitudes in the estimates for the alpha terms for both equations and a clear relationship between the first lagged GDP numbers and consumption along with a very similar covariance structure.

Adding a Bayesian Twist: Hierarchical VARs#

In addition we can add some hierarchical parameters if we want to model multiple countries and the relationship between these economic metrics at the national level. This is a useful technique in the cases where we have reasonably short timeseries data because it allows us to “borrow” information across the countries to inform the estimates of the key parameters.

def make_hierarchical_model(n_lags, n_eqs, df, group_field, prior_checks=True):
    cols = [col for col in df.columns if col != group_field]
    coords = {"lags": np.arange(n_lags) + 1, "equations": cols, "cross_vars": cols}

    groups = df[group_field].unique()

    with pm.Model(coords=coords) as model:
        ## Hierarchical Priors
        rho = pm.Beta("rho", alpha=2, beta=2)
        alpha_hat_location = pm.Normal("alpha_hat_location", 0, 0.1)
        alpha_hat_scale = pm.InverseGamma("alpha_hat_scale", 3, 0.5)
        beta_hat_location = pm.Normal("beta_hat_location", 0, 0.1)
        beta_hat_scale = pm.InverseGamma("beta_hat_scale", 3, 0.5)
        omega_global, _, _ = pm.LKJCholeskyCov(
            "omega_global", n=n_eqs, eta=1.0, sd_dist=pm.Exponential.dist(1)
        )

        for grp in groups:
            df_grp = df[df[group_field] == grp][cols]
            z_scale_beta = pm.InverseGamma(f"z_scale_beta_{grp}", 3, 0.5)
            z_scale_alpha = pm.InverseGamma(f"z_scale_alpha_{grp}", 3, 0.5)
            lag_coefs = pm.Normal(
                f"lag_coefs_{grp}",
                mu=beta_hat_location,
                sigma=beta_hat_scale * z_scale_beta,
                dims=["equations", "lags", "cross_vars"],
            )
            alpha = pm.Normal(
                f"alpha_{grp}",
                mu=alpha_hat_location,
                sigma=alpha_hat_scale * z_scale_alpha,
                dims=("equations",),
            )

            betaX = calc_ar_step(lag_coefs, n_eqs, n_lags, df_grp)
            betaX = pm.Deterministic(f"betaX_{grp}", betaX)
            mean = alpha + betaX

            n = df_grp.shape[1]
            noise_chol, _, _ = pm.LKJCholeskyCov(
                f"noise_chol_{grp}", eta=10, n=n, sd_dist=pm.Exponential.dist(1)
            )
            omega = pm.Deterministic(f"omega_{grp}", rho * omega_global + (1 - rho) * noise_chol)
            obs = pm.MvNormal(f"obs_{grp}", mu=mean, chol=omega, observed=df_grp.values[n_lags:])

        if prior_checks:
            idata = pm.sample_prior_predictive()
            return model, idata
        else:
            idata = pm.sample_prior_predictive()
            idata.extend(sample_blackjax_nuts(2000, random_seed=120))
            pm.sample_posterior_predictive(idata, extend_inferencedata=True)
    return model, idata

The model design allows for a non-centred parameterisation of the key likeihood for each of the individual country components by allowing the us to shift the country specific estimates away from the hierarchical mean. This is done by rho * omega_global + (1 - rho) * noise_chol line. The parameter rho determines the share of impact each country’s data contributes to the estimation of the covariance relationship among the economic variables. Similar country specific adjustments are made with the z_alpha_scale and z_beta_scale parameters.

df_final = gdp_hierarchical[["country", "dl_gdp", "dl_cons", "dl_gfcf"]]
model_full_test, idata_full_test = make_hierarchical_model(
    2,
    3,
    df_final,
    "country",
    prior_checks=False,
)
Sampling: [alpha_Australia, alpha_Canada, alpha_Chile, alpha_Ireland, alpha_New Zealand, alpha_South Africa, alpha_United Kingdom, alpha_United States, alpha_hat_location, alpha_hat_scale, beta_hat_location, beta_hat_scale, lag_coefs_Australia, lag_coefs_Canada, lag_coefs_Chile, lag_coefs_Ireland, lag_coefs_New Zealand, lag_coefs_South Africa, lag_coefs_United Kingdom, lag_coefs_United States, noise_chol_Australia, noise_chol_Canada, noise_chol_Chile, noise_chol_Ireland, noise_chol_New Zealand, noise_chol_South Africa, noise_chol_United Kingdom, noise_chol_United States, obs_Australia, obs_Canada, obs_Chile, obs_Ireland, obs_New Zealand, obs_South Africa, obs_United Kingdom, obs_United States, omega_global, rho, z_scale_alpha_Australia, z_scale_alpha_Canada, z_scale_alpha_Chile, z_scale_alpha_Ireland, z_scale_alpha_New Zealand, z_scale_alpha_South Africa, z_scale_alpha_United Kingdom, z_scale_alpha_United States, z_scale_beta_Australia, z_scale_beta_Canada, z_scale_beta_Chile, z_scale_beta_Ireland, z_scale_beta_New Zealand, z_scale_beta_South Africa, z_scale_beta_United Kingdom, z_scale_beta_United States]
INFO:pymc:Sampling: [alpha_Australia, alpha_Canada, alpha_Chile, alpha_Ireland, alpha_New Zealand, alpha_South Africa, alpha_United Kingdom, alpha_United States, alpha_hat_location, alpha_hat_scale, beta_hat_location, beta_hat_scale, lag_coefs_Australia, lag_coefs_Canada, lag_coefs_Chile, lag_coefs_Ireland, lag_coefs_New Zealand, lag_coefs_South Africa, lag_coefs_United Kingdom, lag_coefs_United States, noise_chol_Australia, noise_chol_Canada, noise_chol_Chile, noise_chol_Ireland, noise_chol_New Zealand, noise_chol_South Africa, noise_chol_United Kingdom, noise_chol_United States, obs_Australia, obs_Canada, obs_Chile, obs_Ireland, obs_New Zealand, obs_South Africa, obs_United Kingdom, obs_United States, omega_global, rho, z_scale_alpha_Australia, z_scale_alpha_Canada, z_scale_alpha_Chile, z_scale_alpha_Ireland, z_scale_alpha_New Zealand, z_scale_alpha_South Africa, z_scale_alpha_United Kingdom, z_scale_alpha_United States, z_scale_beta_Australia, z_scale_beta_Canada, z_scale_beta_Chile, z_scale_beta_Ireland, z_scale_beta_New Zealand, z_scale_beta_South Africa, z_scale_beta_United Kingdom, z_scale_beta_United States]
Compiling...
Compilation time =  0:00:41.131670
Sampling...
Sampling time =  0:02:47.507208
Transforming variables...
Transformation time =  0:00:46.927859
Computing Log Likelihood...
Log Likelihood time =  0:00:06.982440
Sampling: [obs_Australia, obs_Canada, obs_Chile, obs_Ireland, obs_New Zealand, obs_South Africa, obs_United Kingdom, obs_United States]
INFO:pymc:Sampling: [obs_Australia, obs_Canada, obs_Chile, obs_Ireland, obs_New Zealand, obs_South Africa, obs_United Kingdom, obs_United States]
100.00% [8000/8000 14:45<00:00]
idata_full_test
arviz.InferenceData
    • <xarray.Dataset>
      Dimensions:                               (chain: 4, draw: 2000, equations: 3,
                                                 lags: 2, cross_vars: 3,
                                                 omega_global_dim_0: 6,
                                                 noise_chol_Australia_dim_0: 6,
                                                 noise_chol_Canada_dim_0: 6,
                                                 noise_chol_Chile_dim_0: 6,
                                                 ...
                                                 betaX_United States_dim_1: 1,
                                                 noise_chol_United States_corr_dim_0: 3,
                                                 noise_chol_United States_corr_dim_1: 3,
                                                 noise_chol_United States_stds_dim_0: 3,
                                                 omega_United States_dim_0: 3,
                                                 omega_United States_dim_1: 3)
      Coordinates: (12/73)
        * chain                                 (chain) int64 0 1 2 3
        * draw                                  (draw) int64 0 1 2 ... 1997 1998 1999
        * equations                             (equations) <U7 'dl_gdp' ... 'dl_gfcf'
        * lags                                  (lags) int64 1 2
        * cross_vars                            (cross_vars) <U7 'dl_gdp' ... 'dl_g...
        * omega_global_dim_0                    (omega_global_dim_0) int64 0 1 2 3 4 5
          ...                                    ...
        * betaX_United States_dim_1             (betaX_United States_dim_1) int64 0
        * noise_chol_United States_corr_dim_0   (noise_chol_United States_corr_dim_0) int64 ...
        * noise_chol_United States_corr_dim_1   (noise_chol_United States_corr_dim_1) int64 ...
        * noise_chol_United States_stds_dim_0   (noise_chol_United States_stds_dim_0) int64 ...
        * omega_United States_dim_0             (omega_United States_dim_0) int64 0...
        * omega_United States_dim_1             (omega_United States_dim_1) int64 0...
      Data variables: (12/80)
          alpha_hat_location                    (chain, draw) float64 0.02064 ... 0...
          beta_hat_location                     (chain, draw) float64 0.03119 ... 0...
          lag_coefs_Australia                   (chain, draw, equations, lags, cross_vars) float64 ...
          alpha_Australia                       (chain, draw, equations) float64 0....
          lag_coefs_Canada                      (chain, draw, equations, lags, cross_vars) float64 ...
          alpha_Canada                          (chain, draw, equations) float64 0....
          ...                                    ...
          noise_chol_United Kingdom_stds        (chain, draw, noise_chol_United Kingdom_stds_dim_0) float64 ...
          omega_United Kingdom                  (chain, draw, omega_United Kingdom_dim_0, omega_United Kingdom_dim_1) float64 ...
          betaX_United States                   (chain, draw, betaX_United States_dim_0, betaX_United States_dim_1) float64 ...
          noise_chol_United States_corr         (chain, draw, noise_chol_United States_corr_dim_0, noise_chol_United States_corr_dim_1) float64 ...
          noise_chol_United States_stds         (chain, draw, noise_chol_United States_stds_dim_0) float64 ...
          omega_United States                   (chain, draw, omega_United States_dim_0, omega_United States_dim_1) float64 ...
      Attributes:
          created_at:     2022-12-10T19:12:26.413672
          arviz_version:  0.12.1

    • <xarray.Dataset>
      Dimensions:                   (chain: 4, draw: 2000, obs_Australia_dim_0: 49,
                                     obs_Australia_dim_1: 3, obs_Canada_dim_0: 22,
                                     obs_Canada_dim_1: 3, obs_Chile_dim_0: 49,
                                     obs_Chile_dim_1: 3, obs_Ireland_dim_0: 49,
                                     obs_Ireland_dim_1: 3, obs_New Zealand_dim_0: 41,
                                     obs_New Zealand_dim_1: 3,
                                     obs_South Africa_dim_0: 49,
                                     obs_South Africa_dim_1: 3,
                                     obs_United Kingdom_dim_0: 49,
                                     obs_United Kingdom_dim_1: 3,
                                     obs_United States_dim_0: 46,
                                     obs_United States_dim_1: 3)
      Coordinates: (12/18)
        * chain                     (chain) int64 0 1 2 3
        * draw                      (draw) int64 0 1 2 3 4 ... 1996 1997 1998 1999
        * obs_Australia_dim_0       (obs_Australia_dim_0) int64 0 1 2 3 ... 46 47 48
        * obs_Australia_dim_1       (obs_Australia_dim_1) int64 0 1 2
        * obs_Canada_dim_0          (obs_Canada_dim_0) int64 0 1 2 3 4 ... 18 19 20 21
        * obs_Canada_dim_1          (obs_Canada_dim_1) int64 0 1 2
          ...                        ...
        * obs_South Africa_dim_0    (obs_South Africa_dim_0) int64 0 1 2 ... 46 47 48
        * obs_South Africa_dim_1    (obs_South Africa_dim_1) int64 0 1 2
        * obs_United Kingdom_dim_0  (obs_United Kingdom_dim_0) int64 0 1 2 ... 47 48
        * obs_United Kingdom_dim_1  (obs_United Kingdom_dim_1) int64 0 1 2
        * obs_United States_dim_0   (obs_United States_dim_0) int64 0 1 2 ... 43 44 45
        * obs_United States_dim_1   (obs_United States_dim_1) int64 0 1 2
      Data variables:
          obs_Australia             (chain, draw, obs_Australia_dim_0, obs_Australia_dim_1) float64 ...
          obs_Canada                (chain, draw, obs_Canada_dim_0, obs_Canada_dim_1) float64 ...
          obs_Chile                 (chain, draw, obs_Chile_dim_0, obs_Chile_dim_1) float64 ...
          obs_Ireland               (chain, draw, obs_Ireland_dim_0, obs_Ireland_dim_1) float64 ...
          obs_New Zealand           (chain, draw, obs_New Zealand_dim_0, obs_New Zealand_dim_1) float64 ...
          obs_South Africa          (chain, draw, obs_South Africa_dim_0, obs_South Africa_dim_1) float64 ...
          obs_United Kingdom        (chain, draw, obs_United Kingdom_dim_0, obs_United Kingdom_dim_1) float64 ...
          obs_United States         (chain, draw, obs_United States_dim_0, obs_United States_dim_1) float64 ...
      Attributes:
          created_at:                 2022-12-10T19:34:46.232070
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:                   (chain: 4, draw: 2000, obs_Australia_dim_0: 49,
                                     obs_Canada_dim_0: 22, obs_Chile_dim_0: 49,
                                     obs_Ireland_dim_0: 49,
                                     obs_New Zealand_dim_0: 41,
                                     obs_South Africa_dim_0: 49,
                                     obs_United Kingdom_dim_0: 49,
                                     obs_United States_dim_0: 46)
      Coordinates:
        * chain                     (chain) int64 0 1 2 3
        * draw                      (draw) int64 0 1 2 3 4 ... 1996 1997 1998 1999
        * obs_Australia_dim_0       (obs_Australia_dim_0) int64 0 1 2 3 ... 46 47 48
        * obs_Canada_dim_0          (obs_Canada_dim_0) int64 0 1 2 3 4 ... 18 19 20 21
        * obs_Chile_dim_0           (obs_Chile_dim_0) int64 0 1 2 3 4 ... 45 46 47 48
        * obs_Ireland_dim_0         (obs_Ireland_dim_0) int64 0 1 2 3 ... 45 46 47 48
        * obs_New Zealand_dim_0     (obs_New Zealand_dim_0) int64 0 1 2 3 ... 38 39 40
        * obs_South Africa_dim_0    (obs_South Africa_dim_0) int64 0 1 2 ... 46 47 48
        * obs_United Kingdom_dim_0  (obs_United Kingdom_dim_0) int64 0 1 2 ... 47 48
        * obs_United States_dim_0   (obs_United States_dim_0) int64 0 1 2 ... 43 44 45
      Data variables:
          obs_Australia             (chain, draw, obs_Australia_dim_0) float64 8.23...
          obs_Canada                (chain, draw, obs_Canada_dim_0) float64 8.129 ....
          obs_Chile                 (chain, draw, obs_Chile_dim_0) float64 4.683 .....
          obs_Ireland               (chain, draw, obs_Ireland_dim_0) float64 5.325 ...
          obs_New Zealand           (chain, draw, obs_New Zealand_dim_0) float64 6....
          obs_South Africa          (chain, draw, obs_South Africa_dim_0) float64 6...
          obs_United Kingdom        (chain, draw, obs_United Kingdom_dim_0) float64 ...
          obs_United States         (chain, draw, obs_United States_dim_0) float64 ...
      Attributes:
          created_at:     2022-12-10T19:12:26.462020
          arviz_version:  0.12.1

    • <xarray.Dataset>
      Dimensions:                               (chain: 1, draw: 500,
                                                 noise_chol_United States_corr_dim_0: 3,
                                                 noise_chol_United States_corr_dim_1: 3,
                                                 equations: 3,
                                                 noise_chol_Ireland_stds_dim_0: 3,
                                                 noise_chol_Chile_dim_0: 6, lags: 2,
                                                 ...
                                                 noise_chol_Chile_stds_dim_0: 3,
                                                 noise_chol_South Africa_corr_dim_0: 3,
                                                 noise_chol_South Africa_corr_dim_1: 3,
                                                 noise_chol_United Kingdom_corr_dim_0: 3,
                                                 noise_chol_United Kingdom_corr_dim_1: 3,
                                                 noise_chol_Australia_stds_dim_0: 3)
      Coordinates: (12/73)
        * chain                                 (chain) int64 0
        * draw                                  (draw) int64 0 1 2 3 ... 497 498 499
        * noise_chol_United States_corr_dim_0   (noise_chol_United States_corr_dim_0) int64 ...
        * noise_chol_United States_corr_dim_1   (noise_chol_United States_corr_dim_1) int64 ...
        * equations                             (equations) <U7 'dl_gdp' ... 'dl_gfcf'
        * noise_chol_Ireland_stds_dim_0         (noise_chol_Ireland_stds_dim_0) int64 ...
          ...                                    ...
        * noise_chol_Chile_stds_dim_0           (noise_chol_Chile_stds_dim_0) int64 ...
        * noise_chol_South Africa_corr_dim_0    (noise_chol_South Africa_corr_dim_0) int64 ...
        * noise_chol_South Africa_corr_dim_1    (noise_chol_South Africa_corr_dim_1) int64 ...
        * noise_chol_United Kingdom_corr_dim_0  (noise_chol_United Kingdom_corr_dim_0) int64 ...
        * noise_chol_United Kingdom_corr_dim_1  (noise_chol_United Kingdom_corr_dim_1) int64 ...
        * noise_chol_Australia_stds_dim_0       (noise_chol_Australia_stds_dim_0) int64 ...
      Data variables: (12/80)
          noise_chol_United States_corr         (chain, draw, noise_chol_United States_corr_dim_0, noise_chol_United States_corr_dim_1) float64 ...
          alpha_South Africa                    (chain, draw, equations) float64 -0...
          z_scale_alpha_United States           (chain, draw) float64 0.4974 ... 0....
          noise_chol_Ireland_stds               (chain, draw, noise_chol_Ireland_stds_dim_0) float64 ...
          noise_chol_Chile                      (chain, draw, noise_chol_Chile_dim_0) float64 ...
          lag_coefs_Australia                   (chain, draw, equations, lags, cross_vars) float64 ...
          ...                                    ...
          noise_chol_Chile_stds                 (chain, draw, noise_chol_Chile_stds_dim_0) float64 ...
          noise_chol_South Africa_corr          (chain, draw, noise_chol_South Africa_corr_dim_0, noise_chol_South Africa_corr_dim_1) float64 ...
          lag_coefs_South Africa                (chain, draw, equations, lags, cross_vars) float64 ...
          noise_chol_United Kingdom_corr        (chain, draw, noise_chol_United Kingdom_corr_dim_0, noise_chol_United Kingdom_corr_dim_1) float64 ...
          noise_chol_Australia_stds             (chain, draw, noise_chol_Australia_stds_dim_0) float64 ...
          alpha_Chile                           (chain, draw, equations) float64 -0...
      Attributes:
          created_at:                 2022-12-10T19:08:03.439050
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:                   (chain: 1, draw: 500, obs_Ireland_dim_0: 49,
                                     obs_Ireland_dim_1: 3, obs_New Zealand_dim_0: 41,
                                     obs_New Zealand_dim_1: 3,
                                     obs_Australia_dim_0: 49, obs_Australia_dim_1: 3,
                                     obs_South Africa_dim_0: 49,
                                     obs_South Africa_dim_1: 3,
                                     obs_United States_dim_0: 46,
                                     obs_United States_dim_1: 3,
                                     obs_Canada_dim_0: 22, obs_Canada_dim_1: 3,
                                     obs_Chile_dim_0: 49, obs_Chile_dim_1: 3,
                                     obs_United Kingdom_dim_0: 49,
                                     obs_United Kingdom_dim_1: 3)
      Coordinates: (12/18)
        * chain                     (chain) int64 0
        * draw                      (draw) int64 0 1 2 3 4 5 ... 495 496 497 498 499
        * obs_Ireland_dim_0         (obs_Ireland_dim_0) int64 0 1 2 3 ... 45 46 47 48
        * obs_Ireland_dim_1         (obs_Ireland_dim_1) int64 0 1 2
        * obs_New Zealand_dim_0     (obs_New Zealand_dim_0) int64 0 1 2 3 ... 38 39 40
        * obs_New Zealand_dim_1     (obs_New Zealand_dim_1) int64 0 1 2
          ...                        ...
        * obs_Canada_dim_0          (obs_Canada_dim_0) int64 0 1 2 3 4 ... 18 19 20 21
        * obs_Canada_dim_1          (obs_Canada_dim_1) int64 0 1 2
        * obs_Chile_dim_0           (obs_Chile_dim_0) int64 0 1 2 3 4 ... 45 46 47 48
        * obs_Chile_dim_1           (obs_Chile_dim_1) int64 0 1 2
        * obs_United Kingdom_dim_0  (obs_United Kingdom_dim_0) int64 0 1 2 ... 47 48
        * obs_United Kingdom_dim_1  (obs_United Kingdom_dim_1) int64 0 1 2
      Data variables:
          obs_Ireland               (chain, draw, obs_Ireland_dim_0, obs_Ireland_dim_1) float64 ...
          obs_New Zealand           (chain, draw, obs_New Zealand_dim_0, obs_New Zealand_dim_1) float64 ...
          obs_Australia             (chain, draw, obs_Australia_dim_0, obs_Australia_dim_1) float64 ...
          obs_South Africa          (chain, draw, obs_South Africa_dim_0, obs_South Africa_dim_1) float64 ...
          obs_United States         (chain, draw, obs_United States_dim_0, obs_United States_dim_1) float64 ...
          obs_Canada                (chain, draw, obs_Canada_dim_0, obs_Canada_dim_1) float64 ...
          obs_Chile                 (chain, draw, obs_Chile_dim_0, obs_Chile_dim_1) float64 ...
          obs_United Kingdom        (chain, draw, obs_United Kingdom_dim_0, obs_United Kingdom_dim_1) float64 ...
      Attributes:
          created_at:                 2022-12-10T19:08:03.508786
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2

    • <xarray.Dataset>
      Dimensions:                   (obs_Australia_dim_0: 49, obs_Australia_dim_1: 3,
                                     obs_Canada_dim_0: 22, obs_Canada_dim_1: 3,
                                     obs_Chile_dim_0: 49, obs_Chile_dim_1: 3,
                                     obs_Ireland_dim_0: 49, obs_Ireland_dim_1: 3,
                                     obs_New Zealand_dim_0: 41,
                                     obs_New Zealand_dim_1: 3,
                                     obs_South Africa_dim_0: 49,
                                     obs_South Africa_dim_1: 3,
                                     obs_United Kingdom_dim_0: 49,
                                     obs_United Kingdom_dim_1: 3,
                                     obs_United States_dim_0: 46,
                                     obs_United States_dim_1: 3)
      Coordinates: (12/16)
        * obs_Australia_dim_0       (obs_Australia_dim_0) int64 0 1 2 3 ... 46 47 48
        * obs_Australia_dim_1       (obs_Australia_dim_1) int64 0 1 2
        * obs_Canada_dim_0          (obs_Canada_dim_0) int64 0 1 2 3 4 ... 18 19 20 21
        * obs_Canada_dim_1          (obs_Canada_dim_1) int64 0 1 2
        * obs_Chile_dim_0           (obs_Chile_dim_0) int64 0 1 2 3 4 ... 45 46 47 48
        * obs_Chile_dim_1           (obs_Chile_dim_1) int64 0 1 2
          ...                        ...
        * obs_South Africa_dim_0    (obs_South Africa_dim_0) int64 0 1 2 ... 46 47 48
        * obs_South Africa_dim_1    (obs_South Africa_dim_1) int64 0 1 2
        * obs_United Kingdom_dim_0  (obs_United Kingdom_dim_0) int64 0 1 2 ... 47 48
        * obs_United Kingdom_dim_1  (obs_United Kingdom_dim_1) int64 0 1 2
        * obs_United States_dim_0   (obs_United States_dim_0) int64 0 1 2 ... 43 44 45
        * obs_United States_dim_1   (obs_United States_dim_1) int64 0 1 2
      Data variables:
          obs_Australia             (obs_Australia_dim_0, obs_Australia_dim_1) float64 ...
          obs_Canada                (obs_Canada_dim_0, obs_Canada_dim_1) float64 0....
          obs_Chile                 (obs_Chile_dim_0, obs_Chile_dim_1) float64 -0.0...
          obs_Ireland               (obs_Ireland_dim_0, obs_Ireland_dim_1) float64 ...
          obs_New Zealand           (obs_New Zealand_dim_0, obs_New Zealand_dim_1) float64 ...
          obs_South Africa          (obs_South Africa_dim_0, obs_South Africa_dim_1) float64 ...
          obs_United Kingdom        (obs_United Kingdom_dim_0, obs_United Kingdom_dim_1) float64 ...
          obs_United States         (obs_United States_dim_0, obs_United States_dim_1) float64 ...
      Attributes:
          created_at:                 2022-12-10T19:08:03.519442
          arviz_version:              0.12.1
          inference_library:          pymc
          inference_library_version:  4.2.2