From ef4cfa8389f602674bfbe881d54b56b9e536e9f7 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 29 Sep 2026 10:24:50 -0400 Subject: [PATCH 01/12] v1 files for time-varying ascertainment --- docs/tutorials/ascertainment.qmd | 148 ++++++- .../tutorials/building_multisignal_models.qmd | 4 +- .../observation_processes_counts.qmd | 65 ++- .../observation_processes_measurements.qmd | 4 +- pyrenew/ascertainment/__init__.py | 2 + pyrenew/ascertainment/base.py | 60 +++ pyrenew/ascertainment/timevarying.py | 207 +++++++++ pyrenew/convolve.py | 31 ++ pyrenew/model/multisignal_model.py | 96 ++++- pyrenew/model/pyrenew_builder.py | 12 +- pyrenew/observation/base.py | 19 +- pyrenew/observation/count_observations.py | 77 +++- test/test_ascertainment.py | 381 +++++++++++++++++ test/test_convolve.py | 28 ++ test/test_helpers.py | 7 +- test/test_observation_counts.py | 167 ++++++++ test/test_pyrenew_builder.py | 402 +++++++++++++++++- 17 files changed, 1635 insertions(+), 75 deletions(-) create mode 100644 pyrenew/ascertainment/timevarying.py diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index bba714614..07feb18b1 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -24,8 +24,13 @@ jupyter: import jax.numpy as jnp import numpyro.distributions as dist -from pyrenew.ascertainment import JointAscertainment, RatioLinkedAscertainment +from pyrenew.ascertainment import ( + JointAscertainment, + RatioLinkedAscertainment, + TimeVaryingAscertainment, +) from pyrenew.deterministic import DeterministicPMF, DeterministicVariable +from pyrenew.latent import RandomWalk, WeeklyTemporalProcess from pyrenew.model import PyrenewBuilder from pyrenew.observation import NegativeBinomialNoise, PopulationCounts from pyrenew.randomvariable import DistributionalVariable @@ -51,13 +56,14 @@ where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate, and $\pi_ If a model treats the ascertainment rates for different signals as unrelated, each observation process can receive its own `RandomVariable` directly. An ascertainment component is needed when the model represents a relationship between the ascertainment rates of multiple observation processes. -PyRenew provides two such components: +PyRenew provides two scalar models and one time-varying wrapper: - `JointAscertainment` specifies a joint prior for two or more rates. - `RatioLinkedAscertainment` specifies one rate as a multiple of another rate. +- `TimeVaryingAscertainment` uses either scalar model as a baseline and adds an independent temporal deviation for each signal. -The choice between them is part of the statistical model, not only an API choice. -Each component expresses a different assumption about how information is shared across signals. +The choice among them is part of the statistical model, not only an API choice. +Each component expresses a different assumption about how information is shared across signals and over time. ## Independent ascertainment rates @@ -206,7 +212,127 @@ Unlike `JointAscertainment`, this component does not automatically constrain the If both rates should be $<1$ (for instance because both are probabilities), make sure to choose priors that make $\alpha_{\mathrm{base}}\rho > 1$ negligibly likely. The component does not clip an invalid linked rate, because clipping would change the specified statistical model. -## Building a model an ascertainment component +## Related baselines with time-varying rates + +`TimeVaryingAscertainment` combines a scalar baseline model with one temporal process per signal. +If $b_s$ is the sampled baseline for signal $s$ and $z_s(t)$ is its temporal deviation, the rate trajectory is + +$$ + \alpha_s(t) + = + \operatorname{logit}^{-1} + \left(\operatorname{logit}(b_s) + z_s(t)\right). +$$ + +The baseline is the rate obtained when the deviation is zero. +It is not necessarily the first trajectory value: some temporal processes, including `AR1`, sample their initial state around zero. + +The baseline model relates only the scalar reference rates. +After those rates are sampled, each signal receives its own temporal process. +A joint covariance or a baseline ratio is therefore not reimposed at every timepoint. +This lets hospital and ED ascertainment trajectories move independently around related reference levels. + +The time-varying rate applies on the potential-observation date, after the infection-to-observation delay: + +```python +baseline_rates = baseline_model.sample() + +potential_counts = convolve_delay( + infections, + clinical_delay_pmf, +) + +rate_trajectory = sigmoid( + logit(baseline_rates[signal]) + + temporal_deviation[signal] +) + +predicted_counts = potential_counts * rate_trajectory +``` + +This is different from multiplying infections by a changing rate before convolution. +It represents ascertainment conditions on the day an event could be observed. + +### Joint scalar baselines + +```python +baseline = JointAscertainment( + name="he_baseline", + signals=("hospital", "ed_visits"), + baseline_rates=jnp.array([0.01, 0.02]), + scale_tril=jnp.array([[0.3, 0.0], [0.15, 0.26]]), +) + +ascertainment = TimeVaryingAscertainment( + name="he_ascertainment", + baseline_model=baseline, + processes={ + "hospital": RandomWalk( + DeterministicVariable("hospital_ascertainment_sd", 0.05) + ), + "ed_visits": RandomWalk( + DeterministicVariable("ed_ascertainment_sd", 0.05) + ), + }, +) +``` + +### Ratio-linked scalar baselines + +```python +baseline = RatioLinkedAscertainment( + name="he_baseline", + base_signal="ed_visits", + linked_signal="hospital", + base_rate_rv=DistributionalVariable("iedr", dist.Beta(1, 100)), + ratio_rv=DistributionalVariable( + "ihr_rel_iedr", + dist.LogNormal(jnp.log(0.5), 0.35), + ), +) + +ascertainment = TimeVaryingAscertainment( + name="he_ascertainment", + baseline_model=baseline, + processes={ + "ed_visits": RandomWalk( + DeterministicVariable("ed_ascertainment_sd", 0.05) + ), + "hospital": RandomWalk( + DeterministicVariable("hospital_ascertainment_sd", 0.05) + ), + }, +) +``` + +Register only the outer model, because it owns and samples its baseline. +Use the outer accessors in observation processes: + +```python +builder.add_ascertainment(ascertainment) + +hospital_obs = PopulationCounts( + name="hospital", + ascertainment_rate_rv=ascertainment.for_signal("hospital"), + ... +) +``` + +Daily temporal processes such as `RandomWalk` do not require a date. +Calendar-aligned processes do. +For example, wrap a process in `WeeklyTemporalProcess` when deviations should be shared within calendar weeks: + +```python +weekly_hospital_process = WeeklyTemporalProcess( + RandomWalk(DeterministicVariable("hospital_weekly_sd", 0.05)), + start_dow=MMWR_WEEK, +) +``` + +When any registered ascertainment model contains a weekly process, supply `obs_start_date` when sampling or fitting the model. +PyRenew uses it to align weekly blocks to the padded shared model axis. + +## Building a model with an ascertainment component The complete builder pattern is the same for either ascertainment component - here we use a `RatioLinkedAscertainment` but the sequence of definitions is the same for `JointAscertainment`. @@ -257,9 +383,9 @@ model = builder.build() ## Shared assumptions and limitations -Both ascertainment components produce one scalar ascertainment rate for each signal over the entire model period. -Consequently, neither component allows the rates or their implied relationship to change over time. -Whether this assumption is reasonable should be evaluated in the context of the observation processes and study period. +`JointAscertainment` and `RatioLinkedAscertainment` each produce one scalar rate per signal over the entire model period. +`TimeVaryingAscertainment` adds signal-specific deviations to scalar rates from either model. +Its baseline relationship is not a claim that the same covariance or ratio holds pointwise. The choice of ascertainment model should consider: @@ -267,8 +393,8 @@ The choice of ascertainment model should consider: - whether a covariance structure or a relative multiplier is easier to justify - whether the priors keep all rates, including derived linked rates, within scientifically plausible ranges - whether treating ascertainment rates as constant over the model period is scientifically reasonable +- whether daily, stepwise, or calendar-week deviations match the scientific question Ascertainment rates are generally difficult to distinguish from the scale of latent infections using surveillance counts alone. -Allowing the rates to vary over time introduces additional confounding unless external information or strong structural assumptions identify that variation. -The built-in ascertainment components therefore use scalar rates over the model period. -When this assumption is questionable, shorter modeling windows and sensitivity analyses may be more appropriate than introducing a weakly identified time-varying process. +Flexible infection and ascertainment trajectories can be weakly identified without informative priors, external information, or strong structural assumptions. +Use time variation only when the available information can support it, and assess sensitivity to both baseline and temporal-process priors. diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 6d5139e1e..4eb2fbbdd 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -454,8 +454,8 @@ class Wastewater(MeasurementObservation): log10_genome = self.log10_genome_per_infection_rv() def convolve_site(site_infections): - convolved, _ = self._convolve_with_alignment( - site_infections, shedding_pmf, p_observed=1.0 + convolved, _ = self._convolve_delay_with_alignment( + site_infections, shedding_pmf ) return convolved diff --git a/docs/tutorials/observation_processes_counts.qmd b/docs/tutorials/observation_processes_counts.qmd index 392e79fc3..32e2faa9c 100644 --- a/docs/tutorials/observation_processes_counts.qmd +++ b/docs/tutorials/observation_processes_counts.qmd @@ -56,16 +56,19 @@ Observed data can be aggregated or available as subpopulation-level counts, whic ### Observation equation -The deterministic transformation is given by the observation equation: +First convolve infections with the delay distribution to calculate potential observations, then apply ascertainment on the potential-observation date: $$ - \mu(t) = \alpha \sum_{d=0}^{D} I(t-d)\, \pi_d + C(t) = \sum_{d=0}^{D} I(t-d)\, \pi_d, + \qquad + \mu(t) = \alpha(t) C(t) $$ where: - $I(t-d)$ is the number of incident (new) infections on day $t-d$ -- $\alpha$ is the **ascertainment rate**, the probability that an infection results in an observed event (e.g., hospitalization) +- $C(t)$ is the number of potential observations on day $t$ +- $\alpha(t)$ is the **ascertainment rate** on day $t$, the probability that a potential event is observed (e.g., a hospitalization is recorded) - $\pi_d$ is the delay distribution from infection to observation, conditional on an infection leading to an observed event The delay distribution has finite support over lags $d = 0, \dots, D$ and satisfies @@ -75,6 +78,30 @@ $$ $$ This equation is a discrete convolution: observations at time $t$ arise from infections that occurred $d$ days earlier, weighted by the delay distribution. +A scalar ascertainment rate is the special case where $\alpha(t)=\alpha$ for the full model period. +With a scalar rate, multiplication before and after the convolution is numerically equivalent. +With a time-varying rate, PyRenew uses the ordering above; it does not multiply incident infections by the trajectory before convolution. + +For a time-varying component with related scalar baselines, the computation is: + +```python +baseline_rates = baseline_model.sample() + +potential_counts = convolve_delay( + infections, + clinical_delay_pmf, +) + +rate_trajectory = sigmoid( + logit(baseline_rates[signal]) + + temporal_deviation[signal] +) + +predicted_counts = potential_counts * rate_trajectory +``` + +Flexible infection and ascertainment trajectories can be weakly identified from count data alone. +Use informative priors or external information when estimating both over time. ### Stochastic observation model @@ -306,22 +333,19 @@ This gives the predicted number of observations $\mu(t)$. ```{python} #| label: predicted-no-noise -# Compute predicted admissions (convolution only, no observation noise) -from pyrenew.convolve import compute_delay_ascertained_incidence +# Compute predicted admissions (convolution and ascertainment, no noise) +from pyrenew.convolve import compute_delay_incidence -# Scale infections by IHR (ascertainment rate) -infections_scaled = infections * float(ihr_rv.sample()) - -predicted_admissions, offset = compute_delay_ascertained_incidence( - p_observed_given_incident=1.0, - latent_incidence=infections_scaled, +potential_admissions, offset = compute_delay_incidence( + latent_incidence=infections, delay_incidence_to_observation_pmf=hosp_delay_pmf, pad=True, ) +predicted_admissions = potential_admissions * float(ihr_rv.sample()) ``` -*Note:* in the above implementation, the ascertainment rate is applied by scaling infections before convolution. -This is equivalent to applying $\alpha$ after the convolution in the observation equation. +The rate is applied after convolution. +This ordering works for both scalar and time-varying ascertainment and makes the rate correspond to the potential-observation date. ```{python} #| label: plot-predicted-no-noise @@ -500,16 +524,15 @@ ihr_values = [0.005, 0.025] peak_value = 3000 # Peak infections infections_decay = peak_value * jnp.exp(-jnp.arange(n_days) / 20.0) -# Compute predicted hospital admissions (no noise) for each IHR +# Compute potential hospital admissions once, then apply each scalar IHR +potential_hosp, _ = compute_delay_incidence( + latent_incidence=infections_decay, + delay_incidence_to_observation_pmf=hosp_delay_pmf, + pad=True, +) results_list = [] for ihr_val in ihr_values: - infections_scaled = infections_decay * ihr_val - predicted_hosp, _ = compute_delay_ascertained_incidence( - p_observed_given_incident=1.0, - latent_incidence=infections_scaled, - delay_incidence_to_observation_pmf=hosp_delay_pmf, - pad=True, - ) + predicted_hosp = potential_hosp * ihr_val for i, admit in enumerate(predicted_hosp[day_one:]): results_list.append( diff --git a/docs/tutorials/observation_processes_measurements.qmd b/docs/tutorials/observation_processes_measurements.qmd index a8316b8ca..1f6f3716c 100644 --- a/docs/tutorials/observation_processes_measurements.qmd +++ b/docs/tutorials/observation_processes_measurements.qmd @@ -299,8 +299,8 @@ class Wastewater(MeasurementObservation): # Convolve each site's infections with shedding kinetics def convolve_site(site_infections): - convolved, _ = self._convolve_with_alignment( - site_infections, shedding_pmf, p_observed=1.0 + convolved, _ = self._convolve_delay_with_alignment( + site_infections, shedding_pmf ) return convolved diff --git a/pyrenew/ascertainment/__init__.py b/pyrenew/ascertainment/__init__.py index cc6160600..ad8a488bd 100644 --- a/pyrenew/ascertainment/__init__.py +++ b/pyrenew/ascertainment/__init__.py @@ -6,10 +6,12 @@ from pyrenew.ascertainment.base import AscertainmentModel, AscertainmentSignal from pyrenew.ascertainment.joint import JointAscertainment from pyrenew.ascertainment.linked import RatioLinkedAscertainment +from pyrenew.ascertainment.timevarying import TimeVaryingAscertainment __all__ = [ "AscertainmentModel", "AscertainmentSignal", "JointAscertainment", "RatioLinkedAscertainment", + "TimeVaryingAscertainment", ] diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index 9c0c823f2..de1f6d71e 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -8,6 +8,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import Mapping +import jax.numpy as jnp from jax.typing import ArrayLike from pyrenew.ascertainment.context import get_ascertainment_value @@ -168,6 +169,65 @@ def for_signal(self, signal_name: str) -> AscertainmentSignal: signal_name=signal_name, ) + def requires_calendar_anchor(self) -> bool: + """Return whether sampling requires the model-axis day-of-week. + + Returns + ------- + bool + ``True`` when the model needs a calendar anchor; otherwise + ``False``. + """ + return False + + def validate_sampled_values( + self, + values: Mapping[str, ArrayLike], + n_timepoints: int, + ) -> None: + """Validate sampled signal names and output shapes. + + Parameters + ---------- + values + Sampled ascertainment values keyed by signal name. + n_timepoints + Length of the shared model time axis. + + Raises + ------ + TypeError + If ``values`` is not a mapping. + ValueError + If signal names do not match this model or a value is neither a + scalar nor a full-axis trajectory. + """ + if not isinstance(values, Mapping): + raise TypeError( + f"Ascertainment model {self.name!r} must return a mapping, " + f"got {type(values).__name__}." + ) + + actual_signals = set(values) + expected_signals = set(self.signals) + if actual_signals != expected_signals: + missing = tuple(signal for signal in self.signals if signal not in values) + extra = tuple(signal for signal in values if signal not in expected_signals) + raise ValueError( + f"Ascertainment model {self.name!r} must return exactly signals " + f"{self.signals}. Missing: {missing}. Extra: {extra}." + ) + + allowed_shapes = ((), (n_timepoints,)) + for signal in self.signals: + actual_shape = jnp.asarray(values[signal]).shape + if actual_shape not in allowed_shapes: + raise ValueError( + f"Ascertainment model {self.name!r}, signal {signal!r}, " + f"returned shape {actual_shape}; allowed shapes are () and " + f"({n_timepoints},)." + ) + @abstractmethod def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: """ diff --git a/pyrenew/ascertainment/timevarying.py b/pyrenew/ascertainment/timevarying.py new file mode 100644 index 000000000..ee738b4d0 --- /dev/null +++ b/pyrenew/ascertainment/timevarying.py @@ -0,0 +1,207 @@ +# numpydoc ignore=GL08 +""" +Time-varying ascertainment models. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import jax.nn as jnn +import jax.numpy as jnp +import numpyro +from jax.scipy.special import logit +from jax.typing import ArrayLike +from numpyro.util import not_jax_tracer + +from pyrenew.ascertainment.base import AscertainmentModel +from pyrenew.latent import TemporalProcess + + +class TimeVaryingAscertainment(AscertainmentModel): + """ + Add independent temporal deviations to related ascertainment baselines. + + The wrapped baseline model samples one scalar rate per signal. Each scalar + is converted to the logit scale, combined with a signal-specific temporal + deviation, and returned as a probability trajectory on the shared model + time axis. Baseline relationships therefore describe the reference rates, + not a pointwise constraint on the resulting trajectories. + + Register only this outer model with ``PyrenewBuilder``. Use accessors from + this model, such as ``ascertainment.for_signal("hospital")``, in count + observation processes. + """ + + def __init__( + self, + name: str, + baseline_model: AscertainmentModel, + processes: Mapping[str, TemporalProcess], + ) -> None: + """ + Initialize time-varying ascertainment. + + Parameters + ---------- + name + Name of the outer ascertainment model. It must differ from the + wrapped baseline model's name. + baseline_model + Ascertainment model that samples one scalar reference rate for + each signal. + processes + One temporal process per baseline signal. Each process generates + an independent logit-scale deviation trajectory. + + Raises + ------ + TypeError + If the baseline or temporal processes have invalid types. + ValueError + If names conflict or process signal names do not match the + baseline signals. + """ + if not isinstance(baseline_model, AscertainmentModel): + raise TypeError( + "baseline_model must be an AscertainmentModel, " + f"got {type(baseline_model).__name__}." + ) + if not isinstance(processes, Mapping): + raise TypeError( + f"processes must be a mapping, got {type(processes).__name__}." + ) + if name == baseline_model.name: + raise ValueError( + "Time-varying ascertainment and its baseline model must have " + f"different names; both were {name!r}." + ) + + super().__init__(name=name, signals=baseline_model.signals) + expected_signals = set(self.signals) + process_signals = set(processes) + if process_signals != expected_signals: + missing = tuple( + signal for signal in self.signals if signal not in processes + ) + extra = tuple( + signal for signal in processes if signal not in expected_signals + ) + raise ValueError( + "processes must contain exactly the baseline model signals " + f"{self.signals}. Missing: {missing}. Extra: {extra}." + ) + + ordered_processes: dict[str, TemporalProcess] = {} + for signal in self.signals: + process = processes[signal] + if not isinstance(process, TemporalProcess): + raise TypeError( + f"process for signal {signal!r} must satisfy the " + f"TemporalProcess protocol, got {type(process).__name__}." + ) + ordered_processes[signal] = process + + self.baseline_model = baseline_model + self.processes = ordered_processes + + def requires_calendar_anchor(self) -> bool: + """Return whether the baseline or any temporal process needs a date. + + Returns + ------- + bool + ``True`` when sampling needs the model-axis day-of-week; otherwise + ``False``. + """ + return self.baseline_model.requires_calendar_anchor() or any( + getattr(process, "requires_calendar_anchor", False) + for process in self.processes.values() + ) + + def sample( + self, + n_timepoints: int, + first_day_dow: int | None = None, + **kwargs: object, + ) -> Mapping[str, ArrayLike]: + """ + Sample full-axis ascertainment-rate trajectories. + + Parameters + ---------- + n_timepoints + Number of timepoints on the shared model axis. + first_day_dow + Day-of-week for the first model-axis timepoint. Required by + calendar-aligned temporal processes. + **kwargs + Additional model context, ignored. + + Returns + ------- + Mapping[str, ArrayLike] + One probability trajectory of shape ``(n_timepoints,)`` per + signal. + + Raises + ------ + ValueError + If the baseline does not return scalar probabilities in ``(0, 1)`` + or a temporal process returns an unexpected shape. + """ + baseline_values = self.baseline_model.sample( + n_timepoints=n_timepoints, + first_day_dow=first_day_dow, + ) + self.baseline_model.validate_sampled_values( + baseline_values, + n_timepoints=n_timepoints, + ) + + scalar_baselines: dict[str, ArrayLike] = {} + for signal in self.signals: + baseline_value = jnp.asarray(baseline_values[signal]) + if baseline_value.shape != (): + raise ValueError( + f"Time-varying ascertainment model {self.name!r} requires " + f"a scalar baseline for signal {signal!r}; got shape " + f"{baseline_value.shape}, required shape ()." + ) + invalid_baseline = ( + ~jnp.isfinite(baseline_value) + | (baseline_value <= 0) + | (baseline_value >= 1) + ) + if not_jax_tracer(invalid_baseline) and bool(invalid_baseline): + raise ValueError( + f"Time-varying ascertainment model {self.name!r} requires " + f"the baseline for signal {signal!r} to be finite and " + f"strictly inside (0, 1); got {baseline_value}." + ) + scalar_baselines[signal] = baseline_value + + result: dict[str, ArrayLike] = {} + for signal in self.signals: + deviation = jnp.asarray( + self.processes[signal].sample( + n_timepoints=n_timepoints, + initial_value=0.0, + n_processes=1, + name_prefix=f"{self.name}_{signal}", + first_day_dow=first_day_dow, + ) + ) + required_shape = (n_timepoints, 1) + if deviation.shape != required_shape: + raise ValueError( + f"Time-varying ascertainment model {self.name!r}, signal " + f"{signal!r}, received temporal-process shape " + f"{deviation.shape}; required shape is {required_shape}." + ) + deviation = jnp.squeeze(deviation, axis=-1) + rate = jnn.sigmoid(logit(scalar_baselines[signal]) + deviation) + numpyro.deterministic(f"{self.name}_{signal}", rate) + result[signal] = rate + + return result diff --git a/pyrenew/convolve.py b/pyrenew/convolve.py index 0a9129f93..10d12999f 100755 --- a/pyrenew/convolve.py +++ b/pyrenew/convolve.py @@ -243,6 +243,37 @@ def compute_delay_ascertained_incidence( return (delay_obs_incidence, offset) +def compute_delay_incidence( + latent_incidence: ArrayLike, + delay_incidence_to_observation_pmf: ArrayLike, + pad: bool = False, +) -> tuple[ArrayLike, int]: + """Convolve incidence with a delay PMF without ascertainment. + + Parameters + ---------- + latent_incidence + Incidence values generated by the latent process. + delay_incidence_to_observation_pmf + Probability mass function for the delay from incidence to potential + observation. + pad + Whether to prepend ``NaN`` values so the result remains aligned with + the input time axis. + + Returns + ------- + tuple[ArrayLike, int] + Delayed potential observations and their offset from the input axis. + """ + return compute_delay_ascertained_incidence( + latent_incidence=latent_incidence, + delay_incidence_to_observation_pmf=delay_incidence_to_observation_pmf, + p_observed_given_incident=1, + pad=pad, + ) + + def compute_prop_already_reported( reporting_delay_pmf: ArrayLike, n_timepoints: int, diff --git a/pyrenew/model/multisignal_model.py b/pyrenew/model/multisignal_model.py index cd8efa144..e20e95d5a 100644 --- a/pyrenew/model/multisignal_model.py +++ b/pyrenew/model/multisignal_model.py @@ -7,6 +7,7 @@ from __future__ import annotations import datetime as dt +from collections.abc import Mapping import jax.numpy as jnp import numpy as np @@ -14,7 +15,7 @@ import numpyro.handlers from jax.typing import ArrayLike -from pyrenew.ascertainment import AscertainmentModel +from pyrenew.ascertainment import AscertainmentModel, TimeVaryingAscertainment from pyrenew.ascertainment.context import ascertainment_context from pyrenew.latent.base import BaseLatentInfectionProcess from pyrenew.metaclass import Model @@ -124,6 +125,39 @@ def validate(self) -> None: f"the model name {ascertainment_model.name!r}." ) + namespace_owners: dict[str, list[str]] = {} + for name, ascertainment_model in self.ascertainment_models.items(): + namespace_owners.setdefault(name, []).append( + f"registered ascertainment model {name!r}" + ) + if not isinstance(ascertainment_model, TimeVaryingAscertainment): + continue + + baseline = ascertainment_model.baseline_model + registered_baseline_names = [ + registered_name + for registered_name, registered_model in self.ascertainment_models.items() + if registered_model is baseline + ] + if registered_baseline_names: + raise ValueError( + f"Time-varying ascertainment model {name!r} owns baseline " + f"model {baseline.name!r}; that baseline object must not " + f"also be registered as {registered_baseline_names[0]!r}." + ) + namespace_owners.setdefault(baseline.name, []).append( + f"baseline model {baseline.name!r} owned by time-varying " + f"ascertainment model {name!r}" + ) + + for namespace, owners in namespace_owners.items(): + if len(owners) > 1: + raise ValueError( + f"Ascertainment NumPyro namespace {namespace!r} is used by " + f"{' and '.join(owners)}. Registered models and wrapped " + "baselines must have unique names." + ) + def _validate_run_args( self, n_days_post_init: int, @@ -240,6 +274,13 @@ def _check_obs_start_date( f"obs_start_date is required when any observation uses " f"a day-of-week effect; observation '{name}' does." ) + for name, ascertainment_model in self.ascertainment_models.items(): + if ascertainment_model.requires_calendar_anchor(): + raise ValueError( + "obs_start_date is required when a registered ascertainment " + f"model uses a calendar-aligned temporal process; " + f"ascertainment model {name!r} does." + ) if self.latent.requires_calendar_anchor(): raise ValueError( "obs_start_date is required when the latent process uses a " @@ -299,9 +340,10 @@ def validate_data( Date of the first observation day. Required when any observation uses ``aggregation="weekly"`` or a day-of-week effect, or when a latent temporal process uses - calendar-week alignment. Converted once to the axis-origin - ``first_day_dow`` and forwarded to the latent process and - every observation. + calendar-week alignment, including a registered ascertainment + model. Converted once to the axis-origin ``first_day_dow`` and + forwarded to the latent process, ascertainment models, and every + observation. **observation_data Data for each observation process, keyed by observation name. Each value should be a dict of kwargs for that observation's sample(). @@ -344,6 +386,40 @@ def validate_data( **obs_data, ) + def _sample_ascertainment_models( + self, + n_timepoints: int, + first_day_dow: int | None, + ) -> dict[str, Mapping[str, ArrayLike]]: + """Sample and validate every registered ascertainment model. + + Parameters + ---------- + n_timepoints + Length of the padded shared model axis. + first_day_dow + Day-of-week for the first point on that axis, or ``None`` when no + calendar anchor was supplied. + + Returns + ------- + dict[str, Mapping[str, ArrayLike]] + Sampled ascertainment values keyed first by registered model name + and then by signal name. + """ + sampled_values: dict[str, Mapping[str, ArrayLike]] = {} + for name, ascertainment_model in self.ascertainment_models.items(): + values = ascertainment_model.sample( + n_timepoints=n_timepoints, + first_day_dow=first_day_dow, + ) + ascertainment_model.validate_sampled_values( + values, + n_timepoints=n_timepoints, + ) + sampled_values[name] = values + return sampled_values + def sample( self, n_days_post_init: int, @@ -374,7 +450,8 @@ def sample( to the latent process and every observation. Required when any observation uses ``aggregation="weekly"`` or a day-of-week effect, or when a latent temporal process uses - calendar-week alignment. + calendar-week alignment, including a registered ascertainment + model. **observation_data Data for each observation process, keyed by observation name (the ``name`` attribute of each observation process). @@ -412,10 +489,11 @@ def sample( "subpop": inf_all, } - ascertainment_values = { - name: ascertainment_model.sample() - for name, ascertainment_model in self.ascertainment_models.items() - } + n_timepoints = self.latent.n_initialization_points + n_days_post_init + ascertainment_values = self._sample_ascertainment_models( + n_timepoints=n_timepoints, + first_day_dow=first_day_dow, + ) with ascertainment_context(ascertainment_values): # Apply each observation process diff --git a/pyrenew/model/pyrenew_builder.py b/pyrenew/model/pyrenew_builder.py index 438013739..9a5ccef62 100644 --- a/pyrenew/model/pyrenew_builder.py +++ b/pyrenew/model/pyrenew_builder.py @@ -162,8 +162,10 @@ def add_ascertainment( process as ordinary ``RandomVariable`` objects. A registered ascertainment model is sampled once per model execution, - before observation processes run. Observation processes receive - signal-specific accessors from ``ascertainment_model.for_signal(...)``: + before observation processes run. It may return a scalar rate for each + signal or a trajectory spanning the full padded model axis. + Observation processes receive signal-specific accessors from + ``ascertainment_model.for_signal(...)``: ```python ascertainment = JointAscertainment( @@ -183,6 +185,10 @@ def add_ascertainment( ) ``` + ``TimeVaryingAscertainment`` owns and samples its scalar baseline + model. Register only the outer time-varying model, and use that outer + model's ``for_signal(...)`` accessors in observation processes. + The ascertainment model's ``name`` attribute is used as the unique identifier in the built ``MultiSignalModel``. @@ -190,7 +196,7 @@ def add_ascertainment( ---------- ascertainment_model Configured ascertainment model instance, such as - ``JointAscertainment``. + ``JointAscertainment`` or ``TimeVaryingAscertainment``. Returns ------- diff --git a/pyrenew/observation/base.py b/pyrenew/observation/base.py index 6fdfab166..c72097483 100644 --- a/pyrenew/observation/base.py +++ b/pyrenew/observation/base.py @@ -15,7 +15,7 @@ from jax.typing import ArrayLike from pyrenew.arrayutils import require_shape -from pyrenew.convolve import compute_delay_ascertained_incidence +from pyrenew.convolve import compute_delay_incidence from pyrenew.metaclass import RandomVariable from pyrenew.time import validate_dow @@ -51,7 +51,7 @@ class BaseObservationProcess(RandomVariable): See Also -------- - pyrenew.convolve.compute_delay_ascertained_incidence : + pyrenew.convolve.compute_delay_incidence : Underlying convolution function pyrenew.metaclass.RandomVariable : Base class for all random variables @@ -270,16 +270,14 @@ def _compute_period_offset( raise ValueError("first_day_dow is required when start_dow is not None") return (start_dow - first_day_dow) % 7 - def _convolve_with_alignment( + def _convolve_delay_with_alignment( self, latent_incidence: ArrayLike, pmf: ArrayLike, - p_observed: float = 1.0, ) -> tuple[ArrayLike, int]: - """ - Convolve latent incidence with PMF while maintaining timeline alignment. + """Apply delay-only convolution while preserving timeline alignment. - This is a wrapper around ``compute_delay_ascertained_incidence`` that + This is a wrapper around ``compute_delay_incidence`` that always uses ``pad=True`` to ensure day t in the output corresponds to day t in the input. The first ``len(pmf) - 1`` days will be NaN. @@ -290,8 +288,6 @@ def _convolve_with_alignment( Shape: (n_days,) pmf Delay or shedding PMF. Shape: (n_pmf,) - p_observed - Observation probability multiplier. Scales the convolution result. Returns ------- @@ -309,13 +305,12 @@ def _convolve_with_alignment( See Also -------- - pyrenew.convolve.compute_delay_ascertained_incidence : + pyrenew.convolve.compute_delay_incidence : Underlying function """ - return compute_delay_ascertained_incidence( + return compute_delay_incidence( latent_incidence=latent_incidence, delay_incidence_to_observation_pmf=pmf, - p_observed_given_incident=p_observed, pad=True, # Maintains timeline alignment ) diff --git a/pyrenew/observation/count_observations.py b/pyrenew/observation/count_observations.py index 3aaa7a342..31871df77 100644 --- a/pyrenew/observation/count_observations.py +++ b/pyrenew/observation/count_observations.py @@ -195,7 +195,7 @@ def _predicted_obs( infections: ArrayLike, ) -> ArrayLike: """ - Compute predicted counts via ascertainment x delay convolution. + Compute predicted counts using delay then ascertainment. Parameters ---------- @@ -212,19 +212,74 @@ def _predicted_obs( First len(delay_pmf)-1 days are NaN. """ delay_pmf = self.temporal_pmf_rv() - ascertainment_rate = self.ascertainment_rate_rv() if infections.ndim == 1: - return self._convolve_with_alignment( - infections, delay_pmf, ascertainment_rate + potential_counts = self._convolve_delay_with_alignment( + infections, + delay_pmf, )[0] - return jax.vmap( - lambda col: self._convolve_with_alignment( - col, delay_pmf, ascertainment_rate - )[0], - in_axes=1, - out_axes=1, - )(infections) + else: + potential_counts = jax.vmap( + lambda col: self._convolve_delay_with_alignment( + col, + delay_pmf, + )[0], + in_axes=1, + out_axes=1, + )(infections) + + ascertainment_rate = self.ascertainment_rate_rv() + return self._apply_ascertainment(potential_counts, ascertainment_rate) + + def _apply_ascertainment( + self, + potential_counts: ArrayLike, + ascertainment_rate: ArrayLike, + ) -> ArrayLike: + """Scale potential counts by scalar or shared time-varying ascertainment. + + Parameters + ---------- + potential_counts + Delayed potential counts with shape ``(n_timepoints,)`` or + ``(n_timepoints, n_subpopulations)``. + ascertainment_rate + A scalar rate or one shared rate per timepoint. + + Returns + ------- + ArrayLike + Ascertainment-scaled counts with the same shape as + ``potential_counts``. + + Raises + ------ + ValueError + If counts are not one- or two-dimensional or the rate is neither a + scalar nor an exact full-axis trajectory. + """ + potential_counts_array = jnp.asarray(potential_counts) + if potential_counts_array.ndim not in (1, 2): + raise ValueError( + "potential_counts must have shape (n_timepoints,) or " + "(n_timepoints, n_subpopulations); got shape " + f"{potential_counts_array.shape}." + ) + + rate = jnp.asarray(ascertainment_rate) + expected_shape = (potential_counts_array.shape[0],) + if rate.shape not in ((), expected_shape): + raise ValueError( + f"ascertainment_rate must have shape () or {expected_shape}; " + f"got shape {rate.shape}." + ) + if rate.shape == expected_shape and potential_counts_array.ndim == 2: + rate = rate[:, None] + + finite_counts = ~jnp.isnan(potential_counts_array) + safe_counts = jnp.where(finite_counts, potential_counts_array, 0.0) + scaled_counts = safe_counts * rate + return jnp.where(finite_counts, scaled_counts, potential_counts_array) def _apply_right_truncation( self, diff --git a/test/test_ascertainment.py b/test/test_ascertainment.py index 6cf77cd2b..8019f69f3 100644 --- a/test/test_ascertainment.py +++ b/test/test_ascertainment.py @@ -2,24 +2,405 @@ Tests for ascertainment models. """ +from collections.abc import Mapping + import jax.numpy as jnp import numpyro import numpyro.distributions as dist import pytest +from jax.scipy.special import logit +from jax.typing import ArrayLike from pyrenew.ascertainment import ( + AscertainmentModel, AscertainmentSignal, JointAscertainment, RatioLinkedAscertainment, + TimeVaryingAscertainment, ) from pyrenew.ascertainment.context import ( ascertainment_context, get_ascertainment_value, ) from pyrenew.deterministic import DeterministicVariable +from pyrenew.latent import WeeklyTemporalProcess from pyrenew.randomvariable import DistributionalVariable +class DeterministicTemporalProcess: + """Generate a predictable two-dimensional temporal trajectory for tests.""" + + step_size = 1 + + def __init__( + self, + initial: float = 0.0, + increment: float = 0.0, + result: ArrayLike | None = None, + ) -> None: + """Initialize a deterministic temporal-process test double.""" + self.initial = initial + self.increment = increment + self.result = result + self.sample_calls: list[dict[str, object]] = [] + + def sample( + self, + n_timepoints: int, + initial_value: float | ArrayLike | None = None, + n_processes: int = 1, + name_prefix: str = "temporal", + *, + first_day_dow: int | None = None, + ) -> ArrayLike: + """Return a deterministic trajectory with the requested shape. + + Returns + ------- + ArrayLike + Configured result or generated two-dimensional trajectory. + """ + self.sample_calls.append( + { + "n_timepoints": n_timepoints, + "initial_value": initial_value, + "n_processes": n_processes, + "name_prefix": name_prefix, + "first_day_dow": first_day_dow, + } + ) + if self.result is not None: + return jnp.asarray(self.result) + values = self.initial + self.increment * jnp.arange(n_timepoints) + return jnp.broadcast_to(values[:, None], (n_timepoints, n_processes)) + + +class FixedAscertainmentModel(AscertainmentModel): + """Return fixed ascertainment values for baseline-contract tests.""" + + def __init__( + self, + name: str, + values: Mapping[str, ArrayLike], + requires_calendar_anchor: bool = False, + ) -> None: + """Initialize a fixed ascertainment test double.""" + super().__init__(name=name, signals=tuple(values)) + self.values = dict(values) + self._requires_calendar_anchor = requires_calendar_anchor + self.n_sample_calls = 0 + + def requires_calendar_anchor(self) -> bool: + """Return the configured calendar requirement. + + Returns + ------- + bool + Configured calendar requirement. + """ + return self._requires_calendar_anchor + + def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: + """Return configured values and record the sampling call. + + Returns + ------- + Mapping[str, ArrayLike] + Configured signal values. + """ + self.n_sample_calls += 1 + return self.values + + +class TestAscertainmentModelContract: + """Test validation shared by all ascertainment models.""" + + def test_accepts_scalar_and_full_axis_values(self) -> None: + """Test scalar and exact full-axis trajectories are valid.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2, "ed": 0.3}, + ) + + model.validate_sampled_values( + {"hospital": jnp.array(0.2), "ed": jnp.arange(4)}, + n_timepoints=4, + ) + + def test_rejects_non_mapping_values(self) -> None: + """Test sampled values must be returned as a mapping.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + with pytest.raises(TypeError, match="ascertainment.*mapping"): + model.validate_sampled_values([0.2], n_timepoints=4) # type: ignore[arg-type] + + @pytest.mark.parametrize( + "values", + [ + {}, + {"hospital": 0.2, "ed": 0.3}, + ], + ) + def test_rejects_missing_or_extra_signals( + self, + values: Mapping[str, ArrayLike], + ) -> None: + """Test sampled signal names must exactly match the model signals.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + with pytest.raises(ValueError, match="ascertainment.*signals"): + model.validate_sampled_values(values, n_timepoints=4) + + @pytest.mark.parametrize( + "value", + [ + jnp.ones(1), + jnp.ones(3), + jnp.ones((4, 1)), + jnp.ones((4, 1, 1)), + ], + ) + def test_rejects_invalid_sample_shapes(self, value: ArrayLike) -> None: + """Test only scalar or exact full-axis values are accepted.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + with pytest.raises( + ValueError, + match=r"ascertainment.*hospital.*shape.*\(\).*\(4,\)", + ): + model.validate_sampled_values({"hospital": value}, n_timepoints=4) + + def test_scalar_models_do_not_require_a_calendar_anchor(self) -> None: + """Test the base calendar declaration defaults to false.""" + model = JointAscertainment( + name="ascertainment", + signals=("hospital",), + baseline_rates=jnp.array([0.2]), + scale_tril=jnp.eye(1), + ) + + assert not model.requires_calendar_anchor() + + +class TestTimeVaryingAscertainmentConstruction: + """Test time-varying ascertainment constructor validation.""" + + def test_requires_an_ascertainment_baseline(self) -> None: + """Test the baseline must implement the ascertainment contract.""" + with pytest.raises(TypeError, match="baseline_model"): + TimeVaryingAscertainment( + name="outer", + baseline_model=object(), # type: ignore[arg-type] + processes={"hospital": DeterministicTemporalProcess()}, + ) + + def test_requires_distinct_outer_and_baseline_names(self) -> None: + """Test outer and baseline NumPyro namespaces must differ.""" + baseline = FixedAscertainmentModel("same", {"hospital": 0.2}) + + with pytest.raises(ValueError, match="different names"): + TimeVaryingAscertainment( + name="same", + baseline_model=baseline, + processes={"hospital": DeterministicTemporalProcess()}, + ) + + def test_requires_a_process_mapping(self) -> None: + """Test temporal processes must be provided in a mapping.""" + baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + + with pytest.raises(TypeError, match="processes must be a mapping"): + TimeVaryingAscertainment( + name="outer", + baseline_model=baseline, + processes=[], # type: ignore[arg-type] + ) + + @pytest.mark.parametrize( + "processes", + [ + {}, + { + "hospital": DeterministicTemporalProcess(), + "ed": DeterministicTemporalProcess(), + }, + ], + ) + def test_requires_exact_baseline_signal_keys( + self, + processes: Mapping[str, DeterministicTemporalProcess], + ) -> None: + """Test process signal names must exactly match baseline signals.""" + baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + + with pytest.raises(ValueError, match="baseline model signals"): + TimeVaryingAscertainment("outer", baseline, processes) + + def test_requires_temporal_process_protocol(self) -> None: + """Test every configured process must satisfy TemporalProcess.""" + baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + + with pytest.raises(TypeError, match="TemporalProcess"): + TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": object()}, # type: ignore[dict-item] + ) + + def test_orders_processes_like_baseline_signals(self) -> None: + """Test deterministic process order follows the baseline model.""" + baseline = FixedAscertainmentModel( + "baseline", + {"hospital": 0.2, "ed": 0.3}, + ) + processes = { + "ed": DeterministicTemporalProcess(), + "hospital": DeterministicTemporalProcess(), + } + + model = TimeVaryingAscertainment("outer", baseline, processes) + + assert tuple(model.processes) == ("hospital", "ed") + + +class TestTimeVaryingAscertainmentSampling: + """Test baseline and temporal sampling semantics.""" + + def test_joint_baseline_sites_and_independent_trajectories(self) -> None: + """Test a joint baseline is sampled once before independent deviations.""" + baseline_rates = jnp.array([0.2, 0.4]) + baseline = JointAscertainment( + name="he_baseline", + signals=("hospital", "ed"), + baseline_rates=baseline_rates, + scale_tril=jnp.eye(2), + ) + model = TimeVaryingAscertainment( + name="he_ascertainment", + baseline_model=baseline, + processes={ + "hospital": DeterministicTemporalProcess(increment=0.3), + "ed": DeterministicTemporalProcess(increment=-0.2), + }, + ) + + with numpyro.handlers.substitute( + data={"he_baseline_eta": logit(baseline_rates)} + ): + with numpyro.handlers.trace() as trace: + values = model.sample(n_timepoints=4) + + assert trace["he_baseline_eta"]["type"] == "sample" + assert trace["he_baseline_eta"]["value"].shape == (2,) + assert trace["he_baseline_hospital"]["value"].shape == () + assert trace["he_baseline_ed"]["value"].shape == () + assert trace["he_ascertainment_hospital"]["value"].shape == (4,) + assert trace["he_ascertainment_ed"]["value"].shape == (4,) + assert jnp.allclose(values["hospital"][0], baseline_rates[0]) + assert jnp.allclose(values["ed"][0], baseline_rates[1]) + assert not jnp.allclose( + values["hospital"][1:] / values["ed"][1:], + baseline_rates[0] / baseline_rates[1], + ) + + def test_ratio_linked_baseline_relationship_is_not_pointwise(self) -> None: + """Test linked baseline sites persist without constraining trajectories.""" + baseline = RatioLinkedAscertainment( + name="he_baseline", + base_signal="ed", + linked_signal="hospital", + base_rate_rv=DistributionalVariable("iedr", dist.Delta(0.4)), + ratio_rv=DistributionalVariable("ihr_rel_iedr", dist.Delta(0.5)), + ) + model = TimeVaryingAscertainment( + name="he_ascertainment", + baseline_model=baseline, + processes={ + "ed": DeterministicTemporalProcess(increment=-0.1), + "hospital": DeterministicTemporalProcess(increment=0.2), + }, + ) + + with numpyro.handlers.seed(rng_seed=42): + with numpyro.handlers.trace() as trace: + values = model.sample(n_timepoints=4) + + assert trace["iedr"]["type"] == "sample" + assert trace["ihr_rel_iedr"]["type"] == "sample" + assert jnp.allclose(trace["he_baseline_ed"]["value"], 0.4) + assert jnp.allclose(trace["he_baseline_hospital"]["value"], 0.2) + assert values["ed"].shape == (4,) + assert values["hospital"].shape == (4,) + trajectory_ratios = values["hospital"] / values["ed"] + assert jnp.allclose(trajectory_ratios[0], 0.5) + assert not jnp.allclose(trajectory_ratios[1:], 0.5) + + def test_calendar_requirement_and_weekly_alignment(self) -> None: + """Test weekly processes declare and use the model-axis weekday.""" + baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + weekly = WeeklyTemporalProcess( + DeterministicTemporalProcess(increment=1.0), + start_dow=0, + ) + model = TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": weekly}, + ) + + values = model.sample(n_timepoints=10, first_day_dow=2) + + assert model.requires_calendar_anchor() + assert jnp.allclose(values["hospital"][:5], 0.2) + assert jnp.all(values["hospital"][5:] > 0.2) + + @pytest.mark.parametrize("baseline_value", [0.0, 1.0, -0.1, 1.1, jnp.nan]) + def test_rejects_invalid_concrete_baselines( + self, + baseline_value: float, + ) -> None: + """Test concrete baselines must be finite probabilities.""" + baseline = FixedAscertainmentModel( + "baseline", + {"hospital": baseline_value}, + ) + model = TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": DeterministicTemporalProcess()}, + ) + + with pytest.raises(ValueError, match=r"hospital.*\(0, 1\)"): + model.sample(n_timepoints=4) + + def test_rejects_trajectory_baseline(self) -> None: + """Test nested time-varying baselines are rejected.""" + baseline = FixedAscertainmentModel( + "baseline", + {"hospital": jnp.full(4, 0.2)}, + ) + model = TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": DeterministicTemporalProcess()}, + ) + + with pytest.raises(ValueError, match=r"scalar baseline.*\(4,\).*"): + model.sample(n_timepoints=4) + + def test_rejects_malformed_temporal_output_before_squeezing(self) -> None: + """Test temporal output must retain the singleton process dimension.""" + baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + model = TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": DeterministicTemporalProcess(result=jnp.zeros(4))}, + ) + + with pytest.raises(ValueError, match=r"outer.*hospital.*\(4,\).*\(4, 1\)"): + model.sample(n_timepoints=4) + + class TestJointAscertainmentValidation: """Test JointAscertainment constructor validation.""" diff --git a/test/test_convolve.py b/test/test_convolve.py index a9211f85f..977422518 100644 --- a/test/test_convolve.py +++ b/test/test_convolve.py @@ -407,6 +407,34 @@ def test_compute_delay_ascertained_incidence_manual( assert_array_equal(result_padded, expected_padded) +def test_compute_delay_incidence_delegates_without_ascertainment() -> None: + """Test delay-only convolution, offset, and aligned NaN padding.""" + latent_incidence = jnp.array([1.0, 2.0, 3.0]) + delay_pmf = jnp.array([0.25, 0.75]) + + result, offset = pc.compute_delay_incidence( + latent_incidence, + delay_pmf, + ) + padded, padded_offset = pc.compute_delay_incidence( + latent_incidence, + delay_pmf, + pad=True, + ) + existing, existing_offset = pc.compute_delay_ascertained_incidence( + latent_incidence, + delay_pmf, + p_observed_given_incident=1, + ) + + assert_array_equal(result, jnp.array([1.25, 2.25])) + assert offset == 1 + assert_array_equal(padded, jnp.array([jnp.nan, 1.25, 2.25])) + assert padded_offset == 0 + assert_array_equal(result, existing) + assert offset == existing_offset + + @pytest.mark.parametrize( ["reporting_delay_pmf", "n_timepoints", "right_truncation_offset", "expected"], [ diff --git a/test/test_helpers.py b/test/test_helpers.py index 4aa381086..3a7221365 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -7,6 +7,7 @@ import jax import jax.numpy as jnp +from jax.typing import ArrayLike from pyrenew.deterministic import DeterministicVariable from pyrenew.latent import AR1, DifferencedAR1, RandomWalk @@ -135,7 +136,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: """ Simple predicted signal: log(convolution * scale). @@ -149,8 +150,8 @@ def _predicted_obs(self, infections): if infections.ndim == 1: infections = infections[:, jnp.newaxis] - def convolve_col(col): # numpydoc ignore=GL08 - return self._convolve_with_alignment(col, pmf, 1.0)[0] + def convolve_col(col: ArrayLike) -> ArrayLike: # numpydoc ignore=GL08 + return self._convolve_delay_with_alignment(col, pmf)[0] predicted = jax.vmap(convolve_col, in_axes=1, out_axes=1)(infections) diff --git a/test/test_observation_counts.py b/test/test_observation_counts.py index 5fd5b6abb..ac87d7d0d 100644 --- a/test/test_observation_counts.py +++ b/test/test_observation_counts.py @@ -7,9 +7,12 @@ import numpyro import numpyro.distributions as dist import pytest +from jax.typing import ArrayLike from numpyro.infer.util import log_density +from pyrenew.convolve import compute_delay_incidence from pyrenew.deterministic import DeterministicPMF, DeterministicVariable +from pyrenew.metaclass import RandomVariable from pyrenew.observation import ( NegativeBinomialNoise, PoissonNoise, @@ -21,6 +24,27 @@ from test.test_helpers import create_mock_infections +class CountingVariable(RandomVariable): + """Return a fixed value while counting calls for sampling-order tests.""" + + def __init__(self, name: str, value: ArrayLike) -> None: + """Initialize a fixed counting random variable.""" + super().__init__(name=name) + self.value = value + self.n_calls = 0 + + def sample(self, **kwargs: object) -> ArrayLike: + """Return the fixed value and increment the call count. + + Returns + ------- + ArrayLike + Configured value. + """ + self.n_calls += 1 + return self.value + + class TestCountsBasics: """Test basic functionality of aggregated count observation process.""" @@ -145,6 +169,149 @@ def test_observation_passthrough(self, counts_process): assert jnp.allclose(result.observed, known_obs) +class TestTimeVaryingAscertainment: + """Test ascertainment applied on the potential-observation date.""" + + def test_delay_precedes_time_varying_ascertainment(self) -> None: + """Test a rate trajectory scales delayed rather than incident counts.""" + infections = jnp.array([0.0, 10.0, 30.0, 0.0, 0.0]) + delay_pmf = jnp.array([0.25, 0.75]) + rate = jnp.array([0.1, 0.2, 0.3, 0.4, 0.5]) + process = PopulationCounts( + name="hospital", + ascertainment_rate_rv=DeterministicVariable("rate", rate), + delay_distribution_rv=DeterministicPMF("delay", delay_pmf), + noise=PoissonNoise(), + ) + + predicted = process._predicted_obs(infections) + potential_counts, _ = compute_delay_incidence( + infections, + delay_pmf, + pad=True, + ) + incident_date_counts, _ = compute_delay_incidence( + infections * rate, + delay_pmf, + pad=True, + ) + + assert jnp.allclose(predicted, potential_counts * rate, equal_nan=True) + assert not jnp.allclose(predicted, incident_date_counts, equal_nan=True) + + def test_scalar_rate_is_sampled_once_and_matches_old_calculation(self) -> None: + """Test scalar ascertainment preserves prior results and samples once.""" + rate = CountingVariable("rate", 0.25) + infections = jnp.array([1.0, 2.0, 4.0, 8.0]) + delay_pmf = jnp.array([0.2, 0.8]) + process = PopulationCounts( + name="hospital", + ascertainment_rate_rv=rate, + delay_distribution_rv=DeterministicPMF("delay", delay_pmf), + noise=PoissonNoise(), + ) + + predicted = process._predicted_obs(infections) + old_calculation, _ = compute_delay_incidence( + infections * 0.25, + delay_pmf, + pad=True, + ) + + assert rate.n_calls == 1 + assert jnp.allclose(predicted, old_calculation, equal_nan=True) + + @pytest.mark.parametrize( + "rate", + [ + jnp.ones(1), + jnp.ones(3), + jnp.ones((4, 1)), + jnp.ones((4, 2)), + jnp.ones((4, 1, 1)), + ], + ) + def test_rejects_malformed_rate_trajectories(self, rate: ArrayLike) -> None: + """Test a trajectory must exactly match the shared time axis.""" + process = PopulationCounts( + name="hospital", + ascertainment_rate_rv=DeterministicVariable("rate", 0.2), + delay_distribution_rv=DeterministicPMF("delay", jnp.array([1.0])), + noise=PoissonNoise(), + ) + + with pytest.raises(ValueError, match=r"shape \(\) or \(4,\)"): + process._apply_ascertainment(jnp.ones(4), rate) + + def test_trajectory_is_shared_across_subpopulations(self) -> None: + """Test one full-axis trajectory broadcasts across all columns.""" + potential_counts = jnp.array( + [ + [1.0, 3.0], + [2.0, 6.0], + [4.0, 12.0], + ] + ) + rate = jnp.array([0.1, 0.2, 0.4]) + process = SubpopulationCounts( + name="hospital", + ascertainment_rate_rv=DeterministicVariable("rate", rate), + delay_distribution_rv=DeterministicPMF("delay", jnp.array([1.0])), + noise=PoissonNoise(), + ) + + predicted = process._apply_ascertainment(potential_counts, rate) + + assert jnp.allclose(predicted, potential_counts * rate[:, None]) + assert jnp.allclose(predicted[:, 1] / predicted[:, 0], 3.0) + + def test_delay_padding_is_preserved(self) -> None: + """Test ascertainment leaves leading delay-padding values as NaN.""" + process = PopulationCounts( + name="hospital", + ascertainment_rate_rv=DeterministicVariable( + "rate", + jnp.linspace(0.1, 0.5, 5), + ), + delay_distribution_rv=DeterministicPMF( + "delay", + jnp.array([0.2, 0.3, 0.5]), + ), + noise=PoissonNoise(), + ) + + predicted = process._predicted_obs(jnp.ones(5)) + + assert jnp.all(jnp.isnan(predicted[:2])) + assert jnp.all(jnp.isfinite(predicted[2:])) + + def test_delay_padding_does_not_create_nan_rate_gradients(self) -> None: + """Test padded potential counts have a finite ascertainment gradient.""" + process = PopulationCounts( + name="hospital", + ascertainment_rate_rv=DeterministicVariable("rate", 0.2), + delay_distribution_rv=DeterministicPMF("delay", jnp.array([1.0])), + noise=PoissonNoise(), + ) + potential_counts = jnp.array([jnp.nan, 2.0, 4.0]) + + def finite_sum(rate: ArrayLike) -> ArrayLike: + """Return the sum after replacing padding with a finite value. + + Returns + ------- + ArrayLike + Sum of finite scaled counts. + """ + scaled = process._apply_ascertainment(potential_counts, rate) + return jnp.sum(jnp.where(jnp.isnan(scaled), 0.0, scaled)) + + gradient = jax.grad(finite_sum)(jnp.array(0.2)) + + assert jnp.isfinite(gradient) + assert jnp.isclose(gradient, 6.0) + + class TestCountsWithPriors: """Test aggregated count observation with uncertain parameters.""" diff --git a/test/test_pyrenew_builder.py b/test/test_pyrenew_builder.py index a03af53ea..39b08edb1 100644 --- a/test/test_pyrenew_builder.py +++ b/test/test_pyrenew_builder.py @@ -2,14 +2,21 @@ Tests for PyrenewBuilder and MultiSignalModel. """ +from collections.abc import Mapping from datetime import date, timedelta import jax.numpy as jnp import numpyro import numpyro.distributions as dist import pytest +from jax.typing import ArrayLike -from pyrenew.ascertainment import JointAscertainment, RatioLinkedAscertainment +from pyrenew.ascertainment import ( + AscertainmentModel, + JointAscertainment, + RatioLinkedAscertainment, + TimeVaryingAscertainment, +) from pyrenew.deterministic import DeterministicPMF, DeterministicVariable from pyrenew.latent import ( InfectionsWithFeedback, @@ -33,6 +40,48 @@ SUBPOP_FRACTIONS = jnp.array([0.3, 0.25, 0.45]) +class RecordingAscertainmentModel(AscertainmentModel): + """Return configured values while recording axis context for tests.""" + + def __init__( + self, + name: str, + values: Mapping[str, ArrayLike], + requires_calendar_anchor: bool = False, + ) -> None: + """Initialize a recording ascertainment model.""" + super().__init__(name=name, signals=tuple(values)) + self.values = dict(values) + self._requires_calendar_anchor = requires_calendar_anchor + self.sample_calls: list[tuple[int, int | None]] = [] + + def requires_calendar_anchor(self) -> bool: + """Return the configured calendar requirement. + + Returns + ------- + bool + Configured calendar requirement. + """ + return self._requires_calendar_anchor + + def sample( + self, + n_timepoints: int, + first_day_dow: int | None = None, + **kwargs: object, + ) -> Mapping[str, ArrayLike]: + """Return configured values and record shared-axis arguments. + + Returns + ------- + Mapping[str, ArrayLike] + Configured signal values. + """ + self.sample_calls.append((n_timepoints, first_day_dow)) + return self.values + + def _obs_date_for_dow(target_first_day_dow: int, n_init: int) -> date: """ Return an ``obs_start_date`` whose axis-origin day-of-week matches. @@ -58,6 +107,42 @@ def _obs_date_for_dow(target_first_day_dow: int, n_init: int) -> date: return date(2024, 1, 1) + timedelta(days=obs_dow) +def _builder_with_time_varying_ascertainment( + ascertainment: TimeVaryingAscertainment, +) -> PyrenewBuilder: + """Build a two-count-signal model using one outer ascertainment model. + + Returns + ------- + PyrenewBuilder + Configured builder with hospital and ED count observations. + """ + builder = PyrenewBuilder() + builder.configure_latent( + SubpopulationInfections, + gen_int_rv=DeterministicPMF( + "gen_int", + jnp.array([0.2, 0.5, 0.3]), + ), + I0_rv=DeterministicVariable("I0", 0.001), + log_rt_time_0_rv=DeterministicVariable("initial_log_rt", 0.0), + baseline_rt_process=fixed_random_walk(innovation_sd=0.1), + subpop_rt_deviation_process=fixed_random_walk(innovation_sd=0.1), + ) + builder.add_ascertainment(ascertainment) + delay = DeterministicPMF("delay", jnp.array([0.1, 0.3, 0.4, 0.2])) + for signal in ("hospital", "ed"): + builder.add_observation( + PopulationCounts( + name=signal, + ascertainment_rate_rv=ascertainment.for_signal(signal), + delay_distribution_rv=delay, + noise=PoissonNoise(), + ) + ) + return builder + + @pytest.fixture def simple_builder(): """ @@ -504,6 +589,99 @@ def test_prior_predictive_with_ratio_linked_ascertainment(self): assert "hospital_predicted" in prior_samples assert "ed_predicted" in prior_samples + def test_prior_predictive_with_time_varying_joint_baseline(self) -> None: + """Test joint baseline sites and full-axis outer count trajectories.""" + import jax.random + from numpyro.infer import Predictive + + baseline = JointAscertainment( + name="he_baseline", + signals=("hospital", "ed"), + baseline_rates=jnp.array([0.01, 0.02]), + scale_tril=jnp.eye(2) * 0.1, + ) + ascertainment = TimeVaryingAscertainment( + "he_ascertainment", + baseline, + { + "hospital": fixed_random_walk(innovation_sd=0.05), + "ed": fixed_random_walk(innovation_sd=0.05), + }, + ) + model = _builder_with_time_varying_ascertainment(ascertainment).build() + n_days_post_init = 10 + n_total = model.latent.n_initialization_points + n_days_post_init + + samples = Predictive(model.sample, num_samples=2)( + jax.random.PRNGKey(42), + n_days_post_init=n_days_post_init, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ed={"obs": None}, + ) + + assert set(model.ascertainment_models) == {"he_ascertainment"} + assert samples["he_baseline_eta"].shape == (2, 2) + assert samples["he_baseline_hospital"].shape == (2,) + assert samples["he_baseline_ed"].shape == (2,) + assert samples["he_ascertainment_hospital"].shape == (2, n_total) + assert samples["he_ascertainment_ed"].shape == (2, n_total) + assert samples["hospital_predicted"].shape == (2, n_total) + assert samples["ed_predicted"].shape == (2, n_total) + assert samples["latent_infections"].shape == (2, n_total) + + def test_prior_predictive_with_time_varying_ratio_baseline(self) -> None: + """Test ratio sites persist beneath independently varying trajectories.""" + import jax.random + from numpyro.infer import Predictive + + baseline = RatioLinkedAscertainment( + name="he_baseline", + base_signal="ed", + linked_signal="hospital", + base_rate_rv=DistributionalVariable("iedr", dist.Beta(2, 98)), + ratio_rv=DistributionalVariable( + "ihr_rel_iedr", + dist.Beta(2, 2), + ), + ) + ascertainment = TimeVaryingAscertainment( + "he_ascertainment", + baseline, + { + "hospital": fixed_random_walk(innovation_sd=0.05), + "ed": fixed_random_walk(innovation_sd=0.05), + }, + ) + model = _builder_with_time_varying_ascertainment(ascertainment).build() + n_days_post_init = 10 + n_total = model.latent.n_initialization_points + n_days_post_init + + samples = Predictive(model.sample, num_samples=2)( + jax.random.PRNGKey(42), + n_days_post_init=n_days_post_init, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ed={"obs": None}, + ) + + assert samples["iedr"].shape == (2,) + assert samples["ihr_rel_iedr"].shape == (2,) + assert jnp.allclose(samples["he_baseline_ed"], samples["iedr"]) + assert jnp.allclose( + samples["he_baseline_hospital"], + samples["iedr"] * samples["ihr_rel_iedr"], + ) + assert samples["he_ascertainment_hospital"].shape == (2, n_total) + assert samples["he_ascertainment_ed"].shape == (2, n_total) + assert not jnp.allclose( + samples["he_ascertainment_hospital"] / samples["he_ascertainment_ed"], + samples["ihr_rel_iedr"][:, None], + ) + assert samples["latent_infections"].shape == (2, n_total) + def test_prior_predictive_reuses_same_ascertainment_signal(self): """Test two observations can reuse one signal accessor without site conflicts.""" import jax.random @@ -597,6 +775,228 @@ def test_manual_model_rejects_mismatched_ascertainment_key(self, simple_builder) ascertainment_models={"wrong_name": ascertainment}, ) + def test_rejects_registering_time_varying_baseline_object( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test a wrapped baseline object cannot also be registered.""" + built = simple_builder.build() + baseline = RecordingAscertainmentModel("baseline", {"hospital": 0.2}) + outer = TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": fixed_random_walk(innovation_sd=0.1)}, + ) + + with pytest.raises(ValueError, match=r"outer.*baseline.*also be registered"): + MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={"outer": outer, "baseline": baseline}, + ) + + def test_rejects_registered_model_with_wrapped_baseline_name( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test distinct registered and baseline objects cannot share a name.""" + built = simple_builder.build() + baseline = RecordingAscertainmentModel("baseline", {"hospital": 0.2}) + outer = TimeVaryingAscertainment( + "outer", + baseline, + {"hospital": fixed_random_walk(innovation_sd=0.1)}, + ) + registered = RecordingAscertainmentModel("baseline", {"ed": 0.3}) + + with pytest.raises( + ValueError, + match=r"namespace 'baseline'.*baseline model.*outer.*registered", + ): + MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={"outer": outer, "baseline": registered}, + ) + + def test_rejects_two_wrapped_baselines_with_the_same_name( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test independently wrapped baselines need distinct namespaces.""" + built = simple_builder.build() + first_baseline = RecordingAscertainmentModel( + "shared_baseline", + {"hospital": 0.2}, + ) + second_baseline = RecordingAscertainmentModel( + "shared_baseline", + {"ed": 0.3}, + ) + first = TimeVaryingAscertainment( + "first_outer", + first_baseline, + {"hospital": fixed_random_walk(innovation_sd=0.1)}, + ) + second = TimeVaryingAscertainment( + "second_outer", + second_baseline, + {"ed": fixed_random_walk(innovation_sd=0.1)}, + ) + + with pytest.raises( + ValueError, + match=r"shared_baseline.*first_outer.*second_outer", + ): + MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={ + "first_outer": first, + "second_outer": second, + }, + ) + + def test_rejects_outer_name_colliding_with_another_baseline( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test outer and owned baseline namespaces are globally unique.""" + built = simple_builder.build() + first = TimeVaryingAscertainment( + "second_baseline", + RecordingAscertainmentModel("first_baseline", {"hospital": 0.2}), + {"hospital": fixed_random_walk(innovation_sd=0.1)}, + ) + second = TimeVaryingAscertainment( + "second_outer", + RecordingAscertainmentModel("second_baseline", {"ed": 0.3}), + {"ed": fixed_random_walk(innovation_sd=0.1)}, + ) + + with pytest.raises( + ValueError, + match=r"second_baseline.*registered.*second_outer", + ): + MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={ + "second_baseline": first, + "second_outer": second, + }, + ) + + def test_forwards_full_axis_length_and_padded_axis_weekday( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test ascertainment sampling receives the shared padded-axis context.""" + built = simple_builder.build() + ascertainment = RecordingAscertainmentModel( + "recording", + {"hospital": 0.2}, + ) + model = MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={"recording": ascertainment}, + ) + first_day_dow = 4 + obs_start_date = _obs_date_for_dow( + first_day_dow, + model.latent.n_initialization_points, + ) + + with numpyro.handlers.seed(rng_seed=42): + model.sample( + n_days_post_init=5, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + obs_start_date=obs_start_date, + hospital={"obs": None}, + ) + + assert ascertainment.sample_calls == [ + (model.latent.n_initialization_points + 5, first_day_dow) + ] + + def test_calendar_ascertainment_fails_early_and_names_model( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test registered calendar-aligned ascertainment requires a date.""" + built = simple_builder.build() + ascertainment = RecordingAscertainmentModel( + "weekly_ascertainment", + {"hospital": 0.2}, + requires_calendar_anchor=True, + ) + model = MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={"weekly_ascertainment": ascertainment}, + ) + + with pytest.raises( + ValueError, + match=r"obs_start_date.*weekly_ascertainment", + ): + model.sample( + n_days_post_init=5, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ) + + def test_unregistered_calendar_ascertainment_does_not_require_date( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test only registered ascertainment contributes date requirements.""" + built = simple_builder.build() + RecordingAscertainmentModel( + "unregistered_weekly", + {"hospital": 0.2}, + requires_calendar_anchor=True, + ) + + with numpyro.handlers.seed(rng_seed=42): + built.sample( + n_days_post_init=5, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ) + + def test_malformed_ascertainment_fails_before_observation_sampling( + self, + simple_builder: PyrenewBuilder, + ) -> None: + """Test output validation completes before any observation runs.""" + built = simple_builder.build() + malformed = RecordingAscertainmentModel( + "malformed", + {"hospital": jnp.ones((2, 1))}, + ) + model = MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={"malformed": malformed}, + ) + + with numpyro.handlers.seed(rng_seed=42): + with numpyro.handlers.trace() as trace: + with pytest.raises(ValueError, match=r"malformed.*hospital.*shape"): + model.sample( + n_days_post_init=5, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ) + + assert "hospital_obs" not in trace + def test_first_day_dow_reaches_calendar_aligned_latent_process(self): """MultiSignalModel forwards model-axis day of week to the latent process.""" latent = PopulationInfections( From 7bb72c69b177e3dd7fb8304ac0990e5ae4acb463 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 30 Sep 2026 11:37:03 -0400 Subject: [PATCH 02/12] checkpointing - revise API; remove TimeVaryingAscertainment --- docs/tutorials/ascertainment.qmd | 8 +- pyrenew/ascertainment/__init__.py | 4 +- pyrenew/ascertainment/base.py | 241 +++++- pyrenew/ascertainment/independent.py | 87 +++ pyrenew/ascertainment/joint.py | 34 +- pyrenew/ascertainment/linked.py | 36 +- pyrenew/ascertainment/timevarying.py | 207 ----- pyrenew/model/multisignal_model.py | 35 +- pyrenew/model/pyrenew_builder.py | 48 +- test/integration/conftest.py | 19 +- ...nfections_he_weekly_joint_ascertainment.py | 13 +- ...ns_he_weekly_ratio_linked_ascertainment.py | 16 +- test/test_ascertainment.py | 719 +++++++++++++----- test/test_pyrenew_builder.py | 398 +++++----- 14 files changed, 1170 insertions(+), 695 deletions(-) create mode 100644 pyrenew/ascertainment/independent.py delete mode 100644 pyrenew/ascertainment/timevarying.py diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 07feb18b1..2023f144c 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -53,6 +53,12 @@ $$ where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate, and $\pi_{s,d}$ is its delay distribution. +## Individual and related ascertainment rates for multi-signal models + +## Fixed and time-varying ascertainment + +******** stopped here ****** + If a model treats the ascertainment rates for different signals as unrelated, each observation process can receive its own `RandomVariable` directly. An ascertainment component is needed when the model represents a relationship between the ascertainment rates of multiple observation processes. @@ -60,7 +66,7 @@ PyRenew provides two scalar models and one time-varying wrapper: - `JointAscertainment` specifies a joint prior for two or more rates. - `RatioLinkedAscertainment` specifies one rate as a multiple of another rate. -- `TimeVaryingAscertainment` uses either scalar model as a baseline and adds an independent temporal deviation for each signal. + The choice among them is part of the statistical model, not only an API choice. Each component expresses a different assumption about how information is shared across signals and over time. diff --git a/pyrenew/ascertainment/__init__.py b/pyrenew/ascertainment/__init__.py index ad8a488bd..3f584a5e0 100644 --- a/pyrenew/ascertainment/__init__.py +++ b/pyrenew/ascertainment/__init__.py @@ -4,14 +4,14 @@ """ from pyrenew.ascertainment.base import AscertainmentModel, AscertainmentSignal +from pyrenew.ascertainment.independent import IndependentAscertainment from pyrenew.ascertainment.joint import JointAscertainment from pyrenew.ascertainment.linked import RatioLinkedAscertainment -from pyrenew.ascertainment.timevarying import TimeVaryingAscertainment __all__ = [ "AscertainmentModel", "AscertainmentSignal", + "IndependentAscertainment", "JointAscertainment", "RatioLinkedAscertainment", - "TimeVaryingAscertainment", ] diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index de1f6d71e..a4ff9cd50 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -8,10 +8,16 @@ from abc import ABCMeta, abstractmethod from collections.abc import Mapping +import jax.nn as jnn import jax.numpy as jnp +import numpyro +from jax import Array +from jax.scipy.special import logit from jax.typing import ArrayLike +from numpyro.util import not_jax_tracer from pyrenew.ascertainment.context import get_ascertainment_value +from pyrenew.latent import TemporalProcess from pyrenew.metaclass import RandomVariable @@ -101,15 +107,21 @@ class AscertainmentModel(metaclass=ABCMeta): ) ``` - Subclasses own any NumPyro sites needed for the shared structure. - Accessors returned by ``for_signal()`` read the sampled values from the - active model context and do not sample independently. + The model samples one scalar baseline rate per signal and can add an + optional signal-specific temporal deviation on the logit scale. Subclasses + implement ``_sample_baseline_rates()`` to define relationships among the + scalar baselines. The base class owns validation, temporal sampling, and + the standard deterministic sites. + + Accessors returned by ``for_signal()`` read the final sampled values from + the active model context and do not sample independently. """ def __init__( self, name: str, signals: tuple[str, ...], + temporal_processes: Mapping[str, TemporalProcess] | None = None, ) -> None: """ Initialize an ascertainment model. @@ -120,6 +132,17 @@ def __init__( A non-empty string identifying the ascertainment model. signals Unique signal names produced by this model. + temporal_processes + Optional temporal processes keyed by signal name. Signals without + a process retain their scalar baseline rate. + + Raises + ------ + TypeError + If ``temporal_processes`` is not a mapping or a configured process + does not satisfy the ``TemporalProcess`` protocol. + ValueError + If a temporal process is configured for an unknown signal. """ if not isinstance(name, str) or len(name) == 0: raise ValueError( @@ -135,6 +158,37 @@ def __init__( self.name = name self.signals = signals + if temporal_processes is None: + temporal_processes = {} + if not isinstance(temporal_processes, Mapping): + raise TypeError( + "temporal_processes must be a mapping, " + f"got {type(temporal_processes).__name__}." + ) + + expected_signals = set(signals) + unknown_signals = tuple( + signal for signal in temporal_processes if signal not in expected_signals + ) + if unknown_signals: + raise ValueError( + f"temporal_processes contains unknown signals {unknown_signals} for " + f"ascertainment model {name!r}. Available signals: {signals}." + ) + + ordered_processes: dict[str, TemporalProcess] = {} + for signal in signals: + if signal not in temporal_processes: + continue + process = temporal_processes[signal] + if not isinstance(process, TemporalProcess): + raise TypeError( + f"temporal process for signal {signal!r} must satisfy the " + f"TemporalProcess protocol, got {type(process).__name__}." + ) + ordered_processes[signal] = process + self.temporal_processes = ordered_processes + def for_signal(self, signal_name: str) -> AscertainmentSignal: """ Return an observation-process accessor for one signal. @@ -170,7 +224,7 @@ def for_signal(self, signal_name: str) -> AscertainmentSignal: ) def requires_calendar_anchor(self) -> bool: - """Return whether sampling requires the model-axis day-of-week. + """Return whether any temporal process requires a calendar anchor. Returns ------- @@ -178,7 +232,165 @@ def requires_calendar_anchor(self) -> bool: ``True`` when the model needs a calendar anchor; otherwise ``False``. """ - return False + return any( + getattr(process, "requires_calendar_anchor", False) + for process in self.temporal_processes.values() + ) + + def _validate_baseline_rates( + self, + baseline_rates: Mapping[str, ArrayLike], + ) -> dict[str, Array]: + """Validate and convert scalar baseline ascertainment rates. + + Parameters + ---------- + baseline_rates + Sampled scalar baseline rates keyed by signal name. + + Returns + ------- + dict[str, Array] + Validated scalar arrays in signal order. + + Raises + ------ + TypeError + If ``baseline_rates`` is not a mapping. + ValueError + If signal names do not match, a baseline is not scalar, or a + concrete baseline is non-finite or outside its allowed interval. + """ + if not isinstance(baseline_rates, Mapping): + raise TypeError( + f"Ascertainment model {self.name!r} must return baseline rates " + f"as a mapping, got {type(baseline_rates).__name__}." + ) + + expected_signals = set(self.signals) + actual_signals = set(baseline_rates) + if actual_signals != expected_signals: + missing = tuple( + signal for signal in self.signals if signal not in baseline_rates + ) + extra = tuple( + signal for signal in baseline_rates if signal not in expected_signals + ) + raise ValueError( + f"Ascertainment model {self.name!r} must return baseline rates for " + f"exactly signals {self.signals}. Missing: {missing}. Extra: {extra}." + ) + + validated_rates: dict[str, Array] = {} + for signal in self.signals: + baseline_rate = jnp.asarray(baseline_rates[signal]) + if baseline_rate.shape != (): + raise ValueError( + f"Ascertainment model {self.name!r}, signal {signal!r}, " + f"returned baseline shape {baseline_rate.shape}; required shape " + "is ()." + ) + + if signal in self.temporal_processes: + invalid_rate = ( + ~jnp.isfinite(baseline_rate) + | (baseline_rate <= 0) + | (baseline_rate >= 1) + ) + allowed_interval = "strictly inside (0, 1)" + else: + invalid_rate = ( + ~jnp.isfinite(baseline_rate) + | (baseline_rate < 0) + | (baseline_rate > 1) + ) + allowed_interval = "inside [0, 1]" + + if not_jax_tracer(invalid_rate) and bool(invalid_rate): + raise ValueError( + f"Ascertainment model {self.name!r} requires the baseline for " + f"signal {signal!r} to be finite and {allowed_interval}; got " + f"{baseline_rate}." + ) + validated_rates[signal] = baseline_rate + + return validated_rates + + def sample( + self, + n_timepoints: int, + first_day_dow: int | None = None, + ) -> Mapping[str, ArrayLike]: + """Sample baseline rates and optional temporal rate trajectories. + + Parameters + ---------- + n_timepoints + Positive number of timepoints on the shared model axis. + first_day_dow + Day of week for the first model-axis timepoint. Calendar-aligned + temporal processes use this value. + + Returns + ------- + Mapping[str, ArrayLike] + Final scalar rates or full-axis trajectories in signal order. + + Raises + ------ + TypeError + If ``n_timepoints`` is not a built-in integer. + ValueError + If ``n_timepoints`` is less than one, baseline rates are invalid, + or a temporal process returns an unexpected shape. + """ + if type(n_timepoints) is not int: + raise TypeError( + "n_timepoints must be a positive integer, " + f"got {type(n_timepoints).__name__}." + ) + if n_timepoints < 1: + raise ValueError( + f"n_timepoints must be a positive integer, got {n_timepoints}." + ) + + baseline_rates = self._validate_baseline_rates(self._sample_baseline_rates()) + + result: dict[str, ArrayLike] = {} + for signal in self.signals: + baseline_rate = baseline_rates[signal] + numpyro.deterministic( + f"{self.name}_baseline_{signal}", + baseline_rate, + ) + + process = self.temporal_processes.get(signal) + if process is None: + rate = baseline_rate + else: + deviation = jnp.asarray( + process.sample( + n_timepoints=n_timepoints, + initial_value=0.0, + n_processes=1, + name_prefix=f"{self.name}_{signal}", + first_day_dow=first_day_dow, + ) + ) + required_shape = (n_timepoints, 1) + if deviation.shape != required_shape: + raise ValueError( + f"Ascertainment model {self.name!r}, signal {signal!r}, " + f"received temporal-process shape {deviation.shape}; " + f"required shape is {required_shape}." + ) + deviation = jnp.squeeze(deviation, axis=-1) + rate = jnn.sigmoid(logit(baseline_rate) + deviation) + + numpyro.deterministic(f"{self.name}_{signal}", rate) + result[signal] = rate + + return result def validate_sampled_values( self, @@ -229,19 +441,18 @@ def validate_sampled_values( ) @abstractmethod - def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: - """ - Sample all signal-specific ascertainment values owned by this model. - - Parameters - ---------- - **kwargs - Additional model-context arguments supplied by ``MultiSignalModel``. - Subclasses may ignore unused values. + def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + """Sample one scalar baseline ascertainment rate per signal. Returns ------- Mapping[str, ArrayLike] - Mapping from signal name to sampled ascertainment rate. + Scalar baseline rates keyed by the model's signal names. + + Notes + ----- + Custom subclasses should implement only baseline sampling here. The + base ``sample()`` method validates baselines, applies any configured + temporal processes, and records the standard deterministic sites. """ pass # pragma: no cover diff --git a/pyrenew/ascertainment/independent.py b/pyrenew/ascertainment/independent.py new file mode 100644 index 000000000..34fc6b730 --- /dev/null +++ b/pyrenew/ascertainment/independent.py @@ -0,0 +1,87 @@ +# numpydoc ignore=GL08 +""" +Independent ascertainment models. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import cast + +from jax.typing import ArrayLike + +from pyrenew.ascertainment.base import AscertainmentModel +from pyrenew.latent import TemporalProcess +from pyrenew.metaclass import RandomVariable + + +class IndependentAscertainment(AscertainmentModel): + """Ascertainment rates with independently sampled scalar baselines. + + Use this model for one signal or for multiple signals whose baseline rates + do not share a joint prior. Add temporal processes only for signals whose + rates should vary over time. + """ + + def __init__( + self, + name: str, + rate_rvs: Mapping[str, RandomVariable], + temporal_processes: Mapping[str, TemporalProcess] | None = None, + ) -> None: + """Initialize independent ascertainment rates. + + Parameters + ---------- + name + Name of the ascertainment model. + rate_rvs + Non-empty mapping from signal names to random variables for their + scalar baseline rates. Mapping order determines signal order. + temporal_processes + Optional temporal processes keyed by signal name. Signals without + a process retain their scalar baseline rate. + + Raises + ------ + TypeError + If ``rate_rvs`` is not a mapping or a value is not a + ``RandomVariable``. + ValueError + If ``rate_rvs`` is empty or contains an invalid signal name. + """ + if not isinstance(rate_rvs, Mapping): + raise TypeError( + f"rate_rvs must be a mapping, got {type(rate_rvs).__name__}." + ) + if len(rate_rvs) == 0: + raise ValueError("rate_rvs must be a non-empty mapping.") + + for signal, rate_rv in rate_rvs.items(): + if not isinstance(signal, str) or len(signal) == 0: + raise ValueError("rate_rvs keys must be non-empty strings.") + if not isinstance(rate_rv, RandomVariable): + raise TypeError( + f"rate_rv for signal {signal!r} must be a RandomVariable, " + f"got {type(rate_rv).__name__}." + ) + + self.rate_rvs = dict(rate_rvs) + super().__init__( + name=name, + signals=tuple(self.rate_rvs), + temporal_processes=temporal_processes, + ) + + def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + """Sample one scalar baseline rate from each random variable. + + Returns + ------- + Mapping[str, ArrayLike] + Independently sampled scalar rates in signal order. + """ + return { + signal: cast(ArrayLike, rate_rv()) + for signal, rate_rv in self.rate_rvs.items() + } diff --git a/pyrenew/ascertainment/joint.py b/pyrenew/ascertainment/joint.py index 2fc42d74a..0796c0f74 100644 --- a/pyrenew/ascertainment/joint.py +++ b/pyrenew/ascertainment/joint.py @@ -15,11 +15,11 @@ from jax.typing import ArrayLike from pyrenew.ascertainment.base import AscertainmentModel +from pyrenew.latent import TemporalProcess class JointAscertainment(AscertainmentModel): - """ - Joint prior for scalar ascertainment rates across multiple signals. + """Joint prior for baseline ascertainment rates across multiple signals. This model is useful when multiple observation streams have distinct but related probabilities of observing latent incidence. For example, hospital @@ -35,7 +35,8 @@ class JointAscertainment(AscertainmentModel): ascertainment_rate_j = sigmoid(eta_j) ``` - Each returned rate is scalar and constant over the model time axis. + Each sampled baseline is scalar. Optional temporal processes can vary any + subset of the final rates over the model time axis. """ def __init__( @@ -46,6 +47,7 @@ def __init__( scale_tril: ArrayLike | None = None, covariance_matrix: ArrayLike | None = None, precision_matrix: ArrayLike | None = None, + temporal_processes: Mapping[str, TemporalProcess] | None = None, ) -> None: """ Initialize a joint scalar ascertainment model. @@ -72,8 +74,15 @@ def __init__( precision_matrix Precision matrix for the multivariate normal on the logit scale. Exactly one covariance parameter must be supplied. + temporal_processes + Optional temporal processes keyed by signal name. Signals without + a process retain their scalar baseline rate. """ - super().__init__(name=name, signals=signals) + super().__init__( + name=name, + signals=signals, + temporal_processes=temporal_processes, + ) baseline_rates_array = jnp.asarray(baseline_rates) scale_tril_array = self._optional_array(scale_tril) covariance_matrix_array = self._optional_array(covariance_matrix) @@ -126,23 +135,17 @@ def _validate_parameters(self, baseline_rates: Array) -> None: ) if jnp.any(baseline_rates <= 0) or jnp.any(baseline_rates >= 1): raise ValueError( - "baseline_rates must contain probabilities in [0, 1], " + "baseline_rates must contain probabilities in (0, 1), " f"got {baseline_rates}." ) - def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: - """ - Sample jointly distributed scalar ascertainment rates. - - Parameters - ---------- - **kwargs - Additional model-context arguments, ignored. + def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + """Sample jointly distributed scalar baseline rates. Returns ------- Mapping[str, ArrayLike] - Mapping from signal name to sampled scalar ascertainment rate. + Sampled scalar baseline rates keyed by signal name. """ eta = numpyro.sample( f"{self.name}_eta", @@ -150,9 +153,8 @@ def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: ) rates = expit(eta) - result = {} + result: dict[str, ArrayLike] = {} for signal, rate in zip(self.signals, rates): - numpyro.deterministic(f"{self.name}_{signal}", rate) result[signal] = rate return result diff --git a/pyrenew/ascertainment/linked.py b/pyrenew/ascertainment/linked.py index b209cc27f..c4579fed0 100644 --- a/pyrenew/ascertainment/linked.py +++ b/pyrenew/ascertainment/linked.py @@ -7,10 +7,11 @@ from collections.abc import Mapping -import numpyro +import jax.numpy as jnp from jax.typing import ArrayLike from pyrenew.ascertainment.base import AscertainmentModel +from pyrenew.latent import TemporalProcess from pyrenew.metaclass import RandomVariable @@ -19,7 +20,8 @@ class RatioLinkedAscertainment(AscertainmentModel): Two ascertainment rates expressed as a base rate and a ratio. The linked ascertainment rate is the sampled base rate multiplied by the - sampled ratio. + sampled ratio. Optional temporal processes act on the resulting scalar + baselines and do not preserve the ratio pointwise over time. """ def __init__( @@ -29,6 +31,7 @@ def __init__( linked_signal: str, base_rate_rv: RandomVariable, ratio_rv: RandomVariable, + temporal_processes: Mapping[str, TemporalProcess] | None = None, ) -> None: """ Initialize a ratio-linked ascertainment model. @@ -47,35 +50,32 @@ def __init__( ratio_rv Random variable for the ratio of the linked signal's ascertainment rate to the base signal's ascertainment rate. + temporal_processes + Optional temporal processes keyed by signal name. Signals without + a process retain their scalar baseline rate. """ - super().__init__(name=name, signals=(base_signal, linked_signal)) + super().__init__( + name=name, + signals=(base_signal, linked_signal), + temporal_processes=temporal_processes, + ) self.base_signal = base_signal self.linked_signal = linked_signal self.base_rate_rv = base_rate_rv self.ratio_rv = ratio_rv - def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: - """ - Sample the base rate and ratio and calculate the linked rate. - - Parameters - ---------- - **kwargs - Additional model-context arguments, ignored. + def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + """Sample the base rate and ratio and calculate the linked baseline. Returns ------- Mapping[str, ArrayLike] - Mapping from the base and linked signal names to their sampled - ascertainment rates. + Scalar baseline rates for the base and linked signals. """ - base_rate = self.base_rate_rv() - ratio = self.ratio_rv() + base_rate = jnp.asarray(self.base_rate_rv()) + ratio = jnp.asarray(self.ratio_rv()) linked_rate = base_rate * ratio - numpyro.deterministic(f"{self.name}_{self.base_signal}", base_rate) - numpyro.deterministic(f"{self.name}_{self.linked_signal}", linked_rate) - return { self.base_signal: base_rate, self.linked_signal: linked_rate, diff --git a/pyrenew/ascertainment/timevarying.py b/pyrenew/ascertainment/timevarying.py deleted file mode 100644 index ee738b4d0..000000000 --- a/pyrenew/ascertainment/timevarying.py +++ /dev/null @@ -1,207 +0,0 @@ -# numpydoc ignore=GL08 -""" -Time-varying ascertainment models. -""" - -from __future__ import annotations - -from collections.abc import Mapping - -import jax.nn as jnn -import jax.numpy as jnp -import numpyro -from jax.scipy.special import logit -from jax.typing import ArrayLike -from numpyro.util import not_jax_tracer - -from pyrenew.ascertainment.base import AscertainmentModel -from pyrenew.latent import TemporalProcess - - -class TimeVaryingAscertainment(AscertainmentModel): - """ - Add independent temporal deviations to related ascertainment baselines. - - The wrapped baseline model samples one scalar rate per signal. Each scalar - is converted to the logit scale, combined with a signal-specific temporal - deviation, and returned as a probability trajectory on the shared model - time axis. Baseline relationships therefore describe the reference rates, - not a pointwise constraint on the resulting trajectories. - - Register only this outer model with ``PyrenewBuilder``. Use accessors from - this model, such as ``ascertainment.for_signal("hospital")``, in count - observation processes. - """ - - def __init__( - self, - name: str, - baseline_model: AscertainmentModel, - processes: Mapping[str, TemporalProcess], - ) -> None: - """ - Initialize time-varying ascertainment. - - Parameters - ---------- - name - Name of the outer ascertainment model. It must differ from the - wrapped baseline model's name. - baseline_model - Ascertainment model that samples one scalar reference rate for - each signal. - processes - One temporal process per baseline signal. Each process generates - an independent logit-scale deviation trajectory. - - Raises - ------ - TypeError - If the baseline or temporal processes have invalid types. - ValueError - If names conflict or process signal names do not match the - baseline signals. - """ - if not isinstance(baseline_model, AscertainmentModel): - raise TypeError( - "baseline_model must be an AscertainmentModel, " - f"got {type(baseline_model).__name__}." - ) - if not isinstance(processes, Mapping): - raise TypeError( - f"processes must be a mapping, got {type(processes).__name__}." - ) - if name == baseline_model.name: - raise ValueError( - "Time-varying ascertainment and its baseline model must have " - f"different names; both were {name!r}." - ) - - super().__init__(name=name, signals=baseline_model.signals) - expected_signals = set(self.signals) - process_signals = set(processes) - if process_signals != expected_signals: - missing = tuple( - signal for signal in self.signals if signal not in processes - ) - extra = tuple( - signal for signal in processes if signal not in expected_signals - ) - raise ValueError( - "processes must contain exactly the baseline model signals " - f"{self.signals}. Missing: {missing}. Extra: {extra}." - ) - - ordered_processes: dict[str, TemporalProcess] = {} - for signal in self.signals: - process = processes[signal] - if not isinstance(process, TemporalProcess): - raise TypeError( - f"process for signal {signal!r} must satisfy the " - f"TemporalProcess protocol, got {type(process).__name__}." - ) - ordered_processes[signal] = process - - self.baseline_model = baseline_model - self.processes = ordered_processes - - def requires_calendar_anchor(self) -> bool: - """Return whether the baseline or any temporal process needs a date. - - Returns - ------- - bool - ``True`` when sampling needs the model-axis day-of-week; otherwise - ``False``. - """ - return self.baseline_model.requires_calendar_anchor() or any( - getattr(process, "requires_calendar_anchor", False) - for process in self.processes.values() - ) - - def sample( - self, - n_timepoints: int, - first_day_dow: int | None = None, - **kwargs: object, - ) -> Mapping[str, ArrayLike]: - """ - Sample full-axis ascertainment-rate trajectories. - - Parameters - ---------- - n_timepoints - Number of timepoints on the shared model axis. - first_day_dow - Day-of-week for the first model-axis timepoint. Required by - calendar-aligned temporal processes. - **kwargs - Additional model context, ignored. - - Returns - ------- - Mapping[str, ArrayLike] - One probability trajectory of shape ``(n_timepoints,)`` per - signal. - - Raises - ------ - ValueError - If the baseline does not return scalar probabilities in ``(0, 1)`` - or a temporal process returns an unexpected shape. - """ - baseline_values = self.baseline_model.sample( - n_timepoints=n_timepoints, - first_day_dow=first_day_dow, - ) - self.baseline_model.validate_sampled_values( - baseline_values, - n_timepoints=n_timepoints, - ) - - scalar_baselines: dict[str, ArrayLike] = {} - for signal in self.signals: - baseline_value = jnp.asarray(baseline_values[signal]) - if baseline_value.shape != (): - raise ValueError( - f"Time-varying ascertainment model {self.name!r} requires " - f"a scalar baseline for signal {signal!r}; got shape " - f"{baseline_value.shape}, required shape ()." - ) - invalid_baseline = ( - ~jnp.isfinite(baseline_value) - | (baseline_value <= 0) - | (baseline_value >= 1) - ) - if not_jax_tracer(invalid_baseline) and bool(invalid_baseline): - raise ValueError( - f"Time-varying ascertainment model {self.name!r} requires " - f"the baseline for signal {signal!r} to be finite and " - f"strictly inside (0, 1); got {baseline_value}." - ) - scalar_baselines[signal] = baseline_value - - result: dict[str, ArrayLike] = {} - for signal in self.signals: - deviation = jnp.asarray( - self.processes[signal].sample( - n_timepoints=n_timepoints, - initial_value=0.0, - n_processes=1, - name_prefix=f"{self.name}_{signal}", - first_day_dow=first_day_dow, - ) - ) - required_shape = (n_timepoints, 1) - if deviation.shape != required_shape: - raise ValueError( - f"Time-varying ascertainment model {self.name!r}, signal " - f"{signal!r}, received temporal-process shape " - f"{deviation.shape}; required shape is {required_shape}." - ) - deviation = jnp.squeeze(deviation, axis=-1) - rate = jnn.sigmoid(logit(scalar_baselines[signal]) + deviation) - numpyro.deterministic(f"{self.name}_{signal}", rate) - result[signal] = rate - - return result diff --git a/pyrenew/model/multisignal_model.py b/pyrenew/model/multisignal_model.py index e20e95d5a..97f07bb7a 100644 --- a/pyrenew/model/multisignal_model.py +++ b/pyrenew/model/multisignal_model.py @@ -15,7 +15,7 @@ import numpyro.handlers from jax.typing import ArrayLike -from pyrenew.ascertainment import AscertainmentModel, TimeVaryingAscertainment +from pyrenew.ascertainment import AscertainmentModel from pyrenew.ascertainment.context import ascertainment_context from pyrenew.latent.base import BaseLatentInfectionProcess from pyrenew.metaclass import Model @@ -125,39 +125,6 @@ def validate(self) -> None: f"the model name {ascertainment_model.name!r}." ) - namespace_owners: dict[str, list[str]] = {} - for name, ascertainment_model in self.ascertainment_models.items(): - namespace_owners.setdefault(name, []).append( - f"registered ascertainment model {name!r}" - ) - if not isinstance(ascertainment_model, TimeVaryingAscertainment): - continue - - baseline = ascertainment_model.baseline_model - registered_baseline_names = [ - registered_name - for registered_name, registered_model in self.ascertainment_models.items() - if registered_model is baseline - ] - if registered_baseline_names: - raise ValueError( - f"Time-varying ascertainment model {name!r} owns baseline " - f"model {baseline.name!r}; that baseline object must not " - f"also be registered as {registered_baseline_names[0]!r}." - ) - namespace_owners.setdefault(baseline.name, []).append( - f"baseline model {baseline.name!r} owned by time-varying " - f"ascertainment model {name!r}" - ) - - for namespace, owners in namespace_owners.items(): - if len(owners) > 1: - raise ValueError( - f"Ascertainment NumPyro namespace {namespace!r} is used by " - f"{' and '.join(owners)}. Registered models and wrapped " - "baselines must have unique names." - ) - def _validate_run_args( self, n_days_post_init: int, diff --git a/pyrenew/model/pyrenew_builder.py b/pyrenew/model/pyrenew_builder.py index 9a5ccef62..045bc9ab2 100644 --- a/pyrenew/model/pyrenew_builder.py +++ b/pyrenew/model/pyrenew_builder.py @@ -153,42 +153,45 @@ def add_ascertainment( self, ascertainment_model: AscertainmentModel, ) -> PyrenewBuilder: - """ - Add shared ascertainment structure to the model. + """Add ascertainment rates for one or more observation signals. - Use this method when observation probabilities are related across - signals. Independent scalar ascertainment rates do not require an - ascertainment model; those can be passed directly to an observation - process as ordinary ``RandomVariable`` objects. + Choose the ascertainment class based on the relationship among scalar + baseline rates. ``IndependentAscertainment`` is the usual choice for a + single signal or independent baselines. ``JointAscertainment`` and + ``RatioLinkedAscertainment`` provide related baselines. Any class can + receive optional temporal processes for selected signals. A registered ascertainment model is sampled once per model execution, - before observation processes run. It may return a scalar rate for each - signal or a trajectory spanning the full padded model axis. - Observation processes receive signal-specific accessors from - ``ascertainment_model.for_signal(...)``: + before observation processes run. Signals without a temporal process + provide a scalar rate; signals with one provide a trajectory spanning + the full padded model axis. + + Fixed and time-varying models use the same registration and observation + wiring. For example, omit ``temporal_processes`` for a fixed ED rate or + supply it as shown for a time-varying rate: ```python - ascertainment = JointAscertainment( - name="he_ascertainment", - signals=("hospital", "ed_visits"), - baseline_rates=..., - scale_tril=..., + ascertainment = IndependentAscertainment( + name="ed_ascertainment", + rate_rvs={"ed_visits": iedr_rv}, + temporal_processes={"ed_visits": ed_temporal_process}, ) builder.add_ascertainment(ascertainment) + ``` + Observation processes receive signal-specific accessors from + ``ascertainment_model.for_signal(...)``: + + ```python builder.add_observation( PopulationCounts( - name="hospital", - ascertainment_rate_rv=ascertainment.for_signal("hospital"), + name="ed_visits", + ascertainment_rate_rv=ascertainment.for_signal("ed_visits"), ... ) ) ``` - ``TimeVaryingAscertainment`` owns and samples its scalar baseline - model. Register only the outer time-varying model, and use that outer - model's ``for_signal(...)`` accessors in observation processes. - The ascertainment model's ``name`` attribute is used as the unique identifier in the built ``MultiSignalModel``. @@ -196,7 +199,8 @@ def add_ascertainment( ---------- ascertainment_model Configured ascertainment model instance, such as - ``JointAscertainment`` or ``TimeVaryingAscertainment``. + ``IndependentAscertainment``, ``JointAscertainment``, or + ``RatioLinkedAscertainment``. Returns ------- diff --git a/test/integration/conftest.py b/test/integration/conftest.py index 6d8f264bd..b8f0d8b4c 100644 --- a/test/integration/conftest.py +++ b/test/integration/conftest.py @@ -12,7 +12,11 @@ import polars as pl import pytest -from pyrenew.ascertainment import JointAscertainment, RatioLinkedAscertainment +from pyrenew.ascertainment import ( + IndependentAscertainment, + JointAscertainment, + RatioLinkedAscertainment, +) from pyrenew.datasets import ( load_example_infection_admission_interval, load_synthetic_daily_ed_visits, @@ -172,6 +176,15 @@ def _build_he_population_model( # numpydoc ignore=RT01 infection_process=infection_process, ) + ascertainment = IndependentAscertainment( + name="he_ascertainment", + rate_rvs={ + "hospital": DistributionalVariable("ihr", dist.Beta(1, 100)), + "ed": DistributionalVariable("iedr", dist.Beta(1, 100)), + }, + ) + builder.add_ascertainment(ascertainment) + hospital_kwargs = {} if hospital_weekly: hospital_kwargs = { @@ -183,7 +196,7 @@ def _build_he_population_model( # numpydoc ignore=RT01 builder.add_observation( PopulationCounts( name="hospital", - ascertainment_rate_rv=DistributionalVariable("ihr", dist.Beta(1, 100)), + ascertainment_rate_rv=ascertainment.for_signal("hospital"), delay_distribution_rv=DeterministicPMF("hosp_delay", hosp_delay_pmf), noise=NegativeBinomialNoise( DistributionalVariable("hosp_conc", dist.LogNormal(5.0, 1.0)) @@ -194,7 +207,7 @@ def _build_he_population_model( # numpydoc ignore=RT01 builder.add_observation( PopulationCounts( name="ed", - ascertainment_rate_rv=DistributionalVariable("iedr", dist.Beta(1, 100)), + ascertainment_rate_rv=ascertainment.for_signal("ed"), delay_distribution_rv=DeterministicPMF("ed_delay", ed_delay_pmf), noise=NegativeBinomialNoise( DistributionalVariable("ed_conc", dist.LogNormal(4.0, 1.0)) diff --git a/test/integration/test_population_infections_he_weekly_joint_ascertainment.py b/test/integration/test_population_infections_he_weekly_joint_ascertainment.py index 53b34f91f..936fc126c 100644 --- a/test/integration/test_population_infections_he_weekly_joint_ascertainment.py +++ b/test/integration/test_population_infections_he_weekly_joint_ascertainment.py @@ -190,8 +190,17 @@ def test_joint_ascertainment_sites_are_sampled_once( assert trace["he_ascertainment_eta"]["type"] == "sample" assert trace["he_ascertainment_eta"]["value"].shape == (2,) - assert trace["he_ascertainment_hospital"]["type"] == "deterministic" - assert trace["he_ascertainment_ed_visits"]["type"] == "deterministic" + + for signal in ("hospital", "ed_visits"): + baseline_site = trace[f"he_ascertainment_baseline_{signal}"] + final_site = trace[f"he_ascertainment_{signal}"] + + assert baseline_site["type"] == "deterministic" + assert baseline_site["value"].shape == () + assert final_site["type"] == "deterministic" + assert final_site["value"].shape == () + assert jnp.allclose(final_site["value"], baseline_site["value"]) + assert trace["hospital_predicted"]["value"].shape == (n_periods,) assert trace["hospital_predicted_daily"]["value"].shape == (n_total,) assert trace["ed_visits_predicted"]["value"].shape == (n_total,) diff --git a/test/integration/test_population_infections_he_weekly_ratio_linked_ascertainment.py b/test/integration/test_population_infections_he_weekly_ratio_linked_ascertainment.py index 64b051cd3..42114c2df 100644 --- a/test/integration/test_population_infections_he_weekly_ratio_linked_ascertainment.py +++ b/test/integration/test_population_infections_he_weekly_ratio_linked_ascertainment.py @@ -139,15 +139,27 @@ def test_ratio_linked_rates_and_prediction_shapes( base_rate = trace["iedr"]["value"] ratio = trace["ihr_rel_iedr"]["value"] + ed_baseline = trace["he_ascertainment_baseline_ed_visits"]["value"] + hospital_baseline = trace["he_ascertainment_baseline_hospital"]["value"] ed_rate = trace["he_ascertainment_ed_visits"]["value"] hospital_rate = trace["he_ascertainment_hospital"]["value"] assert trace["iedr"]["type"] == "sample" + assert base_rate.shape == () assert trace["ihr_rel_iedr"]["type"] == "sample" + assert ratio.shape == () + assert trace["he_ascertainment_baseline_ed_visits"]["type"] == "deterministic" + assert ed_baseline.shape == () + assert trace["he_ascertainment_baseline_hospital"]["type"] == "deterministic" + assert hospital_baseline.shape == () assert trace["he_ascertainment_ed_visits"]["type"] == "deterministic" + assert ed_rate.shape == () assert trace["he_ascertainment_hospital"]["type"] == "deterministic" - assert jnp.allclose(ed_rate, base_rate) - assert jnp.allclose(hospital_rate, base_rate * ratio) + assert hospital_rate.shape == () + assert jnp.allclose(ed_baseline, base_rate) + assert jnp.allclose(hospital_baseline, base_rate * ratio) + assert jnp.allclose(ed_rate, ed_baseline) + assert jnp.allclose(hospital_rate, hospital_baseline) n_total = model.latent.n_initialization_points + N_DAYS_FIT hospital = model.observations["hospital"] diff --git a/test/test_ascertainment.py b/test/test_ascertainment.py index 8019f69f3..f45a11d7d 100644 --- a/test/test_ascertainment.py +++ b/test/test_ascertainment.py @@ -4,6 +4,7 @@ from collections.abc import Mapping +import jax import jax.numpy as jnp import numpyro import numpyro.distributions as dist @@ -14,16 +15,17 @@ from pyrenew.ascertainment import ( AscertainmentModel, AscertainmentSignal, + IndependentAscertainment, JointAscertainment, RatioLinkedAscertainment, - TimeVaryingAscertainment, ) from pyrenew.ascertainment.context import ( ascertainment_context, get_ascertainment_value, ) from pyrenew.deterministic import DeterministicVariable -from pyrenew.latent import WeeklyTemporalProcess +from pyrenew.latent import TemporalProcess, WeeklyTemporalProcess +from pyrenew.metaclass import RandomVariable from pyrenew.randomvariable import DistributionalVariable @@ -37,11 +39,13 @@ def __init__( initial: float = 0.0, increment: float = 0.0, result: ArrayLike | None = None, + requires_calendar_anchor: bool = False, ) -> None: """Initialize a deterministic temporal-process test double.""" self.initial = initial self.increment = increment self.result = result + self.requires_calendar_anchor = requires_calendar_anchor self.sample_calls: list[dict[str, object]] = [] def sample( @@ -75,6 +79,27 @@ def sample( return jnp.broadcast_to(values[:, None], (n_timepoints, n_processes)) +class CountingVariable(RandomVariable): + """Return a configured value and count sampling calls.""" + + def __init__(self, name: str, value: ArrayLike) -> None: + """Initialize a counting random variable.""" + super().__init__(name=name) + self.value = value + self.n_calls = 0 + + def sample(self, **kwargs: object) -> ArrayLike: + """Return the configured value and increment the call count. + + Returns + ------- + ArrayLike + Configured value. + """ + self.n_calls += 1 + return self.value + + class FixedAscertainmentModel(AscertainmentModel): """Return fixed ascertainment values for baseline-contract tests.""" @@ -82,31 +107,24 @@ def __init__( self, name: str, values: Mapping[str, ArrayLike], - requires_calendar_anchor: bool = False, + temporal_processes: Mapping[str, TemporalProcess] | None = None, ) -> None: """Initialize a fixed ascertainment test double.""" - super().__init__(name=name, signals=tuple(values)) + super().__init__( + name=name, + signals=tuple(values), + temporal_processes=temporal_processes, + ) self.values = dict(values) - self._requires_calendar_anchor = requires_calendar_anchor self.n_sample_calls = 0 - def requires_calendar_anchor(self) -> bool: - """Return the configured calendar requirement. - - Returns - ------- - bool - Configured calendar requirement. - """ - return self._requires_calendar_anchor - - def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: - """Return configured values and record the sampling call. + def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + """Return configured scalar baselines and record the sampling call. Returns ------- Mapping[str, ArrayLike] - Configured signal values. + Configured signal baselines. """ self.n_sample_calls += 1 return self.values @@ -115,6 +133,286 @@ def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]: class TestAscertainmentModelContract: """Test validation shared by all ascertainment models.""" + def test_baseline_sampling_is_the_only_abstract_method(self) -> None: + """Test custom subclasses only need to implement baseline sampling.""" + assert AscertainmentModel.__abstractmethods__ == frozenset( + {"_sample_baseline_rates"} + ) + + def test_temporal_processes_default_to_an_empty_dictionary(self) -> None: + """Test models without temporal processes store an empty dictionary.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + assert model.temporal_processes == {} + + def test_accepts_an_empty_temporal_process_mapping(self) -> None: + """Test an explicitly empty temporal mapping is valid.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2}, + temporal_processes={}, + ) + + assert model.temporal_processes == {} + + def test_temporal_processes_are_copied_in_signal_order(self) -> None: + """Test process storage follows signal order and does not alias input.""" + hospital_process = DeterministicTemporalProcess() + ed_process = DeterministicTemporalProcess() + processes: dict[str, TemporalProcess] = { + "ed": ed_process, + "hospital": hospital_process, + } + + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2, "ed": 0.3}, + temporal_processes=processes, + ) + processes.clear() + + assert tuple(model.temporal_processes) == ("hospital", "ed") + assert model.temporal_processes["hospital"] is hospital_process + assert model.temporal_processes["ed"] is ed_process + + def test_temporal_processes_may_cover_a_signal_subset(self) -> None: + """Test only selected signals need temporal processes.""" + process = DeterministicTemporalProcess() + + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2, "ed": 0.3}, + temporal_processes={"ed": process}, + ) + + assert model.temporal_processes == {"ed": process} + + def test_rejects_non_mapping_temporal_processes(self) -> None: + """Test temporal process configuration must be a mapping.""" + with pytest.raises(TypeError, match="temporal_processes.*mapping"): + FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2}, + temporal_processes=[], # type: ignore[arg-type] + ) + + def test_rejects_unknown_temporal_process_signal(self) -> None: + """Test processes cannot be configured for undeclared signals.""" + with pytest.raises(ValueError, match="unknown signals.*ed"): + FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2}, + temporal_processes={"ed": DeterministicTemporalProcess()}, + ) + + def test_rejects_invalid_temporal_process(self) -> None: + """Test every process must satisfy the runtime protocol.""" + with pytest.raises(TypeError, match="hospital.*TemporalProcess"): + FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2}, + temporal_processes={"hospital": object()}, # type: ignore[dict-item] + ) + + def test_calendar_requirement_checks_all_temporal_processes(self) -> None: + """Test any calendar-aligned process makes the model require a date.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2, "ed": 0.3}, + temporal_processes={ + "hospital": DeterministicTemporalProcess(), + "ed": DeterministicTemporalProcess(requires_calendar_anchor=True), + }, + ) + + assert model.requires_calendar_anchor() + + def test_daily_process_does_not_require_a_calendar_anchor(self) -> None: + """Test ordinary daily processes do not require a date.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2}, + temporal_processes={"hospital": DeterministicTemporalProcess()}, + ) + + assert not model.requires_calendar_anchor() + + @pytest.mark.parametrize("n_timepoints", [True, 1.0, jnp.array(1)]) + def test_rejects_non_integer_n_timepoints(self, n_timepoints: object) -> None: + """Test only built-in integers are accepted for the axis length.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + with pytest.raises(TypeError, match="n_timepoints.*positive integer"): + model.sample(n_timepoints=n_timepoints) # type: ignore[arg-type] + + assert model.n_sample_calls == 0 + + @pytest.mark.parametrize("n_timepoints", [0, -1]) + def test_rejects_non_positive_n_timepoints(self, n_timepoints: int) -> None: + """Test the shared axis must contain at least one timepoint.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + with pytest.raises(ValueError, match="n_timepoints.*positive integer"): + model.sample(n_timepoints=n_timepoints) + + assert model.n_sample_calls == 0 + + def test_rejects_non_mapping_baseline_rates(self) -> None: + """Test baseline sampling must return a mapping.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + model.values = [0.2] # type: ignore[assignment] + + with pytest.raises(TypeError, match="baseline rates.*mapping"): + model.sample(n_timepoints=4) + + @pytest.mark.parametrize( + "values, expected_missing, expected_extra", + [ + ({}, ("hospital",), ()), + ({"hospital": 0.2, "ed": 0.3}, (), ("ed",)), + ({"ed": 0.3}, ("hospital",), ("ed",)), + ], + ) + def test_rejects_missing_or_extra_baseline_signals( + self, + values: Mapping[str, ArrayLike], + expected_missing: tuple[str, ...], + expected_extra: tuple[str, ...], + ) -> None: + """Test sampled baseline keys must exactly match declared signals.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + model.values = dict(values) + + with pytest.raises(ValueError) as error: + model.sample(n_timepoints=4) + + assert f"Missing: {expected_missing}" in str(error.value) + assert f"Extra: {expected_extra}" in str(error.value) + + @pytest.mark.parametrize("value", [jnp.ones(1), jnp.ones((1, 1))]) + def test_rejects_non_scalar_baselines(self, value: ArrayLike) -> None: + """Test baseline samplers cannot return vectors or trajectories.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": value}) + + with pytest.raises( + ValueError, + match=r"hospital.*baseline shape.*required shape.*\(\)", + ): + model.sample(n_timepoints=4) + + @pytest.mark.parametrize("value", [0.0, 1.0]) + def test_fixed_baselines_allow_closed_interval_endpoints( + self, + value: float, + ) -> None: + """Test fixed scalar baselines may equal zero or one.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": value}) + + result = model.sample(n_timepoints=4) + + assert result["hospital"] == value + + @pytest.mark.parametrize("value", [-0.1, 1.1, jnp.nan, jnp.inf]) + def test_rejects_invalid_fixed_baselines(self, value: float) -> None: + """Test fixed baselines must be finite probabilities.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": value}) + + with pytest.raises(ValueError, match=r"hospital.*finite.*\[0, 1\]"): + model.sample(n_timepoints=4) + + @pytest.mark.parametrize("value", [0.0, 1.0, -0.1, 1.1, jnp.nan]) + def test_temporal_baselines_require_open_interval( + self, + value: float, + ) -> None: + """Test logit-transformed baselines must be strictly between bounds.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": value}, + temporal_processes={"hospital": DeterministicTemporalProcess()}, + ) + + with pytest.raises(ValueError, match=r"hospital.*finite.*\(0, 1\)"): + model.sample(n_timepoints=4) + + def test_traced_baseline_values_do_not_trigger_python_validation(self) -> None: + """Test baseline value checks remain compatible with JAX tracing.""" + model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) + + @jax.jit + def validate(value: ArrayLike) -> ArrayLike: + """Validate and return a traced baseline value.""" + return model._validate_baseline_rates({"hospital": value})["hospital"] + + assert jnp.isnan(validate(jnp.nan)) + + def test_validates_all_baselines_before_recording_or_temporal_sampling( + self, + ) -> None: + """Test invalid baselines fail before deterministic or process sites.""" + process = DeterministicTemporalProcess() + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2, "ed": 1.1}, + temporal_processes={"hospital": process}, + ) + + with numpyro.handlers.trace() as trace: + with pytest.raises(ValueError, match="ed"): + model.sample(n_timepoints=4) + + assert trace == {} + assert process.sample_calls == [] + + def test_samples_fixed_and_temporal_rates_in_signal_order(self) -> None: + """Test fixed scalars and varying trajectories share one interface.""" + process = DeterministicTemporalProcess(increment=0.2) + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2, "ed": 0.3}, + temporal_processes={"ed": process}, + ) + + with numpyro.handlers.trace() as trace: + values = model.sample(n_timepoints=4, first_day_dow=3) + + assert tuple(values) == ("hospital", "ed") + assert values["hospital"].shape == () + assert values["ed"].shape == (4,) + assert set(trace) == { + "ascertainment_baseline_hospital", + "ascertainment_hospital", + "ascertainment_baseline_ed", + "ascertainment_ed", + } + assert process.sample_calls == [ + { + "n_timepoints": 4, + "initial_value": 0.0, + "n_processes": 1, + "name_prefix": "ascertainment_ed", + "first_day_dow": 3, + } + ] + expected_ed = jax.nn.sigmoid(logit(0.3) + 0.2 * jnp.arange(4)) + assert jnp.allclose(values["ed"], expected_ed) + + def test_rejects_malformed_temporal_output(self) -> None: + """Test temporal output keeps its singleton process dimension.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.2}, + temporal_processes={ + "hospital": DeterministicTemporalProcess(result=jnp.zeros(4)) + }, + ) + + with pytest.raises( + ValueError, + match=r"ascertainment.*hospital.*\(4,\).*\(4, 1\)", + ): + model.sample(n_timepoints=4) + def test_accepts_scalar_and_full_axis_values(self) -> None: """Test scalar and exact full-axis trajectories are valid.""" model = FixedAscertainmentModel( @@ -170,130 +468,227 @@ def test_rejects_invalid_sample_shapes(self, value: ArrayLike) -> None: ): model.validate_sampled_values({"hospital": value}, n_timepoints=4) - def test_scalar_models_do_not_require_a_calendar_anchor(self) -> None: - """Test the base calendar declaration defaults to false.""" - model = JointAscertainment( - name="ascertainment", - signals=("hospital",), - baseline_rates=jnp.array([0.2]), - scale_tril=jnp.eye(1), - ) - assert not model.requires_calendar_anchor() +def _make_public_ascertainment( + model_kind: str, + temporal_processes: Mapping[str, TemporalProcess] | None, +) -> AscertainmentModel: + """Construct a public ascertainment class with two scalar baselines.""" + if model_kind == "independent": + return IndependentAscertainment( + name="common", + rate_rvs={ + "hospital": DeterministicVariable("common_ihr", 0.2), + "ed": DeterministicVariable("common_iedr", 0.3), + }, + temporal_processes=temporal_processes, + ) + if model_kind == "joint": + return JointAscertainment( + name="common", + signals=("hospital", "ed"), + baseline_rates=jnp.array([0.2, 0.3]), + scale_tril=jnp.eye(2) * 0.1, + temporal_processes=temporal_processes, + ) + if model_kind == "ratio": + return RatioLinkedAscertainment( + name="common", + base_signal="ed", + linked_signal="hospital", + base_rate_rv=DeterministicVariable("common_iedr", 0.3), + ratio_rv=DeterministicVariable("common_ratio", 2.0 / 3.0), + temporal_processes=temporal_processes, + ) + raise ValueError(f"Unknown model kind {model_kind!r}.") -class TestTimeVaryingAscertainmentConstruction: - """Test time-varying ascertainment constructor validation.""" +class TestPublicAscertainmentBehavior: + """Test common fixed and temporal behavior for every public class.""" - def test_requires_an_ascertainment_baseline(self) -> None: - """Test the baseline must implement the ascertainment contract.""" - with pytest.raises(TypeError, match="baseline_model"): - TimeVaryingAscertainment( - name="outer", - baseline_model=object(), # type: ignore[arg-type] - processes={"hospital": DeterministicTemporalProcess()}, - ) + @pytest.mark.parametrize("model_kind", ["independent", "joint", "ratio"]) + @pytest.mark.parametrize( + "varying_signals", + [(), ("ed",), ("hospital", "ed")], + ids=["fixed", "partial", "full"], + ) + def test_fixed_partial_and_fully_varying_rates( + self, + model_kind: str, + varying_signals: tuple[str, ...], + ) -> None: + """Test common output shapes and deterministic names across classes.""" + temporal_processes = { + signal: DeterministicTemporalProcess(increment=0.1) + for signal in varying_signals + } + ascertainment = _make_public_ascertainment( + model_kind, + temporal_processes or None, + ) - def test_requires_distinct_outer_and_baseline_names(self) -> None: - """Test outer and baseline NumPyro namespaces must differ.""" - baseline = FixedAscertainmentModel("same", {"hospital": 0.2}) + with numpyro.handlers.seed(rng_seed=42): + with numpyro.handlers.trace() as trace: + values = ascertainment.sample(n_timepoints=4) - with pytest.raises(ValueError, match="different names"): - TimeVaryingAscertainment( - name="same", - baseline_model=baseline, - processes={"hospital": DeterministicTemporalProcess()}, + assert set(trace) >= { + "common_baseline_hospital", + "common_baseline_ed", + "common_hospital", + "common_ed", + } + assert trace["common_baseline_hospital"]["value"].shape == () + assert trace["common_baseline_ed"]["value"].shape == () + assert trace["common_hospital"]["type"] == "deterministic" + assert trace["common_ed"]["type"] == "deterministic" + for signal in ("hospital", "ed"): + expected_shape = (4,) if signal in varying_signals else () + assert values[signal].shape == expected_shape + assert trace[f"common_{signal}"]["value"].shape == expected_shape + + +class TestIndependentAscertainmentValidation: + """Test IndependentAscertainment constructor validation.""" + + def test_requires_a_rate_rv_mapping(self) -> None: + """Test rate_rvs must be supplied as a mapping.""" + with pytest.raises(TypeError, match="rate_rvs.*mapping"): + IndependentAscertainment( + name="ascertainment", + rate_rvs=[], # type: ignore[arg-type] ) - def test_requires_a_process_mapping(self) -> None: - """Test temporal processes must be provided in a mapping.""" - baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + def test_requires_at_least_one_rate_rv(self) -> None: + """Test at least one signal must be configured.""" + with pytest.raises(ValueError, match="rate_rvs.*non-empty"): + IndependentAscertainment(name="ascertainment", rate_rvs={}) + + @pytest.mark.parametrize("signal", ["", None]) + def test_requires_non_empty_string_signals(self, signal: object) -> None: + """Test rate_rv keys must be non-empty strings.""" + with pytest.raises(ValueError, match="keys.*non-empty strings"): + IndependentAscertainment( + name="ascertainment", + rate_rvs={signal: DeterministicVariable("rate", 0.2)}, # type: ignore[dict-item] + ) - with pytest.raises(TypeError, match="processes must be a mapping"): - TimeVaryingAscertainment( - name="outer", - baseline_model=baseline, - processes=[], # type: ignore[arg-type] + def test_requires_random_variable_values(self) -> None: + """Test each baseline sampler must be a RandomVariable.""" + with pytest.raises(TypeError, match="hospital.*RandomVariable"): + IndependentAscertainment( + name="ascertainment", + rate_rvs={"hospital": object()}, # type: ignore[dict-item] ) - @pytest.mark.parametrize( - "processes", - [ - {}, - { - "hospital": DeterministicTemporalProcess(), - "ed": DeterministicTemporalProcess(), - }, - ], - ) - def test_requires_exact_baseline_signal_keys( - self, - processes: Mapping[str, DeterministicTemporalProcess], - ) -> None: - """Test process signal names must exactly match baseline signals.""" - baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + def test_preserves_mapping_order_and_copies_rate_rvs(self) -> None: + """Test insertion order defines signals without aliasing the input.""" + ed_rv = DeterministicVariable("iedr", 0.2) + hospital_rv = DeterministicVariable("ihr", 0.1) + rate_rvs: dict[str, RandomVariable] = { + "ed": ed_rv, + "hospital": hospital_rv, + } - with pytest.raises(ValueError, match="baseline model signals"): - TimeVaryingAscertainment("outer", baseline, processes) + ascertainment = IndependentAscertainment( + name="ascertainment", + rate_rvs=rate_rvs, + ) + rate_rvs.clear() - def test_requires_temporal_process_protocol(self) -> None: - """Test every configured process must satisfy TemporalProcess.""" - baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) + assert ascertainment.signals == ("ed", "hospital") + assert ascertainment.rate_rvs == { + "ed": ed_rv, + "hospital": hospital_rv, + } - with pytest.raises(TypeError, match="TemporalProcess"): - TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": object()}, # type: ignore[dict-item] - ) - def test_orders_processes_like_baseline_signals(self) -> None: - """Test deterministic process order follows the baseline model.""" - baseline = FixedAscertainmentModel( - "baseline", - {"hospital": 0.2, "ed": 0.3}, +class TestIndependentAscertainmentSampling: + """Test IndependentAscertainment sampling behavior.""" + + def test_calls_each_random_variable_once(self) -> None: + """Test every configured baseline variable is sampled exactly once.""" + ed_rv = CountingVariable("iedr", 0.2) + hospital_rv = CountingVariable("ihr", 0.1) + ascertainment = IndependentAscertainment( + name="ascertainment", + rate_rvs={"ed": ed_rv, "hospital": hospital_rv}, ) - processes = { - "ed": DeterministicTemporalProcess(), - "hospital": DeterministicTemporalProcess(), + + with numpyro.handlers.trace() as trace: + values = ascertainment.sample(n_timepoints=4) + + assert tuple(values) == ("ed", "hospital") + assert ed_rv.n_calls == 1 + assert hospital_rv.n_calls == 1 + assert set(trace) == { + "ascertainment_baseline_ed", + "ascertainment_ed", + "ascertainment_baseline_hospital", + "ascertainment_hospital", } - model = TimeVaryingAscertainment("outer", baseline, processes) + def test_one_signal_trace_retains_underlying_rv_name(self) -> None: + """Test the user-named sample site and standard sites are recorded.""" + ascertainment = IndependentAscertainment( + name="ascertainment", + rate_rvs={"ed": DistributionalVariable("iedr", dist.Delta(0.2))}, + ) + + with numpyro.handlers.seed(rng_seed=42): + with numpyro.handlers.trace() as trace: + values = ascertainment.sample(n_timepoints=4) + + assert ascertainment.signals == ("ed",) + assert trace["iedr"]["type"] == "sample" + assert trace["ascertainment_baseline_ed"]["type"] == "deterministic" + assert trace["ascertainment_ed"]["type"] == "deterministic" + assert jnp.array_equal(values["ed"], trace["iedr"]["value"]) + + def test_supports_calendar_aligned_temporal_process(self) -> None: + """Test independent rates can use calendar-aligned trajectories.""" + weekly = WeeklyTemporalProcess( + DeterministicTemporalProcess(increment=1.0), + start_dow=0, + ) + ascertainment = IndependentAscertainment( + name="ascertainment", + rate_rvs={"hospital": DeterministicVariable("ihr", 0.2)}, + temporal_processes={"hospital": weekly}, + ) + + values = ascertainment.sample(n_timepoints=10, first_day_dow=2) - assert tuple(model.processes) == ("hospital", "ed") + assert ascertainment.requires_calendar_anchor() + assert jnp.allclose(values["hospital"][:5], 0.2) + assert jnp.all(values["hospital"][5:] > 0.2) -class TestTimeVaryingAscertainmentSampling: - """Test baseline and temporal sampling semantics.""" +class TestConcreteTemporalAscertainment: + """Test temporal behavior supplied by concrete baseline relationships.""" - def test_joint_baseline_sites_and_independent_trajectories(self) -> None: - """Test a joint baseline is sampled once before independent deviations.""" + def test_joint_baselines_support_independent_trajectories(self) -> None: + """Test joint baselines are sampled before separate deviations.""" baseline_rates = jnp.array([0.2, 0.4]) - baseline = JointAscertainment( - name="he_baseline", + ascertainment = JointAscertainment( + name="he_ascertainment", signals=("hospital", "ed"), baseline_rates=baseline_rates, scale_tril=jnp.eye(2), - ) - model = TimeVaryingAscertainment( - name="he_ascertainment", - baseline_model=baseline, - processes={ + temporal_processes={ "hospital": DeterministicTemporalProcess(increment=0.3), "ed": DeterministicTemporalProcess(increment=-0.2), }, ) with numpyro.handlers.substitute( - data={"he_baseline_eta": logit(baseline_rates)} + data={"he_ascertainment_eta": logit(baseline_rates)} ): with numpyro.handlers.trace() as trace: - values = model.sample(n_timepoints=4) + values = ascertainment.sample(n_timepoints=4) - assert trace["he_baseline_eta"]["type"] == "sample" - assert trace["he_baseline_eta"]["value"].shape == (2,) - assert trace["he_baseline_hospital"]["value"].shape == () - assert trace["he_baseline_ed"]["value"].shape == () + assert trace["he_ascertainment_eta"]["type"] == "sample" + assert trace["he_ascertainment_baseline_hospital"]["value"].shape == () + assert trace["he_ascertainment_baseline_ed"]["value"].shape == () assert trace["he_ascertainment_hospital"]["value"].shape == (4,) assert trace["he_ascertainment_ed"]["value"].shape == (4,) assert jnp.allclose(values["hospital"][0], baseline_rates[0]) @@ -304,18 +699,14 @@ def test_joint_baseline_sites_and_independent_trajectories(self) -> None: ) def test_ratio_linked_baseline_relationship_is_not_pointwise(self) -> None: - """Test linked baseline sites persist without constraining trajectories.""" - baseline = RatioLinkedAscertainment( - name="he_baseline", + """Test temporal deviations need not preserve the baseline ratio.""" + ascertainment = RatioLinkedAscertainment( + name="he_ascertainment", base_signal="ed", linked_signal="hospital", base_rate_rv=DistributionalVariable("iedr", dist.Delta(0.4)), ratio_rv=DistributionalVariable("ihr_rel_iedr", dist.Delta(0.5)), - ) - model = TimeVaryingAscertainment( - name="he_ascertainment", - baseline_model=baseline, - processes={ + temporal_processes={ "ed": DeterministicTemporalProcess(increment=-0.1), "hospital": DeterministicTemporalProcess(increment=0.2), }, @@ -323,83 +714,22 @@ def test_ratio_linked_baseline_relationship_is_not_pointwise(self) -> None: with numpyro.handlers.seed(rng_seed=42): with numpyro.handlers.trace() as trace: - values = model.sample(n_timepoints=4) + values = ascertainment.sample(n_timepoints=4) assert trace["iedr"]["type"] == "sample" assert trace["ihr_rel_iedr"]["type"] == "sample" - assert jnp.allclose(trace["he_baseline_ed"]["value"], 0.4) - assert jnp.allclose(trace["he_baseline_hospital"]["value"], 0.2) - assert values["ed"].shape == (4,) - assert values["hospital"].shape == (4,) + assert jnp.allclose( + trace["he_ascertainment_baseline_ed"]["value"], + 0.4, + ) + assert jnp.allclose( + trace["he_ascertainment_baseline_hospital"]["value"], + 0.2, + ) trajectory_ratios = values["hospital"] / values["ed"] assert jnp.allclose(trajectory_ratios[0], 0.5) assert not jnp.allclose(trajectory_ratios[1:], 0.5) - def test_calendar_requirement_and_weekly_alignment(self) -> None: - """Test weekly processes declare and use the model-axis weekday.""" - baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) - weekly = WeeklyTemporalProcess( - DeterministicTemporalProcess(increment=1.0), - start_dow=0, - ) - model = TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": weekly}, - ) - - values = model.sample(n_timepoints=10, first_day_dow=2) - - assert model.requires_calendar_anchor() - assert jnp.allclose(values["hospital"][:5], 0.2) - assert jnp.all(values["hospital"][5:] > 0.2) - - @pytest.mark.parametrize("baseline_value", [0.0, 1.0, -0.1, 1.1, jnp.nan]) - def test_rejects_invalid_concrete_baselines( - self, - baseline_value: float, - ) -> None: - """Test concrete baselines must be finite probabilities.""" - baseline = FixedAscertainmentModel( - "baseline", - {"hospital": baseline_value}, - ) - model = TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": DeterministicTemporalProcess()}, - ) - - with pytest.raises(ValueError, match=r"hospital.*\(0, 1\)"): - model.sample(n_timepoints=4) - - def test_rejects_trajectory_baseline(self) -> None: - """Test nested time-varying baselines are rejected.""" - baseline = FixedAscertainmentModel( - "baseline", - {"hospital": jnp.full(4, 0.2)}, - ) - model = TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": DeterministicTemporalProcess()}, - ) - - with pytest.raises(ValueError, match=r"scalar baseline.*\(4,\).*"): - model.sample(n_timepoints=4) - - def test_rejects_malformed_temporal_output_before_squeezing(self) -> None: - """Test temporal output must retain the singleton process dimension.""" - baseline = FixedAscertainmentModel("baseline", {"hospital": 0.2}) - model = TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": DeterministicTemporalProcess(result=jnp.zeros(4))}, - ) - - with pytest.raises(ValueError, match=r"outer.*hospital.*\(4,\).*\(4, 1\)"): - model.sample(n_timepoints=4) - class TestJointAscertainmentValidation: """Test JointAscertainment constructor validation.""" @@ -576,11 +906,13 @@ def test_sample_creates_one_joint_sample_site_and_signal_deterministics(self): with numpyro.handlers.seed(rng_seed=42): with numpyro.handlers.trace() as trace: - values = ascertainment.sample() + values = ascertainment.sample(n_timepoints=4) assert set(values) == {"hospital", "ed"} assert trace["he_ascertainment_eta"]["type"] == "sample" assert trace["he_ascertainment_eta"]["value"].shape == (2,) + assert trace["he_ascertainment_baseline_hospital"]["type"] == ("deterministic") + assert trace["he_ascertainment_baseline_ed"]["type"] == "deterministic" assert trace["he_ascertainment_hospital"]["type"] == "deterministic" assert trace["he_ascertainment_ed"]["type"] == "deterministic" assert jnp.array_equal( @@ -602,7 +934,7 @@ def test_sample_accepts_covariance_matrix_parameterization(self): ) with numpyro.handlers.seed(rng_seed=42): - values = ascertainment.sample() + values = ascertainment.sample(n_timepoints=4) assert set(values) == {"hospital", "ed"} @@ -616,7 +948,7 @@ def test_sample_accepts_precision_matrix_parameterization(self): ) with numpyro.handlers.seed(rng_seed=42): - values = ascertainment.sample() + values = ascertainment.sample(n_timepoints=4) assert set(values) == {"hospital", "ed"} @@ -757,13 +1089,20 @@ def test_sample_returns_base_rate_and_scaled_linked_rate(self): ) with numpyro.handlers.trace() as trace: - values = ascertainment.sample() + values = ascertainment.sample(n_timepoints=4) assert set(values) == {"ed_visits", "hospital"} assert jnp.array_equal(values["ed_visits"], jnp.array(0.2)) assert jnp.array_equal(values["hospital"], jnp.array(0.1)) + assert trace["he_ascertainment_baseline_ed_visits"]["type"] == ("deterministic") + assert trace["he_ascertainment_baseline_hospital"]["type"] == ("deterministic") assert trace["he_ascertainment_ed_visits"]["type"] == "deterministic" assert trace["he_ascertainment_hospital"]["type"] == "deterministic" + assert jnp.allclose( + trace["he_ascertainment_baseline_hospital"]["value"] + / trace["he_ascertainment_baseline_ed_visits"]["value"], + 0.5, + ) assert jnp.array_equal( values["ed_visits"], trace["he_ascertainment_ed_visits"]["value"], @@ -785,7 +1124,7 @@ def test_sample_calls_each_random_variable_once(self): with numpyro.handlers.seed(rng_seed=42): with numpyro.handlers.trace() as trace: - values = ascertainment.sample() + values = ascertainment.sample(n_timepoints=4) assert trace["iedr"]["type"] == "sample" assert trace["ihr_rel_iedr"]["type"] == "sample" diff --git a/test/test_pyrenew_builder.py b/test/test_pyrenew_builder.py index 39b08edb1..6963c085a 100644 --- a/test/test_pyrenew_builder.py +++ b/test/test_pyrenew_builder.py @@ -13,9 +13,9 @@ from pyrenew.ascertainment import ( AscertainmentModel, + IndependentAscertainment, JointAscertainment, RatioLinkedAscertainment, - TimeVaryingAscertainment, ) from pyrenew.deterministic import DeterministicPMF, DeterministicVariable from pyrenew.latent import ( @@ -40,46 +40,91 @@ SUBPOP_FRACTIONS = jnp.array([0.3, 0.25, 0.45]) -class RecordingAscertainmentModel(AscertainmentModel): - """Return configured values while recording axis context for tests.""" +class RecordingTemporalProcess: + """Return zero deviations while recording temporal-process arguments.""" + + step_size = 1 def __init__( self, - name: str, - values: Mapping[str, ArrayLike], requires_calendar_anchor: bool = False, ) -> None: - """Initialize a recording ascertainment model.""" - super().__init__(name=name, signals=tuple(values)) - self.values = dict(values) - self._requires_calendar_anchor = requires_calendar_anchor - self.sample_calls: list[tuple[int, int | None]] = [] + """Initialize a recording temporal process.""" + self.requires_calendar_anchor = requires_calendar_anchor + self.sample_calls: list[dict[str, object]] = [] - def requires_calendar_anchor(self) -> bool: - """Return the configured calendar requirement. + def sample( + self, + n_timepoints: int, + initial_value: float | ArrayLike | None = None, + n_processes: int = 1, + name_prefix: str = "temporal", + *, + first_day_dow: int | None = None, + ) -> ArrayLike: + """Record sampling arguments and return zero deviations. + + Parameters + ---------- + n_timepoints + Number of timepoints to return. + initial_value + Initial value supplied by the ascertainment model. + n_processes + Number of process columns to return. + name_prefix + NumPyro name prefix supplied by the ascertainment model. + first_day_dow + Day of week for the first model-axis timepoint. Returns ------- - bool - Configured calendar requirement. + ArrayLike + Zero deviations with shape ``(n_timepoints, n_processes)``. """ - return self._requires_calendar_anchor + self.sample_calls.append( + { + "n_timepoints": n_timepoints, + "initial_value": initial_value, + "n_processes": n_processes, + "name_prefix": name_prefix, + "first_day_dow": first_day_dow, + } + ) + return jnp.zeros((n_timepoints, n_processes)) + + +class MalformedAscertainmentModel(AscertainmentModel): + """Override the common sampler with malformed output for boundary tests.""" + + def __init__(self, name: str) -> None: + """Initialize a deliberately nonconforming ascertainment model.""" + super().__init__(name=name, signals=("hospital",)) + + def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + """Return a valid baseline that the overridden sampler does not use.""" + return {"hospital": 0.2} def sample( self, n_timepoints: int, first_day_dow: int | None = None, - **kwargs: object, ) -> Mapping[str, ArrayLike]: - """Return configured values and record shared-axis arguments. + """Return an invalid two-dimensional final ascertainment value. + + Parameters + ---------- + n_timepoints + Shared-axis length, ignored. + first_day_dow + Shared-axis day of week, ignored. Returns ------- Mapping[str, ArrayLike] - Configured signal values. + Malformed signal output. """ - self.sample_calls.append((n_timepoints, first_day_dow)) - return self.values + return {"hospital": jnp.ones((2, 1))} def _obs_date_for_dow(target_first_day_dow: int, n_init: int) -> date: @@ -107,10 +152,10 @@ def _obs_date_for_dow(target_first_day_dow: int, n_init: int) -> date: return date(2024, 1, 1) + timedelta(days=obs_dow) -def _builder_with_time_varying_ascertainment( - ascertainment: TimeVaryingAscertainment, +def _builder_with_ascertainment( + ascertainment: AscertainmentModel, ) -> PyrenewBuilder: - """Build a two-count-signal model using one outer ascertainment model. + """Build a two-count-signal model using one ascertainment model. Returns ------- @@ -455,6 +500,85 @@ def test_prior_predictive_multi_signal(self, simple_builder): # All prior predictive infections should be positive assert jnp.all(prior_samples["latent_infections"] > 0) + @pytest.mark.parametrize("varying", [False, True]) + def test_prior_predictive_with_independent_ascertainment( + self, + varying: bool, + ) -> None: + """Test fixed and fully varying independent ascertainment models.""" + import jax.random + from numpyro.infer import Predictive + + temporal_processes = None + if varying: + temporal_processes = { + "hospital": fixed_random_walk(innovation_sd=0.05), + "ed": fixed_random_walk(innovation_sd=0.05), + } + ascertainment = IndependentAscertainment( + name="he_ascertainment", + rate_rvs={ + "hospital": DistributionalVariable("ihr", dist.Beta(2, 198)), + "ed": DistributionalVariable("iedr", dist.Beta(2, 98)), + }, + temporal_processes=temporal_processes, + ) + model = _builder_with_ascertainment(ascertainment).build() + n_days_post_init = 10 + n_total = model.latent.n_initialization_points + n_days_post_init + + samples = Predictive(model.sample, num_samples=2)( + jax.random.PRNGKey(42), + n_days_post_init=n_days_post_init, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ed={"obs": None}, + ) + + assert samples["ihr"].shape == (2,) + assert samples["iedr"].shape == (2,) + assert samples["he_ascertainment_baseline_hospital"].shape == (2,) + assert samples["he_ascertainment_baseline_ed"].shape == (2,) + expected_final_shape = (2, n_total) if varying else (2,) + assert samples["he_ascertainment_hospital"].shape == expected_final_shape + assert samples["he_ascertainment_ed"].shape == expected_final_shape + + def test_prior_predictive_with_mixed_independent_ascertainment(self) -> None: + """Test one registered model can mix fixed and varying signals.""" + import jax.random + from numpyro.infer import Predictive + + ascertainment = IndependentAscertainment( + name="he_ascertainment", + rate_rvs={ + "hospital": DeterministicVariable("ihr", 0.01), + "ed": DeterministicVariable("iedr", 0.02), + }, + temporal_processes={ + "ed": fixed_random_walk(innovation_sd=0.05), + }, + ) + model = _builder_with_ascertainment(ascertainment).build() + n_days_post_init = 10 + n_total = model.latent.n_initialization_points + n_days_post_init + + samples = Predictive(model.sample, num_samples=2)( + jax.random.PRNGKey(42), + n_days_post_init=n_days_post_init, + population_size=1_000_000, + subpop_fractions=SUBPOP_FRACTIONS, + hospital={"obs": None}, + ed={"obs": None}, + ) + + assert samples["he_ascertainment_baseline_hospital"].shape == (2,) + assert samples["he_ascertainment_baseline_ed"].shape == (2,) + assert samples["he_ascertainment_hospital"].shape == (2,) + assert samples["he_ascertainment_ed"].shape == (2, n_total) + assert samples["hospital_predicted"].shape == (2, n_total) + assert samples["ed_predicted"].shape == (2, n_total) + def test_prior_predictive_with_joint_ascertainment(self): """Test model-scoped sampling for shared joint ascertainment.""" import jax.random @@ -511,6 +635,8 @@ def test_prior_predictive_with_joint_ascertainment(self): ) assert prior_samples["he_ascertainment_eta"].shape == (3, 2) + assert prior_samples["he_ascertainment_baseline_hospital"].shape == (3,) + assert prior_samples["he_ascertainment_baseline_ed"].shape == (3,) assert prior_samples["he_ascertainment_hospital"].shape == (3,) assert prior_samples["he_ascertainment_ed"].shape == (3,) assert "hospital_predicted" in prior_samples @@ -576,39 +702,37 @@ def test_prior_predictive_with_ratio_linked_ascertainment(self): assert model.ascertainment_models["he_ascertainment"] is ascertainment assert prior_samples["iedr"].shape == (3,) assert prior_samples["ihr_rel_iedr"].shape == (3,) + assert prior_samples["he_ascertainment_baseline_ed"].shape == (3,) + assert prior_samples["he_ascertainment_baseline_hospital"].shape == (3,) assert prior_samples["he_ascertainment_ed"].shape == (3,) assert prior_samples["he_ascertainment_hospital"].shape == (3,) assert jnp.allclose( - prior_samples["he_ascertainment_ed"], + prior_samples["he_ascertainment_baseline_ed"], prior_samples["iedr"], ) assert jnp.allclose( - prior_samples["he_ascertainment_hospital"], + prior_samples["he_ascertainment_baseline_hospital"], prior_samples["iedr"] * prior_samples["ihr_rel_iedr"], ) assert "hospital_predicted" in prior_samples assert "ed_predicted" in prior_samples - def test_prior_predictive_with_time_varying_joint_baseline(self) -> None: - """Test joint baseline sites and full-axis outer count trajectories.""" + def test_prior_predictive_with_time_varying_joint_ascertainment(self) -> None: + """Test joint baseline sites and full-axis final trajectories.""" import jax.random from numpyro.infer import Predictive - baseline = JointAscertainment( - name="he_baseline", + ascertainment = JointAscertainment( + name="he_ascertainment", signals=("hospital", "ed"), baseline_rates=jnp.array([0.01, 0.02]), scale_tril=jnp.eye(2) * 0.1, - ) - ascertainment = TimeVaryingAscertainment( - "he_ascertainment", - baseline, - { + temporal_processes={ "hospital": fixed_random_walk(innovation_sd=0.05), "ed": fixed_random_walk(innovation_sd=0.05), }, ) - model = _builder_with_time_varying_ascertainment(ascertainment).build() + model = _builder_with_ascertainment(ascertainment).build() n_days_post_init = 10 n_total = model.latent.n_initialization_points + n_days_post_init @@ -622,22 +746,22 @@ def test_prior_predictive_with_time_varying_joint_baseline(self) -> None: ) assert set(model.ascertainment_models) == {"he_ascertainment"} - assert samples["he_baseline_eta"].shape == (2, 2) - assert samples["he_baseline_hospital"].shape == (2,) - assert samples["he_baseline_ed"].shape == (2,) + assert samples["he_ascertainment_eta"].shape == (2, 2) + assert samples["he_ascertainment_baseline_hospital"].shape == (2,) + assert samples["he_ascertainment_baseline_ed"].shape == (2,) assert samples["he_ascertainment_hospital"].shape == (2, n_total) assert samples["he_ascertainment_ed"].shape == (2, n_total) assert samples["hospital_predicted"].shape == (2, n_total) assert samples["ed_predicted"].shape == (2, n_total) assert samples["latent_infections"].shape == (2, n_total) - def test_prior_predictive_with_time_varying_ratio_baseline(self) -> None: - """Test ratio sites persist beneath independently varying trajectories.""" + def test_prior_predictive_with_time_varying_ratio_ascertainment(self) -> None: + """Test ratio baselines persist beneath independently varying rates.""" import jax.random from numpyro.infer import Predictive - baseline = RatioLinkedAscertainment( - name="he_baseline", + ascertainment = RatioLinkedAscertainment( + name="he_ascertainment", base_signal="ed", linked_signal="hospital", base_rate_rv=DistributionalVariable("iedr", dist.Beta(2, 98)), @@ -645,16 +769,12 @@ def test_prior_predictive_with_time_varying_ratio_baseline(self) -> None: "ihr_rel_iedr", dist.Beta(2, 2), ), - ) - ascertainment = TimeVaryingAscertainment( - "he_ascertainment", - baseline, - { + temporal_processes={ "hospital": fixed_random_walk(innovation_sd=0.05), "ed": fixed_random_walk(innovation_sd=0.05), }, ) - model = _builder_with_time_varying_ascertainment(ascertainment).build() + model = _builder_with_ascertainment(ascertainment).build() n_days_post_init = 10 n_total = model.latent.n_initialization_points + n_days_post_init @@ -669,9 +789,12 @@ def test_prior_predictive_with_time_varying_ratio_baseline(self) -> None: assert samples["iedr"].shape == (2,) assert samples["ihr_rel_iedr"].shape == (2,) - assert jnp.allclose(samples["he_baseline_ed"], samples["iedr"]) assert jnp.allclose( - samples["he_baseline_hospital"], + samples["he_ascertainment_baseline_ed"], + samples["iedr"], + ) + assert jnp.allclose( + samples["he_ascertainment_baseline_hospital"], samples["iedr"] * samples["ihr_rel_iedr"], ) assert samples["he_ascertainment_hospital"].shape == (2, n_total) @@ -737,6 +860,7 @@ def test_prior_predictive_reuses_same_ascertainment_signal(self): ) assert prior_samples["shared_ascertainment_eta"].shape == (2, 2) + assert prior_samples["shared_ascertainment_baseline_hospital"].shape == (2,) assert prior_samples["shared_ascertainment_hospital"].shape == (2,) assert prior_samples["hospital_a_predicted"].shape == ( 2, @@ -775,127 +899,17 @@ def test_manual_model_rejects_mismatched_ascertainment_key(self, simple_builder) ascertainment_models={"wrong_name": ascertainment}, ) - def test_rejects_registering_time_varying_baseline_object( - self, - simple_builder: PyrenewBuilder, - ) -> None: - """Test a wrapped baseline object cannot also be registered.""" - built = simple_builder.build() - baseline = RecordingAscertainmentModel("baseline", {"hospital": 0.2}) - outer = TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": fixed_random_walk(innovation_sd=0.1)}, - ) - - with pytest.raises(ValueError, match=r"outer.*baseline.*also be registered"): - MultiSignalModel( - latent_process=built.latent, - observations=built.observations, - ascertainment_models={"outer": outer, "baseline": baseline}, - ) - - def test_rejects_registered_model_with_wrapped_baseline_name( - self, - simple_builder: PyrenewBuilder, - ) -> None: - """Test distinct registered and baseline objects cannot share a name.""" - built = simple_builder.build() - baseline = RecordingAscertainmentModel("baseline", {"hospital": 0.2}) - outer = TimeVaryingAscertainment( - "outer", - baseline, - {"hospital": fixed_random_walk(innovation_sd=0.1)}, - ) - registered = RecordingAscertainmentModel("baseline", {"ed": 0.3}) - - with pytest.raises( - ValueError, - match=r"namespace 'baseline'.*baseline model.*outer.*registered", - ): - MultiSignalModel( - latent_process=built.latent, - observations=built.observations, - ascertainment_models={"outer": outer, "baseline": registered}, - ) - - def test_rejects_two_wrapped_baselines_with_the_same_name( - self, - simple_builder: PyrenewBuilder, - ) -> None: - """Test independently wrapped baselines need distinct namespaces.""" - built = simple_builder.build() - first_baseline = RecordingAscertainmentModel( - "shared_baseline", - {"hospital": 0.2}, - ) - second_baseline = RecordingAscertainmentModel( - "shared_baseline", - {"ed": 0.3}, - ) - first = TimeVaryingAscertainment( - "first_outer", - first_baseline, - {"hospital": fixed_random_walk(innovation_sd=0.1)}, - ) - second = TimeVaryingAscertainment( - "second_outer", - second_baseline, - {"ed": fixed_random_walk(innovation_sd=0.1)}, - ) - - with pytest.raises( - ValueError, - match=r"shared_baseline.*first_outer.*second_outer", - ): - MultiSignalModel( - latent_process=built.latent, - observations=built.observations, - ascertainment_models={ - "first_outer": first, - "second_outer": second, - }, - ) - - def test_rejects_outer_name_colliding_with_another_baseline( - self, - simple_builder: PyrenewBuilder, - ) -> None: - """Test outer and owned baseline namespaces are globally unique.""" - built = simple_builder.build() - first = TimeVaryingAscertainment( - "second_baseline", - RecordingAscertainmentModel("first_baseline", {"hospital": 0.2}), - {"hospital": fixed_random_walk(innovation_sd=0.1)}, - ) - second = TimeVaryingAscertainment( - "second_outer", - RecordingAscertainmentModel("second_baseline", {"ed": 0.3}), - {"ed": fixed_random_walk(innovation_sd=0.1)}, - ) - - with pytest.raises( - ValueError, - match=r"second_baseline.*registered.*second_outer", - ): - MultiSignalModel( - latent_process=built.latent, - observations=built.observations, - ascertainment_models={ - "second_baseline": first, - "second_outer": second, - }, - ) - def test_forwards_full_axis_length_and_padded_axis_weekday( self, simple_builder: PyrenewBuilder, ) -> None: - """Test ascertainment sampling receives the shared padded-axis context.""" + """Test temporal sampling receives full padded-axis context.""" built = simple_builder.build() - ascertainment = RecordingAscertainmentModel( - "recording", - {"hospital": 0.2}, + process = RecordingTemporalProcess() + ascertainment = IndependentAscertainment( + name="recording", + rate_rvs={"hospital": DeterministicVariable("ihr_recording", 0.2)}, + temporal_processes={"hospital": process}, ) model = MultiSignalModel( latent_process=built.latent, @@ -917,8 +931,14 @@ def test_forwards_full_axis_length_and_padded_axis_weekday( hospital={"obs": None}, ) - assert ascertainment.sample_calls == [ - (model.latent.n_initialization_points + 5, first_day_dow) + assert process.sample_calls == [ + { + "n_timepoints": model.latent.n_initialization_points + 5, + "initial_value": 0.0, + "n_processes": 1, + "name_prefix": "recording_hospital", + "first_day_dow": first_day_dow, + } ] def test_calendar_ascertainment_fails_early_and_names_model( @@ -927,10 +947,15 @@ def test_calendar_ascertainment_fails_early_and_names_model( ) -> None: """Test registered calendar-aligned ascertainment requires a date.""" built = simple_builder.build() - ascertainment = RecordingAscertainmentModel( - "weekly_ascertainment", - {"hospital": 0.2}, - requires_calendar_anchor=True, + inner_process = RecordingTemporalProcess() + weekly_process = WeeklyTemporalProcess( + inner_process, + start_dow=0, + ) + ascertainment = IndependentAscertainment( + name="weekly_ascertainment", + rate_rvs={"hospital": DeterministicVariable("weekly_ihr", 0.2)}, + temporal_processes={"hospital": weekly_process}, ) model = MultiSignalModel( latent_process=built.latent, @@ -949,36 +974,43 @@ def test_calendar_ascertainment_fails_early_and_names_model( hospital={"obs": None}, ) - def test_unregistered_calendar_ascertainment_does_not_require_date( + assert inner_process.sample_calls == [] + + def test_daily_ascertainment_process_works_without_obs_start_date( self, simple_builder: PyrenewBuilder, ) -> None: - """Test only registered ascertainment contributes date requirements.""" + """Test a registered daily process does not require a calendar date.""" built = simple_builder.build() - RecordingAscertainmentModel( - "unregistered_weekly", - {"hospital": 0.2}, - requires_calendar_anchor=True, + process = RecordingTemporalProcess() + ascertainment = IndependentAscertainment( + name="daily_ascertainment", + rate_rvs={"hospital": DeterministicVariable("daily_ihr", 0.2)}, + temporal_processes={"hospital": process}, + ) + model = MultiSignalModel( + latent_process=built.latent, + observations=built.observations, + ascertainment_models={"daily_ascertainment": ascertainment}, ) with numpyro.handlers.seed(rng_seed=42): - built.sample( + model.sample( n_days_post_init=5, population_size=1_000_000, subpop_fractions=SUBPOP_FRACTIONS, hospital={"obs": None}, ) + assert process.sample_calls[0]["first_day_dow"] is None + def test_malformed_ascertainment_fails_before_observation_sampling( self, simple_builder: PyrenewBuilder, ) -> None: """Test output validation completes before any observation runs.""" built = simple_builder.build() - malformed = RecordingAscertainmentModel( - "malformed", - {"hospital": jnp.ones((2, 1))}, - ) + malformed = MalformedAscertainmentModel("malformed") model = MultiSignalModel( latent_process=built.latent, observations=built.observations, From 852b1ed104840b41c90351d42f4072fbe47ae7f6 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Thu, 1 Oct 2026 11:12:19 -0400 Subject: [PATCH 03/12] lint fix --- docs/tutorials/ascertainment.qmd | 159 ++++--------------------------- pyrenew/ascertainment/base.py | 42 ++++---- test/test_ascertainment.py | 6 +- test/test_pyrenew_builder.py | 4 +- 4 files changed, 48 insertions(+), 163 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 2023f144c..73f8c1c98 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -25,9 +25,11 @@ import jax.numpy as jnp import numpyro.distributions as dist from pyrenew.ascertainment import ( + AscertainmentModel, + AscertainmentSignal, + IndependentAscertainment, JointAscertainment, RatioLinkedAscertainment, - TimeVaryingAscertainment, ) from pyrenew.deterministic import DeterministicPMF, DeterministicVariable from pyrenew.latent import RandomWalk, WeeklyTemporalProcess @@ -42,7 +44,6 @@ from pyrenew.time import MMWR_WEEK An ascertainment rate is the rate at which an infection results in an observed event. For hospital admissions, this is often called the infection-hospitalization rate (IHR). For emergency department visits, it can be called the infection-to-ED-visit rate (IEDR). - For a count observation process, the ascertainment rate scales delayed latent infections to obtain the expected number of observed events: $$ @@ -53,28 +54,30 @@ $$ where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate, and $\pi_{s,d}$ is its delay distribution. -## Individual and related ascertainment rates for multi-signal models +From a modeling perspective, an ascertainment rate may be considered constant over the observation period or allowed to vary over time. +A time-varying ascertainment component combines a scalar baseline rate with a temporal process that describes deviations from that baseline over time as follows: -## Fixed and time-varying ascertainment +$\mathrm{b}_{\mathrm{s}}$ is the baseline rate for signal $\mathrm{s}$, $\mathrm{z}_{\mathrm{s}}(\mathrm{t})$ is its temporal deviation, and -******** stopped here ****** +$$ + \alpha\_s(t) = \operatorname{logit}^{-1} \left(\operatorname{logit}(b\_s) + z\_s(t)\right) +$$ -If a model treats the ascertainment rates for different signals as unrelated, each observation process can receive its own `RandomVariable` directly. -An ascertainment component is needed when the model represents a relationship between the ascertainment rates of multiple observation processes. +Thus, $\mathrm{b}_{\mathrm{s}}$ is the ascertainment rate when the temporal deviation is zero while $\mathrm{z}_{\mathrm{s}}(\mathrm{t})$ determines how the rate changes during the observation period. -PyRenew provides two scalar models and one time-varying wrapper: +The `AscertainmentModel` component of a PyRenew renewal process model provides the ascertainment rate for one or more observation signals, which can either be constant or time-varying. -- `JointAscertainment` specifies a joint prior for two or more rates. -- `RatioLinkedAscertainment` specifies one rate as a multiple of another rate. +- For a renewal process model for a single signal or for a multisignal model where the signals measure unrelated phenomena, the `IndependentAscertainment` component models the signal baseline and optional temporal process. +- For a multisignal model where the signals are measuring a related phenomenon, such as hospitalizations and ED visits, which provide 2 views on health-care seeking behaviors, these signals can be modeled either by a `JointAscertainment` component which assigns the baseline rates a joint prior or as a `RatioLinkedAscertainment` which expresses one baseline rate relative to another. -The choice among them is part of the statistical model, not only an API choice. -Each component expresses a different assumption about how information is shared across signals and over time. +These components reflect fundamental modeling assumptions about how information is shared across signals and over time. ## Independent ascertainment rates -An ascertainment component is not required when the signal rates are modeled independently. -Pass each rate directly to its observation process: +The `IndependentAscertainment` component models a signal baseline in the case where the ascertainment rate is constant, or a signal baseline plus temporal process is the ascertainment rate varies over time. + +If the model has one or more signals which measure unrelated phenomena and whose ascertainment rate is constant, then an `AscertainmentModel` component is equivalent to a `RandomVariable` which can be specified directly on the observation process component's `ascertainment_rate_rv`. ```python hospital_obs = PopulationCounts( @@ -84,16 +87,8 @@ hospital_obs = PopulationCounts( noise=hosp_noise, ) -ed_obs = PopulationCounts( - name="ed_visits", - ascertainment_rate_rv=DistributionalVariable("iedr", dist.Beta(2, 98)), - delay_distribution_rv=ed_delay_rv, - noise=ed_noise, -) ``` -This specification does not encode prior dependence between IHR and IEDR; each observation process has its own ascertainment-rate prior. - ## Ascertainment components and the builder An ascertainment component is a model-level component shared by multiple observation processes. @@ -218,126 +213,6 @@ Unlike `JointAscertainment`, this component does not automatically constrain the If both rates should be $<1$ (for instance because both are probabilities), make sure to choose priors that make $\alpha_{\mathrm{base}}\rho > 1$ negligibly likely. The component does not clip an invalid linked rate, because clipping would change the specified statistical model. -## Related baselines with time-varying rates - -`TimeVaryingAscertainment` combines a scalar baseline model with one temporal process per signal. -If $b_s$ is the sampled baseline for signal $s$ and $z_s(t)$ is its temporal deviation, the rate trajectory is - -$$ - \alpha_s(t) - = - \operatorname{logit}^{-1} - \left(\operatorname{logit}(b_s) + z_s(t)\right). -$$ - -The baseline is the rate obtained when the deviation is zero. -It is not necessarily the first trajectory value: some temporal processes, including `AR1`, sample their initial state around zero. - -The baseline model relates only the scalar reference rates. -After those rates are sampled, each signal receives its own temporal process. -A joint covariance or a baseline ratio is therefore not reimposed at every timepoint. -This lets hospital and ED ascertainment trajectories move independently around related reference levels. - -The time-varying rate applies on the potential-observation date, after the infection-to-observation delay: - -```python -baseline_rates = baseline_model.sample() - -potential_counts = convolve_delay( - infections, - clinical_delay_pmf, -) - -rate_trajectory = sigmoid( - logit(baseline_rates[signal]) - + temporal_deviation[signal] -) - -predicted_counts = potential_counts * rate_trajectory -``` - -This is different from multiplying infections by a changing rate before convolution. -It represents ascertainment conditions on the day an event could be observed. - -### Joint scalar baselines - -```python -baseline = JointAscertainment( - name="he_baseline", - signals=("hospital", "ed_visits"), - baseline_rates=jnp.array([0.01, 0.02]), - scale_tril=jnp.array([[0.3, 0.0], [0.15, 0.26]]), -) - -ascertainment = TimeVaryingAscertainment( - name="he_ascertainment", - baseline_model=baseline, - processes={ - "hospital": RandomWalk( - DeterministicVariable("hospital_ascertainment_sd", 0.05) - ), - "ed_visits": RandomWalk( - DeterministicVariable("ed_ascertainment_sd", 0.05) - ), - }, -) -``` - -### Ratio-linked scalar baselines - -```python -baseline = RatioLinkedAscertainment( - name="he_baseline", - base_signal="ed_visits", - linked_signal="hospital", - base_rate_rv=DistributionalVariable("iedr", dist.Beta(1, 100)), - ratio_rv=DistributionalVariable( - "ihr_rel_iedr", - dist.LogNormal(jnp.log(0.5), 0.35), - ), -) - -ascertainment = TimeVaryingAscertainment( - name="he_ascertainment", - baseline_model=baseline, - processes={ - "ed_visits": RandomWalk( - DeterministicVariable("ed_ascertainment_sd", 0.05) - ), - "hospital": RandomWalk( - DeterministicVariable("hospital_ascertainment_sd", 0.05) - ), - }, -) -``` - -Register only the outer model, because it owns and samples its baseline. -Use the outer accessors in observation processes: - -```python -builder.add_ascertainment(ascertainment) - -hospital_obs = PopulationCounts( - name="hospital", - ascertainment_rate_rv=ascertainment.for_signal("hospital"), - ... -) -``` - -Daily temporal processes such as `RandomWalk` do not require a date. -Calendar-aligned processes do. -For example, wrap a process in `WeeklyTemporalProcess` when deviations should be shared within calendar weeks: - -```python -weekly_hospital_process = WeeklyTemporalProcess( - RandomWalk(DeterministicVariable("hospital_weekly_sd", 0.05)), - start_dow=MMWR_WEEK, -) -``` - -When any registered ascertainment model contains a weekly process, supply `obs_start_date` when sampling or fitting the model. -PyRenew uses it to align weekly blocks to the padded shared model axis. - ## Building a model with an ascertainment component The complete builder pattern is the same for either ascertainment component - here we use a `RatioLinkedAscertainment` but the sequence of definitions is the same for `JointAscertainment`. diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index a4ff9cd50..6c004b89c 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -84,17 +84,32 @@ def sample(self, **kwargs: object) -> ArrayLike: class AscertainmentModel(metaclass=ABCMeta): """ - Base class for shared ascertainment structure. + An ``AscertainmentModel`` is a component of a PyRenew renewal process + model. It provides an ascertainment rate for one or more observation signals. - An ascertainment rate is the probability that latent incidence is observed - in a particular data stream. Examples include an infection-hospitalization - ratio for hospital admissions or an infection-ED-visit ratio for emergency - department visits. + Each signal has a scalar baseline rate. If a temporal process is specified + for that signal, the component combines the baseline rate with the temporal + process to produce a rate trajectory over the model period. Otherwise, the + baseline rate applies throughout the model period. - ``AscertainmentModel`` objects make shared structure explicit in a model - specification. A user defines the shared model once, registers it with - ``PyrenewBuilder.add_ascertainment(...)``, and passes signal-specific - accessors into observation processes: + The model samples one scalar baseline rate per signal and can add an + optional signal-specific temporal deviation on the logit scale. Subclasses + implement ``_sample_baseline_rates()`` to define relationships among the + scalar baselines. The base class owns validation, temporal sampling, and + the standard deterministic sites. + + Concrete subclasses determine how the baseline rates are specified. + ``IndependentAscertainment`` specifies each baseline separately, + ``JointAscertainment`` assigns the baselines a joint distribution, and + ``RatioLinkedAscertainment`` defines one baseline relative to another. + Both ``IndependentAscertainment`` and ``JointAscertainment`` may be used + with a single signal. + + Register the component with ``PyrenewBuilder.add_ascertainment(...)``. + Use ``for_signal(...)`` to obtain the signal-specific accessor passed + to an observation process as ``ascertainment_rate_rv``. Accessors + returned by ``for_signal()`` read the final sampled values from + the active model context and do not sample independently. ```python ascertainment = JointAscertainment(...) @@ -106,15 +121,6 @@ class AscertainmentModel(metaclass=ABCMeta): ... ) ``` - - The model samples one scalar baseline rate per signal and can add an - optional signal-specific temporal deviation on the logit scale. Subclasses - implement ``_sample_baseline_rates()`` to define relationships among the - scalar baselines. The base class owns validation, temporal sampling, and - the standard deterministic sites. - - Accessors returned by ``for_signal()`` read the final sampled values from - the active model context and do not sample independently. """ def __init__( diff --git a/test/test_ascertainment.py b/test/test_ascertainment.py index f45a11d7d..916085f92 100644 --- a/test/test_ascertainment.py +++ b/test/test_ascertainment.py @@ -340,7 +340,9 @@ def test_traced_baseline_values_do_not_trigger_python_validation(self) -> None: model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) @jax.jit - def validate(value: ArrayLike) -> ArrayLike: + def validate( # numpydoc ignore=RT01 + value: ArrayLike, + ) -> ArrayLike: """Validate and return a traced baseline value.""" return model._validate_baseline_rates({"hospital": value})["hospital"] @@ -469,7 +471,7 @@ def test_rejects_invalid_sample_shapes(self, value: ArrayLike) -> None: model.validate_sampled_values({"hospital": value}, n_timepoints=4) -def _make_public_ascertainment( +def _make_public_ascertainment( # numpydoc ignore=RT01 model_kind: str, temporal_processes: Mapping[str, TemporalProcess] | None, ) -> AscertainmentModel: diff --git a/test/test_pyrenew_builder.py b/test/test_pyrenew_builder.py index 6963c085a..6d1835c22 100644 --- a/test/test_pyrenew_builder.py +++ b/test/test_pyrenew_builder.py @@ -101,7 +101,9 @@ def __init__(self, name: str) -> None: """Initialize a deliberately nonconforming ascertainment model.""" super().__init__(name=name, signals=("hospital",)) - def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: + def _sample_baseline_rates( # numpydoc ignore=RT01 + self, + ) -> Mapping[str, ArrayLike]: """Return a valid baseline that the overridden sampler does not use.""" return {"hospital": 0.2} From b84fc89a73a93718f2aaba16ba16540502752dde Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Fri, 2 Oct 2026 10:48:32 -0400 Subject: [PATCH 04/12] tutorial edits --- docs/tutorials/ascertainment.qmd | 20 +-- .../tutorials/building_multisignal_models.qmd | 1 + pyrenew/ascertainment/base.py | 2 - test/integration/conftest.py | 131 +++++++++++++++++- 4 files changed, 141 insertions(+), 13 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 73f8c1c98..89e06073b 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -60,7 +60,7 @@ A time-varying ascertainment component combines a scalar baseline rate with a te $\mathrm{b}_{\mathrm{s}}$ is the baseline rate for signal $\mathrm{s}$, $\mathrm{z}_{\mathrm{s}}(\mathrm{t})$ is its temporal deviation, and $$ - \alpha\_s(t) = \operatorname{logit}^{-1} \left(\operatorname{logit}(b\_s) + z\_s(t)\right) + \alpha_s(t) = \operatorname{logit}^{-1} \left(\operatorname{logit}(\mathrm{b}_{\mathrm{s}}) + \mathrm{z}_{\mathrm{s}}(t)\right) $$ Thus, $\mathrm{b}_{\mathrm{s}}$ is the ascertainment rate when the temporal deviation is zero while $\mathrm{z}_{\mathrm{s}}(\mathrm{t})$ determines how the rate changes during the observation period. @@ -75,9 +75,11 @@ These components reflect fundamental modeling assumptions about how information ## Independent ascertainment rates -The `IndependentAscertainment` component models a signal baseline in the case where the ascertainment rate is constant, or a signal baseline plus temporal process is the ascertainment rate varies over time. +The `IndependentAscertainment` component models a signal baseline in the case where the ascertainment rate is constant, or a signal baseline plus temporal process if the ascertainment rate varies over time. -If the model has one or more signals which measure unrelated phenomena and whose ascertainment rate is constant, then an `AscertainmentModel` component is equivalent to a `RandomVariable` which can be specified directly on the observation process component's `ascertainment_rate_rv`. +An `AscertainmentModel` component is a subclass of a [Random Variable](random_variables.md). +If modeling a single signal with a constant ascertainment rate or some number of independent signals with constant ascertainment rates, +instead of using an `IndependentAscertainment` component, it is possible to use a `RandomVariable` directly. ```python hospital_obs = PopulationCounts( @@ -264,17 +266,17 @@ model = builder.build() ## Shared assumptions and limitations -`JointAscertainment` and `RatioLinkedAscertainment` each produce one scalar rate per signal over the entire model period. -`TimeVaryingAscertainment` adds signal-specific deviations to scalar rates from either model. -Its baseline relationship is not a claim that the same covariance or ratio holds pointwise. - +- 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. + The choice of ascertainment model should consider: -- whether the rates should be related at all -- whether a covariance structure or a relative multiplier is easier to justify - whether the priors keep all rates, including derived linked rates, within scientifically plausible ranges - whether treating ascertainment rates as constant over the model period is scientifically reasonable - whether daily, stepwise, or calendar-week deviations match the scientific question +- whether the rates should be related at all +- whether a covariance structure or a relative multiplier is easier to justify Ascertainment rates are generally difficult to distinguish from the scale of latent infections using surveillance counts alone. Flexible infection and ascertainment trajectories can be weakly identified without informative priors, external information, or strong structural assumptions. diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 4eb2fbbdd..79cdce229 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -1086,3 +1086,4 @@ That same edge uncertainty is what carries forward when the model is used for fo - Explore different temporal processes for $\mathcal{R}(t)$ in the [Latent Infections](latent_infections.md) and [Latent Subpopulation Infections](latent_subpopulation_infections.md) tutorials - Learn about count-based observation models in [Observation Processes: Counts](observation_processes_counts.md) - Learn about continuous measurement models in [Observation Processes: Measurements](observation_processes_measurements.md) +- Learn about ascertainment models in [Ascertainment Models for Related Observation Processes](ascertainment.md) diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index 6c004b89c..65744ab83 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -102,8 +102,6 @@ class AscertainmentModel(metaclass=ABCMeta): ``IndependentAscertainment`` specifies each baseline separately, ``JointAscertainment`` assigns the baselines a joint distribution, and ``RatioLinkedAscertainment`` defines one baseline relative to another. - Both ``IndependentAscertainment`` and ``JointAscertainment`` may be used - with a single signal. Register the component with ``PyrenewBuilder.add_ascertainment(...)``. Use ``for_signal(...)`` to obtain the signal-specific accessor passed diff --git a/test/integration/conftest.py b/test/integration/conftest.py index b8f0d8b4c..85dc7a647 100644 --- a/test/integration/conftest.py +++ b/test/integration/conftest.py @@ -12,6 +12,7 @@ import polars as pl import pytest +import pyrenew.transformation as transformation from pyrenew.ascertainment import ( IndependentAscertainment, JointAscertainment, @@ -26,12 +27,17 @@ load_synthetic_weekly_hospital_admissions, ) from pyrenew.deterministic import DeterministicPMF, DeterministicVariable -from pyrenew.latent import InfectionsWithFeedback, WeeklyTemporalProcess +from pyrenew.latent import ( + AR1, + InfectionsWithFeedback, + RandomWalk, + WeeklyTemporalProcess, +) from pyrenew.latent.infection_process import InfectionProcess from pyrenew.latent.population_infections import PopulationInfections from pyrenew.model import MultiSignalModel, PyrenewBuilder from pyrenew.observation import NegativeBinomialNoise, PopulationCounts -from pyrenew.randomvariable import DistributionalVariable +from pyrenew.randomvariable import DistributionalVariable, TransformedVariable from pyrenew.time import MMWR_WEEK from test.test_helpers import fixed_ar1, fixed_ar1_state, fixed_differenced_ar1_state @@ -156,6 +162,127 @@ def ed_day_of_week_effects(true_params: dict) -> jnp.ndarray: return jnp.array(true_params["ed_visits"]["day_of_week_effects"]) +@pytest.fixture(scope="module") +def e_rtdaily_rw_timevary_model( + ed_delay_pmf: jnp.ndarray, +) -> MultiSignalModel: + """Build an ED model with daily Rt and time-varying ascertainment. + + This fixture reproduces the PyRenew structure of the + ``e_rtdaily_rw_timevary`` model without depending on the external ARM + model builder. The ascertainment baseline is sampled independently and a + calendar-aligned weekly AR(1) process supplies logit-scale deviations. + + Parameters + ---------- + ed_delay_pmf + Infection-to-ED-visit delay PMF from the synthetic dataset. + + Returns + ------- + MultiSignalModel + ED-only model ready for integration testing. + """ + rt_process = RandomWalk( + innovation_sd_rv=DistributionalVariable( + "eta_sd", + dist.TruncatedNormal( + 0.15 / jnp.sqrt(7.0), + 0.05 / jnp.sqrt(7.0), + low=0, + ), + ), + parameterization="innovation", + ) + infection_process = InfectionsWithFeedback( + name="infections_with_feedback", + infection_feedback_strength=TransformedVariable( + "inf_feedback", + DistributionalVariable( + "inf_feedback_raw", + dist.LogNormal(jnp.log(50.0), jnp.log(1.5)), + ), + transforms=transformation.AffineTransform(loc=0, scale=-1), + ), + infection_feedback_pmf=DeterministicPMF( + "infection_feedback_pmf", + _GEN_INT_PMF, + ), + ) + ascertainment = IndependentAscertainment( + name="ed_ascertainment", + rate_rvs={ + "ed": TransformedVariable( + "p_ed_visit", + DistributionalVariable( + "p_ed_visit_mean", + dist.Normal( + transformation.SigmoidTransform().inv(0.005), + 0.3, + ), + ), + transforms=transformation.SigmoidTransform(), + ) + }, + temporal_processes={ + "ed": WeeklyTemporalProcess( + AR1( + autoreg_rv=DistributionalVariable( + "autoreg_p_ed_visit", + dist.Beta(1, 100), + ), + innovation_sd_rv=DistributionalVariable( + "p_ed_visit_w_sd", + dist.TruncatedNormal(0.0, 0.01, low=0), + ), + parameterization="innovation", + ), + start_dow=MMWR_WEEK, + ) + }, + ) + + builder = PyrenewBuilder() + builder.configure_latent( + PopulationInfections, + gen_int_rv=DeterministicPMF("generation_interval_pmf", _GEN_INT_PMF), + I0_rv=DistributionalVariable("i0_first_obs_n_rv", dist.Beta(1, 10)), + log_rt_time_0_rv=DistributionalVariable( + "log_r_mu_intercept_rv", + dist.Normal(jnp.log(1.2), jnp.log(jnp.sqrt(2.0))), + ), + single_rt_process=rt_process, + infection_process=infection_process, + ) + builder.add_ascertainment(ascertainment) + builder.add_observation( + PopulationCounts( + name="ed", + ascertainment_rate_rv=ascertainment.for_signal("ed"), + delay_distribution_rv=DeterministicPMF("inf_to_ed", ed_delay_pmf), + noise=NegativeBinomialNoise( + DistributionalVariable( + "ed_visit_neg_bin_concentration", + dist.LogNormal(4.0, 1.0), + ) + ), + right_truncation_rv=DeterministicPMF( + "right_truncation_pmf", + jnp.array([1.0]), + ), + day_of_week_rv=TransformedVariable( + "ed_visit_wday_effect", + DistributionalVariable( + "ed_visit_wday_effect_raw", + dist.Dirichlet(jnp.full(7, 5.0)), + ), + transforms=transformation.AffineTransform(loc=0, scale=7), + ), + ) + ) + return builder.build() + + def _build_he_population_model( # numpydoc ignore=RT01 *, single_rt_process: object, From e41a06df85d935308b3fd471f1e0b6e689ed599e Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Fri, 2 Oct 2026 12:49:20 -0400 Subject: [PATCH 05/12] docs cleanup --- docs/tutorials/ascertainment.qmd | 6 ++---- docs/tutorials/building_multisignal_models.qmd | 2 +- docs/tutorials/latent_infections.qmd | 2 +- docs/tutorials/observation_processes_counts.qmd | 5 +---- panache.toml | 2 -- 5 files changed, 5 insertions(+), 12 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index 89e06073b..f37b009bf 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -78,8 +78,7 @@ These components reflect fundamental modeling assumptions about how information The `IndependentAscertainment` component models a signal baseline in the case where the ascertainment rate is constant, or a signal baseline plus temporal process if the ascertainment rate varies over time. An `AscertainmentModel` component is a subclass of a [Random Variable](random_variables.md). -If modeling a single signal with a constant ascertainment rate or some number of independent signals with constant ascertainment rates, -instead of using an `IndependentAscertainment` component, it is possible to use a `RandomVariable` directly. +If modeling a single signal with a constant ascertainment rate or some number of independent signals with constant ascertainment rates, instead of using an `IndependentAscertainment` component, it is possible to use a `RandomVariable` directly. ```python hospital_obs = PopulationCounts( @@ -88,7 +87,6 @@ hospital_obs = PopulationCounts( delay_distribution_rv=hosp_delay_rv, noise=hosp_noise, ) - ``` ## Ascertainment components and the builder @@ -269,7 +267,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. - + The choice of ascertainment model should consider: - whether the priors keep all rates, including derived linked rates, within scientifically plausible ranges diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd index 79cdce229..6e64e2cbf 100644 --- a/docs/tutorials/building_multisignal_models.qmd +++ b/docs/tutorials/building_multisignal_models.qmd @@ -786,7 +786,7 @@ obs_data_90 = { The observation dictionary has the following structure: -```python +```text { "obs_start_date": ..., "hospital": { diff --git a/docs/tutorials/latent_infections.qmd b/docs/tutorials/latent_infections.qmd index 8a774f607..3408011f7 100644 --- a/docs/tutorials/latent_infections.qmd +++ b/docs/tutorials/latent_infections.qmd @@ -259,7 +259,7 @@ n_init = len(gen_int_pmf) n_samples = 200 log_rt_time_0 = 0.5 I0_val = 0.001 -rt_cap = 3. +rt_cap = 3.0 def sample_process(rt_process, label): diff --git a/docs/tutorials/observation_processes_counts.qmd b/docs/tutorials/observation_processes_counts.qmd index 32e2faa9c..9d9a603fc 100644 --- a/docs/tutorials/observation_processes_counts.qmd +++ b/docs/tutorials/observation_processes_counts.qmd @@ -92,10 +92,7 @@ potential_counts = convolve_delay( clinical_delay_pmf, ) -rate_trajectory = sigmoid( - logit(baseline_rates[signal]) - + temporal_deviation[signal] -) +rate_trajectory = sigmoid(logit(baseline_rates[signal]) + temporal_deviation[signal]) predicted_counts = potential_counts * rate_trajectory ``` diff --git a/panache.toml b/panache.toml index 798a6847f..bca7d5378 100644 --- a/panache.toml +++ b/panache.toml @@ -3,5 +3,3 @@ flavor = "gfm" wrap = "sentence" [formatters] python = "ruff" -[linters] -python = "ruff" From d14ed735f10ced9b2ae723e79e6f4554a1dcd3c0 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 6 Oct 2026 15:52:58 -0400 Subject: [PATCH 06/12] more unit tests --- ...lation_infections_e_rtdaily_rw_timevary.py | 167 ++++++++++++++++++ 1 file changed, 167 insertions(+) create mode 100644 test/integration/test_population_infections_e_rtdaily_rw_timevary.py diff --git a/test/integration/test_population_infections_e_rtdaily_rw_timevary.py b/test/integration/test_population_infections_e_rtdaily_rw_timevary.py new file mode 100644 index 000000000..84f8566a1 --- /dev/null +++ b/test/integration/test_population_infections_e_rtdaily_rw_timevary.py @@ -0,0 +1,167 @@ +"""Integration test for daily Rt and time-varying ED ascertainment.""" + +from __future__ import annotations + +from datetime import date + +import jax +import jax.numpy as jnp +import jax.random as random +import numpyro +import polars as pl +import pytest + +from pyrenew.ascertainment import AscertainmentSignal, IndependentAscertainment +from pyrenew.latent import WeeklyTemporalProcess +from pyrenew.model import MultiSignalModel + +pytestmark = pytest.mark.integration + + +N_DAYS_FIT = 126 +N_DAYS_MCMC_SMOKE = 28 +NUM_WARMUP = 2 +NUM_SAMPLES = 2 +OBS_START_DATE = date(2023, 11, 5) + + +def _padded_ed_observations( + model: MultiSignalModel, + daily_ed: pl.DataFrame, + n_days: int, +) -> jax.Array: + """Return synthetic ED observations on the model's padded time axis. + + Parameters + ---------- + model + Model providing the required initialization padding. + daily_ed + Synthetic daily ED visit data. + n_days + Number of observation days to include. + + Returns + ------- + jax.Array + Padded ED observation vector. + """ + values = jnp.array(daily_ed["ed_visits"][:n_days].to_numpy(), dtype=jnp.float32) + return model.pad_observations(values) + + +class TestModelStructure: + """Verify the production-like ascertainment structure.""" + + def test_time_varying_independent_ascertainment_is_registered( + self, + e_rtdaily_rw_timevary_model: MultiSignalModel, + ) -> None: + """Verify the ED observation uses the registered ascertainment signal.""" + model = e_rtdaily_rw_timevary_model + + assert set(model.ascertainment_models) == {"ed_ascertainment"} + ascertainment = model.ascertainment_models["ed_ascertainment"] + assert isinstance(ascertainment, IndependentAscertainment) + assert ascertainment.signals == ("ed",) + assert isinstance(ascertainment.temporal_processes["ed"], WeeklyTemporalProcess) + assert ascertainment.requires_calendar_anchor() + + rate = model.observations["ed"].ascertainment_rate_rv + assert isinstance(rate, AscertainmentSignal) + assert rate.ascertainment_name == "ed_ascertainment" + assert rate.signal_name == "ed" + + +class TestSyntheticExecution: + """Run the model against the existing constant-ascertainment data.""" + + def test_full_graph_records_time_varying_ascertainment( + self, + e_rtdaily_rw_timevary_model: MultiSignalModel, + daily_ed: pl.DataFrame, + true_params: dict, + ) -> None: + """Verify one conditioned execution records valid ascertainment sites.""" + model = e_rtdaily_rw_timevary_model + ed_obs = _padded_ed_observations(model, daily_ed, N_DAYS_FIT) + population_size = float(true_params["population"]) + + model.validate_data( + n_days_post_init=N_DAYS_FIT, + obs_start_date=OBS_START_DATE, + ed={"obs": ed_obs}, + ) + with numpyro.handlers.seed(rng_seed=0): + with numpyro.handlers.trace() as trace: + model.sample( + n_days_post_init=N_DAYS_FIT, + population_size=population_size, + obs_start_date=OBS_START_DATE, + ed={"obs": ed_obs}, + ) + + n_total = model.latent.n_initialization_points + N_DAYS_FIT + baseline = trace["ed_ascertainment_baseline_ed"]["value"] + weekly = trace["ed_ascertainment_ed_weekly"]["value"] + trajectory = trace["ed_ascertainment_ed"]["value"] + predicted = trace["ed_predicted"]["value"] + + assert trace["p_ed_visit_mean"]["type"] == "sample" + assert baseline.shape == () + assert weekly.ndim == 2 + assert weekly.shape[1] == 1 + assert weekly.shape[0] < n_total + assert trajectory.shape == (n_total,) + assert jnp.isfinite(baseline) + assert 0 < baseline < 1 + assert jnp.all(jnp.isfinite(trajectory)) + assert jnp.all((trajectory > 0) & (trajectory < 1)) + assert predicted.shape == (n_total,) + assert jnp.all(jnp.isfinite(predicted[model.latent.n_initialization_points :])) + assert trace["ed_obs"]["is_observed"] + + def test_short_mcmc_fit_runs( + self, + e_rtdaily_rw_timevary_model: MultiSignalModel, + daily_ed: pl.DataFrame, + true_params: dict, + ) -> None: + """Verify NUTS can initialize and sample the integrated model.""" + model = e_rtdaily_rw_timevary_model + ed_obs = _padded_ed_observations(model, daily_ed, N_DAYS_MCMC_SMOKE) + + model.run( + num_warmup=NUM_WARMUP, + num_samples=NUM_SAMPLES, + rng_key=random.PRNGKey(42), + mcmc_args={"progress_bar": False}, + n_days_post_init=N_DAYS_MCMC_SMOKE, + population_size=float(true_params["population"]), + obs_start_date=OBS_START_DATE, + ed={"obs": ed_obs}, + ) + + assert model.mcmc is not None + samples = model.mcmc.get_samples() + jax.block_until_ready(samples) + + required_sites = { + "p_ed_visit_mean", + "autoreg_p_ed_visit", + "p_ed_visit_w_sd", + "ed_ascertainment_ed", + "ed_predicted", + } + assert required_sites <= set(samples) + for site in required_sites - {"ed_predicted"}: + assert jnp.all(jnp.isfinite(samples[site])) + assert jnp.all( + jnp.isfinite( + samples["ed_predicted"][:, model.latent.n_initialization_points :] + ) + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-m", "integration"]) From 65a61f0a45dc8bbc28faf5bf45de0b292b6081ec Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 7 Oct 2026 14:31:38 -0400 Subject: [PATCH 07/12] changes per code review --- docs/tutorials/ascertainment.qmd | 204 +++++++++++++----- pyrenew/ascertainment/base.py | 6 + test/integration/conftest.py | 20 +- ...lation_infections_e_rtdaily_rw_timevary.py | 4 +- 4 files changed, 166 insertions(+), 68 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index f37b009bf..e4dbb8b5b 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -17,28 +17,6 @@ jupyter: name: python3 --- -```{python} -#| label: setup -#| output: false - -import jax.numpy as jnp -import numpyro.distributions as dist - -from pyrenew.ascertainment import ( - AscertainmentModel, - AscertainmentSignal, - IndependentAscertainment, - JointAscertainment, - RatioLinkedAscertainment, -) -from pyrenew.deterministic import DeterministicPMF, DeterministicVariable -from pyrenew.latent import RandomWalk, WeeklyTemporalProcess -from pyrenew.model import PyrenewBuilder -from pyrenew.observation import NegativeBinomialNoise, PopulationCounts -from pyrenew.randomvariable import DistributionalVariable -from pyrenew.time import MMWR_WEEK -``` - ## Overview An ascertainment rate is the rate at which an infection results in an observed event. @@ -52,42 +30,32 @@ $$ \alpha_s \sum_{d=0}^{D_s} I(t-d)\,\pi_{s,d}, $$ -where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate, and $\pi_{s,d}$ is its delay distribution. +where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate and $\pi_{s,d}$ is its delay distribution. From a modeling perspective, an ascertainment rate may be considered constant over the observation period or allowed to vary over time. A time-varying ascertainment component combines a scalar baseline rate with a temporal process that describes deviations from that baseline over time as follows: -$\mathrm{b}_{\mathrm{s}}$ is the baseline rate for signal $\mathrm{s}$, $\mathrm{z}_{\mathrm{s}}(\mathrm{t})$ is its temporal deviation, and +$b_s$ is the baseline rate for signal $s$, $z_s(t)$ is its temporal deviation, and $$ - \alpha_s(t) = \operatorname{logit}^{-1} \left(\operatorname{logit}(\mathrm{b}_{\mathrm{s}}) + \mathrm{z}_{\mathrm{s}}(t)\right) + \alpha_s(t) = \operatorname{logit}^{-1} \left(\operatorname{logit}(b_s) + z_s(t)\right) $$ -Thus, $\mathrm{b}_{\mathrm{s}}$ is the ascertainment rate when the temporal deviation is zero while $\mathrm{z}_{\mathrm{s}}(\mathrm{t})$ determines how the rate changes during the observation period. +Thus, $b_s$ is the ascertainment rate when the temporal deviation is zero while $z_s(t)$ determines how the rate changes during the observation period. -The `AscertainmentModel` component of a PyRenew renewal process model provides the ascertainment rate for one or more observation signals, which can either be constant or time-varying. +We use `AscertainmentModel` classes to assign ascertainment rates to each signal in a PyRenew model. +These components reflect fundamental modeling assumptions about how information is shared across signals and over time: -- For a renewal process model for a single signal or for a multisignal model where the signals measure unrelated phenomena, the `IndependentAscertainment` component models the signal baseline and optional temporal process. +When a signal's ascertainment rate has little to do with ascertainment rates for other signals, it can be modeled as statistically *independent* of other ascertainment rates. +PyRenew provides an `IndependentAscertainment` class. +Alternatively, PyRenew allows the user to use a raw `RandomVariable` to specify any independent ascertainment rates. -- For a multisignal model where the signals are measuring a related phenomenon, such as hospitalizations and ED visits, which provide 2 views on health-care seeking behaviors, these signals can be modeled either by a `JointAscertainment` component which assigns the baseline rates a joint prior or as a `RatioLinkedAscertainment` which expresses one baseline rate relative to another. - -These components reflect fundamental modeling assumptions about how information is shared across signals and over time. - -## Independent ascertainment rates - -The `IndependentAscertainment` component models a signal baseline in the case where the ascertainment rate is constant, or a signal baseline plus temporal process if the ascertainment rate varies over time. - -An `AscertainmentModel` component is a subclass of a [Random Variable](random_variables.md). -If modeling a single signal with a constant ascertainment rate or some number of independent signals with constant ascertainment rates, instead of using an `IndependentAscertainment` component, it is possible to use a `RandomVariable` directly. - -```python -hospital_obs = PopulationCounts( - name="hospital", - ascertainment_rate_rv=DistributionalVariable("ihr", dist.Beta(2, 198)), - delay_distribution_rv=hosp_delay_rv, - noise=hosp_noise, -) -``` +Sometimes, however, ascertainment rates for two or more signals in a multi-signal model should be modeled as statistically *dependent*. +For example, hospital admissions and emergency department visits both reflect the degree to which infected individuals need to and are able to obtain care. +The two rates might be expected to scale up or down together. +In PyRenew, use the `JointAscertainment` and `RatioLinkedAscertainment` classes to model dependent ascertainment rates. +The `JointAscertainment` class assigns the baseline rates a joint prior. +The `RatioLinkedAscertainment` expresses one baseline rate relative to another. ## Ascertainment components and the builder @@ -123,15 +91,43 @@ It does not have to match the observation-process name, but matching names usual ## Choosing an ascertainment model -`JointAscertainment` and `RatioLinkedAscertainment` both relate signal rates, but they represent the relationship differently. +All three ascertainment components can produce constant rates or add signal-specific temporal processes. +They differ primarily in how the scalar baseline rates are specified and whether relationships are imposed among signals. + + | Component | Baseline | Relationship | Probability bounds | + | ----------------------------------------------------------------- | -------------------------------------------------------------- | ------------------------------------ | ------------------------------------------------------ | + | IndependentAscertainment | One `RandomVariable` per signal | No relationship imposed | Rates must be valid probabilities | + | `JointAscertainment` | Logit-scale multivariate normal | Prior covariance | Inverse logit keeps every rate in `(0, 1)` | + | `RatioLinkedAscertainment` | A base rate `RandomVariable`
and a ratio `RandomVariable` | Linked rate = base rate x multiplier | The derived linked rate is not bounded by construction | + +Use `IndependentAscertainment` when signal rates are not related through a joint prior or deterministic relationship. +It is particularly useful when one or more signals need temporal processes or when a named model-level component makes the model easier to organize. + +Use `JointAscertainment` when prior dependence among baseline rates is naturally represented by a covariance structure on the logit scale. - | Model | Parameters with priors | Relationship represented | Probability bounds | - | -------------------------- | ------------------------------------- | -------------------------------------------------------- | -------------------------------------------------------------------- | - | `JointAscertainment` | A vector of logit-scale rates | Prior covariance controls how the rates vary together | Every rate is transformed to `(0, 1)` | - | `RatioLinkedAscertainment` | A base rate and a relative multiplier | The linked rate is the base rate multiplied by the ratio | The derived linked rate is not automatically constrained to `(0, 1)` | +Use `RatioLinkedAscertainment` when a relative quantity, such as IHR / IEDR, is scientifically meaningful and can be assigned an interpretable prior. -Use `JointAscertainment` when prior dependence between the rates is most naturally described through a covariance structure. -Use `RatioLinkedAscertainment` when the relative quantity, such as IHR / IEDR, is itself 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. + +## Independent ascertainment rates + +The `IndependentAscertainment` component models a signal baseline in the case where the ascertainment rate is constant, or a signal baseline plus temporal process if the ascertainment rate varies over time. +Use `IndependentAscertainment` when a signal needs a temporal process or when representing the rates through a named model-level component makes the model easier to organize. + +For constant rate that is not related to another signal's rate, an ascertainment component is not required. +In this situation, the ascertainment component would only be used once, when specifying the `ascertainment_rate_rv` on the observation process corresponding to that signal via `ascertainment.for_signal()` method. +Since the `ascertainment.for_signal()` method returns a `RandomVariable` accessor, specifying the `RandomVariable` directly, simplifies the model specification: + +``` +hospital_obs = PopulationCounts( + name="hospital", + ascertainment_rate_rv=DistributionalVariable("ihr", dist.Beta(2, 198)), + delay_distribution_rv=hosp_delay_rv, + noise=hosp_noise, +) +``` + +For multiple independent signals with constant rates, each observation process can similarly receive its own `RandomVariable`. ## Joint ascertainment @@ -211,7 +207,105 @@ The priors are assigned to the base rate and the ratio, and together they induce The ratio may be any positive value. Unlike `JointAscertainment`, this component does not automatically constrain the derived linked rate to the unit interval. If both rates should be $<1$ (for instance because both are probabilities), make sure to choose priors that make $\alpha_{\mathrm{base}}\rho > 1$ negligibly likely. -The component does not clip an invalid linked rate, because clipping would change the specified statistical model. + +## Time-varying ascertainment + +Any ascertainment component can add a temporal process to one or more signals. +The temporal process generates the deviation $z_s(t)$, and the ascertainment component combines it with the signal's baseline rate: Because this calculation takes the logit of $b_s$, a time-varying signal requires $0 < b_s < 1$. +A constant ascertainment rate may equal $0$ or $1$, but those endpoints cannot be combined with a temporal process because their logits are infinite and thus a finite temporal deviation would leave the rate fixed at the endpoint. + +$$ + \alpha_s(t) + = + \operatorname{logit}^{-1} + \left( + \operatorname{logit}(b_s) + z_s(t) + \right). +$$ + +For example, the following component gives the ED ascertainment rate a logit-normal baseline and a weekly AR(1) deviation. + +First, define the scalar baseline rate: + +```python +ed_baseline_rv = LogitNormalVariable( + name="p_ed_visit", + median=0.005, + scale=0.3, +) +``` + +The baseline has median 0.005 on the probability scale. +Its `scale` parameter is the standard deviation on the logit scale. + +Next, define a temporal process for deviations from that baseline: + +```python +ed_deviation_process = WeeklyTemporalProcess( + inner=AR1( + autoreg_rv=DistributionalVariable( + "autoreg_p_ed_visit", + dist.Beta(1, 100), + ), + innovation_sd_rv=DistributionalVariable( + "p_ed_visit_w_sd", + dist.TruncatedNormal(0.0, 0.01, low=0), + ), + parameterization="innovation", + ), + start_dow=MMWR_WEEK, +) +``` + +The `AR1` process describes how the deviation evolves from one week to the next. +The autoregressive coefficient controls its persistence, and the innovation standard deviation controls the size of weekly changes. +`WeeklyTemporalProcess` aligns those changes with MMWR weeks and expands the weekly values over the model's daily time axis. + +Finally, associate both objects with the ED signal: + +```python +ascertainment = IndependentAscertainment( + name="ed_ascertainment", + rate_rvs={"ed": ed_baseline_rv}, + temporal_processes={"ed": ed_deviation_process}, +) +``` + +The keys in `rate_rvs` and `temporal_processes` identify the signal to which each object applies. +For a constant ascertainment rate, omit the `temporal_processes` argument. +For a component with multiple signals, include entries only for the signals whose rates should vary over time. +For example: + +``` +# Constant rate +IndependentAscertainment( + name="ed_ascertainment", + rate_rvs={"ed": ed_baseline_rv}, +) +``` + +``` +# Time-varying ED rate +IndependentAscertainment( + name="ed_ascertainment", + rate_rvs={"ed": ed_baseline_rv}, + temporal_processes={"ed": ed_deviation_process}, +) +``` + +Register and connect the component using the same builder pattern as for a constant ascertainment component: + +```python +builder.add_ascertainment(ascertainment) + +ed_obs = PopulationCounts( + name="ed", + ascertainment_rate_rv=ascertainment.for_signal("ed"), + ... +) +``` + +The same pattern applies to temporal processes supplied to `JointAscertainment` and `RatioLinkedAscertainment`. ## Building a model with an ascertainment component diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index 65744ab83..aa5c998d1 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -92,6 +92,12 @@ class AscertainmentModel(metaclass=ABCMeta): process to produce a rate trajectory over the model period. Otherwise, the baseline rate applies throughout the model period. + A baseline rate may lie in [0, 1] when no temporal process is configured. + A signal with a temporal process requires a baseline rate strictly inside + (0, 1), because the temporal deviation is added on the logit scale. + At either endpoint, the logit is infinite and finite deviations cannot + change the rate. + The model samples one scalar baseline rate per signal and can add an optional signal-specific temporal deviation on the logit scale. Subclasses implement ``_sample_baseline_rates()`` to define relationships among the diff --git a/test/integration/conftest.py b/test/integration/conftest.py index 85dc7a647..0d4ea6fce 100644 --- a/test/integration/conftest.py +++ b/test/integration/conftest.py @@ -37,7 +37,11 @@ from pyrenew.latent.population_infections import PopulationInfections from pyrenew.model import MultiSignalModel, PyrenewBuilder from pyrenew.observation import NegativeBinomialNoise, PopulationCounts -from pyrenew.randomvariable import DistributionalVariable, TransformedVariable +from pyrenew.randomvariable import ( + DistributionalVariable, + LogitNormalVariable, + TransformedVariable, +) from pyrenew.time import MMWR_WEEK from test.test_helpers import fixed_ar1, fixed_ar1_state, fixed_differenced_ar1_state @@ -212,16 +216,10 @@ def e_rtdaily_rw_timevary_model( ascertainment = IndependentAscertainment( name="ed_ascertainment", rate_rvs={ - "ed": TransformedVariable( - "p_ed_visit", - DistributionalVariable( - "p_ed_visit_mean", - dist.Normal( - transformation.SigmoidTransform().inv(0.005), - 0.3, - ), - ), - transforms=transformation.SigmoidTransform(), + "ed": LogitNormalVariable( + name="p_ed_visit", + median=0.005, + scale=0.3, ) }, temporal_processes={ diff --git a/test/integration/test_population_infections_e_rtdaily_rw_timevary.py b/test/integration/test_population_infections_e_rtdaily_rw_timevary.py index 84f8566a1..e7ce7dcf0 100644 --- a/test/integration/test_population_infections_e_rtdaily_rw_timevary.py +++ b/test/integration/test_population_infections_e_rtdaily_rw_timevary.py @@ -107,7 +107,7 @@ def test_full_graph_records_time_varying_ascertainment( trajectory = trace["ed_ascertainment_ed"]["value"] predicted = trace["ed_predicted"]["value"] - assert trace["p_ed_visit_mean"]["type"] == "sample" + assert trace["logit_p_ed_visit"]["type"] == "sample" assert baseline.shape == () assert weekly.ndim == 2 assert weekly.shape[1] == 1 @@ -147,7 +147,7 @@ def test_short_mcmc_fit_runs( jax.block_until_ready(samples) required_sites = { - "p_ed_visit_mean", + "logit_p_ed_visit", "autoreg_p_ed_visit", "p_ed_visit_w_sd", "ed_ascertainment_ed", From d7b494a7deb1c0f65154979ef6e20ff5fe07eda8 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 7 Oct 2026 15:13:15 -0400 Subject: [PATCH 08/12] changes per code review --- docs/tutorials/ascertainment.qmd | 8 +++++--- pyrenew/ascertainment/base.py | 6 +++--- test/test_ascertainment.py | 17 +++++++++++++++++ 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index e4dbb8b5b..d51706dad 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -27,7 +27,7 @@ For a count observation process, the ascertainment rate scales delayed latent in $$ \mu_s(t) = - \alpha_s \sum_{d=0}^{D_s} I(t-d)\,\pi_{s,d}, + \alpha_s \sum_{d=0}^{D_s} I(t-d)\,\pi_{s,d} $$ where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate and $\pi_{s,d}$ is its delay distribution. @@ -211,7 +211,8 @@ If both rates should be $<1$ (for instance because both are probabilities), make ## Time-varying ascertainment Any ascertainment component can add a temporal process to one or more signals. -The temporal process generates the deviation $z_s(t)$, and the ascertainment component combines it with the signal's baseline rate: Because this calculation takes the logit of $b_s$, a time-varying signal requires $0 < b_s < 1$. +The temporal process generates the deviation $z_s(t)$, and the ascertainment component combines it with the signal's baseline rate. +Because this calculation takes the logit of $b_s$, a time-varying signal requires $0 < b_s < 1$. A constant ascertainment rate may equal $0$ or $1$, but those endpoints cannot be combined with a temporal process because their logits are infinite and thus a finite temporal deviation would leave the rate fixed at the endpoint. $$ @@ -309,7 +310,8 @@ The same pattern applies to temporal processes supplied to `JointAscertainment` ## Building a model with an ascertainment component -The complete builder pattern is the same for either ascertainment component - here we use a `RatioLinkedAscertainment` but the sequence of definitions is the same for `JointAscertainment`. +The complete builder pattern is the same for any ascertainment component - here we use a `RatioLinkedAscertainment` +but the sequence of definitions is the same for `JointAscertainment` or `IndividualAscertainment`. ```python builder = PyrenewBuilder() diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index aa5c998d1..235e6bd09 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -8,12 +8,11 @@ from abc import ABCMeta, abstractmethod from collections.abc import Mapping -import jax.nn as jnn import jax.numpy as jnp import numpyro from jax import Array -from jax.scipy.special import logit from jax.typing import ArrayLike +from numpyro.distributions.transforms import SigmoidTransform from numpyro.util import not_jax_tracer from pyrenew.ascertainment.context import get_ascertainment_value @@ -365,6 +364,7 @@ def sample( ) baseline_rates = self._validate_baseline_rates(self._sample_baseline_rates()) + logit_transform = SigmoidTransform().inv result: dict[str, ArrayLike] = {} for signal in self.signals: @@ -395,7 +395,7 @@ def sample( f"required shape is {required_shape}." ) deviation = jnp.squeeze(deviation, axis=-1) - rate = jnn.sigmoid(logit(baseline_rate) + deviation) + rate = logit_transform.inv(logit_transform(baseline_rate) + deviation) numpyro.deterministic(f"{self.name}_{signal}", rate) result[signal] = rate diff --git a/test/test_ascertainment.py b/test/test_ascertainment.py index 916085f92..69c87a4f7 100644 --- a/test/test_ascertainment.py +++ b/test/test_ascertainment.py @@ -399,6 +399,23 @@ def test_samples_fixed_and_temporal_rates_in_signal_order(self) -> None: expected_ed = jax.nn.sigmoid(logit(0.3) + 0.2 * jnp.arange(4)) assert jnp.allclose(values["ed"], expected_ed) + def test_extreme_temporal_deviations_remain_inside_open_interval(self) -> None: + """Test transformed temporal rates do not reach probability endpoints.""" + model = FixedAscertainmentModel( + "ascertainment", + {"hospital": 0.5}, + temporal_processes={ + "hospital": DeterministicTemporalProcess( + result=jnp.array([[-100.0], [100.0]]) + ) + }, + ) + + rates = model.sample(n_timepoints=2)["hospital"] + + assert jnp.all(rates > 0.0) + assert jnp.all(rates < 1.0) + def test_rejects_malformed_temporal_output(self) -> None: """Test temporal output keeps its singleton process dimension.""" model = FixedAscertainmentModel( From 836e3813e91a47607f82b3d27d819c51c82c09ec Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Wed, 7 Oct 2026 15:13:57 -0400 Subject: [PATCH 09/12] changes per code review --- docs/tutorials/ascertainment.qmd | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index d51706dad..d42597b5a 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -310,8 +310,7 @@ The same pattern applies to temporal processes supplied to `JointAscertainment` ## Building a model with an ascertainment component -The complete builder pattern is the same for any ascertainment component - here we use a `RatioLinkedAscertainment` -but the sequence of definitions is the same for `JointAscertainment` or `IndividualAscertainment`. +The complete builder pattern is the same for any ascertainment component - here we use a `RatioLinkedAscertainment` but the sequence of definitions is the same for `JointAscertainment` or `IndividualAscertainment`. ```python builder = PyrenewBuilder() From 86ebb851e94ab8d0232a7dcef0039279f9eafd65 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 14:15:46 -0400 Subject: [PATCH 10/12] changes per code review --- docs/tutorials/ascertainment.qmd | 9 ++--- pyrenew/ascertainment/base.py | 53 +++----------------------- pyrenew/model/multisignal_model.py | 46 ++++------------------ test/integration/conftest.py | 4 +- test/test_ascertainment.py | 55 --------------------------- test/test_observation_counts.py | 12 +++++- test/test_pyrenew_builder.py | 61 ------------------------------ 7 files changed, 28 insertions(+), 212 deletions(-) diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd index d42597b5a..f135f8db6 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -48,7 +48,7 @@ These components reflect fundamental modeling assumptions about how information When a signal's ascertainment rate has little to do with ascertainment rates for other signals, it can be modeled as statistically *independent* of other ascertainment rates. PyRenew provides an `IndependentAscertainment` class. -Alternatively, PyRenew allows the user to use a raw `RandomVariable` to specify any independent ascertainment rates. +Alternatively, PyRenew allows the user to use a raw `RandomVariable` to specify any independent ascertainment rate. Sometimes, however, ascertainment rates for two or more signals in a multi-signal model should be modeled as statistically *dependent*. For example, hospital admissions and emergency department visits both reflect the degree to which infected individuals need to and are able to obtain care. @@ -100,14 +100,13 @@ They differ primarily in how the scalar baseline rates are specified and whether | `JointAscertainment` | Logit-scale multivariate normal | Prior covariance | Inverse logit keeps every rate in `(0, 1)` | | `RatioLinkedAscertainment` | A base rate `RandomVariable`
and a ratio `RandomVariable` | Linked rate = base rate x multiplier | The derived linked rate is not bounded by construction | -Use `IndependentAscertainment` when signal rates are not related through a joint prior or deterministic relationship. -It is particularly useful when one or more signals need temporal processes or when a named model-level component makes the model easier to organize. +Use `IndependentAscertainment` when a signal or signals have time-varying ascertainment rates that are not related through a joint prior or deterministic relationship. Use `JointAscertainment` when prior dependence among baseline rates is naturally represented by a covariance structure on the logit scale. 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 @@ -116,7 +115,7 @@ Use `IndependentAscertainment` when a signal needs a temporal process or when re For constant rate that is not related to another signal's rate, an ascertainment component is not required. In this situation, the ascertainment component would only be used once, when specifying the `ascertainment_rate_rv` on the observation process corresponding to that signal via `ascertainment.for_signal()` method. -Since the `ascertainment.for_signal()` method returns a `RandomVariable` accessor, specifying the `RandomVariable` directly, simplifies the model specification: +Since the `ascertainment.for_signal()` method returns a `RandomVariable` accessor specifying the `RandomVariable` directly, simplifies the model specification: ``` hospital_obs = PopulationCounts( diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py index 235e6bd09..a81a977eb 100644 --- a/pyrenew/ascertainment/base.py +++ b/pyrenew/ascertainment/base.py @@ -252,6 +252,11 @@ def _validate_baseline_rates( ) -> dict[str, Array]: """Validate and convert scalar baseline ascertainment rates. + A signal without a temporal process may have a baseline in [0, 1]. + A signal with a temporal process requires a baseline strictly inside + (0, 1), because the temporal deviation is added on the logit scale + and the logit of 0 or 1 is infinite. + Parameters ---------- baseline_rates @@ -402,54 +407,6 @@ def sample( return result - def validate_sampled_values( - self, - values: Mapping[str, ArrayLike], - n_timepoints: int, - ) -> None: - """Validate sampled signal names and output shapes. - - Parameters - ---------- - values - Sampled ascertainment values keyed by signal name. - n_timepoints - Length of the shared model time axis. - - Raises - ------ - TypeError - If ``values`` is not a mapping. - ValueError - If signal names do not match this model or a value is neither a - scalar nor a full-axis trajectory. - """ - if not isinstance(values, Mapping): - raise TypeError( - f"Ascertainment model {self.name!r} must return a mapping, " - f"got {type(values).__name__}." - ) - - actual_signals = set(values) - expected_signals = set(self.signals) - if actual_signals != expected_signals: - missing = tuple(signal for signal in self.signals if signal not in values) - extra = tuple(signal for signal in values if signal not in expected_signals) - raise ValueError( - f"Ascertainment model {self.name!r} must return exactly signals " - f"{self.signals}. Missing: {missing}. Extra: {extra}." - ) - - allowed_shapes = ((), (n_timepoints,)) - for signal in self.signals: - actual_shape = jnp.asarray(values[signal]).shape - if actual_shape not in allowed_shapes: - raise ValueError( - f"Ascertainment model {self.name!r}, signal {signal!r}, " - f"returned shape {actual_shape}; allowed shapes are () and " - f"({n_timepoints},)." - ) - @abstractmethod def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]: """Sample one scalar baseline ascertainment rate per signal. diff --git a/pyrenew/model/multisignal_model.py b/pyrenew/model/multisignal_model.py index 97f07bb7a..3e339547e 100644 --- a/pyrenew/model/multisignal_model.py +++ b/pyrenew/model/multisignal_model.py @@ -7,7 +7,6 @@ from __future__ import annotations import datetime as dt -from collections.abc import Mapping import jax.numpy as jnp import numpy as np @@ -353,40 +352,6 @@ def validate_data( **obs_data, ) - def _sample_ascertainment_models( - self, - n_timepoints: int, - first_day_dow: int | None, - ) -> dict[str, Mapping[str, ArrayLike]]: - """Sample and validate every registered ascertainment model. - - Parameters - ---------- - n_timepoints - Length of the padded shared model axis. - first_day_dow - Day-of-week for the first point on that axis, or ``None`` when no - calendar anchor was supplied. - - Returns - ------- - dict[str, Mapping[str, ArrayLike]] - Sampled ascertainment values keyed first by registered model name - and then by signal name. - """ - sampled_values: dict[str, Mapping[str, ArrayLike]] = {} - for name, ascertainment_model in self.ascertainment_models.items(): - values = ascertainment_model.sample( - n_timepoints=n_timepoints, - first_day_dow=first_day_dow, - ) - ascertainment_model.validate_sampled_values( - values, - n_timepoints=n_timepoints, - ) - sampled_values[name] = values - return sampled_values - def sample( self, n_days_post_init: int, @@ -457,10 +422,13 @@ def sample( } n_timepoints = self.latent.n_initialization_points + n_days_post_init - ascertainment_values = self._sample_ascertainment_models( - n_timepoints=n_timepoints, - first_day_dow=first_day_dow, - ) + ascertainment_values = { + name: ascertainment_model.sample( + n_timepoints=n_timepoints, + first_day_dow=first_day_dow, + ) + for name, ascertainment_model in self.ascertainment_models.items() + } with ascertainment_context(ascertainment_values): # Apply each observation process diff --git a/test/integration/conftest.py b/test/integration/conftest.py index 0d4ea6fce..96507723e 100644 --- a/test/integration/conftest.py +++ b/test/integration/conftest.py @@ -172,9 +172,7 @@ def e_rtdaily_rw_timevary_model( ) -> MultiSignalModel: """Build an ED model with daily Rt and time-varying ascertainment. - This fixture reproduces the PyRenew structure of the - ``e_rtdaily_rw_timevary`` model without depending on the external ARM - model builder. The ascertainment baseline is sampled independently and a + The ascertainment baseline is sampled independently and a calendar-aligned weekly AR(1) process supplies logit-scale deviations. Parameters diff --git a/test/test_ascertainment.py b/test/test_ascertainment.py index 69c87a4f7..5c8a755f9 100644 --- a/test/test_ascertainment.py +++ b/test/test_ascertainment.py @@ -432,61 +432,6 @@ def test_rejects_malformed_temporal_output(self) -> None: ): model.sample(n_timepoints=4) - def test_accepts_scalar_and_full_axis_values(self) -> None: - """Test scalar and exact full-axis trajectories are valid.""" - model = FixedAscertainmentModel( - "ascertainment", - {"hospital": 0.2, "ed": 0.3}, - ) - - model.validate_sampled_values( - {"hospital": jnp.array(0.2), "ed": jnp.arange(4)}, - n_timepoints=4, - ) - - def test_rejects_non_mapping_values(self) -> None: - """Test sampled values must be returned as a mapping.""" - model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) - - with pytest.raises(TypeError, match="ascertainment.*mapping"): - model.validate_sampled_values([0.2], n_timepoints=4) # type: ignore[arg-type] - - @pytest.mark.parametrize( - "values", - [ - {}, - {"hospital": 0.2, "ed": 0.3}, - ], - ) - def test_rejects_missing_or_extra_signals( - self, - values: Mapping[str, ArrayLike], - ) -> None: - """Test sampled signal names must exactly match the model signals.""" - model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) - - with pytest.raises(ValueError, match="ascertainment.*signals"): - model.validate_sampled_values(values, n_timepoints=4) - - @pytest.mark.parametrize( - "value", - [ - jnp.ones(1), - jnp.ones(3), - jnp.ones((4, 1)), - jnp.ones((4, 1, 1)), - ], - ) - def test_rejects_invalid_sample_shapes(self, value: ArrayLike) -> None: - """Test only scalar or exact full-axis values are accepted.""" - model = FixedAscertainmentModel("ascertainment", {"hospital": 0.2}) - - with pytest.raises( - ValueError, - match=r"ascertainment.*hospital.*shape.*\(\).*\(4,\)", - ): - model.validate_sampled_values({"hospital": value}, n_timepoints=4) - def _make_public_ascertainment( # numpydoc ignore=RT01 model_kind: str, diff --git a/test/test_observation_counts.py b/test/test_observation_counts.py index ac87d7d0d..bac63ddd3 100644 --- a/test/test_observation_counts.py +++ b/test/test_observation_counts.py @@ -286,7 +286,14 @@ def test_delay_padding_is_preserved(self) -> None: assert jnp.all(jnp.isfinite(predicted[2:])) def test_delay_padding_does_not_create_nan_rate_gradients(self) -> None: - """Test padded potential counts have a finite ascertainment gradient.""" + """Test padded counts stay NaN and have a finite ascertainment gradient. + + The gradient is finite because ``_apply_ascertainment`` multiplies the + rate only by NaN-free counts. The ``jnp.where`` in ``finite_sum`` only + removes NaN values from the summed output; it would not prevent a NaN + gradient if ``_apply_ascertainment`` multiplied the rate by the padded + counts directly. + """ process = PopulationCounts( name="hospital", ascertainment_rate_rv=DeterministicVariable("rate", 0.2), @@ -306,8 +313,11 @@ def finite_sum(rate: ArrayLike) -> ArrayLike: scaled = process._apply_ascertainment(potential_counts, rate) return jnp.sum(jnp.where(jnp.isnan(scaled), 0.0, scaled)) + scaled = process._apply_ascertainment(potential_counts, jnp.array(0.2)) gradient = jax.grad(finite_sum)(jnp.array(0.2)) + assert jnp.isnan(scaled[0]) + assert jnp.all(jnp.isfinite(scaled[1:])) assert jnp.isfinite(gradient) assert jnp.isclose(gradient, 6.0) diff --git a/test/test_pyrenew_builder.py b/test/test_pyrenew_builder.py index 6d1835c22..e8faea662 100644 --- a/test/test_pyrenew_builder.py +++ b/test/test_pyrenew_builder.py @@ -2,7 +2,6 @@ Tests for PyrenewBuilder and MultiSignalModel. """ -from collections.abc import Mapping from datetime import date, timedelta import jax.numpy as jnp @@ -94,41 +93,6 @@ def sample( return jnp.zeros((n_timepoints, n_processes)) -class MalformedAscertainmentModel(AscertainmentModel): - """Override the common sampler with malformed output for boundary tests.""" - - def __init__(self, name: str) -> None: - """Initialize a deliberately nonconforming ascertainment model.""" - super().__init__(name=name, signals=("hospital",)) - - def _sample_baseline_rates( # numpydoc ignore=RT01 - self, - ) -> Mapping[str, ArrayLike]: - """Return a valid baseline that the overridden sampler does not use.""" - return {"hospital": 0.2} - - def sample( - self, - n_timepoints: int, - first_day_dow: int | None = None, - ) -> Mapping[str, ArrayLike]: - """Return an invalid two-dimensional final ascertainment value. - - Parameters - ---------- - n_timepoints - Shared-axis length, ignored. - first_day_dow - Shared-axis day of week, ignored. - - Returns - ------- - Mapping[str, ArrayLike] - Malformed signal output. - """ - return {"hospital": jnp.ones((2, 1))} - - def _obs_date_for_dow(target_first_day_dow: int, n_init: int) -> date: """ Return an ``obs_start_date`` whose axis-origin day-of-week matches. @@ -1006,31 +970,6 @@ def test_daily_ascertainment_process_works_without_obs_start_date( assert process.sample_calls[0]["first_day_dow"] is None - def test_malformed_ascertainment_fails_before_observation_sampling( - self, - simple_builder: PyrenewBuilder, - ) -> None: - """Test output validation completes before any observation runs.""" - built = simple_builder.build() - malformed = MalformedAscertainmentModel("malformed") - model = MultiSignalModel( - latent_process=built.latent, - observations=built.observations, - ascertainment_models={"malformed": malformed}, - ) - - with numpyro.handlers.seed(rng_seed=42): - with numpyro.handlers.trace() as trace: - with pytest.raises(ValueError, match=r"malformed.*hospital.*shape"): - model.sample( - n_days_post_init=5, - population_size=1_000_000, - subpop_fractions=SUBPOP_FRACTIONS, - hospital={"obs": None}, - ) - - assert "hospital_obs" not in trace - def test_first_day_dow_reaches_calendar_aligned_latent_process(self): """MultiSignalModel forwards model-axis day of week to the latent process.""" latent = PopulationInfections( From b9749559fcabb1b394c5b28505992680c2871058 Mon Sep 17 00:00:00 2001 From: Mitzi Morris Date: Tue, 13 Oct 2026 14:22:32 -0400 Subject: [PATCH 11/12] ascertainment class name fix --- 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 f135f8db6..0ad43a5fe 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -309,7 +309,7 @@ The same pattern applies to temporal processes supplied to `JointAscertainment` ## Building a model with an ascertainment component -The complete builder pattern is the same for any ascertainment component - here we use a `RatioLinkedAscertainment` but the sequence of definitions is the same for `JointAscertainment` or `IndividualAscertainment`. +The complete builder pattern is the same for any ascertainment component - here we use a `RatioLinkedAscertainment` but the sequence of definitions is the same for `JointAscertainment` or `IndependentAscertainment`. ```python builder = PyrenewBuilder() From ff8ed3899c8a67548394f0038d13568a6a01dcc0 Mon Sep 17 00:00:00 2001 From: "Dylan H. Morris" Date: Tue, 13 Oct 2026 15:24:33 -0400 Subject: [PATCH 12/12] Update docs/tutorials/ascertainment.qmd --- 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 0ad43a5fe..1f01422e8 100644 --- a/docs/tutorials/ascertainment.qmd +++ b/docs/tutorials/ascertainment.qmd @@ -115,7 +115,7 @@ Use `IndependentAscertainment` when a signal needs a temporal process or when re For constant rate that is not related to another signal's rate, an ascertainment component is not required. In this situation, the ascertainment component would only be used once, when specifying the `ascertainment_rate_rv` on the observation process corresponding to that signal via `ascertainment.for_signal()` method. -Since the `ascertainment.for_signal()` method returns a `RandomVariable` accessor specifying the `RandomVariable` directly, simplifies the model specification: +Since the `ascertainment.for_signal()` method returns a `RandomVariable` accessor, specifying the `RandomVariable` directly simplifies the model specification: ``` hospital_obs = PopulationCounts(