Building Multi-Signal Renewal Models

Code

import numpyro

# to run samplers in parallel you must run `set_host_device_count` before importing jax
numpyro.set_host_device_count(4)
numpyro.enable_x64()

Code

import arviz as az
import jax
import jax.numpy as jnp
import jax.random as random
import numpy as np
import numpyro.distributions as dist
import plotnine as p9
import pandas as pd
import time
import warnings
from datetime import date

warnings.filterwarnings("ignore")

from _tutorial_theme import theme_tutorial

Code

from jax.typing import ArrayLike

from pyrenew import datasets
from pyrenew.deterministic import (
    DeterministicPMF,
    DeterministicVariable,
)
from pyrenew.metaclass import RandomVariable
from pyrenew.randomvariable import DistributionalVariable

from pyrenew.latent import (
    SubpopulationInfections,
    DifferencedAR1,
    RandomWalk,
    StepwiseTemporalProcess,
    WeeklyTemporalProcess,
    GammaGroupSdPrior,
    HierarchicalNormalPrior,
)
from pyrenew.model import PyrenewBuilder

from pyrenew.observation import (
    PopulationCounts,
    HierarchicalNormalNoise,
    MeasurementObservation,
    MeasurementNoise,
    NegativeBinomialNoise,
)
from pyrenew.time import MMWR_WEEK

Overview

Renewal models in PyRenew combine two types of components:

1. Latent infection process: Generates unobserved infections via the renewal equation, driven by a time-varying reproduction number \(\mathcal{R}(t)\)

2. Observation processes: Transform latent infections into observable signals (hospital admissions, wastewater concentrations, etc.) by applying delays, ascertainment, and noise

A multi-signal model combines multiple observation processes—each representing a different data stream, e.g., hospital admissions, emergency department visits, wastewater concentrations, which stem from the same underlying latent infection process. By jointly modeling these signals, we can improve estimation and prediction of the time-varying reproduction number \(\mathcal{R}(t)\). Such a model must:

• Generate a single coherent infection trajectory (or set of trajectories for subpopulations)
• Route those infections to each observation process appropriately
• Handle the initialization period required by delay distributions

The PyrenewBuilder class handles this plumbing. You specify:

1. A single latent process (e.g., SubpopulationInfections) that defines how infections evolve.
2. One or more observation processes (e.g., PopulationCounts, MeasurementObservation) that define how infections become data.

The builder computes initialization requirements, wires components together, and produces a model ready for inference.

Related Tutorials

Before diving into multi-signal models, you may want to review these foundational tutorials:

• Latent Infections and Latent Subpopulation Infections: Understanding temporal process choices for \(\mathcal{R}(t)\)
• Observation Processes: Counts: Modeling count data (admissions, deaths)
• Observation Processes: Measurements: Modeling continuous data (wastewater)

This tutorial shows how to combine these components into a complete multi-signal model.

What This Tutorial Covers

This tutorial demonstrates building a multi-signal renewal model using:

• SubpopulationInfections — subpopulations share a jurisdiction-level baseline \(\mathcal{R}(t)\) with subpopulation-specific deviations
• PopulationCounts — hospital admissions (jurisdiction-level)
• A custom Wastewater class — viral concentrations (subpopulation-level)

Model Structure

In this tutorial, we build a model that jointly fits two data streams to a shared latent infection process:

• Hospital admissions — jurisdiction-level counts that reflect a delayed and partially observed subset of total infections across all subpopulations
• Wastewater concentrations — site-level measurements from a subset of subpopulations (catchment areas), reflecting viral shedding and dilution

The diagram below shows how data flows through the model. The latent process generates infection trajectories for all subpopulations. Each observation process receives the infections it needs — aggregated totals or per-subpopulation arrays — and transforms them into predicted observations via delays, ascertainment, shedding kinetics, and noise.

flowchart TB

    subgraph Latent["Latent Infection Process"]
        L["Renewal equation<br/>(SubpopulationInfections)"]
    end

    subgraph Infections["Infection Trajectories"]
        J["Jurisdiction total<br/>(summed across subpopulations)"]
        S["Per-subpopulation infections<br/>(all subpopulations)"]
    end

    subgraph Obs["Observation Processes"]
        C["Hospital admissions<br/>(PopulationCounts)"]
        W["Wastewater concentrations<br/>(MeasurementObservation)"]
    end

    subgraph Data["Observed Data"]
        HA["Reported admissions"]
        WW["Measured viral concentrations"]
    end

    L --> S
    S -->|"weighted sum"| J
    J --> C
    S -->|"select monitored subpopulations"| W
    C -->|"delay + ascertainment + noise"| HA
    W -->|"shedding + dilution + noise"| WW

Infection Resolution

Different observation processes observe different levels of the model hierarchy. Each observation process declares an infection resolution that determines what infection data it receives:

Resolution	Receives	Example signals
"aggregate"	Infection counts for the whole jurisdiction, shape (T,)	Hospital admissions, case counts
"subpop"	Infection counts in each subpopulation, shape (T, n_subpops)	Regional case counts
"subpop_per_capita"	Infections per member of each subpopulation, shape (T, n_subpops)	Wastewater concentrations

The model routes latent infections to observation processes based on each process’s declared resolution.

The latent infection process generates infections per individual: \(I_k(t)\) is the number of new infections among members of subpopulation \(k\) on day \(t\), divided by the number of members of \(k\). The model converts these to the scale each kind of signal needs, using population_size (\(N\)) and subpop_fractions (\(p_k\)):

• Count signals need infection counts. Subpopulation \(k\) has \(N p_k\) members, so it has \(N p_k I_k(t)\) infections, and these counts sum to the jurisdiction total \(N \sum_k p_k I_k(t)\).
• Some measurement signals provide information about infections per member of the subpopulation. For example, the number genomes shed into wastewater scales with the number of infected individuals, and the total volume of wastewater scales with the total population, so it can be reasonable to model wastewater concentrations purely as a function of per-capita incidence \(I_k(t)\).

For subpopulation-level observations, the observation process selects which subpopulations it observes using subpop_indices provided at sample/fit time. This allows flexible observation patterns—for example, wastewater samples might only cover 5 of 6 subpopulations (catchment areas), while the 6th represents areas without wastewater monitoring.

With this structure in mind, we’ll now define each component following the generative direction: first the latent infection process, then the observation processes.

Latent Infection Process

Latent infection processes implement the renewal equation to generate infection trajectories. All latent processes share common components:

• Generation interval: PMF for secondary infection timing
• Initial infections (\(I(0)\)): Starting condition for the renewal process
• Temporal dynamics: How \(\mathcal{R}(t)\) evolves over time

Generation Interval

The generation interval PMF specifies the probability that a secondary infection occurs \(\tau\) days after the primary infection.

Code

covid_gen_int = [0.16, 0.32, 0.25, 0.14, 0.07, 0.04, 0.02]
gen_int_pmf = jnp.array(covid_gen_int)
gen_int_rv = DeterministicPMF("gen_int", gen_int_pmf)

days = np.arange(len(gen_int_pmf))
print(f"Generation interval: {gen_int_pmf}")

Generation interval: [0.16 0.32 0.25 0.14 0.07 0.04 0.02]

I0: Initial Infections

