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

Generate negative binomial random variables using approximate inverse CDF methods for the gamma and Poisson distributions (fast_gamma and fast_poisson) in order to run fast on GPUs.

The Negative Binomial distribution NB(n, p) represents the number of failures before n successes, where p is the probability of success. Alternatively, it can be parameterized by n (size) and mu (mean).

NB(n, p) has mean mu = n * (1-p) / p.

Parameters:

• key – PRNG key used as the random key.

• n – Number of successes (size parameter). Must be positive.

• p – Probability of success (0 < p <= 1). Exactly one of p or mu must be provided.

• mu – Mean of the distribution. Exactly one of p or mu must be provided.

• dtype – optional, a float or integer dtype for the returned values (default float64 if jax_enable_x64 is true, otherwise float32).

• gamma_newton_loops – Cap on iterations for the Newton-Raphson method used for the gamma sampler.

• poisson_newton_loops – Cap on iterations for the Newton-Raphson method used for the Poisson sampler.

• poisson_inverse_cdf_loops – Cap on iterations for the exact inverse CDF method used for the Poisson sampler.

• gamma_adjustment_size – Size of uniform adjustments to apply for the gamma sampler (handles small shape values).

Returns:

Negative binomial random variables with the same broadcast shape as the inputs.