zoo

//
// Zoo: Infer prevalence of 3 animals after a zoo visit
//
// Observed: 3 lions, 2 tigers and 1 bear
//
// Questions:
// - Prevalence of each species
// - Probability of seeing a bear next time
//

var animals = ['lion', 'tiger', 'bear']
var observations = ['lion', 'lion', 'lion', 'tiger', 'tiger', 'bear']

var model = function() {
	// set Dirichlet prior
	var a = [uniform(0,5), uniform(0,5), uniform(0,5)]
	var p = dirichlet(Vector(a))

	// set Categorical distribution
	var cdist = Categorical({ps: p, vs: animals})

	// observe every data point
	var obsFn = function(x) {observe(cdist, x)}
	mapData({data: observations}, obsFn)

	// get animal prevalencies
	return {p_lion: p.data[0], p_tiger:p.data[1], p_bear:p.data[2]};
}


var g = Infer({method: 'MCMC', samples: 1000000, burn: 5000}, model);

viz.marginals(g);