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,
    SeasonRange,
    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,
    SeasonRange,
    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")

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_6142/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.parameters
curve = LogisticCurve()
curve.parameters

Out[5]:

('m', 'r', 's')

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!





season_bounds = (
    observations.data.groupby("season", as_index=False)
    .agg(
        season_start_date=("season_start_date", "min"),
        season_end_date=("season_end_date", "max"),
    )
    .sort_values("season")
)
seasons = [
    SeasonRange(
        season=row.season,
        start_date=row.season_start_date.date(),
        end_date=row.season_end_date.date(),
    )
    for row in season_bounds.itertuples(index=False)
]

age_covariate_categories = CovariateCategories(
    covariate="age",
    categories=tuple(observations.data["age"].drop_duplicates().tolist()),
)
season_bounds = (
    observations.data.groupby("season", as_index=False)
    .agg(
        season_start_date=("season_start_date", "min"),
        season_end_date=("season_end_date", "max"),
    )
    .sort_values("season")
)
seasons = [
    SeasonRange(
        season=row.season,
        start_date=row.season_start_date.date(),
        end_date=row.season_end_date.date(),
    )
    for row in season_bounds.itertuples(index=False)
]

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 0x139180d10>

Run Prior Predictive Sampling¶

Draw 500 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-03-31	2021-08-01	2021-10-30	2021-10-30	18-49 years	incidence	0.122936
1	0	0	2021/22	2021-08-01	2022-03-31	2021-10-30	2021-11-27	2021-11-27	18-49 years	incidence	0.040334
2	0	0	2021/22	2021-08-01	2022-03-31	2021-11-27	2021-12-31	2021-12-31	18-49 years	incidence	0.033967
3	0	0	2021/22	2021-08-01	2022-03-31	2021-12-31	2022-01-29	2022-01-29	18-49 years	incidence	0.017063
4	0	0	2021/22	2021-08-01	2022-03-31	2022-01-29	2022-02-26	2022-02-26	18-49 years	incidence	0.009292
...	...	...	...	...	...	...	...	...	...	...	...
142495	0	499	2025/26	2025-08-01	2026-03-31	2026-01-17	2026-01-24	2026-01-24	65+ years	incidence	0.004391
142496	0	499	2025/26	2025-08-01	2026-03-31	2026-01-24	2026-01-31	2026-01-31	65+ years	incidence	0.003748
142497	0	499	2025/26	2025-08-01	2026-03-31	2026-01-31	2026-02-07	2026-02-07	65+ years	incidence	0.003195
142498	0	499	2025/26	2025-08-01	2026-03-31	2026-02-07	2026-02-14	2026-02-14	65+ years	incidence	0.002719
142499	0	499	2025/26	2025-08-01	2026-03-31	2026-02-14	2026-02-21	2026-02-21	65+ years	incidence	0.002311

142500 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 4 chains, 500 warmup steps, and 500 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-03-31	2021-10-30	2021-11-27	2021-11-27	18-49 years	incidence	0.100080
1	0	0	2021/22	2021-08-01	2022-03-31	2022-02-26	2022-03-26	2022-03-26	18-49 years	incidence	0.001350
2	0	0	2021/22	2021-08-01	2022-03-31	2022-01-29	2022-02-26	2022-02-26	18-49 years	incidence	0.004571
3	0	0	2021/22	2021-08-01	2022-03-31	2021-08-01	2021-10-30	2021-10-30	18-49 years	incidence	0.171427
4	0	0	2021/22	2021-08-01	2022-03-31	2021-11-27	2021-12-31	2021-12-31	18-49 years	incidence	0.059675
...	...	...	...	...	...	...	...	...	...	...	...
569995	3	499	2025/26	2025-08-01	2026-03-31	2025-08-01	2025-10-11	2025-10-11	65+ years	incidence	0.321021
569996	3	499	2025/26	2025-08-01	2026-03-31	2025-12-20	2025-12-27	2025-12-27	65+ years	incidence	0.005776
569997	3	499	2025/26	2025-08-01	2026-03-31	2026-02-14	2026-02-21	2026-02-21	65+ years	incidence	0.000411
569998	3	499	2025/26	2025-08-01	2026-03-31	2026-01-10	2026-01-17	2026-01-17	65+ years	incidence	0.002174
569999	3	499	2025/26	2025-08-01	2026-03-31	2026-01-17	2026-01-24	2026-01-24	65+ years	incidence	0.001561

