diff --git a/README.md b/README.md
index 924dccac..25f2a897 100755
--- a/README.md
+++ b/README.md
@@ -20,9 +20,10 @@ where $\alpha$ is the ascertainment rate and $\pi(s)$ is the delay distribution
The Pyrenew package provides configurable classes which encapsulate these components and methods to orchestrate the configuration and composition of these processes
resulting in programs which clearly express the model structure and choices, allowing for both ease of model specification and dissemination.
-The fundamental building blocks are the `Model` metaclass, from which we can draw samples,
-and the `RandomVariable` metaclass which has been abstracted to allow for sampling from distributions, computing a mechanistic equation, or simply returning a fixed value.
-The `PyrenewBuilder` class orchestrates the composition process.
+The fundamental building block is the `RandomVariable` abstract base class, which allows for sampling from distributions, computing a mechanistic equation, or simply returning a fixed value.
+
+We use `RandomVariable`s to build probabilistic models. We represent complete models as concrete subclasses of the `Model` abstract base class.
+The `PyrenewBuilder` class orchestrates the composition of `RandomVariables` into a `Model`.
PyRenew's strength lies in multi-signal integration for information pooling across diverse observed data streams
such as hospital admissions, wastewater concentrations, and emergency department visits
diff --git a/docs/index.md b/docs/index.md
index 612c7a5e..a202efe8 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1 +1,69 @@
---8<-- "README.md"
+# PyRenew
+
+**Bayesian renewal modeling with JAX and NumPyro.**
+
+PyRenew is a flexible tool for simulation and statistical inference of epidemiological models, emphasizing hierarchical multi-signal renewal models.
+Built on top of [NumPyro](https://num.pyro.ai/), PyRenew provides configurable classes that encapsulate the components of a renewal model, and methods to orchestrate their composition into programs that clearly express model structure and choices.
+
+## Renewal models
+
+A renewal model estimates new infections from recent past infections using a generation interval distribution $w(s)$: the probability that $s$ time units separate infection in an index case and a secondary case.
+The core renewal equation is:
+
+$$I(t) = R_t \sum_{s} I(t-s) \, w(s)$$
+
+where $R_t$ is the time-varying reproduction number.
+
+Inference is complicated by the fact that observational data require their own models ([Bhatt et al., 2023, S2](https://doi.org/10.1093/jrsssa/qnad030)).
+The observation equation links infections to expected observations:
+
+$$\mu(t) = \alpha \sum_{s} I(t-s) \, \pi(s)$$
+
+where $\alpha$ is the ascertainment rate and $\pi(s)$ is the delay distribution from infection to observation.
+
+## Design
+
+PyRenew's building blocks are:
+
+- **`RandomVariable`** -- an abstract base class for quantities that models can sample from, including stochastic draws via `numpyro.sample()`, mechanistic computations, and fixed values.
+- **`Model`** -- an abstract base class that defines model structure through a `sample()` method and provides functionality for fitting and simulation.
+- **`PyrenewBuilder`** -- an orchestrator that composes `RandomVariable` components into a complete renewal model, auto-computing derived quantities like the number of initialization points.
+
+Components (generation interval, reproduction number process, observation process) are specified independently, so each can be swapped without changing the rest of the model.
+This makes it straightforward to move a quantity between "known" and "inferred" and keeps modeling choices explicit and reviewable.
+
+
+## Multi-signal models
+
+PyRenew's strength lies in multi-signal integration: pooling information across diverse observed data streams such as hospital admissions, wastewater concentrations, and emergency department visits, where each signal has distinct observation delays, noise characteristics, and spatial resolutions.
+For single-signal renewal models, we recommend the excellent R package [EpiNow2](https://epiforecasts.io/EpiNow2/).
+
+## Installation
+
+```bash
+pip install git+https://github.com/CDCgov/PyRenew@main
+```
+
+## Tutorials
+
+- [The RandomVariable abstract base class](tutorials/random_variables.md) -- PyRenew's core abstraction and its concrete implementations.
+- [Building multi-signal models](tutorials/building_multisignal_models.md) -- composing a renewal model from PyRenew components using `PyrenewBuilder`.
+- [Observation processes: count data](tutorials/observation_processes_counts.md) -- connecting latent infections to observed counts.
+- [Observation processes: measurements](tutorials/observation_processes_measurements.md) -- connecting latent infections to continuous measurements.
+- [Latent hierarchical infections](tutorials/latent_hierarchical_infections.md) -- modeling infections with subpopulation structure.
+
+## Resources
+
+- [Model equations](https://github.com/CDCgov/PyRenew/blob/main/equations.md) -- the mathematics of the multi-signal renewal processes PyRenew supports.
+- [API reference](reference/index.md)
+- [Developer documentation](developer_documentation.md)
+
+### Models implemented with PyRenew
+
+- [pyrenew-covid-wastewater](https://github.com/CDCgov/pyrenew-covid-wastewater) -- forecasting COVID-19 hospitalizations using wastewater data.
+- [pyrenew-flu-light](https://github.com/CDCgov/pyrenew-flu-light/) -- an influenza forecasting model used in the 2023-24 respiratory season.
+
+### Further reading
+
+- [Semi-mechanistic Bayesian modelling of COVID-19 with renewal processes](https://academic.oup.com/jrsssa/article-pdf/186/4/601/54770289/qnad030.pdf) (Bhatt et al., 2023)
+- [Unifying incidence and prevalence under a time-varying general branching process](https://link.springer.com/content/pdf/10.1007/s00285-023-01958-w.pdf)
diff --git a/docs/tutorials/.pages b/docs/tutorials/.pages
index 81b29109..163f0e05 100644
--- a/docs/tutorials/.pages
+++ b/docs/tutorials/.pages
@@ -1,11 +1,7 @@
nav:
- - getting_started.md
+ - random_variables.md
- building_multisignal_models.md
- observation_processes_counts.md
- observation_processes_measurements.md
- latent_hierarchical_infections.md
- - custom_randomvariables.md
- - basic_renewal_model.md
- - hospital_admissions_model.md
- - day_of_the_week.md
- periodic_effects.md
diff --git a/docs/tutorials/basic_renewal_model.qmd b/docs/tutorials/basic_renewal_model.qmd
deleted file mode 100644
index 8070259c..00000000
--- a/docs/tutorials/basic_renewal_model.qmd
+++ /dev/null
@@ -1,390 +0,0 @@
----
-title: Fitting a basic renewal model
-format: gfm
-engine: jupyter
----
-
-This tutorial shows the steps to build a simple renewal model featuring a latent infection process, a random walk $\mathcal{R}(t)$ process, and an observation process for the reported infections.
-
-We start by loading the needed components to build a basic renewal model:
-
-```{python}
-# | label: loading-pkgs
-# | output: false
-# | warning: false
-import jax.numpy as jnp
-import numpy as np
-import numpyro
-import numpyro.distributions as dist
-from pyrenew.process import RandomWalk
-from pyrenew.latent import (
- Infections,
- InfectionInitializationProcess,
- InitializeInfectionsZeroPad,
-)
-from pyrenew.observation import PoissonObservation
-from pyrenew.deterministic import DeterministicPMF
-from pyrenew.model import RtInfectionsRenewalModel
-from pyrenew.metaclass import RandomVariable
-from pyrenew.randomvariable import DistributionalVariable, TransformedVariable
-import pyrenew.transformation as t
-from numpyro.infer.reparam import LocScaleReparam
-```
-By default, XLA (which is used by JAX for compilation) considers all CPU cores as one device. Depending on your system's configuration, we recommend using numpyro's [set_host_device_count()](https://num.pyro.ai/en/stable/utilities.html#set-host-device-count) function to set the number of devices available for parallel computing. Here, we set the device count to be 2.
-```{python}
-# | label: set-device-count
-numpyro.set_host_device_count(2)
-```
-
-## `pyrenew` core metaclasses
-
-Most objects in `pyrenew` can be treated as random variables we can sample. At the top-level `pyrenew` has two metaclasses from which most objects inherit: `RandomVariable` and `Model`.
-These, in turn, provide the following sub-modules.
-
-- `deterministic` - `RandomVariable` classes for fixed transformations and non-random parameters
-- `latent` - `RandomVariable` classes for hidden states requiring statistical inference
-- `observation` - `RandomVariable` classes for observed data and likelihoods
-- `process` - `RandomVariable` classes for stochastic processes for disease dynamics
-- `models` - `Model` classes for specific epidemiological models and workflows
-
-## Architecture of `RtInfectionsRenewalModel`
-
-The class `RtInfectionsRenewalModel` is basic renewal model consisting of infections and reproduction numbers.
-It models the real-time reproductive number $\mathcal{R}(t)$, the average number of secondary infections caused by an infected individual, as a renewal process model
-defined in terms of:
-
-1. generation interval, the times between infections
-
-2. initial infections, occurring prior to time $t = 0$
-
-3. $\mathcal{R}(t)$, the time-varying reproductive number,
-
-4. latent infections, i.e., those infections which are known to exist but are not observed (or not observable), and
-
-5. observed infections, a subset of underlying true infections that are reported, perhaps via hospital admissions, physician's office visits, or routine biosurveillance.
-
-The following diagram shows a detailed view of how metaclasses, modules, and classes are used to in the `RtInfectionsRenewalModel`.
-```mermaid
-flowchart LR
-
- rand("(RandomVariable metaclass)")
- models("(Model metaclass)")
-
- subgraph observations["Observations module"]
- obs["infection_obs_process_rv
(PoissonObservation)"]
- end
-
- subgraph latent["Latent module"]
- inf["latent_infections_rv
(Infections)"]
- i0["I0_rv
(DistributionalVariable)"]
- end
-
- subgraph process["Process module"]
- rt["Rt_process_rv
(Custom class built using RandomWalk)"]
- end
-
- subgraph deterministic["Deterministic module"]
- detpmf["gen_int_rv
(DeterministicPMF)"]
- end
-
- subgraph model[Model module]
- model1["model1
(RtInfectionsRenewalModel)"]
- end
-
- rand-->|Inherited by|observations
- rand-->|Inherited by|process
- rand-->|Inherited by|latent
- rand-->|Inherited by|deterministic
- models-->|Inherited by|model
-
- detpmf-->|Composes|model1
- i0-->|Composes|model1
- rt-->|Composes|model1
- obs-->|Composes|model1
- inf-->|Composes|model1
-```
-
-## Implementing a `RtInfectionsRenewalModel`
-
-In this example we specify a renewal model with the following components.
-
-1. The generation interval is provided as a deterministic instance of `RandomVariable`
-
-2. The `InfectionInitializationProcess` specifies the number of latent infections immediately before the renewal process begins as following a log-normal distribution with mean = 0 and standard deviation = 1. By specifying `InitializeInfectionsZeroPad`, the latent infections before this time are assumed to be 0.
-
-3. A process to represent $\mathcal{R}(t)$ as a random walk on the log scale, with an inferred initial value and a fixed Normal step-size distribution. For this, we construct a custom `RandomVariable`, `RtRandomWalk`.
-
-4. An instance of the `Infections` class with default values.
-
-5. An instance of the `PoissonObservation` class with default values
-
-```{python}
-# | label: creating-elements
-# (1) The generation interval (deterministic)
-pmf_array = jnp.array([0.4, 0.3, 0.2, 0.1])
-gen_int = DeterministicPMF(name="gen_int", value=pmf_array)
-
-# (2) Initial infections (inferred with a prior)
-I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(2.5, 1)),
- InitializeInfectionsZeroPad(pmf_array.size),
-)
-
-
-# (3) The random walk process.
-# We construct a custom RandomVariable on log Rt, with an inferred s.d.
-class RtRandomWalk(RandomVariable):
- def validate(self):
- pass
-
- def sample(self, n: int, **kwargs) -> tuple:
- sd_rt = numpyro.sample("Rt_random_walk_sd", dist.HalfNormal(0.025))
-
- rt_rv = TransformedVariable(
- name="log_rt_random_walk",
- base_rv=RandomWalk(
- name="log_rt",
- step_rv=DistributionalVariable(
- name="rw_step_rv", distribution=dist.Normal(0, sd_rt)
- ),
- ),
- transforms=t.ExpTransform(),
- )
- rt_init_rv = DistributionalVariable(
- name="init_log_rt", distribution=dist.Normal(0, 0.2)
- )
- init_rt = rt_init_rv.sample()
-
- return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
-
-
-rt_proc = RtRandomWalk()
-
-# (4) Latent infection process (which will use 1 and 2)
-latent_infections = Infections()
-
-# (5) The observed infections process (with mean at the latent infections)
-observation_process = PoissonObservation("poisson_rv")
-```
-
-With these five pieces, we can build the basic renewal model as an instance of the `RtInfectionsRenewalModel` class:
-
-```{python}
-# | label: model-creation
-model1 = RtInfectionsRenewalModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- Rt_process_rv=rt_proc,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observation_process,
-)
-```
-
-Using `numpyro`, we can simulate data using the `sample()` member function of `RtInfectionsRenewalModel`. The `sample()` method of the `RtInfectionsRenewalModel` class returns a list composed of the `Rt` and `infections` sequences, called `sim_data`:
-
-```{python}
-# | label: simulate
-with numpyro.handlers.seed(rng_seed=53):
- sim_data = model1.sample(n_datapoints=40)
-
-sim_data
-```
-
-To understand what has been accomplished here, visualize an $\mathcal{R}(t)$ sample path (left panel) and infections over time (right panel):
-
-```{python}
-# | label: fig-basic
-# | fig-cap: Rt and Infections
-import matplotlib.pyplot as plt
-
-fig, axs = plt.subplots(1, 2)
-
-# Rt plot
-axs[0].plot(sim_data.Rt)
-axs[0].set_ylabel("Rt")
-
-# Infections plot
-axs[1].plot(sim_data.observed_infections)
-axs[1].set_ylabel("Infections")
-
-fig.suptitle("Basic renewal model")
-fig.supxlabel("Time")
-plt.tight_layout()
-plt.show()
-```
-
-To fit the model, we can use the `run()` method of the `RtInfectionsRenewalModel` class (an inherited method from the metaclass `Model`). `model1.run()` will call the `run` method of the `model1` object, which will generate an instance of model MCMC simulation, with 2000 warm-up iterations for the MCMC algorithm, used to tune the parameters of the MCMC algorithm to improve efficiency of the sampling process. From the posterior distribution of the model parameters, 1000 samples will be drawn and used to estimate the posterior distributions and compute summary statistics. Observed data is provided to the model using the `sim_data` object previously generated. `mcmc_args` provides additional arguments for the MCMC algorithm.
-
-```{python}
-# | label: model-fit
-import jax
-
-model1.run(
- num_warmup=2000,
- num_samples=1000,
- data_observed_infections=sim_data.observed_infections,
- rng_key=jax.random.key(54),
- mcmc_args=dict(progress_bar=False, num_chains=2),
-)
-```
-
-Now, let's investigate the output, particularly the posterior distribution of the $\mathcal{R}(t)$ estimates:
-
-```{python}
-# | label: fig-output-rt
-# | fig-cap: Rt posterior distribution
-import arviz as az
-
-# Create arviz inference data object
-idata = az.from_numpyro(
- posterior=model1.mcmc,
-)
-
-# Extract Rt signal samples across chains
-rt = az.extract(idata.posterior["Rt"], num_samples=100)["Rt"].values
-
-
-# Plot Rt signal
-fig, ax = plt.subplots(1, 1, figsize=(8, 6))
-ax.plot(
- np.arange(rt.shape[0]),
- rt,
- color="skyblue",
- alpha=0.10,
-)
-ax.plot([], [], color="skyblue", alpha=0.05, label="Rt Posterior Samples")
-ax.plot(
- np.arange(rt.shape[0]),
- rt.mean(axis=1),
- color="black",
- linewidth=2.0,
- linestyle="--",
- label="Sample Mean",
-)
-ax.legend(loc="best")
-ax.set_ylabel(r"$\mathscr{R}_t$ Signal", fontsize=20)
-ax.set_xlabel("Days", fontsize=20)
-plt.show()
-```
-We can use the `get_samples` method to extract samples from the model
-```{python}
-# | label: get-model-fitted-samples
-Rt_samp = model1.mcmc.get_samples()["Rt"]
-latent_infection_samp = model1.mcmc.get_samples()["all_latent_infections"]
-```
-
-We can also convert the fitted model to [ArviZ](https://www.arviz.org/) InferenceData object and use ArviZ package to extarct samples, calculate statistics, create model diagnostics and visualizations.
-
-```{python}
-# | label: convert-inference-data
-import arviz as az
-
-idata = az.from_numpyro(model1.mcmc)
-```
-
-and use the InferenceData to compute the model-fit diagnostics. Here, we show diagnostic summary for the first 10 effective reproduction number $\mathcal{R}(t)$.
-
-```{python}
-# | label: diagnostics
-diagnostic_stats_summary = az.summary(
- idata.posterior["Rt"][::, ::, 4:], # ignore nan padding
- kind="diagnostics",
-)
-
-print(diagnostic_stats_summary)
-```
-
-Below we use `plot_trace` to inspect the trace of the first 10 inferred $\mathcal{R}(t)$ values.
-
-```{python}
-# | label: fig-trace-Rt
-# | fig-cap: Trace plot of Rt posterior distribution
-plt.rcParams["figure.constrained_layout.use"] = True
-
-az.plot_trace(
- idata.posterior,
- var_names=["Rt"],
- coords={"Rt_dim_0": np.arange(4, 14)},
- compact=False,
-)
-plt.show()
-```
-
-
-We inspect the posterior distribution of $\mathcal{R}(t)$ by plotting the 90% and 50% highest density intervals:
-
-```{python}
-# | label: fig-hdi-Rt
-# | fig-cap: High density interval for Effective Reproduction Number
-x_data = idata.posterior["Rt_dim_0"][4:]
-y_data = idata.posterior["Rt"][::, ::, 4:]
-
-fig, axes = plt.subplots(figsize=(6, 5))
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.9,
- color="C0",
- fill_kwargs={"alpha": 0.3},
- ax=axes,
-)
-
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.5,
- color="C0",
- fill_kwargs={"alpha": 0.6},
- ax=axes,
-)
-
-# Add mean of the posterior to the figure
-median_ts = y_data.median(dim=["chain", "draw"])
-axes.plot(x_data, median_ts, color="C0", label="Median")
-axes.legend()
-axes.set_title("Posterior Effective Reproduction Number", fontsize=10)
-axes.set_xlabel("Time", fontsize=10)
-axes.set_ylabel("$\\mathcal{R}(t)$", fontsize=10)
-plt.show()
-```
-
-and latent infections:
-
-```{python}
-# | label: fig-hdi-latent-infections
-# | fig-cap: High density interval for Latent Infections
-x_data = idata.posterior["all_latent_infections_dim_0"]
-y_data = idata.posterior["all_latent_infections"]
-
-fig, axes = plt.subplots(figsize=(6, 5))
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.9,
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.3},
- ax=axes,
-)
-
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.5,
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.6},
- ax=axes,
-)
-
-# plot the posterior median
-median_ts = y_data.median(dim=["chain", "draw"])
-axes.plot(x_data, median_ts, color="C0", label="Median")
-
-axes.legend()
-axes.set_title("Posterior Latent Infections", fontsize=10)
-axes.set_xlabel("Time", fontsize=10)
-axes.set_ylabel("Latent Infections", fontsize=10)
-plt.show()
-```
diff --git a/docs/tutorials/building_multisignal_models.qmd b/docs/tutorials/building_multisignal_models.qmd
index 6031e5de..fdf9d544 100644
--- a/docs/tutorials/building_multisignal_models.qmd
+++ b/docs/tutorials/building_multisignal_models.qmd
@@ -106,9 +106,9 @@ The builder computes initialization requirements, wires components together, and
Before diving into multi-signal models, you may want to review these foundational tutorials:
-- **[Hierarchical Latent Infections](latent_hierarchical_infections.qmd)**: Understanding temporal process choices for $\mathcal{R}(t)$
-- **[Observation Processes: Counts](observation_processes_counts.qmd)**: Modeling count data (admissions, deaths)
-- **[Observation Processes: Measurements](observation_processes_measurements.qmd)**: Modeling continuous data (wastewater)
+- **[Hierarchical Latent Infections](latent_hierarchical_infections.md)**: Understanding temporal process choices for $\mathcal{R}(t)$
+- **[Observation Processes: Counts](observation_processes_counts.md)**: Modeling count data (admissions, deaths)
+- **[Observation Processes: Measurements](observation_processes_measurements.md)**: Modeling continuous data (wastewater)
This tutorial shows how to combine these components into a complete multi-signal model.
@@ -359,7 +359,7 @@ print(f" Delay PMF length: {len(inf_to_hosp_pmf)} days")
The `Measurements` base class handles continuous observation processes. Domain-specific
implementations subclass it and implement `_predicted_obs()` to transform infections
-into predicted values. See `observation_processes_measurements.qmd` for a detailed tutorial.
+into predicted values. See `observation_processes_measurements.md` for a detailed tutorial.
```{python}
# | label: wastewater-class
@@ -1119,6 +1119,6 @@ This tutorial demonstrated composing a multi-signal renewal model using `Pyrenew
### Next Steps
-- Explore different temporal processes for $\mathcal{R}(t)$ in the [Hierarchical Latent Infections](latent_hierarchical_infections.qmd) tutorial
-- Learn about count-based observation models in [Observation Processes: Counts](observation_processes_counts.qmd)
-- Learn about continuous measurement models in [Observation Processes: Measurements](observation_processes_measurements.qmd)
+- Explore different temporal processes for $\mathcal{R}(t)$ in the [Hierarchical Latent Infections](latent_hierarchical_infections.md) tutorial
+- Learn about count-based observation models in [Observation Processes: Counts](observation_processes_counts.md)
+- Learn about continuous measurement models in [Observation Processes: Measurements](observation_processes_measurements.md)
diff --git a/docs/tutorials/custom_randomvariables.qmd b/docs/tutorials/custom_randomvariables.qmd
deleted file mode 100644
index 53ece1ee..00000000
--- a/docs/tutorials/custom_randomvariables.qmd
+++ /dev/null
@@ -1,320 +0,0 @@
----
-title: Defining Custom PyRenew RandomVariables
-format: gfm
-engine: jupyter
-jupyter:
- jupytext:
- text_representation:
- extension: .qmd
- format_name: quarto
- format_version: '1.0'
- jupytext_version: 1.18.1
- kernelspec:
- display_name: Python 3 (ipykernel)
- language: python
- name: python3
----
-
-```{python}
-# | label: setup
-import jax
-import jax.numpy as jnp
-import numpy as np
-import numpyro
-import numpyro.distributions as dist
-from pyrenew.deterministic import DeterministicPMF, DeterministicVariable
-from pyrenew.latent import InfectionsWithFeedback
-from pyrenew.model import RtInfectionsRenewalModel
-from pyrenew.process import RandomWalk
-from pyrenew.metaclass import RandomVariable
-from pyrenew.randomvariable import DistributionalVariable, TransformedVariable
-from pyrenew.latent import (
- InfectionInitializationProcess,
- InitializeInfectionsExponentialGrowth,
-)
-import pyrenew.transformation as t
-```
-
-## Expected Behavior of the Feedback Mechanism
-
-This section demonstrates how to use the `InfectionsWithFeedback` class to model the number of infections
-by running a baseline simulation.
-We simulate the model without an observation process which
-lets us see how the feedback mechanism affects the reproduction number and subsequent infections over time.
-
-The following code-chunks define the model components in the following order.
-
-* We define a deterministic PMF `gen_int_array`, the discretized fraction of infectiousness per day.
-For this example, we have a 4-day PMF of equal probabilities with feedback strength ($\gamma$)
-set to 0.01 to create moderate damping.
-
-* We specify the `InfectionInitializationProcess` class, where the number of latent infections immediately before the renewal process begins follows a log-normal distribution with mean = 0 and standard deviation = 1.
-
-* We specify the `InfectionsWithFeedback` component using the same distribution as the generation interval.
-
-* The latent random variable $\mathcal{R}^\mathrm{u}(t)$ is estimated via random walk
-on $\log(\mathcal{R}^\mathrm{u}(t))$, (as in the basic renewal model).
-
-```{python}
-# | label: model-components
-gen_int_array = jnp.array([0.25, 0.25, 0.25, 0.25])
-gen_int = DeterministicPMF(name="gen_int", value=gen_int_array)
-feedback_strength = DeterministicVariable(name="feedback_strength", value=0.01)
-
-# | label: model-components
-I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsExponentialGrowth(
- gen_int_array.size,
- DeterministicVariable(name="rate", value=0.05),
- ),
-)
-
-latent_infections = InfectionsWithFeedback(
- infection_feedback_strength=feedback_strength,
- infection_feedback_pmf=gen_int,
-)
-
-
-class MyRt(RandomVariable):
- def validate(self):
- pass
-
- def sample(self, n: int, **kwargs) -> tuple:
- sd_rt = numpyro.sample("Rt_random_walk_sd", dist.HalfNormal(0.025))
-
- rt_rv = TransformedVariable(
- name="log_rt_random_walk",
- base_rv=RandomWalk(
- name="log_rt",
- step_rv=DistributionalVariable(
- name="rw_step_rv", distribution=dist.Normal(0, sd_rt)
- ),
- ),
- transforms=t.ExpTransform(),
- )
- rt_init_rv = DistributionalVariable(
- name="init_log_rt", distribution=dist.Normal(0, 0.2)
- )
- init_rt = rt_init_rv.sample()
-
- return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
-```
-
-We build the model from these components and then simulate data from it without observed infections.
-
-```{python}
-# | label: build1
-model0 = RtInfectionsRenewalModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- latent_infections_rv=latent_infections,
- Rt_process_rv=MyRt(),
- infection_obs_process_rv=None, # no observed infections
-)
-```
-
-```{python}
-# | label: simulate1
-# Sampling and fitting model 0 (with no obs for infections)
-with numpyro.handlers.seed(rng_seed=223):
- model0_samp = model0.sample(n_datapoints=30)
-```
-
-```{python}
-# | label: fig-simulate1
-# | fig-cap: Simulated infections with no observation process
-import matplotlib.pyplot as plt
-
-fig, ax = plt.subplots()
-ax.plot(model0_samp.latent_infections)
-ax.set_xlabel("Time")
-ax.set_ylabel("Infections")
-plt.show()
-```
-
-With feedback strength $\gamma = 0.01$, we see:
-
-* Initial exponential-like growth when infection history is minimal
-* Damping effects becoming visible as infections accumulate
-
-## PyRenew's `RandomVariable` class
-
-### Fundamentals
-
-All instances of PyRenew's `RandomVariable` should have at least three functions: `__init__()`, `validate()`, and `sample()`. The `__init__()` function is the constructor and initializes the class. The `validate()` function checks if the class is correctly initialized. Finally, the `sample()` method contains the core of the class; it should return a `tuple` or `namedtuple`. The following is a minimal example of a `RandomVariable` class based on `numpyro.distributions.Normal`:
-
-```{python}
-from pyrenew.metaclass import RandomVariable
-
-
-class MyNormal(RandomVariable):
- def __init__(self, loc, scale):
- self.validate(scale)
- self.loc = loc
- self.scale = scale
- return None
-
- @staticmethod
- def validate(self):
- if self.scale <= 0:
- raise ValueError("Scale must be positive")
- return None
-
- def sample(self, **kwargs):
- return (dist.Normal(loc=self.loc, scale=self.scale),)
-```
-
-The `@staticmethod` decorator exposes the `validate` function to be used outside the class.
-
-## Essential elements of class `InfectionsWithFeedback`
-
-As an exercise we define a new class `InfFeedback`, which exactly follows the `InfectionsWithFeedback` class,
-but omits the tedious checks on its inputs.
-
-Although returning a `namedtuple` is not strictly required, they are the recommended return type, as they make the code more readable. The following code-chunk shows how to create a `namedtuple` for the `InfectionsWithFeedback` class:
-
-```{python}
-# | label: data-class
-from collections import namedtuple
-
-# Creating a tuple to store the output
-InfFeedbackSample = namedtuple(
- typename="InfFeedbackSample",
- field_names=["post_initialization_infections", "rt"],
- defaults=(None, None),
-)
-```
-
-The next step is to create the actual class. The bulk of its implementation lies in the function `pyrenew.latent.compute_infections_from_rt_with_feedback()`. We will also use the `pyrenew.arrayutils.pad_edges_to_match()` function to ensure the passed vectors match their lengths. The following code-chunk shows most of the implementation of the `InfectionsWithFeedback` class:
-
-```{python}
-# | label: new-model-def
-# | code-line-numbers: true
-# Creating the class
-from pyrenew.metaclass import RandomVariable
-from pyrenew.latent import compute_infections_from_rt_with_feedback
-from pyrenew import arrayutils as au
-from jax.typing import ArrayLike
-import jax.numpy as jnp
-
-
-class InfFeedback(RandomVariable):
- """Latent infections"""
-
- def __init__(
- self,
- infection_feedback_strength: RandomVariable,
- infection_feedback_pmf: RandomVariable,
- ) -> None:
- """Constructor"""
-
- self.infection_feedback_strength = infection_feedback_strength
- self.infection_feedback_pmf = infection_feedback_pmf
-
- return None
-
- def validate(self):
- """
- Generally, this method should be more meaningful, but we will skip it for now
- """
- return None
-
- def sample(
- self,
- Rt: ArrayLike,
- I0: ArrayLike,
- gen_int: ArrayLike,
- **kwargs,
- ) -> tuple:
- """Sample infections with feedback"""
-
- # Generation interval
- gen_int_rev = jnp.flip(gen_int)
-
- # Baseline infections
- I0_vec = I0[-gen_int_rev.size :]
-
- # Sampling inf feedback strength and adjusting the shape
- inf_feedback_strength = self.infection_feedback_strength(
- **kwargs,
- )
-
- inf_feedback_strength = jnp.atleast_1d(inf_feedback_strength)
-
- inf_feedback_strength, _ = au.pad_edges_to_match(
- x=inf_feedback_strength,
- y=Rt,
- )
-
- # Sampling inf feedback and adjusting the shape
- inf_feedback_pmf = self.infection_feedback_pmf(**kwargs)
- inf_fb_pmf_rev = jnp.flip(inf_feedback_pmf)
-
- # Generating the infections with feedback
- all_infections, Rt_adj = compute_infections_from_rt_with_feedback(
- I0=I0_vec,
- Rt_raw=Rt,
- infection_feedback_strength=inf_feedback_strength,
- reversed_generation_interval_pmf=gen_int_rev,
- reversed_infection_feedback_pmf=inf_fb_pmf_rev,
- )
-
- # Storing adjusted Rt for future use
- numpyro.deterministic("Rt_adjusted", Rt_adj)
-
- # Preparing the output
-
- return InfFeedbackSample(
- post_initialization_infections=all_infections,
- rt=Rt_adj,
- )
-```
-
-The core of the class is implemented in the `sample()` method. Things to highlight from the above code:
-
-1. **Arguments of `sample`**: The `InfFeedback` class will be used within `RtInfectionsRenewalModel` to generate latent infections. During the sampling process, `InfFeedback()` will receive the reproduction number, the initial number of infections, and the generation interval. `RandomVariable()` calls are expected to include the `**kwargs` argument, even if unused.
-
-2. **Saving computed quantities**: Since `Rt_adj` is not generated via `numpyro.sample()`, we use `numpyro.deterministic()` to record the quantity to a site; allowing us to access it later.
-
-3. **Return type of `InfFeedback()`**: As said before, the `sample()` method should return a `tuple` or `namedtuple`. In our case, we return a `namedtuple` `InfFeedbackSample` with two fields: `infections` and `rt`.
-
-To check our work, we build a second model using our new class `InfFeedback` and compare it to our first model, which uses the built-in `InfectionsWithFeedback` class.
-Since class `InfFeedback` exactly follows the `InfectionsWithFeedback` class, if correctly implemented, the two models should provide exactly the same predictions.
-
-```{python}
-# | label: simulation2
-latent_infections2 = InfFeedback(
- infection_feedback_strength=feedback_strength,
- infection_feedback_pmf=gen_int,
-)
-
-model1 = RtInfectionsRenewalModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- latent_infections_rv=latent_infections2,
- Rt_process_rv=MyRt(),
- infection_obs_process_rv=None,
-)
-
-# Sampling and fitting model 0 (with no obs for infections)
-with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model1.sample(n_datapoints=30)
-```
-
-Comparing `model0` with `model1`, these two should match:
-
-```{python}
-# | label: fig-model0-vs-model1
-# | fig-cap: Comparing latent infections from model 0 and model 1
-import matplotlib.pyplot as plt
-
-fig, ax = plt.subplots(ncols=2)
-ax[0].plot(model0_samp.latent_infections)
-ax[1].plot(model1_samp.latent_infections)
-ax[0].set_xlabel("Time (model 0)")
-ax[1].set_xlabel("Time (model 1)")
-ax[0].set_ylabel("Infections")
-plt.show()
-```
diff --git a/docs/tutorials/day_of_the_week.qmd b/docs/tutorials/day_of_the_week.qmd
deleted file mode 100644
index 62f374ec..00000000
--- a/docs/tutorials/day_of_the_week.qmd
+++ /dev/null
@@ -1,447 +0,0 @@
----
-title: Implementing a Day of the Week Effect
-format: gfm
-engine: jupyter
----
-
-This document illustrates how to leverage the time-aware arrays to create a a day of the week effect. We use the same model designed in the hospital admissions-only tutorial.
-
-## Recap: Hospital Admissions Model
-
-In the ["Fitting a hospital admissions-only model" tutorial](hospital_admissions_model.md), we built a fairly complicated model that included pre-defined random variables as well as a custom random variable representing the reproductive number.
-
-
-```mermaid
-flowchart LR
-
- %% Elements
- rt_proc["Random walk RT process
(latent)"];
- latent_inf["Latent infections"]
- latent_ihr["Infection to hospitalization rate
(latent)"]
- neg_binom["Observation process
(hospitalizations)"]
- latent_hosp["Latent hospitalizations"];
- i0["Initial infections
(latent)"];
- gen_int["Generation interval
(fixed)"];
- hosp_int["Hospitalization interval
(fixed)"];
-
- %% Models
- basic_model(("Infections model"));
- admin_model(("Hospital admissions model"));
-
- %% Latent infections
- rt_proc --> latent_inf;
- i0 --> latent_inf;
- gen_int --> latent_inf;
- latent_inf --> basic_model
-
- %% Hospitalizations
- hosp_int --> latent_hosp
-
- neg_binom --> admin_model;
- latent_ihr --> latent_hosp;
- basic_model --> admin_model;
- latent_hosp --> admin_model;
-```
-
-In this tutorial, we will focus on adding a new component: the day-of-the-week effect.
-
-#### Setup: import all packages and load the simulated data
-
-```{python}
-# | label: setup
-import numpyro
-import polars as pl
-import jax.numpy as jnp
-import numpyro.distributions as dist
-from pyrenew import (
- datasets,
- deterministic,
- latent,
- metaclass,
- model,
- observation,
- process,
- randomvariable,
- transformation,
-)
-from pyrenew.latent import (
- InfectionInitializationProcess,
- InitializeInfectionsExponentialGrowth,
-)
-
-numpyro.set_host_device_count(2)
-
-# Loading and processing the data
-dat = (
- datasets.load_wastewater()
- .group_by("date")
- .first()
- .select(["date", "daily_hosp_admits"])
- .sort("date")
- .head(90)
-)
-daily_hosp_admits = dat["daily_hosp_admits"].to_numpy()
-dates = dat["date"].to_numpy()
-
-# Loading additional datasets
-gen_int = datasets.load_generation_interval()
-inf_hosp_int = datasets.load_infection_admission_interval()
-
-# We only need the probability_mass column of each dataset
-gen_int_array = gen_int["probability_mass"].to_numpy()
-inf_hosp_int_array = inf_hosp_int["probability_mass"].to_numpy()
-```
-
-#### Define all model components
-
-We start by (re)defining the hospital admissions model without a day-of-the-week effect.
-
-```{python}
-# | label: hospital-admissions-model
-
-# Hospital admissions
-inf_hosp_int = deterministic.DeterministicPMF(
- name="inf_hosp_int", value=inf_hosp_int_array
-)
-
-hosp_rate = randomvariable.DistributionalVariable(
- name="IHR", distribution=dist.LogNormal(jnp.log(0.05), jnp.log(1.1))
-)
-
-latent_hosp = latent.HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp_int,
- infection_hospitalization_ratio_rv=hosp_rate,
-)
-
-
-# Infection process
-latent_inf = latent.Infections()
-n_initialization_points = max(gen_int_array.size, inf_hosp_int_array.size) - 1
-
-I0 = InfectionInitializationProcess(
- "I0_initialization",
- randomvariable.DistributionalVariable(
- name="I0",
- distribution=dist.LogNormal(loc=jnp.log(100), scale=jnp.log(1.75)),
- ),
- InitializeInfectionsExponentialGrowth(
- n_initialization_points,
- deterministic.DeterministicVariable(name="rate", value=0.05),
- ),
-)
-
-# Generation interval
-gen_int_pmf = deterministic.DeterministicPMF(
- name="gen_int", value=gen_int_array
-)
-
-
-class RtRandomWalk(metaclass.RandomVariable):
- def validate(self):
- pass
-
- def sample(self, n: int, **kwargs) -> tuple:
- sd_rt = numpyro.sample("Rt_random_walk_sd", dist.HalfNormal(0.025))
-
- # Random walk step
- step_rv = randomvariable.DistributionalVariable(
- name="rw_step_rv", distribution=dist.Normal(0, sd_rt)
- )
-
- rt_init_rv = randomvariable.DistributionalVariable(
- name="init_log_rt", distribution=dist.Normal(0, 0.2)
- )
-
- # Random walk process
- base_rv = process.RandomWalk(
- name="log_rt",
- step_rv=step_rv,
- )
-
- # Transforming the random walk to the Rt scale
- rt_rv = randomvariable.TransformedVariable(
- name="Rt_rv",
- base_rv=base_rv,
- transforms=transformation.ExpTransform(),
- )
- init_rt = rt_init_rv.sample()
-
- return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
-
-
-rtproc = RtRandomWalk()
-
-# we place a log-Normal prior on the concentration
-# parameter of the negative binomial.
-nb_conc_rv = randomvariable.TransformedVariable(
- "concentration",
- randomvariable.DistributionalVariable(
- name="concentration_raw",
- distribution=dist.TruncatedNormal(loc=0, scale=1, low=0.01),
- ),
- transformation.PowerTransform(-2),
-)
-
-# now we define the observation process
-obs = observation.NegativeBinomialObservation(
- "negbinom_rv",
- concentration_rv=nb_conc_rv,
-)
-
-# now we build the model
-hosp_model = model.HospitalAdmissionsModel(
- latent_infections_rv=latent_inf,
- latent_hosp_admissions_rv=latent_hosp,
- I0_rv=I0,
- gen_int_rv=gen_int_pmf,
- Rt_process_rv=rtproc,
- hosp_admission_obs_process_rv=obs,
-)
-```
-
-We run the model and then plot the posterior samples.
-
-```{python}
-import jax
-import numpy as np
-
-# Model without weekday effect
-hosp_model.run(
- num_samples=2000,
- num_warmup=2000,
- data_observed_hosp_admissions=daily_hosp_admits,
- rng_key=jax.random.key(54),
- mcmc_args=dict(progress_bar=False),
-)
-```
-
-```{python}
-# | code-fold: true
-# | label: fig-output-admissions-padding-and-weekday
-# | fig-cap: Hospital Admissions posterior distribution without weekday effect
-import arviz as az
-import matplotlib.pyplot as plt
-
-ppc_samples = hosp_model.posterior_predictive(
- n_datapoints=daily_hosp_admits.size
-)
-idata = az.from_numpyro(
- posterior=hosp_model.mcmc,
- posterior_predictive=ppc_samples,
-)
-
-# Use a time series plot (plot_ts) from arviz for plotting
-axes = az.plot_ts(
- idata,
- y="negbinom_rv",
- y_hat="negbinom_rv",
- num_samples=200,
- y_kwargs={
- "color": "blue",
- "linewidth": 1.0,
- "marker": "o",
- "linestyle": "solid",
- },
- y_hat_plot_kwargs={"color": "skyblue", "alpha": 0.05},
- y_mean_plot_kwargs={"color": "black", "linestyle": "--", "linewidth": 2.5},
- backend_kwargs={"figsize": (8, 6)},
- textsize=15.0,
-)
-ax = axes[0][0]
-ax.set_xlabel("Time", fontsize=20)
-ax.set_ylabel("Hospital Admissions", fontsize=20)
-handles, labels = ax.get_legend_handles_labels()
-ax.legend(
- handles, ["Observed", "Sample Mean", "Posterior Samples"], loc="best"
-)
-plt.show()
-```
-
-
-## Round 2: Incorporating day-of-the-week effects
-
-To add a day-of-the-week effect to the previous model we add two additional arguments to the latent hospital admissions random variable:
-
-* `day_of_the_week_rv` (a `RandomVariable`)
-* `obs_data_first_day_of_the_week` (an `int` mapping days of the week from 0:6, zero being Monday).
-
-The `day_of_the_week_rv`'s sample method should return a vector of length seven; those values are then broadcasted to match the length of the dataset.
-Since the observed data may start in a weekday other than Monday, the `obs_data_first_day_of_the_week` argument is used to offset the day-of-the-week effect.
-
-For this effect we use a scaled Dirichlet distribution, and define
-a `TransformedVariable` that samples an array of length seven from numpyro's `distributions.Dirichlet` and applies a `transformation.AffineTransform` to scale it by seven. [^note-other-examples]:
-
-[^note-other-examples]: A similar weekday effect is implemented in its own module, with example code in the [Periodic Effects](periodic_effects.md) tutorial.
-
-```{python}
-# | label: weekly-effect
-# Instantiating the day-of-the-week effect
-dayofweek_effect = randomvariable.TransformedVariable(
- name="dayofweek_effect",
- base_rv=randomvariable.DistributionalVariable(
- name="dayofweek_effect_raw",
- distribution=dist.Dirichlet(jnp.ones(7)),
- ),
- transforms=transformation.AffineTransform(
- loc=0, scale=7, domain=jnp.array([0, 1])
- ),
-)
-```
-
-We use this day-of-the-week effect as the latent hospital admissions random variable.
-
-Since the day-of-the-week effect takes into account the first day of the dataset, we need to determine the day of the week of the first observation.
-To do this, we convert the first date in the dataset to a `datetime` object and extract the day of the week:
-
-```{python}
-# | label: latent-hosp-weekday
-# Figuring out the day of the week of the first observation
-import datetime as dt
-
-first_dow_in_data = dates[0].astype(dt.datetime).weekday()
-first_dow_in_data # zero
-
-# Define the latent hospital admissions RV with the day-of-the-week effect
-latent_hosp_wday_effect = latent.HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp_int,
- infection_hospitalization_ratio_rv=hosp_rate,
- day_of_week_effect_rv=dayofweek_effect,
- # Concidirently, this is zero
- obs_data_first_day_of_the_week=first_dow_in_data,
-)
-
-# New model with day-of-the-week effect
-hosp_model_dow = model.HospitalAdmissionsModel(
- latent_infections_rv=latent_inf,
- latent_hosp_admissions_rv=latent_hosp_wday_effect,
- I0_rv=I0,
- gen_int_rv=gen_int_pmf,
- Rt_process_rv=rtproc,
- hosp_admission_obs_process_rv=obs,
-)
-```
-
-Running the model:
-
-```{python}
-# | label: model-2-run-weekday
-# Model with weekday effect
-hosp_model_dow.run(
- num_samples=2000,
- num_warmup=2000,
- data_observed_hosp_admissions=daily_hosp_admits,
- rng_key=jax.random.key(54),
- mcmc_args=dict(progress_bar=False),
-)
-```
-
-As a result, we can see the posterior distribution of our novel day-of-the-week effect:
-
-```{python}
-# | code-fold: true
-# | label: fig-output-day-of-week
-# | fig-cap: Day of the week effect
-# Create an InferenceData object from hosp_model_dow
-dow_idata = az.from_numpyro(
- posterior=hosp_model_dow.mcmc,
-)
-
-# Extract day of week effect (DOW)
-dow_effect_raw = dow_idata.posterior["dayofweek_effect_raw"].squeeze().T
-indices = np.random.choice(dow_effect_raw.shape[1], size=200, replace=False)
-dow_plot_samples = dow_effect_raw[:, indices]
-fig, ax = plt.subplots(1, 1, figsize=(8, 6))
-ax.plot(
- np.arange(dow_effect_raw.shape[0]),
- dow_plot_samples,
- color="skyblue",
- alpha=0.10,
-)
-ax.plot([], [], color="skyblue", alpha=0.10, label="DOW Posterior Samples")
-ax.plot(
- np.arange(dow_effect_raw.shape[0]),
- dow_plot_samples.mean(dim="draw"),
- color="black",
- linewidth=2.0,
- linestyle="--",
- label="Sample Mean",
-)
-ax.legend(loc="best")
-ax.set_ylabel("Effect", fontsize=20)
-ax.set_xlabel("Day Of Week", fontsize=20)
-plt.show()
-```
-
-The new model with the day-of-the-week effect can be compared to the previous model, which doesn't distinguish the day of the week, by plotting the posterior distributions from each.
-
-```{python}
-# | code-fold: true
-# | label: fig-output-admissions-original
-# | fig-cap: Hospital Admissions posterior distribution without weekday effect
-# Without weekday effect (from earlier)
-axes = az.plot_ts(
- idata,
- y="negbinom_rv",
- y_hat="negbinom_rv",
- num_samples=200,
- y_kwargs={
- "color": "blue",
- "linewidth": 1.0,
- "marker": "o",
- "linestyle": "solid",
- },
- y_hat_plot_kwargs={"color": "skyblue", "alpha": 0.05},
- y_mean_plot_kwargs={"color": "black", "linestyle": "--", "linewidth": 2.5},
- backend_kwargs={"figsize": (8, 6)},
- textsize=15.0,
-)
-ax = axes[0][0]
-ax.set_xlabel("Time", fontsize=20)
-ax.set_ylabel("Hospital Admissions", fontsize=20)
-handles, labels = ax.get_legend_handles_labels()
-ax.legend(
- handles,
- ["Observed", "Posterior Predictive", "Samples wo/ WDE"],
- loc="best",
-)
-plt.show()
-```
-
-```{python}
-# | code-fold: true
-# | label: fig-output-admissions-wof
-# | fig-cap: Hospital Admissions posterior distribution with weekday effect
-# Figure with weekday effect
-ppc_samples = hosp_model_dow.posterior_predictive(
- n_datapoints=daily_hosp_admits.size
-)
-idata = az.from_numpyro(
- posterior=hosp_model_dow.mcmc,
- posterior_predictive=ppc_samples,
-)
-
-axes = az.plot_ts(
- idata,
- y="negbinom_rv",
- y_hat="negbinom_rv",
- num_samples=200,
- y_kwargs={
- "color": "blue",
- "linewidth": 1.0,
- "marker": "o",
- "linestyle": "solid",
- },
- y_hat_plot_kwargs={"color": "skyblue", "alpha": 0.05},
- y_mean_plot_kwargs={"color": "black", "linestyle": "--", "linewidth": 2.5},
- backend_kwargs={"figsize": (8, 6)},
- textsize=15.0,
-)
-ax = axes[0][0]
-ax.set_xlabel("Time", fontsize=20)
-ax.set_ylabel("Hospital Admissions", fontsize=20)
-handles, labels = ax.get_legend_handles_labels()
-ax.legend(
- handles, ["Observed", "Posterior Predictive", "Samples w/ WDE"], loc="best"
-)
-plt.show()
-```
diff --git a/docs/tutorials/getting_started.qmd b/docs/tutorials/getting_started.qmd
deleted file mode 100644
index 58d391c1..00000000
--- a/docs/tutorials/getting_started.qmd
+++ /dev/null
@@ -1,64 +0,0 @@
----
-title: "Getting started with pyrenew"
-format: gfm
-engine: jupyter
----
-
-`pyrenew` is a flexible tool for simulating and making statistical inferences from epidemiologic models, with an emphasis on renewal models. Built on `numpyro`, `pyrenew` provides core components for model building and pre-defined models for processing various observational processes.
-
-## Prerequisites
-
-The following tutorials assume that you are familiar with both infectious disease dynamics and Python programming.
-
- - You must have a Python3 installation (use tools like [pyenv](https://realpython.com/intro-to-pyenv/) or [compile and install](https://ubuntuhandbook.org/index.php/2023/05/install-python-3-12-ubuntu/) from the [release page](https://www.python.org/downloads/)).
- - You should be able to install packages, load modules in python, and manage virtual environments (we recommend [uv](https://docs.astral.sh/uv/)).
- - You should understand the concept of a [class](https://realpython.com/python-classes/) and [metaclass](https://realpython.com/courses/python-metaclasses/) in python.
- - You should be familiar with Bayesian inference and have a working understanding of MCMC methods used to fit Bayesian models to data (some resources are available [here](https://mc-stan.org/docs/2_18/reference-manual/effective-sample-size-section.html), and [here](https://xcelab.net/rm/)).
-
-## Installing pyrenew
-
-You can install `pyrenew` using either uv or pip. To install `pyrenew` using `uv`, run the following command from within the directory containing the `pyrenew` project:
-
-```bash
-uv sync
-```
-
-To install `pyrenew` using pip, run the following command:
-
-```bash
-pip install git+https://github.com/CDCgov/PyRenew@main
-```
-
-
-## The fundamentals
-
-The `pyrenew` package leverages `numpyro`'s flexibility to build models via composition.
-It's core components are the metaclasses `RandomVariable` and `Model` (in Python, a _metaclass_ is a class whose instances are also classes, where a _class_ is a template for making objects). Within the `pyrenew` package, a `RandomVariable` is a quantity that models can estimate and sample from, **including deterministic quantities**. The benefit of this design is that the definition of the `sample()` function can be arbitrary, allowing the user to either sample from a distribution using `numpyro.sample()`, compute fixed quantities (like a mechanistic equation), or return a fixed value (like a pre-computed PMF.) For instance, when estimating a PMF, the `RandomVariable` sampling function may roughly be defined as:
-
-```python
-# define a new class called MyRandVar that inherits from the RandomVariable class
-class MyRandVar(RandomVariable):
- #define a method called sample that returns an object of type ArrayLike
- def sample(...) -> ArrayLike:
- # calls sample function from NumPyro package
- return numpyro.sample(...)
-```
-
-We can instead use a fixed quantity for that variable (like a pre-computed PMF), where the `RandomVariable`'s sample function could be defined as:
-
-```python
-# instead define MyRandVar to still inherit from the RandVariable class
-class MyRandVar(RandomVariable):
- #define sample method that still returns an ArrayLike object
- def sample(...) -> ArrayLike:
- #sampling method is a pre-computed PMF, a JAX NumPy array with explicit elements
- return jax.numpy.array([0.2, 0.7, 0.1])
-```
-
-Thus, when a `Model` samples from `MyRandVar`, it could be either adding random variables to be estimated (first example) or just retrieving some quantity needed for other calculations (second example).
-
-The `Model` metaclass provides basic functionality for estimating and simulation. Like `RandomVariable`, the `Model` metaclass has a `sample()` method that defines the model structure. Ultimately, models can be nested (or inherited), providing a straightforward way to add layers of complexity.
-Currently the `Model` metaclass contains of two model classes
-
-* `RtInfectionsRenewalModel` - a basic renewal model consisting of infections and reproduction numbers.
-* `HospitalAdmissionsModel`- a model which extends a basic renewal model to account for latent and observed hospital admissions.
diff --git a/docs/tutorials/hospital_admissions_model.qmd b/docs/tutorials/hospital_admissions_model.qmd
deleted file mode 100644
index 25109bdc..00000000
--- a/docs/tutorials/hospital_admissions_model.qmd
+++ /dev/null
@@ -1,630 +0,0 @@
----
-title: "Fitting a hospital admissions-only model"
-format: gfm
-engine: jupyter
----
-
-This document illustrates how a hospital admissions-only model can be fitted using simulated example data.
-
-We begin by loading `numpyro` and configuring the device count to 2 to enable running MCMC chains in parallel. By default, XLA (which is used by JAX for compilation) considers all CPU cores as one device. Depending on your system's configuration, we recommend using numpyro's [set_host_device_count()](https://num.pyro.ai/en/stable/utilities.html#set-host-device-count) function to set the number of devices available for parallel computing.
-
-
-```{python}
-# | label: numpyro setup
-import numpyro
-
-numpyro.set_host_device_count(2)
-```
-
-## Model definition
-
-The hospital admissions model is a semi-mechanistic model that describes the number of observed hospital admissions, a positive integer, as a discretely distributed variable whose location is the number of latent hospital admissions
-$$
-h(t) \sim \text{HospDist}\left(H(t)\right),
-$$
-where $h(t)$ is the observed number of hospital admissions at time $t$, and $H(t)$ is the number of latent hospital admissions at time $t$.
-For this example $\text{HospDist}$ is a negative binomial distribution with an inferred concentration $k$.
-$$
-h(t) \sim \mathrm{NegativeBinomial}\left(\mathrm{mean} = H(t), \mathrm{concentration} = k\right)
-$$
-
-The number of latent hospital admissions at time $t$ is a function of the number of latent infections at time $t$ and the infection to admission rate
-$$
-H(t) = p_\mathrm{hosp}(t) \sum_{\tau = 0}^{T_d} d(\tau) I(t-\tau)
-$$
-where $d(\tau)$ is the infection to hospital admission interval, $I(t)$ is the number of latent infections at time $t$, and $p_\mathrm{hosp}(t)$ is the infection to admission rate.
-
-The priors on $p_\mathrm{hosp}(t)$ and $k$ reflect our knowledge, or lack thereof, of these variables.
-
-$$
-\begin{align*}
-\log(p_{\mathrm{hosp}(t)}) & \sim \mathrm{Normal}(\mu=\log(0.05), \sigma=\log(1.1)) \\
-\log(k) & \sim \mathrm{Normal}(\mu=\log(1), \sigma=\log(10))
-\end{align*}
-$$
-
-Without much information, the prior on $k$ allows for any amount of overdispersion.
-
-The latent infections are modeled as a renewal process:
-
-$$
-\begin{align*}
-I(t) &= \mathcal{R}(t) \times \sum_{\tau < t} I(\tau) g(t - \tau) \\
-\log[I(0)] &\sim \text{Normal}(\mu=\log(100), \sigma=\log(1.75))
-\end{align*}
-$$
-
-The reproduction number $\mathcal{R}(t)$ is modeled as a random walk in logarithmic space, i.e.:
-
-$$
-\begin{align*}
-\log[\mathcal{R}(t)] & = \log[\mathcal{R}(t-1)] + \epsilon\\
-\epsilon & \sim \text{Normal}(\mu=0, \sigma=0.025) \\
-\mathcal{R}(0) &\sim \text{TruncatedNormal}(\text{loc}=1.2, \text{scale}=0.2, \text{min}=0)
-\end{align*}
-$$
-
-
-## Data processing
-
-We start by loading the data and inspecting the first five rows.
-
-```{python}
-# | label: data-inspect
-import polars as pl
-from pyrenew import datasets
-
-dat = datasets.load_wastewater()
-dat[["date", "site", "daily_hosp_admits"]].head(5)
-```
-
-The data shows one entry per site. In this simulated dataset, all sites have the same number of admissions each day; therefore, we only use the first observation per day, keeping the date and number of admissions. Furthermore, we only use the first 90 days' worth of data.
-
-```{python}
-# | label: aggregation
-# Keeping the first observation of each date, sort by date, first 90 days only
-dat = (
- dat.group_by("date")
- .first()
- .select(["date", "daily_hosp_admits"])
- .sort("date")
- .head(90)
-)
-dat.head(5)
-```
-
-Let's take a look at the daily prevalence of hospital admissions.
-
-```{python}
-# | label: fig-plot-hospital-admissions
-# | fig-cap: Daily hospital admissions from the simulated data
-import matplotlib.pyplot as plt
-import matplotlib.dates as mdates
-
-daily_hosp_admits = dat["daily_hosp_admits"].to_numpy()
-dates = dat["date"].to_numpy()
-ax = plt.gca()
-ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m-%d"))
-ax.xaxis.set_major_locator(mdates.DayLocator(interval=7))
-ax.set_xlim(dates[0], dates[-1])
-plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
-plt.plot(dates, daily_hosp_admits, "-o")
-plt.xlabel("Date")
-plt.ylabel("Admissions")
-plt.show()
-```
-
-## Building the model
-
-
-This model uses the following imports from `pyrenew`, `jax`, and `numpyro`.
-
-```{python}
-from pyrenew import (
- deterministic,
- latent,
- metaclass,
- model,
- observation,
- process,
- randomvariable,
- transformation,
-)
-from pyrenew.latent import (
- InfectionInitializationProcess,
- InitializeInfectionsExponentialGrowth,
-)
-import jax.numpy as jnp
-import numpyro.distributions as dist
-```
-
-The generation interval and the infection to hospital admission interval are `RandomVariable`s.
-Here we load two more datasets which provide deterministic quantities for these variables.
-
-```{python}
-# | label: data-extract-intervals
-gen_int = datasets.load_generation_interval()
-inf_hosp_int = datasets.load_infection_admission_interval()
-
-# We only need the probability_mass column of each dataset
-gen_int_array = gen_int["probability_mass"].to_numpy()
-inf_hosp_int_array = inf_hosp_int["probability_mass"].to_numpy()
-
-# Taking a peek at the first 5 elements of each
-print("Generation interval\n", gen_int[:5])
-print("Hospital admission interval\n", inf_hosp_int_array[:5])
-```
-
-We plot the full distribution curves side-by-side.
-
-```{python}
-# | label: fig-data-plot
-# | fig-cap: Generation interval and infection to hospital admission interval
-fig, axs = plt.subplots(1, 2)
-
-axs[0].plot(gen_int_array)
-axs[0].set_title("Generation interval")
-axs[1].plot(inf_hosp_int_array)
-axs[1].set_title("Infection to hospital admission interval")
-fig.supxlabel("Days")
-plt.show()
-```
-
-With these arrays in hand, we define the necessary components of the model.
-
-* `inf_hosp_int` is a `DeterministicPMF` object that takes the infection to hospital admission interval as input.
-
-* `hosp_rate` is a `DistributionalVariable` object that takes a numpyro distribution to represent the infection to hospital admission rate
-
-* `latent_hosp` is an instance of the `HospitalAdmissions` class. It is a `RandomVariable` that takes two distributions as inputs: the infection to admission interval and the infection to hospital admission rate.
-
-
-```{python}
-# | label: latent-hosp
-inf_hosp_int = deterministic.DeterministicPMF(
- name="inf_hosp_int", value=inf_hosp_int_array
-)
-
-hosp_rate = randomvariable.DistributionalVariable(
- name="IHR", distribution=dist.LogNormal(jnp.log(0.05), jnp.log(1.1))
-)
-
-latent_hosp = latent.HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp_int,
- infection_hospitalization_ratio_rv=hosp_rate,
-)
-```
-
-Next we define infection process `IO` and generation interval `gen_int`.
-
-```{python}
-# | label: initializing-rest-of-model
-# Infection process
-latent_inf = latent.Infections()
-n_initialization_points = max(gen_int_array.size, inf_hosp_int_array.size) - 1
-I0 = InfectionInitializationProcess(
- "I0_initialization",
- randomvariable.DistributionalVariable(
- name="I0",
- distribution=dist.LogNormal(loc=jnp.log(100), scale=jnp.log(1.75)),
- ),
- InitializeInfectionsExponentialGrowth(
- n_initialization_points,
- deterministic.DeterministicVariable(name="rate", value=0.05),
- ),
-)
-
-# Generation interval
-gen_int_pmf = deterministic.DeterministicPMF(
- name="gen_int", value=gen_int_array
-)
-```
-
-Then we define the $\mathcal{R}(t)$ random walk process `rtproc`.
-
-```{python}
-class RtRandomWalk(metaclass.RandomVariable):
- def validate(self):
- pass
-
- def sample(self, n: int, **kwargs) -> tuple:
- sd_rt = numpyro.sample("Rt_random_walk_sd", dist.HalfNormal(0.025))
-
- # Random walk step
- step_rv = randomvariable.DistributionalVariable(
- name="rw_step_rv", distribution=dist.Normal(0, sd_rt)
- )
-
- rt_init_rv = randomvariable.DistributionalVariable(
- name="init_log_rt", distribution=dist.Normal(0, 0.2)
- )
-
- # Random walk process
- base_rv = process.RandomWalk(
- name="log_rt",
- step_rv=step_rv,
- )
-
- # Transforming the random walk to the Rt scale
- rt_rv = randomvariable.TransformedVariable(
- name="Rt_rv",
- base_rv=base_rv,
- transforms=transformation.ExpTransform(),
- )
- init_rt = rt_init_rv.sample()
-
- return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
-
-
-rtproc = RtRandomWalk()
-```
-
-The final model component is the observation process `obs`.
-
-```{python}
-# Put a log-Normal prior on the concentration parameter
-# of the negative binomial.
-nb_conc_rv = randomvariable.DistributionalVariable(
- "concentration",
- distribution=dist.LogNormal(loc=0.0, scale=jnp.log(10.0)),
-)
-
-# The observation process
-obs = observation.NegativeBinomialObservation(
- "negbinom_rv",
- concentration_rv=nb_conc_rv,
-)
-```
-
-We can now build the model.
-
-```{python}
-# | label: init-model
-hosp_model = model.HospitalAdmissionsModel(
- latent_infections_rv=latent_inf,
- latent_hosp_admissions_rv=latent_hosp,
- I0_rv=I0,
- gen_int_rv=gen_int_pmf,
- Rt_process_rv=rtproc,
- hosp_admission_obs_process_rv=obs,
-)
-```
-
-Let's simulate from the prior predictive distribution to check that the model is working:
-
-```{python}
-# | label: simulation
-import numpy as np
-
-timeframe = 120
-
-with numpyro.handlers.seed(rng_seed=223):
- simulated_data = hosp_model.sample(n_datapoints=timeframe)
-```
-
-```{python}
-# | label: fig-basic
-# | fig-cap: Simulated Rt and Admissions
-import matplotlib.pyplot as plt
-
-fig, axs = plt.subplots(1, 2)
-
-# Rt plot
-axs[0].plot(simulated_data.Rt)
-axs[0].set_ylabel("Simulated Rt")
-
-# Admissions plot
-axs[1].plot(simulated_data.observed_hosp_admissions, "-o")
-axs[1].set_ylabel("Simulated Admissions")
-
-fig.suptitle("Basic renewal model")
-fig.supxlabel("Time")
-plt.tight_layout()
-plt.show()
-```
-
-## Fitting the model
-
-We now fit the model, not to these simulated data, but rather to the dataset we retrieved above. We use the `run` method of the `Model` object:
-
-```{python}
-# | label: model-fit
-import jax
-
-hosp_model.run(
- num_samples=1000,
- num_warmup=1000,
- data_observed_hosp_admissions=daily_hosp_admits,
- rng_key=jax.random.key(54),
- mcmc_args=dict(progress_bar=False, num_chains=2),
-)
-```
-
-We use [`arviz`](https://www.arviz.org/) to visualize the model fit.
-First we load the module and converting the fitted model to an `arviz` `InferenceData` object,
-then we plot the observed values against the inferred latent values (i.e. the mean of the negative binomial observation process).
-
-```{python}
-# | label: fig-output-hospital-admissions
-# | fig-cap: Latent hospital admissions posterior samples (gray) and observed admissions timeseries (red).
-import arviz as az
-
-ppc_samples = hosp_model.posterior_predictive(
- n_datapoints=daily_hosp_admits.size
-)
-idata = az.from_numpyro(
- posterior=hosp_model.mcmc,
- posterior_predictive=ppc_samples,
-)
-
-axes = az.plot_ts(
- idata,
- y="negbinom_rv",
- y_hat="negbinom_rv",
- num_samples=200,
- y_kwargs={
- "color": "blue",
- "linewidth": 1.0,
- "marker": "o",
- "linestyle": "solid",
- },
- y_hat_plot_kwargs={"color": "skyblue", "alpha": 0.05},
- y_mean_plot_kwargs={"color": "black", "linestyle": "--", "linewidth": 2.5},
- backend_kwargs={"figsize": (8, 6)},
- textsize=15.0,
-)
-ax = axes[0][0]
-ax.set_xlabel("Time", fontsize=20)
-ax.set_ylabel("Hospital Admissions", fontsize=20)
-handles, labels = ax.get_legend_handles_labels()
-ax.legend(
- handles, ["Observed", "Sample Mean", "Posterior Samples"], loc="best"
-)
-plt.show()
-```
-
-## Results exploration and MCMC diagnostics
-
-To explore further, we use `arviz` to visualize the results.
-We obtain the summary of model diagnostics and print the diagnostics for `latent_hospital_admissions[1]`
-
-```{python}
-# | label: mcmc-diagnostics
-# | warning: false
-
-idata = az.from_numpyro(hosp_model.mcmc)
-
-diagnostic_stats_summary = az.summary(
- idata.posterior,
- kind="diagnostics",
-)
-
-print(diagnostic_stats_summary.loc["latent_hospital_admissions[1]"])
-```
-
-Below we plot 90% and 50% highest density intervals for latent hospital admissions using [plot_hdi](https://python.arviz.org/en/stable/api/generated/arviz.plot_hdi.html):
-
-```{python}
-# | label: fig-output-admission-distribution
-# | fig-cap: Hospital Admissions posterior distribution
-x_data = idata.posterior["latent_hospital_admissions_dim_0"]
-y_data = idata.posterior["latent_hospital_admissions"]
-
-
-fig, axes = plt.subplots(figsize=(6, 5))
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.9,
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.3},
- ax=axes,
-)
-
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.5,
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.6},
- ax=axes,
-)
-
-# Add the posterior median to the figure
-median_ts = y_data.median(dim=["chain", "draw"])
-
-axes.plot(x_data, median_ts, color="C0", label="Median")
-axes.legend()
-axes.set_title("Posterior Hospital Admissions", fontsize=10)
-axes.set_xlabel("Time", fontsize=10)
-axes.set_ylabel("Hospital Admissions", fontsize=10)
-plt.show()
-```
-
-We can also look at credible intervals for the posterior distribution of latent infections:
-
-```{python}
-# | label: fig-output-infections-distribution
-# | fig-cap: Posterior Latent Infections
-x_data = (
- idata.posterior["all_latent_infections_dim_0"] - n_initialization_points
-)
-y_data = idata.posterior["all_latent_infections"]
-
-fig, axes = plt.subplots(figsize=(6, 5))
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.9,
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.3},
- ax=axes,
-)
-
-az.plot_hdi(
- x_data,
- y_data,
- hdi_prob=0.5,
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.6},
- ax=axes,
-)
-
-# Add the posterior median to the figure
-median_ts = y_data.median(dim=["chain", "draw"])
-axes.plot(x_data, median_ts, color="C0", label="Median")
-axes.legend()
-```
-
-## Predictive checks and forecasting
-We can use the `Model`'s `posterior_predictive` and `prior_predictive` methods to generate posterior and prior predictive samples for observed admissions.
-
-```{python}
-# | label: demonstrate-use-of-predictive-methods
-
-idata = az.from_numpyro(
- hosp_model.mcmc,
- posterior_predictive=hosp_model.posterior_predictive(
- n_datapoints=len(daily_hosp_admits)
- ),
- prior=hosp_model.prior_predictive(
- n_datapoints=len(daily_hosp_admits),
- numpyro_predictive_args={"num_samples": 1000},
- ),
-)
-```
-
-We will use `plot_lm` method from `arviz` to plot the posterior predictive distribution against the actual observed data below:
-
-```{python}
-# | label: fig-posterior-predictive
-# | fig-cap: Hospital Admissions posterior distribution with plot_lm
-fig, ax = plt.subplots()
-az.plot_lm(
- "negbinom_rv",
- idata=idata,
- kind_pp="hdi",
- y_kwargs={"color": "black"},
- y_hat_fill_kwargs={"color": "C0"},
- axes=ax,
-)
-
-ax.set_title("Posterior Predictive Plot")
-ax.set_ylabel("Hospital Admissions")
-ax.set_xlabel("Days")
-plt.show()
-```
-
-By increasing `n_datapoints`, we can perform forecasting using the posterior predictive distribution.
-
-```{python}
-# | label: posterior-predictive-distribution
-n_forecast_points = 28
-idata = az.from_numpyro(
- hosp_model.mcmc,
- posterior_predictive=hosp_model.posterior_predictive(
- n_datapoints=len(daily_hosp_admits) + n_forecast_points,
- ),
- prior=hosp_model.prior_predictive(
- n_datapoints=len(daily_hosp_admits),
- numpyro_predictive_args={"num_samples": 1000},
- ),
-)
-```
-
-Below we plot the prior predictive distributions using equal tailed Bayesian credible intervals:
-
-```{python}
-# | label: fig-output-prior-predictive
-# | fig-cap: Prior Predictive Admissions
-def compute_eti(dataset, eti_prob):
- eti_bdry = dataset.quantile(
- ((1 - eti_prob) / 2, 1 / 2 + eti_prob / 2), dim=("chain", "draw")
- )
- return eti_bdry.T
-
-
-fig, axes = plt.subplots(figsize=(6, 5))
-az.plot_hdi(
- idata.prior_predictive["negbinom_rv_dim_0"],
- hdi_data=compute_eti(idata.prior_predictive["negbinom_rv"], 0.9),
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.3},
- ax=axes,
-)
-
-az.plot_hdi(
- idata.prior_predictive["negbinom_rv_dim_0"],
- hdi_data=compute_eti(idata.prior_predictive["negbinom_rv"], 0.5),
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.6},
- ax=axes,
-)
-
-plt.scatter(
- idata.observed_data["negbinom_rv_dim_0"],
- idata.observed_data["negbinom_rv"],
- color="black",
-)
-
-axes.set_title("Prior Predictive Admissions", fontsize=10)
-axes.set_xlabel("Time", fontsize=10)
-axes.set_ylabel("Observed Admissions", fontsize=10)
-plt.yscale("log")
-plt.show()
-```
-
-And now we plot the posterior predictive distributions with a `{python} n_forecast_points`-day-ahead forecast:
-
-```{python}
-# | label: fig-output-posterior-predictive-forecast
-# | fig-cap: Posterior predictive admissions, including a forecast.
-x_data = idata.posterior_predictive["negbinom_rv_dim_0"]
-y_data = idata.posterior_predictive["negbinom_rv"]
-fig, axes = plt.subplots(figsize=(6, 5))
-az.plot_hdi(
- x_data,
- hdi_data=compute_eti(y_data, 0.9),
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.3},
- ax=axes,
-)
-
-az.plot_hdi(
- x_data,
- hdi_data=compute_eti(y_data, 0.5),
- color="C0",
- smooth=False,
- fill_kwargs={"alpha": 0.6},
- ax=axes,
-)
-
-# Add median of the posterior to the figure
-median_ts = y_data.median(dim=["chain", "draw"])
-
-plt.plot(
- x_data,
- median_ts,
- color="C0",
- label="Median",
-)
-plt.scatter(
- idata.observed_data["negbinom_rv_dim_0"],
- idata.observed_data["negbinom_rv"],
- color="black",
-)
-axes.legend()
-axes.set_title(
- "Posterior Predictive Admissions, including a forecast", fontsize=10
-)
-axes.set_xlabel("Time", fontsize=10)
-axes.set_ylabel("Hospital Admissions", fontsize=10)
-plt.show()
-```
diff --git a/docs/tutorials/latent_hierarchical_infections.qmd b/docs/tutorials/latent_hierarchical_infections.qmd
index 397c05b9..82b5a75f 100644
--- a/docs/tutorials/latent_hierarchical_infections.qmd
+++ b/docs/tutorials/latent_hierarchical_infections.qmd
@@ -909,4 +909,4 @@ and call the `configure_latent` method
to specify the latent process model structure: generation interval, initial infections, and temporal dynamics.
See tutorial:
-- [Building Multi-Signal Renewal Models](building_multisignal_models.qmd)
+- [Building Multi-Signal Renewal Models](building_multisignal_models.md)
diff --git a/docs/tutorials/random_variables.qmd b/docs/tutorials/random_variables.qmd
new file mode 100644
index 00000000..342eb092
--- /dev/null
+++ b/docs/tutorials/random_variables.qmd
@@ -0,0 +1,526 @@
+---
+title: "PyRenew's RandomVariable abstract base class"
+format:
+ gfm:
+ fig-width: 16
+ fig-height: 10
+ html:
+ toc: true
+ embed-resources: true
+ self-contained-math: true
+ code-fold: true
+ code-tools: true
+engine: jupyter
+jupyter:
+ jupytext:
+ text_representation:
+ extension: .qmd
+ format_name: quarto
+ format_version: '1.0'
+ jupytext_version: 1.18.1
+ kernelspec:
+ display_name: Python 3 (ipykernel)
+ language: python
+ name: python3
+---
+
+## Design principle: all quantities are RandomVariables
+
+In a Bayesian model, all quantities — data, parameters, hyperparameters, derived computations —
+are random variables living in a single joint probability model (Gelman et al., BDA3 §1.3).
+The only distinction is whether a given quantity is known (observed, conditioned on) or unknown (to be inferred).
+A fixed constant is just a degenerate random variable; an estimated rate is a draw from a prior.
+Both participate in the same joint distribution.
+
+PyRenew embodies this through its `RandomVariable` abstract base class.
+All model components implement the same `sample()` interface,
+so you can swap a fixed quantity for an estimated one — or vice versa —
+without changing any surrounding model code.
+
+
+```{python}
+# | label: setup
+# | output: false
+import jax.numpy as jnp
+import numpy as np
+import numpyro
+import numpyro.distributions as dist
+import pandas as pd
+import plotnine as p9
+import warnings
+from plotnine.exceptions import PlotnineWarning
+
+warnings.filterwarnings("ignore", category=PlotnineWarning)
+
+from pyrenew.metaclass import RandomVariable
+from pyrenew.deterministic import DeterministicVariable, DeterministicPMF
+from pyrenew.randomvariable import (
+ DistributionalVariable,
+ StaticDistributionalVariable,
+ DynamicDistributionalVariable,
+ TransformedVariable,
+)
+from pyrenew.observation import Counts, NegativeBinomialNoise
+import pyrenew.transformation as transformation
+from pyrenew import datasets
+```
+
+## Concrete implementations
+
+PyRenew provides four `RandomVariable` implementations that cover most modeling needs.
+
+### DeterministicVariable
+
+A degenerate random variable that returns a fixed value.
+Its `sample()` method simply returns the stored value, unchanged.
+
+```{python}
+# | label: deterministic-example
+ihr_fixed = DeterministicVariable("ihr", 0.01)
+
+with numpyro.handlers.seed(rng_seed=0):
+ value = ihr_fixed()
+ print(f"IHR (fixed): {value}")
+```
+
+`DeterministicPMF` specializes this for probability mass functions,
+validating at construction time that the values sum to 1:
+
+```{python}
+# | label: deterministic-pmf-example
+delay_pmf = DeterministicPMF(
+ "delay",
+ jnp.array([0.0, 0.1, 0.3, 0.3, 0.2, 0.1]),
+)
+
+with numpyro.handlers.seed(rng_seed=0):
+ pmf = delay_pmf()
+ print(f"Delay PMF: {np.round(pmf, 2)}, sum: {pmf.sum():.1f}")
+```
+
+### DistributionalVariable
+
+A random variable that draws from a numpyro distribution via `numpyro.sample()`.
+The `DistributionalVariable` factory function dispatches to one of two classes
+depending on whether the distribution is known at construction time or built at sample time.
+
+**StaticDistributionalVariable**: the distribution is fully specified at construction time.
+
+```{python}
+# | label: static-distributional-example
+# IHR with a Beta(2, 198) prior: mean ~1%, moderate uncertainty
+ihr_estimated = DistributionalVariable("ihr", dist.Beta(2, 198))
+print(f"Type: {type(ihr_estimated).__name__}")
+
+with numpyro.handlers.seed(rng_seed=0):
+ print(f"IHR (sampled): {ihr_estimated():.4f}")
+```
+
+**DynamicDistributionalVariable**: the distribution is constructed at sample time
+from a callable. This is useful when distribution parameters depend on
+other sampled quantities.
+
+```{python}
+# | label: dynamic-distributional-example
+# A Normal whose location is determined at sample time
+dynamic_rv = DistributionalVariable(
+ "dynamic_normal",
+ lambda loc: dist.Normal(loc, 0.1),
+)
+print(f"Type: {type(dynamic_rv).__name__}")
+
+with numpyro.handlers.seed(rng_seed=0):
+ print(f"Sample (loc=2.0): {dynamic_rv.sample(2.0):.3f}")
+ print(f"Sample (loc=5.0): {dynamic_rv.sample(5.0):.3f}")
+```
+
+### TransformedVariable
+
+Wraps another `RandomVariable` and applies a deterministic transformation to its output.
+This is useful for reparameterizations and derived quantities.
+
+```{python}
+# | label: transformed-example
+# Day-of-week effect: Dirichlet draw scaled by 7
+# so effects are multiplicative and preserve weekly totals
+dow_effect = TransformedVariable(
+ name="dow_effect",
+ base_rv=DistributionalVariable(
+ name="dow_raw",
+ distribution=dist.Dirichlet(jnp.ones(7)),
+ ),
+ transforms=transformation.AffineTransform(loc=0, scale=7),
+)
+
+with numpyro.handlers.seed(rng_seed=0):
+ effect = dow_effect()
+ print(f"Day-of-week multipliers: {np.round(effect, 3)}")
+ print(f"Sum: {effect.sum():.1f}")
+```
+
+### Interchangeability
+
+Because all implementations share the `sample()` interface,
+you can swap them freely. For example, the `Counts` observation process
+accepts any `RandomVariable` as its `ascertainment_rate_rv`.
+The model code is identical whether the rate is fixed or estimated:
+
+```{python}
+# | label: interchangeability
+hosp_delay_pmf = jnp.array(
+ datasets.load_infection_admission_interval()["probability_mass"].to_numpy()
+)
+
+# Fixed ascertainment rate
+hosp_fixed = Counts(
+ name="hospital",
+ ascertainment_rate_rv=DeterministicVariable("ihr", 0.01),
+ delay_distribution_rv=DeterministicPMF("delay", hosp_delay_pmf),
+ noise=NegativeBinomialNoise(DeterministicVariable("conc", 10.0)),
+)
+
+# Estimated ascertainment rate --- same model structure, different RV
+hosp_estimated = Counts(
+ name="hospital",
+ ascertainment_rate_rv=DistributionalVariable("ihr", dist.Beta(2, 198)),
+ delay_distribution_rv=DeterministicPMF("delay", hosp_delay_pmf),
+ noise=NegativeBinomialNoise(DeterministicVariable("conc", 10.0)),
+)
+```
+
+## The RandomVariable public API
+
+The `RandomVariable` metaclass (defined in `pyrenew.metaclass`) requires
+subclasses to implement:
+
+| Method | Signature | Purpose |
+|--------|-----------|---------|
+| `sample` | `sample(**kwargs) -> tuple` | Core computation: return a value, draw from a distribution, or perform a calculation |
+| `validate` | `validate(**kwargs) -> None` | Check that parameters are well-formed; raise an error if not |
+
+The metaclass also provides:
+
+| Method | Behavior |
+|--------|----------|
+| `__call__(**kwargs)` | Alias for `sample(**kwargs)`, so `my_rv()` is equivalent to `my_rv.sample()` |
+
+The `**kwargs` pattern is central to composability: a `RandomVariable`
+accepts whatever arguments its `sample()` method needs, and passes
+through any additional keyword arguments to internal calls.
+
+## Writing a custom RandomVariable
+
+The built-in classes handle most cases, but sometimes you need a component with
+custom logic---for instance, one that makes multiple `numpyro.sample()` calls,
+performs domain-specific validation, or records derived quantities
+via `numpyro.deterministic()`.
+
+### Example: ascertainment with day-of-week effects
+
+Hospital admissions data typically shows day-of-week reporting patterns:
+fewer admissions are reported on weekends, more on weekdays.
+We can model this as a multiplicative adjustment to a baseline ascertainment rate.
+
+The predicted hospital admissions on day $t$ are:
+
+$$\lambda_t = \alpha \cdot w_{t \bmod 7} \cdot \sum_{d=0}^{D} I_{t-d} \cdot p_d$$
+
+where $\alpha$ is the baseline ascertainment rate,
+$w_j$ for $j = 0, \ldots, 6$ are day-of-week multipliers
+(positive, summing to 7 so that the weekly total is preserved),
+and the summation is the delay convolution.
+
+We define a custom `RandomVariable` that bundles the
+ascertainment rate and day-of-week effect into a single component,
+returning a daily rate vector.
+
+```{python}
+# | label: custom-rv-definition
+from jax.typing import ArrayLike
+
+
+class AscertainmentWithDayOfWeek(RandomVariable):
+ """
+ Ascertainment rate modulated by day-of-week reporting effects.
+
+ Combines a scalar ascertainment rate with a 7-element
+ day-of-week multiplier to produce a daily rate vector.
+ """
+
+ def __init__(
+ self,
+ name: str,
+ baseline_rate_rv: RandomVariable,
+ dow_concentration: ArrayLike,
+ first_day_offset: int = 0,
+ ) -> None:
+ """
+ Parameters
+ ----------
+ name
+ Name prefix for numpyro sample sites.
+ baseline_rate_rv
+ RandomVariable for the baseline ascertainment rate.
+ dow_concentration
+ Dirichlet concentration parameters, shape (7,).
+ Larger values produce less day-to-day variation.
+ first_day_offset
+ Day of the week for the first day of the timeseries
+ (0 = Monday, 6 = Sunday).
+ """
+ self.validate(dow_concentration, first_day_offset)
+ self.name = name
+ self.baseline_rate_rv = baseline_rate_rv
+ self.dow_concentration = jnp.asarray(dow_concentration)
+ self.first_day_offset = first_day_offset
+
+ @staticmethod
+ def validate(dow_concentration, first_day_offset) -> None:
+ """Check that parameters are well-formed."""
+ dow_concentration = jnp.asarray(dow_concentration)
+ if dow_concentration.shape != (7,):
+ raise ValueError(
+ f"dow_concentration must have shape (7,), "
+ f"got {dow_concentration.shape}"
+ )
+ if jnp.any(dow_concentration <= 0):
+ raise ValueError("dow_concentration values must be positive")
+ if not (0 <= first_day_offset <= 6):
+ raise ValueError(
+ f"first_day_offset must be 0-6, got {first_day_offset}"
+ )
+
+ def sample(self, n_days: int, **kwargs) -> tuple:
+ """
+ Sample a daily ascertainment rate vector.
+
+ Parameters
+ ----------
+ n_days
+ Number of days in the timeseries.
+
+ Returns
+ -------
+ tuple
+ Containing a single array of daily ascertainment rates,
+ shape (n_days,).
+ """
+ # Sample baseline rate (delegates to another RandomVariable)
+ baseline = self.baseline_rate_rv()
+ numpyro.deterministic(f"{self.name}_baseline_rate", baseline)
+
+ # Sample day-of-week effects from Dirichlet, scale to sum to 7
+ dow_raw = numpyro.sample(
+ f"{self.name}_dow_raw",
+ dist.Dirichlet(self.dow_concentration),
+ )
+ dow_effect = dow_raw * 7.0
+ numpyro.deterministic(f"{self.name}_dow_effect", dow_effect)
+
+ # Tile the 7-element vector across the timeseries
+ full_cycle = jnp.tile(dow_effect, (n_days // 7) + 1)
+ daily_dow = full_cycle[
+ self.first_day_offset : self.first_day_offset + n_days
+ ]
+
+ # Combine: daily rate = baseline * day-of-week multiplier
+ daily_rate = baseline * daily_dow
+
+ return (daily_rate,)
+```
+
+This class bundles three things that belong together:
+
+1. **Two sample statements**: one for the baseline rate (delegated to `baseline_rate_rv`) and one for the Dirichlet day-of-week effects
+2. **Derived computation**: scaling, tiling, offsetting, and combining into a daily vector
+3. **Validation**: checking that concentration has shape `(7,)` and the offset is valid
+
+### Sampling from the custom RV
+
+```{python}
+# | label: sample-custom-rv
+ascertainment_rv = AscertainmentWithDayOfWeek(
+ name="hosp",
+ baseline_rate_rv=DistributionalVariable("ihr", dist.Beta(2, 198)),
+ dow_concentration=jnp.array([5, 5, 5, 5, 5, 2, 2]), # lower on weekends
+ first_day_offset=0, # timeseries starts on Monday
+)
+
+n_days = 28
+
+with numpyro.handlers.seed(rng_seed=42):
+ (daily_rates,) = ascertainment_rv(n_days=n_days)
+```
+
+```{python}
+# | label: plot-custom-rv
+day_labels = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
+
+rates_df = pd.DataFrame(
+ {
+ "day": np.arange(n_days),
+ "rate": np.array(daily_rates),
+ "dow": [day_labels[i % 7] for i in range(n_days)],
+ "day_type": [
+ "Weekend" if i % 7 >= 5 else "Weekday" for i in range(n_days)
+ ],
+ }
+)
+
+(
+ p9.ggplot(rates_df, p9.aes(x="day", y="rate", fill="day_type"))
+ + p9.geom_col(alpha=0.7, color="black", size=0.3)
+ + p9.scale_fill_manual(
+ values={"Weekday": "lightblue", "Weekend": "orange"}
+ )
+ + p9.labs(
+ x="Day",
+ y="Daily Ascertainment Rate",
+ title="Ascertainment Rate with Day-of-Week Effects",
+ fill="",
+ )
+ + p9.theme_grey()
+ + p9.theme(plot_title=p9.element_text(size=14, weight="bold"))
+)
+```
+
+The weekly pattern is clear: weekday rates are higher than weekend rates,
+reflecting reduced reporting on weekends.
+
+### Applying to a simulated observation process
+
+We now use the `AscertainmentWithDayOfWeek` instance defined above
+to generate day-of-week multipliers within a realistic observation
+pipeline, showing how a custom `RandomVariable` integrates into a
+data-generating process.
+
+In a hospital surveillance system, the observation pipeline is:
+
+1. **Infections** occur over time (the latent process).
+2. **Delay**: each infection leads to a hospital admission after some delay, modeled as a convolution with a delay PMF.
+3. **Reporting**: the admission is *reported* with day-of-week effects — fewer reports on weekends, more on weekdays.
+4. **Noise**: observed counts include measurement noise.
+
+The day-of-week effect is a *reporting* artifact, so it is
+applied **after** the delay convolution.
+We walk through each step explicitly to make the pipeline clear.
+
+```{python}
+# | label: simulate-infections
+# Epidemic curve: exponential growth then decay
+n_days = 60
+infections = 10000.0 * jnp.exp(-((jnp.arange(n_days) - 25.0) ** 2) / 200.0)
+```
+
+We convolve infections with the delay PMF and apply a baseline
+ascertainment rate to get the expected admissions before noise:
+
+```{python}
+# | label: delay-convolve
+baseline_rate = 0.01
+expected_admissions = jnp.convolve(
+ baseline_rate * infections,
+ hosp_delay_pmf,
+ mode="full",
+)[:n_days]
+```
+
+Now we simulate observed admissions with and without day-of-week effects.
+Without effects, we add noise directly to the smooth expected curve.
+With effects, we first call our custom `ascertainment_rv` to get a
+day-of-week multiplier, apply it to the expected admissions
+(post-convolution, at the reporting stage), then add noise.
+
+```{python}
+# | label: simulate-with-dow
+n_samples = 20
+concentration = 50.0
+
+results = []
+for seed in range(n_samples):
+ with numpyro.handlers.seed(rng_seed=seed):
+ # Without day-of-week: noise on the smooth expected curve
+ obs_no_dow = numpyro.sample(
+ "obs_no_dow",
+ dist.NegativeBinomial2(expected_admissions, concentration),
+ )
+
+ with numpyro.handlers.seed(rng_seed=seed + 1000):
+ # Sample day-of-week multiplier from the custom RV
+ (daily_rates,) = ascertainment_rv(n_days=n_days)
+ dow_multiplier = daily_rates / daily_rates.mean()
+
+ # Apply day-of-week *after* the delay convolution
+ obs_with_dow = numpyro.sample(
+ "obs_with_dow",
+ dist.NegativeBinomial2(
+ expected_admissions * dow_multiplier, concentration
+ ),
+ )
+
+ for i in range(n_days):
+ results.append(
+ {
+ "day": i,
+ "admissions": float(obs_no_dow[i]),
+ "type": "No day-of-week effect",
+ "sample": seed,
+ }
+ )
+ results.append(
+ {
+ "day": i,
+ "admissions": float(obs_with_dow[i]),
+ "type": "With day-of-week effect",
+ "sample": seed,
+ }
+ )
+```
+
+```{python}
+# | label: plot-comparison
+results_df = pd.DataFrame(results)
+
+(
+ p9.ggplot(
+ results_df,
+ p9.aes(x="day", y="admissions", group="sample"),
+ )
+ + p9.geom_line(alpha=0.3, size=0.5, color="steelblue")
+ + p9.facet_wrap("~ type", ncol=2)
+ + p9.labs(
+ x="Day",
+ y="Hospital Admissions",
+ title="Simulated Hospital Admissions",
+ )
+ + p9.theme_grey()
+ + p9.theme(plot_title=p9.element_text(size=14, weight="bold"))
+)
+```
+
+The left panel shows smooth variation from noise alone.
+The right panel shows the additional weekly oscillation introduced by
+day-of-week reporting effects — the sawtooth pattern of weekend dips
+and weekday peaks that is characteristic of real hospital admissions data.
+
+## Summary
+
+### Choosing a RandomVariable implementation
+
+| Need | Use |
+|------|-----|
+| Fixed known value | `DeterministicVariable` |
+| Fixed known PMF | `DeterministicPMF` |
+| Sample from a fixed distribution | `DistributionalVariable` (static) |
+| Sample from a distribution parameterized at sample time | `DistributionalVariable` (dynamic, pass a callable) |
+| Deterministic transformation of another RV | `TransformedVariable` |
+| Multiple sample statements, custom validation, or derived computation | Custom `RandomVariable` subclass |
+
+### Writing a custom RandomVariable
+
+1. Subclass `RandomVariable` from `pyrenew.metaclass`
+2. Implement `validate()` as a `@staticmethod`; call it in `__init__`
+3. Implement `sample(**kwargs)` returning a `tuple`
+4. Use `numpyro.sample()` for quantities to be estimated
+5. Use `numpyro.deterministic()` to record derived quantities in the trace
+6. Accept `**kwargs` in `sample()` for composability
diff --git a/pyrenew/latent/__init__.py b/pyrenew/latent/__init__.py
index 433d13e9..ead837f4 100644
--- a/pyrenew/latent/__init__.py
+++ b/pyrenew/latent/__init__.py
@@ -11,7 +11,6 @@
HierarchicalNormalPrior,
StudentTGroupModePrior,
)
-from pyrenew.latent.hospitaladmissions import HospitalAdmissions
from pyrenew.latent.infection_functions import (
compute_infections_from_rt,
compute_infections_from_rt_with_feedback,
@@ -36,7 +35,6 @@
)
__all__ = [
- "HospitalAdmissions",
"Infections",
"logistic_susceptibility_adjustment",
"compute_infections_from_rt",
@@ -51,7 +49,7 @@
"BaseLatentInfectionProcess",
"LatentSample",
"PopulationStructure",
- # Hierarchical infection processes
+ # Hierarchical infect1`_).
-
- Following other semi-mechanistic renewal frameworks, we model the *expected*
- hospital admissions per capita $H(t)$ as a convolution of the
- *expected* latent incident infections per capita $I(t)$, and a
- discrete infection to hospitalization distribution $d(\tau)$, scaled
- by the probability of being hospitalized $p_\mathrm{hosp}(t)$.
-
- To account for day-of-week effects in hospital reporting, we use an
- estimated *day of the week effect* $\omega(t)$. If $t$ and $t'$
- are the same day of the week, $\omega(t) = \omega(t')$. The seven
- values that $\omega(t)$ takes on are constrained to have mean 1.
-
- ```math
- H(t) = \omega(t) p_\mathrm{hosp}(t) \sum_{\tau = 0}^{T_d} d(\tau) I(t-\tau)
- ```
-
- Where $T_d$ is the maximum delay from infection to hospitalization
- that we consider.
- """
-
- def __init__(
- self,
- infection_to_admission_interval_rv: RandomVariable,
- infection_hospitalization_ratio_rv: RandomVariable,
- day_of_week_effect_rv: RandomVariable | None = None,
- hospitalization_reporting_ratio_rv: RandomVariable | None = None,
- obs_data_first_day_of_the_week: int = 0,
- ) -> None:
- """
- Default constructor
-
- Parameters
- ----------
- infection_to_admission_interval_rv
- pmf for reporting (informing) hospital admissions (see
- pyrenew.observations.Deterministic).
- infection_hospitalization_ratio_rv
- Infection to hospitalization rate random variable.
- day_of_week_effect_rv
- Day of the week effect. Should return a ArrayLike with 7
- values. Defaults to a deterministic variable with
- jax.numpy.ones(7) (no effect).
- hospitalization_reporting_ratio_rv
- Random variable for the hospital admission reporting
- probability. Defaults to 1 (full reporting).
- obs_data_first_day_of_the_week
- The day of the week that the first day of the observation data
- corresponds to. Valid values are 0-6, where 0 is Monday and 6 is
- Sunday. Defaults to 0.
-
- Returns
- -------
- None
- """
-
- if day_of_week_effect_rv is None:
- day_of_week_effect_rv = DeterministicVariable(
- name="weekday_effect", value=jnp.ones(7)
- )
- if hospitalization_reporting_ratio_rv is None:
- hospitalization_reporting_ratio_rv = DeterministicVariable(
- name="hosp_report_prob", value=1.0
- )
-
- HospitalAdmissions.validate(
- infection_to_admission_interval_rv,
- infection_hospitalization_ratio_rv,
- day_of_week_effect_rv,
- hospitalization_reporting_ratio_rv,
- obs_data_first_day_of_the_week,
- )
-
- self.infection_to_admission_interval_rv = infection_to_admission_interval_rv
- self.infection_hospitalization_ratio_rv = infection_hospitalization_ratio_rv
- self.day_of_week_effect_rv = day_of_week_effect_rv
- self.hospitalization_reporting_ratio_rv = hospitalization_reporting_ratio_rv
- self.obs_data_first_day_of_the_week = obs_data_first_day_of_the_week
-
- @staticmethod
- def validate(
- infection_to_admission_interval_rv: Any,
- infection_hospitalization_ratio_rv: Any,
- day_of_week_effect_rv: Any,
- hospitalization_reporting_ratio_rv: Any,
- obs_data_first_day_of_the_week: Any,
- ) -> None:
- """
- Validates that the IHR, weekday effects, probability of being
- reported hospitalized distributions, and infection to
- hospital admissions reporting delay pmf are RandomVariable types
-
- Parameters
- ----------
- infection_to_admission_interval_rv
- Possibly incorrect input for the infection to hospitalization
- interval distribution.
- infection_hospitalization_ratio_rv
- Possibly incorrect input for infection to hospitalization rate distribution.
- day_of_week_effect_rv
- Possibly incorrect input for day of the week effect.
- hospitalization_reporting_ratio_rv
- Possibly incorrect input for distribution or fixed value for the
- hospital admission reporting probability.
- obs_data_first_day_of_the_week
- Possibly incorrect input for the day of the week that the first day
- of the observation data corresponds to. Valid values are 0-6, where
- 0 is Monday and 6 is Sunday.
-
- Returns
- -------
- None
-
- Raises
- ------
- AssertionError
- If any of the random variables are not of the correct type, or if
- the day of the week is not within the valid range.
- """
- assert isinstance(infection_to_admission_interval_rv, RandomVariable)
- assert isinstance(infection_hospitalization_ratio_rv, RandomVariable)
- assert isinstance(day_of_week_effect_rv, RandomVariable)
- assert isinstance(hospitalization_reporting_ratio_rv, RandomVariable)
- assert isinstance(obs_data_first_day_of_the_week, int)
- assert 0 <= obs_data_first_day_of_the_week <= 6
-
- return None
-
- def sample(
- self,
- latent_infections: ArrayLike,
- **kwargs,
- ) -> HospitalAdmissionsSample:
- """
- Samples from the observation process
-
- Parameters
- ----------
- latent_infections
- Latent infections.
- **kwargs
- Additional keyword arguments passed through to
- internal `sample()` calls,
- should there be any.
-
- Returns
- -------
- HospitalAdmissionsSample
- """
-
- infection_hosp_rate = self.infection_hospitalization_ratio_rv(**kwargs)
-
- infection_to_admission_interval = self.infection_to_admission_interval_rv(
- **kwargs
- )
-
- latent_hospital_admissions, _ = compute_delay_ascertained_incidence(
- latent_infections,
- infection_to_admission_interval,
- infection_hosp_rate,
- )
-
- # Applying the day of the week effect. For this we need to:
- # 1. Get the day of the week effect
- # 2. Identify the offset of the latent_infections
- # 3. Apply the day of the week effect to the
- # latent_hospital_admissions
- dow_effect_sampled = self.day_of_week_effect_rv(**kwargs)
-
- if dow_effect_sampled.size != 7:
- raise ValueError(
- "Day of the week effect should have 7 values. "
- f"Got {dow_effect_sampled.size} instead."
- )
-
- inf_offset = self.obs_data_first_day_of_the_week % 7
-
- # Replicating the day of the week effect to match the number of
- # timepoints
- dow_effect = au.tile_until_n(
- data=dow_effect_sampled,
- n_timepoints=latent_hospital_admissions.size,
- offset=inf_offset,
- )
-
- latent_hospital_admissions = latent_hospital_admissions * dow_effect
-
- # Applying reporting probability
- latent_hospital_admissions = (
- latent_hospital_admissions
- * self.hospitalization_reporting_ratio_rv(**kwargs)
- )
-
- numpyro.deterministic("latent_hospital_admissions", latent_hospital_admissions)
-
- return HospitalAdmissionsSample(
- infection_hosp_rate=infection_hosp_rate,
- latent_hospital_admissions=latent_hospital_admissions,
- multiplier=dow_effect,
- )
diff --git a/pyrenew/model/__init__.py b/pyrenew/model/__init__.py
index 302e6424..c2972079 100644
--- a/pyrenew/model/__init__.py
+++ b/pyrenew/model/__init__.py
@@ -1,13 +1,9 @@
# numpydoc ignore=GL08
-from pyrenew.model.admissionsmodel import HospitalAdmissionsModel
from pyrenew.model.multisignal_model import MultiSignalModel
from pyrenew.model.pyrenew_builder import PyrenewBuilder
-from pyrenew.model.rtinfectionsrenewalmodel import RtInfectionsRenewalModel
__all__ = [
- "HospitalAdmissionsModel",
"PyrenewBuilder",
"MultiSignalModel",
- "RtInfectionsRenewalModel",
]
diff --git a/pyrenew/model/admissionsmodel.py b/pyrenew/model/admissionsmodel.py
deleted file mode 100644
index 62ea760b..00000000
--- a/pyrenew/model/admissionsmodel.py
+++ /dev/null
@@ -1,202 +0,0 @@
-# numpydoc ignore=GL08
-
-from __future__ import annotations
-
-from typing import NamedTuple
-
-from jax.typing import ArrayLike
-
-from pyrenew.deterministic import NullObservation
-from pyrenew.metaclass import Model, RandomVariable
-from pyrenew.model.rtinfectionsrenewalmodel import RtInfectionsRenewalModel
-
-
-class HospModelSample(NamedTuple):
- """
- A container for holding the output from `model.HospitalAdmissionsModel.sample()`.
-
- Attributes
- ----------
- Rt : ArrayLike | None, optional
- The reproduction number over time. Defaults to None.
- latent_infections : ArrayLike | None, optional
- The estimated number of new infections over time. Defaults to None.
- infection_hosp_rate : ArrayLike | None, optional
- The infected hospitalization rate. Defaults to None.
- latent_hosp_admissions : ArrayLike | None, optional
- The estimated latent hospitalizations. Defaults to None.
- observed_hosp_admissions : ArrayLike | None, optional
- The sampled or observed hospital admissions. Defaults to None.
- """
-
- Rt: ArrayLike | None = None
- latent_infections: ArrayLike | None = None
- infection_hosp_rate: ArrayLike | None = None
- latent_hosp_admissions: ArrayLike | None = None
- observed_hosp_admissions: ArrayLike | None = None
-
- def __repr__(self):
- return (
- f"HospModelSample(Rt={self.Rt}, "
- f"latent_infections={self.latent_infections}, "
- f"infection_hosp_rate={self.infection_hosp_rate}, "
- f"latent_hosp_admissions={self.latent_hosp_admissions}, "
- f"observed_hosp_admissions={self.observed_hosp_admissions}"
- )
-
-
-class HospitalAdmissionsModel(Model):
- """
- Hospital Admissions Model (BasicRenewal + HospitalAdmissions)
-
- This class inherits from pyrenew.models.Model. It extends the
- basic renewal model by adding a hospital admissions module, e.g.,
- pyrenew.observations.HospitalAdmissions.
- """
-
- def __init__(
- self,
- latent_hosp_admissions_rv: RandomVariable,
- latent_infections_rv: RandomVariable,
- gen_int_rv: RandomVariable,
- I0_rv: RandomVariable,
- Rt_process_rv: RandomVariable,
- hosp_admission_obs_process_rv: RandomVariable,
- ) -> None: # numpydoc ignore=PR04
- """
- Default constructor
-
- Parameters
- ----------
- latent_hosp_admissions_rv
- Latent process for the hospital admissions.
- latent_infections_rv
- The infections latent process (passed to RtInfectionsRenewalModel).
- gen_int_rv
- Generation time (passed to RtInfectionsRenewalModel)
- I0_rv
- Initial infections (passed to RtInfectionsRenewalModel)
- Rt_process_rv
- Rt process (passed to RtInfectionsRenewalModel).
- hosp_admission_obs_process_rv
- Observation process for the hospital admissions.
-
- Returns
- -------
- None
- """
- self.basic_renewal = RtInfectionsRenewalModel(
- gen_int_rv=gen_int_rv,
- I0_rv=I0_rv,
- latent_infections_rv=latent_infections_rv,
- infection_obs_process_rv=None, # why is this None?
- Rt_process_rv=Rt_process_rv,
- )
-
- HospitalAdmissionsModel.validate(
- latent_hosp_admissions_rv, hosp_admission_obs_process_rv
- )
-
- self.latent_hosp_admissions_rv = latent_hosp_admissions_rv
- if hosp_admission_obs_process_rv is None:
- hosp_admission_obs_process_rv = NullObservation()
-
- self.hosp_admission_obs_process_rv = hosp_admission_obs_process_rv
-
- @staticmethod
- def validate(latent_hosp_admissions_rv, hosp_admission_obs_process_rv) -> None:
- """
- Verifies types and status (RV) of latent and observed hospital admissions
-
- Parameters
- ----------
- latent_hosp_admissions_rv
- The latent process for the hospital admissions.
- hosp_admission_obs_process_rv
- The observed hospital admissions.
-
- Returns
- -------
- None
- """
- assert isinstance(latent_hosp_admissions_rv, RandomVariable)
- if hosp_admission_obs_process_rv is not None:
- assert isinstance(hosp_admission_obs_process_rv, RandomVariable)
-
- return None
-
- def sample(
- self,
- n_datapoints: int | None = None,
- data_observed_hosp_admissions: ArrayLike | None = None,
- padding: int = 0,
- **kwargs,
- ) -> HospModelSample:
- """
- Sample from the HospitalAdmissions model
-
- Parameters
- ----------
- n_datapoints
- Number of timepoints to sample (passed to the basic renewal model).
- data_observed_hosp_admissions
- The observed hospitalization data (passed to the basic renewal
- model). Defaults to None (simulation, rather than fit).
- padding
- Number of padding timepoints to add to the beginning of the
- simulation. Defaults to 0.
- **kwargs
- Additional keyword arguments passed through to internal sample()
- calls, should there be any.
-
- Returns
- -------
- HospModelSample
-
- See Also
- --------
- basic_renewal.sample : For sampling the basic renewal model
- sample_observed_admissions : For sampling observed hospital admissions
- """
- if n_datapoints is None and data_observed_hosp_admissions is None:
- raise ValueError(
- "Either n_datapoints or data_observed_hosp_admissions must be passed."
- )
- elif n_datapoints is not None and data_observed_hosp_admissions is not None:
- raise ValueError(
- "Cannot pass both n_datapoints and data_observed_hosp_admissions."
- )
- elif n_datapoints is None:
- n_datapoints = len(data_observed_hosp_admissions)
- else:
- n_datapoints = n_datapoints
-
- # Getting the initial quantities from the basic model
- basic_model = self.basic_renewal.sample(
- n_datapoints=n_datapoints,
- data_observed_infections=None,
- padding=padding,
- **kwargs,
- )
- # Sampling the latent hospital admissions
- (
- infection_hosp_rate,
- latent_hosp_admissions,
- *_,
- ) = self.latent_hosp_admissions_rv(
- latent_infections=basic_model.latent_infections,
- **kwargs,
- )
- observed_hosp_admissions = self.hosp_admission_obs_process_rv(
- mu=latent_hosp_admissions[-n_datapoints:],
- obs=data_observed_hosp_admissions,
- **kwargs,
- )
-
- return HospModelSample(
- Rt=basic_model.Rt,
- latent_infections=basic_model.latent_infections,
- infection_hosp_rate=infection_hosp_rate,
- latent_hosp_admissions=latent_hosp_admissions,
- observed_hosp_admissions=observed_hosp_admissions,
- )
diff --git a/pyrenew/model/rtinfectionsrenewalmodel.py b/pyrenew/model/rtinfectionsrenewalmodel.py
deleted file mode 100644
index 1f1be0c4..00000000
--- a/pyrenew/model/rtinfectionsrenewalmodel.py
+++ /dev/null
@@ -1,218 +0,0 @@
-# numpydoc ignore=GL08
-
-from __future__ import annotations
-
-from typing import NamedTuple
-
-import jax.numpy as jnp
-import numpyro
-from numpy.typing import ArrayLike
-
-from pyrenew.deterministic import NullObservation
-from pyrenew.metaclass import Model, RandomVariable
-
-
-# Output class of the RtInfectionsRenewalModel
-class RtInfectionsRenewalSample(NamedTuple):
- """
- A container for holding the output from `model.RtInfectionsRenewalModel.sample()`.
-
- Attributes
- ----------
- Rt : ArrayLike | None, optional
- The reproduction number over time. Defaults to None.
- latent_infections : ArrayLike | None, optional
- The estimated latent infections. Defaults to None.
- observed_infections : ArrayLike | None, optional
- The sampled infections. Defaults to None.
- """
-
- Rt: ArrayLike | None = None
- latent_infections: ArrayLike | None = None
- observed_infections: ArrayLike | None = None
-
- def __repr__(self):
- return (
- f"RtInfectionsRenewalSample("
- f"Rt={self.Rt}, "
- f"latent_infections={self.latent_infections}, "
- f"observed_infections={self.observed_infections})"
- )
-
-
-class RtInfectionsRenewalModel(Model):
- """
- Basic Renewal Model (Infections + Rt)
-
- The basic renewal model consists of a sampler of two steps: Sample from
- Rt and then used that to sample the infections.
- """
-
- def __init__(
- self,
- latent_infections_rv: RandomVariable,
- gen_int_rv: RandomVariable,
- I0_rv: RandomVariable,
- Rt_process_rv: RandomVariable,
- infection_obs_process_rv: RandomVariable = None,
- ) -> None:
- """
- Default constructor
-
- Parameters
- ----------
- latent_infections_rv
- Infections latent process (e.g.,
- pyrenew.latent.Infections.).
- gen_int_rv
- The generation interval.
- I0_rv
- The initial infections.
- Rt_process_rv
- The sample function of the process should return a tuple where the
- first element is the drawn Rt.
- infection_obs_process_rv
- Infections observation process (e.g.,
- pyrenew.observations.Poisson.).
-
- Returns
- -------
- None
- """
-
- if infection_obs_process_rv is None:
- infection_obs_process_rv = NullObservation()
-
- RtInfectionsRenewalModel.validate(
- gen_int_rv=gen_int_rv,
- I0_rv=I0_rv,
- latent_infections_rv=latent_infections_rv,
- infection_obs_process_rv=infection_obs_process_rv,
- Rt_process_rv=Rt_process_rv,
- )
-
- self.gen_int_rv = gen_int_rv
- self.I0_rv = I0_rv
- self.latent_infections_rv = latent_infections_rv
- self.infection_obs_process_rv = infection_obs_process_rv
- self.Rt_process_rv = Rt_process_rv
-
- @staticmethod
- def validate(
- gen_int_rv: any,
- I0_rv: any,
- latent_infections_rv: any,
- infection_obs_process_rv: any,
- Rt_process_rv: any,
- ) -> None:
- """
- Verifies types and status (RV) of the generation interval, initial
- infections, latent and observed infections, and the Rt process.
-
- Parameters
- ----------
- gen_int_rv
- The generation interval. Expects RandomVariable.
- I0_rv
- The initial infections. Expects RandomVariable.
- latent_infections_rv
- Infections latent process. Expects RandomVariable.
- infection_obs_process_rv
- Infections observation process. Expects RandomVariable.
- Rt_process_rv
- The sample function of the process should return a tuple where the
- first element is the drawn Rt. Expects RandomVariable.
-
- Returns
- -------
- None
- """
- assert isinstance(gen_int_rv, RandomVariable)
- assert isinstance(I0_rv, RandomVariable)
- assert isinstance(latent_infections_rv, RandomVariable)
- assert isinstance(infection_obs_process_rv, RandomVariable)
- assert isinstance(Rt_process_rv, RandomVariable)
- return None
-
- def sample(
- self,
- n_datapoints: int | None = None,
- data_observed_infections: ArrayLike | None = None,
- padding: int = 0,
- **kwargs,
- ) -> RtInfectionsRenewalSample:
- """
- Sample from the Basic Renewal Model
-
- Parameters
- ----------
- n_datapoints
- Number of timepoints to sample.
- data_observed_infections
- Observed infections. Defaults to None.
- padding
- Number of padding timepoints to add to the beginning of the
- simulation. Defaults to 0.
- **kwargs
- Additional keyword arguments passed through to internal sample()
- calls, if any
-
- Notes
- -----
- Either `data_observed_infections` or `n_datapoints`
- must be specified, not both.
-
- Returns
- -------
- RtInfectionsRenewalSample
- """
-
- if n_datapoints is None and data_observed_infections is None:
- raise ValueError(
- "Either n_datapoints or data_observed_infections must be passed."
- )
- elif n_datapoints is not None and data_observed_infections is not None:
- raise ValueError(
- "Cannot pass both n_datapoints and data_observed_infections."
- )
- elif n_datapoints is None:
- n_timepoints = len(data_observed_infections) + padding
- else:
- n_timepoints = n_datapoints + padding
- # Sampling from Rt (possibly with a given Rt, depending on
- # the Rt_process (RandomVariable) object.)
-
- Rt = self.Rt_process_rv(
- n=n_timepoints,
- **kwargs,
- )
-
- # Getting the generation interval
- gen_int = self.gen_int_rv(**kwargs)
-
- # Sampling initial infections
- I0 = self.I0_rv(**kwargs)
-
- # Sampling from the latent process
- post_initialization_latent_infections = self.latent_infections_rv(
- Rt=Rt,
- gen_int=gen_int,
- I0=I0,
- **kwargs,
- ).post_initialization_infections
- observed_infections = self.infection_obs_process_rv(
- mu=post_initialization_latent_infections[padding:],
- obs=data_observed_infections,
- **kwargs,
- )
-
- all_latent_infections = jnp.hstack([I0, post_initialization_latent_infections])
- numpyro.deterministic("all_latent_infections", all_latent_infections)
-
- numpyro.deterministic("Rt", Rt)
-
- return RtInfectionsRenewalSample(
- Rt=Rt,
- latent_infections=all_latent_infections,
- observed_infections=observed_infections,
- )
diff --git a/pyrenew/observation/__init__.py b/pyrenew/observation/__init__.py
index 84224e16..50616364 100644
--- a/pyrenew/observation/__init__.py
+++ b/pyrenew/observation/__init__.py
@@ -30,13 +30,11 @@
PoissonNoise,
VectorizedRV,
)
-from pyrenew.observation.poisson import PoissonObservation
from pyrenew.observation.types import ObservationSample
__all__ = [
# Existing (kept for backward compatibility)
"NegativeBinomialObservation",
- "PoissonObservation",
# Base classes and types
"BaseObservationProcess",
"ObservationSample",
diff --git a/pyrenew/observation/poisson.py b/pyrenew/observation/poisson.py
deleted file mode 100644
index a29f6408..00000000
--- a/pyrenew/observation/poisson.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# numpydoc ignore=GL08
-
-from __future__ import annotations
-
-import numpyro
-import numpyro.distributions as dist
-from jax.typing import ArrayLike
-
-from pyrenew.metaclass import RandomVariable
-
-
-class PoissonObservation(RandomVariable):
- """
- Poisson observation process
- """
-
- def __init__(
- self,
- name: str,
- eps: float = 1e-8,
- ) -> None:
- """
- Default Constructor
-
- Parameters
- ----------
- name
- Passed to [`numpyro.primitives.sample`][].
- eps
- Small value added to the rate parameter to avoid zero values.
- Defaults to 1e-8.
-
- Returns
- -------
- None
- """
-
- self.name = name
- self.eps = eps
-
- return None
-
- @staticmethod
- def validate(): # numpydoc ignore=GL08
- None
-
- def sample(
- self,
- mu: ArrayLike,
- obs: ArrayLike | None = None,
- **kwargs,
- ) -> ArrayLike:
- """
- Sample from the Poisson process
-
- Parameters
- ----------
- mu
- Rate parameter of the Poisson distribution.
- obs
- Observed data. Defaults to None.
- **kwargs
- Additional keyword arguments passed through to internal sample calls, should there be any.
-
- Returns
- -------
- ArrayLike
- """
-
- poisson_sample = numpyro.sample(
- name=self.name,
- fn=dist.Poisson(rate=mu + self.eps),
- obs=obs,
- )
- return poisson_sample
diff --git a/test/test_deterministic.py b/test/test_deterministic.py
index 2df30227..5c8b0e76 100644
--- a/test/test_deterministic.py
+++ b/test/test_deterministic.py
@@ -10,6 +10,7 @@
from pyrenew.deterministic import (
DeterministicPMF,
DeterministicVariable,
+ NullObservation,
NullVariable,
)
@@ -89,3 +90,21 @@ def test_deterministic_validation():
for arraylike_val in some_array_likes:
DeterministicVariable.validate(arraylike_val)
DeterministicVariable(value=arraylike_val, name="valid_variable")
+
+
+def test_null_observation():
+ """
+ Test that NullObservation can be constructed and
+ that its sample method returns None.
+ """
+ null_obs = NullObservation()
+ assert null_obs.sample(mu=jnp.array([1.0, 2.0]), obs=None) is None
+ assert null_obs.sample(mu=jnp.array([1.0]), obs=jnp.array([1.0])) is None
+
+
+def test_deterministic_pmf_size():
+ """
+ Test that DeterministicPMF.size() returns the correct size.
+ """
+ pmf = DeterministicPMF(name="test_pmf", value=jnp.array([0.25, 0.25, 0.2, 0.3]))
+ assert pmf.size() == 4
diff --git a/test/test_forecast.py b/test/test_forecast.py
deleted file mode 100644
index d58462b1..00000000
--- a/test/test_forecast.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-import jax.random as jr
-import numpyro
-import numpyro.distributions as dist
-from numpy.testing import assert_array_equal
-
-from pyrenew.deterministic import DeterministicPMF
-from pyrenew.latent import (
- InfectionInitializationProcess,
- Infections,
- InitializeInfectionsZeroPad,
-)
-from pyrenew.model import RtInfectionsRenewalModel
-from pyrenew.observation import PoissonObservation
-from pyrenew.randomvariable import DistributionalVariable
-from test.utils import SimpleRt
-
-
-def test_forecast():
- """
- Check that forecasts are the right length and
- match the posterior up until forecast begins.
- """
- pmf_array = jnp.array([0.25, 0.25, 0.25, 0.25])
- gen_int = DeterministicPMF(name="gen_int", value=pmf_array)
- I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
- latent_infections = Infections()
- observed_infections = PoissonObservation(name="poisson_rv")
- rt = SimpleRt()
-
- model = RtInfectionsRenewalModel(
- I0_rv=I0,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
- )
-
- n_datapoints = 30
- n_forecast_points = 10
- with numpyro.handlers.seed(rng_seed=223):
- model_sample = model.sample(n_datapoints=n_datapoints)
-
- model.run(
- num_warmup=5,
- num_samples=5,
- data_observed_infections=model_sample.observed_infections,
- rng_key=jr.key(54),
- )
-
- posterior_predictive_samples = model.posterior_predictive(
- n_datapoints=n_datapoints + n_forecast_points,
- )
-
- # Check the length of the predictive distribution
- assert (
- len(posterior_predictive_samples["poisson_rv"][0])
- == n_datapoints + n_forecast_points
- )
-
- # Check the first elements of the posterior predictive Rt
- # are the same as the posterior Rt
- assert_array_equal(
- model.mcmc.get_samples()["Rt"][0],
- posterior_predictive_samples["Rt"][0][: len(model.mcmc.get_samples()["Rt"][0])],
- )
diff --git a/test/test_latent_admissions.py b/test/test_latent_admissions.py
deleted file mode 100644
index 053e1545..00000000
--- a/test/test_latent_admissions.py
+++ /dev/null
@@ -1,130 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-import numpy.testing as testing
-import numpyro
-import numpyro.distributions as dist
-
-from pyrenew.deterministic import DeterministicPMF, DeterministicVariable
-from pyrenew.latent import HospitalAdmissions, Infections
-from pyrenew.randomvariable import DistributionalVariable
-from test.utils import SimpleRt
-
-
-def test_admissions_sample():
- """
- Check that a HospitalAdmissions latent process
- can be initialized and sampled from.
- """
-
- # Generating Rt and Infections to compute the hospital admissions
-
- rt = SimpleRt()
- n_steps = 30
-
- with numpyro.handlers.seed(rng_seed=223):
- sim_rt = rt(n=n_steps)
-
- gen_int = jnp.array([0.5, 0.1, 0.1, 0.2, 0.1])
- inf_hosp_int_array = jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ]
- )
- i0 = 10 * jnp.ones_like(inf_hosp_int_array)
- inf1 = Infections()
-
- with numpyro.handlers.seed(rng_seed=223):
- inf_sampled1 = inf1(Rt=sim_rt, gen_int=gen_int, I0=i0)
-
- # Testing the hospital admissions
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=inf_hosp_int_array,
- )
-
- hosp1 = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR", distribution=dist.LogNormal(jnp.log(0.05), 0.05)
- ),
- )
-
- with numpyro.handlers.seed(rng_seed=223):
- sim_hosp_1 = hosp1(
- latent_infections=jnp.hstack(
- [i0, inf_sampled1.post_initialization_infections]
- )
- )
-
- testing.assert_array_less(
- sim_hosp_1.latent_hospital_admissions[-n_steps:],
- inf_sampled1.post_initialization_infections,
- )
- inf_hosp2 = jnp.ones(30)
- inf_hosp2 = DeterministicPMF("i2h", inf_hosp2 / sum(inf_hosp2))
-
- dow_effect = jnp.array([1, 1, 1, 1, 0.5, 0.5, 0.5])
- dow_effect = DeterministicPMF(
- name="dow_effect",
- value=dow_effect / sum(dow_effect),
- )
-
- dow_effect_wrong = DeterministicPMF(
- name="dow_effect",
- value=jnp.array([0.3, 0.3, 1 - 0.6]),
- )
- hosp2a = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp2,
- infection_hospitalization_ratio_rv=DeterministicVariable("ihr", 1),
- day_of_week_effect_rv=dow_effect,
- obs_data_first_day_of_the_week=0,
- )
-
- hosp2b = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp2,
- infection_hospitalization_ratio_rv=DeterministicVariable("ihr", 1),
- day_of_week_effect_rv=dow_effect,
- obs_data_first_day_of_the_week=2,
- )
-
- hosp3b = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp2,
- infection_hospitalization_ratio_rv=DeterministicVariable("ihr", 1),
- day_of_week_effect_rv=dow_effect_wrong,
- obs_data_first_day_of_the_week=2,
- )
-
- inf_sampled2 = jnp.ones(30)
-
- with numpyro.handlers.seed(rng_seed=223):
- sim_hosp_2a = hosp2a(latent_infections=inf_sampled2).multiplier
-
- with numpyro.handlers.seed(rng_seed=223):
- sim_hosp_2b = hosp2b(latent_infections=inf_sampled2).multiplier
-
- with numpyro.handlers.seed(rng_seed=223):
- with testing.assert_raises(ValueError):
- hosp3b(latent_infections=inf_sampled2).multiplier
-
- testing.assert_array_equal(
- sim_hosp_2a[2 : (sim_hosp_2b.size - 2)],
- sim_hosp_2b[: (sim_hosp_2b.size - 4)],
- )
diff --git a/test/test_model_basic_renewal.py b/test/test_model_basic_renewal.py
deleted file mode 100644
index 973acf61..00000000
--- a/test/test_model_basic_renewal.py
+++ /dev/null
@@ -1,242 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-import jax.random as jr
-import numpy as np
-import numpyro
-import numpyro.distributions as dist
-import pytest
-
-from pyrenew.deterministic import DeterministicPMF, NullObservation
-from pyrenew.latent import (
- InfectionInitializationProcess,
- Infections,
- InitializeInfectionsZeroPad,
-)
-from pyrenew.model import RtInfectionsRenewalModel
-from pyrenew.observation import PoissonObservation
-from pyrenew.randomvariable import DistributionalVariable
-from test.utils import SimpleRt
-
-
-def test_model_basicrenewal_no_timepoints_or_observations():
- """
- Test that the basic renewal model does not run
- without either n_datapoints or
- observed_admissions
- """
-
- gen_int = DeterministicPMF(
- name="gen_int", value=jnp.array([0.25, 0.25, 0.25, 0.25])
- )
-
- I0_init_rv = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
-
- latent_infections = Infections()
-
- observed_infections = PoissonObservation("poisson_rv")
-
- rt = SimpleRt()
-
- model1 = RtInfectionsRenewalModel(
- I0_rv=I0_init_rv,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
- )
-
- with numpyro.handlers.seed(rng_seed=223):
- with pytest.raises(ValueError, match="Either"):
- model1.sample(n_datapoints=None, data_observed_infections=None)
-
-
-def test_model_basicrenewal_both_timepoints_and_observations():
- """
- Test that the basic renewal model does not run with both n_datapoints and observed_admissions passed
- """
-
- gen_int = DeterministicPMF(
- name="gen_int",
- value=jnp.array([0.25, 0.25, 0.25, 0.25]),
- )
-
- I0_init_rv = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
-
- latent_infections = Infections()
-
- observed_infections = PoissonObservation("possion_rv")
-
- rt = SimpleRt()
-
- model1 = RtInfectionsRenewalModel(
- I0_rv=I0_init_rv,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
- )
-
- with numpyro.handlers.seed(rng_seed=223):
- with pytest.raises(ValueError, match="Cannot pass both"):
- model1.sample(
- n_datapoints=30,
- data_observed_infections=jnp.repeat(jnp.nan, 30),
- )
-
-
-def test_model_basicrenewal_no_obs_model():
- """
- Test the basic semi-deterministic renewal model runs. Semi-deterministic
- from the perspective of the infections. It returns expected, not sampled.
- """
-
- gen_int = DeterministicPMF(
- name="gen_int",
- value=jnp.array([0.25, 0.25, 0.25, 0.25]),
- )
-
- with pytest.raises(ValueError):
- _ = DistributionalVariable(name="I0", distribution=1)
-
- I0_init_rv = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
-
- latent_infections = Infections()
-
- rt = SimpleRt()
-
- model0 = RtInfectionsRenewalModel(
- gen_int_rv=gen_int,
- I0_rv=I0_init_rv,
- latent_infections_rv=latent_infections,
- Rt_process_rv=rt,
- # Explicitly use None, this should call the NullObservation
- infection_obs_process_rv=None,
- )
-
- # Sampling and fitting model 0 (with no obs for infections)
- with numpyro.handlers.seed(rng_seed=223):
- model0_samp = model0.sample(n_datapoints=30)
- model0_samp.Rt
- model0_samp.latent_infections
- model0_samp.observed_infections
-
- # Generating
- model0.infection_obs_process_rv = NullObservation()
- with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model0.sample(n_datapoints=30)
-
- np.testing.assert_array_equal(model0_samp.Rt, model1_samp.Rt)
- np.testing.assert_array_equal(
- model0_samp.latent_infections,
- model1_samp.latent_infections,
- )
- np.testing.assert_array_equal(
- model0_samp.observed_infections,
- model1_samp.observed_infections,
- )
-
- model0.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(272),
- data_observed_infections=model0_samp.latent_infections,
- )
-
-
-def test_model_basicrenewal_with_obs_model():
- """
- Test the basic random renewal model runs. Random
- from the perspective of the infections. It returns sampled, not expected.
- """
-
- gen_int = DeterministicPMF(
- name="gen_int", value=jnp.array([0.25, 0.25, 0.25, 0.25])
- )
-
- I0_init_rv = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
-
- latent_infections = Infections()
-
- observed_infections = PoissonObservation("poisson_rv")
-
- rt = SimpleRt()
-
- model1 = RtInfectionsRenewalModel(
- I0_rv=I0_init_rv,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
- )
-
- # Sampling and fitting model 1 (with obs infections)
- with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model1.sample(n_datapoints=30)
-
- print(model1_samp)
- print(model1_samp.Rt.size)
- print(model1_samp.latent_infections.size)
- print(model1_samp.observed_infections.size)
-
- model1.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(22),
- data_observed_infections=model1_samp.observed_infections,
- )
-
-
-def test_model_basicrenewal_padding() -> None: # numpydoc ignore=GL08
- gen_int = DeterministicPMF(
- name="gen_int", value=jnp.array([0.25, 0.25, 0.25, 0.25])
- )
-
- I0_init_rv = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
-
- latent_infections = Infections()
-
- observed_infections = PoissonObservation("poisson_rv")
-
- rt = SimpleRt()
-
- model1 = RtInfectionsRenewalModel(
- I0_rv=I0_init_rv,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
- )
-
- pad_size = 5
-
- with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model1.sample(n_datapoints=30, padding=pad_size)
-
- model1.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(22),
- data_observed_infections=model1_samp.observed_infections,
- padding=pad_size,
- )
diff --git a/test/test_model_hosp_admissions.py b/test/test_model_hosp_admissions.py
deleted file mode 100644
index be041cee..00000000
--- a/test/test_model_hosp_admissions.py
+++ /dev/null
@@ -1,523 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-import jax.random as jr
-import numpy as np
-import numpyro
-import numpyro.distributions as dist
-import pytest
-
-from pyrenew.deterministic import (
- DeterministicPMF,
- DeterministicVariable,
- NullObservation,
-)
-from pyrenew.latent import (
- HospitalAdmissions,
- InfectionInitializationProcess,
- Infections,
- InitializeInfectionsZeroPad,
-)
-from pyrenew.model import HospitalAdmissionsModel
-from pyrenew.observation import PoissonObservation
-from pyrenew.randomvariable import DistributionalVariable
-from test.utils import SimpleRt
-
-
-def test_model_hosp_no_timepoints_or_observations():
- """
- Checks that the hospital admissions model does not run
- without either n_datapoints or observed_admissions
- """
-
- gen_int = DeterministicPMF(
- name="gen_int", value=jnp.array([0.25, 0.25, 0.25, 0.25])
- )
-
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ],
- ),
- )
-
- I0 = DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1))
-
- latent_infections = Infections()
- Rt_process = SimpleRt()
-
- observed_admissions = PoissonObservation("poisson_rv")
-
- latent_admissions = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR", distribution=dist.LogNormal(jnp.log(0.05), 0.05)
- ),
- )
-
- model1 = HospitalAdmissionsModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- Rt_process_rv=Rt_process,
- latent_infections_rv=latent_infections,
- latent_hosp_admissions_rv=latent_admissions,
- hosp_admission_obs_process_rv=observed_admissions,
- )
-
- with numpyro.handlers.seed(rng_seed=223):
- with pytest.raises(ValueError, match="Either"):
- model1.sample(n_datapoints=None, data_observed_admissions=None)
-
-
-def test_model_hosp_both_timepoints_and_observations():
- """
- Checks that the hospital admissions model does not run with
- both n_datapoints and observed_admissions passed
- """
-
- gen_int = DeterministicPMF(
- name="gen_int",
- value=jnp.array([0.25, 0.25, 0.25, 0.25]),
- )
-
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ],
- ),
- )
-
- I0 = DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1))
-
- latent_infections = Infections()
- Rt_process = SimpleRt()
- observed_admissions = PoissonObservation("poisson_rv")
-
- latent_admissions = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR", distribution=dist.LogNormal(jnp.log(0.05), 0.05)
- ),
- )
-
- model1 = HospitalAdmissionsModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- Rt_process_rv=Rt_process,
- latent_infections_rv=latent_infections,
- latent_hosp_admissions_rv=latent_admissions,
- hosp_admission_obs_process_rv=observed_admissions,
- )
-
- with numpyro.handlers.seed(rng_seed=223):
- with pytest.raises(ValueError, match="Cannot pass both"):
- model1.sample(
- n_datapoints=30,
- data_observed_hosp_admissions=jnp.repeat(jnp.nan, 30),
- )
-
-
-def test_model_hosp_no_obs_model():
- """
- Checks that the partially deterministic
- Hospitalization model runs
- """
-
- gen_int = DeterministicPMF(
- name="gen_int",
- value=jnp.array([0.25, 0.25, 0.25, 0.25]),
- )
-
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ]
- ),
- )
- n_initialization_points = max(gen_int.size(), inf_hosp.size()) - 1
-
- I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=n_initialization_points),
- )
-
- latent_infections = Infections()
- Rt_process = SimpleRt()
-
- latent_admissions = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR",
- distribution=dist.LogNormal(jnp.log(0.05), 0.05),
- ),
- )
-
- model0 = HospitalAdmissionsModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- Rt_process_rv=Rt_process,
- latent_infections_rv=latent_infections,
- latent_hosp_admissions_rv=latent_admissions,
- hosp_admission_obs_process_rv=None,
- )
-
- # Sampling and fitting model 0 (with no obs for infections)
- with numpyro.handlers.seed(rng_seed=223):
- model0_samp = model0.sample(n_datapoints=30)
-
- model0.hosp_admission_obs_process_rv = NullObservation()
-
- with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model0.sample(n_datapoints=30)
-
- np.testing.assert_array_almost_equal(model0_samp.Rt, model1_samp.Rt)
- np.testing.assert_array_equal(
- model0_samp.latent_infections,
- model1_samp.latent_infections,
- )
- np.testing.assert_array_equal(
- model0_samp.infection_hosp_rate,
- model1_samp.infection_hosp_rate,
- )
- np.testing.assert_array_equal(
- model0_samp.latent_hosp_admissions,
- model1_samp.latent_hosp_admissions,
- )
-
- # These are supposed to be none, both
- assert model0_samp.observed_hosp_admissions is None
- assert model1_samp.observed_hosp_admissions is None
-
- model0.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(272),
- data_observed_hosp_admissions=model0_samp.latent_hosp_admissions,
- )
-
-
-def test_model_hosp_with_obs_model():
- """
- Checks that the random Hospitalization model runs
- """
-
- gen_int = DeterministicPMF(
- name="gen_int", value=jnp.array([0.25, 0.25, 0.25, 0.25])
- )
-
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ]
- ),
- )
-
- n_initialization_points = max(gen_int.size(), inf_hosp.size()) - 1
-
- I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=n_initialization_points),
- )
-
- latent_infections = Infections()
- Rt_process = SimpleRt()
- observed_admissions = PoissonObservation("poisson_rv")
-
- latent_admissions = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR",
- distribution=dist.LogNormal(jnp.log(0.05), 0.05),
- ),
- )
-
- model1 = HospitalAdmissionsModel(
- gen_int_rv=gen_int,
- I0_rv=I0,
- Rt_process_rv=Rt_process,
- latent_infections_rv=latent_infections,
- latent_hosp_admissions_rv=latent_admissions,
- hosp_admission_obs_process_rv=observed_admissions,
- )
-
- # Sampling and fitting model 0 (with no obs for infections)
- with numpyro.handlers.seed(rng_seed=233):
- model1_samp = model1.sample(n_datapoints=30)
-
- model1.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(272),
- data_observed_hosp_admissions=model1_samp.observed_hosp_admissions,
- )
-
-
-def test_model_hosp_with_obs_model_weekday_phosp_2():
- """
- Checks that the random Hospitalization model runs
- """
-
- gen_int = DeterministicPMF(
- name="gen_int",
- value=jnp.array([0.25, 0.25, 0.25, 0.25]),
- )
-
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ]
- ),
- )
-
- n_initialization_points = max(gen_int.size(), inf_hosp.size()) - 1
-
- I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=n_initialization_points),
- )
-
- latent_infections = Infections()
- Rt_process = SimpleRt()
- observed_admissions = PoissonObservation("poisson_rv")
-
- hosp_report_prob_dist = DistributionalVariable(
- "hosp_report_prob_dist", dist.Uniform()
- )
- weekday = DistributionalVariable("weekday", dist.Uniform()).expand_by((7,))
-
- latent_admissions = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- day_of_week_effect_rv=weekday,
- hospitalization_reporting_ratio_rv=hosp_report_prob_dist,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR", distribution=dist.LogNormal(jnp.log(0.05), 0.05)
- ),
- )
-
- model1 = HospitalAdmissionsModel(
- I0_rv=I0,
- gen_int_rv=gen_int,
- Rt_process_rv=Rt_process,
- latent_infections_rv=latent_infections,
- latent_hosp_admissions_rv=latent_admissions,
- hosp_admission_obs_process_rv=observed_admissions,
- )
-
- # Sampling and fitting model 0 (with no obs for admissions)
- with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model1.sample(n_datapoints=30)
-
- model1.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(272),
- data_observed_hosp_admissions=model1_samp.observed_hosp_admissions,
- )
-
-
-def test_model_hosp_with_obs_model_weekday_phosp():
- """
- Checks that the random Hospitalization model runs
- """
-
- gen_int = DeterministicPMF(
- name="gen_int",
- value=jnp.array([0.25, 0.25, 0.25, 0.25]),
- )
- n_obs_to_generate = 30
- pad_size = 5
-
- inf_hosp = DeterministicPMF(
- name="inf_hosp",
- value=jnp.array(
- [
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0,
- 0.25,
- 0.5,
- 0.1,
- 0.1,
- 0.05,
- ]
- ),
- )
-
- n_initialization_points = max(gen_int.size(), inf_hosp.size())
-
- I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=n_initialization_points),
- )
-
- latent_infections = Infections()
- Rt_process = SimpleRt()
-
- observed_admissions = PoissonObservation("poisson_rv")
-
- # Other random components
- total_length = n_obs_to_generate + pad_size
- total_length = n_obs_to_generate + pad_size + 1 # gen_int.size()
- weekday = jnp.array([1, 1, 1, 1, 2, 2, 2])
- weekday = weekday / weekday.sum()
-
- weekday = DeterministicVariable(name="weekday", value=weekday)
-
- hosp_report_prob_dist = jnp.array([0.9, 0.8, 0.7, 0.7, 0.6, 0.4])
- hosp_report_prob_dist = jnp.tile(hosp_report_prob_dist, 10)
- hosp_report_prob_dist = hosp_report_prob_dist[:total_length]
- hosp_report_prob_dist = hosp_report_prob_dist / hosp_report_prob_dist.sum()
-
- hosp_report_prob_dist = DeterministicVariable(
- name="hosp_report_prob_dist",
- value=hosp_report_prob_dist,
- )
-
- latent_admissions = HospitalAdmissions(
- infection_to_admission_interval_rv=inf_hosp,
- day_of_week_effect_rv=weekday,
- hospitalization_reporting_ratio_rv=hosp_report_prob_dist,
- infection_hospitalization_ratio_rv=DistributionalVariable(
- name="IHR",
- distribution=dist.LogNormal(jnp.log(0.05), 0.05),
- ),
- )
-
- model1 = HospitalAdmissionsModel(
- I0_rv=I0,
- gen_int_rv=gen_int,
- Rt_process_rv=Rt_process,
- latent_infections_rv=latent_infections,
- latent_hosp_admissions_rv=latent_admissions,
- hosp_admission_obs_process_rv=observed_admissions,
- )
-
- # Sampling and fitting model 0 (with no obs for infections)
- with numpyro.handlers.seed(rng_seed=223):
- model1_samp = model1.sample(n_datapoints=n_obs_to_generate, padding=pad_size)
-
- # Showed during merge conflict, but unsure if it will be needed
- # pad_size = 5
- # obs = jnp.hstack(
- # [
- # jnp.repeat(jnp.nan, pad_size),
- # model1_samp.observed_hosp_admissions[pad_size:],
- # ]
- # )
- # Running with padding
- model1.run(
- num_warmup=500,
- num_samples=500,
- rng_key=jr.key(272),
- data_observed_hosp_admissions=model1_samp.observed_hosp_admissions,
- padding=pad_size,
- )
diff --git a/test/test_observation_poisson.py b/test/test_observation_poisson.py
deleted file mode 100644
index 8b9c0716..00000000
--- a/test/test_observation_poisson.py
+++ /dev/null
@@ -1,29 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-import numpy as np
-import numpy.testing as testing
-import numpyro
-
-from pyrenew.observation import PoissonObservation
-
-
-def test_poisson_obs():
- """
- Check that an PoissonObservation can sample
- """
-
- pois = PoissonObservation("rv")
-
- rates = np.random.randint(1, 5, size=10)
- with numpyro.handlers.seed(rng_seed=223):
- sim_pois = pois(mu=rates)
-
- testing.assert_array_equal(sim_pois, jnp.ceil(sim_pois))
-
-
-def test_poisson_validate():
- """
- Check that PoissonObservation.validate() runs without error.
- """
- PoissonObservation.validate()
diff --git a/test/test_predictive.py b/test/test_predictive.py
deleted file mode 100644
index 4619a3aa..00000000
--- a/test/test_predictive.py
+++ /dev/null
@@ -1,47 +0,0 @@
-"""
-Ensures that posterior predictive samples are not generated
-when no posterior samples are available.
-"""
-
-import jax.numpy as jnp
-import numpyro.distributions as dist
-import pytest
-
-from pyrenew.deterministic import DeterministicPMF
-from pyrenew.latent import (
- InfectionInitializationProcess,
- Infections,
- InitializeInfectionsZeroPad,
-)
-from pyrenew.model import RtInfectionsRenewalModel
-from pyrenew.observation import PoissonObservation
-from pyrenew.randomvariable import DistributionalVariable
-from test.utils import SimpleRt
-
-pmf_array = jnp.array([0.25, 0.1, 0.2, 0.45])
-gen_int = DeterministicPMF(name="gen_int", value=pmf_array)
-I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
-)
-latent_infections = Infections()
-observed_infections = PoissonObservation("poisson_rv")
-rt = SimpleRt()
-
-model = RtInfectionsRenewalModel(
- I0_rv=I0,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
-)
-
-
-def test_posterior_predictive_no_posterior():
- """
- Tests that posterior predictive samples are not generated when
- no posterior samples are available.
- """
- with pytest.raises(ValueError, match="No posterior"):
- model.posterior_predictive(n_datapoints=10)
diff --git a/test/test_random_key.py b/test/test_random_key.py
deleted file mode 100644
index a7a26bfd..00000000
--- a/test/test_random_key.py
+++ /dev/null
@@ -1,191 +0,0 @@
-"""
-Ensures that models created with the same or
-with different random keys behave appropriately.
-"""
-
-import jax.numpy as jnp
-import jax.random as jr
-import numpyro
-import numpyro.distributions as dist
-from numpy.testing import assert_array_equal, assert_raises
-
-from pyrenew.deterministic import DeterministicPMF
-from pyrenew.latent import (
- InfectionInitializationProcess,
- Infections,
- InitializeInfectionsZeroPad,
-)
-from pyrenew.model import RtInfectionsRenewalModel
-from pyrenew.observation import PoissonObservation
-from pyrenew.randomvariable import DistributionalVariable
-from test.utils import SimpleRt
-
-
-def create_test_model(): # numpydoc ignore=GL08
- pmf_array = jnp.array([0.25, 0.25, 0.25, 0.25])
- gen_int = DeterministicPMF(name="gen_int", value=pmf_array)
- I0 = InfectionInitializationProcess(
- "I0_initialization",
- DistributionalVariable(name="I0", distribution=dist.LogNormal(0, 1)),
- InitializeInfectionsZeroPad(n_timepoints=gen_int.size()),
- )
- latent_infections = Infections()
- observed_infections = PoissonObservation("poisson_rv")
- rt = SimpleRt()
- model = RtInfectionsRenewalModel(
- I0_rv=I0,
- gen_int_rv=gen_int,
- latent_infections_rv=latent_infections,
- infection_obs_process_rv=observed_infections,
- Rt_process_rv=rt,
- )
- return model
-
-
-def run_test_model(test_model, observed_infections, rng_key): # numpydoc ignore=GL08
- test_model.run(
- num_warmup=50,
- num_samples=50,
- data_observed_infections=observed_infections,
- rng_key=rng_key,
- mcmc_args=dict(progress_bar=True),
- )
-
-
-def prior_predictive_test_model(
- test_model, n_datapoints, rng_key
-): # numpydoc ignore=GL08
- prior_predictive_samples = test_model.prior_predictive(
- rng_key=rng_key,
- numpyro_predictive_args={"num_samples": 20},
- n_datapoints=n_datapoints,
- )
- return prior_predictive_samples
-
-
-def posterior_predictive_test_model(
- test_model, n_datapoints, rng_key
-): # numpydoc ignore=GL08
- posterior_predictive_samples = test_model.posterior_predictive(
- rng_key=rng_key, n_datapoints=n_datapoints
- )
- return posterior_predictive_samples
-
-
-def test_rng_keys_produce_correct_samples():
- """
- Tests that the random keys specified for
- MCMC sampling produce appropriate
- output if left to None or specified directly.
- """
-
- # set up singular epidemiological process
-
- # set up base models for testing
- models = [create_test_model() for _ in range(5)]
- n_datapoints = [30] * len(models)
- n_timepoints_posterior_predictive = [
- x + models[0].gen_int_rv.size() for x in n_datapoints
- ]
- # sample only a single model and use that model's samples
- # as the observed_infections for the rest of the models
- with numpyro.handlers.seed(rng_seed=223):
- model_sample = models[0].sample(n_datapoints=n_datapoints[0])
- obs_infections = [model_sample.observed_infections] * len(models)
- rng_keys = [jr.key(54), jr.key(54), None, None, jr.key(74)]
-
- # run test models with the different keys
-
- for elt in list(zip(models, obs_infections, rng_keys)):
- run_test_model(*elt)
-
- prior_predictive_list = [
- prior_predictive_test_model(*elt)
- for elt in list(zip(models, n_datapoints, rng_keys))
- ]
-
- posterior_predictive_list = [
- posterior_predictive_test_model(*elt)
- for elt in list(zip(models, n_timepoints_posterior_predictive, rng_keys))
- ]
- # using same rng_key should get same run samples
- assert_array_equal(
- models[0].mcmc.get_samples()["Rt"][0],
- models[1].mcmc.get_samples()["Rt"][0],
- )
-
- assert_array_equal(
- prior_predictive_list[0]["Rt"][0],
- prior_predictive_list[1]["Rt"][0],
- )
-
- assert_array_equal(
- posterior_predictive_list[0]["Rt"][0],
- posterior_predictive_list[1]["Rt"][0],
- )
-
- # using None for rng_key should get different run samples
- assert_raises( # negate assert_array_equal
- AssertionError,
- assert_array_equal,
- models[2].mcmc.get_samples()["Rt"][0],
- models[3].mcmc.get_samples()["Rt"][0],
- )
-
- assert_raises(
- AssertionError,
- assert_array_equal,
- prior_predictive_list[2]["Rt"][0],
- prior_predictive_list[3]["Rt"][0],
- )
-
- assert_raises(
- AssertionError,
- assert_array_equal,
- posterior_predictive_list[2]["Rt"][0],
- posterior_predictive_list[3]["Rt"][0],
- )
-
- # using None vs preselected rng_key should get different samples
- assert_raises(
- AssertionError,
- assert_array_equal,
- models[0].mcmc.get_samples()["Rt"][0],
- models[2].mcmc.get_samples()["Rt"][0],
- )
-
- assert_raises(
- AssertionError,
- assert_array_equal,
- prior_predictive_list[0]["Rt"][0],
- prior_predictive_list[2]["Rt"][0],
- )
-
- assert_raises(
- AssertionError,
- assert_array_equal,
- posterior_predictive_list[0]["Rt"][0],
- posterior_predictive_list[2]["Rt"][0],
- )
-
- # using two different non-None rng keys should get different samples
- assert_raises(
- AssertionError,
- assert_array_equal,
- models[1].mcmc.get_samples()["Rt"][0],
- models[4].mcmc.get_samples()["Rt"][0],
- )
-
- assert_raises(
- AssertionError,
- assert_array_equal,
- prior_predictive_list[1]["Rt"][0],
- prior_predictive_list[4]["Rt"][0],
- )
-
- assert_raises(
- AssertionError,
- assert_array_equal,
- posterior_predictive_list[1]["Rt"][0],
- posterior_predictive_list[4]["Rt"][0],
- )
diff --git a/test/test_transformed_rv_class.py b/test/test_transformed_rv_class.py
index cf306c8b..73961940 100644
--- a/test/test_transformed_rv_class.py
+++ b/test/test_transformed_rv_class.py
@@ -228,3 +228,44 @@ def test_transforms_variable_naming():
assert "transformed_rv_0" in mymodel3.mcmc.get_samples()
assert "transformed_rv_1" in mymodel3.mcmc.get_samples()
+
+
+def test_posterior_predictive():
+ """
+ Test that Model.posterior_predictive() returns samples
+ after running MCMC, and raises before running.
+ """
+ rv = DistributionalVariable(name="test_normal", distribution=dist.Normal(0, 1))
+ model = MyModel(rv)
+
+ with pytest.raises(ValueError, match="No posterior samples available"):
+ model.posterior_predictive(rng_key=jax.random.key(0))
+
+ model.run(num_samples=5, num_warmup=5, rng_key=jax.random.key(1))
+ result = model.posterior_predictive(rng_key=jax.random.key(2))
+ assert isinstance(result, dict)
+
+ # Also test with rng_key=None (auto-generated)
+ result_auto_key = model.posterior_predictive()
+ assert isinstance(result_auto_key, dict)
+
+
+def test_prior_predictive():
+ """
+ Test that Model.prior_predictive() returns samples
+ without needing to run MCMC first.
+ """
+ rv = DistributionalVariable(name="test_normal", distribution=dist.Normal(0, 1))
+ model = MyModel(rv)
+
+ result = model.prior_predictive(
+ rng_key=jax.random.key(0),
+ numpyro_predictive_args={"num_samples": 5},
+ )
+ assert isinstance(result, dict)
+
+ # Also test with rng_key=None (auto-generated)
+ result_auto_key = model.prior_predictive(
+ numpyro_predictive_args={"num_samples": 5},
+ )
+ assert isinstance(result_auto_key, dict)