Merge pull request #23497 from QuLogic/less-gridspec-demo

timhoffm · web-flow · commit 37300c359533 · 2022-08-12T08:32:45.000+02:00
Avoid gridspec in more examples
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()
 
diff --git a/examples/text_labels_and_annotations/demo_text_rotation_mode.py b/examples/text_labels_and_annotations/demo_text_rotation_mode.py
@@ -20,23 +20,14 @@
 """
 
 import matplotlib.pyplot as plt
-import numpy as np
 
 
-def test_rotation_mode(fig, mode, subplot_location):
+def test_rotation_mode(fig, mode):
     ha_list = ["left", "center", "right"]
     va_list = ["top", "center", "baseline", "bottom"]
-    axs = np.empty((len(va_list), len(ha_list)), object)
-    gs = subplot_location.subgridspec(*axs.shape, hspace=0, wspace=0)
-    axs[0, 0] = fig.add_subplot(gs[0, 0])
-    for i in range(len(va_list)):
-        for j in range(len(ha_list)):
-            if (i, j) == (0, 0):
-                continue  # Already set.
-            axs[i, j] = fig.add_subplot(
-                gs[i, j], sharex=axs[0, 0], sharey=axs[0, 0])
-    for ax in axs.flat:
-        ax.set(aspect=1)
+    axs = fig.subplots(len(va_list), len(ha_list), sharex=True, sharey=True,
+                       subplot_kw=dict(aspect=1),
+                       gridspec_kw=dict(hspace=0, wspace=0))
 
     # labels and title
     for ha, ax in zip(ha_list, axs[-1, :]):
@@ -50,6 +41,8 @@ def test_rotation_mode(fig, mode, subplot_location):
         {"bbox": dict(boxstyle="square,pad=0.", ec="none", fc="C1", alpha=0.3)}
     )
 
+    texts = {}
+
     # use a different text alignment in each axes
     for i, va in enumerate(va_list):
         for j, ha in enumerate(ha_list):
@@ -64,22 +57,22 @@ def test_rotation_mode(fig, mode, subplot_location):
                          size="x-large", rotation=40,
                          horizontalalignment=ha, verticalalignment=va,
                          rotation_mode=mode, **kw)
+            texts[ax] = tx
 
     if mode == "default":
         # highlight bbox
         fig.canvas.draw()
-        for ax in axs.flat:
-            text, = ax.texts
+        for ax, text in texts.items():
             bb = text.get_window_extent().transformed(ax.transData.inverted())
             rect = plt.Rectangle((bb.x0, bb.y0), bb.width, bb.height,
                                  facecolor="C1", alpha=0.3, zorder=2)
             ax.add_patch(rect)
 
 
 fig = plt.figure(figsize=(8, 5))
-gs = fig.add_gridspec(1, 2)
-test_rotation_mode(fig, "default", gs[0])
-test_rotation_mode(fig, "anchor", gs[1])
+subfigs = fig.subfigures(1, 2)
+test_rotation_mode(subfigs[0], "default")
+test_rotation_mode(subfigs[1], "anchor")
 plt.show()
 
 