570000 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.319810	0.336123
1	2021/22	50-64 years	low	0.05	0.30	0.449283	0.468149
2	2021/22	65+ years	low	0.05	0.30	0.629529	0.650314
3	2022/23	18-49 years	low	0.05	0.30	0.334149	0.340351
4	2022/23	50-64 years	low	0.05	0.30	0.487632	0.500387
5	2022/23	65+ years	low	0.05	0.30	0.659659	0.666099
6	2023/24	18-49 years	low	0.05	0.30	0.315925	0.340030
7	2023/24	50-64 years	low	0.05	0.30	0.446859	0.461671
8	2023/24	65+ years	low	0.05	0.30	0.646534	0.671252
9	2024/25	18-49 years	low	0.05	0.30	0.308783	0.336346
10	2024/25	50-64 years	low	0.05	0.30	0.421356	0.447332
11	2024/25	65+ years	low	0.05	0.30	0.641780	0.675849
12	2025/26	18-49 years	low	0.05	0.30	0.305617	0.331458
13	2025/26	50-64 years	low	0.05	0.30	0.416439	0.439399
14	2025/26	65+ years	low	0.05	0.30	0.620012	0.650571
15	2021/22	18-49 years	typical	0.30	0.70	0.336123	0.347397
16	2021/22	50-64 years	typical	0.30	0.70	0.468149	0.481144
17	2021/22	65+ years	typical	0.30	0.70	0.650314	0.661678
18	2022/23	18-49 years	typical	0.30	0.70	0.340351	0.353813
19	2022/23	50-64 years	typical	0.30	0.70	0.500387	0.511207
20	2022/23	65+ years	typical	0.30	0.70	0.666099	0.679575
21	2023/24	18-49 years	typical	0.30	0.70	0.340030	0.356214
22	2023/24	50-64 years	typical	0.30	0.70	0.461671	0.489910
23	2023/24	65+ years	typical	0.30	0.70	0.671252	0.696902
24	2024/25	18-49 years	typical	0.30	0.70	0.336346	0.363077
25	2024/25	50-64 years	typical	0.30	0.70	0.447332	0.480310
26	2024/25	65+ years	typical	0.30	0.70	0.675849	0.708066
27	2025/26	18-49 years	typical	0.30	0.70	0.331458	0.347512
28	2025/26	50-64 years	typical	0.30	0.70	0.439399	0.470341
29	2025/26	65+ years	typical	0.30	0.70	0.650571	0.666772
30	2021/22	18-49 years	high	0.70	0.95	0.347397	0.363827
31	2021/22	50-64 years	high	0.70	0.95	0.481144	0.502936
32	2021/22	65+ years	high	0.70	0.95	0.661678	0.681220
33	2022/23	18-49 years	high	0.70	0.95	0.353813	0.366116
34	2022/23	50-64 years	high	0.70	0.95	0.511207	0.520574
35	2022/23	65+ years	high	0.70	0.95	0.679575	0.692981
36	2023/24	18-49 years	high	0.70	0.95	0.356214	0.372592
37	2023/24	50-64 years	high	0.70	0.95	0.489910	0.517859
38	2023/24	65+ years	high	0.70	0.95	0.696902	0.711046
39	2024/25	18-49 years	high	0.70	0.95	0.363077	0.373036
40	2024/25	50-64 years	high	0.70	0.95	0.480310	0.512963
41	2024/25	65+ years	high	0.70	0.95	0.708066	0.721521
42	2025/26	18-49 years	high	0.70	0.95	0.347512	0.367037
43	2025/26	50-64 years	high	0.70	0.95	0.470341	0.500462
44	2025/26	65+ years	high	0.70	0.95	0.666772	0.692194

In [ ]: