sankey: grouping with animations

Author: antoinerg

src/snapshot/helpers.js

@@ -15,6 +15,7 @@ exports.getDelay = function(fullLayout) {
     return (
         fullLayout._has('gl3d') ||
         fullLayout._has('gl2d') ||
+        fullLayout._has('sankey') ||
         fullLayout._has('mapbox')
     ) ? 500 : 0;
 };

src/traces/sankey/attributes.js

@@ -89,6 +89,16 @@ var attrs = module.exports = overrideAll({
             role: 'info',
             description: 'The shown name of the node.'
         },
+        groups: {
+            valType: 'data_array',
+            dflt: [],
+            role: 'calc',
+            description: [
+                'Groups of nodes.',
+                'Each group is defined by an array with the indices of the nodes it contains.',
+                'Multiple groups can be specified.'
+            ].join(' ')
+        },
         color: {
             valType: 'color',
             role: 'style',

src/traces/sankey/calc.js

@@ -17,8 +17,10 @@ var isIndex = Lib.isIndex;
 var Colorscale = require('../../components/colorscale');
 
 function convertToD3Sankey(trace) {
-    var nodeSpec = trace.node;
-    var linkSpec = trace.link;
+    // var nodeSpec = trace.node;
+    // var linkSpec = trace.link;
+    var nodeSpec = Lib.extendDeep({}, trace.node);
+    var linkSpec = Lib.extendDeep({}, trace.link);
 
     var links = [];
     var hasLinkColorArray = isArrayOrTypedArray(linkSpec.color);
@@ -34,7 +36,32 @@ function convertToD3Sankey(trace) {
         components[cscale.label] = scale;
     }
 
-    var nodeCount = nodeSpec.label.length;
+    var maxNodeId = 0;
+    for(i = 0; i < linkSpec.value.length; i++) {
+        if(linkSpec.source[i] > maxNodeId) maxNodeId = linkSpec.source[i];
+        if(linkSpec.target[i] > maxNodeId) maxNodeId = linkSpec.target[i];
+    }
+    var nodeCount = maxNodeId + 1;
+
+    // Group nodes
+    var j;
+    var groups = trace.node.groups;
+    var groupLookup = {};
+    for(i = 0; i < groups.length; i++) {
+        var group = groups[i];
+        // Build a lookup table to quickly find in which group a node is
+        if(Array.isArray(group)) {
+            for(j = 0; j < group.length; j++) {
+                var nodeIndex = group[j];
+                var groupIndex = nodeCount + i;
+                groupLookup[nodeIndex] = groupIndex;
+            }
+        } else {
+            Lib.warn('node.groups must be an array, default to empty array []');
+        }
+    }
+
+    // Process links
     for(i = 0; i < linkSpec.value.length; i++) {
         var val = linkSpec.value[i];
         // remove negative values, but keep zeros with special treatment
@@ -44,6 +71,22 @@ function convertToD3Sankey(trace) {
             continue;
         }
 
+        // Remove links that are within the same group
+        if(groupLookup.hasOwnProperty(source) && groupLookup.hasOwnProperty(target) && groupLookup[source] === groupLookup[target]) {
+            continue;
+        }
+
+        // if link targets a node in the group, relink target to that group
+        if(groupLookup.hasOwnProperty(target)) {
+            target = groupLookup[target];
+        }
+
+        // if link originates from a node in a group, relink source to that group
+        // if(group.indexOf(source) !== -1) {
+        if(groupLookup.hasOwnProperty(source)) {
+            source = groupLookup[source];
+        }
+
         source = +source;
         target = +target;
         linkedNodes[source] = linkedNodes[target] = true;
@@ -65,34 +108,29 @@ function convertToD3Sankey(trace) {
         });
     }
 
+    // Process nodes
+    var totalCount = nodeCount + groups.length;
     var hasNodeColorArray = isArrayOrTypedArray(nodeSpec.color);
     var nodes = [];
-    var removedNodes = false;
-    var nodeIndices = {};
-
-    for(i = 0; i < nodeCount; i++) {
-        if(linkedNodes[i]) {
-            var l = nodeSpec.label[i];
-            nodeIndices[i] = nodes.length;
-            nodes.push({
-                pointNumber: i,
-                label: l,
-                color: hasNodeColorArray ? nodeSpec.color[i] : nodeSpec.color
-            });
-        } else removedNodes = true;
-    }
+    for(i = 0; i < totalCount; i++) {
+        if(!linkedNodes[i]) continue;
+        var l = nodeSpec.label[i];
 
-    // need to re-index links now, since we didn't put all the nodes in
-    if(removedNodes) {
-        for(i = 0; i < links.length; i++) {
-            links[i].source = nodeIndices[links[i].source];
-            links[i].target = nodeIndices[links[i].target];
-        }
+        nodes.push({
+            group: (i > nodeCount - 1),
+            pointNumber: i,
+            label: l,
+            color: hasNodeColorArray ? nodeSpec.color[i] : nodeSpec.color
+        });
     }
 
     return {
         links: links,
-        nodes: nodes
+        nodes: nodes,
+
+        // Data structure for groups
+        groups: groups,
+        groupLookup: groupLookup
     };
 }
 
@@ -130,6 +168,10 @@ module.exports = function calc(gd, trace) {
     return wrap({
         circular: circular,
         _nodes: result.nodes,
-        _links: result.links
+        _links: result.links,
+
+        // Data structure for grouping
+        _groups: result.groups,
+        _groupLookup: result.groupLookup,
     });
 };

src/traces/sankey/constants.js

@@ -15,8 +15,8 @@ module.exports = {
     sankeyIterations: 50,
     forceIterations: 5,
     forceTicksPerFrame: 10,
-    duration: 500,
-    ease: 'cubic-in-out',
+    duration: 350,
+    ease: 'quart-in-out',
     cn: {
         sankey: 'sankey',
         sankeyLinks: 'sankey-links',

src/traces/sankey/defaults.js

@@ -32,6 +32,7 @@ module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout
         return Lib.coerce(nodeIn, nodeOut, attributes.node, attr, dflt);
     }
     coerceNode('label');
+    coerceNode('groups');
     coerceNode('pad');
     coerceNode('thickness');
     coerceNode('line.color');

src/traces/sankey/render.js

@@ -45,10 +45,7 @@ function sankeyModel(layout, d, traceIndex) {
     if(circular) {
         sankey = d3SankeyCircular
             .sankeyCircular()
-            .circularLinkGap(0)
-            .nodeId(function(d) {
-                return d.pointNumber;
-            });
+            .circularLinkGap(0);
     } else {
         sankey = d3Sankey.sankey();
     }
@@ -58,6 +55,9 @@ function sankeyModel(layout, d, traceIndex) {
       .size(horizontal ? [width, height] : [height, width])
       .nodeWidth(nodeThickness)
       .nodePadding(nodePad)
+      .nodeId(function(d) {
+          return d.pointNumber;
+      })
       .nodes(nodes)
       .links(links);
 
@@ -67,6 +67,30 @@ function sankeyModel(layout, d, traceIndex) {
         Lib.warn('node.pad was reduced to ', sankey.nodePadding(), ' to fit within the figure.');
     }
 
+    // Create transient nodes for animations
+    Object.keys(calcData._groupLookup).forEach(function(nodePointNumber) {
+        var groupIndex = parseInt(calcData._groupLookup[nodePointNumber]);
+
+        var groupingNode;
+        for(var i = 0; i < graph.nodes.length; i++) {
+            if(graph.nodes[i].pointNumber === groupIndex) {
+                groupingNode = graph.nodes[i];
+                break;
+            }
+        }
+
+        graph.nodes.push({
+            pointNumber: parseInt(nodePointNumber),
+            x0: groupingNode.x0,
+            x1: groupingNode.x1,
+            y0: groupingNode.y0,
+            y1: groupingNode.y1,
+            partOfGroup: true,
+            sourceLinks: [],
+            targetLinks: []
+        });
+    });
+
     function computeLinkConcentrations() {
         var i, j, k;
         for(i = 0; i < graph.nodes.length; i++) {
@@ -343,7 +367,7 @@ function linkPath() {
     return path;
 }
 
-function nodeModel(d, n, i) {
+function nodeModel(d, n) {
     var tc = tinycolor(n.color);
     var zoneThicknessPad = c.nodePadAcross;
     var zoneLengthPad = d.nodePad / 2;
@@ -352,8 +376,11 @@ function nodeModel(d, n, i) {
     var visibleThickness = n.dx;
     var visibleLength = Math.max(0.5, n.dy);
 
-    var basicKey = n.label;
-    var key = basicKey + '__' + i;
+    var key = 'node_' + n.pointNumber;
+    // If it's a group, it's mutable and should be unique
+    if(n.group) {
+        key = 'group_' + Math.floor(1e12 * (1 + Math.random()));
+    }
 
     // for event data
     n.trace = d.trace;
@@ -362,6 +389,8 @@ function nodeModel(d, n, i) {
     return {
         index: n.pointNumber,
         key: key,
+        partOfGroup: n.partOfGroup || false,
+        group: n.group,
         traceId: d.key,
         node: n,
         nodePad: d.nodePad,
@@ -540,7 +569,10 @@ function attachDragHandler(sankeyNode, sankeyLink, callbacks) {
 function attachForce(sankeyNode, forceKey, d) {
     // Attach force to nodes in the same column (same x coordinate)
     switchToForceFormat(d.graph.nodes);
-    var nodes = d.graph.nodes.filter(function(n) {return n.originalX === d.node.originalX;});
+    var nodes = d.graph.nodes
+        .filter(function(n) {return n.originalX === d.node.originalX;})
+        // Filter out children
+        .filter(function(n) {return !n.partOfGroup;});
     d.forceLayouts[forceKey] = d3Force.forceSimulation(nodes)
         .alphaDecay(0)
         .force('collide', d3Force.forceCollide()
@@ -683,7 +715,6 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
     sankeyLink
           .enter().append('path')
           .classed(c.cn.sankeyLink, true)
-          .attr('d', linkPath())
           .call(attachPointerEvents, sankey, callbacks.linkEvents);
 
     sankeyLink
@@ -701,13 +732,17 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
         })
         .style('stroke-width', function(d) {
             return salientEnough(d) ? d.linkLineWidth : 1;
-        });
+        })
+        .attr('d', linkPath());
 
-    sankeyLink.transition()
-       .ease(c.ease).duration(c.duration)
-       .attr('d', linkPath());
+    sankeyLink
+        .style('opacity', 0)
+        .transition()
+        .ease(c.ease).duration(c.duration)
+        .style('opacity', 1);
 
-    sankeyLink.exit().transition()
+    sankeyLink.exit()
+        .transition()
         .ease(c.ease).duration(c.duration)
         .style('opacity', 0)
         .remove();
@@ -733,24 +768,26 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
             var nodes = d.graph.nodes;
             persistOriginalPlace(nodes);
             return nodes
-                .filter(function(n) {return n.value;})
-                .map(nodeModel.bind(null, d));
+              .map(nodeModel.bind(null, d));
         }, keyFun);
 
     sankeyNode.enter()
         .append('g')
         .classed(c.cn.sankeyNode, true)
         .call(updateNodePositions)
-        .call(attachPointerEvents, sankey, callbacks.nodeEvents);
+        .style('opacity', 0);
 
     sankeyNode
+        .call(attachPointerEvents, sankey, callbacks.nodeEvents)
         .call(attachDragHandler, sankeyLink, callbacks); // has to be here as it binds sankeyLink
 
     sankeyNode.transition()
         .ease(c.ease).duration(c.duration)
-        .call(updateNodePositions);
+        .call(updateNodePositions)
+        .style('opacity', function(n) { return n.partOfGroup ? 0 : 1;});
 
-    sankeyNode.exit().transition()
+    sankeyNode.exit()
+        .transition()
         .ease(c.ease).duration(c.duration)
         .style('opacity', 0)
         .remove();