All examples By author By category About

rpgove

Full window Github gist

Speeding up KDE heatmaps with Barnes-Hut quadtrees

Performance comparison of a brute force kernel density estimation (KDE) heatmap vs. a Barnes-Hut quadtree KDE heatmap.

This example creates an array of node positions at every tick of a force-directed algorithm. Then it generates two heatmaps from all of the collected node positions.

The Barnes-Hut approximation is markedly faster and looks the same when there are many points. When there is a very small number of points, the extra cost of constructing the quadtree makes the Barnes-Hut heatmap slower. But after just a few ticks, there are enough points that the Barnes-Hut heatmap is much, much faster.

In data visualization, the Barnes-Hut approximation is often used for speeding up force-directed graph layouts, but it has many other potential uses, such as this example here.

See also:

<!DOCTYPE html>
<meta charset="utf-8">
<style>

body {
  background: #fff;
  font: 10px sans-serif;
  margin: 0;
}

.container {
  width: 945px;
  height: 315px;
}

.container div {
  display: inline-block;
  width: 315px;
  height: 315px;
  overflow: hidden;
}

.force svg {
  background: rgba(0,0,0,0);
  position: absolute;
  top: 0;
  left: 0;
}

p {
  color: black;
}

.container div p {
  font-size: 12px;
  font-weight: bold;
  text-align: center;
  margin: 0;
  position: relative;
  bottom: 25px;
}

.brute p {
  color: steelblue;
}

.barneshut p {
  color: maroon;
}

svg {
  background: #fff;
}

.node {
  fill: #555;
  stroke: #fff;
  stroke-width: 1px;
}

.link {
  stroke: #555;
  stroke-opacity: 0.6;
}

.point {
  fill: black;
}

</style>
<body>
<div class="container"><div class="force"><canvas></canvas><svg></svg><p>Graph layout node positions</p></div><div class="brute"><canvas></canvas><p>Brute force heatmap</p></div><div class="barneshut"><canvas></canvas><p>Barnes-Hut heatmap</p></div></div>
<script src="https://d3js.org/d3.v4.min.js"></script>
<script>

var kdeWidth = 315;
var kdeHeight = 315;

var chartMargin = {top: 10, right: 20, bottom: 30, left: 50};
var chartWidth = 945 - chartMargin.left - chartMargin.right;
var chartHeight = 175 - chartMargin.top - chartMargin.bottom;

var computeTimes = [
  {name: 'Brute force', values: []},
  {name: 'Barnes-Hut', values: []}
];

var gridSize = 9;

var xyPoints = [];
var qt = d3.quadtree()
  .extent([[0, 0], [kdeWidth, kdeHeight]])
  .visitAfter(accumulate);
var h, h2, iqr, xPoints;

// Array of grid cell points. Each point is the center of the grid cell.
var grid = d3.merge(d3.range(0, kdeHeight/gridSize).map(function(i) {
  return d3.range(0, kdeWidth/gridSize).map(function(j) { return [j*gridSize + gridSize/2, i*gridSize + gridSize/2] });
}));

var x = d3.scaleLinear().rangeRound([0, chartWidth]);
var y = d3.scaleLinear().rangeRound([chartHeight, 0]);

var outerScale = d3.scalePow()
    .exponent(0.4)
    .domain([0,1])
    .range([0,1]);
var heatmapColor = d3.scaleLinear()
    .clamp(true)
    .domain([0, 0.1111111111111111, 0.2222222222222222, 0.3333333333333333, 0.4444444444444444, 0.5555555555555555, 0.6666666666666666, 0.7777777777777777, 0.8888888888888888, 1])
    .range(["#fff", "#dadada", "#b9b9b9", "#9b9b9b", "#7f7e7e", "#646363", "#4a4a4a", "#313030", "#161616", "#000"]);

var force = d3.forceSimulation()
  .alphaDecay(0.03)
  .force('link', d3.forceLink().id(function(d, i) { return i; }).distance(10))
  .force('charge', d3.forceManyBody().strength(-40))
  .force('x', d3.forceX(kdeWidth/2))
  .force('y', d3.forceY(kdeHeight/2))
  .force('center', d3.forceCenter(kdeWidth/2, kdeHeight/2));

var tickCount = 0;

var forceSvg = d3.select(".force svg")
    .attr("width", kdeWidth)
    .attr("height", kdeHeight)
    .append("g");

var pointsContext = d3.select('.force canvas')
    .attr("width", kdeWidth)
    .attr("height", kdeHeight)
    .node()
    .getContext('2d');

