Vaxflux Model Demo¶

This notebook extends the season-only demo by adding age group covariates. The structure is the same, but now the model learns how vaccination uptake differs across age groups (youth, adult, elderly) in addition to varying across seasons. It's a natural next step once you're comfortable with the basics.

In [1]:

Copied!





from pprint import pprint

import arviz as az
import matplotlib.pyplot as plt
import numpyro
import pandas as pd

from vaxflux import (
    CovariateCategories,
    LogisticCurve,
    PartiallyPooledGaussianCovariate,
    SeasonRange,
    SeasonVaryingPartiallyPooledGaussianCovariate,
    VaxfluxModel,
    VaxfluxObservations,
    daily_date_ranges,
)
from vaxflux.data import sample_dataset

numpyro.set_host_device_count(8)
az.style.use("arviz-docgrid")
from pprint import pprint

import arviz as az
import matplotlib.pyplot as plt
import numpyro
import pandas as pd

from vaxflux import (
    CovariateCategories,
    LogisticCurve,
    PartiallyPooledGaussianCovariate,
    SeasonRange,
    SeasonVaryingPartiallyPooledGaussianCovariate,
    VaxfluxModel,
    VaxfluxObservations,
    daily_date_ranges,
)
from vaxflux.data import sample_dataset

numpyro.set_host_device_count(8)
az.style.use("arviz-docgrid")

Instantiate the Logistic Curve¶

LogisticCurve is the functional form vaxflux uses to model vaccination uptake over time within a season. Displaying it here confirms that it loaded correctly and shows its parameter names.

In [2]:

Copied!

curve = LogisticCurve()
curve
curve = LogisticCurve()
curve

Out[2]:

$$f(t \mid m, r, s) = \operatorname{expit}(m)\operatorname{expit}\left(e^r(t-s)\right)$$

Define Seasons and Observation Dates¶

Each SeasonRange specifies the name and calendar window for a flu season. daily_date_ranges then generates weekly observation intervals (7-day windows, i.e. range_days=6) across all seasons — these become the time points at which the model expects observations.

In [3]:

Copied!





seasons = [
    SeasonRange(season="2022/23", start_date="2022-10-03", end_date="2023-02-05"),
    SeasonRange(season="2023/24", start_date="2023-10-02", end_date="2024-02-04"),
    SeasonRange(season="2024/25", start_date="2024-10-07", end_date="2025-02-02"),
]
dates = daily_date_ranges(seasons, range_days=6)
seasons = [
    SeasonRange(season="2022/23", start_date="2022-10-03", end_date="2023-02-05"),
    SeasonRange(season="2023/24", start_date="2023-10-02", end_date="2024-02-04"),
    SeasonRange(season="2024/25", start_date="2024-10-07", end_date="2025-02-02"),
]
dates = daily_date_ranges(seasons, range_days=6)

Define Age Covariate Categories¶

CovariateCategories declares the demographic groupings the model will track. Here we have three age groups: youth, adult, and elderly. These labels will later be used to assign age-specific parameter offsets.

In [4]:

Copied!





age_covariate_categories = CovariateCategories(
    covariate="age",
    categories=["youth", "adult", "elderly"],
)
age_covariate_categories = CovariateCategories(
    covariate="age",
    categories=["youth", "adult", "elderly"],
)

Define Prior Covariates¶

PartiallyPooledGaussianCovariate defines a hierarchical (partially pooled) prior for a logistic curve parameter. Each season gets its own value drawn from a shared population distribution, whose mean (mu_mu, mu_sigma) and spread (sigma) are specified here.

SeasonVaryingPartiallyPooledGaussianCovariate extends this to allow the age-group offset to vary across seasons — so the model can learn that, say, the elderly advantage in uptake timing differs between 2022/23 and 2024/25.

The three parameters of the logistic curve are:

m: the midpoint (location of the uptake peak on the logistic scale)
r: the rate (steepness of the curve)
s: the shift (horizontal offset in days from season start)

In [5]:

Copied!





