diff --git a/.gitignore b/.gitignore index 14a31920..d93f2737 100755 --- a/.gitignore +++ b/.gitignore @@ -332,3 +332,5 @@ replay_pid* /.quarto/ *_files + +**/*.quarto_ipynb diff --git a/docs/.gitignore b/docs/.gitignore index 075b2542..0e3521a7 100644 --- a/docs/.gitignore +++ b/docs/.gitignore @@ -1 +1,3 @@ /.quarto/ + +**/*.quarto_ipynb diff --git a/docs/index.md b/docs/index.md index 213754e9..f9aec02b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -105,6 +105,7 @@ pip install git+https://github.com/CDCgov/PyRenew@main - [Building multi-signal models](tutorials/building_multisignal_models.md) -- composing a renewal model from PyRenew components using `PyrenewBuilder`. - [Latent infections](tutorials/latent_infections.md) -- modeling latent infection trajectories over time. - [Latent subpopulation infections](tutorials/latent_subpopulation_infections.md) -- modeling latent infections with subpopulation structure. +- [Latent infections with mobility](tutorials/latent_mobility_infections.md) -- Modeling subpopulations coupled by a mobility matrix. - [Observation processes: count data](tutorials/observation_processes_counts.md) -- connecting latent infections to observed counts. - [Observation processes: measurements](tutorials/observation_processes_measurements.md) -- connecting latent infections to continuous measurements. - [Joint ascertainment](tutorials/ascertainment.md) -- sharing ascertainment structure across count signals. diff --git a/docs/tutorials/.pages b/docs/tutorials/.pages index fd693b8d..6aea842f 100644 --- a/docs/tutorials/.pages +++ b/docs/tutorials/.pages @@ -3,6 +3,7 @@ nav: - building_multisignal_models.md - latent_infections.md - latent_subpopulation_infections.md + - latent_mobility_infections.md - observation_processes_counts.md - observation_processes_measurements.md - right_truncation.md diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd new file mode 100644 index 00000000..631aac33 --- /dev/null +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -0,0 +1,456 @@ +--- +title: Latent Mobility Infections +format: + gfm: + code-fold: true +engine: jupyter +jupyter: + jupytext: + text_representation: + extension: .qmd + format_name: quarto + format_version: "1.0" + jupytext_version: 1.18.1 + kernelspec: + display_name: Python 3 (ipykernel) + language: python + name: python3 +--- + +```{python} +#| label: setup +#| output: false + +import numpyro + +numpyro.enable_x64() + +import jax.numpy as jnp +import jax.random as random +import numpy as np +from numpyro.infer import Predictive +import pandas as pd +import plotnine as p9 +import warnings +from plotnine.exceptions import PlotnineWarning + +warnings.filterwarnings("ignore", category=PlotnineWarning) + +from _tutorial_theme import theme_tutorial +``` + +```{python} +#| label: imports + +from pyrenew.latent import MobilitySubpopulationInfections, RandomWalk +from pyrenew.deterministic import DeterministicPMF, DeterministicVariable +``` + +## Overview + +`SubpopulationInfections` gives each subpopulation its own renewal equation, and infections in one subpopulation do not depend on infections in any other. +`MobilitySubpopulationInfections` keeps the same $\mathcal{R}(t)$ structure, a shared baseline plus deviations that sum to zero, but treats subpopulations as home regions whose residents spend part of each day in other regions. +Infections among residents of a region then depend on transmission in the regions they visit and on the infectious visitors their own region receives. + +The class implements Equation 2 of Roy, Clapham, and Mishra (2025), [Incorporating human mobility to enhance epidemic response and estimate real-time reproduction numbers](https://doi.org/10.1371/journal.pcbi.1013642). +For home region $i$ on day $t$: + +$$ + I_{t,i} = \sum_{j=1}^{n} \frac{m_{ji}\, S_{t,i}}{\sum_{l=1}^{n} m_{jl} N_l}\; R_j(t) \left[ \sum_{k=1}^{n} m_{jk} \sum_{\tau} I_{\tau,k}\, g_{t-\tau} \right] +$$ + +where: + +- $I_{t,i}$ is the modeled infection count +- $m_{jk}$ is the fraction of residents of home region $k$ present at location $j$ +- $N_k$ is the number of residents of region $k$ +- $S_{t,i}$ is the number of residents of region $i$ not infected before day $t$ +- $R_j(t)$ is the reproduction number at location $j$, where transmission happens +- $g$ is the generation interval + +PyRenew works with infections per resident. +With depletion enabled, PyRenew adds a ceiling to Equation 2: each day's new infections cannot exceed the remaining susceptible population. +With `M` the mobility matrix and `p` the population fractions, one day of the computation with depletion enabled is: + +```python +present = M @ p +share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None] +infectiousness = reversed_gen_int @ history +susceptible = 1.0 - cumulative +new_infections_raw = M.T @ (rt * (share_present @ infectiousness)) +new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible) +``` + +- `share_present[j, k]` is the share of the people present at location `j` who live in region `k`. +- `reversed_gen_int` is the generation interval in reverse order and `history` holds the new infections per resident on each of the last `len(reversed_gen_int)` days, oldest first, with one column per home region. + With this ordering, the vector-matrix product pairs each day's infections with the generation-interval probability for that many days back, so `infectiousness` is the per-region weighted sum of recent new infections per resident. +- `cumulative` is the share of each region's residents infected before today, summed over every earlier day starting with the first day of the initialization period. +- `new_infections` is capped at `susceptible`. +- `rt` multiplies at each location, because $R_j(t)$ describes transmission where it happens. +- `M.T` assigns infections acquired at each location back to the home region of the person infected. +- `susceptible` scales infections after they are assigned to home regions, because residents of a region lose susceptibility wherever they were infected. + +## The Mobility Matrix + +This tutorial uses the three regions of Roy et al. +(Fig 2), with populations of 20, 10, and 15 million, and the mobility matrix of their Fig 2b. + +```{python} +#| label: matrix + +M = np.array( + [ + [0.9, 0.07, 0.21], + [0.05, 0.85, 0.09], + [0.05, 0.08, 0.7], + ] +) +p = np.array([20.0, 10.0, 15.0]) / 45.0 +region_names = ["Region 1", "Region 2", "Region 3"] + +matrix_df = pd.DataFrame( + { + "location": np.repeat(region_names, len(region_names)), + "home_region": np.tile(region_names, len(region_names)), + "fraction": M.ravel(), + } +) + +( + p9.ggplot(matrix_df, p9.aes(x="home_region", y="location", fill="fraction")) + + p9.geom_tile(color="white", size=1) + + p9.geom_text( + p9.aes(label="fraction", color="fraction > 0.5"), + format_string="{:.2f}", + size=12, + show_legend=False, + ) + + p9.scale_fill_cmap(cmap_name="Blues", limits=(0, 1)) + + p9.scale_color_manual(values=["black", "white"]) + + p9.scale_y_discrete(limits=region_names[::-1]) + + p9.coord_equal() + + p9.labs( + x="Home region", + y="Daytime location", + fill="Fraction of residents", + title="Mobility Matrix", + ) + + theme_tutorial +) +``` + +A mobility matrix is a square matrix with one row and one column per region. +`M[j, k]` is the fraction of residents of home region `k` present at location `j`. +Columns are home regions and each column sums to $1$. +Rows are the same regions, viewed as the locations where residents spend the day; the diagonal is the fraction of residents who stay home. +In region 3, 70% of residents stay home on a typical day, 21% spend it in region 1, and 9% in region 2. + +- The matrix must represent home regions by column, not rows. + A matrix with home regions in its rows, so that each row sums to 1, must be transposed before use, otherwise `MobilitySubpopulationInfections` rejects it which an error message that says to pass the transpose. +- A matrix stored as a data frame must be converted with `.to_numpy()`, with rows and columns in that order. +- Regions must be in the same order in the matrix, in `subpop_fractions`, in a vector-valued `I0_rv` or `log_rt_time_0_rv`, and in every observation's `subpop_indices`. + +You may want to model only some of the areas represented in your raw mobility matrix, or your matrix may include an "everywhere else" bin. +To use it in this model, drop any flows to un-modeled locations, then renormalize by columns, so columns once again sum to 1. +Roy et al. +do this for the twelve local authorities of North East England. + +```{python} +#| label: partial-area + +M_full = np.array( + [ + [0.85, 0.05, 0.10, 0.02], + [0.05, 0.90, 0.05, 0.03], + [0.05, 0.03, 0.80, 0.05], + [0.05, 0.02, 0.05, 0.90], + ] +) +M_area = M_full[:3, :3] +M_area = M_area / M_area.sum(axis=0, keepdims=True) +print(f"Column sums after renormalizing: {M_area.sum(axis=0)}") +``` + +## Building the Latent Process + +The latent process takes the matrix as a `DeterministicVariable`. +The matrix is data: it is fixed when the process is constructed and is never estimated. + +`I0_rv` is the number of new incident infections per resident on the first day of the initialization period, the earliest day on the time axis. +Infections grow from that value over the initialization period, and every initialization day counts toward susceptible depletion. + +Each region starts at its own $\mathcal{R}(t)$, given as one value per region by `log_rt_time_0_rv`. +Region 2 starts below 1. +The random walks have a very small step size, so $\mathcal{R}(t)$ stays close to its starting value in every region. + +```{python} +#| label: configure + +gen_int_rv = DeterministicPMF( + "gen_int", jnp.array([0.16, 0.32, 0.25, 0.14, 0.07, 0.04, 0.02]) +) +n_init = 7 +I0 = 1e-6 +start_rt = jnp.array([1.3, 0.8, 1.2]) + + +def simulate(matrix, susceptible_depletion=True, n_days=120): + """Simulate one draw of the mobility latent process; return its recorded values.""" + process = MobilitySubpopulationInfections( + name="mobility", + gen_int_rv=gen_int_rv, + n_initialization_points=n_init, + I0_rv=DeterministicVariable("I0", I0), + log_rt_time_0_rv=DeterministicVariable("log_rt_time_0", jnp.log(start_rt)), + baseline_rt_process=RandomWalk( + innovation_sd_rv=DeterministicVariable("innovation_sd", 0.001), + ), + subpop_rt_deviation_process=RandomWalk( + innovation_sd_rv=DeterministicVariable("innovation_sd", 0.001), + ), + mobility_matrix_rv=DeterministicVariable("mobility_matrix", matrix), + susceptible_depletion=susceptible_depletion, + ) + + def sampler(): + """Sample the latent process and record infections per resident.""" + latent = process.sample(n_days_post_init=n_days, subpop_fractions=p) + numpyro.deterministic("infections", latent.all_subpops) + + samples = Predictive(sampler, num_samples=1)(random.PRNGKey(0)) + return { + name.removeprefix("mobility::"): np.asarray(value[0]) + for name, value in samples.items() + } + + +def to_long(values, run): + """Convert a (days, regions) array to a long data frame, after initialization.""" + values = values[n_init:] + n_days = values.shape[0] + return pd.DataFrame( + { + "day": np.repeat(np.arange(n_days), values.shape[1]), + "region": np.tile(region_names, n_days), + "value": values.ravel(), + "run": run, + } + ) +``` + +Every call to `simulate` uses the same random key, so the $\mathcal{R}(t)$ trajectories are the same whatever the matrix. + +## Mobility Versus No Mobility + +With the identity matrix, every resident stays home, and each region follows its own renewal equation. +With depletion on, this is the model Roy et al. +compare against for regions not linked by mobility. + +```{python} +#| label: compare-matrices + +isolated = simulate(np.eye(3)) +linked = simulate(M) + +infections_df = pd.concat( + [ + to_long(isolated["infections"], "Identity matrix"), + to_long(linked["infections"], "Fig 2b matrix"), + ] +) + +( + p9.ggplot(infections_df, p9.aes(x="day", y="value", color="run")) + + p9.geom_line(size=1) + + p9.facet_wrap("~ region", ncol=3) + + p9.scale_y_log10() + + p9.labs( + x="Days", + y="New infections per resident", + title="Infections With and Without Mobility", + color="", + ) + + theme_tutorial + + p9.theme(legend_position="bottom") +) +``` + +On its own, region 2, with $\mathcal{R}(t)$ below 1, has a declining epidemic. +Linked to the other regions, it has an outbreak: its residents are infected while visiting regions 1 and 3, and infectious visitors from those regions bring infections to region 2. +Roy et al. +show the same effect in their S2 Fig. + +## Where $\mathcal{R}(t)$ Applies + +`rt_subpop[t, j]` is the reproduction number at location `j`, where transmission happens. +It is not the reproduction number of the residents of region `j`, who are also infected in the regions they visit. + +With depletion enabled, `rt_subpop_adjusted` is the reproduction number at each location after susceptibility adjustment. +After the initialization period it is `rt_subpop * (share_present @ susceptible_fraction)`. + +```{python} +#| label: rt-at-location + +rt_df = pd.concat( + [ + to_long(linked["rt_subpop"], "rt_subpop"), + to_long(linked["rt_subpop_adjusted"], "rt_subpop_adjusted"), + ] +) + +( + p9.ggplot(rt_df, p9.aes(x="day", y="value", color="run")) + + p9.geom_line(size=1) + + p9.geom_hline(yintercept=1.0, linetype="dashed", color="grey") + + p9.facet_wrap("~ region", ncol=3) + + p9.labs( + x="Days", + y=r"$\mathcal{R}(t)$ at location", + title="Reproduction Number Before and After Depletion", + color="", + ) + + theme_tutorial + + p9.theme(legend_position="bottom") +) +``` + +## Susceptible Depletion + +Susceptible depletion is on by default, as in Eq 2. + +```{python} +#| label: depletion + +depleted = simulate(M, susceptible_depletion=True, n_days=250) +not_depleted = simulate(M, susceptible_depletion=False, n_days=250) + +susceptible_df = pd.concat( + [ + to_long(depleted["susceptible_fraction"], "Depletion on"), + to_long(not_depleted["susceptible_fraction"], "Depletion off"), + ] +) + +( + p9.ggplot(susceptible_df, p9.aes(x="day", y="value", color="run")) + + p9.geom_line(size=1) + + p9.facet_grid("run ~ region", scales="free_y") + + p9.labs( + x="Days", + y="1 - cumulative infections per resident", + title="Susceptibility With and Without Depletion", + ) + + p9.guides(color="none") + + theme_tutorial + + p9.theme(figure_size=(10, 7)) +) +``` + +With depletion off, the recorded susceptible fraction becomes negative, indicating that cumulative infections exceed the population. +All residents are susceptible at the start of the initialization period, as in Roy et al. + +## Sources of Infection + +Roy et al. +(Eqs 3 to 5) split each region's new infections into three sources: + +- **local**: residents infected at home by infectious residents of the same region; +- **incoming**: residents infected at home by infectious visitors from other regions; +- **outgoing**: residents infected while visiting other regions. + +PyRenew does not provide a built-in function to compute this split, but it can be computed from recorded values, the matrix, the population fractions, and the generation interval. +The generation interval must be fixed, because a sampled one is not recorded. + +```{python} +#| label: sources + +per_resident = linked["infections"] +rt_subpop = linked["rt_subpop"] +susceptible_fraction = linked["susceptible_fraction"] +reversed_gen_int = np.flip(np.asarray(gen_int_rv())) +n_gen = reversed_gen_int.size +present = M @ p +share_present = M * p[None, :] / np.where(present > 0, present, 1.0)[:, None] + +sources = [] +for t in range(n_init, per_resident.shape[0]): + infectiousness = reversed_gen_int @ per_resident[t - n_gen : t] + at_location = rt_subpop[t] * (share_present @ infectiousness) + susceptible = susceptible_fraction[t] + local = ( + susceptible + * np.diag(M) + * rt_subpop[t] + * np.diag(share_present) + * infectiousness + ) + incoming = susceptible * np.diag(M) * at_location - local + new_infections_raw = M.T @ at_location + outgoing = susceptible * new_infections_raw - local - incoming + ceiling_scale = 1.0 / np.maximum(new_infections_raw, 1.0) + sources.append(np.stack([local, incoming, outgoing]) * ceiling_scale) +sources = np.stack(sources) + +total = sources.sum(axis=1) +print( + "Largest absolute difference between the summed sources and the infections: " + f"{np.max(np.abs(total - per_resident[n_init:])):.1e}" +) +``` + +The three sources add up to each day's new infections, to floating-point precision. + +```{python} +#| label: plot-sources + +source_names = ["Local", "Incoming", "Outgoing"] +share = np.divide( + sources, + total[:, np.newaxis, :], + out=np.zeros_like(sources), + where=total[:, np.newaxis, :] > 0, +) +n_days = share.shape[0] +sources_df = pd.DataFrame( + { + "day": np.repeat(np.arange(n_days), 9), + "source": np.tile(np.repeat(source_names, 3), n_days), + "region": np.tile(region_names, 3 * n_days), + "share": share.ravel(), + } +) +sources_df["source"] = pd.Categorical( + sources_df["source"], categories=source_names, ordered=True +) + +( + p9.ggplot(sources_df, p9.aes(x="day", y="share", fill="source")) + + p9.geom_area() + + p9.facet_wrap("~ region", ncol=3) + + p9.labs( + x="Days", + y="Share of new infections", + title="Sources of New Infections", + fill="", + ) + + theme_tutorial + + p9.theme(legend_position="bottom") +) +``` + +More than half of region 2's infections come from mobility: its residents are infected while visiting regions with $\mathcal{R}(t)$ above 1, or at home by visitors from those regions. + +## Connecting to Observation Processes + +`MobilitySubpopulationInfections` is configured with `PyrenewBuilder` like `SubpopulationInfections`, and `model.run()` takes the same arguments. +The matrix is part of the model definition, not run-time data, so fitting and prediction use the same matrix, and no run-time argument can replace it. +To model another jurisdiction, build a model with that jurisdiction's matrix. + +Before fitting, `run()` checks that the matrix has one row and one column per entry of `subpop_fractions`. +Empty locations are allowed and contribute no infections. + +For building complete models with observation processes, see: + +- [Latent Subpopulation Infections](latent_subpopulation_infections.md) +- [Observation Processes: Counts](observation_processes_counts.md) +- [Building Multi-Signal Models](building_multisignal_models.md) diff --git a/docs/tutorials/latent_subpopulation_infections.qmd b/docs/tutorials/latent_subpopulation_infections.qmd index e3e8b975..2f3e2a74 100644 --- a/docs/tutorials/latent_subpopulation_infections.qmd +++ b/docs/tutorials/latent_subpopulation_infections.qmd @@ -798,3 +798,5 @@ To learn how to build complete models with observation processes, see: - [Observation Processes: Counts](observation_processes_counts.md) - [Observation Processes: Measurements](observation_processes_measurements.md) - [Building Multi-Signal Models](building_multisignal_models.md) + +For subpopulations linked by mobility, where residents spend part of each day in other regions, see [Latent Mobility Infections](latent_mobility_infections.md). diff --git a/pyrenew/latent/__init__.py b/pyrenew/latent/__init__.py index 4203bf40..4c5e9769 100644 --- a/pyrenew/latent/__init__.py +++ b/pyrenew/latent/__init__.py @@ -13,6 +13,7 @@ from pyrenew.latent.infection_functions import ( compute_infections_from_rt, compute_infections_from_rt_with_feedback, + compute_infections_from_rt_with_mobility, logistic_susceptibility_adjustment, ) from pyrenew.latent.infection_initialization_method import ( @@ -26,8 +27,14 @@ ) from pyrenew.latent.infections import InfectionProcessSample, Infections from pyrenew.latent.infectionswithfeedback import InfectionsWithFeedback +from pyrenew.latent.mobility_subpopulation_infections import ( + MobilitySubpopulationInfections, +) from pyrenew.latent.population_infections import PopulationInfections -from pyrenew.latent.subpopulation_infections import SubpopulationInfections +from pyrenew.latent.subpopulation_infections import ( + BaseSubpopulationInfections, + SubpopulationInfections, +) from pyrenew.latent.temporal_processes import ( AR1, DifferencedAR1, @@ -42,6 +49,7 @@ "logistic_susceptibility_adjustment", "compute_infections_from_rt", "compute_infections_from_rt_with_feedback", + "compute_infections_from_rt_with_mobility", "InfectionProcessSample", "InfectionInitializationMethod", "InitializeInfectionsExponentialGrowth", @@ -51,9 +59,11 @@ "InfectionsWithFeedback", # Base classes and types "BaseLatentInfectionProcess", + "BaseSubpopulationInfections", "LatentSample", "PopulationStructure", "SubpopulationInfections", + "MobilitySubpopulationInfections", "PopulationInfections", # Hierarchical priors "HierarchicalNormalPrior", diff --git a/pyrenew/latent/infection_functions.py b/pyrenew/latent/infection_functions.py index 86b15b39..471fc2fd 100755 --- a/pyrenew/latent/infection_functions.py +++ b/pyrenew/latent/infection_functions.py @@ -2,9 +2,12 @@ from __future__ import annotations +import functools + import jax import jax.numpy as jnp from jax.typing import ArrayLike +from numpyro.util import not_jax_tracer from pyrenew.convolve import new_convolve_scanner, new_double_convolve_scanner from pyrenew.transformation import ExpTransform, IdentityTransform @@ -183,3 +186,186 @@ def compute_infections_from_rt_with_feedback( infections, R_adjustment = infs_and_R_adj Rt_adjusted = R_adjustment * Rt_raw return infections, Rt_adjusted + + +def _mobility_renewal_step( + state: tuple[ArrayLike, ArrayLike], + rt_raw: ArrayLike, + mobility_matrix: ArrayLike, + share_present: ArrayLike, + reversed_generation_interval_pmf: ArrayLike, + susceptible_depletion: bool, +) -> tuple[tuple[jnp.ndarray, jnp.ndarray], tuple[jnp.ndarray, jnp.ndarray]]: + """ + Advance the mobility-coupled renewal process by one day. + + Parameters + ---------- + state + Recent new infections per resident, shape ``(n_gen, n)``, and + cumulative new infections per resident, shape ``(n,)``. + rt_raw + Reproduction number at each location on this day, shape ``(n,)``. + mobility_matrix + Fraction of residents of each home region present at each location, + shape ``(n, n)``. + share_present + Share of the people present at each location who live in each home + region, shape ``(n, n)``. + reversed_generation_interval_pmf + Generation interval PMF in reverse order. + susceptible_depletion + Whether new infections are scaled by, and capped at, the + susceptible fraction. + + Returns + ------- + tuple[tuple[jnp.ndarray, jnp.ndarray], tuple[jnp.ndarray, jnp.ndarray]] + The updated state, and this day's new infections per resident and + reproduction number at each location after susceptibility + adjustment, before the incidence ceiling. + """ + history, cumulative = state + infectiousness = reversed_generation_interval_pmf @ history + new_infections_raw = mobility_matrix.T @ (rt_raw * (share_present @ infectiousness)) + if susceptible_depletion: + susceptible = 1.0 - cumulative + new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible) + rt_adjusted = rt_raw * (share_present @ susceptible) + else: + new_infections = new_infections_raw + rt_adjusted = rt_raw + history = jnp.concatenate([history[1:], new_infections[jnp.newaxis]], axis=0) + return (history, cumulative + new_infections), (new_infections, rt_adjusted) + + +def compute_infections_from_rt_with_mobility( + I0: ArrayLike, + Rt_raw: ArrayLike, + reversed_generation_interval_pmf: ArrayLike, + mobility_matrix: ArrayLike, + subpop_fractions: ArrayLike, + susceptible_depletion: bool = True, +) -> tuple[jnp.ndarray, jnp.ndarray]: + r""" + Generate infections from a renewal process coupled by a fixed mobility matrix. + + Residents of each home region spend part of each day at other locations. + Transmission happens at each location, at that location's reproduction + number, among the people present there. Each infection is assigned to + the home region of the person infected. + + Parameters + ---------- + I0 + New infections per resident on each day of the initialization + period, shape ``(n_init, n)``, with ``n_init`` at least the length of + the generation interval. Every row counts toward susceptible + depletion. Only the last ``len(reversed_generation_interval_pmf)`` + rows enter the generation-interval sum. Entries must be finite and + non-negative. With depletion enabled, ``I0.sum(axis=0)`` must be + at most 1 in every region. + Rt_raw + Finite, non-negative reproduction number at each location on each + day after the initialization period, shape ``(n_days, n)``. + reversed_generation_interval_pmf + Generation interval PMF in reverse order: the last entry is one day + back, the second-to-last two days back, and so on. + mobility_matrix + Matrix used on every day, shape ``(n, n)``. + ``mobility_matrix[j, k]`` is the fraction of residents of home + region ``k`` present at location ``j``. Each column sums to 1. + subpop_fractions + Each home region's share of the total population, shape ``(n,)``. + Sums to 1. + susceptible_depletion + If True (default), new infections among residents of each region + are scaled by, and capped at, the fraction of those residents not + yet infected. If False, neither adjustment is applied. + + Returns + ------- + tuple[jnp.ndarray, jnp.ndarray] + New infections per resident, shape ``(n_days, n)``, and the + reproduction number at each location after susceptibility + adjustment, the same shape. The second array does not account for + the incidence ceiling. Without depletion, it is ``Rt_raw``. + + Raises + ------ + ValueError + If concrete ``I0`` values are non-finite or negative, or if their + sum exceeds 1 in any region with depletion enabled. Value checks + are skipped during JAX tracing; callers must ensure traced inputs + meet the same requirements. + + Notes + ----- + This implements Equation 2 of Roy, Clapham, and Mishra (2025), + . In infection counts (before + converting to infections per resident and applying the incidence + ceiling), the equation is: + + ```math + I_{t,i} = \sum_{j=1}^{n} \frac{m_{ji} S_{t,i}}{\sum_{l=1}^{n} m_{jl} N_l} + R_{t,j} \left[ \sum_{k=1}^{n} m_{jk} \sum_{\tau} I_{\tau,k}\, g_{t-\tau} \right] + ``` + + where $m_{jk}$ is ``mobility_matrix[j, k]``, $N_k$ is the number of + residents of region $k$, $R_{t,j}$ is the reproduction number at + location $j$, $g$ is the generation interval, and $S_{t,i}$ is the + number of residents of region $i$ not infected before day $t$. + + With ``M = mobility_matrix`` and ``p = subpop_fractions``, one day with + depletion enabled is: + + ```python + present = M @ p + share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None] + infectiousness = reversed_generation_interval_pmf @ history + susceptible = 1.0 - cumulative + new_infections_raw = M.T @ (rt * (share_present @ infectiousness)) + new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible) + ``` + + ``share_present[j, k]`` is the share of the people present at location + ``j`` who live in region ``k``. ``cumulative`` counts every day of + ``I0`` and every earlier day of the result. + + A location with no one present contributes no infections. With + depletion enabled, its adjusted reproduction number is 0. + + The incidence ceiling is an addition to Roy et al.'s Equation 2. + With valid initialization, it keeps cumulative infections at most 1 + per resident, so the susceptible fraction needs no floor. Without + depletion, cumulative infections may exceed 1 per resident. + """ + I0 = jnp.asarray(I0) + invalid_initial = jnp.any(~jnp.isfinite(I0) | (I0 < 0)) + if not_jax_tracer(invalid_initial) and invalid_initial: + raise ValueError("I0 must contain only finite, non-negative infections.") + cumulative = jnp.sum(I0, axis=0) + if susceptible_depletion: + exceeds_population = jnp.any(cumulative > 1.0) + if not_jax_tracer(exceeds_population) and exceeds_population: + raise ValueError( + "I0.sum(axis=0) must be <= 1 in every region when " + "susceptible depletion is enabled." + ) + present = mobility_matrix @ subpop_fractions + share_present = ( + mobility_matrix + * subpop_fractions[jnp.newaxis, :] + / jnp.where(present > 0, present, 1.0)[:, jnp.newaxis] + ) + n_gen = reversed_generation_interval_pmf.size + initial_state = (I0[-n_gen:], cumulative) + step = functools.partial( + _mobility_renewal_step, + mobility_matrix=mobility_matrix, + share_present=share_present, + reversed_generation_interval_pmf=reversed_generation_interval_pmf, + susceptible_depletion=susceptible_depletion, + ) + _, (infections, Rt_adjusted) = jax.lax.scan(step, initial_state, Rt_raw) + return infections, Rt_adjusted diff --git a/pyrenew/latent/mobility_subpopulation_infections.py b/pyrenew/latent/mobility_subpopulation_infections.py new file mode 100644 index 00000000..474ccf24 --- /dev/null +++ b/pyrenew/latent/mobility_subpopulation_infections.py @@ -0,0 +1,293 @@ +""" +Subpopulation latent infections coupled by a fixed mobility matrix. +""" + +from __future__ import annotations + +import jax.numpy as jnp +import numpyro +from jax.typing import ArrayLike + +from pyrenew.deterministic import DeterministicVariable +from pyrenew.latent.base import PopulationStructure +from pyrenew.latent.infection_functions import ( + compute_infections_from_rt_with_mobility, +) +from pyrenew.latent.subpopulation_infections import BaseSubpopulationInfections +from pyrenew.latent.temporal_processes import TemporalProcess +from pyrenew.metaclass import RandomVariable + + +class MobilitySubpopulationInfections(BaseSubpopulationInfections): + r""" + Subpopulation renewal model coupled by a fixed mobility matrix. + + Subpopulations are home regions whose residents spend part of each day + at other locations. Transmission happens at each location, at that + location's reproduction number, among the people present there, and + each infection is assigned to the home region of the person infected. + This implements Equation 2 of Roy, Clapham, and Mishra (2025), + [Incorporating human mobility to enhance epidemic response and estimate + real-time reproduction numbers](https://doi.org/10.1371/journal.pcbi.1013642). + With depletion enabled, PyRenew additionally caps each day's new + infections at the remaining susceptible population. + + Infections are per resident: ``infections[t, k]`` is new infections + among residents of home region ``k`` on day ``t``, divided by the number + of residents of ``k``. With ``M`` the mobility matrix and + ``p = subpop_fractions``, one day with depletion enabled is: + + ```python + present = M @ p + share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None] + infectiousness = reversed_gen_int @ history + susceptible = 1.0 - cumulative + new_infections_raw = M.T @ (rt * (share_present @ infectiousness)) + new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible) + ``` + + ``share_present[j, k]`` is the share of the people present at location + ``j`` who live in region ``k``, and ``cumulative`` counts every infection + since the first day of the initialization period. Initial infections + must be finite and non-negative. With depletion enabled, their sum + over all initialization days must be at most 1 per resident in each + region. Concrete invalid values raise ``ValueError``; during JAX + tracing, callers must ensure sampled initialization meets these bounds. + + ``M[j, k]`` is the fraction of residents of home region ``k`` present at + location ``j``: columns are home regions, rows are locations, and each + column sums to 1. Regions must be in the same order in the matrix, + ``subpop_fractions``, a vector-valued ``I0_rv``, a vector-valued + ``log_rt_time_0_rv``, and every observation's ``subpop_indices``. + + The recorded values are those of + [pyrenew.latent.SubpopulationInfections][], plus + ``susceptible_fraction``. Under mobility: + + - ``rt_subpop[t, j]`` is the reproduction number at location ``j``, + where transmission happens, when everyone is susceptible. + - ``rt_subpop_adjusted`` is + ``rt_subpop * (share_present @ susceptible_fraction)`` after the + initialization period, with depletion enabled. This adjusts for + susceptibility but does not account for the incidence ceiling. It + equals new infections per unit of infectiousness present only when + the ceiling does not bind. It is 0 at a location with no one present. + - ``susceptible_fraction[t, i]`` is the share of residents of home + region ``i`` not infected before day ``t``. It is recorded as 1 minus + cumulative infections and can become negative with depletion off. + + Susceptible depletion is on by default, as in Eq 2. With it off, + neither susceptibility scaling nor the incidence ceiling is applied, + and inferred raw values of ``rt_subpop`` will need to directly absorb + any reductions in transmission due to susceptible depletion. + + The matrix is fixed. Fitting and prediction use the matrix the model + was built with; it is available as ``mobility_matrix_rv.value``. + """ + + def __init__( + self, + *, + name: str, + gen_int_rv: RandomVariable, + n_initialization_points: int, + I0_rv: RandomVariable, + baseline_rt_process: TemporalProcess, + subpop_rt_deviation_process: TemporalProcess, + log_rt_time_0_rv: RandomVariable, + mobility_matrix_rv: DeterministicVariable, + susceptible_depletion: bool = True, + ) -> None: + """ + Initialize the mobility-coupled infections process. + + Parameters + ---------- + name + Name prefix for numpyro sample sites. All deterministic + quantities are recorded under this scope (e.g., + ``"{name}::rt_subpop"``, ``"{name}::rt_subpop_adjusted"``, and + ``"{name}::susceptible_fraction"``). + gen_int_rv + Generation interval PMF + n_initialization_points + Number of initialization days before day 0. + I0_rv + New infections per resident on the first day of the + initialization period, the earliest day on the time axis. This + is not prevalence. Returns either a scalar, one value for all + regions, or an array of shape ``(n_subpops,)``, one value per + region in the order of ``subpop_fractions``. Every + initialization day counts toward susceptible depletion. With + depletion enabled, the sum over all initialization days must + be at most 1 per resident in each region. + baseline_rt_process + Temporal process for baseline $\\mathcal{R}(t)$ dynamics + subpop_rt_deviation_process + Temporal process for subpopulation deviations + log_rt_time_0_rv + Starting log($\\mathcal{R}(t)$) on the first day of the + initialization period. Returns either a scalar, one value for + all regions, or an array of shape ``(n_subpops,)``, one value per + region in the order of ``subpop_fractions``, which gives each + region its own starting $\\mathcal{R}(t)$, as in Roy et al. + mobility_matrix_rv + The mobility matrix used on every day, as a + ``DeterministicVariable`` whose value has shape + ``(n_subpops, n_subpops)``. ``mobility_matrix[j, k]`` is the + fraction of residents of home region ``k`` present at location + ``j``; each column sums to 1. Rows and columns are in the order + of ``subpop_fractions``. Convert a data frame with + ``.to_numpy()``. Transpose a matrix whose rows sum to 1. + susceptible_depletion + If True (default), new infections among residents of each region + are scaled by the fraction of those residents not yet infected, + as in Roy et al. Eq 2, and capped at that fraction. If False, + neither scaling nor the ceiling is applied, and ``rt_subpop`` + absorbs depletion. + + Raises + ------ + TypeError + If mobility_matrix_rv is not a DeterministicVariable. + ValueError + If required parameters are missing or invalid, or if the + mobility matrix is not a valid column-stochastic matrix. + """ + super().__init__( + name=name, + gen_int_rv=gen_int_rv, + n_initialization_points=n_initialization_points, + I0_rv=I0_rv, + baseline_rt_process=baseline_rt_process, + subpop_rt_deviation_process=subpop_rt_deviation_process, + log_rt_time_0_rv=log_rt_time_0_rv, + ) + if not isinstance(mobility_matrix_rv, DeterministicVariable): + raise TypeError( + "mobility_matrix_rv must be a DeterministicVariable; the " + "mobility matrix is data and cannot be random. Got " + f"{type(mobility_matrix_rv).__name__}." + ) + self._validate_mobility_matrix(jnp.asarray(mobility_matrix_rv.value)) + self.mobility_matrix_rv = mobility_matrix_rv + self.susceptible_depletion = susceptible_depletion + + @staticmethod + def _validate_mobility_matrix(matrix: ArrayLike) -> None: + """ + Check that the mobility matrix is a valid column-stochastic matrix. + + Parameters + ---------- + matrix + The mobility matrix, as a JAX array. + + Raises + ------ + ValueError + If the matrix is not square and two-dimensional, has non-finite + or negative entries, or has a column that does not sum to 1. + """ + if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]: + raise ValueError( + "mobility_matrix_rv must be a single square (n, n) matrix; " + f"got shape {matrix.shape}." + ) + if not jnp.all(jnp.isfinite(matrix)): + raise ValueError("mobility_matrix_rv must have only finite entries.") + if jnp.any(matrix < 0): + raise ValueError("mobility_matrix_rv must have only non-negative entries.") + column_sums = jnp.sum(matrix, axis=0) + if not jnp.all(jnp.isclose(column_sums, 1.0, atol=1e-6)): + message = ( + "Each column of mobility_matrix_rv must sum to 1; got column " + f"sums {column_sums.tolist()}." + ) + if jnp.all(jnp.isclose(jnp.sum(matrix, axis=1), 1.0, atol=1e-6)): + message += ( + " The rows sum to 1; mobility_matrix[j, k] must be the " + "fraction of residents of k present at j. Did you mean to pass the transpose?" + ) + raise ValueError(message) + + def _parse_and_validate_fractions( + self, + subpop_fractions: ArrayLike | None = None, + ) -> PopulationStructure: + """ + Parse population fractions and check them against the mobility matrix. + + Parameters + ---------- + subpop_fractions + Population fractions for all subpopulations. + + Returns + ------- + PopulationStructure + Parsed and validated population structure. + + Raises + ------ + ValueError + If the fractions are invalid, or if the matrix does not have one + row and one column per subpopulation. + """ + pop = super()._parse_and_validate_fractions(subpop_fractions) + matrix = jnp.asarray(self.mobility_matrix_rv.value) + if matrix.shape != (pop.n_subpops, pop.n_subpops): + raise ValueError( + f"mobility_matrix_rv has shape {matrix.shape}, but " + f"subpop_fractions has {pop.n_subpops} entries; the matrix " + f"must have shape ({pop.n_subpops}, {pop.n_subpops})." + ) + return pop + + def _infections_from_rt( + self, + rt_raw: ArrayLike, + I0: ArrayLike, + gen_int: ArrayLike, + pop: PopulationStructure, + ) -> tuple[ArrayLike, ArrayLike]: + """ + Compute post-initialization infections and adjusted Rt with mobility coupling. + + Parameters + ---------- + rt_raw + Raw Rt at each location after the initialization period. + I0 + Infections per resident on every initialization day. + gen_int + Generation interval PMF. + pop + Parsed population structure. + + Returns + ------- + tuple[ArrayLike, ArrayLike] + Post-initialization infections per resident and the reproduction + number at each location after susceptibility adjustment, + before the incidence ceiling. + """ + post_init_infections, rt_adjusted = compute_infections_from_rt_with_mobility( + I0=I0, + Rt_raw=rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int), + mobility_matrix=jnp.asarray(self.mobility_matrix_rv()), + subpop_fractions=pop.fractions, + susceptible_depletion=self.susceptible_depletion, + ) + infections_all = jnp.vstack([I0, post_init_infections]) + infected_before = jnp.vstack( + [ + jnp.zeros((1, infections_all.shape[1])), + jnp.cumsum(infections_all, axis=0)[:-1], + ] + ) + susceptible_fraction = 1.0 - infected_before + with numpyro.handlers.scope(prefix=self.name, divider="::"): + numpyro.deterministic("susceptible_fraction", susceptible_fraction) + return post_init_infections, rt_adjusted diff --git a/pyrenew/latent/subpopulation_infections.py b/pyrenew/latent/subpopulation_infections.py index 5d7eb9b2..8482ea97 100644 --- a/pyrenew/latent/subpopulation_infections.py +++ b/pyrenew/latent/subpopulation_infections.py @@ -4,6 +4,8 @@ from __future__ import annotations +from abc import abstractmethod + import jax.numpy as jnp import numpyro from jax.typing import ArrayLike @@ -23,58 +25,29 @@ from pyrenew.metaclass import RandomVariable -class SubpopulationInfections(BaseLatentInfectionProcess): +class BaseSubpopulationInfections(BaseLatentInfectionProcess): r""" - Multi-subpopulation renewal model with hierarchical $\mathcal{R}(t)$ structure. - - Each subpopulation has its own renewal equation with $\mathcal{R}(t)$ deviating from a - shared baseline. Suitable when transmission dynamics vary substantially - across subpopulations. - - Mathematical form: - - Baseline $\mathcal{R}(t)$: log[R_baseline(t)] ~ TemporalProcess - - Subpopulation $\mathcal{R}(t)$: log R_k(t) = log[R_baseline(t)] + delta_k(t) - - Deviations: delta_k(t) ~ TemporalProcess with sum-to-zero constraint - - Starting values on the first day of the initialization period, from - log_rt_time_0_rv: a scalar s gives log[R_baseline(0)] = s and - delta_k(0) = 0; one value v_k per subpopulation gives - log[R_baseline(0)] = mean(v) and delta_k(0) = v_k - mean(v), so - log R_k(0) = v_k. - - Renewal per subpop: I_k(t) = R_k(t) * sum_tau I_k(t-tau) * g(tau) - - Aggregate total: I_aggregate(t) = sum_k p_k * I_k(t) + Base class for subpopulation renewal models with hierarchical $\mathcal{R}(t)$. - Infection feedback is optional. It is on when both - ``infection_feedback_strength_rv`` and ``infection_feedback_pmf_rv`` - are given. On each day after the initialization period, it multiplies - the raw $\mathcal{R}(t)$ of each subpopulation by a factor that depends - on recent infections in that subpopulation: + Subpopulations share a baseline $\mathcal{R}(t)$. The log + $\mathcal{R}(t)$ of each subpopulation is the log of the baseline plus + a deviation, and the deviations sum to zero on each day, so the + baseline is the geometric mean of the subpopulation values. - ```math - \mathcal{R}^{\mathrm{adj}}_k(t) = \mathcal{R}_k(t) - \exp\left(\gamma_k(t) \sum_{\tau=1}^{T_g} I_k(t - \tau) f(\tau)\right) - ``` - - where $I_k(t)$ is new infections per resident of subpopulation $k$, - $\gamma_k(t)$ is the feedback strength, $f$ is the feedback PMF, and - $T_g$ is the length of the generation interval. The renewal equation - of subpopulation $k$ then uses the adjusted - $\mathcal{R}^{\mathrm{adj}}_k(t)$ in place of $\mathcal{R}_k(t)$. - Negative strength lowers $\mathcal{R}_k(t)$ when recent infections are - high, and positive strength raises it. With both arguments left out, - infections follow the plain renewal equation. + The starting log $\mathcal{R}(t)$ and the initial infections are each + either one value for all subpopulations or one value per subpopulation. + Initial infections grow exponentially over the initialization period + at the rate implied by each subpopulation's starting $\mathcal{R}(t)$. - The constructor specifies model structure (priors, temporal processes). - Population structure (subpop_fractions) is provided at sample time, - allowing a single model to be fit to multiple jurisdictions. + Each subclass defines how infections are generated from + $\mathcal{R}(t)$. Notes ----- - Sum-to-zero constraint on deviations ensures R_baseline(t) is the geometric - mean of subpopulation $\mathcal{R}(t)$ values, providing identifiability. - ``rt_baseline`` and ``rt_subpop`` are raw reproduction numbers generated by - the temporal-process hierarchy. ``rt_subpop_adjusted`` is the reproduction - number used to generate infections. It equals ``rt_subpop`` unless - infection feedback is on. + ``rt_baseline`` and ``rt_subpop`` are raw reproduction numbers generated + by the temporal-process hierarchy. ``rt_subpop_adjusted`` is the + reproduction number used to generate infections, as each subclass + defines it. """ def __init__( @@ -87,11 +60,9 @@ def __init__( baseline_rt_process: TemporalProcess, subpop_rt_deviation_process: TemporalProcess, log_rt_time_0_rv: RandomVariable, - infection_feedback_strength_rv: RandomVariable | None = None, - infection_feedback_pmf_rv: RandomVariable | None = None, ) -> None: """ - Initialize hierarchical infections process. + Initialize the shared subpopulation structure. Parameters ---------- @@ -105,50 +76,25 @@ def __init__( n_initialization_points Number of initialization days before day 0. I0_rv - Initial infection prevalence (proportion of population). Returns - either a scalar, one value for all subpopulations, or an array of - shape ``(n_subpops,)``, one value per subpopulation in the order of - ``subpop_fractions``. + New infections per resident on the first day of the + initialization period, the earliest day on the time axis. This + is not prevalence. Returns either a scalar, one value for all + subpopulations, or an array of shape ``(n_subpops,)``, one value + per subpopulation in the order of ``subpop_fractions``. baseline_rt_process Temporal process for baseline $\\mathcal{R}(t)$ dynamics subpop_rt_deviation_process Temporal process for subpopulation deviations log_rt_time_0_rv - Value(s) for subpopulation-level $\\log(\\mathcal{R}_i(t))$ on the first day of the - initialization period. Should return either a scalar (which will be - broadcast to all subpopulations) or an array of shape ``(n_subpops,)`` - containing subpopultion-specific $\log(\\mathcal{R}_i(t))$ values, - assigned in order of ``subpop_fractions``. - Each subpopulation's growth rate for infection seeding will - be calculated from its initial $\mathcal{R}(t))$ value. - The first value for the baseline time-varying $\\mathcal{R}(t)$ - (from which subpopulation-specific values deviate) is set as - the mean of the initial per-subpopulation values. - infection_feedback_strength_rv - Infection feedback strength $\\gamma$. Returns an array that - broadcasts to ``(n_days_post_init, n_subpops)``: - - - a scalar: one strength for all subpopulations and days; - - shape ``(n_subpops,)``: one strength per subpopulation; - - shape ``(n_days_post_init, 1)``: one strength per day; - - shape ``(n_days_post_init, n_subpops)``: one strength per day - and subpopulation. - - A 1-D array is read as one strength per subpopulation, so a - strength that varies by day must have shape - ``(n_days_post_init, 1)``. Give together with - ``infection_feedback_pmf_rv``, or leave both out for plain renewal. - infection_feedback_pmf_rv - Feedback PMF $f$, which weights recent infections in the feedback - factor. Element 0 weights infections on the previous day. Must have - the same length as the generation interval. Give together with - ``infection_feedback_strength_rv``. + Starting log($\\mathcal{R}(t)$) on the first day of the + initialization period. Returns either a scalar, one value for all + subpopulations, or an array of shape ``(n_subpops,)``, one value per + subpopulation in the order of ``subpop_fractions``. Raises ------ ValueError - If required parameters are missing or invalid, or if only one of - the two infection feedback arguments is given. + If required parameters are missing or invalid. """ super().__init__( name=name, @@ -175,12 +121,6 @@ def __init__( raise ValueError("subpop_rt_deviation_process is required") self.subpop_rt_deviation_process = subpop_rt_deviation_process - _validate_feedback_options( - infection_feedback_strength_rv, infection_feedback_pmf_rv - ) - self.infection_feedback_strength_rv = infection_feedback_strength_rv - self.infection_feedback_pmf_rv = infection_feedback_pmf_rv - def _validate_and_prepare_I0( self, I0: ArrayLike, @@ -195,7 +135,8 @@ def _validate_and_prepare_I0( Parameters ---------- I0 - Initial infection prevalence from I0_rv, as a JAX array. + New infections per resident on the first day of the + initialization period, from I0_rv, as a JAX array. pop Parsed population structure. @@ -272,6 +213,38 @@ def validate(self) -> None: """ validate_discrete_dist_vector(self.gen_int_rv()) + @abstractmethod + def _infections_from_rt( + self, + rt_raw: ArrayLike, + I0: ArrayLike, + gen_int: ArrayLike, + pop: PopulationStructure, + ) -> tuple[ArrayLike, ArrayLike]: + """ + Compute post-initialization infections and adjusted Rt for all subpopulations. + + Parameters + ---------- + rt_raw + Raw subpopulation Rt after the initialization period, + shape ``(n_days_post_init, n_subpops)``. + I0 + Infections per resident on every initialization day, + shape ``(n_initialization_points, n_subpops)``. + gen_int + Generation interval PMF. + pop + Parsed population structure. + + Returns + ------- + tuple[ArrayLike, ArrayLike] + Post-initialization infections per resident and the adjusted Rt + used to generate them, both with the shape of ``rt_raw``. + """ + pass # pragma: no cover + def sample( self, n_days_post_init: int, @@ -280,16 +253,16 @@ def sample( **kwargs: object, ) -> LatentSample: """ - Sample hierarchical infections for all subpopulations. + Sample infections for all subpopulations. Generates daily baseline $\\mathcal{R}(t)$, daily subpopulation - deviations with a sum-to-zero constraint, initial infections, and - computes infections with the renewal equations, with infection - feedback if it is on. The temporal processes may sample parameters at - any supported cadence (for example daily or weekly/stepwise), but they - must return daily-length trajectories before the renewal equations are - evaluated. The raw subpopulation Rt is recorded as ``rt_subpop``; the - adjusted Rt used to generate infections is recorded as + deviations with a sum-to-zero constraint, and initial infections, + then computes infections after the initialization period with the + subclass's renewal. The temporal processes may sample parameters at + any supported cadence (for example daily or weekly/stepwise), but + they must return daily-length trajectories before the renewal is + evaluated. The raw subpopulation Rt is recorded as ``rt_subpop``; + the Rt used to generate infections is recorded as ``rt_subpop_adjusted``. Parameters @@ -316,11 +289,9 @@ def sample( Raises ------ ValueError - If ``subpop_fractions`` is invalid, if ``log_rt_time_0_rv`` or - ``I0_rv`` returns neither a scalar nor an array of shape - ``(n_subpops,)``, if the infection feedback strength does not - broadcast to shape ``(n_days_post_init, n_subpops)``, or if the - feedback PMF length differs from the generation interval length. + If ``subpop_fractions`` is invalid, or if ``log_rt_time_0_rv`` + or ``I0_rv`` returns neither a scalar nor an array of shape + ``(n_subpops,)``. A subclass may raise for its own inputs. """ pop = self._parse_and_validate_fractions( subpop_fractions=subpop_fractions, @@ -373,12 +344,11 @@ def sample( initial_r_subpop[jnp.newaxis, :] * time_indices[:, jnp.newaxis] ) - post_init_infections_all, rt_adjusted = _compute_renewal( - I0=I0_all, + post_init_infections_all, rt_adjusted = self._infections_from_rt( rt_raw=rt_subpop[self.n_initialization_points :, :], + I0=I0_all, gen_int=gen_int, - infection_feedback_strength_rv=self.infection_feedback_strength_rv, - infection_feedback_pmf_rv=self.infection_feedback_pmf_rv, + pop=pop, ) rt_subpop_adjusted = jnp.vstack( [rt_subpop[: self.n_initialization_points, :], rt_adjusted] @@ -410,3 +380,187 @@ def sample( aggregate=infections_aggregate, all_subpops=infections_all, ) + + +class SubpopulationInfections(BaseSubpopulationInfections): + r""" + Multi-subpopulation renewal model with hierarchical $\mathcal{R}(t)$ structure. + + Each subpopulation has its own renewal equation with $\mathcal{R}(t)$ deviating from a + shared baseline. Suitable when transmission dynamics vary substantially + across subpopulations. Each subpopulation renews independently: + infections in one subpopulation do not depend on infections in any other. + + Mathematical form: + - Baseline $\mathcal{R}(t)$: log[R_baseline(t)] ~ TemporalProcess + - Subpopulation $\mathcal{R}(t)$: log R_k(t) = log[R_baseline(t)] + delta_k(t) + - Deviations: delta_k(t) ~ TemporalProcess with sum-to-zero constraint + - Starting values on the first day of the initialization period, from + log_rt_time_0_rv: a scalar s gives log[R_baseline(0)] = s and + delta_k(0) = 0; one value v_k per subpopulation gives + log[R_baseline(0)] = mean(v) and delta_k(0) = v_k - mean(v), so + log R_k(0) = v_k. + - Renewal per subpop: I_k(t) = R_k(t) * sum_tau I_k(t-tau) * g(tau) + - Aggregate total: I_aggregate(t) = sum_k p_k * I_k(t) + + Infection feedback is optional. It is on when both + ``infection_feedback_strength_rv`` and ``infection_feedback_pmf_rv`` + are given. On each day after the initialization period, it multiplies + the raw $\mathcal{R}(t)$ of each subpopulation by a factor that depends + on recent infections in that subpopulation: + + ```math + \mathcal{R}^{\mathrm{adj}}_k(t) = \mathcal{R}_k(t) + \exp\left(\gamma_k(t) \sum_{\tau=1}^{T_g} I_k(t - \tau) f(\tau)\right) + ``` + + where $I_k(t)$ is new infections per resident of subpopulation $k$, + $\gamma_k(t)$ is the feedback strength, $f$ is the feedback PMF, and + $T_g$ is the length of the generation interval. The renewal equation + of subpopulation $k$ then uses the adjusted + $\mathcal{R}^{\mathrm{adj}}_k(t)$ in place of $\mathcal{R}_k(t)$. + Negative strength lowers $\mathcal{R}_k(t)$ when recent infections are + high, and positive strength raises it. With both arguments left out, + infections follow the plain renewal equation. Sampling raises + ``ValueError`` if the feedback strength does not broadcast to shape + ``(n_days_post_init, n_subpops)``, or if the feedback PMF length + differs from the generation interval length. + + The constructor specifies model structure (priors, temporal processes). + Population structure (subpop_fractions) is provided at sample time, + allowing a single model to be fit to multiple jurisdictions. + + Notes + ----- + Sum-to-zero constraint on deviations ensures R_baseline(t) is the geometric + mean of subpopulation $\mathcal{R}(t)$ values, providing identifiability. + ``rt_baseline`` and ``rt_subpop`` are raw reproduction numbers generated by + the temporal-process hierarchy. ``rt_subpop_adjusted`` is the reproduction + number used to generate infections. It equals ``rt_subpop`` unless + infection feedback is on. + """ + + def __init__( + self, + *, + name: str, + gen_int_rv: RandomVariable, + n_initialization_points: int, + I0_rv: RandomVariable, + baseline_rt_process: TemporalProcess, + subpop_rt_deviation_process: TemporalProcess, + log_rt_time_0_rv: RandomVariable, + infection_feedback_strength_rv: RandomVariable | None = None, + infection_feedback_pmf_rv: RandomVariable | None = None, + ) -> None: + """ + Initialize hierarchical infections process. + + Parameters + ---------- + name + Name prefix for numpyro sample sites. All deterministic + quantities are recorded under this scope (e.g., + ``"{name}::rt_baseline"``, ``"{name}::rt_subpop"``, and + ``"{name}::rt_subpop_adjusted"``). + gen_int_rv + Generation interval PMF + n_initialization_points + Number of initialization days before day 0. + I0_rv + New infections per resident on the first day of the + initialization period, the earliest day on the time axis. This + is not prevalence. Returns either a scalar, one value for all + subpopulations, or an array of shape ``(n_subpops,)``, one value + per subpopulation in the order of ``subpop_fractions``. + baseline_rt_process + Temporal process for baseline $\\mathcal{R}(t)$ dynamics + subpop_rt_deviation_process + Temporal process for subpopulation deviations + log_rt_time_0_rv + Starting log($\\mathcal{R}(t)$) on the first day of the + initialization period. Returns either a scalar, one value for all + subpopulations, or an array of shape ``(n_subpops,)``, one value per + subpopulation in the order of ``subpop_fractions``. Per-subpopulation + values also give each subpopulation its own initial growth rate. + The baseline starts at the mean of the per-subpopulation values, so + giving every subpopulation the same prior narrows the prior on the + baseline's starting value: with + ``dist.Normal(0.0, s).expand([n_subpops])``, its standard deviation + is ``s / sqrt(n_subpops)``, not ``s``. + infection_feedback_strength_rv + Infection feedback strength $\\gamma$. Returns an array that + broadcasts to ``(n_days_post_init, n_subpops)``: + + - a scalar: one strength for all subpopulations and days; + - shape ``(n_subpops,)``: one strength per subpopulation; + - shape ``(n_days_post_init, 1)``: one strength per day; + - shape ``(n_days_post_init, n_subpops)``: one strength per day + and subpopulation. + + A 1-D array is read as one strength per subpopulation, so a + strength that varies by day must have shape + ``(n_days_post_init, 1)``. Give together with + ``infection_feedback_pmf_rv``, or leave both out for plain renewal. + infection_feedback_pmf_rv + Feedback PMF $f$, which weights recent infections in the feedback + factor. Element 0 weights infections on the previous day. Must have + the same length as the generation interval. Give together with + ``infection_feedback_strength_rv``. + + Raises + ------ + ValueError + If required parameters are missing or invalid, or if only one of + the two infection feedback arguments is given. + """ + super().__init__( + name=name, + gen_int_rv=gen_int_rv, + n_initialization_points=n_initialization_points, + I0_rv=I0_rv, + baseline_rt_process=baseline_rt_process, + subpop_rt_deviation_process=subpop_rt_deviation_process, + log_rt_time_0_rv=log_rt_time_0_rv, + ) + + _validate_feedback_options( + infection_feedback_strength_rv, infection_feedback_pmf_rv + ) + self.infection_feedback_strength_rv = infection_feedback_strength_rv + self.infection_feedback_pmf_rv = infection_feedback_pmf_rv + + def _infections_from_rt( + self, + rt_raw: ArrayLike, + I0: ArrayLike, + gen_int: ArrayLike, + pop: PopulationStructure, + ) -> tuple[ArrayLike, ArrayLike]: + """ + Compute post-initialization infections and adjusted Rt for each subpopulation independently. + + Parameters + ---------- + rt_raw + Raw subpopulation Rt after the initialization period. + I0 + Infections per resident on every initialization day. + gen_int + Generation interval PMF. + pop + Parsed population structure (unused). + + Returns + ------- + tuple[ArrayLike, ArrayLike] + Post-initialization infections per resident and the adjusted Rt, + with infection feedback if it is on. + """ + return _compute_renewal( + I0=I0, + rt_raw=rt_raw, + gen_int=gen_int, + infection_feedback_strength_rv=self.infection_feedback_strength_rv, + infection_feedback_pmf_rv=self.infection_feedback_pmf_rv, + ) diff --git a/test/conftest.py b/test/conftest.py index cabb5d90..aa6aff1c 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -12,16 +12,20 @@ from collections.abc import Callable +import jax import jax.numpy as jnp import numpyro.distributions as dist import pytest +from jax.typing import ArrayLike from pyrenew.deterministic import DeterministicPMF, DeterministicVariable from pyrenew.latent import ( + MobilitySubpopulationInfections, PopulationInfections, SubpopulationInfections, ) from pyrenew.metaclass import RandomVariable +from pyrenew.model import MultiSignalModel, PyrenewBuilder from pyrenew.observation import ( HierarchicalNormalNoise, NegativeBinomialNoise, @@ -308,6 +312,274 @@ def make(log_rt_time_0_rv: RandomVariable) -> SubpopulationInfections: return make +# ============================================================================= +# Mobility Fixtures +# ============================================================================= + + +@pytest.fixture +def roy_mobility_matrix() -> jnp.ndarray: + """ + Mobility matrix of Roy et al. (2025) Fig 2b. + + Returns + ------- + jnp.ndarray + Column-stochastic matrix; entry [j, k] is the fraction of + residents of region k present at location j. + """ + return jnp.array([[0.9, 0.07, 0.21], [0.05, 0.85, 0.09], [0.05, 0.08, 0.7]]) + + +@pytest.fixture +def roy_subpop_fractions() -> jnp.ndarray: + """ + Population fractions of the three regions of Roy et al. (2025) Fig 2. + + Returns + ------- + jnp.ndarray + Fractions for populations of 20, 10, and 15 million. + """ + return jnp.array([20.0, 10.0, 15.0]) / 45.0 + + +@pytest.fixture +def mobility_initial_infections() -> jnp.ndarray: + """ + Initialization-period infections per resident for three regions. + + Returns + ------- + jnp.ndarray + Shape (10, 3): ten days, longer than the 7-day generation + interval, with values that differ by region. + """ + return jnp.tile(jnp.array([1e-5, 2e-5, 5e-6]), (10, 1)) + + +@pytest.fixture +def mobility_rt_raw() -> jnp.ndarray: + """ + Unequal random-walk Rt for three locations over 150 days. + + Returns + ------- + jnp.ndarray + Shape (150, 3), starting at 2.0, 1.6, and 1.4. + """ + steps = 0.01 * jax.random.normal(jax.random.key(901), (150, 3)) + return jnp.exp(jnp.log(jnp.array([2.0, 1.6, 1.4])) + jnp.cumsum(steps, axis=0)) + + +@pytest.fixture(params=["empty_row", "zero_population"]) +def empty_location_mobility( + request: pytest.FixtureRequest, +) -> tuple[jnp.ndarray, jnp.ndarray]: + """ + Mobility matrix and population fractions with no one present at location 2. + + Parameters + ---------- + request + Selects the case: row 2 of the matrix is zero, or the only people + at location 2 live in region 2, which has no population. + + Returns + ------- + tuple[jnp.ndarray, jnp.ndarray] + The mobility matrix and the population fractions. + """ + if request.param == "empty_row": + return ( + jnp.array([[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 0.0]]), + jnp.array([0.4, 0.3, 0.3]), + ) + return jnp.eye(3), jnp.array([0.5, 0.5, 0.0]) + + +def _mobility_latent_params( + gen_int_rv: DeterministicPMF, + mobility_matrix: ArrayLike, + log_rt_time_0_rv: RandomVariable, + susceptible_depletion: bool, +) -> dict[str, object]: + """ + Constructor arguments shared by the mobility latent-process fixtures. + + Matches random_walk_subpopulation_infections_factory in name, initial + infections, and temporal processes. + + Parameters + ---------- + gen_int_rv + Generation interval PMF. + mobility_matrix + Mobility matrix value. + log_rt_time_0_rv + Starting log Rt, a scalar or one value per region. + susceptible_depletion + Whether susceptible depletion is on. + + Returns + ------- + dict[str, object] + Keyword arguments for MobilitySubpopulationInfections, without + n_initialization_points. + """ + return { + "name": "subpopulation", + "gen_int_rv": gen_int_rv, + "I0_rv": DeterministicVariable("I0", 0.001), + "log_rt_time_0_rv": log_rt_time_0_rv, + "baseline_rt_process": fixed_random_walk(innovation_sd=0.05), + "subpop_rt_deviation_process": fixed_random_walk(innovation_sd=0.025), + "mobility_matrix_rv": DeterministicVariable("mobility_matrix", mobility_matrix), + "susceptible_depletion": susceptible_depletion, + } + + +@pytest.fixture +def mobility_subpopulation_infections_factory( + gen_int_rv: DeterministicPMF, +) -> Callable[..., MobilitySubpopulationInfections]: + """ + Factory for MobilitySubpopulationInfections. + + Returns + ------- + Callable[..., MobilitySubpopulationInfections] + Function that builds the process with the given matrix, starting + log Rt, and depletion switch. + """ + + def make( + mobility_matrix: ArrayLike, + log_rt_time_0_rv: RandomVariable, + susceptible_depletion: bool = True, + ) -> MobilitySubpopulationInfections: + """ + Build MobilitySubpopulationInfections. + + Parameters + ---------- + mobility_matrix + Mobility matrix value. + log_rt_time_0_rv + Starting log Rt, a scalar or one value per region. + susceptible_depletion + Whether susceptible depletion is on. + + Returns + ------- + MobilitySubpopulationInfections + Configured infection process. + """ + return MobilitySubpopulationInfections( + n_initialization_points=7, + **_mobility_latent_params( + gen_int_rv, + mobility_matrix, + log_rt_time_0_rv, + susceptible_depletion, + ), + ) + + return make + + +@pytest.fixture +def mobility_subpopulation_infections( + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + roy_mobility_matrix: jnp.ndarray, +) -> MobilitySubpopulationInfections: + """ + MobilitySubpopulationInfections with the Fig 2b matrix. + + Returns + ------- + MobilitySubpopulationInfections + Process with a scalar starting log Rt of 0 and depletion on. + """ + return mobility_subpopulation_infections_factory( + roy_mobility_matrix, DeterministicVariable("log_rt_time_0", 0.0) + ) + + +@pytest.fixture +def mobility_model_factory( + gen_int_rv: DeterministicPMF, +) -> Callable[[ArrayLike], MultiSignalModel]: + """ + Factory for a builder-constructed model with mobility coupling. + + Returns + ------- + Callable[[ArrayLike], MultiSignalModel] + Function that builds the model with the given mobility matrix. + """ + + def make(mobility_matrix: ArrayLike) -> MultiSignalModel: + """ + Build a model with one subpopulation count observation. + + Parameters + ---------- + mobility_matrix + Mobility matrix value. + + Returns + ------- + MultiSignalModel + Model with the latent configuration of the mobility factory + and a SubpopulationCounts observation named "hospital_subpop". + """ + builder = PyrenewBuilder() + builder.configure_latent( + MobilitySubpopulationInfections, + **_mobility_latent_params( + gen_int_rv, + mobility_matrix, + DeterministicVariable("log_rt_time_0", 0.0), + True, + ), + ) + builder.add_observation( + SubpopulationCounts( + name="hospital_subpop", + ascertainment_rate_rv=DeterministicVariable("ihr", 0.01), + delay_distribution_rv=DeterministicPMF( + "delay", jnp.array([0.1, 0.3, 0.4, 0.2]) + ), + noise=NegativeBinomialNoise(DeterministicVariable("conc", 10.0)), + ) + ) + return builder.build() + + return make + + +@pytest.fixture +def mobility_model_args(roy_subpop_fractions: jnp.ndarray) -> dict[str, object]: + """ + Keyword arguments for sampling the mobility model. + + Returns + ------- + dict[str, object] + Arguments for model.sample: 30 days, the Fig 2 populations, and + prior predictive counts for all three regions. + """ + return { + "n_days_post_init": 30, + "population_size": 45_000_000, + "subpop_fractions": roy_subpop_fractions, + "hospital_subpop": {"obs": None, "subpop_indices": jnp.arange(3)}, + } + + @pytest.fixture def population_infections_factory( gen_int_rv: DeterministicPMF, diff --git a/test/test_infection_functions.py b/test/test_infection_functions.py index 752f99c3..46a4e9ab 100644 --- a/test/test_infection_functions.py +++ b/test/test_infection_functions.py @@ -3,10 +3,13 @@ submodule """ +import jax import jax.numpy as jnp +import numpy as np import pytest -from numpy.testing import assert_array_equal +from numpy.testing import assert_allclose, assert_array_equal +from pyrenew.deterministic import DeterministicPMF from pyrenew.latent import infection_functions as inf @@ -95,3 +98,435 @@ def test_compute_infections_from_rt_with_feedback_2d(I0, gen_int, inf_pmf, Rt_ra assert_array_equal(Rt_adj, Rt_raw) return None + + +def _roy_stan_loop( + I0_counts: np.ndarray, + Rt: np.ndarray, + gen_int: np.ndarray, + mobility_matrix: np.ndarray, + pop: np.ndarray, +) -> np.ndarray: + """ + Transcribe the time loop of the Roy et al. (2025) Stan model, in counts. + + Follows ``uk_spatial/simulated_data/simulated_fitting_region.stan`` in + https://github.com/mlgh-sg/Spatial_transmission. The generation-interval + sum runs over the full history, and the susceptibles of each region are + its population minus every earlier day's infections. + + Parameters + ---------- + I0_counts + Infections on each initialization day, shape ``(n_init, n)``. + Rt + Reproduction number at each location after initialization, + shape ``(n_days, n)``. + gen_int + Generation interval PMF; element 0 is one day back. + mobility_matrix + Column-stochastic mobility matrix, shape ``(n, n)``. + pop + Population of each region, shape ``(n,)``. + + Returns + ------- + np.ndarray + Infections after initialization, shape ``(n_days, n)``. + """ + n_init, n = I0_counts.shape + n_total = n_init + Rt.shape[0] + infections = np.zeros((n_total, n)) + infections[:n_init] = I0_counts + g = np.zeros(n_total + 1) + g[1 : gen_int.size + 1] = gen_int + eff_pop = mobility_matrix @ pop + for t in range(n_init, n_total): + convolution_inf = g[t - np.arange(t)] @ infections[:t] + total_inf = mobility_matrix @ convolution_inf + for m in range(n): + sus = pop[m] - infections[:t, m].sum() + infections[t, m] = np.dot( + mobility_matrix[:, m], (sus / eff_pop) * Rt[t - n_init] * total_inf + ) + return infections[n_init:] + + +def _eq1_loop(I0: np.ndarray, Rt: np.ndarray, gen_int: np.ndarray) -> np.ndarray: + """ + Renew each region on its own with susceptible depletion, per resident. + + This is Eq 1 of Roy et al. (2025) for each region. + + Parameters + ---------- + I0 + Infections per resident on each initialization day, + shape ``(n_init, n)``. + Rt + Reproduction number of each region after initialization, + shape ``(n_days, n)``. + gen_int + Generation interval PMF; element 0 is one day back. + + Returns + ------- + np.ndarray + Infections per resident after initialization, shape ``(n_days, n)``. + """ + n_init, n = I0.shape + n_total = n_init + Rt.shape[0] + infections = np.zeros((n_total, n)) + infections[:n_init] = I0 + g = np.zeros(n_total + 1) + g[1 : gen_int.size + 1] = gen_int + for t in range(n_init, n_total): + susceptible = 1.0 - infections[:t].sum(axis=0) + infections[t] = ( + susceptible * Rt[t - n_init] * (g[t - np.arange(t)] @ infections[:t]) + ) + return infections[n_init:] + + +def test_mobility_matches_roy_stan_loop( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """Infections match the Roy et al. Stan time loop through heavy depletion.""" + pop = np.asarray(roy_subpop_fractions) * 45e6 + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=roy_mobility_matrix, + subpop_fractions=roy_subpop_fractions, + ) + expected = ( + _roy_stan_loop( + np.asarray(mobility_initial_infections) * pop, + np.asarray(mobility_rt_raw), + np.asarray(gen_int_rv()), + np.asarray(roy_mobility_matrix), + pop, + ) + / pop + ) + + assert_allclose(infections, expected, rtol=1e-10) + depletion = mobility_initial_infections.sum(axis=0) + infections.sum(axis=0) + assert jnp.all(depletion > 0.5) + + +def test_mobility_identity_without_depletion_matches_uncoupled( + gen_int_rv: DeterministicPMF, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """Identity matrix without depletion gives independent plain renewal.""" + reversed_gen_int = jnp.flip(gen_int_rv()) + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=reversed_gen_int, + mobility_matrix=jnp.eye(3), + subpop_fractions=roy_subpop_fractions, + susceptible_depletion=False, + ) + expected = inf.compute_infections_from_rt( + I0=mobility_initial_infections[-reversed_gen_int.size :], + Rt=mobility_rt_raw, + reversed_generation_interval_pmf=reversed_gen_int, + ) + + assert_allclose(infections, expected, rtol=1e-12) + + +def test_mobility_identity_with_depletion_matches_eq1( + gen_int_rv: DeterministicPMF, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """Identity matrix with depletion gives Eq 1 for each region.""" + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=jnp.eye(3), + subpop_fractions=roy_subpop_fractions, + ) + expected = _eq1_loop( + np.asarray(mobility_initial_infections), + np.asarray(mobility_rt_raw), + np.asarray(gen_int_rv()), + ) + + assert_allclose(infections, expected, rtol=1e-10) + + +def test_mobility_empty_location( + empty_location_mobility: tuple[jnp.ndarray, jnp.ndarray], + gen_int_rv: DeterministicPMF, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """A location with no one present adds no infections, whatever its Rt, and its adjusted Rt is 0.""" + mobility_matrix, subpop_fractions = empty_location_mobility + infections, rt_adjusted = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=mobility_matrix, + subpop_fractions=subpop_fractions, + ) + infections_other_rt, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw.at[:, 2].multiply(3.0), + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=mobility_matrix, + subpop_fractions=subpop_fractions, + ) + + assert jnp.all(jnp.isfinite(infections)) + assert_array_equal(infections_other_rt, infections) + assert_array_equal(rt_adjusted[:, 2], 0.0) + + +def test_mobility_zero_population_region_matches_eq1( + gen_int_rv: DeterministicPMF, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """With the identity matrix, a region with no population leaves the others at Eq 1 and has no new infections.""" + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=jnp.eye(3), + subpop_fractions=jnp.array([0.5, 0.5, 0.0]), + ) + expected = _eq1_loop( + np.asarray(mobility_initial_infections), + np.asarray(mobility_rt_raw), + np.asarray(gen_int_rv()), + ) + + assert_allclose(infections[:, :2], expected[:, :2], rtol=1e-10) + assert_array_equal(infections[:, 2], 0.0) + + +def test_mobility_adjusted_rt_without_depletion( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """Without depletion, the adjusted Rt is the raw Rt.""" + _, rt_adjusted = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=roy_mobility_matrix, + subpop_fractions=roy_subpop_fractions, + susceptible_depletion=False, + ) + + assert_array_equal(rt_adjusted, mobility_rt_raw) + + +def test_mobility_conservation( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, + mobility_rt_raw: jnp.ndarray, +) -> None: + """Total new infections equal adjusted Rt times infectiousness present, summed over locations.""" + reversed_gen_int = np.asarray(jnp.flip(gen_int_rv())) + infections, rt_adjusted = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=mobility_rt_raw, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=roy_mobility_matrix, + subpop_fractions=roy_subpop_fractions, + ) + n_init = mobility_initial_infections.shape[0] + n_gen = reversed_gen_int.size + history = np.vstack([mobility_initial_infections, infections]) + infectiousness = np.stack( + [ + reversed_gen_int @ history[n_init + t - n_gen : n_init + t] + for t in range(infections.shape[0]) + ] + ) + p = np.asarray(roy_subpop_fractions) + present = (infectiousness * p) @ np.asarray(roy_mobility_matrix).T + + assert_allclose( + np.asarray(infections) @ p, + np.sum(np.asarray(rt_adjusted) * present, axis=1), + rtol=1e-12, + ) + + +def test_mobility_counts_all_initialization_days( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, +) -> None: + """Initialization days older than the generation interval count toward depletion.""" + reversed_gen_int = jnp.flip(gen_int_rv()) + recent = jnp.full((reversed_gen_int.size, 3), 1e-4) + I0_low = jnp.vstack([jnp.full((13, 3), 1e-4), recent]) + I0_high = jnp.vstack([jnp.full((13, 3), 0.02), recent]) + + first_day = [] + for I0 in (I0_low, I0_high): + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=I0, + Rt_raw=jnp.full((5, 3), 1.5), + reversed_generation_interval_pmf=reversed_gen_int, + mobility_matrix=roy_mobility_matrix, + subpop_fractions=roy_subpop_fractions, + ) + first_day.append(infections[0]) + + assert_allclose( + first_day[1] / first_day[0], + (1.0 - I0_high.sum(axis=0)) / (1.0 - I0_low.sum(axis=0)), + rtol=1e-12, + ) + + +def test_mobility_extreme_rt_caps_incidence( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, +) -> None: + """At extreme Rt, incidence stays within the remaining susceptible fraction.""" + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=jnp.full((60, 3), 30.0), + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=roy_mobility_matrix, + subpop_fractions=roy_subpop_fractions, + ) + initial_cumulative = mobility_initial_infections.sum(axis=0) + cumulative = initial_cumulative + jnp.cumsum(infections, axis=0) + susceptible = 1.0 - jnp.vstack([initial_cumulative, cumulative[:-1]]) + + assert jnp.all(infections >= 0) + assert jnp.all(jnp.isfinite(infections)) + assert jnp.all(infections <= susceptible + 1e-12) + assert jnp.all(cumulative <= 1.0 + 1e-12) + assert_allclose(cumulative[-1], 1.0, atol=1e-12) + + +@pytest.mark.parametrize("susceptible_depletion", [True, False]) +@pytest.mark.parametrize("rt", [0.0, 1.0, 10.0]) +@pytest.mark.parametrize("initial", [0.0, 0.2, 1.0]) +def test_mobility_incidence_boundary( + initial: float, rt: float, susceptible_depletion: bool +) -> None: + """Jitted renewal caps only depleted incidence, including zero and full depletion.""" + compute = jax.jit( + inf.compute_infections_from_rt_with_mobility, + static_argnames="susceptible_depletion", + ) + infections, rt_adjusted = compute( + I0=jnp.array([[initial]]), + Rt_raw=jnp.full((2, 1), rt), + reversed_generation_interval_pmf=jnp.ones(1), + mobility_matrix=jnp.ones((1, 1)), + subpop_fractions=jnp.ones(1), + susceptible_depletion=susceptible_depletion, + ) + first = rt * initial + if susceptible_depletion: + first = min((1.0 - initial) * first, 1.0 - initial) + susceptible = 1.0 - initial - first + second = min(susceptible * rt * first, susceptible) + expected_rt = [rt * (1.0 - initial), rt * susceptible] + else: + second = rt * first + expected_rt = [rt, rt] + + assert_allclose(infections[:, 0], [first, second], atol=1e-12) + assert_allclose(rt_adjusted[:, 0], expected_rt, atol=1e-12) + + +@pytest.mark.parametrize("susceptible_depletion", [True, False]) +@pytest.mark.parametrize("invalid", [-0.1, jnp.nan, jnp.inf]) +def test_mobility_rejects_invalid_initial_infections( + invalid: float, susceptible_depletion: bool +) -> None: + """Initialization must contain only finite, non-negative incident infections.""" + with pytest.raises(ValueError, match="finite, non-negative"): + inf.compute_infections_from_rt_with_mobility( + I0=jnp.array([[invalid], [0.1]]), + Rt_raw=jnp.ones((1, 1)), + reversed_generation_interval_pmf=jnp.ones(1), + mobility_matrix=jnp.ones((1, 1)), + subpop_fractions=jnp.ones(1), + susceptible_depletion=susceptible_depletion, + ) + + +def test_mobility_initialization_cannot_exceed_population() -> None: + """Every initialization day counts toward the per-region population bound.""" + kwargs = { + "I0": jnp.array([[0.6, 0.1], [0.5, 0.1]]), + "Rt_raw": jnp.ones((1, 2)), + "reversed_generation_interval_pmf": jnp.ones(1), + "mobility_matrix": jnp.eye(2), + "subpop_fractions": jnp.full(2, 0.5), + } + with pytest.raises(ValueError, match="<= 1 in every region"): + inf.compute_infections_from_rt_with_mobility(**kwargs) + + infections, _ = inf.compute_infections_from_rt_with_mobility( + **kwargs, susceptible_depletion=False + ) + assert_allclose(infections, [[0.5, 0.1]]) + + +def test_mobility_gradient_finite( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, +) -> None: + """Gradients with respect to Rt stay finite when the incidence ceiling binds.""" + + def total_infections(rt: jnp.ndarray) -> jnp.ndarray: + """ + Sum infections per resident over all days and regions. + + Parameters + ---------- + rt + Reproduction number at each location after initialization. + + Returns + ------- + jnp.ndarray + Scalar total. + """ + infections, _ = inf.compute_infections_from_rt_with_mobility( + I0=mobility_initial_infections, + Rt_raw=rt, + reversed_generation_interval_pmf=jnp.flip(gen_int_rv()), + mobility_matrix=roy_mobility_matrix, + subpop_fractions=roy_subpop_fractions, + ) + return jnp.sum(infections) + + grad = jax.grad(total_infections)(jnp.full((60, 3), 30.0)) + + assert jnp.all(jnp.isfinite(grad)) diff --git a/test/test_interface_coverage.py b/test/test_interface_coverage.py index c03ccb95..24084e35 100644 --- a/test/test_interface_coverage.py +++ b/test/test_interface_coverage.py @@ -24,6 +24,7 @@ HierarchicalNormalPrior, Infections, InfectionsWithFeedback, + MobilitySubpopulationInfections, StudentTGroupModePrior, SubpopulationInfections, ) @@ -241,6 +242,16 @@ def test_get_required_lookback(gen_int_rv): assert infections.get_required_lookback() == expected_length +def test_mobility_get_required_lookback( + mobility_subpopulation_infections: MobilitySubpopulationInfections, + gen_int_rv: DeterministicPMF, +) -> None: + """MobilitySubpopulationInfections lookback is the generation interval length.""" + assert mobility_subpopulation_infections.get_required_lookback() == len( + gen_int_rv() + ) + + # ============================================================================= # SubpopulationInfections.validate() coverage # ============================================================================= @@ -260,6 +271,13 @@ def test_subpopulation_infections_validate(gen_int_rv): infections.validate() +def test_mobility_subpopulation_infections_validate( + mobility_subpopulation_infections: MobilitySubpopulationInfections, +) -> None: + """MobilitySubpopulationInfections.validate() runs without error on valid PMF.""" + mobility_subpopulation_infections.validate() + + # ============================================================================= # MultiSignalModel._validate_observation_resolutions error path # ============================================================================= @@ -451,6 +469,13 @@ def test_subpopulation_infections_name(gen_int_rv): assert infections.name == "test_hi" +def test_mobility_subpopulation_infections_name( + mobility_subpopulation_infections: MobilitySubpopulationInfections, +) -> None: + """MobilitySubpopulationInfections.name is correctly set during construction.""" + assert mobility_subpopulation_infections.name == "subpopulation" + + # ============================================================================= # Internal auto-generated RV names (composition checks) # ============================================================================= diff --git a/test/test_mobility_subpopulation_infections.py b/test/test_mobility_subpopulation_infections.py new file mode 100644 index 00000000..b93bda4e --- /dev/null +++ b/test/test_mobility_subpopulation_infections.py @@ -0,0 +1,415 @@ +""" +Unit tests for MobilitySubpopulationInfections. +""" + +from collections.abc import Callable + +import jax +import jax.numpy as jnp +import numpyro +import numpyro.distributions as dist +import pytest +from numpy.testing import assert_allclose, assert_array_equal +from numpyro.infer import Predictive + +from pyrenew.deterministic import DeterministicPMF, DeterministicVariable +from pyrenew.latent import ( + LatentSample, + MobilitySubpopulationInfections, + SubpopulationInfections, +) +from pyrenew.latent.base import BaseLatentInfectionProcess +from pyrenew.metaclass import RandomVariable +from pyrenew.model import MultiSignalModel +from pyrenew.randomvariable import DistributionalVariable +from test.test_helpers import fixed_random_walk + + +def _sample_with_trace( + process: BaseLatentInfectionProcess, + subpop_fractions: jnp.ndarray, +) -> tuple[dict, LatentSample]: + """ + Sample a latent process for 30 days with a fixed seed, recording its sites. + + Parameters + ---------- + process + Latent infection process to sample. + subpop_fractions + Population fractions for all subpopulations. + + Returns + ------- + tuple[dict, LatentSample] + The NumPyro trace and the sampled latent infections. + """ + with numpyro.handlers.seed(rng_seed=42): + with numpyro.handlers.trace() as trace: + latent = process.sample( + n_days_post_init=30, subpop_fractions=subpop_fractions + ) + return trace, latent + + +class TestMobilitySample: + """Test sampling the mobility-coupled latent process.""" + + @pytest.mark.parametrize("per_region", [False, True]) + def test_identity_without_depletion_matches_subpopulation_infections( + self, + per_region: bool, + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + random_walk_subpopulation_infections_factory: Callable[ + [RandomVariable], SubpopulationInfections + ], + roy_subpop_fractions: jnp.ndarray, + subpop_log_rt_time_0: jnp.ndarray, + ) -> None: + """Identity matrix without depletion records what SubpopulationInfections records.""" + log_rt_time_0 = subpop_log_rt_time_0 if per_region else 0.0 + uncoupled = random_walk_subpopulation_infections_factory( + DeterministicVariable("log_rt_time_0", log_rt_time_0) + ) + coupled = mobility_subpopulation_infections_factory( + jnp.eye(3), + DeterministicVariable("log_rt_time_0", log_rt_time_0), + susceptible_depletion=False, + ) + + trace_uncoupled, latent_uncoupled = _sample_with_trace( + uncoupled, roy_subpop_fractions + ) + trace_coupled, latent_coupled = _sample_with_trace( + coupled, roy_subpop_fractions + ) + + for site_name, site in trace_uncoupled.items(): + assert_allclose( + trace_coupled[site_name]["value"], site["value"], rtol=1e-12 + ) + assert_allclose( + latent_coupled.aggregate, latent_uncoupled.aggregate, rtol=1e-12 + ) + assert_allclose( + latent_coupled.all_subpops, latent_uncoupled.all_subpops, rtol=1e-12 + ) + + def test_aggregate_is_weighted_sum( + self, + mobility_subpopulation_infections: MobilitySubpopulationInfections, + roy_subpop_fractions: jnp.ndarray, + ) -> None: + """Aggregate infections are the population-weighted sum over regions.""" + _, latent = _sample_with_trace( + mobility_subpopulation_infections, roy_subpop_fractions + ) + + assert_allclose( + latent.aggregate, latent.all_subpops @ roy_subpop_fractions, rtol=1e-12 + ) + + @pytest.mark.parametrize("susceptible_depletion", [True, False]) + def test_records_susceptible_fraction( + self, + susceptible_depletion: bool, + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + ) -> None: + """The susceptible fraction is recorded on the full axis, starts at 1, and never increases.""" + process = mobility_subpopulation_infections_factory( + roy_mobility_matrix, + DeterministicVariable("log_rt_time_0", 0.0), + susceptible_depletion=susceptible_depletion, + ) + trace, latent = _sample_with_trace(process, roy_subpop_fractions) + susceptible_fraction = trace["subpopulation::susceptible_fraction"]["value"] + + assert susceptible_fraction.shape == (process.n_initialization_points + 30, 3) + assert jnp.all(susceptible_fraction[0] == 1.0) + assert jnp.all(jnp.diff(susceptible_fraction, axis=0) <= 0.0) + assert_allclose( + susceptible_fraction[1:], + 1.0 - jnp.cumsum(latent.all_subpops, axis=0)[:-1], + atol=1e-12, + ) + + def test_adjusted_rt_uses_susceptible_fraction( + self, + mobility_subpopulation_infections: MobilitySubpopulationInfections, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + ) -> None: + """With depletion, adjusted Rt is raw Rt times the susceptible share present.""" + trace, _ = _sample_with_trace( + mobility_subpopulation_infections, roy_subpop_fractions + ) + n_init = mobility_subpopulation_infections.n_initialization_points + susceptible_fraction = trace["subpopulation::susceptible_fraction"]["value"] + rt_subpop = trace["subpopulation::rt_subpop"]["value"] + rt_subpop_adjusted = trace["subpopulation::rt_subpop_adjusted"]["value"] + share_present = ( + roy_mobility_matrix + * roy_subpop_fractions[jnp.newaxis, :] + / (roy_mobility_matrix @ roy_subpop_fractions)[:, jnp.newaxis] + ) + + assert_allclose( + rt_subpop_adjusted[n_init:], + rt_subpop[n_init:] * (susceptible_fraction[n_init:] @ share_present.T), + rtol=1e-10, + ) + assert_array_equal(rt_subpop_adjusted[:n_init], rt_subpop[:n_init]) + + def test_susceptible_fraction_without_depletion_is_not_floored( + self, + mobility_subpopulation_infections: MobilitySubpopulationInfections, + roy_subpop_fractions: jnp.ndarray, + ) -> None: + """Without depletion, the recorded fraction shows cumulative overshoot.""" + process = mobility_subpopulation_infections + process.susceptible_depletion = False + process.I0_rv = DeterministicVariable("I0", 0.2) + trace, latent = _sample_with_trace(process, roy_subpop_fractions) + susceptible = trace["subpopulation::susceptible_fraction"]["value"] + + assert jnp.any(susceptible < 0.0) + assert_allclose( + susceptible[1:], + 1.0 - jnp.cumsum(latent.all_subpops, axis=0)[:-1], + ) + + def test_adjusted_rt_without_depletion( + self, + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + ) -> None: + """Without depletion, adjusted Rt equals raw Rt.""" + process = mobility_subpopulation_infections_factory( + roy_mobility_matrix, + DeterministicVariable("log_rt_time_0", 0.0), + susceptible_depletion=False, + ) + trace, _ = _sample_with_trace(process, roy_subpop_fractions) + + assert_array_equal( + trace["subpopulation::rt_subpop_adjusted"]["value"], + trace["subpopulation::rt_subpop"]["value"], + ) + + +class TestMobilityValidation: + """Test validation of the mobility matrix and its fit to the population.""" + + def test_rejects_excessive_initialization( + self, + mobility_subpopulation_infections: MobilitySubpopulationInfections, + roy_subpop_fractions: jnp.ndarray, + ) -> None: + """Sampling rejects initialization whose total exceeds a region's population.""" + mobility_subpopulation_infections.I0_rv = DeterministicVariable("I0", 0.2) + with pytest.raises(ValueError, match="<= 1 in every region"): + _sample_with_trace(mobility_subpopulation_infections, roy_subpop_fractions) + + @pytest.mark.parametrize( + ("case", "match"), + [ + ("non_square", "square"), + ("three_dimensional", "square"), + ("negative", "non-negative"), + ("nan", "finite"), + ("row_stochastic", "transpose"), + ("columns_not_one", "sum to 1"), + ], + ) + def test_rejects_invalid_matrix( + self, + case: str, + match: str, + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + roy_mobility_matrix: jnp.ndarray, + ) -> None: + """Invalid matrices are rejected at construction with a message naming the problem.""" + matrices = { + "non_square": jnp.ones((2, 3)) / 2, + "three_dimensional": jnp.stack([roy_mobility_matrix, roy_mobility_matrix]), + "negative": jnp.array([[1.1, 0.0, 0.0], [-0.1, 1.0, 0.0], [0.0, 0.0, 1.0]]), + "nan": roy_mobility_matrix.at[0, 0].set(jnp.nan), + "row_stochastic": roy_mobility_matrix.T, + "columns_not_one": 0.5 * jnp.eye(3), + } + + with pytest.raises(ValueError, match=match): + mobility_subpopulation_infections_factory( + matrices[case], DeterministicVariable("log_rt_time_0", 0.0) + ) + + def test_column_sum_message_without_transpose( + self, + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + ) -> None: + """The transpose hint is given only when the rows sum to 1.""" + with pytest.raises(ValueError) as excinfo: + mobility_subpopulation_infections_factory( + 0.5 * jnp.eye(3), DeterministicVariable("log_rt_time_0", 0.0) + ) + + assert "transpose" not in str(excinfo.value) + + def test_rejects_random_matrix_rv(self, gen_int_rv: DeterministicPMF) -> None: + """A random mobility matrix is rejected.""" + with pytest.raises(TypeError, match="DeterministicVariable"): + MobilitySubpopulationInfections( + name="subpopulation", + gen_int_rv=gen_int_rv, + n_initialization_points=7, + I0_rv=DeterministicVariable("I0", 0.001), + log_rt_time_0_rv=DeterministicVariable("log_rt_time_0", 0.0), + baseline_rt_process=fixed_random_walk(innovation_sd=0.05), + subpop_rt_deviation_process=fixed_random_walk(innovation_sd=0.025), + mobility_matrix_rv=DistributionalVariable( + "mobility_matrix", dist.Dirichlet(jnp.ones(3)).expand([3]) + ), + ) + + def test_rejects_feedback_arguments( + self, + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + ) -> None: + """The class has no infection feedback option.""" + with pytest.raises(TypeError, match="infection_feedback_strength_rv"): + MobilitySubpopulationInfections( + name="subpopulation", + gen_int_rv=gen_int_rv, + n_initialization_points=7, + I0_rv=DeterministicVariable("I0", 0.001), + log_rt_time_0_rv=DeterministicVariable("log_rt_time_0", 0.0), + baseline_rt_process=fixed_random_walk(innovation_sd=0.05), + subpop_rt_deviation_process=fixed_random_walk(innovation_sd=0.025), + mobility_matrix_rv=DeterministicVariable( + "mobility_matrix", roy_mobility_matrix + ), + infection_feedback_strength_rv=DeterministicVariable("strength", -1.0), + ) + + def test_rejects_size_mismatch( + self, + mobility_subpopulation_infections: MobilitySubpopulationInfections, + ) -> None: + """A matrix that does not match the number of subpopulations is rejected.""" + with pytest.raises(ValueError) as excinfo: + mobility_subpopulation_infections._parse_and_validate_fractions( + jnp.array([0.5, 0.5]) + ) + + message = str(excinfo.value) + assert "(3, 3)" in message + assert "2 entries" in message + + def test_accepts_location_without_population( + self, + empty_location_mobility: tuple[jnp.ndarray, jnp.ndarray], + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + ) -> None: + """A location with no one present samples finite infections and adjusted Rt 0 there.""" + mobility_matrix, subpop_fractions = empty_location_mobility + process = mobility_subpopulation_infections_factory( + mobility_matrix, DeterministicVariable("log_rt_time_0", 0.0) + ) + trace, latent = _sample_with_trace(process, subpop_fractions) + n_init = process.n_initialization_points + rt_subpop_adjusted = trace["subpopulation::rt_subpop_adjusted"]["value"] + + assert jnp.all(jnp.isfinite(latent.all_subpops)) + assert_array_equal(rt_subpop_adjusted[n_init:, 2], 0.0) + + +class TestMobilityModel: + """Test the mobility latent process inside a builder-constructed model.""" + + def test_builder_prior_predictive( + self, + mobility_model_factory: Callable[[jnp.ndarray], MultiSignalModel], + roy_mobility_matrix: jnp.ndarray, + mobility_model_args: dict[str, object], + ) -> None: + """Prior predictive sampling through PyrenewBuilder gives the expected shapes.""" + model = mobility_model_factory(roy_mobility_matrix) + samples = Predictive(model.sample, num_samples=5)( + jax.random.PRNGKey(0), **mobility_model_args + ) + n_total = model.latent.n_initialization_points + 30 + + assert samples["latent_infections_by_subpop"].shape == (5, n_total, 3) + assert jnp.all(samples["latent_infections_by_subpop"] > 0) + assert samples["subpopulation::susceptible_fraction"].shape == (5, n_total, 3) + + def test_validate_data_rejects_before_mcmc( + self, + mobility_model_factory: Callable[[jnp.ndarray], MultiSignalModel], + roy_mobility_matrix: jnp.ndarray, + ) -> None: + """A matrix that does not match the number of subpopulations is rejected before MCMC starts.""" + model = mobility_model_factory(roy_mobility_matrix) + fractions = jnp.array([0.5, 0.5]) + + with pytest.raises(ValueError, match=r"\(3, 3\)"): + model.validate_data(n_days_post_init=30, subpop_fractions=fractions) + with pytest.raises(ValueError, match=r"\(3, 3\)"): + model.run( + num_warmup=10, + num_samples=10, + n_days_post_init=30, + population_size=45_000_000, + subpop_fractions=fractions, + hospital_subpop={ + "obs": None, + "subpop_indices": jnp.arange(fractions.size), + }, + ) + assert model.mcmc is None + + def test_counterfactual_with_another_matrix( + self, + mobility_model_factory: Callable[[jnp.ndarray], MultiSignalModel], + roy_mobility_matrix: jnp.ndarray, + mobility_model_args: dict[str, object], + ) -> None: + """Draws replayed through a model with another matrix give different infections.""" + model = mobility_model_factory(roy_mobility_matrix) + identity_model = mobility_model_factory(jnp.eye(3)) + draws = Predictive(model.sample, num_samples=4)( + jax.random.PRNGKey(0), **mobility_model_args + ) + + replay = Predictive(model.sample, posterior_samples=draws)( + jax.random.PRNGKey(1), **mobility_model_args + ) + counterfactual = Predictive(identity_model.sample, posterior_samples=draws)( + jax.random.PRNGKey(1), **mobility_model_args + ) + + assert_allclose( + replay["latent_infections_by_subpop"], + draws["latent_infections_by_subpop"], + rtol=1e-12, + ) + assert not jnp.allclose( + counterfactual["latent_infections_by_subpop"], + draws["latent_infections_by_subpop"], + )