Avoid gridspec in more examples

QuLogic · QuLogic · commit 0046b3854fcc · 2022-08-12T01:11:25.000-04:00
diff --git a/examples/images_contours_and_fields/plot_streamplot.py b/examples/images_contours_and_fields/plot_streamplot.py
@@ -16,63 +16,58 @@
 """
 import numpy as np
 import matplotlib.pyplot as plt
-import matplotlib.gridspec as gridspec
+
 
 w = 3
 Y, X = np.mgrid[-w:w:100j, -w:w:100j]
 U = -1 - X**2 + Y
 V = 1 + X - Y**2
 speed = np.sqrt(U**2 + V**2)
 
-fig = plt.figure(figsize=(7, 9))
-gs = gridspec.GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2])
+fig, axs = plt.subplots(3, 2, figsize=(7, 9), height_ratios=[1, 1, 2])
+axs = axs.flat
 
 #  Varying density along a streamline
-ax0 = fig.add_subplot(gs[0, 0])
-ax0.streamplot(X, Y, U, V, density=[0.5, 1])
-ax0.set_title('Varying Density')
+axs[0].streamplot(X, Y, U, V, density=[0.5, 1])
+axs[0].set_title('Varying Density')
 
 # Varying color along a streamline
-ax1 = fig.add_subplot(gs[0, 1])
-strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn')
+strm = axs[1].streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn')
 fig.colorbar(strm.lines)
-ax1.set_title('Varying Color')
+axs[1].set_title('Varying Color')
 
 #  Varying line width along a streamline
-ax2 = fig.add_subplot(gs[1, 0])
 lw = 5*speed / speed.max()
-ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw)
-ax2.set_title('Varying Line Width')
+axs[2].streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw)
+axs[2].set_title('Varying Line Width')
 
 # Controlling the starting points of the streamlines
 seed_points = np.array([[-2, -1, 0, 1, 2, -1], [-2, -1,  0, 1, 2, 2]])
 
-ax3 = fig.add_subplot(gs[1, 1])
-strm = ax3.streamplot(X, Y, U, V, color=U, linewidth=2,
-                      cmap='autumn', start_points=seed_points.T)
+strm = axs[3].streamplot(X, Y, U, V, color=U, linewidth=2,
+                         cmap='autumn', start_points=seed_points.T)
 fig.colorbar(strm.lines)
-ax3.set_title('Controlling Starting Points')
+axs[3].set_title('Controlling Starting Points')
 
 # Displaying the starting points with blue symbols.
-ax3.plot(seed_points[0], seed_points[1], 'bo')
-ax3.set(xlim=(-w, w), ylim=(-w, w))
+axs[3].plot(seed_points[0], seed_points[1], 'bo')
+axs[3].set(xlim=(-w, w), ylim=(-w, w))
 
 # Create a mask
 mask = np.zeros(U.shape, dtype=bool)
 mask[40:60, 40:60] = True
 U[:20, :20] = np.nan
 U = np.ma.array(U, mask=mask)
 
-ax4 = fig.add_subplot(gs[2, 0])
-ax4.streamplot(X, Y, U, V, color='r')
-ax4.set_title('Streamplot with Masking')
+axs[4].streamplot(X, Y, U, V, color='r')
+axs[4].set_title('Streamplot with Masking')
 
-ax4.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray', aspect='auto')
-ax4.set_aspect('equal')
+axs[4].imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray',
+              aspect='auto')
+axs[4].set_aspect('equal')
 
-ax5 = fig.add_subplot(gs[2, 1])
-ax5.streamplot(X, Y, U, V, broken_streamlines=False)
-ax5.set_title('Streamplot with unbroken streamlines')
+axs[5].streamplot(X, Y, U, V, broken_streamlines=False)
+axs[5].set_title('Streamplot with unbroken streamlines')
 
 plt.tight_layout()
 plt.show()
diff --git a/examples/lines_bars_and_markers/linestyles.py b/examples/lines_bars_and_markers/linestyles.py
@@ -65,9 +65,7 @@ def plot_linestyles(ax, linestyles, title):
                     color="blue", fontsize=8, ha="right", family="monospace")
 
 
-ax0, ax1 = (plt.figure(figsize=(10, 8))
-            .add_gridspec(2, 1, height_ratios=[1, 3])
-            .subplots())
+fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(10, 8), height_ratios=[1, 3])
 
 plot_linestyles(ax0, linestyle_str[::-1], title='Named linestyles')
 plot_linestyles(ax1, linestyle_tuple[::-1], title='Parametrized linestyles')
diff --git a/examples/lines_bars_and_markers/psd_demo.py b/examples/lines_bars_and_markers/psd_demo.py
@@ -12,7 +12,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import matplotlib.mlab as mlab
-import matplotlib.gridspec as gridspec
 
 # Fixing random state for reproducibility
 np.random.seed(19680801)
@@ -58,39 +57,44 @@
 y = y + np.random.randn(*t.shape)
 
 # Plot the raw time series
-fig = plt.figure(constrained_layout=True)
-gs = gridspec.GridSpec(2, 3, figure=fig)
-ax = fig.add_subplot(gs[0, :])
-ax.plot(t, y)
-ax.set_xlabel('time [s]')
-ax.set_ylabel('signal')
+fig, axs = plt.subplot_mosaic([
+    ['signal', 'signal', 'signal'],
+    ['zero padding', 'block size', 'overlap'],
+], layout='constrained')
+
+axs['signal'].plot(t, y)
+axs['signal'].set_xlabel('time [s]')
+axs['signal'].set_ylabel('signal')
 
 # Plot the PSD with different amounts of zero padding. This uses the entire
 # time series at once
-ax2 = fig.add_subplot(gs[1, 0])
-ax2.psd(y, NFFT=len(t), pad_to=len(t), Fs=fs)
-ax2.psd(y, NFFT=len(t), pad_to=len(t) * 2, Fs=fs)
-ax2.psd(y, NFFT=len(t), pad_to=len(t) * 4, Fs=fs)
-ax2.set_title('zero padding')
+axs['zero padding'].psd(y, NFFT=len(t), pad_to=len(t), Fs=fs)
+axs['zero padding'].psd(y, NFFT=len(t), pad_to=len(t) * 2, Fs=fs)
+axs['zero padding'].psd(y, NFFT=len(t), pad_to=len(t) * 4, Fs=fs)
 
 # Plot the PSD with different block sizes, Zero pad to the length of the
 # original data sequence.
-ax3 = fig.add_subplot(gs[1, 1], sharex=ax2, sharey=ax2)
-ax3.psd(y, NFFT=len(t), pad_to=len(t), Fs=fs)
-ax3.psd(y, NFFT=len(t) // 2, pad_to=len(t), Fs=fs)
-ax3.psd(y, NFFT=len(t) // 4, pad_to=len(t), Fs=fs)
-ax3.set_ylabel('')
-ax3.set_title('block size')
+axs['block size'].psd(y, NFFT=len(t), pad_to=len(t), Fs=fs)
+axs['block size'].psd(y, NFFT=len(t) // 2, pad_to=len(t), Fs=fs)
+axs['block size'].psd(y, NFFT=len(t) // 4, pad_to=len(t), Fs=fs)
+axs['block size'].set_ylabel('')
 
 # Plot the PSD with different amounts of overlap between blocks
-ax4 = fig.add_subplot(gs[1, 2], sharex=ax2, sharey=ax2)
-ax4.psd(y, NFFT=len(t) // 2, pad_to=len(t), noverlap=0, Fs=fs)
-ax4.psd(y, NFFT=len(t) // 2, pad_to=len(t),
-        noverlap=int(0.05 * len(t) / 2.), Fs=fs)
-ax4.psd(y, NFFT=len(t) // 2, pad_to=len(t),
-        noverlap=int(0.2 * len(t) / 2.), Fs=fs)
-ax4.set_ylabel('')
-ax4.set_title('overlap')
+axs['overlap'].psd(y, NFFT=len(t) // 2, pad_to=len(t), noverlap=0, Fs=fs)
+axs['overlap'].psd(y, NFFT=len(t) // 2, pad_to=len(t),
+                   noverlap=int(0.025 * len(t)), Fs=fs)
+axs['overlap'].psd(y, NFFT=len(t) // 2, pad_to=len(t),
+                   noverlap=int(0.1 * len(t)), Fs=fs)
+axs['overlap'].set_ylabel('')
+axs['overlap'].set_title('overlap')
+
+for title, ax in axs.items():
+    if title == 'signal':
+        continue
+
+    ax.set_title(title)
+    ax.sharex(axs['zero padding'])
+    ax.sharey(axs['zero padding'])
 
 plt.show()
 
