# Copyright Contributors to the Pyro project.

# SPDX-License-Identifier: Apache-2.0

r"""

Example: VAR(2) process

=======================

In this example, we demonstrate how to implement and perform Bayesian inference for a

Vector Autoregressive process of order 2 (VAR(2)). VAR models are widely used in

time series analysis, especially for capturing the dynamics between multiple variables.

A VAR(2) process for a multivariate time series :math:`y_t` with :math:`K` variables is defined as:

.. math::

y_t = c + \Phi_1 y_{t-1} + \Phi_2 y_{t-2} + \epsilon_t

Here, :math:`c` is a constant vector, :math:`\Phi_1` and :math:`\Phi_2` are coefficient matrices for lag 1

and lag 2, respectively, and :math:`\epsilon_t` is a Gaussian noise term with zero mean and a

covariance matrix :math:`\Sigma`.

This example uses NumPyro's `scan` utility to efficiently model the temporal dependencies without

explicit Python loops.

For more general time series forecasting techniques and examples, refer to the

`Time Series Forecasting` tutorial:

https://num.pyro.ai/en/stable/tutorials/time_series_forecasting.html#Forecasting

Reference

---------

For more information on Vector Autoregressive models, see:

https://otexts.com/fpp2/VAR.html

.. image:: ../_static/img/examples/var2.png

:align: center

"""

import argparse

import os

import time

import matplotlib.pyplot as plt

import numpy as np

from jax import random

import jax.numpy as jnp

import numpyro

from numpyro.contrib.control_flow import scan

import numpyro.distributions as dist

def var2_scan(y):

T, K = y.shape # Number of time steps and number of variables

# Priors for constants and coefficients

c = numpyro.sample("c", dist.Normal(0, 1).expand([K])) # Constants vector of size K

Phi1 = numpyro.sample(

"Phi1", dist.Normal(0, 1).expand([K, K]).to_event(2)

) # Coefficients for lag 1

Phi2 = numpyro.sample(

"Phi2", dist.Normal(0, 1).expand([K, K]).to_event(2)

) # Coefficients for lag 2

# Priors for error terms

sigma = numpyro.sample("sigma", dist.HalfNormal(1.0).expand([K]).to_event(1))

L_omega = numpyro.sample(

"L_omega", dist.LKJCholesky(dimension=K, concentration=1.0)

)

L_Sigma = (

sigma[..., None] * L_omega

) # Alternative: jnp.einsum("...i,...ij->...ij", sigma, L_omega)

def transition(carry, t):

y_prev1, y_prev2, y_obs = carry # Previous two observations and observed data

m_t = c + jnp.dot(Phi1, y_prev1) + jnp.dot(Phi2, y_prev2) # Mean prediction

# Conditioned on observed y

y_t = numpyro.sample(

f"y_{t}",

dist.MultivariateNormal(loc=m_t, scale_tril=L_Sigma),

obs=y_obs[t],

)

new_carry = (y_t, y_prev1, y_obs)

return new_carry, m_t

# Initial carry: observations at time steps 1 and 0

init_carry = (y[1], y[0], y[2:])

# Time indices starting from time step 2

time_indices = jnp.arange(T - 2)

# Run the scan

_, mu = scan(transition, init_carry, time_indices)

# Store the mean trajectory as a deterministic variable

numpyro.deterministic("mu", mu)

def generate_var2_data(T, K, c, Phi1, Phi2, sigma):

"""

Generate time series data from a VAR(2) process.

Args:

T (int): Number of time steps.

K (int): Number of variables in the time series.

c (array): Constants (shape: (K,)).

Phi1 (array): Coefficients for lag 1 (shape: (K, K)).

Phi2 (array): Coefficients for lag 2 (shape: (K, K)).

sigma (array): Covariance matrix for the noise (shape: (K, K)).

Returns:

np.ndarray: Generated time series data (shape: (T, K)).

"""

# Initialize time series with random values

y = np.zeros((T, K))

y[:2] = np.random.multivariate_normal(mean=np.zeros(K), cov=sigma, size=2)

# Generate the time series

for t in range(2, T):

y[t] = (

c

+ Phi1 @ y[t - 1]

+ Phi2 @ y[t - 2]

+ np.random.multivariate_normal(mean=np.zeros(K), cov=sigma)

)

return y

def run_inference(model, args, rng_key, y):

"""

Run MCMC inference for the given model.

Args:

model: The probabilistic model to infer.

args: Command-line arguments.

rng_key: PRNG key for randomness.

y: Observed time series data.

"""

start = time.time()

sampler = numpyro.infer.NUTS(model)

mcmc = numpyro.infer.MCMC(

sampler,

num_warmup=args.num_warmup,

num_samples=args.num_samples,

num_chains=args.num_chains,

progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,

)

mcmc.run(rng_key, y=y)

mcmc.print_summary()

print("\nMCMC elapsed time:", time.time() - start)

return mcmc.get_samples()

def main(args):

# Generate artificial dataset

T = args.num_data # Number of time steps

K = 2 # Number of variables

c_true = jnp.array([0.5, -0.3]) # Constants

Phi1_true = jnp.array([[0.7, 0.1], [0.2, 0.5]]) # Coefficients for lag 1

Phi2_true = jnp.array([[0.2, -0.1], [-0.1, 0.2]]) # Coefficients for lag 2

sigma_true = jnp.array([[0.1, 0.02], [0.02, 0.1]]) # Covariance matrix

rng_key = random.key(0)

y = generate_var2_data(T, K, c_true, Phi1_true, Phi2_true, sigma_true)

# Perform inference

samples = run_inference(var2_scan, args, rng_key, y)

# Prediction

mean_prediction = samples["mu"].mean(axis=0)

lower_bound = jnp.percentile(samples["mu"], 2.5, axis=0) # 2.5th percentile

upper_bound = jnp.percentile(samples["mu"], 97.5, axis=0) # 97.5th percentile

# Plot results

fig, axes = plt.subplots(K, 1, figsize=(10, 6), sharex=True)

time_steps = jnp.arange(T)

for i in range(K):

# True values

axes[i].plot(time_steps, y[:, i], label=f"True Variable {i + 1}", color="blue")

# Posterior mean prediction

axes[i].plot(

time_steps[2:],

mean_prediction[:, i],

label=f"Predicted Mean Variable {i + 1}",

color="orange",

)

# 95% confidence interval

axes[i].fill_between(

time_steps[2:],

lower_bound[:, i],

upper_bound[:, i],

color="orange",

alpha=0.2,

label="95% CI",

)

axes[i].set_title(f"Variable {i + 1}")

axes[i].legend()

axes[i].grid(True)

plt.xlabel("Time Steps")

plt.tight_layout()

plt.savefig("var2.png")

if __name__ == "__main__":

assert numpyro.__version__.startswith("0.21.0")

parser = argparse.ArgumentParser(description="VAR(2) example")

parser.add_argument("--num-data", nargs="?", default=100, type=int)

parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)

parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)

parser.add_argument("--num-chains", nargs="?", default=1, type=int)

parser.add_argument("--device", default="cpu", type=str, help='use "cpu" or "gpu".')

args = parser.parse_args()

numpyro.set_platform(args.device)

numpyro.set_host_device_count(args.num_chains)

main(args)