Skip to content

Latent

HospitalAdmissions 

HospitalAdmissions(
    infection_to_admission_interval_rv: RandomVariable,
    infection_hospitalization_ratio_rv: RandomVariable,
    day_of_week_effect_rv: RandomVariable | None = None,
    hospitalization_reporting_ratio_rv: RandomVariable | None = None,
    obs_data_first_day_of_the_week: int = 0,
)

Bases: RandomVariable

Latent hospital admissions

Implements a renewal process for the expected number of hospital admissions.

Notes

The following text was directly extracted from the wastewater model documentation (link <https://github.com/CDCgov/ww-inference-model/blob/main/model_definition.md#hospital-admissions-component>_).

Following other semi-mechanistic renewal frameworks, we model the expected hospital admissions per capita \(H(t)\) as a convolution of the expected latent incident infections per capita \(I(t)\), and a discrete infection to hospitalization distribution \(d(\tau)\), scaled by the probability of being hospitalized \(p_\mathrm{hosp}(t)\).

To account for day-of-week effects in hospital reporting, we use an estimated day of the week effect \(\omega(t)\). If \(t\) and \(t'\) are the same day of the week, \(\omega(t) = \omega(t')\). The seven values that \(\omega(t)\) takes on are constrained to have mean 1.

\[ H(t) = \omega(t) p_\mathrm{hosp}(t) \sum_{\tau = 0}^{T_d} d(\tau) I(t-\tau) \]

Where \(T_d\) is the maximum delay from infection to hospitalization that we consider.

Default constructor

Parameters:

Name Type Description Default
infection_to_admission_interval_rv RandomVariable

pmf for reporting (informing) hospital admissions (see pyrenew.observations.Deterministic).

 required
infection_hospitalization_ratio_rv RandomVariable

Infection to hospitalization rate random variable.

 required
day_of_week_effect_rv RandomVariable | None

Day of the week effect. Should return a ArrayLike with 7 values. Defaults to a deterministic variable with jax.numpy.ones(7) (no effect).

 None
hospitalization_reporting_ratio_rv RandomVariable | None

Random variable for the hospital admission reporting probability. Defaults to 1 (full reporting).

 None
obs_data_first_day_of_the_week int

The day of the week that the first day of the observation data corresponds to. Valid values are 0-6, where 0 is Monday and 6 is Sunday. Defaults to 0.

 0

Returns:

Type Description
None
Source code in pyrenew/latent/hospitaladmissions.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def __init__(
    self,
    infection_to_admission_interval_rv: RandomVariable,
    infection_hospitalization_ratio_rv: RandomVariable,
    day_of_week_effect_rv: RandomVariable | None = None,
    hospitalization_reporting_ratio_rv: RandomVariable | None = None,
    obs_data_first_day_of_the_week: int = 0,
) -> None:
    """
    Default constructor

    Parameters
    ----------
    infection_to_admission_interval_rv
        pmf for reporting (informing) hospital admissions (see
        pyrenew.observations.Deterministic).
    infection_hospitalization_ratio_rv
        Infection to hospitalization rate random variable.
    day_of_week_effect_rv
        Day of the week effect. Should return a ArrayLike with 7
        values. Defaults to a deterministic variable with
        jax.numpy.ones(7) (no effect).
    hospitalization_reporting_ratio_rv
        Random variable for the hospital admission reporting
        probability. Defaults to 1 (full reporting).
    obs_data_first_day_of_the_week
        The day of the week that the first day of the observation data
        corresponds to. Valid values are 0-6, where 0 is Monday and 6 is
        Sunday. Defaults to 0.

    Returns
    -------
    None
    """

    if day_of_week_effect_rv is None:
        day_of_week_effect_rv = DeterministicVariable(
            name="weekday_effect", value=jnp.ones(7)
        )
    if hospitalization_reporting_ratio_rv is None:
        hospitalization_reporting_ratio_rv = DeterministicVariable(
            name="hosp_report_prob", value=1.0
        )

    HospitalAdmissions.validate(
        infection_to_admission_interval_rv,
        infection_hospitalization_ratio_rv,
        day_of_week_effect_rv,
        hospitalization_reporting_ratio_rv,
        obs_data_first_day_of_the_week,
    )

    self.infection_to_admission_interval_rv = infection_to_admission_interval_rv
    self.infection_hospitalization_ratio_rv = infection_hospitalization_ratio_rv
    self.day_of_week_effect_rv = day_of_week_effect_rv
    self.hospitalization_reporting_ratio_rv = hospitalization_reporting_ratio_rv
    self.obs_data_first_day_of_the_week = obs_data_first_day_of_the_week

sample 

sample(latent_infections: ArrayLike, **kwargs) -> HospitalAdmissionsSample

Samples from the observation process

Parameters:

Name Type Description Default
latent_infections ArrayLike

Latent infections.

 required
**kwargs

Additional keyword arguments passed through to internal sample() calls, should there be any.

 {}

Returns:

Type Description
HospitalAdmissionsSample
Source code in pyrenew/latent/hospitaladmissions.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
def sample(
    self,
    latent_infections: ArrayLike,
    **kwargs,
) -> HospitalAdmissionsSample:
    """
    Samples from the observation process

    Parameters
    ----------
    latent_infections
        Latent infections.
    **kwargs
        Additional keyword arguments passed through to
        internal `sample()` calls,
        should there be any.

    Returns
    -------
    HospitalAdmissionsSample
    """

    infection_hosp_rate = self.infection_hospitalization_ratio_rv(**kwargs)

    infection_to_admission_interval = self.infection_to_admission_interval_rv(
        **kwargs
    )

    latent_hospital_admissions, _ = compute_delay_ascertained_incidence(
        latent_infections,
        infection_to_admission_interval,
        infection_hosp_rate,
    )

    # Applying the day of the week effect. For this we need to:
    # 1. Get the day of the week effect
    # 2. Identify the offset of the latent_infections
    # 3. Apply the day of the week effect to the
    # latent_hospital_admissions
    dow_effect_sampled = self.day_of_week_effect_rv(**kwargs)

    if dow_effect_sampled.size != 7:
        raise ValueError(
            "Day of the week effect should have 7 values. "
            f"Got {dow_effect_sampled.size} instead."
        )

    inf_offset = self.obs_data_first_day_of_the_week % 7

    # Replicating the day of the week effect to match the number of
    # timepoints
    dow_effect = au.tile_until_n(
        data=dow_effect_sampled,
        n_timepoints=latent_hospital_admissions.size,
        offset=inf_offset,
    )

    latent_hospital_admissions = latent_hospital_admissions * dow_effect

    # Applying reporting probability
    latent_hospital_admissions = (
        latent_hospital_admissions
        * self.hospitalization_reporting_ratio_rv(**kwargs)
    )

    numpyro.deterministic("latent_hospital_admissions", latent_hospital_admissions)

    return HospitalAdmissionsSample(
        infection_hosp_rate=infection_hosp_rate,
        latent_hospital_admissions=latent_hospital_admissions,
        multiplier=dow_effect,
    )

validate staticmethod 

validate(
    infection_to_admission_interval_rv: Any,
    infection_hospitalization_ratio_rv: Any,
    day_of_week_effect_rv: Any,
    hospitalization_reporting_ratio_rv: Any,
    obs_data_first_day_of_the_week: Any,
) -> None

Validates that the IHR, weekday effects, probability of being reported hospitalized distributions, and infection to hospital admissions reporting delay pmf are RandomVariable types

Parameters:

Name Type Description Default
infection_to_admission_interval_rv Any

Possibly incorrect input for the infection to hospitalization interval distribution.

 required
infection_hospitalization_ratio_rv Any

Possibly incorrect input for infection to hospitalization rate distribution.

 required
day_of_week_effect_rv Any

Possibly incorrect input for day of the week effect.

 required
hospitalization_reporting_ratio_rv Any

Possibly incorrect input for distribution or fixed value for the hospital admission reporting probability.

 required
obs_data_first_day_of_the_week Any

Possibly incorrect input for the day of the week that the first day of the observation data corresponds to. Valid values are 0-6, where 0 is Monday and 6 is Sunday.

 required

Returns:

Type Description
None

Raises:

Type Description
AssertionError

If any of the random variables are not of the correct type, or if the day of the week is not within the valid range.
Source code in pyrenew/latent/hospitaladmissions.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@staticmethod
def validate(
    infection_to_admission_interval_rv: Any,
    infection_hospitalization_ratio_rv: Any,
    day_of_week_effect_rv: Any,
    hospitalization_reporting_ratio_rv: Any,
    obs_data_first_day_of_the_week: Any,
) -> None:
    """
    Validates that the IHR, weekday effects, probability of being
    reported hospitalized distributions, and infection to
    hospital admissions reporting delay pmf are RandomVariable types

    Parameters
    ----------
    infection_to_admission_interval_rv
        Possibly incorrect input for the infection to hospitalization
        interval distribution.
    infection_hospitalization_ratio_rv
        Possibly incorrect input for infection to hospitalization rate distribution.
    day_of_week_effect_rv
        Possibly incorrect input for day of the week effect.
    hospitalization_reporting_ratio_rv
        Possibly incorrect input for distribution or fixed value for the
        hospital admission reporting probability.
    obs_data_first_day_of_the_week
        Possibly incorrect input for the day of the week that the first day
        of the observation data corresponds to. Valid values are 0-6, where
        0 is Monday and 6 is Sunday.

    Returns
    -------
    None

    Raises
    ------
    AssertionError
        If any of the random variables are not of the correct type, or if
        the day of the week is not within the valid range.
    """
    assert isinstance(infection_to_admission_interval_rv, RandomVariable)
    assert isinstance(infection_hospitalization_ratio_rv, RandomVariable)
    assert isinstance(day_of_week_effect_rv, RandomVariable)
    assert isinstance(hospitalization_reporting_ratio_rv, RandomVariable)
    assert isinstance(obs_data_first_day_of_the_week, int)
    assert 0 <= obs_data_first_day_of_the_week <= 6

    return None

InfectionInitializationMethod 

InfectionInitializationMethod(n_timepoints: int)

Method for initializing infections in a renewal process.

Default constructor for pyrenew.latent.infection_initialization_method.InfectionInitializationMethod.

Parameters:

Name Type Description Default
n_timepoints int

the number of time points for which to generate initial infections

 required

Returns:

Type Description
None
Source code in pyrenew/latent/infection_initialization_method.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def __init__(self, n_timepoints: int):
    """Default constructor for
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][].

    Parameters
    ----------
    n_timepoints
        the number of time points for which to
        generate initial infections

    Returns
    -------
    None
    """
    self.validate(n_timepoints)
    self.n_timepoints = n_timepoints

initialize_infections abstractmethod 

initialize_infections(I_pre_init: ArrayLike)

Generate the number of initialized infections at each time point.

Parameters:

Name Type Description Default
I_pre_init ArrayLike

An array representing some number of latent infections to be used with the specified [pyrenew.latent.infection_initialization_method.InfectionInitializationMethod][].

 required

Returns:

Type Description
ArrayLike

An array of length n_timepoints with the number of initialized infections at each time point.
Source code in pyrenew/latent/infection_initialization_method.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
@abstractmethod
def initialize_infections(self, I_pre_init: ArrayLike):
    """Generate the number of initialized infections at each time point.

    Parameters
    ----------
    I_pre_init
        An array representing some number of latent infections to be used with the specified `[`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][]`.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of initialized infections at each time point.
    """

validate staticmethod 

validate(n_timepoints: int) -> None

Validate inputs to the pyrenew.latent.infection_initialization_method.InfectionInitializationMethod constructor.

Parameters:

Name Type Description Default
n_timepoints int

the number of time points to generate initial infections for

 required

Returns:

Type Description
None
Source code in pyrenew/latent/infection_initialization_method.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
@staticmethod
def validate(n_timepoints: int) -> None:
    """
    Validate inputs to the
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][]
    constructor.

    Parameters
    ----------
    n_timepoints
        the number of time points to generate initial infections for

    Returns
    -------
    None
    """
    if not isinstance(n_timepoints, int):
        raise TypeError(
            f"n_timepoints must be an integer. Got {type(n_timepoints)}"
        )
    if n_timepoints <= 0:
        raise ValueError(f"n_timepoints must be positive. Got {n_timepoints}")

InfectionInitializationProcess 

InfectionInitializationProcess(
    name,
    I_pre_init_rv: RandomVariable,
    infection_init_method: InfectionInitializationMethod,
)

Bases: RandomVariable

Generate an initial infection history

Default class constructor for InfectionInitializationProcess

Parameters:

Name Type Description Default
name

A name to assign to the RandomVariable.

 required
I_pre_init_rv RandomVariable

A RandomVariable representing the number of infections that occur at some time before the renewal process begins. Each infection_init_method uses this random variable in different ways.

 required
infection_init_method InfectionInitializationMethod

An pyrenew.latent.infection_initialization_method.InfectionInitializationMethod that generates the initial infections for the renewal process.

 required

Returns:

Type Description
None
Source code in pyrenew/latent/infection_initialization_process.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def __init__(
    self,
    name,
    I_pre_init_rv: RandomVariable,
    infection_init_method: InfectionInitializationMethod,
) -> None:
    """Default class constructor for InfectionInitializationProcess

    Parameters
    ----------
    name
        A name to assign to the RandomVariable.
    I_pre_init_rv
        A RandomVariable representing the number of infections that occur at some time before the renewal process begins. Each `infection_init_method` uses this random variable in different ways.
    infection_init_method
        An [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][] that generates the initial infections for the renewal process.

    Returns
    -------
    None
    """
    InfectionInitializationProcess.validate(I_pre_init_rv, infection_init_method)

    self.I_pre_init_rv = I_pre_init_rv
    self.infection_init_method = infection_init_method
    self.name = name

sample 

sample() -> ArrayLike

Sample the Infection Initialization Process.

Returns:

Type Description
ArrayLike

the number of initialized infections at each time point.
Source code in pyrenew/latent/infection_initialization_process.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def sample(self) -> ArrayLike:
    """Sample the Infection Initialization Process.

    Returns
    -------
    ArrayLike
        the number of initialized infections at each time point.
    """

    I_pre_init = self.I_pre_init_rv()

    infection_initialization = self.infection_init_method(
        I_pre_init,
    )

    return infection_initialization

validate staticmethod 

validate(
    I_pre_init_rv: RandomVariable,
    infection_init_method: InfectionInitializationMethod,
) -> None

Validate the input arguments to the InfectionInitializationProcess class constructor

Parameters:

Name Type Description Default
I_pre_init_rv RandomVariable

A random variable representing the number of infections that occur at some time before the renewal process begins.

 required
infection_init_method InfectionInitializationMethod

An method to generate the initial infections.

 required

Returns:

Type Description
None
Source code in pyrenew/latent/infection_initialization_process.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
@staticmethod
def validate(
    I_pre_init_rv: RandomVariable,
    infection_init_method: InfectionInitializationMethod,
) -> None:
    """Validate the input arguments to the InfectionInitializationProcess class constructor

    Parameters
    ----------
    I_pre_init_rv
        A random variable representing the number of infections that occur at some time before the renewal process begins.
    infection_init_method
        An method to generate the initial infections.

    Returns
    -------
    None
    """
    _assert_type("I_pre_init_rv", I_pre_init_rv, RandomVariable)
    _assert_type(
        "infection_init_method",
        infection_init_method,
        InfectionInitializationMethod,
    )

Infections 

Infections(**kwargs)

Bases: RandomVariable

Latent infections

This class samples infections given \(\mathcal{R}(t)\), initial infections, and generation interval.

Notes

The mathematical model is given by:

\[ I(t) = R(t) \times \sum_{\tau < t} I(\tau) g(t-\tau) \]

where \(I(t)\) is the number of infections at time \(t\), \(R(t)\) is the reproduction number at time \(t\), and \(g(t-\tau)\) is the generation interval.

Source code in pyrenew/metaclass.py 
47
48
49
50
51
def __init__(self, **kwargs):
    """
    Default constructor
    """
    pass

sample 

sample(
    Rt: ArrayLike, I0: ArrayLike, gen_int: ArrayLike, **kwargs
) -> InfectionsSample

Sample infections given \(\mathcal{R}(t)\), initial infections, and generation interval.

Parameters:

Name Type Description Default
Rt ArrayLike

Reproduction number.

 required
I0 ArrayLike

Initial infections vector of the same length as the generation interval.

 required
gen_int ArrayLike

Generation interval pmf vector.

 required
**kwargs

Additional keyword arguments passed through to internal sample calls, should there be any.

 {}

Returns:

Type Description
InfectionsSample

A named tuple with a post_initialization_infections field.
Source code in pyrenew/latent/infections.py 
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def sample(
    self,
    Rt: ArrayLike,
    I0: ArrayLike,
    gen_int: ArrayLike,
    **kwargs,
) -> InfectionsSample:
    r"""
    Sample infections given
    $\mathcal{R}(t)$, initial infections,
    and generation interval.

    Parameters
    ----------
    Rt
        Reproduction number.
    I0
        Initial infections vector
        of the same length as the
        generation interval.
    gen_int
        Generation interval pmf vector.
    **kwargs
        Additional keyword arguments passed through to internal
        sample calls, should there be any.

    Returns
    -------
    InfectionsSample
        A named tuple with a
        `post_initialization_infections` field.
    """
    if I0.shape[0] < gen_int.size:
        raise ValueError(
            "Initial infections vector must be at least as long as "
            "the generation interval. "
            f"Initial infections vector length: {I0.shape[0]}, "
            f"generation interval length: {gen_int.size}."
        )

    if I0.shape[1:] != Rt.shape[1:]:
        raise ValueError(
            "Initial infections and Rt must have the "
            "same batch shapes. "
            f"Got initial infections of batch shape {I0.shape[1:]} "
            f"and Rt of batch shape {Rt.shape[1:]}."
        )

    gen_int_rev = jnp.flip(gen_int)
    recent_I0 = I0[-gen_int_rev.size :]

    post_initialization_infections = inf.compute_infections_from_rt(
        I0=recent_I0,
        Rt=Rt,
        reversed_generation_interval_pmf=gen_int_rev,
    )

    return InfectionsSample(post_initialization_infections)

InfectionsWithFeedback 

InfectionsWithFeedback(
    infection_feedback_strength: RandomVariable,
    infection_feedback_pmf: RandomVariable,
)

Bases: RandomVariable

Latent infections

This class computes infections, given Rt, initial infections, and generation interval.

Parameters:

Name Type Description Default
infection_feedback_strength RandomVariable

Infection feedback strength.

 required
infection_feedback_pmf RandomVariable

Infection feedback pmf.

 required
Notes

This function implements the following renewal process (reproduced from pyrenew.latent.infection_functions.compute_infections_from_rt_with_feedback):

\[ I(t) & = \mathcal{R}(t)\sum_{\tau=1}^{T_g}I(t - \tau)g(\tau) \mathcal{R}(t) & = \mathcal{R}^u(t)\exp\left(-\gamma(t)\ \sum_{\tau=1}^{T_f}I(t - \tau)f(\tau)\right) \]

where \(\mathcal{R}(t)\) is the reproductive number, \(\gamma(t)\) is the infection feedback strength, \(T_g\) is the max-length of the generation interval, \(\mathcal{R}^u(t)\) is the raw reproduction number, \(f(t)\) is the infection feedback pmf, and \(T_f\) is the max-length of the infection feedback pmf.

Default constructor for Infections class.

Parameters:

Name Type Description Default
infection_feedback_strength RandomVariable

Infection feedback strength.

 required
infection_feedback_pmf RandomVariable

Infection feedback pmf.

 required

Returns:

Type Description
None
Source code in pyrenew/latent/infectionswithfeedback.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def __init__(
    self,
    infection_feedback_strength: RandomVariable,
    infection_feedback_pmf: RandomVariable,
) -> None:
    """
    Default constructor for Infections class.

    Parameters
    ----------
    infection_feedback_strength
        Infection feedback strength.
    infection_feedback_pmf
        Infection feedback pmf.

    Returns
    -------
    None
    """

    self.validate(infection_feedback_strength, infection_feedback_pmf)

    self.infection_feedback_strength = infection_feedback_strength
    self.infection_feedback_pmf = infection_feedback_pmf

    return None

sample 

sample(
    Rt: ArrayLike, I0: ArrayLike, gen_int: ArrayLike, **kwargs
) -> InfectionsRtFeedbackSample

Samples infections given Rt, initial infections, and generation interval.

Parameters:

Name Type Description Default
Rt ArrayLike

Reproduction number.

 required
I0 ArrayLike

Initial infections, as an array at least as long as the generation interval PMF.

 required
gen_int ArrayLike

Generation interval PMF.

 required
**kwargs

Additional keyword arguments passed through to internal sample calls, should there be any.

 {}

Returns:

Type Description
InfectionsWithFeedback

Named tuple with "infections".
Source code in pyrenew/latent/infectionswithfeedback.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
def sample(
    self,
    Rt: ArrayLike,
    I0: ArrayLike,
    gen_int: ArrayLike,
    **kwargs,
) -> InfectionsRtFeedbackSample:
    """
    Samples infections given Rt, initial infections, and generation
    interval.

    Parameters
    ----------
    Rt
        Reproduction number.
    I0
        Initial infections, as an array
        at least as long as the generation
        interval PMF.
    gen_int
        Generation interval PMF.
    **kwargs
        Additional keyword arguments passed through to internal
        sample calls, should there be any.

    Returns
    -------
    InfectionsWithFeedback
        Named tuple with "infections".
    """
    if I0.shape[0] < gen_int.size:
        raise ValueError(
            "Initial infections must be at least as long as the "
            f"generation interval. Got initial infections length {I0.shape[0]}"
            f"and generation interval length {gen_int.size}."
        )

    if I0.shape[1:] != Rt.shape[1:]:
        raise ValueError(
            "Initial infections and Rt must have the same batch shapes. "
            f"Got initial infections of batch shape {I0.shape[1:]} "
            f"and Rt of batch shape {Rt.shape[1:]}."
        )

    gen_int_rev = jnp.flip(gen_int)

    I0 = I0[-gen_int_rev.size :]

    # Sampling inf feedback strength
    inf_feedback_strength = jnp.atleast_1d(
        self.infection_feedback_strength(
            **kwargs,
        )
    )

    try:
        inf_feedback_strength = jnp.broadcast_to(inf_feedback_strength, Rt.shape)
    except Exception as e:
        raise ValueError(
            "Could not broadcast inf_feedback_strength "
            f"(shape {inf_feedback_strength.shape}) "
            "to the shape of Rt"
            f"{Rt.shape}"
        ) from e

    # Sampling inf feedback pmf
    inf_feedback_pmf = self.infection_feedback_pmf(**kwargs)

    inf_fb_pmf_rev = jnp.flip(inf_feedback_pmf)

    (
        post_initialization_infections,
        Rt_adj,
    ) = inf.compute_infections_from_rt_with_feedback(
        I0=I0,
        Rt_raw=Rt,
        infection_feedback_strength=inf_feedback_strength,
        reversed_generation_interval_pmf=gen_int_rev,
        reversed_infection_feedback_pmf=inf_fb_pmf_rev,
    )

    return InfectionsRtFeedbackSample(
        post_initialization_infections=post_initialization_infections,
        rt=Rt_adj,
    )

validate staticmethod 

validate(inf_feedback_strength: any, inf_feedback_pmf: any) -> None

Validates the input parameters.

Parameters:

Name Type Description Default
inf_feedback_strength any

Infection feedback strength.

 required
inf_feedback_pmf any

Infection feedback pmf.

 required

Returns:

Type Description
None
Source code in pyrenew/latent/infectionswithfeedback.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@staticmethod
def validate(
    inf_feedback_strength: any,
    inf_feedback_pmf: any,
) -> None:  # numpydoc ignore=GL08
    """
    Validates the input parameters.

    Parameters
    ----------
    inf_feedback_strength
        Infection feedback strength.
    inf_feedback_pmf
        Infection feedback pmf.

    Returns
    -------
    None
    """
    assert isinstance(inf_feedback_strength, RandomVariable)
    assert isinstance(inf_feedback_pmf, RandomVariable)

    return None

InitializeInfectionsExponentialGrowth 

InitializeInfectionsExponentialGrowth(
    n_timepoints: int, rate_rv: RandomVariable, t_pre_init: int | None = None
)

Bases: InfectionInitializationMethod

Generate initial infections according to exponential growth.

Notes

The number of incident infections at time t is given by:

\[ I(t) = I_p \exp \left( r (t - t_p) \right) \]

Where \(I_p\) is I_pre_init, \(r\) is rate, and \(t_p\) is t_pre_init. This ensures that \(I(t_p) = I_p\). We default to t_pre_init = n_timepoints - 1, so that I_pre_init represents the number of incident infections immediately before the renewal process begins.

Default constructor for the pyrenew.latent.infection_initialization_method.InitializeInfectionsExponentialGrowth class.

Parameters:

Name Type Description Default
n_timepoints int

the number of time points to generate initial infections for

 required
rate_rv RandomVariable

A random variable representing the rate of exponential growth

 required
t_pre_init int | None

The time point whose number of infections is described by I_pre_init. Defaults to n_timepoints - 1.

 None
Source code in pyrenew/latent/infection_initialization_method.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __init__(
    self,
    n_timepoints: int,
    rate_rv: RandomVariable,
    t_pre_init: int | None = None,
):
    """Default constructor for the [`pyrenew.latent.infection_initialization_method.InitializeInfectionsExponentialGrowth`][] class.

    Parameters
    ----------
    n_timepoints
        the number of time points to generate initial infections for
    rate_rv
        A random variable representing the rate of exponential growth
    t_pre_init
         The time point whose number of infections is described by ``I_pre_init``. Defaults to ``n_timepoints - 1``.
    """
    super().__init__(n_timepoints)
    self.rate_rv = rate_rv
    if t_pre_init is None:
        t_pre_init = n_timepoints - 1
    self.t_pre_init = t_pre_init

initialize_infections 

initialize_infections(I_pre_init: ArrayLike)

Generate initial infections according to exponential growth.

Parameters:

Name Type Description Default
I_pre_init ArrayLike

An array of size 1 representing the number of infections at time t_pre_init.

 required

Returns:

Type Description
ArrayLike

An array of length n_timepoints with the number of initialized infections at each time point.
Source code in pyrenew/latent/infection_initialization_method.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def initialize_infections(self, I_pre_init: ArrayLike):
    """Generate initial infections according to exponential growth.

    Parameters
    ----------
    I_pre_init
        An array of size 1 representing the number of infections at time ``t_pre_init``.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of initialized infections at each time point.
    """
    I_pre_init = jnp.array(I_pre_init)
    rate = jnp.array(self.rate_rv())
    initial_infections = I_pre_init * jnp.exp(
        rate * (jnp.arange(self.n_timepoints)[:, jnp.newaxis] - self.t_pre_init)
    )
    return jnp.squeeze(initial_infections)

InitializeInfectionsFromVec 

InitializeInfectionsFromVec(n_timepoints: int)

Bases: InfectionInitializationMethod

Create initial infections from a vector of infections.

Source code in pyrenew/latent/infection_initialization_method.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def __init__(self, n_timepoints: int):
    """Default constructor for
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][].

    Parameters
    ----------
    n_timepoints
        the number of time points for which to
        generate initial infections

    Returns
    -------
    None
    """
    self.validate(n_timepoints)
    self.n_timepoints = n_timepoints

initialize_infections 

initialize_infections(I_pre_init: ArrayLike) -> ArrayLike

Create initial infections from a vector of infections.

Parameters:

Name Type Description Default
I_pre_init ArrayLike

An array with the same length as n_timepoints to be used as the initial infections.

 required

Returns:

Type Description
ArrayLike

An array of length n_timepoints with the number of initialized infections at each time point.
Source code in pyrenew/latent/infection_initialization_method.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def initialize_infections(self, I_pre_init: ArrayLike) -> ArrayLike:
    """Create initial infections from a vector of infections.

    Parameters
    ----------
    I_pre_init
        An array with the same length as ``n_timepoints`` to be
        used as the initial infections.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of
        initialized infections at each time point.
    """
    I_pre_init = jnp.array(I_pre_init)
    if I_pre_init.size != self.n_timepoints:
        raise ValueError(
            "I_pre_init must have the same size as n_timepoints. "
            f"Got I_pre_init of size {I_pre_init.size} "
            f"and n_timepoints of size {self.n_timepoints}."
        )
    return I_pre_init

InitializeInfectionsZeroPad 

InitializeInfectionsZeroPad(n_timepoints: int)

Bases: InfectionInitializationMethod

Create an initial infection vector of specified length by padding a shorter vector with an appropriate number of zeros at the beginning of the time series.

Source code in pyrenew/latent/infection_initialization_method.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def __init__(self, n_timepoints: int):
    """Default constructor for
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][].

    Parameters
    ----------
    n_timepoints
        the number of time points for which to
        generate initial infections

    Returns
    -------
    None
    """
    self.validate(n_timepoints)
    self.n_timepoints = n_timepoints

initialize_infections 

initialize_infections(I_pre_init: ArrayLike)

Pad the initial infections with zeros at the beginning of the time series.

Parameters:

Name Type Description Default
I_pre_init ArrayLike

An array with initialized infections to be padded with zeros.

 required

Returns:

Type Description
ArrayLike

An array of length n_timepoints with the number of initialized infections at each time point.
Source code in pyrenew/latent/infection_initialization_method.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def initialize_infections(self, I_pre_init: ArrayLike):
    """Pad the initial infections with zeros at the beginning of the time series.

    Parameters
    ----------
    I_pre_init
        An array with initialized infections to be padded with zeros.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of initialized infections at each time point.
    """
    I_pre_init = jnp.atleast_1d(I_pre_init)
    if self.n_timepoints < I_pre_init.size:
        raise ValueError(
            "I_pre_init must be no longer than n_timepoints. "
            f"Got I_pre_init of size {I_pre_init.size} and "
            f" n_timepoints of size {self.n_timepoints}."
        )
    return jnp.pad(I_pre_init, (self.n_timepoints - I_pre_init.size, 0))

compute_infections_from_rt 

compute_infections_from_rt(
    I0: ArrayLike, Rt: ArrayLike, reversed_generation_interval_pmf: ArrayLike
) -> ndarray

Generate infections according to a renewal process with a time-varying reproduction number \(\mathcal{R}(t)\)

Parameters:

Name Type Description Default
I0 ArrayLike

Array of initial infections of the same length as the generation interval pmf vector.

 required
Rt ArrayLike

Timeseries of \(\mathcal{R}(t)\) values

 required
reversed_generation_interval_pmf ArrayLike

discrete probability mass vector representing the generation interval of the infection process, where the final entry represents an infection 1 time unit in the past, the second-to-last entry represents an infection two time units in the past, etc.

 required

Returns:

Type Description
ndarray

The timeseries of infections.
Source code in pyrenew/latent/infection_functions.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def compute_infections_from_rt(
    I0: ArrayLike,
    Rt: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
) -> jnp.ndarray:
    """
    Generate infections according to a
    renewal process with a time-varying
    reproduction number $\\mathcal{R}(t)$

    Parameters
    ----------
    I0
        Array of initial infections of the
        same length as the generation interval
        pmf vector.
    Rt
        Timeseries of $\\mathcal{R}(t)$ values
    reversed_generation_interval_pmf
        discrete probability mass vector
        representing the generation interval
        of the infection process, where the final
        entry represents an infection 1 time unit in the
        past, the second-to-last entry represents
        an infection two time units in the past, etc.

    Returns
    -------
    jnp.ndarray
        The timeseries of infections.
    """
    incidence_func = new_convolve_scanner(
        reversed_generation_interval_pmf, IdentityTransform()
    )

    latest, all_infections = jax.lax.scan(f=incidence_func, init=I0, xs=Rt)

    return all_infections

compute_infections_from_rt_with_feedback 

compute_infections_from_rt_with_feedback(
    I0: ArrayLike,
    Rt_raw: ArrayLike,
    infection_feedback_strength: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
    reversed_infection_feedback_pmf: ArrayLike,
) -> tuple

