d3 fmm

ball-tree.js

// class BallTree {
//   constructor (points) {
//     const X = new Array(points.length);
//     const Y = new Array(points.length);

//     for (let i = 0; i < points.length; i++) X[i] = Y[i] = i;

//     X.sort((a, b) => points[a].x - points[b].x);
//     Y.sort((a, b) => points[a].y - points[b].y);

//     this._root = this.buildTree(points, X, Y);
//   }

//   buildTree(pts, byX, bY) {
//     const dx = Math.abs(byX[byX.length - 1] - byX[0]);
//     const dy = Math.abs(byY[byY.length - 1] - byY[0]);
//     if (dx > dy) { // split by x
//       const med = Math.floor(pts.length / 2);
//       const left = new Array(med), right = new Array(points.length - med);
//       for (let i = 0; i < points.length; i++) {
//         if (i < med) {
//           left[i] = pts[byX[i]];
//         } else {
//           right[i - med] = pts[byX[i]];
//         }
//       }
//     } else { // split by y

//     }
//   }
// }



// function construct_balltree is
//        input:
//            D, an array of data points
//        output:
//            B, the root of a constructed ball tree
//        if a single point remains then
//            create a leaf B containing the single point in D
//            return B
//        else
//            let c be the dimension of greatest spread
//            let p be the central point selected considering c
//            let L,R be the sets of points lying to the left and right of the median along dimension c
//            create B with two children:
//                B.pivot = p
//                B.child1 = construct_balltree(L),
//                B.child2 = construct_balltree(R),
//                let B.radius be maximum distance from p among children
//            return B
//        end if
//    end function


class BallTree {
  constructor (data, leaf_size = 1) {
    this.data = data;
    this.leaf_size = leaf_size;

    this.n_samples = this.data.length;
    this.n_features = 2;

    // determine number of levels in the tree, and from this
    // the number of nodes in the tree.  This results in leaf nodes
    // with numbers of points betweeen leaf_size and 2 * leaf_size
    this.n_levels = 1 + Math.log2(Math.max(1, Math.floor(this.n_samples - 1) / this.leaf_size));
    this.n_nodes = (1 << this.n_levels) - 1;

    // allocate arrays for storage
    this.idx_array = data.map((_, i) => i);
    this.node_radius = new Float32Array(data.length);
    this.node_idx_start = new Uint32Array(this.n_nodes);
    this.node_idx_end = new Uint32Array(this.n_nodes);
    this.node_is_leaf = new Uint8Array(this.n_nodes);
    this.node_centroidsx = new Float32Array(data.length);
    this.node_centroidsy = new Float32Array(data.length);

    // Allocate tree-specific data from TreeBase
    this._recursive_build(0, 0, this.n_samples);
  }

  _recursive_build (i_node, idx_start, idx_end) {
    //initialize node data
    this.init_node(i_node, idx_start, idx_end);

    if (2 * i_node + 1 >= this.n_nodes) {
      this.node_is_leaf[i_node] = true;
      if (idx_end - idx_start > 2 * this.leaf_size) {
        // this shouldn't happen if our memory allocation is correct
        // we'll proactively prevent memory errors, but raise a
        // warning saying we're doing so.
        //console.warn("Internal: memory layout is flawed: not enough nodes allocated");
      } else if (idx_end - idx_start < 2) {
        // again, this shouldn't happen if our memory allocation
        // is correct.  Raise a warning.
        //console.warn("Internal: memory layout is flawed: too many nodes allocated");
        this.node_is_leaf[i_node] = true;
      }
    } else {
      // split node and recursively construct child nodes.
      this.node_is_leaf[i_node] = false;
      const n_mid = Math.floor((idx_end + idx_start) / 2);
      this._partition_indices(this.data, this.idx_array, idx_start, idx_end, n_mid);
      this._recursive_build(2 * i_node + 1, idx_start, n_mid);
      this._recursive_build(2 * i_node + 2, n_mid, idx_end);
    }
  }

