Merge pull request #3712 from plotly/sankey-selection2

antoinerg · web-flow · commit 8fd213311815 · 2019-04-08T18:48:15.000-04:00
sankey: implement node grouping via mouse selection
diff --git a/src/components/modebar/buttons.js b/src/components/modebar/buttons.js
@@ -513,6 +513,26 @@ function toggleHover(gd) {
     Registry.call('_guiRelayout', gd, 'hovermode', newHover);
 }
 
+modeBarButtons.resetViewSankey = {
+    name: 'resetSankeyGroup',
+    title: function(gd) { return _(gd, 'Reset view'); },
+    icon: Icons.home,
+    click: function(gd) {
+        var aObj = {
+            'node.groups': [],
+            'node.x': [],
+            'node.y': []
+        };
+        for(var i = 0; i < gd._fullData.length; i++) {
+            var viewInitial = gd._fullData[i]._viewInitial;
+            aObj['node.groups'].push(viewInitial.node.groups.slice());
+            aObj['node.x'].push(viewInitial.node.x.slice());
+            aObj['node.y'].push(viewInitial.node.y.slice());
+        }
+        Registry.call('restyle', gd, aObj);
+    }
+};
+
 // buttons when more then one plot types are present
 
 modeBarButtons.toggleHover = {
diff --git a/src/components/modebar/manage.js b/src/components/modebar/manage.js
@@ -142,6 +142,7 @@ function getButtonGroups(gd, buttonsToRemove, buttonsToAdd, showSendToCloud) {
     }
     else if(hasSankey) {
         hoverGroup = ['hoverClosestCartesian', 'hoverCompareCartesian'];
+        resetGroup = ['resetViewSankey'];
     }
     else { // hasPolar, hasTernary
         // always show at least one hover icon.
diff --git a/src/plots/cartesian/select.js b/src/plots/cartesian/select.js
@@ -49,6 +49,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;
 
@@ -285,6 +287,8 @@ function prepSelect(e, startX, startY, dragOptions, mode) {
                 dragOptions.mergedPolygons.length = 0;
                 [].push.apply(dragOptions.mergedPolygons, mergedPolygons);
             }
+
+            doneFnCompleted(selection);
         });
     };
 }
@@ -520,6 +524,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;
diff --git a/src/traces/sankey/base_plot.js b/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;
@@ -24,6 +30,7 @@ exports.baseLayoutAttrOverrides = overrideAll({
 exports.plot = function(gd) {
     var calcData = getModuleCalcData(gd.calcdata, SANKEY)[0];
     plot(gd, calcData);
+    exports.updateFx(gd);
 };
 
 exports.clean = function(newFullData, newFullLayout, oldFullData, oldFullLayout) {
@@ -32,5 +39,99 @@ exports.clean = function(newFullData, newFullLayout, oldFullData, oldFullLayout)
 
     if(hadPlot && !hasPlot) {
         oldFullLayout._paperdiv.selectAll('.sankey').remove();
+        oldFullLayout._paperdiv.selectAll('.bgsankey').remove();
     }
 };
+
+exports.updateFx = function(gd) {
+    for(var i = 0; i < gd._fullData.length; i++) {
+        subplotUpdateFx(gd, i);
+    }
+};
+
+function subplotUpdateFx(gd, index) {
+    var trace = gd._fullData[index];
+    var fullLayout = gd._fullLayout;
+
+    var dragMode = fullLayout.dragmode;
+    var cursor = fullLayout.dragmode === 'pan' ? 'move' : 'crosshair';
+    var bgRect = trace._bgRect;
+
+    if(dragMode === 'pan' || dragMode === 'zoom') return;
+
+    setCursor(bgRect, cursor);
+
+    var xaxis = {
+        _id: 'x',
+        c2p: Lib.identity,
+        _offset: trace._sankey.translateX,
+        _length: trace._sankey.width
+    };
+    var yaxis = {
+        _id: 'y',
+        c2p: Lib.identity,
+        _offset: trace._sankey.translateY,
+        _length: trace._sankey.height
+    };
+
+    // Note: dragOptions is needed to be declared for all dragmodes because
+    // it's the object that holds persistent selection state.
+    var dragOptions = {
+        gd: gd,
+        element: bgRect.node(),
+        plotinfo: {
+            id: index,
+            xaxis: xaxis,
+            yaxis: yaxis,
+            fillRangeItems: Lib.noop
+        },
+        subplot: index,
+        // create mock x/y axes for hover routine
+        xaxes: [xaxis],
+        yaxes: [yaxis],
+        doneFnCompleted: function(selection) {
+            var traceNow = gd._fullData[index];
+            var newGroups;
+            var oldGroups = traceNow.node.groups.slice();
+            var newGroup = [];
+
+            function findNode(pt) {
+                var nodes = traceNow._sankey.graph.nodes;
+                for(var i = 0; i < nodes.length; i++) {
+                    if(nodes[i].pointNumber === pt) return nodes[i];
+                }
+            }
+
+            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 - traceNow.node._count] = false;
+                } else {
+                    newGroup.push(node.pointNumber);
+                }
+            }
+
+            newGroups = oldGroups
+                .filter(Boolean)
+                .concat([newGroup]);
+
+            Registry.call('_guiRestyle', gd, {
+                'node.groups': [ newGroups ]
+            }, index);
+        }
+    };
+
+    dragOptions.prepFn = function(e, startX, startY) {
+        prepSelect(e, startX, startY, dragOptions, dragMode);
+    };
+
+    dragElement.init(dragOptions);
+}
diff --git a/src/traces/sankey/calc.js b/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;
diff --git a/src/traces/sankey/index.js b/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: [
diff --git a/src/traces/sankey/plot.js b/src/traces/sankey/plot.js
@@ -127,6 +127,20 @@ module.exports = function plot(gd, calcData) {
     var svg = fullLayout._paper;
     var size = fullLayout._size;
 
+    // stash initial view
+    for(var i = 0; i < calcData.length; i++) {
+        if(!gd._fullData[i]._viewInitial) {
+            var node = gd._fullData[i].node;
+            gd._fullData[i]._viewInitial = {
+                node: {
+                    groups: node.groups.slice(),
+                    x: node.x.slice(),
+                    y: node.y.slice()
+                }
+            };
+        }
+    }
+
     var linkSelect = function(element, d) {
         var evt = d.link;
         evt.originalEvent = d3.event;
diff --git a/src/traces/sankey/render.js b/src/traces/sankey/render.js
@@ -836,6 +836,25 @@ 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;
+        // Create dragbox if missing
+        var dragboxClassName = 'bgsankey-' + d.trace.uid + '-' + i;
+        Lib.ensureSingle(gd._fullLayout._draggers, 'rect', dragboxClassName);
+
+        gd._fullData[i]._bgRect = d3.select('.' + dragboxClassName);
+
+        // Style dragbox
+        gd._fullData[i]._bgRect
+          .style('pointer-events', 'all')
+          .attr('width', d.width)
+          .attr('height', d.height)
+          .attr('x', d.translateX)
+          .attr('y', d.translateY)
+          .classed('bgsankey', true)
+          .style({fill: 'transparent', 'stroke-width': 0});
+    });
+
     sankey.transition()
         .ease(c.ease).duration(c.duration)
         .attr('transform', sankeyTransform);
@@ -925,7 +944,8 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
         .call(attachPointerEvents, sankey, callbacks.nodeEvents)
         .call(attachDragHandler, sankeyLink, callbacks, gd); // has to be here as it binds sankeyLink
 
-    sankeyNode.transition()
+    sankeyNode
+        .transition()
         .ease(c.ease).duration(c.duration)
         .call(updateNodePositions)
         .style('opacity', function(n) { return n.partOfGroup ? 0 : 1;});
diff --git a/src/traces/sankey/select.js b/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
+
+        // Position of node's centroid
+        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 && selectionTester.contains(pos, false, i, searchInfo)) {
+            selection.push({
+                pointNumber: node.pointNumber
+                // TODO: add eventData
+            });
+        }
+    }
+    return selection;
+};
diff --git a/test/jasmine/tests/sankey_test.js b/test/jasmine/tests/sankey_test.js
@@ -23,6 +23,7 @@ var defaultColors = require('@src/components/color/attributes').defaults;
 var drag = require('../assets/drag');
 var checkOverlap = require('../assets/check_overlap');
 var delay = require('../assets/delay');
