Skip to content

Commit 391c24a

Browse files
committed
Return arrow collection as 2nd argument of streamplot.
1 parent cadd152 commit 391c24a

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

lib/matplotlib/pyplot.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3055,7 +3055,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
30553055
draw_if_interactive()
30563056
finally:
30573057
ax.hold(washold)
3058-
sci(ret)
3058+
sci(ret[0])
30593059
return ret
30603060

30613061
# This function was autogenerated by boilerplate.py. Do not edit as

lib/matplotlib/streamplot.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
108108
cmap = cm.get_cmap(cmap)
109109

110110
streamlines = []
111+
arrows = []
111112
for t in trajectories:
112113
tgx = np.array(t[0])
113114
tgy = np.array(t[1])
@@ -139,6 +140,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
139140
transform=transform,
140141
**arrow_kw)
141142
axes.add_patch(p)
143+
arrows.append(p)
142144

143145
lc = mcollections.LineCollection(streamlines,
144146
transform=transform,
@@ -151,7 +153,9 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
151153

152154
axes.update_datalim(((x.min(), y.min()), (x.max(), y.max())))
153155
axes.autoscale_view(tight=True)
154-
return lc
156+
157+
arrow_collection = matplotlib.collections.PatchCollection(arrows)
158+
return lc, arrow_collection
155159

156160

157161
# Coordinate definitions

0 commit comments

Comments
 (0)