  init_node (i_node, idx_start, idx_end) {
    // determine Node centroid
    this.node_centroidsx[i_node] = 0;
    this.node_centroidsy[i_node] = 0;
    for (let i = idx_start; i < idx_end; i++) {
      this.node_centroidsx[i_node] += this.data[this.idx_array[i]].x;
      this.node_centroidsy[i_node] += this.data[this.idx_array[i]].y;
    }
    this.node_centroidsx[i_node] /= (idx_end - idx_start);
    this.node_centroidsy[i_node] /= (idx_end - idx_start);

    // determine Node radius
    let sq_radius = 0;
    for (let i = idx_start; i < idx_end; i++) {
      const x1 = this.node_centroidsx[i_node], x2 = this.data[this.idx_array[i]].x;
      const y1 = this.node_centroidsy[i_node], y2 = this.data[this.idx_array[i]].y;
      const dx = x2 - x1, dy = y2 - y1;
      const sq_dist = dx * dx + dy * dy;
      sq_radius = Math.max(sq_radius, sq_dist);
    }

    //this.node_radius[i_node] = np.sqrt(sq_radius);
    this.node_idx_start[i_node] = idx_start;
    this.node_idx_end[i_node] = idx_end;
  }

  rdist (X1, i1, X2, i2) {
    let d = 0;

    tmp = (X1[i1, k] - X2[i2, k])
    d += tmp * tmp
    tmp = (X1[i1, k] - X2[i2, k])
    d += tmp * tmp
    return d;
  }

  min_rdist(i_node, X, j) {
    d = this.rdist(this.node_centroids, i_node, X, j);
    return max(0, np.sqrt(d) - this.node_radius[i_node]) ** 2;
  }


  _partition_indices (data, idx_array, idx_start, idx_end, split_index) {
    // Find the split dimension
    let max_spread = 0;

    let minx = Infinity, miny = Infinity, maxx = -Infinity, maxy = -Infinity;
    for (let i = idx_start; i < idx_end; i++) {
      const { x, y } = this.data[i];
      if (x < minx) minx = x;
      if (y < miny) miny = y;
      if (x > maxx) maxx = x;
      if (y > maxy) maxy = y;
    }


    const max_spreadx = maxx - minx;
    const max_spready = maxy - miny;

    const split_dim = (max_spreadx > max_spready) ? 'x' : 'y';

    // Partition using the split dimension
    let left = idx_start;
    let right = idx_end - 1;
    let tmp;

    while (true) {
      let midindex = left;
      for (let i = left; i < right; i++) {
        const d1 = data[idx_array[i]][split_dim];
        const d2 = data[idx_array[right]][split_dim];
        if (d1 < d2) {
          tmp = idx_array[i];
          idx_array[i] = idx_array[midindex]
          idx_array[midindex] = tmp;
          midindex += 1;
        }
      }
      tmp = idx_array[midindex];
      idx_array[midindex] = idx_array[right];
      idx_array[right] = tmp;

      if (midindex === split_index) break;
      else if (midindex < split_index) left = midindex + 1;
      else                             right = midindex - 1;
    }
  }
}

// class NeighborsHeap:
//     def __init__(self, n_pts, n_nbrs):
//         this.distances = np.zeros((n_pts, n_nbrs), dtype=float) + np.inf
//         this.indices = np.zeros((n_pts, n_nbrs), dtype=int)

//     def get_arrays(self, sort=true):
//         if sort:
//             i = np.arange(len(this.distances), dtype=int)[:, None]
//             j = np.argsort(this.distances, 1)
//             return this.distances[i, j], this.indices[i, j]
//         else:
//             return this.distances, this.indices

//     def largest(self, row):
//         return this.distances[row, 0]

//     def push(self, row, val, i_val):
//         size = this.distances.shape[1]

//         # check if val should be in heap
//         if val > this.distances[row, 0]:
//             return

//         # insert val at position zero
//         this.distances[row, 0] = val
//         this.indices[row, 0] = i_val

//         #descend the heap, swapping values until the max heap criterion is met
//         i = 0
//         while true:
//             ic1 = 2 * i + 1
//             ic2 = ic1 + 1

//             if ic1 >= size:
//                 break
//             elif ic2 >= size:
//                 if this.distances[row, ic1] > val:
//                     i_swap = ic1
//                 else:
//                     break
//             elif this.distances[row, ic1] >= this.distances[row, ic2]:
//                 if val < this.distances[row, ic1]:
//                     i_swap = ic1
//                 else:
//                     break
//             else:
//                 if val < this.distances[row, ic2]:
//                     i_swap = ic2
//                 else:
//                     break

//             this.distances[row, i] = this.distances[row, i_swap]
//             this.indices[row, i] = this.indices[row, i_swap]

//             i = i_swap

//         this.distances[row, i] = val
//         this.indices[row, i] = i_val


(function (N, D, LS) {
  console.log("-------------------------------------------------------")
  console.log(`${N} points in ${D} dimensions`);

  const X = new Array(N).fill(0).map(() => ({ x: Math.random() * N, y: Math.random() * N }));
  console.time('build');
  const bt = new BallTree(X, LS);
  console.timeEnd('build');
})(10000, 3, 1)

