(forecasting_with_ar)=

Forecasting with Structural AR Timeseries¶

:::{post} Oct 20, 2022 :tags: forecasting, autoregressive, bayesian structural timeseries :category: intermediate :author: Nathaniel Forde :::

Bayesian structural timeseries models are an interesting way to learn about the structure inherent in any observed timeseries data. It also gives us the ability to project forward the implied predictive distribution granting us another view on forecasting problems. We can treat the learned characteristics of the timeseries data observed to-date as informative about the structure of the unrealised future state of the same measure.

In this notebook we'll see how to fit and predict a range of auto-regressive structural timeseries models and, importantly, how to predict future observations of the models.

In [ ]:
import arviz as az
import numpy as np
import pandas as pd
import pymc as pm

from matplotlib import pyplot as plt
In [ ]:
RANDOM_SEED = 8929
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

Generate Fake Autoregressive Data¶

First we will generate a simple autoregressive timeseries. We will show how to specify a model to fit this data and then add a number of complexities to the data and show how they too can be captured with an autoregressive model and used to predict the shape of the future.

In [ ]:
def simulate_ar(intercept, coef1, coef2, noise=0.3, *, warmup=10, steps=200):
    # We sample some extra warmup steps, to let the AR process stabilize
    draws = np.zeros(warmup + steps)
    # Initialize first draws at intercept
    draws[:2] = intercept
    for step in range(2, warmup + steps):
        draws[step] = (
            intercept
            + coef1 * draws[step - 1]
            + coef2 * draws[step - 2]
            + np.random.normal(0, noise)
        )
    # Discard the warmup draws
    return draws[warmup:]


# True parameters of the AR process
ar1_data = simulate_ar(10, -0.9, 0)

fig, ax = plt.subplots(figsize=(10, 3))
ax.set_title("Generated Autoregressive Timeseries", fontsize=15)
ax.plot(ar1_data);

Specifying the Model¶

We'll walk through the model step by step and then generalise the pattern into a function that can be used to take increasingly complex structural combinations of components.

In [ ]:
## Set up a dictionary for the specification of our priors
## We set up the dictionary to specify size of the AR coefficients in
## case we want to vary the AR lags.
priors = {
    "coefs": {"mu": [10, 0.2], "sigma": [0.1, 0.1], "size": 2},
    "sigma": 8,
    "init": {"mu": 9, "sigma": 0.1, "size": 1},
}

## Initialise the model
with pm.Model() as AR:
    pass

## Define the time interval for fitting the data
t_data = list(range(len(ar1_data)))
## Add the time interval as a mutable coordinate to the model to allow for future predictions
AR.add_coord("obs_id", t_data, mutable=True)

with AR:
    ## Data containers to enable prediction
    t = pm.MutableData("t", t_data, dims="obs_id")
    y = pm.MutableData("y", ar1_data, dims="obs_id")

    # The first coefficient will be the constant term but we need to set priors for each coefficient in the AR process
    coefs = pm.Normal("coefs", priors["coefs"]["mu"], priors["coefs"]["sigma"])
    sigma = pm.HalfNormal("sigma", priors["sigma"])
    # We need one init variable for each lag, hence size is variable too
    init = pm.Normal.dist(
        priors["init"]["mu"], priors["init"]["sigma"], size=priors["init"]["size"]
    )
    # Steps of the AR model minus the lags required
    ar1 = pm.AR(
        "ar",
        coefs,
        sigma=sigma,
        init_dist=init,
        constant=True,
        steps=t.shape[0] - (priors["coefs"]["size"] - 1),
        dims="obs_id",
    )

    # The Likelihood
    outcome = pm.Normal("likelihood", mu=ar1, sigma=sigma, observed=y, dims="obs_id")
    ## Sampling
    idata_ar = pm.sample_prior_predictive()
    idata_ar.extend(pm.sample(2000, random_seed=100, target_accept=0.95))
    idata_ar.extend(pm.sample_posterior_predictive(idata_ar))

Sampling: [ar, coefs, likelihood, sigma]
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [coefs, sigma, ar]
100.00% [12000/12000 00:26<00:00 Sampling 4 chains, 0 divergences] 
Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 46 seconds.
Sampling: [likelihood]
100.00% [8000/8000 00:00<00:00]
In [ ]:
idata_ar
Out[ ]:
arviz.InferenceData
    • <xarray.Dataset>
Dimensions:      (chain: 4, draw: 2000, coefs_dim_0: 2, obs_id: 200)
Coordinates:
  * chain        (chain) int64 0 1 2 3
  * draw         (draw) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
  * coefs_dim_0  (coefs_dim_0) int64 0 1
  * obs_id       (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
Data variables:
    coefs        (chain, draw, coefs_dim_0) float64 9.784 -0.8044 ... -0.7937
    ar           (chain, draw, obs_id) float64 8.744 5.527 4.382 ... 5.203 5.425
    sigma        (chain, draw) float64 0.4881 0.5189 0.4896 ... 0.4919 0.536
Attributes:
    created_at:                 2022-11-09T13:37:35.440374
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
    sampling_time:              46.495951890945435
    tuning_steps:               1000
      xarray.Dataset
        • chain: 4
        • draw: 2000
        • coefs_dim_0: 2
        • obs_id: 200
        • chain
          (chain)
          int64
          0 1 2 3
          array([0, 1, 2, 3])
        • draw
          (draw)
          int64
          0 1 2 3 4 ... 1996 1997 1998 1999
          array([   0,    1,    2, ..., 1997, 1998, 1999])
        • coefs_dim_0
          (coefs_dim_0)
          int64
          0 1
          array([0, 1])
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • coefs
          (chain, draw, coefs_dim_0)
          float64
          9.784 -0.8044 ... 9.728 -0.7937
          array([[[ 9.78424718, -0.80438722],
        [ 9.74703572, -0.83706764],
        [ 9.79944405, -0.79930181],
        ...,
        [ 9.53893097, -0.79144711],
        [ 9.59709183, -0.78652045],
        [ 9.73692629, -0.81651835]],

       [[ 9.74268028, -0.80803045],
        [ 9.79526179, -0.81136089],
        [ 9.7890239 , -0.82699725],
        ...,
        [ 9.86478399, -0.81825659],
        [ 9.56231758, -0.7842574 ],
        [ 9.91745608, -0.85332664]],

       [[ 9.86287461, -0.80656223],
        [ 9.7503848 , -0.78838031],
        [ 9.73456476, -0.82507864],
        ...,
        [ 9.67317157, -0.79964398],
        [ 9.72808985, -0.81919129],
        [ 9.79594168, -0.8316699 ]],

       [[ 9.66388511, -0.77070848],
        [ 9.81792066, -0.86152692],
        [ 9.76193626, -0.80400692],
        ...,
        [ 9.82540452, -0.82762375],
        [ 9.68635269, -0.81628275],
        [ 9.72849693, -0.79365451]]])
        • ar
          (chain, draw, obs_id)
          float64
          8.744 5.527 4.382 ... 5.203 5.425
          array([[[8.74352758, 5.52741943, 4.38191711, ..., 5.14480198,
         5.24617995, 5.41743901],
        [8.77550506, 4.90547857, 4.15303127, ..., 4.60138867,
         5.93864166, 4.59592892],
        [8.71329607, 5.52044291, 4.52511514, ..., 5.06972045,
         5.27857848, 5.00547222],
        ...,
        [8.88502664, 5.57377862, 3.89409635, ..., 5.59909617,
         5.32225401, 5.20482076],
        [8.55664857, 4.88659312, 4.38079346, ..., 4.83263795,
         5.33348824, 5.05874786],
        [8.76757844, 4.90809102, 4.54824247, ..., 5.25881292,
         6.12158758, 5.20176982]],

       [[8.86872656, 5.0022022 , 4.58149104, ..., 5.78242943,
         5.27274199, 5.02613374],
        [8.60545624, 5.36327142, 3.87549526, ..., 4.78169205,
         5.73215524, 4.94515273],
        [8.72934777, 5.23184664, 4.26631068, ..., 5.43608436,
         5.69109134, 4.9359198 ],
...
        [8.51872267, 5.67053018, 3.68091342, ..., 5.38074753,
         5.28068957, 4.81562403],
        [8.84368012, 5.00970368, 4.76789871, ..., 5.08216918,
         5.81855931, 5.10367318],
        [8.5910571 , 4.86528321, 4.36407615, ..., 5.52091691,
         5.53541407, 4.56811379]],

       [[8.8404785 , 5.25125451, 3.84738158, ..., 5.62313306,
         4.79819979, 5.02215709],
        [8.58596763, 5.2463655 , 4.40220118, ..., 4.7902282 ,
         5.66964292, 4.48441869],
        [8.60738243, 5.61050158, 4.28259516, ..., 5.1362292 ,
         5.91231786, 5.18238246],
        ...,
        [8.60479962, 5.23713291, 4.55945121, ..., 5.27873283,
         4.82834783, 6.02702215],
        [8.86453311, 5.38558531, 4.04103726, ..., 4.84883412,
         5.98583658, 4.13783256],
        [8.62106165, 5.49281595, 3.85385124, ..., 5.77014132,
         5.20270335, 5.42464754]]])
        • sigma
          (chain, draw)
          float64
          0.4881 0.5189 ... 0.4919 0.536
          array([[0.48813477, 0.51890268, 0.48960848, ..., 0.51885326, 0.46876196,
        0.48209071],
       [0.52594341, 0.5199723 , 0.51310773, ..., 0.51786577, 0.51361691,
        0.54464582],
       [0.53240365, 0.51134871, 0.48844114, ..., 0.47774974, 0.4976345 ,
        0.51907641],
       [0.57185232, 0.50741003, 0.52455992, ..., 0.49067704, 0.49186708,
        0.53600122]])
      • created_at :
        2022-11-09T13:37:35.440374
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2
        sampling_time :
        46.495951890945435
        tuning_steps :
        1000

    • <xarray.Dataset>
Dimensions:     (chain: 4, draw: 2000, obs_id: 200)
Coordinates:
  * chain       (chain) int64 0 1 2 3
  * draw        (draw) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
  * obs_id      (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
Data variables:
    likelihood  (chain, draw, obs_id) float64 8.641 5.752 4.043 ... 5.464 5.011
Attributes:
    created_at:                 2022-11-09T13:37:51.112643
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • chain: 4
        • draw: 2000
        • obs_id: 200
        • chain
          (chain)
          int64
          0 1 2 3
          array([0, 1, 2, 3])
        • draw
          (draw)
          int64
          0 1 2 3 4 ... 1996 1997 1998 1999
          array([   0,    1,    2, ..., 1997, 1998, 1999])
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • likelihood
          (chain, draw, obs_id)
          float64
          8.641 5.752 4.043 ... 5.464 5.011
          array([[[8.64144021, 5.75203576, 4.04303489, ..., 5.3221561 ,
         5.40895724, 6.37164037],
        [9.20737422, 4.77367821, 3.79719738, ..., 4.46492958,
         5.71891587, 5.20647785],
        [8.79904259, 5.56006589, 4.66385612, ..., 5.00503684,
         4.73687063, 3.91708441],
        ...,
        [9.02951775, 5.1090761 , 4.01542758, ..., 5.58777459,
         4.99099305, 5.63332393],
        [8.28701055, 5.0233726 , 4.39672476, ..., 4.91416348,
         5.00382789, 5.08662577],
        [8.9022953 , 4.21894133, 4.39714918, ..., 6.12681816,
         6.37711376, 5.08892117]],

       [[8.90659297, 5.2929161 , 5.2026002 , ..., 5.66422753,
         6.46146195, 4.96200898],
        [7.86936615, 4.71119665, 4.23075325, ..., 3.85400236,
         5.44458598, 3.88321931],
        [8.55068507, 5.08472439, 4.63502134, ..., 5.79729243,
         4.90870811, 4.79077975],
...
        [8.60001121, 5.82915763, 3.56456982, ..., 4.72691436,
         4.40132837, 5.11605809],
        [8.14890069, 5.16532458, 5.14091545, ..., 5.38339109,
         6.27748992, 5.65373504],
        [8.74895254, 5.05543356, 4.57462896, ..., 6.51428266,
         4.76162078, 4.14433726]],

       [[8.45070324, 5.85182674, 4.11871452, ..., 5.48546099,
         4.51368516, 3.61546662],
        [8.6352959 , 4.88094404, 4.17249571, ..., 4.8559512 ,
         5.60901423, 3.71047229],
        [8.37373307, 5.92288004, 4.81028947, ..., 5.65188383,
         5.09270627, 4.74995286],
        ...,
        [8.66215207, 6.10395075, 3.98724745, ..., 4.89645146,
         4.25671533, 6.76667397],
        [8.17018768, 5.84520123, 3.1958053 , ..., 5.40706817,
         5.52490914, 3.75315205],
        [8.13533292, 4.99707698, 4.77117672, ..., 5.46631615,
         5.46421884, 5.01056552]]])
      • created_at :
        2022-11-09T13:37:51.112643
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

    • <xarray.Dataset>
Dimensions:     (chain: 4, draw: 2000, obs_id: 200)
Coordinates:
  * chain       (chain) int64 0 1 2 3
  * draw        (draw) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
  * obs_id      (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
Data variables:
    likelihood  (chain, draw, obs_id) float64 -60.87 -3.722 ... -0.4692 -1.227
Attributes:
    created_at:                 2022-11-09T13:37:35.999780
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • chain: 4
        • draw: 2000
        • obs_id: 200
        • chain
          (chain)
          int64
          0 1 2 3
          array([0, 1, 2, 3])
        • draw
          (draw)
          int64
          0 1 2 3 4 ... 1996 1997 1998 1999
          array([   0,    1,    2, ..., 1997, 1998, 1999])
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • likelihood
          (chain, draw, obs_id)
          float64
          -60.87 -3.722 ... -0.4692 -1.227
          array([[[-60.86542114,  -3.7223262 ,  -1.84617856, ...,  -0.20197954,
          -0.35776789,  -1.3033495 ],
        [-54.58626866,  -7.08844946,  -1.06286115, ...,  -0.83136725,
          -0.59016851,  -0.28036118],
        [-59.82761828,  -3.74197706,  -2.41088758, ...,  -0.21984483,
          -0.32518437,  -0.40857451],
        ...,
        [-56.8192156 ,  -3.15976965,  -0.55614316, ...,  -0.62963094,
          -0.33457504,  -0.74953226],
        [-61.44944842,  -8.69068665,  -1.93988539, ...,  -0.39726649,
          -0.23944424,  -0.46583361],
        [-62.94125141,  -8.07551173,  -2.568245  , ...,  -0.21264417,
          -0.97093642,  -0.746406  ]],

       [[-54.99352111,  -6.26693   ,  -2.40353173, ...,  -0.98867984,
          -0.38584232,  -0.47709643],
        [-51.01674802,  -4.20380108,  -0.53033473, ...,  -0.52223463,
          -0.34911562,  -0.38263294],
        [-54.86481123,  -5.05792678,  -1.37657487, ...,  -0.40205764,
          -0.30802281,  -0.36382854],
...
        [-58.32492137,  -3.08826499,  -0.25462351, ...,  -0.29222589,
          -0.30450405,  -0.21326374],
        [-60.78544609,  -6.85752476,  -3.48383595, ...,  -0.23166495,
          -0.40243526,  -0.56173413],
        [-50.91078934,  -7.37322541,  -1.65941333, ...,  -0.51213834,
          -0.26374465,  -0.29213058]],

       [[-46.17031693,  -4.13573953,  -0.54811451, ...,  -0.69559651,
          -1.15407877,  -0.52582118],
        [-53.14058874,  -5.06600445,  -1.83288695, ...,  -0.4984505 ,
          -0.28467215,  -0.32491124],
        [-50.17837711,  -2.94379157,  -1.39609645, ...,  -0.27436152,
          -0.55508837,  -0.70910831],
        ...,
        [-57.18552426,  -5.42782435,  -2.55258599, ...,  -0.23892836,
          -1.19708281,  -3.90329025],
        [-62.67539028,  -4.47764883,  -0.82179919, ...,  -0.40271297,
          -0.66012368,  -0.84613764],
        [-48.34204689,  -3.37324921,  -0.51733771, ...,  -0.95455446,
          -0.46924544,  -1.22720162]]])
      • created_at :
        2022-11-09T13:37:35.999780
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

    • <xarray.Dataset>
Dimensions:              (chain: 4, draw: 2000)
Coordinates:
  * chain                (chain) int64 0 1 2 3
  * draw                 (draw) int64 0 1 2 3 4 5 ... 1995 1996 1997 1998 1999
Data variables: (12/16)
    acceptance_rate      (chain, draw) float64 0.9895 0.8118 ... 0.9199 0.8375
    n_steps              (chain, draw) float64 31.0 31.0 31.0 ... 31.0 31.0 31.0
    tree_depth           (chain, draw) int64 5 5 5 5 5 5 5 5 ... 5 5 5 5 5 5 5 5
    energy_error         (chain, draw) float64 -0.08384 0.09877 ... 0.2174
    max_energy_error     (chain, draw) float64 -0.3237 0.3343 ... 0.1872 0.3201
    process_time_diff    (chain, draw) float64 0.005672 0.005849 ... 0.005637
    ...                   ...
    lp                   (chain, draw) float64 -350.1 -362.7 ... -352.8 -366.3
    smallest_eigval      (chain, draw) float64 nan nan nan nan ... nan nan nan
    step_size_bar        (chain, draw) float64 0.1601 0.1601 ... 0.1662 0.1662
    diverging            (chain, draw) bool False False False ... False False
    step_size            (chain, draw) float64 0.1495 0.1495 ... 0.1676 0.1676
    largest_eigval       (chain, draw) float64 nan nan nan nan ... nan nan nan
Attributes:
    created_at:                 2022-11-09T13:37:35.452666
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
    sampling_time:              46.495951890945435
    tuning_steps:               1000
      xarray.Dataset
        • chain: 4
        • draw: 2000
        • chain
          (chain)
          int64
          0 1 2 3
          array([0, 1, 2, 3])
        • draw
          (draw)
          int64
          0 1 2 3 4 ... 1996 1997 1998 1999
          array([   0,    1,    2, ..., 1997, 1998, 1999])
        • acceptance_rate
          (chain, draw)
          float64
          0.9895 0.8118 ... 0.9199 0.8375
          array([[0.98949688, 0.81181998, 0.93384018, ..., 0.98307339, 0.96983487,
        0.9753301 ],
       [0.83695072, 0.99086602, 0.98335682, ..., 0.95839766, 0.98665995,
        0.96171918],
       [0.98013982, 1.        , 0.94590947, ..., 0.70126769, 0.9770592 ,
        0.97624725],
       [0.94769058, 0.99594134, 0.87097286, ..., 0.97516364, 0.91986975,
        0.83751788]])
        • n_steps
          (chain, draw)
          float64
          31.0 31.0 31.0 ... 31.0 31.0 31.0
          array([[31., 31., 31., ..., 31., 31., 31.],
       [31., 31., 31., ..., 31., 31., 31.],
       [31., 31., 31., ..., 31., 31., 31.],
       [31., 31., 31., ..., 31., 31., 31.]])
        • tree_depth
          (chain, draw)
          int64
          5 5 5 5 5 5 5 5 ... 5 5 5 5 5 5 5 5
          array([[5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5]])
        • energy_error
          (chain, draw)
          float64
          -0.08384 0.09877 ... 0.1187 0.2174
          array([[-0.08383905,  0.09877008,  0.10360686, ...,  0.08640426,
        -0.02942147,  0.13710444],
       [ 0.29973228, -0.02229062,  0.01491618, ...,  0.10387911,
        -0.06748885,  0.10912723],
       [ 0.04044541, -0.23421695,  0.09694502, ...,  0.34785799,
         0.12383348, -0.08932485],
       [ 0.03315935, -0.22976407,  0.11084977, ..., -0.14849933,
         0.11874378,  0.21740001]])
        • max_energy_error
          (chain, draw)
          float64
          -0.3237 0.3343 ... 0.1872 0.3201
          array([[-0.32374934,  0.33432771,  0.22142569, ..., -0.17662432,
        -0.28160666, -0.31878279],
       [ 0.41461154, -0.1309749 , -0.22603568, ..., -0.26374009,
        -0.21565493,  0.13201054],
       [ 0.0612734 , -0.25813869,  0.18569643, ...,  0.75338588,
        -0.34793549, -0.28261466],
       [ 0.11330144, -0.3627097 ,  0.24214139, ..., -0.18257251,
         0.18716959,  0.32013469]])
        • process_time_diff
          (chain, draw)
          float64
          0.005672 0.005849 ... 0.005637
          array([[0.005672, 0.005849, 0.00564 , ..., 0.00568 , 0.005698, 0.005707],
       [0.005691, 0.005698, 0.005651, ..., 0.006401, 0.006081, 0.005596],
       [0.00568 , 0.005666, 0.005591, ..., 0.006672, 0.005902, 0.005704],
       [0.006102, 0.005791, 0.005657, ..., 0.005625, 0.005599, 0.005637]])
        • perf_counter_start
          (chain, draw)
          float64
          30.63 30.64 30.65 ... 45.87 45.88
          array([[30.63440198, 30.64022168, 30.64623307, ..., 42.8325296 ,
        42.8383483 , 42.84420278],
       [30.79673515, 30.80257273, 30.80840904, ..., 43.00991387,
        43.01650429, 43.02274264],
       [30.72152346, 30.7273735 , 30.73319075, ..., 42.88104742,
        42.88788403, 42.89393226],
       [33.69940925, 33.70568169, 33.71165544, ..., 45.86818541,
        45.87395314, 45.87968532]])
        • perf_counter_diff
          (chain, draw)
          float64
          0.005671 0.005848 ... 0.005636
          array([[0.0056714 , 0.00584828, 0.00563893, ..., 0.00567942, 0.00569751,
        0.00570738],
       [0.00568993, 0.00569761, 0.00565047, ..., 0.00640032, 0.00608116,
        0.00559526],
       [0.00569556, 0.00566596, 0.00559076, ..., 0.00668251, 0.00590671,
        0.00570238],
       [0.00610157, 0.0058214 , 0.00565711, ..., 0.00562505, 0.00559811,
        0.00563632]])
        • energy
          (chain, draw)
          float64
          444.7 483.6 463.6 ... 448.2 466.1
          array([[444.68972844, 483.5695198 , 463.55971002, ..., 444.57757127,
        447.96787097, 431.73910582],
       [464.65544859, 466.9452041 , 463.05623033, ..., 442.79489645,
        453.3358193 , 453.46156508],
       [475.2418436 , 451.65355904, 443.98860678, ..., 421.51260822,
        423.31378223, 450.35141672],
       [496.02305635, 473.48968202, 477.72154185, ..., 443.6172865 ,
        448.22309191, 466.10519362]])
        • index_in_trajectory
          (chain, draw)
          int64
          -17 -26 20 14 -18 ... 7 16 25 -15
          array([[-17, -26,  20, ..., -13,  16, -10],
       [ 22,  17, -13, ...,  18, -20, -25],
       [-14, -10,  23, ..., -28, -17,  15],
       [ 16, -22,  12, ...,  16,  25, -15]])
        • lp
          (chain, draw)
          float64
          -350.1 -362.7 ... -352.8 -366.3
          array([[-350.08090895, -362.66549122, -357.07212596, ..., -353.87023145,
        -338.88169682, -341.80997914],
       [-367.54496137, -365.96190233, -359.90856048, ..., -351.78816851,
        -345.68808173, -358.68915983],
       [-369.11506089, -342.62999597, -354.44561405, ..., -325.0752659 ,
        -339.09205768, -353.71032042],
       [-382.71459084, -359.86492696, -363.72296097, ..., -342.89580393,
        -352.80227461, -366.33444894]])
        • smallest_eigval
          (chain, draw)
          float64
          nan nan nan nan ... nan nan nan nan
          array([[nan, nan, nan, ..., nan, nan, nan],
       [nan, nan, nan, ..., nan, nan, nan],
       [nan, nan, nan, ..., nan, nan, nan],
       [nan, nan, nan, ..., nan, nan, nan]])
        • step_size_bar
          (chain, draw)
          float64
          0.1601 0.1601 ... 0.1662 0.1662
          array([[0.16014141, 0.16014141, 0.16014141, ..., 0.16014141, 0.16014141,
        0.16014141],
       [0.16079267, 0.16079267, 0.16079267, ..., 0.16079267, 0.16079267,
        0.16079267],
       [0.16011546, 0.16011546, 0.16011546, ..., 0.16011546, 0.16011546,
        0.16011546],
       [0.16618295, 0.16618295, 0.16618295, ..., 0.16618295, 0.16618295,
        0.16618295]])
        • diverging
          (chain, draw)
          bool
          False False False ... False False
          array([[False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False]])
        • step_size
          (chain, draw)
          float64
          0.1495 0.1495 ... 0.1676 0.1676
          array([[0.14954232, 0.14954232, 0.14954232, ..., 0.14954232, 0.14954232,
        0.14954232],
       [0.16448243, 0.16448243, 0.16448243, ..., 0.16448243, 0.16448243,
        0.16448243],
       [0.17057224, 0.17057224, 0.17057224, ..., 0.17057224, 0.17057224,
        0.17057224],
       [0.16763304, 0.16763304, 0.16763304, ..., 0.16763304, 0.16763304,
        0.16763304]])
        • largest_eigval
          (chain, draw)
          float64
          nan nan nan nan ... nan nan nan nan
          array([[nan, nan, nan, ..., nan, nan, nan],
       [nan, nan, nan, ..., nan, nan, nan],
       [nan, nan, nan, ..., nan, nan, nan],
       [nan, nan, nan, ..., nan, nan, nan]])
      • created_at :
        2022-11-09T13:37:35.452666
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2
        sampling_time :
        46.495951890945435
        tuning_steps :
        1000

    • <xarray.Dataset>
Dimensions:      (chain: 1, draw: 500, coefs_dim_0: 2, obs_id: 200)
Coordinates:
  * chain        (chain) int64 0
  * draw         (draw) int64 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
  * coefs_dim_0  (coefs_dim_0) int64 0 1
  * obs_id       (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
Data variables:
    sigma        (chain, draw) float64 16.77 6.404 5.943 ... 5.954 13.04 4.308
    coefs        (chain, draw, coefs_dim_0) float64 10.2 0.01543 ... 0.3662
    ar           (chain, draw, obs_id) float64 9.003 -18.85 ... 14.98 15.32
Attributes:
    created_at:                 2022-11-09T13:36:41.366276
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • chain: 1
        • draw: 500
        • coefs_dim_0: 2
        • obs_id: 200
        • chain
          (chain)
          int64
          0
          array([0])
        • draw
          (draw)
          int64
          0 1 2 3 4 5 ... 495 496 497 498 499
          array([  0,   1,   2, ..., 497, 498, 499])
        • coefs_dim_0
          (coefs_dim_0)
          int64
          0 1
          array([0, 1])
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • sigma
          (chain, draw)
          float64
          16.77 6.404 5.943 ... 13.04 4.308
          array([[16.76843906,  6.40411708,  5.94291271,  6.16181643,  2.47538257,
         3.62704806,  4.08747738,  3.05154676,  9.98411565, 10.22012446,
         4.95067989,  4.74501328,  9.85943452, 11.94628761,  2.41601397,
         8.71346319, 13.13855773,  1.0601317 ,  2.21669304, 14.96364015,
         1.37556993,  5.18216339,  8.77430372,  1.17204321,  0.03186847,
        10.08246153,  2.09745603,  2.77245341,  3.14897039, 12.59529963,
         8.55964877,  8.22785164,  7.56273695, 12.87647578,  2.40645396,
         5.90681503, 10.5035026 ,  6.32313202,  7.4668807 ,  6.84945782,
        17.94909433,  1.0308095 ,  6.0482961 ,  7.61469857, 11.01262292,
         0.15081356,  1.21012571,  6.4473999 ,  7.92874495,  8.75355943,
         4.96667642,  0.14523657,  8.55426643,  2.24921563,  1.25806684,
         0.27619812, 13.44382133,  9.25850098,  9.79599173,  0.6672222 ,
        12.05348543,  5.92001852,  0.80672269,  0.93566355,  5.59091408,
         6.29986579,  9.7177924 ,  4.86035198,  1.12038787, 13.48032932,
         6.18407487, 16.28252479,  1.48399394,  1.07747448,  5.65697727,
         2.81119188,  0.92584337, 25.98667756,  3.59242822, 14.28534513,
         2.76029648,  2.51794623,  7.60043114,  3.46047742,  7.90877254,
         5.65940604,  3.72743813,  1.06058534,  6.59035297,  3.71149226,
        20.13797437,  3.29191708,  4.51498548,  0.50688542,  5.00019235,
        17.83737287,  2.70596516,  1.25279284,  0.80422021,  3.50313686,
...
        17.96512064,  5.07288842,  7.37048254,  6.59084309,  8.35662642,
         3.1228615 ,  1.43030014,  1.98235744,  8.09340368,  9.54951366,
         7.94913832,  8.05049286,  4.4664681 , 13.66716976,  0.4319542 ,
         7.81165264, 13.07872709,  3.45355066,  9.14849984,  6.06336408,
         1.95074813,  0.68419654,  4.67810159,  7.97706402,  6.44474432,
         5.76473098,  0.0817305 ,  8.29923585,  1.95538943,  0.62189289,
         2.40151006,  4.83413595,  5.32816377,  8.11257452, 13.61203394,
         3.58400336,  2.74650475, 13.93278234,  5.61353355,  7.39171159,
         0.74204684,  1.95517276,  2.52728228, 14.91513793,  1.34110589,
         3.47079636, 10.08241886, 11.102863  ,  3.03700871,  1.24477134,
         3.89339953,  4.1232135 ,  3.64064671,  5.91050689,  3.89746377,
         2.63544878,  7.30505574,  0.52221113,  5.04649846,  6.65574185,
        10.22539428,  5.42955168,  4.71980875,  1.22224769,  8.22603356,
        11.50819064, 11.47901484,  5.33599859,  1.84813256,  5.0322994 ,
         1.51924399,  1.4478636 ,  2.74324755,  0.79407895,  2.72371208,
         0.38521757,  9.73540043,  9.89166482,  1.18204207,  3.83214476,
         2.64988781,  3.86276387, 17.34555872,  5.20432375,  3.93386763,
         2.12652191, 10.82108038,  1.76419649,  7.92825281,  0.09060821,
         7.39175656,  3.52224885,  1.52515064, 14.24381387,  6.15655092,
         4.05759568,  2.03310616,  5.9535346 , 13.04491808,  4.30809941]])
        • coefs
          (chain, draw, coefs_dim_0)
          float64
          10.2 0.01543 10.09 ... 10.08 0.3662
          array([[[ 1.02039330e+01,  1.54297312e-02],
        [ 1.00938385e+01,  3.13731916e-01],
        [ 1.00057961e+01,  1.37232981e-01],
        [ 9.95214242e+00,  2.35381613e-01],
        [ 1.00660870e+01,  3.74344308e-02],
        [ 1.00405283e+01, -6.18654848e-02],
        [ 9.98597008e+00,  2.37025717e-01],
        [ 9.91189827e+00, -5.32000367e-02],
        [ 9.94953246e+00,  3.49698134e-01],
        [ 9.99156183e+00,  2.87964182e-01],
        [ 9.74941640e+00,  1.82074548e-01],
        [ 9.95952955e+00,  3.41724581e-01],
        [ 1.00531934e+01, -8.46082088e-03],
        [ 1.00366640e+01,  3.05965720e-01],
        [ 1.00968753e+01,  1.94996336e-01],
        [ 9.96058027e+00,  3.97799625e-01],
        [ 1.02390532e+01,  2.83879757e-01],
        [ 1.01409765e+01,  2.86228897e-01],
        [ 9.96172666e+00,  3.06176422e-01],
        [ 9.88131487e+00,  1.47223295e-01],
...
        [ 1.00701898e+01,  2.79532854e-01],
        [ 1.01310345e+01,  8.54532070e-02],
        [ 9.77100828e+00,  2.64596899e-01],
        [ 9.97567160e+00,  2.31133292e-01],
        [ 9.86563621e+00,  2.07563173e-01],
        [ 9.98229572e+00,  3.63847661e-01],
        [ 1.02755464e+01,  2.21790826e-01],
        [ 1.01532330e+01,  1.92084682e-01],
        [ 1.00084068e+01,  9.49213643e-03],
        [ 9.95043381e+00,  7.42764610e-02],
        [ 1.01031381e+01,  2.69489783e-01],
        [ 1.01005543e+01,  2.38210860e-01],
        [ 1.00319015e+01,  2.91596405e-01],
        [ 1.01362719e+01,  1.18845843e-01],
        [ 9.87177697e+00,  3.45208445e-01],
        [ 9.99246974e+00,  1.15384365e-01],
        [ 9.81693170e+00,  2.11180153e-01],
        [ 9.90388553e+00,  2.04965340e-01],
        [ 1.00303113e+01,  3.41089129e-01],
        [ 1.00756179e+01,  3.66170303e-01]]])
        • ar
          (chain, draw, obs_id)
          float64
          9.003 -18.85 -4.548 ... 14.98 15.32
          array([[[  9.00299151, -18.8485096 ,  -4.54752041, ...,  10.94949945,
           4.16275636,  34.84306712],
        [  8.82045826,  14.2681331 ,  26.43646958, ...,  23.55809525,
          32.46011812,  15.9262657 ],
        [  8.96405557,  12.68371253,   0.61840562, ...,   9.79001285,
           7.938978  ,   8.52595322],
        ...,
        [  9.00497583,  19.18717977,  12.13418556, ...,   7.99254007,
          12.15391017,   8.09397143],
        [  9.01389166,  13.07090743,   9.1545123 , ...,  16.4515381 ,
          24.39032085,  18.45802279],
        [  8.87730016,  10.97257878,  18.45178715, ...,  15.21527822,
          14.9834655 ,  15.31628038]]])
      • created_at :
        2022-11-09T13:36:41.366276
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

    • <xarray.Dataset>
Dimensions:     (chain: 1, draw: 500, obs_id: 200)
Coordinates:
  * chain       (chain) int64 0
  * draw        (draw) int64 0 1 2 3 4 5 6 7 ... 492 493 494 495 496 497 498 499
  * obs_id      (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
Data variables:
    likelihood  (chain, draw, obs_id) float64 17.04 -1.075 -15.33 ... 25.43 17.2
Attributes:
    created_at:                 2022-11-09T13:36:41.368764
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • chain: 1
        • draw: 500
        • obs_id: 200
        • chain
          (chain)
          int64
          0
          array([0])
        • draw
          (draw)
          int64
          0 1 2 3 4 5 ... 495 496 497 498 499
          array([  0,   1,   2, ..., 497, 498, 499])
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • likelihood
          (chain, draw, obs_id)
          float64
          17.04 -1.075 -15.33 ... 25.43 17.2
          array([[[ 17.03846974,  -1.07466242, -15.33174563, ...,  -1.52121242,
          17.77951906,  35.77846349],
        [ -2.84662079,  12.39775329,  22.96673596, ...,  25.81453948,
          35.72581656,   5.01295186],
        [ 14.29887991,  10.45486191,  -3.43067872, ...,  -2.09864025,
           7.30810655,  13.45918065],
        ...,
        [  7.72346751,  27.22335768,   8.32598136, ...,   6.83689315,
           9.35751716,   7.55726229],
        [ -3.77649705,  13.44666845,  -6.05583262, ...,  21.25178371,
          45.73048555,   6.76886275],
        [ 12.3604314 ,   8.45027733,  16.38837641, ...,  16.62578364,
          25.42655466,  17.19886967]]])
      • created_at :
        2022-11-09T13:36:41.368764
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

    • <xarray.Dataset>
Dimensions:     (obs_id: 200)
Coordinates:
  * obs_id      (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
Data variables:
    likelihood  (obs_id) float64 3.367 6.823 3.497 7.269 ... 5.155 5.519 4.693
Attributes:
    created_at:                 2022-11-09T13:36:41.369822
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • obs_id: 200
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • likelihood
          (obs_id)
          float64
          3.367 6.823 3.497 ... 5.519 4.693
          array([3.36678808, 6.82268877, 3.49668131, 7.26876957, 3.79726746,
       6.59548238, 3.85896672, 6.5615307 , 4.05187049, 6.39620385,
       5.02195835, 5.31900993, 5.28465527, 5.20073556, 5.49522023,
       5.25258134, 4.78561231, 5.44677169, 5.16980359, 5.44528742,
       4.88947788, 5.77533767, 4.73572778, 6.36542966, 4.41177264,
       6.70727576, 3.98831032, 6.21671786, 4.40674358, 5.78290024,
       4.21893995, 6.33334036, 4.01557432, 6.62857595, 4.21017154,
       6.5386415 , 4.05395515, 6.41760073, 4.50837529, 5.91609097,
       5.27493145, 5.58677319, 5.11850318, 5.14996592, 5.53025179,
       4.70241421, 5.64903475, 5.09930691, 5.3286019 , 5.67767762,
       5.17115473, 5.6031667 , 4.84977124, 5.33115466, 5.14301679,
       5.23124031, 5.46781857, 4.72940137, 5.58059212, 4.97770599,
       5.34919693, 5.55250843, 4.82338116, 5.80944196, 4.55894121,
       6.17651742, 4.33505687, 5.65113914, 5.14290722, 5.36211065,
       5.38583573, 5.61634357, 4.65378798, 5.32347245, 5.482609  ,
       4.92004293, 5.63395489, 5.64054565, 4.92332935, 5.2539528 ,
       5.22183762, 5.18365224, 5.1375205 , 5.41105977, 5.70106931,
       4.93409239, 5.26799831, 5.25082489, 4.85024757, 5.10428416,
       5.32210233, 4.93611418, 5.41159451, 5.19494374, 5.56520868,
       4.68153102, 5.93229857, 4.23313192, 6.13822931, 4.60389058,
       6.15254809, 4.73332254, 5.99038712, 4.83631964, 5.56771104,
       5.26638219, 5.15806332, 5.36293655, 5.41918158, 5.18915235,
       5.46906279, 4.31710722, 5.63757617, 4.74755102, 6.07273406,
       3.99768924, 6.69052524, 3.85380521, 6.38528013, 4.81101254,
       5.87277283, 4.68711718, 5.73039993, 4.88056962, 5.45045625,
       4.87487659, 5.49853415, 5.47817712, 4.93098895, 5.57935722,
       4.8261684 , 5.52227139, 4.93651788, 5.7265195 , 4.84740354,
       5.64811706, 4.68137641, 5.75039632, 4.93424071, 5.76059072,
       4.45359485, 6.15746325, 4.90744046, 5.64399384, 5.15348945,
       5.24356944, 5.63829457, 4.81623659, 6.09475212, 4.35197041,
       6.07276028, 3.879273  , 6.82881183, 3.44011781, 6.7827564 ,
       4.37400916, 6.21361693, 4.49974697, 6.52664073, 3.97976962,
       6.39735001, 4.05252036, 6.20867084, 4.3600925 , 5.91438246,
       4.56348411, 5.96917553, 4.37056863, 5.78234943, 5.06369301,
       5.39549948, 5.73906972, 5.0505406 , 5.05771741, 5.44805861,
       5.15773415, 4.79321612, 6.13843482, 4.62551841, 5.9491786 ,
       4.87962984, 5.52865342, 5.16212009, 5.45923183, 5.05984681,
       5.49351839, 4.9232828 , 5.31009998, 5.34136063, 5.13034039,
       5.12328974, 5.49532282, 4.82028118, 5.93108626, 5.07202009,
       4.93355709, 5.35331517, 5.15467991, 5.51883117, 4.69290009])
      • created_at :
        2022-11-09T13:36:41.369822
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

    • <xarray.Dataset>
Dimensions:  (obs_id: 200)
Coordinates:
  * obs_id   (obs_id) int64 0 1 2 3 4 5 6 7 ... 192 193 194 195 196 197 198 199
Data variables:
    t        (obs_id) int32 0 1 2 3 4 5 6 7 ... 192 193 194 195 196 197 198 199
    y        (obs_id) float64 3.367 6.823 3.497 7.269 ... 5.155 5.519 4.693
Attributes:
    created_at:                 2022-11-09T13:36:41.370951
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • obs_id: 200
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • t
          (obs_id)
          int32
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,
        13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,
        26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,
        39,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,
        52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,
        65,  66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,
        78,  79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,
        91,  92,  93,  94,  95,  96,  97,  98,  99, 100, 101, 102, 103,
       104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116,
       117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129,
       130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142,
       143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155,
       156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168,
       169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194,
       195, 196, 197, 198, 199], dtype=int32)
        • y
          (obs_id)
          float64
          3.367 6.823 3.497 ... 5.519 4.693
          array([3.36678808, 6.82268877, 3.49668131, 7.26876957, 3.79726746,
       6.59548238, 3.85896672, 6.5615307 , 4.05187049, 6.39620385,
       5.02195835, 5.31900993, 5.28465527, 5.20073556, 5.49522023,
       5.25258134, 4.78561231, 5.44677169, 5.16980359, 5.44528742,
       4.88947788, 5.77533767, 4.73572778, 6.36542966, 4.41177264,
       6.70727576, 3.98831032, 6.21671786, 4.40674358, 5.78290024,
       4.21893995, 6.33334036, 4.01557432, 6.62857595, 4.21017154,
       6.5386415 , 4.05395515, 6.41760073, 4.50837529, 5.91609097,
       5.27493145, 5.58677319, 5.11850318, 5.14996592, 5.53025179,
       4.70241421, 5.64903475, 5.09930691, 5.3286019 , 5.67767762,
       5.17115473, 5.6031667 , 4.84977124, 5.33115466, 5.14301679,
       5.23124031, 5.46781857, 4.72940137, 5.58059212, 4.97770599,
       5.34919693, 5.55250843, 4.82338116, 5.80944196, 4.55894121,
       6.17651742, 4.33505687, 5.65113914, 5.14290722, 5.36211065,
       5.38583573, 5.61634357, 4.65378798, 5.32347245, 5.482609  ,
       4.92004293, 5.63395489, 5.64054565, 4.92332935, 5.2539528 ,
       5.22183762, 5.18365224, 5.1375205 , 5.41105977, 5.70106931,
       4.93409239, 5.26799831, 5.25082489, 4.85024757, 5.10428416,
       5.32210233, 4.93611418, 5.41159451, 5.19494374, 5.56520868,
       4.68153102, 5.93229857, 4.23313192, 6.13822931, 4.60389058,
       6.15254809, 4.73332254, 5.99038712, 4.83631964, 5.56771104,
       5.26638219, 5.15806332, 5.36293655, 5.41918158, 5.18915235,
       5.46906279, 4.31710722, 5.63757617, 4.74755102, 6.07273406,
       3.99768924, 6.69052524, 3.85380521, 6.38528013, 4.81101254,
       5.87277283, 4.68711718, 5.73039993, 4.88056962, 5.45045625,
       4.87487659, 5.49853415, 5.47817712, 4.93098895, 5.57935722,
       4.8261684 , 5.52227139, 4.93651788, 5.7265195 , 4.84740354,
       5.64811706, 4.68137641, 5.75039632, 4.93424071, 5.76059072,
       4.45359485, 6.15746325, 4.90744046, 5.64399384, 5.15348945,
       5.24356944, 5.63829457, 4.81623659, 6.09475212, 4.35197041,
       6.07276028, 3.879273  , 6.82881183, 3.44011781, 6.7827564 ,
       4.37400916, 6.21361693, 4.49974697, 6.52664073, 3.97976962,
       6.39735001, 4.05252036, 6.20867084, 4.3600925 , 5.91438246,
       4.56348411, 5.96917553, 4.37056863, 5.78234943, 5.06369301,
       5.39549948, 5.73906972, 5.0505406 , 5.05771741, 5.44805861,
       5.15773415, 4.79321612, 6.13843482, 4.62551841, 5.9491786 ,
       4.87962984, 5.52865342, 5.16212009, 5.45923183, 5.05984681,
       5.49351839, 4.9232828 , 5.31009998, 5.34136063, 5.13034039,
       5.12328974, 5.49532282, 4.82028118, 5.93108626, 5.07202009,
       4.93355709, 5.35331517, 5.15467991, 5.51883117, 4.69290009])
      • created_at :
        2022-11-09T13:36:41.370951
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

Lets check the model structure with plate notation and then examine the convergence diagnostics.

In [ ]:
az.plot_trace(idata_ar, figsize=(10, 6), kind="rank_vlines");

Next we'll check the summary estimates for the to AR coefficients and the sigma term.

In [ ]:
az.summary(idata_ar, var_names=["~ar"])
Out[ ]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
coefs[0] 9.744 0.094 9.555 9.911 0.001 0.001 13107.0 6200.0 1.0
coefs[1] -0.815 0.021 -0.854 -0.776 0.000 0.000 9141.0 5792.0 1.0
sigma 0.506 0.028 0.453 0.557 0.001 0.000 2965.0 5254.0 1.0

We can see here that the model fit has fairly correctly estimated the true parameters of the data generating process. We can also see this if we plot the posterior ar distribution against our observed data.

In [ ]:
fig, ax = plt.subplots(figsize=(10, 4))
idata_ar.posterior.ar.mean(["chain", "draw"]).plot(ax=ax, label="Posterior Mean AR level")
ax.plot(ar1_data, "o", color="black", markersize=2, label="Observed Data")
ax.legend()
ax.set_title("Fitted AR process\nand observed data");

Prediction Step¶

The next step works somewhat differently from generating posterior predictive observations for new data in a GLM model. Since we are forecasting from a learned posterior distribution of structural parameters we must condition on the learned parameters. Or put another way, we have to tell the model how many prediction steps we want to impute with the model we have just fit and from what basis to impute those values.

So for the shape handling purposes we have to feed our model new data for prediction and specify how to incorporate the learned parameters of the AR process. To do so, we initialise a new AR process for the future and feed it a set of initialising values we have learned when fitting our model to data. To make this as precise as can be use use the Dirac distribution to constrain the initial AR values very tightly around the learned posterior parameters.

In [ ]:
prediction_length = 250
n = prediction_length - ar1_data.shape[0]
obs = list(range(prediction_length))
with AR:
    ## We need to have coords for the observations minus the lagged term to correctly centre the prediction step
    AR.add_coords({"obs_id_fut_1": range(ar1_data.shape[0] - 1, 250, 1)})
    AR.add_coords({"obs_id_fut": range(ar1_data.shape[0], 250, 1)})
    # condition on the learned values of the AR process
    # initialise the future AR process precisely at the last observed value in the AR process
    # using the special feature of the dirac delta distribution to be 0 everywhere else.
    ar1_fut = pm.AR(
        "ar1_fut",
        init_dist=pm.DiracDelta.dist(ar1[..., -1]),
        rho=coefs,
        sigma=sigma,
        constant=True,
        dims="obs_id_fut_1",
    )
    yhat_fut = pm.Normal("yhat_fut", mu=ar1_fut[1:], sigma=sigma, dims="obs_id_fut")
    # use the updated values and predict outcomes and probabilities:
    idata_preds = pm.sample_posterior_predictive(
        idata_ar, var_names=["likelihood", "yhat_fut"], predictions=True, random_seed=100
    )

Sampling: [ar1_fut, likelihood, yhat_fut]
100.00% [8000/8000 00:31<00:00]

It's important to understand the conditional nature of the autoregressive forecast and the manner in which it depends on the observed data. In our two-step model fit and predict process we have learned the posterior distribution for the parameters of an AR process, and then used those parameters to centre our forecasts.

In [ ]:
pm.model_to_graphviz(AR)
Out[ ]:
clusterobs_id (200) obs_id (200) cluster2 2 clusterobs_id_fut_1 (51) obs_id_fut_1 (51) clusterobs_id_fut (50) obs_id_fut (50) y y ~ MutableData ar ar ~ Deterministic likelihood likelihood ~ Normal ar->likelihood ar1_fut ar1_fut ~ Deterministic ar->ar1_fut t t ~ MutableData t->ar likelihood->y coefs coefs ~ Normal coefs->ar coefs->ar1_fut sigma sigma ~ HalfNormal sigma->ar sigma->likelihood sigma->ar1_fut yhat_fut yhat_fut ~ Normal sigma->yhat_fut ar1_fut->yhat_fut
In [ ]:
idata_preds
Out[ ]:
arviz.InferenceData
    • <xarray.Dataset>
Dimensions:     (chain: 4, draw: 2000, obs_id: 200, obs_id_fut: 50)
Coordinates:
  * chain       (chain) int64 0 1 2 3
  * draw        (draw) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
  * obs_id      (obs_id) int64 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
  * obs_id_fut  (obs_id_fut) int64 200 201 202 203 204 ... 245 246 247 248 249
Data variables:
    likelihood  (chain, draw, obs_id) float64 8.599 6.006 3.882 ... 5.206 5.22
    yhat_fut    (chain, draw, obs_id_fut) float64 6.495 5.331 ... 3.495 6.905
Attributes:
    created_at:                 2022-11-09T13:38:50.219240
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • chain: 4
        • draw: 2000
        • obs_id: 200
        • obs_id_fut: 50
        • chain
          (chain)
          int64
          0 1 2 3
          array([0, 1, 2, 3])
        • draw
          (draw)
          int64
          0 1 2 3 4 ... 1996 1997 1998 1999
          array([   0,    1,    2, ..., 1997, 1998, 1999])
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • obs_id_fut
          (obs_id_fut)
          int64
          200 201 202 203 ... 246 247 248 249
          array([200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213,
       214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227,
       228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241,
       242, 243, 244, 245, 246, 247, 248, 249])
        • likelihood
          (chain, draw, obs_id)
          float64
          8.599 6.006 3.882 ... 5.206 5.22
          array([[[8.59884897, 6.00604341, 3.88206887, ..., 4.98169029,
         5.40359092, 5.68191078],
        [8.79670654, 5.93529041, 4.40429365, ..., 5.01546634,
         4.43474002, 3.60012164],
        [8.56847869, 5.39800394, 3.88356292, ..., 4.63002488,
         5.01572423, 4.43078707],
        ...,
        [9.20136104, 4.57883509, 4.48042098, ..., 5.31498913,
         4.44754853, 4.55047726],
        [9.00546596, 5.50518055, 6.01138556, ..., 4.70408191,
         5.11675198, 5.62857239],
        [8.88946697, 5.24419495, 4.00156719, ..., 4.20165221,
         5.82695311, 5.21254244]],

       [[9.58466508, 4.26734586, 4.27939642, ..., 5.14275023,
         5.82883516, 5.43938218],
        [9.6388725 , 5.1134545 , 4.08002374, ..., 4.39881783,
         5.77437617, 5.18038996],
        [8.10630832, 5.11994266, 4.43082836, ..., 4.67106398,
         5.45678235, 4.86631852],
...
        [8.15595583, 5.20394562, 2.951609  , ..., 5.3213695 ,
         6.18673762, 4.1615559 ],
        [8.80889118, 4.55812225, 5.01817234, ..., 5.01598758,
         4.9526523 , 5.56593906],
        [8.61649003, 5.09371976, 4.18762361, ..., 5.47849819,
         5.13814072, 5.95369104]],

       [[9.16858942, 5.12279053, 3.36495293, ..., 6.11475246,
         4.69878594, 4.39136645],
        [8.4054303 , 4.92610756, 4.21464288, ..., 4.26666647,
         5.84868038, 4.81156007],
        [8.67098501, 5.30548722, 4.40671805, ..., 5.16418269,
         5.67730921, 5.38076556],
        ...,
        [8.274733  , 4.58119292, 5.12602147, ..., 5.09634608,
         4.50401847, 5.5991206 ],
        [8.64207991, 5.23009314, 3.9992984 , ..., 4.43087845,
         6.05891835, 4.42354604],
        [8.16730802, 5.70020461, 4.63431287, ..., 5.95958104,
         5.20641388, 5.21960772]]])
        • yhat_fut
          (chain, draw, obs_id_fut)
          float64
          6.495 5.331 6.987 ... 3.495 6.905
          array([[[6.4953486 , 5.33147583, 6.98718248, ..., 6.36780178,
         5.19327288, 5.55648387],
        [6.35693501, 4.72419911, 6.34624456, ..., 6.34945251,
         2.44257016, 6.26709538],
        [6.41316592, 4.07332944, 6.75438993, ..., 3.94588935,
         5.45166734, 5.85294635],
        ...,
        [5.53022118, 3.81466774, 5.38812548, ..., 4.30277357,
         5.58011075, 4.3171974 ],
        [5.94471308, 3.53128884, 5.05963497, ..., 4.53589301,
         4.88987214, 5.09390923],
        [4.3616181 , 4.95844304, 5.64978963, ..., 5.81761355,
         5.16029723, 4.70581458]],

       [[4.78506671, 4.28035922, 4.18195535, ..., 6.33435425,
         5.15477639, 5.63097243],
        [5.6816181 , 5.56311808, 4.40708914, ..., 5.73869459,
         5.34695395, 4.84345455],
        [5.84512397, 5.51856095, 6.70363985, ..., 5.61478793,
         5.61052781, 4.7551204 ],
...
        [5.65597631, 4.89324145, 4.57705603, ..., 4.83832309,
         6.64658325, 3.62682651],
        [5.82433233, 4.10600282, 6.54642625, ..., 4.26288255,
         5.16459549, 5.82842815],
        [5.46029854, 5.90144111, 6.09288155, ..., 6.81420745,
         5.6433365 , 5.80063532]],

       [[5.46329009, 4.32768119, 7.01060984, ..., 7.33606912,
         4.79217401, 6.80963022],
        [5.64209241, 3.92499247, 6.18934625, ..., 5.4748427 ,
         4.45647414, 5.62512031],
        [4.40948812, 4.93112986, 6.26481676, ..., 5.87055775,
         3.87404467, 5.58124849],
        ...,
        [6.3038289 , 5.12722683, 5.24307923, ..., 4.31622093,
         5.74889202, 5.40140064],
        [6.05972035, 4.98366893, 5.07665905, ..., 6.3783536 ,
         4.17927265, 5.84074104],
        [6.19750378, 5.32128673, 5.88932402, ..., 5.47006247,
         3.49463187, 6.90531056]]])
      • created_at :
        2022-11-09T13:38:50.219240
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

    • <xarray.Dataset>
Dimensions:  (obs_id: 200)
Coordinates:
  * obs_id   (obs_id) int64 0 1 2 3 4 5 6 7 ... 192 193 194 195 196 197 198 199
Data variables:
    t        (obs_id) int32 0 1 2 3 4 5 6 7 ... 192 193 194 195 196 197 198 199
    y        (obs_id) float64 3.367 6.823 3.497 7.269 ... 5.155 5.519 4.693
Attributes:
    created_at:                 2022-11-09T13:38:50.221893
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.2.2
      xarray.Dataset
        • obs_id: 200
        • obs_id
          (obs_id)
          int64
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
        28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
        56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
        70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
        84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
        98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
       112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
       126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
       154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
       168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
       196, 197, 198, 199])
        • t
          (obs_id)
          int32
          0 1 2 3 4 5 ... 195 196 197 198 199
          array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,
        13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,
        26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,
        39,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,
        52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,
        65,  66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,
        78,  79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,
        91,  92,  93,  94,  95,  96,  97,  98,  99, 100, 101, 102, 103,
       104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116,
       117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129,
       130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142,
       143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155,
       156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168,
       169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194,
       195, 196, 197, 198, 199], dtype=int32)
        • y
          (obs_id)
          float64
          3.367 6.823 3.497 ... 5.519 4.693
          array([3.36678808, 6.82268877, 3.49668131, 7.26876957, 3.79726746,
       6.59548238, 3.85896672, 6.5615307 , 4.05187049, 6.39620385,
       5.02195835, 5.31900993, 5.28465527, 5.20073556, 5.49522023,
       5.25258134, 4.78561231, 5.44677169, 5.16980359, 5.44528742,
       4.88947788, 5.77533767, 4.73572778, 6.36542966, 4.41177264,
       6.70727576, 3.98831032, 6.21671786, 4.40674358, 5.78290024,
       4.21893995, 6.33334036, 4.01557432, 6.62857595, 4.21017154,
       6.5386415 , 4.05395515, 6.41760073, 4.50837529, 5.91609097,
       5.27493145, 5.58677319, 5.11850318, 5.14996592, 5.53025179,
       4.70241421, 5.64903475, 5.09930691, 5.3286019 , 5.67767762,
       5.17115473, 5.6031667 , 4.84977124, 5.33115466, 5.14301679,
       5.23124031, 5.46781857, 4.72940137, 5.58059212, 4.97770599,
       5.34919693, 5.55250843, 4.82338116, 5.80944196, 4.55894121,
       6.17651742, 4.33505687, 5.65113914, 5.14290722, 5.36211065,
       5.38583573, 5.61634357, 4.65378798, 5.32347245, 5.482609  ,
       4.92004293, 5.63395489, 5.64054565, 4.92332935, 5.2539528 ,
       5.22183762, 5.18365224, 5.1375205 , 5.41105977, 5.70106931,
       4.93409239, 5.26799831, 5.25082489, 4.85024757, 5.10428416,
       5.32210233, 4.93611418, 5.41159451, 5.19494374, 5.56520868,
       4.68153102, 5.93229857, 4.23313192, 6.13822931, 4.60389058,
       6.15254809, 4.73332254, 5.99038712, 4.83631964, 5.56771104,
       5.26638219, 5.15806332, 5.36293655, 5.41918158, 5.18915235,
       5.46906279, 4.31710722, 5.63757617, 4.74755102, 6.07273406,
       3.99768924, 6.69052524, 3.85380521, 6.38528013, 4.81101254,
       5.87277283, 4.68711718, 5.73039993, 4.88056962, 5.45045625,
       4.87487659, 5.49853415, 5.47817712, 4.93098895, 5.57935722,
       4.8261684 , 5.52227139, 4.93651788, 5.7265195 , 4.84740354,
       5.64811706, 4.68137641, 5.75039632, 4.93424071, 5.76059072,
       4.45359485, 6.15746325, 4.90744046, 5.64399384, 5.15348945,
       5.24356944, 5.63829457, 4.81623659, 6.09475212, 4.35197041,
       6.07276028, 3.879273  , 6.82881183, 3.44011781, 6.7827564 ,
       4.37400916, 6.21361693, 4.49974697, 6.52664073, 3.97976962,
       6.39735001, 4.05252036, 6.20867084, 4.3600925 , 5.91438246,
       4.56348411, 5.96917553, 4.37056863, 5.78234943, 5.06369301,
       5.39549948, 5.73906972, 5.0505406 , 5.05771741, 5.44805861,
       5.15773415, 4.79321612, 6.13843482, 4.62551841, 5.9491786 ,
       4.87962984, 5.52865342, 5.16212009, 5.45923183, 5.05984681,
       5.49351839, 4.9232828 , 5.31009998, 5.34136063, 5.13034039,
       5.12328974, 5.49532282, 4.82028118, 5.93108626, 5.07202009,
       4.93355709, 5.35331517, 5.15467991, 5.51883117, 4.69290009])
      • created_at :
        2022-11-09T13:38:50.221893
        arviz_version :
        0.12.1
        inference_library :
        pymc
        inference_library_version :
        4.2.2

Inspecting model fit and forecast¶

We can look at the standard posterior predictive fits but since our data is timeseries data we have to also look how draws from the posterior predictive distribution vary over time.

In [ ]:
def plot_fits(idata_ar, idata_preds):
    palette = "plasma"
    cmap = plt.get_cmap(palette)
    percs = np.linspace(51, 99, 100)
    colors = (percs - np.min(percs)) / (np.max(percs) - np.min(percs))
    mosaic = """AABB
                CCCC"""
    fig, axs = plt.subplot_mosaic(mosaic, sharex=False, figsize=(20, 10))
    axs = [axs[k] for k in axs.keys()]
    for i, p in enumerate(percs[::-1]):
        upper = np.percentile(
            az.extract_dataset(idata_ar, group="prior_predictive", num_samples=1000)["likelihood"],
            p,
            axis=1,
        )
        lower = np.percentile(
            az.extract_dataset(idata_ar, group="prior_predictive", num_samples=1000)["likelihood"],
            100 - p,
            axis=1,
        )
        color_val = colors[i]
        axs[0].fill_between(
            x=idata_ar["constant_data"]["t"],
            y1=upper.flatten(),
            y2=lower.flatten(),
            color=cmap(color_val),
            alpha=0.1,
        )

    axs[0].plot(
        az.extract_dataset(idata_ar, group="prior_predictive", num_samples=1000)["likelihood"].mean(
            axis=1
        ),
        color="cyan",
        label="Prior Predicted Mean Realisation",
    )

    axs[0].scatter(
        x=idata_ar["constant_data"]["t"],
        y=idata_ar["constant_data"]["y"],
        color="k",
        label="Observed Data points",
    )
    axs[0].set_title("Prior Predictive Fit", fontsize=20)
    axs[0].legend()

    for i, p in enumerate(percs[::-1]):
        upper = np.percentile(
            az.extract_dataset(idata_preds, group="predictions", num_samples=1000)["likelihood"],
            p,
            axis=1,
        )
        lower = np.percentile(
            az.extract_dataset(idata_preds, group="predictions", num_samples=1000)["likelihood"],
            100 - p,
            axis=1,
        )
        color_val = colors[i]
        axs[2].fill_between(
            x=idata_preds["predictions_constant_data"]["t"],
            y1=upper.flatten(),
            y2=lower.flatten(),
            color=cmap(color_val),
            alpha=0.1,
        )

        upper = np.percentile(
            az.extract_dataset(idata_preds, group="predictions", num_samples=1000)["yhat_fut"],
            p,
            axis=1,
        )
        lower = np.percentile(
            az.extract_dataset(idata_preds, group="predictions", num_samples=1000)["yhat_fut"],
            100 - p,
            axis=1,
        )
        color_val = colors[i]
        axs[2].fill_between(
            x=idata_preds["predictions"].coords["obs_id_fut"].data,
            y1=upper.flatten(),
            y2=lower.flatten(),
            color=cmap(color_val),
            alpha=0.1,
        )

    axs[2].plot(
        az.extract_dataset(idata_preds, group="predictions", num_samples=1000)["likelihood"].mean(
            axis=1
        ),
        color="cyan",
    )
    idata_preds.predictions.yhat_fut.mean(["chain", "draw"]).plot(
        ax=axs[2], color="cyan", label="Predicted Mean Realisation"
    )
    axs[2].scatter(
        x=idata_ar["constant_data"]["t"],
        y=idata_ar["constant_data"]["y"],
        color="k",
        label="Observed Data",
    )
    axs[2].set_title("Posterior Predictions Plotted", fontsize=20)
    axs[2].axvline(np.max(idata_ar["constant_data"]["t"]), color="black")
    axs[2].legend()
    axs[2].set_xlabel("Time in Days")
    axs[0].set_xlabel("Time in Days")
    az.plot_ppc(idata_ar, ax=axs[1])


plot_fits(idata_ar, idata_preds)

Here we can see that although the model converged and ends up with a reasonable fit to the existing data, and a plausible projection for future values. However, we have set the prior specification very poorly in allowing an absurdly broad range of values due to the kind of compounding logic of the auto-regressive function. For this reason it's very important to be able to inspect and tailor your model with prior predictive checks.

Secondly, the mean forecast fails to capture any long lasting structure, quickly dying down to a stable baseline. To account for these kind of short-lived forecasts, we can add more structure to our model, but first, let's complicate the picture.

Complicating the Picture¶

Often our data will involve more than one latent process, and might have more complex factors which drive the outcomes. To see one such complication let's add a trend to our data. By adding more structure to our forecast we are telling our model that we expect certain patterns or trends to remain in the data out into the future. The choice of which structures to add are at the discretion of the creative modeller - here we'll demonstrate some simple examples.

In [ ]:
y_t = -0.3 + np.arange(200) * -0.2 + np.random.normal(0, 10, 200)
y_t = y_t + ar1_data

fig, ax = plt.subplots(figsize=(10, 4))
ax.plot(y_t)
ax.set_title("AR Process + Trend data");

Wrapping our model into a function¶

In [ ]:
def make_latent_AR_model(ar_data, priors, prediction_steps=250, full_sample=True, samples=2000):
    with pm.Model() as AR:
        pass

    t_data = list(range(len(ar_data)))
    AR.add_coord("obs_id", t_data, mutable=True)

    with AR:
        ## Data containers to enable prediction
        t = pm.MutableData("t", t_data, dims="obs_id")
        y = pm.MutableData("y", ar_data, dims="obs_id")
        # The first coefficient will be the intercept term
        coefs = pm.Normal("coefs", priors["coefs"]["mu"], priors["coefs"]["sigma"])
        sigma = pm.HalfNormal("sigma", priors["sigma"])
        # We need one init variable for each lag, hence size is variable too
        init = pm.Normal.dist(
            priors["init"]["mu"], priors["init"]["sigma"], size=priors["init"]["size"]
        )
        # Steps of the AR model minus the lags required given specification
        ar1 = pm.AR(
            "ar",
            coefs,
            sigma=sigma,
            init_dist=init,
            constant=True,
            steps=t.shape[0] - (priors["coefs"]["size"] - 1),
            dims="obs_id",
        )

        # The Likelihood
        outcome = pm.Normal("likelihood", mu=ar1, sigma=sigma, observed=y, dims="obs_id")
        ## Sampling
        idata_ar = pm.sample_prior_predictive()
        if full_sample:
            idata_ar.extend(pm.sample(samples, random_seed=100, target_accept=0.95))
            idata_ar.extend(pm.sample_posterior_predictive(idata_ar))
        else:
            return idata_ar

    n = prediction_steps - ar_data.shape[0]

    with AR:
        AR.add_coords({"obs_id_fut_1": range(ar1_data.shape[0] - 1, 250, 1)})
        AR.add_coords({"obs_id_fut": range(ar1_data.shape[0], 250, 1)})
        # condition on the learned values of the AR process
        # initialise the future AR process precisely at the last observed value in the AR process
        # using the special feature of the dirac delta distribution to be 0 probability everywhere else.
        ar1_fut = pm.AR(
            "ar1_fut",
            init_dist=pm.DiracDelta.dist(ar1[..., -1]),
            rho=coefs,
            sigma=sigma,
            constant=True,
            dims="obs_id_fut_1",
        )
        yhat_fut = pm.Normal("yhat_fut", mu=ar1_fut[1:], sigma=sigma, dims="obs_id_fut")
        # use the updated values and predict outcomes and probabilities:
        idata_preds = pm.sample_posterior_predictive(
            idata_ar, var_names=["likelihood", "yhat_fut"], predictions=True, random_seed=100
        )

    return idata_ar, idata_preds, AR

Next we'll cycle through a number of prior specifications to show how that impacts the prior predictive distribution i.e. the implied distribution of our outcome if we were to forward sample from the model specified by our priors.

In [ ]:
priors_0 = {
    "coefs": {"mu": [-4, 0.2], "sigma": 0.1, "size": 2},
    "sigma": 8,
    "init": {"mu": 9, "sigma": 0.1, "size": 1},
}

priors_1 = {
    "coefs": {"mu": [-2, 0.2], "sigma": 0.1, "size": 2},
    "sigma": 12,
    "init": {"mu": 8, "sigma": 0.1, "size": 1},
}

priors_2 = {
    "coefs": {"mu": [0, 0.2], "sigma": 0.1, "size": 2},
    "sigma": 15,
    "init": {"mu": 8, "sigma": 0.1, "size": 1},
}

models = {}
for i, p in enumerate([priors_0, priors_1, priors_2]):
    models[i] = {}
    idata = make_latent_AR_model(y_t, p, full_sample=False)
    models[i]["idata"] = idata

Sampling: [ar, coefs, likelihood, sigma]
Sampling: [ar, coefs, likelihood, sigma]
Sampling: [ar, coefs, likelihood, sigma]
In [ ]:
fig, axs = plt.subplots(1, 3, figsize=(10, 4), sharey=True)
axs = axs.flatten()
for i, p in zip(range(3), [priors_0, priors_1, priors_2]):
    axs[i].plot(
        az.extract_dataset(models[i]["idata"], group="prior_predictive", num_samples=100)[
            "likelihood"
        ],
        color="blue",
        alpha=0.1,
    )
    axs[i].plot(y_t, "o", color="black", markersize=2)
    axs[i].set_title(
        "$y_{t+1}$" + f'= N({p["coefs"]["mu"][0]} + {p["coefs"]["mu"][1]}y$_t$, {p["sigma"]})'
    )
plt.suptitle("Prior Predictive Specifications", fontsize=20);

We can see the manner in which the model struggles to capture the trend line. Increasing the variability of the model will never capture the directional pattern we know to be in the data.

In [ ]:
priors_0 = {
    "coefs": {"mu": [-4, 0.2], "sigma": [0.5, 0.03], "size": 2},
    "sigma": 8,
    "init": {"mu": -4, "sigma": 0.1, "size": 1},
}

idata_no_trend, preds_no_trend, model = make_latent_AR_model(y_t, priors_0)

Sampling: [ar, coefs, likelihood, sigma]
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [coefs, sigma, ar]
100.00% [12000/12000 00:22<00:00 Sampling 4 chains, 0 divergences] 
Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 49 seconds.
Sampling: [likelihood]
100.00% [8000/8000 00:00<00:00] 
Sampling: [ar1_fut, likelihood, yhat_fut]
100.00% [8000/8000 00:29<00:00]
In [ ]:
plot_fits(idata_no_trend, preds_no_trend)

Forecasting with this model is somewhat hopeless because, while the model fit adjusts well with observed data, but it completely fails to capture the structural trend in the data. So without some structural constraint when we seek to make predictions with this simple AR model, it reverts to the mean level forecast very quickly.

Specifying a Trend Model¶

