Sankey: colorscales per component, linked to concentration

src/traces/sankey/attributes.js

@@ -14,6 +14,8 @@ var colorAttrs = require('../../components/color/attributes');
 var fxAttrs = require('../../components/fx/attributes');
 var domainAttrs = require('../../plots/domain').attributes;
 var hovertemplateAttrs = require('../../components/fx/hovertemplate_attributes');
+var colorAttributes = require('../../components/colorscale/attributes');
+var templatedArray = require('../../plot_api/plot_template').templatedArray;
 
 var extendFlat = require('../../lib/extend').extendFlat;
 var overrideAll = require('../../plot_api/edit_types').overrideAll;
@@ -225,7 +227,37 @@ var attrs = module.exports = overrideAll({
             description: 'Variables `source` and `target` are node objects.',
             keys: ['value', 'label']
         }),
-        description: 'The links of the Sankey plot.'
+        colorscales: templatedArray('concentrationscales', {
+            editType: 'calc',
+            label: {
+                valType: 'string',
+                role: 'info',
+                editType: 'calc',
+                description: 'The label of the links to color based on their concentration within a flow.',
+                dflt: ''
+            },
+            cmax: {
+                valType: 'number',
+                role: 'info',
+                editType: 'calc',
+                dflt: 1,
+                description: [
+                    'Sets the upper bound of the color domain.'
+                ].join('')
+            },
+            cmin: {
+                valType: 'number',
+                role: 'info',
+                editType: 'calc',
+                dflt: 0,
+                description: [
+                    'Sets the lower bound of the color domain.'
+                ].join('')
+            },
+            colorscale: extendFlat(colorAttributes().colorscale, {dflt: [[0, 'white'], [1, 'black']]})
+        }),
+        description: 'The links of the Sankey plot.',
+        role: 'info'
     }
 }, 'calc', 'nested');
 attrs.transforms = undefined;

src/traces/sankey/calc.js

@@ -14,7 +14,7 @@ var wrap = require('../../lib/gup').wrap;
 
 var isArrayOrTypedArray = Lib.isArrayOrTypedArray;
 var isIndex = Lib.isIndex;
