GARCH models in Numpyro

I was curious how best to implement GARCH models in numpyro. I tried reading https://num.pyro.ai/en/stable/tutorials/time_series_forecasting.html but found it generally unclear (the model notation and variable names are not easy to map, the model is overly complicated for a basic intro)

I wrote the following code to estimate a GARCH-M model, but the performance doesn't seem great and I'm curious to see if others have suggestions.

The model is

but the question applies quite generally. I chose a GARCH-M model because it requires looping, you can't infer the entire sequence of epsilon's at once because you need to infer time-varying volatility.

# preliminaries
import numpyro
import numpyro.distributions    as dist
import jax.numpy as jnp
from numpyro.infer import MCMC, NUTS
from numpyro.contrib.control_flow import scan
import numpy                    as np

def my_model(y=None):

    μ = numpyro.sample("μ", dist.Normal(0, 4))
    ω = numpyro.sample("ω", dist.HalfCauchy(2))
    α = numpyro.sample("α", dist.Uniform())
    β = numpyro.sample("β", dist.Uniform())
    λ = numpyro.sample("λ", dist.Normal(0, 4))

    σ2_0 = ω / (1 - α - β)

    def gjr_var(state, new):
        # state is past variance and past shock
        σ2_t, r_t = state

        ɛ_t = (r_t - μ - λ*σ2_t**0.5)/σ2_t**0.5

        σ2_tp = ω + α * ɛ_t**2 + β * σ2_t 
        r_tp = numpyro.sample('r', dist.Normal(μ + λ*σ2_tp**0.5, σ2_tp**0.5), obs=new)

        return (σ2_tp, r_tp), r_tp

    _, r = scan(gjr_var, (σ2_0, y[0]), y[1:])

    return r
    

I re-visited this after several months and came up with some tweaks that resolved the problem.

The main problem seems to be that if alpha + beta >= 1, then the volatility process is non-stationary. Worse yet, it leads to an negative unconditional variance due to the formula σ2_0 = ω / (1 - α - β)

I imposed stationary by drawing alpha and beta from a 3d Dirichlet distribution, and now the estimates match what is produced by the arch Python package

def my_model(y: np.ndarray):

    μ = numpyro.sample("μ", dist.ImproperUniform(dist.constraints.real, (), ()))
    ω = numpyro.sample("ω", dist.ImproperUniform(dist.constraints.positive, (), ()))
    β_vec = numpyro.sample("β_intermediate", dist.Dirichlet(jnp.ones(3)))
    α = numpyro.deterministic("α", β_vec[0])
    β = numpyro.deterministic("β", β_vec[1])

    λ = numpyro.sample("λ", dist.ImproperUniform(dist.constraints.real, (), ()))

    σ2_0 = ω / (1 - α - β)

    def gjr_var(state, new):
        # state is past variance and past shock
        σ2_t, r_t = state

        ɛ_t = r_t - μ - λ * σ2_t**0.5

        σ2_tp = numpyro.deterministic(
            "σ2", ω + α * ɛ_t**2 + β * σ2_t
        )

        r_tp = numpyro.sample(
            "r", dist.Normal(μ + λ * σ2_tp**0.5, σ2_tp**0.5), obs=new
        )

        ɛ_tp = numpyro.deterministic("ɛ", r_tp - μ - λ * σ2_tp**0.5)

        return (σ2_tp, r_tp), (r_tp, σ2_tp, ɛ_tp)

    scan(gjr_var, (σ2_0, y[0]), y[1:])