Test convolve mode in hospitaladmissionspy

docs/source/tutorials/day_of_the_week.qmd

@@ -43,7 +43,7 @@ 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()
 gen_int = gen_int_array
-inf_hosp_int = inf_hosp_int["probability_mass"].to_numpy()
+inf_hosp_int_array = inf_hosp_int["probability_mass"].to_numpy()
 ```
 
 2. Next, we defined the model's components:
@@ -56,7 +56,7 @@ import jax.numpy as jnp
 import numpyro.distributions as dist
 
 inf_hosp_int = deterministic.DeterministicPMF(
-    name="inf_hosp_int", value=inf_hosp_int
+    name="inf_hosp_int", value=inf_hosp_int_array
 )
 
 hosp_rate = metaclass.DistributionalRV(
@@ -84,7 +84,7 @@ I0 = InfectionInitializationProcess(
         distribution=dist.LogNormal(loc=jnp.log(100), scale=jnp.log(1.75)),
     ),
     InitializeInfectionsExponentialGrowth(
-        gen_int_array.size,
+        inf_hosp_int_array.size,
         deterministic.DeterministicVariable(name="rate", value=0.05),
     ),
     t_unit=1,
@@ -201,11 +201,7 @@ hosp_model.run(
 out = hosp_model.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
-    obs_signal=np.pad(
-        daily_hosp_admits.astype(float),
-        (gen_int_array.size, 0),
-        constant_values=np.nan,
-    ),
+    obs_signal=daily_hosp_admits.astype(float),
 )
 ```
 
@@ -299,11 +295,7 @@ The new model with the day-of-the-week effect can be compared to the previous mo
 out = hosp_model.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
-    obs_signal=np.pad(
-        daily_hosp_admits.astype(float),
-        (gen_int_array.size, 0),
-        constant_values=np.nan,
-    ),
+    obs_signal=daily_hosp_admits.astype(float),
 )
 ```
 
@@ -314,10 +306,6 @@ out = hosp_model.plot_posterior(
 out_dow = hosp_model_dow.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
-    obs_signal=np.pad(
-        daily_hosp_admits.astype(float),
-        (gen_int_array.size, 0),
-        constant_values=np.nan,
-    ),
+    obs_signal=daily_hosp_admits.astype(float),
 )
 ```

docs/source/tutorials/hospital_admissions_model.qmd

@@ -118,17 +118,17 @@ 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()
 gen_int = gen_int_array
-inf_hosp_int = inf_hosp_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
-gen_int[:5], inf_hosp_int[:5]
+gen_int[:5], inf_hosp_int_array[:5]
 
 # Visualizing both quantities side by side
 fig, axs = plt.subplots(1, 2)
 
 axs[0].plot(gen_int)
 axs[0].set_title("Generation interval")
-axs[1].plot(inf_hosp_int)
+axs[1].plot(inf_hosp_int_array)
 axs[1].set_title("Infection to hospital admission interval")
 plt.show()
 ```
@@ -142,7 +142,7 @@ import jax.numpy as jnp
 import numpyro.distributions as dist
 
 inf_hosp_int = deterministic.DeterministicPMF(
-    name="inf_hosp_int", value=inf_hosp_int
+    name="inf_hosp_int", value=inf_hosp_int_array
 )
 
 hosp_rate = metaclass.DistributionalRV(
@@ -175,7 +175,7 @@ I0 = InfectionInitializationProcess(
         distribution=dist.LogNormal(loc=jnp.log(100), scale=jnp.log(1.75)),
     ),
     InitializeInfectionsExponentialGrowth(
-        gen_int_array.size,
+        inf_hosp_int_array.size,
         deterministic.DeterministicVariable(name="rate", value=0.05),
     ),
     t_unit=1,
@@ -313,11 +313,7 @@ We can use the `Model` object's `plot_posterior` method to visualize the model f
 out = hosp_model.plot_posterior(
     var="latent_hospital_admissions",
     ylab="Hospital Admissions",
-    obs_signal=np.pad(
-        daily_hosp_admits.astype(float),
-        (gen_int_array.size, 0),
-        constant_values=np.nan,
-    ),
+    obs_signal=daily_hosp_admits.astype(float),
 )
 ```
 
