pypomp.random.fast_nbinomial

pypomp.random.fast_nbinomial(key: Array, n: Array, p: Array | None = None, mu: Array | None = None, dtype: dtype | None = None) → Array[source]

Generate negative binomial random variables using approximate inverse CDF methods 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).

Returns:

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