Random Number Generators Benchmark & Comparison

Introduction

This file contains performance benchmarks and statistical validation tests for the random number generators in pypomp.

We compare pypomp’s fast approximate inverse CDF samplers against their counterpart samplers in standard libraries (jax.random and scipy.stats).

Environment Setup & Imports

Environment Setup & Imports Log

import os
import json

USE_CPU = os.environ.get("USE_CPU", "false").lower() == "true"
if USE_CPU:
    os.environ["JAX_PLATFORMS"] = "cpu"

import jax
import jax.numpy as jnp
import numpy as np
import time
import matplotlib.pyplot as plt
import pypomp.random as ppr
import warnings
from scipy import stats
from jax.scipy.stats import binom as jax_binom
from scipy.stats import binom as scipy_binom

# Determine benchmarking modes
BENCHMARK_GPU = os.environ.get("BENCHMARK_GPU", "false").lower() == "true"
BENCHMARK_CPU = os.environ.get("BENCHMARK_CPU", "false").lower() == "true"

print("JAX Devices:", jax.devices())
print("Using CPU Mode (forced):", USE_CPU)
print("BENCHMARK_GPU Mode:", BENCHMARK_GPU)
print("BENCHMARK_CPU Mode:", BENCHMARK_CPU)

# Detect if JAX is running on GPU
HAS_GPU = not USE_CPU and len(jax.devices()) > 0 and jax.devices()[0].platform == "gpu"
if HAS_GPU:
    print("GPU Model Used:", jax.devices()[0].device_kind)

# Load cached execution times if they exist
CACHE_FILE = "benchmark_results.json"
benchmark_results = {}
if os.path.exists(CACHE_FILE):
    try:
        with open(CACHE_FILE, "r") as f:
            benchmark_results = json.load(f)
        print("Loaded cached benchmark results from", CACHE_FILE)
        
        # Convert older format cache where GPU times were stored under 't_pp' / 't_jax'
        for category in list(benchmark_results.keys()):
            if isinstance(benchmark_results[category], dict):
                for name in list(benchmark_results[category].keys()):
                    entry = benchmark_results[category][name]
                    if isinstance(entry, dict):
                        if "t_pp" in entry and "t_pp_gpu" not in entry:
                            entry["t_pp_gpu"] = entry.pop("t_pp")
                        if "t_jax" in entry and "t_jax_gpu" not in entry:
                            entry["t_jax_gpu"] = entry.pop("t_jax")
    except Exception as e:
        print(f"Error loading cached benchmark results: {e}")

ERROR:2026-07-06 15:40:30,755:jax._src.xla_bridge:491: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/home/aaronabk/research/quant/.venv/lib64/python3.12/site-packages/jax/_src/xla_bridge.py", line 489, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/home/aaronabk/research/quant/.venv/lib64/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 328, in initialize
    _check_cuda_versions(raise_on_first_error=True)
  File "/home/aaronabk/research/quant/.venv/lib64/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 285, in _check_cuda_versions
    local_device_count = cuda_versions.cuda_device_count()
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: jaxlib/cuda/versions_helpers.cc:113: operation cuInit(0) failed: Unknown CUDA error 303; cuGetErrorName failed. This probably means that JAX was unable to load the CUDA libraries.

JAX Devices: [CpuDevice(id=0)]
Using CPU Mode (forced): True
BENCHMARK_GPU Mode: False
BENCHMARK_CPU Mode: True
Loaded cached benchmark results from benchmark_results.json

Performance Benchmarks

We measure execution times for generating samples across multiple parameter ranges using both JAX built-ins and pypomp’s fast approximations. We generate samples in batches of 1,000,000 for 20 replications, then record the average time. The number 1,000,000 was chosen because it’s a realistic number of parallel draws that could be made in a particle filtering step across particles and method replications. In our experience, increasing the number of samples further causes the speedup ratio to increase significantly, so this benchmark could be conservative in some circumstances. We also implemented the 20 replications as a JIT-compiled scan to reduce overhead and better approximate how the samplers would be used in practice.

Poisson Performance

n = 1_000_000 # number of draws per rep

def poisson_performance():
    key = jax.random.key(42)
    reps = 20

    regimes = {
        "Low Rate (lambda=0.01)": jnp.full((n,), 0.01, dtype=jnp.float32),
        "Before pypomp Switch (lambda=3.99)": jnp.full((n,), 3.99, dtype=jnp.float32),
        "After pypomp Switch (lambda=4.01)": jnp.full((n,), 4.01, dtype=jnp.float32),
        "Before JAX Switch (lambda=9.99)": jnp.full((n,), 9.99, dtype=jnp.float32),
        "After JAX Switch (lambda=10.01)": jnp.full((n,), 10.01, dtype=jnp.float32),
        "High Rate (lambda=100.0)": jnp.full((n,), 100.0, dtype=jnp.float32),
        "Mixed Rates (around switches)": jnp.repeat(
            jnp.array(
                [0.01, 3.99, 4.01, 9.99, 10.01, 20.0, 50.0, 100.0], dtype=jnp.float32
            ),
            n // 8,
        ),
    }

    if "poisson" not in benchmark_results:
        benchmark_results["poisson"] = {}

    # Run/update GPU benchmark if BENCHMARK_GPU is requested
    if BENCHMARK_GPU:
        for name, lam_samples in regimes.items():
            key1, key2 = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_poisson, key1, lam_samples, reps=reps) * 1000
            t_jax = run_benchmark(jax.random.poisson, key2, lam_samples, reps=reps) * 1000
            if name not in benchmark_results["poisson"]:
                benchmark_results["poisson"][name] = {}
            benchmark_results["poisson"][name]["t_pp_gpu"] = t_pp
            benchmark_results["poisson"][name]["t_jax_gpu"] = t_jax

    # Run/update CPU benchmark if BENCHMARK_CPU is requested
    if BENCHMARK_CPU:
        for name, lam_samples in regimes.items():
            key1, key2 = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_poisson, key1, lam_samples, reps=reps) * 1000
            t_jax = run_benchmark(jax.random.poisson, key2, lam_samples, reps=reps) * 1000
            if name not in benchmark_results["poisson"]:
                benchmark_results["poisson"][name] = {}
            benchmark_results["poisson"][name]["t_pp_cpu"] = t_pp
            benchmark_results["poisson"][name]["t_jax_cpu"] = t_jax

    # Save cache if benchmarks were executed
    if BENCHMARK_GPU or BENCHMARK_CPU:
        try:
            with open(CACHE_FILE, "w") as f:
                json.dump(benchmark_results, f, indent=2)
        except Exception as e:
            print(f"Error saving cache: {e}")

    # Display GPU table if GPU times are present in cache
    has_gpu = all("t_pp_gpu" in benchmark_results["poisson"].get(name, {}) for name in regimes)
    if has_gpu:
        print("#### GPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_poisson (ms) | jax.random.poisson (ms) | Speedup Ratio |")
        print("|---|---|---|---|")
        for name in regimes:
            entry = benchmark_results["poisson"][name]
            t_pp = entry["t_pp_gpu"]
            t_jax = entry["t_jax_gpu"]
            ratio = t_jax / t_pp
            print(f"| {name} | {t_pp:.2f} ms | {t_jax:.2f} ms | **{ratio:.2f}x** |")
        print("\n")

    # Display CPU table if CPU times are present in cache
    has_cpu = all("t_pp_cpu" in benchmark_results["poisson"].get(name, {}) for name in regimes)
    if has_cpu:
        print("#### CPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_poisson (ms) | jax.random.poisson (ms) | Speedup Ratio |")
        print("|---|---|---|---|")
        for name in regimes:
            entry = benchmark_results["poisson"][name]
            t_pp = entry["t_pp_cpu"]
            t_jax = entry["t_jax_cpu"]
            ratio = t_jax / t_pp
            print(f"| {name} | {t_pp:.2f} ms | {t_jax:.2f} ms | **{ratio:.2f}x** |")
        print("\n")


