Merge pull request #14079 from timhoffm/subplots-flat

jklymak · web-flow · commit dc043e159387 · 2019-04-29T14:02:08.000-07:00
Consistently use axs.flat instead of axs.flatten()
diff --git a/examples/specialty_plots/radar_chart.py b/examples/specialty_plots/radar_chart.py
@@ -168,7 +168,7 @@ def example_data():
 
     colors = ['b', 'r', 'g', 'm', 'y']
     # Plot the four cases from the example data on separate axes
-    for ax, (title, case_data) in zip(axes.flatten(), data):
+    for ax, (title, case_data) in zip(axes.flat, data):
         ax.set_rgrids([0.2, 0.4, 0.6, 0.8])
         ax.set_title(title, weight='bold', size='medium', position=(0.5, 1.1),
                      horizontalalignment='center', verticalalignment='center')
diff --git a/examples/statistics/boxplot.py b/examples/statistics/boxplot.py
@@ -48,7 +48,7 @@
 axes[1, 2].boxplot(data, labels=labels, showfliers=False)
 axes[1, 2].set_title('showfliers=False', fontsize=fs)
 
-for ax in axes.flatten():
+for ax in axes.flat:
     ax.set_yscale('log')
     ax.set_yticklabels([])
 
@@ -88,7 +88,7 @@
 axes[1, 2].boxplot(data, whis=[15, 85])
 axes[1, 2].set_title('whis=[15, 85]\n#percentiles', fontsize=fs)
 
-for ax in axes.flatten():
+for ax in axes.flat:
     ax.set_yscale('log')
     ax.set_yticklabels([])
 
diff --git a/examples/statistics/bxp.py b/examples/statistics/bxp.py
@@ -62,7 +62,7 @@
 axes[1, 2].bxp(stats, showfliers=False)
 axes[1, 2].set_title('showfliers=False', fontsize=fs)
 
-for ax in axes.flatten():
+for ax in axes.flat:
     ax.set_yscale('log')
     ax.set_yticklabels([])
 
@@ -95,7 +95,7 @@
                showmeans=True)
 axes[1, 1].set_title('Custom mean\nas line', fontsize=fs)
 
-for ax in axes.flatten():
+for ax in axes.flat:
     ax.set_yscale('log')
     ax.set_yticklabels([])
 
diff --git a/examples/statistics/violinplot.py b/examples/statistics/violinplot.py
@@ -57,7 +57,7 @@
                       bw_method=0.5)
 axes[1, 2].set_title('Custom violinplot 6', fontsize=fs)
 
-for ax in axes.flatten():
+for ax in axes.flat:
     ax.set_yticklabels([])
 
 fig.suptitle("Violin Plotting Examples")
diff --git a/examples/subplots_axes_and_figures/demo_constrained_layout.py b/examples/subplots_axes_and_figures/demo_constrained_layout.py
@@ -27,15 +27,15 @@ def example_plot(ax):
 
 fig, axs = plt.subplots(nrows=2, ncols=2, constrained_layout=False)
 
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax)
 
 ###############################################################################
 # adding ``constrained_layout=True`` automatically adjusts.
 
 fig, axs = plt.subplots(nrows=2, ncols=2, constrained_layout=True)
 
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax)
 
 ###############################################################################
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -275,11 +275,10 @@ def test_minorticks_on_rcParams_both(fig_test, fig_ref):
 @image_comparison(baseline_images=["autoscale_tiny_range"], remove_text=True)
 def test_autoscale_tiny_range():
     # github pull #904
-    fig, ax = plt.subplots(2, 2)
-    ax = ax.flatten()
-    for i in range(4):
+    fig, axs = plt.subplots(2, 2)
+    for i, ax in enumerate(axs.flat):
         y1 = 10**(-11 - i)
-        ax[i].plot([0, 1], [1, 1 + y1])
+        ax.plot([0, 1], [1, 1 + y1])
 
 
 @pytest.mark.style('default')
@@ -358,7 +357,7 @@ def test_arrow_simple():
     kwargs = product(length_includes_head, shape, head_starts_at_zero)
 
     fig, axs = plt.subplots(3, 4)
-    for i, (ax, kwarg) in enumerate(zip(axs.flatten(), kwargs)):
+    for i, (ax, kwarg) in enumerate(zip(axs.flat, kwargs)):
         ax.set_xlim(-2, 2)
         ax.set_ylim(-2, 2)
         # Unpack kwargs
diff --git a/lib/matplotlib/tests/test_constrainedlayout.py b/lib/matplotlib/tests/test_constrainedlayout.py
@@ -47,7 +47,7 @@ def test_constrained_layout1():
 def test_constrained_layout2():
     'Test constrained_layout for 2x2 subplots'
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         example_plot(ax, fontsize=24)
 
 
@@ -56,7 +56,7 @@ def test_constrained_layout2():
 def test_constrained_layout3():
     'Test constrained_layout for colorbars with subplots'
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for nn, ax in enumerate(axs.flatten()):
+    for nn, ax in enumerate(axs.flat):
         pcm = example_pcolor(ax, fontsize=24)
         if nn == 3:
             pad = 0.08
