Season-Only Vaxflux Model Demo¶

This notebook is a good place to start if you're new to vaxflux. It walks through the core modeling workflow from start to finish using synthetic data, with variation across flu seasons but no demographic covariates. By the end you'll have fit a model, checked that it sampled well, and seen how its predictions compare to the data.

In [1]:

Copied!





from pprint import pprint

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

from vaxflux import (
    LogisticCurve,
    PartiallyPooledGaussianCovariate,
    SeasonRange,
    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 arviz_plots
import matplotlib.pyplot as plt
import numpyro
import pandas as pd

from vaxflux import (
    LogisticCurve,
    PartiallyPooledGaussianCovariate,
    SeasonRange,
    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)$$

Visualize the Logistic Curve Parameters¶

Before generating data, it helps to build intuition for what each logistic curve parameter does. These plots sweep one parameter at a time while holding the others fixed, showing how m shifts the peak timing, r controls steepness, and s shifts the horizontal position of the uptake curve.

In [3]:

Copied!





for title, parameter_set in [
    (
        "Logistic curve varying $m$",
        [
            (-1.0, -1.0, 30.0),
            (0.0, -1.0, 30.0),
            (1.0, -1.0, 30.0),
        ],
    ),
    (
        "Logistic curve varying $r$",
        [
            (-0.75, -2.5, 40.0),
            (-0.75, -1.75, 40.0),
            (-0.75, -1.0, 40.0),
        ],
    ),
    (
        "Logistic curve varying $s$",
        [
            (0.0, -2.0, 20.0),
            (0.0, -2.0, 40.0),
            (0.0, -2.0, 60.0),
        ],
    ),
]:
    curve.plot(
        [float(day) for day in range(85)],
        parameter_sets=[
            dict(zip(curve.parameters, values, strict=True)) for values in parameter_set
        ],
        title=title,
    )
plt.show()
for title, parameter_set in [
    (
        "Logistic curve varying $m$",
        [
            (-1.0, -1.0, 30.0),
            (0.0, -1.0, 30.0),
            (1.0, -1.0, 30.0),
        ],
    ),
    (
        "Logistic curve varying $r$",
        [
            (-0.75, -2.5, 40.0),
            (-0.75, -1.75, 40.0),
            (-0.75, -1.0, 40.0),
        ],
    ),
    (
        "Logistic curve varying $s$",
        [
            (0.0, -2.0, 20.0),
            (0.0, -2.0, 40.0),
            (0.0, -2.0, 60.0),
        ],
    ),
]:
    curve.plot(
        [float(day) for day in range(85)],
        parameter_sets=[
            dict(zip(curve.parameters, values, strict=True)) for values in parameter_set
        ],
        title=title,
    )
plt.show()

No description has been provided for this image

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 [4]:

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 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.

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,
    ),
]
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,
    ),
]

Specify True Parameters for Synthetic Data¶

This defines the ground-truth logistic curve parameters used to generate the synthetic observations. Since this is a season-only model (no age groups), each season gets a single value for each parameter. These are the values the model should recover during inference.

In [6]:

Copied!





params = []

params_grid = {
    "2022/23": {
        "m": -0.5,
        "r": -3.2,
        "s": 40.0,
    },
    "2023/24": {
        "m": -0.6,
        "r": -3.1,
        "s": 42.5,
    },
    "2024/25": {
        "m": -0.4,
        "r": -3.3,
        "s": 41.0,
    },
}

for curve_param in curve.parameters:
    for season_range in seasons:
        season_params = params_grid[season_range.season]
        params.append(
            (
                curve_param,
                season_range.season,
                season_params[curve_param],
            ),
        )

pprint(params)
params = []

params_grid = {
    "2022/23": {
        "m": -0.5,
        "r": -3.2,
        "s": 40.0,
    },
    "2023/24": {
        "m": -0.6,
        "r": -3.1,
        "s": 42.5,
    },
    "2024/25": {
        "m": -0.4,
        "r": -3.3,
        "s": 41.0,
    },
}

for curve_param in curve.parameters:
    for season_range in seasons:
        season_params = params_grid[season_range.season]
        params.append(
            (
                curve_param,
                season_range.season,
                season_params[curve_param],
            ),
        )

pprint(params)

[('m', '2022/23', -0.5),
 ('m', '2023/24', -0.6),
 ('m', '2024/25', -0.4),
 ('r', '2022/23', -3.2),
 ('r', '2023/24', -3.1),
 ('r', '2024/25', -3.3),
 ('s', '2022/23', 40.0),
 ('s', '2023/24', 42.5),
 ('s', '2024/25', 41.0)]

Generate Synthetic Observations¶

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

In [7]:

Copied!





observations = VaxfluxObservations.from_dataframe(
    sample_dataset(
        curve,
        seasons,
        dates,
        [],
        params,
        epsilon=0.001,
        noise="normal",
    )
)

