Additional modules

Metaclass Module

Convolution Utility Module

convolve

Factory functions for calculating convolutions of timeseries with discrete distributions of times-to-event using jax.lax.scan(). Factories generate functions that can be passed to jax.lax.scan() or numpyro.contrib.control_flow.scan() with an appropriate array to scan along.

compute_delay_ascertained_incidence(latent_incidence, delay_incidence_to_observation_pmf, p_observed_given_incident=1, pad=False)

Computes incidences observed according to a given observation rate and based on a delay interval.

In addition to the output array, returns the offset (number of time units) separating the first entry of the the input latent_incidence array from the first entry of the output (delay ascertained incidence) array. Note that if the pad keyword argument is True, the offset will be always 0.

Parameters:

• p_observed_given_incident (ArrayLike) – The rate at which latent incident counts translate into observed counts. For example, setting p_observed_given_incident=0.001 when the incident counts are infections and the observed counts are reported hospital admissions could be used to model disease and population for which the probability of a latent infection leading to a reported hospital admission is 0.001.

• latent_incidence (ArrayLike) – Incidence values based on the true underlying process.

• delay_incidence_to_observation_pmf (ArrayLike) – Probability mass function of delay interval from incidence to observation with support on the interval 0 to the length of the array’s first dimension. The \(i\) h entry represents the probability mass for a delay of \(i\) time units, i.e delay_incidence_to_observation_pmf[0] represents the fraction of observations that are delayed 0 time unit, delay_incidence_to_observation_pmf[1] represents the fraction that are delayed 1 time units, et cetera.

• pad (bool) – Return an output array that has been nan-padded so that its first entry represents the same timepoint as the first timepoint of the input latent_incidence array? Boolean, default False.

Returns:

Tuple whose first entry is the predicted timeseries of delayed observations and whose second entry is the offset.

Return type:

tuple[ArrayLike, int]

new_convolve_scanner(array_to_convolve, transform)

Factory function to create a “scanner” function that can be used with jax.lax.scan() or numpyro.contrib.control_flow.scan() to construct an array via backward-looking iterative convolution.

Parameters:

• array_to_convolve (ArrayLike) – A 1D jax array to convolve with subsets of the iteratively constructed history array.

• transform (Callable) – A transformation to apply to the result of the dot product and multiplication.

Returns:

A scanner function that can be used with jax.lax.scan() or numpyro.contrib.control_flow.scan() for convolution. This function takes a history subset array and a scalar, computes the dot product of the supplied convolution array with the history subset array, multiplies by the scalar, and returns the resulting value and a new history subset array formed by the 2nd-through-last entries of the old history subset array followed by that same resulting value.

Return type:

Callable

Notes

The following iterative operation is found often in renewal processes:

\[X(t) = f\left(m(t) \begin{bmatrix} X(t - n) \\ X(t - n + 1) \\ \vdots{} \\ X(t - 1)\end{bmatrix} \cdot{} \mathbf{d} \right) \]

Where \(\mathbf{d}\) is a vector of length \(n\), \(m(t)\) is a scalar for each value of time \(t\), and \(f\) is a scalar-valued function.

Given \(\mathbf{d}\), and optionally \(f\), this factory function returns a new function that performs one step of this process while scanning along an array of multipliers (i.e. an array giving the values of \(m(t)\)) using jax.lax.scan().

new_double_convolve_scanner(arrays_to_convolve, transforms)

Factory function to create a scanner function that iteratively constructs arrays by applying the dot-product/multiply/transform operation twice per history subset, with the first yielding operation yielding an additional scalar multiplier for the second.

Parameters:

• arrays_to_convolve (tuple[ArrayLike, ArrayLike]) – A tuple of two 1D jax arrays, one for each of the two stages of convolution. The first entry in the arrays_to_convolve tuple will be convolved with the current history subset array first, the the second entry will be convolved with it second.

• transforms (tuple[Callable, Callable]) – A tuple of two functions, each transforming the output of the dot product at each convolution stage. The first entry in the transforms tuple will be applied first, then the second will be applied.

Returns:

A scanner function that applies two sets of convolution, multiply, and transform operations in sequence to construct a new array by scanning along a pair of input arrays that are equal in length to each other.

Return type:

Callable

Notes

Using the same notation as in the documentation for new_convolve_scanner(), this function aids in applying the iterative operation:

\[\begin{aligned} Y(t) &= f_1 \left(m_1(t) \begin{bmatrix} X(t - n) \\ X(t - n + 1) \\ \vdots{} \\ X(t - 1) \end{bmatrix} \cdot{} \mathbf{d}_1 \right) \\ \\ X(t) &= f_2 \left( m_2(t) Y(t) \begin{bmatrix} X(t - n) \\ X(t - n + 1) \\ \vdots{} \\ X(t - 1)\end{bmatrix} \cdot{} \mathbf{d}_2 \right) \end{aligned} \]

Where \(\mathbf{d}_1\) and \(\mathbf{d}_2\) are vectors of length \(n\), \(m_1(t)\) and \(m_2(t)\) are scalars for each value of time \(t\), and \(f_1\) and \(f_2\) are scalar-valued functions.

