tweak poisson functions

statistics.lua

@@ -3,36 +3,36 @@
 -- https://en.wikipedia.org/wiki/Error_function
 
 local statistics = {}
-local ROOT_2 = math.sqrt(2.0)
-local erf
-local erf_inv
+local random, floor, ceil = math.random, math.floor, math.ceil
+local exp, log, sqrt = math.exp, math.log, math.sqrt
+local ROOT_2 = sqrt(2.0)
 local A = 8 * (math.pi - 3.0) / (3.0 * math.pi * (4.0 - math.pi))
 local B = 4.0 / math.pi
 local C = 2.0 / (math.pi * A)
 local D = 1.0 / A
 
 
-erf = function(x)
+local function erf(x)
 
-	if x == 0 then return 0; end
+	if x == 0 then return 0 end
 
 	local xSq  = x * x
 	local aXSq = A * xSq
-	local v = math.sqrt(1.0 - math.exp(-xSq * (B + aXSq) / (1.0 + aXSq)))
+	local v = sqrt(1.0 - exp(-xSq * (B + aXSq) / (1.0 + aXSq)))
 
 	return (x > 0 and v) or -v
 end
 
 
-erf_inv = function(x)
+local function erf_inv(x)
 
-	if x == 0 then return 0; end
+	if x == 0 then return 0 end
 
-	if x <= -1 or x >= 1 then return nil; end
+	if x <= -1 or x >= 1 then return nil end
 
-	local y = math.log(1 - x * x)
+	local y = log(1 - x * x)
 	local u = C + 0.5 * y
-	local v = math.sqrt(math.sqrt(u * u - D * y) - u)
+	local v = sqrt(sqrt(u * u - D * y) - u)
 
 	return (x > 0 and v) or -v
 end
@@ -43,14 +43,10 @@ local function std_normal(u)
 end
 
 
-local poisson
-local cdf_table = {}
-
-
 local function generate_cdf(lambda_index, lambda)
 
-	local max = math.ceil(4 * lambda)
-	local pdf = math.exp(-lambda)
+	local max = ceil(4 * lambda)
+	local pdf = exp(-lambda)
 	local cdf = pdf
 	local t = { [0] = pdf }
 
@@ -64,30 +60,32 @@ local function generate_cdf(lambda_index, lambda)
 end
 
 
+local cdf_table = {}
+
 for li = 1, 100 do
 	cdf_table[li] = generate_cdf(li, 0.25 * li)
 end
 
 
-poisson = function(lambda, max)
+local function poisson(lambda, max)
 
 	if max < 2 then
-		return (math.random() < math.exp(-lambda) and 0) or 1
+		return (random() < exp(-lambda) and 0) or 1
 	elseif lambda >= 2 * max then
 		return max
 	end
 
-	local u = math.random()
-	local lambda_index = math.floor(4 * lambda + 0.5)
+	local u = random()
+	local lambda_index = floor(4 * lambda + 0.5)
 	local cdfs = cdf_table[lambda_index]
 
 	if cdfs then
 
 		lambda = 0.25 * lambda_index
 
-		if u < cdfs[0] then return 0; end
-		if max > #cdfs then max = #cdfs + 1 else max = math.floor(max); end
-		if u >= cdfs[max - 1] then return max; end
+		if u < cdfs[0] then return 0 end
+		if max > #cdfs then max = #cdfs + 1 else max = floor(max) end
+		if u >= cdfs[max - 1] then return max end
 
 		if max > 4 then  -- Binary search
 
@@ -95,30 +93,27 @@ poisson = function(lambda, max)
 
 			while s + 1 < max do
 
-				local m = math.floor(0.5 * (s + max))
+				local m = floor(0.5 * (s + max))
 
-				if u < cdfs[m] then max = m; else s = m; end
+				if u < cdfs[m] then max = m else s = m end
 			end
 		else
 			for i = 1, max - 1 do
-				if u < cdfs[i] then return i; end
+				if u < cdfs[i] then return i end
 			end
 		end
 
 		return max
 	else
-		local x = lambda + math.sqrt(lambda) * std_normal(u)
+		local x = lambda + sqrt(lambda) * std_normal(u)
 
-		return (x < 0.5 and 0) or (x >= max - 0.5 and max) or math.floor(x + 0.5)
+		return (x < 0.5 and 0) or (x >= max - 0.5 and max) or floor(x + 0.5)
 	end
 end
 
--- Error function.
+-- Error and Inverse error functions
 
 statistics.erf = erf
-
--- Inverse error function.
-
 statistics.erf_inv = erf_inv
 
 --- Standard normal distribution function (mean 0, standard deviation 1).
@@ -126,13 +121,9 @@ statistics.erf_inv = erf_inv
 
 statistics.std_normal = function()
 
-	local u = math.random()
+	local u = random()
 
-	if u < 0.001 then
-		return -3.0
-	elseif u > 0.999 then
-		return 3.0
-	end
+	if u < 0.001 then return -3.0 elseif u > 0.999 then return 3.0 end
 
 	return std_normal(u)
 end
@@ -144,7 +135,7 @@ end
 
 statistics.normal = function(mu, sigma)
 
-	local u = math.random()
+	local u = random()
 
 	if u < 0.001 then
 		return mu - 3.0 * sigma
@@ -164,7 +155,7 @@ statistics.poisson = function(lambda, max)
 
 	lambda, max = tonumber(lambda), tonumber(max)
 
-	if not lambda or not max or lambda <= 0 or max < 1 then return 0; end
+	if not lambda or not max or lambda <= 0 or max < 1 then return 0 end
 
 	return poisson(lambda, max)
 end