Mem bugfix 903 subpop scaling

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:
 

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

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:

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
 

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

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

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

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

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"``:

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

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
 # =============================================================================