Merge pull request #876 from null-a/numeric

Pull more fns into numeric module.

src/aggregation/ScoreAggregator.ad.js

@@ -5,18 +5,7 @@ var assert = require('assert');
 var _ = require('lodash');
 var dists = require('../dists');
 var util = require('../util');
-
-function logaddexp(a, b) {
-  if (a === -Infinity) {
-    return b;
-  } else if (b === -Infinity) {
-    return a;
-  } else if (a > b) {
-    return Math.log(1 + Math.exp(b - a)) + a;
-  } else {
-    return Math.log(1 + Math.exp(a - b)) + b;
-  }
-}
+var numeric = require('../math/numeric');
 
 var ScoreAggregator = function() {
   this.dist = {};
@@ -34,13 +23,13 @@ ScoreAggregator.prototype.add = function(value, score) {
   if (this.dist[key] === undefined) {
     this.dist[key] = { score: -Infinity, val: value };
   }
-  this.dist[key].score = logaddexp(this.dist[key].score, score);
+  this.dist[key].score = numeric.logaddexp(this.dist[key].score, score);
 };
 
 function normalize(dist) {
   // Note, this also maps dist from log space into probability space.
   var logNorm = _.reduce(dist, function(acc, obj) {
-    return logaddexp(acc, obj.score);
+    return numeric.logaddexp(acc, obj.score);
   }, -Infinity);
   return _.mapValues(dist, function(obj) {
     return { val: obj.val, prob: Math.exp(obj.score - logNorm) };

src/dists/discrete.ad.js

@@ -9,7 +9,7 @@ var numeric = require('../math/numeric');
 var T = ad.tensor;
 
 function sample(theta) {
-  var thetaSum = util.sum(theta);
+  var thetaSum = numeric._sum(theta);
   var x = util.random() * thetaSum;
   var k = theta.length;
   var probAccum = 0;

src/dists/kde.ad.js

@@ -6,6 +6,7 @@ var types = require('../types');
 var util = require('../util');
 var Tensor = require('../tensor');
 var stats = require('../math/statistics');
+var numeric = require('../math/numeric');
 var gaussian = require('./gaussian');
 var diagCovGaussian = require('./diagCovGaussian');
 
@@ -102,7 +103,7 @@ var KDE = base.makeDistributionType({
     var kernel = this.kernel;
     return data.reduce(
       function(acc, x) {
-        return util.logaddexp(acc, kernel.score(x, width, val));
+        return numeric.logaddexp(acc, kernel.score(x, width, val));
       },
       -Infinity) - Math.log(n);
   }

src/dists/multinomial.ad.js

@@ -16,7 +16,7 @@ function zeros(n) {
 }
 
 function sample(theta, n) {
-  // var thetaSum = util.sum(theta);
+  // var thetaSum = numeric._sum(theta);
   var a = zeros(theta.length);
   for (var i = 0; i < n; i++) {
     a[discrete.sample(theta)]++;

src/inference/asyncpf.js

@@ -8,6 +8,7 @@
 
 var _ = require('lodash');
 var util = require('../util');
+var numeric = require('../math/numeric');
 var CountAggregator = require('../aggregation/CountAggregator');
 
 module.exports = function(env) {
@@ -114,8 +115,8 @@ module.exports = function(env) {
       var currWeight = this.activeParticle.weight;
       var denom = lk.length + currMultiplicity; // k - 1 + Ckn
       var prevWBar = lk[lk.length - 1].wbar;
-      var wbar = -Math.log(denom) + util.logsumexp([Math.log(lk.length) + prevWBar,
-                                                    Math.log(currMultiplicity) + currWeight]);
+      var wbar = -Math.log(denom) + numeric._logsumexp([Math.log(lk.length) + prevWBar,
+                                                        Math.log(currMultiplicity) + currWeight]);
       if (wbar > 0) throw new Error('Positive weight!!'); // sanity check
       var logRatio = currWeight - wbar;
       var numChildrenAndWeight = [];

src/inference/forwardSample.js

@@ -5,6 +5,7 @@
 
 var _ = require('lodash');
 var util = require('../util');
+var numeric = require('../math/numeric');
 var CountAggregator = require('../aggregation/CountAggregator');
 var ad = require('../ad');
 var guide = require('../guide');
@@ -104,7 +105,7 @@ module.exports = function(env) {
         function() {
           var dist = hist.toDist();
           if (!opts.guide) {
-            dist.normalizationConstant = util.logsumexp(logWeights) - Math.log(opts.samples);
+            dist.normalizationConstant = numeric._logsumexp(logWeights) - Math.log(opts.samples);
           }
           return k(s, dist);
         }

src/inference/mhkernel.js

@@ -3,6 +3,7 @@
 var _ = require('lodash');
 var assert = require('assert');
 var util = require('../util');
+var numeric = require('../math/numeric');
 var ad = require('../ad');
 
 module.exports = function(env) {
@@ -215,7 +216,7 @@ module.exports = function(env) {
     var score = ad.value(proposalDist.score(regenChoice.val));
 
     // Rest of the trace.
-    score += util.sum(toTrace.choices.slice(this.regenFrom + 1).map(function(choice) {
+    score += numeric._sum(toTrace.choices.slice(this.regenFrom + 1).map(function(choice) {
       return this.reused.hasOwnProperty(choice.address) ? 0 : ad.value(choice.dist.score(choice.val));
     }, this));
 

src/inference/smc.js

@@ -2,6 +2,7 @@
 
 var _ = require('lodash');
 var util = require('../util');
+var numeric = require('../math/numeric');
 var discrete = require('../dists/discrete');
 var Trace = require('../trace');
 
@@ -154,7 +155,7 @@ module.exports = function(env) {
     }
     // Residual resampling following Liu 2008; p. 72, section 3.4.4
     var m = particles.length;
-    var logW = util.logsumexp(_.map(particles, 'logWeight'));
+    var logW = numeric._logsumexp(_.map(particles, 'logWeight'));
     var logAvgW = logW - Math.log(m);
     if (logAvgW === -Infinity) {
       // do not return, execution continues

src/math/numeric.ad.js

@@ -1,15 +1,40 @@
 'use strict';
 
+var _ = require('lodash');
 var ad = require('../ad');
 
 var LOG_PI = 1.1447298858494002;
 var LOG_2PI = 1.8378770664093453;
 
+// By convention, non-adified versions of functions are prefixed with
+// an underscore.
+
 function sum(xs) {
   'use ad';
   return xs.reduce(function(a, b) { return a + b; }, 0);
 }
 
+function _sum(xs) {
+  if (xs.length === 0) {
+    return 0.0;
+  } else {
+    var total = _.reduce(xs,
+        function(a, b) {
+          return a + b;
+        });
+    return total;
+  }
+}
+
+function product(xs) {
+  'use ad';
+  var result = 1;
+  for (var i = 0, n = xs.length; i < n; i++) {
+    result *= xs[i];
+  }
+  return result;
+}
+
 function fact(x) {
   'use ad';
   var t = 1;
@@ -47,11 +72,37 @@ function squishToProbSimplex(x) {
   return ad.tensor.softmax(u);
 }
 
+function logaddexp(a, b) {
+  'use ad';
+  if (a === -Infinity) {
+    return b;
+  } else if (b === -Infinity) {
+    return a;
+  } else if (a > b) {
+    return Math.log(1 + Math.exp(b - a)) + a;
+  } else {
+    return Math.log(1 + Math.exp(a - b)) + b;
+  }
+}
+
+function _logsumexp(a) {
+  var m = Math.max.apply(null, a);
+  var sum = 0;
+  for (var i = 0; i < a.length; ++i) {
+    sum += (a[i] === -Infinity ? 0 : Math.exp(a[i] - m));
+  }
+  return m + Math.log(sum);
+}
+
 module.exports = {
   LOG_PI: LOG_PI,
   LOG_2PI: LOG_2PI,
   sum: sum,
+  _sum: _sum,
+  product: product,
   fact: fact,
   lnfact: lnfact,
-  squishToProbSimplex: squishToProbSimplex
+  squishToProbSimplex: squishToProbSimplex,
+  logaddexp: logaddexp,
+  _logsumexp: _logsumexp
 };

src/transforms/adify.js

@@ -8,7 +8,6 @@ var generate = require('escodegen').generate;
 var build = require('ast-types').builders;
 var _ = require('lodash');
 var ad = require('./ad').ad;
-var util = require('../util');
 
 function isMarkedForGlobalTransform(ast) {
   assert.ok(ast.type === 'Program');

src/types.js

@@ -8,6 +8,7 @@
 
 var _ = require('lodash');
 var util = require('./util');
+var numeric = require('./math/numeric');
 var interval = require('./math/interval');
 
 var isInterval = interval.isInterval;
@@ -136,7 +137,7 @@ var probabilityArray = function() {
     name: 'probabilityArray',
     desc: 'real array with elements that sum to one',
     check: function(val) {
-      return baseType.check(val) && Math.abs(1 - util.sum(val)) < tol;
+      return baseType.check(val) && Math.abs(1 - numeric._sum(val)) < tol;
     }
   };
 };

src/util.js

@@ -5,9 +5,13 @@ var assert = require('assert');
 var seedrandom = require('seedrandom');
 var ad = require('./ad');
 var Tensor = require('./tensor');
+var numeric = require('./math/numeric');
 
 var rng = Math.random;
 
+// Re-export sum from this module, as expected by webppl-viz.
+var sum = numeric._sum;
+
 var trampolineRunners = {
   web: function(yieldEvery) {
     yieldEvery = yieldEvery || 100;
@@ -83,47 +87,6 @@ function asArray(arg) {
   return arg ? [].concat(arg) : [];
 }
 
-function sum(xs) {
-  if (xs.length === 0) {
-    return 0.0;
-  } else {
-    var total = _.reduce(xs,
-        function(a, b) {
-          return a + b;
-        });
-    return total;
-  }
-}
-
-function product(xs) {
-  var result = 1;
-  for (var i = 0, n = xs.length; i < n; i++) {
-    result *= xs[i];
-  }
-  return result;
-}
-
-function logsumexp(a) {
-  var m = Math.max.apply(null, a);
-  var sum = 0;
-  for (var i = 0; i < a.length; ++i) {
-    sum += (a[i] === -Infinity ? 0 : Math.exp(a[i] - m));
-  }
-  return m + Math.log(sum);
-}
-
-function logaddexp(a, b) {
-  if (a === -Infinity) {
-    return b;
-  } else if (b === -Infinity) {
-    return a;
-  } else if (a > b) {
-    return Math.log(1 + Math.exp(b - a)) + a;
-  } else {
-    return Math.log(1 + Math.exp(a - b)) + b;
-  }
-}
-
 var deleteIndex = function(arr, i) {
   return arr.slice(0, i).concat(arr.slice(i + 1))
 }
@@ -396,16 +359,13 @@ module.exports = {
   histStd: histStd,
   histsApproximatelyEqual: histsApproximatelyEqual,
   gensym: gensym,
-  logsumexp: logsumexp,
-  logaddexp: logaddexp,
   deleteIndex: deleteIndex,
   makeGensym: makeGensym,
   prettyJSON: prettyJSON,
   runningInBrowser: runningInBrowser,
   mergeDefaults: mergeDefaults,
   getValAndOpts: getValAndOpts,
   sum: sum,
-  product: product,
   asArray: asArray,
   serialize: serialize,
   deserialize: deserialize,

tests/test-data/sampler/beta.js

@@ -3,7 +3,7 @@
 var _ = require('lodash');
 var assert = require('assert');
 var beta = require('../../../src/dists/beta');
-var util = require('../../../src/util');
+var numeric = require('../../../src/math/numeric');
 var statistics = require('../../../src/math/statistics');
 
 var ln = Math.log,
@@ -42,7 +42,7 @@ module.exports = {
     var a = params[0];
     var b = params[1];
     // https://en.wikipedia.org/wiki/Beta_distribution#Higher_moments
-    return util.product(_.range(0, n - 1).map(function(k) { return (a + k) / (a + b + k) }))
+    return numeric.product(_.range(0, n - 1).map(function(k) { return (a + k) / (a + b + k) }))
   },
   // mostly HT https://en.wikipedia.org/wiki/Gamma_distribution
   populationStatisticFunctions: {

tests/test-data/sampler/gamma.js

@@ -3,7 +3,7 @@
 var _ = require('lodash');
 var assert = require('assert');
 var gamma = require('../../../src/dists/gamma');
-var util = require('../../../src/util');
+var numeric = require('../../../src/math/numeric');
 var statistics = require('../../../src/math/statistics');
 
 var ln = Math.log,
@@ -41,7 +41,7 @@ module.exports = {
     // HT http://ocw.mit.edu/courses/mathematics/
     // 18-443-statistics-for-applications-fall-2006/lecture-notes/lecture6.pdf
     // (but NB: they use shape, rate whereas we have shape, scale)
-    return util.product(_.range(0, n - 1).map(function(k) { return shape + k })) * pow(scale, n)
+    return numeric.product(_.range(0, n - 1).map(function(k) { return shape + k })) * pow(scale, n)
   },
   // mostly HT https://en.wikipedia.org/wiki/Gamma_distribution
   populationStatisticFunctions: {