NIS-FRVM Survey Demo¶

This notebook fits a vaxflux model to real flu vaccination coverage data from the CDC's National Immunization Survey (NIS-FluVaxView / FRVM), covering five seasons and three age groups. Unlike the synthetic demos, the data and season structure come directly from the survey, so this is closer to what a real analysis workflow looks like.

In [1]:

Copied!





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

from vaxflux import (
    CovariateCategories,
    LogisticCurve,
    PartiallyPooledGaussianCovariate,
    SeasonVaryingPartiallyPooledGaussianCovariate,
    VaxfluxModel,
    VaxfluxObservations,
)
from vaxflux.data import get_ncird_nis_frvm_flu_vaccination_coverage

numpyro.set_host_device_count(8)
az.style.use("arviz-docgrid")
import arviz as az
import matplotlib.pyplot as plt
import numpyro
import pandas as pd

from vaxflux import (
    CovariateCategories,
    LogisticCurve,
    PartiallyPooledGaussianCovariate,
    SeasonVaryingPartiallyPooledGaussianCovariate,
    VaxfluxModel,
    VaxfluxObservations,
)
from vaxflux.data import get_ncird_nis_frvm_flu_vaccination_coverage

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

/Users/twillard/Desktop/GitHub/ACCIDDA/vaxflux/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Load Real Survey Data¶

get_ncird_nis_frvm_flu_vaccination_coverage fetches CDC NIS-FRVM flu vaccination coverage estimates. Passing include_age_groups=True returns data stratified by age group (18–49, 50–64, 65+). The result is wrapped in VaxfluxObservations. Displaying the head confirms the expected long-form structure: one row per (season, age, time window) with an incidence value.

In [2]:

Copied!





observations = VaxfluxObservations.from_dataframe(
    get_ncird_nis_frvm_flu_vaccination_coverage(include_age_groups=True)
)
observations.data.head()
observations = VaxfluxObservations.from_dataframe(
    get_ncird_nis_frvm_flu_vaccination_coverage(include_age_groups=True)
)
observations.data.head()

Out[2]:

	season	season_start_date	season_end_date	start_date	end_date	report_date	age	type	value
0	2021/22	2021-08-01	2022-03-31	2021-08-01	2021-10-30	2021-10-30	18-49 years	incidence	0.170
1	2021/22	2021-08-01	2022-03-31	2021-10-30	2021-11-27	2021-11-27	18-49 years	incidence	0.109
2	2021/22	2021-08-01	2022-03-31	2021-11-27	2021-12-31	2021-12-31	18-49 years	incidence	0.043
3	2021/22	2021-08-01	2022-03-31	2021-12-31	2022-01-29	2022-01-29	18-49 years	incidence	0.024
4	2021/22	2021-08-01	2022-03-31	2022-01-29	2022-02-26	2022-02-26	18-49 years	incidence	0.004

Inspect Total Season Prevalence¶

A quick sanity check: summing incidence within each (season, age) group gives cumulative season prevalence — the total vaccination coverage by season end. The 65+ group consistently shows the highest coverage (~65–72%), as expected for flu vaccination.

In [3]:

Copied!

observations.data.groupby(["season", "age"], as_index=False)["value"].sum()
observations.data.groupby(["season", "age"], as_index=False)["value"].sum()

Out[3]:

	season	age	value
0	2021/22	18-49 years	0.363
1	2021/22	50-64 years	0.488
2	2021/22	65+ years	0.655
3	2022/23	18-49 years	0.358
4	2022/23	50-64 years	0.503
5	2022/23	65+ years	0.684
6	2023/24	18-49 years	0.366
7	2023/24	50-64 years	0.502
8	2023/24	65+ years	0.715
9	2024/25	18-49 years	0.363
10	2024/25	50-64 years	0.483
11	2024/25	65+ years	0.717
12	2025/26	18-49 years	0.376
13	2025/26	50-64 years	0.454
14	2025/26	65+ years	0.690

Plot the Observed Data¶

A visual overview of the NIS-FRVM data. The top panel shows bi-weekly incidence (new vaccinations) and the bottom shows cumulative prevalence, broken out by season and age group. The early-season surge pattern and age-group ordering are clearly visible.

In [4]:

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_37897/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

Instantiate the Logistic Curve¶

LogisticCurve is the functional form used to model vaccination uptake over time within a season.

In [5]:

Copied!

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

Out[5]:

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

Extract Season Ranges and Age Categories from the Data¶

