Refine AscertainmentModel component to allow for time-varying component

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                                     |
+  | ----------------------------------------------------------------- | -------------------------------------------------------------- | ------------------------------------ | ------------------------------------------------------ |
+  | <code style="white-space: nowrap">IndependentAscertainment</code> | 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` <br> 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.

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)

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