Forecasting Interface

docs/pyproject.toml

@@ -9,8 +9,8 @@ package-mode = false
 [tool.poetry.dependencies]
 python = "^3.12"
 sphinx = "^7.2.6"
-jax = "^0.4.25"
-jaxlib = "^0.4.25"
+jax = "0.4.29"
+jaxlib = "^0.4.29"
 numpyro = "^0.15.0"
 sphinxcontrib-mermaid = "^0.9.2"
 polars = "^0.20.16"

docs/source/tutorials/basic_renewal_model.qmd

@@ -214,7 +214,9 @@ model1.run(
     num_samples=1000,
     data_observed_infections=sim_data.observed_infections,
     rng_key=jax.random.PRNGKey(54),
-    mcmc_args=dict(progress_bar=False, num_chains=2, chain_method="sequential"),
+    mcmc_args=dict(
+        progress_bar=False, num_chains=2, chain_method="sequential"
+    ),
 )
 ```
 
@@ -335,7 +337,9 @@ az.plot_hdi(
 )
 
 # Add mean of the posterior to the figure
-mean_latent_infection = np.mean(idata.posterior["all_latent_infections"], axis=1)
+mean_latent_infection = np.mean(
+    idata.posterior["all_latent_infections"], axis=1
+)
 axes.plot(x_data, mean_latent_infection[0], color="C0", label="Mean")
 axes.legend()
 axes.set_title("Posterior Latent Infections", fontsize=10)

docs/source/tutorials/hospital_admissions_model.qmd

@@ -81,11 +81,12 @@ Let's take a look at the daily prevalence of hospital admissions.
 # | fig-cap: Daily hospital admissions from the simulated data
 import matplotlib.pyplot as plt
 
+daily_hosp_admits = dat["daily_hosp_admits"].to_numpy()
 # Rotating the x-axis labels, and only showing ~10 labels
 ax = plt.gca()
 ax.xaxis.set_major_locator(plt.MaxNLocator(nbins=10))
 ax.xaxis.set_tick_params(rotation=45)
-plt.plot(dat["date"].to_numpy(), dat["daily_hosp_admits"].to_numpy())
+plt.plot(dat["date"].to_numpy(), daily_hosp_admits)
 plt.xlabel("Date")
 plt.ylabel("Admissions")
 plt.show()
@@ -147,7 +148,10 @@ The `inf_hosp_int` is a `DeterministicPMF` object that takes the infection to ho
 ```{python}
 # | label: initializing-rest-of-model
 from pyrenew import model, process, observation, metaclass, transformation
-from pyrenew.latent import InfectionSeedingProcess, SeedInfectionsExponentialGrowth
+from pyrenew.latent import (
+    InfectionSeedingProcess,
+    SeedInfectionsExponentialGrowth,
+)
 
 # Infection process
 latent_inf = latent.Infections()
@@ -237,9 +241,11 @@ npro.set_host_device_count(jax.local_device_count())
 hosp_model.run(
     num_samples=1000,
     num_warmup=1000,
-    data_observed_hosp_admissions=dat["daily_hosp_admits"].to_numpy(),
+    data_observed_hosp_admissions=daily_hosp_admits,
     rng_key=jax.random.PRNGKey(54),
-    mcmc_args=dict(progress_bar=False, num_chains=2, chain_method="sequential"),
+    mcmc_args=dict(
+        progress_bar=False, num_chains=2, chain_method="sequential"
+    ),
 )
 ```
 
@@ -254,7 +260,7 @@ out = hosp_model.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
     obs_signal=np.pad(
-        dat["daily_hosp_admits"].to_numpy().astype(float),
+        daily_hosp_admits.astype(float),
         (gen_int_array.size, 0),
         constant_values=np.nan,
     ),
@@ -270,10 +276,10 @@ import arviz as az
 idata = az.from_numpyro(
     hosp_model.mcmc,
     posterior_predictive=hosp_model.posterior_predictive(
-        n_timepoints_to_simulate=len(dat["daily_hosp_admits"])
+        n_timepoints_to_simulate=len(daily_hosp_admits)
     ),
     prior=hosp_model.prior_predictive(
-        n_timepoints_to_simulate=len(dat["daily_hosp_admits"]),
+        n_timepoints_to_simulate=len(daily_hosp_admits),
         numpyro_predictive_args={"num_samples": 1000},
     ),
 )
@@ -308,22 +314,15 @@ We can use the padding argument to solve the overestimation of hospital admissio
 
 ```{python}
 # | label: model-fit-padding
