<commit_before>local S = terralib.require("qs.lib.std")
local erp = terralib.require("qs.erp")
local distrib = terralib.require("qs.distrib")
local tmath = terralib.require("qs.lib.tmath")
local qs = terralib.require("qs.globals")



local ERPs = {}

--------------------------------------------

ERPs.flip = erp.makeRandomChoice(
	distrib.bernoulli,
	S.memoize(function(real)
		return terra(currval: bool, p: real) : {bool, real, real}
			if currval then
				return false, real(0.0), real(0.0)
			else
				return true, real(0.0), real(0.0)
			end
		end
	end)
)

--------------------------------------------

ERPs.uniform = erp.makeRandomChoice(
	distrib.uniform,
	nil,
	erp.Bounds.LowerUpper(function(real)
		return terra(lo: real, hi: real)
			return lo, hi
		end
	end)
)

--------------------------------------------

ERPs.gaussian = erp.makeRandomChoice(
	distrib.gaussian,
	-- Drift kernel
	S.memoize(function(real)
		return terra(currval: real, mu: real, sigma: real)
			var newval = [distrib.gaussian(real)].sample(currval, sigma)
			var fwdlp = [distrib.gaussian(real)].logprob(newval, currval, sigma)
			var rvslp = [distrib.gaussian(real)].logprob(currval, newval, sigma)
			return newval, fwdlp, rvslp
		end
	end)
)

--------------------------------------------

ERPs.gamma = erp.makeRandomChoice(
	distrib.gamma,
	nil,
	erp.Bounds.Lower(function(real)
		return terra(shape: real, scale: real)
			return 0.0
		end
	end)
)

-- Often more convenient to specify a gamma in terms of mean and variance
ERPs.gammamv = macro(function(m, v, opts)
	opts = opts or `{}
	return quote
		var shape = m*m / v
		var scale = v / m
	in
		ERPs.gamma(shape, scale, opts)
	end
end)
ERPs.gammamv.observe = macro(function(val, m, v)
	return quote
		var shape = m*m / v
		var scale = v / m
	in
		ERPs.gamma.observe(val, shape, scale)
	end
end)

--------------------------------------------

ERPs.beta = erp.makeRandomChoice(
	distrib.beta,
	nil,
	erp.Bounds.LowerUpper(function(real)
		return terra(a: real, b: real)
			return 0.0, 1.0
		end
	end)
)

-- Often more convenient to specify a beta in terms of mean and variance
ERPs.betamv = macro(function(m, v, opts)
	opts = opts or `{}
	return quote
		var tmp = m*m - m + v
		var a = - m * tmp / v
		var b = (m - 1) * tmp / v
	in
		ERPs.beta(a, b, opts)
	end
end)<commit_msg>Bounds-checking assert for betamv.<commit_after>local S = terralib.require("qs.lib.std")
local erp = terralib.require("qs.erp")
local distrib = terralib.require("qs.distrib")
local tmath = terralib.require("qs.lib.tmath")
local qs = terralib.require("qs.globals")



local ERPs = {}

--------------------------------------------

ERPs.flip = erp.makeRandomChoice(
	distrib.bernoulli,
	S.memoize(function(real)
		return terra(currval: bool, p: real) : {bool, real, real}
			if currval then
				return false, real(0.0), real(0.0)
			else
				return true, real(0.0), real(0.0)
			end
		end
	end)
)

--------------------------------------------

ERPs.uniform = erp.makeRandomChoice(
	distrib.uniform,
	nil,
	erp.Bounds.LowerUpper(function(real)
		return terra(lo: real, hi: real)
			return lo, hi
		end
	end)
)

--------------------------------------------

ERPs.gaussian = erp.makeRandomChoice(
	distrib.gaussian,
	-- Drift kernel
	S.memoize(function(real)
		return terra(currval: real, mu: real, sigma: real)
			var newval = [distrib.gaussian(real)].sample(currval, sigma)
			var fwdlp = [distrib.gaussian(real)].logprob(newval, currval, sigma)
			var rvslp = [distrib.gaussian(real)].logprob(currval, newval, sigma)
			return newval, fwdlp, rvslp
		end
	end)
)

--------------------------------------------

ERPs.gamma = erp.makeRandomChoice(
	distrib.gamma,
	nil,
	erp.Bounds.Lower(function(real)
		return terra(shape: real, scale: real)
			return 0.0
		end
	end)
)

-- Often more convenient to specify a gamma in terms of mean and variance
ERPs.gammamv = macro(function(m, v, opts)
	opts = opts or `{}
	return quote
		var shape = m*m / v
		var scale = v / m
	in
		ERPs.gamma(shape, scale, opts)
	end
end)
ERPs.gammamv.observe = macro(function(val, m, v)
	return quote
		var shape = m*m / v
		var scale = v / m
	in
		ERPs.gamma.observe(val, shape, scale)
	end
end)

--------------------------------------------

ERPs.beta = erp.makeRandomChoice(
	distrib.beta,
	nil,
	erp.Bounds.LowerUpper(function(real)
		return terra(a: real, b: real)
			return 0.0, 1.0
		end
	end)
)

-- Often more convenient to specify a beta in terms of mean and variance
ERPs.betamv = macro(function(m, v, opts)
	opts = opts or `{}
	return quote
		var tmp = m*m - m + v
		-- v must be less than m - m^2, otherwise we get NaNs
		-- (it is impossible to construct a beta distribution
		--    with variance that high)
		S.assert(tmp < 0.0)
		var a = - m * tmp / v
		var b = (m - 1) * tmp / v
	in
		ERPs.beta(a, b, opts)
	end
end)
{'fraction_non_alphanumeric': 0.20947630922693267, 'fraction_numerical': 0.00685785536159601, 'mean_word_length': 2.3563458856345885, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2427857', 'n_tokens_mistral': 1880, 'n_tokens_neox': 1760, 'n_words': 469}