d3-spring-model/src/barnesHut.js

import constant from "./constant";
import jiggle from "./jiggle";
import {x, y} from "./xy";
import {quadtree} from "d3-quadtree";

/**
 * The refinement of the existing Barnes-Hut implementation in D3
 * to fit the use case of the project. Previously the algorithm stored
 * strength as internal node, now the random child is stored as internal
 * node and the force calculations are done between the node and that internal
 * object if they are sufficiently far away.
 * The check to see if the nodes are far away was also changed to the one described in original Barnes-Hut paper.
 * @return {force} calculated forces.
 */
export default function() {
  var nodes,
      node,
      alpha,
      distance = constant(300),
      theta = 0.5;

  /**
   * Constructs a quadtree at every iteration and apply the forces by visiting
   * each node in a tree.
   * @param  {number} _ - controls the stopping of the
   *     particle simulations.
   */
  function force(_) {
    var i, n = nodes.length, tree = quadtree(nodes, x, y).visitAfter(accumulate);
    for (alpha = _, i = 0; i < n; ++i) {
      node = nodes[i], tree.visit(apply);
    }
  }

  /**
   * Function used during the tree construction to fill out the nodes with
   * correct data. Internal nodes acquire the random child while the leaf
   * nodes accumulate forces from coincident quadrants.
   * @param  {quadrant} quad - node representing the quadrant in quadtree.
   */
  function accumulate(quad) {
    var q, d, children = [];

    // For internal nodes, accumulate forces from child quadrants.
    if (quad.length) {
      for (var i = 0; i < 4; ++i) {
        if ((q = quad[i]) && (d = q.data)) {
          children.push(d);
        }
      }
      // Choose a random child.
      quad.data = children[Math.floor(Math.random() * children.length)];
      quad.x = quad.data.x;
      quad.y = quad.data.y;
    }

    // For leaf nodes, accumulate forces from coincident quadrants.
    else {
      q = quad;
      q.x = q.data.x;
      q.y = q.data.y;
    }
  }

  /**
   * Function that applies the forces for each node. If the objects are
   * far away, the approximation is made. Otherwise, forces are calculated
   * directly between the nodes.
   * @param  {quadrant} quad - node representing the quadrant in quadtree.
   * @param  {number} x1   - lower x bound of the node.
   * @param  {number} _    - lower y bound of the node.
   * @param  {number} x2   - upper x bound of the node.
   * @return {boolean}     - true if the approximation was applied.
   */
  function apply(quad, x1, _, x2) {

    var x = quad.data.x + quad.data.vx - node.x - node.vx,
        y = quad.data.y + quad.data.vy - node.y - node.vy,
        w = x2 - x1,
        l = Math.sqrt(x * x + y * y);

    // Apply the Barnes-Hut approximation if possible.
    // Limit forces for very close nodes; randomize direction if coincident.
    if (w / l < theta) {
      if (x === 0) x = jiggle(), l += x * x;
      if (y === 0) y = jiggle(), l += y * y;
      if (quad.data) {
        l = (l - +distance(node, quad.data)) / l * alpha;
        x *= l, y *= l;
        quad.data.vx -= x;
        quad.data.vy -= y;
        node.vx += x;
        node.vy += y;
      }
      return true;
    }

    // Otherwise, process points directly.
    else if (quad.length) return;

    // Limit forces for very close nodes; randomize direction if coincident.
    if (quad.data !== node || quad.next) {
      if (x === 0) x = jiggle(), l += x * x;
      if (y === 0) y = jiggle(), l += y * y;
    }

    do if (quad.data !== node) {
      l = (l - +distance(node, quad.data)) / l * alpha;
      x *= l, y *= l;
      quad.data.vx -= x;
      quad.data.vy -= y;
      node.vx += x;
      node.vy += y;
    } while (quad = quad.next);
  }

   /**
    * Calculates the stress. Basically, it computes the difference between
    * high dimensional distance and real distance. The lower the stress is,
    * the better layout.
    * @return {number} - stress of the layout.
    */
  function getStress() {
    var totalDiffSq = 0, totalHighDistSq = 0;
    for (var i = 0, source, target, realDist, highDist; i < nodes.length; i++) {
      for (var j = 0; j < nodes.length; j++) {
        if (i !== j) {
          source = nodes[i], target = nodes[j];
          realDist = Math.hypot(target.x-source.x, target.y-source.y);
          highDist = +distance(nodes[i], nodes[j]);
          totalDiffSq += Math.pow(realDist-highDist, 2);
          totalHighDistSq += highDist * highDist;
        }
      }
    }
    return Math.sqrt(totalDiffSq/totalHighDistSq);
  }

  // API for initializing the algorithm, setting parameters and querying
  // metrics.
  force.initialize = function(_) {
    nodes = _;
  };

  force.distance = function(_) {
    return arguments.length ? (distance = typeof _ === "function" ? _ : constant(+_), force) : distance;
  };

  force.theta = function(_) {
    return arguments.length ? (theta = _, force) : theta;
  };

  force.stress = function() {
    return getStress();
  };

  return force;
}