Add hosp_only model and demo

notebooks/hosp_only_ww_model.qmd

@@ -0,0 +1,121 @@
+---
+title: "Replicating Hospital Only Model from ww-inference-model"
+format: gfm
+engine: jupyter
+---
+
+```{python}
+# | label: setup
+import jax
+import numpyro
+import arviz as az
+import pyrenew_covid_wastewater.plotting as plotting
+from pyrenew_covid_wastewater.hosp_only_ww_model import (
+    create_hosp_only_ww_model_from_stan_data,
+)
+
+numpyro.set_host_device_count(4)
+```
+
+## Background
+
+This tutorial provides a demonstration of our reimplementation of "Model 2" from the [`ww-inference-model` project](https://github.com/CDCgov/ww-inference-model).
+The model is described [here](https://github.com/CDCgov/ww-inference-model/blob/main/model_definition.md).
+Stan code for the model is [here](https://github.com/CDCgov/ww-inference-model/blob/main/inst/stan/wwinference.stan).
+
+The model we provide is designed to be fully-compatible with the stan_data generated in the that project.
+We provide the stan data used in the `wwinference` [vignette](https://github.com/CDCgov/ww-inference-model/blob/main/vignettes/wwinference.Rmd) in the [`ww-inference-model` project](https://github.com/CDCgov/ww-inference-model).
+The data is available in `notebooks/data/fit_hosp_only/stan_data.json`.
+This data was generated by running `notebooks/wwinference.Rmd`, which replicates the original vignette and saves the relevant data.
+This script also saves the posterior samples from the model for comparison to our own model.
+
+## Load Data and Create Priors
+
+We begin by loading the Stan data, converting it the correct inputs for our model, and definitng the model.
+
+```{python}
+# | label: create model
+my_hosp_only_ww_model, data_observed_hospital_admissions = (
+    create_hosp_only_ww_model_from_stan_data(
+        "data/fit_hosp_only/stan_data.json"
+    )
+)
+```
+
+# Simulate from the model
+
+We check that we can simulate from the prior predictive
+```{python}
+# | label: prior predictive
+n_forecast_days = 28
+prior_predictive = my_hosp_only_ww_model.prior_predictive(
+    n_datapoints=len(data_observed_hospital_admissions) + n_forecast_days,
+    numpyro_predictive_args={"num_samples": 200},
+)
+```
+
+# Fit the model
+
+Now we can fit the model to the observed data:
+```{python}
+# | label: fit the model
+my_hosp_only_ww_model.run(
+    num_warmup=500,
+    num_samples=500,
+    rng_key=jax.random.key(200),
+    data_observed_hospital_admissions=data_observed_hospital_admissions,
+    mcmc_args=dict(num_chains=4, progress_bar=False),
+    nuts_args=dict(find_heuristic_step_size=True),
+)
+```
+
+Create the posterior predictive and forecast:
+
+```{python}
+# | label: posterior predictive
+posterior_predictive = my_hosp_only_ww_model.posterior_predictive(
+    n_datapoints=len(data_observed_hospital_admissions) + n_forecast_days
+)
+```
+
+## Prepare for plotting
+
+```{python}
+# | label: prepare for plotting
+import arviz as az
+
+idata = az.from_numpyro(
+    my_hosp_only_ww_model.mcmc,
+    posterior_predictive=posterior_predictive,
+    prior=prior_predictive,
+)
+```
+
+## Plot Predictive Distributions
+
+```{python}
+# | label: plot prior preditive
+plotting.plot_predictive(idata, prior=True)
+```
+
+```{python}
+# | label: plot posterior preditive
+plotting.plot_predictive(idata)
+```
+
+## Plot all posteriors
+
+```{python}
+# | label: plot all posteriors
+for key in list(idata.posterior.keys()):
+    try:
+        plotting.plot_posterior(idata, key)
+    except Exception as e:
+        print(f"An error occurred while plotting {key}: {e}")
+```
+
+## Save for Post-Processing
+
+```{python}
+idata.to_dataframe().to_csv("data/fit_hosp_only/pyrenew_inference_data.csv")
+```

pyproject.toml

@@ -8,6 +8,12 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.12"
 pyrenew = {git = "https://github.com/CDCgov/PyRenew"}
+ipywidgets = "^8.1.5"
+arviz = "^0.19.0"
+pyyaml = "^6.0.2"
+jupyter = "^1.0.0"
+ipykernel = "^6.29.5"
+polars = "^1.5.0"
 
 
 [build-system]

pyrenew-covid-wastewater.Rproj

@@ -0,0 +1,16 @@
+Version: 1.0
+
+RestoreWorkspace: Default
+SaveWorkspace: Default
+AlwaysSaveHistory: Default
+
+EnableCodeIndexing: Yes
+UseSpacesForTab: Yes
+NumSpacesForTab: 2
+Encoding: UTF-8
+
+RnwWeave: Sweave
+LaTeX: pdfLaTeX
+
+AutoAppendNewline: Yes
+StripTrailingWhitespace: Yes