diff --git a/notebooks/hosp_only_ww_model.qmd b/notebooks/hosp_only_ww_model.qmd
index 33e42602..778aca0e 100644
--- a/notebooks/hosp_only_ww_model.qmd
+++ b/notebooks/hosp_only_ww_model.qmd
@@ -35,7 +35,7 @@ We begin by loading the Stan data, converting it the correct inputs for our mode
 
 ```{python}
 # | label: create model
-my_hosp_only_ww_model, data_observed_hospital_admissions = (
+my_hosp_only_ww_model, data_observed_disease_hospital_admissions = (
     create_hosp_only_ww_model_from_stan_data(
         "data/fit_hosp_only/stan_data.json"
     )
@@ -50,7 +50,7 @@ We check that we can simulate from the prior predictive
 n_forecast_days = 35
 
 prior_predictive = my_hosp_only_ww_model.prior_predictive(
-    n_datapoints=len(data_observed_hospital_admissions) + n_forecast_days,
+    n_datapoints=len(data_observed_disease_hospital_admissions) + n_forecast_days,
     numpyro_predictive_args={"num_samples": 200},
 )
 ```
@@ -64,7 +64,7 @@ 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,
+    data_observed_disease_hospital_admissions=data_observed_disease_hospital_admissions,
     mcmc_args=dict(num_chains=4, progress_bar=False),
     nuts_args=dict(find_heuristic_step_size=True),
 )