poisson_performance()

GPU Performance Benchmarks

Parameter Regime	ppr.fast_poisson (ms)	jax.random.poisson (ms)	Speedup Ratio
Low Rate (lambda=0.01)	0.02 ms	0.41 ms	17.31x
Before pypomp Switch (lambda=3.99)	0.02 ms	0.77 ms	31.47x
After pypomp Switch (lambda=4.01)	0.03 ms	0.78 ms	25.94x
Before JAX Switch (lambda=9.99)	0.03 ms	1.08 ms	43.09x
After JAX Switch (lambda=10.01)	0.02 ms	0.47 ms	19.84x
High Rate (lambda=100.0)	0.02 ms	0.36 ms	14.65x
Mixed Rates (around switches)	0.02 ms	1.08 ms	44.33x

CPU Performance Benchmarks

Parameter Regime	ppr.fast_poisson (ms)	jax.random.poisson (ms)	Speedup Ratio
Low Rate (lambda=0.01)	27.57 ms	396.82 ms	14.40x
Before pypomp Switch (lambda=3.99)	27.61 ms	615.85 ms	22.30x
After pypomp Switch (lambda=4.01)	27.59 ms	616.28 ms	22.33x
Before JAX Switch (lambda=9.99)	27.61 ms	789.31 ms	28.59x
After JAX Switch (lambda=10.01)	27.58 ms	519.97 ms	18.85x
High Rate (lambda=100.0)	27.52 ms	367.52 ms	13.35x
Mixed Rates (around switches)	27.52 ms	854.55 ms	31.05x

Binomial Performance

def binomial_performance():
    key = jax.random.key(43)
    reps = 20

    # Mixed parameter grid setup
    trials = jnp.array([10, 40, 100, 250], dtype=jnp.float32)
    p = jnp.array([0.01, 0.039, 0.041, 0.099, 0.101, 0.499, 0.501, 0.99], dtype=jnp.float32)
    trial_grid, p_grid = jnp.meshgrid(trials, p, indexing="ij")
    trial_flat = trial_grid.reshape(-1)
    p_flat = p_grid.reshape(-1)
    n_repeat = n // len(trial_flat)
    trial_samples = jnp.tile(trial_flat, n_repeat)
    p_samples = jnp.tile(p_flat, n_repeat)

    regimes = {
        "Before pypomp Switch (n=100, p=0.039)": (
            jnp.full((n,), 100, dtype=jnp.int32),
            jnp.full((n,), 0.039, dtype=jnp.float32),
        ),
        "After pypomp Switch (n=100, p=0.041)": (
            jnp.full((n,), 100, dtype=jnp.int32),
            jnp.full((n,), 0.041, dtype=jnp.float32),
        ),
        "Before JAX Switch (n=100, p=0.099)": (
            jnp.full((n,), 100, dtype=jnp.int32),
            jnp.full((n,), 0.099, dtype=jnp.float32),
        ),
        "After JAX Switch (n=100, p=0.101)": (
            jnp.full((n,), 100, dtype=jnp.int32),
            jnp.full((n,), 0.101, dtype=jnp.float32),
        ),
        "Before Flip Switch (n=100, p=0.499)": (
            jnp.full((n,), 100, dtype=jnp.int32),
            jnp.full((n,), 0.499, dtype=jnp.float32),
        ),
        "After Flip Switch (n=100, p=0.501)": (
            jnp.full((n,), 100, dtype=jnp.int32),
            jnp.full((n,), 0.501, dtype=jnp.float32),
        ),
        "Mixed Grid (around switches)": (trial_samples, p_samples),
    }

    if "binomial" not in benchmark_results:
        benchmark_results["binomial"] = {}

    # Run/update GPU benchmark if BENCHMARK_GPU is requested
    if BENCHMARK_GPU:
        for name, (n_arr, p_arr) in regimes.items():
            key1, key2 = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_binomial, key1, n_arr, p_arr, reps=reps) * 1000
            t_jax = run_benchmark(jax.random.binomial, key2, n_arr, p_arr, reps=reps) * 1000
            if name not in benchmark_results["binomial"]:
                benchmark_results["binomial"][name] = {}
            benchmark_results["binomial"][name]["t_pp_gpu"] = t_pp
            benchmark_results["binomial"][name]["t_jax_gpu"] = t_jax

    # Run/update CPU benchmark if BENCHMARK_CPU is requested
    if BENCHMARK_CPU:
        for name, (n_arr, p_arr) in regimes.items():
            key1, key2 = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_binomial, key1, n_arr, p_arr, reps=reps) * 1000
            t_jax = run_benchmark(jax.random.binomial, key2, n_arr, p_arr, reps=reps) * 1000
            if name not in benchmark_results["binomial"]:
                benchmark_results["binomial"][name] = {}
            benchmark_results["binomial"][name]["t_pp_cpu"] = t_pp
            benchmark_results["binomial"][name]["t_jax_cpu"] = t_jax

    # Save cache if benchmarks were executed
    if BENCHMARK_GPU or BENCHMARK_CPU:
        try:
            with open(CACHE_FILE, "w") as f:
                json.dump(benchmark_results, f, indent=2)
        except Exception as e:
            print(f"Error saving cache: {e}")

    # Display GPU table if GPU times are present in cache
    has_gpu = all("t_pp_gpu" in benchmark_results["binomial"].get(name, {}) for name in regimes)
    if has_gpu:
        print("#### GPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_binomial (ms) | jax.random.binomial (ms) | Speedup Ratio |")
        print("|---|---|---|---|")
        for name in regimes:
            entry = benchmark_results["binomial"][name]
            t_pp = entry["t_pp_gpu"]
            t_jax = entry["t_jax_gpu"]
            ratio = t_jax / t_pp
            print(f"| {name} | {t_pp:.2f} ms | {t_jax:.2f} ms | **{ratio:.2f}x** |")
        print("\n")

    # Display CPU table if CPU times are present in cache
    has_cpu = all("t_pp_cpu" in benchmark_results["binomial"].get(name, {}) for name in regimes)
    if has_cpu:
        print("#### CPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_binomial (ms) | jax.random.binomial (ms) | Speedup Ratio |")
        print("|---|---|---|---|")
        for name in regimes:
            entry = benchmark_results["binomial"][name]
            t_pp = entry["t_pp_cpu"]
            t_jax = entry["t_jax_cpu"]
            ratio = t_jax / t_pp
            print(f"| {name} | {t_pp:.2f} ms | {t_jax:.2f} ms | **{ratio:.2f}x** |")
        print("\n")


binomial_performance()

GPU Performance Benchmarks

Parameter Regime	ppr.fast_binomial (ms)	jax.random.binomial (ms)	Speedup Ratio
Before pypomp Switch (n=100, p=0.039)	0.02 ms	0.77 ms	31.77x
After pypomp Switch (n=100, p=0.041)	0.02 ms	0.78 ms	32.11x
Before JAX Switch (n=100, p=0.099)	0.02 ms	1.04 ms	43.64x
After JAX Switch (n=100, p=0.101)	0.02 ms	0.51 ms	21.73x
Before Flip Switch (n=100, p=0.499)	0.03 ms	0.47 ms	18.58x
After Flip Switch (n=100, p=0.501)	0.02 ms	0.47 ms	19.22x
Mixed Grid (around switches)	0.02 ms	1.05 ms	43.88x