-days_to_impute = 21
-
-# Add 21 Nas to the beginning of dat_w_padding
-dat_w_padding = np.pad(
-    dat["daily_hosp_admits"].to_numpy().astype(float),
-    (days_to_impute, 0),
-    constant_values=np.nan,
-)
+pad_size = 21
 
 hosp_model.run(
     num_samples=1000,
     num_warmup=1000,
-    data_observed_hosp_admissions=dat_w_padding,
+    data_observed_hosp_admissions=daily_hosp_admits,
     rng_key=jax.random.PRNGKey(54),
     mcmc_args=dict(progress_bar=False, num_chains=2),
-    padding=days_to_impute,  # Padding the model
+    padding=pad_size,  # Padding the model
 )
 ```
 
@@ -336,7 +335,9 @@ out = hosp_model.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
     obs_signal=np.pad(
-        dat_w_padding, (gen_int_array.size, 0), constant_values=np.nan
+        daily_hosp_admits.astype(float),
+        (gen_int_array.size + pad_size, 0),
+        constant_values=np.nan,
     ),
 )
 ```
@@ -407,7 +408,9 @@ We can look at individual draws from the posterior distribution of latent infect
 ```{python}
 # | label: fig-output-infections-with-padding
 # | fig-cap: Latent infections
-out2 = hosp_model.plot_posterior(var="all_latent_infections", ylab="Latent Infections")
+out2 = hosp_model.plot_posterior(
+    var="all_latent_infections", ylab="Latent Infections"
+)
 ```
 
 We can also look at credible intervals for the posterior distribution of latent infections:
@@ -440,7 +443,9 @@ az.plot_hdi(
 )
 
 # Add mean of the posterior to the figure
-mean_latent_infection = np.mean(idata.posterior["all_latent_infections"], axis=1)
+mean_latent_infection = np.mean(
+    idata.posterior["all_latent_infections"], axis=1
+)
 axes.plot(x_data, mean_latent_infection[0], color="C0", label="Mean")
 axes.legend()
 axes.set_title("Posterior Latent Infections", fontsize=10)
@@ -520,10 +525,10 @@ Running the model (with the same padding as before):
 hosp_model_weekday.run(
     num_samples=2000,
     num_warmup=2000,
-    data_observed_hosp_admissions=dat_w_padding,
+    data_observed_hosp_admissions=daily_hosp_admits,
     rng_key=jax.random.PRNGKey(54),
     mcmc_args=dict(progress_bar=False),
-    padding=days_to_impute,
+    padding=pad_size,
 )
 ```
 
@@ -535,21 +540,29 @@ And plotting the results:
 out = hosp_model_weekday.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
-    obs_signal=np.pad(dat_w_padding, (gen_int_array.size, 0), constant_values=np.nan),
+    obs_signal=np.pad(
+        daily_hosp_admits.astype(float),
+        (gen_int_array.size + pad_size, 0),
+        constant_values=np.nan,
+    ),
 )
 ```
 
 We will use ArviZ to visualize the posterior/prior predictive distributions.
