BUG: Mask NaNs in color arrays passed to streamplot

tonysyu · tonysyu · commit 8350febf8b46 · 2012-11-16T16:29:24.000-05:00
Bug report and fix provided by Jon Ramsey.
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -82,6 +82,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
     if use_multicolor_lines:
         assert color.shape == grid.shape
         line_colors = []
+        if np.any(np.isnan(color)):
+            color = np.ma.array(color, mask=np.isnan(color))
     else:
         line_kw['color'] = color
         arrow_kw['color'] = color