CPU Performance Benchmarks

Parameter Regime	ppr.fast_binomial (ms)	jax.random.binomial (ms)	Speedup Ratio
Before pypomp Switch (n=100, p=0.039)	15.36 ms	640.05 ms	41.68x
After pypomp Switch (n=100, p=0.041)	15.15 ms	649.51 ms	42.88x
Before JAX Switch (n=100, p=0.099)	15.15 ms	806.00 ms	53.19x
After JAX Switch (n=100, p=0.101)	15.29 ms	581.09 ms	38.00x
Before Flip Switch (n=100, p=0.499)	15.23 ms	522.80 ms	34.33x
After Flip Switch (n=100, p=0.501)	15.16 ms	518.39 ms	34.18x
Mixed Grid (around switches)	15.25 ms	877.63 ms	57.54x

Gamma Performance

def gamma_performance():
    key = jax.random.key(44)
    reps = 20

    alpha = jnp.array([0.1, 0.5, 0.99, 1.01, 2.0, 10.0, 50.0, 100.0], dtype=jnp.float32)
    alpha_samples = jnp.repeat(alpha, n // len(alpha))

    regimes = {
        "Before JAX Switch (alpha=0.99)": jnp.full((n,), 0.99, dtype=jnp.float32),
        "After JAX Switch (alpha=1.01)": jnp.full((n,), 1.01, dtype=jnp.float32),
        "Mixed Shapes (around switch)": alpha_samples,
    }

    if "gamma" not in benchmark_results:
        benchmark_results["gamma"] = {}

    # Run/update GPU benchmark if BENCHMARK_GPU is requested
    if BENCHMARK_GPU:
        for name, alpha_arr in regimes.items():
            key1, key2 = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_gamma, key1, alpha_arr, reps=reps) * 1000
            t_jax = run_benchmark(jax.random.gamma, key2, alpha_arr, reps=reps) * 1000
            if name not in benchmark_results["gamma"]:
                benchmark_results["gamma"][name] = {}
            benchmark_results["gamma"][name]["t_pp_gpu"] = t_pp
            benchmark_results["gamma"][name]["t_jax_gpu"] = t_jax

    # Run/update CPU benchmark if BENCHMARK_CPU is requested
    if BENCHMARK_CPU:
        for name, alpha_arr in regimes.items():
            key1, key2 = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_gamma, key1, alpha_arr, reps=reps) * 1000
            t_jax = run_benchmark(jax.random.gamma, key2, alpha_arr, reps=reps) * 1000
            if name not in benchmark_results["gamma"]:
                benchmark_results["gamma"][name] = {}
            benchmark_results["gamma"][name]["t_pp_cpu"] = t_pp
            benchmark_results["gamma"][name]["t_jax_cpu"] = t_jax

    # Save cache if benchmarks were executed
    if BENCHMARK_GPU or BENCHMARK_CPU:
        try:
            with open(CACHE_FILE, "w") as f:
                json.dump(benchmark_results, f, indent=2)
        except Exception as e:
            print(f"Error saving cache: {e}")

    # Display GPU table if GPU times are present in cache
    has_gpu = all("t_pp_gpu" in benchmark_results["gamma"].get(name, {}) for name in regimes)
    if has_gpu:
        print("#### GPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_gamma (ms) | jax.random.gamma (ms) | Speedup Ratio |")
        print("|---|---|---|---|")
        for name in regimes:
            entry = benchmark_results["gamma"][name]
            t_pp = entry["t_pp_gpu"]
            t_jax = entry["t_jax_gpu"]
            ratio = t_jax / t_pp
            print(f"| {name} | {t_pp:.2f} ms | {t_jax:.2f} ms | **{ratio:.2f}x** |")
        print("\n")

    # Display CPU table if CPU times are present in cache
    has_cpu = all("t_pp_cpu" in benchmark_results["gamma"].get(name, {}) for name in regimes)
    if has_cpu:
        print("#### CPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_gamma (ms) | jax.random.gamma (ms) | Speedup Ratio |")
        print("|---|---|---|---|")
        for name in regimes:
            entry = benchmark_results["gamma"][name]
            t_pp = entry["t_pp_cpu"]
            t_jax = entry["t_jax_cpu"]
            ratio = t_jax / t_pp
            print(f"| {name} | {t_pp:.2f} ms | {t_jax:.2f} ms | **{ratio:.2f}x** |")
        print("\n")


gamma_performance()

GPU Performance Benchmarks

Parameter Regime	ppr.fast_gamma (ms)	jax.random.gamma (ms)	Speedup Ratio
Before JAX Switch (alpha=0.99)	0.03 ms	0.61 ms	18.25x
After JAX Switch (alpha=1.01)	0.03 ms	0.95 ms	30.42x
Mixed Shapes (around switch)	0.03 ms	0.84 ms	25.46x

CPU Performance Benchmarks

Parameter Regime	ppr.fast_gamma (ms)	jax.random.gamma (ms)	Speedup Ratio
Before JAX Switch (alpha=0.99)	90.59 ms	1782.78 ms	19.68x
After JAX Switch (alpha=1.01)	92.64 ms	2884.97 ms	31.14x
Mixed Shapes (around switch)	90.58 ms	2493.77 ms	27.53x

Negative Binomial Performance

def rnbinom_performance():
    key = jax.random.key(45)
    reps = 20

    size = jnp.array([1.0, 10.0, 100.0], dtype=jnp.float32)
    p = jnp.array([0.1, 0.4998, 0.5003, 0.7138, 0.7148, 0.9], dtype=jnp.float32)
    size_grid, p_grid = jnp.meshgrid(size, p, indexing="ij")
    size_flat = size_grid.reshape(-1)
    p_flat = p_grid.reshape(-1)
    n_repeat = max(1, n // len(size_flat))
    size_samples = jnp.tile(size_flat, n_repeat)
    p_samples = jnp.tile(p_flat, n_repeat)

    regimes = {
        "Before Poisson Switch (size=10.0, p=0.7148)": (
            jnp.full((n,), 10.0, dtype=jnp.float32),
            jnp.full((n,), 10.0 / 13.99, dtype=jnp.float32),
        ),
        "After Poisson Switch (size=10.0, p=0.7138)": (
            jnp.full((n,), 10.0, dtype=jnp.float32),
            jnp.full((n,), 10.0 / 14.01, dtype=jnp.float32),
        ),
        "Before JAX Poisson Switch (size=10.0, p=0.5003)": (
            jnp.full((n,), 10.0, dtype=jnp.float32),
            jnp.full((n,), 10.0 / 19.99, dtype=jnp.float32),
        ),
        "After JAX Poisson Switch (size=10.0, p=0.4998)": (
            jnp.full((n,), 10.0, dtype=jnp.float32),
            jnp.full((n,), 10.0 / 20.01, dtype=jnp.float32),
        ),
        "Mixed Grid (around switches)": (size_samples, p_samples),
    }

    if "nbinomial" not in benchmark_results:
        benchmark_results["nbinomial"] = {}

    # Run/update GPU benchmark if BENCHMARK_GPU is requested
    if BENCHMARK_GPU:
        for name, (size_arr, p_arr) in regimes.items():
            key1, _ = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_nbinomial, key1, size_arr, p=p_arr, reps=reps) * 1000
            if name not in benchmark_results["nbinomial"]:
                benchmark_results["nbinomial"][name] = {}
            benchmark_results["nbinomial"][name]["t_pp_gpu"] = t_pp

    # Run/update CPU benchmark if BENCHMARK_CPU is requested
    if BENCHMARK_CPU:
        for name, (size_arr, p_arr) in regimes.items():
            key1, _ = jax.random.split(key)
            t_pp = run_benchmark(ppr.fast_nbinomial, key1, size_arr, p=p_arr, reps=reps) * 1000
            if name not in benchmark_results["nbinomial"]:
                benchmark_results["nbinomial"][name] = {}
            benchmark_results["nbinomial"][name]["t_pp_cpu"] = t_pp

    # Save cache if benchmarks were executed
    if BENCHMARK_GPU or BENCHMARK_CPU:
        try:
            with open(CACHE_FILE, "w") as f:
                json.dump(benchmark_results, f, indent=2)
        except Exception as e:
            print(f"Error saving cache: {e}")

    # Display GPU table if GPU times are present in cache
    has_gpu = all("t_pp_gpu" in benchmark_results["nbinomial"].get(name, {}) for name in regimes)
    if has_gpu:
        print("#### GPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_nbinomial (ms) |")
        print("|---|---|")
        for name in regimes:
            t_pp = benchmark_results["nbinomial"][name]["t_pp_gpu"]
            print(f"| {name} | {t_pp:.2f} ms |")
        print("\n")

    # Display CPU table if CPU times are present in cache
    has_cpu = all("t_pp_cpu" in benchmark_results["nbinomial"].get(name, {}) for name in regimes)
    if has_cpu:
        print("#### CPU Performance Benchmarks")
        print("| Parameter Regime | ppr.fast_nbinomial (ms) |")
        print("|---|---|")
        for name in regimes:
            t_pp = benchmark_results["nbinomial"][name]["t_pp_cpu"]
            print(f"| {name} | {t_pp:.2f} ms |")
        print("\n")


