Probability.lua

--[=[
	Probability utility functions
	@class Probability
]=]

local Probability = {}

--[=[
	Returns a boxMuller random distribution
	@return number
]=]
function Probability.boxMuller()
	return math.sqrt(-2 * math.log(math.random())) * math.cos(2 * math.pi * math.random()) / 2
end

--[=[
	Returns a normal distribution
	@param mean number
	@param standardDeviation number
	@return number
]=]
function Probability.normal(mean, standardDeviation)
	return mean + Probability.boxMuller() * standardDeviation
end

--[=[
	Returns a bounded normal, clamping the normal value
	@param mean number
	@param standardDeviation number
	@param hardMin number
	@param hardMax number
	@return number
]=]
function Probability.boundedNormal(mean, standardDeviation, hardMin, hardMax)
	return math.clamp(Probability.normal(mean, standardDeviation), hardMin, hardMax)
end

--[=[
	Approximation of the error function (erf) using the Abramowitz and Stegun formula

	https://en.wikipedia.org/wiki/Error_function

	@param x number
	@return number
]=]
function Probability.erf(x)
	local t = 1 / (1 + 0.5 * math.abs(x))
	local erf_approx = 1 - t * math.exp(-x * x - 1.26551223 +
			t * (1.00002368 +
			t * (0.37409196 +
			t * (0.09678418 +
			t * (-0.18628806 +
			t * (0.27886807 +
			t * (-1.13520398 +
			t * (1.48851587 +
			t * (-0.82215223 +
			t * (0.17087277))))))))))
	if x >= 0 then
		return erf_approx
	else
		return -erf_approx
	end
end

--[=[
	Standard normal cumulative distribution function. Returns the value from 0 to 1.

	This is also known as percentile!

	@param zScore number
	@return number
]=]
function Probability.cdf(zScore)
	assert(type(zScore) == "number", "Bad zScore")

	return 0.5 * (1 + Probability.erf(zScore / math.sqrt(2)))
end

--[=[
	Function to calculate the inverse error function (erfinv) using Newton's method

	@param x number
	@return number
]=]
function Probability.erfinv(x)
	assert(type(x) == "number", "Bad x")

	if x < -1 or x > 1 then
		return nil
	elseif x == -1 then
		return -math.huge
	elseif x == 1 then
		return math.huge
	end

	local tolerance = 1e-15
	local maxIterations = 1000

	local function derivative(y)
		return 2 * math.exp(-y * y) / math.sqrt(math.pi)
	end

	local y = 0
	for _ = 1, maxIterations do
		local error = Probability.erf(y) - x
		if math.abs(error) < tolerance then
			return y
		end
		y = y - error / derivative(y)
	end

	-- If Newton's method fails to converge, return nil
	return nil
end

--[=[
	Standard normal cumulative distribution function. Returns the value from 0 to 1.

	This is also known as percentile!

	@param percentile number
	@return number
]=]
function Probability.percentileToZScore(percentile)
	assert(type(percentile) == "number" and percentile >= 0 and percentile <= 1, "Bad percentile")

	-- Calculate z-score using inverse cumulative distribution function (norm.ppf)
	local zScore = math.sqrt(2) * Probability.erfinv(2 * percentile - 1)

	return zScore
end

return Probability