Mathematics Utilities Module

Helper functions for doing analytical and/or numerical calculations.

get_asymptotic_growth_rate(R, generation_interval_pmf)

Get the asymptotic per timestep growth rate for a renewal process with a given value of \(\mathcal{R}\) and a given discrete generation interval probability mass vector.

This function computes that growth rate finding the dominant eigenvalue of the renewal process’s Leslie matrix.

Parameters:

• R (float) – The reproduction number of the renewal process

• generation_interval_pmf (ArrayLike) – The discrete generation interval probability mass vector of the renewal process

Returns:

The asymptotic growth rate of the renewal process, as a jax float.

Return type:

float

get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf)

Get the asymptotic per-timestep growth rate of the renewal process (the dominant eigenvalue of its Leslie matrix) and the associated stable age distribution (a normalized eigenvector associated to that eigenvalue).

Parameters:

• R (float) – The reproduction number of the renewal process

• generation_interval_pmf (ArrayLike) – The discrete generation interval probability mass vector of the renewal process

Returns:

A tuple consisting of the asymptotic growth rate of the process, as jax float, and the stable age distribution of the process, as a jax array probability vector of the same shape as the generation interval probability vector.

Return type:

tuple[float, ArrayLike]

Raises:

ValueError – If an age distribution vector with non-zero imaginary part is produced.

get_leslie_matrix(R, generation_interval_pmf)

Create the Leslie matrix corresponding to a basic renewal process with the given \(\mathcal{R}\) value and discrete generation interval pmf vector.

Parameters:

• R (float) – The reproduction number of the renewal process

• generation_interval_pmf (ArrayLike) – The discrete generation interval probability mass vector of the renewal process

Returns:

The Leslie matrix for the renewal process, as a jax array.

Return type:

ArrayLike

get_stable_age_distribution(R, generation_interval_pmf)

Get the stable age distribution for a renewal process with a given value of R and a given discrete generation interval probability mass vector.

This function computes that stable age distribution by finding and then normalizing an eigenvector associated to the dominant eigenvalue of the renewal process’s Leslie matrix.

Parameters:

• R (float) – The reproduction number of the renewal process

• generation_interval_pmf (ArrayLike) – The discrete generation interval probability mass vector of the renewal process

Returns:

The stable age distribution for the process, as a jax array probability vector of the same shape as the generation interval probability vector.

Return type:

ArrayLike

integrate_discrete(init_diff_vals, highest_order_diff_vals)

Integrate (de-difference) the differenced process, obtaining the process values \(X(t=0), X(t=1), ..., X(t)\) from the \(n^{th}\) differences and a set of initial process / difference values \(X(t=0), X^1(t=1), X^2(t=2), ..., X^{(n-1)}(t=n-1)\), where \(X^k(t)\) is the value of the \(n^{th}\) difference at index \(t\) of the process, obtaining a sequence of length equal to the length of the provided highest_order_diff_vals vector plus the order of the process.

Parameters:

• init_diff_vals (ArrayLike) – Values of \(X(t=0), X^1(t=1), X^2(t=2), ..., X^{(n-1)}(t=n-1)\).

• highest_order_diff_vals (ArrayLike) – Array of differences at the highest order of differencing, i.e. the order of the overall process, starting with \(X^{n}(t=n)\)

Returns:

The integrated (de-differenced) sequence of values, of length n_diffs + order, where n_diffs is the number of highest_order_diff_vals and order is the order of the process.

Return type:

ArrayLike

neg_MGF(r, w)

Compute the negative moment generating function (MGF) for a given rate r and weights w.

Parameters:

• r (float) – The rate parameter.

• w (ArrayLike) – An array of weights.

Returns:

The value of the negative MGF evaluated at r and w.

Return type:

float

Notes

For a finite discrete random variable \(X\) supported on the first \(n\) positive integers (\(\{1, 2, ..., n \}\)), the moment generating function (MGF) \(M_+(r)\) is defined as the expected value of \(\exp(rX)\). Similarly, the negative moment generating function \(M_-(r)\) is the expected value of \(\exp(-rX)\). So if we represent the PMF of \(X\) as a “weights” vector \(w\) of length \(n\), the negative MGF \(M_-(r)\) is given by:

\[M_-(r) = \sum_{t = 1}^{n} w_i \exp(-rt) \]

neg_MGF_del_r(r, w)

Compute the value of the partial deriative of neg_MGF() with respect to r evaluated at a particular r and w pair.

Parameters:

• r (float) – The rate parameter.

• w (ArrayLike) – An array of weights.

Returns:

The value of the partial derivative evaluated at r and w.

Return type:

float

r_approx_from_R(R, g, n_newton_steps)

Get the approximate asymptotic geometric growth rate r for a renewal process with a fixed reproduction number R and discrete generation interval PMF g.

Uses Newton’s method with a fixed number of steps.

Parameters:

• R (float) – The reproduction number

