From 39061ce0ac96423eef1903cca8c4f4cbb772a3fd Mon Sep 17 00:00:00 2001 From: "Dylan H. Morris" Date: Mon, 9 Feb 2026 18:49:46 -0500 Subject: [PATCH 1/4] Class skeleton --- pyrenew/randomvariable/pmf.py | 238 ++++++++++++++++++++++++++++++++++ 1 file changed, 238 insertions(+) create mode 100644 pyrenew/randomvariable/pmf.py diff --git a/pyrenew/randomvariable/pmf.py b/pyrenew/randomvariable/pmf.py new file mode 100644 index 00000000..7e2661ad --- /dev/null +++ b/pyrenew/randomvariable/pmf.py @@ -0,0 +1,238 @@ +""" +Class for vectors that represent discrete probability +mass functions. +""" + +from abc import abstractmethod + +import jax.numpy as jnp +from jax.typing import ArrayLike + +from pyrenew.metaclass import RandomVariable + + +class PMFVector(RandomVariable): + """ + Abstract [`pyrenew.metaclass.RandomVariable`][] that + represents a probability mass function (PMF) as a + vector of probabilities that sums to 1. + + These vectors of probabilities can be deterministic + or stochastic in concrete subclasses. + """ + + def __init__(self, values: ArrayLike, **kwargs) -> None: + """ + Default constructor. + + Parameters + ---------- + values + Vector of values of the same shape as the + output of [`self.sample`][], representing the + values of the variable to which those probabilities + correspond. + + **kwargs + Additional keyword arguments passed to the parent + constructor. + + Returns + ------- + None + """ + self.values = values + super().__init__(**kwargs) + + @abstractmethod + def sample(self, **kwargs) -> ArrayLike: + """ + Sample a vector of probabilities. + """ + pass + + +class DelayPMF(PMFVector): + """ + Subclass of [`pyrenew.randomvariable.PMFVector`] that + represents a discrete time delay PMF. + + Discrete delay PMFs are fundamental to discrete-time + renewal modeling. They are used to represent generation + interval distributions (minimum value 1 time unit), as well + as delays between infectious and various downstream events + (e.g. an infection-to-hospital-admission delay distribution, + minimum value 0 time units). + + Enforces continguousness. [`self.values`][] must be + an array of consecutive integers representing time units. + + Enforces either 0 or 1 indexing. Shortest represented delay must + be either 0 or 1 time unit. + """ + + def __init__(self, min_delay: int, max_delay: int, **kwargs) -> None: + """ + Default constructor + + Parameters + ---------- + min_delay + Shortest possible delay in time units. + Will become the first value of [`self.values`][] + (corresponding to the zeroth entry of the probability + vector returned by [`self.sample`][]). Must be an integer + greater than or equal to 0. + + max_delay + Longest possible delay in time units. + Will become the final value of [`self.values`][] + (corresponding to the final entry of the probability + vector returned by [`self.sample`][]). Must be an + integer greater than or equal to `min_delay`. + + **kwargs + Additional keyword arguments passed to the parent + constructor. + + Returns + ------- + None + + Raises + ------ + ValueError + If min_delay and max_delay do not satisfy the specified + constraints. + """ + if not all([isinstance(x, int) for x in [min_delay, max_delay]]): + raise ValueError("min_delay and max_delay must be integers.") + + if not min_delay > 0: + raise ValueError("min_delay must be greater than or equal to 0.") + if not max_delay >= min_delay: + raise ValueError("max_delay must be greater than or equal to min_delay") + + super().__init__(values=jnp.arange(min_delay, max_delay + 1)) + + @property + def min_delay(self) -> int: + """ + The minimum possible delay in integer time units. + Corresponds to the zeroth entry of the probability vector + returned by [`self.sample`][]. + + Returns + ------- + int + The value of the minimum possible delay. + """ + return self.values[0] + + @property + def max_delay(self) -> int: + """ + The maximum possible delay in integer time units. + Corresponds to the final entry of the probability vector + returned by [`self.sample`][]. + + Returns + ------- + int + The value of the maximum possible delay. + """ + return self.values[-1] + + +class NonnegativeDelayPMF(DelayPMF): + """ + Subclass of [`pyrenew.randomvariable.DelayPMF`] that + represents the PMF of a delay that can possibly be + 0 time units (i.e. no delay). + + Enforces a `min_delay` value of 0. + + In PyRenew, we have a convention of using + `NonnegativeDelayPMF`s to represent discrete-time delays + from infection to ascertained observation. This + simplifies the computation of predicted observations. + """ + + def __init__(self, max_delay: int) -> None: + """ + Default constructor. + + Parameters + ---------- + max_delay + Longest possible delay in time units. + Will become the final value of [`self.values`][] + (corresponding to the final entry of the probability + vector returned by [`self.sample`][]). Must be an + integer greater than or equal to 0. + + Returns + ------- + None + + Raises + ------ + ValueError + If max_delay does not satisfy the specified constraints. + """ + super().__init__(min_delay=0, max_delay=max_delay) + + +class PositiveDelayPMF(DelayPMF): + """ + Subclass of [`pyrenew.randomvariable.DelayPMF`] that + represents the PMF of a strictly positive discrete time + delay (i.e. of at least 1 time unit). + + Enforces a `min_delay` value of 1. + + In PyRenew, we have a convention of using + `PositiveDelayPMF`s to represent generation interval + distributions. This simplifies the computation of the + renewal equation. + """ + + def __init__(self, max_delay: int) -> None: + """ + Default constructor. + + Parameters + ---------- + max_delay + Longest possible delay in time units. + Will become the final value of [`self.values`][] + (corresponding to the final entry of the probability + vector returned by [`self.sample`][]). Must be an + integer greater than or equal to 1. + + Returns + ------- + None + + Raises + ------ + ValueError + If max_delay does not satisfy the specified constraints. + """ + + super().__init__(min_delay=1, max_delay=max_delay) + + +class GenerationIntervalPMF(PositiveDelayPMF): + """ + Subclass of [`pyrenew.randomvariable.PositiveDelayPMF`] that + represents the PMF of a generation interval distribution. + """ + + +class AscertainmentDelayPMF(NonnegativeDelayPMF): + """ + Subclass of [`pyrenew.randomvariable.NonnegativeDelayPMF`] that + represents the PMF of a delay from an event to when it is + ascertained + """ From 6506c241c0599d978c2de763b30df34caa3a117e Mon Sep 17 00:00:00 2001 From: "Dylan H. Morris" Date: Tue, 10 Feb 2026 09:41:29 -0500 Subject: [PATCH 2/4] Expose new classes --- pyrenew/randomvariable/__init__.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pyrenew/randomvariable/__init__.py b/pyrenew/randomvariable/__init__.py index c599d101..8275ead0 100644 --- a/pyrenew/randomvariable/__init__.py +++ b/pyrenew/randomvariable/__init__.py @@ -5,6 +5,14 @@ DynamicDistributionalVariable, StaticDistributionalVariable, ) +from pyrenew.randomvariable.pmf import ( + AscertainmentDelayPMF, + DelayPMF, + GenerationIntervalPMF, + NonnegativeDelayPMF, + PMFVector, + PositiveDelayPMF, +) from pyrenew.randomvariable.transformedvariable import TransformedVariable __all__ = [ @@ -12,4 +20,10 @@ "StaticDistributionalVariable", "DynamicDistributionalVariable", "TransformedVariable", + "PMFVector", + "DelayPMF", + "PositiveDelayPMF", + "NonnegativeDelayPMF", + "GenerationIntervalPMF", + "AscertainmentDelayPMF", ] From 26a39d45d924b6c34ec04ae17b285ec354a78f70 Mon Sep 17 00:00:00 2001 From: "Dylan H. Morris" Date: Wed, 25 Feb 2026 22:19:48 -0500 Subject: [PATCH 3/4] Add names --- pyrenew/randomvariable/pmf.py | 86 ++++++++++++++++++++++++++++++++--- 1 file changed, 79 insertions(+), 7 deletions(-) diff --git a/pyrenew/randomvariable/pmf.py b/pyrenew/randomvariable/pmf.py index 7e2661ad..c409f462 100644 --- a/pyrenew/randomvariable/pmf.py +++ b/pyrenew/randomvariable/pmf.py @@ -8,6 +8,7 @@ import jax.numpy as jnp from jax.typing import ArrayLike +from pyrenew.deterministic import DeterministicVariable from pyrenew.metaclass import RandomVariable @@ -21,12 +22,15 @@ class PMFVector(RandomVariable): or stochastic in concrete subclasses. """ - def __init__(self, values: ArrayLike, **kwargs) -> None: + def __init__(self, name: str, values: ArrayLike, **kwargs) -> None: """ Default constructor. Parameters ---------- + name + Name for the random variable. + values Vector of values of the same shape as the output of [`self.sample`][], representing the @@ -41,6 +45,7 @@ def __init__(self, values: ArrayLike, **kwargs) -> None: ------- None """ + self.name = name self.values = values super().__init__(**kwargs) @@ -71,12 +76,15 @@ class DelayPMF(PMFVector): be either 0 or 1 time unit. """ - def __init__(self, min_delay: int, max_delay: int, **kwargs) -> None: + def __init__(self, name: str, min_delay: int, max_delay: int, **kwargs) -> None: """ Default constructor Parameters ---------- + name + Name for the random variable. + min_delay Shortest possible delay in time units. Will become the first value of [`self.values`][] @@ -113,7 +121,7 @@ def __init__(self, min_delay: int, max_delay: int, **kwargs) -> None: if not max_delay >= min_delay: raise ValueError("max_delay must be greater than or equal to min_delay") - super().__init__(values=jnp.arange(min_delay, max_delay + 1)) + super().__init__(name=name, values=jnp.arange(min_delay, max_delay + 1)) @property def min_delay(self) -> int: @@ -158,12 +166,15 @@ class NonnegativeDelayPMF(DelayPMF): simplifies the computation of predicted observations. """ - def __init__(self, max_delay: int) -> None: + def __init__(self, name: str, max_delay: int) -> None: """ Default constructor. Parameters ---------- + name + Name for the random variable. + max_delay Longest possible delay in time units. Will become the final value of [`self.values`][] @@ -180,7 +191,7 @@ def __init__(self, max_delay: int) -> None: ValueError If max_delay does not satisfy the specified constraints. """ - super().__init__(min_delay=0, max_delay=max_delay) + super().__init__(name=name, min_delay=0, max_delay=max_delay) class PositiveDelayPMF(DelayPMF): @@ -197,12 +208,15 @@ class PositiveDelayPMF(DelayPMF): renewal equation. """ - def __init__(self, max_delay: int) -> None: + def __init__(self, name: str, max_delay: int) -> None: """ Default constructor. Parameters ---------- + name + Name for the random variable. + max_delay Longest possible delay in time units. Will become the final value of [`self.values`][] @@ -220,7 +234,7 @@ def __init__(self, max_delay: int) -> None: If max_delay does not satisfy the specified constraints. """ - super().__init__(min_delay=1, max_delay=max_delay) + super().__init__(name=name, min_delay=1, max_delay=max_delay) class GenerationIntervalPMF(PositiveDelayPMF): @@ -236,3 +250,61 @@ class AscertainmentDelayPMF(NonnegativeDelayPMF): represents the PMF of a delay from an event to when it is ascertained """ + + +class DeterministicGenerationIntervalPMF(GenerationIntervalPMF): + """ + Subclass of [`pyrenew.randomvariable.GenerationIntervalPMF`] + where the PMF is treated as fixed. + """ + + def __init__(self, name: str, probabilities: ArrayLike, max_delay: int) -> None: + """ + Default constructor. + + Parameters + ---------- + name + Name for the random variable. + + probabilities + Vector of probabilities representing the pmf + + max_delay + Longest possible delay in time units. + Will become the final value of [`self.values`][] + (corresponding to the final entry of the probability + vector returned by [`self.sample`][]). Must be an + integer greater than or equal to 1. + + Returns + ------- + None + + Raises + ------ + ValueError + If max_delay does not satisfy the specified constraints. + """ + + self.base_variable_ = DeterministicVariable( + name="base_variable_", value=probabilities + ) + super().__init__(name=name, max_delay=max_delay) + + def sample(self, **kwargs) -> ArrayLike: + """ + Retrieve the probability vector representing + the deterministic PMF. + + Parameters + ---------- + **kwargs + Keyword arguments passed to `self.base_variable_.sample()`. + + Returns + ------- + ArrayLike + The probability vector. + """ + return self.base_variable_.sample(**kwargs) From ca6cd0e811fbe0c9077f8a0267c1dcf5cc7c754a Mon Sep 17 00:00:00 2001 From: "Dylan H. Morris" Date: Thu, 26 Feb 2026 13:25:37 -0500 Subject: [PATCH 4/4] Enforce probability dimension at construction --- pyrenew/randomvariable/pmf.py | 39 +++++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/pyrenew/randomvariable/pmf.py b/pyrenew/randomvariable/pmf.py index c409f462..568c80b2 100644 --- a/pyrenew/randomvariable/pmf.py +++ b/pyrenew/randomvariable/pmf.py @@ -22,7 +22,7 @@ class PMFVector(RandomVariable): or stochastic in concrete subclasses. """ - def __init__(self, name: str, values: ArrayLike, **kwargs) -> None: + def __init__(self, name: str, values: ArrayLike) -> None: """ Default constructor. @@ -37,24 +37,19 @@ def __init__(self, name: str, values: ArrayLike, **kwargs) -> None: values of the variable to which those probabilities correspond. - **kwargs - Additional keyword arguments passed to the parent - constructor. - Returns ------- None """ - self.name = name self.values = values - super().__init__(**kwargs) + super().__init__(name=name) @abstractmethod - def sample(self, **kwargs) -> ArrayLike: + def sample(self, **kwargs: object) -> ArrayLike: """ Sample a vector of probabilities. """ - pass + raise NotImplementedError() class DelayPMF(PMFVector): @@ -76,7 +71,7 @@ class DelayPMF(PMFVector): be either 0 or 1 time unit. """ - def __init__(self, name: str, min_delay: int, max_delay: int, **kwargs) -> None: + def __init__(self, name: str, min_delay: int, max_delay: int) -> None: """ Default constructor @@ -99,10 +94,6 @@ def __init__(self, name: str, min_delay: int, max_delay: int, **kwargs) -> None: vector returned by [`self.sample`][]). Must be an integer greater than or equal to `min_delay`. - **kwargs - Additional keyword arguments passed to the parent - constructor. - Returns ------- None @@ -268,7 +259,8 @@ def __init__(self, name: str, probabilities: ArrayLike, max_delay: int) -> None: Name for the random variable. probabilities - Vector of probabilities representing the pmf + Vector of probabilities representing the pmf. + Must have a first dimension of length max_delay Longest possible delay in time units. @@ -286,13 +278,28 @@ def __init__(self, name: str, probabilities: ArrayLike, max_delay: int) -> None: ValueError If max_delay does not satisfy the specified constraints. """ + probabilities = jnp.array(probabilities) + if not probabilities.shape[0] == max_delay: + raise ValueError( + f"When `max_delay` = {max_delay}, " + "first dimension of `probabilities` " + f"must be of length {max_delay}, with " + "one entry for each possible delay in " + f"{jnp.arange(1, max_delay + 1)}" + ) self.base_variable_ = DeterministicVariable( name="base_variable_", value=probabilities ) super().__init__(name=name, max_delay=max_delay) - def sample(self, **kwargs) -> ArrayLike: + def validate(self, **kwargs: object) -> None: + """ + Empty validation + """ + pass + + def sample(self, **kwargs: object) -> ArrayLike: """ Retrieve the probability vector representing the deterministic PMF.