Fix selection of arrow's linewidth and color

tonysyu · tonysyu · commit b884fc813db4 · 2012-01-16T22:07:27.000-05:00
diff --git a/examples/pylab_examples/streamplot_demo.py b/examples/pylab_examples/streamplot_demo.py
@@ -13,7 +13,7 @@
 ax1.streamplot(X, Y, U, V)
 
 lw = 5*speed/speed.max()
-ax2.streamplot(X, Y, U, V, density=0.6, color=U, linewidth=lw)
+ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw)
 
 plt.show()
 
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -98,39 +98,41 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=1, color='k', cmap=None,
     line_kw = {}
     arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10*arrowsize)
 
+    if not type(linewidth) == np.ndarray:
+        line_kw['linewidth'] = linewidth
+        arrow_kw['linewidth'] = linewidth
+
+    if not type(color) == np.ndarray:
+        line_kw['color'] = color
+        arrow_kw['color'] = color
+
     for t in trajectories:
         tgx = np.array(t[0])
         tgy = np.array(t[1])
-
-        if type(linewidth) == np.ndarray:
-            line_kw['linewidth'] = interpgrid(linewidth, tgx, tgy)[:-1]
-            arrow_kw['linewidth'] = line_kw['linewidth'][len(tgx) / 2]
-        else:
-            line_kw['linewidth'] = linewidth
-            arrow_kw['linewidth'] = linewidth
-
-        if type(color) == np.ndarray:
-            line_kw['color'] = cmap(norm(interpgrid(color, tgx, tgy)[:-1]))
-            arrow_kw['color'] = line_kw['color'][len(tgx) / 2]
-        else:
-            line_kw['color'] = color
-            arrow_kw['color'] = color
-
         # Rescale from grid-coordinates to data-coordinates.
         tx = np.array(t[0]) * grid.dx + grid.x_origin
         ty = np.array(t[1]) * grid.dy + grid.y_origin
 
         points = np.transpose([tx, ty]).reshape(-1, 1, 2)
         segments = np.concatenate([points[:-1], points[1:]], axis=1)
 
-        lc = matplotlib.collections.LineCollection(segments, **line_kw)
-        axes.add_collection(lc)
-
         ## Add arrows half way along each trajectory.
         s = np.cumsum(np.sqrt(np.diff(tx)**2 + np.diff(ty)**2))
         n = np.searchsorted(s, s[-1] / 2.)
         arrow_tail = (tx[n], ty[n])
         arrow_head = (np.mean(tx[n:n+2]), np.mean(ty[n:n+2]))
+
+        if type(linewidth) == np.ndarray:
+            line_kw['linewidth'] = interpgrid(linewidth, tgx, tgy)[:-1]
+            arrow_kw['linewidth'] = line_kw['linewidth'][n]
+
+        if type(color) == np.ndarray:
+            line_kw['color'] = cmap(norm(interpgrid(color, tgx, tgy)[:-1]))
+            arrow_kw['color'] = line_kw['color'][n]
+
+        lc = matplotlib.collections.LineCollection(segments, **line_kw)
+        axes.add_collection(lc)
+
         p = mpp.FancyArrowPatch(arrow_tail, arrow_head, **arrow_kw)
         axes.add_patch(p)
 
