End-to-end support for coarser-than-daily count signals #794

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

cdc-mitzimorris wants to merge 53 commits into main from mem_789_temporal_aggregation

+4,868 −724

_typos.toml

-Original file line number
+Diff line change
@@ Expand Up / @@ -3,3 +3,4 @@ @@
     # words that should not be corrected
     arange = "arange"
     lod = "lod"
+    dows = "dows"

docs/tutorials/building_multisignal_models.qmd

-Original file line number
+Diff line change
@@ Expand Up / @@ -39,6 +39,7 @@ import plotnine as p9 @@
     import pandas as pd
     import time
     import warnings
+    from datetime import date
     warnings.filterwarnings("ignore")
@@ Expand Down Expand Up / @@ -67,6 +68,7 @@ from pyrenew.latent import ( @@
         SubpopulationInfections,
         AR1,
         RandomWalk,
+        StepwiseTemporalProcess,
         GammaGroupSdPrior,
         HierarchicalNormalPrior,
     )
@@ Expand All / @@ -79,6 +81,7 @@ from pyrenew.observation import ( @@
         MeasurementNoise,
         NegativeBinomialNoise,
     )
+    from pyrenew.time import MMWR_WEEK
     ```
     ## Overview
@@ Expand Down Expand Up / @@ -237,6 +240,42 @@ baseline_rt_process = AR1(autoreg=0.9, innovation_sd=0.05) @@
     subpop_rt_deviation_process = RandomWalk(innovation_sd=0.025)
     ```
+    ### Choosing the Rt Parameter Cadence
+    The renewal equation is evaluated on the model's daily time axis, but the temporal process for $\mathcal{R}(t)$ does not have to sample a new parameter every day.
+    This separates three model choices:
+    -   **Parameter cadence**: how often the $\mathcal{R}(t)$ temporal process samples a new latent value
+    -   **Model time axis**: the daily axis used by the renewal equation and delay convolutions
+    -   **Observation cadence**: the temporal granularity for each signal, such as daily ED visits or weekly hospital admissions
+    The AR(1) process above samples one value per model day:
+    ```python
+    baseline_rt_process = AR1(autoreg=0.9, innovation_sd=0.05)
+    ```
+    To use a  weekly $\mathcal{R}(t)$ while still running the renewal equation daily, wrap the temporal process in `StepwiseTemporalProcess`.
+    The wrapper samples a coarse trajectory and broadcasts it to the daily model axis before the latent infection process uses it.
+    ```python
+    weekly_baseline_rt_process = StepwiseTemporalProcess(
+        inner=AR1(autoreg=0.9, innovation_sd=0.05),
+        step_size=7,
+        alignment="calendar_week",
+        week=MMWR_WEEK,
+    )
+    ```
+    Use `alignment="calendar_week"` when the weekly Rt blocks should align to a calendar week.
+    PyRenew provides class `WeekCycle` which identifies the start day of a calendar week by ISO convention, e.g. `0 == Monday`, `6 == Sunday`
+    and two pre-defined instances: `MMWR_WEEK: WeekCycle = WeekCycle(start_dow=6)` and `ISO_WEEK: WeekCycle = WeekCycle(start_dow=0)`.
+    `pyrenew.time` exports `MMWR_WEEK` (Sunday-Saturday epiweeks) and `ISO_WEEK` (Monday-Sunday); use `WeekCycle(start_dow=k)` for any other convention.
+    At sample or run time, pass the date of the first observation day as `obs_start_date`.
+    Argument `obs_start_date` is required whenever any component performs calendar-aligned work: a latent temporal process with `alignment="calendar_week"`, a count observation with `aggregation="weekly"`, or any observation with a day-of-week effect.
+    For `alignment="model_index"` and daily observations with no day-of-week effect, `obs_start_date` can be omitted.
+    The model handles the calendar bookkeeping and forwards the day-of-week information to every component that needs it.
     ## Observation Processes
     Observation processes transform latent infections into observable signals and define the statistical model linking predictions to data. Each observation process:
@@ Expand Down Expand Up / @@ -268,7 +307,7 @@ for the first 10 months of 2023 (as reported to the CDC). @@
     # | label: load-hospital-data
     # Load daily hospital admissions for California
     ca_hosp_data = datasets.load_hospital_data_for_state("CA", "2023-11-06.csv")
+    obs_start_date = ca_hosp_data["dates"][0]
     hosp_admits = ca_hosp_data["daily_admits"]
     population_size = ca_hosp_data["population"]
     n_hosp_days = ca_hosp_data["n_days"]
@@ Expand Down Expand Up / @@ -658,18 +697,31 @@ model.run( @@
         n_days_post_init=n_days,
         population_size=population_size,
         subpop_fractions=subpop_fractions,
-        hospital={"obs": model.pad_observations(hosp_counts)},
-        wastewater={
-            "obs": ww_conc,
-            "times": model.shift_times(ww_times),
-            "subpop_indices": ww_subpop_indices,
-            "sensor_indices": ww_sensor_indices,
-            "n_sensors": n_ww_sensors,
-        },
+        **obs_data,
     )
     samples = model.mcmc.get_samples()
     ```
+    where `**obs_data` is a data dictionary which supplies the observation start date as `obs_start_date`  and a data dictionary for each set of signal data, where the name of the data dictionary corresponds to the signal name registered on the builder.
+    For this example,  such a data dictionary would have the following structure:
+    ```python
+    {
+        "obs_start_date": ...
+        "hospital": {
+            "obs": ...
+        },
+        "wastewater": {
+            "obs": ...
+            "times": ...
+            "subpop_indices": ...
+            "sensor_indices": ...
+            "n_sensors": ...
+        },
+    }
+    ```
     ## Running the Model
     First we declare the population structure. We have 6 subpopulations, where 5 have wastewater monitoring and 1 does not. The subpopulations with wastewater monitoring need not be contiguous indices—they could be any subset of {0, 1, ..., n_subpops-1}.
@@ Expand Down Expand Up @@
     ```
     We define a function to prepare observation data using the model's helper methods to align with the shared time axis.
