diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 1f01422e..82deb3b6 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/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 6e64e2cb..ae51c99a 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -170,12 +170,21 @@ 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 member 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 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$ 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. 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. @@ -297,7 +306,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 +434,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 member of each subpopulation into + predicted log-concentrations via shedding kinetics convolution and + genome/volume scaling. """ def __init__( @@ -524,7 +534,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/latent_subpopulation_infections.qmd b/docs/tutorials/latent_subpopulation_infections.qmd index 39c2e099..33f9643e 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 9d9a603f..108473a8 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 1f6f3716..f8e40783 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 on day $t-d$ +- $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,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 member 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,11 @@ 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 Subpopulation member", + title="Infections (Input)", + ) + theme_tutorial + p9.annotate( "text", @@ -714,8 +719,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/pyproject.toml b/pyproject.toml index 302a8726..245c3921 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/pyrenew/model/multisignal_model.py b/pyrenew/model/multisignal_model.py index 3e339547..0284062d 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 member 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 b38ab471..40b304ed 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 member 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 c7209748..b1a37e9a 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 population. + Shape: ``(n_total,)`` + - ``"subpop"``: infection counts in each subpopulation. + Shape: ``(n_total, n_subpops)``, indexed via ``subpop_indices``. + - ``"subpop_per_capita"``: infections per member 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 57f0881a..dcd41940 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 03fe13e2..ca4c5d34 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 member of each subpopulation. See Also -------- @@ -92,17 +93,17 @@ 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 member of each + subpopulation. Returns ------- str - ``"subpop"`` + ``"subpop_per_capita"`` """ - return "subpop" + return "subpop_per_capita" def validate_data( self, @@ -175,8 +176,8 @@ def sample( Parameters ---------- infections - Infections from the infection process. - Shape: (n_total, n_subpops) + Infections per member 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 2ba07556..59407c84 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -4,21 +4,12 @@ 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). +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.setdefault("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 diff --git a/test/integration/test_subpopulation_counts_scaling.py b/test/integration/test_subpopulation_counts_scaling.py new file mode 100644 index 00000000..e1b646d6 --- /dev/null +++ b/test/integration/test_subpopulation_counts_scaling.py @@ -0,0 +1,287 @@ +""" +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 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 __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.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 +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_CAPITA = 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 per-capita infections and observed counts for each subpopulation. + + 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 + per-capita infections 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) + ) + 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]), + ) + 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 - offset :]) + 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_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( + 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_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-6, + atol=0, + ) + + 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 3a722136..2d006d8d 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: ArrayLike) -> ArrayLike: + """ + 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 fad284d0..c03ccb95 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 00000000..da1588b0 --- /dev/null +++ b/test/test_multisignal_infection_scaling.py @@ -0,0 +1,330 @@ +""" +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-6, + atol=0, + ) + + +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-6, atol=0) + np.testing.assert_allclose( + sites["latent_infections_by_subpop"], expected, rtol=1e-6, atol=0 + ) + + +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-6, + atol=0, + ) + + +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-6, atol=0) + assert sites["counts_infections"].shape == (expected.size, 1) + np.testing.assert_allclose( + sites["counts_infections"][:, 0], expected, rtol=1e-6, atol=0 + ) + + +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 25d79bd2..f24e7413 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 diff --git a/uv.lock b/uv.lock index 4c3122ab..9e69f7b6 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 = [