covariates = [
    PartiallyPooledGaussianCovariate(
        parameter="m",
        covariate=None,
        mu_mu=0.0,
        mu_sigma=0.75,
        sigma=0.75,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="r",
        covariate=None,
        mu_mu=-2.5,
        mu_sigma=1.5,
        sigma=1.5,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="s",
        covariate=None,
        mu_mu=35.0,
        mu_sigma=20.0,
        sigma=20.0,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="m",
        covariate="age",
        mu_mu=1.0,
        mu_sigma=1.0,
        sigma=1.0,
    ),
]
covariates = [
    PartiallyPooledGaussianCovariate(
        parameter="m",
        covariate=None,
        mu_mu=0.0,
        mu_sigma=0.75,
        sigma=0.75,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="r",
        covariate=None,
        mu_mu=-2.5,
        mu_sigma=1.5,
        sigma=1.5,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="s",
        covariate=None,
        mu_mu=35.0,
        mu_sigma=20.0,
        sigma=20.0,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="m",
        covariate="age",
        mu_mu=1.0,
        mu_sigma=1.0,
        sigma=1.0,
    ),
]

Specify True Parameters for Synthetic Data¶

This defines the ground-truth logistic curve parameters used to generate the synthetic observations. Each age group within each season gets its own m value (via an additive offset), while r and s are shared within a season. These parameter values are what the model should recover during inference.

In [6]:

Copied!





params = []

params_grid = {
    "2022/23": {
        "m": -0.5,
        "r": -3.2,
        "s": 40.0,
        "adult": {
            "m": 1.0,
        },
        "elderly": {
            "m": 1.8,
        },
    },
    "2023/24": {
        "m": -0.6,
        "r": -3.1,
        "s": 42.5,
        "adult": {
            "m": 1.05,
        },
        "elderly": {
            "m": 1.78,
        },
    },
    "2024/25": {
        "m": -0.4,
        "r": -3.3,
        "s": 41.0,
        "adult": {
            "m": 1.03,
        },
        "elderly": {
            "m": 1.81,
        },
    },
}

for curve_param in curve.parameters:
    for season_range in seasons:
        for age_category in age_covariate_categories.categories:
            season_params = params_grid[season_range.season]
            param_value = season_params[curve_param] + season_params.get(
                age_category, {}
            ).get(curve_param, 0)
            params.append(
                (
                    curve_param,
                    season_range.season,
                    age_category,
                    param_value,
                ),
            )

pprint(params)
params = []

params_grid = {
    "2022/23": {
        "m": -0.5,
        "r": -3.2,
        "s": 40.0,
        "adult": {
            "m": 1.0,
        },
        "elderly": {
            "m": 1.8,
        },
    },
    "2023/24": {
        "m": -0.6,
        "r": -3.1,
        "s": 42.5,
        "adult": {
            "m": 1.05,
        },
        "elderly": {
            "m": 1.78,
        },
    },
    "2024/25": {
        "m": -0.4,
        "r": -3.3,
        "s": 41.0,
        "adult": {
            "m": 1.03,
        },
        "elderly": {
            "m": 1.81,
        },
    },
}

for curve_param in curve.parameters:
    for season_range in seasons:
        for age_category in age_covariate_categories.categories:
            season_params = params_grid[season_range.season]
            param_value = season_params[curve_param] + season_params.get(
                age_category, {}
            ).get(curve_param, 0)
            params.append(
                (
                    curve_param,
                    season_range.season,
                    age_category,
                    param_value,
                ),
            )

pprint(params)

[('m', '2022/23', 'youth', -0.5),
 ('m', '2022/23', 'adult', 0.5),
 ('m', '2022/23', 'elderly', 1.3),
 ('m', '2023/24', 'youth', -0.6),
 ('m', '2023/24', 'adult', 0.45000000000000007),
 ('m', '2023/24', 'elderly', 1.1800000000000002),
 ('m', '2024/25', 'youth', -0.4),
 ('m', '2024/25', 'adult', 0.63),
 ('m', '2024/25', 'elderly', 1.4100000000000001),
 ('r', '2022/23', 'youth', -3.2),
 ('r', '2022/23', 'adult', -3.2),
 ('r', '2022/23', 'elderly', -3.2),
 ('r', '2023/24', 'youth', -3.1),
 ('r', '2023/24', 'adult', -3.1),
 ('r', '2023/24', 'elderly', -3.1),
 ('r', '2024/25', 'youth', -3.3),
 ('r', '2024/25', 'adult', -3.3),
 ('r', '2024/25', 'elderly', -3.3),
 ('s', '2022/23', 'youth', 40.0),
 ('s', '2022/23', 'adult', 40.0),
 ('s', '2022/23', 'elderly', 40.0),
 ('s', '2023/24', 'youth', 42.5),
 ('s', '2023/24', 'adult', 42.5),
 ('s', '2023/24', 'elderly', 42.5),
 ('s', '2024/25', 'youth', 41.0),
 ('s', '2024/25', 'adult', 41.0),
 ('s', '2024/25', 'elderly', 41.0)]

Generate Synthetic Observations¶

