Sankey: group nodes

src/traces/sankey/attributes.js

@@ -89,6 +89,19 @@ var attrs = module.exports = overrideAll({
             role: 'info',
             description: 'The shown name of the node.'
         },
+        groups: {
+            valType: 'info_array',
+            dimensions: 2,
+            freeLength: true,
+            dflt: [],
+            items: {valType: 'number', editType: 'calc'},
+            role: 'info',
+            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

@@ -34,7 +34,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
+        for(j = 0; j < group.length; j++) {
+            var nodeIndex = group[j];
+            var groupIndex = nodeCount + i;
+            if(groupLookup.hasOwnProperty(nodeIndex)) {
+                Lib.warn('Node ' + nodeIndex + ' is already part of a group.');
+            } else {
+                groupLookup[nodeIndex] = groupIndex;
+            }
+        }
+    }
+
+    // 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 +69,21 @@ 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(groupLookup.hasOwnProperty(source)) {
+            source = groupLookup[source];
+        }
+
         source = +source;
         target = +target;
         linkedNodes[source] = linkedNodes[target] = true;
@@ -65,34 +105,30 @@ 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),
+            childrenNodes: [],
+            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 +166,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

@@ -16,7 +16,7 @@ module.exports = {
     forceIterations: 5,
     forceTicksPerFrame: 10,
     duration: 500,
-    ease: 'cubic-in-out',
+    ease: 'linear',
     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,33 @@ function sankeyModel(layout, d, traceIndex) {
         Lib.warn('node.pad was reduced to ', sankey.nodePadding(), ' to fit within the figure.');
     }
 
+    // Create transient nodes for animations
+    for(var nodePointNumber in calcData._groupLookup) {
+        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;
+            }
+        }
+
+        var child = {
+            pointNumber: parseInt(nodePointNumber),
+            x0: groupingNode.x0,
+            x1: groupingNode.x1,
+            y0: groupingNode.y0,
+            y1: groupingNode.y1,
+            partOfGroup: true,
+            sourceLinks: [],
+            targetLinks: []
+        };
+
+        graph.nodes.unshift(child);
+        groupingNode.childrenNodes.unshift(child);
+    }
+
     function computeLinkConcentrations() {
         var i, j, k;
         for(i = 0; i < graph.nodes.length; i++) {
@@ -137,7 +164,7 @@ function sankeyModel(layout, d, traceIndex) {
         circular: circular,
         key: traceIndex,
         trace: trace,
-        guid: Math.floor(1e12 * (1 + Math.random())),
+        guid: Lib.randstr(),
         horizontal: horizontal,
         width: width,
         height: height,
@@ -343,7 +370,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 +379,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 = Lib.randstr();
+    }
 
     // for event data
     n.trace = d.trace;
@@ -362,6 +392,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,
@@ -445,19 +477,19 @@ function attachPointerEvents(selection, sankey, eventSet) {
     selection
         .on('.basic', null) // remove any preexisting handlers
         .on('mouseover.basic', function(d) {
-            if(!d.interactionState.dragInProgress) {
+            if(!d.interactionState.dragInProgress && !d.partOfGroup) {
                 eventSet.hover(this, d, sankey);
                 d.interactionState.hovered = [this, d];
             }
         })
         .on('mousemove.basic', function(d) {
-            if(!d.interactionState.dragInProgress) {
+            if(!d.interactionState.dragInProgress && !d.partOfGroup) {
                 eventSet.follow(this, d);
                 d.interactionState.hovered = [this, d];
             }
         })
         .on('mouseout.basic', function(d) {
-            if(!d.interactionState.dragInProgress) {
+            if(!d.interactionState.dragInProgress && !d.partOfGroup) {
                 eventSet.unhover(this, d, sankey);
                 d.interactionState.hovered = false;
             }
@@ -467,7 +499,7 @@ function attachPointerEvents(selection, sankey, eventSet) {
                 eventSet.unhover(this, d, sankey);
                 d.interactionState.hovered = false;
             }
-            if(!d.interactionState.dragInProgress) {
+            if(!d.interactionState.dragInProgress && !d.partOfGroup) {
                 eventSet.select(this, d, sankey);
             }
         });
@@ -530,6 +562,10 @@ function attachDragHandler(sankeyNode, sankeyLink, callbacks) {
 
         .on('dragend', function(d) {
             d.interactionState.dragInProgress = false;
+            for(var i = 0; i < d.node.childrenNodes.length; i++) {
+                d.node.childrenNodes[i].x = d.node.x;
+                d.node.childrenNodes[i].y = d.node.y;
+            }
         });
 
     sankeyNode
@@ -540,7 +576,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()
@@ -639,6 +678,11 @@ function switchToSankeyFormat(nodes) {
 
 // scene graph
 module.exports = function(gd, svg, calcData, layout, callbacks) {
+    // To prevent animation on first render
+    var firstRender = false;
+    Lib.ensureSingle(gd._fullLayout._infolayer, 'g', 'first-render', function() {
+        firstRender = true;
+    });
 
     var styledData = calcData
             .filter(function(d) {return unwrap(d).trace.visible;})
@@ -683,7 +727,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 +744,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', function() { return (gd._context.staticPlot || firstRender) ? 1 : 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 +780,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', function(n) { return ((gd._context.staticPlot || firstRender) && !n.partOfGroup) ? 1 : 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();