#!/usr/bin/env python
# coding: utf-8

# DO NOT EDIT
# Autogenerated from the notebook statespace_tvpvar_mcmc_cfa.ipynb.
# Edit the notebook and then sync the output with this file.
#
# flake8: noqa
# DO NOT EDIT

# ## TVP-VAR, MCMC, and sparse simulation smoothing

from importlib import reload
import numpy as np
import pandas as pd
import statsmodels.api as sm
import matplotlib.pyplot as plt

from scipy.stats import invwishart, invgamma

# Get the macro dataset
dta = sm.datasets.macrodata.load_pandas().data
dta.index = pd.date_range('1959Q1', '2009Q3', freq='QS')

# ### Background
#
# Bayesian analysis of linear Gaussian state space models via Markov chain
# Monte Carlo (MCMC) methods has become both commonplace and relatively
# straightforward in recent years, due especially to advances in sampling
# from the joint posterior of the unobserved state vector conditional on the
# data and model parameters (see especially Carter and Kohn (1994), de Jong
# and Shephard (1995), and Durbin and Koopman (2002)). This is particularly
# useful for Gibbs sampling MCMC approaches.
#
# While these procedures make use of the forward/backward application of
# the recursive Kalman filter and smoother, another recent line of research
# takes a different approach and constructs the posterior joint distribution
# of the entire vector of states at once - see in particular Chan and
# Jeliazkov (2009) for an econometric time series treatment and McCausland
# et al. (2011) for a more general survey. In particular, the posterior mean
# and precision matrix are constructed explicitly, with the latter a sparse
# band matrix. Advantage is then taken of efficient algorithms for Cholesky
# factorization of sparse band matrices; this reduces memory costs and can
# improve performance. Following McCausland et al. (2011), we refer to this
# method as the "Cholesky Factor Algorithm" (CFA) approach.
#
# The CFA-based simulation smoother has some advantages and some drawbacks
# compared to that based on the more typical Kalman filter and smoother
# (KFS).
#
# **Advantages of CFA**:
#
# - Derivation of the joint posterior distribution is relatively
# straightforward and easy to understand.
# - In some cases can be both faster and less memory-intensive than the
# KFS approach
#     - In the Appendix at the end of this notebook, we briefly discuss
# the performance of the two simulation smoothers for the TVP-VAR model. In
# summary: simple tests on a single machine suggest that for the TVP-VAR
# model, the CFA and KFS implementations in Statsmodels have about the same
# runtimes, while both implementations are about twice as fast as the
# replication code, written in Matlab, provided by Chan and Jeliazkov
# (2009).
#
# **Drawbacks of CFA**:
#
# The main drawback is that this method has not (at least so far) reached
# the generality of the KFS approach. For example:
#
# - It can not be used with models that have reduced-rank error terms in
# the observation or state equations.
#     - One implication of this is that the typical state space model
# trick of including identities in the state equation to accommodate, for
# example, higher-order lags in autoregressive models is not applicable.
# These models can still be handled by the CFA approach, but at the cost of
# requiring a slightly different implementation for each lag that is
# included.
#     - As an example, standard ways of representing ARMA and VARMA
# processes in state space form do include identities in the observation
# and/or state equations, and so the basic formulas presented in Chan and
# Jeliazkov (2009) do not apply immediately to these models.
# - Less flexibility is available in the state initialization / prior.

# ### Implementation in Statsmodels
#
# A CFA simulation smoother along the lines of the basic formulas
# presented in Chan and Jeliazkov (2009) has been implemented in
# Statsmodels.
#
# **Notes**:
#
# - Therefore, the CFA simulation smoother in Statsmodels so-far only
# supports the case that the state transition is truly a first-order Markov
# process (i.e. it does not support a p-th order Markov process that has
# been stacked using identities into a first-order process).
# - By contrast, the KFS smoother in Statsmodels is fully general any can
# be used for any state space model, including those with stacked p-th order
# Markov processes or other identities in the observation and state
# equations.
#
# Either a KFS or the CFA simulation smoothers can be constructed from a
# state space model using the `simulation_smoother` method. To show the
# basic idea, we first consider a simple example.

