diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py index 17a3f4076407..df293fd44aaf 100644 --- a/lib/matplotlib/axes.py +++ b/lib/matplotlib/axes.py @@ -6389,13 +6389,15 @@ def quiver(self, *args, **kw): quiver.__doc__ = mquiver.Quiver.quiver_doc def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None, - cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1): + cmap=None, norm=None, arrowsize=1, arrowstyle='-|>', + minlength=0.1): 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) diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index 4c507c6bbb1b..e06721bd2da1 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -2658,7 +2658,7 @@ def step(x, y, *args, **kwargs): # This function was autogenerated by boilerplate.py. Do not edit as # changes will be lost @autogen_docstring(Axes.streamplot) -def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.10000000000000001, hold=None): +def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, norm=None, arrowsize=1, arrowstyle='-|>', minlength=0.10000000000000001, hold=None): ax = gca() # allow callers to override the hold state by passing hold=True|False washold = ax.ishold() @@ -2666,7 +2666,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, arr if hold is not None: ax.hold(hold) try: - ret = ax.streamplot(x, y, u, v, density, linewidth, color, cmap, arrowsize, arrowstyle, minlength) + ret = ax.streamplot(x, y, u, v, density, linewidth, color, cmap, norm, arrowsize, arrowstyle, minlength) draw_if_interactive() finally: ax.hold(washold) diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py index c70eb76edda7..92a43b430539 100644 --- a/lib/matplotlib/streamplot.py +++ b/lib/matplotlib/streamplot.py @@ -11,7 +11,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, - cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1): + cmap=None, norm=None, arrowsize=1, arrowstyle='-|>', + minlength=0.1): """Draws streamlines of a vector flow. Parameters @@ -31,19 +32,23 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, *color* : matplotlib color code, or 2d array Streamline color. When given an array with the same shape as velocities, *color* values are converted to colors using *cmap*. - *cmap* : Colormap + *cmap* : :class:`~matplotlib.colors.Colormap` Colormap used to plot streamlines and arrows. Only necessary when using an array input for *color*. + *norm* : :class:`~matplotlib.colors.Normalize` + Normalize object used to scale luminance data to 0, 1. If None, stretch + (min, max) to (0, 1). Only necessary when *color* is an array. *arrowsize* : float Factor scale arrow size. *arrowstyle* : str - Arrow style specification. See `matplotlib.patches.FancyArrowPatch`. + Arrow style specification. + See :class:`~matplotlib.patches.FancyArrowPatch`. *minlength* : float Minimum length of streamline in axes coordinates. Returns ------- - *streamlines* : `matplotlib.collections.LineCollection` + *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 @@ -62,7 +67,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, line_kw = {} arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10*arrowsize) - if isinstance(color, np.ndarray): + use_multicolor_lines = isinstance(color, np.ndarray) + if use_multicolor_lines: assert color.shape == grid.shape line_colors = [] else: @@ -95,12 +101,11 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, if t != None: trajectories.append(t) - # Load up the defaults - needed to get the color right. - if isinstance(color, np.ndarray): - norm = matplotlib.colors.normalize(color.min(), color.max()) - if cmap == None: cmap = matplotlib.cm.get_cmap( - matplotlib.rcParams['image.cmap']) - + if use_multicolor_lines: + if norm is None: + norm = matplotlib.colors.normalize(color.min(), color.max()) + if cmap is None: + cmap = matplotlib.cm.get_cmap(matplotlib.rcParams['image.cmap']) streamlines = [] for t in trajectories: @@ -113,7 +118,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, points = np.transpose([tx, ty]).reshape(-1, 1, 2) streamlines.extend(np.hstack([points[:-1], points[1:]])) - ## Add arrows half way along each trajectory. + # Add arrows half way along each trajectory. s = np.cumsum(np.sqrt(np.diff(tx)**2 + np.diff(ty)**2)) n = np.searchsorted(s, s[-1] / 2.) arrow_tail = (tx[n], ty[n]) @@ -124,7 +129,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, line_kw['linewidth'].extend(line_widths) arrow_kw['linewidth'] = line_widths[n] - if isinstance(color, np.ndarray): + if use_multicolor_lines: color_values = interpgrid(color, tgx, tgy)[:-1] line_colors.extend(color_values) arrow_kw['color'] = cmap(norm(color_values[n])) @@ -133,7 +138,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None, axes.add_patch(p) lc = matplotlib.collections.LineCollection(streamlines, **line_kw) - if isinstance(color, np.ndarray): + if use_multicolor_lines: lc.set_array(np.asarray(line_colors)) lc.set_cmap(cmap) lc.set_norm(norm)