|
16 | 16 | """ |
17 | 17 | import numpy as np |
18 | 18 | import matplotlib.pyplot as plt |
19 | | -import matplotlib.gridspec as gridspec |
| 19 | + |
20 | 20 |
|
21 | 21 | w = 3 |
22 | 22 | Y, X = np.mgrid[-w:w:100j, -w:w:100j] |
23 | 23 | U = -1 - X**2 + Y |
24 | 24 | V = 1 + X - Y**2 |
25 | 25 | speed = np.sqrt(U**2 + V**2) |
26 | 26 |
|
27 | | -fig = plt.figure(figsize=(7, 9)) |
28 | | -gs = gridspec.GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2]) |
| 27 | +fig, axs = plt.subplots(3, 2, figsize=(7, 9), height_ratios=[1, 1, 2]) |
| 28 | +axs = axs.flat |
29 | 29 |
|
30 | 30 | # Varying density along a streamline |
31 | | -ax0 = fig.add_subplot(gs[0, 0]) |
32 | | -ax0.streamplot(X, Y, U, V, density=[0.5, 1]) |
33 | | -ax0.set_title('Varying Density') |
| 31 | +axs[0].streamplot(X, Y, U, V, density=[0.5, 1]) |
| 32 | +axs[0].set_title('Varying Density') |
34 | 33 |
|
35 | 34 | # Varying color along a streamline |
36 | | -ax1 = fig.add_subplot(gs[0, 1]) |
37 | | -strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn') |
| 35 | +strm = axs[1].streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn') |
38 | 36 | fig.colorbar(strm.lines) |
39 | | -ax1.set_title('Varying Color') |
| 37 | +axs[1].set_title('Varying Color') |
40 | 38 |
|
41 | 39 | # Varying line width along a streamline |
42 | | -ax2 = fig.add_subplot(gs[1, 0]) |
43 | 40 | lw = 5*speed / speed.max() |
44 | | -ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw) |
45 | | -ax2.set_title('Varying Line Width') |
| 41 | +axs[2].streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw) |
| 42 | +axs[2].set_title('Varying Line Width') |
46 | 43 |
|
47 | 44 | # Controlling the starting points of the streamlines |
48 | 45 | seed_points = np.array([[-2, -1, 0, 1, 2, -1], [-2, -1, 0, 1, 2, 2]]) |
49 | 46 |
|
50 | | -ax3 = fig.add_subplot(gs[1, 1]) |
51 | | -strm = ax3.streamplot(X, Y, U, V, color=U, linewidth=2, |
52 | | - cmap='autumn', start_points=seed_points.T) |
| 47 | +strm = axs[3].streamplot(X, Y, U, V, color=U, linewidth=2, |
| 48 | + cmap='autumn', start_points=seed_points.T) |
53 | 49 | fig.colorbar(strm.lines) |
54 | | -ax3.set_title('Controlling Starting Points') |
| 50 | +axs[3].set_title('Controlling Starting Points') |
55 | 51 |
|
56 | 52 | # Displaying the starting points with blue symbols. |
57 | | -ax3.plot(seed_points[0], seed_points[1], 'bo') |
58 | | -ax3.set(xlim=(-w, w), ylim=(-w, w)) |
| 53 | +axs[3].plot(seed_points[0], seed_points[1], 'bo') |
| 54 | +axs[3].set(xlim=(-w, w), ylim=(-w, w)) |
59 | 55 |
|
60 | 56 | # Create a mask |
61 | 57 | mask = np.zeros(U.shape, dtype=bool) |
62 | 58 | mask[40:60, 40:60] = True |
63 | 59 | U[:20, :20] = np.nan |
64 | 60 | U = np.ma.array(U, mask=mask) |
65 | 61 |
|
66 | | -ax4 = fig.add_subplot(gs[2, 0]) |
67 | | -ax4.streamplot(X, Y, U, V, color='r') |
68 | | -ax4.set_title('Streamplot with Masking') |
| 62 | +axs[4].streamplot(X, Y, U, V, color='r') |
| 63 | +axs[4].set_title('Streamplot with Masking') |
69 | 64 |
|
70 | | -ax4.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray', aspect='auto') |
71 | | -ax4.set_aspect('equal') |
| 65 | +axs[4].imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray', |
| 66 | + aspect='auto') |
| 67 | +axs[4].set_aspect('equal') |
72 | 68 |
|
73 | | -ax5 = fig.add_subplot(gs[2, 1]) |
74 | | -ax5.streamplot(X, Y, U, V, broken_streamlines=False) |
75 | | -ax5.set_title('Streamplot with unbroken streamlines') |
| 69 | +axs[5].streamplot(X, Y, U, V, broken_streamlines=False) |
| 70 | +axs[5].set_title('Streamplot with unbroken streamlines') |
76 | 71 |
|
77 | 72 | plt.tight_layout() |
78 | 73 | plt.show() |
|
0 commit comments