Merge pull request #773 from null-a/only-map-fwd-sample

Add onlyMAP option to forward sampling method

docs/inference/methods.rst

@@ -329,6 +329,13 @@ Optimization
 
          Default: ``1``
 
+      .. describe:: onlyMAP
+
+         When ``true``, only the sample with the highest score is
+         retained. The marginal is a delta distribution on this value.
+
+         Default: ``false``
+
    In addition, all of the options supported by :ref:`Optimize
    <optimize>` are also supported here.
 
@@ -368,6 +375,13 @@ Forward Sampling
 
       Default: ``false``
 
+   .. describe:: onlyMAP
+
+      When ``true``, only the sample with the highest score is
+      retained. The marginal is a delta distribution on this value.
+
+      Default: ``false``
+
    Example usage::
 
      Infer({method: 'forward', model: model});

src/header.wppl

@@ -474,7 +474,7 @@ var SampleGuide = function(wpplFn, options) {
 
 var OptimizeThenSample = function(wpplFn, options) {
   Optimize(wpplFn, _.omit(options, 'samples'));
-  var opts = _.pick(options, 'samples', 'verbose');
+  var opts = _.pick(options, 'samples', 'onlyMAP', 'verbose');
   return SampleGuide(wpplFn, opts);
 };
 

src/inference/forwardSample.js

@@ -15,6 +15,7 @@ module.exports = function(env) {
     this.opts = util.mergeDefaults(options, {
       samples: 1,
       guide: false, // true = sample guide, false = sample target
+      onlyMAP: false,
       verbose: false
     });
 
@@ -34,18 +35,19 @@ module.exports = function(env) {
 
     run: function() {
 
-      var hist = new CountAggregator();
+      var hist = new CountAggregator(this.opts.onlyMAP);
       var logWeights = [];   // Save total factor weights
 
       return util.cpsLoop(
           this.opts.samples,
 
           // Loop body.
           function(i, next) {
+            this.score = 0;
             this.logWeight = 0;
             return this.wpplFn(_.clone(this.s), function(s, val) {
               logWeights.push(this.logWeight);
-              hist.add(val);
+              hist.add(val, this.score);
               return next();
             }.bind(this), this.a);
           }.bind(this),
@@ -64,16 +66,21 @@ module.exports = function(env) {
     },
 
     sample: function(s, k, a, dist, options) {
+      var cont = function(s, dist) {
+        var val = dist.sample();
+        this.score += dist.score(val);
+        return k(s, val);
+      }.bind(this);
+
       if (this.opts.guide) {
         options = options || {};
         return guide.getDist(
             options.guide, options.noAutoGuide, dist, env, s, a,
             function(s, maybeGuideDist) {
-              var d = maybeGuideDist || dist;
-              return k(s, d.sample());
+              return cont(s, maybeGuideDist || dist);
             });
       } else {
-        return k(s, dist.sample());
+        return cont(s, dist);
       }
     },
 

tests/test-data/stochastic/expected/onlyMAP.json

@@ -1,3 +1,4 @@
 {
-  "mean": 1
+  "mean": 1,
+  "std": 0
 }

tests/test-inference.js

@@ -49,7 +49,12 @@ var tests = [
       mixed4: true,
       bivariateGaussian: true,
       indirectDependency: true,
-      guidedFlip: true
+      guidedFlip: true,
+      onlyMAP: {
+        mean: { tol: 0.1 },
+        std: { tol: 0 },
+        args: { samples: 150, onlyMAP: true }
+      }
     }
   },
   {
@@ -112,7 +117,11 @@ var tests = [
       variableSupport: true,
       query: true,
       query2: { hist: { tol: 0.1, exactSupport: true } },
-      onlyMAP: { mean: { tol: 0.1 }, args: { samples: 150, onlyMAP: true } },
+      onlyMAP: {
+        mean: { tol: 0.1 },
+        std: { tol: 0 },
+        args: { samples: 150, onlyMAP: true }
+      },
       nestedEnum1: { mean: { tol: 0.1 }, std: { tol: 0.075 } },
       nestedEnum2: { mean: { tol: 0.1 }, std: { tol: 0.075 } },
       nestedEnum3: { mean: { tol: 0.1 }, std: { tol: 0.075 } },
@@ -222,7 +231,11 @@ var tests = [
       store: { hist: { tol: 0 }, args: { particles: 100 } },
       store2: { hist: { tol: 0 }, args: { particles: 100 } },
       notapes: { hist: { tol: 0 }, args: { samples: 100 } },
-      onlyMAP: { mean: { tol: 0.1 }, args: { particles: 150, onlyMAP: true } },
+      onlyMAP: {
+        mean: { tol: 0.1 },
+        std: { tol: 0 },
+        args: { particles: 150, onlyMAP: true }
+      },
       gaussianMean: { mean: { tol: 0.3 }, std: { tol: 0.3 }, args: { particles: 10000 } },
       varFactors1: { args: { particles: 5000 } },
       varFactors2: true,
@@ -372,7 +385,11 @@ var tests = [
       variableSupport: true,
       query: true,
       query2: { hist: { tol: 0.1, exactSupport: true } },
-      onlyMAP: { mean: { tol: 0.1 }, args: { samples: 150, onlyMAP: true } },
+      onlyMAP: {
+        mean: { tol: 0.1 },
+        std: { tol: 0 },
+        args: { samples: 150, onlyMAP: true }
+      },
       nestedEnum1: { mean: { tol: 0.1 }, std: { tol: 0.075 } },
       nestedEnum2: { mean: { tol: 0.1 }, std: { tol: 0.075 } },
       nestedEnum3: { mean: { tol: 0.1 }, std: { tol: 0.075 } },
@@ -408,7 +425,11 @@ var tests = [
       variableSupport: true,
       query: true,
       query2: { hist: { tol: 0.1, exactSupport: true } },
-      onlyMAP: { mean: { tol: 0.1 }, args: { samples: 150, kernel: 'HMC', onlyMAP: true } },
+      onlyMAP: {
+        mean: { tol: 0.1 },
+        std: { tol: 0 },
+        args: { samples: 150, kernel: 'HMC', onlyMAP: true }
+      },
       mixed1: true,
       mixed1Factor: true,
       mixed2: {
@@ -631,7 +652,17 @@ var tests = [
           verbose: false
         }
       },
-      mapData: true
+      mapData: true,
+      onlyMAP: {
+        mean: { tol: 0.1 },
+        std: { tol: 0 },
+        args: {
+          steps: 10000,
+          samples: 150,
+          onlyMAP: true,
+          verbose: false
+        }
+      }
     }
   }
 ];