var bruteCanvas = d3.select(".brute canvas")
    .attr('id', 'brute')
    .attr("width", kdeWidth)
    .attr("height", kdeHeight);

var barneshutCanvas = d3.select(".barneshut canvas")
    .attr('id', 'barneshut')
    .attr("width", kdeWidth)
    .attr("height", kdeHeight);

var chartSvg = d3.select("body").append("svg")
    .attr("width", chartWidth + chartMargin.left + chartMargin.right)
    .attr("height", chartHeight + chartMargin.top + chartMargin.bottom)
    .append('g')
    .attr("transform", "translate(" + chartMargin.left + "," + chartMargin.top + ")");

var line = d3.line()
    .x(function (d) { return x(d.numPoints); })
    .y(function (d) { return y(d.time); });

var bruteContext = bruteCanvas.node().getContext("2d");
var barneshutContext = barneshutCanvas.node().getContext("2d");

d3.json("miserables.json", function(error, graph) {
  if (error) throw error;

var link = forceSvg.selectAll(".link")
      .data(graph.links)
    .enter().append("line")
      .attr("class", "link")
      .style("stroke-width", function(d) { return Math.sqrt(d.value); });

var node = forceSvg.selectAll(".node")
      .data(graph.nodes)
    .enter().append("circle")
      .attr("class", "node")
      .attr("r", 3);

point = forceSvg.selectAll('.point')
      .data(xyPoints)
    .enter().append('circle')
      .attr('class', 'point');

node.append("title")
      .text(function(d) { return d.name; });

force.nodes(graph.nodes)
    .on("end", function(event) {
      computeDensities();
      updateChart();
    })
    .on('tick', function(event) {
      link.attr("x1", function(d) { d.source.x = boundedCoords(d.source.x, kdeWidth); return d.source.x; })
          .attr("y1", function(d) { d.source.y = boundedCoords(d.source.y, kdeHeight); return d.source.y; })
          .attr("x2", function(d) { d.target.x = boundedCoords(d.target.x, kdeWidth); return d.target.x; })
          .attr("y2", function(d) { d.target.y = boundedCoords(d.target.y, kdeHeight); return d.target.y; });

node.attr("cx", function(d) { d.x = boundedCoords(d.x, kdeWidth); return d.x; })
          .attr("cy", function(d) { d.y = boundedCoords(d.y, kdeHeight); return d.y; });

// Track new node positions
      node.each(function(d) {
        xyPoints.push([d.x,d.y]);

pointsContext.fillStyle = '#c9c9c9';
        pointsContext.beginPath();
        pointsContext.arc(d.x, d.y, 3, 0, 2*Math.PI);
        pointsContext.fill();
      });

if (++tickCount == 50) {
        force.stop();
      }

computeDensities();
      updateChart();
    });

force.force('link')
    .links(graph.links);
});

function updateChart () {
  x.domain([0, xyPoints.length]);
  y.domain([0, d3.max(computeTimes, function (d) { return d3.max(d.values.map(function(v) { return v.time; })); })]);

chartSvg.selectAll('g').remove();

chartSvg.append('g')
      .attr('transform', 'translate(0,' + chartHeight + ')')
      .call(d3.axisBottom(x))
    .append("text")
      .attr("fill", "#000")
      .attr('transform', 'translate(' + chartMargin.left + ',' + 0 + ')')
      .attr("y", 19)
      .attr("dy", "0.71em")
      .attr("text-anchor", "begin")
      .text("Number of points");

chartSvg.append('g')
      .call(d3.axisLeft(y))
    .append("text")
      .attr("fill", "#000")
      .attr("transform", "rotate(-90)")
      .attr("y", 6)
      .attr("dy", "0.71em")
      .attr("text-anchor", "end")
      .text("Compute time (ms)");

var lineG = chartSvg.selectAll('g.line-g')
      .data(computeTimes)
    .enter().append('g')
      .attr('class', 'line-g');

lineG.append('path')
      .attr("fill", "none")
      .attr("stroke", function (d) { return d.name === 'Brute force' ? 'steelblue' : 'maroon'; })
      .attr("stroke-linejoin", "round")
      .attr("stroke-linecap", "round")
      .attr("stroke-width", 1.5)
      .attr("d", function (d) { return line(d.values); });

lineG.append("text")
      .datum(function(d) { return {name: d.name, value: d.values[d.values.length - 1]}; })
      .attr("transform", function(d) { return "translate(" + x(d.value.numPoints) + "," + y(d.value.time) + ")"; })
      .attr("x", 3)
      .attr('y', -6)
      .attr("dy", "0.35em")
      .attr('text-anchor', 'end')
      .style("font", "10px sans-serif")
      .text(function(d) { return d.name; });
}

function computeDensities () {
  // Density at each (x,y) coordinate in the grid
  var densities = [];
  var startTime;

// Update bandwidth
  // Use same bandwidth for each dimension
  xPoints = xyPoints.map(function(d) { return d[0] }).sort(function(a,b) { return a - b });
  iqr = d3.quantile(xPoints, 0.75) - d3.quantile(xPoints, 0.25);
  h = 1.06 * Math.min(d3.deviation(xPoints), iqr / 1.34) * Math.pow(xyPoints.length, -0.2);
  h2 = h*h;

// Compute KDE for each (x,y) pair and update the color scale
  startTime = new Date();
  densities = grid.map(function(point) { return kdeBrute(point); });
  computeTimes[0].values.push({numPoints: xyPoints.length, time: new Date() - startTime})
  outerScale.domain([0, d3.max(densities)]);
  drawHeatmap(bruteContext, densities);

startTime = new Date();
  qt = d3.quadtree(xyPoints)
    .extent([[0, 0], [kdeWidth, kdeHeight]])
    .visitAfter(accumulate);
  densities = grid.map(function(point) { return kdeBarneshut(point); });
  computeTimes[1].values.push({numPoints: xyPoints.length, time: new Date() - startTime})
  outerScale.domain([0, d3.max(densities)]);
  drawHeatmap(barneshutContext, densities);
}

function drawHeatmap(context, densities) {
  // Draw the grid
  grid.forEach(function(point, idx) {
    context.beginPath();
    context.fillStyle = heatmapColor(outerScale(densities[idx]));
    // Subtract to get the corner of the grid cell
    context.rect(point[0] - gridSize/2, point[1] - gridSize/2, gridSize, gridSize);
    context.fill();
    context.closePath();
  });
}

function boundedCoords(pos, maxPos) {
  return Math.min(maxPos, Math.max(0, pos));
}

function accumulate(quad) {
  // Sum of x/y coordinates
  var sx = 0;
  var sy = 0;

var count = 0, q;

// This quadrant has children.
  if (quad.length) {
    var i = -1;
    var c;
    while (++i < 4) {
      c = quad[i];
      if (c == null) continue;
      count += c.count;
      sx += c.sx;
      sy += c.sy;
    }
  }

// This is a leaf node.
  else {
    q = quad;
    do {
      count += 1;
      sx += q.data[0];
      sy += q.data[1];
    } while (q = q.next);
  }

quad.sx = sx;
  quad.sy = sy;
  quad.count = count;
}

// Use same bandwidth for each dimension
function kdeBrute(gridPoint) {
  return d3.mean(xyPoints, function(p) { return gaussian(norm(p, gridPoint) / h) }) / h2;
}

function kdeBarneshut(gridPoint) {
  //return d3.mean(xyPoints, function(p) { return gaussian(norm(p, gridPoint) / h) }) / h2;
  var estimate = 0;
  var theta2 = 0.64;

function recurse(quadNode, x1, y1, x2, y2) {
    // Average x/y coordinates in this quadrant.
    var p = [quadNode.sx/quadNode.count, quadNode.sy/quadNode.count];
    var dx = p[0] - gridPoint[0];
    var dy = p[1] - gridPoint[1];
    var dw = x2 - x1;
    var dist = dx * dx + dy * dy;

// Barnes-Hut: If the quadrant size is small relative to its
    // distance from the grid point, then aggregate the points in
    // that quadrant and treat them as a single large point (or many
    // individual points all in the same location).
    if (dw * dw / theta2 < dist && !(gridPoint[0] >= x1 && gridPoint[0] < x2 && gridPoint[1] >= y1 && gridPoint[1] < y2)) {
      estimate += quadNode.count * gaussian(norm(p, gridPoint) / h);
      return true;
    }

if (quadNode.point) {
      estimate += gaussian(norm(quadNode.point, gridPoint) / h);
    }

// Stop recursing if there are no more nodes.
    return !quadNode.count;
  }

qt.visit(recurse);

return estimate / xyPoints.length;
}

// Norm of 2D arrays/vectors
function norm(v1, v2) {
  return Math.sqrt((v1[0] - v2[0]) * (v1[0] - v2[0]) + (v1[1] - v2[1]) * (v1[1] - v2[1]));
}

function gaussian(x) {
  // sqrt(2 * PI) is approximately 2.5066
  return Math.exp(-x * x / 2) / 2.5066;
}

</script>

https://d3js.org/d3.v4.min.js