Add support for coloring by an independent scalar array

Author: degzhaus

src/traces/quiver/attributes.js

@@ -33,6 +33,13 @@ var attrs = {
         anim: true,
         description: 'Sets the y components of the arrow vectors.'
     },
+    // Optional scalar field used for colorscale mapping. If omitted, |(u,v)| is used.
+    c: {
+        valType: 'data_array',
+        editType: 'calc',
+        anim: true,
+        description: 'Sets the scalar values used to map colors via the colorscale. If not provided, the magnitude sqrt(u^2 + v^2) is used.'
+    },
     sizemode: {
         valType: 'enumerated',
         values: ['scaled', 'absolute', 'raw'],

src/traces/quiver/calc.js

@@ -5,6 +5,7 @@ var Axes = require('../../plots/cartesian/axes');
 var isNumeric = require('fast-isnumeric');
 var BADNUM = require('../../constants/numerical').BADNUM;
 var scatterCalc = require('../scatter/calc');
+var colorscaleCalc = require('../../components/colorscale/calc');
 
 /**
  * Main calculation function for quiver trace
@@ -24,6 +25,12 @@ module.exports = function calc(gd, trace) {
     trace._length = len;
     var cd = new Array(len);
 
+    var normMin = Infinity;
+    var normMax = -Infinity;
+    var cMin = Infinity;
+    var cMax = -Infinity;
+    var hasC = Array.isArray(trace.c);
+
     for(var i = 0; i < len; i++) {
         var cdi = cd[i] = { i: i };
         var xValid = isNumeric(xVals[i]);
@@ -36,11 +43,36 @@ module.exports = function calc(gd, trace) {
             cdi.x = BADNUM;
             cdi.y = BADNUM;
         }
+
+        // track ranges for colorscale
+        if(hasC) {
+            var ci = trace.c[i];
+            if(isNumeric(ci)) {
+                if(ci < cMin) cMin = ci;
+                if(ci > cMax) cMax = ci;
+            }
+        } else {
+            var ui = (trace.u && trace.u[i]) || 0;
+            var vi = (trace.v && trace.v[i]) || 0;
+            var n = Math.sqrt(ui * ui + vi * vi);
+            if(isFinite(n)) {
+                if(n < normMin) normMin = n;
+                if(n > normMax) normMax = n;
+            }
+        }
     }
 
     // Ensure axes are expanded and categories registered like scatter traces do
     scatterCalc.calcAxisExpansion(gd, trace, xa, ya, xVals, yVals);
 
+    // Colorscale cmin/cmax computation: prefer provided c, else magnitude
+    var vals = hasC ? [cMin, cMax] : [normMin, normMax];
+    colorscaleCalc(gd, trace, {
+        vals: vals,
+        containerStr: '',
+        cLetter: 'c'
+    });
+
     return cd;
 };
 

src/traces/quiver/defaults.js

@@ -64,23 +64,8 @@ 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.)
+    // Colorscale defaults (adds colorscale, showscale, colorbar, etc.)
+    // Keep colorscale enabled by default for quiver
     traceOut._hasColorscale = hasColorscale(traceIn) || true;
     colorscaleDefaults(traceIn, traceOut, layout, coerce, { prefix: '', cLetter: 'c' });
 

src/traces/quiver/plot.js

@@ -189,10 +189,16 @@ function plotOne(gd, idx, plotinfo, cdscatter, cdscatterAll, element, transition
     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);
+            var cArr = trace.c;
+            var value;
+            if (Array.isArray(cArr) && cArr.length > cdi.i && isFinite(cArr[cdi.i])) {
+                value = cArr[cdi.i];
+            } else {
+                var uVal = (trace.u && trace.u[cdi.i]) || 0;
+                var vVal = (trace.v && trace.v[cdi.i]) || 0;
+                value = Math.sqrt(uVal * uVal + vVal * vVal);
+            }
+            return colorFunc(value);
         });
     }