From 391c24ae0dde742715f24ba1f13581c930e62c8e Mon Sep 17 00:00:00 2001 From: Tony S Yu Date: Tue, 27 Mar 2012 21:08:19 -0400 Subject: [PATCH 1/4] Return arrow collection as 2nd argument of streamplot. --- lib/matplotlib/pyplot.py | 2 +- lib/matplotlib/streamplot.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index effc065083f2..77ca38873d93 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -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[0]) return ret # This function was autogenerated by boilerplate.py. Do not edit as diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py index a3a349a74fee..60bdc1d7b329 100644 --- a/lib/matplotlib/streamplot.py +++ b/lib/matplotlib/streamplot.py @@ -108,6 +108,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 +140,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 +153,9 @@ 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 + + arrow_collection = matplotlib.collections.PatchCollection(arrows) + return lc, arrow_collection # Coordinate definitions From 57a7d3de898553132a416bfb712cc03f3ab4ed3b Mon Sep 17 00:00:00 2001 From: Tony S Yu Date: Sat, 28 Apr 2012 14:22:32 -0400 Subject: [PATCH 2/4] Return container object containing lines and arrows. --- lib/matplotlib/axes.py | 22 +++++++++++----------- lib/matplotlib/pyplot.py | 2 +- lib/matplotlib/streamplot.py | 30 +++++++++++++++++++++++------- 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py index 21e09b811957..16aaa134e05a 100644 --- a/lib/matplotlib/axes.py +++ b/lib/matplotlib/axes.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 diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index 77ca38873d93..b1444f85d9ea 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -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[0]) + sci(ret.lines) return ret # This function was autogenerated by boilerplate.py. Do not edit as diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py index 60bdc1d7b329..ac29c3eff237 100644 --- a/lib/matplotlib/streamplot.py +++ b/lib/matplotlib/streamplot.py @@ -8,6 +8,8 @@ import matplotlib.colors as mcolors import matplotlib.collections as mcollections import matplotlib.patches as patches +import matplotlib.container as container + __all__ = ['streamplot'] @@ -46,11 +48,13 @@ 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* : StreamplotContainer + Container object with attributes + lines : `matplotlib.collections.LineCollection` of streamlines + arrows : collection of `matplotlib.patches.FancyArrowPatch` objects + repesenting arrows half-way along stream lines. """ grid = Grid(x, y) mask = StreamMask(density) @@ -154,8 +158,20 @@ 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) - arrow_collection = matplotlib.collections.PatchCollection(arrows) - return lc, arrow_collection + ac = matplotlib.collections.PatchCollection(arrows) + stream_container = StreamplotContainer(lc, arrows=ac) + return stream_container + + +class StreamplotContainer(container.Container): + + def __new__(cls, *kl, **kwargs): + return tuple.__new__(cls) + + def __init__(self, lines, arrows=None, **kwargs): + self.lines = lines + self.arrows = arrows + container.Container.__init__(self, lines, **kwargs) # Coordinate definitions From 9875323ed79767751f05a109e44c5da7556e6ce2 Mon Sep 17 00:00:00 2001 From: Tony S Yu Date: Sun, 29 Apr 2012 11:46:30 -0400 Subject: [PATCH 3/4] Change return value of `streamplot`. - Return object that derives from `object` since deriving from `Container` was not beneficial. - This return value is a stopgap; the final return value should allow users to set the colormap, alpha, etc. for both lines and arrows. --- lib/matplotlib/streamplot.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py index ac29c3eff237..2edd56229ef0 100644 --- a/lib/matplotlib/streamplot.py +++ b/lib/matplotlib/streamplot.py @@ -8,7 +8,6 @@ import matplotlib.colors as mcolors import matplotlib.collections as mcollections import matplotlib.patches as patches -import matplotlib.container as container __all__ = ['streamplot'] @@ -50,7 +49,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, Returns ------- - *stream_container* : StreamplotContainer + *stream_container* : StreamplotSet Container object with attributes lines : `matplotlib.collections.LineCollection` of streamlines arrows : collection of `matplotlib.patches.FancyArrowPatch` objects @@ -159,19 +158,15 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, axes.autoscale_view(tight=True) ac = matplotlib.collections.PatchCollection(arrows) - stream_container = StreamplotContainer(lc, arrows=ac) + stream_container = StreamplotSet(lc, ac) return stream_container -class StreamplotContainer(container.Container): +class StreamplotSet(object): - def __new__(cls, *kl, **kwargs): - return tuple.__new__(cls) - - def __init__(self, lines, arrows=None, **kwargs): + def __init__(self, lines, arrows, **kwargs): self.lines = lines self.arrows = arrows - container.Container.__init__(self, lines, **kwargs) # Coordinate definitions From a7248d989bf17c93afc2c956ef566ab5eede3e2a Mon Sep 17 00:00:00 2001 From: Tony S Yu Date: Sun, 2 Sep 2012 09:25:49 -0400 Subject: [PATCH 4/4] DOC: Add note about future changes to return value --- lib/matplotlib/streamplot.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py index 2edd56229ef0..ff8e343fe26c 100644 --- a/lib/matplotlib/streamplot.py +++ b/lib/matplotlib/streamplot.py @@ -54,6 +54,9 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, 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)