+var selectButton = require('../assets/modebar_button');
 
 describe('sankey tests', function() {
     'use strict';
@@ -379,6 +380,20 @@ describe('sankey tests', function() {
                 });
         });
 
+        it('Plotly.deleteTraces removes draggers', function(done) {
+            var mockCopy = Lib.extendDeep({}, mock);
+            Plotly.plot(gd, mockCopy)
+                .then(function() {
+                    expect(document.getElementsByClassName('bgsankey').length).toBe(1);
+                    return Plotly.deleteTraces(gd, [0]);
+                })
+                .then(function() {
+                    expect(document.getElementsByClassName('bgsankey').length).toBe(0);
+                })
+                .catch(failTest)
+                .then(done);
+        });
+
         it('Plotly.plot does not show Sankey if \'visible\' is false', function(done) {
             var mockCopy = Lib.extendDeep({}, mock);
 
@@ -546,6 +561,52 @@ describe('sankey tests', function() {
             .catch(failTest)
             .then(done);
         });
+
+        it('resets each subplot to its initial view (ie. x, y groups) via modebar button', function(done) {
+            var mockCopy = Lib.extendDeep({}, require('@mocks/sankey_subplots_circular'));
+
+            // Set initial view
+            mockCopy.data[0].node.x = [0.25];
+            mockCopy.data[0].node.y = [0.25];
+
+            mockCopy.data[0].node.groups = [];
+            mockCopy.data[1].node.groups = [[2, 3]];
+
+            Plotly.plot(gd, mockCopy)
+            .then(function() {
+                expect(gd._fullData[0].node.groups).toEqual([]);
+                expect(gd._fullData[1].node.groups).toEqual([[2, 3]]);
+
+                // Change groups
+                return Plotly.restyle(gd, {
+                    'node.groups': [[[1, 2]], [[]]],
+                    'node.x': [[0.1]],
+                    'node.y': [[0.1]]
+                });
+            })
+            .then(function() {
+                // Check current state
+                expect(gd._fullData[0].node.x).toEqual([0.1]);
+                expect(gd._fullData[0].node.y).toEqual([0.1]);
+
+                expect(gd._fullData[0].node.groups).toEqual([[1, 2]]);
+                expect(gd._fullData[1].node.groups).toEqual([[]]);
+
+                // Click reset
+                var resetButton = selectButton(gd._fullLayout._modeBar, 'resetViewSankey');
+                resetButton.click();
+            })
+            .then(function() {
+                // Check we are back to initial view
+                expect(gd._fullData[0].node.x).toEqual([0.25]);
+                expect(gd._fullData[0].node.y).toEqual([0.25]);
+
+                expect(gd._fullData[0].node.groups).toEqual([]);
+                expect(gd._fullData[1].node.groups).toEqual([[2, 3]]);
+            })
+            .catch(failTest)
+            .then(done);
+        });
     });
 
     describe('Test hover/click interactions:', function() {
diff --git a/test/jasmine/tests/select_test.js b/test/jasmine/tests/select_test.js

Original file line number	Diff line number	Diff line change
`@@ -142,6 +142,7 @@ function getButtonGroups(gd, buttonsToRemove, buttonsToAdd, showSendToCloud) {`
`142`	`142`	`}`
`143`	`143`	`else if(hasSankey) {`
`144`	`144`	`hoverGroup = ['hoverClosestCartesian', 'hoverCompareCartesian'];`
	`145`	`+ resetGroup = ['resetViewSankey'];`
`145`	`146`	`}`
`146`	`147`	`else { // hasPolar, hasTernary`
`147`	`148`	`// always show at least one hover icon.`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ function convertToD3Sankey(trace) {`
`40`	`40`	`if(linkSpec.target[i] > maxNodeId) maxNodeId = linkSpec.target[i];`
`41`	`41`	`}`
`42`	`42`	`var nodeCount = maxNodeId + 1;`
	`43`	`+ trace.node._count = nodeCount;`
`43`	`44`
`44`	`45`	`// Group nodes`
`45`	`46`	`var j;`