-
+var Colorscale = require('../../components/colorscale');
 
 function convertToD3Sankey(trace) {
     var nodeSpec = trace.node;
@@ -24,8 +24,17 @@ function convertToD3Sankey(trace) {
     var hasLinkColorArray = isArrayOrTypedArray(linkSpec.color);
     var linkedNodes = {};
 
-    var nodeCount = nodeSpec.label.length;
+    var components = {};
+    var componentCount = linkSpec.colorscales.length;
     var i;
+    for(i = 0; i < componentCount; i++) {
+        var cscale = linkSpec.colorscales[i];
+        var specs = Colorscale.extractScale(cscale, {cLetter: 'c'});
+        var scale = Colorscale.makeColorScaleFunc(specs);
+        components[cscale.label] = scale;
+    }
+
+    var nodeCount = nodeSpec.label.length;
     for(i = 0; i < linkSpec.value.length; i++) {
         var val = linkSpec.value[i];
         // remove negative values, but keep zeros with special treatment
@@ -42,10 +51,14 @@ function convertToD3Sankey(trace) {
         var label = '';
         if(linkSpec.label && linkSpec.label[i]) label = linkSpec.label[i];
 
+        var concentrationscale = null;
+        if(label && components.hasOwnProperty(label)) concentrationscale = components[label];
+
         links.push({
             pointNumber: i,
             label: label,
             color: hasLinkColorArray ? linkSpec.color[i] : linkSpec.color,
+            concentrationscale: concentrationscale,
             source: source,
             target: target,
             value: +val

src/traces/sankey/defaults.js

@@ -15,6 +15,7 @@ var tinycolor = require('tinycolor2');
 var handleDomainDefaults = require('../../plots/domain').defaults;
 var handleHoverLabelDefaults = require('../../components/fx/hoverlabel_defaults');
 var Template = require('../../plot_api/plot_template');
+var handleArrayContainerDefaults = require('../../plots/array_container_defaults');
 
 module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout) {
     function coerce(attr, dflt) {
@@ -48,7 +49,7 @@ module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout
     }));
 
     // link attributes
-    var linkIn = traceIn.link;
+    var linkIn = traceIn.link || {};
     var linkOut = Template.newContainer(traceOut, 'link');
 
     function coerceLink(attr, dflt) {
@@ -70,6 +71,11 @@ module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout
 
     coerceLink('color', Lib.repeat(defaultLinkColor, linkOut.value.length));
 
+    handleArrayContainerDefaults(linkIn, linkOut, {
+        name: 'colorscales',
+        handleItemDefaults: concentrationscalesDefaults
+    });
+
     handleDomainDefaults(traceOut, layout, coerce);
 
     coerce('orientation');
@@ -83,3 +89,14 @@ module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout
     // don't match, between nodes and links
     traceOut._length = null;
 };
+
+function concentrationscalesDefaults(In, Out) {
+    function coerce(attr, dflt) {
+        return Lib.coerce(In, Out, attributes.link.colorscales, attr, dflt);
+    }
+
+    coerce('label');
+    coerce('cmin');
+    coerce('cmax');
+    coerce('colorscale');
+}

src/traces/sankey/plot.js

@@ -288,6 +288,7 @@ module.exports = function plot(gd, calcData) {
     };
 
     render(
+        gd,
         svg,
         calcData,
         {

src/traces/sankey/render.js

@@ -10,6 +10,7 @@
 
 var c = require('./constants');
 var d3 = require('d3');
+var sum = require('d3-array').sum;
 var tinycolor = require('tinycolor2');
 var Color = require('../../components/color');
 var Drawing = require('../../components/drawing');
@@ -67,6 +68,57 @@ function sankeyModel(layout, d, traceIndex) {
         Lib.warn('node.pad was reduced to ', sankey.nodePadding(), ' to fit within the figure.');
     }
 
+    function computeLinkConcentrations() {
+        graph.nodes.forEach(function(node) {
+            // Links connecting the same two nodes are part of a flow
+            var flows = {};
+            node.targetLinks.forEach(function(link) {
+                var flowKey = link.source.pointNumber + ':' + link.target.pointNumber;
+                if(!flows.hasOwnProperty(flowKey)) flows[flowKey] = [];
+                flows[flowKey].push(link);
+            });
+
+            // Compute statistics for each flow
+            Object.keys(flows).forEach(function(flowKey) {
+                var flowLinks = flows[flowKey];
+
+                // Find the total size of the flow and total size per label
+                var total = 0;
+                var totalPerLabel = {};
+                flowLinks.forEach(function(link) {
+                    if(!totalPerLabel[link.label]) totalPerLabel[link.label] = 0;
+                    totalPerLabel[link.label] += link.value;
+                    total += link.value;
+                });
+
+                // Find the ratio of the link's value and the size of the flow
+                flowLinks.forEach(function(link) {
+                    link.flow = {
+                        value: total,
+                        labelConcentration: totalPerLabel[link.label] / total,
+                        concentration: link.value / total,
+                        links: flowLinks
+                    };
+                });
+            });
+
+            // Gather statistics of all links at current node
+            var totalOutflow = sum(node.sourceLinks, function(n) {
+                return n.value;
+            });
+            node.sourceLinks.forEach(function(link) {
+                link.concentrationOut = link.value / totalOutflow;
+            });
+            var totalInflow = sum(node.targetLinks, function(n) {
+                return n.value;
+            });
+            node.targetLinks.forEach(function(link) {
+                link.concenrationIn = link.value / totalInflow;
+            });
+        });
+    }
+    computeLinkConcentrations();
+
     return {
         circular: circular,
         key: traceIndex,
@@ -100,6 +152,9 @@ function sankeyModel(layout, d, traceIndex) {
 
 function linkModel(d, l, i) {
     var tc = tinycolor(l.color);
+    if(l.concentrationscale) {
+        tc = tinycolor(l.concentrationscale(l.flow.labelConcentration));
+    }
     var basicKey = l.source.label + '|' + l.target.label;
     var key = basicKey + '__' + i;
 
@@ -121,7 +176,8 @@ function linkModel(d, l, i) {
         valueSuffix: d.valueSuffix,
         sankey: d.sankey,
         parent: d,
-        interactionState: d.interactionState
+        interactionState: d.interactionState,
+        flow: l.flow
     };
 }
 
@@ -568,7 +624,7 @@ function switchToSankeyFormat(nodes) {
 }
 
 // scene graph
-module.exports = function(svg, calcData, layout, callbacks) {
+module.exports = function(gd, svg, calcData, layout, callbacks) {
 
     var styledData = calcData
             .filter(function(d) {return unwrap(d).trace.visible;})
@@ -616,6 +672,7 @@ module.exports = function(svg, calcData, layout, callbacks) {
           .attr('d', linkPath())
           .call(attachPointerEvents, sankey, callbacks.linkEvents);
 
+
     sankeyLink
         .style('stroke', function(d) {
             return salientEnough(d) ? Color.tinyRGB(tinycolor(d.linkLineColor)) : d.tinyColorHue;

test/image/mocks/sankey_link_concentration.json

@@ -0,0 +1,74 @@
+{
+    "data": [
+        {
+            "type": "sankey",
+            "node": {
+                "pad": 25,
+                "line": {
+                  "color": "white",
+                  "width": 2
+                },
+                "label": ["process0", "process1", "process2", "process3", "process4"]
+            },
+            "link": {
+                "source": [
+                    0, 0, 0, 0,
+                    1, 1, 1, 1,
+                    1, 1, 1, 1,
+                    1, 1,
+                    2
+                ],
+                "target": [
+                    1, 1, 1, 1,
+                    2, 2, 2, 2,
+                    3, 3, 3, 3,
+                    4, 4,
+                    0
+                ],
+                "value": [
+                    10, 20, 40, 30,
+                    10, 5, 10, 20,
+                    0, 10, 10, 10,
+                    15, 5,
+                    5
+
+                ],
+                "label": [
+                    "elementA", "elementB", "elementC", "elementD",
+                    "elementA", "elementB", "elementC", "elementD",
+                    "elementA", "elementB", "elementC", "elementD",
+                    "elementC", "elementC",
+                    "elementA"
+                ],
+                "line": {
+                    "color": "white",
+                    "width": 2
+                },
+                "colorscales": [
+                  {
+                    "label": "elementA",
+                    "colorscale": [[0, "white"], [1, "blue"]]
+                  },
+                  {
+                    "label": "elementB",
+                    "colorscale": [[0, "white"], [1, "red"]]
+                  },
+                  {
+                    "label": "elementC",
+                    "colorscale": [[0, "white"], [1, "green"]]
+                  },
+                  {
+                    "label": "elementD"
+                  }
+                ],
+
+                "hovertemplate": "%{label}<br><b>flow.labelConcentration</b>: %{flow.labelConcentration:%0.2f}<br><b>flow.concentration</b>: %{flow.concentration:%0.2f}<br><b>flow.value</b>: %{flow.value}"
+            }
+
+        }],
+    "layout": {
+        "title": "Sankey with links colored based on its concentration within a flow",
+        "width": 800,
+        "height": 800
+    }
+}

test/jasmine/tests/sankey_test.js

@@ -120,6 +120,9 @@ describe('sankey tests', function() {
 
                 expect(fullTrace.link.label)
                     .toEqual([], 'presence of link target array is guaranteed');
+
+                expect(fullTrace.link.colorscales)
+                    .toEqual([], 'presence of link colorscales array is guaranteed');
             });
 
         it('\'Sankey\' specification should have proper types',
@@ -784,6 +787,7 @@ describe('sankey tests', function() {
                 var pt = d.points[0];
                 expect(pt.hasOwnProperty('source')).toBeTruthy();
                 expect(pt.hasOwnProperty('target')).toBeTruthy();
+                expect(pt.hasOwnProperty('flow')).toBeTruthy();
             })
             .then(function() { return _unhover('node'); })
             .then(function(d) {