diff --git a/docs/tutorials/ascertainment.qmd b/docs/tutorials/ascertainment.qmd
index bba71461..1f01422e 100644
--- a/docs/tutorials/ascertainment.qmd
+++ b/docs/tutorials/ascertainment.qmd
@@ -17,70 +17,45 @@ jupyter:
name: python3
---
-```{python}
-#| label: setup
-#| output: false
-
-import jax.numpy as jnp
-import numpyro.distributions as dist
-
-from pyrenew.ascertainment import JointAscertainment, RatioLinkedAscertainment
-from pyrenew.deterministic import DeterministicPMF, DeterministicVariable
-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.
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:
$$
\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.
+where $s$ identifies the signal, $\alpha_s$ is its ascertainment rate and $\pi_{s,d}$ is its delay distribution.
-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.
+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:
-PyRenew provides two such components:
+$b_s$ is the baseline rate for signal $s$, $z_s(t)$ is its temporal deviation, and
-- `JointAscertainment` specifies a joint prior for two or more rates.
-- `RatioLinkedAscertainment` specifies one rate as a multiple of another rate.
-
-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.
-
-## Independent ascertainment rates
+$$
+ \alpha_s(t) = \operatorname{logit}^{-1} \left(\operatorname{logit}(b_s) + z_s(t)\right)
+$$
-An ascertainment component is not required when the signal rates are modeled independently.
-Pass each rate directly to its observation process:
+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.
-```python
-hospital_obs = PopulationCounts(
- name="hospital",
- ascertainment_rate_rv=DistributionalVariable("ihr", dist.Beta(2, 198)),
- delay_distribution_rv=hosp_delay_rv,
- noise=hosp_noise,
-)
+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:
-ed_obs = PopulationCounts(
- name="ed_visits",
- ascertainment_rate_rv=DistributionalVariable("iedr", dist.Beta(2, 98)),
- delay_distribution_rv=ed_delay_rv,
- noise=ed_noise,
-)
-```
+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 rate.
-This specification does not encode prior dependence between IHR and IEDR; each observation process has its own ascertainment-rate prior.
+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
@@ -116,15 +91,42 @@ 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 a signal or signals have time-varying ascertainment rates that are not related through a joint prior or deterministic relationship.
- | 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 `JointAscertainment` when prior dependence among baseline rates is naturally represented by a covariance structure on the logit scale.
-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.
+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.
+
+## 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
@@ -204,11 +206,110 @@ 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.
-## Building a model an ascertainment component
+## 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.
-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`.
+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
+
+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()
@@ -257,18 +358,18 @@ 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.
+- 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.
-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 6d5139e1..6e64e2cb 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
@@ -786,7 +786,7 @@ obs_data_90 = {
The observation dictionary has the following structure:
-```python
+```text
{
"obs_start_date": ...,
"hospital": {
@@ -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/docs/tutorials/latent_infections.qmd b/docs/tutorials/latent_infections.qmd
index 8a774f60..3408011f 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 392e79fc..9d9a603f 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,27 @@ $$
$$
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 +330,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
-
-# Scale infections by IHR (ascertainment rate)
-infections_scaled = infections * float(ihr_rv.sample())
+# Compute predicted admissions (convolution and ascertainment, no noise)
+from pyrenew.convolve import compute_delay_incidence
-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 +521,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 a8316b8c..1f6f3716 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/panache.toml b/panache.toml
index 798a6847..bca7d537 100644
--- a/panache.toml
+++ b/panache.toml
@@ -3,5 +3,3 @@ flavor = "gfm"
wrap = "sentence"
[formatters]
python = "ruff"
-[linters]
-python = "ruff"
diff --git a/pyrenew/ascertainment/__init__.py b/pyrenew/ascertainment/__init__.py
index cc616060..3f584a5e 100644
--- a/pyrenew/ascertainment/__init__.py
+++ b/pyrenew/ascertainment/__init__.py
@@ -4,12 +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
__all__ = [
"AscertainmentModel",
"AscertainmentSignal",
+ "IndependentAscertainment",
"JointAscertainment",
"RatioLinkedAscertainment",
]
diff --git a/pyrenew/ascertainment/base.py b/pyrenew/ascertainment/base.py
index 9c0c823f..a81a977e 100644
--- a/pyrenew/ascertainment/base.py
+++ b/pyrenew/ascertainment/base.py
@@ -8,9 +8,15 @@
from abc import ABCMeta, abstractmethod
from collections.abc import Mapping
+import jax.numpy as jnp
+import numpyro
+from jax import Array
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
+from pyrenew.latent import TemporalProcess
from pyrenew.metaclass import RandomVariable
@@ -77,17 +83,36 @@ 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:
+ 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
+ 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.
+
+ 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(...)
@@ -99,16 +124,13 @@ 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.
"""
def __init__(
self,
name: str,
signals: tuple[str, ...],
+ temporal_processes: Mapping[str, TemporalProcess] | None = None,
) -> None:
"""
Initialize an ascertainment model.
@@ -119,6 +141,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(
@@ -134,6 +167,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.
@@ -168,20 +232,194 @@ def for_signal(self, signal_name: str) -> AscertainmentSignal:
signal_name=signal_name,
)
- @abstractmethod
- def sample(self, **kwargs: object) -> Mapping[str, ArrayLike]:
+ def requires_calendar_anchor(self) -> bool:
+ """Return whether any temporal process requires a calendar anchor.
+
+ Returns
+ -------
+ bool
+ ``True`` when the model needs a calendar anchor; otherwise
+ ``False``.
"""
- Sample all signal-specific ascertainment values owned by this model.
+ 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.
+
+ 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
----------
- **kwargs
- Additional model-context arguments supplied by ``MultiSignalModel``.
- Subclasses may ignore unused values.
+ 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())
+ logit_transform = SigmoidTransform().inv
+
+ 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 = logit_transform.inv(logit_transform(baseline_rate) + deviation)
+
+ numpyro.deterministic(f"{self.name}_{signal}", rate)
+ result[signal] = rate
+
+ return result
+
+ @abstractmethod
+ 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 00000000..34fc6b73
--- /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 2fc42d74..0796c0f7 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 b209cc27..c4579fed 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/convolve.py b/pyrenew/convolve.py
index 0a9129f9..10d12999 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 cd8efa14..3e339547 100644
--- a/pyrenew/model/multisignal_model.py
+++ b/pyrenew/model/multisignal_model.py
@@ -240,6 +240,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 +306,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().
@@ -374,7 +382,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,8 +421,12 @@ def sample(
"subpop": inf_all,
}
+ n_timepoints = self.latent.n_initialization_points + n_days_post_init
ascertainment_values = {
- name: ascertainment_model.sample()
+ name: ascertainment_model.sample(
+ n_timepoints=n_timepoints,
+ first_day_dow=first_day_dow,
+ )
for name, ascertainment_model in self.ascertainment_models.items()
}
diff --git a/pyrenew/model/pyrenew_builder.py b/pyrenew/model/pyrenew_builder.py
index 43801373..045bc9ab 100644
--- a/pyrenew/model/pyrenew_builder.py
+++ b/pyrenew/model/pyrenew_builder.py
@@ -153,31 +153,40 @@ 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. 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"),
...
)
)
@@ -190,7 +199,8 @@ def add_ascertainment(
----------
ascertainment_model
Configured ascertainment model instance, such as
- ``JointAscertainment``.
+ ``IndependentAscertainment``, ``JointAscertainment``, or
+ ``RatioLinkedAscertainment``.
Returns
-------
diff --git a/pyrenew/observation/base.py b/pyrenew/observation/base.py
index 6fdfab16..c7209748 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 3aaa7a34..31871df7 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/integration/conftest.py b/test/integration/conftest.py
index 6d8f264b..96507723 100644
--- a/test/integration/conftest.py
+++ b/test/integration/conftest.py
@@ -12,7 +12,12 @@
import polars as pl
import pytest
-from pyrenew.ascertainment import JointAscertainment, RatioLinkedAscertainment
+import pyrenew.transformation as transformation
+from pyrenew.ascertainment import (
+ IndependentAscertainment,
+ JointAscertainment,
+ RatioLinkedAscertainment,
+)
from pyrenew.datasets import (
load_example_infection_admission_interval,
load_synthetic_daily_ed_visits,
@@ -22,12 +27,21 @@
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,
+ LogitNormalVariable,
+ TransformedVariable,
+)
from pyrenew.time import MMWR_WEEK
from test.test_helpers import fixed_ar1, fixed_ar1_state, fixed_differenced_ar1_state
@@ -152,6 +166,119 @@ 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.
+
+ 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": LogitNormalVariable(
+ name="p_ed_visit",
+ median=0.005,
+ scale=0.3,
+ )
+ },
+ 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,
@@ -172,6 +299,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 +319,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 +330,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_e_rtdaily_rw_timevary.py b/test/integration/test_population_infections_e_rtdaily_rw_timevary.py
new file mode 100644
index 00000000..e7ce7dcf
--- /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["logit_p_ed_visit"]["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 = {
+ "logit_p_ed_visit",
+ "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"])
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 53b34f91..936fc126 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 64b051cd..42114c2d 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 6cf77cd2..5c8a755f 100644
--- a/test/test_ascertainment.py
+++ b/test/test_ascertainment.py
@@ -2,13 +2,20 @@
Tests for ascertainment models.
"""
+from collections.abc import Mapping
+
+import jax
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,
+ IndependentAscertainment,
JointAscertainment,
RatioLinkedAscertainment,
)
@@ -17,9 +24,677 @@
get_ascertainment_value,
)
from pyrenew.deterministic import DeterministicVariable
+from pyrenew.latent import TemporalProcess, WeeklyTemporalProcess
+from pyrenew.metaclass import RandomVariable
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,
+ 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(
+ 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 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."""
+
+ def __init__(
+ self,
+ name: str,
+ values: Mapping[str, ArrayLike],
+ temporal_processes: Mapping[str, TemporalProcess] | None = None,
+ ) -> None:
+ """Initialize a fixed ascertainment test double."""
+ super().__init__(
+ name=name,
+ signals=tuple(values),
+ temporal_processes=temporal_processes,
+ )
+ self.values = dict(values)
+ self.n_sample_calls = 0
+
+ def _sample_baseline_rates(self) -> Mapping[str, ArrayLike]:
+ """Return configured scalar baselines and record the sampling call.
+
+ Returns
+ -------
+ Mapping[str, ArrayLike]
+ Configured signal baselines.
+ """
+ self.n_sample_calls += 1
+ return self.values
+
+
+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( # numpydoc ignore=RT01
+ 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_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(
+ "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 _make_public_ascertainment( # numpydoc ignore=RT01
+ 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 TestPublicAscertainmentBehavior:
+ """Test common fixed and temporal behavior for every public class."""
+
+ @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,
+ )
+
+ with numpyro.handlers.seed(rng_seed=42):
+ with numpyro.handlers.trace() as trace:
+ values = ascertainment.sample(n_timepoints=4)
+
+ 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_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]
+ )
+
+ 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]
+ )
+
+ 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,
+ }
+
+ ascertainment = IndependentAscertainment(
+ name="ascertainment",
+ rate_rvs=rate_rvs,
+ )
+ rate_rvs.clear()
+
+ assert ascertainment.signals == ("ed", "hospital")
+ assert ascertainment.rate_rvs == {
+ "ed": ed_rv,
+ "hospital": hospital_rv,
+ }
+
+
+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},
+ )
+
+ 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",
+ }
+
+ 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 ascertainment.requires_calendar_anchor()
+ assert jnp.allclose(values["hospital"][:5], 0.2)
+ assert jnp.all(values["hospital"][5:] > 0.2)
+
+
+class TestConcreteTemporalAscertainment:
+ """Test temporal behavior supplied by concrete baseline relationships."""
+
+ 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])
+ ascertainment = JointAscertainment(
+ name="he_ascertainment",
+ signals=("hospital", "ed"),
+ baseline_rates=baseline_rates,
+ scale_tril=jnp.eye(2),
+ temporal_processes={
+ "hospital": DeterministicTemporalProcess(increment=0.3),
+ "ed": DeterministicTemporalProcess(increment=-0.2),
+ },
+ )
+
+ with numpyro.handlers.substitute(
+ data={"he_ascertainment_eta": logit(baseline_rates)}
+ ):
+ with numpyro.handlers.trace() as trace:
+ values = ascertainment.sample(n_timepoints=4)
+
+ 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])
+ 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 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)),
+ temporal_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 = ascertainment.sample(n_timepoints=4)
+
+ assert trace["iedr"]["type"] == "sample"
+ assert trace["ihr_rel_iedr"]["type"] == "sample"
+ 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)
+
+
class TestJointAscertainmentValidation:
"""Test JointAscertainment constructor validation."""
@@ -195,11 +870,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(
@@ -221,7 +898,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"}
@@ -235,7 +912,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"}
@@ -376,13 +1053,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"],
@@ -404,7 +1088,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_convolve.py b/test/test_convolve.py
index a9211f85..97742251 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 4aa38108..3a722136 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 5fd5b6ab..bac63ddd 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,159 @@ 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 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),
+ 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))
+
+ 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)
+
+
class TestCountsWithPriors:
"""Test aggregated count observation with uncertain parameters."""
diff --git a/test/test_pyrenew_builder.py b/test/test_pyrenew_builder.py
index a03af53e..e8faea66 100644
--- a/test/test_pyrenew_builder.py
+++ b/test/test_pyrenew_builder.py
@@ -8,8 +8,14 @@
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,
+ IndependentAscertainment,
+ JointAscertainment,
+ RatioLinkedAscertainment,
+)
from pyrenew.deterministic import DeterministicPMF, DeterministicVariable
from pyrenew.latent import (
InfectionsWithFeedback,
@@ -33,6 +39,60 @@
SUBPOP_FRACTIONS = jnp.array([0.3, 0.25, 0.45])
+class RecordingTemporalProcess:
+ """Return zero deviations while recording temporal-process arguments."""
+
+ step_size = 1
+
+ def __init__(
+ self,
+ requires_calendar_anchor: bool = False,
+ ) -> None:
+ """Initialize a recording temporal process."""
+ self.requires_calendar_anchor = requires_calendar_anchor
+ 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:
+ """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
+ -------
+ ArrayLike
+ Zero deviations with shape ``(n_timepoints, n_processes)``.
+ """
+ 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))
+
+
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 +118,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_ascertainment(
+ ascertainment: AscertainmentModel,
+) -> PyrenewBuilder:
+ """Build a two-count-signal model using one 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():
"""
@@ -370,6 +466,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
@@ -426,6 +601,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
@@ -491,19 +668,109 @@ 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_ascertainment(self) -> None:
+ """Test joint baseline sites and full-axis final trajectories."""
+ import jax.random
+ from numpyro.infer import Predictive
+
+ ascertainment = JointAscertainment(
+ name="he_ascertainment",
+ signals=("hospital", "ed"),
+ baseline_rates=jnp.array([0.01, 0.02]),
+ scale_tril=jnp.eye(2) * 0.1,
+ temporal_processes={
+ "hospital": fixed_random_walk(innovation_sd=0.05),
+ "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 set(model.ascertainment_models) == {"he_ascertainment"}
+ 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_ascertainment(self) -> None:
+ """Test ratio baselines persist beneath independently varying rates."""
+ import jax.random
+ from numpyro.infer import Predictive
+
+ ascertainment = RatioLinkedAscertainment(
+ name="he_ascertainment",
+ 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),
+ ),
+ temporal_processes={
+ "hospital": fixed_random_walk(innovation_sd=0.05),
+ "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["iedr"].shape == (2,)
+ assert samples["ihr_rel_iedr"].shape == (2,)
+ assert jnp.allclose(
+ 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)
+ 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
@@ -559,6 +826,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,
@@ -597,6 +865,111 @@ def test_manual_model_rejects_mismatched_ascertainment_key(self, simple_builder)
ascertainment_models={"wrong_name": ascertainment},
)
+ def test_forwards_full_axis_length_and_padded_axis_weekday(
+ self,
+ simple_builder: PyrenewBuilder,
+ ) -> None:
+ """Test temporal sampling receives full padded-axis context."""
+ built = simple_builder.build()
+ process = RecordingTemporalProcess()
+ ascertainment = IndependentAscertainment(
+ name="recording",
+ rate_rvs={"hospital": DeterministicVariable("ihr_recording", 0.2)},
+ temporal_processes={"hospital": process},
+ )
+ 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 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(
+ self,
+ simple_builder: PyrenewBuilder,
+ ) -> None:
+ """Test registered calendar-aligned ascertainment requires a date."""
+ built = simple_builder.build()
+ 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,
+ 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},
+ )
+
+ assert inner_process.sample_calls == []
+
+ def test_daily_ascertainment_process_works_without_obs_start_date(
+ self,
+ simple_builder: PyrenewBuilder,
+ ) -> None:
+ """Test a registered daily process does not require a calendar date."""
+ built = simple_builder.build()
+ 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):
+ 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_first_day_dow_reaches_calendar_aligned_latent_process(self):
"""MultiSignalModel forwards model-axis day of week to the latent process."""
latent = PopulationInfections(