Generate infections according to a renewal process with infection feedback (generalizing Asher 2018 <https://doi.org/10.1016/j.epidem.2017.02.009>_).

Parameters:

Name Type Description Default
I0 ArrayLike

Array of initial infections of the same length as the generation interval pmf vector.

 required
Rt_raw ArrayLike

Timeseries of raw \(\mathcal{R}(t)\) values not adjusted by infection feedback

 required
infection_feedback_strength ArrayLike

Strength of the infection feedback. Either a scalar (constant feedback strength in time) or a vector representing the infection feedback strength at a given point in time.

 required
reversed_generation_interval_pmf ArrayLike

discrete probability mass vector representing the generation interval of the infection process, where the final entry represents an infection 1 time unit in the past, the second-to-last entry represents an infection two time units in the past, etc.

 required
reversed_infection_feedback_pmf ArrayLike

discrete probability mass vector representing the infection feedback process, where the final entry represents the relative contribution to infection feedback from infections that occurred 1 time unit in the past, the second-to-last entry represents the contribution from infections that occurred 2 time units in the past, etc.

 required

Returns:

Type Description
tuple

A tuple (infections, Rt_adjusted), where Rt_adjusted is the infection-feedback-adjusted timeseries of the reproduction number \(\mathcal{R}(t)\) and infections is the incident infection timeseries.
Notes

This function implements the following renewal process:

\[ \begin{aligned} I(t) & = \mathcal{R}(t)\sum_{\tau=1}^{T_g}I(t - \tau)g(\tau) \\ \mathcal{R}(t) & = \mathcal{R}^u(t)\exp\left(\gamma(t)\ \sum_{\tau=1}^{T_f}I(t - \tau)f(\tau)\right) \end{aligned} \]

where \(\mathcal{R}(t)\) is the reproductive number, \(\gamma(t)\) is the infection feedback strength, \(T_g\) is the max-length of the generation interval, \(\mathcal{R}^u(t)\) is the raw reproduction number, \(f(t)\) is the infection feedback pmf, and \(T_f\) is the max-length of the infection feedback pmf.

Note that negative \(\gamma(t)\) implies that recent incident infections reduce \(\mathcal{R}(t)\) below its raw value in the absence of feedback, while positive \(\gamma\) implies that recent incident infections increase \(\mathcal{R}(t)\) above its raw value, and \(\gamma(t)=0\) implies no feedback.

In general, negative \(\gamma\) is the more common modeling choice, as it can be used to model susceptible depletion, reductions in contact rate due to awareness of high incidence, et cetera.

Source code in pyrenew/latent/infection_functions.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def compute_infections_from_rt_with_feedback(
    I0: ArrayLike,
    Rt_raw: ArrayLike,
    infection_feedback_strength: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
    reversed_infection_feedback_pmf: ArrayLike,
) -> tuple:
    r"""
    Generate infections according to
    a renewal process with infection
    feedback (generalizing `Asher 2018
    <https://doi.org/10.1016/j.epidem.2017.02.009>`_).

    Parameters
    ----------
    I0
        Array of initial infections of the
        same length as the generation interval
        pmf vector.
    Rt_raw
        Timeseries of raw $\mathcal{R}(t)$ values not
        adjusted by infection feedback
    infection_feedback_strength
        Strength of the infection feedback.
        Either a scalar (constant feedback
        strength in time) or a vector representing
        the infection feedback strength at a
        given point in time.
    reversed_generation_interval_pmf
        discrete probability mass vector
        representing the generation interval
        of the infection process, where the final
        entry represents an infection 1 time unit in the
        past, the second-to-last entry represents
        an infection two time units in the past, etc.
    reversed_infection_feedback_pmf
        discrete probability mass vector
        representing the infection feedback
        process, where the final entry represents
        the relative contribution to infection
        feedback from infections that occurred
        1 time unit in the past, the second-to-last
        entry represents the contribution from infections
        that occurred 2 time units in the past, etc.

    Returns
    -------
    tuple
        A tuple ``(infections, Rt_adjusted)``,
        where `Rt_adjusted` is the infection-feedback-adjusted
        timeseries of the reproduction number $\mathcal{R}(t)$
        and `infections` is the incident infection timeseries.

    Notes
    -----
    This function implements the following renewal process:

    ```math
    \begin{aligned}
    I(t) & = \mathcal{R}(t)\sum_{\tau=1}^{T_g}I(t - \tau)g(\tau) \\
    \mathcal{R}(t) & = \mathcal{R}^u(t)\exp\left(\gamma(t)\
        \sum_{\tau=1}^{T_f}I(t - \tau)f(\tau)\right)
    \end{aligned}
    ```

    where $\mathcal{R}(t)$ is the reproductive number,
    $\gamma(t)$ is the infection feedback strength,
    $T_g$ is the max-length of the
    generation interval, $\mathcal{R}^u(t)$ is the raw reproduction
    number, $f(t)$ is the infection feedback pmf, and $T_f$
    is the max-length of the infection feedback pmf.

    Note that negative $\gamma(t)$ implies
    that recent incident infections reduce $\mathcal{R}(t)$
    below its raw value in the absence of feedback, while
    positive $\gamma$ implies that recent incident infections
    *increase* $\mathcal{R}(t)$ above its raw value, and
    $\gamma(t)=0$ implies no feedback.

    In general, negative $\gamma$ is the more common modeling
    choice, as it can be used to model susceptible depletion,
    reductions in contact rate due to awareness of high incidence,
    et cetera.
    """
    feedback_scanner = new_double_convolve_scanner(
        arrays_to_convolve=(
            reversed_infection_feedback_pmf,
            reversed_generation_interval_pmf,
        ),
        transforms=(ExpTransform(), IdentityTransform()),
    )
    latest, infs_and_R_adj = jax.lax.scan(
        f=feedback_scanner,
        init=I0,
        xs=(infection_feedback_strength, Rt_raw),
    )

    infections, R_adjustment = infs_and_R_adj
    Rt_adjusted = R_adjustment * Rt_raw
    return infections, Rt_adjusted

logistic_susceptibility_adjustment 

logistic_susceptibility_adjustment(
    I_raw_t: float, frac_susceptible: float, n_population: float
) -> float

Apply the logistic susceptibility adjustment to a potential new incidence I_raw_t proposed in equation 6 of Bhatt et al 2023 <https://doi.org/10.1093/jrsssa/qnad030>_.

Parameters:

Name Type Description Default
I_raw_t float

The "unadjusted" incidence at time t, i.e. the incidence given an infinite number of available susceptible individuals.

 required
frac_susceptible float

fraction of remaining susceptible individuals in the population

 required
n_population float

Total size of the population.

 required

Returns:

Type Description
float

The adjusted value of \(I(t)\).
Source code in pyrenew/latent/infection_functions.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def logistic_susceptibility_adjustment(
    I_raw_t: float,
    frac_susceptible: float,
    n_population: float,
) -> float:
    """
    Apply the logistic susceptibility
    adjustment to a potential new
    incidence `I_raw_t` proposed in
    equation 6 of `Bhatt et al 2023
    <https://doi.org/10.1093/jrsssa/qnad030>`_.

    Parameters
    ----------
    I_raw_t
        The "unadjusted" incidence at time t,
        i.e. the incidence given an infinite
        number of available susceptible individuals.
    frac_susceptible
        fraction of remaining susceptible individuals
        in the population
    n_population
        Total size of the population.

    Returns
    -------
    float
        The adjusted value of $I(t)$.
    """
    approx_frac_infected = 1 - jnp.exp(-I_raw_t / n_population)
    return n_population * frac_susceptible * approx_frac_infected