We will define a model to account for the trend in our data and combine this trend in an additive model with the autoregressive components. Again the model is much as before, but now we add additional latent features. These are to be combined in a simple additive combination but we can be more creative here if it would suit our model.

In [ ]:
def make_latent_AR_trend_model(
    ar_data, priors, prediction_steps=250, full_sample=True, samples=2000
):
    with pm.Model() as AR:
        pass

    t_data = list(range(len(ar_data)))
    AR.add_coord("obs_id", t_data, mutable=True)

    with AR:
        ## Data containers to enable prediction
        t = pm.MutableData("t", t_data, dims="obs_id")
        y = pm.MutableData("y", ar_data, dims="obs_id")
        # The first coefficient will be the intercept term
        coefs = pm.Normal("coefs", priors["coefs"]["mu"], priors["coefs"]["sigma"])
        sigma = pm.HalfNormal("sigma", priors["sigma"])
        # We need one init variable for each lag, hence size is variable too
        init = pm.Normal.dist(
            priors["init"]["mu"], priors["init"]["sigma"], size=priors["init"]["size"]
        )
        # Steps of the AR model minus the lags required given specification
        ar1 = pm.AR(
            "ar",
            coefs,
            sigma=sigma,
            init_dist=init,
            constant=True,
            steps=t.shape[0] - (priors["coefs"]["size"] - 1),
            dims="obs_id",
        )

        ## Priors for the linear trend component
        alpha = pm.Normal("alpha", priors["alpha"]["mu"], priors["alpha"]["sigma"])
        beta = pm.Normal("beta", priors["beta"]["mu"], priors["beta"]["sigma"])
        trend = pm.Deterministic("trend", alpha + beta * t, dims="obs_id")

        mu = ar1 + trend

        # The Likelihood
        outcome = pm.Normal("likelihood", mu=mu, sigma=sigma, observed=y, dims="obs_id")
        ## Sampling
        idata_ar = pm.sample_prior_predictive()
        if full_sample:
            idata_ar.extend(pm.sample(samples, random_seed=100, target_accept=0.95))
            idata_ar.extend(pm.sample_posterior_predictive(idata_ar))
        else:
            return idata_ar

    n = prediction_steps - ar_data.shape[0]

    with AR:
        AR.add_coords({"obs_id_fut_1": range(ar1_data.shape[0] - 1, prediction_steps, 1)})
        AR.add_coords({"obs_id_fut": range(ar1_data.shape[0], prediction_steps, 1)})
        t_fut = pm.MutableData("t_fut", list(range(ar1_data.shape[0], prediction_steps, 1)))
        # condition on the learned values of the AR process
        # initialise the future AR process precisely at the last observed value in the AR process
        # using the special feature of the dirac delta distribution to be 0 probability everywhere else.
        ar1_fut = pm.AR(
            "ar1_fut",
            init_dist=pm.DiracDelta.dist(ar1[..., -1]),
            rho=coefs,
            sigma=sigma,
            constant=True,
            dims="obs_id_fut_1",
        )
        trend = pm.Deterministic("trend_fut", alpha + beta * t_fut, dims="obs_id_fut")
        mu = ar1_fut[1:] + trend

        yhat_fut = pm.Normal("yhat_fut", mu=mu, sigma=sigma, dims="obs_id_fut")
        # use the updated values and predict outcomes and probabilities:
        idata_preds = pm.sample_posterior_predictive(
            idata_ar, var_names=["likelihood", "yhat_fut"], predictions=True, random_seed=100
        )

    return idata_ar, idata_preds, AR

We will fit this model by specifying priors on the negative trend and the range of the standard deviation to respect the direction of the data drift.

In [ ]:
priors_0 = {
    "coefs": {"mu": [0.2, 0.2], "sigma": [0.5, 0.03], "size": 2},
    "alpha": {"mu": -4, "sigma": 0.1},
    "beta": {"mu": -0.1, "sigma": 0.2},
    "sigma": 8,
    "init": {"mu": -4, "sigma": 0.1, "size": 1},
}


idata_trend, preds_trend, model = make_latent_AR_trend_model(y_t, priors_0, full_sample=True)

Sampling: [alpha, ar, beta, coefs, likelihood, sigma]
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [coefs, sigma, ar, alpha, beta]
100.00% [12000/12000 00:33<00:00 Sampling 4 chains, 0 divergences] 
Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 57 seconds.
Sampling: [likelihood]
100.00% [8000/8000 00:00<00:00] 
Sampling: [ar1_fut, likelihood, yhat_fut]
100.00% [8000/8000 00:30<00:00]
In [ ]:
pm.model_to_graphviz(model)
Out[ ]:
clusterobs_id (200) obs_id (200) cluster2 2 cluster50 50 clusterobs_id_fut_1 (51) obs_id_fut_1 (51) clusterobs_id_fut (50) obs_id_fut (50) y y ~ MutableData ar ar ~ Deterministic likelihood likelihood ~ Normal ar->likelihood ar1_fut ar1_fut ~ Deterministic ar->ar1_fut t t ~ MutableData t->ar trend trend ~ Deterministic t->trend trend->likelihood likelihood->y coefs coefs ~ Normal coefs->ar coefs->ar1_fut sigma sigma ~ HalfNormal sigma->ar sigma->likelihood sigma->ar1_fut yhat_fut yhat_fut ~ Normal sigma->yhat_fut alpha alpha ~ Normal alpha->trend trend_fut trend_fut ~ Deterministic alpha->trend_fut beta beta ~ Normal beta->trend beta->trend_fut t_fut t_fut ~ MutableData t_fut->trend_fut ar1_fut->yhat_fut trend_fut->yhat_fut

We can see the structure more clearly with the plate notation, and this additional structure has helped to appropriately fit the directional trend of the timeseries data.

In [ ]:
plot_fits(idata_trend, preds_trend);
In [ ]:
az.summary(idata_trend, var_names=["coefs", "sigma", "alpha", "beta"])
Out[ ]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
coefs[0] 1.086 0.475 0.220 2.006 0.010 0.007 2482.0 4278.0 1.0
coefs[1] 0.210 0.029 0.156 0.264 0.000 0.000 7296.0 5931.0 1.0
sigma 7.557 0.389 6.841 8.284 0.007 0.005 3332.0 4817.0 1.0
alpha -3.969 0.099 -4.171 -3.794 0.001 0.001 9220.0 5606.0 1.0
beta -0.145 0.009 -0.161 -0.128 0.000 0.000 1785.0 3995.0 1.0

Complicating the picture further¶

Next we'll add a seasonal component to our data and see how we can recover this aspect of the data with a bayesian structural timeseries model. Again, this is is because in reality our data is often the result of multiple converging influences. These influences can be capture in an additive bayesian structural model where our inferential model ensures that we allocate appropriate weight to each of the components.

In [ ]:
t_data = list(range(200))
n_order = 10
periods = np.array(t_data) / 7

fourier_features = pd.DataFrame(
    {
        f"{func}_order_{order}": getattr(np, func)(2 * np.pi * periods * order)
        for order in range(1, n_order + 1)
        for func in ("sin", "cos")
    }
)

y_t_s = y_t + 20 * fourier_features["sin_order_1"]

fig, ax = plt.subplots(figsize=(10, 4))
ax.plot(y_t_s)
ax.set_title("AR + Trend + Seasonality");

The key to fitting this model is to understand that we're now passing in synthetic fourier features to help account for seasonality effects. This works because (roughly speaking) we're trying to fit a complex oscillating phenomena using a weighted combination of sine and cosine waves. So we add these sine waves and consine waves like we would add any other feature variables in a regression model.

However, since we're using this weighted sum to fit the observed data, the model now expects a linear combination of those synthetic features also in the prediction step. As such we need to be able to supply those features even out into the future. This fact remains key for any other type of predictive feature we might want to add e.g. day of the week, holiday dummy variable or any other. If a feature is required to fit the observed data the feature must be available in the prediction step too.

Specifying the Trend + Seasonal Model¶