• g (ArrayLike) – The probability mass function of the generation interval.

• n_newton_steps (int) – Number of steps to take when performing Newton’s method.

Returns:

The approximate value of r.

Return type:

float

Notes

For a fixed value of \(\mathcal{R}\), a renewal process has an asymptotic geometric growth rate \(r\) that satisfies

\[M_{-}(r) - \frac{1}{\mathcal{R}} = 0 \]

where \(M_-(r)\) is the negative moment generating function for a random variable \(\tau\) representing the (discrete) generation interval. See neg_MGF() for details.

We obtain a value for \(r\) via approximate numerical solution of this implicit equation.

We first make an initial guess based on the mean generation interval \(\bar{\tau} = \mathbb{E}(\tau)\):

\[r \approx \frac{\mathcal{R} - 1}{\mathcal{R} \bar{\tau}} \]

We then refine this approximation by applying Newton’s method for a fixed number of steps.

Transformations Module

This module exposes numpyro’s transformations module to the user, and defines and adds additional custom transformations

ScaledLogitTransform(x_max)

Scaled logistic transformation from the interval (0, X_max) to the interval (-infinity, +infinity).

Parameters:

x_max (float) – Maximum value of the untransformed scale (will be transformed to +infinity).

Returns:

A composition of the following transformations: - numpyro.distributions.transforms.AffineTransform(0.0, 1.0/x_max) - numpyro.distributions.transforms.SigmoidTransform().inv

Return type:

nt.ComposeTransform

Regression Module

Helper classes for regression problems

class AbstractRegressionPrediction

Bases: object

abstractmethod predict()

Make a regression prediction

abstractmethod sample(obs=None)

Observe or sample from the regression problem according to the specified distributions

class GLMPrediction(name, intercept_prior, coefficient_priors, transform=None, intercept_suffix='_intercept', coefficient_suffix='_coefficients')

Bases: AbstractRegressionPrediction

Generalized linear model regression predictions

Parameters:

• name (str)

• intercept_prior (dist.Distribution)

• coefficient_priors (dist.Distribution)

• transform (t.Transform)

predict(intercept, coefficients, predictor_values)

Generates a transformed prediction w/ intercept, coefficients, and predictor values

Parameters:

• intercept (ArrayLike) – Sampled numpyro distribution generated from intercept priors.

• coefficients (ArrayLike) – Sampled prediction coefficients distribution generated from coefficients priors.

• predictor_values (ArrayLike(n_predictors, n_observations)) – Matrix of predictor variables (covariates) for the regression problem. Each row should represent the predictor values corresponding to an observation; each column should represent a predictor variable. You do not include values of 1 for the intercept; these will be added automatically.

Returns:

Array of transformed predictions.

Return type:

ArrayLike

sample(predictor_values)

Sample generalized linear model

Parameters:

predictor_values (ArrayLike(n_predictors, n_observations)) – Matrix of predictor variables (covariates) for the regression problem. Each row should represent the predictor values corresponding to an observation; each column should represent a predictor variable. Do not include values of 1 for the intercept; these will be added automatically. Passed as the predictor_values argument to GLMPrediction.predict()

Return type:

GLMPredictionSample

class GLMPredictionSample(prediction=None, intercept=None, coefficients=None)

Bases: NamedTuple

A container for holding the output from GLMPrediction.sample().

Parameters:

• prediction (ArrayLike | None)

• intercept (ArrayLike | None)

• coefficients (ArrayLike | None)

prediction

Transformed predictions. Defaults to None.

Type:

ArrayLike | None, optional

intercept

Sampled intercept from intercept priors. Defaults to None.

Type:

ArrayLike | None, optional

coefficients

Prediction coefficients generated from coefficients priors. Defaults to None.

Type:

ArrayLike | None, optional

coefficients: ArrayLike | None

Alias for field number 2

intercept: ArrayLike | None

Alias for field number 1

prediction: ArrayLike | None

Alias for field number 0

Distributions Utility Module

distutil

Utilities for working with commonly- encountered probability distributions found in renewal equation modeling, such as discrete time-to-event distributions

reverse_discrete_dist_vector(dist)

Reverse a discrete distribution vector (useful for discrete time-to-event distributions).

Parameters:

dist (ArrayLike) – A discrete distribution vector (likely discrete time-to-event distribution)

Returns:

A reversed (jnp.flip) discrete distribution vector

Return type:

ArrayLike

validate_discrete_dist_vector(discrete_dist, tol=1e-05)

Validate that a vector represents a discrete probability distribution to within a specified tolerance, raising a ValueError if not.

Parameters:

• discrete_dist (ArrayLike) – An jax array containing non-negative values that represent a discrete probability distribution. The values must sum to 1 within the specified tolerance.

• tol (float, optional) – The tolerance within which the sum of the distribution must be 1. Defaults to 1e-5.

Returns:

The normalized distribution array if the input is valid.

Return type:

ArrayLike

Raises:

ValueError – If any value in discrete_dist is negative or if the sum of the distribution does not equal 1 within the specified tolerance.