@@ -69,7 +69,7 @@ def test_constrained_layout3():
 def test_constrained_layout4():
     'Test constrained_layout for a single colorbar with subplots'
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         pcm = example_pcolor(ax, fontsize=24)
     fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
 
@@ -82,7 +82,7 @@ def test_constrained_layout5():
     colorbar bottom
     '''
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         pcm = example_pcolor(ax, fontsize=24)
     fig.colorbar(pcm, ax=axs,
                  use_gridspec=False, pad=0.01, shrink=0.6,
@@ -162,7 +162,7 @@ def test_constrained_layout9():
     'Test for handling suptitle and for sharex and sharey'
     fig, axs = plt.subplots(2, 2, constrained_layout=True,
                             sharex=False, sharey=False)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         pcm = example_pcolor(ax, fontsize=24)
         ax.set_xlabel('')
         ax.set_ylabel('')
@@ -176,7 +176,7 @@ def test_constrained_layout9():
 def test_constrained_layout10():
     'Test for handling legend outside axis'
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         ax.plot(np.arange(12), label='This is a label')
     ax.legend(loc='center left', bbox_to_anchor=(0.8, 0.5))
 
@@ -250,7 +250,7 @@ def test_constrained_layout12():
 def test_constrained_layout13():
     'Test that padding works.'
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         pcm = example_pcolor(ax, fontsize=12)
         fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
     fig.set_constrained_layout_pads(w_pad=24./72., h_pad=24./72.)
@@ -261,7 +261,7 @@ def test_constrained_layout13():
 def test_constrained_layout14():
     'Test that padding works.'
     fig, axs = plt.subplots(2, 2, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         pcm = example_pcolor(ax, fontsize=12)
         fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
     fig.set_constrained_layout_pads(
@@ -275,7 +275,7 @@ def test_constrained_layout15():
     'Test that rcparams work.'
     rcParams['figure.constrained_layout.use'] = True
     fig, axs = plt.subplots(2, 2)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         example_plot(ax, fontsize=12)
 
 
@@ -393,7 +393,7 @@ def test_colorbar_location():
     """
 
     fig, axs = plt.subplots(4, 5, constrained_layout=True)
-    for ax in axs.flatten():
+    for ax in axs.flat:
         pcm = example_pcolor(ax)
         ax.set_xlabel('')
         ax.set_ylabel('')
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -206,7 +206,7 @@ def test_given_colors_levels_and_extends():
     colors = ['red', 'yellow', 'pink', 'blue', 'black']
     levels = [2, 4, 8, 10]
 
-    for i, ax in enumerate(axes.flatten()):
+    for i, ax in enumerate(axes.flat):
         filled = i % 2 == 0.
         extend = ['neither', 'min', 'max', 'both'][i // 2]
 
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -360,23 +360,23 @@ def test_invalid_figure_add_axes():
 
 
 def test_subplots_shareax_loglabels():
-    fig, ax_arr = plt.subplots(2, 2, sharex=True, sharey=True, squeeze=False)
-    for ax in ax_arr.flatten():
+    fig, axs = plt.subplots(2, 2, sharex=True, sharey=True, squeeze=False)
+    for ax in axs.flat:
         ax.plot([10, 20, 30], [10, 20, 30])
 
     ax.set_yscale("log")
     ax.set_xscale("log")
 
-    for ax in ax_arr[0, :]:
+    for ax in axs[0, :]:
         assert 0 == len(ax.xaxis.get_ticklabels(which='both'))
 
-    for ax in ax_arr[1, :]:
+    for ax in axs[1, :]:
         assert 0 < len(ax.xaxis.get_ticklabels(which='both'))
 
-    for ax in ax_arr[:, 1]:
+    for ax in axs[:, 1]:
         assert 0 == len(ax.yaxis.get_ticklabels(which='both'))
 
-    for ax in ax_arr[:, 0]:
+    for ax in axs[:, 0]:
         assert 0 < len(ax.yaxis.get_ticklabels(which='both'))
 
 
diff --git a/tutorials/intermediate/constrainedlayout_guide.py b/tutorials/intermediate/constrainedlayout_guide.py
@@ -90,15 +90,15 @@ def example_plot(ax, fontsize=12, nodec=False):
 # axes overlapping each other.
 
 fig, axs = plt.subplots(2, 2, constrained_layout=False)
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax)
 
 ###############################################################################
 # Specifying ``constrained_layout=True`` in the call to ``plt.subplots``
 # causes the layout to be properly constrained.
 
 fig, axs = plt.subplots(2, 2, constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax)
 
 ###############################################################################
@@ -132,7 +132,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 # the specified axes.
 
 fig, axs = plt.subplots(2, 2, figsize=(4, 4), constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     im = ax.pcolormesh(arr, **pc_kwargs)
 fig.colorbar(im, ax=axs, shrink=0.6)
 
@@ -142,7 +142,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 # still be the same size.
 
 fig, axs = plt.subplots(3, 3, figsize=(4, 4), constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     im = ax.pcolormesh(arr, **pc_kwargs)
 fig.colorbar(im, ax=axs[1:, ][:, 1], shrink=0.8)
 fig.colorbar(im, ax=axs[:, -1], shrink=0.6)
@@ -179,7 +179,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 # ``constrained_layout`` can also make room for `~.figure.Figure.suptitle`.
 
 fig, axs = plt.subplots(2, 2, figsize=(4, 4), constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     im = ax.pcolormesh(arr, **pc_kwargs)
 fig.colorbar(im, ax=axs, shrink=0.6)
 fig.suptitle('Big Suptitle')
@@ -265,15 +265,15 @@ def example_plot(ax, fontsize=12, nodec=False):
 # `~.figure.Figure.set_constrained_layout_pads`:
 
 fig, axs = plt.subplots(2, 2, constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax, nodec=True)
     ax.set_xticklabels('')
     ax.set_yticklabels('')
 fig.set_constrained_layout_pads(w_pad=4./72., h_pad=4./72.,
         hspace=0., wspace=0.)
 
 fig, axs = plt.subplots(2, 2, constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax, nodec=True)
     ax.set_xticklabels('')
     ax.set_yticklabels('')
@@ -288,7 +288,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 # the above, but the space between subplots does.
 
 fig, axs = plt.subplots(2, 2, constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax, nodec=True)
     ax.set_xticklabels('')
     ax.set_yticklabels('')
@@ -308,7 +308,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 # of the axis it is attached to.
 
 fig, axs = plt.subplots(2, 2, constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     pc = ax.pcolormesh(arr, **pc_kwargs)
     fig.colorbar(pc, ax=ax, shrink=0.6, pad=0)
     ax.set_xticklabels('')
@@ -322,7 +322,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 # for ``pad`` to be non-zero.
 
 fig, axs = plt.subplots(2, 2, constrained_layout=True)
-for ax in axs.flatten():
+for ax in axs.flat:
     pc = ax.pcolormesh(arr, **pc_kwargs)
     fig.colorbar(im, ax=ax, shrink=0.6, pad=0.05)
     ax.set_xticklabels('')
@@ -347,7 +347,7 @@ def example_plot(ax, fontsize=12, nodec=False):
 
 plt.rcParams['figure.constrained_layout.use'] = True
 fig, axs = plt.subplots(2, 2, figsize=(3, 3))
-for ax in axs.flatten():
+for ax in axs.flat:
     example_plot(ax)
 
 #############################
@@ -739,10 +739,10 @@ def docomplicated(suptitle=None):
 # layoutboxes in the gridspec, and it is made to be centered between
 # those two points.
 
-fig, ax = plt.subplots(2, 2, constrained_layout=True)
-for a in ax.flatten():
-    im = a.pcolormesh(arr, **pc_kwargs)
-fig.colorbar(im, ax=ax, shrink=0.6)
+fig, axs = plt.subplots(2, 2, constrained_layout=True)
+for ax in axs.flat:
+    im = ax.pcolormesh(arr, **pc_kwargs)
+fig.colorbar(im, ax=axs, shrink=0.6)
 plot_children(fig, fig._layoutbox, printit=False)
 
 #######################################################################
diff --git a/tutorials/text/text_intro.py b/tutorials/text/text_intro.py
@@ -329,24 +329,23 @@
 # labels fit in the right-hand plot.
 
 fig, axs = plt.subplots(2, 2, figsize=(8, 5), tight_layout=True)
-axs = axs.flatten()
-for n, ax in enumerate(axs):
+for n, ax in enumerate(axs.flat):
     ax.plot(x1*10., y1)
 
 formatter = matplotlib.ticker.FormatStrFormatter('%1.1f')
 locator = matplotlib.ticker.MaxNLocator(nbins='auto', steps=[1, 4, 10])
-axs[1].xaxis.set_major_locator(locator)
-axs[1].xaxis.set_major_formatter(formatter)
+axs[0, 1].xaxis.set_major_locator(locator)
+axs[0, 1].xaxis.set_major_formatter(formatter)
 
 formatter = matplotlib.ticker.FormatStrFormatter('%1.5f')
 locator = matplotlib.ticker.AutoLocator()
-axs[2].xaxis.set_major_formatter(formatter)
-axs[2].xaxis.set_major_locator(locator)
+axs[1, 0].xaxis.set_major_formatter(formatter)
+axs[1, 0].xaxis.set_major_locator(locator)
 
 formatter = matplotlib.ticker.FormatStrFormatter('%1.5f')
 locator = matplotlib.ticker.MaxNLocator(nbins=4)
-axs[3].xaxis.set_major_formatter(formatter)
-axs[3].xaxis.set_major_locator(locator)
+axs[1, 1].xaxis.set_major_formatter(formatter)
+axs[1, 1].xaxis.set_major_locator(locator)
 
 plt.show()
 