# #### Local level model
#
# A local level model decomposes an observed series $y_t$ into a
# persistent trend $\mu_t$ and a transitory error component
#
# $$
# \begin{aligned}
# y_t & = \mu_t + \varepsilon_t, \qquad \varepsilon_t \sim N(0,
# \sigma_\text{irregular}^2) \\
# \mu_t & = \mu_{t-1} + \eta_t, \quad ~ \eta_t \sim N(0,
# \sigma_\text{level}^2)
# \end{aligned}
# $$
#
# This model satisfies the requirements of the CFA simulation smoother
# because both the observation error term $\varepsilon_t$ and the state
# innovation term $\eta_t$ are non-degenerate - that is, their covariance
# matrices are full rank.
#
# We apply this model to inflation, and consider simulating draws from the
# posterior of the joint state vector. That is, we are interested in
# sampling from
#
# $$p(\mu^t \mid y^t, \sigma_\text{irregular}^2, \sigma_\text{level}^2)$$
#
# where we define $\mu^t \equiv (\mu_1, \dots, \mu_T)'$ and $y^t \equiv
# (y_1, \dots, y_T)'$.
#
# In Statsmodels, the local level model falls into the more general class
# of "unobserved components" models, and can be constructed as follows:

# Construct a local level model for inflation
mod = sm.tsa.UnobservedComponents(dta.infl, 'llevel')

# Fit the model's parameters (sigma2_varepsilon and sigma2_eta)
# via maximum likelihood
res = mod.fit()
print(res.params)

# Create simulation smoother objects
sim_kfs = mod.simulation_smoother()  # default method is KFS
sim_cfa = mod.simulation_smoother(method='cfa')  # can specify CFA method

# The simulation smoother objects `sim_kfs` and `sim_cfa` have `simulate`
# methods that perform simulation smoothing. Each time that `simulate` is
# called, the `simulated_state` attribute will be re-populated with a new
# simulated draw from the posterior.
#
# Below, we construct 20 simulated paths for the trend, using the KFS and
# CFA approaches, where the simulation is at the maximum likelihood
# parameter estimates.

nsimulations = 20
simulated_state_kfs = pd.DataFrame(np.zeros((mod.nobs, nsimulations)),
                                   index=dta.index)
simulated_state_cfa = pd.DataFrame(np.zeros((mod.nobs, nsimulations)),
                                   index=dta.index)

for i in range(nsimulations):
    # Apply KFS simulation smoothing
    sim_kfs.simulate()
    # Save the KFS simulated state
    simulated_state_kfs.iloc[:, i] = sim_kfs.simulated_state[0]

    # Apply CFA simulation smoothing
    sim_cfa.simulate()
    # Save the CFA simulated state
    simulated_state_cfa.iloc[:, i] = sim_cfa.simulated_state[0]

# Plotting the observed data and the simulations created using each method
# below, it is not too hard to see that these two methods are doing the same
# thing.

# Plot the inflation data along with simulated trends
fig, axes = plt.subplots(2, figsize=(15, 6))

# Plot data and KFS simulations
dta.infl.plot(ax=axes[0], color='k')
axes[0].set_title('Simulations based on KFS approach, MLE parameters')
simulated_state_kfs.plot(ax=axes[0], color='C0', alpha=0.25, legend=False)

# Plot data and CFA simulations
dta.infl.plot(ax=axes[1], color='k')
axes[1].set_title('Simulations based on CFA approach, MLE parameters')
simulated_state_cfa.plot(ax=axes[1], color='C0', alpha=0.25, legend=False)

# Add a legend, clean up layout
handles, labels = axes[0].get_legend_handles_labels()
axes[0].legend(handles[:2], ['Data', 'Simulated state'])
fig.tight_layout()

# #### Updating the model's parameters
#
# The simulation smoothers are tied to the model instance, here the
# variable `mod`. Whenever the model instance is updated with new
# parameters, the simulation smoothers will take those new parameters into
# account in future calls to the `simulate` method.
#
# This is convenient for MCMC algorithms, which repeatedly (a) update the
# model's parameters, (b) draw a sample of the state vector, and then (c)
# draw new values for the model's parameters.
#
# Here we will change the model to a different parameterization that
# yields a smoother trend, and show how the simulated values change (for
# brevity we only show the simulations from the KFS approach, but
# simulations from the CFA approach would be the same).

fig, ax = plt.subplots(figsize=(15, 3))

# Update the model's parameterization to one that attributes more
# variation in inflation to the observation error and so has less
# variation in the trend component
mod.update([4, 0.05])

# Plot simulations
for i in range(nsimulations):
    sim_kfs.simulate()
    ax.plot(dta.index,
            sim_kfs.simulated_state[0],
            color='C0',
            alpha=0.25,
            label='Simulated state')

# Plot data
dta.infl.plot(ax=ax, color='k', label='Data', zorder=-1)

# Add title, legend, clean up layout
ax.set_title(
    'Simulations with alternative parameterization yielding a smoother trend')
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles[-2:], labels[-2:])
fig.tight_layout()

# ### Application: Bayesian analysis of a TVP-VAR model by MCMC
#
# One of the applications that Chan and Jeliazkov (2009) consider is the
# time-varying parameters vector autoregression (TVP-VAR) model, estimated
# with Bayesian Gibb sampling (MCMC) methods. They apply this to model the
# co-movements in four macroeconomic time series:
#
# - Real GDP growth
# - Inflation
# - Unemployment rate
# - Short-term interest rates
#
# We will replicate their example, using a very similar dataset that is
# included in Statsmodels.

# Subset to the four variables of interest
y = dta[['realgdp', 'cpi', 'unemp', 'tbilrate']].copy()
y.columns = ['gdp', 'inf', 'unemp', 'int']

# Convert to real GDP growth and CPI inflation rates
y[['gdp', 'inf']] = np.log(y[['gdp', 'inf']]).diff() * 100
y = y.iloc[1:]

fig, ax = plt.subplots(figsize=(15, 5))
y.plot(ax=ax)
ax.set_title(
    'Evolution of macroeconomic variables included in TVP-VAR exercise')

# #### TVP-VAR model
#
# **Note**: this section is based on Chan and Jeliazkov (2009) section
# 3.1, which can be consulted for additional details.
#
# The usual (time-invariant) VAR(1) model is typically written:
#
# $$
# \begin{aligned}
# y_t & = \mu + \Phi y_{t-1} + \varepsilon_t, \qquad \varepsilon_t \sim
# N(0, H)
# \end{aligned}
# $$
#
# where $y_t$ is a $p \times 1$ vector of variables observed at time $t$
# and $H$ is a covariance matrix.
#
# The TVP-VAR(1) model generalizes this to allow the coefficients to vary
# over time according. Stacking all the parameters into a vector according
# to $\alpha_t = \text{vec}([\mu_t : \Phi_t])$, where $\text{vec}$ denotes
# the operation that stacks columns of a matrix into a vector, we model
# their evolution over time according to:
#
# $$\alpha_{i,t+1} = \alpha_{i, t} + \eta_{i,t}, \qquad \eta_{i, t} \sim
# N(0, \sigma_i^2)$$
#
# In other words, each parameter evolves independently according to a
# random walk.
#
# Note that there are $p$ coefficients in $\mu_t$ and $p^2$ coefficients
# in $\Phi_t$, so the full state vector $\alpha$ is shaped $p * (p + 1)
# \times 1$.

# Putting the TVP-VAR(1) model into state-space form is relatively
# straightforward, and in fact we just have to re-write the observation
# equation into SUR form:
#
# $$
# \begin{aligned}
# y_t & = Z_t \alpha_t + \varepsilon_t, \qquad \varepsilon_t \sim N(0, H)
# \\
# \alpha_{t+1} & = \alpha_t + \eta_t, \qquad \eta_t \sim N(0,
# \text{diag}(\{\sigma_i^2\}))
# \end{aligned}
# $$
#
# where
#
# $$
# Z_t = \begin{bmatrix}
# 1 & y_{t-1}' & 0 & \dots & &  0 \\
# 0 & 0 & 1 & y_{t-1}' &  & 0 \\
# \vdots & & & \ddots & \ddots & \vdots \\
# 0 & 0 & 0 & 0 & 1 & y_{t-1}'  \\
# \end{bmatrix}
# $$
#
# As long as $H$ is full rank and each of the variances $\sigma_i^2$ is
# non-zero, the model satisfies the requirements of the CFA simulation
# smoother.
#
# We also need to specify the initialization / prior for the initial
# state, $\alpha_1$. Here we will follow Chan and Jeliazkov (2009) in using
# $\alpha_1 \sim N(0, 5 I)$, although we could also model it as diffuse.

# Aside from the time-varying coefficients $\alpha_t$, the other
# parameters that we will need to estimate are terms in the covariance
# matrix $H$ and the random walk variances $\sigma_i^2$.

# #### TVP-VAR model in Statsmodels
#
# Constructing this model programatically in Statsmodels is also
# relatively straightforward, since there are basically four steps:
#
# 1. Create a new `TVPVAR` class as a subclass of
# `sm.tsa.statespace.MLEModel`
# 2. Fill in the fixed values of the state space system matrices
# 3. Specify the initialization of $\alpha_1$
# 4. Create a method for updating the state space system matrices with new
# values of the covariance matrix $H$ and the random walk variances
# $\sigma_i^2$.
#
# To do this, first note that the general state space representation used
# by Statsmodels is:
#
# $$
# \begin{aligned}
# y_t & = d_t + Z_t \alpha_t + \varepsilon_t, \qquad \varepsilon_t \sim
# N(0, H_t) \\
# \alpha_{t+1} & = c_t + T_t \alpha_t + R_t \eta_t, \qquad \eta_t \sim
# N(0, Q_t) \\
# \end{aligned}
# $$
#
# Then the TVP-VAR(1) model implies the following specializations:
#
# - The intercept terms are zero, i.e. $c_t = d_t = 0$
# - The design matrix $Z_t$ is time-varying but its values are fixed as
# described above (i.e. its values contain ones and lags of $y_t$)
# - The observation covariance matrix is not time-varying, i.e. $H_t =
# H_{t+1} = H$
# - The transition matrix is not time-varying and is equal to the identity
# matrix, i.e. $T_t = T_{t+1} = I$
# - The selection matrix $R_t$ is not time-varying and is also equal to
# the identity matrix, i.e. $R_t = R_{t+1} = I$
# - The state covariance matrix $Q_t$ is not time-varying and is diagonal,
# i.e. $Q_t = Q_{t+1} = \text{diag}(\{\sigma_i^2\})$


# 1. Create a new TVPVAR class as a subclass of sm.tsa.statespace.MLEModel
class TVPVAR(sm.tsa.statespace.MLEModel):
    # Steps 2-3 are best done in the class "constructor", i.e. the __init__ method
    def __init__(self, y):
        # Create a matrix with [y_t' : y_{t-1}'] for t = 2, ..., T
        augmented = sm.tsa.lagmat(y,
                                  1,
                                  trim='both',
                                  original='in',
                                  use_pandas=True)
        # Separate into y_t and z_t = [1 : y_{t-1}']
        p = y.shape[1]
        y_t = augmented.iloc[:, :p]
        z_t = sm.add_constant(augmented.iloc[:, p:])

        # Recall that the length of the state vector is p * (p + 1)
        k_states = p * (p + 1)
        super().__init__(y_t, exog=z_t, k_states=k_states)

        # Note that the state space system matrices default to contain zeros,
        # so we don't need to explicitly set c_t = d_t = 0.

        # Construct the design matrix Z_t
        # Notes:
        # -> self.k_endog = p is the dimension of the observed vector
        # -> self.k_states = p * (p + 1) is the dimension of the observed vector
        # -> self.nobs = T is the number of observations in y_t
        self['design'] = np.zeros((self.k_endog, self.k_states, self.nobs))
        for i in range(self.k_endog):
            start = i * (self.k_endog + 1)
            end = start + self.k_endog + 1
            self['design', i, start:end, :] = z_t.T

        # Construct the transition matrix T = I
        self['transition'] = np.eye(k_states)

        # Construct the selection matrix R = I
        self['selection'] = np.eye(k_states)

        # Step 3: Initialize the state vector as alpha_1 ~ N(0, 5I)
        self.ssm.initialize('known', stationary_cov=5 * np.eye(self.k_states))

    # Step 4. Create a method that we can call to update H and Q
    def update_variances(self, obs_cov, state_cov_diag):
        self['obs_cov'] = obs_cov
        self['state_cov'] = np.diag(state_cov_diag)

    # Finally, it can be convenient to define human-readable names for
    # each element of the state vector. These will be available in output
    @property
    def state_names(self):
        state_names = np.empty((self.k_endog, self.k_endog + 1), dtype=object)
        for i in range(self.k_endog):
            endog_name = self.endog_names[i]
            state_names[i] = (['intercept.%s' % endog_name] + [
                'L1.%s->%s' % (other_name, endog_name)
                for other_name in self.endog_names
            ])
        return state_names.ravel().tolist()


# The above class defined the state space model for any given dataset. Now
# we need to create a specific instance of it with the dataset that we
# created earlier containing real GDP growth, inflation, unemployment, and
# interest rates.

# Create an instance of our TVPVAR class with our observed dataset y
mod = TVPVAR(y)

# #### Preliminary investigation with ad-hoc parameters in H, Q

# In our analysis below, we will need to begin our MCMC iterations with
# some initial parameterization. Following Chan and Jeliazkov (2009) we will
# set $H$ to be the sample covariance matrix of our dataset, and we will set
# $\sigma_i^2 = 0.01$ for each $i$.
#
# Before discussing the MCMC scheme that will allow us to make inferences
# about the model, first we can consider the output of the model when simply
# plugging in these initial parameters. To fill in these parameters, we use
# the `update_variances` method that we defined earlier and then perform
# Kalman filtering and smoothing conditional on those parameters.
#
# **Warning: This exercise is just by way of explanation - we must wait
# for the output of the MCMC exercise to study the actual implications of
# the model in a meaningful way**.

initial_obs_cov = np.cov(y.T)
initial_state_cov_diag = [0.01] * mod.k_states

# Update H and Q
mod.update_variances(initial_obs_cov, initial_state_cov_diag)

# Perform Kalman filtering and smoothing
# (the [] is just an empty list that in some models might contain
# additional parameters. Here, we don't have any additional parameters
# so we just pass an empty list)
initial_res = mod.smooth([])

# The `initial_res` variable contains the output of Kalman filtering and
# smoothing, conditional on those initial parameters. In particular, we may
# be interested in the "smoothed states", which are $E[\alpha_t \mid y^t, H,
# \{\sigma_i^2\}]$.
#
# First, lets create a function that graphs the coefficients over time,
# separated into the equations for equation of the observed variables.


def plot_coefficients_by_equation(states):
    fig, axes = plt.subplots(2, 2, figsize=(15, 8))

    # The way we defined Z_t implies that the first 5 elements of the
    # state vector correspond to the first variable in y_t, which is GDP growth
    ax = axes[0, 0]
    states.iloc[:, :5].plot(ax=ax)
    ax.set_title('GDP growth')
    ax.legend()

    # The next 5 elements correspond to inflation
    ax = axes[0, 1]
    states.iloc[:, 5:10].plot(ax=ax)
    ax.set_title('Inflation rate')
    ax.legend()

    # The next 5 elements correspond to unemployment
    ax = axes[1, 0]
    states.iloc[:, 10:15].plot(ax=ax)
    ax.set_title('Unemployment equation')
    ax.legend()

    # The last 5 elements correspond to the interest rate
    ax = axes[1, 1]
    states.iloc[:, 15:20].plot(ax=ax)
    ax.set_title('Interest rate equation')
    ax.legend()

    return ax


# Now, we are interested in the smoothed states, which are available in
# the `states.smoothed` attribute out our results object `initial_res`.
#
# As the graph below shows, the initial parameterization implies
# substantial time-variation in some of the coefficients.

# Here, for illustration purposes only, we plot the time-varying
# coefficients conditional on an ad-hoc parameterization

# Recall that `initial_res` contains the Kalman filtering and smoothing,
# and the `states.smoothed` attribute contains the smoothed states
plot_coefficients_by_equation(initial_res.states.smoothed)

# #### Bayesian estimation via MCMC
#
# We will now implement the Gibbs sampler scheme described in Chan and
# Jeliazkov (2009), Algorithm 2.
#
#
# We use the following (conditionally conjugate) priors:
#
# $$
# \begin{aligned}
# H & \sim \mathcal{IW}(\nu_1^0, S_1^0) \\
# \sigma_i^2 & \sim \mathcal{IG} \left ( \frac{\nu_{i2}^0}{2},
# \frac{S_{i2}^0}{2} \right )
# \end{aligned}
# $$
#
# where $\mathcal{IW}$ denotes the inverse-Wishart distribution and
# $\mathcal{IG}$ denotes the inverse-Gamma distribution. We set the prior
# hyperparameters as:
#
# $$
# \begin{aligned}
# v_1^0 = T + 3, & \quad S_1^0 = I \\
# v_{i2}^0 = 6, & \quad S_{i2}^0 = 0.01 \qquad \text{for each} ~ i\\
# \end{aligned}
# $$

# Prior hyperparameters

# Prior for obs. cov. is inverse-Wishart(v_1^0=k + 3, S10=I)
v10 = mod.k_endog + 3
S10 = np.eye(mod.k_endog)

# Prior for state cov. variances is inverse-Gamma(v_{i2}^0 / 2 = 3,
# S+{i2}^0 / 2 = 0.005)
vi20 = 6
Si20 = 0.01

# Before running the MCMC iterations, there are a couple of practical
# steps:
#
# 1. Create arrays to store the draws of our state vector, observation
# covariance matrix, and state error variances.
# 2. Put the initial values for H and Q (described above) into the storage
# vectors
# 3. Construct the simulation smoother object associated with our `TVPVAR`
# instance to make draws of the state vector

# Gibbs sampler setup
niter = 11000
nburn = 1000

# 1. Create storage arrays
store_states = np.zeros((niter + 1, mod.nobs, mod.k_states))
store_obs_cov = np.zeros((niter + 1, mod.k_endog, mod.k_endog))
store_state_cov = np.zeros((niter + 1, mod.k_states))

# 2. Put in the initial values
store_obs_cov[0] = initial_obs_cov
store_state_cov[0] = initial_state_cov_diag
mod.update_variances(store_obs_cov[0], store_state_cov[0])

# 3. Construct posterior samplers
sim = mod.simulation_smoother(method='cfa')

# As before, we could have used either the simulation smoother based on
# the Kalman filter and smoother or that based on the Cholesky Factor
# Algorithm.

for i in range(niter):
    mod.update_variances(store_obs_cov[i], store_state_cov[i])
    sim.simulate()

    # 1. Sample states
    store_states[i + 1] = sim.simulated_state.T

    # 2. Simulate obs cov
    fitted = np.matmul(mod['design'].transpose(2, 0, 1),
                       store_states[i + 1][..., None])[..., 0]
    resid = mod.endog - fitted
    store_obs_cov[i + 1] = invwishart.rvs(v10 + mod.nobs,
                                          S10 + resid.T @ resid)

    # 3. Simulate state cov variances
    resid = store_states[i + 1, 1:] - store_states[i + 1, :-1]
    sse = np.sum(resid**2, axis=0)

    for j in range(mod.k_states):
        rv = invgamma.rvs((vi20 + mod.nobs - 1) / 2, scale=(Si20 + sse[j]) / 2)
        store_state_cov[i + 1, j] = rv

# After removing a number of initial draws, the remaining draws from the
# posterior allow us to conduct inference. Below, we plot the posterior mean
# of the time-varying regression coefficients.
#
# (**Note**: these plots are different from those in Figure 1 of the
# published version of Chan and Jeliazkov (2009), but they are very similar
# to those produced by the Matlab replication code available at
# http://joshuachan.org/code/code_TVPVAR.html)

# Collect the posterior means of each time-varying coefficient
states_posterior_mean = pd.DataFrame(np.mean(store_states[nburn + 1:], axis=0),
                                     index=mod._index,
                                     columns=mod.state_names)

# Plot these means over time
plot_coefficients_by_equation(states_posterior_mean)

# Python also has a number of libraries to assist with exploring Bayesian
# models. Here we'll just use the [arviz](https://arviz-
# devs.github.io/arviz/index.html) package to explore the credible intervals
# of each of the covariance and variance parameters, although it makes
# available a much wider set of tools for analysis.

import arviz as az

# Collect the observation error covariance parameters
az_obs_cov = az.convert_to_inference_data({
    ('Var[%s]' % mod.endog_names[i] if i == j else 'Cov[%s, %s]' %
     (mod.endog_names[i], mod.endog_names[j])): store_obs_cov[nburn + 1:, i, j]
    for i in range(mod.k_endog) for j in range(i, mod.k_endog)
})

# Plot the credible intervals
az.plot_forest(az_obs_cov, figsize=(8, 7))

# Collect the state innovation variance parameters
az_state_cov = az.convert_to_inference_data({
    r'$\sigma^2$[%s]' % mod.state_names[i]: store_state_cov[nburn + 1:, i]
    for i in range(mod.k_states)
})

# Plot the credible intervals
az.plot_forest(az_state_cov, figsize=(8, 7))

# ### Appendix: performance
#
# Finally, we run a few simple tests to compare the performance of the KFS
# and CFA simulation smoothers by using the `%timeit` Jupyter notebook
# magic.
#
# One caveat is that the KFS simulation smoother can produce a variety of
# output beyond just simulations of the posterior state vector, and these
# additional computations could bias the results. To make the results
# comparable, we will tell the KFS simulation smoother to only compute
# simulations of the state by using the `simulation_output` argument.

from statsmodels.tsa.statespace.simulation_smoother import SIMULATION_STATE

sim_cfa = mod.simulation_smoother(method='cfa')
sim_kfs = mod.simulation_smoother(simulation_output=SIMULATION_STATE)

# Then we can use the following code to perform a basic timing exercise:
#
# ```python
# %timeit -n 10000 -r 3 sim_cfa.simulate()
# %timeit -n 10000 -r 3 sim_kfs.simulate()
# ```
#
# On the machine this was tested on, this resulted in the following:
#
# ```
# 2.06 ms ± 26.5 µs per loop (mean ± std. dev. of 3 runs, 10000 loops
# each)
# 2.02 ms ± 68.4 µs per loop (mean ± std. dev. of 3 runs, 10000 loops
# each)
# ```

# These results suggest that - at least for this model - there are not
# noticeable computational gains from the CFA approach relative to the KFS
# approach. However, this does not rule out the following:
#
# 1. The Statsmodels implementation of the CFA simulation smoother could
# possibly be further optimized
# 2. The CFA approach may only show improvement for certain models (for
# example with a large number of `endog` variables)
#
# One simple way to take a first pass at assessing the first possibility
# is to compare the runtime of the Statsmodels implementation of the CFA
# simulation smoother to the Matlab implementation in the replication codes
# of Chan and Jeliazkov (2009), available at
# http://joshuachan.org/code/code_TVPVAR.html.
#
# While the Statsmodels version of the CFA simulation smoother is written
# in Cython and compiled to C code, the Matlab version takes advantage of
# the Matlab's sparse matrix capabilities. As a result, even though it is
# not compiled code,  we might expect it to have relatively good
# performance.
#
# On the machine this was tested on, the Matlab version typically ran the
# MCMC loop with 11,000 iterations in 70-75 seconds, while the MCMC loop in
# this notebook using the Statsmodels CFA simulation smoother (see above),
# also with 11,0000 iterations, ran in 40-45 seconds. This is some evidence
# that the Statsmodels implementation of the CFA smoother already performs
# relatively well (although it does not rule out that there are additional
# gains possible).

# ### Bibliography
#
# Carter, Chris K., and Robert Kohn. "On Gibbs sampling for state space
# models." Biometrika 81, no. 3 (1994): 541-553.
#
# Chan, Joshua CC, and Ivan Jeliazkov. "Efficient simulation and
# integrated likelihood estimation in state space models." International
# Journal of Mathematical Modelling and Numerical Optimisation 1, no. 1-2
# (2009): 101-120.
#
# De Jong, Piet, and Neil Shephard. "The simulation smoother for time
# series models." Biometrika 82, no. 2 (1995): 339-350.
#
# Durbin, James, and Siem Jan Koopman. "A simple and efficient simulation
# smoother for state space time series analysis." Biometrika 89, no. 3
# (2002): 603-616.
#
# McCausland, William J., Shirley Miller, and Denis Pelletier. "Simulation
# smoothing for state–space models: A computational efficiency analysis."
# Computational Statistics & Data Analysis 55, no. 1 (2011): 199-212.