From 2b530dd15c643f2e931d95f318a5c5c7ab2c1a3c Mon Sep 17 00:00:00 2001
From: kczimm <kevinczimmerman@gmail.com>
Date: Tue, 12 Dec 2017 11:47:55 -0600
Subject: [PATCH] stat: implement Poisson rejection sampling. (#333)

* stat: implement Poisson rejection sampling.
---
 stat/distuv/poisson.go      | 56 +++++++++++++++++++++++++++----------
 stat/distuv/poisson_test.go |  4 +++
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/stat/distuv/poisson.go b/stat/distuv/poisson.go
index 56c63184..9316ff3f 100644
--- a/stat/distuv/poisson.go
+++ b/stat/distuv/poisson.go
@@ -66,26 +66,52 @@ func (p Poisson) Prob(x float64) float64 {
 
 // Rand returns a random sample drawn from the distribution.
 func (p Poisson) Rand() float64 {
-	// poisson generator based upon the multiplication of
-	// uniform random variates.
-	// see:
-	//  Non-Uniform Random Variate Generation,
-	//  Luc Devroye (p504)
-	//  http://www.eirene.de/Devroye.pdf
-	x := 0.0
-	prod := 1.0
-	exp := math.Exp(-p.Lambda)
-	rnd := rand.Float64
+	// NUMERICAL RECIPES IN C: THE ART OF SCIENTIFIC COMPUTING (ISBN 0-521-43108-5)
+	// p. 294
+	// <http://www.aip.de/groups/soe/local/numres/bookcpdf/c7-3.pdf>
+
+	rnd := rand.ExpFloat64
+	if p.Source != nil {
+		rnd = p.Source.ExpFloat64
+	}
+
+	if p.Lambda < 10.0 {
+		// Use direct method.
+		var em float64
+		t := 0.0
+		for {
+			t += rnd()
+			if t >= p.Lambda {
+				break
+			}
+			em++
+		}
+		return em
+	}
+	// Use rejection method.
+	rnd = rand.Float64
 	if p.Source != nil {
 		rnd = p.Source.Float64
 	}
-
+	sq := math.Sqrt(2.0 * p.Lambda)
+	alxm := math.Log(p.Lambda)
+	lg, _ := math.Lgamma(p.Lambda + 1)
+	g := p.Lambda*alxm - lg
 	for {
-		prod *= rnd()
-		if prod <= exp {
-			return x
+		var em, y float64
+		for {
+			y = math.Tan(math.Pi * rnd())
+			em = sq*y + p.Lambda
+			if em >= 0 {
+				break
+			}
+		}
+		em = math.Floor(em)
+		lg, _ = math.Lgamma(em + 1)
+		t := 0.9 * (1.0 + y*y) * math.Exp(em*alxm-lg-g)
+		if rnd() <= t {
+			return em
 		}
-		x++
 	}
 }
 
diff --git a/stat/distuv/poisson_test.go b/stat/distuv/poisson_test.go
index f80ca235..fd8bb27e 100644
--- a/stat/distuv/poisson_test.go
+++ b/stat/distuv/poisson_test.go
@@ -101,6 +101,10 @@ func TestPoissonCDF(t *testing.T) {
 func TestPoisson(t *testing.T) {
 	src := rand.New(rand.NewSource(1))
 	for i, b := range []Poisson{
+		{100, src},
+		{15, src},
+		{10, src},
+		{9.9, src},
 		{3, src},
 		{1.5, src},
 		{0.9, src},