sankey.js 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388
  1. import { sum, max, min } from 'd3-array';
  2. import { justify } from './align';
  3. import { constant } from './constant';
  4. function ascendingSourceBreadth(a, b) {
  5. return ascendingBreadth(a.source, b.source) || a.index - b.index;
  6. }
  7. function ascendingTargetBreadth(a, b) {
  8. return ascendingBreadth(a.target, b.target) || a.index - b.index;
  9. }
  10. function ascendingBreadth(a, b) {
  11. return a.y0 - b.y0;
  12. }
  13. function value(d) {
  14. return d.value;
  15. }
  16. function defaultId(d) {
  17. return d.index;
  18. }
  19. function defaultNodes(graph) {
  20. return graph.nodes;
  21. }
  22. function defaultLinks(graph) {
  23. return graph.links;
  24. }
  25. function find(nodeById, id) {
  26. const node = nodeById.get(id);
  27. if (!node)
  28. throw new Error('missing: ' + id);
  29. return node;
  30. }
  31. function computeLinkBreadths({ nodes }) {
  32. for (const node of nodes) {
  33. let y0 = node.y0;
  34. let y1 = y0;
  35. for (const link of node.sourceLinks) {
  36. link.y0 = y0 + link.width / 2;
  37. y0 += link.width;
  38. }
  39. for (const link of node.targetLinks) {
  40. link.y1 = y1 + link.width / 2;
  41. y1 += link.width;
  42. }
  43. }
  44. }
  45. export function Sankey() {
  46. let x0 = 0, y0 = 0, x1 = 1, y1 = 1; // extent
  47. let dx = 24; // nodeWidth
  48. let dy = 8, py; // nodePadding
  49. let id = defaultId;
  50. let align = justify;
  51. let depth;
  52. let sort;
  53. let linkSort;
  54. let nodes = defaultNodes;
  55. let links = defaultLinks;
  56. let iterations = 6;
  57. function sankey(arg) {
  58. const graph = {
  59. nodes: nodes(arg),
  60. links: links(arg),
  61. };
  62. computeNodeLinks(graph);
  63. computeNodeValues(graph);
  64. computeNodeDepths(graph);
  65. computeNodeHeights(graph);
  66. computeNodeBreadths(graph);
  67. computeLinkBreadths(graph);
  68. return graph;
  69. }
  70. sankey.update = function (graph) {
  71. computeLinkBreadths(graph);
  72. return graph;
  73. };
  74. sankey.nodeId = function (_) {
  75. return arguments.length
  76. ? ((id = typeof _ === 'function' ? _ : constant(_)), sankey)
  77. : id;
  78. };
  79. sankey.nodeAlign = function (_) {
  80. return arguments.length
  81. ? ((align = typeof _ === 'function' ? _ : constant(_)), sankey)
  82. : align;
  83. };
  84. sankey.nodeDepth = function (_) {
  85. return arguments.length
  86. ? ((depth = typeof _ === 'function' ? _ : _), sankey)
  87. : depth;
  88. };
  89. sankey.nodeSort = function (_) {
  90. return arguments.length ? ((sort = _), sankey) : sort;
  91. };
  92. sankey.nodeWidth = function (_) {
  93. return arguments.length ? ((dx = +_), sankey) : dx;
  94. };
  95. sankey.nodePadding = function (_) {
  96. return arguments.length ? ((dy = py = +_), sankey) : dy;
  97. };
  98. sankey.nodes = function (_) {
  99. return arguments.length
  100. ? ((nodes = typeof _ === 'function' ? _ : constant(_)), sankey)
  101. : nodes;
  102. };
  103. sankey.links = function (_) {
  104. return arguments.length
  105. ? ((links = typeof _ === 'function' ? _ : constant(_)), sankey)
  106. : links;
  107. };
  108. sankey.linkSort = function (_) {
  109. return arguments.length ? ((linkSort = _), sankey) : linkSort;
  110. };
  111. sankey.size = function (_) {
  112. return arguments.length
  113. ? ((x0 = y0 = 0), (x1 = +_[0]), (y1 = +_[1]), sankey)
  114. : [x1 - x0, y1 - y0];
  115. };
  116. sankey.extent = function (_) {
  117. return arguments.length
  118. ? ((x0 = +_[0][0]),
  119. (x1 = +_[1][0]),
  120. (y0 = +_[0][1]),
  121. (y1 = +_[1][1]),
  122. sankey)
  123. : [
  124. [x0, y0],
  125. [x1, y1],
  126. ];
  127. };
  128. sankey.iterations = function (_) {
  129. return arguments.length ? ((iterations = +_), sankey) : iterations;
  130. };
  131. function computeNodeLinks({ nodes, links }) {
  132. nodes.forEach((node, idx) => {
  133. node.index = idx;
  134. node.sourceLinks = [];
  135. node.targetLinks = [];
  136. });
  137. const nodeById = new Map(nodes.map((d) => [id(d), d]));
  138. links.forEach((link, idx) => {
  139. link.index = idx;
  140. let { source, target } = link;
  141. if (typeof source !== 'object')
  142. source = link.source = find(nodeById, source);
  143. if (typeof target !== 'object')
  144. target = link.target = find(nodeById, target);
  145. source.sourceLinks.push(link);
  146. target.targetLinks.push(link);
  147. });
  148. if (linkSort != null) {
  149. for (const { sourceLinks, targetLinks } of nodes) {
  150. sourceLinks.sort(linkSort);
  151. targetLinks.sort(linkSort);
  152. }
  153. }
  154. }
  155. function computeNodeValues({ nodes }) {
  156. for (const node of nodes) {
  157. node.value =
  158. node.fixedValue === undefined
  159. ? Math.max(sum(node.sourceLinks, value), sum(node.targetLinks, value))
  160. : node.fixedValue;
  161. }
  162. }
  163. function computeNodeDepths({ nodes }) {
  164. const n = nodes.length;
  165. let current = new Set(nodes);
  166. let next = new Set();
  167. let x = 0;
  168. while (current.size) {
  169. current.forEach((node) => {
  170. node.depth = x;
  171. for (const { target } of node.sourceLinks) {
  172. next.add(target);
  173. }
  174. });
  175. if (++x > n)
  176. throw new Error('circular link');
  177. current = next;
  178. next = new Set();
  179. }
  180. // 如果配置了 depth,则设置自定义 depth
  181. if (depth) {
  182. const maxDepth = Math.max(max(nodes, (d) => d.depth) + 1, 0);
  183. let node;
  184. for (let i = 0; i < nodes.length; i++) {
  185. node = nodes[i];
  186. node.depth = depth.call(null, node, maxDepth);
  187. }
  188. }
  189. }
  190. function computeNodeHeights({ nodes }) {
  191. const n = nodes.length;
  192. let current = new Set(nodes);
  193. let next = new Set();
  194. let x = 0;
  195. while (current.size) {
  196. current.forEach((node) => {
  197. node.height = x;
  198. for (const { source } of node.targetLinks) {
  199. next.add(source);
  200. }
  201. });
  202. if (++x > n)
  203. throw new Error('circular link');
  204. current = next;
  205. next = new Set();
  206. }
  207. }
  208. function computeNodeLayers({ nodes }) {
  209. const x = Math.max(max(nodes, (d) => d.depth) + 1, 0);
  210. const kx = (x1 - x0 - dx) / (x - 1);
  211. const columns = new Array(x).fill(0).map(() => []);
  212. for (const node of nodes) {
  213. const i = Math.max(0, Math.min(x - 1, Math.floor(align.call(null, node, x))));
  214. node.layer = i;
  215. node.x0 = x0 + i * kx;
  216. node.x1 = node.x0 + dx;
  217. if (columns[i])
  218. columns[i].push(node);
  219. else
  220. columns[i] = [node];
  221. }
  222. if (sort)
  223. for (const column of columns) {
  224. column.sort(sort);
  225. }
  226. return columns;
  227. }
  228. function initializeNodeBreadths(columns) {
  229. const ky = min(columns, (c) => (y1 - y0 - (c.length - 1) * py) / sum(c, value));
  230. for (const nodes of columns) {
  231. let y = y0;
  232. for (const node of nodes) {
  233. node.y0 = y;
  234. node.y1 = y + node.value * ky;
  235. y = node.y1 + py;
  236. for (const link of node.sourceLinks) {
  237. link.width = link.value * ky;
  238. }
  239. }
  240. y = (y1 - y + py) / (nodes.length + 1);
  241. for (let i = 0; i < nodes.length; ++i) {
  242. const node = nodes[i];
  243. node.y0 += y * (i + 1);
  244. node.y1 += y * (i + 1);
  245. }
  246. reorderLinks(nodes);
  247. }
  248. }
  249. function computeNodeBreadths(graph) {
  250. const columns = computeNodeLayers(graph);
  251. py = Math.min(dy, (y1 - y0) / (max(columns, (c) => c.length) - 1));
  252. initializeNodeBreadths(columns);
  253. for (let i = 0; i < iterations; ++i) {
  254. const alpha = Math.pow(0.99, i);
  255. const beta = Math.max(1 - alpha, (i + 1) / iterations);
  256. relaxRightToLeft(columns, alpha, beta);
  257. relaxLeftToRight(columns, alpha, beta);
  258. }
  259. }
  260. // Reposition each node based on its incoming (target) links.
  261. function relaxLeftToRight(columns, alpha, beta) {
  262. for (let i = 1, n = columns.length; i < n; ++i) {
  263. const column = columns[i];
  264. for (const target of column) {
  265. let y = 0;
  266. let w = 0;
  267. for (const { source, value } of target.targetLinks) {
  268. const v = value * (target.layer - source.layer);
  269. y += targetTop(source, target) * v;
  270. w += v;
  271. }
  272. if (!(w > 0))
  273. continue;
  274. const dy = (y / w - target.y0) * alpha;
  275. target.y0 += dy;
  276. target.y1 += dy;
  277. reorderNodeLinks(target);
  278. }
  279. if (sort === undefined)
  280. column.sort(ascendingBreadth);
  281. if (column.length)
  282. resolveCollisions(column, beta);
  283. }
  284. }
  285. // Reposition each node based on its outgoing (source) links.
  286. function relaxRightToLeft(columns, alpha, beta) {
  287. for (let n = columns.length, i = n - 2; i >= 0; --i) {
  288. const column = columns[i];
  289. for (const source of column) {
  290. let y = 0;
  291. let w = 0;
  292. for (const { target, value } of source.sourceLinks) {
  293. const v = value * (target.layer - source.layer);
  294. y += sourceTop(source, target) * v;
  295. w += v;
  296. }
  297. if (!(w > 0))
  298. continue;
  299. const dy = (y / w - source.y0) * alpha;
  300. source.y0 += dy;
  301. source.y1 += dy;
  302. reorderNodeLinks(source);
  303. }
  304. if (sort === undefined)
  305. column.sort(ascendingBreadth);
  306. if (column.length)
  307. resolveCollisions(column, beta);
  308. }
  309. }
  310. function resolveCollisions(nodes, alpha) {
  311. const i = nodes.length >> 1;
  312. const subject = nodes[i];
  313. resolveCollisionsBottomToTop(nodes, subject.y0 - py, i - 1, alpha);
  314. resolveCollisionsTopToBottom(nodes, subject.y1 + py, i + 1, alpha);
  315. resolveCollisionsBottomToTop(nodes, y1, nodes.length - 1, alpha);
  316. resolveCollisionsTopToBottom(nodes, y0, 0, alpha);
  317. }
  318. // Push any overlapping nodes down.
  319. function resolveCollisionsTopToBottom(nodes, y, i, alpha) {
  320. for (; i < nodes.length; ++i) {
  321. const node = nodes[i];
  322. const dy = (y - node.y0) * alpha;
  323. if (dy > 1e-6)
  324. (node.y0 += dy), (node.y1 += dy);
  325. y = node.y1 + py;
  326. }
  327. }
  328. // Push any overlapping nodes up.
  329. function resolveCollisionsBottomToTop(nodes, y, i, alpha) {
  330. for (; i >= 0; --i) {
  331. const node = nodes[i];
  332. const dy = (node.y1 - y) * alpha;
  333. if (dy > 1e-6)
  334. (node.y0 -= dy), (node.y1 -= dy);
  335. y = node.y0 - py;
  336. }
  337. }
  338. function reorderNodeLinks({ sourceLinks, targetLinks }) {
  339. if (linkSort === undefined) {
  340. for (const { source: { sourceLinks }, } of targetLinks) {
  341. sourceLinks.sort(ascendingTargetBreadth);
  342. }
  343. for (const { target: { targetLinks }, } of sourceLinks) {
  344. targetLinks.sort(ascendingSourceBreadth);
  345. }
  346. }
  347. }
  348. function reorderLinks(nodes) {
  349. if (linkSort === undefined) {
  350. for (const { sourceLinks, targetLinks } of nodes) {
  351. sourceLinks.sort(ascendingTargetBreadth);
  352. targetLinks.sort(ascendingSourceBreadth);
  353. }
  354. }
  355. }
  356. // Returns the target.y0 that would produce an ideal link from source to target.
  357. function targetTop(source, target) {
  358. let y = source.y0 - ((source.sourceLinks.length - 1) * py) / 2;
  359. for (const { target: node, width } of source.sourceLinks) {
  360. if (node === target)
  361. break;
  362. y += width + py;
  363. }
  364. for (const { source: node, width } of target.targetLinks) {
  365. if (node === source)
  366. break;
  367. y -= width;
  368. }
  369. return y;
  370. }
  371. // Returns the source.y0 that would produce an ideal link from source to target.
  372. function sourceTop(source, target) {
  373. let y = target.y0 - ((target.targetLinks.length - 1) * py) / 2;
  374. for (const { source: node, width } of target.targetLinks) {
  375. if (node === source)
  376. break;
  377. y += width + py;
  378. }
  379. for (const { target: node, width } of source.sourceLinks) {
  380. if (node === target)
  381. break;
  382. y -= width;
  383. }
  384. return y;
  385. }
  386. return sankey;
  387. }
  388. //# sourceMappingURL=sankey.js.map