Add latent infection class "MobilitySubpopulationInfections" for contact-mixing models

Start in pyrenew/latent by reading SubpopulationInfections and BaseLatentInfectionProcess to learn the latent-process contract, then trace PyrenewBuilder.configure_latent and MultiSignalModel.sample() to see where runtime latent_data and population values must flow. Implement the mixing equations as written, add the runtime validation hook, and wire counts conversion in MultiSignalModel. Done when the required checks pass: identity mobility reproduces SubpopulationInfections, non-column-stochastic matrices are rejected before sampling, and fitting/prediction accept different mobility matrices.

Estimate time-varying $\mathcal{R}(t)$ for each subpopulation when residents of one subpopulation spend part of each day in others. 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 new class implements Equation 2 of Roy, Clapham, and Mishra (2025),
Incorporating human mobility to enhance epidemic response and estimate real-time reproduction numbers.
This issue states the computation the class performs, the inputs it accepts, and what it does not do. The tests listed at the end check that contract.

Prerequisites

These merge in this order:

• PR #897: time-varying ascertainment. It changes no code in pyrenew/latent/.
• #903, fixed by PR #904: subpopulation count observations are scaled by the subpopulation's population (see Infection counts).
• #902, fixed by PR #906: removes infection_process. Infection feedback becomes a constructor option of PopulationInfections and SubpopulationInfections.
• #905, fixed by PR #907: log_rt_time_0_rv can give each subpopulation its own starting $\mathcal{R}(t)$.

Background

SubpopulationInfections gives each subpopulation its own renewal equation.
Subpopulation $\mathcal{R}(t)$ trajectories share a baseline, but infections in one subpopulation do not depend on infections in any other.

MobilitySubpopulationInfections keeps the same $\mathcal{R}(t)$ structure and couples the renewal equations through a fixed mobility matrix,
such as one derived from census commuting data.

The model: Roy et al. Equation 2

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=0}^{t-1} I_{\tau,k}, g_{t-\tau} \right]
$$

with $S_{t,i} = N_i - \sum_{s<t} I_{s,i}$, the susceptible residents of region $i$ (Eq 1 applied to each region). This assumes infection confers lasting immunity (Roy et al. section 2.2).

The paper writes the Eq 1 sum from $s = 1$ while its renewal sum starts at $\tau = 0$. The authors' code subtracts every day's infections, including the first, and this class does the same.

• $m_{jk}$: fraction of residents of home region $k$ present at location $j$
• $N_k$: number of residents of region $k$
• $R_{t,j}$: reproduction number at location $j$, where transmission happens
• $g$: generation interval

The paper writes $M$ for the number of regions and $m_{ji}(t)$ for a matrix that may change over time. Here M is the matrix, n is the number of regions, and in this version the matrix is fixed.

Mobility matrix convention

M[j, k] is the fraction of residents of home region k who are present at location j. Columns are home regions (origins), rows are locations (destinations), and each column sums to 1. This is the paper's $m_{jk}$.

A matrix with home regions in its rows (rows sum to 1) must be transposed before use.

Regions must be in the same order in M, subpop_fractions, a vector-valued I0_rv, a vector-valued log_rt_time_0_rv, and every observation's subpop_indices.

The computation, one day at a time

PyRenew's latent infections are per resident: infections[t, k] is the number of new infections among residents of home region k on day t, divided by the number of residents of k. Dividing Eq 2 by $N_i$ and writing $N_k$ = population_size * p[k], where p is subpop_fractions, the total population cancels.

The share of the people present at each location who live in each home region depends only on M and p, so it is computed once, before the daily loop:

share_present = M * p[None, :] / (M @ p)[:, None]

One day of the renewal is:

infectiousness = reversed_gen_int @ history
infectiousness_at_location = share_present @ infectiousness
susceptible = 1 - cumulative
new_infections = susceptible * (M.T @ (rt * infectiousness_at_location))