In [ ]:
def make_latent_AR_trend_seasonal_model(
    ar_data, ff, priors, prediction_steps=250, full_sample=True, samples=2000
):
    with pm.Model() as AR:
        pass

    ff = ff.to_numpy().T
    t_data = list(range(len(ar_data)))
    AR.add_coord("obs_id", t_data, mutable=True)
    ## The fourier features must be mutable to allow for addition fourier features to be
    ## passed in the prediction step.
    AR.add_coord("fourier_features", np.arange(len(ff)), mutable=True)

    with AR:
        ## Data containers to enable prediction
        t = pm.MutableData("t", t_data, dims="obs_id")
        y = pm.MutableData("y", ar_data, dims="obs_id")
        # The first coefficient will be the intercept term
        coefs = pm.Normal("coefs", priors["coefs"]["mu"], priors["coefs"]["sigma"])
        sigma = pm.HalfNormal("sigma", priors["sigma"])
        # We need one init variable for each lag, hence size is variable too
        init = pm.Normal.dist(
            priors["init"]["mu"], priors["init"]["sigma"], size=priors["init"]["size"]
        )
        # Steps of the AR model minus the lags required given specification
        ar1 = pm.AR(
            "ar",
            coefs,
            sigma=sigma,
            init_dist=init,
            constant=True,
            steps=t.shape[0] - (priors["coefs"]["size"] - 1),
            dims="obs_id",
        )

        ## Priors for the linear trend component
        alpha = pm.Normal("alpha", priors["alpha"]["mu"], priors["alpha"]["sigma"])
        beta = pm.Normal("beta", priors["beta"]["mu"], priors["beta"]["sigma"])
        trend = pm.Deterministic("trend", alpha + beta * t, dims="obs_id")

        ## Priors for seasonality
        beta_fourier = pm.Normal(
            "beta_fourier",
            mu=priors["beta_fourier"]["mu"],
            sigma=priors["beta_fourier"]["sigma"],
            dims="fourier_features",
        )
        fourier_terms = pm.MutableData("fourier_terms", ff)
        seasonality = pm.Deterministic(
            "seasonality", pm.math.dot(beta_fourier, fourier_terms), dims="obs_id"
        )

        mu = ar1 + trend + seasonality

        # The Likelihood
        outcome = pm.Normal("likelihood", mu=mu, sigma=sigma, observed=y, dims="obs_id")
        ## Sampling
        idata_ar = pm.sample_prior_predictive()
        if full_sample:
            idata_ar.extend(pm.sample(samples, random_seed=100, target_accept=0.95))
            idata_ar.extend(pm.sample_posterior_predictive(idata_ar))
        else:
            return idata_ar

    n = prediction_steps - ar_data.shape[0]
    n_order = 10
    periods = (ar_data.shape[0] + np.arange(n)) / 7

    fourier_features_new = pd.DataFrame(
        {
            f"{func}_order_{order}": getattr(np, func)(2 * np.pi * periods * order)
            for order in range(1, n_order + 1)
            for func in ("sin", "cos")
        }
    )

    with AR:
        AR.add_coords({"obs_id_fut_1": range(ar1_data.shape[0] - 1, prediction_steps, 1)})
        AR.add_coords({"obs_id_fut": range(ar1_data.shape[0], prediction_steps, 1)})
        t_fut = pm.MutableData(
            "t_fut", list(range(ar1_data.shape[0], prediction_steps, 1)), dims="obs_id_fut"
        )
        ff_fut = pm.MutableData("ff_fut", fourier_features_new.to_numpy().T)
        # condition on the learned values of the AR process
        # initialise the future AR process precisely at the last observed value in the AR process
        # using the special feature of the dirac delta distribution to be 0 probability everywhere else.
        ar1_fut = pm.AR(
            "ar1_fut",
            init_dist=pm.DiracDelta.dist(ar1[..., -1]),
            rho=coefs,
            sigma=sigma,
            constant=True,
            dims="obs_id_fut_1",
        )
        trend = pm.Deterministic("trend_fut", alpha + beta * t_fut, dims="obs_id_fut")
        seasonality = pm.Deterministic(
            "seasonality_fut", pm.math.dot(beta_fourier, ff_fut), dims="obs_id_fut"
        )
        mu = ar1_fut[1:] + trend + seasonality

        yhat_fut = pm.Normal("yhat_fut", mu=mu, sigma=sigma, dims="obs_id_fut")
        # use the updated values and predict outcomes and probabilities:
        idata_preds = pm.sample_posterior_predictive(
            idata_ar, var_names=["likelihood", "yhat_fut"], predictions=True, random_seed=743
        )

    return idata_ar, idata_preds, AR
In [ ]:
priors_0 = {
    "coefs": {"mu": [0.2, 0.2], "sigma": [0.5, 0.03], "size": 2},
    "alpha": {"mu": -4, "sigma": 0.1},
    "beta": {"mu": -0.1, "sigma": 0.2},
    "beta_fourier": {"mu": 0, "sigma": 2},
    "sigma": 8,
    "init": {"mu": -4, "sigma": 0.1, "size": 1},
}


idata_t_s, preds_t_s, model = make_latent_AR_trend_seasonal_model(y_t_s, fourier_features, priors_0)

Sampling: [alpha, ar, beta, beta_fourier, coefs, likelihood, sigma]
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [coefs, sigma, ar, alpha, beta, beta_fourier]
100.00% [12000/12000 00:43<00:00 Sampling 4 chains, 0 divergences] 
Sampling 4 chains for 1_000 tune and 2_000 draw iterations (4_000 + 8_000 draws total) took 70 seconds.
Sampling: [likelihood]
100.00% [8000/8000 00:01<00:00] 
Sampling: [ar1_fut, likelihood, yhat_fut]
100.00% [8000/8000 00:31<00:00]
In [ ]:
pm.model_to_graphviz(model)
Out[ ]:
clusterobs_id (200) obs_id (200) cluster2 2 clusterfourier_features (20) fourier_features (20) cluster20 x 200 20 x 200 clusterobs_id_fut (50) obs_id_fut (50) cluster20 x 50 20 x 50 clusterobs_id_fut_1 (51) obs_id_fut_1 (51) y y ~ MutableData ar ar ~ Deterministic likelihood likelihood ~ Normal ar->likelihood ar1_fut ar1_fut ~ Deterministic ar->ar1_fut seasonality seasonality ~ Deterministic seasonality->likelihood t t ~ MutableData t->ar trend trend ~ Deterministic t->trend trend->likelihood likelihood->y coefs coefs ~ Normal coefs->ar coefs->ar1_fut sigma sigma ~ HalfNormal sigma->ar sigma->likelihood yhat_fut yhat_fut ~ Normal sigma->yhat_fut sigma->ar1_fut alpha alpha ~ Normal alpha->trend trend_fut trend_fut ~ Deterministic alpha->trend_fut beta beta ~ Normal beta->trend beta->trend_fut beta_fourier beta_fourier ~ Normal beta_fourier->seasonality seasonality_fut seasonality_fut ~ Deterministic beta_fourier->seasonality_fut fourier_terms fourier_terms ~ MutableData fourier_terms->seasonality trend_fut->yhat_fut t_fut t_fut ~ MutableData t_fut->trend_fut seasonality_fut->yhat_fut ff_fut ff_fut ~ MutableData ff_fut->seasonality_fut ar1_fut->yhat_fut
In [ ]:
az.summary(idata_t_s, var_names=["alpha", "beta", "coefs", "beta_fourier"])
Out[ ]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
alpha -3.986 0.099 -4.172 -3.797 0.001 0.001 13106.0 6145.0 1.0
beta -0.182 0.012 -0.204 -0.159 0.000 0.000 3522.0 5420.0 1.0
coefs[0] 0.595 0.484 -0.287 1.536 0.009 0.006 3177.0 5203.0 1.0
coefs[1] 0.195 0.029 0.137 0.246 0.000 0.000 11284.0 6668.0 1.0
beta_fourier[0] 5.700 1.681 2.558 8.785 0.018 0.013 8413.0 5505.0 1.0
beta_fourier[1] 0.275 1.647 -2.913 3.280 0.018 0.018 8501.0 5712.0 1.0
beta_fourier[2] 0.188 1.641 -2.862 3.241 0.017 0.017 9114.0 5857.0 1.0
beta_fourier[3] -0.066 1.669 -3.209 3.026 0.018 0.018 8683.0 6748.0 1.0
beta_fourier[4] 0.172 1.667 -2.771 3.382 0.017 0.018 9684.0 6683.0 1.0
beta_fourier[5] -0.143 1.668 -3.285 2.910 0.017 0.018 9728.0 5846.0 1.0
beta_fourier[6] -0.176 1.624 -3.215 2.860 0.017 0.018 9020.0 5924.0 1.0
beta_fourier[7] -0.151 1.672 -3.222 3.015 0.017 0.018 9259.0 6636.0 1.0
beta_fourier[8] -0.174 1.659 -3.277 2.944 0.017 0.018 9347.0 6390.0 1.0
beta_fourier[9] -0.021 1.637 -3.198 2.878 0.018 0.017 8138.0 6374.0 1.0
beta_fourier[10] -5.691 1.683 -8.789 -2.487 0.018 0.013 8219.0 6287.0 1.0
beta_fourier[11] 0.298 1.652 -2.769 3.415 0.018 0.018 8453.0 5843.0 1.0
beta_fourier[12] 0.016 2.041 -3.673 3.914 0.016 0.026 15708.0 5857.0 1.0
beta_fourier[13] 5.603 1.333 3.018 7.990 0.023 0.016 3460.0 5277.0 1.0
beta_fourier[14] 5.691 1.650 2.570 8.754 0.018 0.013 8261.0 6341.0 1.0
beta_fourier[15] 0.289 1.634 -2.897 3.245 0.018 0.018 8464.0 6278.0 1.0
beta_fourier[16] 0.214 1.654 -2.989 3.200 0.017 0.017 9169.0 6719.0 1.0
beta_fourier[17] -0.031 1.663 -3.224 2.996 0.018 0.019 8464.0 5444.0 1.0
beta_fourier[18] 0.177 1.683 -2.978 3.316 0.018 0.020 9126.0 5907.0 1.0
beta_fourier[19] -0.146 1.669 -3.339 2.925 0.018 0.019 8771.0 6316.0 1.0
In [ ]:
plot_fits(idata_t_s, preds_t_s)

We can see here how the model fit again recovers the broad structure and trend of the data, but in addition we have captured the oscillation of the seasonal effect and projected that into the future.

Closing Remarks¶

The strength of a Bayesian model is largely the flexibility it offers for each modelling task. Hopefully this notebook gives a flavour of the variety of combinations worth considering when building a model to suit your use-case. We've seen how the Bayesian structural timeseries approach to forecasting can reveal the structure underlying our data, and be used to project that structure forward in time. We've seen how to encode different assumptions in the data generating model and calibrate our models against the observed data with posterior predictive checks.

Notably in the case of Auto-regressive modelling we've explicitly relied on the learned posterior distribution of the structural components. In this aspect we think the above is a kind of pure (neatly contained) example of Bayesian learning.

Authors¶

Adapted from Nathaniel Forde's Examined Algorithms Blog by Nathaniel Forde in Oct 2022.

Watermark¶

In [ ]:
%load_ext watermark
%watermark -n -u -v -iv -w -p pytensor,aeppl,xarray

Last updated: Wed Nov 09 2022

Python implementation: CPython
Python version       : 3.9.0
IPython version      : 8.4.0

pytensor: 2.8.7
aeppl : 0.0.38
xarray: 2022.10.0

arviz     : 0.12.1
pymc      : 4.2.2
numpy     : 1.23.4
pandas    : 1.5.1
sys       : 3.9.0 (default, Nov 15 2020, 06:25:35) 
[Clang 10.0.0 ]
matplotlib: 3.6.1

Watermark: 2.3.1

:::{include} ../page_footer.md :::