header.wppl

'no caching';

resetEnv();

// Distributions (note: this gets macro transformed)

defineDistConstructors(
    Bernoulli
    MultivariateBernoulli
    Uniform
    UniformDrift
    RandomInteger
    Gaussian
    GaussianDrift
    MultivariateGaussian
    DiagCovGaussian
    TensorGaussian
    Cauchy
    Discrete
    Gamma
    Exponential
    Beta
    Binomial
    Multinomial
    Poisson
    Dirichlet
    DirichletDrift
    LogisticNormal
    Categorical
    Delta);

// Helpers to sample from distributions.

var bernoulli = function(arg) {
  var params = util.isObject(arg) ? arg : {p: arg};
  return sample(Bernoulli(params));
};

var randomInteger = function(arg) {
  var params = util.isObject(arg) ? arg : {n: arg};
  return sample(RandomInteger(params));
};

var discrete = function(arg) {
  var params = util.isObject(arg) ? arg : {ps: arg};
  return sample(Discrete(params));
};

var multinomial = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {ps: arg1, n: arg2};
  return sample(Multinomial(params));
};

var gaussian = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {mu: arg1, sigma: arg2};
  return sample(Gaussian(params));
};

var multivariateGaussian = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {mu: arg1, cov: arg2};
  return sample(MultivariateGaussian(params));
};

var cauchy = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {location: arg1, scale: arg2};
  return sample(Cauchy(params));
};

var uniform = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {a: arg1, b: arg2};
  return sample(Uniform(params));
};

var dirichlet = function(arg) {
  var params = util.isObject(arg) ? arg : {alpha: arg};
  return sample(Dirichlet(params));
};

var poisson = function(arg) {
  var params = util.isObject(arg) ? arg : {mu: arg};
  return sample(Poisson(params));
};

var binomial = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {p: arg1, n: arg2};
  return sample(Binomial(params));
};

var beta = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {a: arg1, b: arg2};
  return sample(Beta(params));
};

var exponential = function(arg) {
  var params = util.isObject(arg) ? arg : {a: arg};
  return sample(Exponential(params));
};

var gamma = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {shape: arg1, scale: arg2};
  return sample(Gamma(params));
};

// Other distribution helpers

var flip = function(p) {
  var params = {p: (p !== undefined) ? p : .5};
  return sample(Bernoulli(params));
};

var categorical = function(arg1, arg2) {
  var params = util.isObject(arg1) ? arg1 : {ps: arg1, vs: arg2};
  return params.vs[discrete(params.ps)];
};

var delta = function(arg) {
  var params = util.isObject(arg) ? arg : {v: arg};
  return sample(Delta(params));
};

var uniformDraw = function(arg) {
  var vs = util.isObject(arg) ? arg.vs : arg;
  return vs[sample(RandomInteger({n: vs.length}))];
};

var serializeDist = function(d) {
  return dists.serialize(d);
};

var deserializeDist = function(JSONString) {
  return dists.deserialize(JSONString);
};

// XRPs

var makeBetaBernoulli = function(pseudocounts) {
  globalStore.BBindex = 1 + (globalStore.BBindex || 0);
  var bbname = 'BB' + globalStore.BBindex;
  globalStore[bbname] = pseudocounts;
  return function() {
    var pc = globalStore[bbname];  // get current sufficient stats
    var val = sample(Bernoulli({p: pc[0] / (pc[0] + pc[1])}));  // sample from predictive.
    globalStore[bbname] = [pc[0] + val, pc[1] + !val];  // update sufficient stats
    return val;
  };
};

var makeDirichletDiscrete = function(pseudocounts) {
  var addCount = function(a, i, j) {
    var j = j || 0;
    if (a.length === 0) {
      return [];
    } else {
      return [a[0] + (i === j)].concat(addCount(a.slice(1), i, j + 1));
    }
  };
  globalStore.DDindex = 1 + (globalStore.DDindex || 0);
  var ddname = 'DD' + globalStore.DDindex;
  globalStore[ddname] = pseudocounts;
  return function() {
    var pc = globalStore[ddname];  // get current sufficient stats
    var val = sample(Discrete({ps: pc}));  // sample from predictive. (doesn't need to be normalized.)
    globalStore[ddname] = addCount(pc, val); // update sufficient stats
    return val;
  };
};

var isEven = function(v) {
  return v % 2 === 0;
};
var isOdd = function(v) {
  return v % 2 != 0;
};

var idF = function(x) {
  return x;
};
var constF = function(f) {
  return function() {
    return f;
  };
};
var falseF = function() {
  return false;
};
var trueF = function() {
  return true;
};

// Probability computations & calculations

var MAP = function(dist) {
  return dist.MAP();
};

var expectation = function(dist, func) {
  var f = func || idF;
  var supp = dist.support();
  return reduce(function(s, acc) {
    return acc + Math.exp(dist.score(s)) * f(s);
  }, 0, supp);
};

var entropy = function(dist) {
  return dist.entropy();
};

// Data structures & higher-order functions

var append = function(a, b) {
  return a.concat(b);
};
var cons = function(a, b) {
  return [a].concat(b);
};
var snoc = function(a, b) {
  return a.concat([b]);
};

var first = function(xs) {
  return xs[0];
};
var second = function(xs) {
  return xs[1];
};
var third = function(xs) {
  return xs[2];
};
var fourth = function(xs) {
  return xs[3];
};
var secondLast = function(xs) {
  return xs[xs.length - 2];
};
var last = function(xs) {
  return xs[xs.length - 1];
};

var most = function(xs) {
  return xs.slice(0, xs.length - 1);
}

var rest = function(xs) {
  return xs.slice(1);
};

var map_helper = function(i, j, f) {
  var n = j - i + 1;
  if (n == 0) {
    return []
  } else if (n == 1) {
    return [f(i)];
  } else {
    var n1 = Math.ceil(n / 2);
    return map_helper(i, i + n1 - 1, f).concat(map_helper(i + n1, j, f));
  }
}

// recursively split input array so that we only call
// concat a logarithmic number of times
var map = function(fn, l) {
  return map_helper(0, l.length - 1, function(i) { return fn(l[i]) })
}

// assumes that length l1 == length l2
var map2 = function(fn, l1, l2) {
  return map_helper(0, l1.length - 1, function(i) { return fn(l1[i], l2[i]) })
}

// sugar for map(f, [0,1,...,n-1])
var mapN = function(fn, n) {
  return map_helper(0, n - 1, function(i) { return fn(i) })
}

var mapIndexed = function(fn, l) {
  return map_helper(0, l.length - 1, function(i) { return fn(i, l[i]) })
}

var mapObject = function(fn, obj) {
  return _.object(
      map(
      function(kv) {
        return [kv[0], fn(kv[0], kv[1])]
      },
      _.pairs(obj))
  );
};

var reduce = function(fn, init, ar) {
  var n = ar.length;
  var helper = function(i) {
    if (i === n) {
      return init
    } else {
      return fn(ar[i], helper(i + 1))
    }
  }

  return helper(0);
};

var sum = function(l) {
  return reduce(function(a, b) { return a + b; }, 0, l);
};
var product = function(l) {
  return reduce(function(a, b) { return a * b; }, 1, l);
};

var listMean = function(l) {
  return sum(l) / l.length;
};
var listVar = function(l, mu) {
  var mu = mu === undefined ? listMean(l) : mu;
  return reduce(function(a, acc) {
    return acc + (a - mu) * (a - mu);
  }, 0, l) / l.length;

};
var listStdev = function(l, mu) {
  return Math.sqrt(listVar(l, mu));
};

var normalize = function(xs) {
  var Z = sum(xs);
  return map(function(x) {
    return x / Z;
  }, xs);
};

var all = function(p, l) {
  return reduce(function(x, acc) {
    return acc && p(x);
  }, true, l);
};

var any = function(p, l) {
  return reduce(function(x, acc) {
    return acc || p(x);
  }, false, l);
};

var zip = function(xs, ys) {
  return map2(function(x, y) { return [x, y]},
              xs,
              ys)
};

var filter = function(fn, ar) {
  var helper = function(i, j) {
    var n = j - i + 1;
    if (n == 0) {
      return [];
    } else if (n == 1) {
      return (fn(ar[i]) ? [ar[i]] : []);
    } else {
      var n1 = Math.ceil(n / 2);
      return helper(i, i + n1 - 1).concat(helper(i + n1, j));
    }
  }

  return helper(0, ar.length - 1)
};

var find = function(f, ar) {
  var n = ar.length;
  var helper = function(i) {
    if (i === n) {
      return undefined;
    } else if (f(ar[i])) {
      return ar[i];
    } else {
      return helper(i + 1);
    }
  }
  return helper(0);
};

var remove = function(a, ar) {
  return filter(function(e) {
    return a != e;
  }, ar);
};


var minWith = function(f, ar) {
  var fn = function(_ar, _best) {
    if (_ar.length === 0) {
      return _best;
    } else if (_ar[0][1] < _best[1]) {
      return fn(_ar.slice(1), _ar[0]);
    } else {
      return fn(_ar.slice(1), _best);
    }
  };
  return fn(zip(ar, map(f, ar)), [Infinity, Infinity]);
};

var maxWith = function(f, ar) {
  var fn = function(_ar, _best) {
    if (_ar.length === 0) {
      return _best;
    } else if (_ar[0][1] > _best[1]) {
      return fn(_ar.slice(1), _ar[0]);
    } else {
      return fn(_ar.slice(1), _best);
    }
  };
  return fn(zip(ar, map(f, ar)), [-Infinity, -Infinity]);
};

