Add "RatioLinkedAscertainment" component

[cdc-mitzimorris]

Add a new class "RatioLinkedAscertainment" which relates the ascertainment rate of two observation processes where the ascertainment rate for one process is the base rate and the ascertainment rate for the other is some ratio of the two signals.

This follows pyrenew-hew which parameterizes the ascertainment for the ED visits and hospitalizations as a base rate (IEDR) and a multiplier (IHR / IEDR).

This can be implemented straightforwardly using the base Ascertainment component:

"""Linked ascertainment components."""

from __future__ import annotations

from collections.abc import Mapping

import jax
import numpyro
from pyrenew.ascertainment import AscertainmentModel
from pyrenew.metaclass import RandomVariable


class RatioLinkedAscertainment(AscertainmentModel):
    """Two ascertainment rates expressed as a base rate and a ratio."""

    def __init__(
        self,
        name: str,
        base_signal: str,
        linked_signal: str,
        base_rate_rv: RandomVariable,
        ratio_rv: RandomVariable,
    ) -> None:
        """Create a linked ascertainment model for two named signals."""
        super().__init__(name=name, signals=(base_signal, linked_signal))
        self.base_signal = base_signal
        self.linked_signal = linked_signal
        self.base_rate_rv = base_rate_rv
        self.ratio_rv = ratio_rv

    def sample(self, **kwargs: object) -> Mapping[str, jax.Array]:
        """Sample the base rate and ratio, then compute the linked rate."""
        base_rate = self.base_rate_rv()
        ratio = self.ratio_rv()
        linked_rate = base_rate * ratio
        numpyro.deterministic(f"{self.name}_{self.base_signal}", base_rate)
        numpyro.deterministic(f"{self.name}_{self.linked_signal}", linked_rate)
        return {self.base_signal: base_rate, self.linked_signal: linked_rate}

The corresponding Python code to specify a model using the PyrenewBuilder class is:

iedr_rv = DistributionalVariable(name="iedr",  dist.LogNormal(0.0, jnp.log(jnp.sqrt(2.0))))

ratio_rv = DistributionalVariable(name="ihr_rel_iedr",  dist.LogNormal(0.0, jnp.log(jnp.sqrt(2.0))))

ascertainment = RatioLinkedAscertainment(
    name="he_ascertainment",
    base_signal="ed_visits",
    linked_signal="hospital",
    base_rate_rv=iedr_rv,
    ratio_rv=ratio_rv,
)

hospital_obs = PopulationCounts(
    name="hospital",
    ascertainment_rate_rv=ascertainment.for_signal("hospital"),
    ...
)

ed_obs = PopulationCounts(
    name="ed_visits",
    ascertainment_rate_rv=ascertainment.for_signal("ed_visits"),
    ...
)