index.html

<!doctype html>
<html>
  <head>
    <title>Graph playground</title>
    <script src="https://unpkg.com/d3@5.5.0/dist/d3.min.js"></script>
    <style>

    html, body {
      margin: 0;
      padding: 0;
      width: 100%;
      height: 100%;
      overflow: hidden;
    }

    .control {
      display: none;
      position: absolute;
      top: 20px;
      right: 20px;
      padding: 10px 20px;
      font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
      font-weight: 300;
      font-size: 13px;
      background: #fff;
      border-radius: 4px;
      box-shadow: 0 0 5px rgba(0,0,0,0.5);
    }
      .control label {
        display: block;
      }

    #range {
      position: absolute;
      bottom: 20px;
      left: 20px;
      width: 90%;
    }

    #indicator {
      font-family: Georgia, "Times New Roman", serif;
      position: absolute;
      bottom: 50px;
      right: 20px;
      left: 20px;
      font-size: 1.5em;
      font-weight: 300;
      color: #222222;
    }

    #center-view {
      display: block;
      position: absolute;
      bottom: 20px;
      right: 20px;
      cursor: pointer;
      border-radius: 4px;
      box-shadow: 0 0 5px rgba(0,0,0,0.5);
      width: 20px;
      height: 20px;
      text-align: center;
      line-height: 20px;
    }
      #center-view:active {
        background: #aaa;
      }
      #center-view:hover {
        background: #ddd;
      }
    </style>
  </head>
  <body>
    <canvas id="canvas"></canvas>
    <form class="control" id="params">
      <label><input name="input1" value="0" type="number"/> input 1</label>
      <label><input name="input2" value="1" type="number"> input 2</label>
    </form>

    <input type="range" id="range" min="10" max="10000" value="250" />
    <div id="indicator">250</div>
    <div id="center-view" title="center view">&bullet;</div>
    <script src="utils.js"></script>
    <script src="index.js"></script>
  </body>
</html>

index.js

const screenWidth  = document.documentElement.clientWidth;
const screenHeight = document.documentElement.clientHeight;

const pxRatio = window.devicePixelRatio;

const canvas = document.getElementById('canvas');
const ctx = canvas.getContext('2d');

canvas.style.width  = screenWidth + 'px';
canvas.style.height = screenHeight + 'px';

// global width/height to use for rendering, accunting for retina screens
const w = canvas.width = screenWidth * pxRatio;
const h = canvas.height = screenHeight * pxRatio;


let data = { nodes: [], edges: [] };  // store graph here
let nodesMap = {};

// store highlighted nodes and edges there to render
const highlightedEdges = [];
const highlightedNodes = [];

// central node to be highlighted differently, use as you want
let selectedNode = null;
let q = d3.quadtree([], n => n.x, n => n.y);

let transform = { x: 0, y: 0, k: 1 };

function render() {
  if (d3.event) transform = d3.event.transform;
  ctx.save();
  ctx.clearRect(0, 0, w, h);
  ctx.translate(transform.x * pxRatio, transform.y * pxRatio);
  ctx.scale(transform.k, transform.k);

  const Nn = data.nodes.length;
  const { nodes, edges } = data;

  // edges
  ctx.lineWidth = 3;
  ctx.beginPath();
  ctx.strokeStyle = '#707070';
  edges.forEach((e) => {
    const s = nodesMap[e.source],
          t = nodesMap[e.target];
    ctx.moveTo(s.x, s.y);
    ctx.lineTo(t.x, t.y);
  });
  ctx.stroke();
  ctx.closePath();

  // highlighted edges
  ctx.beginPath();
  ctx.fillStyle   = 'red';
  ctx.strokeStyle = '#505050';
  ctx.lineWidth = 5;
  highlightedEdges.forEach((e) => {
    const s = nodesMap[e.source],
          t = nodesMap[e.target];
    ctx.moveTo(s.x, s.y);
    ctx.lineTo(t.x, t.y);
  });
  ctx.closePath();
  ctx.stroke();

  // nodes;
  ctx.beginPath();
  ctx.fillStyle   = 'orange';
  ctx.strokeStyle = '#333333';
  ctx.lineWidth = 1;
  nodes.forEach((n) => {
    ctx.moveTo(n.x + n.r, n.y);
    ctx.arc(n.x, n.y, n.r, 0, 2 * Math.PI, false);
  });
  ctx.stroke();
  ctx.fill();

  // selected node
  if (selectedNode) {
    ctx.beginPath();
    ctx.fillStyle   = 'red';
    ctx.strokeStyle = '#333333';
    ctx.moveTo(selectedNode.x + selectedNode.r, selectedNode.y);
    ctx.arc(selectedNode.x, selectedNode.y, selectedNode.r, 0, 2 * Math.PI, false);
    ctx.stroke();
    ctx.fill();
  }

  // highlighted nodes
  ctx.beginPath();
  ctx.fillStyle   = 'red';
  ctx.strokeStyle = '#333333';
  ctx.lineWidth = 15;
  highlightedNodes.forEach((n) => {
    ctx.moveTo(n.x + n.r, n.y);
    ctx.arc(n.x, n.y, n.r, 0, 2 * Math.PI, false);
  });
  ctx.stroke();
  ctx.fill();

  drawQuadTree(ctx);
  ctx.restore();
}


function drawQuadTree(ctx) {
  q = d3.quadtree([], n => n.x, n => n.y);
  q.addAll(data.nodes);
  q.visit((n) => {
    if (n.length) {
      for (let i = 0; i < 4; i++) if (n[i]) n[i].parent = n;
    }
  });
  q.visitAfter((quad) => {
    var strength = 0, q, c, weight = 0, x, y, i;
    const out = [];
    // For internal nodes, accumulate forces from child quadrants.
    if (quad.length) {
      const ins = [];
      for (x = y = i = 0; i < 4; ++i) {
        if ((q = quad[i]) && q.med) ins.push(q.med);
      }
      out.push(quad.med = d3.packEnclose(ins));
    }

    // For leaf nodes, accumulate forces from coincident quadrants.
    else {
      out.push(quad.med = quad.data);
      // do strength += strengths[q.data.index]; while (q = q.next);
    }

    quad.med = d3.packEnclose(out);
  });
  ctx.beginPath();
  ctx.globalAlpha = 0.75;
  ctx.lineWidth = 0.5;
  q.visit((n, x0, y0, x1, y1) => {
    ctx.rect(x0, y0, x1 - x0, y1 - y0);
    if (n.length) {
      ctx.moveTo(n.med.x + n.med.r, n.med.y);
      ctx.arc(n.med.x, n.med.y, n.med.r, 0, 2 * Math.PI, false);
    }
  });
  ctx.stroke();
  ctx.globalAlpha = 1;
  ctx.closePath();
}


d3.select('canvas')
  .call(d3.zoom().scaleExtent([0.1, 10]).on('zoom', function() {
    transform = d3.event.transform;
    requestAnimationFrame(render);
  }));
// render on mouse move
canvas.addEventListener('mousemove', ({ clientX, clientY }) => {
  const x = clientX * pxRatio;
  const y = clientY * pxRatio;

  const offset = 10 * pxRatio;

//   qBounds[0] = x - offset;
//   qBounds[1] = y - offset;
//   qBounds[2] = x + offset;
//   qBounds[3] = y + offset;

  requestAnimationFrame(render);
});

const form = document.querySelector('#params');
function onChange() {
  // read controls, apply actions
  console.log(form['input1'].value);
  console.log(form['input2'].value);
}

// listen to form changes
form.addEventListener('change', onChange);

// graph size control
const indicator  = document.querySelector('#indicator');
const rangeInput = document.querySelector('#range');

let populateTimer = 0;
let N = 550;
function onGraphChanged() {
  clearTimeout(populateTimer);
  populateTimer = requestAnimationFrame(() => {
    N = parseInt(rangeInput.value);
    indicator.innerHTML = `G = (${N}<sub>V</sub> &times; ${N}<sub>E</sub>)`;
    init(N);
  });
}
rangeInput.addEventListener('change', onGraphChanged);

// center view;
document.querySelector('#center-view').addEventListener('click', () => {
  const padding = 20 * pxRatio;
  scaleGraphToFitBounds(data, [
    0 + padding, 0 + padding,
    w - padding, h - padding
  ]);
});

function init (N) {
  onChange();
  data = generateTree(N);
  nodesMap = data.nodes.reduce((a, n) => {
    a[n.id] = n;
    return a;
  }, {});
  layout(data, render).then(render);
}

onGraphChanged();

utils.js

/**
 * @param {object} G
 * @param {Array<Object>} G.nodes
 * @param {Array<Object>} G.edges
 * @param {Array<Number>} [xmin,ymin,xmax,ymax]
 */
