'use strict';

MCMC.registerAlgorithm('SVGD', {

  description: 'Stein Variational Gradient Descent',

  about: function() {
    window.open('http://www.cs.dartmouth.edu/~dartml/project.html?p=vgd');
  },

  init: function(self) {
    self.chain = [ ];
    self.n = 125; // number of particlese
    self.epsilon = 0.25; // step size
    self.h = 0.15; // bandwidth
    self.reset(self);
  },

  reset: function(self) {
    // initialize chain with samples from standard normal
    self.chain = [ ];
    for (var i = 0; i < self.n; i++) {
      self.chain.push(MultivariateNormal.getSample(self.dim));
    }
  },

  attachUI: function(self, folder) {
    folder.add(self, 'h', 0.05, 1).step(0.05).name('bandwidth');
    folder.add(self, 'epsilon', 0.05, 0.5).step(0.05).name('stepsize');
    folder.add(self, 'n', 10, 400).step(1).name('numParticles').onChange(function(value) {
      self.n = parseInt(value);
      if (self.n > self.chain.length) {
        for (var i = 0; i < self.chain.length - self.n; i++) {
          self.chain.push(MultivariateNormal.getSample(self.dim));
        }
      } else if (self.n < self.chain.length) {
        self.chain = self.chain.slice(0, self.n);
      }
      console.log(self.n, self.chain.length);
    });
    folder.open();
  },

  step: function(self, visualizer) {

    var gradx = [];

    var x = self.chain;
    var h = self.h;

    // precompute log densities
    var gradLogDensities = [ ];
    x.forEach(function(x_i, i) {
      gradLogDensities.push(self.gradLogDensity(x_i));
    });

    // compute gradient of x as empirical average of gradx_i for i = 1...n
    x.forEach(function(x_i, i) {
      var gradx_i = Float64Array.zeros(self.dim);
      x.forEach(function(x_j, j) {
        // compute rbf and gradient
        var rbf = Math.exp(-x_i.dist2(x_j) / (2 * h * h));
        var grad_rbf = x_i.subtract(x_j).scale(2 * rbf / (h * h));
        gradx_i.increment(gradLogDensities[j].scale(rbf).add(grad_rbf));
      });
      gradx.push(gradx_i.scale(1 / self.n)); // average over x_j
    });

    visualizer.queue.push({ type: 'svgd-step', x: x, gradx: gradx, h: h});

    // update particles
    for (var i = 0; i < self.n; i++) {
      x[i].increment(gradx[i].scale(self.epsilon));
    }

  }

});

@@ -847,70 +847,130 @@
 es')
-.onChange(function(value
+;%0A    folder.open();%0A  %7D,%0A%0A  step: function(self, visualizer
 ) %7B%0A
- 
+%0A
     
- self.n = parseInt(value);%0A  
+var gradx = %5B%5D;%0A%0A    // resize samples appropriately%0A
     
@@ -1009,18 +1009,16 @@
 %7B%0A      
-  
 for (var
@@ -1024,24 +1024,33 @@
 r i = 0; i %3C
+ self.n -
  self.chain.
@@ -1059,17 +1059,8 @@
 ngth
- - self.n
 ; i+
@@ -1060,26 +1060,24 @@
 gth; i++) %7B%0A
-  
         self
@@ -1131,30 +1131,26 @@
 im));%0A      
-  %7D%0A  
+%7D%0A
     %7D else i
@@ -1180,34 +1180,32 @@
 length) %7B%0A      
-  
 self.chain = sel
@@ -1238,149 +1238,14 @@
     
-  
 %7D%0A    
-  console.log(self.n, self.chain.length);%0A    %7D);%0A    folder.open();%0A  %7D,%0A%0A  step: function(self, visualizer) %7B%0A%0A    var gradx = %5B%5D;%0A
 %0A   
@@ -2080,33 +2080,44 @@
 i = 0; i %3C self.
-n
+chain.length
 ; i++) %7B%0A      x

Update SVGD.js