// bin array into items satisfying a predicate p and items not satifying it
var span = function(p, ar) {
  var n = ar.length;
  var helper = function(i, _ts, _fs) {
    return (i === n ?
            [_ts, _fs] :
            (p(ar[i]) ?
             helper(i + 1, snoc(_ts, ar[i]), _fs) :
             helper(i + 1, _ts, snoc(_fs, ar[i]))));
  };
  return helper(0, [], []);
};

// group array items by a comparator
// NB: there is still room for optimization here
var groupBy = function(cmp, ar) {
  if (ar.length === 0) {
    return [];
  } else {
    var x = ar[0];
    var sp = span(function(b) { return cmp(x, b); }, ar.slice(1));
    return [cons(x, sp[0])].concat(groupBy(cmp, sp[1]));
  }
};

var repeat = function(n, fn) {
  var helper = function(m) {
    if (m == 0) {
      return [];
    } else if (m == 1) {
      return [fn()];
    } else {
      var m1 = Math.ceil(m / 2),
          m2 = m - m1;
      return helper(m1).concat(helper(m2));
    }
  }

  return helper(n);
}

var compose = function(f, g) {
  return function(x) {
    return f(g(x));
  };
};

var everyOther = function(l) {
  return l.length <= 1 ? l : [l[0]].concat(everyOther(l.slice(2)));
};

var _merge = function(l1, l2, pred, key) {
  return (l1.length === 0 ?
          l2 :
          (l2.length === 0 ?
           l1 :
           (pred(key(l1[0]), key(l2[0])) ?
            [l1[0]].concat(_merge(l1.slice(1), l2, pred, key)) :
            [l2[0]].concat(_merge(l1, l2.slice(1), pred, key)))));
};

var _sort = function(l, pred, key) {
  return ((l.length <= 1) ?
          l :
          _merge(_sort(everyOther(l), pred, key),
                 _sort(everyOther(l.slice(1)), pred, key),
                 pred,
                 key));
};

var lt = function(a, b) {
  return a < b;
};

var gt = function(a, b) {
  return a > b;
};

var sort = function(l, pred, key) {
  return _sort(l, (pred || lt), (key || idF));
};

var sortOn = function(l, key, pred) {
  return _sort(l, (pred || lt), key);
};

var condition = function(bool) {
  factor(bool ? 0 : -Infinity);
};

var MH = function(wpplFn, samples, burn) {
  return MCMC(wpplFn, { samples: samples, burn: burn });
};

var ParticleFilter = function(wpplFn, particles) {
  return SMC(wpplFn, { particles: particles, rejuvSteps: 0 });
};

var ParticleFilterRejuv = function(wpplFn, particles, rejuvSteps) {
  return SMC(wpplFn, { particles: particles, rejuvSteps: rejuvSteps });
};

var SampleGuide = function(wpplFn, options) {
  return ForwardSample(wpplFn, _.extendOwn({guide: true}, _.omit(options, 'guide')));
};

var OptimizeThenSample = function(wpplFn, options) {
  var params = Optimize(wpplFn, _.omit(options, 'samples'));
  var opts = _.extendOwn({params: params}, _.pick(options, 'samples', 'verbose'));
  return SampleGuide(wpplFn, opts);
};

var Infer = function(options, wpplFn) {
  if (!util.isObject(options)) {
    util.fatal('Infer: expected first argument to be an options object.');
  }
  if (!_.isFunction(wpplFn)) {
    util.fatal('Infer: expected second argument to be a function.');
  }

  // Map from camelCase options to PascalCase coroutine names. Also
  // used to ensure the supplied method name is a valid inference
  // routine.
  var methodMap = {
    SMC: SMC,
    MCMC: MCMC,
    PMCMC: PMCMC,
    asyncPF: AsyncPF,
    rejection: Rejection,
    enumerate: Enumerate,
    incrementalMH: IncrementalMH,
    forward: ForwardSample,
    optimize: OptimizeThenSample
  };

  var methodName = options.method;
  if (methodName === undefined) {
    util.fatal('Infer: the \'method\' option must be specified.');
  }
  if (!_.has(methodMap, methodName)) {
    util.fatal('Infer: unknown method \'' + methodName + '\'.');
  }
  var method = methodMap[methodName];
  return method(wpplFn, _.omit(options, 'method'));
};

// Convenience function to create scalar valued parameters.
var scalarParam = function(mean, sd) {
  return ad.tensorEntry(tensorParam([1], mean, sd), 0);
};

var param = function(a1, a2, a3) {
  return _.isArray(a1) ? tensorParam(a1, a2, a3) : scalarParam(a1, a2);
};