Code	Roy et al. Eq 2
history	recent $I_{\tau,k} / N_k$, one row per day of the generation interval
infectiousness[k]	$\sum_{\tau} I_{\tau,k}, g_{t-\tau} / N_k$
share_present[j, k]	$m_{jk} N_k / \sum_l m_{jl} N_l$: share of the people present at $j$ who live in $k$
infectiousness_at_location[j]	infectiousness per person present at $j$: the bracketed sum divided by $\sum_l m_{jl} N_l$
rt[j]	$R_{t,j}$
M.T @ (...)	the sum over $j$ weighted by $m_{ji}$
cumulative[i]	$\sum_{s<t} I_{s,i} / N_i$
susceptible[i]	$S_{t,i} / N_i$
new_infections[i]	$I_{t,i} / N_i$

The order of operations follows Eq 2:

• rt multiplies at the location, before M.T, 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 after M.T, because $S_{t,i}$ does not depend on $j$: residents of region $i$ lose susceptibility wherever they were infected.
• The population present at a location counts all residents present, susceptible or not, as in the denominator of Eq 2.
• cumulative counts every latent infection since the first day of the initialization period, as in $\sum_{s<t}$.

The aggregate is the population-weighted sum, aggregate = p @ new_infections, which is total new infections divided by the total population.

When M is the identity matrix, each region follows Eq 1 on its own. With susceptible depletion turned off, the result is identical to SubpopulationInfections.

Reproduction numbers

Subpopulation $\mathcal{R}(t)$ has the same structure as in SubpopulationInfections:

log_rt_subpop = log_rt_baseline + deviations

where deviations sum to zero across subpopulations on each day. Under mobility, index j of rt_subpop is a location: rt_subpop[t, j] is $R_{t,j}$.

The recorded values (under the latent process name prefix) are those of SubpopulationInfections, plus susceptible_fraction. Under mobility they mean:

Name	Meaning
rt_baseline	geometric mean of rt_subpop across locations
rt_subpop	$R_{t,j}$: reproduction number at location j when everyone is susceptible
rt_subpop_adjusted	rt_subpop * (share_present @ susceptible_fraction) after the initialization period, and rt_subpop during it: reproduction number at location j after depletion, which is new infections caused per unit of infectiousness present at j; equal to rt_subpop when depletion is off
infections_aggregate	aggregate new infections per resident of the jurisdiction
susceptible_fraction	susceptible_fraction[t, i] is $S_{t,i} / N_i$, the share of residents of home region i not infected before day t, on the full time axis; recorded with depletion on or off, and used by the renewal only with depletion on

With susceptible depletion on, rt_subpop is the paper's $R_{t,j}$. With depletion off, rt_subpop absorbs depletion, as in SubpopulationInfections.

Public API

from pyrenew.deterministic import DeterministicVariable
from pyrenew.latent import MobilitySubpopulationInfections

builder.configure_latent(
    MobilitySubpopulationInfections,
    gen_int_rv=gen_int_rv,
    I0_rv=I0_rv,
    log_rt_time_0_rv=log_rt_time_0_rv,
    baseline_rt_process=baseline_rt_process,
    subpop_rt_deviation_process=subpop_rt_deviation_process,
    mobility_matrix_rv=DeterministicVariable("mobility_matrix", mobility_matrix),
    susceptible_depletion=True,
)
model = builder.build()

model.run(
    ...,
    population_size=population_size,
    subpop_fractions=subpop_fractions,
    **signal_data,
)

