Implements SampledValue

docs/source/tutorials/time.qmd

@@ -10,7 +10,7 @@ The fundamental time unit should represent a period of fixed (or approximately f
 
 For many infectious disease renewal models of interest, the fundamental time unit will be days, and we will proceed with this tutorial treating days as our fundamental unit.
 
- `pyrenew` deals with time having `RandomVariable`s carry information about (i) their own time unit expressed relative to the fundamental unit (`t_unit`) and (ii) the starting time, `t_start`, measured relative to `t = 0` in model time in fundamental time units.
+ `pyrenew` deals with time having `RandomVariable`s carry information about (i) their own time unit expressed relative to the fundamental unit (`t_unit`) and (ii) the starting time, `t_start`, measured relative to `t = 0` in model time in fundamental time units. Moreover, return values from `RandomVariable.sample()` are namedtuples with `TimeArray` objects that carry the same information.
 
 The tuple `(t_unit, t_start)` can encode different types of time series data. For example:
 
@@ -31,14 +31,14 @@ The `PeriodicBroadcaster()` class provides a way of tiling and repeating data ac
 
 The following section describes some preliminary design principles that may be included in future versions of `pyrenew`.
 
-### Validation
-
-With random variables possibly spanning different time scales, *e.g.*, weekly, daily, hourly, the metaclass `Model` should ensure random variables within the model share the same time unit.
-
 ### Array alignment
 
 Using `t_unit` and `t_start`, random variables should be able to align input and output data. For example, in the case of the `RtInfectionsRenewalModel()`, the computed values of `Rt` and `infections` are padded left with `nan` values to account for the seeding process. Instead, we expect to either pre-process the padding leveraging the `t_start` information of the involved variables or simplify the process via a function call that aligns the arrays. A possible implementation could be a method `align()` that takes a list of random variables and aligns them based on the `t_unit` and `t_start` information, e.g.:
 
 ```python
 Rt_aligned, infections_aligned = align([Rt, infections])
 ```
+
+### Retrieving time information from sites
+
+Since numpyro only stores Jax arrays, we cannot store the time information in the arrays themselves. Next iterations of `pyrenew` should include a way to retrieve the time information from the sites of the model after running them.

model/src/pyrenew/deterministic/deterministic.py

@@ -6,7 +6,7 @@
 import jax.numpy as jnp
 import numpyro as npro
 from jax.typing import ArrayLike
-from pyrenew.metaclass import RandomVariable
+from pyrenew.metaclass import RandomVariable, TimeArray
 
 
 class DeterministicVariable(RandomVariable):
@@ -19,6 +19,8 @@ def __init__(
         self,
         vars: ArrayLike,
         name: str,
+        t_start: int | None = None,
+        t_unit: int | None = None,
     ) -> None:
         """Default constructor
 
@@ -28,13 +30,18 @@ def __init__(
             A tuple with arraylike objects.
         name : str, optional
             A name to assign to the process.
+        t_start : int, optional
+            The start time of the process.
+        t_unit : int, optional
+            The unit of time relative to the model's fundamental (smallest) time unit.
 
         Returns
         -------
         None
         """
 
         self.validate(vars)
+        self.set_timeseries(t_start, t_unit)
         self.vars = jnp.atleast_1d(vars)
         self.name = name
 
@@ -83,8 +90,13 @@ def sample(
         Returns
         -------
         tuple
-            Containing the stored values during construction.
+            Containing the stored values during construction wrapped in a TimeArrayß.
         """
         if record:
             npro.deterministic(self.name, self.vars)
-        return (self.vars,)
+        return (
+            TimeArray(
+            array = self.vars,
+            t_start=self.t_start,
+            t_unit=self.t_unit,
+            ),)

model/src/pyrenew/deterministic/deterministicpmf.py

@@ -18,6 +18,8 @@ def __init__(
         vars: ArrayLike,
         name: str,
         tol: float = 1e-5,
+        t_start: int | None = None,
+        t_unit: int | None = None,
     ) -> None:
         """
         Default constructor
@@ -36,7 +38,12 @@ def __init__(
         tol : float, optional
             Passed to pyrenew.distutil.validate_discrete_dist_vector. Defaults
             to 1e-5.
-
+        t_start : int, optional
+            The start time of the process.
+        t_unit : int, optional
+            The unit of time relative to the model's fundamental (smallest)
+            time unit.    
+
         Returns
         -------
         None
@@ -46,7 +53,12 @@ def __init__(
             tol=tol,
         )
 
-        self.basevar = DeterministicVariable(vars, name)
+        self.basevar = DeterministicVariable(
+            vars=vars,
+            name=name,
+            t_start=t_start,
+            t_unit=t_unit,
+            )
 
         return None
 
@@ -82,7 +94,7 @@ def sample(
         Returns
         -------
         tuple
-            Containing the stored values during construction.
+            Containing the stored values during construction wrapped in a TimeArray.
         """
 
         return self.basevar.sample(**kwargs)

model/src/pyrenew/deterministic/nullrv.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from jax.typing import ArrayLike
+from pyrenew.metaclass import TimeArray
 from pyrenew.deterministic.deterministic import DeterministicVariable
 
 
@@ -46,10 +47,10 @@ def sample(
         Returns
         -------
         tuple
-            Containing None.
+            Containing a TimeArray with None.
         """
 
-        return (None,)
+        return (TimeArray(None),)
 
 
 class NullProcess(NullVariable):
@@ -95,10 +96,10 @@ def sample(
         Returns
         -------
         tuple
-            Containing None.
+            Containing a TimeArray with None.
         """
 
-        return (None,)
+        return (TimeArray(None),)
 
 
 class NullObservation(NullVariable):
@@ -151,7 +152,7 @@ def sample(
         Returns
         -------
         tuple
-            Containing None.
+            Containing a TimeArray with None.
         """
 
-        return (None,)
+        return (TimeArray(None),)

model/src/pyrenew/deterministic/process.py

@@ -1,6 +1,7 @@
 # numpydoc ignore=GL08
 
 import jax.numpy as jnp
+from pyrenew.metaclass import TimeArray
 from pyrenew.deterministic.deterministic import DeterministicVariable
 
 
@@ -29,14 +30,25 @@ def sample(
         Returns
         -------
         tuple
-            Containing the stored values during construction.
+            Containing the stored values during construction wrapped in a TimeArray.
         """
 
         res, *_ = super().sample(**kwargs)
 
-        dif = duration - res.shape[0]
+        dif = duration - res.array.shape[0]
 
         if dif > 0:
-            return (jnp.hstack([res, jnp.repeat(res[-1], dif)]),)
-
-        return (res[:duration],)
+            return (
+                TimeArray(
+                    jnp.hstack([res.array, jnp.repeat(res.array[-1], dif)]), t_start=self.t_start,
+                    t_unit=self.t_unit,
+                    ),
+                    )
+
+        return (
+            TimeArray(
+                array=res.array[:duration],
+                t_start=self.t_start,
+                t_unit=self.t_unit
+            ),
+        )

model/src/pyrenew/latent/hospitaladmissions.py

@@ -9,7 +9,7 @@
 import numpyro as npro
 from jax.typing import ArrayLike
 from pyrenew.deterministic import DeterministicVariable
-from pyrenew.metaclass import RandomVariable
+from pyrenew.metaclass import RandomVariable, TimeArray
 
 
 class HospitalAdmissionsSample(NamedTuple):
@@ -20,12 +20,12 @@ class HospitalAdmissionsSample(NamedTuple):
     ----------
     infection_hosp_rate : float, optional
         The infection-to-hospitalization rate. Defaults to None.
-    latent_hospital_admissions : ArrayLike or None
+    latent_hospital_admissions : TimeArray or None
         The computed number of hospital admissions. Defaults to None.
     """
 
     infection_hosp_rate: float | None = None
-    latent_hospital_admissions: ArrayLike | None = None
+    latent_hospital_admissions: TimeArray | None = None
 
     def __repr__(self):
         return f"HospitalAdmissionsSample(infection_hosp_rate={self.infection_hosp_rate}, latent_hospital_admissions={self.latent_hospital_admissions})"
@@ -162,7 +162,7 @@ def sample(
 
         Parameters
         ----------
-        latent : ArrayLike
+        latent : ArrayLike or TimeArray
             Latent infections.
         **kwargs : dict, optional
             Additional keyword arguments passed through to internal `sample()`
@@ -175,7 +175,7 @@ def sample(
 
         infection_hosp_rate, *_ = self.infect_hosp_rate_rv(**kwargs)
 
-        infection_hosp_rate_t = infection_hosp_rate * latent_infections
+        infection_hosp_rate_t = infection_hosp_rate.array * latent_infections
 
         (
             infection_to_admission_interval,
@@ -184,25 +184,31 @@ def sample(
 
         latent_hospital_admissions = jnp.convolve(
             infection_hosp_rate_t,
-            infection_to_admission_interval,
+            infection_to_admission_interval.array,
             mode="full",
         )[: infection_hosp_rate_t.shape[0]]
 
         # Applying the day of the week effect
         latent_hospital_admissions = (
             latent_hospital_admissions
-            * self.day_of_week_effect_rv(**kwargs)[0]
+            * self.day_of_week_effect_rv(**kwargs)[0].array
         )
 
         # Applying probability of hospitalization effect
         latent_hospital_admissions = (
-            latent_hospital_admissions * self.hosp_report_prob_rv(**kwargs)[0]
+            latent_hospital_admissions *
+            self.hosp_report_prob_rv(**kwargs)[0].array
         )
 
         npro.deterministic(
             self.latent_hospital_admissions_varname, latent_hospital_admissions
         )
 
         return HospitalAdmissionsSample(
-            infection_hosp_rate, latent_hospital_admissions
+            infection_hosp_rate=infection_hosp_rate,
+            latent_hospital_admissions=TimeArray(
+                array=latent_hospital_admissions,
+                t_start=self.infection_to_admission_interval_rv.t_start,
+                t_unit=self.infection_to_admission_interval_rv.t_unit,
+                )
         )

model/src/pyrenew/latent/infection_initialization_method.py

@@ -177,6 +177,7 @@ def seed_infections(self, I_pre_seed: ArrayLike):
                 f"I_pre_seed must be an array of size 1. Got size {I_pre_seed.size}."
             )
         (rate,) = self.rate()
+        rate = rate.array
         if rate.size != 1:
             raise ValueError(
                 f"rate must be an array of size 1. Got size {rate.size}."

model/src/pyrenew/latent/infection_initialization_process.py

@@ -4,7 +4,7 @@
 from pyrenew.latent.infection_initialization_method import (
     InfectionInitializationMethod,
 )
-from pyrenew.metaclass import RandomVariable
+from pyrenew.metaclass import RandomVariable, TimeArray
 
 
 class InfectionInitializationProcess(RandomVariable):
@@ -96,7 +96,13 @@ def sample(self) -> tuple:
             a tuple where the only element is an array with the number of seeded infections at each time point.
         """
         (I_pre_seed,) = self.I_pre_seed_rv()
-        infection_seeding = self.infection_seed_method(I_pre_seed)
+        infection_seeding = self.infection_seed_method(I_pre_seed.array)
         npro.deterministic(self.name, infection_seeding)
 
-        return (infection_seeding,)
+        return (
+            TimeArray(
+                array=infection_seeding,
+                t_start=self.t_start,
+                t_unit=self.t_unit,
+            ),
+            )

model/src/pyrenew/latent/infections.py

@@ -8,7 +8,7 @@
 import jax.numpy as jnp
 import pyrenew.latent.infection_functions as inf
 from jax.typing import ArrayLike
-from pyrenew.metaclass import RandomVariable
+from pyrenew.metaclass import RandomVariable, TimeArray
 
 
 class InfectionsSample(NamedTuple):
@@ -17,7 +17,7 @@ class InfectionsSample(NamedTuple):
 
     Attributes
     ----------
-    post_initialization_infections : ArrayLike | None, optional
+    post_initialization_infections : TimeArray | None, optional
         The estimated latent infections. Defaults to None.
     """
 
@@ -97,4 +97,4 @@ def sample(
             reversed_generation_interval_pmf=gen_int_rev,
         )
 
-        return InfectionsSample(post_initialization_infections)
+        return InfectionsSample(TimeArray(post_initialization_infections))

model/src/pyrenew/latent/infectionswithfeedback.py

@@ -8,7 +8,9 @@
 import pyrenew.arrayutils as au
 import pyrenew.latent.infection_functions as inf
 from numpy.typing import ArrayLike
-from pyrenew.metaclass import RandomVariable, _assert_sample_and_rtype
+from pyrenew.metaclass import (
+    RandomVariable, _assert_sample_and_rtype, TimeArray
+)
 
 
 class InfectionsRtFeedbackSample(NamedTuple):
@@ -159,6 +161,7 @@ def sample(
         inf_feedback_strength, *_ = self.infection_feedback_strength(
             **kwargs,
         )
+        inf_feedback_strength = inf_feedback_strength.array
 
         # Making sure inf_feedback_strength spans the Rt length
         if inf_feedback_strength.size == 1:
@@ -177,7 +180,7 @@ def sample(
         # Sampling inf feedback pmf
         inf_feedback_pmf, *_ = self.infection_feedback_pmf(**kwargs)
 
-        inf_fb_pmf_rev = jnp.flip(inf_feedback_pmf)
+        inf_fb_pmf_rev = jnp.flip(inf_feedback_pmf.array)
 
         (
             post_initialization_infections,
@@ -195,6 +198,6 @@ def sample(
         npro.deterministic("Rt_adjusted", Rt_adj)
 
         return InfectionsRtFeedbackSample(
-            post_initialization_infections=post_initialization_infections,
-            rt=Rt_adj,
+            post_initialization_infections=TimeArray(post_initialization_infections),
+            rt=TimeArray(Rt_adj),
         )