From f08062fb16966e5819d60fa5a6d680a7ea13f92f Mon Sep 17 00:00:00 2001
From: Mike Bostock <mbostock@gmail.com>
Date: Fri, 8 Apr 2016 08:32:24 -0700
Subject: [PATCH] Assign node.index in d3.forceSimulation.

Also get ready for d3-quadtree 0.3 and per-node forces.
---
 package.json      |  2 +-
 src/constant.js   |  5 +++++
 src/link.js       | 17 ++++-------------
 src/manyBody.js   | 40 +++++++++++++++++++++-------------------
 src/simulation.js | 10 +++++++++-
 5 files changed, 40 insertions(+), 34 deletions(-)
 create mode 100644 src/constant.js

diff --git a/package.json b/package.json
index d7d593f..2a57ff9 100644
--- a/package.json
+++ b/package.json
@@ -32,7 +32,7 @@
   "dependencies": {
     "d3-collection": "0.1",
     "d3-dispatch": "0.4",
-    "d3-quadtree": "0.2",
+    "d3-quadtree": "0.3",
     "d3-timer": "0.4"
   },
   "devDependencies": {
diff --git a/src/constant.js b/src/constant.js
new file mode 100644
index 0000000..b7d42e7
--- /dev/null
+++ b/src/constant.js
@@ -0,0 +1,5 @@
+export default function(x) {
+  return function() {
+    return x;
+  };
+}
diff --git a/src/link.js b/src/link.js
index deb677f..6ce2dea 100644
--- a/src/link.js
+++ b/src/link.js
@@ -22,19 +22,10 @@ export default function(links) {
   }
 
   function initialize() {
-    var i, n, link, count;
-
-    for (i = 0, n = nodes.length, count = new Array(n); i < n; ++i) {
-      count[nodes[i].index = i] = 0;
-    }
-
-    for (i = 0, n = links.length; i < n; ++i) {
-      link = links[i], ++count[link.source.index], ++count[link.target.index];
-    }
-
-    for (i = 0, bias = new Array(n); i < n; ++i) {
-      link = links[i], bias[i] = count[link.source.index] / (count[link.source.index] + count[link.target.index]);
-    }
+    var i, n = nodes.length, m = links.length, k = new Array(n), l;
+    for (i = 0; i < n; ++i) k[i] = 0;
+    for (i = 0, bias = new Array(m); i < m; ++i) l = links[i], ++k[l.source.index], ++k[l.target.index];
+    for (i = 0; i < m; ++i) l = links[i], bias[i] = k[l.source.index] / (k[l.source.index] + k[l.target.index]);
   }
 
   force.nodes = function(_) {
diff --git a/src/manyBody.js b/src/manyBody.js
index 9b7944e..0cec6f1 100644
--- a/src/manyBody.js
+++ b/src/manyBody.js
@@ -1,4 +1,5 @@
 import {quadtree} from "d3-quadtree";
+import constant from "./constant";
 
 var nodeQuadtree = quadtree()
     .x(function(d) { return d.x; })
@@ -7,21 +8,19 @@ var nodeQuadtree = quadtree()
 function accumulate(quad, k) {
   var cx = 0, cy = 0, charge = 0;
 
-  if (!quad.leaf) {
-    for (var i = 0, child; i < 4; ++i) {
-      child = quad.nodes[i];
-      if (child == null) continue;
-      accumulate(child, k);
-      charge += child.charge;
-      cx += child.charge * child.cx;
-      cy += child.charge * child.cy;
-    }
+  if (quad.point) { // TODO handle coincident nodes
+    charge += quad.selfCharge = k;
+    cx += k * quad.point[0];
+    cy += k * quad.point[1];
   }
 
-  if (quad.data) {
-    charge += quad.selfCharge = k;
-    cx += k * quad.data.x;
-    cy += k * quad.data.y;
+  else for (var i = 0, child; i < 4; ++i) {
+    child = quad[i];
+    if (child == null) continue;
+    accumulate(child, k);
+    charge += child.charge;
+    cx += child.charge * child.cx;
+    cy += child.charge * child.cy;
   }
 
   quad.cx = cx / charge;
@@ -31,19 +30,22 @@ function accumulate(quad, k) {
 
 export default function() {
   var nodes,
-      strength = -100, // TODO compute per-node charge on initialization
+      strength = constant(-100),
       distance2 = Infinity,
       theta2 = 0.64,
       target;
 
   function force(alpha) {
     var root = nodeQuadtree(nodes);
-    accumulate(root, alpha * strength);
-    for (var i = 0, n = nodes.length; i < n; ++i) target = nodes[i], root.visit(apply);
+    accumulate(root, alpha * strength()); // TODO
+    for (var i = 0, n = nodes.length; i < n; ++i) {
+      target = nodes[i];
+      root.each(apply);
+    }
   }
 
   function apply(quad, x1, _, x2) {
-    if (quad.data !== target) {
+    if (!quad.point || quad.point.data !== target) {
       var dx = quad.cx - target.x,
           dy = quad.cy - target.y,
           dw = x2 - x1,
@@ -59,7 +61,7 @@ export default function() {
         return true;
       }
 
-      if (quad.data && dn && dn < distance2) {
+      if (quad.point && dn && dn < distance2) { // TODO handle coincident nodes
         k = quad.selfCharge / dn;
         target.vx += dx * k;
         target.vy += dy * k;
@@ -74,7 +76,7 @@ export default function() {
   };
 
   force.strength = function(_) {
-    return arguments.length ? (strength = +_, force) : strength;
+    return arguments.length ? (strength = typeof _ === "function" ? _ : constant(+_), force) : strength;
   };
 
   force.distance = function(_) {
diff --git a/src/simulation.js b/src/simulation.js
index 6522e1a..0bcded3 100644
--- a/src/simulation.js
+++ b/src/simulation.js
@@ -47,6 +47,12 @@ export default function(nodes) {
     dispatch.call("tick", simulation);
   }
 
+  function index() {
+    for (var i = 0, n = nodes.length; i < n; ++i) {
+      nodes[i].index = i;
+    }
+  }
+
   function initialize(force) {
     force.nodes(nodes);
   }
@@ -55,12 +61,14 @@ export default function(nodes) {
     force(alpha);
   }
 
+  index();
+
   return simulation = {
     start: start,
     stop: stop,
     tick: tick,
     nodes: function(_) {
-      return arguments.length ? (nodes = _, force.each(initialize), simulation) : nodes;
+      return arguments.length ? (nodes = _, index(), force.each(initialize), simulation) : nodes;
     },
     alphaMin: function(_) {
       return arguments.length ? (alphaMin = _, simulation) : alphaMin;