[Matplotlib-users] Using an array which includes NaNs for color in streamplot

[Jon R. <jon...@gm...>]

Hi Everyone,

I have just upgraded to matplotlib 1.2.0 so that I can use the streamplot
module, which I'm quite happy about!
However, I've noticed that when one tries to color the streamlines using a
2-D array which contains NaNs, streamlines of only one color are shown!  I
have appended example code below which reproduces the problem.

Meanwhile, if the following two lines are inserted inside an "if
use_multicolor_lines:" region within streamplot.py, then the problem goes
away (for example, after line 84 or line 115):
    if np.any(np.isnan(color)):
        color = np.ma.array(color, mask=np.isnan(color))

This check already exists on the input arrays U and V, but not for color.
I am also not sure this issue will persist when a normalize object is
explicitly specified.

Example code (derived from streamplot_demo.py):

import numpy as np
import matplotlib.pyplot as plt

Y, X = np.mgrid[-3:3:100j, -3:3:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

m = np.sqrt(X**2 + Y**2) < 1.0
speed[m] = np.nan

plt.streamplot(X, Y, U, V, color=speed, linewidth=2, cmap=plt.cm.autumn)
plt.colorbar()
plt.show()

Additional info:
Linux 2.6.38-16-generic #67-Ubuntu SMP Thu Sep 6 17:58:38 UTC 2012 x86_64
x86_64 x86_64 GNU/Linux
Matplotlib v1.2.0


Best Regards,

Jon Ramsey

P.S. Long time reader, first time poster.