retrospective mcmc crp

gistfile1.js

var rDPmem = $b({
    name: 'rDPmem',
    desc: 'Retrospective implementation of DPmem',
    params: [{name: 'alpha', type: 'positive real', desc: 'Concentration parameter'},
             {name: 'f', type: 'function', desc: 'Function to stochastically memoize'}],
    fn: function(alpha, f) {
        var allStickSets = {};
        var cumulativeStickSums = {};
        var caches = {};
        
        return function() {
            var args = args_to_array(arguments);

            var extractingTables = false
            if (args[0] == 'get-tables') {
                extractingTables = true;
                args.shift(); 
            }

            var argsHash = JSON.stringify(args);

            if (!allStickSets[argsHash]) {
                allStickSets[argsHash] = [];
                caches[argsHash] = {}
                cumulativeStickSums[argsHash] = 0;
            }
            
            
            var sticks = allStickSets[argsHash];
            var cache = caches[argsHash]; 
            var stickSumSoFar = cumulativeStickSums[argsHash];

            if (extractingTables) {
                var ret = [];
                for(var i = 0; i < sticks.length; i++) {
                    var key = i + "";
                    if (cache[key]) {
                        var table = cache[key];
                        ret.push([table.label,table.count])
                    }
                }
                return arrayToList(ret);
            }
            
            // sample a probability for this customer
            var customerLength = wrapped_uniform(0,1);
            

            // create new sticks until we have enough to exceed customerLength 
            while (customerLength > stickSumSoFar) {
                var next = {
                    beta: wrapped_beta(1.0, alpha)
                };
                if (sticks.length == 0) {
                    next.length = next.beta;
                    next.endpoint = next.beta
                } else {
                    var prev = sticks[sticks.length - 1];
                    next.length = prev.length * (1-prev.beta)/prev.beta * next.beta;
                    next.endpoint = prev.endpoint + next.length;
                }

                stickSumSoFar += next.length;
                sticks.push(next);

            }

            cumulativeStickSums[argsHash] = stickSumSoFar;

            var customerStick; 

            for(var i = 0, ii = sticks.length; i < ii; i++) {
                if (customerLength < sticks[i].endpoint) {
                    customerStick = i + "";
                    break;
                }
            }

            if (cache[customerStick] === undefined) {
                cache[customerStick] = {label: f.apply(null, arguments),
                                        count: 1};
            } else {
                cache[customerStick].count++;
            }

            return cache[customerStick].label;

        };
    }
})