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