observations.data.head()
observations = VaxfluxObservations.from_dataframe(
    sample_dataset(
        curve,
        seasons,
        dates,
        [],
        params,
        epsilon=0.001,
        noise="normal",
    )
)

observations.data.head()

Out[7]:

	season	season_start_date	season_end_date	start_date	end_date	report_date	type	value
0	2022/23	2022-10-03	2023-02-05	2022-10-03	2022-10-09	2022-10-09	incidence	0.017206
1	2022/23	2022-10-03	2023-02-05	2022-10-10	2022-10-16	2022-10-16	incidence	0.020050
2	2022/23	2022-10-03	2023-02-05	2022-10-17	2022-10-23	2022-10-23	incidence	0.020793
3	2022/23	2022-10-03	2023-02-05	2022-10-24	2022-10-30	2022-10-30	incidence	0.021211
4	2022/23	2022-10-03	2023-02-05	2022-10-31	2022-11-06	2022-11-06	incidence	0.029847

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. The seasonal variation in timing and height 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"])["value"].cumsum()

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

for season, group in plot_data.groupby(["season"]):
    label = f"{season}"
    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")
ax_inc.set_ylabel("Incidence")

ax_prev.set_title("Cumulative prevalence by season")
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"])["value"].cumsum()

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

for season, group in plot_data.groupby(["season"]):
    label = f"{season}"
    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")
ax_inc.set_ylabel("Incidence")

ax_prev.set_title("Cumulative prevalence by season")
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_64365/1021347413.py:28: UserWarning: The figure layout has changed to tight
  fig.tight_layout(rect=(0, 0, 0.85, 1))

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_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_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_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 0x11ee2b8f0>

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) 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	type	value
0	0	0	2022/23	2022-10-03	2023-02-05	2023-01-23	2023-01-29	2023-01-29	incidence	0.008205
1	0	0	2022/23	2022-10-03	2023-02-05	2023-01-16	2023-01-22	2023-01-22	incidence	0.010159
2	0	0	2022/23	2022-10-03	2023-02-05	2022-10-10	2022-10-16	2022-10-16	incidence	0.009531
3	0	0	2022/23	2022-10-03	2023-02-05	2022-10-31	2022-11-06	2022-11-06	incidence	0.016378
4	0	0	2022/23	2022-10-03	2023-02-05	2023-01-02	2023-01-08	2023-01-08	incidence	0.014722
...	...	...	...	...	...	...	...	...	...	...
26495	0	499	2024/25	2024-10-07	2025-02-02	2024-12-23	2024-12-29	2024-12-29	incidence	0.002985
26496	0	499	2024/25	2024-10-07	2025-02-02	2025-01-06	2025-01-12	2025-01-12	incidence	0.001987
26497	0	499	2024/25	2024-10-07	2025-02-02	2024-10-28	2024-11-03	2024-11-03	incidence	0.011937
26498	0	499	2024/25	2024-10-07	2025-02-02	2024-12-09	2024-12-15	2024-12-15	incidence	0.004416
26499	0	499	2024/25	2024-10-07	2025-02-02	2025-01-27	2025-02-02	2025-02-02	incidence	0.001060

26500 rows × 10 columns

Inspect Prior Parameter Distributions¶

arviz_plots.plot_dist shows the marginal prior distribution for each model parameter across all 500 prior draws. This helps verify that the specified hyperpriors produce plausible parameter ranges before seeing any data.

In [13]:

Copied!





arviz_plots.plot_dist(
    vaxflux_prior_idata,
    group="prior",
    var_names=[
        "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",
    ],
    col_wrap=1,
)
plt.show()
arviz_plots.plot_dist(
    vaxflux_prior_idata,
    group="prior",
    var_names=[
        "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",
    ],
    col_wrap=1,
)
plt.show()

$No description has been provided for this image$

Plot Prior Predictive Curves¶

plot_prior_predictive overlays the prior predictive draws against the observed data. The goal is for the prior to cover the observed data without being so diffuse that it is uninformative about the scale.

In [14]:

Copied!

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

Plot Prior Predictive with Observation Noise¶

Passing predictive="observation" adds the observation-level noise on top of the latent curve draws. This is the quantity actually compared against the data during inference, so it should look consistent with the observed incidence values.

In [15]:

Copied!

vaxflux_prior_idata.plot_prior_predictive(predictive="observation")
plt.show()
vaxflux_prior_idata.plot_prior_predictive(predictive="observation")
plt.show()

Summarize Prior Prevalence Scenarios¶

prior_prevalence_scenarios summarizes what total-season prevalence (cumulative uptake) the prior implies for a set of quantile-based scenarios. The "low", "typical", and "high" bands show the range of plausible season-end vaccination coverage under the prior alone, before any data is used.

In [16]:

Copied!

