Likelihood function in the scan function (AR2 model)

1

Hi all,

I want to implement a state space model.
However, it does not seem to work correctly as shown in the following code.

I would be thankful if you could help me with this question.

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


# 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

# 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)
    
    def transition(y_obs, 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 = m_t
        y_t = pm.Normal("y_t", mu=m_t, sigma=sigma, observed=y_obs)
        return y_t, y_prev

    y_pt = pt.as_tensor_variable(y)

    n_steps = num_data - 2
    inits = [y_pt[1], y_pt[0]]

    # Pytensor scan is a looping function
    result, updates = pytensor.scan(
        fn=transition,  # function
        outputs_info=inits,  # initial conditions
        sequences=y_pt[2:],
        non_sequences=[alpha_1, alpha_2, const, sigma],  # parameters
        n_steps=n_steps, # number of loops
    ) 
    
with ar2_model:
    trace_scan = pm.sample(tune=500, draws=500)

error

MissingInputError: Input 1 () of the graph (indices start from 0), used to compute Mul(alpha_2, ), was not provided and not given a value. Use the PyTensor flag exception_verbosity='high', for more information on this error.
 
Backtrace when that variable is created:

  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/ipykernel/zmqshell.py", line 549, in run_cell
    return super().run_cell(*args, **kwargs)
  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3048, in run_cell
    result = self._run_cell(
  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3103, in _run_cell
    result = runner(coro)
  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
    coro.send(None)
  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3308, in run_cell_async
    has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3490, in run_ast_nodes
    if await self.run_code(code, result, async_=asy):
  File "/Users/yyamaguchi/Desktop/programming/05_PyMC/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py", line 3550, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/var/folders/57/xr8vd64j2_qd3ckpnwhjpw_h0000gn/T/ipykernel_4769/3370695671.py", line 23, in 
    result, updates = pytensor.scan(
2

Your model is incomplete. You should wrap the transition and scan into a single helper function and pass it to a pm.CustomDist. Here’s an example.

You could also give the SARIMAX model in pymc-experimental a try. It’s under active development, so if you do and you hit any problems, don’t hesitate to open an issue/ask here.

3

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()

スクリーンショット 2024-01-12 20.56.41
スクリーンショット 2024-01-12 20.56.411696×1156 119 KB

4

No, just the ones inside the scan. The job of collect_deafult_updates is to explain to pytensor how to update random generators inside the scan loop. Since the non_sequence variables are sampled only once, outside the scan, they don’t need to be updated inside the loop.

It looks like the posterior of the AR parameters includes combinations of parameters that induce non-stationary dynamics into the system. You can see this notebook in the Samuelson Multiplier-Accelerator section for some discussion specifically in the context of an AR(2). You can see that the posterior mean gives damped oscillations, but there are explosive oscillations in the tails.

Unsolicited opinion: Your data don’t look like an AR(2) at all. They exhibit strong seasonal behavior, so you should model that first, then consider AR dynamics if it seems like the residuals are still auto-correlated. The model is learning parameter combinations with complex eigenvalues to try to capture the seasonal pattern.

2 Likes
5

No, just the ones inside the scan.

I Understood. Thank you very much.

You can see this notebook in the Samuelson Multiplier-Accelerator section for some discussion specifically in the context of an AR(2). You can see that the posterior mean gives damped oscillations, but there are explosive oscillations in the tails.
Unsolicited opinion: Your data don’t look like an AR(2) at all. They exhibit strong seasonal behavior, so you should model that first, then consider AR dynamics if it seems like the residuals are still auto-correlated. The model is learning parameter combinations with complex eigenvalues to try to capture the seasonal pattern.

Thanks for the advice.
I used NumPyro’s sample as a reference. I will check with other datasets.

https://num.pyro.ai/en/latest/examples/ar2.html

6

The other difference comes because by default PyMC gives you the unconditional posterior trajectory (that is, sequences of length T starting from t=0 simulated by model dynamics), while the numpyro example is showing the conditional posterior trajectory (that is, a sequences of one-step-ahead predictions, each one starting from datapoint t-1 and generating t from the AR process).

The SARIMAX notebook I linked goes into the difference (and illustrates it) more.

7

Let me ask one more question.
What should I do when I have several different likelihoods, for example?

        y_t = pm.Normal.dist(mu=m_t, sigma=sigma)
        y_p = pm.Poisson(mu=m_t)

Do I just give a (num_data, 2)-shaped matrix with normal distribution y and Poisson distribution y as the observed argument?

   ar_innov = pm.CustomDist(
        name,
        params,
        observed=y_obs,
        shape=(n_steps,)
    )
8

You want to model something like a discrete jump diffusion process? A statespace where one latent variable has poisson innovations and the other is normal?

9

I considered the case of multiple observations.

10

If they aren’t jointly distributed you can just make two entirely separate CustomDists. If they are jointly distributed, your ar2_dist needs to return a vector of length 2, and you pass in all the parameters each distribution requires.

2 Likes
11

I appreciate your correspondence. It was very helpful.

2 Likes