The initial infections RV I0_rv specifies the proportion of the population infected at the first observation time. This must be a value in the interval (0, 1]. We use a Beta prior centered near a small value:

Code

I0_rv = DistributionalVariable("I0", dist.Beta(1, 100))

Log Rt at time \(0\)

We place a prior on the log \(\mathcal{R}(t)\) at time \(0\), centered at \(0.0\) (\(\mathcal{R}(t) = 1.0\)) with moderate uncertainty:

Code

log_rt_time_0_rv = DistributionalVariable("log_rt_time_0", dist.Normal(0.0, 0.5))

Temporal Processes for \(\mathcal{R}(t)\)

We configure two temporal processes:

• Jurisdiction-level (baseline_rt_process): DifferencedAR(1) process for the baseline \(\mathcal{R}(t)\)
• Subpopulation-level (subpop_rt_deviation_process): RandomWalk for subpopulation deviations

The RandomWalk allows flexible evolution of subpopulation-specific transmission without mean reversion.

Code

# DifferencedAR1 allows persistent trends while stabilizing the growth rate.
baseline_rt_process = DifferencedAR1(
    autoreg_rv=DeterministicVariable("autoreg", 0.5),
    innovation_sd_rv=DeterministicVariable("innovation_sd", 0.01),
)

# RandomWalk allows flexible subpopulation deviations
subpop_rt_deviation_process = RandomWalk(
    innovation_sd_rv=DeterministicVariable("innovation_sd", 0.025),
)

Choosing the \(\mathcal{R}(t)\) Parameter Cadence

The renewal equation is evaluated on the model’s daily time axis, but the temporal process for \(\mathcal{R}(t)\) does not have to sample a new parameter every day. This separates three model choices:

• Parameter cadence: how often the \(\mathcal{R}(t)\) temporal process samples a new latent value
• Model time axis: the daily axis used by the renewal equation and delay convolutions
• Observation cadence: the temporal granularity for each signal, such as daily ED visits or weekly hospital admissions

The DifferencedAR1 process above samples one value per model day:

baseline_rt_process = DifferencedAR1(
    autoreg_rv=DeterministicVariable("autoreg", 0.5),
    innovation_sd_rv=DeterministicVariable("innovation_sd", 0.01),
)

To use a weekly \(\mathcal{R}(t)\) while still running the renewal equation daily, wrap the temporal process in WeeklyTemporalProcess. The wrapper samples a weekly trajectory and broadcasts it to the daily model axis before the latent infection process uses it.

weekly_baseline_rt_process = WeeklyTemporalProcess(
    inner=DifferencedAR1(
        autoreg_rv=DeterministicVariable("autoreg", 0.5),
        innovation_sd_rv=DeterministicVariable("innovation_sd", 0.01),
    ),
    start_dow=MMWR_WEEK,
)

Use WeeklyTemporalProcess when the weekly Rt blocks should align to a calendar week. The start_dow parameter is the day of week on which the calendar-week cycle begins (0=Monday, 6=Sunday, ISO convention). pyrenew.time exports MMWR_WEEK = 6 (Sunday-Saturday epiweeks) and ISO_WEEK = 0 (Monday-Sunday); pass any integer in [0, 6] for other conventions. At sample or run time, pass the date of the first observation day as obs_start_date. Argument obs_start_date is required whenever any component performs calendar-aligned work: a WeeklyTemporalProcess, a count observation with aggregation="weekly", or any observation with a day-of-week effect. For StepwiseTemporalProcess and daily observations with no day-of-week effect, obs_start_date can be omitted. The model handles the calendar bookkeeping and forwards the day-of-week information to every component that needs it.

Observation Processes

Observation processes transform latent infections into observable signals and define the statistical model linking predictions to data. Each observation process:

• Has a unique name that identifies the signal in model outputs
• Declares what infection resolution it needs ("aggregate", "subpop", or "subpop_per_capita")
• Applies signal-specific transformations (ascertainment, delay convolution, shedding kinetics)
• Defines the noise model

Signal Naming

Each observation process requires a name parameter—a short, meaningful identifier like "hospital" or "wastewater". This name serves as the single identifier for the signal throughout the model:

• Numpyro sites: Prefixes all sample and deterministic sites (e.g., hospital_obs, hospital_predicted)
• Data binding: Becomes the keyword argument for passing data to model.run() (e.g., hospital={...})

This unified naming provides several benefits:

• Interpretable outputs: When examining MCMC samples or posterior diagnostics, site names like hospital_predicted immediately indicate which signal each quantity refers to
• Multiple signals of the same type: You can include multiple count observations (e.g., hospital admissions and deaths) by giving each a distinct name
• Clearer debugging: Error messages and trace inspection show meaningful signal names rather than generic identifiers

Hospital Admissions

In this example we use a dataset consisting of hospital admissions for COVID-19 across California for the first 10 months of 2023 (as reported to the CDC).

Code

# Load daily hospital admissions for California
ca_hosp_data = datasets.load_hospital_data_for_state("CA", "2023-11-06.csv")
obs_start_date = ca_hosp_data["dates"][0]
hosp_admits = ca_hosp_data["daily_admits"]
population_size = ca_hosp_data["population"]
n_hosp_days = ca_hosp_data["n_days"]

print("State: California")
print(f"Population: {population_size:,}")
print(f"Date range: {ca_hosp_data['dates'][0]} to {ca_hosp_data['dates'][-1]}")
print(f"Number of days: {n_hosp_days}")
print(f"Admissions range: {int(hosp_admits.min())} to {int(hosp_admits.max())}")

State: California
Population: 39,512,223
Date range: 2023-01-01 to 2023-11-06
Number of days: 310
Admissions range: 2 to 2574

The hospital admissions data is aggregated at the jurisdiction level, therefore we specify a PopulationCounts observation process.

Code

# Infection-to-hospitalization delay (COVID-19, from literature)
inf_to_hosp_pmf = jnp.array(
    [
        0,
        0.00469,
        0.01452,
        0.02786,
        0.04237,
        0.05581,
        0.06657,
        0.07379,
        0.07729,
        0.07737,
        0.07465,
        0.06988,
        0.06377,
        0.05696,
        0.04996,
        0.04315,
        0.03677,
        0.03097,
        0.02583,
        0.02135,
        0.01751,
        0.01427,
        0.01156,
        0.00931,
        0.00746,
        0.00596,
        0.00474,
        0.00375,
        0.00296,
        0.00233,
        0.00183,
        0.00143,
        0.00107,
        0.00077,
        0.00054,
        0.00036,
        0.00024,
        0.00015,
        0.00009,
        0.00005,
        0.00003,
        0.00002,
        0.00001,
    ]
)
hosp_delay_rv = DeterministicPMF("inf_to_hosp_delay", inf_to_hosp_pmf)

# IHR: ~1% of infections lead to hospitalization
ihr_rv = DeterministicVariable("ihr", 0.01)

# Negative binomial concentration (moderate overdispersion)
hosp_concentration_rv = DeterministicVariable("hosp_concentration", 10.0)

# Create hospital observation process
hosp_obs = PopulationCounts(
    name="hospital",
    ascertainment_rate_rv=ihr_rv,
    delay_distribution_rv=hosp_delay_rv,
    noise=NegativeBinomialNoise(hosp_concentration_rv),
)