rnbinom_performance()

GPU Performance Benchmarks

Parameter Regime	ppr.fast_nbinomial (ms)
Before Poisson Switch (size=10.0, p=0.7148)	0.07 ms
After Poisson Switch (size=10.0, p=0.7138)	0.07 ms
Before JAX Poisson Switch (size=10.0, p=0.5003)	0.07 ms
After JAX Poisson Switch (size=10.0, p=0.4998)	0.07 ms
Mixed Grid (around switches)	0.07 ms

CPU Performance Benchmarks

Parameter Regime	ppr.fast_nbinomial (ms)
Before Poisson Switch (size=10.0, p=0.7148)	147.36 ms
After Poisson Switch (size=10.0, p=0.7138)	147.28 ms
Before JAX Poisson Switch (size=10.0, p=0.5003)	147.29 ms
After JAX Poisson Switch (size=10.0, p=0.4998)	147.21 ms
Mixed Grid (around switches)	147.71 ms

Statistical Distance Benchmarks

To quantify the statistical accuracy of the fast approximate samplers relative to the reference JAX/SciPy samplers, we compute the Wasserstein Distance (1st Wasserstein metric) across a range of parameter values using 1,000,000 generated samples for each.

We plot both:

1. The Absolute Wasserstein Distance (transport cost in the units of the random variable).
2. The Relative Wasserstein Distance (normalized by the standard deviation of the reference distribution, \(W_1 / \sigma\)), providing a scale-invariant measure of the approximation error relative to the natural spread of the distribution.