@@ -75,7 +75,7 @@ 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
+    n_datapoints=len(data_observed_disease_hospital_admissions) + n_forecast_days
 )
 ```
 
diff --git a/nssp_demo/build_model.py b/nssp_demo/build_model.py
index 78099ed2..e14576e7 100644
--- a/nssp_demo/build_model.py
+++ b/nssp_demo/build_model.py
@@ -45,8 +45,8 @@ def build_model_from_dir(model_dir):
         jnp.array(model_data["generation_interval_pmf"]),
     )  # check if off by 1 or reversed
 
-    data_observed_hospital_admissions = jnp.array(
-        model_data["data_observed_hospital_admissions"]
+    data_observed_disease_hospital_admissions = jnp.array(
+        model_data["data_observed_disease_hospital_admissions"]
     )
     state_pop = jnp.array(model_data["state_pop"])
 
@@ -84,4 +84,8 @@ def build_model_from_dir(model_dir):
         n_initialization_points=uot,
     )
 
-    return my_model, data_observed_hospital_admissions, right_truncation_offset
+    return (
+        my_model,
+        data_observed_disease_hospital_admissions,
+        right_truncation_offset,
+    )
diff --git a/nssp_demo/fit_model.py b/nssp_demo/fit_model.py
index d72dbd9a..b2f57cab 100644
--- a/nssp_demo/fit_model.py
+++ b/nssp_demo/fit_model.py
@@ -21,14 +21,16 @@
 args = parser.parse_args()
 model_dir = args.model_dir
 
-my_model, data_observed_hospital_admissions, right_truncation_offset = (
-    build_model_from_dir(model_dir)
-)
+(
+    my_model,
+    data_observed_disease_hospital_admissions,
+    right_truncation_offset,
+) = build_model_from_dir(model_dir)
 my_model.run(
     num_warmup=500,
     num_samples=500,
     rng_key=jax.random.key(200),
-    data_observed_hospital_admissions=data_observed_hospital_admissions,
+    data_observed_disease_hospital_admissions=data_observed_disease_hospital_admissions,
     right_truncation_offset=right_truncation_offset,
     mcmc_args=dict(num_chains=n_chains, progress_bar=True),
     nuts_args=dict(find_heuristic_step_size=True),
diff --git a/nssp_demo/generate_predictive.py b/nssp_demo/generate_predictive.py
index a4b6bbe9..fcca649b 100644
--- a/nssp_demo/generate_predictive.py
+++ b/nssp_demo/generate_predictive.py
@@ -27,9 +27,11 @@
 args = parser.parse_args()
 model_dir = args.model_dir
 n_forecast_points = args.n_forecast_points
-my_model, data_observed_hospital_admissions, right_truncation_offset = (
-    build_model_from_dir(model_dir)
-)
+(
+    my_model,
+    data_observed_disease_hospital_admissions,
+    right_truncation_offset,
+) = build_model_from_dir(model_dir)
 
 my_model._init_model(1, 1)
 fresh_sampler = my_model.mcmc.sampler
@@ -47,12 +49,13 @@
 #         "num_samples": my_model.mcmc.num_samples * my_model.mcmc.num_chains,
 #         "batch_ndims":1
 #     },
-#     n_datapoints=len(data_observed_hospital_admissions) + n_forecast_points,
+#     n_datapoints=len(data_observed_disease_hospital_admissions) + n_forecast_points,
 # )
 # need to figure out a way to generate these as distinct chains, so that the result of the to_datarame method is more compact
 
 posterior_predictive = my_model.posterior_predictive(
-    n_datapoints=len(data_observed_hospital_admissions) + n_forecast_points
+    n_datapoints=len(data_observed_disease_hospital_admissions)
+    + n_forecast_points
 )
 
 idata = az.from_numpyro(
diff --git a/nssp_demo/other_ed_admissions_forecast.R b/nssp_demo/other_ed_admissions_forecast.R
new file mode 100644
index 00000000..02fa76b1
--- /dev/null
+++ b/nssp_demo/other_ed_admissions_forecast.R
@@ -0,0 +1,101 @@
+library(tidyverse)
+library(fs)
+library(fable)
+library(jsonlite)
+library(glue)
+library(argparser)
+library(arrow)
+
+p <- arg_parser("Forecast other (non-target-disease) ED admissions") %>%
+  add_argument(p, "--model_dir",
+    help = "Directory containing the model data",
+    required = TRUE
+  ) %>%
+  add_argument("--n_forecast_days",
+    help = "Number of days to forecast",
+    default = 28L
+  ) %>%
+  add_argument("--n_samples",
+    help = "Number of samples to generate",
+    default = 2000L
+  )
+
+argv <- parse_args(p)
+model_dir <- path(argv$model_dir)
+n_forecast_days <- argv$n_forecast_days
+n_samples <- arv$n_samples
+
+disease_name_nssp_map <- c(
+  "covid-19" = "COVID-19/Omicron",
+  "influenza" = "Influenza"
+)
+
+base_dir <- path_dir(model_dir)
+
+disease_name_raw <- base_dir %>%
+  path_file() %>%
+  str_extract("^.+(?=_r_)")
+
+disease_name_nssp <- unname(disease_name_nssp_map[disease_name_raw])
+
+fit_and_forecast <- function(other_data,
+                             n_forecast_days = 28,
+                             n_samples = 2000) {
+  forecast_horizon <- glue("{n_forecast_days} days")
+
+  fit <-
+    other_data %>%
+    filter(data_type == "train") %>%
+    model(
+      comb_model = combination_ensemble(
+        ETS(log(ED_admissions) ~ trend(method = c("N", "M", "A"))),
+        ARIMA(log(ED_admissions))
+      ),
+      arima = ARIMA(log(ED_admissions)),
+      ets = ETS(log(ED_admissions) ~ trend(method = c("N", "M", "A")))
+    )
+
+  forecast_samples <- fit |>
+    generate(h = forecast_horizon, times = n_samples) |>
+    as_tibble() %>%
+    mutate(ED_admissions = .sim, .draw = as.integer(.rep)) |>
+    filter(.model == "comb_model") %>%
+    select(date, .draw, other_ED_admissions = ED_admissions)
+
+  forecast_samples
+}
+
+main <- function(model_dir, n_forecast_days = 28, n_samples = 2000) {
+  # to do: do this with json data that has dates
+  data_path <- path(model_dir, "data", ext = "csv")
+
+  other_data <- read_csv(data_path) %>%
+    mutate(disease = if_else(disease == disease_name_nssp,
+      "Disease", disease
+    )) %>%
+    pivot_wider(names_from = disease, values_from = ED_admissions) %>%
+    mutate(Other = Total - Disease) %>%
+    select(date, ED_admissions = Other, data_type) %>%
+    as_tsibble(index = date)
+
+  forecast_samples <- fit_and_forecast(other_data, n_forecast_days, n_samples)
+
+  save_path <- path(model_dir, "other_ed_admissions_forecast", ext = "parquet")
+  write_parquet(forecast_samples, save_path)
+}
+
+
+main(model_dir, n_forecast_days, n_samples)
+# File will end here once command line version is working
+# Temp code to run for all states while command line version doesn't work
+# Command line version is dependent on https://github.com/rstudio/renv/pull/2018
+
+base_dir <- path(
+  "nssp_demo/private_data/influenza_r_2024-10-21_f_2024-07-16_t_2024-10-13"
+)
+
+dir_ls(base_dir, type = "dir") %>%
+  walk(.f = function(model_dir) {
+    print(path_file(model_dir))
+    main(model_dir, n_forecast_days = 28, n_samples = 2000)
+  })
diff --git a/nssp_demo/post_process.R b/nssp_demo/post_process.R
index f768a39d..d1bb0cdd 100644
--- a/nssp_demo/post_process.R
+++ b/nssp_demo/post_process.R
@@ -6,10 +6,15 @@ library(glue)
 library(scales)
 library(here)
 library(argparser)
+library(arrow)
 
 theme_set(theme_minimal_grid())
 
 disease_name_formatter <- c("covid-19" = "COVID-19", "influenza" = "Flu")
+disease_name_nssp_map <- c(
+  "covid-19" = "COVID-19/Omicron",
+  "influenza" = "Influenza"
+)
 
 # Create a parser
 p <- arg_parser("Generate forecast figures") %>%
@@ -23,7 +28,7 @@ p <- arg_parser("Generate forecast figures") %>%
   ) %>%
   add_argument(p, "--good_chain_tol",
     help = "Tolerance level for determining good chains",
-    default = 2
+    default = 2L
   )
 
 argv <- parse_args(p)
@@ -33,8 +38,14 @@ good_chain_tol <- argv$good_chain_tol
 
 base_dir <- path_dir(model_dir)
 
+disease_name_raw <- base_dir %>%
+  path_file() %>%
+  str_extract("^.+(?=_r_)")
 
+disease_name_nssp <- unname(disease_name_nssp_map[disease_name_raw])
+disease_name_pretty <- unname(disease_name_formatter[disease_name_raw])
 
+# To be replaced with reading tidy data from forecasttools
 read_pyrenew_samples <- function(inference_data_path,
                                  filter_bad_chains = TRUE,
                                  good_chain_tol = 2) {
@@ -84,13 +95,67 @@ read_pyrenew_samples <- function(inference_data_path,
   good_pyrenew_samples
 }
 
-make_forecast_fig <- function(model_dir,
-                              filter_bad_chains = TRUE,
-                              good_chain_tol = 2) {
-  disease_name_raw <- base_dir %>%
-    path_file() %>%
-    str_extract("^.+(?=_r_)")
+make_one_forecast_fig <- function(target_disease,
+                                  dat,
+                                  last_training_date,
+                                  last_data_date,
+                                  posterior_predictive_ci,
+                                  state_abb) {
+  y_scale <- if (str_starts(target_disease, "prop")) {
+    scale_y_continuous("Proportion of Emergency Department Admissions",
+      labels = percent
+    )
+  } else {
+    scale_y_continuous("Emergency Department Admissions", labels = comma)
+  }
+
+  title <- if (target_disease == "Other") {
+    glue("Other ED Admissions in {state_abb}")
+  } else {
+    glue("{disease_name_pretty} ED Admissions in {state_abb}")
+  }
 
+  ggplot(mapping = aes(date, .value)) +
+    geom_lineribbon(
+      data = posterior_predictive_ci %>% filter(disease == target_disease),
+      mapping = aes(ymin = .lower, ymax = .upper),
+      color = "#08519c", key_glyph = draw_key_rect, step = "mid"
+    ) +
+    geom_point(
+      mapping = aes(shape = data_type),
+      data = dat %>% filter(disease == target_disease)
+    ) +
+    geom_vline(xintercept = last_training_date, linetype = "dashed") +
+    annotate(
+      geom = "text",
+      x = last_training_date,
+      y = -Inf,
+      label = "Fit Period ←\n",
+      hjust = "right",
+      vjust = "bottom"
+    ) +
+    annotate(
+      geom = "text",
+      x = last_training_date,
+      y = -Inf, label = "→ Forecast Period\n",
+      hjust = "left",
+      vjust = "bottom",
+    ) +
+    ggtitle(title, subtitle = glue("as of {last_data_date}")) +
+    y_scale +
+    scale_x_date("Date") +
+    scale_shape_discrete("Data Type", labels = str_to_title) +
+    scale_fill_brewer(
+      name = "Credible Interval Width",
+      labels = ~ percent(as.numeric(.))
+    ) +
+    theme(legend.position = "bottom")
+}
+
+
+make_forecast_figs <- function(model_dir,
+                               filter_bad_chains = TRUE,
+                               good_chain_tol = 2) {
   state_abb <- model_dir %>%
     path_split() %>%
     pluck(1) %>%
@@ -100,11 +165,27 @@ make_forecast_fig <- function(model_dir,
   inference_data_path <- path(model_dir, "inference_data",
     ext = "csv"
   )
+  other_ed_admissions_path <- path(model_dir, "other_ed_admissions_forecast",
+    ext = "parquet"
+  )
 
-  dat <- read_csv(data_path) %>%
-    arrange(date) %>%
-    mutate(time = row_number() - 1) %>%
-    rename(.value = ED_admissions)
+  dat <-
+    read_csv(data_path) %>%
+    mutate(disease = if_else(disease == disease_name_nssp,
+      "Disease", # assign a common name for use in plotting functions
+      disease
+    )) %>%
+    pivot_wider(names_from = disease, values_from = ED_admissions) %>%
+    mutate(
+      Other = Total - Disease,
+      prop_disease_ed_admissions = Disease / Total
+    ) %>%
+    select(-Total) %>%
+    mutate(time = dense_rank(date)) %>%
+    pivot_longer(c(Disease, Other, prop_disease_ed_admissions),
+      names_to = "disease",
+      values_to = ".value"
+    )
 
   last_training_date <- dat %>%
     filter(data_type == "train") %>%
@@ -115,70 +196,76 @@ make_forecast_fig <- function(model_dir,
     pull(date) %>%
     max()
 
-
   pyrenew_samples <- read_pyrenew_samples(inference_data_path,
     filter_bad_chains = filter_bad_chains,
     good_chain_tol = good_chain_tol
   )
 
-  hosp_ci <-
+  other_ed_admission_forecast <-
+    read_parquet(other_ed_admissions_path) %>%
+    rename(Other = other_ED_admissions)
+
+
+  other_ed_admission_samples <-
+    bind_rows(
+      dat %>%
+        filter(
+          disease == "Other",
+          date <= last_training_date
+        ) %>%
+        select(date, Other = .value) %>%
+        expand_grid(.draw = 1:max(other_ed_admission_forecast$.draw)),
+      other_ed_admission_forecast
+    )
+
+  posterior_predictive_samples <-
     pyrenew_samples$posterior_predictive %>%
     gather_draws(observed_hospital_admissions[time]) %>%
-    median_qi(.width = c(0.5, 0.8, 0.95)) %>%
-    mutate(date = min(dat$date) + time)
+    pivot_wider(names_from = .variable, values_from = .value) %>%
+    rename(Disease = observed_hospital_admissions) %>%
+    ungroup() %>%
+    mutate(date = min(dat$date) + time) %>%
+    left_join(other_ed_admission_samples) %>%
+    mutate(prop_disease_ed_admissions = Disease / (Disease + Other)) %>%
+    pivot_longer(c(Other, Disease, prop_disease_ed_admissions),
+      names_to = "disease",
+      values_to = ".value"
+    )
 
+  posterior_predictive_ci <-
+    posterior_predictive_samples %>%
+    select(date, disease, .value) %>%
+    group_by(date, disease) %>%
+    median_qi(.width = c(0.5, 0.8, 0.95))
 
 
-  forecast_plot <-
-    ggplot(mapping = aes(date, .value)) +
-    geom_lineribbon(
-      data = hosp_ci,
-      mapping = aes(ymin = .lower, ymax = .upper),
-      color = "#08519c", key_glyph = draw_key_rect, step = "mid"
-    ) +
-    geom_point(mapping = aes(shape = data_type), data = dat) +
-    scale_y_continuous("Emergency Department Admissions") +
-    scale_x_date("Date") +
-    scale_fill_brewer(
-      name = "Credible Interval Width",
-      labels = ~ percent(as.numeric(.))
-    ) +
-    scale_shape_discrete("Data Type", labels = str_to_title) +
-    geom_vline(xintercept = last_training_date, linetype = "dashed") +
-    annotate(
-      geom = "text",
-      x = last_training_date,
-      y = -Inf,
-      label = "Fit Period ←\n",
-      hjust = "right",
-      vjust = "bottom"
-    ) +
-    annotate(
-      geom = "text",
-      x = last_training_date,
-      y = -Inf, label = "→ Forecast Period\n",
-      hjust = "left",
-      vjust = "bottom",
-    ) +
-    ggtitle(
-      glue(
-        "{disease_name_formatter[disease_name_raw]} ",
-        "NSSP-based forecast for {state_abb}"
-      ),
-      subtitle = glue("as of {last_data_date}")
-    ) +
-    theme(legend.position = "bottom")
+  all_forecast_plots <- map(
+    set_names(unique(dat$disease)),
+    ~ make_one_forecast_fig(
+      .x,
+      dat,
+      last_training_date,
+      last_data_date,
+      posterior_predictive_ci,
+      state_abb
+    )
+  )
 
-  forecast_plot
+  all_forecast_plots
 }
 
-forecast_fig <- make_forecast_fig(model_dir, filter_bad_chains, good_chain_tol)
+forecast_figs <- make_forecast_figs(
+  model_dir,
+  filter_bad_chains,
+  good_chain_tol
+)
 
-save_plot(
-  filename = path(model_dir, "forecast_plot", ext = "pdf"),
-  plot = forecast_fig,
+iwalk(forecast_figs, ~ save_plot(
+  filename = path(model_dir, glue("{.y}_forecast_plot"), ext = "pdf"),
+  plot = .x,
   device = cairo_pdf, base_height = 6
-)
+))
+
 
 # File will end here once command line version is working
 # Temp code to run for all states while command line version doesn't work
@@ -186,26 +273,39 @@ save_plot(
 base_dir <- path(
   "nssp_demo",
   "private_data",
-  "influenza_r_2024-10-10_f_2024-04-12_l_2024-10-09_t_2024-10-05"
+  "influenza_r_2024-10-21_f_2024-07-16_t_2024-10-13"
 )
 
-walk(dir_ls(base_dir), function(model_dir) {
-  forecast_fig <- make_forecast_fig(model_dir)
+# Save all figures for each state
+walk(dir_ls(base_dir, type = "dir"), function(model_dir) {
+  print(model_dir)
+  forecast_figs <- make_forecast_figs(model_dir,
+    filter_bad_chains = TRUE,
+    good_chain_tol = 2
+  )
 
-  save_plot(
-    filename = path(model_dir, "forecast_plot", ext = "pdf"),
-    plot = forecast_fig,
+  iwalk(forecast_figs, ~ save_plot(
+    filename = path(model_dir, glue("{.y}_forecast_plot"), ext = "pdf"),
+    plot = .x,
     device = cairo_pdf, base_height = 6
-  )
+  ))
 })
 
-path(dir_ls(base_dir, type = "directory"), "forecast_plot", ext = "pdf") %>%
-  str_c(collapse = " ") %>%
-  str_c(
-    path(base_dir,
-      glue("{path_file(base_dir)}_all_forecasts"),
-      ext = "pdf"
-    ),
-    sep = " "
-  ) %>%
-  system2("pdfunite", args = .)
+# Combine figures across states
+tibble(
+  full_path = dir_ls(base_dir,
+    type = "file",
+    glob = "*_forecast_plot.pdf",
+    recurse = TRUE
+  ),
+  plot_type = path_file(full_path)
+) %>%
+  group_by(plot_type) %>%
+  summarize(all_fig_paths = str_c(full_path, collapse = " ")) %>%
+  mutate(combined_plot_path = path(
+    base_dir,
+    glue("{path_file(base_dir)}_{plot_type}")
+  )) %>%
+  select(-plot_type) %>%
+  as.list() %>%
+  pwalk(~ system2("pdfunite", args = glue("{.x} {.y}")))
diff --git a/nssp_demo/prep_data.py b/nssp_demo/prep_data.py
index fb078094..cebd6bba 100644
--- a/nssp_demo/prep_data.py
+++ b/nssp_demo/prep_data.py
@@ -6,7 +6,6 @@
 from pathlib import Path
 
 import polars as pl
-import pyarrow.parquet as pq
 
 
 def process_and_save_state(
@@ -25,16 +24,17 @@ def process_and_save_state(
         "COVID-19": "COVID-19/Omicron",
         "Influenza": "Influenza",
         "RSV": "RSV",
+        "Total": "Total",
     }
 
     data_to_save = (
         nssp_data.filter(
-            (pl.col("disease") == disease_map[disease])
+            (pl.col("disease").is_in([disease_map[disease], "Total"]))
             & (pl.col("metric") == "count_ed_visits")
             & (pl.col("geo_value") == state_abb)
             & (pl.col("reference_date") >= first_training_date)
         )
-        .group_by(["report_date", "reference_date"])
+        .group_by(["reference_date", "disease"])
         .agg(pl.col("value").sum().alias("ED_admissions"))
         .with_columns(
             pl.when(pl.col("reference_date") <= last_training_date)
@@ -42,7 +42,8 @@ def process_and_save_state(
             .otherwise(pl.lit("test"))
             .alias("data_type")
         )
-        .sort(["report_date", "reference_date"])
+        .rename({"reference_date": "date"})
+        .sort(["date", "disease"])
     )
 
     state_pop = (
@@ -95,15 +96,39 @@ def process_and_save_state(
 
     right_truncation_offset = (report_date - last_training_date).days
 
-    train_ed_admissions = (
-        data_to_save.filter(pl.col("data_type") == "train")
+    train_disease_ed_admissions = (
+        data_to_save.filter(
+            (pl.col("data_type") == "train")
+            & (pl.col("disease") == disease_map[disease])
+        )
+        .collect()
+        .get_column("ED_admissions")
+        .to_list()
+    )
+
+    test_disease_ed_admissions = (
+        data_to_save.filter(
+            (pl.col("data_type") == "test")
+            & (pl.col("disease") == disease_map[disease])
+        )
         .collect()
         .get_column("ED_admissions")
         .to_list()
     )
 
-    test_ed_admissions = (
-        data_to_save.filter(pl.col("data_type") == "test")
+    train_total_ed_admissions = (
+        data_to_save.filter(
+            (pl.col("data_type") == "train") & (pl.col("disease") == "Total")
+        )
+        .collect()
+        .get_column("ED_admissions")
+        .to_list()
+    )
+
+    test_total_ed_admissions = (
+        data_to_save.filter(
+            (pl.col("data_type") == "test") & (pl.col("disease") == "Total")
+        )
         .collect()
         .get_column("ED_admissions")
         .to_list()
@@ -113,8 +138,10 @@ def process_and_save_state(
         "inf_to_hosp_pmf": delay_pmf,
         "generation_interval_pmf": generation_interval_pmf,
         "right_truncation_pmf": right_truncation_pmf,
-        "data_observed_hospital_admissions": train_ed_admissions,
-        "test_ed_admissions": test_ed_admissions,
+        "data_observed_disease_hospital_admissions": train_disease_ed_admissions,
+        "data_observed_disease_hospital_admissions_test": test_disease_ed_admissions,
+        "data_observed_total_hospital_admissions": train_total_ed_admissions,
+        "data_observed_total_hospital_admissions_test": test_total_ed_admissions,
         "state_pop": state_pop,
         "right_truncation_offset": right_truncation_offset,
     }
@@ -124,7 +151,8 @@ def process_and_save_state(
 
     if logger is not None:
         logger.info(f"Saving {state_abb} to {state_dir}")
-    data_to_save.sink_csv(Path(state_dir, "data.csv"))
+    # data_to_save.sink_csv(Path(state_dir, "data.csv")) # Not yet supported
+    data_to_save.collect().write_csv(Path(state_dir, "data.csv"))
 
     with open(Path(state_dir, "data_for_model_fit.json"), "w") as json_file:
         json.dump(data_for_model_fit, json_file)
diff --git a/pyrenew_hew/hosp_only_ww_model.py b/pyrenew_hew/hosp_only_ww_model.py
index 77f4a2e7..6ef5dd2f 100644
--- a/pyrenew_hew/hosp_only_ww_model.py
+++ b/pyrenew_hew/hosp_only_ww_model.py
@@ -85,23 +85,26 @@ def validate(self):  # numpydoc ignore=GL08
     def sample(
         self,
         n_datapoints=None,
-        data_observed_hospital_admissions=None,
+        data_observed_disease_hospital_admissions=None,
         right_truncation_offset=None,
     ):  # numpydoc ignore=GL08
-        if n_datapoints is None and data_observed_hospital_admissions is None:
+        if (
+            n_datapoints is None
+            and data_observed_disease_hospital_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_hospital_admissions is not None
+            and data_observed_disease_hospital_admissions is not None
         ):
             raise ValueError(
-                "Cannot pass both n_datapoints and data_observed_hospital_admissions."
+                "Cannot pass both n_datapoints and data_observed_disease_hospital_admissions."
             )
         elif n_datapoints is None:
-            n_datapoints = len(data_observed_hospital_admissions)
+            n_datapoints = len(data_observed_disease_hospital_admissions)
         else:
             n_datapoints = n_datapoints
 
@@ -225,7 +228,7 @@ def sample(
 
         observed_hospital_admissions = hospital_admission_obs_rv(
             mu=latent_hospital_admissions_now,
-            obs=data_observed_hospital_admissions,
+            obs=data_observed_disease_hospital_admissions,
         )
 
         return observed_hospital_admissions
@@ -360,7 +363,7 @@ def create_hosp_only_ww_model_from_stan_data(stan_data_file):
     uot = len(jnp.array(stan_data["inf_to_hosp"]))
     state_pop = stan_data["state_pop"]
 
-    data_observed_hospital_admissions = jnp.array(stan_data["hosp"])
+    data_observed_disease_hospital_admissions = jnp.array(stan_data["hosp"])
     right_truncation_pmf_rv = DeterministicVariable(
         "right_truncation_pmf", jnp.array(1)
     )
@@ -384,4 +387,4 @@ def create_hosp_only_ww_model_from_stan_data(stan_data_file):
         n_initialization_points=uot,
     )
 
-    return my_model, data_observed_hospital_admissions
+    return my_model, data_observed_disease_hospital_admissions
diff --git a/renv.lock b/renv.lock
index bf93dc94..0ec74e61 100644
--- a/renv.lock
+++ b/renv.lock
@@ -9,6 +9,13 @@
     ]
   },
   "Packages": {
+    "BH": {
+      "Package": "BH",
+      "Version": "1.84.0-0",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Hash": "a8235afbcd6316e6e91433ea47661013"
+    },
     "DBI": {
       "Package": "DBI",
       "Version": "1.2.3",
@@ -106,6 +113,28 @@
       ],
       "Hash": "4f57884290cc75ab22f4af9e9d4ca862"
     },
+    "anytime": {
+      "Package": "anytime",
+      "Version": "0.3.9",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "BH",
+        "R",
+        "Rcpp"
+      ],
+      "Hash": "74a64813f17b492da9c6afda6b128e3d"
+    },
+    "argparser": {
+      "Package": "argparser",
+      "Version": "0.7.2",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "methods"
+      ],
+      "Hash": "6ae51dd1cd6a2f5784273ab70d41c66c"
+    },
     "arrayhelpers": {
       "Package": "arrayhelpers",
       "Version": "1.1-0",
@@ -118,6 +147,28 @@
       ],
       "Hash": "3d4e52d458784c335af3846f2de64f75"
     },
+    "arrow": {
+      "Package": "arrow",
+      "Version": "17.0.0.1",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "R",
+        "R6",
+        "assertthat",
+        "bit64",
+        "cpp11",
+        "glue",
+        "methods",
+        "purrr",
+        "rlang",
+        "stats",
+        "tidyselect",
+        "utils",
+        "vctrs"
+      ],
+      "Hash": "14af96cb2973f6a6c220ce9c3e5b02cd"
+    },
     "askpass": {
       "Package": "askpass",
       "Version": "1.2.0",
@@ -128,6 +179,16 @@
       ],
       "Hash": "cad6cf7f1d5f6e906700b9d3e718c796"
     },
+    "assertthat": {
+      "Package": "assertthat",
+      "Version": "0.2.1",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "tools"
+      ],
+      "Hash": "50c838a310445e954bc13f26f26a6ecf"
+    },
     "backports": {
       "Package": "backports",
       "Version": "1.5.0",
@@ -530,6 +591,17 @@
       ],
       "Hash": "54ed3ea01b11e81a86544faaecfef8e2"
     },
+    "ellipsis": {
+      "Package": "ellipsis",
+      "Version": "0.3.2",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "R",
+        "rlang"
+      ],
+      "Hash": "bb0eec2fe32e88d9e2836c2f73ea2077"
+    },
     "evaluate": {
       "Package": "evaluate",
       "Version": "0.24.0",
@@ -541,6 +613,53 @@
       ],
       "Hash": "a1066cbc05caee9a4bf6d90f194ff4da"
     },
+    "fable": {
+      "Package": "fable",
+      "Version": "0.4.0",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "R",
+        "Rcpp",
+        "distributional",
+        "dplyr",
+        "fabletools",
+        "rlang",
+        "stats",
+        "tibble",
+        "tidyr",
+        "tsibble",
+        "utils"
+      ],
+      "Hash": "c0d483fecb4ea17420795160dc0bdd92"
+    },
+    "fabletools": {
+      "Package": "fabletools",
+      "Version": "0.5.0",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "R",
+        "R6",
+        "distributional",
+        "dplyr",
+        "generics",
+        "ggdist",
+        "ggplot2",
+        "lifecycle",
+        "progressr",
+        "rlang",
+        "scales",
+        "stats",
+        "tibble",
+        "tidyr",
+        "tidyselect",
+        "tsibble",
+        "utils",
+        "vctrs"
+      ],
+      "Hash": "0242939c5bf3a090df0b8877c7a6c97a"
+    },
     "fansi": {
       "Package": "fansi",
       "Version": "1.0.6",
@@ -1171,6 +1290,18 @@
       ],
       "Hash": "f4625e061cb2865f111b47ff163a5ca6"
     },
+    "progressr": {
+      "Package": "progressr",
+      "Version": "0.14.0",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "R",
+        "digest",
+        "utils"
+      ],
+      "Hash": "ac50c4ffa8f6a46580dd4d7813add3c4"
+    },
     "ps": {
       "Package": "ps",
       "Version": "1.7.7",
@@ -1660,6 +1791,27 @@
       ],
       "Hash": "cfbad971a71f0e27cec22e544a08bc3b"
     },
+    "tsibble": {
+      "Package": "tsibble",
+      "Version": "1.1.5",
+      "Source": "Repository",
+      "Repository": "CRAN",
+      "Requirements": [
+        "R",
+        "anytime",
+        "dplyr",
+        "ellipsis",
+        "generics",
+        "lifecycle",
+        "lubridate",
+        "methods",
+        "rlang",
+        "tibble",
+        "tidyselect",
+        "vctrs"
+      ],
+      "Hash": "a75e397766b45996310908b5b32557ba"
+    },
     "tzdb": {
       "Package": "tzdb",
       "Version": "0.4.0",