From 7869052fb8daaec665459b4112346a2ec654c140 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 21 Oct 2026 16:55:43 -0400 Subject: [PATCH 01/21] merge fix --- .gitignore | 2 + docs/.gitignore | 2 + docs/index.md | 1 + docs/tutorials/.pages | 1 + docs/tutorials/latent_mobility_infections.qmd | 430 ++++++++++++++++++ .../latent_subpopulation_infections.qmd | 2 + pyrenew/latent/__init__.py | 12 +- pyrenew/latent/infection_functions.py | 151 ++++++ .../mobility_subpopulation_infections.py | 294 ++++++++++++ pyrenew/latent/subpopulation_infections.py | 368 ++++++++++----- test/conftest.py | 245 ++++++++++ test/test_infection_functions.py | 314 ++++++++++++- test/test_interface_coverage.py | 25 + .../test_mobility_subpopulation_infections.py | 389 ++++++++++++++++ 14 files changed, 2127 insertions(+), 109 deletions(-) create mode 100644 docs/tutorials/latent_mobility_infections.qmd create mode 100644 pyrenew/latent/mobility_subpopulation_infections.py create mode 100644 test/test_mobility_subpopulation_infections.py diff --git a/.gitignore b/.gitignore index 14a319203..d93f2737e 100755 --- a/.gitignore +++ b/.gitignore @@ -332,3 +332,5 @@ replay_pid* /.quarto/ *_files + +**/*.quarto_ipynb diff --git a/docs/.gitignore b/docs/.gitignore index 075b2542a..0e3521a7d 100644 --- a/docs/.gitignore +++ b/docs/.gitignore @@ -1 +1,3 @@ /.quarto/ + +**/*.quarto_ipynb diff --git a/docs/index.md b/docs/index.md index 213754e9a..e4a5ea4e5 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 mobility infections](tutorials/latent_mobility_infections.md) -- 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 fd693b8db..6aea842fa 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 000000000..2437b599e --- /dev/null +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -0,0 +1,430 @@ +--- +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_{t,j} \left[ \sum_{k=1}^{n} m_{jk} \sum_{\tau} I_{\tau,k}\, g_{t-\tau} \right] +$$ + +where: + +- $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_{t,j}$ is the reproduction number at location $j$, where transmission happens +- $g$ is the generation interval + +PyRenew works with infections per resident. +With `M` the mobility matrix and `p` the population fractions, one day of the computation is: + +```python +share_present = M * p[None, :] / (M @ p)[:, None] +infectiousness = reversed_gen_int @ history +susceptible = 1 - cumulative +new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) +``` + +`share_present[j, k]` is the share of the people present at location `j` who live in region `k`. +`rt` multiplies at each location, because $R_{t,j}$ describes transmission where it happens. +`M.T` assigns infections acquired at each location back to the home region of the person infected. +`susceptible` multiplies last, 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"] + +print(f"Column sums: {M.sum(axis=0)}") +print(f"Row sums: {M.sum(axis=1)}") +``` + +`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. +In region 3, 70% of residents stay home on a typical day, 21% spend it in region 1, and 9% in region 2. + +A matrix with home regions in its rows, so that each row sums to 1, must be transposed before use. +`MobilitySubpopulationInfections` rejects it, and the error message says to pass the transpose. + +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`. +A matrix stored as a data frame must be converted with `.to_numpy()`, with rows and columns in that order. + +A matrix for part of a larger area needs one more step. +After dropping the flows to places outside the area, the columns no longer sum to 1, so each column is renormalized. +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 infections per resident on the first day of the initialization period, the earliest day on the time axis. +It is not prevalence. +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. + +This comparison is a counterfactual: the same $\mathcal{R}(t)$ under two matrices. +A fitted model gives one the same way. +Build a second model with the other matrix and pass the fitted posterior draws to NumPyro's `Predictive`: + +```python +Predictive(model2.model, posterior_samples=model1.mcmc.get_samples()) +``` + +`model2.posterior_predictive()` does not work for this, because `model2` has not been fitted. + +## 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. + +`rt_subpop_adjusted` is the reproduction number at each location after depletion: new infections caused per unit of infectiousness present there. +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. +With `susceptible_depletion=False`, the $S_{t,i}$ factor is left out, and over a long epidemic `rt_subpop` absorbs the depletion instead. + +`susceptible_fraction[t, i]` is the share of residents of region `i` not infected before day `t`. +It is recorded with depletion on or off. +With depletion off, the renewal does not use it, and it shows how much depletion the model leaves out. + +```{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_wrap("~ region", ncol=3) + + p9.labs( + x="Days", + y="Susceptible fraction", + title="Share of Residents Not Yet Infected", + color="", + ) + + theme_tutorial + + p9.theme(legend_position="bottom") +) +``` + +With depletion off, the susceptible fraction reaches zero, which means that the model infects more residents than there are. +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 has no function for this split. +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. +In this simulation of the latent process alone, the infections per resident are the values the process returns. +After fitting a model, they are `latent_infections_by_subpop / (population_size * p)`. + +```{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 +share_present = M * p[None, :] / (M @ p)[:, 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 + outgoing = susceptible * (M.T @ at_location) - local - incoming + sources.append(np.stack([local, incoming, outgoing])) +sources = np.stack(sources) + +total = sources.sum(axis=1) +print( + "Largest relative difference between the summed sources and the infections: " + f"{np.max(np.abs(total - per_resident[n_init:]) / 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 = sources / total[:, np.newaxis, :] +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`, and that every location has people present. + +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 e3e8b975d..2f3e2a74b 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 4203bf40e..4c5e97694 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 86b15b398..394dac147 100755 --- a/pyrenew/latent/infection_functions.py +++ b/pyrenew/latent/infection_functions.py @@ -2,6 +2,8 @@ from __future__ import annotations +import functools + import jax import jax.numpy as jnp from jax.typing import ArrayLike @@ -183,3 +185,152 @@ 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 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 + adjusted reproduction number at each location. + """ + history, cumulative = state + infectiousness = reversed_generation_interval_pmf @ history + if susceptible_depletion: + susceptible = jnp.maximum(1.0 - cumulative, 0.0) + rt_adjusted = rt_raw * (share_present @ susceptible) + else: + susceptible = jnp.ones_like(cumulative) + rt_adjusted = rt_raw + new_infections = susceptible * ( + mobility_matrix.T @ (rt_raw * (share_present @ infectiousness)) + ) + 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. + Rt_raw + 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 the fraction of those residents not yet infected. + + Returns + ------- + tuple[jnp.ndarray, jnp.ndarray] + New infections per resident, shape ``(n_days, n)``, and the + reproduction number at each location after depletion, the same + shape. Without depletion, the second array is ``Rt_raw``. + + Notes + ----- + This is Equation 2 of Roy, Clapham, and Mishra (2025), + , divided by the number + of residents of region $i$: + + ```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 is: + + ```python + share_present = M * p[None, :] / (M @ p)[:, None] + infectiousness = reversed_generation_interval_pmf @ history + susceptible = jnp.maximum(1 - cumulative, 0) + new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) + ``` + + ``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. + + The susceptible fraction is floored at zero. On a day with extreme + reproduction numbers, new infections can exceed the residents still + susceptible. That day is not capped, but later days stay non-negative. + """ + share_present = ( + mobility_matrix + * subpop_fractions[jnp.newaxis, :] + / (mobility_matrix @ subpop_fractions)[:, jnp.newaxis] + ) + n_gen = reversed_generation_interval_pmf.size + initial_state = (I0[-n_gen:], jnp.sum(I0, axis=0)) + 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 000000000..72563d6d3 --- /dev/null +++ b/pyrenew/latent/mobility_subpopulation_infections.py @@ -0,0 +1,294 @@ +""" +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 numpyro.util import not_jax_tracer + +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 is 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). + + 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 is: + + ```python + share_present = M * p[None, :] / (M @ p)[:, None] + infectiousness = reversed_gen_int @ history + susceptible = 1 - cumulative + new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) + ``` + + ``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. + + ``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: new infections caused per unit of + infectiousness present at each location, after depletion. + - ``susceptible_fraction[t, i]`` is the share of residents of home + region ``i`` not infected before day ``t``. + + Susceptible depletion is on by default, as in Eq 2. With it off, + ``susceptible`` is 1 and ``rt_subpop`` absorbs depletion. There is no + infection feedback option: Eq 2 accounts for depletion directly. + + 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. + 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. If False, that factor is left out 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, has a column that does not sum to 1, or has + a row with no positive entry. + """ + 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. Pass the transpose." + ) + raise ValueError(message) + empty_rows = jnp.flatnonzero(jnp.sum(matrix, axis=1) == 0).tolist() + if empty_rows: + raise ValueError( + f"Locations {empty_rows} have no residents of any region " + "present: their rows of mobility_matrix_rv are all zero." + ) + + 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, if the matrix does not have one row + and one column per subpopulation, or if a location has no + population present. + """ + 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})." + ) + present = matrix @ pop.fractions + empty_check = jnp.any(present <= 0) + if not_jax_tracer(empty_check) and empty_check: + empty = jnp.flatnonzero(present <= 0).tolist() + raise ValueError( + f"Locations {empty} have no population present: " + "mobility_matrix @ subpop_fractions must be positive." + ) + 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 depletion. + """ + 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 = jnp.maximum(1.0 - infected_before, 0.0) + 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 5d7eb9b20..8482ea97f 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 cabb5d904..102554a95 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -15,13 +15,16 @@ 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 +311,248 @@ 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)) + + +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 752f99c36..8c8db2619 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,312 @@ 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_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_nonnegative( + gen_int_rv: DeterministicPMF, + roy_mobility_matrix: jnp.ndarray, + roy_subpop_fractions: jnp.ndarray, + mobility_initial_infections: jnp.ndarray, +) -> None: + """At extreme Rt, infections overshoot the susceptibles but stay non-negative and finite.""" + 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, + ) + cumulative = mobility_initial_infections.sum(axis=0) + infections.sum(axis=0) + + assert jnp.all(infections >= 0) + assert jnp.all(jnp.isfinite(infections)) + assert jnp.any(cumulative > 1.0) + + +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 susceptible floor is reached.""" + + 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 c03ccb956..24084e358 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 000000000..d22b24ee2 --- /dev/null +++ b/test/test_mobility_subpopulation_infections.py @@ -0,0 +1,389 @@ +""" +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, _ = _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) + + 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_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.""" + + @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"), + ("empty_row", "no residents"), + ], + ) + 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), + "empty_row": jnp.array([[1.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + } + + 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_rejects_location_without_population( + self, + mobility_subpopulation_infections_factory: Callable[ + ..., MobilitySubpopulationInfections + ], + ) -> None: + """A location with no population present is rejected.""" + process = mobility_subpopulation_infections_factory( + jnp.eye(3), DeterministicVariable("log_rt_time_0", 0.0) + ) + + with pytest.raises(ValueError, match=r"Locations \[2\]"): + process._parse_and_validate_fractions(jnp.array([0.5, 0.5, 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) + + @pytest.mark.parametrize( + ("case", "fractions", "match"), + [ + ("roy", jnp.array([0.5, 0.5]), r"\(3, 3\)"), + ("identity", jnp.array([0.5, 0.5, 0.0]), r"Locations \[2\]"), + ], + ) + def test_validate_data_rejects_before_mcmc( + self, + case: str, + fractions: jnp.ndarray, + match: str, + mobility_model_factory: Callable[[jnp.ndarray], MultiSignalModel], + roy_mobility_matrix: jnp.ndarray, + ) -> None: + """Size mismatches and empty locations are rejected before MCMC starts.""" + matrices = {"roy": roy_mobility_matrix, "identity": jnp.eye(3)} + model = mobility_model_factory(matrices[case]) + + with pytest.raises(ValueError, match=match): + model.validate_data(n_days_post_init=30, subpop_fractions=fractions) + with pytest.raises(ValueError, match=match): + 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"], + ) From ce9bbd6af22e2ce55fc03408ea5873ee327cbf7b Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 21 Oct 2026 17:08:22 -0400 Subject: [PATCH 02/21] merge fix --- test/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/conftest.py b/test/conftest.py index 102554a95..c66f68ffa 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -12,6 +12,7 @@ from collections.abc import Callable +import jax import jax.numpy as jnp import numpyro.distributions as dist import pytest From aca9ae760a85e4e0468b68eedcc2246676396e2b Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 21 Oct 2026 18:01:11 -0400 Subject: [PATCH 03/21] Revise notes and citation in infection_functions.py Updated citation format and clarified equation context. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- pyrenew/latent/infection_functions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyrenew/latent/infection_functions.py b/pyrenew/latent/infection_functions.py index 394dac147..8913c8d67 100755 --- a/pyrenew/latent/infection_functions.py +++ b/pyrenew/latent/infection_functions.py @@ -288,8 +288,8 @@ def compute_infections_from_rt_with_mobility( Notes ----- This is Equation 2 of Roy, Clapham, and Mishra (2025), - , divided by the number - of residents of region $i$: + . In infection counts (before + converting to infections per resident), 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} From b4cd9466d986b164de4949fd39f902ac313ad579 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 09:38:53 -0400 Subject: [PATCH 04/21] Update docs/index.md Co-authored-by: Dylan H. Morris --- docs/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/index.md b/docs/index.md index e4a5ea4e5..f9aec02b4 100644 --- a/docs/index.md +++ b/docs/index.md @@ -105,7 +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 mobility infections](tutorials/latent_mobility_infections.md) -- subpopulations coupled by a mobility matrix. +- [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. From ea85bafd2ba6b9ca86b6b2d5bfc6179eaa69151b Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 09:48:09 -0400 Subject: [PATCH 05/21] changes per code review --- docs/tutorials/latent_mobility_infections.qmd | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 2437b599e..4574d08bf 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -61,6 +61,7 @@ $$ where: +- $I_{t,i} is the infection count (a non-negative integer) - $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$ From 560ce4504c60f7830a717e56fc6660e34f2b076c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:48:28 +0000 Subject: [PATCH 06/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/tutorials/latent_mobility_infections.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 4574d08bf..3d2a22f2b 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -61,7 +61,7 @@ $$ where: -- $I_{t,i} is the infection count (a non-negative integer) +- $I\_{t,i} is the infection count (a non-negative integer) - $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$ From 61f1c28361914ea2915a5c88c55afb5b55efa926 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:08:48 -0400 Subject: [PATCH 07/21] doc changes per code review --- docs/tutorials/latent_mobility_infections.qmd | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 4574d08bf..821bf4983 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -61,7 +61,7 @@ $$ where: -- $I_{t,i} is the infection count (a non-negative integer) +- $I_{t,i}$ is the infection count (a non-negative integer) - $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$ @@ -74,11 +74,15 @@ With `M` the mobility matrix and `p` the population fractions, one day of the co ```python share_present = M * p[None, :] / (M @ p)[:, None] infectiousness = reversed_gen_int @ history -susceptible = 1 - cumulative +susceptible = jnp.maximum(1 - cumulative, 0) new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) ``` `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. +`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. +`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. +`susceptible` is floored at zero, so that a day of extreme infections cannot exceed the remaining susceptibles. `rt` multiplies at each location, because $R_{t,j}$ describes transmission where it happens. `M.T` assigns infections acquired at each location back to the home region of the person infected. `susceptible` multiplies last, because residents of a region lose susceptibility wherever they were infected. From 41d23f4a11678d2ac1f944b69c059c78413da8cc Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:10:18 -0400 Subject: [PATCH 08/21] Update docs/tutorials/latent_mobility_infections.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/latent_mobility_infections.qmd | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 821bf4983..1f1bf9c2b 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -146,7 +146,6 @@ 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 infections per resident on the first day of the initialization period, the earliest day on the time axis. -It is not prevalence. 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`. From 11977e42627736eb8cb85aa85d7dec5974335743 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:10:46 -0400 Subject: [PATCH 09/21] Update docs/tutorials/latent_mobility_infections.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/latent_mobility_infections.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 1f1bf9c2b..dec04b5a0 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -145,7 +145,7 @@ print(f"Column sums after renormalizing: {M_area.sum(axis=0)}") 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 infections per resident on the first day of the initialization period, the earliest day on the time axis. +`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`. From 77bb0df9d7b544216578fc60ee74d232df14a9a8 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:11:37 -0400 Subject: [PATCH 10/21] Update docs/tutorials/latent_mobility_infections.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/latent_mobility_infections.qmd | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index dec04b5a0..c49bc1948 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -119,8 +119,7 @@ A matrix with home regions in its rows, so that each row sums to 1, must be tran 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`. A matrix stored as a data frame must be converted with `.to_numpy()`, with rows and columns in that order. -A matrix for part of a larger area needs one more step. -After dropping the flows to places outside the area, the columns no longer sum to 1, so each column is renormalized. +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. From b1662e961cd8f6f229edc0dfda69ebd9b7420db4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:11:59 +0000 Subject: [PATCH 11/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/tutorials/latent_mobility_infections.qmd | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index c49bc1948..9eb089bd3 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -119,7 +119,8 @@ A matrix with home regions in its rows, so that each row sums to 1, must be tran 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`. A matrix stored as a data frame must be converted with `.to_numpy()`, with rows and columns in that order. -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. +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. From 353d84cdd855acc381f307bb71c39a5d5d4bcdab Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:13:46 -0400 Subject: [PATCH 12/21] Update docs/tutorials/latent_mobility_infections.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/latent_mobility_infections.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 9eb089bd3..6fa70e480 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -300,7 +300,7 @@ With `susceptible_depletion=False`, the $S_{t,i}$ factor is left out, and over a `susceptible_fraction[t, i]` is the share of residents of region `i` not infected before day `t`. It is recorded with depletion on or off. -With depletion off, the renewal does not use it, and it shows how much depletion the model leaves out. +With `susceptible_depletion=False`, the renewal equation does not consider it, and the epidemic rapidly reaches 100% of the population infected. ```{python} #| label: depletion From 4e317f5cb728e3aaaf37f0b396830e88574a1f65 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:14:01 -0400 Subject: [PATCH 13/21] Update docs/tutorials/latent_mobility_infections.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/latent_mobility_infections.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 6fa70e480..c2e42f3a7 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -342,7 +342,7 @@ Roy et al. - **incoming**: residents infected at home by infectious visitors from other regions; - **outgoing**: residents infected while visiting other regions. -PyRenew has no function for this split. +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. In this simulation of the latent process alone, the infections per resident are the values the process returns. From d02cdb1237ae48592f8021b67ac52baecea2fe98 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:14:25 +0000 Subject: [PATCH 14/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/tutorials/latent_mobility_infections.qmd | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index c2e42f3a7..74af6b494 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -342,8 +342,7 @@ Roy et al. - **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. +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. In this simulation of the latent process alone, the infections per resident are the values the process returns. After fitting a model, they are `latent_infections_by_subpop / (population_size * p)`. From 4a4f107dd9107a6f4e83c6c7d50c55041884ec48 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 10:47:59 -0400 Subject: [PATCH 15/21] doc changes per code review --- docs/tutorials/latent_mobility_infections.qmd | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 74af6b494..5e2028a03 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -249,16 +249,6 @@ Linked to the other regions, it has an outbreak: its residents are infected whil Roy et al. show the same effect in their S2 Fig. -This comparison is a counterfactual: the same $\mathcal{R}(t)$ under two matrices. -A fitted model gives one the same way. -Build a second model with the other matrix and pass the fitted posterior draws to NumPyro's `Predictive`: - -```python -Predictive(model2.model, posterior_samples=model1.mcmc.get_samples()) -``` - -`model2.posterior_predictive()` does not work for this, because `model2` has not been fitted. - ## Where $\mathcal{R}(t)$ Applies `rt_subpop[t, j]` is the reproduction number at location `j`, where transmission happens. @@ -296,7 +286,6 @@ rt_df = pd.concat( ## Susceptible Depletion Susceptible depletion is on by default, as in Eq 2. -With `susceptible_depletion=False`, the $S_{t,i}$ factor is left out, and over a long epidemic `rt_subpop` absorbs the depletion instead. `susceptible_fraction[t, i]` is the share of residents of region `i` not infected before day `t`. It is recorded with depletion on or off. @@ -342,10 +331,8 @@ Roy et al. - **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. +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. -In this simulation of the latent process alone, the infections per resident are the values the process returns. -After fitting a model, they are `latent_infections_by_subpop / (population_size * p)`. ```{python} #| label: sources From 33a586722b59e97d943d43fdeb442f305e8b1c6e Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 11:25:37 -0400 Subject: [PATCH 16/21] tutorial edits --- docs/tutorials/latent_mobility_infections.qmd | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 5e2028a03..b4ff61c3a 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -56,7 +56,7 @@ The class implements Equation 2 of Roy, Clapham, and Mishra (2025), [Incorporati 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_{t,j} \left[ \sum_{k=1}^{n} m_{jk} \sum_{\tau} I_{\tau,k}\, g_{t-\tau} \right] + 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: @@ -65,7 +65,7 @@ where: - $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_{t,j}$ is the reproduction number at location $j$, where transmission happens +- $R_j(t)$ is the reproduction number at location $j$, where transmission happens - $g$ is the generation interval PyRenew works with infections per resident. @@ -78,14 +78,14 @@ susceptible = jnp.maximum(1 - cumulative, 0) new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) ``` -`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. -`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. -`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. -`susceptible` is floored at zero, so that a day of extreme infections cannot exceed the remaining susceptibles. -`rt` multiplies at each location, because $R_{t,j}$ describes transmission where it happens. -`M.T` assigns infections acquired at each location back to the home region of the person infected. -`susceptible` multiplies last, because residents of a region lose susceptibility wherever they were infected. +- `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. +- `susceptible` is floored at zero, so that a day of extreme infections cannot exceed the remaining susceptibles. +- `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` multiplies last, because residents of a region lose susceptibility wherever they were infected. ## The Mobility Matrix From ca1d47a0ba338186403cbaef26c02ae48f103b03 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 12:11:12 -0400 Subject: [PATCH 17/21] Update pyrenew/latent/mobility_subpopulation_infections.py Co-authored-by: Dylan H. Morris --- pyrenew/latent/mobility_subpopulation_infections.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyrenew/latent/mobility_subpopulation_infections.py b/pyrenew/latent/mobility_subpopulation_infections.py index 72563d6d3..d0f97fd2f 100644 --- a/pyrenew/latent/mobility_subpopulation_infections.py +++ b/pyrenew/latent/mobility_subpopulation_infections.py @@ -67,8 +67,8 @@ class MobilitySubpopulationInfections(BaseSubpopulationInfections): region ``i`` not infected before day ``t``. Susceptible depletion is on by default, as in Eq 2. With it off, - ``susceptible`` is 1 and ``rt_subpop`` absorbs depletion. There is no - infection feedback option: Eq 2 accounts for depletion directly. + ``susceptible`` is 1 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``. From 308dce69c8d44f0c4e778eb940539d21087611f1 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 12:11:33 -0400 Subject: [PATCH 18/21] Update pyrenew/latent/mobility_subpopulation_infections.py Co-authored-by: Dylan H. Morris --- pyrenew/latent/mobility_subpopulation_infections.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/latent/mobility_subpopulation_infections.py b/pyrenew/latent/mobility_subpopulation_infections.py index d0f97fd2f..c696e8d6b 100644 --- a/pyrenew/latent/mobility_subpopulation_infections.py +++ b/pyrenew/latent/mobility_subpopulation_infections.py @@ -194,7 +194,7 @@ def _validate_mobility_matrix(matrix: ArrayLike) -> None: 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. Pass the transpose." + "fraction of residents of k present at j. Did you mean to pass the transpose?" ) raise ValueError(message) empty_rows = jnp.flatnonzero(jnp.sum(matrix, axis=1) == 0).tolist() From c2c413ff9a287b1fb290a2360cf4763a8e294b53 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 12:14:34 -0400 Subject: [PATCH 19/21] tutorial tweaks; plot mobility matrix --- docs/tutorials/latent_mobility_infections.qmd | 44 +++++++++++++++---- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index b4ff61c3a..2c6f5b795 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -105,19 +105,47 @@ M = np.array( p = np.array([20.0, 10.0, 15.0]) / 45.0 region_names = ["Region 1", "Region 2", "Region 3"] -print(f"Column sums: {M.sum(axis=0)}") -print(f"Row sums: {M.sum(axis=1)}") +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, rows are locations, and each column sums to 1. +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. -A matrix with home regions in its rows, so that each row sums to 1, must be transposed before use. -`MobilitySubpopulationInfections` rejects it, and the error message says to pass the transpose. - -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`. -A matrix stored as a data frame must be converted with `.to_numpy()`, with rows and columns in that order. +- 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. From 6ed31626e54a181f1648404daca54ade04ef79ed Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 13:20:53 -0400 Subject: [PATCH 20/21] changes per code review --- pyrenew/latent/infection_functions.py | 6 ++- .../mobility_subpopulation_infections.py | 26 ++-------- test/conftest.py | 26 ++++++++++ test/test_infection_functions.py | 51 +++++++++++++++++++ .../test_mobility_subpopulation_infections.py | 37 ++++++-------- 5 files changed, 102 insertions(+), 44 deletions(-) diff --git a/pyrenew/latent/infection_functions.py b/pyrenew/latent/infection_functions.py index 8913c8d67..c9649efd1 100755 --- a/pyrenew/latent/infection_functions.py +++ b/pyrenew/latent/infection_functions.py @@ -314,14 +314,18 @@ def compute_infections_from_rt_with_mobility( ``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, and its + reproduction number after depletion is 0. + The susceptible fraction is floored at zero. On a day with extreme reproduction numbers, new infections can exceed the residents still susceptible. That day is not capped, but later days stay non-negative. """ + present = mobility_matrix @ subpop_fractions share_present = ( mobility_matrix * subpop_fractions[jnp.newaxis, :] - / (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:], jnp.sum(I0, axis=0)) diff --git a/pyrenew/latent/mobility_subpopulation_infections.py b/pyrenew/latent/mobility_subpopulation_infections.py index c696e8d6b..c5672c1c8 100644 --- a/pyrenew/latent/mobility_subpopulation_infections.py +++ b/pyrenew/latent/mobility_subpopulation_infections.py @@ -7,7 +7,6 @@ import jax.numpy as jnp import numpyro from jax.typing import ArrayLike -from numpyro.util import not_jax_tracer from pyrenew.deterministic import DeterministicVariable from pyrenew.latent.base import PopulationStructure @@ -62,7 +61,8 @@ class MobilitySubpopulationInfections(BaseSubpopulationInfections): - ``rt_subpop_adjusted`` is ``rt_subpop * (share_present @ susceptible_fraction)`` after the initialization period: new infections caused per unit of - infectiousness present at each location, after depletion. + infectiousness present at each location, after depletion. 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``. @@ -173,8 +173,7 @@ def _validate_mobility_matrix(matrix: ArrayLike) -> None: ------ ValueError If the matrix is not square and two-dimensional, has non-finite - or negative entries, has a column that does not sum to 1, or has - a row with no positive entry. + 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( @@ -197,12 +196,6 @@ def _validate_mobility_matrix(matrix: ArrayLike) -> None: "fraction of residents of k present at j. Did you mean to pass the transpose?" ) raise ValueError(message) - empty_rows = jnp.flatnonzero(jnp.sum(matrix, axis=1) == 0).tolist() - if empty_rows: - raise ValueError( - f"Locations {empty_rows} have no residents of any region " - "present: their rows of mobility_matrix_rv are all zero." - ) def _parse_and_validate_fractions( self, @@ -224,9 +217,8 @@ def _parse_and_validate_fractions( Raises ------ ValueError - If the fractions are invalid, if the matrix does not have one row - and one column per subpopulation, or if a location has no - population present. + 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) @@ -236,14 +228,6 @@ def _parse_and_validate_fractions( f"subpop_fractions has {pop.n_subpops} entries; the matrix " f"must have shape ({pop.n_subpops}, {pop.n_subpops})." ) - present = matrix @ pop.fractions - empty_check = jnp.any(present <= 0) - if not_jax_tracer(empty_check) and empty_check: - empty = jnp.flatnonzero(present <= 0).tolist() - raise ValueError( - f"Locations {empty} have no population present: " - "mobility_matrix @ subpop_fractions must be positive." - ) return pop def _infections_from_rt( diff --git a/test/conftest.py b/test/conftest.py index c66f68ffa..aa6aff1cb 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -372,6 +372,32 @@ def mobility_rt_raw() -> jnp.ndarray: 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, diff --git a/test/test_infection_functions.py b/test/test_infection_functions.py index 8c8db2619..eda5fb092 100644 --- a/test/test_infection_functions.py +++ b/test/test_infection_functions.py @@ -268,6 +268,57 @@ def test_mobility_identity_with_depletion_matches_eq1( 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, diff --git a/test/test_mobility_subpopulation_infections.py b/test/test_mobility_subpopulation_infections.py index d22b24ee2..b3a2e05b1 100644 --- a/test/test_mobility_subpopulation_infections.py +++ b/test/test_mobility_subpopulation_infections.py @@ -195,7 +195,6 @@ class TestMobilityValidation: ("nan", "finite"), ("row_stochastic", "transpose"), ("columns_not_one", "sum to 1"), - ("empty_row", "no residents"), ], ) def test_rejects_invalid_matrix( @@ -215,7 +214,6 @@ def test_rejects_invalid_matrix( "nan": roy_mobility_matrix.at[0, 0].set(jnp.nan), "row_stochastic": roy_mobility_matrix.T, "columns_not_one": 0.5 * jnp.eye(3), - "empty_row": jnp.array([[1.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), } with pytest.raises(ValueError, match=match): @@ -288,19 +286,24 @@ def test_rejects_size_mismatch( assert "(3, 3)" in message assert "2 entries" in message - def test_rejects_location_without_population( + 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 population present is rejected.""" + """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( - jnp.eye(3), DeterministicVariable("log_rt_time_0", 0.0) + 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"] - with pytest.raises(ValueError, match=r"Locations \[2\]"): - process._parse_and_validate_fractions(jnp.array([0.5, 0.5, 0.0])) + assert jnp.all(jnp.isfinite(latent.all_subpops)) + assert_array_equal(rt_subpop_adjusted[n_init:, 2], 0.0) class TestMobilityModel: @@ -323,28 +326,18 @@ def test_builder_prior_predictive( assert jnp.all(samples["latent_infections_by_subpop"] > 0) assert samples["subpopulation::susceptible_fraction"].shape == (5, n_total, 3) - @pytest.mark.parametrize( - ("case", "fractions", "match"), - [ - ("roy", jnp.array([0.5, 0.5]), r"\(3, 3\)"), - ("identity", jnp.array([0.5, 0.5, 0.0]), r"Locations \[2\]"), - ], - ) def test_validate_data_rejects_before_mcmc( self, - case: str, - fractions: jnp.ndarray, - match: str, mobility_model_factory: Callable[[jnp.ndarray], MultiSignalModel], roy_mobility_matrix: jnp.ndarray, ) -> None: - """Size mismatches and empty locations are rejected before MCMC starts.""" - matrices = {"roy": roy_mobility_matrix, "identity": jnp.eye(3)} - model = mobility_model_factory(matrices[case]) + """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=match): + with pytest.raises(ValueError, match=r"\(3, 3\)"): model.validate_data(n_days_post_init=30, subpop_fractions=fractions) - with pytest.raises(ValueError, match=match): + with pytest.raises(ValueError, match=r"\(3, 3\)"): model.run( num_warmup=10, num_samples=10, From c69747b7de3f341ba096397c9e8f33cfda4b6987 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 27 Oct 2026 22:49:46 -0400 Subject: [PATCH 21/21] changes per code review; cap cumulative, remove floor on susceptible --- docs/tutorials/latent_mobility_infections.qmd | 58 +++++++------ pyrenew/latent/infection_functions.py | 81 ++++++++++++------ .../mobility_subpopulation_infections.py | 47 +++++++---- test/test_infection_functions.py | 82 +++++++++++++++++-- .../test_mobility_subpopulation_infections.py | 35 +++++++- 5 files changed, 231 insertions(+), 72 deletions(-) diff --git a/docs/tutorials/latent_mobility_infections.qmd b/docs/tutorials/latent_mobility_infections.qmd index 2c6f5b795..631aac33f 100644 --- a/docs/tutorials/latent_mobility_infections.qmd +++ b/docs/tutorials/latent_mobility_infections.qmd @@ -61,7 +61,7 @@ $$ where: -- $I_{t,i}$ is the infection count (a non-negative integer) +- $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$ @@ -69,23 +69,26 @@ where: - $g$ is the generation interval PyRenew works with infections per resident. -With `M` the mobility matrix and `p` the population fractions, one day of the computation is: +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 -share_present = M * p[None, :] / (M @ p)[:, None] +present = M @ p +share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None] infectiousness = reversed_gen_int @ history -susceptible = jnp.maximum(1 - cumulative, 0) -new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) +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. -- `susceptible` is floored at zero, so that a day of extreme infections cannot exceed the remaining susceptibles. +- `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` multiplies last, because residents of a region lose susceptibility wherever they were 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 @@ -282,7 +285,7 @@ show the same effect in their S2 Fig. `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. -`rt_subpop_adjusted` is the reproduction number at each location after depletion: new infections caused per unit of infectiousness present there. +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} @@ -315,10 +318,6 @@ rt_df = pd.concat( Susceptible depletion is on by default, as in Eq 2. -`susceptible_fraction[t, i]` is the share of residents of region `i` not infected before day `t`. -It is recorded with depletion on or off. -With `susceptible_depletion=False`, the renewal equation does not consider it, and the epidemic rapidly reaches 100% of the population infected. - ```{python} #| label: depletion @@ -335,19 +334,19 @@ susceptible_df = pd.concat( ( p9.ggplot(susceptible_df, p9.aes(x="day", y="value", color="run")) + p9.geom_line(size=1) - + p9.facet_wrap("~ region", ncol=3) + + p9.facet_grid("run ~ region", scales="free_y") + p9.labs( x="Days", - y="Susceptible fraction", - title="Share of Residents Not Yet Infected", - color="", + y="1 - cumulative infections per resident", + title="Susceptibility With and Without Depletion", ) + + p9.guides(color="none") + theme_tutorial - + p9.theme(legend_position="bottom") + + p9.theme(figure_size=(10, 7)) ) ``` -With depletion off, the susceptible fraction reaches zero, which means that the model infects more residents than there are. +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 @@ -370,7 +369,8 @@ 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 -share_present = M * p[None, :] / (M @ p)[:, None] +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]): @@ -385,14 +385,16 @@ for t in range(n_init, per_resident.shape[0]): * infectiousness ) incoming = susceptible * np.diag(M) * at_location - local - outgoing = susceptible * (M.T @ at_location) - local - incoming - sources.append(np.stack([local, incoming, outgoing])) + 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 relative difference between the summed sources and the infections: " - f"{np.max(np.abs(total - per_resident[n_init:]) / per_resident[n_init:]):.1e}" + "Largest absolute difference between the summed sources and the infections: " + f"{np.max(np.abs(total - per_resident[n_init:])):.1e}" ) ``` @@ -402,7 +404,12 @@ The three sources add up to each day's new infections, to floating-point precisi #| label: plot-sources source_names = ["Local", "Incoming", "Outgoing"] -share = sources / total[:, np.newaxis, :] +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( { @@ -439,7 +446,8 @@ More than half of region 2's infections come from mobility: its residents are in 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`, and that every location has people present. +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: diff --git a/pyrenew/latent/infection_functions.py b/pyrenew/latent/infection_functions.py index c9649efd1..471fc2fd7 100755 --- a/pyrenew/latent/infection_functions.py +++ b/pyrenew/latent/infection_functions.py @@ -7,6 +7,7 @@ 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 @@ -214,25 +215,26 @@ def _mobility_renewal_step( reversed_generation_interval_pmf Generation interval PMF in reverse order. susceptible_depletion - Whether new infections are scaled by the susceptible fraction. + 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 - adjusted reproduction number at each location. + 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 = jnp.maximum(1.0 - cumulative, 0.0) + susceptible = 1.0 - cumulative + new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible) rt_adjusted = rt_raw * (share_present @ susceptible) else: - susceptible = jnp.ones_like(cumulative) + new_infections = new_infections_raw rt_adjusted = rt_raw - new_infections = susceptible * ( - mobility_matrix.T @ (rt_raw * (share_present @ infectiousness)) - ) history = jnp.concatenate([history[1:], new_infections[jnp.newaxis]], axis=0) return (history, cumulative + new_infections), (new_infections, rt_adjusted) @@ -260,10 +262,12 @@ def compute_infections_from_rt_with_mobility( 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. + 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 - Reproduction number at each location on each day after the - initialization period, shape ``(n_days, n)``. + 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. @@ -276,20 +280,31 @@ def compute_infections_from_rt_with_mobility( Sums 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. + 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 depletion, the same - shape. Without depletion, the second array is ``Rt_raw``. + 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 is Equation 2 of Roy, Clapham, and Mishra (2025), + This implements Equation 2 of Roy, Clapham, and Mishra (2025), . In infection counts (before - converting to infections per resident), the equation is: + 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} @@ -301,26 +316,42 @@ def compute_infections_from_rt_with_mobility( 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 is: + With ``M = mobility_matrix`` and ``p = subpop_fractions``, one day with + depletion enabled is: ```python - share_present = M * p[None, :] / (M @ p)[:, None] + present = M @ p + share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None] infectiousness = reversed_generation_interval_pmf @ history - susceptible = jnp.maximum(1 - cumulative, 0) - new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) + 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, and its - reproduction number after depletion is 0. + A location with no one present contributes no infections. With + depletion enabled, its adjusted reproduction number is 0. - The susceptible fraction is floored at zero. On a day with extreme - reproduction numbers, new infections can exceed the residents still - susceptible. That day is not capped, but later days stay non-negative. + 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 @@ -328,7 +359,7 @@ def compute_infections_from_rt_with_mobility( / jnp.where(present > 0, present, 1.0)[:, jnp.newaxis] ) n_gen = reversed_generation_interval_pmf.size - initial_state = (I0[-n_gen:], jnp.sum(I0, axis=0)) + initial_state = (I0[-n_gen:], cumulative) step = functools.partial( _mobility_renewal_step, mobility_matrix=mobility_matrix, diff --git a/pyrenew/latent/mobility_subpopulation_infections.py b/pyrenew/latent/mobility_subpopulation_infections.py index c5672c1c8..474ccf24a 100644 --- a/pyrenew/latent/mobility_subpopulation_infections.py +++ b/pyrenew/latent/mobility_subpopulation_infections.py @@ -26,25 +26,33 @@ class MobilitySubpopulationInfections(BaseSubpopulationInfections): 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 is Equation 2 of Roy, Clapham, and Mishra (2025), + 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 is: + ``p = subpop_fractions``, one day with depletion enabled is: ```python - share_present = M * p[None, :] / (M @ p)[:, None] + present = M @ p + share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None] infectiousness = reversed_gen_int @ history - susceptible = 1 - cumulative - new_infections = susceptible * (M.T @ (rt * (share_present @ infectiousness))) + 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. + 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 @@ -60,14 +68,17 @@ class MobilitySubpopulationInfections(BaseSubpopulationInfections): where transmission happens, when everyone is susceptible. - ``rt_subpop_adjusted`` is ``rt_subpop * (share_present @ susceptible_fraction)`` after the - initialization period: new infections caused per unit of - infectiousness present at each location, after depletion. It is 0 at - a location with no one present. + 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``. + 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, - ``susceptible`` is 1 and inferred raw values of ``rt_subpop`` will need to directly absorb + 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 @@ -107,7 +118,9 @@ def __init__( 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. + 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 @@ -129,8 +142,9 @@ def __init__( 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. If False, that factor is left out and - ``rt_subpop`` absorbs depletion. + 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 ------ @@ -255,7 +269,8 @@ def _infections_from_rt( ------- tuple[ArrayLike, ArrayLike] Post-initialization infections per resident and the reproduction - number at each location after depletion. + number at each location after susceptibility adjustment, + before the incidence ceiling. """ post_init_infections, rt_adjusted = compute_infections_from_rt_with_mobility( I0=I0, @@ -272,7 +287,7 @@ def _infections_from_rt( jnp.cumsum(infections_all, axis=0)[:-1], ] ) - susceptible_fraction = jnp.maximum(1.0 - infected_before, 0.0) + 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/test/test_infection_functions.py b/test/test_infection_functions.py index eda5fb092..46a4e9ab8 100644 --- a/test/test_infection_functions.py +++ b/test/test_infection_functions.py @@ -403,13 +403,13 @@ def test_mobility_counts_all_initialization_days( ) -def test_mobility_extreme_rt_nonnegative( +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, infections overshoot the susceptibles but stay non-negative and finite.""" + """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), @@ -417,11 +417,83 @@ def test_mobility_extreme_rt_nonnegative( mobility_matrix=roy_mobility_matrix, subpop_fractions=roy_subpop_fractions, ) - cumulative = mobility_initial_infections.sum(axis=0) + infections.sum(axis=0) + 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.any(cumulative > 1.0) + 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( @@ -430,7 +502,7 @@ def test_mobility_gradient_finite( roy_subpop_fractions: jnp.ndarray, mobility_initial_infections: jnp.ndarray, ) -> None: - """Gradients with respect to Rt stay finite when the susceptible floor is reached.""" + """Gradients with respect to Rt stay finite when the incidence ceiling binds.""" def total_infections(rt: jnp.ndarray) -> jnp.ndarray: """ diff --git a/test/test_mobility_subpopulation_infections.py b/test/test_mobility_subpopulation_infections.py index b3a2e05b1..b93bda4e2 100644 --- a/test/test_mobility_subpopulation_infections.py +++ b/test/test_mobility_subpopulation_infections.py @@ -127,12 +127,17 @@ def test_records_susceptible_fraction( DeterministicVariable("log_rt_time_0", 0.0), susceptible_depletion=susceptible_depletion, ) - trace, _ = _sample_with_trace(process, roy_subpop_fractions) + 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, @@ -161,6 +166,24 @@ def test_adjusted_rt_uses_susceptible_fraction( ) 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[ @@ -186,6 +209,16 @@ def test_adjusted_rt_without_depletion( 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"), [