Permalink
Browse files

Merge pull request #803 from tonysyu/streamplot-return-arrows

Return arrow collection as 2nd argument of streamplot.
  • Loading branch information...
2 parents 4c1d072 + a7248d9 commit 7a61238c930d725f5cd67b457e61b7034fcecd78 @efiring efiring committed Sep 2, 2012
Showing with 36 additions and 18 deletions.
  1. +11 −11 lib/matplotlib/axes.py
  2. +1 −1 lib/matplotlib/pyplot.py
  3. +24 −6 lib/matplotlib/streamplot.py
@@ -6601,17 +6601,17 @@ def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
minlength=0.1, transform=None):
if not self._hold: self.cla()
- lines = mstream.streamplot(self, x, y, u, v,
- density=density,
- linewidth=linewidth,
- color=color,
- cmap=cmap,
- norm=norm,
- arrowsize=arrowsize,
- arrowstyle=arrowstyle,
- minlength=minlength,
- transform=transform)
- return lines
+ stream_container = mstream.streamplot(self, x, y, u, v,
+ density=density,
+ linewidth=linewidth,
+ color=color,
+ cmap=cmap,
+ norm=norm,
+ arrowsize=arrowsize,
+ arrowstyle=arrowstyle,
+ minlength=minlength,
+ transform=transform)
+ return stream_container
streamplot.__doc__ = mstream.streamplot.__doc__
@docstring.dedent_interpd
@@ -3055,7 +3055,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
draw_if_interactive()
finally:
ax.hold(washold)
- sci(ret)
+ sci(ret.lines)
return ret
# This function was autogenerated by boilerplate.py. Do not edit as
@@ -9,6 +9,7 @@
import matplotlib.collections as mcollections
import matplotlib.patches as patches
+
__all__ = ['streamplot']
@@ -46,11 +47,16 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
*minlength* : float
Minimum length of streamline in axes coordinates.
- Returns *streamlines* : :class:`~matplotlib.collections.LineCollection`
- Line collection with all streamlines as a series of line segments.
- Currently, there is no way to differentiate between line segments
- on different streamlines (other than manually checking that segments
- are connected).
+ Returns
+ -------
+ *stream_container* : StreamplotSet
+ Container object with attributes
+ lines : `matplotlib.collections.LineCollection` of streamlines
+ arrows : collection of `matplotlib.patches.FancyArrowPatch` objects
+ repesenting arrows half-way along stream lines.
+ This container will probably change in the future to allow changes to
+ the colormap, alpha, etc. for both lines and arrows, but these changes
+ should be backward compatible.
"""
grid = Grid(x, y)
mask = StreamMask(density)
@@ -108,6 +114,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
cmap = cm.get_cmap(cmap)
streamlines = []
+ arrows = []
for t in trajectories:
tgx = np.array(t[0])
tgy = np.array(t[1])
@@ -139,6 +146,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
transform=transform,
**arrow_kw)
axes.add_patch(p)
+ arrows.append(p)
lc = mcollections.LineCollection(streamlines,
transform=transform,
@@ -151,7 +159,17 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
axes.update_datalim(((x.min(), y.min()), (x.max(), y.max())))
axes.autoscale_view(tight=True)
- return lc
+
+ ac = matplotlib.collections.PatchCollection(arrows)
+ stream_container = StreamplotSet(lc, ac)
+ return stream_container
+
+
+class StreamplotSet(object):
+
+ def __init__(self, lines, arrows, **kwargs):
+ self.lines = lines
+ self.arrows = arrows
# Coordinate definitions

0 comments on commit 7a61238

Please sign in to comment.