vaxflux_prior_idata.prior_prevalence_scenarios(
    {"low": (0.025, 0.25), "typical": (0.25, 0.75), "high": (0.75, 0.975)}
)
vaxflux_prior_idata.prior_prevalence_scenarios(
    {"low": (0.025, 0.25), "typical": (0.25, 0.75), "high": (0.75, 0.975)}
)

Out[16]:

	season	scenario	low_quantile	high_quantile	low_prevalence	high_prevalence
0	2022/23	low	0.025	0.250	0.004458	0.159307
1	2023/24	low	0.025	0.250	0.001449	0.156781
2	2024/25	low	0.025	0.250	0.007901	0.134623
3	2022/23	typical	0.250	0.750	0.159307	0.517974
4	2023/24	typical	0.250	0.750	0.156781	0.504385
5	2024/25	typical	0.250	0.750	0.134623	0.498256
6	2022/23	high	0.750	0.975	0.517974	0.828640
7	2023/24	high	0.750	0.975	0.504385	0.800235
8	2024/25	high	0.750	0.975	0.498256	0.830469

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) combination. This makes it easy to compute credible intervals and compare directly against the observed incidence.

In [18]:

Copied!

vaxflux_idata.posterior_observations
vaxflux_idata.posterior_observations

Out[18]:

	chain	draw	season	season_start_date	season_end_date	start_date	end_date	report_date	type	value
0	0	0	2022/23	2022-10-03	2023-02-05	2023-01-23	2023-01-29	2023-01-29	incidence	0.005217
1	0	0	2022/23	2022-10-03	2023-02-05	2023-01-16	2023-01-22	2023-01-22	incidence	0.006592
2	0	0	2022/23	2022-10-03	2023-02-05	2022-10-10	2022-10-16	2022-10-16	incidence	0.019198
3	0	0	2022/23	2022-10-03	2023-02-05	2022-10-31	2022-11-06	2022-11-06	incidence	0.025160
4	0	0	2022/23	2022-10-03	2023-02-05	2023-01-02	2023-01-08	2023-01-08	incidence	0.010216
...	...	...	...	...	...	...	...	...	...	...
52995	3	249	2024/25	2024-10-07	2025-02-02	2024-12-23	2024-12-29	2024-12-29	incidence	0.015442
52996	3	249	2024/25	2024-10-07	2025-02-02	2025-01-06	2025-01-12	2025-01-12	incidence	0.010374
52997	3	249	2024/25	2024-10-07	2025-02-02	2024-10-28	2024-11-03	2024-11-03	incidence	0.023840
52998	3	249	2024/25	2024-10-07	2025-02-02	2024-12-09	2024-12-15	2024-12-15	incidence	0.021023
52999	3	249	2024/25	2024-10-07	2025-02-02	2025-01-27	2025-02-02	2025-02-02	incidence	0.005136

53000 rows × 10 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 [19]:

Copied!





az.plot_trace(
    vaxflux_idata,
    var_names=[
        "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=[
        "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()

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 [20]:

Copied!

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

Plot Posterior Predictive with Observation Noise¶

Passing predictive="observation" includes the modeled observation noise in the posterior predictive draws. This represents the full uncertainty — both parameter uncertainty and measurement noise — and is the appropriate comparison against individual data points.

In [21]:

Copied!

vaxflux_idata.plot_posterior_predictive(predictive="observation")
plt.show()
vaxflux_idata.plot_posterior_predictive(predictive="observation")
plt.show()

Summarize Posterior Prevalence Scenarios¶

posterior_prevalence_scenarios reports season-end vaccination coverage (cumulative incidence) under the posterior, grouped into "low", "typical", and "high" quantile bands. These are the posterior estimates of how much of the population got vaccinated by the end of each season.

In [22]:

Copied!

vaxflux_idata.posterior_prevalence_scenarios(
    {"low": (0.025, 0.25), "typical": (0.25, 0.75), "high": (0.75, 0.975)}
)
vaxflux_idata.posterior_prevalence_scenarios(
    {"low": (0.025, 0.25), "typical": (0.25, 0.75), "high": (0.75, 0.975)}
)

Out[22]:

	season	scenario	low_quantile	high_quantile	low_prevalence	high_prevalence
0	2022/23	low	0.025	0.250	0.281009	0.294154
1	2023/24	low	0.025	0.250	0.273335	0.285584
2	2024/25	low	0.025	0.250	0.278789	0.290287
3	2022/23	typical	0.250	0.750	0.294154	0.306520
4	2023/24	typical	0.250	0.750	0.285584	0.299913
5	2024/25	typical	0.250	0.750	0.290287	0.301608
6	2022/23	high	0.750	0.975	0.306520	0.318617
7	2023/24	high	0.750	0.975	0.299913	0.312330
8	2024/25	high	0.750	0.975	0.301608	0.312601

In [ ]: