sankey: implement node grouping via mouse selection

src/components/modebar/buttons.js

@@ -513,6 +513,17 @@ function toggleHover(gd) {
     Registry.call('_guiRelayout', gd, 'hovermode', newHover);
 }
 
+modeBarButtons.resetSankeyGroup = {
+    name: 'resetSankeyGroup',
+    title: function(gd) { return _(gd, 'Ungroup all nodes'); },
+    icon: Icons.home,
+    click: function(gd) {
+        Registry.call('restyle', gd, {
+            'node.groups': [[]],
+        });
+    }
+};
+
 // buttons when more then one plot types are present
 
 modeBarButtons.toggleHover = {

src/components/modebar/manage.js

@@ -86,6 +86,7 @@ function getButtonGroups(gd, buttonsToRemove, buttonsToAdd, showSendToCloud) {
     var hasTernary = fullLayout._has('ternary');
     var hasMapbox = fullLayout._has('mapbox');
     var hasPolar = fullLayout._has('polar');
+    var hasSankey = fullLayout._has('sankey');
     var allAxesFixed = areAllAxesFixed(fullLayout);
 
     var groups = [];
@@ -139,6 +140,9 @@ function getButtonGroups(gd, buttonsToRemove, buttonsToAdd, showSendToCloud) {
     else if(hasPie) {
         hoverGroup = ['hoverClosestPie'];
     }
+    else if(hasSankey) {
+        hoverGroup = ['resetSankeyGroup'];
+    }
     else { // hasPolar, hasTernary
         // always show at least one hover icon.
         hoverGroup = ['toggleHover'];

src/plots/cartesian/select.js

@@ -48,6 +48,8 @@ function prepSelect(e, startX, startY, dragOptions, mode) {
     var allAxes = dragOptions.xaxes.concat(dragOptions.yaxes);
     var subtract = e.altKey;
 
+    var doneFnCompleted = dragOptions.doneFnCompleted;
+
     var filterPoly, selectionTester, mergedPolygons, currentPolygon;
     var i, searchInfo, eventData;
 
@@ -284,6 +286,8 @@ function prepSelect(e, startX, startY, dragOptions, mode) {
                 dragOptions.mergedPolygons.length = 0;
                 [].push.apply(dragOptions.mergedPolygons, mergedPolygons);
             }
+
+            doneFnCompleted(selection);
         });
     };
 }
@@ -519,6 +523,11 @@ function determineSearchTraces(gd, xAxes, yAxes, subplot) {
             var info = createSearchInfo(trace._module, cd, xAxes[0], yAxes[0]);
             info.scene = gd._fullLayout._splomScenes[trace.uid];
             searchTraces.push(info);
+        } else if(
+          trace.type === 'sankey'
+        ) {
+            var sankeyInfo = createSearchInfo(trace._module, cd, xAxes[0], yAxes[0]);
+            searchTraces.push(sankeyInfo);
         } else {
             if(xAxisIds.indexOf(trace.xaxis) === -1) continue;
             if(yAxisIds.indexOf(trace.yaxis) === -1) continue;

src/traces/sankey/base_plot.js

@@ -13,6 +13,12 @@ var getModuleCalcData = require('../../plots/get_data').getModuleCalcData;
 var plot = require('./plot');
 var fxAttrs = require('../../components/fx/layout_attributes');
 
+var setCursor = require('../../lib/setcursor');
+var dragElement = require('../../components/dragelement');
+var prepSelect = require('../../plots/cartesian/select').prepSelect;
+var Lib = require('../../lib');
+var Registry = require('../../registry');
+
 var SANKEY = 'sankey';
 
 exports.name = SANKEY;
@@ -34,3 +40,94 @@ exports.clean = function(newFullData, newFullLayout, oldFullData, oldFullLayout)
         oldFullLayout._paperdiv.selectAll('.sankey').remove();
     }
 };
