laser.core.distributions

laser.core.distributions

Various probability distributions implemented using NumPy and Numba.

LASER based models generally move from pure NumPy to Numba-accelerated version of the core dynamics, e.g., transmission.

It would be a hassle to re-implement these functions for each desired distribution, so we provide these Numba-wrapped distributions here which can be passed in to other Numba compiled functions.

For example, a simple SIR model may want to parameterize the infectious period using a distribution. By passing in a Numba-wrapped distribution function, we can pick and parameterize a distribution based on configuration and sample from that distribution within the Numba-compiled SIR model without needing to re-implement the distribution logic within the SIR model itself.

A simple example of usage::

import laser.core.distributions as dist

# Create a Numba-wrapped beta distribution
beta_dist = dist.beta(2.0, 5.0)

# Assign the distribution to the model's infectious period distribution
# so the transmission component can sample from it during simulation
model.infectious_period_dist = beta_dist

Note that the distribution functions take two parameters, tick and node, which are currently unused but match the desired signature for disease model components that may need to sample from distributions based on the current simulation tick or node index. In other words, distributions with spatial or temporal variation could be implemented in the future.

Here are examples of Numba-wrapped distribution functions that could vary based on tick or tick + node::

# temporal variation only
cosine = np.cos(np.linspace(0, np.pi * 2, 365))

@nb.njit(nogil=True, cache=True)
def seasonal_distribution(tick: int, node: int) -> np.float32:
    # ignore node for this seasonal distribution
    day_of_year = tick % 365
    base_value = 42.0 + 3.14159 * cosine[day_of_year]
    # parameterize normal() with seasonal factor
    return np.float32(np.random.normal(base_value, 2.0))

# additional spatial variation
ramp = np.linspace(0, 2, 42)

@nb.njit(nogil=True, cache=True)
def ramped_distribution(tick: int, node: int) -> np.float32:
    day_of_year = tick % 365
    # use seasonal factor
    base_value = 42.0 + 3.14159 * cosine[day_of_year]
    # apply spatial ramp based on node index
    base_value *= ramp[node]
    # parameterize normal() with seasonal + spatial factor
    return np.float32(np.random.normal(base_value, 1.0))

Normally, these distributions—built in or custom—will be used once per agent as above. However, the sample_ints() and sample_floats() functions can be used to efficiently sample large arrays using multiple CPU cores in parallel.

laser.core.distributions.beta(a, b) cached

Beta distribution.

\(f(x; a, b) = \frac {x^{a-1} (1-x)^{b-1}} {B(a, b)}\)

where \(B(a, b)\) is the beta function.

laser.core.distributions.binomial(n, p) cached

Binomial distribution.

\(f(k,n,p) = Pr(X = k) = \binom {n} {k} p^k (1-p)^{n-k}\)

where \(n\) is the number of trials and \(p\) is the probability of success [0, 1].

laser.core.distributions.constant_float(value) cached

Constant distribution. Always returns the same floating point value.

laser.core.distributions.constant_int(value) cached

Constant distribution. Always returns the same integer value.

laser.core.distributions.exponential(scale) cached

Exponential distribution.

\(f(x; \frac {1} {\beta}) = \frac {1} {\beta} e^{-\frac {x} {\beta}}\)

where \(\beta\) is the scale parameter (\(\beta = 1 / \lambda\)).

laser.core.distributions.gamma(shape, scale) cached

Gamma distribution.

\(p(x) = x^{k-1} \frac {e^{- x / \theta}}{\theta^k \Gamma(k)}\)

where \(k\) is the shape, \(\theta\) is the scale, and \(\Gamma(k)\) is the gamma function.

laser.core.distributions.logistic(loc, scale) cached

Logistic distribution.

\(P(x) = \frac {e^{-(x - \mu) / s}} {s (1 + e^{-(x - \mu) / s})^2}\)

where \(\mu\) is the location parameter and \(s\) is the scale parameter.

laser.core.distributions.lognormal(mean, sigma) cached

Log-normal distribution.

\(p(x) = \frac {1} {\sigma x \sqrt {2 \pi}} e^{- \frac {(\ln x - \mu)^2} {2 \sigma^2}}\)

where \(\mu\) is the mean and \(\sigma\) is the standard deviation of the underlying normal distribution.

laser.core.distributions.negative_binomial(n, p) cached

Negative binomial distribution.

\(P(N; n, p) = \frac {\Gamma (N + n)} {N! \Gamma (n)} p^n (1 - p)^N\)

where \(n\) is the number of successes, \(p\) is the probability of success on each trial, \(N + n\) is the number of trials, and \(\Gamma()\) is the gamma function. When \(n\) is an integer,

\(\frac {\Gamma (N + n)} {N! \Gamma (n)} = \binom {N + n - 1} {n - 1}\)

which is the more common form of this term.

laser.core.distributions.normal(loc, scale) cached

Normal (Gaussian) distribution.

\(p(x) = \frac {1} {\sqrt {2 \pi \sigma^2}} e^{- \frac {(x - \mu)^2} {2 \sigma^2}}\)

where \(\mu\) is the mean and \(\sigma\) is the standard deviation.

laser.core.distributions.poisson(lam) cached

Poisson distribution.

\(f( k ; \lambda ) = \frac {\lambda^k e^{- \lambda}} {k!}\)

where \(\lambda\) is the expected number of events in the given interval.

laser.core.distributions.sample_floats(fn, dest, tick=0, node=0)

Fill an array with floating point values sampled from a Numba-wrapped distribution function.

Parameters:

Name	Type	Description	Default
fn	function	Numba-wrapped distribution function returning float32 values.	required
dest	ndarray	Pre-allocated destination float32 array to store samples.	required
tick	int	Current simulation tick (default is 0). Passed through to the distribution function.	0
node	int	Current node index (default is 0). Passed through to the distribution function.	0

Returns:

Name	Type	Description
dest	ndarray	The destination array filled with sampled values.

laser.core.distributions.sample_ints(fn, dest, tick=0, node=0)

Fill an array with integer values sampled from a Numba-wrapped distribution function.

Parameters:

Name	Type	Description	Default
fn	function	Numba-wrapped distribution function returning int32 values.	required
dest	ndarray	Pre-allocated destination int32 array to store samples.	required
tick	int	Current simulation tick (default is 0). Passed through to the distribution function.	0
node	int	Current node index (default is 0). Passed through to the distribution function.	0

Returns:

Name	Type	Description
dest	ndarray	The destination array filled with sampled values.

laser.core.distributions.uniform(low, high) cached

Uniform distribution.

\(p(x) = \frac {1} {b - a}\)

where \(a\) is the lower bound and \(b\) is the upper bound, [\(a\), \(b\)).

laser.core.distributions.weibull(a, lam) cached

Weibull distribution.

\(X = \lambda (- \ln ( U ))^{1 / a}\)

where \(a\) is the shape parameter and \(\lambda\) is the scale parameter.