print("Hospital observation:")
print(f"  Infection resolution: {hosp_obs.infection_resolution()}")
print(f"  Delay PMF length: {len(inf_to_hosp_pmf)} days")

Hospital observation:
  Infection resolution: aggregate
  Delay PMF length: 43 days

Wastewater Concentrations

Wastewater Observation Process

The MeasurementObservation base class handles continuous observation processes. Domain-specific implementations subclass it and implement _predicted_obs() to transform infections into predicted values. See observation_processes_measurements.md for a detailed tutorial.

Code

class Wastewater(MeasurementObservation):
    """
    Wastewater viral concentration observation process.

    Transforms infections per member of each subpopulation into
    predicted log-concentrations via shedding kinetics convolution and
    genome/volume scaling.
    """

    def __init__(
        self,
        name: str,
        shedding_kinetics_rv: RandomVariable,
        log10_genome_per_infection_rv: RandomVariable,
        ml_per_person_per_day: float,
        noise: MeasurementNoise,
    ) -> None:
        super().__init__(name=name, temporal_pmf_rv=shedding_kinetics_rv, noise=noise)
        self.log10_genome_per_infection_rv = log10_genome_per_infection_rv
        self.ml_per_person_per_day = ml_per_person_per_day

    def validate(self) -> None:
        shedding_pmf = self.temporal_pmf_rv()
        self._validate_pmf(shedding_pmf, "shedding_kinetics_rv")
        self.noise.validate()

    def lookback_days(self) -> int:
        return len(self.temporal_pmf_rv()) - 1

    def _predicted_obs(self, infections: ArrayLike) -> ArrayLike:
        shedding_pmf = self.temporal_pmf_rv()
        log10_genome = self.log10_genome_per_infection_rv()

        def convolve_site(site_infections):
            convolved, _ = self._convolve_delay_with_alignment(
                site_infections, shedding_pmf
            )
            return convolved

        shedding_signal = jax.vmap(convolve_site, in_axes=1, out_axes=1)(infections)
        genome_copies = 10**log10_genome
        concentration = shedding_signal * genome_copies / self.ml_per_person_per_day
        return jnp.log(concentration)

Wastewater Data

For the wastewater data, we use a simulated dataset for California with realistic noise patterns that covers the same time period.

Code

# Load wastewater data for California
ca_ww_data = datasets.load_wastewater_data_for_state("CA", "fake_nwss.csv")

ww_conc = ca_ww_data["observed_conc"]  # log concentrations
ww_site_ids = ca_ww_data["site_ids"]
ww_time_indices = ca_ww_data["time_indices"]
ww_n_sites = ca_ww_data["n_sites"]
ww_n_obs = ca_ww_data["n_obs"]
ww_wwtp_names = ca_ww_data["wwtp_names"]

print("State: California")
print(f"Number of sites: {ww_n_sites}")
print(f"Number of observations: {ww_n_obs}")
print(f"Date range: {ca_ww_data['dates'][0]} to {ca_ww_data['dates'][-1]}")
print(f"Time index range: {int(ww_time_indices.min())} to {int(ww_time_indices.max())}")
print("\nSites:")
for i, name in enumerate(ww_wwtp_names[:5]):
    print(f"  {i}: {name}")
if ww_n_sites > 5:
    print(f"  ... and {ww_n_sites - 5} more")

State: California
Number of sites: 5
Number of observations: 1495
Date range: 2023-01-01 to 2023-11-06
Time index range: 0 to 309

Sites:
  0: 1
  1: 2
  2: 3
  3: 4
  4: 5

Wastewater observations are site-level: each measurement is associated with a specific measurement site. The Wastewater observation process uses LogNormalNoise, which takes hierarchical priors for the site-level mode and standard deviation parameters. This enables partial pooling across measurement sites.

Here we specify HierarchicalNormalPrior for the site-level mode and GammaGroupSdPrior for the standard deviation.

Code

# Viral shedding kinetics PMF (days post-infection)
shedding_pmf = jnp.array(
    [
        0.0,
        0.02,
        0.08,
        0.15,
        0.20,
        0.18,
        0.14,
        0.10,
        0.06,
        0.04,
        0.02,
        0.01,
    ]
)
shedding_pmf = shedding_pmf / shedding_pmf.sum()  # normalize
shedding_rv = DeterministicPMF("shedding_kinetics", shedding_pmf)

# Log10 genomes shed per infection
log10_genome_rv = DeterministicVariable("log10_genome_per_inf", 4.1)

# Wastewater volume per person per day (mL)
ml_per_person_per_day = 1000.0

# Hierarchical priors for site-level effects
site_mode_prior = HierarchicalNormalPrior(
    "ww_site_mode", sd_rv=DeterministicVariable("site_mode_sd", 0.5)
)
site_sd_prior = GammaGroupSdPrior(
    "ww_site_sd",
    sd_mean_rv=DeterministicVariable("site_sd_mean", 0.3),
    sd_concentration_rv=DeterministicVariable("site_sd_conc", 4.0),
)

# Create wastewater observation process
ww_obs = Wastewater(
    name="wastewater",
    shedding_kinetics_rv=shedding_rv,
    log10_genome_per_infection_rv=log10_genome_rv,
    ml_per_person_per_day=ml_per_person_per_day,
    noise=HierarchicalNormalNoise(site_mode_prior, site_sd_prior),
)

print("Wastewater observation:")
print(f"  Infection resolution: {ww_obs.infection_resolution()}")
print(f"  Shedding PMF length: {len(shedding_pmf)} days")

Wastewater observation:
  Infection resolution: subpop_per_capita
  Shedding PMF length: 12 days

Model Building

We instantiate a PyrenewBuilder object which handles the composition of the latent infection process and the observation process.

Code

# Build the multi-signal model
builder = PyrenewBuilder()

The PyrenewBuilder object has 3 key methods:

• configure_latent
• add_observation
• build

Methods configure_latent and add_observation can be called in any order. Method build is called once all processes have been specified in the model.

Configuring the Latent Process

We use configure_latent to specify the model structure: generation interval, initial infections, and temporal dynamics.

Code

print("Latent process configuration:")
print(f"  Generation interval length: {len(gen_int_rv())} days")

builder.configure_latent(
    SubpopulationInfections,
    gen_int_rv=gen_int_rv,
    I0_rv=I0_rv,
    log_rt_time_0_rv=log_rt_time_0_rv,
    baseline_rt_process=baseline_rt_process,
    subpop_rt_deviation_process=subpop_rt_deviation_process,
)

Latent process configuration:
  Generation interval length: 7 days

<pyrenew.model.pyrenew_builder.pyrenewbuilder 0x7f82da6e1400="" at="">

Specifying the Observation Processes, Data

Each observation process’s name attribute becomes the keyword used to pass that observation’s data to model.run() (e.g., hospital={...}, wastewater={...}).

Code

builder.add_observation(hosp_obs)  # Uses hosp_obs.name = "hospital"
builder.add_observation(ww_obs)  # Uses ww_obs.name = "wastewater"
model = builder.build()

n_init = model.latent.n_initialization_points
print("Model built successfully")
print(f"  n_initialization_points: {n_init}")
print(f"  Latent process: {type(model.latent).__name__}")
print(f"  Observation processes: {list(model.observations.keys())}")

