dirichlet process

dp.js

// dirichlet process implementation

// determine which portion of the stick it landed on
var portion_hit = function (stick_location, cumulative_sum) {
    var finder = function(current) {
        if (current == stick_location.length) {
            return -1
        } else if (cumulative_sum[current] > stick_location) {
            return current
        } else {
            return finder(current + 1)
        }
    }

    return finder(0)
}

var DP_inner = function (config, state) {
    display("DP_inner(): " + state.i + " / " + config.n)

    if (state.i == config.n) {
        return state.samples
    } else {
        var stick_location = uniform(0, 1)

        // random draw landed somewhere on existing stick, draw a corresponding sample
        if (stick_location <= last(state.cumulative_sum)) {
            var index = portion_hit(stick_location, state.cumulative_sum)
            var params = state.params[index]
            var new_sample = config.sampler(params)

            return DP_inner(config, {
                i: state.i + 1,
                pis: state.pis,
                betas: state.betas,
                params: state.params,
                samples: append(state.samples, new_sample),
                cumulative_sum: state.cumulative_sum,
                beta_product: state.beta_product
            })
        // break another portion of the stick
        } else {
            var beta_draw = beta(1, config.alpha)
            var pi = beta_draw * state.beta_product

            return DP_inner(config, {
                i: state.i,
                pis: append(state.pis, pi),
                betas: append(state.betas, beta_draw),
                params: append(state.params, config.gen_params()),
                samples: state.samples,
                cumulative_sum: append(state.cumulative_sum, pi + last(state.cumulative_sum)),
                beta_product: state.beta_product * (1 - beta_draw)
            })
        }

    }
}

var DP = function(config) {
    var betas = [beta(1, config.alpha)]
    var params = [config.gen_params()]
    return DP_inner(config, {
        i: 1,
        pis: [betas[0]],
        betas: betas,
        params: params,
        samples: [config.sampler(params)],
        cumulative_sum: [betas[0]],
        beta_product: (1 - betas[0])
    })
}

var model = function () {
    var DP_samples = DP({
        alpha: 1,
        n: 5,

        // function called for each time a uniform draw lands on the existing stick
        sampler: function(params) {
            gaussian(params.mu, params.sigma)
        },

        // each group of params is associated with
        gen_params: function () {
            return {
                mu: uniform(0, 100),
                sigma: uniform(0, 3)
            }
        }
    })

    return DP_samples
}

display('Samples:')
sample(Infer({
    method: 'forward',
    samples: 1
}, model))

/*
Output:

Samples:
DP_inner(): 1 / 5
DP_inner(): 1 / 5
DP_inner(): 2 / 5
DP_inner(): 2 / 5
DP_inner(): 3 / 5
DP_inner(): 4 / 5
DP_inner(): 5 / 5
[ [Function], [Function], [Function], [Function], [Function] ]

*/