d3-spring-model/src/neighbourSampling.js

import constant from "./constant";
import jiggle from "./jiggle";
/**
 * An implementation of Chalmers' 1996 Neighbour and Sampling algorithm.
 * It uses random sampling to find the most suited neighbours from the
 * data set.
 */

function sortDistances(a, b) {
  return b[1] - a[1];
}

export default function () {
  var neighbours = [],
    distance = constant(300),
    nodes,
    neighbourSize = 10,
    sampleSize = 10,
    stableVelocity = 0,
    stableVeloHandler = null,
    dataSizeFactor,
    latestVelocityDiff = 0;

    /**
   * Apply spring forces at each simulation iteration.
   * @param  {number} alpha - multiplier for amount of force applied
   */
  function force(alpha) {
    let n = nodes.length;
    // Cache old velocity for comparison later
    if (stableVeloHandler!==null && stableVelocity>=0) {
      for (let i = n-1, node; i>=0; i--) {
        node = nodes[i];
        node.oldvx = node.vx;
        node.oldvy = node.vy;
      }
    }

    for (let i = n-1, node, samples; i>=0; i--) {
      node = nodes[i];
      samples = createRandomSamples(i);

      for (let [neighbourID, highDDist] of neighbours[i]) {
        setVelocity(node, nodes[neighbourID], highDDist, alpha);
      }

      for (let [sampleID, highDDist] of samples) {
        setVelocity(node, nodes[sampleID], highDDist, alpha);
      }

      neighbours[i] = findNewNeighbours(neighbours[i], samples);
    }

    // Calculate velocity changes, aka force applied.
    if (stableVeloHandler!==null && stableVelocity>=0) {
      let velocityDiff = 0;
      for (let i = n-1, node; i>=0; i--) {
        node = nodes[i];
        velocityDiff += Math.abs(Math.hypot(node.vx-node.oldvx, node.vy-node.oldvy));
      }
      velocityDiff /= n;
      latestVelocityDiff = velocityDiff;

      if(velocityDiff<stableVelocity){
        stableVeloHandler();
      }
    }
  }

  /**
   * Apply force to both source and target nodes.
   * @param {number} source  - source node object
   * @param {number} target  - target node object
   * @param {number} dist    - high dimensional distance between the two nodes
   * @param {number} alpha   - multiplier for the amount of force applied
   */
  function setVelocity(source, target, dist, alpha) {
    let x, y, l;
    // jiggle so l won't be zero and divide by zero error after this
    x = target.x + target.vx - source.x - source.vx || jiggle();
    y = target.y + target.vy - source.y - source.vy || jiggle();
    l = Math.sqrt(x * x + y * y);
    l = (l - dist) / l * dataSizeFactor * alpha;
    x *= l, y *= l;
    // Set the calculated velocites for both nodes.
    target.vx -= x;
    target.vy -= y;
    source.vx += x;
    source.vy += y;
  }

  // Called on nodes change and added to a simulation
  function initialize() {
    if (!nodes) return;

    // Initialize for each node some random neighbours.
    for (let i = nodes.length-1; i>=0; i--) {
      let neighbs = pickRandomNodesFor(i, [i], neighbourSize);
      // Sort the neighbour set by the distances.
      neighbours[i] = new Map(neighbs.sort(sortDistances));
    }

    initDataSizeFactor();
  }

  function initDataSizeFactor(){
    dataSizeFactor = 0.5/(neighbourSize+sampleSize);
  }

  /**
   * Generates an array of array[index, high-d distance to the node of index]
   * where all indices are different and the size is as specified unless
   * impossible (may be due to too large size requested)
   * @param  {number} index  - index of a node to calculate distance against
   * @param  {array} exclude - indices of the nodes to ignore.
   * @param  {number} size   - max number of elements in the map to return.
   * @return {array}
   */
  function pickRandomNodesFor(index, exclude, size) {
    let randElements = [];
    let max = nodes.length;

    for (let i = 0; i < size; i++) {
      // Stop when no new elements can be found.
      if (randElements.length + exclude.length >= nodes.length) {
        break;
      }

      let rand = Math.floor(Math.random() * max);
      // Re-random until suitable value is found.
      while (randElements.includes(rand) || exclude.includes(rand)) {
        rand = Math.floor(Math.random() * max);
      }
      randElements.push(rand);
    }
    for(let i=randElements.length-1, rand; i>=0; i--){
      rand = randElements[i];
      randElements[i] = [rand, distance(nodes[index], nodes[rand])];
    }
    return randElements;
  }

  /**
   * Generates a map {index: high-dimensional distance to the node of index}
   * to be used as samples set for the node of the specified index.
   * @param  {number} index - index of the node to generate sample for
   * @return {map}
   */
  function createRandomSamples(index) {
    // Ignore the current neighbours of the node and itself.
    let exclude = [index];
    exclude = exclude.concat(Array.from(neighbours[index].keys()));
    return new Map(pickRandomNodesFor(index, exclude, sampleSize));
  }

  /**
   * Compares the elements from sample set to the neighbour set and replaces the
   * elements in the neighbour set if any better neighbours are found.
   * @param  {map} neighbours - map of neighbours
   * @param  {map} samples    - map of samples
   * @return {map}            - new map of neighbours
   */
  function findNewNeighbours(neighbours, samples) {
    let combined = [...neighbours.entries()].concat([...samples.entries()]);
    combined = combined.sort(sortDistances);
    return new Map(combined.slice(0, neighbourSize));
  }


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

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

  force.neighbours = function () {
    return neighbours;
  };

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

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

  force.latestAccel = function () {
    return latestVelocityDiff;
  };

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

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