def parameter_wise_wasserstein_benchmarks():
    n_samples = 1_000_000
    key = jax.random.key(200)
    
    # We will check cache first to avoid recomputing on every render
    grids_data = benchmark_results.get("grids", {})
    
    # Define grids for each distribution
    poisson_lambdas = np.geomspace(0.1, 200.0, 30)
    binom_ps = np.linspace(0.01, 0.99, 30)
    binom_n = 100
    gamma_alphas = np.geomspace(0.1, 100.0, 30)
    nbinom_ps = np.linspace(0.05, 0.95, 30)
    nbinom_size = 20
    
    dirty = False
    
    # --- Poisson ---
    if "poisson" not in grids_data:
        poisson_w = []
        poisson_w_norm = []
        for lam in poisson_lambdas:
            key, key_pp, key_ref = jax.random.split(key, 3)
            lam_arr = jnp.full((n_samples,), lam, dtype=jnp.float32)
            pp_samples = ppr.fast_poisson(key_pp, lam_arr)
            ref_samples = jax.random.poisson(key_ref, lam_arr)
            w_dist = float(stats.wasserstein_distance(np.array(pp_samples), np.array(ref_samples)))
            std = np.sqrt(lam)
            poisson_w.append(w_dist)
            poisson_w_norm.append(w_dist / std)
        grids_data["poisson"] = {
            "lambdas": poisson_lambdas.tolist(),
            "w_dist": poisson_w,
            "w_dist_norm": poisson_w_norm
        }
        dirty = True
        
    # --- Binomial ---
    if "binomial" not in grids_data:
        binom_w = []
        binom_w_norm = []
        for p_val in binom_ps:
            key, key_pp, key_ref = jax.random.split(key, 3)
            n_arr = jnp.full((n_samples,), binom_n, dtype=jnp.int32)
            p_arr = jnp.full((n_samples,), p_val, dtype=jnp.float32)
            pp_samples = ppr.fast_binomial(key_pp, n_arr, p_arr)
            ref_samples = jax.random.binomial(key_ref, n=n_arr, p=p_arr)
            w_dist = float(stats.wasserstein_distance(np.array(pp_samples), np.array(ref_samples)))
            std = np.sqrt(binom_n * p_val * (1.0 - p_val))
            binom_w.append(w_dist)
            binom_w_norm.append(w_dist / std)
        grids_data["binomial"] = {
            "ps": binom_ps.tolist(),
            "w_dist": binom_w,
            "w_dist_norm": binom_w_norm,
            "n_fixed": binom_n
        }
        dirty = True
        
    # --- Gamma ---
    if "gamma" not in grids_data:
        gamma_w = []
        gamma_w_norm = []
        for alpha in gamma_alphas:
            key, key_pp, key_ref = jax.random.split(key, 3)
            alpha_arr = jnp.full((n_samples,), alpha, dtype=jnp.float32)
            pp_samples = ppr.fast_gamma(key_pp, alpha_arr)
            ref_samples = jax.random.gamma(key_ref, alpha_arr)
            w_dist = float(stats.wasserstein_distance(np.array(pp_samples), np.array(ref_samples)))
            std = np.sqrt(alpha)
            gamma_w.append(w_dist)
            gamma_w_norm.append(w_dist / std)
        grids_data["gamma"] = {
            "alphas": gamma_alphas.tolist(),
            "w_dist": gamma_w,
            "w_dist_norm": gamma_w_norm
        }
        dirty = True
        
    # --- Negative Binomial ---
    if "nbinomial" not in grids_data:
        nbinom_w = []
        nbinom_w_norm = []
        for p_val in nbinom_ps:
            key, key_pp, key_ref = jax.random.split(key, 3)
            size_arr = jnp.full((n_samples,), nbinom_size, dtype=jnp.float32)
            p_arr = jnp.full((n_samples,), p_val, dtype=jnp.float32)
            pp_samples = ppr.fast_nbinomial(key_pp, size_arr, p=p_arr)
            size_np = np.array(size_arr)
            p_np = np.array(p_arr)
            ref_samples = stats.nbinom.rvs(size_np, p_np)
            w_dist = float(stats.wasserstein_distance(np.array(pp_samples), np.array(ref_samples)))
            std = np.sqrt(nbinom_size * (1.0 - p_val)) / p_val
            nbinom_w.append(w_dist)
            nbinom_w_norm.append(w_dist / std)
        grids_data["nbinomial"] = {
            "ps": nbinom_ps.tolist(),
            "w_dist": nbinom_w,
            "w_dist_norm": nbinom_w_norm,
            "size_fixed": nbinom_size
        }
        dirty = True
        
    # Save cache if dirty
    if dirty:
        benchmark_results["grids"] = grids_data
        try:
            with open(CACHE_FILE, "w") as f:
                json.dump(benchmark_results, f, indent=2)
        except Exception as e:
            print(f"Error saving cache: {e}")
            
    # Set up matplotlib style
    plt.rcParams.update({
        'font.size': 10,
        'axes.labelsize': 11,
        'axes.titlesize': 12,
        'xtick.labelsize': 9,
        'ytick.labelsize': 9,
        'grid.color': '#DDDDDD',
        'grid.linestyle': ':',
        'grid.linewidth': 0.8
    })
    
    fig, axes = plt.subplots(2, 2, figsize=(11, 9))
    
    # 1. Poisson Plot
    ax = axes[0, 0]
    data = grids_data["poisson"]
    ax.plot(data["lambdas"], data["w_dist"], color="tab:blue", marker="o", label="Absolute W1")
    ax.set_xscale("log")
    ax.set_xlabel("Rate (lambda)")
    ax.set_ylabel("Absolute Wasserstein Distance", color="tab:blue")
    ax.tick_params(axis='y', labelcolor="tab:blue")
    ax.grid(True)
    
    ax_twin = ax.twinx()
    ax_twin.plot(data["lambdas"], data["w_dist_norm"], color="tab:orange", marker="s", linestyle="--", label="Relative W1 / std")
    ax_twin.set_ylabel("Relative Wasserstein (W1 / std)", color="tab:orange")
    ax_twin.tick_params(axis='y', labelcolor="tab:orange")
    ax.set_title("Poisson Approximation Error vs. Lambda")
    
    # 2. Binomial Plot
    ax = axes[0, 1]
    data = grids_data["binomial"]
    ax.plot(data["ps"], data["w_dist"], color="tab:blue", marker="o")
    ax.set_xlabel(f"Success Probability (p) [n={data['n_fixed']}]")
    ax.set_ylabel("Absolute Wasserstein Distance", color="tab:blue")
    ax.tick_params(axis='y', labelcolor="tab:blue")
    ax.grid(True)
    
    ax_twin = ax.twinx()
    ax_twin.plot(data["ps"], data["w_dist_norm"], color="tab:orange", marker="s", linestyle="--")
    ax_twin.set_ylabel("Relative Wasserstein (W1 / std)", color="tab:orange")
    ax_twin.tick_params(axis='y', labelcolor="tab:orange")
    ax.set_title("Binomial Approximation Error vs. p")
    
    # 3. Gamma Plot
    ax = axes[1, 0]
    data = grids_data["gamma"]
    ax.plot(data["alphas"], data["w_dist"], color="tab:blue", marker="o")
    ax.set_xscale("log")
    ax.set_xlabel("Shape (alpha)")
    ax.set_ylabel("Absolute Wasserstein Distance", color="tab:blue")
    ax.tick_params(axis='y', labelcolor="tab:blue")
    ax.grid(True)
    
    ax_twin = ax.twinx()
    ax_twin.plot(data["alphas"], data["w_dist_norm"], color="tab:orange", marker="s", linestyle="--")
    ax_twin.set_ylabel("Relative Wasserstein (W1 / std)", color="tab:orange")
    ax_twin.tick_params(axis='y', labelcolor="tab:orange")
    ax.set_title("Gamma Approximation Error vs. Alpha")
    
    # 4. Negative Binomial Plot
    ax = axes[1, 1]
    data = grids_data["nbinomial"]
    ax.plot(data["ps"], data["w_dist"], color="tab:blue", marker="o")
    ax.set_xlabel(f"Probability (p) [size={data['size_fixed']}]")
    ax.set_ylabel("Absolute Wasserstein Distance", color="tab:blue")
    ax.tick_params(axis='y', labelcolor="tab:blue")
    ax.grid(True)
    
    ax_twin = ax.twinx()
    ax_twin.plot(data["ps"], data["w_dist_norm"], color="tab:orange", marker="s", linestyle="--")
    ax_twin.set_ylabel("Relative Wasserstein (W1 / std)", color="tab:orange")
    ax_twin.tick_params(axis='y', labelcolor="tab:orange")
    ax.set_title("Negative Binomial Approximation Error vs. p")
    
    plt.tight_layout()
    plt.show()

parameter_wise_wasserstein_benchmarks()

Absolute and relative (scale-normalized) Wasserstein distance across parameter ranges.

Quantile Function Miss Rates

To give a sense of how often pypomp’s approximate inverse CDF implementation errs, we compare the approximate quantile functions poissoninv and binominv against the exact quantile functions from SciPy (scipy.stats.poisson.ppf and scipy.stats.binom.ppf). We evaluate the miss rate, defined as the proportion of \(1,000,000\) random queries \(u \in [0.0001, 0.9999]\) for which the approximate quantile differs from the exact integer quantile.

import scipy.stats as scipy_stats

poisson_results = []
lam_vals = [0.1, 1.0, 5.0, 10.0, 50.0, 100.0, 500.0, 1000.0]
num_samples = 1_000_000
key = jax.random.key(42)
u = jax.random.uniform(key, (num_samples,), minval=0.0001, maxval=0.9999)

# --- Poisson Miss Rates ---
for lam in lam_vals:
    approx_q = ppr.poissoninv(u, lam)
    exact_q = scipy_stats.poisson.ppf(np.array(u), lam)
    misses = np.sum(np.array(approx_q) != exact_q)
    miss_rate = misses / num_samples
    poisson_results.append((lam, miss_rate))

# --- Binomial Miss Rates ---
n_vals = [10, 100, 1000]
p_vals = [0.01, 0.1, 0.5, 0.9, 0.99]
binomial_results = {n: [] for n in n_vals}

for n in n_vals:
    for p in p_vals:
        approx_q = ppr.binominv(u, n, p, exact_max=5)
        exact_q = scipy_stats.binom.ppf(np.array(u), n, p)
        misses = np.sum(np.array(approx_q) != exact_q)
        miss_rate = misses / num_samples
        binomial_results[n].append((p, miss_rate))

Quantile Miss Rate Visualizations

# Set up matplotlib style
plt.rcParams.update({
    'font.size': 10,
    'axes.labelsize': 11,
    'axes.titlesize': 12,
    'xtick.labelsize': 9,
    'ytick.labelsize': 9,
    'grid.color': '#DDDDDD',
    'grid.linestyle': ':',
    'grid.linewidth': 0.8
})

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# 1. Poisson Plot
ax1 = axes[0]
lams = [r[0] for r in poisson_results]
p_miss = [r[1] for r in poisson_results]
# For log-scale plotting, clamp 0 to a small value (e.g. 1e-7) so they show up on the bottom line
p_miss_plot = [m if m > 0 else 1e-7 for m in p_miss]

ax1.plot(lams, p_miss_plot, color="tab:blue", marker="o", linewidth=2)
ax1.set_xscale("log")
ax1.set_yscale("log")
ax1.set_ylim(1e-8, 1e-3)
ax1.set_xlabel("Rate (lambda)")
ax1.set_ylabel("Quantile Miss Rate (clamped at 1e-7 for 0.0)")
ax1.set_title("Poisson Quantile Miss Rate vs. Lambda")
ax1.grid(True, which="both", linestyle=":", alpha=0.5)

