Merge pull request matplotlib#1514 from tonysyu/streamplot-nan-colors

Fix streamplot when color argument has NaNs

Author: dmcdougall

lib/matplotlib/__init__.py

@@ -1105,6 +1105,7 @@ def tk_window_focus():
     'matplotlib.tests.test_scale',
     'matplotlib.tests.test_simplification',
     'matplotlib.tests.test_spines',
+    'matplotlib.tests.test_streamplot',
     'matplotlib.tests.test_subplots',
     'matplotlib.tests.test_text',
     'matplotlib.tests.test_ticker',

lib/matplotlib/streamplot.py

@@ -85,6 +85,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