function scaleGraphToFitBounds(G, bounds) {
  var min = Math.min, max = Math.max;
  var bbox = getBounds(G);
  var cx = (bounds[0] + bounds[2]) / 2,
      cy = (bounds[1] + bounds[3]) / 2,
      gcx = (bbox[0] + bbox[2]) / 2,
      gcy = (bbox[1] + bbox[3]) / 2;

  var tx = cx - gcx,
      ty = cy - gcy;

  var scale = min(
    (bounds[2] - bounds[0]) / (bbox[2] - bbox[0]),
    (bounds[3] - bounds[1]) / (bbox[3] - bbox[1])
  );


  G.nodes.forEach(function (n) {
    n.x = cx + ((n.x + tx) - cx) * scale;
    n.y = cy + ((n.y + ty) - cy) * scale;
  });
}

/**
 * G -> [xmin, ymin, xmax, ymax], radii are ignored
 */
function getBounds(G) {
  return G.nodes.reduce((b, n) => {
    b[0] = Math.min(b[0], n.x);
    b[1] = Math.min(b[1], n.y);
    b[2] = Math.max(b[2], n.x);
    b[3] = Math.max(b[3], n.y);
    return b;
  }, [Infinity, Infinity, -Infinity, -Infinity]);
}

/**
 * O(V+E) tree graph generator
 */
function generateTree(N = 550, width = 1000, height = 1000) {
  let counter = 0;
  return {
    nodes: new Array(N).fill(0).map(function(_, i) {
      counter++;
      return {
        id: i,
        text: 'Node #' + i,
        x: Math.random() * width,
        y: Math.random() * height,
        r: 3 + Math.random() * 8
      };
    }),
    edges: new Array(N - 1).fill(0).map(function(_, i) {
      counter++;
      return {
        id: i,
        source: Math.floor(Math.sqrt(i)),
        target: i + 1
      };
    })
  };
}

/**
 * Random RGB color
 */
function getRandomColor() {
  const letters = '0123456789ABCDEF';
  let color = '#';
  for (var i = 0; i < 6; i++) {
    color += letters[Math.floor(Math.random() * 16)];
  }
  return color;
}

let simulation = null;


function fmm ({ nodes }) {

  const getX = n => n.x;
  const getY = n => n.y;

  function force (alpha) {
    let i, n = nodes.length;
    const tree = d3.quadtree(nodes, getX, getY);
    const leafs = [];
    tree.visit((node) => {
      if (node.length) {
        for (let i = 0; i < 4; i++) {
          const child = node[i];
          if (child) child.parent = node;
        }
      } else leafs.push(node);
    });

    //d3.packEnclose([{x: 0, y: 0, r: 200}])
  }

  return force;
}

/**
 * Force layout to get the initial positions
 * @param {Graph} data
 * @param [function] onUpdate - to call on each layout pass
 */
function layout(data, onUpdate = () => {}) {
  return new Promise((resolve) => {
    const m = 2;
    const {
      forceSimulation,
      forceLink,
      forceX, forceY,
      forceCenter,
      forceCollide,
      forceManyBody
    } = d3;
    if (simulation) simulation.stop();
    simulation = forceSimulation(data.nodes)
      .force('charge', forceManyBody())
      .force('link', forceLink(data.edges.map((e) => ({
          id: e.id,
          target: e.target,
          source: e.source
        }))
      )
        .id((d) => d.id))
      .force('collision', forceCollide(10))
      .force('center', forceCenter(w / 2, h / 2))
      .force('y', forceY(0))
      .force('x', forceX(0))
      //.alphaDecay(0.1)  // increase for rougher results
      .on('tick', onUpdate)
      .on('end', resolve);
  });
}


function layout(data, onUpdate = () => {}) {
  return new Promise((resolve) => {
    const m = 2;
    const {
      forceSimulation,
      forceLink,
      forceX, forceY,
      forceCenter,
      forceCollide,
      forceManyBody
    } = d3;
    if (simulation) simulation.stop();
    simulation = forceSimulation(data.nodes)
      .force('fmm', fmm(data))
      .force('charge', forceManyBody())
      .force('link', forceLink(data.edges.map((e) => ({
          id: e.id,
          target: e.target,
          source: e.source
        }))
      )
        .id((d) => d.id))
      //.force('collision', forceCollide(10))
      .force('center', forceCenter(w / 2, h / 2))
      .force('y', forceY(0))
      .force('x', forceX(0))
      //.alphaDecay(0.1)  // increase for rougher results
      .on('tick', onUpdate)
      .on('end', resolve);
  });
}