merge origin main

README.md

@@ -46,6 +46,11 @@ Install via pip with
 pip install git+https://github.com/CDCgov/PyRenew@main
 ```
 
+## Models Implemented With PyRenew
+
+- [CDCgov/pyrenew-covid-wastewater](https://github.com/CDCgov/pyrenew-covid-wastewater): _Models and infrastructure for forecasting COVID-19 hospitalizations using wastewater data with PyRenew._
+- [CDCgov/pyrenew-flu-light](https://github.com/CDCgov/pyrenew-flu-light/): _An instantiation in PyRenew of an influenza forecasting model used in the 2023-24 respiratory season._
+
 ## Resources
 
 * [The MSR Website](https://cdcgov.github.io/PyRenew/tutorials/index.html) provides general documentation and tutorials on using MSR.

docs/source/developer_documentation.rst

@@ -41,6 +41,10 @@ PyRenew Principles
    -  Using default site names in a ``RandomVariable`` is discouraged. Only use default site names at the ``Model`` level.
    -  Use ``DeterministicVariable``\ s instead of constants within a model.
 
+-  Multidimensional array conventions
+
+   -  In a multidimensional array of timeseries, time is always the first dimension. By default, ``numpyro.scan`` builds by augmenting the first dimension, and variables are often scanned over time, making default output of scan over time sensible.
+
 Adding Documentation to Sphinx
 ------------------------------
 

docs/source/msei_reference/distributions.rst

@@ -0,0 +1,7 @@
+Distributions
+===========
+
+.. automodule:: pyrenew.distributions
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/source/tutorials/basic_renewal_model.qmd

@@ -125,7 +125,6 @@ I0 = InfectionInitializationProcess(
     "I0_initialization",
     DistributionalVariable(name="I0", distribution=dist.LogNormal(2.5, 1)),
     InitializeInfectionsZeroPad(pmf_array.size),
-    t_unit=1,
 )
 
 
@@ -152,9 +151,9 @@ class MyRt(RandomVariable):
         rt_init_rv = DistributionalVariable(
             name="init_log_rt", distribution=dist.Normal(0, 0.2)
         )
-        init_rt, *_ = rt_init_rv.sample()
+        init_rt = rt_init_rv.sample()
 
-        return rt_rv.sample(n=n, init_vals=init_rt.value, **kwargs)
+        return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
 
 
 rt_proc = MyRt()
@@ -220,11 +219,11 @@ import matplotlib.pyplot as plt
 fig, axs = plt.subplots(1, 2)
 
 # Rt plot
-axs[0].plot(sim_data.Rt.value)
+axs[0].plot(sim_data.Rt)
 axs[0].set_ylabel("Rt")
 
 # Infections plot
-axs[1].plot(sim_data.observed_infections.value)
+axs[1].plot(sim_data.observed_infections)
 axs[1].set_ylabel("Infections")
 
 fig.suptitle("Basic renewal model")
@@ -242,7 +241,7 @@ import jax
 model1.run(
     num_warmup=2000,
     num_samples=1000,
-    data_observed_infections=sim_data.observed_infections.value,
+    data_observed_infections=sim_data.observed_infections,
     rng_key=jax.random.key(54),
     mcmc_args=dict(progress_bar=False, num_chains=2),
 )

docs/source/tutorials/day_of_the_week.qmd

@@ -89,7 +89,6 @@ I0 = InfectionInitializationProcess(
         n_initialization_points,
         deterministic.DeterministicVariable(name="rate", value=0.05),
     ),
-    t_unit=1,
 )
 
 # Generation interval and Rt
@@ -110,11 +109,11 @@ class MyRt(metaclass.RandomVariable):
 
     def sample(self, n: int, **kwargs) -> tuple:
         # Standard deviation of the random walk
-        sd_rt, *_ = self.sd_rv()
+        sd_rt = self.sd_rv()
 
         # Random walk step
         step_rv = randomvariable.DistributionalVariable(
-            name="rw_step_rv", distribution=dist.Normal(0, sd_rt.value)
+            name="rw_step_rv", distribution=dist.Normal(0, sd_rt)
         )
 
         rt_init_rv = randomvariable.DistributionalVariable(
@@ -133,9 +132,9 @@ class MyRt(metaclass.RandomVariable):
             base_rv=base_rv,
             transforms=transformation.ExpTransform(),
         )
-        init_rt, *_ = rt_init_rv.sample()
+        init_rt = rt_init_rv.sample()
 
-        return rt_rv.sample(n=n, init_vals=init_rt.value, **kwargs)
+        return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
 
 
 rtproc = MyRt(
@@ -168,7 +167,7 @@ obs = observation.NegativeBinomialObservation(
 )
 ```
 
-4. And finally, we built the model:
+4. And finally, we build the model:
 
 ```{python}
 # | label: init-model
@@ -282,8 +281,10 @@ As a result, we can see the posterior distribution of our novel day-of-the-week
 # | label: fig-output-day-of-week
 # | fig-cap: Day of the week effect
 out = hosp_model_dow.plot_posterior(
-    var="dayofweek_effect", ylab="Day of the Week Effect", samples=500
+    var="dayofweek_effect_raw", ylab="Day of the Week Effect", samples=500
 )
+
+sp = hosp_model_dow.spread_draws(["dayofweek_effect"])
 ```
 
 The new model with the day-of-the-week effect can be compared to the previous model without the effect. Finally, let's reproduce the figure without the day-of-the-week effect, and then plot the new model with the effect:

docs/source/tutorials/extending_pyrenew.qmd

@@ -55,7 +55,6 @@ I0 = InfectionInitializationProcess(
         gen_int_array.size,
         DeterministicVariable(name="rate", value=0.05),
     ),
-    t_unit=1,
 )
 
 latent_infections = InfectionsWithFeedback(
@@ -85,9 +84,9 @@ class MyRt(RandomVariable):
         rt_init_rv = DistributionalVariable(
             name="init_log_rt", distribution=dist.Normal(0, 0.2)
         )
-        init_rt, *_ = rt_init_rv.sample()
+        init_rt = rt_init_rv.sample()
 
-        return rt_rv.sample(n=n, init_vals=init_rt.value, **kwargs)
+        return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
 ```
 
 With all the components defined, we can build the model:
@@ -118,7 +117,7 @@ with numpyro.handlers.seed(rng_seed=223):
 import matplotlib.pyplot as plt
 
 fig, ax = plt.subplots()
-ax.plot(model0_samp.latent_infections.value)
+ax.plot(model0_samp.latent_infections)
 ax.set_xlabel("Time")
 ax.set_ylabel("Infections")
 plt.show()
@@ -164,7 +163,7 @@ from collections import namedtuple
 # Creating a tuple to store the output
 InfFeedbackSample = namedtuple(
     typename="InfFeedbackSample",
-    field_names=["infections", "rt"],
+    field_names=["post_initialization_infections", "rt"],
     defaults=(None, None),
 )
 ```
@@ -175,7 +174,7 @@ The next step is to create the actual class. The bulk of its implementation lies
 # | label: new-model-def
 # | code-line-numbers: true
 # Creating the class
-from pyrenew.metaclass import RandomVariable, SampledValue
+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
@@ -219,11 +218,11 @@ class InfFeedback(RandomVariable):
         I0_vec = I0[-gen_int_rev.size :]
 
         # Sampling inf feedback strength and adjusting the shape
-        inf_feedback_strength, *_ = self.infection_feedback_strength(
+        inf_feedback_strength = self.infection_feedback_strength(
             **kwargs,
         )
 
-        inf_feedback_strength = jnp.atleast_1d(inf_feedback_strength.value)
+        inf_feedback_strength = jnp.atleast_1d(inf_feedback_strength)
 
         inf_feedback_strength = au.pad_x_to_match_y(
             x=inf_feedback_strength,
@@ -232,8 +231,8 @@ class InfFeedback(RandomVariable):
         )
 
         # Sampling inf feedback and adjusting the shape
-        inf_feedback_pmf, *_ = self.infection_feedback_pmf(**kwargs)
-        inf_fb_pmf_rev = jnp.flip(inf_feedback_pmf.value)
+        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(
@@ -250,20 +249,18 @@ class InfFeedback(RandomVariable):
         # Preparing theoutput
 
         return InfFeedbackSample(
-            infections=SampledValue(all_infections),
-            rt=SampledValue(Rt_adj),
+            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. **Calls to `RandomVariable()`**: All calls to `RandomVariable()` are expected to return a tuple or named tuple. In our implementation, we capture the output of `infection_feedback_strength()` and `infection_feedback_pmf()` in the variables `inf_feedback_strength` and `inf_feedback_pmf`, respectively, disregarding the other outputs (i.e., using `*_`).
-
-3. **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.
+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.
 
-4. **Return type of `InfFeedback()`**: As said before, the `sample()` method should return a tuple or named tuple. In our case, we return a named tuple `InfFeedbackSample` with two fields: `infections` and `rt`.
+3. **Return type of `InfFeedback()`**: As said before, the `sample()` method should return a tuple or named tuple. In our case, we return a named tuple `InfFeedbackSample` with two fields: `infections` and `rt`.
 
 ```{python}
 # | label: simulation2
@@ -293,8 +290,8 @@ Comparing `model0` with `model1`, these two should match:
 import matplotlib.pyplot as plt
 
 fig, ax = plt.subplots(ncols=2)
-ax[0].plot(model0_samp.latent_infections.value)
-ax[1].plot(model1_samp.latent_infections.value)
+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")

docs/source/tutorials/hospital_admissions_model.qmd

@@ -179,7 +179,6 @@ I0 = InfectionInitializationProcess(
         n_initialization_points,
         deterministic.DeterministicVariable(name="rate", value=0.05),
     ),
-    t_unit=1,
 )
 
 # Generation interval and Rt
@@ -207,9 +206,9 @@ class MyRt(metaclass.RandomVariable):
         rt_init_rv = randomvariable.DistributionalVariable(
             name="init_log_rt", distribution=dist.Normal(0, 0.2)
         )
-        init_rt, *_ = rt_init_rv.sample()
+        init_rt = rt_init_rv.sample()
 
-        return rt_rv.sample(n=n, init_vals=init_rt.value, **kwargs)
+        return rt_rv.sample(n=n, init_vals=init_rt, **kwargs)
 
 
 rtproc = MyRt()
@@ -256,7 +255,6 @@ import numpy as np
 
 timeframe = 120
 
-
 with numpyro.handlers.seed(rng_seed=223):
     simulated_data = hosp_model.sample(n_datapoints=timeframe)
 ```
@@ -269,11 +267,11 @@ import matplotlib.pyplot as plt
 fig, axs = plt.subplots(1, 2)
 
 # Rt plot
-axs[0].plot(simulated_data.Rt.value)
+axs[0].plot(simulated_data.Rt)
 axs[0].set_ylabel("Simulated Rt")
 
 # Admissions plot
-axs[1].plot(simulated_data.observed_hosp_admissions.value, "-o")
+axs[1].plot(simulated_data.observed_hosp_admissions, "-o")
 axs[1].set_ylabel("Simulated Admissions")
 
 fig.suptitle("Basic renewal model")
@@ -483,7 +481,7 @@ 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.values.T
+    return eti_bdry.T
 
 
 fig, axes = plt.subplots(figsize=(6, 5))

docs/source/tutorials/periodic_effects.qmd

@@ -46,7 +46,7 @@ with numpyro.handlers.seed(rng_seed=20):
 # Plotting the Rt values
 import matplotlib.pyplot as plt
 
-plt.step(np.arange(len(sim_data.rt.value)), sim_data.rt.value, where="post")
+plt.step(np.arange(len(sim_data)), sim_data, where="post")
 plt.xlabel("Time")
 plt.ylabel("Rt")
 plt.title("Simulated Rt values")
@@ -79,7 +79,6 @@ dayofweek = process.DayOfWeekEffect(
     quantity_to_broadcast=randomvariable.DistributionalVariable(
         name="simp", distribution=mysimplex
     ),
-    t_start=0,
 )
 ```
 
@@ -92,9 +91,7 @@ with numpyro.handlers.seed(rng_seed=20):
 # Plotting the effect values
 import matplotlib.pyplot as plt
 
-plt.step(
-    np.arange(len(sim_data.value.value)), sim_data.value.value, where="post"
-)
+plt.step(np.arange(len(sim_data)), sim_data, where="post")
 plt.xlabel("Time")
 plt.ylabel("Effect size")
 plt.title("Simulated Day of Week Effect values")

pyrenew/arrayutils.py

@@ -109,7 +109,7 @@ class PeriodicProcessSample(NamedTuple):
     value: ArrayLike | None = None
 
     def __repr__(self):
-        return f"PeriodicProcessSample(value={self.value})"
+        return f"PeriodicProcessSample(value={self})"
 
 
 def tile_until_n(