# 2. Binomial Plot
ax2 = axes[1]
colors = ["tab:blue", "tab:orange", "tab:green"]
markers = ["o", "s", "^"]
for i, n in enumerate(n_vals):
    p_vals_n = [r[0] for r in binomial_results[n]]
    b_miss = [r[1] for r in binomial_results[n]]
    b_miss_plot = [m if m > 0 else 1e-7 for m in b_miss]
    ax2.plot(p_vals_n, b_miss_plot, color=colors[i], marker=markers[i], 
             linewidth=2, label=f"n = {n}")

ax2.set_yscale("log")
ax2.set_ylim(1e-8, 1e-2)
ax2.set_xlabel("Success Probability (p)")
ax2.set_ylabel("Quantile Miss Rate (clamped at 1e-7 for 0.0)")
ax2.set_title("Binomial Quantile Miss Rate vs. p")
ax2.legend()
ax2.grid(True, which="both", linestyle=":", alpha=0.5)

plt.tight_layout()
plt.show()

Quantile miss rates for the Poisson and Binomial fast approximate inverse CDF implementations.

Distribution Comparisons & Visualizations

We validate the statistical accuracy of the fast approximate samplers using histograms/density plots and Quantile-Quantile (Q-Q) plots. We want the histograms/density plots to overlap as much as possible and the Q-Q plots to fall along the line \(y = x\) as much as possible.

For discrete distributions (Poisson, Binomial, and Negative Binomial), the Q-Q plots use a small amount of random jitter. Because discrete variables only take integer values, their quantiles lie on a grid, packing all of the points onto a few spots and making it difficult to tell if a point off of the \(y = x\) line is a handful of points or thousands of them. Adding jitter spreads the points slightly so that the relative density/number of points at each coordinate pair is clearly visible.

Poisson Distribution Comparison

def compare_rpoisson_and_jax_poisson(
    seed=42,
    lam_vals=[0.0001, 0.1, 1.0, 4.0, 4.01, 8.0, 15, 19.9, 20.1, 25, 30, 100.0, 500.0],
    show_labels=True,
):
    key = jax.random.key(seed)
    n_samples = 1_000_000

    fig, axes = plt.subplots(2, len(lam_vals), figsize=(3.5 * len(lam_vals), 7))
    if len(lam_vals) == 1:
        axes = axes.reshape(2, 1)
    hist_axes = axes[0, :]
    qq_axes = axes[1, :]

    for i, lam_val in enumerate(lam_vals):
        lam_arr = jnp.full((n_samples,), lam_val)
        key_rpoisson, key = jax.random.split(key)
        rpoisson_samples = ppr.fast_poisson(key_rpoisson, lam_arr)
        key_jax, key = jax.random.split(key)
        jax_poisson_samples = jax.random.poisson(key_jax, lam_arr)

        ax_hist = hist_axes[i]
        min_bin = min(
            float(jnp.min(jax_poisson_samples)),
            float(jnp.min(rpoisson_samples)),
        )
        max_bin = max(
            float(jnp.max(jax_poisson_samples)),
            float(jnp.max(rpoisson_samples)),
        )
        bin_edges = jnp.arange(jnp.floor(min_bin), jnp.ceil(max_bin) + 1)
        ax_hist.hist(
            jax_poisson_samples,
            bins=bin_edges,
            alpha=0.5,
            label="jax.random.poisson",
            density=True,
            color="C1",
            edgecolor="k",
        )
        ax_hist.hist(
            rpoisson_samples,
            bins=bin_edges,
            alpha=0.5,
            label="ppr.fast_poisson",
            density=True,
            color="C0",
            edgecolor="k",
        )
        ax_hist.set_title(r"$\lambda$ = {:.2f}".format(lam_val), fontsize=9)
        if show_labels:
            ax_hist.set_xlabel("Sample value", fontsize=9)
            if i == 0:
                ax_hist.set_ylabel("Density", fontsize=9)
            ax_hist.tick_params(axis="both", which="major", labelsize=8)
        else:
            ax_hist.tick_params(labelbottom=False, labelleft=False)
        if i == len(lam_vals) - 1:
            ax_hist.legend(fontsize=8)

        ax_qq = qq_axes[i]
        sorted_jax = np.sort(np.array(jax_poisson_samples))
        sorted_rpoisson = np.sort(np.array(rpoisson_samples))
        
        # Add jitter to discrete values for QQ plot to show density
        np.random.seed(seed)
        x_jitter = np.random.normal(0, 0.1, size=sorted_jax.size)
        y_jitter = np.random.normal(0, 0.1, size=sorted_rpoisson.size)
        
        ax_qq.scatter(
            sorted_jax + x_jitter,
            sorted_rpoisson + y_jitter,
            alpha=0.3,
            s=5,
            color="C0",
        )
        
        min_val = min(float(sorted_jax[0]), float(sorted_rpoisson[0]))
        max_val = max(float(sorted_jax[-1]), float(sorted_rpoisson[-1]))
        ax_qq.plot(
            [min_val, max_val],
            [min_val, max_val],
            color="red",
            linestyle="--",
            alpha=0.7,
            lw=1.5,
            label="y=x",
        )
        if show_labels:
            ax_qq.set_xlabel("jax.random.poisson quantiles", fontsize=9)
            if i == 0:
                ax_qq.set_ylabel("ppr.fast_poisson quantiles", fontsize=9)
            ax_qq.tick_params(axis="both", which="major", labelsize=8)
        else:
            ax_qq.tick_params(labelbottom=False, labelleft=False)
        if i == len(lam_vals) - 1:
            ax_qq.legend(fontsize=8)
    plt.tight_layout()
    plt.show()


compare_rpoisson_and_jax_poisson()

Poisson Distribution comparison: Histogram and Q-Q Plot across rates.

Binomial Distribution Comparison

