8000 Add `norm` parameter to `streamplot`. · matplotlib/matplotlib@5e1a266 · GitHub
[go: up one dir, main page]
Skip to content

Commit 5e1a266

Browse files
tonysyutonysyu
committed
Add norm parameter to streamplot.
Also, fixed/enhanced some docstrings.
1 parent 76665d0 commit 5e1a266

File tree

3 files changed

+26
-18
lines changed

3 files changed

+26
-18
lines changed

‎lib/matplotlib/axes.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6389,13 +6389,15 @@ def quiver(self, *args, **kw):
63896389
quiver.__doc__ = mquiver.Quiver.quiver_doc
63906390

63916391
def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
6392-
cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1):
6392+
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
6393+
minlength=0.1):
63936394
if not self._hold: self.cla()
63946395
lines = mstream.streamplot(self, x, y, u, v,
63956396
density=density,
63966397
linewidth=linewidth,
63976398
color=color,
63986399
cmap=cmap,
6400+
norm=norm,
63996401
arrowsize=arrowsize,
64006402
arrowstyle=arrowstyle,
64016403
minlength=minlength)

‎lib/matplotlib/pyplot.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2658,15 +2658,15 @@ def step(x, y, *args, **kwargs):
26582658
# This function was autogenerated by boilerplate.py. Do not edit as
26592659
# changes will be lost
26602660
@autogen_docstring(Axes.streamplot)
2661-
def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.10000000000000001, hold=None):
2661+
def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, norm=None, arrowsize=1, arrowstyle='-|>', minlength=0.10000000000000001, hold=None):
26622662
ax = gca()
26632663
# allow callers to override the hold state by passing hold=True|False
26642664
washold = ax.ishold()
26652665

26662666
if hold is not None:
26672667
ax.hold(hold)
26682668
try:
2669-
ret = ax.streamplot(x, y, u, v, density, linewidth, color, cmap, arrowsize, arrowstyle, minlength)
2669+
ret = ax.streamplot(x, y, u, v, density, linewidth, color, cmap, norm, arrowsize, arrowstyle, minlength)
26702670
draw_if_interactive()
26712671
finally:
26722672
ax.hold(washold)

‎lib/matplotlib/streamplot.py

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111

1212

1313
def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
14-
cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1):
14+
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
15+
minlength=0.1):
1516
"""Draws streamlines of a vector flow.
1617
1718
Parameters
@@ -31,23 +32,28 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
3132
*color* : matplotlib color code, or 2d array
3233
Streamline color. When given an array with the same shape as
3334
velocities, *color* values are converted to colors using *cmap*.
34-
*cmap* : Colormap
35+
*cmap* : :class:`~matplotlib.colors.Colormap`
3536
Colormap used to plot streamlines and arrows. Only necessary when using
3637
an array input for *color*.
38+
*norm* : :class:`~matplotlib.colors.Normalize`
39+
Normalize object used to scale luminance data to 0, 1. If None, stretch
40+
(min, max) to (0, 1). Only necessary when *color* is an array.
3741
*arrowsize* : float
3842
Factor scale arrow size.
3943
*arrowstyle* : str
40-
Arrow style specification. See `matplotlib.patches.FancyArrowPatch`.
44+
Arrow style specification.
45+
See :class:`~matplotlib.patches.FancyArrowPatch`.
4146
*minlength* : float
4247
Minimum length of streamline in axes coordinates.
4348
4449
Returns
4550
-------
46-
*streamlines* : `matplotlib.collections.LineCollection`
51+
*streamlines* : `~matplotlib.collections.LineCollection`
4752
Line collection with all streamlines as a series of line segments.
4853
Currently, there is no way to differentiate between line segments
4954
on different streamlines (other than manually checking that segments
50-
are connected).
55+
are connected). NOTE: this return value is likely to change since it
56+
does not include the arrow patches.
5157
"""
5258
grid = Grid(x, y)
5359
mask = StreamMask(density)
@@ -62,7 +68,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
6268
line_kw = {}
6369
arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10*arrowsize)
6470

65-
if isinstance(color, np.ndarray):
71+
use_multicolor_lines = isinstance(color, np.ndarray)
72+
if use_multicolor_lines:
6673
assert color.shape == grid.shape
6774
line_colors = []
6875
else:
@@ -95,12 +102,11 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
95102
if t != None:
96103
trajectories.append(t)
97104

98-
# Load up the defaults - needed to get the color right.
99-
if isinstance(color, np.ndarray):
100-
norm = matplotlib.colors.normalize(color.min(), color.max())
101-
if cmap == None: cmap = matplotlib.cm.get_cmap(
102-
matplotlib.rcParams['image.cmap'])
103-
105+
if use_multicolor_lines:
106+
if norm is None:
107+
norm = matplotlib.colors.normalize(color.min(), color.max())
108+
if cmap is None:
109+
cmap = matplotlib.cm.get_cmap(matplotlib.rcParams['image.cmap'])
104110

105111
streamlines = []
106112
for t in trajectories:
@@ -113,7 +119,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
113119
points = np.transpose([tx, ty]).reshape(-1, 1, 2)
114120
streamlines.extend(np.hstack([points[:-1], points[1:]]))
115121

116-
## Add arrows half way along each trajectory.
122+
# Add arrows half way along each trajectory.
117123
s = np.cumsum(np.sqrt(np.diff(tx)**2 + np.diff(ty)**2))
118124
n = np.searchsorted(s, s[-1] / 2.)
119125
arrow_tail = (tx[n], ty[n])
@@ -124,7 +130,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
124130
line_kw['linewidth'].extend(line_widths)
125131
arrow_kw['linewidth'] = line_widths[n]
126132

127-
if isinstance(color, np.ndarray):
133+
if use_multicolor_lines:
128134
color_values = interpgrid(color, tgx, tgy)[:-1]
129135
line_colors.extend(color_values)
130136
arrow_kw['color'] = cmap(norm(color_values[n]))
@@ -133,7 +139,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
133139
axes.add_patch(p)
134140

135141
lc = matplotlib.collections.LineCollection(streamlines, **line_kw)
136-
if isinstance(color, np.ndarray):
142+
if use_multicolor_lines:
137143
lc.set_array(np.asarray(line_colors))
138144
lc.set_cmap(cmap)
139145
lc.set_norm(norm)

0 commit comments

Comments
 (0)
0