pypomp.random.fast_multinomial

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

Generate multinomial random variables using the inverse CDF method with fast_binomial in order to run fast on GPUs.

The implementation follows the methodology from Giles and Beentjes (2024). To more accurately handle cases where np is very small or the random draw is expected to be close to 0 or n, we apply the exact inverse CDF method in a manner similar to Giles (2016). Our implementation of the method does not produce exact multinomial random variables, but it is very close to exact.

Parameters:

  • key – PRNG key used as the random key.

  • n – Number of trials for the multinomial distribution. Shape: (…,)

  • p – Probabilities for each category. Shape: (…, k), where k = num categories. Probabilities along the last axis must sum to 1.

  • dtype – optional, a float dtype for the returned values (default float64 if jax_enable_x64 is true, otherwise float32). If integer, returns -1 for invalid inputs instead of nan.

Returns:

Multinomial counts. Same shape as p, but with specified dtype.