sample_dataset evaluates the logistic curve at each observation date using the specified parameters and adds small noise (epsilon). The result is wrapped in a VaxfluxObservations container. Displaying the head confirms the expected data shape: one row per (season, date, age) combination with an incidence value.

In [7]:

Copied!





observations = VaxfluxObservations.from_dataframe(
    sample_dataset(
        curve,
        seasons,
        dates,
        [age_covariate_categories],
        params,
        epsilon=0.0005,
    )
)

observations.data.head()
observations = VaxfluxObservations.from_dataframe(
    sample_dataset(
        curve,
        seasons,
        dates,
        [age_covariate_categories],
        params,
        epsilon=0.0005,
    )
)

observations.data.head()

Out[7]:

	season	season_start_date	season_end_date	start_date	end_date	report_date	age	type	value
0	2022/23	2022-10-03	2023-02-05	2022-10-03	2022-10-09	2022-10-09	youth	incidence	0.016363
1	2022/23	2022-10-03	2023-02-05	2022-10-03	2022-10-09	2022-10-09	adult	incidence	0.025353
2	2022/23	2022-10-03	2023-02-05	2022-10-03	2022-10-09	2022-10-09	elderly	incidence	0.033812
3	2022/23	2022-10-03	2023-02-05	2022-10-10	2022-10-16	2022-10-16	youth	incidence	0.017795
4	2022/23	2022-10-03	2023-02-05	2022-10-10	2022-10-16	2022-10-16	adult	incidence	0.030507

Plot the Synthetic Observations¶

A quick visual sanity check of the generated data. The top panel shows weekly incidence and the bottom shows cumulative prevalence (running sum of incidence), broken down by season and age group. The age-group ordering and seasonal variation should match the parameter grid defined above.

In [8]:

Copied!





plot_data = observations.data.copy()
plot_data["start_date"] = pd.to_datetime(plot_data["start_date"])
plot_data["end_date"] = pd.to_datetime(plot_data["end_date"])
plot_data["mid_date"] = (
    plot_data["start_date"] + (plot_data["end_date"] - plot_data["start_date"]) / 2
)
plot_data = plot_data.sort_values("mid_date")
plot_data["prevalence"] = plot_data.groupby(["season", "age"])["value"].cumsum()

fig, (ax_inc, ax_prev) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)

for (season, age), group in plot_data.groupby(["season", "age"]):
    label = f"{season} | {age}"
    ax_inc.plot(group["mid_date"], group["value"], label=label)
    ax_prev.plot(group["mid_date"], group["prevalence"], label=label)

ax_inc.set_title("Sample observations by season and age")
ax_inc.set_ylabel("Incidence")

ax_prev.set_title("Cumulative prevalence by season and age")
ax_prev.set_ylabel("Prevalence")
ax_prev.set_xlabel("Date")

handles, labels = ax_inc.get_legend_handles_labels()
fig.legend(handles, labels, loc="center left", bbox_to_anchor=(1.0, 0.5))
for label in ax_prev.get_xticklabels():
    label.set_rotation(30)
fig.tight_layout(rect=(0, 0, 0.85, 1))
plt.show()
plot_data = observations.data.copy()
plot_data["start_date"] = pd.to_datetime(plot_data["start_date"])
plot_data["end_date"] = pd.to_datetime(plot_data["end_date"])
plot_data["mid_date"] = (
    plot_data["start_date"] + (plot_data["end_date"] - plot_data["start_date"]) / 2
)
plot_data = plot_data.sort_values("mid_date")
plot_data["prevalence"] = plot_data.groupby(["season", "age"])["value"].cumsum()

fig, (ax_inc, ax_prev) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)

for (season, age), group in plot_data.groupby(["season", "age"]):
    label = f"{season} | {age}"
    ax_inc.plot(group["mid_date"], group["value"], label=label)
    ax_prev.plot(group["mid_date"], group["prevalence"], label=label)

ax_inc.set_title("Sample observations by season and age")
ax_inc.set_ylabel("Incidence")

ax_prev.set_title("Cumulative prevalence by season and age")
ax_prev.set_ylabel("Prevalence")
ax_prev.set_xlabel("Date")

handles, labels = ax_inc.get_legend_handles_labels()
fig.legend(handles, labels, loc="center left", bbox_to_anchor=(1.0, 0.5))
for label in ax_prev.get_xticklabels():
    label.set_rotation(30)
fig.tight_layout(rect=(0, 0, 0.85, 1))
plt.show()