Model built successfully
  n_initialization_points: 42
  Latent process: SubpopulationInfections
  Observation processes: ['hospital', 'wastewater']

Model Identifiability.

The renewal equation is linear in the initial infections I0_rv, so scaling I0_rv by a factor \(c\) scales the entire infection trajectory by \(c\). In practice, I0_rv is weakly identified because each observation process links infections to data through a signal-specific ascertainment rate \(\alpha_s\) — the probability that an infection is observed as an event in signal \(s\). Doubling I0_rv while halving all ascertainment rates produces identical expected observations. Without external information to anchor either the ascertainment rates or the absolute infection level, the data cannot distinguish “more infections, lower ascertainment” from “fewer infections, higher ascertainment.” The priors on I0_rv and on the ascertainment rates resolve this ambiguity.

Fitting the Model to Data: model.run()

When you call model.run(), you supply three types of information:

• Model and population information — the fitting period, total population size, and subpopulation fractions
• Observation data — one data dictionary per registered observation process
• MCMC controls — basic settings for posterior sampling

For this model, a complete call has the following structure:

model.run(
    num_warmup=500,
    num_samples=500,
    rng_key=random.PRNGKey(42),
    mcmc_args={"num_chains": 4, "progress_bar": False},
    n_days_post_init=n_days,
    population_size=population_size,
    subpop_fractions=subpop_fractions,
    **obs_data,
)
samples = model.mcmc.get_samples()

The **obs_data dictionary supplies the observation start date and one data dictionary per observation process. The keys of the observation dictionaries are the names registered on the builder: hospital and wastewater.

Model Time vs. Observation Time

In the raw data, day 0 is the first day with observations. PyRenew adds an initialization period before this date so the renewal process has enough past infections to generate infections during the observed period.

This means the model has two related time scales:

• Observation time starts at the first day in the data.
• Model time starts n_init days earlier.

On the model time axis, indices 0 through n_init - 1 are initialization days. The first observed day is model index n_init, and the last observed day is model index n_init + n_days - 1.

The model provides helper methods to convert observation data onto this model time axis:

• model.pad_observations(obs) prepends n_init NaN values to dense daily observation vectors.
• model.shift_times(times) adds n_init to sparse observation time indices.

The hospital admissions signal is dense: it has one value per observation day, so we pad it. The wastewater signal is sparse: each row has an observation day, so we shift its time index.

Population Structure

First we declare the population structure. We have 6 subpopulations, where 5 have wastewater monitoring and 1 does not. The subpopulations with wastewater monitoring need not be contiguous indices; they can be any subset of {0, 1, …, n_subpops-1}.

Code

# All 6 subpopulations with their population fractions
subpop_fractions = jnp.array([0.10, 0.14, 0.21, 0.22, 0.07, 0.26])

n_subpops = len(subpop_fractions)

# Which subpopulations have wastewater monitoring?
# These indices will be used in subpop_indices for wastewater observations.
# They can be any subset of {0, 1, ..., n_subpops-1}, not necessarily contiguous.
ww_monitored_subpops = jnp.array(
    [0, 1, 2, 3, 4]
)  # subpop 5 has no wastewater monitoring

print(f"Total subpopulations: {n_subpops}")
print(f"Subpopulations with wastewater monitoring: {list(ww_monitored_subpops)}")
print(
    f"Wastewater coverage: {float(jnp.sum(subpop_fractions[ww_monitored_subpops])):.0%}"
)
print(f"Total population: {float(jnp.sum(subpop_fractions)):.0%}")

Total subpopulations: 6
Subpopulations with wastewater monitoring: [Array(0, dtype=int64), Array(1, dtype=int64), Array(2, dtype=int64), Array(3, dtype=int64), Array(4, dtype=int64)]
Wastewater coverage: 74%
Total population: 100%

The subpop_fractions array defines all subpopulations in the latent infection process, and its entries must sum to 1.0. Which subpopulations each observation process sees is determined by the observation data.

Measurement data typically exhibits sensor-level variability: different instruments, labs, or sampling locations may have systematic biases and different precision levels. The wastewater observations record both the wastewater treatment plant from which the sample was collected and the laboratory at which the sample was processed. For wastewater, a “sensor” is a WWTP/lab pair - the combination of treatment plant and laboratory associated with sensor-specific bias and measurement variability. Because the wastewater component’s HierarchicalNormalNoise samples one mode and one standard deviation parameter. Therefore, the wastewater signal’s data dictionary includes:

• subpop_indices links each observation to the appropriate subpopulation
• sensor_indices selects the sensor-specific noise parameters.
• n_sensors the total number of sensors, used to size the sensor-level noise parameters.

Preparing Observation Data

The observation data dictionary is structured to match the keyword arguments of model.run(): obs_start_date at the top level, plus one sub-dictionary per observation process. The entries in the sub-dictionaries are forwarded as arguments to each observation process’s sample method.

[!WARNING]

Nest observation-specific arguments under the signal name

Arguments such as subpop_indices must be placed inside the dictionary for the corresponding registered observation process. Passing them at the top level does not associate them with that signal and raises a validation error before MCMC begins.

# Correct: observation-specific metadata is nested under the signal name.
model.run(
    ...,
    regional_counts={
        "obs": counts,
        "subpop_indices": subpop_indices,
    },
)

# Incorrect: this top-level argument is not associated with regional_counts.
model.run(
    ...,
    regional_counts={"obs": counts},
    subpop_indices=subpop_indices,
)

The two observation streams are aligned differently because they are represented differently. Hospital admissions are a dense daily series: every day in the fitting window has a position in the vector, so model.pad_observations() prepends n_init missing values for the initialization period. Wastewater measurements are sparse rows: only sampled days appear in the table, so the observations themselves are not padded. Instead, model.shift_times() adds n_init to each wastewater time index so those rows point to the correct days on the model’s shared time axis.

Code

n_days_90 = 90

# Hospital: dense, NaN-padded to length n_total
hosp_obs = model.pad_observations(hosp_admits[:n_days_90])

# Wastewater: sparse, times shifted to shared axis
ww_mask = ca_ww_data["time_indices"] < n_days_90
ww_times = model.shift_times(ca_ww_data["time_indices"][ww_mask])
ww_conc = ca_ww_data["observed_conc"][ww_mask]
ww_sites = ca_ww_data["site_ids"][ww_mask]

# Map wastewater sensors to subpopulation indices.
# In practice, this mapping comes from your data
# (which WWTP serves which catchment). For this demo, we cycle
# sensors through the monitored subpopulations.
n_ww_sites = ca_ww_data["n_sites"]
n_monitored = len(ww_monitored_subpops)
sensor_to_subpop = {
    i: int(ww_monitored_subpops[i % n_monitored]) for i in range(n_ww_sites)
}
ww_subpop_indices = jnp.array([sensor_to_subpop[int(s)] for s in ww_sites])

obs_data_90 = {
    "obs_start_date": obs_start_date,
    "hospital": {
        "obs": hosp_obs,
    },
    "wastewater": {
        "obs": ww_conc,
        "times": ww_times,
        "subpop_indices": ww_subpop_indices,
        "sensor_indices": ww_sites,
        "n_sensors": n_ww_sites,
    },
}

