1""" 2========== 3Streamplot 4========== 5 6A stream plot, or streamline plot, is used to display 2D vector fields. This 7example shows a few features of the `~.axes.Axes.streamplot` function: 8 9* Varying the color along a streamline. 10* Varying the density of streamlines. 11* Varying the line width along a streamline. 12* Controlling the starting points of streamlines. 13* Streamlines skipping masked regions and NaN values. 14""" 15import numpy as np 16import matplotlib.pyplot as plt 17import matplotlib.gridspec as gridspec 18 19w = 3 20Y, X = np.mgrid[-w:w:100j, -w:w:100j] 21U = -1 - X**2 + Y 22V = 1 + X - Y**2 23speed = np.sqrt(U**2 + V**2) 24 25fig = plt.figure(figsize=(7, 9)) 26gs = gridspec.GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2]) 27 28# Varying density along a streamline 29ax0 = fig.add_subplot(gs[0, 0]) 30ax0.streamplot(X, Y, U, V, density=[0.5, 1]) 31ax0.set_title('Varying Density') 32 33# Varying color along a streamline 34ax1 = fig.add_subplot(gs[0, 1]) 35strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn') 36fig.colorbar(strm.lines) 37ax1.set_title('Varying Color') 38 39# Varying line width along a streamline 40ax2 = fig.add_subplot(gs[1, 0]) 41lw = 5*speed / speed.max() 42ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw) 43ax2.set_title('Varying Line Width') 44 45# Controlling the starting points of the streamlines 46seed_points = np.array([[-2, -1, 0, 1, 2, -1], [-2, -1, 0, 1, 2, 2]]) 47 48ax3 = fig.add_subplot(gs[1, 1]) 49strm = ax3.streamplot(X, Y, U, V, color=U, linewidth=2, 50 cmap='autumn', start_points=seed_points.T) 51fig.colorbar(strm.lines) 52ax3.set_title('Controlling Starting Points') 53 54# Displaying the starting points with blue symbols. 55ax3.plot(seed_points[0], seed_points[1], 'bo') 56ax3.set(xlim=(-w, w), ylim=(-w, w)) 57 58# Create a mask 59mask = np.zeros(U.shape, dtype=bool) 60mask[40:60, 40:60] = True 61U[:20, :20] = np.nan 62U = np.ma.array(U, mask=mask) 63 64ax4 = fig.add_subplot(gs[2:, :]) 65ax4.streamplot(X, Y, U, V, color='r') 66ax4.set_title('Streamplot with Masking') 67 68ax4.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray', aspect='auto') 69ax4.set_aspect('equal') 70 71plt.tight_layout() 72plt.show() 73############################################################################# 74# 75# .. admonition:: References 76# 77# The use of the following functions, methods, classes and modules is shown 78# in this example: 79# 80# - `matplotlib.axes.Axes.streamplot` / `matplotlib.pyplot.streamplot` 81# - `matplotlib.gridspec.GridSpec` 82