Merge pull request #4714 from Cadair/streamplot_seeds

efiring · efiring · commit 2c92c4c14805 · 2015-07-16T10:31:24.000-10:00
Add an option to streamplot to manually specify the seed points.
diff --git a/examples/images_contours_and_fields/streamplot_demo_start_points.py b/examples/images_contours_and_fields/streamplot_demo_start_points.py
@@ -0,0 +1,31 @@
+"""
+Demo of the `streamplot` function.
+
+A streamplot, or streamline plot, is used to display 2D vector fields. This
+example shows a few features of the stream plot function:
+
+    * Varying the color along a streamline.
+    * Varying the density of streamlines.
+    * Varying the line width along a stream line.
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+plt.ion()
+
+X, Y = (np.linspace(-3, 3, 100),
+        np.linspace(-3, 3, 100))
+
+U, V = np.mgrid[-3:3:100j, 0:0:100j]
+
+seed_points = np.array([[-2, 0, 1], [-2, 0, 1]])
+
+fig0, ax0 = plt.subplots()
+strm = ax0.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn, start_points=seed_points.T)
+fig0.colorbar(strm.lines)
+
+ax0.plot(seed_points[0], seed_points[1], 'bo')
+
+ax0.axis((-3,3,-3,3))
+
+
+#plt.show()
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4342,7 +4342,7 @@ def stackplot(self, x, *args, **kwargs):
 
     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, zorder=1):
+                   minlength=0.1, transform=None, zorder=1, start_points=None):
         if not self._hold:
             self.cla()
         stream_container = mstream.streamplot(self, x, y, u, v,
@@ -4354,6 +4354,7 @@ def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
                                               arrowsize=arrowsize,
                                               arrowstyle=arrowstyle,
                                               minlength=minlength,
+                                              start_points=start_points,
                                               transform=transform,
                                               zorder=zorder)
         return stream_container
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2662,10 +2662,10 @@ def broken_barh(xranges, yrange, hold=None, **kwargs):
 # This function was autogenerated by boilerplate.py.  Do not edit as
 # changes will be lost
 @_autogen_docstring(Axes.boxplot)
-def boxplot(x, notch=False, sym=None, vert=True, whis=1.5, positions=None,
-            widths=None, patch_artist=False, bootstrap=None, usermedians=None,
-            conf_intervals=None, meanline=False, showmeans=False, showcaps=True,
-            showbox=True, showfliers=True, boxprops=None, labels=None,
+def boxplot(x, notch=None, sym=None, vert=None, whis=None, positions=None,
+            widths=None, patch_artist=None, bootstrap=None, usermedians=None,
+            conf_intervals=None, meanline=None, showmeans=None, showcaps=None,
+            showbox=None, showfliers=None, boxprops=None, labels=None,
             flierprops=None, medianprops=None, meanprops=None, capprops=None,
             whiskerprops=None, manage_xticks=True, hold=None):
     ax = gca()
@@ -3338,7 +3338,7 @@ def step(x, y, *args, **kwargs):
 @_autogen_docstring(Axes.streamplot)
 def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
                norm=None, arrowsize=1, arrowstyle='-|>', minlength=0.1,
-               transform=None, zorder=1, hold=None):
+               transform=None, zorder=1, start_points=None, hold=None):
     ax = gca()
     # allow callers to override the hold state by passing hold=True|False
     washold = ax.ishold()
@@ -3350,7 +3350,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
                             color=color, cmap=cmap, norm=norm,
                             arrowsize=arrowsize, arrowstyle=arrowstyle,
                             minlength=minlength, transform=transform,
-                            zorder=zorder)
+                            zorder=zorder, start_points=start_points)
     finally:
         ax.hold(washold)
     sci(ret.lines)
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -21,7 +21,7 @@
 
 def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
                cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
-               minlength=0.1, transform=None, zorder=1):
+               minlength=0.1, transform=None, zorder=1, start_points=None):
     """Draws streamlines of a vector flow.
 
     *x*, *y* : 1d arrays
@@ -52,6 +52,9 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
         See :class:`~matplotlib.patches.FancyArrowPatch`.
     *minlength* : float
         Minimum length of streamline in axes coordinates.
+    *start_points*: Nx2 array
+        Coordinates of starting points for the streamlines.
+        In data coordinates, the same as the ``x`` and ``y`` arrays.
     *zorder* : int
         any number
 
@@ -122,9 +125,22 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
     integrate = get_integrator(u, v, dmap, minlength)
 
     trajectories = []
-    for xm, ym in _gen_starting_points(mask.shape):
-        if mask[ym, xm] == 0:
-            xg, yg = dmap.mask2grid(xm, ym)
+    if start_points is None:
+        for xm, ym in _gen_starting_points(mask.shape):
+            if mask[ym, xm] == 0:
+                xg, yg = dmap.mask2grid(xm, ym)
+                t = integrate(xg, yg)
+                if t is not None:
+                    trajectories.append(t)
+    else:
+        # Convert start_points from data to array coords
+        # Shift the seed points from the bottom left of the data so that
+        # data2grid works properly.
+        sp2 = np.asanyarray(start_points).copy()
+        sp2[:,0] += np.abs(x[0])
+        sp2[:,1] += np.abs(y[0])
+        for xs, ys in sp2:
+            xg, yg = dmap.data2grid(xs, ys)
             t = integrate(xg, yg)
             if t is not None:
                 trajectories.append(t)