The observation dictionary has the following structure:

{
    "obs_start_date": ...,
    "hospital": {
        "obs": ...
    },
    "wastewater": {
        "obs": ...
        "times": ...
        "subpop_indices": ...
        "sensor_indices": ...
        "n_sensors": ...
    },
}

At call time, we unpack this dictionary with **, forwarding each entry as a keyword argument to model.run(). Note that we could also have manually passed each dictionary entry to model.run(). We chose to create (and then unpack) an obs_data dictionary to keep our data organized in one place.

MCMC Controls

model.run() uses NumPyro’s No-U-Turn Sampler (NUTS) through numpyro.infer.MCMC. For this introductory example, the main controls are:

• num_warmup, the number of warmup iterations
• num_samples, the number of posterior samples to keep
• rng_key, the random seed for sampling
• mcmc_args, optional NumPyro MCMC settings

For this fit, we use 4 chains and turn off the progress bar:

mcmc_args = {"num_chains": 4, "progress_bar": False}

See the NumPyro MCMC reference for more advanced MCMC options.

Running the Model

We now unpack the observation dictionary prepared above into model.run(). The fit uses 90 days of data.

Fit: 90 Days

Code

# Clear JAX caches to avoid interference from earlier cells
jax.clear_caches()

print(f"Fitting model with {n_days_90} days of data...")
print("  This may take a few minutes...")

start_time = time.time()
model.run(
    num_warmup=500,
    num_samples=500,
    rng_key=random.PRNGKey(42),
    mcmc_args={"num_chains": 4, "progress_bar": False},
    n_days_post_init=n_days_90,
    population_size=population_size,
    subpop_fractions=subpop_fractions,
    **obs_data_90,
)
# JAX uses asynchronous dispatch, so we must block until sampling completes
# to get accurate timing
samples_90 = model.mcmc.get_samples()
jax.block_until_ready(samples_90)
elapsed_90 = time.time() - start_time
print(f"Elapsed time: {elapsed_90:.1f} seconds")

Fitting model with 90 days of data...
  This may take a few minutes...
Elapsed time: 76.4 seconds

We use ArviZ to assess MCMC convergence and mixing via the \(\hat{R}\) statistic and effective sample size (ESS). Before running these diagnostics, we trim the first n_init time steps from all time-series variables. Since the model cannot estimate latent infections until it has seen a full generation interval’s worth of data, these early time steps have no meaningful epidemiological interpretation and therefore should be excluded from summaries and visualizations.

Code

def trim_time(ds):
    """Trim first n_init entries from time dimension and reindex."""
    if "time" in ds.dims:
        ds = ds.isel(time=slice(n_init, None))
        ds = ds.assign_coords(time=range(ds.sizes["time"]))
    return ds

We first label the time axis of all time-series variables using the dims argument to az.from_numpyro. Then we trim the first n_init time steps from all time-series variables. Finally we call the az.summary report.

Code

idata_90 = az.from_numpyro(
    model.mcmc,
    dims={
        "latent_infections": ["time"],
        "SubpopulationInfections::infections_aggregate": ["time"],
        "SubpopulationInfections::log_rt_baseline": ["time", "dummy"],
        "SubpopulationInfections::rt_baseline": ["time", "dummy"],
        "SubpopulationInfections::rt_subpop": ["time", "subpop"],
        "SubpopulationInfections::subpop_deviations": ["time", "subpop"],
        "latent_infections_by_subpop": ["time", "subpop"],
        "hospital_predicted": ["time"],
        "wastewater_predicted": ["time", "subpop"],
    },
)

idata_90_trimmed = idata_90.map_over_datasets(trim_time)
az.summary(idata_90_trimmed, var_names=["latent_infections", "hospital_predicted"])

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
latent_infections[0]	1.17e+05	1.04e+04	100000	130000	1598	1798	1.00	260	190
latent_infections[1]	1.16e+05	1.06e+04	100000	130000	1584	1922	1.00	270	200
latent_infections[2]	1.16e+05	1.02e+04	100000	130000	1632	1875	1.00	250	180
latent_infections[3]	1.16e+05	9.85e+03	100000	130000	1728	1735	1.00	240	170
latent_infections[4]	1.17e+05	9.74e+03	100000	130000	1859	1872	1.00	230	160
latent_infections[5]	1.17e+05	9.70e+03	100000	130000	2041	1982	1.00	220	160
latent_infections[6]	1.18e+05	9.81e+03	100000	130000	2253	1959	1.00	210	170
latent_infections[7]	1.19e+05	9.88e+03	100000	140000	2434	1915	1.00	200	180
latent_infections[8]	1.20e+05	1.00e+04	110000	140000	2519	1910	1.00	200	200
latent_infections[9]	1.22e+05	1.01e+04	110000	140000	2554	1867	1.00	200	200
latent_infections[10]	1.23e+05	1.02e+04	110000	140000	2621	1895	1.00	200	190
latent_infections[11]	1.24e+05	1.03e+04	110000	140000	2672	1866	1.00	200	200
latent_infections[12]	1.26e+05	1.04e+04	110000	140000	2702	1920	1.00	200	190
latent_infections[13]	1.27e+05	1.04e+04	110000	140000	2653	1869	1.00	200	190
latent_infections[14]	1.29e+05	1.05e+04	110000	150000	2625	1976	1.00	210	190
latent_infections[15]	1.31e+05	1.06e+04	110000	150000	2531	1847	1.00	210	190
latent_infections[16]	1.32e+05	1.05e+04	120000	150000	2468	2035	1.00	210	180
latent_infections[17]	1.34e+05	1.06e+04	120000	150000	2374	1885	1.00	220	170
latent_infections[18]	1.35e+05	1.07e+04	120000	150000	2342	1831	1.00	220	170
latent_infections[19]	1.36e+05	1.07e+04	120000	150000	2319	1837	1.00	220	170
latent_infections[20]	1.37e+05	1.07e+04	120000	160000	2286	1702	1.00	220	170
latent_infections[21]	1.38e+05	1.06e+04	120000	160000	2255	1795	1.00	220	170
latent_infections[22]	1.38e+05	1.05e+04	120000	160000	2192	1796	1.00	230	180
latent_infections[23]	1.39e+05	1.05e+04	120000	160000	2177	1923	1.00	230	180
latent_infections[24]	1.39e+05	1.06e+04	120000	160000	2149	1820	1.00	230	190
latent_infections[25]	1.38e+05	1.06e+04	120000	160000	2115	1852	1.00	230	190
latent_infections[26]	1.37e+05	1.05e+04	120000	160000	2073	1833	1.00	230	200
latent_infections[27]	1.36e+05	1.05e+04	120000	150000	2114	1778	1.00	230	200
latent_infections[28]	1.35e+05	1.04e+04	120000	150000	2096	1913	1.00	230	200
latent_infections[29]	1.33e+05	1.03e+04	120000	150000	2079	1822	1.00	230	200
latent_infections[30]	1.32e+05	1.02e+04	120000	150000	2112	1870	1.00	220	200
latent_infections[31]	1.30e+05	1.01e+04	110000	150000	2174	1857	1.00	220	190
latent_infections[32]	1.27e+05	9.90e+03	110000	140000	2217	1919	1.00	210	180
latent_infections[33]	1.25e+05	9.71e+03	110000	140000	2230	2039	1.00	210	170
latent_infections[34]	1.22e+05	9.49e+03	110000	140000	2359	2028	1.00	200	160
latent_infections[35]	1.20e+05	9.36e+03	110000	140000	2425	2008	1.00	190	160
latent_infections[36]	1.17e+05	9.21e+03	100000	130000	2514	2095	1.00	180	160
latent_infections[37]	1.15e+05	9.07e+03	100000	130000	2605	2075	1.00	180	160
latent_infections[38]	1.13e+05	8.91e+03	99000	130000	2724	2185	1.00	170	160
latent_infections[39]	1.11e+05	8.73e+03	98000	130000	2747	1881	1.00	170	170
latent_infections[40]	1.09e+05	8.57e+03	96000	120000	2744	1704	1.00	170	160
latent_infections[41]	1.07e+05	8.43e+03	95000	120000	2608	1794	1.00	170	150
latent_infections[42]	1.06e+05	8.32e+03	94000	120000	2698	1827	1.00	160	150
latent_infections[43]	1.05e+05	8.18e+03	93000	120000	2796	1835	1.00	160	150
latent_infections[44]	1.04e+05	8.08e+03	92000	120000	2741	1845	1.00	160	150
latent_infections[45]	1.04e+05	8.01e+03	92000	120000	2735	1873	1.00	150	150
latent_infections[46]	1.04e+05	8.03e+03	92000	120000	2675	1930	1.00	160	150
latent_infections[47]	1.04e+05	8.04e+03	92000	120000	2552	1815	1.00	160	150
latent_infections[48]	1.05e+05	8.09e+03	92000	120000	2496	1927	1.00	160	150
latent_infections[49]	1.05e+05	8.18e+03	93000	120000	2511	2004	1.00	160	160
latent_infections[50]	1.06e+05	8.28e+03	93000	120000	2508	2216	1.00	170	160
latent_infections[51]	1.07e+05	8.35e+03	94000	120000	2477	2110	1.00	170	150
latent_infections[52]	1.08e+05	8.49e+03	95000	120000	2539	1887	1.00	170	160
latent_infections[53]	1.09e+05	8.65e+03	96000	120000	2529	2066	1.00	170	160
latent_infections[54]	1.10e+05	8.70e+03	97000	120000	2453	1947	1.00	180	170
latent_infections[55]	1.12e+05	8.84e+03	98000	130000	2380	1852	1.00	180	170
latent_infections[56]	1.13e+05	8.99e+03	99000	130000	2360	1837	1.00	180	170
latent_infections[57]	1.15e+05	9.11e+03	100000	130000	2312	1926	1.00	190	180
latent_infections[58]	1.16e+05	9.24e+03	100000	130000	2242	1733	1.00	200	180
latent_infections[59]	1.18e+05	9.40e+03	100000	130000	2243	1717	1.00	200	180
latent_infections[60]	1.19e+05	9.54e+03	100000	140000	2217	1848	1.00	200	180
latent_infections[61]	1.21e+05	9.77e+03	110000	140000	2249	1882	1.00	210	180
latent_infections[62]	1.23e+05	9.99e+03	110000	140000	2330	1897	1.00	210	190
latent_infections[63]	1.24e+05	1.02e+04	110000	140000	2482	1816	1.00	210	200
latent_infections[64]	1.26e+05	1.04e+04	110000	140000	2586	1885	1.00	210	210
latent_infections[65]	1.28e+05	1.06e+04	110000	150000	2744	1769	1.00	200	210
latent_infections[66]	1.30e+05	1.08e+04	110000	150000	2747	1891	1.00	210	210
latent_infections[67]	1.32e+05	1.10e+04	120000	150000	2739	1984	1.00	210	210
latent_infections[68]	1.34e+05	1.13e+04	120000	150000	2771	1879	1.00	220	220
latent_infections[69]	1.36e+05	1.16e+04	120000	150000	2765	1888	1.00	220	230
latent_infections[70]	1.38e+05	1.19e+04	120000	160000	2868	1707	1.00	220	240
latent_infections[71]	1.40e+05	1.23e+04	120000	160000	2802	1762	1.00	230	250
latent_infections[72]	1.43e+05	1.26e+04	120000	160000	2643	2044	1.00	250	260
latent_infections[73]	1.46e+05	1.31e+04	130000	170000	2480	1975	1.00	260	260
latent_infections[74]	1.49e+05	1.37e+04	130000	170000	2267	1715	1.00	290	260
latent_infections[75]	1.53e+05	1.45e+04	130000	180000	2173	1688	1.00	310	270
latent_infections[76]	1.56e+05	1.57e+04	130000	180000	2030	1743	1.00	350	290
latent_infections[77]	1.60e+05	1.73e+04	130000	190000	1944	1558	1.00	390	330
latent_infections[78]	1.65e+05	1.93e+04	140000	200000	1965	1635	1.00	440	370
latent_infections[79]	1.69e+05	2.18e+04	140000	210000	2022	1709	1.00	490	420
latent_infections[80]	1.75e+05	2.51e+04	140000	220000	2038	1882	1.00	560	470
latent_infections[81]	1.80e+05	2.90e+04	140000	230000	2045	1908	1.00	640	540
latent_infections[82]	1.87e+05	3.37e+04	140000	250000	2146	2053	1.00	740	630
latent_infections[83]	1.94e+05	3.92e+04	140000	260000	2195	1988	1.00	860	720
latent_infections[84]	2.02e+05	4.55e+04	140000	280000	2310	2118	1.00	970	860
latent_infections[85]	2.10e+05	5.30e+04	140000	300000	2343	2203	1.00	1100	1000
latent_infections[86]	2.20e+05	6.17e+04	140000	330000	2404	2129	1.00	1300	1200
latent_infections[87]	2.30e+05	7.16e+04	140000	360000	2433	2006	1.00	1500	1500
latent_infections[88]	2.42e+05	8.33e+04	140000	390000	2524	2091	1.00	1700	1800
latent_infections[89]	2.55e+05	9.67e+04	140000	430000	2571	1956	1.00	2000	2200
hospital_predicted[0]	1508	188	1200	1800	2324	1919	1.00	3.9	3.6
hospital_predicted[1]	1467	166	1200	1700	2384	1961	1.00	3.4	3.2
hospital_predicted[2]	1428	147	1200	1700	2428	1943	1.00	3	2.8
hospital_predicted[3]	1391	131	1200	1600	2436	2001	1.00	2.7	2.6
hospital_predicted[4]	1358	119	1200	1600	2409	1894	1.00	2.5	2.4
hospital_predicted[5]	1327	110	1200	1500	2324	1841	1.00	2.3	2.2
hospital_predicted[6]	1301	103	1100	1500	2214	1784	1.00	2.2	2
hospital_predicted[7]	1277	98	1100	1400	2102	1801	1.00	2.2	1.9
hospital_predicted[8]	1257	95	1100	1400	1963	1801	1.00	2.1	1.8
hospital_predicted[9]	1241	92	1100	1400	1788	1783	1.00	2.1	1.7
hospital_predicted[10]	1228	89	1100	1400	1687	1802	1.00	2.2	1.6
hospital_predicted[11]	1218	87	1100	1400	1613	1802	1.00	2.2	1.6
hospital_predicted[12]	1211	86	1100	1400	1600	1881	1.00	2.1	1.5
hospital_predicted[13]	1207	85	1100	1300	1624	1829	1.00	2.1	1.5
hospital_predicted[14]	1205	84	1100	1300	1688	1647	1.00	2	1.4
hospital_predicted[15]	1205	83	1100	1300	1812	1777	1.00	1.9	1.4
hospital_predicted[16]	1208	83	1100	1300	1915	1860	1.00	1.9	1.4
hospital_predicted[17]	1213	83	1100	1400	2055	1902	1.00	1.8	1.4
hospital_predicted[18]	1219	84	1100	1400	2211	1933	1.00	1.8	1.4
hospital_predicted[19]	1227	84	1100	1400	2227	1779	1.00	1.8	1.4
hospital_predicted[20]	1236	85	1100	1400	2300	1841	1.00	1.8	1.5
hospital_predicted[21]	1246	85	1100	1400	2359	1824	1.00	1.8	1.5
hospital_predicted[22]	1257	86	1100	1400	2392	1933	1.00	1.8	1.5
hospital_predicted[23]	1268	87	1100	1400	2397	1927	1.00	1.8	1.5
hospital_predicted[24]	1280	87	1100	1400	2401	1872	1.00	1.8	1.5
hospital_predicted[25]	1291	88	1200	1400	2375	1908	1.00	1.8	1.5
hospital_predicted[26]	1303	88	1200	1400	2334	1871	1.00	1.8	1.5
hospital_predicted[27]	1313	88	1200	1500	2280	1829	1.00	1.8	1.5
hospital_predicted[28]	1323	88	1200	1500	2217	1859	1.00	1.9	1.5
hospital_predicted[29]	1332	88	1200	1500	2163	1788	1.00	1.9	1.5
hospital_predicted[30]	1339	88	1200	1500	2107	1676	1.00	1.9	1.5
hospital_predicted[31]	1345	88	1200	1500	2062	1679	1.00	1.9	1.5
hospital_predicted[32]	1348	87	1200	1500	2015	1898	1.00	2	1.5
hospital_predicted[33]	1350	87	1200	1500	1971	1810	1.00	2	1.6
hospital_predicted[34]	1349	87	1200	1500	1943	1860	1.00	2	1.6
hospital_predicted[35]	1347	86	1200	1500	1921	1819	1.00	2	1.6
hospital_predicted[36]	1341	86	1200	1500	1913	1677	1.00	2	1.6
hospital_predicted[37]	1334	85	1200	1500	1918	1745	1.00	2	1.6
hospital_predicted[38]	1325	85	1200	1500	1921	1700	1.00	1.9	1.6
hospital_predicted[39]	1313	84	1200	1500	1937	1691	1.00	1.9	1.6
hospital_predicted[40]	1299	83	1200	1400	1948	1594	1.00	1.9	1.6
hospital_predicted[41]	1284	82	1200	1400	1975	1617	1.00	1.9	1.5
hospital_predicted[42]	1268	81	1100	1400	2007	1538	1.00	1.8	1.5
hospital_predicted[43]	1250	80	1100	1400	2046	1613	1.00	1.8	1.5
hospital_predicted[44]	1232	79	1100	1400	2089	1751	1.00	1.7	1.4
hospital_predicted[45]	1214	78	1100	1300	2141	1853	1.00	1.7	1.4
hospital_predicted[46]	1196	77	1100	1300	2189	1989	1.00	1.7	1.4
hospital_predicted[47]	1178	76	1100	1300	2246	1962	1.00	1.6	1.4
hospital_predicted[48]	1161	75	1100	1300	2294	1980	1.00	1.6	1.4
hospital_predicted[49]	1145	74	1000	1300	2335	1894	1.00	1.5	1.4
hospital_predicted[50]	1130	73	1000	1200	2377	1963	1.00	1.5	1.4
hospital_predicted[51]	1117	72	1000	1200	2395	1965	1.00	1.5	1.4
hospital_predicted[52]	1106	71	1000	1200	2419	1909	1.00	1.5	1.3
hospital_predicted[53]	1096	71	990	1200	2414	1834	1.00	1.5	1.3
hospital_predicted[54]	1088	70	980	1200	2400	1831	1.00	1.4	1.3
hospital_predicted[55]	1082	70	980	1200	2388	1844	1.00	1.4	1.3
hospital_predicted[56]	1079	70	970	1200	2375	1865	1.00	1.4	1.3
hospital_predicted[57]	1077	70	970	1200	2359	1849	1.00	1.4	1.3
hospital_predicted[58]	1076	70	970	1200	2341	1838	1.00	1.4	1.3
hospital_predicted[59]	1078	70	970	1200	2316	2009	1.00	1.5	1.3
hospital_predicted[60]	1081	71	970	1200	2287	2096	1.00	1.5	1.3
hospital_predicted[61]	1086	71	980	1200	2260	2036	1.00	1.5	1.3
hospital_predicted[62]	1093	72	980	1200	2228	2055	1.00	1.5	1.4
hospital_predicted[63]	1101	72	990	1200	2190	1992	1.00	1.5	1.4
hospital_predicted[64]	1110	73	1000	1200	2160	1843	1.00	1.6	1.4
hospital_predicted[65]	1120	74	1000	1200	2127	1786	1.00	1.6	1.4
hospital_predicted[66]	1132	75	1000	1300	2104	1856	1.00	1.6	1.5
hospital_predicted[67]	1144	76	1000	1300	2093	1779	1.00	1.7	1.5
hospital_predicted[68]	1157	77	1000	1300	2092	1789	1.00	1.7	1.5
hospital_predicted[69]	1171	79	1000	1300	2102	1595	1.00	1.7	1.5
hospital_predicted[70]	1185	80	1100	1300	2114	1732	1.00	1.7	1.5
hospital_predicted[71]	1200	81	1100	1300	2146	1701	1.00	1.8	1.5
hospital_predicted[72]	1216	83	1100	1400	2190	1558	1.00	1.8	1.6
hospital_predicted[73]	1232	85	1100	1400	2244	1619	1.00	1.8	1.6
hospital_predicted[74]	1249	86	1100	1400	2295	1718	1.00	1.8	1.7
hospital_predicted[75]	1267	88	1100	1400	2355	1729	1.00	1.8	1.7
hospital_predicted[76]	1286	90	1100	1400	2402	1714	1.00	1.8	1.8
hospital_predicted[77]	1306	92	1200	1500	2416	1738	1.00	1.9	1.9
hospital_predicted[78]	1327	94	1200	1500	2402	1835	1.00	1.9	1.9
hospital_predicted[79]	1350	96	1200	1500	2372	1786	1.00	2	2
hospital_predicted[80]	1374	99	1200	1500	2308	1813	1.00	2.1	2
hospital_predicted[81]	1399	102	1200	1600	2275	1671	1.00	2.2	2.1
hospital_predicted[82]	1427	107	1300	1600	2145	1544	1.00	2.3	2.2
hospital_predicted[83]	1457	113	1300	1700	2024	1478	1.00	2.5	2.3
hospital_predicted[84]	1489	122	1300	1700	1955	1606	1.00	2.8	2.4
hospital_predicted[85]	1524	134	1300	1800	1946	1699	1.00	3.1	2.6
hospital_predicted[86]	1562	149	1300	1800	1938	1667	1.00	3.4	2.9
hospital_predicted[87]	1603	169	1400	1900	1961	1760	1.00	3.9	3.3
hospital_predicted[88]	1648	193	1400	2000	1989	1946	1.00	4.4	3.7
hospital_predicted[89]	1.70e+03	223	1400	2100	2053	1943	1.00	5	4.2

