From 277a26714fbd3628b7254ecefa6a3a8bf3c1085c Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Sun, 11 Oct 2026 18:05:14 -0400 Subject: [PATCH 01/24] checkpointing; bugfix implemented by claude app --- .../tutorials/building_multisignal_models.qmd | 18 +- .../observation_processes_measurements.qmd | 15 +- pyrenew/model/multisignal_model.py | 35 +- pyrenew/observation/__init__.py | 5 +- pyrenew/observation/base.py | 18 +- pyrenew/observation/count_observations.py | 2 +- .../observation/measurement_observations.py | 18 +- test/conftest.py | 34 +- .../test_subpopulation_counts_scaling.py | 295 ++++++++++++++++ test/test_helpers.py | 32 ++ test/test_interface_coverage.py | 4 +- test/test_multisignal_infection_scaling.py | 326 ++++++++++++++++++ test/test_observation_validation.py | 39 +-- 13 files changed, 748 insertions(+), 93 deletions(-) create mode 100644 test/integration/test_subpopulation_counts_scaling.py create mode 100644 test/test_multisignal_infection_scaling.py diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 6e64e2cbf..4d955b05c 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -170,10 +170,11 @@ flowchart TB Different observation processes observe different levels of the model hierarchy. Each observation process declares an **infection resolution** that determines what infection data it receives: - | Resolution | Receives | Example signals | - | ------------- | ------------------------------------------------------------------- | -------------------------------------- | - | `"aggregate"` | Aggregated infections (sum across all subpopulations), shape `(T,)` | Hospital admissions, case counts | - | `"subpop"` | Infection matrix for all subpopulations, shape `(T, n_subpops)` | Wastewater, site-specific surveillance | + | Resolution | Receives | Example signals | + | --------------------- | --------------------------------------------------------------------- | -------------------------------- | + | `"aggregate"` | Infection counts for the whole jurisdiction, shape `(T,)` | Hospital admissions, case counts | + | `"subpop"` | Infection counts in each subpopulation, shape `(T, n_subpops)` | Regional case counts | + | `"subpop_per_capita"` | Infections per resident of each subpopulation, shape `(T, n_subpops)` | Wastewater concentrations | The `PyrenewBuilder` routes latent infections to observation processes based on each process's declared resolution. @@ -297,7 +298,7 @@ Observation processes transform latent infections into observable signals and de Each observation process: - Has a unique **name** that identifies the signal in model outputs -- Declares what **infection resolution** it needs (`"aggregate"` or `"subpop"`) +- Declares what **infection resolution** it needs (`"aggregate"`, `"subpop"`, or `"subpop_per_capita"`) - Applies signal-specific transformations (ascertainment, delay convolution, shedding kinetics) - Defines the noise model @@ -425,8 +426,9 @@ class Wastewater(MeasurementObservation): """ Wastewater viral concentration observation process. - Transforms site-level infections into predicted log-concentrations - via shedding kinetics convolution and genome/volume scaling. + Transforms infections per resident of each subpopulation into + predicted log-concentrations via shedding kinetics convolution and + genome/volume scaling. """ def __init__( @@ -524,7 +526,7 @@ shedding_pmf = shedding_pmf / shedding_pmf.sum() # normalize shedding_rv = DeterministicPMF("shedding_kinetics", shedding_pmf) # Log10 genomes shed per infection -log10_genome_rv = DeterministicVariable("log10_genome_per_inf", 9.0) +log10_genome_rv = DeterministicVariable("log10_genome_per_inf", 4.1) # Wastewater volume per person per day (mL) ml_per_person_per_day = 1000.0 diff --git a/docs/tutorials/observation_processes_measurements.qmd b/docs/tutorials/observation_processes_measurements.qmd index 1f6f3716c..490a1d350 100644 --- a/docs/tutorials/observation_processes_measurements.qmd +++ b/docs/tutorials/observation_processes_measurements.qmd @@ -203,7 +203,7 @@ $$ where: -- $I(t-d)$ is the number of infections on day $t-d$ +- $I(t-d)$ is the number of infections per resident on day $t-d$ - $\pi_d$ is the shedding kinetics PMF, giving the fraction of total shedding occurring $d$ days after infection, analogous to the delay distribution in count observation models - $G$ is the number of genome copies shed per infection - $V$ is the wastewater volume per person per day @@ -245,8 +245,9 @@ class Wastewater(MeasurementObservation): """ Wastewater viral concentration observation process. - Transforms site-level infections into predicted log-concentrations - via shedding kinetics convolution and genome/volume scaling. + Transforms infections per resident of each subpopulation into + predicted log-concentrations via shedding kinetics convolution and + genome/volume scaling. """ def __init__( @@ -461,7 +462,7 @@ day_one = first_valid_observation_day(ww_process) # Create infections with a spike (shape: n_days x n_subpops) infection_spike_day = day_one + 10 infections = jnp.zeros((n_days, 1)) # 1 subpopulation -infections = infections.at[infection_spike_day, 0].set(2000.0) +infections = infections.at[infection_spike_day, 0].set(0.002) # For plotting rel_spike_day = infection_spike_day - day_one @@ -506,7 +507,7 @@ plot_infections = ( linetype="dashed", alpha=0.5, ) - + p9.labs(x="Day", y="Daily Infections", title="Infections (Input)") + + p9.labs(x="Day", y="Daily Infections per Resident", title="Infections (Input)") + theme_tutorial + p9.annotate( "text", @@ -714,8 +715,8 @@ The observed concentrations reflect both the underlying infection differences AN # Two subpopulations with different infection patterns n_days_mp = 40 -infections_subpop1 = 1000.0 * jnp.exp(-jnp.arange(n_days_mp) / 20.0) # Slow decay -infections_subpop2 = 2000.0 * jnp.exp(-jnp.arange(n_days_mp) / 10.0) # Fast decay +infections_subpop1 = 0.001 * jnp.exp(-jnp.arange(n_days_mp) / 20.0) # Slow decay +infections_subpop2 = 0.002 * jnp.exp(-jnp.arange(n_days_mp) / 10.0) # Fast decay infections_multi = jnp.stack([infections_subpop1, infections_subpop2], axis=1) # Two sensors, each observing a different subpopulation diff --git a/pyrenew/model/multisignal_model.py b/pyrenew/model/multisignal_model.py index 3e339547e..c5e8ef5dd 100644 --- a/pyrenew/model/multisignal_model.py +++ b/pyrenew/model/multisignal_model.py @@ -51,8 +51,16 @@ class MultiSignalModel(Model): ----- The model automatically routes latent infections to observations based on each observation's infection_resolution() method: - - "aggregate" → receives aggregate infections from latent process - - "subpop" → receives all subpopulation infections; observation selects via indices + + - ``"aggregate"``: infection counts for the whole jurisdiction, + shape ``(n_total,)``. + - ``"subpop"``: infection counts in each subpopulation, + shape ``(n_total, n_subpops)``. + - ``"subpop_per_capita"``: infections per resident of each + subpopulation, shape ``(n_total, n_subpops)``. + + Subpopulation observations select their subpopulations via + ``subpop_indices``. """ def __init__( @@ -86,7 +94,7 @@ def __init__( self.ascertainment_models = ascertainment_models self.validate() - _SUPPORTED_RESOLUTIONS = {"aggregate", "subpop"} + _SUPPORTED_RESOLUTIONS = {"aggregate", "subpop", "subpop_per_capita"} def validate(self) -> None: """ @@ -371,8 +379,9 @@ def sample( n_days_post_init Number of days to simulate after initialization period population_size - Total population size. Used to convert infection proportions - (from latent process) to infection counts (for observation processes). + Total population size. Aggregate infections are multiplied by + ``population_size``; infections in subpopulation ``k`` are + multiplied by ``population_size * subpop_fractions[k]``. subpop_fractions Population fractions for all subpopulations. Shape: (n_subpops,). obs_start_date @@ -400,25 +409,27 @@ def sample( self._check_obs_start_date(obs_start_date) first_day_dow = self._resolve_first_day_dow(obs_start_date) - # Generate latent infections (proportions) latent_sample = self.latent.sample( n_days_post_init=n_days_post_init, subpop_fractions=subpop_fractions, first_day_dow=first_day_dow, ) + pop = self.latent._parse_and_validate_fractions( + subpop_fractions=subpop_fractions, + ) - # Scale from proportions to counts inf_aggregate = latent_sample.aggregate * population_size - inf_all = latent_sample.all_subpops * population_size + inf_subpop = ( + latent_sample.all_subpops * population_size * pop.fractions[jnp.newaxis, :] + ) - # Record scaled infections for posterior analysis numpyro.deterministic("latent_infections", inf_aggregate) - numpyro.deterministic("latent_infections_by_subpop", inf_all) + numpyro.deterministic("latent_infections_by_subpop", inf_subpop) - # Map infection resolution to infection arrays latent_map = { "aggregate": inf_aggregate, - "subpop": inf_all, + "subpop": inf_subpop, + "subpop_per_capita": latent_sample.all_subpops, } n_timepoints = self.latent.n_initialization_points + n_days_post_init diff --git a/pyrenew/observation/__init__.py b/pyrenew/observation/__init__.py index b38ab471c..21b51a814 100644 --- a/pyrenew/observation/__init__.py +++ b/pyrenew/observation/__init__.py @@ -11,7 +11,10 @@ All observation processes implement: - ``sample()``: Sample observations given infections -- ``infection_resolution()``: returns ``"aggregate"`` or ``"subpop"`` +- ``infection_resolution()``: returns ``"aggregate"`` (infection counts for + the whole jurisdiction), ``"subpop"`` (infection counts in each + subpopulation), or ``"subpop_per_capita"`` (infections per resident of + each subpopulation) - ``lookback_days()``: returns required infection history length Noise models (``CountNoise``, ``MeasurementNoise``) are composable—pass them diff --git a/pyrenew/observation/base.py b/pyrenew/observation/base.py index c72097483..0fbb2b3d9 100644 --- a/pyrenew/observation/base.py +++ b/pyrenew/observation/base.py @@ -39,7 +39,8 @@ class BaseObservationProcess(RandomVariable): Subclasses must implement: - ``lookback_days()``: Return PMF length for initialization - - ``infection_resolution()``: Return ``"aggregate"`` or ``"subpop"`` + - ``infection_resolution()``: Return ``"aggregate"``, ``"subpop"``, + or ``"subpop_per_capita"`` - ``_predicted_obs()``: Transform infections to predicted values - ``sample()``: Apply noise model to predicted observations @@ -107,19 +108,22 @@ def lookback_days(self) -> int: @abstractmethod def infection_resolution(self) -> str: """ - Return whether this observation uses aggregate or subpop infections. + Return which infections this observation receives. Returns one of: - - ``"aggregate"``: Uses a single aggregated infection trajectory. - Shape: ``(n_days,)`` - - ``"subpop"``: Uses subpopulation-level infection trajectories. - Shape: ``(n_days, n_subpops)``, indexed via ``subpop_indices``. + - ``"aggregate"``: infection counts for the whole jurisdiction. + Shape: ``(n_total,)`` + - ``"subpop"``: infection counts in each subpopulation. + Shape: ``(n_total, n_subpops)``, indexed via ``subpop_indices``. + - ``"subpop_per_capita"``: infections per resident of each + subpopulation. Shape: ``(n_total, n_subpops)``, indexed via + ``subpop_indices``. Returns ------- str - Either ``"aggregate"`` or ``"subpop"`` + ``"aggregate"``, ``"subpop"``, or ``"subpop_per_capita"`` Notes ----- diff --git a/pyrenew/observation/count_observations.py b/pyrenew/observation/count_observations.py index 57f0881a7..dcd419408 100644 --- a/pyrenew/observation/count_observations.py +++ b/pyrenew/observation/count_observations.py @@ -975,7 +975,7 @@ def sample( Parameters ---------- infections - Subpopulation-level infections from the infection process. + Infection counts in each subpopulation. Shape ``(n_total, n_subpops)``. obs Observed counts. For ``reporting_schedule="regular"``: diff --git a/pyrenew/observation/measurement_observations.py b/pyrenew/observation/measurement_observations.py index 03fe13e2f..aa1921ddf 100644 --- a/pyrenew/observation/measurement_observations.py +++ b/pyrenew/observation/measurement_observations.py @@ -36,7 +36,8 @@ class MeasurementObservation(BaseObservationProcess): ----- Subclasses must implement ``_predicted_obs()`` according to their specific signal processing (e.g., wastewater shedding kinetics, - dilution factors, etc.). + dilution factors, etc.). ``_predicted_obs()`` receives infections + per resident of each subpopulation. See Also -------- @@ -92,17 +93,18 @@ def lookback_days(self) -> int: def infection_resolution(self) -> str: """ - Return "subpop" for measurement observations. + Return "subpop_per_capita" for measurement observations. - Measurement observations require subpopulation-level infections - because each measurement corresponds to a specific catchment area. + Measurement observations receive infections per resident of each + subpopulation, because each measurement corresponds to a specific + catchment area and depends on infections per resident. Returns ------- str - ``"subpop"`` + ``"subpop_per_capita"`` """ - return "subpop" + return "subpop_per_capita" def validate_data( self, @@ -175,8 +177,8 @@ def sample( Parameters ---------- infections - Infections from the infection process. - Shape: (n_total, n_subpops) + Infections per resident of each subpopulation. + Shape: (n_total, n_subpops). times Day index for each observation on the shared time axis. Must be in range [0, n_total). Shape: (n_obs,) diff --git a/test/conftest.py b/test/conftest.py index 2ba07556a..9ad236836 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -4,19 +4,20 @@ This module provides reusable fixtures for creating observation processes, test data, and common configurations used across multiple test files. -The module also sets two JAX environment variables before any jax import so -that all tests (unit and integration) run with 64-bit precision and with -four logical host devices available for parallel MCMC chains. JAX reads -these variables at import time, so they must be set before the first -``import jax`` anywhere in the test process; placing them at the top of -this file (loaded by pytest before any test module) satisfies that -requirement. The ``setdefault`` form respects any value the caller already -set at the shell level (e.g., a CI with different configuration). +All tests (unit and integration) run with 64-bit floats. The module sets +``JAX_ENABLE_X64`` before any jax import, overriding any shell value, and +the autouse fixture ``require_float64`` fails any test that runs without +64-bit floats. It also makes four logical host devices available for +parallel MCMC chains, unless the caller already set ``XLA_FLAGS``. JAX +reads these variables at import time, so they must be set before the +first ``import jax`` anywhere in the test process; placing them at the +top of this file (loaded by pytest before any test module) satisfies that +requirement. """ import os -os.environ.setdefault("JAX_ENABLE_X64", "true") +os.environ["JAX_ENABLE_X64"] = "true" os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") import jax.numpy as jnp @@ -39,6 +40,21 @@ from pyrenew.time import MMWR_WEEK from test.test_helpers import fixed_ar1, fixed_random_walk + +@pytest.fixture(autouse=True) +def require_float64() -> None: + """ + Fail any test that would run without 64-bit floats. + + Raises + ------ + pytest.fail.Exception + If JAX creates 32-bit floating point arrays by default. + """ + if jnp.asarray(0.0).dtype != jnp.float64: + pytest.fail("Tests must run with 64-bit floats; JAX_ENABLE_X64 is off.") + + # ============================================================================= # PMF Fixtures # ============================================================================= diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py new file mode 100644 index 000000000..7dcefac8f --- /dev/null +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -0,0 +1,295 @@ +""" +Integration test: subpopulation count scaling in MultiSignalModel (issue 903). + +Four subpopulations share the same true Rt and the same initial infections +per resident, so they differ only in population size. Counts are simulated +as infections per resident times the subpopulation population. A +SubpopulationInfections model fitted to these counts should recover Rt in +every subpopulation, with no dependence on subpopulation size, and its +recorded subpopulation counts should sum to the aggregate count. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import jax.random as random +import numpy as np +import numpyro.distributions as dist +import pytest + +from pyrenew.ascertainment import IndependentAscertainment +from pyrenew.deterministic import DeterministicPMF, DeterministicVariable +from pyrenew.latent import RandomWalk, SubpopulationInfections +from pyrenew.math import r_approx_from_R +from pyrenew.model import MultiSignalModel, PyrenewBuilder +from pyrenew.observation import PoissonNoise, SubpopulationCounts +from pyrenew.randomvariable import DistributionalVariable + +pytestmark = pytest.mark.integration + +POPULATION_SIZE = 1_000_000.0 +SUBPOP_FRACTIONS = np.array([0.5, 0.25, 0.15, 0.10]) +N_SUBPOPS = len(SUBPOP_FRACTIONS) +GEN_INT = np.array([0.2, 0.35, 0.25, 0.15, 0.05]) +DELAY = np.array([0.1, 0.3, 0.3, 0.2, 0.1]) +I0_PER_RESIDENT = 2e-5 +ASCERTAINMENT = 0.3 +N_DAYS_FIT = 70 +N_DAYS_EARLY = 21 +DATA_SEED = 20261011 +NUM_WARMUP = 500 +NUM_SAMPLES = 500 +NUM_CHAINS = 4 + + +def _true_rt(n_total: int) -> np.ndarray: + """ + Return the true Rt shared by all subpopulations on the padded time axis. + + Parameters + ---------- + n_total + Number of days, including the initialization period. + + Returns + ------- + np.ndarray + Rt for each day, falling from about 1.4 to about 0.9. + """ + t = np.arange(n_total) + return 0.9 + 0.5 / (1.0 + np.exp((t - 35.0) / 4.0)) + + +def _simulate(n_init: int) -> dict[str, np.ndarray]: + """ + Simulate infections per resident and observed counts for each subpopulation. + + Initialization-period infections grow exponentially from + ``I0_PER_RESIDENT`` at the rate implied by the first day's Rt, as in + ``SubpopulationInfections``. Counts in subpopulation ``k`` are + infections per resident times ``POPULATION_SIZE * SUBPOP_FRACTIONS[k]``, + delayed by ``DELAY``, scaled by ``ASCERTAINMENT``, and drawn from a + Poisson distribution. + + Parameters + ---------- + n_init + Number of initialization days. + + Returns + ------- + dict[str, np.ndarray] + ``"rt"``: true Rt on the padded axis, shape ``(n_total,)``. + ``"observed"``: observed counts after initialization, + shape ``(N_DAYS_FIT, N_SUBPOPS)``. + """ + n_total = n_init + N_DAYS_FIT + rt = _true_rt(n_total) + initial_growth_rate = float( + r_approx_from_R(R=jnp.asarray(rt[0]), g=jnp.asarray(GEN_INT), n_newton_steps=4) + ) + infections = np.zeros(n_total) + infections[:n_init] = I0_PER_RESIDENT * np.exp( + initial_growth_rate * np.arange(n_init) + ) + reversed_gen_int = GEN_INT[::-1] + for t in range(n_init, n_total): + infections[t] = rt[t] * reversed_gen_int @ infections[t - GEN_INT.size : t] + counts = POPULATION_SIZE * np.outer(infections, SUBPOP_FRACTIONS) + predicted = np.zeros_like(counts) + for d, weight in enumerate(DELAY): + predicted[d:] += ASCERTAINMENT * weight * counts[: n_total - d] + rng = np.random.default_rng(DATA_SEED) + observed = rng.poisson(predicted[n_init:]) + return {"rt": rt, "observed": observed} + + +@pytest.fixture(scope="module") +def scaling_model() -> MultiSignalModel: + """ + Build a SubpopulationInfections model with subpopulation case counts. + + Initial infections per resident are fixed at the true value, and one + ascertainment rate is shared across subpopulations. + + Returns + ------- + MultiSignalModel + The built model. + """ + ascertainment = IndependentAscertainment( + name="case_ascertainment", + rate_rvs={ + "cases": DistributionalVariable("ascertainment_rate", dist.Beta(3.0, 7.0)) + }, + ) + cases = SubpopulationCounts( + name="cases", + ascertainment_rate_rv=ascertainment.for_signal("cases"), + delay_distribution_rv=DeterministicPMF("delay", jnp.asarray(DELAY)), + noise=PoissonNoise(), + ) + builder = PyrenewBuilder() + builder.configure_latent( + SubpopulationInfections, + name="latent", + gen_int_rv=DeterministicPMF("gen_int", jnp.asarray(GEN_INT)), + I0_rv=DeterministicVariable("I0", jnp.full(N_SUBPOPS, I0_PER_RESIDENT)), + log_rt_time_0_rv=DistributionalVariable("log_rt_time_0", dist.Normal(0.0, 0.5)), + baseline_rt_process=RandomWalk(DeterministicVariable("baseline_sd", 0.05)), + subpop_rt_deviation_process=RandomWalk( + DeterministicVariable("deviation_sd", 0.03) + ), + ) + builder.add_observation(cases) + builder.add_ascertainment(ascertainment) + return builder.build() + + +@pytest.fixture(scope="module") +def simulated(scaling_model: MultiSignalModel) -> dict[str, np.ndarray]: + """ + Simulate data on the model's padded time axis. + + Parameters + ---------- + scaling_model + Built model, used for its number of initialization days. + + Returns + ------- + dict[str, np.ndarray] + True Rt and observed counts. + """ + return _simulate(scaling_model.latent.n_initialization_points) + + +@pytest.fixture(scope="module") +def posterior( + scaling_model: MultiSignalModel, + simulated: dict[str, np.ndarray], +) -> dict[str, np.ndarray]: + """ + Fit the model to the simulated counts and return posterior samples. + + Parameters + ---------- + scaling_model + Built model. + simulated + Simulated data. + + Returns + ------- + dict[str, np.ndarray] + Posterior samples, keyed by site name. + """ + scaling_model.run( + num_warmup=NUM_WARMUP, + num_samples=NUM_SAMPLES, + rng_key=random.PRNGKey(1), + mcmc_args={"num_chains": NUM_CHAINS, "progress_bar": False}, + n_days_post_init=N_DAYS_FIT, + population_size=POPULATION_SIZE, + subpop_fractions=jnp.asarray(SUBPOP_FRACTIONS), + cases={ + "obs": scaling_model.pad_observations( + jnp.asarray(simulated["observed"], dtype=float) + ), + "subpop_indices": jnp.arange(N_SUBPOPS), + }, + ) + samples = scaling_model.mcmc.get_samples() + jax.block_until_ready(samples) + return {name: np.asarray(value) for name, value in samples.items()} + + +class TestSubpopulationCountScaling: + """Fitted subpopulation Rt and infections do not depend on subpopulation size.""" + + def test_fit_uses_float64( + self, + posterior: dict[str, np.ndarray], + ) -> None: + """ + The model is fitted in 64-bit floats. + + Parameters + ---------- + posterior + Posterior samples. + """ + assert posterior["latent::rt_subpop"].dtype == np.float64 + assert posterior["latent_infections_by_subpop"].dtype == np.float64 + + def test_subpop_counts_sum_to_aggregate( + self, + posterior: dict[str, np.ndarray], + ) -> None: + """ + Recorded subpopulation infection counts sum to the aggregate count. + + Parameters + ---------- + posterior + Posterior samples. + """ + np.testing.assert_allclose( + posterior["latent_infections_by_subpop"].sum(axis=-1), + posterior["latent_infections"], + rtol=1e-10, + ) + + def test_rt_recovered_in_every_subpop( + self, + scaling_model: MultiSignalModel, + simulated: dict[str, np.ndarray], + posterior: dict[str, np.ndarray], + ) -> None: + """ + Posterior median Rt is close to the true Rt in every subpopulation. + + Checks the mean of log(fitted Rt / true Rt) over the first three + weeks after initialization, where errors in the infection level + show up most strongly. + + Parameters + ---------- + scaling_model + Fitted model. + simulated + Simulated data. + posterior + Posterior samples. + """ + n_init = scaling_model.latent.n_initialization_points + rt_fit = np.median(posterior["latent::rt_subpop"], axis=0)[n_init:] + log_ratio = np.log(rt_fit / simulated["rt"][n_init:, np.newaxis]) + early_bias = log_ratio[:N_DAYS_EARLY].mean(axis=0) + assert np.all(np.abs(early_bias) < 0.06), early_bias + + def test_no_rt_trend_with_subpop_size( + self, + scaling_model: MultiSignalModel, + posterior: dict[str, np.ndarray], + ) -> None: + """ + Subpopulation Rt deviations show no trend with subpopulation size. + + All subpopulations share the same true Rt, so the slope of each + subpopulation's mean deviation against the log of its population + fraction should be close to zero. + + Parameters + ---------- + scaling_model + Fitted model. + posterior + Posterior samples. + """ + n_init = scaling_model.latent.n_initialization_points + deviations = np.median(posterior["latent::subpop_deviations"], axis=0) + mean_deviation = deviations[n_init:].mean(axis=0) + slope = np.polyfit(np.log(SUBPOP_FRACTIONS), mean_deviation, 1)[0] + assert abs(slope) < 0.01, (slope, mean_deviation) diff --git a/test/test_helpers.py b/test/test_helpers.py index 3a7221365..b8364a248 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -160,6 +160,38 @@ def convolve_col(col: ArrayLike) -> ArrayLike: # numpydoc ignore=GL08 return log_predicted +class StubMeasurementObservation(MeasurementObservation): + """Minimal concrete MeasurementObservation whose predicted values are its infections.""" + + def validate(self) -> None: + """ + Validate parameters. + + Raises + ------ + ValueError + If PMF is invalid. + """ + pmf = self.temporal_pmf_rv() + self._validate_pmf(pmf, "temporal_pmf_rv") + + def _predicted_obs(self, infections): + """ + Return infections unchanged (identity transform). + + Parameters + ---------- + infections : ArrayLike + Input infections. + + Returns + ------- + ArrayLike + Same as input. + """ + return infections + + def create_mock_infections( n_days: int, peak_day: int = 10, diff --git a/test/test_interface_coverage.py b/test/test_interface_coverage.py index fad284d0a..c03ccb956 100644 --- a/test/test_interface_coverage.py +++ b/test/test_interface_coverage.py @@ -193,9 +193,9 @@ def test_counts_by_subpop_infection_resolution(): def test_measurements_infection_resolution(): - """ConcreteMeasurementObservation.infection_resolution() returns 'subpop'.""" + """ConcreteMeasurementObservation.infection_resolution() returns 'subpop_per_capita'.""" m = _make_measurements() - assert m.infection_resolution() == "subpop" + assert m.infection_resolution() == "subpop_per_capita" def test_base_count_observation_infection_resolution_raises(): diff --git a/test/test_multisignal_infection_scaling.py b/test/test_multisignal_infection_scaling.py new file mode 100644 index 000000000..b7f82815c --- /dev/null +++ b/test/test_multisignal_infection_scaling.py @@ -0,0 +1,326 @@ +""" +Unit tests for how MultiSignalModel scales latent infections for each +infection resolution. +""" + +from __future__ import annotations + +import jax.numpy as jnp +import numpy as np +import numpyro +import pytest +from jax.typing import ArrayLike + +from pyrenew.deterministic import DeterministicPMF +from pyrenew.latent import ( + LatentSample, + PopulationInfections, + SubpopulationInfections, +) +from pyrenew.model import MultiSignalModel +from pyrenew.observation import HierarchicalNormalNoise +from pyrenew.observation.base import BaseObservationProcess +from test.test_helpers import StubMeasurementObservation + +POPULATION_SIZE = 1_000_000.0 +SUBPOP_FRACTIONS = jnp.array([0.6, 0.3, 0.1]) +N_DAYS_POST_INIT = 10 + + +class RecordingObservation(BaseObservationProcess): + """Observation process that records the infections it receives.""" + + def __init__(self, name: str, resolution: str) -> None: + """ + Initialize the recording observation. + + Parameters + ---------- + name + Observation name; the recorded site is ``{name}_infections``. + resolution + Infection resolution this observation requests. + """ + super().__init__( + name=name, + temporal_pmf_rv=DeterministicPMF(f"{name}_pmf", jnp.array([1.0])), + ) + self.resolution = resolution + + def lookback_days(self) -> int: + """ + Return the number of days of infection history required. + + Returns + ------- + int + Zero. + """ + return 0 + + def infection_resolution(self) -> str: + """ + Return the infection resolution given at construction. + + Returns + ------- + str + The requested infection resolution. + """ + return self.resolution + + def _predicted_obs(self, infections: ArrayLike) -> ArrayLike: + """ + Return infections unchanged. + + Parameters + ---------- + infections + Infections received from the model. + + Returns + ------- + ArrayLike + Same as input. + """ + return infections + + def validate_data( + self, + n_total: int, + n_subpops: int, + **obs_data: dict[str, object], + ) -> None: + """ + Accept any data. + + Parameters + ---------- + n_total + Total number of time steps. + n_subpops + Number of subpopulations. + **obs_data + Ignored. + """ + + def sample(self, infections: ArrayLike, **kwargs: object) -> ArrayLike: + """ + Record the infections received as a deterministic site. + + Parameters + ---------- + infections + Infections routed to this observation by the model. + **kwargs + Ignored. + + Returns + ------- + ArrayLike + The infections received. + """ + self._deterministic("infections", infections) + return infections + + +def _trace_model( + model: MultiSignalModel, + subpop_fractions: ArrayLike | None, + **observation_data: dict[str, object], +) -> tuple[dict[str, np.ndarray], LatentSample]: + """ + Sample the model once and recover the latent sample it used. + + The latent process is re-run with its sampled sites substituted from + the model trace, so the returned ``LatentSample`` is the one the + observations received. + + Parameters + ---------- + model + Model to sample. + subpop_fractions + Subpopulation fractions passed to ``model.sample``. + **observation_data + Per-observation data passed to ``model.sample``. + + Returns + ------- + tuple[dict[str, np.ndarray], LatentSample] + Site values from the model trace and the latent sample. + """ + with numpyro.handlers.trace() as trace, numpyro.handlers.seed(rng_seed=0): + model.sample( + n_days_post_init=N_DAYS_POST_INIT, + population_size=POPULATION_SIZE, + subpop_fractions=subpop_fractions, + **observation_data, + ) + sites = {name: np.asarray(site["value"]) for name, site in trace.items()} + sampled = { + name: site["value"] for name, site in trace.items() if site["type"] == "sample" + } + with ( + numpyro.handlers.substitute(data=sampled), + numpyro.handlers.seed(rng_seed=1), + ): + latent_sample = model.latent.sample( + n_days_post_init=N_DAYS_POST_INIT, + subpop_fractions=subpop_fractions, + ) + return sites, latent_sample + + +@pytest.fixture +def subpop_trace( + subpopulation_infections: SubpopulationInfections, + hierarchical_normal_noise: HierarchicalNormalNoise, +) -> tuple[dict[str, np.ndarray], LatentSample]: + """ + Trace a SubpopulationInfections model with one observation per resolution. + + Parameters + ---------- + subpopulation_infections + Latent process with three subpopulations supplied at sample time. + hierarchical_normal_noise + Noise model for the measurement observation. + + Returns + ------- + tuple[dict[str, np.ndarray], LatentSample] + Site values from the model trace and the latent sample. + """ + model = MultiSignalModel( + latent_process=subpopulation_infections, + observations={ + "agg": RecordingObservation("agg", "aggregate"), + "counts": RecordingObservation("counts", "subpop"), + "ww": StubMeasurementObservation( + name="ww", + temporal_pmf_rv=DeterministicPMF("ww_pmf", jnp.array([1.0])), + noise=hierarchical_normal_noise, + ), + }, + ) + n_subpops = SUBPOP_FRACTIONS.size + return _trace_model( + model, + SUBPOP_FRACTIONS, + ww={ + "times": jnp.full( + n_subpops, subpopulation_infections.n_initialization_points + ), + "subpop_indices": jnp.arange(n_subpops), + "sensor_indices": jnp.arange(n_subpops), + "n_sensors": n_subpops, + }, + ) + + +def test_subpop_counts_sum_to_aggregate( + subpop_trace: tuple[dict[str, np.ndarray], LatentSample], +) -> None: + """ + Recorded subpopulation infection counts sum to the aggregate count. + + Parameters + ---------- + subpop_trace + Traced SubpopulationInfections model. + """ + sites, _ = subpop_trace + np.testing.assert_allclose( + sites["latent_infections_by_subpop"].sum(axis=1), + sites["latent_infections"], + rtol=1e-12, + ) + + +def test_subpop_counts_scaled_by_subpop_population( + subpop_trace: tuple[dict[str, np.ndarray], LatentSample], +) -> None: + """ + A "subpop" observation receives infections per resident times the + population of each subpopulation. + + Parameters + ---------- + subpop_trace + Traced SubpopulationInfections model. + """ + sites, latent_sample = subpop_trace + expected = ( + np.asarray(latent_sample.all_subpops) + * POPULATION_SIZE + * np.asarray(SUBPOP_FRACTIONS) + ) + np.testing.assert_allclose(sites["counts_infections"], expected, rtol=1e-12) + np.testing.assert_allclose( + sites["latent_infections_by_subpop"], expected, rtol=1e-12 + ) + + +def test_measurement_receives_per_capita( + subpop_trace: tuple[dict[str, np.ndarray], LatentSample], +) -> None: + """ + A MeasurementObservation receives infections per resident. + + Parameters + ---------- + subpop_trace + Traced SubpopulationInfections model. + """ + sites, latent_sample = subpop_trace + np.testing.assert_allclose( + sites["ww_predicted"], + np.asarray(latent_sample.all_subpops), + rtol=1e-12, + ) + + +def test_population_infections_unchanged( + population_infections: PopulationInfections, +) -> None: + """ + With PopulationInfections and no subpopulation fractions, aggregate + and subpopulation counts are both infections per resident times the + population size. + + Parameters + ---------- + population_infections + Latent process with a single population. + """ + model = MultiSignalModel( + latent_process=population_infections, + observations={ + "agg": RecordingObservation("agg", "aggregate"), + "counts": RecordingObservation("counts", "subpop"), + }, + ) + sites, latent_sample = _trace_model(model, None) + expected = np.asarray(latent_sample.aggregate) * POPULATION_SIZE + np.testing.assert_allclose(sites["agg_infections"], expected, rtol=1e-12) + assert sites["counts_infections"].shape == (expected.size, 1) + np.testing.assert_allclose(sites["counts_infections"][:, 0], expected, rtol=1e-12) + + +def test_validate_accepts_subpop_per_capita( + subpopulation_infections: SubpopulationInfections, +) -> None: + """ + MultiSignalModel accepts an observation requesting "subpop_per_capita". + + Parameters + ---------- + subpopulation_infections + Latent process with subpopulations. + """ + model = MultiSignalModel( + latent_process=subpopulation_infections, + observations={"ww": RecordingObservation("ww", "subpop_per_capita")}, + ) + model.validate() diff --git a/test/test_observation_validation.py b/test/test_observation_validation.py index 25d79bd23..f24e7413e 100644 --- a/test/test_observation_validation.py +++ b/test/test_observation_validation.py @@ -13,50 +13,13 @@ from pyrenew.deterministic import DeterministicPMF, DeterministicVariable from pyrenew.observation import ( HierarchicalNormalNoise, - MeasurementObservation, PoissonNoise, PopulationCounts, SubpopulationCounts, ) from pyrenew.randomvariable import DistributionalVariable, VectorizedVariable from pyrenew.time import MMWR_WEEK, daily_to_weekly - -# --------------------------------------------------------------------------- -# Helpers – minimal concrete subclass of MeasurementObservation for testing -# --------------------------------------------------------------------------- - - -class StubMeasurementObservation(MeasurementObservation): - """Minimal concrete MeasurementObservation for testing validate_data().""" - - def validate(self) -> None: - """ - Validate parameters. - - Raises - ------ - ValueError - If PMF is invalid. - """ - pmf = self.temporal_pmf_rv() - self._validate_pmf(pmf, "temporal_pmf_rv") - - def _predicted_obs(self, infections): - """ - Return infections unchanged (identity transform). - - Parameters - ---------- - infections : ArrayLike - Input infections. - - Returns - ------- - ArrayLike - Same as input. - """ - return infections - +from test.test_helpers import StubMeasurementObservation # --------------------------------------------------------------------------- # Fixtures From 39508f05fabb74be873df3a07dff87cea12c0505 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Mon, 12 Oct 2026 13:49:37 -0400 Subject: [PATCH 02/24] update tutorials --- docs/tutorials/building_multisignal_models.qmd | 10 +++++++++- docs/tutorials/latent_subpopulation_infections.qmd | 4 +--- docs/tutorials/observation_processes_counts.qmd | 1 + docs/tutorials/observation_processes_measurements.qmd | 4 ++-- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 4d955b05c..2512b0164 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -176,7 +176,15 @@ Each observation process declares an **infection resolution** that determines wh | `"subpop"` | Infection counts in each subpopulation, shape `(T, n_subpops)` | Regional case counts | | `"subpop_per_capita"` | Infections per resident of each subpopulation, shape `(T, n_subpops)` | Wastewater concentrations | -The `PyrenewBuilder` routes latent infections to observation processes based on each process's declared resolution. +The model routes latent infections to observation processes based on each process's declared resolution. + +The latent infection process generates infections per resident: $I_k(t)$ is the number of new infections among residents of subpopulation $k$ on day $t$, divided by the number of residents of $k$. +The model converts these to the scale each kind of signal needs, using `population_size` ($N$) and `subpop_fractions` ($p_k$): + +- **Count signals** need infection counts. + Subpopulation $k$ has $N p_k$ residents, so it has $N p_k I_k(t)$ infections, and these counts sum to the jurisdiction total $N \sum_k p_k I_k(t)$. +- **Concentration signals** need infections per resident. + The amount of virus shed in a catchment area and the volume of wastewater that dilutes it both grow with the number of residents, so the number of residents cancels and the concentration depends only on $I_k(t)$. For subpopulation-level observations, the observation process selects which subpopulations it observes using `subpop_indices` provided at sample/fit time. This allows flexible observation patterns---for example, wastewater samples might only cover 5 of 6 subpopulations (catchment areas), while the 6th represents areas without wastewater monitoring. diff --git a/docs/tutorials/latent_subpopulation_infections.qmd b/docs/tutorials/latent_subpopulation_infections.qmd index 39c2e0995..33f9643e1 100644 --- a/docs/tutorials/latent_subpopulation_infections.qmd +++ b/docs/tutorials/latent_subpopulation_infections.qmd @@ -107,8 +107,6 @@ $$ I_{\text{aggregate}}(t) = \sum_{k=1}^{K} p_k \, I_k(t). $$ -The aggregate trajectory $I_{\text{aggregate}}(t)$ is then passed to observation processes, which transform it into expected observations as described in the observation model tutorials. - ### Relationship to the single-population model This model generalizes the single-population renewal model: @@ -759,7 +757,7 @@ Each observation process selects a subset of subpopulations (via `subpop_indices The `PyrenewBuilder` handles the wiring: 1. **`configure_latent()`** sets the hierarchical infection process (called once) -2. **`add_observation()`** adds an observation process, specifying which subpopulations it observes (called once per data stream) +2. **`add_observation()`** adds an observation process (called once per data stream) 3. **`build()`** computes `n_initialization_points` from all delay distributions and produces a model ready for inference To learn how to build complete models with observation processes, see: diff --git a/docs/tutorials/observation_processes_counts.qmd b/docs/tutorials/observation_processes_counts.qmd index 9d9a603fc..108473a89 100644 --- a/docs/tutorials/observation_processes_counts.qmd +++ b/docs/tutorials/observation_processes_counts.qmd @@ -53,6 +53,7 @@ In the renewal modeling framework, observations are generated in two steps: 2. A **stochastic observation model**, where observed counts are drawn from a distribution with that predicted value. Observed data can be aggregated or available as subpopulation-level counts, which are modeled by the classes `PopulationCounts` and `SubpopulationCounts`, respectively. +See [Infection Resolution](building_multisignal_models.md#infection-resolution) for how a multi-signal model converts latent infections to counts for each. ### Observation equation diff --git a/docs/tutorials/observation_processes_measurements.qmd b/docs/tutorials/observation_processes_measurements.qmd index 490a1d350..f0c8620d5 100644 --- a/docs/tutorials/observation_processes_measurements.qmd +++ b/docs/tutorials/observation_processes_measurements.qmd @@ -90,7 +90,7 @@ The core convolution structure is shared with count observations, but key aspect | Output space | Linear (expected counts) | Often log-transformed | | Noise model | Poisson or Negative Binomial | Normal (often on log scale) | | Scaling | Ascertainment rate $\alpha \in [0,1]$ | Domain-specific | - | Subpop structure | Optional (`CountsBySubpop`) | Inherent (sensor/site effects) | + | Subpop structure | Optional (`SubpopulationCounts`) | Inherent (sensor/site effects) | ### The noise model @@ -203,7 +203,7 @@ $$ where: -- $I(t-d)$ is the number of infections per resident on day $t-d$ +- $I(t-d)$ is the number of infections per resident on day $t-d$ (see [Infection Resolution](building_multisignal_models.md#infection-resolution)) - $\pi_d$ is the shedding kinetics PMF, giving the fraction of total shedding occurring $d$ days after infection, analogous to the delay distribution in count observation models - $G$ is the number of genome copies shed per infection - $V$ is the wastewater volume per person per day From f331b29d1b6158d82a8dd92159916828c424e58a Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 16:32:17 -0400 Subject: [PATCH 03/24] changes per co-pilot review --- docs/tutorials/ascertainment.qmd | 2 +- test/test_helpers.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 1f01422e8..82deb3b60 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -360,7 +360,7 @@ model = builder.build() - Ascertainment models produce one scalar baseline rate per signal. - Supplying `temporal_processes` adds signal-specific deviations to selected baselines to model time-varying ascertainment rates. -- For `JointAscertainment` and `RatioLinkedAscertainemt` models, the temporal deviations do not preserve the baseline covariance or ratio pointwise over time. +- For `JointAscertainment` and `RatioLinkedAscertainment` models, the temporal deviations do not preserve the baseline covariance or ratio pointwise over time. The choice of ascertainment model should consider: diff --git a/test/test_helpers.py b/test/test_helpers.py index b8364a248..2d006d8d4 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -175,7 +175,7 @@ def validate(self) -> None: pmf = self.temporal_pmf_rv() self._validate_pmf(pmf, "temporal_pmf_rv") - def _predicted_obs(self, infections): + def _predicted_obs(self, infections: ArrayLike) -> ArrayLike: """ Return infections unchanged (identity transform). From dc3d81705761932fc90daf67b18e5eaa1e2f7dd0 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 16:33:27 -0400 Subject: [PATCH 04/24] changes per co-pilot review --- docs/tutorials/ascertainment.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 82deb3b60..0dc389f10 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -106,7 +106,7 @@ Use `JointAscertainment` when prior dependence among baseline rates is naturally Use `RatioLinkedAscertainment` when a relative quantity, such as IHR / IEDR, is scientifically meaningful and can be assigned an interpretable prior. -For constant, unrelated rates that do not otherwise need an ascertainment component, each observation process can instead receive a `RandomVariable` directly, as described in the next section. +For constant, unrelated rates that do not otherwise need an ascertainment component each observation process can instead receive a `RandomVariable` directly, as described in the next section. ## Independent ascertainment rates From 31caf255c3783d622465b825c03ef9082e0be36a Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 16:44:00 -0400 Subject: [PATCH 05/24] changes per co-pilot review --- docs/tutorials/ascertainment.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 0dc389f10..82deb3b60 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -106,7 +106,7 @@ Use `JointAscertainment` when prior dependence among baseline rates is naturally Use `RatioLinkedAscertainment` when a relative quantity, such as IHR / IEDR, is scientifically meaningful and can be assigned an interpretable prior. -For constant, unrelated rates that do not otherwise need an ascertainment component each observation process can instead receive a `RandomVariable` directly, as described in the next section. +For constant, unrelated rates that do not otherwise need an ascertainment component, each observation process can instead receive a `RandomVariable` directly, as described in the next section. ## Independent ascertainment rates From 90165b6e10af0ece9f8440726b40640b6d900064 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 18:08:31 -0400 Subject: [PATCH 06/24] Update docs/tutorials/building_multisignal_models.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/building_multisignal_models.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 2512b0164..892e72663 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -183,7 +183,7 @@ The model converts these to the scale each kind of signal needs, using `populati - **Count signals** need infection counts. Subpopulation $k$ has $N p_k$ residents, so it has $N p_k I_k(t)$ infections, and these counts sum to the jurisdiction total $N \sum_k p_k I_k(t)$. -- **Concentration signals** need infections per resident. +- **Some measurement signals** reflection information about infections per resident. The amount of virus shed in a catchment area and the volume of wastewater that dilutes it both grow with the number of residents, so the number of residents cancels and the concentration depends only on $I_k(t)$. For subpopulation-level observations, the observation process selects which subpopulations it observes using `subpop_indices` provided at sample/fit time. From 9ac7e19ccb398a6da84a75c5b24d3ace4ec8d754 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 18:09:11 -0400 Subject: [PATCH 07/24] Update docs/tutorials/building_multisignal_models.qmd Co-authored-by: Dylan H. Morris --- docs/tutorials/building_multisignal_models.qmd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 892e72663..9cfe9e5c6 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -184,7 +184,7 @@ The model converts these to the scale each kind of signal needs, using `populati - **Count signals** need infection counts. Subpopulation $k$ has $N p_k$ residents, so it has $N p_k I_k(t)$ infections, and these counts sum to the jurisdiction total $N \sum_k p_k I_k(t)$. - **Some measurement signals** reflection information about infections per resident. - The amount of virus shed in a catchment area and the volume of wastewater that dilutes it both grow with the number of residents, so the number of residents cancels and the concentration depends only on $I_k(t)$. + For example, the number genomes shed into wastewater scales with the number of infected individuals, and the total volume of wastewater scales with the total population, so it can be reasonable to model wastewater concentrations purely as a function of per-capita incidence $I_k(t)$. For subpopulation-level observations, the observation process selects which subpopulations it observes using `subpop_indices` provided at sample/fit time. This allows flexible observation patterns---for example, wastewater samples might only cover 5 of 6 subpopulations (catchment areas), while the 6th represents areas without wastewater monitoring. From 868196147843e5b3daf33aec37a8d1cc4c08e9f9 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 18:19:47 -0400 Subject: [PATCH 08/24] generalize docs language: use member of subpopulation instead of resident --- docs/tutorials/building_multisignal_models.qmd | 18 +++++++++--------- .../observation_processes_measurements.qmd | 10 +++++++--- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 9cfe9e5c6..ae51c99a1 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -170,20 +170,20 @@ flowchart TB Different observation processes observe different levels of the model hierarchy. Each observation process declares an **infection resolution** that determines what infection data it receives: - | Resolution | Receives | Example signals | - | --------------------- | --------------------------------------------------------------------- | -------------------------------- | - | `"aggregate"` | Infection counts for the whole jurisdiction, shape `(T,)` | Hospital admissions, case counts | - | `"subpop"` | Infection counts in each subpopulation, shape `(T, n_subpops)` | Regional case counts | - | `"subpop_per_capita"` | Infections per resident of each subpopulation, shape `(T, n_subpops)` | Wastewater concentrations | + | Resolution | Receives | Example signals | + | --------------------- | ------------------------------------------------------------------- | -------------------------------- | + | `"aggregate"` | Infection counts for the whole jurisdiction, shape `(T,)` | Hospital admissions, case counts | + | `"subpop"` | Infection counts in each subpopulation, shape `(T, n_subpops)` | Regional case counts | + | `"subpop_per_capita"` | Infections per member of each subpopulation, shape `(T, n_subpops)` | Wastewater concentrations | The model routes latent infections to observation processes based on each process's declared resolution. -The latent infection process generates infections per resident: $I_k(t)$ is the number of new infections among residents of subpopulation $k$ on day $t$, divided by the number of residents of $k$. +The latent infection process generates infections per individual: $I_k(t)$ is the number of new infections among members of subpopulation $k$ on day $t$, divided by the number of members of $k$. The model converts these to the scale each kind of signal needs, using `population_size` ($N$) and `subpop_fractions` ($p_k$): - **Count signals** need infection counts. - Subpopulation $k$ has $N p_k$ residents, so it has $N p_k I_k(t)$ infections, and these counts sum to the jurisdiction total $N \sum_k p_k I_k(t)$. -- **Some measurement signals** reflection information about infections per resident. + Subpopulation $k$ has $N p_k$ members, so it has $N p_k I_k(t)$ infections, and these counts sum to the jurisdiction total $N \sum_k p_k I_k(t)$. +- **Some measurement signals** provide information about infections per member of the subpopulation. For example, the number genomes shed into wastewater scales with the number of infected individuals, and the total volume of wastewater scales with the total population, so it can be reasonable to model wastewater concentrations purely as a function of per-capita incidence $I_k(t)$. For subpopulation-level observations, the observation process selects which subpopulations it observes using `subpop_indices` provided at sample/fit time. @@ -434,7 +434,7 @@ class Wastewater(MeasurementObservation): """ Wastewater viral concentration observation process. - Transforms infections per resident of each subpopulation into + Transforms infections per member of each subpopulation into predicted log-concentrations via shedding kinetics convolution and genome/volume scaling. """ diff --git a/docs/tutorials/observation_processes_measurements.qmd b/docs/tutorials/observation_processes_measurements.qmd index f0c8620d5..f8e407839 100644 --- a/docs/tutorials/observation_processes_measurements.qmd +++ b/docs/tutorials/observation_processes_measurements.qmd @@ -203,7 +203,7 @@ $$ where: -- $I(t-d)$ is the number of infections per resident on day $t-d$ (see [Infection Resolution](building_multisignal_models.md#infection-resolution)) +- $I(t-d)$ is the number of infections per subpopulation member on day $t-d$ (see [Infection Resolution](building_multisignal_models.md#infection-resolution)) - $\pi_d$ is the shedding kinetics PMF, giving the fraction of total shedding occurring $d$ days after infection, analogous to the delay distribution in count observation models - $G$ is the number of genome copies shed per infection - $V$ is the wastewater volume per person per day @@ -245,7 +245,7 @@ class Wastewater(MeasurementObservation): """ Wastewater viral concentration observation process. - Transforms infections per resident of each subpopulation into + Transforms infections per member of each subpopulation into predicted log-concentrations via shedding kinetics convolution and genome/volume scaling. """ @@ -507,7 +507,11 @@ plot_infections = ( linetype="dashed", alpha=0.5, ) - + p9.labs(x="Day", y="Daily Infections per Resident", title="Infections (Input)") + + p9.labs( + x="Day", + y="Daily Infections per Subpopulation member", + title="Infections (Input)", + ) + theme_tutorial + p9.annotate( "text", From cef08e00f96f9ea57e7e162bb04fffb4e7a06a0a Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:25:39 -0400 Subject: [PATCH 09/24] Update pyrenew/model/multisignal_model.py Co-authored-by: Dylan H. Morris --- pyrenew/model/multisignal_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/model/multisignal_model.py b/pyrenew/model/multisignal_model.py index c5e8ef5dd..0284062d4 100644 --- a/pyrenew/model/multisignal_model.py +++ b/pyrenew/model/multisignal_model.py @@ -56,7 +56,7 @@ class MultiSignalModel(Model): shape ``(n_total,)``. - ``"subpop"``: infection counts in each subpopulation, shape ``(n_total, n_subpops)``. - - ``"subpop_per_capita"``: infections per resident of each + - ``"subpop_per_capita"``: infections per member of each subpopulation, shape ``(n_total, n_subpops)``. Subpopulation observations select their subpopulations via From 80db5d3ca28f8af7a70ec2d2885bb0b0924d5dd7 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:25:59 -0400 Subject: [PATCH 10/24] Update pyrenew/observation/__init__.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/observation/__init__.py b/pyrenew/observation/__init__.py index 21b51a814..40b304ed0 100644 --- a/pyrenew/observation/__init__.py +++ b/pyrenew/observation/__init__.py @@ -13,7 +13,7 @@ - ``sample()``: Sample observations given infections - ``infection_resolution()``: returns ``"aggregate"`` (infection counts for the whole jurisdiction), ``"subpop"`` (infection counts in each - subpopulation), or ``"subpop_per_capita"`` (infections per resident of + subpopulation), or ``"subpop_per_capita"`` (infections per member of each subpopulation) - ``lookback_days()``: returns required infection history length From 942bece36429c1e90770663fbe635bbc84c27251 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:26:17 -0400 Subject: [PATCH 11/24] Update pyrenew/observation/base.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/observation/base.py b/pyrenew/observation/base.py index 0fbb2b3d9..fdb7546dc 100644 --- a/pyrenew/observation/base.py +++ b/pyrenew/observation/base.py @@ -112,7 +112,7 @@ def infection_resolution(self) -> str: Returns one of: - - ``"aggregate"``: infection counts for the whole jurisdiction. + - ``"aggregate"``: infection counts for the whole population. Shape: ``(n_total,)`` - ``"subpop"``: infection counts in each subpopulation. Shape: ``(n_total, n_subpops)``, indexed via ``subpop_indices``. From c48737e5e8b0f3c3eec34ba03e5aa8a8724b1c5d Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:26:31 -0400 Subject: [PATCH 12/24] Update pyrenew/observation/base.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/observation/base.py b/pyrenew/observation/base.py index fdb7546dc..b1a37e9a2 100644 --- a/pyrenew/observation/base.py +++ b/pyrenew/observation/base.py @@ -116,7 +116,7 @@ def infection_resolution(self) -> str: Shape: ``(n_total,)`` - ``"subpop"``: infection counts in each subpopulation. Shape: ``(n_total, n_subpops)``, indexed via ``subpop_indices``. - - ``"subpop_per_capita"``: infections per resident of each + - ``"subpop_per_capita"``: infections per member of each subpopulation. Shape: ``(n_total, n_subpops)``, indexed via ``subpop_indices``. From 545640edd16622ce33ed15c6e5ba977047051254 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:26:44 -0400 Subject: [PATCH 13/24] Update pyrenew/observation/measurement_observations.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/measurement_observations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/observation/measurement_observations.py b/pyrenew/observation/measurement_observations.py index aa1921ddf..1865b534b 100644 --- a/pyrenew/observation/measurement_observations.py +++ b/pyrenew/observation/measurement_observations.py @@ -37,7 +37,7 @@ class MeasurementObservation(BaseObservationProcess): Subclasses must implement ``_predicted_obs()`` according to their specific signal processing (e.g., wastewater shedding kinetics, dilution factors, etc.). ``_predicted_obs()`` receives infections - per resident of each subpopulation. + per member of each subpopulation. See Also -------- From 223bf28e085d06ffd73e05f0edca78e03fa43041 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:26:57 -0400 Subject: [PATCH 14/24] Update pyrenew/observation/measurement_observations.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/measurement_observations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/observation/measurement_observations.py b/pyrenew/observation/measurement_observations.py index 1865b534b..68f5e42bd 100644 --- a/pyrenew/observation/measurement_observations.py +++ b/pyrenew/observation/measurement_observations.py @@ -95,7 +95,7 @@ def infection_resolution(self) -> str: """ Return "subpop_per_capita" for measurement observations. - Measurement observations receive infections per resident of each + Measurement observations receive infections per member of each subpopulation, because each measurement corresponds to a specific catchment area and depends on infections per resident. From e6081a51448fd06ecd76ee43d4989587d406a29c Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:27:40 -0400 Subject: [PATCH 15/24] Update pyrenew/observation/measurement_observations.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/measurement_observations.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pyrenew/observation/measurement_observations.py b/pyrenew/observation/measurement_observations.py index 68f5e42bd..7e6a5cd0f 100644 --- a/pyrenew/observation/measurement_observations.py +++ b/pyrenew/observation/measurement_observations.py @@ -96,8 +96,7 @@ def infection_resolution(self) -> str: Return "subpop_per_capita" for measurement observations. Measurement observations receive infections per member of each - subpopulation, because each measurement corresponds to a specific - catchment area and depends on infections per resident. + subpopulation. Returns ------- From a4f88bed5d61bce31bc10f17f11140e5f86f87b2 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:27:55 -0400 Subject: [PATCH 16/24] Update pyrenew/observation/measurement_observations.py Co-authored-by: Dylan H. Morris --- pyrenew/observation/measurement_observations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrenew/observation/measurement_observations.py b/pyrenew/observation/measurement_observations.py index 7e6a5cd0f..ca4c5d348 100644 --- a/pyrenew/observation/measurement_observations.py +++ b/pyrenew/observation/measurement_observations.py @@ -176,7 +176,7 @@ def sample( Parameters ---------- infections - Infections per resident of each subpopulation. + Infections per member of each subpopulation. Shape: (n_total, n_subpops). times Day index for each observation on the shared time axis. From 19ac501ae578284af852b671002b229041339867 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:28:11 -0400 Subject: [PATCH 17/24] Update test/integration/test_subpopulation_counts_scaling.py Co-authored-by: Dylan H. Morris --- test/integration/test_subpopulation_counts_scaling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index 7dcefac8f..aa7f54b24 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -3,7 +3,7 @@ Four subpopulations share the same true Rt and the same initial infections per resident, so they differ only in population size. Counts are simulated -as infections per resident times the subpopulation population. A +as infections per capita times the subpopulation population. A SubpopulationInfections model fitted to these counts should recover Rt in every subpopulation, with no dependence on subpopulation size, and its recorded subpopulation counts should sum to the aggregate count. From 4b70a8b6112f1ea1651f6e4f1903249b9902d5e9 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:28:24 -0400 Subject: [PATCH 18/24] Update test/integration/test_subpopulation_counts_scaling.py Co-authored-by: Dylan H. Morris --- test/integration/test_subpopulation_counts_scaling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index aa7f54b24..5dfc6bf43 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -33,7 +33,7 @@ N_SUBPOPS = len(SUBPOP_FRACTIONS) GEN_INT = np.array([0.2, 0.35, 0.25, 0.15, 0.05]) DELAY = np.array([0.1, 0.3, 0.3, 0.2, 0.1]) -I0_PER_RESIDENT = 2e-5 +I0_PER_CAPITA = 2e-5 ASCERTAINMENT = 0.3 N_DAYS_FIT = 70 N_DAYS_EARLY = 21 From de99c05f223f38f75b20b02392943e292c5d7265 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:30:14 -0400 Subject: [PATCH 19/24] Update test/integration/test_subpopulation_counts_scaling.py Co-authored-by: Dylan H. Morris --- test/integration/test_subpopulation_counts_scaling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index 5dfc6bf43..672c55cb5 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -63,7 +63,7 @@ def _true_rt(n_total: int) -> np.ndarray: def _simulate(n_init: int) -> dict[str, np.ndarray]: """ - Simulate infections per resident and observed counts for each subpopulation. + Simulate per-capita infections and observed counts for each subpopulation. Initialization-period infections grow exponentially from ``I0_PER_RESIDENT`` at the rate implied by the first day's Rt, as in From 5120879eb739d203627141747db6cb8018571dac Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:30:30 -0400 Subject: [PATCH 20/24] Update test/integration/test_subpopulation_counts_scaling.py Co-authored-by: Dylan H. Morris --- test/integration/test_subpopulation_counts_scaling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index 672c55cb5..f669a0759 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -66,7 +66,7 @@ def _simulate(n_init: int) -> dict[str, np.ndarray]: Simulate per-capita infections and observed counts for each subpopulation. Initialization-period infections grow exponentially from - ``I0_PER_RESIDENT`` at the rate implied by the first day's Rt, as in + ``I0_PER_CAPITA`` at the rate implied by the first day's Rt, as in ``SubpopulationInfections``. Counts in subpopulation ``k`` are infections per resident times ``POPULATION_SIZE * SUBPOP_FRACTIONS[k]``, delayed by ``DELAY``, scaled by ``ASCERTAINMENT``, and drawn from a From 253aa484a5ac5f01ee5bffe3f5395eb8514eea98 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:31:01 -0400 Subject: [PATCH 21/24] Update test/integration/test_subpopulation_counts_scaling.py Co-authored-by: Dylan H. Morris --- test/integration/test_subpopulation_counts_scaling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index f669a0759..69301fa70 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -68,7 +68,7 @@ def _simulate(n_init: int) -> dict[str, np.ndarray]: Initialization-period infections grow exponentially from ``I0_PER_CAPITA`` at the rate implied by the first day's Rt, as in ``SubpopulationInfections``. Counts in subpopulation ``k`` are - infections per resident times ``POPULATION_SIZE * SUBPOP_FRACTIONS[k]``, + per-capita infections times ``POPULATION_SIZE * SUBPOP_FRACTIONS[k]``, delayed by ``DELAY``, scaled by ``ASCERTAINMENT``, and drawn from a Poisson distribution. From 1fe36eb51fa676407bb4a8003532bcbf7bf7845c Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:40:46 -0400 Subject: [PATCH 22/24] changes per code review --- test/integration/test_subpopulation_counts_scaling.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index 69301fa70..1d93c1a63 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -90,7 +90,7 @@ def _simulate(n_init: int) -> dict[str, np.ndarray]: r_approx_from_R(R=jnp.asarray(rt[0]), g=jnp.asarray(GEN_INT), n_newton_steps=4) ) infections = np.zeros(n_total) - infections[:n_init] = I0_PER_RESIDENT * np.exp( + infections[:n_init] = I0_PER_CAPITA * np.exp( initial_growth_rate * np.arange(n_init) ) reversed_gen_int = GEN_INT[::-1] @@ -135,7 +135,7 @@ def scaling_model() -> MultiSignalModel: SubpopulationInfections, name="latent", gen_int_rv=DeterministicPMF("gen_int", jnp.asarray(GEN_INT)), - I0_rv=DeterministicVariable("I0", jnp.full(N_SUBPOPS, I0_PER_RESIDENT)), + I0_rv=DeterministicVariable("I0", jnp.full(N_SUBPOPS, I0_PER_CAPITA)), log_rt_time_0_rv=DistributionalVariable("log_rt_time_0", dist.Normal(0.0, 0.5)), baseline_rt_process=RandomWalk(DeterministicVariable("baseline_sd", 0.05)), subpop_rt_deviation_process=RandomWalk( From 44081fb669d9d03c621dfa3482dbe543440a4e61 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 11:52:54 -0400 Subject: [PATCH 23/24] changes per code review --- .../test_subpopulation_counts_scaling.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index 1d93c1a63..df38f1f4c 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -19,8 +19,10 @@ import pytest from pyrenew.ascertainment import IndependentAscertainment +from pyrenew.convolve import compute_delay_ascertained_incidence from pyrenew.deterministic import DeterministicPMF, DeterministicVariable from pyrenew.latent import RandomWalk, SubpopulationInfections +from pyrenew.latent.infection_functions import compute_infections_from_rt from pyrenew.math import r_approx_from_R from pyrenew.model import MultiSignalModel, PyrenewBuilder from pyrenew.observation import PoissonNoise, SubpopulationCounts @@ -89,19 +91,23 @@ def _simulate(n_init: int) -> dict[str, np.ndarray]: initial_growth_rate = float( r_approx_from_R(R=jnp.asarray(rt[0]), g=jnp.asarray(GEN_INT), n_newton_steps=4) ) - infections = np.zeros(n_total) - infections[:n_init] = I0_PER_CAPITA * np.exp( - initial_growth_rate * np.arange(n_init) + initial_infections = I0_PER_CAPITA * np.exp(initial_growth_rate * np.arange(n_init)) + post_initialization_infections = compute_infections_from_rt( + I0=jnp.asarray(initial_infections[-GEN_INT.size :]), + Rt=jnp.asarray(rt[n_init:]), + reversed_generation_interval_pmf=jnp.asarray(GEN_INT[::-1]), ) - reversed_gen_int = GEN_INT[::-1] - for t in range(n_init, n_total): - infections[t] = rt[t] * reversed_gen_int @ infections[t - GEN_INT.size : t] - counts = POPULATION_SIZE * np.outer(infections, SUBPOP_FRACTIONS) - predicted = np.zeros_like(counts) - for d, weight in enumerate(DELAY): - predicted[d:] += ASCERTAINMENT * weight * counts[: n_total - d] + infections = jnp.concatenate( + [jnp.asarray(initial_infections), post_initialization_infections] + ) + predicted_per_capita, offset = compute_delay_ascertained_incidence( + latent_incidence=infections, + delay_incidence_to_observation_pmf=jnp.asarray(DELAY), + p_observed_given_incident=ASCERTAINMENT, + ) + predicted = POPULATION_SIZE * np.outer(predicted_per_capita, SUBPOP_FRACTIONS) rng = np.random.default_rng(DATA_SEED) - observed = rng.poisson(predicted[n_init:]) + observed = rng.poisson(predicted[n_init - offset :]) return {"rt": rt, "observed": observed} From aa0773eb30460341d3edcb3a06e88773e2bfd4be Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 14 Oct 2026 12:18:30 -0400 Subject: [PATCH 24/24] refine pytests - use pytest-env --- pyproject.toml | 6 +++ test/conftest.py | 33 ++----------- .../test_subpopulation_counts_scaling.py | 18 +------ test/test_multisignal_infection_scaling.py | 16 ++++--- uv.lock | 48 ++++++++++++++----- 5 files changed, 58 insertions(+), 63 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 302a8726f..245c39211 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dev = [ "pre-commit>=4.6.0", "pytest>=9.0.3", "pytest-cov>=6.3.0", + "pytest-env>=1.2.0", "ruff>=0.15.16", "scipy>=1.16.1", ] @@ -76,10 +77,15 @@ known_first_party = ["pyrenew", "test"] DEP004 = ["arviz", "pytest", "scipy", "bs4"] [tool.pytest.ini_options] +required_plugins = ["pytest-env"] markers = [ "integration: integration tests that fit models via MCMC (deselect with '-m \"not integration\"')", ] +[tool.pytest_env] +JAX_ENABLE_X64 = { value = "true", skip_if_set = true } +XLA_FLAGS = { value = "--xla_force_host_platform_device_count=4", skip_if_set = true } + [tool.ruff] fix = true diff --git a/test/conftest.py b/test/conftest.py index 9ad236836..59407c845 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -4,22 +4,12 @@ This module provides reusable fixtures for creating observation processes, test data, and common configurations used across multiple test files. -All tests (unit and integration) run with 64-bit floats. The module sets -``JAX_ENABLE_X64`` before any jax import, overriding any shell value, and -the autouse fixture ``require_float64`` fails any test that runs without -64-bit floats. It also makes four logical host devices available for -parallel MCMC chains, unless the caller already set ``XLA_FLAGS``. JAX -reads these variables at import time, so they must be set before the -first ``import jax`` anywhere in the test process; placing them at the -top of this file (loaded by pytest before any test module) satisfies that -requirement. +JAX defaults are configured through pytest-env in ``pyproject.toml``: +64-bit floats and four logical host devices for parallel MCMC chains. +Existing environment values take precedence; set ``JAX_ENABLE_X64=false`` +to run tests with 32-bit floats. """ -import os - -os.environ["JAX_ENABLE_X64"] = "true" -os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") - import jax.numpy as jnp import numpyro.distributions as dist import pytest @@ -40,21 +30,6 @@ from pyrenew.time import MMWR_WEEK from test.test_helpers import fixed_ar1, fixed_random_walk - -@pytest.fixture(autouse=True) -def require_float64() -> None: - """ - Fail any test that would run without 64-bit floats. - - Raises - ------ - pytest.fail.Exception - If JAX creates 32-bit floating point arrays by default. - """ - if jnp.asarray(0.0).dtype != jnp.float64: - pytest.fail("Tests must run with 64-bit floats; JAX_ENABLE_X64 is off.") - - # ============================================================================= # PMF Fixtures # ============================================================================= diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py index df38f1f4c..e1b646d6d 100644 --- a/test/integration/test_subpopulation_counts_scaling.py +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -214,21 +214,6 @@ def posterior( class TestSubpopulationCountScaling: """Fitted subpopulation Rt and infections do not depend on subpopulation size.""" - def test_fit_uses_float64( - self, - posterior: dict[str, np.ndarray], - ) -> None: - """ - The model is fitted in 64-bit floats. - - Parameters - ---------- - posterior - Posterior samples. - """ - assert posterior["latent::rt_subpop"].dtype == np.float64 - assert posterior["latent_infections_by_subpop"].dtype == np.float64 - def test_subpop_counts_sum_to_aggregate( self, posterior: dict[str, np.ndarray], @@ -244,7 +229,8 @@ def test_subpop_counts_sum_to_aggregate( np.testing.assert_allclose( posterior["latent_infections_by_subpop"].sum(axis=-1), posterior["latent_infections"], - rtol=1e-10, + rtol=1e-6, + atol=0, ) def test_rt_recovered_in_every_subpop( diff --git a/test/test_multisignal_infection_scaling.py b/test/test_multisignal_infection_scaling.py index b7f82815c..da1588b06 100644 --- a/test/test_multisignal_infection_scaling.py +++ b/test/test_multisignal_infection_scaling.py @@ -234,7 +234,8 @@ def test_subpop_counts_sum_to_aggregate( np.testing.assert_allclose( sites["latent_infections_by_subpop"].sum(axis=1), sites["latent_infections"], - rtol=1e-12, + rtol=1e-6, + atol=0, ) @@ -256,9 +257,9 @@ def test_subpop_counts_scaled_by_subpop_population( * POPULATION_SIZE * np.asarray(SUBPOP_FRACTIONS) ) - np.testing.assert_allclose(sites["counts_infections"], expected, rtol=1e-12) + np.testing.assert_allclose(sites["counts_infections"], expected, rtol=1e-6, atol=0) np.testing.assert_allclose( - sites["latent_infections_by_subpop"], expected, rtol=1e-12 + sites["latent_infections_by_subpop"], expected, rtol=1e-6, atol=0 ) @@ -277,7 +278,8 @@ def test_measurement_receives_per_capita( np.testing.assert_allclose( sites["ww_predicted"], np.asarray(latent_sample.all_subpops), - rtol=1e-12, + rtol=1e-6, + atol=0, ) @@ -303,9 +305,11 @@ def test_population_infections_unchanged( ) sites, latent_sample = _trace_model(model, None) expected = np.asarray(latent_sample.aggregate) * POPULATION_SIZE - np.testing.assert_allclose(sites["agg_infections"], expected, rtol=1e-12) + np.testing.assert_allclose(sites["agg_infections"], expected, rtol=1e-6, atol=0) assert sites["counts_infections"].shape == (expected.size, 1) - np.testing.assert_allclose(sites["counts_infections"][:, 0], expected, rtol=1e-12) + np.testing.assert_allclose( + sites["counts_infections"][:, 0], expected, rtol=1e-6, atol=0 + ) def test_validate_accepts_subpop_per_capita( diff --git a/uv.lock b/uv.lock index 4c3122ab9..9e69f7b66 100644 --- a/uv.lock +++ b/uv.lock @@ -174,7 +174,7 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ @@ -1482,7 +1482,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -1758,6 +1758,7 @@ dev = [ { name = "pre-commit" }, { name = "pytest" }, { name = "pytest-cov" }, + { name = "pytest-env" }, { name = "ruff" }, { name = "scipy" }, ] @@ -1788,6 +1789,7 @@ dev = [ { name = "pre-commit", specifier = ">=4.6.0" }, { name = "pytest", specifier = ">=9.0.3" }, { name = "pytest-cov", specifier = ">=6.3.0" }, + { name = "pytest-env", specifier = ">=1.2.0" }, { name = "ruff", specifier = ">=0.15.16" }, { name = "scipy", specifier = ">=1.16.1" }, ] @@ -1822,6 +1824,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" }, ] +[[package]] +name = "pytest-env" +version = "1.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, + { name = "python-dotenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/32/72/d3e125d18f798b2968430b77dba991c69df354ec649db4c1480ac9bc2ab4/pytest_env-1.7.1.tar.gz", hash = "sha256:f2c5aed2621dbfc73c2866a710e2e456409495c5ac10905612c9509b3ee631a2", size = 17091, upload-time = "2026-09-08T14:49:44.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/82/fd9f2855bb9d48022942c2123f287d006b7569ae9b2cae8939f94a25fdf6/pytest_env-1.7.1-py3-none-any.whl", hash = "sha256:22341b945305b65ef4e53d7c678982cfd64b5c5800afdf7d37caad2362ea002c", size = 10361, upload-time = "2026-09-08T14:49:43.288Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -1846,6 +1861,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/7d/e9ffbadfbf89c93848412d04594135c4ae8c1d37d9e053b9c3ed718fabc4/python_discovery-1.6.1-py3-none-any.whl", hash = "sha256:d43fcdef879fe795352bd13ccf8d185ba5a9f86f36cfcd00529f596e737442b3", size = 38664, upload-time = "2026-09-18T01:31:52.448Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/74/26/2fbeedb218a787a5eea551c7532cac4e009f83d689dd2faa0d0353473f86/python_dotenv-1.2.4.tar.gz", hash = "sha256:f0d53e69935a851c0dcc78f3ab7aaccd8cabef0b92382b576b824212902873c0", size = 60824, upload-time = "2026-10-01T05:36:10Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/60/d1/38f3a3405989a89ac18390803e70c6ad7c7760da4f9b83cbeca0c44a0c72/python_dotenv-1.2.4-py3-none-any.whl", hash = "sha256:42269a8a5b3fd54ffa6f3d84b18abed50064717576b4ecf03dc4a55d8aa04fdc", size = 23266, upload-time = "2026-10-01T05:36:08.633Z" }, +] + [[package]] name = "pyyaml" version = "6.0.3" @@ -2160,11 +2184,11 @@ resolution-markers = [ "python_full_version < '3.14' and sys_platform == 'emscripten'", ] dependencies = [ - { name = "numpy" }, - { name = "packaging" }, - { name = "pandas" }, - { name = "patsy" }, - { name = "scipy" }, + { name = "numpy", marker = "sys_platform == 'emscripten'" }, + { name = "packaging", marker = "sys_platform == 'emscripten'" }, + { name = "pandas", marker = "sys_platform == 'emscripten'" }, + { name = "patsy", marker = "sys_platform == 'emscripten'" }, + { name = "scipy", marker = "sys_platform == 'emscripten'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1f/3b/963a015dd8ea17e10c7b0e2f14d7c4daec903baf60a017e756b57953a4bf/statsmodels-0.14.4.tar.gz", hash = "sha256:5d69e0f39060dc72c067f9bb6e8033b6dccdb0bae101d76a7ef0bcc94e898b67", size = 20354802, upload-time = "2024-10-03T16:15:36.273Z" } @@ -2179,11 +2203,11 @@ resolution-markers = [ "python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy" }, - { name = "packaging" }, - { name = "pandas" }, - { name = "patsy" }, - { name = "scipy" }, + { name = "numpy", marker = "sys_platform != 'emscripten'" }, + { name = "packaging", marker = "sys_platform != 'emscripten'" }, + { name = "pandas", marker = "sys_platform != 'emscripten'" }, + { name = "patsy", marker = "sys_platform != 'emscripten'" }, + { name = "scipy", marker = "sys_platform != 'emscripten'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0d/81/e8d74b34f85285f7335d30c5e3c2d7c0346997af9f3debf9a0a9a63de184/statsmodels-0.14.6.tar.gz", hash = "sha256:4d17873d3e607d398b85126cd4ed7aad89e4e9d89fc744cdab1af3189a996c2a", size = 20689085, upload-time = "2025-12-05T23:08:39.522Z" } wheels = [