Add Ascertainment model for JointAscertainment

[cdc-mitzimorris]

Summary

This PR adds model-level ascertainment support so multiple observation signals can share structured ascertainment rates instead of each observation independently sampling its own rate.

The main motivation is to support H+E-style multisignal models where hospital admissions and ED visits have distinct but related observation probabilities. The new abstraction lets those probabilities be sampled once at the model level, then passed into observation processes through signal-specific accessors without creating duplicate NumPyro sample sites.

Additions / Changes

• Added a new pyrenew.ascertainment package with:

  • AscertainmentModel, the base interface for shared ascertainment models.
  • AscertainmentSignal, a signal-specific accessor used by observation processes.
  • JointAscertainment, which samples correlated scalar ascertainment rates across signals on the logit scale.
  • A context layer that makes sampled ascertainment values available to observation processes during model execution.

• Extended PyrenewBuilder with add_ascertainment(...), allowing shared ascertainment models to be registered alongside latent and observation components.

• Extended MultiSignalModel so registered ascertainment models are sampled once before observation processes run. Observation processes then retrieve their signal-specific values from the active ascertainment context.

• Added unit coverage for validation, sampling behavior, context safety, duplicate-site avoidance, builder registration, and direct MultiSignalModel validation.

• Added integration coverage for mixed-cadence hospital + ED models:

  • weekly hospital admissions plus daily ED visits with JointAscertainment

Suggested Review Order

1. Start with pyrenew/ascertainment/base.py and pyrenew/ascertainment/context.py to understand the core abstraction and how accessors retrieve already-sampled values.
2. Review pyrenew/ascertainment/joint.py
3. Review the integration points in pyrenew/model/pyrenew_builder.py and pyrenew/model/multisignal_model.py.
4. Review tutorial docs/tutorials/ascertainment.py
5. Review test/test_ascertainment.py and the new builder tests in test/test_pyrenew_builder.py.
6. Finish with the integration fixtures and tests in test/integration/conftest.py and the two new H+E integration test files, since those show the intended end-to-end usage.

[dylanhmorris]

Thanks @cdc-mitzimorris. Before I review further, could you explain the decision to use a context layer here versus any alternative approaches you considered?

[cdc-mitzimorris]

Thanks @cdc-mitzimorris. Before I review further, could you explain the decision to use a context layer here versus any alternative approaches you considered?

The main design constraint was that joint ascertainment needs to create one shared NumPyro sample site, while existing observation processes expect signal-specific RandomVariable-like objects.

I considered passing the sampled mapping explicitly through the model stack, but that would require broader plumbing/API changes. I also wanted to avoid having each signal accessor sample independently, since that would either duplicate sites or lose the intended joint dependence. Storing sampled values directly on the model object also seemed fragile for NumPyro execution, since model calls/traces should not depend on mutable object state.

The context layer is intended to keep the public observation API unchanged: the joint model samples once, then the signal accessors retrieve their already-sampled values from execution-scoped context without creating additional NumPyro sites. This keeps ownership of the shared sample site centralized in the ascertainment model, while keeping observation processes focused on consuming signal-specific rates.

[dylanhmorris]

A few minor suggestions. Happy for you to merge without re-review from me if @damonbayer is satisfied that his requested changes have been made. Thanks, @cdc-mitzimorris!

[damonbayer]

I would like to review the tutorial, once it can be rendered.

[cdc-mitzimorris]

@damonbayer - tutorial renders.

[damonbayer]

One recommendation

[damonbayer]

Expect panache to reformat the tutorial doc.