MCMC Sampling

Math
Statistics
Hexbin
Code 
chart = {
  const width = fullWidth,
        height = fullHeight;

  const window = d3.create("svg")
    .attr("width", "100%")
    .attr("height", "100%")
    .attr("viewBox", `0 0 ${fullWidth} ${fullHeight}`)
    .attr("preserveAspectRatio", "xMidYMid meet");

  const svg = window.append("g")
    .attr("class", "mcmc");

  svg.append("defs")
    .append("clipPath")
    .attr("id", "def-mcmc-clipPath")
      .append("rect")
      .attr("x", 0)
      .attr("y", 0)
      .attr("width", width)
      .attr("height", height);

  svg.append("rect")
    .attr("width", width)
    .attr("height", height)
    .style("fill", "white")
    .style("stroke", "none");

  svg.append("g").attr("clip-path", "url(#def-mcmc-clipPath)")
    .attr("class", "hexbins");

  svg.append("g").attr("clip-path", "url(#def-mcmc-clipPath)")
    .attr("class", "contours");

  svg.append("g").attr("clip-path", "url(#def-mcmc-clipPath)")
    .attr("class", "walkers")

  let xc, yc, logLikelihood, abortController;

  const projection = d3.geoIdentity();

  const numBins = 100;
  const x2grid = d3.scaleLinear().range([0, numBins - 1]),
        y2grid = d3.scaleLinear().range([0, numBins - 1]);

  const hexbin = d3.hexbin()
    .x(d => projection([x2grid(d[0]), y2grid(d[1])])[0])
    .y(d => projection([x2grid(d[0]), y2grid(d[1])])[1])
    .radius(5)
    .extent([[0, 0], [width, height]])

  const color = d3.scaleSequential(d3.interpolateYlGn);

  const data = {
    rosenbrock: {
      x0: -2,
      x1: 2,
      y0: -2,
      y1: 3,
      z: [-60, -30, -20, -10],
      logLikelihood: d =>
        -(100 * Math.pow(d[1] - Math.pow(d[0], 2), 2) + Math.pow(d[0] - 1, 2)),
    },
    booth: {
      x0: 1-4,
      x1: 1+4,
      y0: 3-4,
      y1: 3+4,
      z: [-20, -15, -10, -5],
      logLikelihood: d =>
        -(Math.pow(d[0] + 2 * d[1] - 7, 2) + Math.pow(2 * d[0] + d[1] - 5, 2)),
    }
  };

  function linspace(domain, n){
    const dx = (domain[1] - domain[0]) / (n - 1)
    return d3.range(n).map(i => domain[0] + i * dx);
  }

  function change(value) {
    abortController?.abort();

    xc = 0;
    yc = 1;

    const d = data[value];
    logLikelihood = d.logLikelihood;

    svg.selectAll(".contour")
      .remove();

    x2grid.domain([d.x0, d.x1]);
    y2grid.domain([d.y0, d.y1]);

    const values = new Array(numBins * numBins);
    const x = linspace(x2grid.domain(), numBins);
    const y = linspace(y2grid.domain(), numBins);
    for (let j = 0, k = 0; j < numBins; ++j) {
      for (let i = 0; i < numBins; ++i, ++k) {
        values[k] = logLikelihood([x[i], y[j]]);
      }
    }

    const contours = d3.contours()
      .size([numBins, numBins])
      .thresholds(d.z)
      (values);

    projection.fitSize([width, height], contours[0]);

    svg.select(".contours").selectAll(".contour")
      .data(contours.slice(1))
      .join("path")
        .attr("class", "contour")
        .attr("fill", "none")
        .attr("stroke", "#C0C0C0")
        .attr("d", d3.geoPath(projection));

    start();
  }

  function start() {
    svg.selectAll(".walker")
      .remove();

    svg.selectAll(".hexbin")
      .remove();

    const numWalkers = 100,
          numIterations = 10000,
          delay = 100;

    const initialPosition = d3.range(numWalkers)
      .map(() => [xc + (Math.random() - 0.5), yc + (Math.random() - 0.5)]);

    abortController = new AbortController();
    const sample = MCMC(logLikelihood, initialPosition, numIterations)
      .delay(delay)
      .callback(render)
      .signal(abortController.signal);
    sample(d3.interval).then((sampler) => {
      console.log(`a = ${sampler.getAcceptanceFraction().toFixed(2)}`);
    });
  }

  function render(result) {
    svg.select(".walkers").selectAll(".walker")
      .data(result.iteration)
      .join("circle")
        .attr("class", "walker")
        .attr("r", 1.5)
        .attr("fill", "black")
        .attr("stroke", "none")
        .attr("cx", d => projection([x2grid(d[0]), y2grid(d[1])])[0])
        .attr("cy", d => projection([x2grid(d[0]), y2grid(d[1])])[1]);

    const bins = hexbin(result.chain);
    color.domain([0, d3.max(bins, d => d.length)]);

    svg.select(".hexbins").selectAll(".hexbin")
      .data(bins)
      .join("path")
        .attr("class", "hexbin")
        .attr("d", hexbin.hexagon())
        .attr("transform", d => `translate(${d.x},${d.y})`)
        .attr("fill", d => color(d.length));
  }

  svg.on("click", function(event) {
    abortController.abort();

    const pointer = projection.invert(d3.pointer(event, this));
    xc = x2grid.invert(pointer[0]);
    yc = y2grid.invert(pointer[1]);

    start();
  });

  invalidation.then(() => {
    abortController.abort();
  });

  return Object.assign(window.node(), {
    update(values) {
      change(values.function);
    },
  });
}

Affine invariant MCMC sampling is a popular algorithm for sampling hard-to-sample probability density functions, since it works well in many cases without much fine-tuning, using an ensemble of walkers (shown as black points).

Resources

Tools