pypomp.random.fast_nbinomial¶

pypomp.random.fast_nbinomial(key: Array, n: Array, p: Array | None = None, mu: Array | None = None, dtype: dtype | None = None, gamma_newton_loops: int = 3, poisson_newton_loops: int = 5, poisson_inverse_cdf_loops: int = 20, gamma_adjustment_size: int = 3) → Array[source]¶

Sample Negative Binomial random variates using a GPU-optimized algorithm.

Draws from NB(n, p) (number of failures before n successes) via a Gamma-Poisson mixture. Both steps use the approximate GPU-optimized samplers fast_gamma() and fast_poisson().

Parameters:

key (jax.Array) – JAX PRNG key.
n (jax.Array) – Size (number of successes) parameter. Must be positive.
p (jax.Array or None, optional) – Success probability in (0, 1]. Mutually exclusive with mu.
mu (jax.Array or None, optional) – Mean of the distribution: mu = n * (1 - p) / p. Mutually exclusive with p.
dtype (np.dtype or None, optional) – Output dtype (float or integer). Defaults to float64 if jax_enable_x64=True, otherwise float32.
gamma_newton_loops (int, optional) – Newton-Raphson iterations for the Gamma sampler. Defaults to 3.
poisson_newton_loops (int, optional) – Newton-Raphson iterations for the Poisson sampler. Defaults to 5.
poisson_inverse_cdf_loops (int, optional) – Exact inverse CDF iterations for the Poisson sampler. Defaults to 20.
gamma_adjustment_size (int, optional) – Uniform adjustment steps for the Gamma sampler. Defaults to 3.

Returns:

Negative Binomial samples with the broadcast shape of the inputs.

Return type:

jax.Array

Notes

For speed and accuracy metrics, see the Quant Tests.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from pypomp.random import fast_nbinomial
>>> fast_nbinomial(jax.random.key(0), n=jnp.array(5.0), mu=jnp.array(3.0))
Array(2, dtype=int32)