+    The returned dictionary is structured to match the keyword arguments of `model.run()`: `obs_start_date` at the top level, plus one sub-dictionary per observation process, keyed by the name registered on the builder (`hospital`, `wastewater`).
+    At call time the returned dict is unpacked with `**` (for example, `**obs_data_90` in the 90-day fit below), forwarding each entry as a keyword argument to `model.run()`.
     ```{python}
     # | label: prepare-observation-data
     def prepare_observation_data(
-        model, n_days_fit, hosp_admits, ww_data, ww_monitored_subpops
+        model,
+        n_days_fit,
+        obs_start_date,
+        hosp_admits,
+        ww_data,
+        ww_monitored_subpops,
     ):
         """
         Prepare observation data for fitting.
@@ Expand All / @@ -717,6 +776,8 @@ def prepare_observation_data( @@
             The model (provides padding/shifting helpers)
         n_days_fit : int
             Number of days to include in fit
+        obs_start_date : date
+            Date of first observation
         hosp_admits : array
             Hospital admissions time series
         ww_data : dict
@@ Expand Down Expand Up / @@ -749,6 +810,7 @@ def prepare_observation_data( @@
         )
         return {
+            "obs_start_date": obs_start_date,
             "hospital": {
                 "obs": hosp_obs,
             },
@@ Expand All / @@ -774,7 +836,12 @@ jax.clear_caches() @@
     n_days_90 = 90
     obs_data_90 = prepare_observation_data(
-        model, n_days_90, hosp_admits, ca_ww_data, ww_monitored_subpops
+        model,
+        n_days_90,
+        obs_start_date,
+        hosp_admits,
+        ca_ww_data,
+        ww_monitored_subpops,
     )
     print(f"Fitting model with {n_days_90} days of data...")
@@ Expand Down Expand Up / @@ -939,7 +1006,12 @@ jax.clear_caches() @@
     n_days_180 = 180
     obs_data_180 = prepare_observation_data(
-        model, n_days_180, hosp_admits, ca_ww_data, ww_monitored_subpops
+        model,
+        n_days_180,
+        obs_start_date,
+        hosp_admits,
+        ca_ww_data,
+        ww_monitored_subpops,
     )
     print(f"Fitting model with {n_days_180} days of data...")
@@ Expand Down @@

docs/tutorials/day_of_week_effects.qmd

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -257,8 +257,8 @@ offset_df["offset"] = pd.Categorical(
  
    The two curves have the same shape but are phase-shifted: their weekend dips fall on different days.

    Getting `first_day_dow` right matters — a misaligned offset would attribute Monday's high to Sunday or vice versa.

    When using `MultiSignalModel`, the shared time axis starts `n_init` days before the first observation.

    The convenience method `model.compute_first_day_dow(obs_start_dow)` converts the known day of the week of the first observation to the correct offset for element 0 of the time axis.

    When using `MultiSignalModel`, pass `obs_start_date` (the date of the first observation day) to `model.sample()` or `model.run()`.

    The model handles the calendar bookkeeping and forwards the day-of-week information to every component that needs it.

    ## Sampled observations

pyproject.toml

-Original file line number
+Diff line change
@@ Expand Up / @@ -77,6 +77,11 @@ known_first_party = ["pyrenew", "test"] @@
     [tool.deptry.per_rule_ignores]
     DEP004 = ["arviz", "pytest", "scipy", "bs4"]
+    [tool.pytest.ini_options]
+    markers = [
+        "integration: integration tests that fit models via MCMC (deselect with '-m \"not integration\"')",
+    ]
     [tool.ruff]
     fix = true
@@ Expand Down @@

pyrenew/arrayutils.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -24,6 +24,28 @@ def __repr__(self) -> str: @@
             return f"PeriodicProcessSample(value={self})"
+    def require_shape(arr: ArrayLike, expected: tuple[int, ...], label: str) -> None:
+        """
+        Validate that ``arr.shape`` equals ``expected``.
+        Parameters
+        ----------
+        arr
+            Array whose shape is being checked.
+        expected
+            Required shape.
+        label
+            Name of the producing component, used in the error message.
+        Raises
+        ------
+        ValueError
+            If ``arr.shape`` does not equal ``expected``.
+        """
+        if arr.shape != expected:
+            raise ValueError(f"{label} must return shape {expected}; got {arr.shape}")
     def tile_until_n(
         data: ArrayLike,
         n_timepoints: int,
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

End-to-end support for coarser-than-daily count signals #794

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!

End-to-end support for coarser-than-daily count signals #794

Are you sure you want to change the base?

End-to-end support for coarser-than-daily count signals #794

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!