+
+exports.updateFx = function(gd) {
+    for(var i = 0; i < gd._fullData.length; i++) {
+        subplotUpdateFx(gd, i);
+    }
+};
+
+var dragOptions = [];
+function subplotUpdateFx(gd, index) {
+    var i = index;
+    var fullData = gd._fullData[i];
+    var fullLayout = gd._fullLayout;
+
+    var dragMode = fullLayout.dragmode;
+    var cursor = fullLayout.dragmode === 'pan' ? 'move' : 'crosshair';
+    var bgRect = fullData._bgRect;
+
+    setCursor(bgRect, cursor);
+
+    var xaxis = {
+        _id: 'x',
+        c2p: function(v) { return v; },
+        _offset: fullData._sankey.translateX,
+        _length: fullData._sankey.width
+    };
+    var yaxis = {
+        _id: 'y',
+        c2p: function(v) { return v; },
+        _offset: fullData._sankey.translateY,
+        _length: fullData._sankey.height
+    };
+
+    // Note: dragOptions is needed to be declared for all dragmodes because
+    // it's the object that holds persistent selection state.
+    dragOptions[i] = {
+        gd: gd,
+        element: bgRect.node(),
+        plotinfo: {
+            id: i,
+            xaxis: xaxis,
+            yaxis: yaxis,
+            fillRangeItems: Lib.noop
+        },
+        subplot: i,
+        // create mock x/y axes for hover routine
+        xaxes: [xaxis],
+        yaxes: [yaxis],
+        doneFnCompleted: function(selection) {
+            var newGroups;
+            var oldGroups = gd._fullData[i].node.groups.slice();
+            var newGroup = [];
+
+            function findNode(pt) {
+                return gd._fullData[i]._sankey.graph.nodes.find(function(n) {
+                    return n.pointNumber === pt;
+                });
+            }
+
+            for(var j = 0; j < selection.length; j++) {
+                var node = findNode(selection[j].pointNumber);
+                if(!node) continue;
+
+                // If the node represents a group
+                if(node.group) {
+                    // Add all its children to the current selection
+                    for(var k = 0; k < node.childrenNodes.length; k++) {
+                        newGroup.push(node.childrenNodes[k].pointNumber);
+                    }
+                    // Flag group for removal from existing list of groups
+                    oldGroups[node.pointNumber - fullData.node._count] = false;
+                } else {
+                    newGroup.push(node.pointNumber);
+                }
+            }
+
+            newGroups = oldGroups
+                .filter(function(g) { return g; })
+                .concat([newGroup]);
+
+            Registry.call('_guiRestyle', gd, {
+                'node.groups': [ newGroups ]
+            }, i).catch(Lib.noop); // TODO will this ever fail?
+        }
+    };
+
+    dragOptions[i].prepFn = function(e, startX, startY) {
+        prepSelect(e, startX, startY, dragOptions[i], dragMode);
+    };
+
+    dragElement.init(dragOptions[i]);
+}

src/traces/sankey/calc.js

@@ -40,6 +40,7 @@ function convertToD3Sankey(trace) {
         if(linkSpec.target[i] > maxNodeId) maxNodeId = linkSpec.target[i];
     }
     var nodeCount = maxNodeId + 1;
+    trace.node._count = nodeCount;
 
     // Group nodes
     var j;

src/traces/sankey/index.js

@@ -18,6 +18,7 @@ Plot.plot = require('./plot');
 Plot.moduleType = 'trace';
 Plot.name = 'sankey';
 Plot.basePlotModule = require('./base_plot');
+Plot.selectPoints = require('./select.js');
 Plot.categories = ['noOpacity'];
 Plot.meta = {
     description: [

src/traces/sankey/render.js

@@ -836,6 +836,23 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
         .style('pointer-events', 'auto')
         .attr('transform', sankeyTransform);
 
+    sankey.each(function(d, i) {
+        gd._fullData[i]._sankey = d;
+
+        // Draw dragbox
+        Lib.ensureSingle(gd._fullLayout._draggers, 'rect', 'bg-' + i, function(el) {
+            el
+              .style('pointer-events', 'all')
+              .attr('width', d.width)
+              .attr('height', d.height)
+              .attr('x', d.translateX)
+              .attr('y', d.translateY)
+              .style({fill: 'transparent', 'stroke-width': 0});
+
+            gd._fullData[i]._bgRect = el;
+        });
+    });
+
     sankey.transition()
         .ease(c.ease).duration(c.duration)
         .attr('transform', sankeyTransform);

src/traces/sankey/select.js

@@ -0,0 +1,36 @@
+/**
+* Copyright 2012-2019, Plotly, Inc.
+* All rights reserved.
+*
+* This source code is licensed under the MIT license found in the
+* LICENSE file in the root directory of this source tree.
+*/
+
+'use strict';
+
+module.exports = function selectPoints(searchInfo, selectionTester) {
+    var cd = searchInfo.cd;
+    var selection = [];
+    var fullData = cd[0].trace;
+
+    var nodes = fullData._sankey.graph.nodes;
+
+    for(var i = 0; i < nodes.length; i++) {
+        var node = nodes[i];
+        if(node.partOfGroup) continue; // Those are invisible
+
+        // TODO: decide on selection criteria, using centroid for now
+        var pos = [(node.x0 + node.x1) / 2, (node.y0 + node.y1) / 2];
+
+        // Swap x and y if trace is vertical
+        if(fullData.orientation === 'v') pos.reverse();
+
+        if(selectionTester.contains(pos, false, i, searchInfo)) {
+            selection.push({
+                pointNumber: node.pointNumber
+                // TODO: add eventData
+            });
+        }
+    }
+    return selection;
+};