We extract the posterior quantiles and print summary statistics.

Code

latent_inf = idata_90_trimmed.posterior["latent_infections"]

quantiles_90 = {
    "q05": latent_inf.quantile(0.05, dim=["chain", "draw"]).values,
    "q50": latent_inf.quantile(0.50, dim=["chain", "draw"]).values,
    "q95": latent_inf.quantile(0.95, dim=["chain", "draw"]).values,
}

ci_width_90 = quantiles_90["q95"] - quantiles_90["q05"]
print(f"Posterior summary for {n_days_90} days:")
print(f"  Mean 90% CI width: {ci_width_90.mean():,.0f} infections")
print(f"  Median infections (day 45): {quantiles_90['q50'][45]:,.0f}")

Posterior summary for 90 days:
  Mean 90% CI width: 49,190 infections
  Median infections (day 45): 103,626

Finally, we visualize the posterior latent infections alongside observed hospitalizations. Note that hospital admissions lag behind infections by the infection-to-hospitalization delay (mode ~10 days in our delay PMF). When comparing the two panels, peaks in the infection curve should precede corresponding peaks in hospitalizations by roughly 10-14 days.

Code

# Visualize posterior latent infections and observed hospitalizations (90 days)
# Create separate dataframes for faceted plot
infections_df_90 = pd.DataFrame(
    {
        "day": np.arange(n_days_90),
        "median": quantiles_90["q50"],
        "q05": quantiles_90["q05"],
        "q95": quantiles_90["q95"],
        "signal": "Latent Infections",
    }
)

