Fast Random Number Generation on GPUs

The simulators included in the JAX package often suffer from warp divergence due to using rejection sampling. This problem is especially pronounced when running particle filtering methods with hundreds of thousands of particles being run in parallel across method replications. To address this, Pypomp includes replacement functions that use Inverse Transform Sampling to generate random variables. These functions use JAX under the hood, so they can be used in a JIT-compiled context. While the following functions include some branching in order to handle edge cases, the performance loss from warp divergence is minimal.

Random Variate Generators

fast_poisson(key, lam[, dtype, ...])	Sample Poisson random variates using a GPU-optimized inverse CDF algorithm.
fast_binomial(key, n, p[, order, exact_max, ...])	Sample binomial random variates using a GPU-optimized inverse CDF algorithm.
fast_multinomial(key, n, p[, order, ...])	Sample multinomial random variates using a GPU-optimized inverse CDF algorithm.
fast_gamma(key, alpha[, dtype, ...])	Sample Gamma random variates using a GPU-optimized inverse CDF algorithm.
fast_nbinomial(key, n[, p, mu, dtype, ...])	Sample Negative Binomial random variates using a GPU-optimized algorithm.

Inverse Cumulative Distribution Functions (CDFs)

poissoninv(u, lam[, dtype, ...])	Compute the approximate inverse Poisson CDF using JAX primitives.
binominv(u, n, p[, exact_max, order, dtype])	Compute the approximate inverse binomial CDF using JAX primitives.
gammainv(u, alpha[, dtype, newton_steps])	Compute the approximate inverse Gamma CDF using JAX primitives.