1010__all__ = ['streamplot' ]
1111
1212
13- def streamplot (axes , x , y , u , v , density = 1 , linewidth = 1 , color = 'k' , cmap = None ,
14- arrowsize = 1 , arrowstyle = '-|>' , minlength = 0.1 ):
13+ def streamplot (axes , x , y , u , v , density = 1 , linewidth = None , color = None ,
14+ cmap = None , arrowsize = 1 , arrowstyle = '-|>' , minlength = 0.1 ):
1515 """Draws streamlines of a vector flow.
1616
1717 Parameters
@@ -56,13 +56,26 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=1, color='k', cmap=None,
5656
5757 if color is None :
5858 color = matplotlib .rcParams ['lines.color' ]
59- elif isinstance (color , np .ndarray ):
60- assert color .shape == grid .shape
6159
6260 if linewidth is None :
6361 linewidth = matplotlib .rcParams ['lines.linewidth' ]
64- elif isinstance (linewidth , np .ndarray ):
62+
63+ line_kw = {}
64+ arrow_kw = dict (arrowstyle = arrowstyle , mutation_scale = 10 * arrowsize )
65+
66+ if isinstance (color , np .ndarray ):
67+ assert color .shape == grid .shape
68+ line_colors = []
69+ else :
70+ line_kw ['color' ] = color
71+ arrow_kw ['color' ] = color
72+
73+ if isinstance (linewidth , np .ndarray ):
6574 assert linewidth .shape == grid .shape
75+ line_kw ['linewidth' ] = []
76+ else :
77+ line_kw ['linewidth' ] = linewidth
78+ arrow_kw ['linewidth' ] = linewidth
6679
6780 ## Sanity checks.
6881 assert u .shape == grid .shape
@@ -89,20 +102,6 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=1, color='k', cmap=None,
89102 if cmap == None : cmap = matplotlib .cm .get_cmap (
90103 matplotlib .rcParams ['image.cmap' ])
91104
92- line_kw = {}
93- arrow_kw = dict (arrowstyle = arrowstyle , mutation_scale = 10 * arrowsize )
94-
95- if isinstance (linewidth , np .ndarray ):
96- line_kw ['linewidth' ] = []
97- else :
98- line_kw ['linewidth' ] = linewidth
99- arrow_kw ['linewidth' ] = linewidth
100-
101- if isinstance (color , np .ndarray ):
102- line_colors = []
103- else :
104- line_kw ['color' ] = color
105- arrow_kw ['color' ] = color
106105
107106 streamlines = []
108107 for t in trajectories :
0 commit comments