Periodic effects and broadcasting

A common practice in time series forecasting is featuring periodic effects such as day of the week or weekend effects. The arrayutils.PeriodicBroadcaster class provides this functionality. This tutorial shows how to use this to model both repeating elements of a sequence and repeating a sequence as a whole (tiling).

Repeated units

The RtPeriodicDiff and RtWeeklyDiff classes use PeriodicBroadcaster to repeat each vector element for \(\mathcal{R}(t)\) values held constant within a period.

import jax.numpy as jnp
import numpy as np
import numpyro
from pyrenew import process, deterministic

# The random process for Rt
rt_proc = process.RtWeeklyDiffARProcess(
    name="rt_weekly_diff",
    offset=0,
    log_rt_rv=deterministic.DeterministicVariable(
        name="log_rt", value=jnp.array([0.1, 0.2])
    ),
    autoreg_rv=deterministic.DeterministicVariable(
        name="autoreg", value=jnp.array([0.7])
    ),
    periodic_diff_sd_rv=deterministic.DeterministicVariable(
        name="periodic_diff_sd", value=jnp.array([0.1])
    ),
)

with numpyro.handlers.seed(rng_seed=20):
    sim_data = rt_proc(duration=30)

# Plotting the Rt values
import matplotlib.pyplot as plt

plt.step(np.arange(len(sim_data)), sim_data, where="post")
plt.xlabel("Days")
plt.ylabel("Rt")
plt.title("Simulated Rt values")

# Adding bands to mark weeks
for i in range(0, 30, 7):
    plt.axvline(i, color="black", linestyle="--", alpha=0.5)
plt.show()

The implementation of the RtWeeklyDiffARProcess (which is an instance of RtPeriodicDiffARProcess), uses repeat_until_n to repeat values: repeat_until_n(..., period_size=7). The RtWeeklyDiff class is a particular case of RtPeriodicDiff with a period size of seven.

Repeated sequences (tiling)

By specifying broadcast_type = "tile", the PeriodicBroadcaster repeats the sequence as a whole. The DayOfWeekEffect class is a particular case of PeriodicEffect with a period size of seven. We sample from a scaled Dirchlet distribution such that the sum of the samples is 7:

import numpyro.distributions as dist
from pyrenew import transformation, randomvariable

# Building the transformed prior: Dirichlet * 7
mysimplex = dist.TransformedDistribution(
    dist.Dirichlet(concentration=jnp.ones(7)),
    transformation.AffineTransform(loc=0, scale=7.0),
)

# Constructing the day of week effect
dayofweek = process.DayOfWeekEffect(
    offset=0,
    quantity_to_broadcast=randomvariable.DistributionalVariable(
        name="simp", distribution=mysimplex
    ),
)

We use the sample method to generate samples from the day of week effect:

with numpyro.handlers.seed(rng_seed=20):
    sim_data = dayofweek(duration=30)

# Plotting the effect values
import matplotlib.pyplot as plt

plt.step(np.arange(len(sim_data)), sim_data, where="post")
plt.xlabel("Days")
plt.ylabel("Effect size")
plt.title("Simulated Day of Week Effect values")

# Adding bands to mark weeks
for i in range(0, 30, 7):
    plt.axvline(i, color="black", linestyle="--", alpha=0.5)
plt.show()