More PyRenew Rethinking: what does pyrenew-multisignal HEW model need from PyRenew

[cdc-mitzimorris]

I asked claude.ai to review the HEW model in repo pyrenew-multisignal and it produced this analysis of what the HEW model does and how PyRenew could support it.

Architecture

PyrenewHEWModel.sample()
  |
  +-- LatentInfectionProcess.sample()        -> latent infections (aggregate + per-subpop)
  +-- EDVisitObservationProcess.sample()     -> daily ED visit likelihood
  +-- HospAdmitObservationProcess.sample()   -> weekly hospital admission likelihood
  +-- WastewaterObservationProcess.sample()  -> wastewater concentration likelihood (optional)

Subpopulations exist only when wastewater sites are present (each WW
sampling site defines a subpopulation). With hospital + ED only, it is a
single-population model with no deviation parameters.

Latent Infection Process

Rt parameterization

$\log \mathcal{R}(t)$ follows a differenced AR(1) process at weekly
resolution:

• The rate of change of $\log \mathcal{R}(t)$ is AR(1), not
  $\log \mathcal{R}(t)$ itself. This allows persistent upward/downward
  trends while the growth rate stabilizes.
• Weekly values are repeated 7 times to produce daily Rt.
• Parameters: log_r_mu_intercept (initial value), autoreg_rt
  (AR coefficient), eta_sd (innovation SD).

Initialization

• i0_first_obs_n: fraction of population infected at the first
  observation time.
• Exponential growth initialization using r_approx_from_R() applied to
  the initial Rt value, run for n_initialization_points days.

Infection feedback

Uses InfectionsWithFeedback from PyRenew, which modulates Rt based on
recent infection levels (typically a small negative feedback).

Subpopulation structure (when WW is present)

• Subpop 0 is the reference; it uses the main Rt trajectory.
• Non-reference subpops get AR(1) deviations added to the reference
  Rt (no sum-to-zero constraint).
• Initial infection fractions for non-reference subpops are offset from
  the reference via a logit-scale parameter.

ED Visit Observation Process

Pipeline: infections $\to$ delay convolution $\to$ ascertainment $\to$
day-of-week $\to$ right truncation $\to$ NegBin likelihood.

Ascertainment (IEDR)

The infection-to-ED-visit ratio is time-varying:

• Weekly AR(1) process on the logit scale, centered at p_ed_mean.
• Sigmoid transform produces daily IEDR values in (0, 1).
• Weekly values are repeated to daily.

Day-of-week effects

A 7-element multiplicative vector (ed_wday_effect_rv) tiled across the
time series. Applied after ascertainment.

Right truncation

Accounts for reporting delays: recent days have incomplete data. A
right-truncation PMF is converted to a CDF representing the proportion
of cases already reported at each lag. Applied multiplicatively.

Noise

Negative binomial with concentration parameter ed_neg_bin_concentration.

Hospital Admission Observation Process

Pipeline: infections $\to$ delay convolution $\to$ IHR $\to$
MMWR epiweekly aggregation $\to$ NegBin likelihood.

IHR

Two modes depending on whether ED visits are also being fit:

• H+E mode (ihr_rel_iedr_rv): IHR is defined as a ratio relative
  to the IEDR. Specifically, ihr = iedr[0] * ihr_rel_iedr. This
  creates a principled link between the two signals.
• H-only mode (ihr_rv): IHR is an independent parameter.

Delay distribution

When both H and E are fit, the infection-to-hospitalization delay is an
OffsetDiscretizedLognormalPMF with inferred offsets to a reference
lognormal. This allows the model to learn the delay difference between
the two signals.

Weekly aggregation

Daily predicted admissions are aggregated to MMWR epiweek totals via
daily_to_mmwr_epiweekly(). The calculate_weekly_hosp_indices() static
method handles alignment:

1. Find the first complete epiweek (ending Saturday) after the model
   start.
2. Validate that all observed week-ending dates are Saturdays.
3. Compute indices mapping observed weeks to positions in the aggregated
   predictions array.

Noise

Negative binomial with concentration parameter
hosp_admit_neg_bin_concentration.

Initialization Points

n_initialization_points = max(
    generation_interval_pmf.size(),
    infection_feedback_pmf.size(),
    inf_to_ed_pmf.size()          if fitting ED,
    inf_to_hosp_pmf.size() + 6    if fitting hospital,
    max_shed_interval              if fitting wastewater,
) - 1

The + 6 for hospital accounts for the weekly aggregation window.

Key Differences from PyRenew's HierarchicalInfections

Aspect	HierarchicalInfections	pyrenew-multisignal hew
Subpop deviations	Sum-to-zero constraint; always sampled	Reference + offsets; only when WW present
1-subpop behavior	Degenerate (144 nuisance parameters)	Clean single-population model
Rt process	AR(1) on log Rt	Differenced AR(1) on log Rt (weekly)
Rt resolution	Daily	Weekly, repeated to daily
IHR/IEDR link	Independent	IHR = IEDR * ratio (when fitting both)
Infection feedback	Not supported	Supported via InfectionsWithFeedback
Builder	PyrenewBuilder	Custom build_pyrenew_hew_model()

Implications for PyRenew

To support H+E models in PyRenew's PyrenewBuilder framework, the main
gaps are:

1. Single-population latent process: A BaseLatentInfectionProcess
   subclass that uses a temporal process for Rt without subpopulation
   deviations.
2. Linked ascertainment: A mechanism for observation processes to
   share parameters (e.g., IHR defined relative to IEDR).
3. Weekly Rt: Support for temporal processes at coarser-than-daily
   resolution, expanded to daily for the renewal equation.