/var/folders/2z/h3pc0p7s3ng1tvxrgsw5kr680000gp/T/ipykernel_64306/996366010.py:28: UserWarning: The figure layout has changed to tight
  fig.tight_layout(rect=(0, 0, 0.85, 1))

No description has been provided for this image

Build the Model¶

VaxfluxModel is assembled using a fluent builder pattern. Each .add_* call registers a component:

.add_seasons / .add_dates — the temporal structure
.add_covariate_categories — the age groups to track
.add_covariates — the prior distributions for each curve parameter
.add_observations — the data to condition on
.add_observation_process — the likelihood (here a Normal with fixed noise and a prevalence penalty that discourages the cumulative curve from exceeding 1)

In [9]:

Copied!





model = (
    VaxfluxModel(curve=curve)
    .add_seasons(seasons)
    .add_dates(dates)
    .add_covariate_categories(age_covariate_categories)
    .add_covariates(covariates)
    .add_observations(observations)
    .add_observation_process(
        kind="normal",
        noise=0.005,
        partially_pool_by_season=False,
        prevalence_penalty=10.0,
    )
)
model
model = (
    VaxfluxModel(curve=curve)
    .add_seasons(seasons)
    .add_dates(dates)
    .add_covariate_categories(age_covariate_categories)
    .add_covariates(covariates)
    .add_observations(observations)
    .add_observation_process(
        kind="normal",
        noise=0.005,
        partially_pool_by_season=False,
        prevalence_penalty=10.0,
    )
)
model

Out[9]:

<vaxflux.VaxfluxModel object at 0x11e98e0c0>

Render the Model Graph¶

render_model produces a graphviz diagram of the probabilistic model structure. This is a useful check that the hierarchy of parameters, covariates, and observations is wired together as intended before running any sampling.

In [10]:

Copied!

model.render_model(render_params=True)
model.render_model(render_params=True)

Out[10]:

Run Prior Predictive Sampling¶

Calling model.sample with samples=None skips the posterior and draws only from the prior. The result is an InferenceData-like object with a prior group. Inspecting it before running the full MCMC is a prior predictive check — useful for confirming the priors produce plausible observations.

Inspect Prior Predictive Observations¶

prior_observations is a long-form DataFrame of simulated incidence values drawn from the prior — one row per (chain, draw, season, date, age) combination. Browsing this confirms the prior is generating data in a reasonable range before any observations are seen.

In [12]:

Copied!

vaxflux_prior_idata.prior_observations
vaxflux_prior_idata.prior_observations

Out[12]:

	chain	draw	season	season_start_date	season_end_date	start_date	end_date	report_date	age	type	value
0	0	0	2022/23	2022-10-03	2023-02-05	2023-01-16	2023-01-22	2023-01-22	youth	incidence	1.015942e-02
1	0	0	2022/23	2022-10-03	2023-02-05	2022-10-17	2022-10-23	2022-10-23	youth	incidence	1.166365e-02
2	0	0	2022/23	2022-10-03	2023-02-05	2022-12-26	2023-01-01	2023-01-01	youth	incidence	1.709607e-02
3	0	0	2022/23	2022-10-03	2023-02-05	2022-11-28	2022-12-04	2022-12-04	youth	incidence	2.243422e-02
4	0	0	2022/23	2022-10-03	2023-02-05	2022-11-07	2022-11-13	2022-11-13	youth	incidence	1.864388e-02
...	...	...	...	...	...	...	...	...	...	...	...
39745	0	249	2024/25	2024-10-07	2025-02-02	2024-11-25	2024-12-01	2024-12-01	elderly	incidence	6.461699e-03
39746	0	249	2024/25	2024-10-07	2025-02-02	2024-12-02	2024-12-08	2024-12-08	elderly	incidence	1.491969e-03
39747	0	249	2024/25	2024-10-07	2025-02-02	2024-12-09	2024-12-15	2024-12-15	elderly	incidence	3.419897e-04
39748	0	249	2024/25	2024-10-07	2025-02-02	2025-01-20	2025-01-26	2025-01-26	elderly	incidence	2.111469e-08
39749	0	249	2024/25	2024-10-07	2025-02-02	2024-11-11	2024-11-17	2024-11-17	elderly	incidence	1.004628e-01

39750 rows × 11 columns

Plot Prior Predictive Curves¶

This overlays the prior predictive draws against the observed data for a visual prior predictive check. Wide bands indicate an uninformative prior; the goal is for the prior to cover the observed data without being so diffuse that it is uninformative about the scale.

In [13]:

Copied!

