Likelihood function in the scan function (AR2 model)

[Yamaguchi]

Thanks for your reply.

I actually got it to work by changing the code as follows.
However, I have two questions I am unclear about.

1. is it sufficient to apply the collect_default_updates function only to random numbers created inside scan? In other words, do I need to include random numbers created outside of scan (alpha_1, alpha_2, etc.)? 2.
2. the results are strange. The uncertainty seems to increase as time progresses. Is this a mistake in the code?

import os
import random
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import arviz as az
from copy import copy
from scipy import stats

import pymc as pm
import pytensor
import pytensor.tensor as pt
from pymc.pytensorf import collect_default_updates

# Initialize random number generator
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

# create data
num_data = 142
t = np.arange(0, num_data)
y = np.sin(t) + np.random.normal(size=num_data) * 0.1

def ar2_dist(y_prev, y_prev_prev, alpha_1, alpha_2, const, sigma, n_steps, size):
    
    def transition(y_prev, y_prev_prev, alpha_1, alpha_2, const, sigma):
        m_t = const + alpha_1 * y_prev + alpha_2 * y_prev_prev
        y_t = pm.Normal.dist(mu=m_t, sigma=sigma)
        return (y_t, y_prev), collect_default_updates([y_t, y_prev, alpha_1, alpha_2, const, sigma]) 
        #return (y_t, y_prev), collect_default_updates([y_t, y_prev])
    
    # Pytensor scan is a looping function
    result, updates = pytensor.scan(
        fn=transition,  # function
        outputs_info=[y_prev, y_prev_prev],  # initial conditions
        non_sequences=[alpha_1, alpha_2, const, sigma],  # parameters
        n_steps=n_steps, # number of loops
        strict=True,
    ) 
    
    return result[0] 

# ar2 model
with pm.Model() as ar2_model:
    alpha_1 = pm.Normal("alpha_1", mu=0.0, sigma=1.0)
    alpha_2 = pm.Normal("alpha_2", mu=0.0, sigma=1.0)
    const = pm.Normal("const", mu=0.0, sigma=1.0)
    sigma = pm.HalfNormal("sigma", sigma=1.0)

    #y_pt = pt.as_tensor_variable(y)

    n_steps = num_data - 2
    y_prev, y_prev_prev = y[1], y[0]

    ar_innov = pm.CustomDist(
        "ar2_dist",
        y_prev,
        y_prev_prev,
        alpha_1,
        alpha_2,
        const,
        sigma,
        n_steps,
        dist=ar2_dist,
        observed=y[2:],
        shape=(n_steps,)
    )
    
    ar = pm.Deterministic("ar", pt.concatenate([[y_prev_prev, y_prev], ar_innov], axis=-1))
    
with ar2_model:
    trace_scan = pm.sample(tune=500, draws=1000)

az.plot_trace(trace_scan);

with ar2_model:
    post_pred = pm.sample_posterior_predictive(
        trace_scan, 
        random_seed=rng, 
    )
    
post_pred_ar = post_pred.posterior_predictive["ar"]
for hdi_prob in (0.94, 0.64):
    hdi = az.hdi(post_pred_ar, hdi_prob=hdi_prob)["ar"]
    lower = hdi.sel(hdi="lower")
    upper = hdi.sel(hdi="higher")
    plt.fill_between(np.arange(num_data), y1=lower, y2=upper, alpha=.2, color="C0")
plt.plot(post_pred_ar.mean(("chain", "draw")), color="C0")    
plt.plot(y, color="k")
plt.show()