+By increasing `n_timepoints_to_simulate`, we can perform forecasting.
 
 ```{python}
 # | label: posterior-predictive-distribution
+n_forecast_points = 28
 idata_weekday = az.from_numpyro(
     hosp_model_weekday.mcmc,
     posterior_predictive=hosp_model_weekday.posterior_predictive(
-        n_timepoints_to_simulate=len(dat_w_padding)
+        n_timepoints_to_simulate=len(daily_hosp_admits) + n_forecast_points,
+        padding=pad_size,
     ),
     prior=hosp_model_weekday.prior_predictive(
-        n_timepoints_to_simulate=len(dat_w_padding),
+        n_timepoints_to_simulate=len(daily_hosp_admits),
+        padding=pad_size,
         numpyro_predictive_args={"num_samples": 1000},
     ),
 )
@@ -569,7 +582,9 @@ def compute_eti(dataset, eti_prob):
 
 fig, axes = plt.subplots(figsize=(6, 5))
 az.plot_hdi(
-    idata_weekday.prior_predictive["negbinom_rv_dim_0"],
+    idata_weekday.prior_predictive["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
     hdi_data=compute_eti(idata_weekday.prior_predictive["negbinom_rv"], 0.9),
     color="C0",
     smooth=False,
@@ -578,7 +593,9 @@ az.plot_hdi(
 )
 
 az.plot_hdi(
-    idata_weekday.prior_predictive["negbinom_rv_dim_0"],
+    idata_weekday.prior_predictive["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
     hdi_data=compute_eti(idata_weekday.prior_predictive["negbinom_rv"], 0.5),
     color="C0",
     smooth=False,
@@ -587,7 +604,9 @@ az.plot_hdi(
 )
 
 plt.scatter(
-    idata_weekday.observed_data["negbinom_rv_dim_0"] + days_to_impute,
+    idata_weekday.observed_data["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
     idata_weekday.observed_data["negbinom_rv"],
     color="black",
 )
@@ -605,17 +624,25 @@ And now we plot the posterior predictive distributions:
 # | fig-cap: Posterior Predictive Infections
 fig, axes = plt.subplots(figsize=(6, 5))
 az.plot_hdi(
-    idata_weekday.posterior_predictive["negbinom_rv_dim_0"],
-    hdi_data=compute_eti(idata_weekday.posterior_predictive["negbinom_rv"], 0.9),
+    idata_weekday.posterior_predictive["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
+    hdi_data=compute_eti(
+        idata_weekday.posterior_predictive["negbinom_rv"], 0.9
+    ),
     color="C0",
     smooth=False,
     fill_kwargs={"alpha": 0.3},
     ax=axes,
 )
 
 az.plot_hdi(
-    idata_weekday.posterior_predictive["negbinom_rv_dim_0"],
-    hdi_data=compute_eti(idata_weekday.posterior_predictive["negbinom_rv"], 0.5),
+    idata_weekday.posterior_predictive["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
+    hdi_data=compute_eti(
+        idata_weekday.posterior_predictive["negbinom_rv"], 0.5
+    ),
     color="C0",
     smooth=False,
     fill_kwargs={"alpha": 0.6},
@@ -628,13 +655,17 @@ mean_latent_infection = np.mean(
 )
 
 plt.plot(
-    idata_weekday.posterior_predictive["negbinom_rv_dim_0"],
+    idata_weekday.posterior_predictive["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
     mean_latent_infection[0],
     color="C0",
     label="Mean",
 )
 plt.scatter(
-    idata_weekday.observed_data["negbinom_rv_dim_0"] + days_to_impute,
+    idata_weekday.observed_data["negbinom_rv_dim_0"]
+    + pad_size
+    + gen_int.size(),
     idata_weekday.observed_data["negbinom_rv"],
     color="black",
 )

model/pyproject.toml

@@ -12,7 +12,8 @@ exclude = [{path = "datasets/*.rds"}]
 [tool.poetry.dependencies]
 python = "^3.12"
 numpyro = "^0.15.0"
-jax = "^0.4.25"
+jax = "0.4.29"
+jaxlib = "0.4.29"
 numpy = "^1.26.4"
 polars = "^0.20.16"
 pillow = "^10.3.0" # See #56 on CDCgov/multisignal-epi-inference

model/src/pyrenew/model/admissionsmodel.py

@@ -5,9 +5,8 @@
 
 from typing import NamedTuple
 
-import jax.numpy as jnp
-import pyrenew.arrayutils as au
 from jax.typing import ArrayLike
+from pyrenew.deterministic import NullObservation
 from pyrenew.metaclass import Model, RandomVariable, _assert_sample_and_rtype
 from pyrenew.model.rtinfectionsrenewalmodel import RtInfectionsRenewalModel
 
@@ -99,6 +98,9 @@ def __init__(
         )
 
         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
@@ -178,13 +180,13 @@ def sample(
                 "Cannot pass both n_timepoints_to_simulate and data_observed_hosp_admissions."
             )
         elif n_timepoints_to_simulate is None:
-            n_timepoints = len(data_observed_hosp_admissions)
+            n_datapoints = len(data_observed_hosp_admissions)
         else:
-            n_timepoints = n_timepoints_to_simulate
+            n_datapoints = n_timepoints_to_simulate
 
         # Getting the initial quantities from the basic model
         basic_model = self.basic_renewal.sample(
-            n_timepoints_to_simulate=n_timepoints,
+            n_timepoints_to_simulate=n_datapoints,
             data_observed_infections=None,
             padding=padding,
             **kwargs,
@@ -199,35 +201,15 @@ def sample(
             latent_infections=basic_model.latent_infections,
             **kwargs,
         )
-        i0_size = len(latent_hosp_admissions) - n_timepoints
-        if self.hosp_admission_obs_process_rv is None:
-            observed_hosp_admissions = None
-        else:
-            if data_observed_hosp_admissions is None:
-                (
-                    observed_hosp_admissions,
-                    *_,
-                ) = self.hosp_admission_obs_process_rv.sample(
-                    mu=latent_hosp_admissions[i0_size + padding :],
-                    obs=data_observed_hosp_admissions,
-                    **kwargs,
-                )
-            else:
-                data_observed_hosp_admissions = au.pad_x_to_match_y(
-                    data_observed_hosp_admissions,
-                    latent_hosp_admissions,
-                    jnp.nan,
-                    pad_direction="start",
-                )
-
-                (
-                    observed_hosp_admissions,
-                    *_,
-                ) = self.hosp_admission_obs_process_rv.sample(
-                    mu=latent_hosp_admissions[i0_size + padding :],
-                    obs=data_observed_hosp_admissions[i0_size + padding :],
-                    **kwargs,
-                )
+
+        (
+            observed_hosp_admissions,
+            *_,
+        ) = self.hosp_admission_obs_process_rv.sample(
+            mu=latent_hosp_admissions[-n_datapoints:],
+            obs=data_observed_hosp_admissions,
+            **kwargs,
+        )
 
         return HospModelSample(
             Rt=basic_model.Rt,

model/src/pyrenew/model/rtinfectionsrenewalmodel.py

@@ -187,9 +187,9 @@ def sample(
                 "Cannot pass both n_timepoints_to_simulate and data_observed_infections."
             )
         elif n_timepoints_to_simulate is None:
-            n_timepoints = len(data_observed_infections)
+            n_timepoints = len(data_observed_infections) + padding
         else:
-            n_timepoints = n_timepoints_to_simulate
+            n_timepoints = n_timepoints_to_simulate + padding
         # Sampling from Rt (possibly with a given Rt, depending on
         # the Rt_process (RandomVariable) object.)
         Rt, *_ = self.Rt_process_rv.sample(
@@ -210,9 +210,6 @@ def sample(
             **kwargs,
         )
 
-        if data_observed_infections is not None:
-            data_observed_infections = data_observed_infections[padding:]
-
         observed_infections, *_ = self.infection_obs_process_rv.sample(
             mu=post_seed_latent_infections[padding:],
             obs=data_observed_infections,