|
8 | 8 | import matplotlib.colors as mcolors |
9 | 9 | import matplotlib.collections as mcollections |
10 | 10 | import matplotlib.patches as patches |
| 11 | +import matplotlib.container as container |
| 12 | + |
11 | 13 |
|
12 | 14 | __all__ = ['streamplot'] |
13 | 15 |
|
@@ -46,11 +48,13 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, |
46 | 48 | *minlength* : float |
47 | 49 | Minimum length of streamline in axes coordinates. |
48 | 50 |
|
49 | | - Returns *streamlines* : :class:`~matplotlib.collections.LineCollection` |
50 | | - Line collection with all streamlines as a series of line segments. |
51 | | - Currently, there is no way to differentiate between line segments |
52 | | - on different streamlines (other than manually checking that segments |
53 | | - are connected). |
| 51 | + Returns |
| 52 | + ------- |
| 53 | + *stream_container* : StreamplotContainer |
| 54 | + Container object with attributes |
| 55 | + lines : `matplotlib.collections.LineCollection` of streamlines |
| 56 | + arrows : collection of `matplotlib.patches.FancyArrowPatch` objects |
| 57 | + repesenting arrows half-way along stream lines. |
54 | 58 | """ |
55 | 59 | grid = Grid(x, y) |
56 | 60 | mask = StreamMask(density) |
@@ -154,8 +158,20 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, |
154 | 158 | axes.update_datalim(((x.min(), y.min()), (x.max(), y.max()))) |
155 | 159 | axes.autoscale_view(tight=True) |
156 | 160 |
|
157 | | - arrow_collection = matplotlib.collections.PatchCollection(arrows) |
158 | | - return lc, arrow_collection |
| 161 | + ac = matplotlib.collections.PatchCollection(arrows) |
| 162 | + stream_container = StreamplotContainer(lc, arrows=ac) |
| 163 | + return stream_container |
| 164 | + |
| 165 | + |
| 166 | +class StreamplotContainer(container.Container): |
| 167 | + |
| 168 | + def __new__(cls, *kl, **kwargs): |
| 169 | + return tuple.__new__(cls) |
| 170 | + |
| 171 | + def __init__(self, lines, arrows=None, **kwargs): |
| 172 | + self.lines = lines |
| 173 | + self.arrows = arrows |
| 174 | + container.Container.__init__(self, lines, **kwargs) |
159 | 175 |
|
160 | 176 |
|
161 | 177 | # Coordinate definitions |
|
0 commit comments