@@ -504,7 +500,7 @@ def compute_eti(dataset, eti_prob):
 
 fig, axes = plt.subplots(figsize=(6, 5))
 az.plot_hdi(
-    idata.prior_predictive["negbinom_rv_dim_0"] + gen_int.size(),
+    idata.prior_predictive["negbinom_rv_dim_0"],
     hdi_data=compute_eti(idata.prior_predictive["negbinom_rv"], 0.9),
     color="C0",
     smooth=False,
@@ -513,7 +509,7 @@ az.plot_hdi(
 )
 
 az.plot_hdi(
-    idata.prior_predictive["negbinom_rv_dim_0"] + gen_int.size(),
+    idata.prior_predictive["negbinom_rv_dim_0"],
     hdi_data=compute_eti(idata.prior_predictive["negbinom_rv"], 0.5),
     color="C0",
     smooth=False,
@@ -522,7 +518,7 @@ az.plot_hdi(
 )
 
 plt.scatter(
-    idata.observed_data["negbinom_rv_dim_0"] + gen_int.size(),
+    idata.observed_data["negbinom_rv_dim_0"],
     idata.observed_data["negbinom_rv"],
     color="black",
 )
@@ -538,7 +534,7 @@ And now we plot the posterior predictive distributions with a `{python} n_foreca
 ```{python}
 # | label: fig-output-posterior-predictive-forecast
 # | fig-cap: Posterior predictive admissions, including a forecast.
-x_data = idata.posterior_predictive["negbinom_rv_dim_0"] + gen_int.size()
+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(
@@ -569,7 +565,7 @@ plt.plot(
     label="Median",
 )
 plt.scatter(
-    idata.observed_data["negbinom_rv_dim_0"] + gen_int.size(),
+    idata.observed_data["negbinom_rv_dim_0"],
     idata.observed_data["negbinom_rv"],
     color="black",
 )

src/pyrenew/latent/hospitaladmissions.py

@@ -10,7 +10,11 @@
 
 import pyrenew.arrayutils as au
 from pyrenew.deterministic import DeterministicVariable
-from pyrenew.metaclass import RandomVariable, SampledValue
+from pyrenew.metaclass import (
+    RandomVariable,
+    SampledValue,
+    compute_incidence_observed_with_delay,
+)
 
 
 class HospitalAdmissionsSample(NamedTuple):
@@ -210,11 +214,11 @@ def sample(
             *_,
         ) = self.infection_to_admission_interval_rv(**kwargs)
 
-        latent_hospital_admissions = jnp.convolve(
-            infection_hosp_rate.value * latent_infections.value,
+        latent_hospital_admissions = compute_incidence_observed_with_delay(
+            infection_hosp_rate.value,
+            latent_infections.value,
             infection_to_admission_interval.value,
-            mode="full",
-        )[: latent_infections.value.shape[0]]
+        )
 
         # Applying the day of the week effect. For this we need to:
         # 1. Get the day of the week effect

src/pyrenew/metaclass.py

@@ -126,6 +126,38 @@ def _assert_sample_and_rtype(
     return None
 
 
+def compute_incidence_observed_with_delay(
+    incidence_to_observation_rate: ArrayLike,
+    latent_incidence: ArrayLike,
+    incidence_to_observation_delay_interval: ArrayLike,
+) -> ArrayLike:
+    """
+    Computes incidences observed according
+    to a given observation rate and based
+    on a delay interval.
+
+    Parameters
+    ----------
+    incidence_to_observation_rate: ArrayLike
+        The rate at which latent incidences are observed.
+    latent_incidence: ArrayLike
+        Incidence values based on the true underlying process.
+    incidence_to_observation_delay_interval: ArrayLike
+        Pmf of delay interval between incidence to observation.
+
+    Returns
+    --------
+    ArrayLike
+        The incidence after the observation delay.
+    """
+    delay_obs_incidence = jnp.convolve(
+        incidence_to_observation_rate * latent_incidence,
+        incidence_to_observation_delay_interval,
+        mode="valid",
+    )
+    return delay_obs_incidence
+
+
 class SampledValue(NamedTuple):
     """
     A container for a value sampled from a RandomVariable.

src/test/test_incidence_observed_with_delay.py

@@ -0,0 +1,49 @@
+# numpydoc ignore=GL08
+
+import jax.numpy as jnp
+import pytest
+from numpy.testing import assert_array_equal
+
+from pyrenew.metaclass import compute_incidence_observed_with_delay
+
+
+@pytest.mark.parametrize(
+    ["obs_rate", "latent_incidence", "delay_interval", "expected_output"],
+    [
+        [
+            jnp.array([1.0]),
+            jnp.array([1.0, 2.0, 3.0]),
+            jnp.array([1.0]),
+            jnp.array([1.0, 2.0, 3.0]),
+        ],
+        [
+            jnp.array([1.0, 0.1, 1.0]),
+            jnp.array([1.0, 2.0, 3.0]),
+            jnp.array([1.0]),
+            jnp.array([1.0, 0.2, 3.0]),
+        ],
+        [
+            jnp.array([1.0]),
+            jnp.array([1.0, 2.0, 3.0]),
+            jnp.array([0.5, 0.5]),
+            jnp.array([1.5, 2.5]),
+        ],
+        [
+            jnp.array([1.0]),
+            jnp.array([0, 2.0, 4.0]),
+            jnp.array([0.25, 0.5, 0.25]),
+            jnp.array([2]),
+        ],
+    ],
+)
+def test(obs_rate, latent_incidence, delay_interval, expected_output):
+    """
+    Tests for helper function to compute
+    incidence observed with a delay
+    """
+    result = compute_incidence_observed_with_delay(
+        obs_rate,
+        latent_incidence,
+        delay_interval,
+    )
+    assert_array_equal(result, expected_output)