import json
import numpy as np
from polyagamma import random_polyagamma
TEST_CASES = [
(1.0, 0.0, "alternate"),
(1.0, 2.0, "alternate"),
(2.0, 0.0, "alternate"),
(2.0, 4.0, "alternate"),
(3.0, 6.0, "alternate"),
(4.0, 0.0, "alternate"),
(4.0, 8.0, "alternate"),
(7.0, 6.0, "alternate"),
(5.0, 0.0, "saddlepoint"),
(5.0, 3.0, "saddlepoint"),
(7.0, 2.0, "saddlepoint"),
(8.0, 0.0, "saddlepoint"),
(10.0, 5.0, "saddlepoint"),
(20.0, 0.0, "saddlepoint"),
(30.0, 8.0, "saddlepoint"),
]
N_REF = 500_000 N_QUANTILES = 500
probs = np.linspace(0.001, 0.999, N_QUANTILES)
rng = np.random.default_rng(seed=0)
results = []
for h, z, path in TEST_CASES:
samples = random_polyagamma(h, z, size=N_REF, random_state=rng)
quantiles = np.quantile(samples, probs).tolist()
results.append({
"h": h,
"z": z,
"path": path,
"probs": probs.tolist(),
"quantiles": quantiles,
})
print(f"PG({h:.1f}, {z:.1f}) path={path} mean={np.mean(samples):.6f}")
out_path = "tests/pg_reference_quantiles.json"
with open(out_path, "w") as f:
json.dump(results, f, separators=(",", ":"))
print(f"\nWrote {len(results)} cases to {out_path}")