def compare_rbinom2(
    seed=42,
    n_trials_list=[3, 20, 100, 2000],
    prob_vals=[0.02 / 365.25, 0.01, 0.1, 0.3, 0.5, 0.8, 0.95, 0.99],
    compare_to_exact=False,
    jitter_scale=0.01,  # minor jitter to reduce overplotting
    show_labels=True,
):
    key = jax.random.key(seed)
    n_samples = 1000000

    n_rows = len(n_trials_list)
    n_cols = len(prob_vals)

    plot_data = {}

    for row, n in enumerate(n_trials_list):
        for col, p_val in enumerate(prob_vals):
            p_arr = jnp.full((n_samples,), p_val, dtype=jnp.float32)
            n_arr = jnp.full((n_samples,), n, dtype=jnp.int32)

            key_rbinom, key = jax.random.split(key)
            rbinom_samples = ppr.fast_binomial(key_rbinom, n_arr, p_arr, exact_max=5)

            if compare_to_exact:
                plot_data[(row, col)] = (rbinom_samples, None, n, p_val)
            else:
                key_jax, key = jax.random.split(key)
                jax_binom_samples = jax.random.binomial(key_jax, n=n_arr, p=p_arr)
                plot_data[(row, col)] = (rbinom_samples, jax_binom_samples, n, p_val)

    # --- Histogram Figure ---
    fig_hist, hist_axes = plt.subplots(
        n_rows, n_cols, figsize=(3 * n_cols, 1.8 * n_rows), squeeze=False
    )

    for row, n in enumerate(n_trials_list):
        for col, p_val in enumerate(prob_vals):
            rbinom_samples, jax_samples, n, p_val = plot_data[(row, col)]
            ax = hist_axes[row, col]

            min_v = int(jnp.min(rbinom_samples))
            max_v = int(jnp.max(rbinom_samples))
            if jax_samples is not None:
                min_v = min(min_v, int(jnp.min(jax_samples)))
                max_v = max(max_v, int(jnp.max(jax_samples)))

            bins = np.arange(min_v, max_v + 2) - 0.5

            if compare_to_exact:
                x_exact = np.arange(min_v, max_v + 1)
                pmf_exact = jax_binom.pmf(x_exact, n, p_val)
                ax.step(
                    x_exact,
                    pmf_exact,
                    where="mid",
                    color="red",
                    lw=1.5,
                    label="Exact PMF (SciPy)",
                )
                ax.scatter(x_exact, pmf_exact, color="red", s=10)
            else:
                ax.hist(
                    jax_samples,
                    bins=bins,
                    alpha=0.5,
                    label="jax.random.binomial",
                    density=True,
                    color="C1",
                    edgecolor="k",
                )

            ax.hist(
                rbinom_samples,
                bins=bins,
                alpha=0.5,
                label="ppr.fast_binomial",
                density=True,
                color="C0",
                edgecolor="k",
            )

            ax.set_title(f"n={n}, p={p_val:.4f}", fontsize=8)
            if show_labels:
                if col == 0:
                    ax.set_ylabel("Density", fontsize=8)
                if row == n_rows - 1:
                    ax.set_xlabel("Value", fontsize=8)
                ax.tick_params(axis="both", which="major", labelsize=7)
            else:
                ax.tick_params(
                    axis="both", which="both", labelbottom=False, labelleft=False
                )
            if row == 0 and col == n_cols - 1:
                ax.legend(fontsize=7)

    fig_hist.tight_layout()
    plt.show()

    # --- QQ Plot Figure ---
    fig_qq, qq_axes = plt.subplots(
        n_rows, n_cols, figsize=(3 * n_cols, 1.8 * n_rows), squeeze=False
    )

    for row, n in enumerate(n_trials_list):
        for col, p_val in enumerate(prob_vals):
            rbinom_samples, jax_samples, n, p_val = plot_data[(row, col)]
            ax = qq_axes[row, col]

            sorted_approx = np.sort(np.array(rbinom_samples))

            if compare_to_exact:
                quantiles = np.linspace(0 + 1e-5, 1 - 1e-5, n_samples)
                x_data = scipy_binom.ppf(quantiles, n, p_val)
                xlabel = "Theoretical Quantiles"
            else:
                x_data = np.sort(np.array(jax_samples))
                xlabel = "jax.random.binomial quantiles"

            # Add jitter to reduce overplotting for discrete values
            np.random.seed(seed)
            x_jitter = np.random.normal(0, 0.1, size=x_data.size)
            y_jitter = np.random.normal(0, 0.1, size=sorted_approx.size)
            
            ax.scatter(
                x_data + x_jitter,
                sorted_approx + y_jitter,
                alpha=0.3,
                s=5,
                color="C0",
            )
            if show_labels:
                ax.set_xlabel(xlabel, fontsize=8)
                if col == 0:
                    ax.set_ylabel("ppr.fast_binomial quantiles", fontsize=8)
                ax.tick_params(axis="both", which="major", labelsize=7)
            else:
                ax.tick_params(
                    axis="both", which="both", labelbottom=False, labelleft=False
                )

            limit = [
                min(float(x_data[0]), float(sorted_approx[0])),
                max(float(x_data[-1]), float(sorted_approx[-1])),
            ]
            ax.plot(limit, limit, color="red", linestyle="--", alpha=0.7, lw=1.5, label="y=x")
            ax.set_title(f"n={n}, p={p_val:.4f}", fontsize=8)

    fig_qq.tight_layout()
    plt.show()


compare_rbinom2()

Binomial Distribution comparison: Histogram and Q-Q Plot.

Gamma Distribution Comparison

For continuous Gamma distributions, we compare the shape of the density using smooth Kernel Density Estimation (KDE) plots.

Log-KDE and Log-Scale Q-Q Plots for Extreme Skewness (\(\alpha=0.01\))

For extremely low shape parameters, values span dozens of orders of magnitude. A standard linear KDE is heavily distorted, and setting the x-axis to a log-scale mathematically distorts the area under the curve.

To resolve this, we perform KDE on the log-transformed variable \(Y = \log_{10}(X)\) and plot it on a linear scale, relabeling the ticks to \(10^y\). Because \(P(10^a \le X \le 10^b) = \int_a^b f_Y(y) dy\), the visual area under the curve maps exactly to the true probability mass. Monotonicity of \(\log\) preserves the exact quantile mapping in the Q-Q plots.

def compare_rgamma_and_jax_gamma(
    seed=42,
    alpha_vals=[0.01, 0.5, 1.0, 1.5, 2.0, 5.0, 10.0, 50.0, 100.0],
    show_labels=True,
):
    key = jax.random.key(seed)
    n_samples = 1000000

    fig, axes = plt.subplots(2, len(alpha_vals), figsize=(3.5 * len(alpha_vals), 7))
    if len(alpha_vals) == 1:
        axes = axes.reshape(2, 1)
    hist_axes = axes[0, :]
    qq_axes = axes[1, :]

    for i, alpha_val in enumerate(alpha_vals):
        alpha_arr = jnp.full((n_samples,), alpha_val, dtype=jnp.float32)
        key_rgamma, key = jax.random.split(key)
        rgamma_samples = ppr.fast_gamma(key_rgamma, alpha_arr)
        key_jax, key = jax.random.split(key)
        jax_gamma_samples = jax.random.gamma(key_jax, alpha_arr)

        ax_hist = hist_axes[i]

        if alpha_val == 0.01:
            import matplotlib.ticker as ticker

            # Filter out non-positive samples if any exist, then log-transform
            jax_pos = np.array(jax_gamma_samples[jax_gamma_samples > 0])
            rgamma_pos = np.array(rgamma_samples[rgamma_samples > 0])
            
            log10_jax = np.log10(jax_pos)
            log10_rgamma = np.log10(rgamma_pos)

            kde_jax = stats.gaussian_kde(log10_jax[:20000])
            kde_rgamma = stats.gaussian_kde(log10_rgamma[:20000])

            min_log = min(float(np.min(log10_jax)), float(np.min(log10_rgamma)))
            max_log = max(float(np.max(log10_jax)), float(np.max(log10_rgamma)))
            xs = np.linspace(min_log, max_log, 200)

            ax_hist.plot(xs, kde_jax(xs), color="C1", label="jax.random.gamma", lw=2)
            ax_hist.plot(
                xs,
                kde_rgamma(xs),
                color="C0",
                label="ppr.fast_gamma",
                lw=2,
                linestyle="--",
            )
            ax_hist.fill_between(xs, kde_jax(xs), alpha=0.15, color="C1")
            ax_hist.fill_between(xs, kde_rgamma(xs), alpha=0.15, color="C0")
            ax_hist.xaxis.set_major_formatter(
                ticker.FuncFormatter(lambda x, pos: f"$10^{{{int(np.round(x))}}}$")
            )
        else:
            min_bin = min(
                float(jnp.min(jax_gamma_samples)),
                float(jnp.min(rgamma_samples)),
            )
            max_bin = max(
                float(jnp.max(jax_gamma_samples)),
                float(jnp.max(rgamma_samples)),
            )

            # Estimate densities using gaussian_kde on a subset of samples for speed and smooth curves
            kde_samples_jax = np.array(jax_gamma_samples[:20000])
            kde_samples_rgamma = np.array(rgamma_samples[:20000])
            
            # Check variance to avoid singular covariance matrix in gaussian_kde
            if np.var(kde_samples_jax) > 1e-9 and np.var(kde_samples_rgamma) > 1e-9:
                kde_jax = stats.gaussian_kde(kde_samples_jax)
                kde_rgamma = stats.gaussian_kde(kde_samples_rgamma)
                
                xs = np.linspace(min_bin, max_bin, 200)
                ax_hist.plot(
                    xs,
                    kde_jax(xs),
                    color="C1",
                    label="jax.random.gamma",
                    lw=2,
                )
                ax_hist.plot(
                    xs,
                    kde_rgamma(xs),
                    color="C0",
                    label="ppr.fast_gamma",
                    lw=2,
                    linestyle="--",
                )
                ax_hist.fill_between(xs, kde_jax(xs), alpha=0.15, color="C1")
                ax_hist.fill_between(xs, kde_rgamma(xs), alpha=0.15, color="C0")
            else:
                # Fallback to histogram if KDE fails due to low/zero variance
                n_bins = 50
                bin_edges = jnp.linspace(min_bin, max_bin, n_bins + 1)
                ax_hist.hist(
                    jax_gamma_samples,
                    bins=bin_edges,
                    alpha=0.5,
                    label="jax.random.gamma",
                    density=True,
                    color="C1",
                    edgecolor="k",
                )
                ax_hist.hist(
                    rgamma_samples,
                    bins=bin_edges,
                    alpha=0.5,
                    label="ppr.fast_gamma",
                    density=True,
                    color="C0",
                    edgecolor="k",
                )

        ax_hist.set_title(r"$\alpha$ = {:.2f}".format(alpha_val), fontsize=9)
        if show_labels:
            ax_hist.set_xlabel("Sample value", fontsize=9)
            if i == 0:
                ax_hist.set_ylabel("Density", fontsize=9)
            ax_hist.tick_params(axis="both", which="major", labelsize=8)
        else:
            ax_hist.tick_params(labelbottom=False, labelleft=False)
        if i == len(alpha_vals) - 1:
            ax_hist.legend(fontsize=8)

        ax_qq = qq_axes[i]
        sorted_jax = np.sort(np.array(jax_gamma_samples))
        sorted_rgamma = np.sort(np.array(rgamma_samples))
        
        ax_qq.scatter(sorted_jax, sorted_rgamma, alpha=0.3, s=5, color="C0")
        
        min_val = min(float(sorted_jax[0]), float(sorted_rgamma[0]))
        max_val = max(float(sorted_jax[-1]), float(sorted_rgamma[-1]))
        ax_qq.plot(
            [min_val, max_val],
            [min_val, max_val],
            color="red",
            linestyle="--",
            alpha=0.7,
            lw=1.5,
            label="y=x",
        )
        if alpha_val == 0.01:
            ax_qq.set_xscale("log")
            ax_qq.set_yscale("log")
        if show_labels:
            ax_qq.set_xlabel("jax.random.gamma quantiles", fontsize=9)
            if i == 0:
                ax_qq.set_ylabel("ppr.fast_gamma quantiles", fontsize=9)
            ax_qq.tick_params(axis="both", which="major", labelsize=8)
        else:
            ax_qq.tick_params(labelbottom=False, labelleft=False)
        if i == len(alpha_vals) - 1:
            ax_qq.legend(fontsize=8)
    plt.tight_layout()
    plt.show()


compare_rgamma_and_jax_gamma()

Gamma Distribution comparison: Histogram and Q-Q Plot.

Negative Binomial Distribution Comparison

def compare_rnbinom_and_scipy(
    seed=42,
    n_vals=[1.0, 5.0, 20.0, 100.0],
    p_vals=[0.1, 0.5, 0.9],
    show_labels=True,
):
    key = jax.random.key(seed)
    n_samples = 1000000

    fig, axes = plt.subplots(
        2, len(n_vals) * len(p_vals), figsize=(3 * len(n_vals) * len(p_vals), 7)
    )
    hist_axes = axes[0, :]
    qq_axes = axes[1, :]

    plot_idx = 0
    for n_val in n_vals:
        for p_val in p_vals:
            n_arr = jnp.full((n_samples,), n_val, dtype=jnp.float32)
            p_arr = jnp.full((n_samples,), p_val, dtype=jnp.float32)
            key_rnbinom, key = jax.random.split(key)
            rnbinom_samples = ppr.fast_nbinomial(key_rnbinom, n_arr, p=p_arr)

            ax_hist = hist_axes[plot_idx]
            min_bin = int(jnp.min(rnbinom_samples))
            max_bin = int(jnp.max(rnbinom_samples))
            bin_edges = np.arange(min_bin, max_bin + 2) - 0.5

            ax_hist.hist(
                rnbinom_samples,
                bins=bin_edges,
                alpha=0.5,
                label="ppr.fast_nbinomial",
                density=True,
                color="C0",
                edgecolor="k",
            )

            x_theory = np.arange(min_bin, max_bin + 1)
            pmf_theory = stats.nbinom.pmf(x_theory, n_val, p_val)
            ax_hist.step(
                x_theory,
                pmf_theory,
                where="mid",
                color="red",
                alpha=0.8,
                label="Exact PMF (SciPy)",
            )

            ax_hist.set_title(f"n={n_val}, p={p_val}", fontsize=9)
            if show_labels:
                ax_hist.set_xlabel("Sample value", fontsize=9)
                if plot_idx == 0:
                    ax_hist.set_ylabel("Density", fontsize=9)
                ax_hist.tick_params(axis="both", which="major", labelsize=8)
            else:
                ax_hist.tick_params(labelbottom=False, labelleft=False)
            if plot_idx == len(n_vals) * len(p_vals) - 1:
                ax_hist.legend(fontsize=8)

            ax_qq = qq_axes[plot_idx]
            sorted_samples = np.sort(np.array(rnbinom_samples))
            
            quantiles = np.linspace(1e-5, 1 - 1e-5, n_samples)
            theo_quants = stats.nbinom.ppf(quantiles, n_val, p_val)

            # Add jitter to reduce overplotting for discrete values
            np.random.seed(seed)
            x_jitter = np.random.normal(0, 0.1, size=theo_quants.size)
            y_jitter = np.random.normal(0, 0.1, size=sorted_samples.size)

            ax_qq.scatter(
                theo_quants + x_jitter,
                sorted_samples + y_jitter,
                alpha=0.3,
                s=5,
                color="C0",
            )
            
            limit = [
                min(float(theo_quants[0]), float(sorted_samples[0])),
                max(float(theo_quants[-1]), float(sorted_samples[-1])),
            ]
            ax_qq.plot(limit, limit, color="red", linestyle="--", alpha=0.7, lw=1.5, label="y=x")

            if show_labels:
                ax_qq.set_xlabel("scipy.stats.nbinom quantiles", fontsize=9)
                if plot_idx == 0:
                    ax_qq.set_ylabel("ppr.fast_nbinomial quantiles", fontsize=9)
                ax_qq.tick_params(axis="both", which="major", labelsize=8)
            else:
                ax_qq.tick_params(labelbottom=False, labelleft=False)
            if plot_idx == len(n_vals) * len(p_vals) - 1:
                ax_qq.legend(fontsize=8)

            plot_idx += 1

    plt.tight_layout()
    plt.show()


compare_rnbinom_and_scipy()

Negative Binomial Distribution comparison: Histogram and Q-Q Plot against SciPy exact.