# Add 14-day moving average to smooth noisy daily admissions
hosp_raw_90 = np.array(hosp_admits[:n_days_90], dtype=float)
hosp_ma_90 = pd.Series(hosp_raw_90).rolling(window=14, center=True).mean().values

hosp_df_90 = pd.DataFrame(
    {
        "day": np.arange(n_days_90),
        "median": hosp_ma_90,
        "raw": hosp_raw_90,
        "q05": np.nan,
        "q95": np.nan,
        "signal": "Hospital Admissions (14-day MA)",
    }
)

plot_df_90 = pd.concat([infections_df_90, hosp_df_90], ignore_index=True)
plot_df_90["signal"] = pd.Categorical(
    plot_df_90["signal"],
    categories=["Hospital Admissions (14-day MA)", "Latent Infections"],
    ordered=True,
)

(
    p9.ggplot(plot_df_90, p9.aes(x="day"))
    + p9.geom_ribbon(
        p9.aes(ymin="q05", ymax="q95"),
        fill="steelblue",
        alpha=0.3,
    )
    + p9.geom_point(
        p9.aes(y="raw"),
        color="gray",
        alpha=0.3,
        size=1,
    )
    + p9.geom_line(
        p9.aes(y="median"),
        color="darkblue",
        size=1,
    )
    + p9.facet_wrap("~signal", ncol=1, scales="free_y")
    + p9.scale_y_log10()
    + p9.labs(
        x="Day",
        y="Count",
        title="Posterior Latent Infections vs Observed Hospitalizations (90 days)",
    )
    + theme_tutorial
)

Figure 1: Posterior latent infections and observed hospitalizations (90 days).

Month-by-month uncertainty

The posterior ribbon above shows the infection trajectory and its 90% credible interval, but the ribbon width itself is hard to compare across days. To focus directly on uncertainty, we compute the daily 90% interval width and summarize it over each 30-day period.

Code

ci_width_90 = quantiles_90["q95"] - quantiles_90["q05"]

ci_width_df = pd.DataFrame(
    {
        "day": np.arange(n_days_90),
        "ci_width": ci_width_90,
        "period": pd.cut(
            np.arange(n_days_90),
            bins=[-1, 29, 59, 89],
            labels=["Days 0-29", "Days 30-59", "Days 60-89"],
        ),
    }
)

period_summary_df = (
    ci_width_df.groupby("period", observed=True)
    .agg(
        start=("day", "min"),
        end=("day", "max"),
        mean_ci_width=("ci_width", "mean"),
    )
    .reset_index()
)

print("\nMean 90% interval width by time period:")
for row in period_summary_df.itertuples(index=False):
    print(f"  {row.period}: {row.mean_ci_width:,.0f}")

(
    p9.ggplot(ci_width_df, p9.aes(x="day", y="ci_width"))
    + p9.geom_line(color="gray", alpha=0.6, size=0.7)
    + p9.geom_segment(
        data=period_summary_df,
        mapping=p9.aes(
            x="start",
            xend="end",
            y="mean_ci_width",
            yend="mean_ci_width",
            color="period",
        ),
        size=1.6,
    )
    + p9.scale_y_log10()
    + p9.labs(
        x="Day",
        y="90% interval width",
        color="Period",
        title="Posterior uncertainty in latent infections over the 90-day fit",
    )
    + theme_tutorial
)

Mean 90% interval width by time period:
  Days 0-29: 33,773
  Days 30-59: 28,592
  Days 60-89: 85,204

This pattern has a direct implication for forecasting: renewal models are most uncertain at the edge of the observation window. The daily line shows local changes in posterior uncertainty, while the monthly means make the broader edge effect easier to see. Future observations constrain past latent infections through the renewal equation and observation delays, but the most recent days have less future signal available to anchor them. That edge uncertainty is the same uncertainty that carries forward when the model is used for forecasting.

Summary

In this tutorial, we built a renewal model that combines hospital admissions and wastewater concentrations through a shared latent infection process. The latent process describes infections through time and across subpopulations; each observation process describes how one data stream is generated from those infections.

The main workflow was:

1. Configure the latent infection process with configure_latent.
2. Add named observation processes with add_observation.
3. Build the model with build.
4. Prepare observation dictionaries whose keys match the observation process names.
5. Fit the model with model.run.

A few details are especially important when adapting this pattern:

• Observation process names become the data argument names passed to model.run, such as hospital={...} and wastewater={...}.
• Dense daily observations are padded with model.pad_observations.
• Sparse observations use shifted time indices from model.shift_times.
• subpop_fractions defines the latent subpopulations; subpop_indices tells an observation process which subpopulation each measurement observes.
• The model time axis starts before the first observed day because PyRenew adds an initialization period for the renewal process.

The 90-day fit shows why the end of the observation window matters. Future observations help constrain past latent infections, so posterior uncertainty is usually largest near the most recent observed days. That same edge uncertainty is what carries forward when the model is used for forecasting.

Next Steps

• Explore different temporal processes for \(\mathcal{R}(t)\) in the Latent Infections and Latent Subpopulation Infections tutorials
• Learn about count-based observation models in Observation Processes: Counts
• Learn about continuous measurement models in Observation Processes: Measurements
• Learn about ascertainment models in Ascertainment Models for Related Observation Processes