• mobility_matrix_rv is required. It is the matrix used on every day. In this version it must be a DeterministicVariable whose value is an (n, n) numeric array. A data frame must be converted with .to_numpy(), with rows and columns in subpopulation order.
• susceptible_depletion defaults to True, which is Eq 2 as published. False drops the $S_{t,i}$ term.
• I0_rv returns new infections per resident on the first day of the initialization period, the earliest day on the time axis, either one value for all regions or one value per region. It is not prevalence. New infections grow from that value over the initialization period, and every initialization day counts toward depletion.
• log_rt_time_0_rv returns the starting log $\mathcal{R}(t)$, either one value for all regions or one value per region (#905). Roy et al. give each region its own starting value.
• model.run() takes the same arguments as for SubpopulationInfections. The matrix is part of the model definition, not run-time data.
• Fitting and forecasting use the same matrix. Prediction with NumPyro's Predictive, including model.posterior_predictive(), runs the fitted model, so it uses the matrix the model was built with. No run-time argument can replace it.
• To model another jurisdiction, build a model with that jurisdiction's matrix.
• pyrenew.latent also exports BaseSubpopulationInfections (see Class structure) and compute_infections_from_rt_with_mobility, the renewal function the class uses, next to compute_infections_from_rt_with_feedback.

Validation

When the latent process is constructed, the matrix must be:

• two-dimensional and square;
• finite and non-negative;
• column-stochastic: each column sums to 1, within 1e-6. If the columns do not but the rows do, the error message says to pass the transpose;
• such that every location has at least one positive entry in its row.

A three-dimensional array is rejected: mobility_matrix_rv must be a single (n, n) matrix. A mobility_matrix_rv that is not a DeterministicVariable is rejected.

Before fitting (MultiSignalModel.validate_data, which run() calls):

• the matrix has one row and one column per entry of subpop_fractions;
• every location has a positive population present: M @ subpop_fractions > 0.

The size check also runs whenever the model is sampled, including prediction with Predictive. The population check does not: under Predictive and during MCMC the model is traced, and checks on values are skipped.

When the latent process is sampled:

• a vector-valued initial infection level has one value per subpopulation;
• a vector-valued starting log $\mathcal{R}(t)$ has one value per subpopulation (#905).

Infection counts

This class depends on the fix for subpopulation count scaling (#903, PR #904). With that fix:

• count observations of subpopulation k receive infections[:, k] * population_size * subpop_fractions[k];
• measurement observations, such as wastewater concentrations, receive infections per resident;
• aggregate count observations receive aggregate * population_size, as now.

What this version does not do

• Time-varying mobility. One fixed matrix. A later version is planned to accept a list of dated matrices, each replacing the fixed matrix on its dates: for example, school holidays, or the identity during lockdowns as in Roy et al. The matrices are data and are never estimated. That version would run one fixed-matrix renewal per date range, carrying the state from one range to the next.
• Estimating the matrix. The matrix is data and is never estimated. A random matrix is rejected.
• Scenario projections. This version uses one matrix on every day, so a projection that switches to another matrix from the forecast start (for example a lockdown) is not supported. That needs a matrix that changes on a date, as planned for dated matrices. A counterfactual with another matrix on every day is possible: build a second model with that matrix and pass the fitted model's posterior samples to NumPyro's Predictive, as Predictive(model2.model, posterior_samples=model1.mcmc.get_samples()). model2.posterior_predictive() does not work, because model2 has not been fitted.
• Recording the matrix. The matrix is constant, so it is not recorded with every draw. It is available as model.latent.mobility_matrix_rv.value.
• Mobility that responds to the epidemic, for example people travelling less when local transmission is high. That would need its own renewal step, computing each day's matrix from recent infections, and data on mobility: with count data alone, a mobility response trades off against $\mathcal{R}(t)$.
• Infection feedback. This class has no feedback option. In PopulationInfections and SubpopulationInfections (#902), infection feedback is an approximate way to account for susceptible depletion. Here Eq 2 accounts for depletion directly, through susceptible_depletion.
• Prior immunity. All residents are susceptible at the start of the initialization period, as in Roy et al.
• Coupled initialization. Initialization-period infections grow at a rate computed from each region's own initial $\mathcal{R}(t)$, as in SubpopulationInfections. Under coupling this can produce a short adjustment over the first generation interval after initialization.
• Infection sources. The split of infections into local, incoming, and outgoing sources (Roy et al. Eqs 3 to 5) is not recorded. It can be computed after fitting from latent_infections_by_subpop, rt_subpop, and susceptible_fraction, with the matrix, population_size, subpop_fractions, and the generation interval. The generation interval must be fixed, because a sampled one is not recorded. With depletion off, use 1 in place of susceptible_fraction.
• Matrix orientation conversion. A row-stochastic matrix is rejected, not transposed.
• Other kinds of contact mixing, such as age-structured contact matrices. Those are different models.

Deliberate differences from Roy et al.

• $\mathcal{R}(t)$ structure. Roy et al. give each region an independent random walk on $\log R_{t,i}$ (simulation study) or a covariate model (England case study). This class uses PyRenew's shared baseline plus sum-to-zero deviations, with any TemporalProcess. As in Roy et al., each region can start at its own $\mathcal{R}(t)$ by giving log_rt_time_0_rv one value per region (#905).
• Initialization. Roy et al. seed a constant number of infections for six days. This class starts from I0 on the first day of the initialization period and grows infections exponentially over that period, as SubpopulationInfections does.
• Overshoot. On a day with extreme $\mathcal{R}(t)$, Eq 2 as written can infect more residents than remain susceptible. This class does not cap that day's infections, so cumulative infections can exceed the number of residents. It floors the susceptible fraction at zero, so infections on later days stay non-negative. At realistic incidence the floor has no effect.
• Observation model. Delays, ascertainment, and noise are PyRenew's observation processes. They are not part of Eq 2.

Class structure

BaseSubpopulationInfections holds what both subpopulation models share: the hierarchical $\mathcal{R}(t)$, initial infections, aggregation, and recorded values. Each subclass defines its renewal step, and its constructor lists its options (#902):

• SubpopulationInfections: each subpopulation renews independently, with optional infection feedback. Its constructor, behavior, and recorded values do not change.
• MobilitySubpopulationInfections: the coupled renewal above, with optional susceptible depletion.

Susceptible depletion is part of the coupled renewal step in Eq 2, so it is a constructor option of MobilitySubpopulationInfections. PopulationInfections and SubpopulationInfections have no depletion option. Independent regions with depletion (Eq 1 for each region) are MobilitySubpopulationInfections with the identity matrix and depletion on. This is the model Roy et al. compare against for regions not linked by mobility.

Required checks

• Eq 2 reference. Infections match a direct transcription of the time loop in the authors' Stan model (uk_spatial/simulated_data/simulated_fitting_region.stan in mlgh-sg/Spatial_transmission). The comparison uses the Fig 2b mobility matrix, unequal regional $\mathcal{R}(t)$, depletion on, and enough days for substantial depletion, and agrees to floating-point precision.
• Identity, no depletion. With the identity matrix and depletion off, every value SubpopulationInfections records agrees to floating-point precision (allclose) for the same draws, with one starting log $\mathcal{R}(t)$ or one per region.
• Identity, with depletion. With the identity matrix and depletion on, each region follows Eq 1.
• Adjusted Rt. With depletion off, rt_subpop_adjusted equals rt_subpop.
• Conservation. Total new infections equal the sum over locations of rt_subpop_adjusted times the infectiousness present there.
• Aggregate. aggregate equals all_subpops @ p.
• Running total. Every initialization day counts toward depletion, including days older than the generation interval.
• Susceptible fraction. susceptible_fraction starts at 1 and never increases. With depletion on, rt_subpop_adjusted equals rt_subpop * (share_present @ susceptible_fraction) after the initialization period.
• Extreme Rt. Infections stay non-negative and finite, and gradients stay finite.
• Construction checks. Matrices that are non-square, negative, non-finite, row-stochastic but not column-stochastic, three-dimensional, or have an empty row are rejected. So is a non-deterministic mobility_matrix_rv.
• SubpopulationInfections unchanged. Its tests, as #902 and #905 leave them, pass after the base class is extracted, including those for infection feedback and per-region starting $\mathcal{R}(t)$.
• Pre-fit checks. A matrix whose size does not match subpop_fractions, or a location with no population present, is rejected by validate_data before MCMC starts.
• End to end. Prior predictive sampling through PyrenewBuilder and MultiSignalModel works.
• Counterfactual with another matrix. Samples drawn from a model, passed to Predictive for the same model, reproduce its latent infections. Passed to Predictive for a model built with the identity matrix, they give different latent infections.