47. Posterior Distributions for AR(1) Parameters#

GPU

This lecture was built using a machine with the latest CUDA and CUDANN frameworks installed with access to a GPU.

To run this lecture on Google Colab, click on the “play” icon top right, select Colab, and set the runtime environment to include a GPU.

To run this lecture on your own machine, you need to install the software listed following this notice.

!pip install numpyro jax

In addition to what’s included in base Anaconda, we need to install the following packages

!pip install arviz pymc

Show code cell output Hide code cell output

Collecting arviz

  Downloading arviz-0.21.0-py3-none-any.whl.metadata (8.8 kB)

Collecting pymc

  Downloading pymc-5.22.0-py3-none-any.whl.metadata (16 kB)
Requirement already satisfied: setuptools>=60.0.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (75.1.0)
Requirement already satisfied: matplotlib>=3.5 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (3.9.2)
Requirement already satisfied: numpy>=1.23.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (1.26.4)
Requirement already satisfied: scipy>=1.9.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (1.13.1)
Requirement already satisfied: packaging in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (24.1)
Requirement already satisfied: pandas>=1.5.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (2.2.2)
Requirement already satisfied: xarray>=2022.6.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (2023.6.0)
Collecting h5netcdf>=1.0.2 (from arviz)

  Downloading h5netcdf-1.6.1-py3-none-any.whl.metadata (13 kB)
Requirement already satisfied: typing-extensions>=4.1.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from arviz) (4.11.0)
Collecting xarray-einstats>=0.3 (from arviz)

  Downloading xarray_einstats-0.8.0-py3-none-any.whl.metadata (5.8 kB)
Requirement already satisfied: cachetools>=4.2.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from pymc) (5.3.3)
Requirement already satisfied: cloudpickle in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from pymc) (3.0.0)

Collecting pytensor<2.31,>=2.30.2 (from pymc)

  Downloading pytensor-2.30.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (10.0 kB)
Requirement already satisfied: rich>=13.7.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from pymc) (13.7.1)
Requirement already satisfied: threadpoolctl<4.0.0,>=3.1.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from pymc) (3.5.0)
Requirement already satisfied: h5py in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from h5netcdf>=1.0.2->arviz) (3.11.0)
Requirement already satisfied: contourpy>=1.0.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (1.2.0)
Requirement already satisfied: cycler>=0.10 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (0.11.0)
Requirement already satisfied: fonttools>=4.22.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (4.51.0)
Requirement already satisfied: kiwisolver>=1.3.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (1.4.4)
Requirement already satisfied: pillow>=8 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (10.4.0)
Requirement already satisfied: pyparsing>=2.3.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (3.1.2)
Requirement already satisfied: python-dateutil>=2.7 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from matplotlib>=3.5->arviz) (2.9.0.post0)
Requirement already satisfied: pytz>=2020.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from pandas>=1.5.0->arviz) (2024.1)
Requirement already satisfied: tzdata>=2022.7 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from pandas>=1.5.0->arviz) (2023.3)

Collecting filelock>=3.15 (from pytensor<2.31,>=2.30.2->pymc)
  Downloading filelock-3.18.0-py3-none-any.whl.metadata (2.9 kB)

Collecting etuples (from pytensor<2.31,>=2.30.2->pymc)
  Downloading etuples-0.3.9.tar.gz (30 kB)

  Preparing metadata (setup.py) ... ?25l-

 done

?25hCollecting logical-unification (from pytensor<2.31,>=2.30.2->pymc)
  Downloading logical-unification-0.4.6.tar.gz (31 kB)

  Preparing metadata (setup.py) ... ?25l-

 done

?25hCollecting miniKanren (from pytensor<2.31,>=2.30.2->pymc)
  Downloading miniKanren-1.0.3.tar.gz (41 kB)

  Preparing metadata (setup.py) ... ?25l-

 \ done

?25hCollecting cons (from pytensor<2.31,>=2.30.2->pymc)
  Downloading cons-0.4.6.tar.gz (26 kB)

  Preparing metadata (setup.py) ... ?25l-

 done
