[cdc-mitzimorris]

• New latent process MobilitySubpopulationInfections: subpopulations are home regions whose residents spend part of each day at other locations, coupled by a fixed mobility matrix.
• Implements Eq 2 of Roy, Clapham, and Mishra (2025), https://doi.org/10.1371/journal.pcbi.1013642.
• Builds on prerequisites already on this branch:
  • Refine AscertainmentModel component to allow for time-varying component #897
  • Mem bugfix 903 subpop scaling #904
  • Make infection feedback a constructor option of the latent classes; remove infection_process #902
  • Let log_rt_time_0_rv give each subpopulation its own starting Rt #905

Code

• In file pyrenew/latent/infection_functions.py adds a new function compute_infections_from_rt_with_mobility which does a jax.lax.scan over a private _mobility_renewal_step.

  • Returns infections per resident and the Rt at each location after depletion.
  • Does no validation; the class validates.

  With M = mobility_matrix and p = subpop_fractions, one day is:

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)))

• pyrenew/latent/subpopulation_infections.py: abstract BaseSubpopulationInfections extracted from SubpopulationInfections.
  • The base holds the constructor, I0 and starting log Rt handling, validate, and sample.
  • Subclasses implement _infections_from_rt.
  • SubpopulationInfections keeps the infection feedback options. Its behavior and error order are unchanged.
  • I0_rv docstrings now say "new infections per resident on the first day of the initialization period", not "prevalence".
• pyrenew/latent/mobility_subpopulation_infections.py (new):
  • mobility_matrix_rv must be a DeterministicVariable. mobility_matrix[j, k] is the fraction of residents of region k present at location j, so each column sums to 1.
  • susceptible_depletion: bool = True, as in Eq 2. With False, rt_subpop absorbs depletion.
  • No infection feedback option, because Eq 2 handles depletion directly. Passing a feedback argument raises TypeError.
  • rt_subpop[t, j] is the Rt at location j. After initialization, rt_subpop_adjusted = rt_subpop * (share_present @ susceptible_fraction).
  • Records susceptible_fraction, shape (n_initialization_points + n_days_post_init, n_subpops), starting at 1 and non-increasing, whether depletion is on or off.
  • Every initialization day counts toward depletion.
  • The matrix is fixed: fitting and prediction use the matrix the model was built with.
• Validation:
  • At construction: the matrix is square and 2D, finite, non-negative, has columns summing to 1, and has no all-zero rows. If the rows sum to 1 instead, the error message says to pass the transpose.
  • Override of _parse_and_validate_fractions: the matrix shape matches subpop_fractions (always checked), and every location has people present, M @ p > 0 (skipped under tracing). run() performs both checks through validate_data before MCMC starts.
• pyrenew/latent/__init__.py exports BaseSubpopulationInfections, MobilitySubpopulationInfections, and compute_infections_from_rt_with_mobility.

Tests

• test/test_infection_functions.py:
  • Matches an independent NumPy implementation of Roy et al.'s model, written from their published code, to a relative tolerance of 1e-10, with over half of each region infected.
  • With the identity matrix, matches uncoupled renewal (depletion off) and per-region Eq 1 (depletion on).
  • Conservation of infections across locations.
  • Every initialization day counts toward depletion.
  • At extreme Rt, infections stay non-negative and gradients stay finite.
• test/test_mobility_subpopulation_infections.py (new):
  • With the identity matrix and no depletion, every trace site matches SubpopulationInfections, for both a scalar and a per-region starting Rt.
  • The aggregate is the population-weighted sum.
  • susceptible_fraction and rt_subpop_adjusted are correct.
  • Matrix, type, size, and population checks raise the right errors.
  • Builder prior predictive runs.
  • validate_data and run reject bad inputs before MCMC starts.
  • Counterfactual: replaying draws under the same matrix reproduces them, and replaying them under another matrix changes them.
• test/test_interface_coverage.py: three interface tests for the new class.
• test/test_subpopulation_infections.py: not edited by this change; it is the regression check for the base class extraction.
• New fixtures in test/conftest.py, including the Roy et al. Fig 2b matrix and populations.

Docs

• New tutorial docs/tutorials/latent_mobility_infections.qmd, forward simulation only:
  • matrix orientation, and how to subset and renormalize a matrix
  • a counterfactual with and without mobility (a region with Rt below 1 grows only with mobility)
  • where Rt applies
  • susceptible depletion
  • sources of infection (Roy et al. Eqs 3 to 5)
  • connecting to observation processes
• Link from latent_subpopulation_infections.qmd, and navigation entries in docs/index.md and docs/tutorials/.pages.

[dylanhmorris]

Also tagging @KOVALW, @paigemiller, and @confunguido for awareness (no need to review/comment unless you feel inspired)

[cdc-mitzimorris]

@dylanhmorris - ready for re-review.

• accepted all tutorial and docstring changes
• dropped empty-row restriction instead added more safe checks for divide-by-zero
• added mobility-matrix plot to tutorial
• leaving computation susceptible_fraction = jnp.maximum(1.0 - infected_before, 0.0) as is - this too makes biological sense and the code is cleaner (imo).

[dylanhmorris]

\mathcal{R}?

[dylanhmorris]

\mathcal{R}?

[dylanhmorris]

Same again

[dylanhmorris]

LGTM, thanks @cdc-mitzimorris!