Exploring TensorFlow Probability STS Forecasting

Smoothing

We begin the analysis by studying and visualize various level of smoothness using a gaussian_filter. You can think of it as a weighted centered moving average based on a normal distribution. In particular, the level of smoothness is controlled by the parameter sigma which represents the standard deviation.

df_smooth = df_train.assign(
    y_smooth_1=lambda x: gaussian_filter(input=x["y"], sigma=3.5)
).assign(y_smooth_2=lambda x: gaussian_filter(input=x["y"], sigma=15))
fig, ax = plt.subplots()
sns.lineplot(x="date", y="y", label="y_train", data=df_smooth, alpha=0.5, ax=ax)
sns.lineplot(
    x="date", y="y_smooth_1", label="y_train_smooth monthly", data=df_smooth, ax=ax
)
sns.lineplot(
    x="date", y="y_smooth_2", label="y_train_smooth yearly", data=df_smooth, ax=ax
)
ax.legend(loc="upper left")
ax.set(title="Dependent Variable - Smoothing", ylabel="y")

From this plot we see two clear seasonalities: yearly and monthly. The variable y_train_smooth yearly also includes a positive trend component, which does not look linear.

To understand how to model the monthly seasonality we plot the mean of y over the day of the month.

fig, ax = plt.subplots()
df_train.groupby("day_of_month").agg({"y": np.mean}).plot(ax=ax)
ax.set(title="Day of the month mean")

This shows a good first approximation is to model this component via a cyclic variable of the form \(z\longmapsto \sin(z)\).

Remove yearly and monthly seasonality

Let us remove these two seasonal components:

# Remove yearly seasonality.
y_no_year_season = df_smooth["y"] - df_smooth["y_smooth_2"]
# Remove monthly seasonality.
y_no_year_month_season = y_no_year_season - gaussian_filter(
    input=y_no_year_season, sigma=3.5
)

# Plot components.
fig, ax = plt.subplots(
    nrows=3, ncols=1, figsize=(12, 16), sharex=True, sharey=False, layout="constrained"
)
sns.lineplot(x="date", y="y", label="y_train", data=df_smooth, alpha=0.5, ax=ax[0])
sns.lineplot(
    x="date", y="y_smooth_1", label="y_train_smooth monthly", data=df_smooth, ax=ax[0]
)
sns.lineplot(
    x="date", y="y_smooth_2", label="y_train_smooth yearly", data=df_smooth, ax=ax[0]
)
ax[0].set(title="y")

ax[1].plot(df_smooth["date"], y_no_year_season, c="C2")
ax[1].axhline(y_no_year_season.mean(), color="C3", linestyle="--", label="mean")
ax[1].legend()
ax[1].set(title="y no year seasonality")

ax[2].plot(df_smooth["date"], y_no_year_month_season, c="C1")
ax[2].axhline(y_no_year_month_season.mean(), color="C3", linestyle="--", label="mean")
ax[2].legend()
ax[2].set(title="y no year & month seasonality")

AC and PAC

Let us now compute the autocorrelation and partial-autocorrelation of the remainder component.

fig, ax = plt.subplots(2, 1)
_ = plot_acf(x=y_no_year_month_season, ax=ax[0])
_ = plot_pacf(x=y_no_year_month_season, ax=ax[1])

From the autocorrelation plot we observe a clear seasonality component at lag 7. Moreover, looking into the shape of the plot we can also see that the encoding is cyclical.

Correlations

We want to see whether there is a linear effect of x on y. A first good indication can be obtained by looking into correlations. Naively, one would simply compute:

np.corrcoef(df_train["y"], df_train["x"])[0, 1]

0.34364604482253697

However, this correlation does not reflect the potential relation of x on y since we have a clear positive trend and seasonality components. A more meaningful indication is the correlation:

np.corrcoef(y_no_year_month_season, df_train["x"])[0, 1]

0.6635134779159232

Note that we still have not removed the 7-day (weekly) seasonality. This of course has an effect on this correlation. To see this let us consider the scatter plot:

fig, ax = plt.subplots()
sns.scatterplot(
    x=y_no_year_month_season,
    y=df_train["x"],
    hue=df["day_of_week"],
    palette="viridis",
    ax=ax,
)

ax.set(title="y_no_year_month_season vs x", xlabel="x", ylabel="y_no_year_month_season")

We indeed see we have clear clusters on each level of x corresponding to each day of the week.

One way of removing this seasonality is via a linear regression model. We use a one-hot encoding for the day of the week variable.

Remark: As we have seen before a cyclical encoding might be a better choice. Nevertheless, given how the data was generated, we want to test the one-hot encoding approach.

# Prepare model data frame.
dow_df = df_train[["day_of_week"]].copy()
dow_df["y_no_year_month_season"] = y_no_year_month_season
# One-hot encoding of the day of the week.
dow_dummies = pd.get_dummies(dow_df["day_of_week"], drop_first=True)
dow_dummies.columns = ["d" + str(i) for i in dow_dummies.columns]
dow_df = pd.concat([dow_df, dow_dummies], axis=1)
dow_df.head()

Next, we fit the linear model.

dow_mod = smf.ols(
    formula="y_no_year_month_season ~  d1 + d2 + d3 + d4 + d5 + d6", data=dow_df
)

dow_res = dow_mod.fit()

print(dow_res.summary())

                              OLS Regression Results                              
==================================================================================
Dep. Variable:     y_no_year_month_season   R-squared:                       0.425
Model:                                OLS   Adj. R-squared:                  0.423
Method:                     Least Squares   F-statistic:                     156.6
Date:                    Thu, 09 Nov 2023   Prob (F-statistic):          6.44e-149
Time:                            11:05:19   Log-Likelihood:                -2103.6
No. Observations:                    1277   AIC:                             4221.
Df Residuals:                        1270   BIC:                             4257.
Df Model:                               6                                         
Covariance Type:                nonrobust                                         
==============================================================================
                 coef    std err          t      P>|t|      [0.025      0.975]
------------------------------------------------------------------------------
Intercept      1.4894      0.093     15.991      0.000       1.307       1.672
d1[T.True]    -0.5122      0.132     -3.888      0.000      -0.771      -0.254
d2[T.True]    -1.5880      0.132    -12.039      0.000      -1.847      -1.329
d3[T.True]    -2.9809      0.132    -22.599      0.000      -3.240      -2.722
d4[T.True]    -2.9059      0.132    -22.030      0.000      -3.165      -2.647
d5[T.True]    -1.7887      0.132    -13.561      0.000      -2.047      -1.530
d6[T.True]    -0.6685      0.132     -5.075      0.000      -0.927      -0.410
==============================================================================
Omnibus:                      752.611   Durbin-Watson:                   2.159
Prob(Omnibus):                  0.000   Jarque-Bera (JB):             6163.206
Skew:                           2.694   Prob(JB):                         0.00
Kurtosis:                      12.317   Cond. No.                         7.86
==============================================================================

Notes:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

We now compute the correlation of the regressor x with the model residuals:

np.corrcoef(dow_res.resid, df_train["x"])[0, 1]

0.8502148605418938

The correlation increased significantly.

Regressor Effect

A common question in time series analysis is to estimate the effect (ROI) of a specific external regressor (e.g. what is the ROI of media spend on sales?) For this specific case, given the correlation above, we can estimate this for x running a linear model on the residuals of the dow_model.

x_df = df_train[["x"]].copy()
x_df["dow_model_resid"] = dow_res.resid

x_mod = smf.ols(formula="dow_model_resid ~ x", data=x_df)

x_res = x_mod.fit()

print(x_res.summary())

                            OLS Regression Results                            
==============================================================================
Dep. Variable:        dow_model_resid   R-squared:                       0.723
Model:                            OLS   Adj. R-squared:                  0.723
Method:                 Least Squares   F-statistic:                     3326.
Date:                Thu, 09 Nov 2023   Prob (F-statistic):               0.00
Time:                        11:05:19   Log-Likelihood:                -1284.2
No. Observations:                1277   AIC:                             2572.
Df Residuals:                    1275   BIC:                             2583.
Df Model:                           1                                         
Covariance Type:            nonrobust                                         
==============================================================================
                 coef    std err          t      P>|t|      [0.025      0.975]
------------------------------------------------------------------------------
Intercept     -0.2309      0.019    -12.182      0.000      -0.268      -0.194
x              5.3018      0.092     57.668      0.000       5.121       5.482
==============================================================================
Omnibus:                        9.800   Durbin-Watson:                   1.280
Prob(Omnibus):                  0.007   Jarque-Bera (JB):                9.950
Skew:                          -0.215   Prob(JB):                      0.00691
Kurtosis:                       2.960   Cond. No.                         4.97
==============================================================================

Notes:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

The estimates effect is then:

x_res.params["x"]

5.301793890986457

Recall that the true effect is 5.0.

Periodogram

There is another way to extract seasonality patterns: using a periodogram to estimate the spectral density of a signal.

f, Pxx_den = signal.periodogram(x=df_train["y"], detrend="linear", nfft=int(7e2))

fig, ax = plt.subplots()
sns.lineplot(x=f, y=Pxx_den, ax=ax)
ax.set(title="Power Spectral Density")

Each of these peaks represent the seasonal components. To compute the associated frequency wee need to compute their multiplicative inverse.

# Sort to get the peak values.
sort_freq_index = np.argsort(a=Pxx_den)[::-1]

periodogram_df = pd.DataFrame(
    {"sort_freq": f[sort_freq_index], "Pxx_den": Pxx_den[sort_freq_index]}
)

periodogram_df.assign(days=lambda x: 1 / x["sort_freq"]).head(5)

The first three values correspond to days = 30, 350 and 7. Which correspond to monthly, yearly and weekly seasonality.

Time Series Decomposition

statsmodels has an inbuilt decomposition function using moving averages. Let us use it to estimate the effect of the external regressor (which we already know it should the residual component).

decomposition_obj = seasonal_decompose(
    x=df_train[["date", "y"]].set_index("date"), model="additive"
)

fig, ax = plt.subplots(4, 1, figsize=(12, 12))

decomposition_obj.observed.plot(ax=ax[0])
ax[0].set(title="observed")
decomposition_obj.trend.plot(ax=ax[1])
ax[1].set(title="trend")
decomposition_obj.seasonal.plot(ax=ax[2])
ax[2].set(title="seasonal")
decomposition_obj.resid.plot(ax=ax[3])
ax[3].set(title="residual")
plt.tight_layout()

Note that the trend component also includes the monthly and yearly seasonality. This might not always be desired. You can of course use the smoothing technique presented above or keep decomposing the resulting trend by specifying the period parameter in the seasonal_decompose function.

• Correlation

# We remove the first and last 3 entries as they are np.nan coming
# from the moving average method.
np.corrcoef(decomposition_obj.resid[3:-3], df_train["x"][3:-3])[0, 1]

0.8533497914653861

• Effect Estimation

x_df["decomposition_resid"] = decomposition_obj.resid.values

x_mod2 = smf.ols(formula="decomposition_resid ~ x", data=x_df[3:-3])

x_res2 = x_mod2.fit()

print(x_res2.summary())

                             OLS Regression Results                            
===============================================================================
Dep. Variable:     decomposition_resid   R-squared:                       0.728
Model:                             OLS   Adj. R-squared:                  0.728
Method:                  Least Squares   F-statistic:                     3400.
Date:                 Thu, 09 Nov 2023   Prob (F-statistic):               0.00
Time:                         11:05:23   Log-Likelihood:                -1242.5
No. Observations:                 1271   AIC:                             2489.
Df Residuals:                     1269   BIC:                             2499.
Df Model:                            1                                         
Covariance Type:             nonrobust                                         
==============================================================================
                 coef    std err          t      P>|t|      [0.025      0.975]
------------------------------------------------------------------------------
Intercept     -0.2265      0.018    -12.259      0.000      -0.263      -0.190
x              5.2128      0.089     58.309      0.000       5.037       5.388
==============================================================================
Omnibus:                       37.931   Durbin-Watson:                   1.562
Prob(Omnibus):                  0.000   Jarque-Bera (JB):               40.589
Skew:                          -0.432   Prob(JB):                     1.54e-09
Kurtosis:                       3.139   Cond. No.                         4.96
==============================================================================

Notes:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

We get similar results as above.

Remarks:

• In real life applications the effects of external regressors might include lags. A correlation analysis with lags needs to be done before modeling.

• Also, linear relations might serve as a good first approximation. Nevertheless, going into non-linear relations might improve the model significantly. Two concrete methods which appear often is to model via adstock effect and generalized additive models.

Variational Inference (Short Intro)

We follow the strategy of the TensorFlow Probability use cases and fit the model using variational inference. I find the article Variational Inference: A Review for Statisticians very good for an introduction to the subject. The main idea of variational inference is to use optimization methods to minimize the Kullback-Leibler divergence in order to approximate a conditional density within a family of specified parameter densities \(\mathfrak{D}\). Specifically, assume we are given observed variables \(x\) (i.e. data) and let \(z\) be a set of latent variables with joint density \(p(x, z)\). We are interested in computing the conditional density

\[ p(z|x) = \frac{p(x, z)}{p(x)}, \quad\text{where} \quad p(x)=\int p(x, z) dz. \]

This last integral (known as the evidence) is in general very hard to compute. The main idea of variational inference is to solve the optimization problem

\[ q^{*}(z)=\min_{q \in \mathfrak{D}} \text{KL}(q(z)||p(z|x)), \]

where

\[ \begin{align} \text{KL}(q(z)||p(z|x)) =& \text{E}[\log(q(z))] - \text{E}[\log(p(z|x))] \\ =& \text{E}[\log(q(z))] - \text{E}[\log(p(z,x))] - \text{E}[\log(p(x))] \\ =& \text{E}[\log(q(z))] - \text{E}[\log(p(z,x))] - \log(p(x) \end{align} \]

is the Kullback-Leibler (KL) divergence (the expected values are taken with respect to \(q(z)\)). Note that this quantity contains a term \(\text{E}[\log(p(x))]\), which is hard to compute. Because we cannot compute the KL, we optimize an alternative objective that is equivalent to the KL up to an added constant ²,

\[ \text{ELBO}(q) = \text{E}[\log(p(z, x))] - \text{E}[\log(q(z))] \]

ELBO stands for evidence lower bound.

Remark: Note that

\[ \text{ELBO}(q)= -\text{KL}(q(z)||p(z|x)) + \log(p(x)) \]

Moreover, it is easy to see that

\[ \text{ELBO}(q)= \text{E}[\log(p(x|z)]-\text{KL}(q(z)||p(x)) \]

The first term is the expected likelihood and the second term is the negative divergence between the variational density and the prior.

For more details and enlightening comments, please refer to the article mentioned above (from which this short introduction was taken from).

Variational Inference in TensorFlow Probability

First we build the variational surrogate posteriors, which consist of independent normal distributions with two trainable hyper-parameters loc and scale for each parameter in the model. That is, \(\mathfrak{D}\) above consists of normal distributions.

variational_posteriors = tfp.sts.build_factored_surrogate_posterior(
    model=toy_model, seed=seed
)

Let us sample from the prior distributions and plot the corresponding densities:

q_prior_samples = variational_posteriors.sample(1000)

num_parameters = len(toy_model.parameters)

fig, ax = plt.subplots(
    nrows=num_parameters,
    ncols=1,
    figsize=(12, 21),
    sharex=False,
    sharey=False,
    layout="constrained",
)

for i, param in enumerate(toy_model.parameters):
    param_mean = np.mean(q_prior_samples[param.name], axis=0)
    param_median = np.median(q_prior_samples[param.name], axis=0)
    param_std = np.std(q_prior_samples[param.name], axis=0)

    sns.histplot(x=q_prior_samples[param.name].numpy().flatten(), kde=True, ax=ax[i])
    ax[i].set(title=param.name)
    ax[i].axvline(x=param_median, color="C1", linestyle="--", label="median")
    ax[i].axvline(x=param_mean, color="C2", linestyle="--", label="mean")
    ax[i].axvline(
        x=param_mean + 2 * param_std,
        color="C3",
        linestyle="--",
        label=r"$\mu \pm 2\sigma$",
    )
    ax[i].axvline(x=param_mean - param_std, color="C3", linestyle="--")
    ax[i].legend()

fig.suptitle("PRIOR DISTRIBUTIONS", fontweight="bold", fontsize=16, y=1.02)

Next we run the optimization procedure.

num_variational_steps = 500

# Set optimizer.
optimizer = tf.optimizers.Adam(learning_rate=0.1)


# Using fit_surrogate_posterior to build and optimize
# the variational loss function.
@tf.function(experimental_compile=True)
def train():
    # Build the joint density.
    target_log_prob_fn = toy_model.joint_distribution(
        observed_time_series=df_train["y"]
    ).log_prob

    return tfp.vi.fit_surrogate_posterior(
        target_log_prob_fn=target_log_prob_fn,
        surrogate_posterior=variational_posteriors,
        optimizer=optimizer,
        num_steps=num_variational_steps,
        seed=seed,
    )


# Run optimization.
elbo_loss_curve = train()

Let us plot the elbo_loss_curve.

fig, ax = plt.subplots()
ax.plot(elbo_loss_curve, marker=".")
ax.set(title="ELBO Loss Curve", xlabel="iteration")

We see that a minimum has been reached.

Now we sample from the variational posteriors obtained:

q_samples = variational_posteriors.sample(2_000)

Let us plot the parameter posterior distributions:

num_parameters = len(toy_model.parameters)

fig, ax = plt.subplots(
    nrows=num_parameters,
    ncols=1,
    figsize=(12, 21),
    sharex=False,
    sharey=False,
    layout="constrained",
)

for i, param in enumerate(toy_model.parameters):
    param_mean = np.mean(q_samples[param.name], axis=0)
    param_median = np.median(q_samples[param.name], axis=0)
    param_std = np.std(q_samples[param.name], axis=0)

    sns.histplot(x=q_samples[param.name].numpy().flatten(), kde=True, ax=ax[i])
    ax[i].set(title=param.name)
    ax[i].axvline(x=param_median, color="C1", linestyle="--", label="median")
    ax[i].axvline(x=param_mean, color="C2", linestyle="--", label="mean")
    ax[i].axvline(
        x=param_mean + 2 * param_std,
        color="C3",
        linestyle="--",
        label=r"$\mu \pm 2\sigma$",
    )
    ax[i].axvline(x=param_mean - param_std, color="C3", linestyle="--")
    ax[i].legend()

fig.suptitle("POSTERIOR DISTRIBUTIONS", fontweight="bold", fontsize=16, y=1.02)

# Get mean and std for each parameter.
print("Inferred parameters:")
for param in toy_model.parameters:
    print(
        "{}: {} +- {}".format(
            param.name,
            np.mean(q_samples[param.name], axis=0),
            np.std(q_samples[param.name], axis=0),
        )
    )

Inferred parameters:
observation_noise_scale: 0.484006404876709 +- 0.010281678289175034
local_linear_trend/_level_scale: 0.019886409863829613 +- 0.012175584211945534
local_linear_trend/_slope_scale: 0.0008743585785850883 +- 0.0002825716510415077
month_of_year/_drift_scale: 0.014278052374720573 +- 0.04950084537267685
day_of_week/_drift_scale: 0.004724537953734398 +- 0.0029768438544124365
day_of_month/_weights: [1.9374084] +- [0.02063216]
x_var/_weights: [6.03262] +- [0.0631116]

Remark: Observe that the estimated effect of the regressor x is around 6.