?25hRequirement already satisfied: markdown-it-py>=2.2.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from rich>=13.7.1->pymc) (2.2.0)
Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from rich>=13.7.1->pymc) (2.15.1)

Requirement already satisfied: mdurl~=0.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from markdown-it-py>=2.2.0->rich>=13.7.1->pymc) (0.1.0)
Requirement already satisfied: six>=1.5 in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from python-dateutil>=2.7->matplotlib>=3.5->arviz) (1.16.0)
Requirement already satisfied: toolz in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from logical-unification->pytensor<2.31,>=2.30.2->pymc) (0.12.0)
Requirement already satisfied: multipledispatch in /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages (from logical-unification->pytensor<2.31,>=2.30.2->pymc) (0.6.0)
Downloading arviz-0.21.0-py3-none-any.whl (1.7 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/1.7 MB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.7/1.7 MB 47.4 MB/s eta 0:00:00
?25h

Downloading pymc-5.22.0-py3-none-any.whl (517 kB)
Downloading h5netcdf-1.6.1-py3-none-any.whl (49 kB)
Downloading pytensor-2.30.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/2.0 MB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.0/2.0 MB 96.9 MB/s eta 0:00:00
?25hDownloading xarray_einstats-0.8.0-py3-none-any.whl (32 kB)
Downloading filelock-3.18.0-py3-none-any.whl (16 kB)

Building wheels for collected packages: cons, logical-unification, etuples, miniKanren

  Building wheel for cons (setup.py) ... ?25l-

 done
?25h  Created wheel for cons: filename=cons-0.4.6-py3-none-any.whl size=9094 sha256=86b3be2bd10ffb0010144ebd0ac23b60e0eee8c96dff6b1825adb1234a53473e
  Stored in directory: /home/runner/.cache/pip/wheels/91/0c/f7/7ba4ffbf0db6bf272bee59789131f6ec2d26966307a3054267

  Building wheel for logical-unification (setup.py) ... ?25l-

 done
?25h  Created wheel for logical-unification: filename=logical_unification-0.4.6-py3-none-any.whl size=13912 sha256=f4b283880e15295de0109fb4753ec1da19f6544b5912dac266fa3d7defd898cb
  Stored in directory: /home/runner/.cache/pip/wheels/4b/07/cf/3fdaedafc40d7b1861c98be0df1bf3b2b3e87e03187fa95298

  Building wheel for etuples (setup.py) ... ?25l-

 done
?25h  Created wheel for etuples: filename=etuples-0.3.9-py3-none-any.whl size=12619 sha256=58b63de5d1a418cde40e18442aad887a4255ca4ed8c2e4746bae60427f085c1e
  Stored in directory: /home/runner/.cache/pip/wheels/9e/ea/42/9562f37ce4e92913b74da9bcd504a749f0bf5d42d1d148ca76

  Building wheel for miniKanren (setup.py) ... ?25l-

 done
?25h  Created wheel for miniKanren: filename=miniKanren-1.0.3-py3-none-any.whl size=23908 sha256=070773e8cc5fe2b980dd20ea9f982a82422f0af460c0f6783a2fd82074ffdf8e
  Stored in directory: /home/runner/.cache/pip/wheels/57/48/e3/441a42666f9a4ac9ccb6897d6046d81cf970c986be39938c6e
Successfully built cons logical-unification etuples miniKanren

Installing collected packages: filelock, logical-unification, h5netcdf, cons, xarray-einstats, etuples, miniKanren, arviz, pytensor, pymc
  Attempting uninstall: filelock
    Found existing installation: filelock 3.13.1
    Uninstalling filelock-3.13.1:
      Successfully uninstalled filelock-3.13.1

Successfully installed arviz-0.21.0 cons-0.4.6 etuples-0.3.9 filelock-3.18.0 h5netcdf-1.6.1 logical-unification-0.4.6 miniKanren-1.0.3 pymc-5.22.0 pytensor-2.30.3 xarray-einstats-0.8.0

We’ll begin with some Python imports.

import arviz as az
import pymc as pmc
import numpyro
from numpyro import distributions as dist

import numpy as np
import jax.numpy as jnp
from jax import random
import matplotlib.pyplot as plt

import logging
logging.basicConfig()
logger = logging.getLogger('pymc')
logger.setLevel(logging.CRITICAL)

This lecture uses Bayesian methods offered by pymc and numpyro to make statistical inferences about two parameters of a univariate first-order autoregression.

The model is a good laboratory for illustrating consequences of alternative ways of modeling the distribution of the initial $y_{0}$ :

As a fixed number
As a random variable drawn from the stationary distribution of the ${y_{t}}$ stochastic process

The first component of the statistical model is

(47.1)#

y_{t + 1} = ρ y_{t} + σ_{x} ϵ_{t + 1}, t \geq 0

where the scalars $ρ$ and $σ_{x}$ satisfy $| ρ | < 1$ and $σ_{x} > 0$ ; ${ϵ_{t + 1}}$ is a sequence of i.i.d. normal random variables with mean $0$ and variance $1$ .

The second component of the statistical model is

(47.2)#

y_{0} \sim N (μ_{0}, σ_{0}^{2})

Consider a sample ${y_{t}}_{t = 0}^{T}$ governed by this statistical model.

The model implies that the likelihood function of ${y_{t}}_{t = 0}^{T}$ can be factored:

f (y_{T}, y_{T - 1}, \dots, y_{0}) = f (y_{T} | y_{T - 1}) f (y_{T - 1} | y_{T - 2}) \dots f (y_{1} | y_{0}) f (y_{0})

where we use $f$ to denote a generic probability density.

The statistical model (47.1)-(47.2) implies

\begin{array}{r} \begin{aligned} f (y_{t} | y_{t - 1}) & \sim N (ρ y_{t - 1}, σ_{x}^{2}) \\ f (y_{0}) & \sim N (μ_{0}, σ_{0}^{2}) \end{aligned} \end{array}

We want to study how inferences about the unknown parameters $(ρ, σ_{x})$ depend on what is assumed about the parameters $μ_{0}, σ_{0}$ of the distribution of $y_{0}$ .

Below, we study two widely used alternative assumptions:

$(μ_{0}, σ_{0}) = (y_{0}, 0)$ which means that $y_{0}$ is drawn from the distribution $N (y_{0}, 0)$ ; in effect, we are conditioning on an observed initial value.
$μ_{0}, σ_{0}$ are functions of $ρ, σ_{x}$ because $y_{0}$ is drawn from the stationary distribution implied by $ρ, σ_{x}$ .

Note: We do not treat a third possible case in which $μ_{0}, σ_{0}$ are free parameters to be estimated.

Unknown parameters are $ρ, σ_{x}$ .

We have independent prior probability distributions for $ρ, σ_{x}$ and want to compute a posterior probability distribution after observing a sample ${y_{t}}_{t = 0}^{T}$ .

The notebook uses pymc4 and numpyro to compute a posterior distribution of $ρ, σ_{x}$ . We will use NUTS samplers to generate samples from the posterior in a chain. Both of these libraries support NUTS samplers.

NUTS is a form of Monte Carlo Markov Chain (MCMC) algorithm that bypasses random walk behaviour and allows for convergence to a target distribution more quickly. This not only has the advantage of speed, but allows for complex models to be fitted without having to employ specialised knowledge regarding the theory underlying those fitting methods.

Thus, we explore consequences of making these alternative assumptions about the distribution of $y_{0}$ :

A first procedure is to condition on whatever value of $y_{0}$ is observed. This amounts to assuming that the probability distribution of the random variable $y_{0}$ is a Dirac delta function that puts probability one on the observed value of $y_{0}$ .
A second procedure assumes that $y_{0}$ is drawn from the stationary distribution of a process described by (47.1) so that $y_{0} \sim N (0, \frac{σ_{x}^{2}}{(1 - ρ)^{2}})$

When the initial value $y_{0}$ is far out in a tail of the stationary distribution, conditioning on an initial value gives a posterior that is more accurate in a sense that we’ll explain.

Basically, when $y_{0}$ happens to be in a tail of the stationary distribution and we don’t condition on $y_{0}$ , the likelihood function for ${y_{t}}_{t = 0}^{T}$ adjusts the posterior distribution of the parameter pair $ρ, σ_{x}$ to make the observed value of $y_{0}$ more likely than it really is under the stationary distribution, thereby adversely twisting the posterior in short samples.

An example below shows how not conditioning on $y_{0}$ adversely shifts the posterior probability distribution of $ρ$ toward larger values.

We begin by solving a direct problem that simulates an AR(1) process.

How we select the initial value $y_{0}$ matters.

If we think $y_{0}$ is drawn from the stationary distribution $N (0, \frac{σ_{x}^{2}}{1 - ρ^{2}})$ , then it is a good idea to use this distribution as $f (y_{0})$ . Why? Because $y_{0}$ contains information about $ρ, σ_{x}$ .
If we suspect that $y_{0}$ is far in the tails of the stationary distribution – so that variation in early observations in the sample have a significant transient component – it is better to condition on $y_{0}$ by setting $f (y_{0}) = 1$ .

To illustrate the issue, we’ll begin by choosing an initial $y_{0}$ that is far out in a tail of the stationary distribution.

def ar1_simulate(rho, sigma, y0, T):

    # Allocate space and draw epsilons
    y = np.empty(T)
    eps = np.random.normal(0.,sigma,T)

    # Initial condition and step forward
    y[0] = y0
    for t in range(1, T):
        y[t] = rho*y[t-1] + eps[t]

    return y

sigma =  1.
rho = 0.5
T = 50

np.random.seed(145353452)
y = ar1_simulate(rho, sigma, 10, T)

plt.plot(y)
plt.tight_layout()

_images/86fce48cb891af782332a8afe2522140e433e395d6cbb2723aa93a52c5b92105.png

Now we shall use Bayes’ law to construct a posterior distribution, conditioning on the initial value of $y_{0}$ .

(Later we’ll assume that $y_{0}$ is drawn from the stationary distribution, but not now.)

First we’ll use pymc4.

47.1. PyMC Implementation#

For a normal distribution in pymc, $v a r = 1 / τ = σ^{2}$ .

AR1_model = pmc.Model()

with AR1_model:

    # Start with priors
    rho = pmc.Uniform('rho', lower=-1., upper=1.) # Assume stable rho
    sigma = pmc.HalfNormal('sigma', sigma = np.sqrt(10))

    # Expected value of y at the next period (rho * y)
    yhat = rho * y[:-1]

    # Likelihood of the actual realization
    y_like = pmc.Normal('y_obs', mu=yhat, sigma=sigma, observed=y[1:])

pmc.sample by default uses the NUTS samplers to generate samples as shown in the below cell:

with AR1_model:
    trace = pmc.sample(50000, tune=10000, return_inferencedata=True)

                                                                                                                   
  Progress                    Draws   Divergences   Step size   Grad evals   Sampling Speed   Elapsed   Remaining  
 ───────────────────────────────────────────────────────────────────────────────────────────────────────────────── 
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   0             1.15        3            759.50 draws/s   0:01:18   0:00:00    
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   0             1.39        3            744.23 draws/s   0:01:20   0:00:00    
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   0             1.11        3            754.13 draws/s   0:01:19   0:00:00    
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   0             0.89        3            746.14 draws/s   0:01:20   0:00:00

with AR1_model:
    az.plot_trace(trace, figsize=(17,6))

_images/cba5b2b63c97c52ac2ce08fc8135dab597a96ff8ced13f64b59304b1642dca4f.png

Evidently, the posteriors aren’t centered on the true values of $.5, 1$ that we used to generate the data.

This is a symptom of the classic Hurwicz bias for first order autoregressive processes (see Leonid Hurwicz [Hurwicz, 1950].)

The Hurwicz bias is worse the smaller is the sample (see [Orcutt and Winokur, 1969]).

Be that as it may, here is more information about the posterior.

with AR1_model:
    summary = az.summary(trace, round_to=4)

summary

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
rho	0.5363	0.0712	0.4025	0.6707	0.0002	0.0002	164211.1076	124922.4228	1.0001
sigma	1.0103	0.1062	0.8243	1.2166	0.0003	0.0003	168094.1812	135012.0580	1.0000

Now we shall compute a posterior distribution after seeing the same data but instead assuming that $y_{0}$ is drawn from the stationary distribution.

This means that

y_{0} \sim N (0, \frac{σ_{x}^{2}}{1 - ρ^{2}})

We alter the code as follows:

AR1_model_y0 = pmc.Model()

with AR1_model_y0:

    # Start with priors
    rho = pmc.Uniform('rho', lower=-1., upper=1.) # Assume stable rho
    sigma = pmc.HalfNormal('sigma', sigma=np.sqrt(10))

    # Standard deviation of ergodic y
    y_sd = sigma / np.sqrt(1 - rho**2)

    # yhat
    yhat = rho * y[:-1]
    y_data = pmc.Normal('y_obs', mu=yhat, sigma=sigma, observed=y[1:])
    y0_data = pmc.Normal('y0_obs', mu=0., sigma=y_sd, observed=y[0])

with AR1_model_y0:
    trace_y0 = pmc.sample(50000, tune=10000, return_inferencedata=True)

# Grey vertical lines are the cases of divergence

                                                                                                                   
  Progress                    Draws   Divergences   Step size   Grad evals   Sampling Speed   Elapsed   Remaining  
 ───────────────────────────────────────────────────────────────────────────────────────────────────────────────── 
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   0             0.87        3            720.66 draws/s   0:01:23   0:00:00    
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   30            0.58        3            678.53 draws/s   0:01:28   0:00:00    
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   2             1.03        7            667.60 draws/s   0:01:29   0:00:00    
  ━━━━━━━━━━━━━━━━━━━━━━━━━   60000   9             0.94        3            674.17 draws/s   0:01:28   0:00:00

with AR1_model_y0:
    az.plot_trace(trace_y0, figsize=(17,6))

_images/96d06c12c19b9435f6dbbda2e28f1eec86eadcb29fb0a7cf53c3bce0e8944585.png

with AR1_model:
    summary_y0 = az.summary(trace_y0, round_to=4)

summary_y0

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
rho	0.8759	0.0815	0.7321	0.9990	0.0002	0.0002	105808.9127	78800.5559	1.0
sigma	1.4050	0.1478	1.1415	1.6859	0.0004	0.0004	119533.6471	103442.9811	1.0

Please note how the posterior for $ρ$ has shifted to the right relative to when we conditioned on $y_{0}$ instead of assuming that $y_{0}$ is drawn from the stationary distribution.

Think about why this happens.

Hint

It is connected to how Bayes Law (conditional probability) solves an inverse problem by putting high probability on parameter values that make observations more likely.

We’ll return to this issue after we use numpyro to compute posteriors under our two alternative assumptions about the distribution of $y_{0}$ .

We’ll now repeat the calculations using numpyro.

47.2. Numpyro Implementation#

def plot_posterior(sample):
    """
    Plot trace and histogram
    """
    # To np array
    rhos = sample['rho']
    sigmas = sample['sigma']
    rhos, sigmas, = np.array(rhos), np.array(sigmas)

    fig, axs = plt.subplots(2, 2, figsize=(17, 6))
    # Plot trace
    axs[0, 0].plot(rhos)   # rho
    axs[1, 0].plot(sigmas) # sigma

    # Plot posterior
    axs[0, 1].hist(rhos, bins=50, density=True, alpha=0.7)
    axs[0, 1].set_xlim([0, 1])
    axs[1, 1].hist(sigmas, bins=50, density=True, alpha=0.7)

    axs[0, 0].set_title("rho")
    axs[0, 1].set_title("rho")
    axs[1, 0].set_title("sigma")
    axs[1, 1].set_title("sigma")
    plt.show()

def AR1_model(data):
    # set prior
    rho = numpyro.sample('rho', dist.Uniform(low=-1., high=1.))
    sigma = numpyro.sample('sigma', dist.HalfNormal(scale=np.sqrt(10)))

    # Expected value of y at the next period (rho * y)
    yhat = rho * data[:-1]

    # Likelihood of the actual realization.
    y_data = numpyro.sample('y_obs', dist.Normal(loc=yhat, scale=sigma), obs=data[1:])

# Make jnp array
y = jnp.array(y)

# Set NUTS kernal
NUTS_kernel = numpyro.infer.NUTS(AR1_model)

# Run MCMC
mcmc = numpyro.infer.MCMC(NUTS_kernel, num_samples=50000, num_warmup=10000, progress_bar=False)
mcmc.run(rng_key=random.PRNGKey(1), data=y)

plot_posterior(mcmc.get_samples())

_images/86fa4a45d484c2cf869e928255530a9bddabc0c9c1b9260b69d1f4c04d01a185.png

mcmc.print_summary()

                mean       std    median      5.0%     95.0%     n_eff     r_hat
       rho      0.54      0.07      0.54      0.43      0.66  33977.54      1.00
     sigma      1.01      0.11      1.00      0.84      1.18  32265.69      1.00

Number of divergences: 0

Next, we again compute the posterior under the assumption that $y_{0}$ is drawn from the stationary distribution, so that

y_{0} \sim N (0, \frac{σ_{x}^{2}}{1 - ρ^{2}})

Here’s the new code to achieve this.

def AR1_model_y0(data):
    # Set prior
    rho = numpyro.sample('rho', dist.Uniform(low=-1., high=1.))
    sigma = numpyro.sample('sigma', dist.HalfNormal(scale=np.sqrt(10)))

    # Standard deviation of ergodic y
    y_sd = sigma / jnp.sqrt(1 - rho**2)

    # Expected value of y at the next period (rho * y)
    yhat = rho * data[:-1]

    # Likelihood of the actual realization.
    y_data = numpyro.sample('y_obs', dist.Normal(loc=yhat, scale=sigma), obs=data[1:])
    y0_data = numpyro.sample('y0_obs', dist.Normal(loc=0., scale=y_sd), obs=data[0])

# Make jnp array
y = jnp.array(y)

# Set NUTS kernal
NUTS_kernel = numpyro.infer.NUTS(AR1_model_y0)

# Run MCMC
mcmc2 = numpyro.infer.MCMC(NUTS_kernel, num_samples=50000, num_warmup=10000, progress_bar=False)
mcmc2.run(rng_key=random.PRNGKey(1), data=y)

plot_posterior(mcmc2.get_samples())

_images/a700f9d3ad3e123b6b1bc90499c665e50a64b5e690883df1bd03f44dcbad8cda.png

mcmc2.print_summary()

                mean       std    median      5.0%     95.0%     n_eff     r_hat
       rho      0.88      0.08      0.89      0.76      1.00  29481.95      1.00
     sigma      1.40      0.15      1.39      1.17      1.64  24707.84      1.00

Number of divergences: 0

Look what happened to the posterior!

It has moved far from the true values of the parameters used to generate the data because of how Bayes’ Law (i.e., conditional probability) is telling numpyro to explain what it interprets as “explosive” observations early in the sample.

Bayes’ Law is able to generate a plausible likelihood for the first observation by driving $ρ \to 1$ and $σ ↑$ in order to raise the variance of the stationary distribution.

Our example illustrates the importance of what you assume about the distribution of initial conditions.