Rather than hardcoding season boundaries, we derive SeasonRange objects directly from the observations data. Similarly, age categories are extracted from the unique values in the age column. This keeps the model setup in sync with whatever the data actually contains.

In [6]:

Copied!





seasons = observations.to_season_ranges()

age_covariate_categories = CovariateCategories(
    covariate="age",
    categories=tuple(observations.data["age"].drop_duplicates().tolist()),
)
seasons = observations.to_season_ranges()

age_covariate_categories = CovariateCategories(
    covariate="age",
    categories=tuple(observations.data["age"].drop_duplicates().tolist()),
)

Define Prior Covariates¶

Priors are specified for all three logistic curve parameters (m, r, s) at the season level, plus season-varying age offsets for all three parameters. The tighter sigma values here (compared to the synthetic demo) reflect that the NIS-FRVM data is real and the priors should be more constrained around known vaccination patterns.

In [7]:

Copied!





covariates = [
    PartiallyPooledGaussianCovariate(
        parameter="m",
        covariate=None,
        mu_mu=-0.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="r",
        covariate=None,
        mu_mu=-3.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="s",
        covariate=None,
        mu_mu=60.0,
        mu_sigma=10.0,
        sigma=10.0,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="m",
        covariate="age",
        mu_mu=1.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="r",
        covariate="age",
        mu_mu=0.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="s",
        covariate="age",
        mu_mu=0.0,
        mu_sigma=10.0,
        sigma=10.0,
    ),
]
covariates = [
    PartiallyPooledGaussianCovariate(
        parameter="m",
        covariate=None,
        mu_mu=-0.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="r",
        covariate=None,
        mu_mu=-3.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    PartiallyPooledGaussianCovariate(
        parameter="s",
        covariate=None,
        mu_mu=60.0,
        mu_sigma=10.0,
        sigma=10.0,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="m",
        covariate="age",
        mu_mu=1.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="r",
        covariate="age",
        mu_mu=0.5,
        mu_sigma=0.5,
        sigma=0.5,
    ),
    SeasonVaryingPartiallyPooledGaussianCovariate(
        parameter="s",
        covariate="age",
        mu_mu=0.0,
        mu_sigma=10.0,
        sigma=10.0,
    ),
]

Build the Model¶

The model is assembled without .add_dates because the NIS-FRVM data has non-uniform reporting intervals (bi-weekly, monthly), so dates come directly from the observations rather than being generated by daily_date_ranges. A slightly larger observation noise (noise=0.01) is used compared to the synthetic demos, since real survey data carries measurement error.

In [8]:

Copied!





model = (
    VaxfluxModel(curve=curve)
    .add_seasons(seasons)
    .add_covariate_categories(age_covariate_categories)
    .add_covariates(covariates)
    .add_observations(observations)
    .add_observation_process(
        kind="normal",
        noise=0.01,
    )
)
model
model = (
    VaxfluxModel(curve=curve)
    .add_seasons(seasons)
    .add_covariate_categories(age_covariate_categories)
    .add_covariates(covariates)
    .add_observations(observations)
    .add_observation_process(
        kind="normal",
        noise=0.01,
    )
)
model

Out[8]:

<vaxflux.VaxfluxModel object at 0x10acf8c10>

Run Prior Predictive Sampling¶

Draw 100 samples from the prior to perform a prior predictive check before conditioning on the real data.

Inspect Prior Predictive Observations¶

Browse the prior predictive draws in long-form to get a sense of the range of incidence values the prior generates. These should look qualitatively similar to the observed data.

In [10]:

Copied!

vaxflux_prior_idata.prior_observations
vaxflux_prior_idata.prior_observations

Out[10]:

	chain	draw	season	season_start_date	season_end_date	start_date	end_date	report_date	age	type	value
0	0	0	2021/22	2021-08-01	2022-04-30	2021-08-01	2021-10-30	2021-10-30	18-49 years	incidence	0.122936
1	0	0	2021/22	2021-08-01	2022-04-30	2021-10-30	2021-11-27	2021-11-27	18-49 years	incidence	0.040334
2	0	0	2021/22	2021-08-01	2022-04-30	2021-11-27	2021-12-31	2021-12-31	18-49 years	incidence	0.033967
3	0	0	2021/22	2021-08-01	2022-04-30	2021-12-31	2022-01-29	2022-01-29	18-49 years	incidence	0.017063
4	0	0	2021/22	2021-08-01	2022-04-30	2022-01-29	2022-02-26	2022-02-26	18-49 years	incidence	0.009292
...	...	...	...	...	...	...	...	...	...	...	...
28495	0	99	2025/26	2025-08-01	2026-02-21	2026-01-17	2026-01-24	2026-01-24	65+ years	incidence	0.009216
28496	0	99	2025/26	2025-08-01	2026-02-21	2026-01-24	2026-01-31	2026-01-31	65+ years	incidence	0.008976
28497	0	99	2025/26	2025-08-01	2026-02-21	2026-01-31	2026-02-07	2026-02-07	65+ years	incidence	0.008732
28498	0	99	2025/26	2025-08-01	2026-02-21	2026-02-07	2026-02-14	2026-02-14	65+ years	incidence	0.008485
28499	0	99	2025/26	2025-08-01	2026-02-21	2026-02-14	2026-02-21	2026-02-21	65+ years	incidence	0.008236

28500 rows × 11 columns

Plot Prior Predictive Curves¶

Visual prior predictive check: the band of prior draws should envelope the observed data while being sufficiently broad to reflect prior uncertainty.

In [11]:

Copied!

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

Run MCMC Posterior Sampling¶

Run NUTS with 2 chains, 250 warmup steps, and 250 posterior draws. More samples than the synthetic demos are used here given the greater complexity of fitting real data across 5 seasons and 3 age groups.

Inspect Posterior Observations¶

Browse the posterior predictive draws. For real data, these should cluster tightly around the observed values if the model fits well.

In [13]:

Copied!

vaxflux_idata.posterior_observations
vaxflux_idata.posterior_observations

Out[13]:

	chain	draw	season	season_start_date	season_end_date	start_date	end_date	report_date	age	type	value
0	0	0	2021/22	2021-08-01	2022-04-30	2021-08-01	2021-10-30	2021-10-30	18-49 years	incidence	0.173874
1	0	0	2021/22	2021-08-01	2022-04-30	2021-12-31	2022-01-29	2022-01-29	18-49 years	incidence	0.013714
2	0	0	2021/22	2021-08-01	2022-04-30	2022-01-29	2022-02-26	2022-02-26	18-49 years	incidence	0.003825
3	0	0	2021/22	2021-08-01	2022-04-30	2021-10-30	2021-11-27	2021-11-27	18-49 years	incidence	0.097548
4	0	0	2021/22	2021-08-01	2022-04-30	2022-02-26	2022-03-26	2022-03-26	18-49 years	incidence	0.001093
...	...	...	...	...	...	...	...	...	...	...	...
142495	1	249	2025/26	2025-08-01	2026-02-21	2025-10-18	2025-10-25	2025-10-25	65+ years	incidence	0.048348
142496	1	249	2025/26	2025-08-01	2026-02-21	2025-12-20	2025-12-27	2025-12-27	65+ years	incidence	0.005434
142497	1	249	2025/26	2025-08-01	2026-02-21	2025-11-22	2025-11-29	2025-11-29	65+ years	incidence	0.017069
142498	1	249	2025/26	2025-08-01	2026-02-21	2026-02-07	2026-02-14	2026-02-14	65+ years	incidence	0.000622
142499	1	249	2025/26	2025-08-01	2026-02-21	2025-12-06	2025-12-13	2025-12-13	65+ years	incidence	0.009799

142500 rows × 11 columns

Trace Plots for Convergence Diagnostics¶

Check MCMC chain convergence. With real data the parameter space is more complex — all three curve parameters have both season-level and age-level components — so trace plots are especially important for catching any chain mixing issues.

In [14]:

Copied!





az.plot_trace(
    vaxflux_idata,
    var_names=[
        "m_age",
        "m_age_mu",
        "m_age_sigma",
        "r_age",
        "r_age_mu",
        "r_age_sigma",
        "s_age",
        "s_age_mu",
        "s_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=[
        "m_age",
        "m_age_mu",
        "m_age_sigma",
        "r_age",
        "r_age_mu",
        "r_age_sigma",
        "s_age",
        "s_age_mu",
        "s_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()

$No description has been provided for this image$

Plot Posterior Predictive Fit¶

Overlay the posterior credible interval on the observed data. A well-fitting model should show the data within the posterior band with a narrow uncertainty region.

In [15]:

Copied!

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

Plot Posterior Predictive with Observation Noise¶

Including observation noise gives the full predictive distribution — the range of incidence values the model expects to see in future surveys, accounting for both parameter uncertainty and measurement noise.

In [16]:

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 the model's posterior estimates of season-end vaccination coverage, broken out by season and age group. These quantile bands ("low", "typical", "high") summarize how much of the population was vaccinated by the end of each season according to the fitted model.

In [17]:

Copied!

vaxflux_idata.posterior_prevalence_scenarios(
    scenarios={"low": (0.05, 0.3), "typical": (0.3, 0.7), "high": (0.7, 0.95)}
)
vaxflux_idata.posterior_prevalence_scenarios(
    scenarios={"low": (0.05, 0.3), "typical": (0.3, 0.7), "high": (0.7, 0.95)}
)

Out[17]:

	season	age	scenario	low_quantile	high_quantile	low_prevalence	high_prevalence
0	2021/22	18-49 years	low	0.05	0.30	0.322036	0.335361
1	2021/22	50-64 years	low	0.05	0.30	0.450548	0.468673
2	2021/22	65+ years	low	0.05	0.30	0.627752	0.648644
3	2022/23	18-49 years	low	0.05	0.30	0.333732	0.344840
4	2022/23	50-64 years	low	0.05	0.30	0.485962	0.499770
5	2022/23	65+ years	low	0.05	0.30	0.659864	0.671404
6	2023/24	18-49 years	low	0.05	0.30	0.320403	0.338799
7	2023/24	50-64 years	low	0.05	0.30	0.444701	0.474165
8	2023/24	65+ years	low	0.05	0.30	0.643939	0.667762
9	2024/25	18-49 years	low	0.05	0.30	0.307251	0.333225
10	2024/25	50-64 years	low	0.05	0.30	0.420386	0.456986
11	2024/25	65+ years	low	0.05	0.30	0.632744	0.668154
12	2025/26	18-49 years	low	0.05	0.30	0.306299	0.327429
13	2025/26	50-64 years	low	0.05	0.30	0.420436	0.446056
14	2025/26	65+ years	low	0.05	0.30	0.617792	0.644229
15	2021/22	18-49 years	typical	0.30	0.70	0.335361	0.348988
16	2021/22	50-64 years	typical	0.30	0.70	0.468673	0.484836
17	2021/22	65+ years	typical	0.30	0.70	0.648644	0.663516
18	2022/23	18-49 years	typical	0.30	0.70	0.344840	0.354104
19	2022/23	50-64 years	typical	0.30	0.70	0.499770	0.509800
20	2022/23	65+ years	typical	0.30	0.70	0.671404	0.681288
21	2023/24	18-49 years	typical	0.30	0.70	0.338799	0.355123
22	2023/24	50-64 years	typical	0.30	0.70	0.474165	0.497018
23	2023/24	65+ years	typical	0.30	0.70	0.667762	0.687026
24	2024/25	18-49 years	typical	0.30	0.70	0.333225	0.355773
25	2024/25	50-64 years	typical	0.30	0.70	0.456986	0.485124
26	2024/25	65+ years	typical	0.30	0.70	0.668154	0.697703
27	2025/26	18-49 years	typical	0.30	0.70	0.327429	0.346771
28	2025/26	50-64 years	typical	0.30	0.70	0.446056	0.476744
29	2025/26	65+ years	typical	0.30	0.70	0.644229	0.667800
30	2021/22	18-49 years	high	0.70	0.95	0.348988	0.363783
31	2021/22	50-64 years	high	0.70	0.95	0.484836	0.505452
32	2021/22	65+ years	high	0.70	0.95	0.663516	0.680285
33	2022/23	18-49 years	high	0.70	0.95	0.354104	0.366072
34	2022/23	50-64 years	high	0.70	0.95	0.509800	0.521680
35	2022/23	65+ years	high	0.70	0.95	0.681288	0.692829
36	2023/24	18-49 years	high	0.70	0.95	0.355123	0.368728
37	2023/24	50-64 years	high	0.70	0.95	0.497018	0.518690
38	2023/24	65+ years	high	0.70	0.95	0.687026	0.706213
39	2024/25	18-49 years	high	0.70	0.95	0.355773	0.376438
40	2024/25	50-64 years	high	0.70	0.95	0.485124	0.517415
41	2024/25	65+ years	high	0.70	0.95	0.697703	0.720351
42	2025/26	18-49 years	high	0.70	0.95	0.346771	0.366473
43	2025/26	50-64 years	high	0.70	0.95	0.476744	0.506836
44	2025/26	65+ years	high	0.70	0.95	0.667800	0.691789

In [ ]: