Implementing a Day of the Week Effect

This document illustrates how to leverage the time-aware arrays to create a a day of the week effect. We use the same model designed in the hospital admissions-only tutorial.

Recap: Hospital Admissions Model

In the “Fitting a hospital admissions-only model” tutorial, we built a fairly complicated model that included pre-defined random variables as well as a custom random variable representing the reproductive number. In this tutorial, we will focus on adding a new component: the day-of-the-week effect. We start by reproducing the model without the day-of-the-week effect:

1. We load the data:

# Setup
import numpyro
import polars as pl
from pyrenew import datasets

# Setting the number of devices
numpyro.set_host_device_count(2)

# Loading and processing the data
dat = (
    datasets.load_wastewater()
    .group_by("date")
    .first()
    .select(["date", "daily_hosp_admits"])
    .sort("date")
    .head(90)
)

daily_hosp_admits = dat["daily_hosp_admits"].to_numpy()
dates = dat["date"].to_numpy()

# Loading additional datasets
gen_int = datasets.load_generation_interval()
inf_hosp_int = datasets.load_infection_admission_interval()

# We only need the probability_mass column of each dataset
gen_int_array = gen_int["probability_mass"].to_numpy()
gen_int = gen_int_array
inf_hosp_int_array = inf_hosp_int["probability_mass"].to_numpy()

2. Next, we defined the model’s components:

from pyrenew import latent, deterministic, randomvariable
import jax.numpy as jnp
import numpyro.distributions as dist

inf_hosp_int = deterministic.DeterministicPMF(
    name="inf_hosp_int", value=inf_hosp_int_array
)

hosp_rate = randomvariable.DistributionalVariable(
    name="IHR", distribution=dist.LogNormal(jnp.log(0.05), jnp.log(1.1))
)

latent_hosp = latent.HospitalAdmissions(
    infection_to_admission_interval_rv=inf_hosp_int,
    infection_hospitalization_ratio_rv=hosp_rate,
)

from pyrenew import model, process, observation, metaclass, transformation
from pyrenew.latent import (
    InfectionInitializationProcess,
    InitializeInfectionsExponentialGrowth,
)


# Infection process
latent_inf = latent.Infections()
n_initialization_points = max(gen_int_array.size, inf_hosp_int_array.size) - 1

I0 = InfectionInitializationProcess(
    "I0_initialization",
    randomvariable.DistributionalVariable(
        name="I0",
        distribution=dist.LogNormal(loc=jnp.log(100), scale=jnp.log(1.75)),
    ),
    InitializeInfectionsExponentialGrowth(
        n_initialization_points,
        deterministic.DeterministicVariable(name="rate", value=0.05),
    ),
)

# Generation interval and Rt
gen_int = deterministic.DeterministicPMF(name="gen_int", value=gen_int)

including the Rt effect:

class MyRt(metaclass.RandomVariable):
    def __init__(self, sd_rv):
        self.sd_rv = sd_rv

    def validate(self):
        pass

    def sample(self, n: int, **kwargs) -> tuple:
        # Standard deviation of the random walk
        sd_rt = self.sd_rv()

        # Random walk step
        step_rv = randomvariable.DistributionalVariable(
            name="rw_step_rv", distribution=dist.Normal(0, sd_rt)
        )

        rt_init_rv = randomvariable.DistributionalVariable(
            name="init_log_rt", distribution=dist.Normal(0, 0.2)
        )

        # Random walk process
        base_rv = process.RandomWalk(
            name="log_rt",
            step_rv=step_rv,
        )

        # Transforming the random walk to the Rt scale
        rt_rv = randomvariable.TransformedVariable(
            name="Rt_rv",
            base_rv=base_rv,
            transforms=transformation.ExpTransform(),
        )
        init_rt = rt_init_rv.sample()

        return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)


rtproc = MyRt(
    randomvariable.DistributionalVariable(
        name="Rt_random_walk_sd", distribution=dist.HalfNormal(0.025)
    )
)

3. We defined the observation model:

# we place a log-Normal prior on the concentration
# parameter of the negative binomial.
nb_conc_rv = randomvariable.TransformedVariable(
    "concentration",
    randomvariable.DistributionalVariable(
        name="concentration_raw",
        distribution=dist.TruncatedNormal(loc=0, scale=1, low=0.01),
    ),
    transformation.PowerTransform(-2),
)

# now we define the observation process
obs = observation.NegativeBinomialObservation(
    "negbinom_rv",
    concentration_rv=nb_conc_rv,
)

4. And finally, we build the model:

hosp_model = model.HospitalAdmissionsModel(
    latent_infections_rv=latent_inf,
    latent_hosp_admissions_rv=latent_hosp,
    I0_rv=I0,
    gen_int_rv=gen_int,
    Rt_process_rv=rtproc,
    hosp_admission_obs_process_rv=obs,
)

Here is what the model looks like without the day-of-the-week effect:

import jax
import numpy as np

# Model without weekday effect
hosp_model.run(
    num_samples=2000,
    num_warmup=2000,
    data_observed_hosp_admissions=daily_hosp_admits,
    rng_key=jax.random.key(54),
    mcmc_args=dict(progress_bar=False),
)

# Plotting the posterior
out = hosp_model.plot_posterior(
    var="latent_hospital_admissions",
    ylab="Hospital Admissions",
    obs_signal=daily_hosp_admits.astype(float),
)

Round 2: Incorporating day-of-the-week effects

We will re-use the infection to admission interval and infection to hospitalization rate from the previous model. But we will also add a day-of-the-week effect. To do this, we will add two additional arguments to the latent hospital admissions random variable: day_of_the_week_rv (a RandomVariable) and obs_data_first_day_of_the_week (an int mapping days of the week from 0:6, zero being Monday). The day_of_the_week_rv’s sample method should return a vector of length seven; those values are then broadcasted to match the length of the dataset. Moreover, since the observed data may start in a weekday other than Monday, the obs_data_first_day_of_the_week argument is used to offset the day-of-the-week effect.

For this example, the effect will be passed as a scaled Dirichlet distribution. It will consist of a TransformedVariable that samples an array of length seven from numpyro’s distributions.Dirichlet and applies a transformation.AffineTransform to scale it by seven.

[1]:

# Instantiating the day-of-the-week effect
dayofweek_effect = randomvariable.TransformedVariable(
    name="dayofweek_effect",
    base_rv=randomvariable.DistributionalVariable(
        name="dayofweek_effect_raw",
        distribution=dist.Dirichlet(jnp.ones(7)),
    ),
    transforms=transformation.AffineTransform(
        loc=0, scale=7, domain=jnp.array([0, 1])
    ),
)

Now, by re-defining the latent hospital admissions random variable with the day-of-the-week effect, we can build a model that includes this effect. Since the day-of-the-week effect takes into account the first day of the dataset, we need to determine the day of the week of the first observation. We can do this by converting the first date in the dataset to a datetime object and extracting the day of the week:

# Figuring out the day of the week of the first observation
import datetime as dt

first_dow_in_data = dates[0].astype(dt.datetime).weekday()
first_dow_in_data  # zero

# Re-defining the latent hospital admissions RV, now with the
# day-of-the-week effect
latent_hosp_wday_effect = latent.HospitalAdmissions(
    infection_to_admission_interval_rv=inf_hosp_int,
    infection_hospitalization_ratio_rv=hosp_rate,
    day_of_week_effect_rv=dayofweek_effect,
    # Concidirently, this is zero
    obs_data_first_day_of_the_week=first_dow_in_data,
)

# New model with day-of-the-week effect
hosp_model_dow = model.HospitalAdmissionsModel(
    latent_infections_rv=latent_inf,
    latent_hosp_admissions_rv=latent_hosp_wday_effect,
    I0_rv=I0,
    gen_int_rv=gen_int,
    Rt_process_rv=rtproc,
    hosp_admission_obs_process_rv=obs,
)

Running the model:

# Model with weekday effect
hosp_model_dow.run(
    num_samples=2000,
    num_warmup=2000,
    data_observed_hosp_admissions=daily_hosp_admits,
    rng_key=jax.random.key(54),
    mcmc_args=dict(progress_bar=False),
)

As a result, we can see the posterior distribution of our novel day-of-the-week effect:

out = hosp_model_dow.plot_posterior(
    var="dayofweek_effect_raw", ylab="Day of the Week Effect", samples=500
)

sp = hosp_model_dow.spread_draws(["dayofweek_effect_raw"])
# dayofweek_effect is not recorded

The new model with the day-of-the-week effect can be compared to the previous model without the effect. Finally, let’s reproduce the figure without the day-of-the-week effect, and then plot the new model with the effect:

# Figure without weekday effect
out = hosp_model.plot_posterior(
    var="latent_hospital_admissions",
    ylab="Hospital Admissions",
    obs_signal=daily_hosp_admits.astype(float),
)

# Figure with weekday effect
out_dow = hosp_model_dow.plot_posterior(
    var="latent_hospital_admissions",
    ylab="Hospital Admissions",
    obs_signal=daily_hosp_admits.astype(float),
)

[1]

A similar weekday effect is implemented in its own module, with example code here.