Include colorscale attributes in quiver

Author: degzhaus

src/traces/quiver/attributes.js

@@ -6,6 +6,7 @@ var fontAttrs = require('../../plots/font_attributes');
 var dash = require('../../components/drawing/attributes').dash;
 
 var extendFlat = require('../../lib/extend').extendFlat;
+var colorScaleAttrs = require('../../components/colorscale/attributes');
 
 var attrs = {
     x: {
@@ -223,6 +224,16 @@ var attrs = {
 // Extend with base attributes (includes hoverinfo, etc.)
 extendFlat(attrs, baseAttrs);
 
+// Colorscale attributes to color arrows by |(u,v)| magnitude
+extendFlat(
+    attrs,
+    colorScaleAttrs('', {
+        colorAttr: 'u/v norm',
+        showScaleDflt: true,
+        editTypeOverride: 'calc'
+    })
+);
+
 // Add hoverinfo with proper flags for quiver
 // We need to create a new object to avoid mutating the shared base attributes
 attrs.hoverinfo = extendFlat({}, baseAttrs.hoverinfo, {

src/traces/quiver/defaults.js

@@ -2,6 +2,9 @@
 
 var Lib = require('../../lib');
 var attributes = require('./attributes');
+var Colorscale = require('../../components/colorscale');
+var colorscaleDefaults = Colorscale.handleDefaults;
+var hasColorscale = Colorscale.hasColorscale;
 
 module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout) {
     // Selection styling - use coerce to set proper defaults
@@ -61,6 +64,26 @@ module.exports = function supplyDefaults(traceIn, traceOut, defaultColor, layout
     // traceOut.hoverinfo will be set by Lib.coerceHoverinfo in plots.js
     traceOut.hovertemplate = traceIn.hovertemplate;
 
+    // Colorscale for magnitude coloring: compute cmin/cmax from |(u,v)|
+    var cmin = Infinity;
+    var cmax = -Infinity;
+    for (var k = 0; k < len; k++) {
+        var uu = (traceOut.u && traceOut.u[k]) || (traceIn.u && traceIn.u[k]) || 0;
+        var vv = (traceOut.v && traceOut.v[k]) || (traceIn.v && traceIn.v[k]) || 0;
+        var nrm = Math.sqrt(uu * uu + vv * vv);
+        if (isFinite(nrm)) {
+            if (nrm < cmin) cmin = nrm;
+            if (nrm > cmax) cmax = nrm;
+        }
+    }
+    if (!isFinite(cmin)) cmin = 0;
+    if (!isFinite(cmax)) cmax = 1;
+    if (traceIn.cmin === undefined && traceOut.cmin === undefined) traceOut.cmin = cmin;
+    if (traceIn.cmax === undefined && traceOut.cmax === undefined) traceOut.cmax = cmax;
+    // Flag colorscale and apply defaults (adds colorscale, showscale, colorbar, etc.)
+    traceOut._hasColorscale = hasColorscale(traceIn) || true;
+    colorscaleDefaults(traceIn, traceOut, layout, coerce, { prefix: '', cLetter: 'c' });
+
     // Text
     traceOut.text = traceIn.text;
     traceOut.textposition = traceIn.textposition || 'middle center';

src/traces/quiver/plot.js

@@ -5,6 +5,7 @@ var d3 = require('@plotly/d3');
 var Registry = require('../../registry');
 var Lib = require('../../lib');
 var Drawing = require('../../components/drawing');
+var Colorscale = require('../../components/colorscale');
 
 module.exports = function plot(gd, plotinfo, cdscatter, scatterLayer, transitionOpts, makeOnCompleteCallback) {
     var join, onComplete;
@@ -184,6 +185,17 @@ function plotOne(gd, idx, plotinfo, cdscatter, cdscatterAll, element, transition
     // Apply styling using Plotly's standard styling system
     Drawing.lineGroupStyle(lineSegments, trace.line && trace.line.width, trace.line && trace.line.color, trace.line && trace.line.dash);
 
+    // If colorscale present, color arrows by magnitude |(u,v)|
+    if (trace._hasColorscale) {
+        var colorFunc = Colorscale.makeColorScaleFuncFromTrace(trace);
+        lineSegments.style('stroke', function(cdi) {
+            var uVal = (trace.u && trace.u[cdi.i]) || 0;
+            var vVal = (trace.v && trace.v[cdi.i]) || 0;
+            var nVal = Math.sqrt(uVal * uVal + vVal * vVal);
+            return colorFunc(nVal);
+        });
+    }
+
     // Handle transitions
     if(transitionOpts && transitionOpts.duration > 0) {
         var transition = d3.transition()