vaxflux_prior_idata.plot_prior_predictive()
plt.show()
vaxflux_prior_idata.plot_prior_predictive()
plt.show()

Run MCMC Posterior Sampling¶

Now we condition on the observed data. model.sample runs NUTS (No-U-Turn Sampler) via NumPyro with 4 parallel chains, each running 250 warmup steps followed by 250 posterior draws. The resulting InferenceData object contains the posterior, posterior predictive, log likelihood, and sample diagnostics.

Inspect Posterior Observations¶

posterior_observations returns the posterior predictive draws in the same long-form DataFrame format as the input observations — one row per (chain, draw, season, date, age). This makes it easy to compute credible intervals and compare directly against the observed incidence.

In [15]:

Copied!

vaxflux_idata.posterior_observations
vaxflux_idata.posterior_observations

Out[15]:

	chain	draw	season	season_start_date	season_end_date	start_date	end_date	report_date	age	type	value
0	0	0	2022/23	2022-10-03	2023-02-05	2023-01-16	2023-01-22	2023-01-22	youth	incidence	0.005906
1	0	0	2022/23	2022-10-03	2023-02-05	2022-10-17	2022-10-23	2022-10-23	youth	incidence	0.021702
2	0	0	2022/23	2022-10-03	2023-02-05	2022-12-26	2023-01-01	2023-01-01	youth	incidence	0.012082
3	0	0	2022/23	2022-10-03	2023-02-05	2022-11-28	2022-12-04	2022-12-04	youth	incidence	0.023557
4	0	0	2022/23	2022-10-03	2023-02-05	2023-01-02	2023-01-08	2023-01-08	youth	incidence	0.009651
...	...	...	...	...	...	...	...	...	...	...	...
158995	3	249	2024/25	2024-10-07	2025-02-02	2024-11-25	2024-12-01	2024-12-01	elderly	incidence	0.049539
158996	3	249	2024/25	2024-10-07	2025-02-02	2024-12-02	2024-12-08	2024-12-08	elderly	incidence	0.046355
158997	3	249	2024/25	2024-10-07	2025-02-02	2024-12-09	2024-12-15	2024-12-15	elderly	incidence	0.042154
158998	3	249	2024/25	2024-10-07	2025-02-02	2025-01-20	2025-01-26	2025-01-26	elderly	incidence	0.015311
158999	3	249	2024/25	2024-10-07	2025-02-02	2024-11-11	2024-11-17	2024-11-17	elderly	incidence	0.051573

159000 rows × 11 columns

Trace Plots for Convergence Diagnostics¶

az.plot_trace shows the MCMC trace (sample values over iterations) and marginal posterior distribution for each named parameter. Healthy chains should look like "fuzzy caterpillars" with all chains overlapping — divergences or drifting traces indicate sampling problems.

In [16]:

Copied!





az.plot_trace(
    vaxflux_idata,
    var_names=[
        "covariate_values_m_age",
        "covariate_values_m_season",
        "covariate_values_r_season",
        "covariate_values_s_season",
        "m_age",
        "m_age_mu",
        "m_age_sigma",
        "m_season",
        "m_season_mu",
        "m_season_sigma",
        "r_season",
        "r_season_mu",
        "r_season_sigma",
        "s_season",
        "s_season_mu",
        "s_season_sigma",
        "observation_sigma",
    ],
)
plt.show()
az.plot_trace(
    vaxflux_idata,
    var_names=[
        "covariate_values_m_age",
        "covariate_values_m_season",
        "covariate_values_r_season",
        "covariate_values_s_season",
        "m_age",
        "m_age_mu",
        "m_age_sigma",
        "m_season",
        "m_season_mu",
        "m_season_sigma",
        "r_season",
        "r_season_mu",
        "r_season_sigma",
        "s_season",
        "s_season_mu",
        "s_season_sigma",
        "observation_sigma",
    ],
)
plt.show()

/Users/twillard/Desktop/GitHub/ACCIDDA/vaxflux/.venv/lib/python3.12/site-packages/arviz/stats/density_utils.py:488: UserWarning: Your data appears to have a single value or no finite values
  warnings.warn("Your data appears to have a single value or no finite values")

Plot Posterior Predictive Fit¶

plot_posterior_predictive overlays the posterior predictive credible interval against the observed incidence data. A good fit shows the observed data falling within the posterior band, with the band being reasonably narrow.

In [17]:

Copied!

vaxflux_idata.plot_posterior_predictive()
plt.show()
vaxflux_idata.plot_posterior_predictive()
plt.show()

$No description has been provided for this image$

In [ ]: