Skip to content

Conversation

@cdc-mitzimorris
Copy link
Collaborator

@cdc-mitzimorris cdc-mitzimorris commented Jan 15, 2026

Adds infrastructure for modeling infections across multiple subpopulations with hierarchical Rt structure, plus a builder pattern for composing multi-signal renewal models.

Key additions:

  • HierarchicalInfections: Latent process with shared baseline Rt and subpopulation-specific deviations
  • TemporalProcess protocol with AR1, DifferencedAR1, RandomWalk implementations
  • ModelBuilder + MultiSignalModel for composing latent + observation processes
  • Hierarchical priors (HierarchicalNormalPrior, GammaGroupSdPrior)
  • Sample datasets (NHSN hospital admissions, fake NWSS wastewater)
  • Two tutorials demonstrating the API

Architecture

Latent layer:

  • TemporalProcess (protocol)
    • AR1, DifferencedAR1, RandomWalk (implementations)
    • Used by HierarchicalInfections for baseline Rt and subpopulation deviations
  • HierarchicalNormalPrior, GammaGroupSdPrior
    • Used by HierarchicalInfections for I0 and other group-level effects
  • BaseLatentInfectionProcess → HierarchicalInfections
    • Outputs LatentSample: (aggregate, all_subpops, observed, unobserved)

Model layer:

  • ModelBuilder
    • Configures latent process + observation processes
    • Computes n_initialization_points from all delay PMFs
    • Builds MultiSignalModel
  • MultiSignalModel
    • Routes infections to observations based on infection_resolution()
    • "aggregate" → summed infections
    • "subpop" → per-subpopulation infections

Suggested Review Order

  1. latent/temporal_processes.py — TemporalProcess protocol and implementations (AR1, DifferencedAR1, RandomWalk). Module docstring explains relationship to pyrenew.process.
  2. latent/base.py — BaseLatentInfectionProcess base class defining the LatentSample output structure
  3. latent/hierarchical_infections.py — HierarchicalInfections implementation
  4. model/multisignal_model.py → model/model_builder.py — Composition layer
  5. Tutorials — Demonstrate the API in action:
    - latent_hierarchical_infections.qmd: temporal process choices
    - building_multisignal_models.qmd: full inference workflow
  6. Tests and datasets — Can skim; mostly mechanical

Review Scope

Feedback requested on:

  • API design: Are the interfaces intuitive? Would you change any method signatures?
  • Correctness: Does the hierarchical Rt decomposition and sum-to-zero constraint make epidemiological sense?

Deferred to follow-up PRs:


Known Limitations / Future Work

  • TemporalProcess implementations don't yet support periodic effects (available in pyrenew.process)
  • Tested with synthetic data; real-data validation ongoing
  • ModelBuilder API may evolve based on user feedback

cdc-mitzimorris and others added 30 commits September 15, 2025 18:24
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Copy link
Collaborator

@dylanhmorris dylanhmorris left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @cdc-mitzimorris. Have now reviewed everything except the tutorials and tests. Flagged a bunch of small things and tried to be clear where changes are optional. My main big questions (open to pushback on both):

  • Not sure specific distributional priors are needed / in scope. Currently prefer to show the user how to create a needed prior via VectorizedRV in the tutorials.
  • TemporalProcess RandomWalk should wrap the process module's random walk rather than reimplementing.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • Suggest log y axis for all infection and admission plots
  • Why not show the predicted latent (and/or posterior predictive) admissions on top of the observed?
  • Would be nice for completeness to have a wastewater predicted/observed plot as well.

Copy link
Collaborator

@dylanhmorris dylanhmorris left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have now reviewed everything but tests. @cdc-mitzimorris I propose you make and/or push back all proposed changes and then I will re-review. I'd like to wait to review tests since some changes could require them to be reworked or all some to be eliminated (e.g. the hierarchical priors).

cdc-mitzimorris and others added 2 commits February 5, 2026 23:34
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
cdc-mitzimorris and others added 7 commits February 5, 2026 23:53
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
…n, remove fit(), remove sample() return

- Rename baseline_temporal -> baseline_rt_process, subpop_temporal -> subpop_rt_deviation_process
- Refactor RandomWalk temporal process to wrap pyrenew.process.RandomWalk
- Replace custom PMF validation with validate_discrete_dist_vector
- Remove MultiSignalModel.fit() in favor of Model.run()
- Remove return value from MultiSignalModel.sample() (use Predictive/deterministic sites)
- Remove unreachable lookback_days defensive check (enforced by ABC)
- Update tests and tutorials

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
n_processes=pop.K,
initial_value=jnp.zeros(pop.K),
n_processes=pop.n_subpops,
initial_value=jnp.zeros(pop.n_subpops),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To fix (not in this PR): initial deviations should be confiurable, not fixed.

Copy link
Collaborator

@dylanhmorris dylanhmorris left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, @cdc-mitzimorris!

@dylanhmorris dylanhmorris merged commit 265b2bf into main Feb 6, 2026
8 checks passed
@dylanhmorris dylanhmorris deleted the mem_hier_latent_processes branch February 6, 2026 17:14
@cdc-mitzimorris
Copy link
Collaborator Author

thank you @dylanhmorris for your careful and patient review!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants