如何在 python 中为子图运行智能循环

Posted 2023-03-27

【中文标题】如何在 python 中为子图运行智能循环

【英文标题】：How to run a smart loop for subplots in python

【发布时间】：2017-08-08 07:47:12

【问题描述】：

所以我正在尝试创建一个由 9 个图表组成的 3x3 框，虽然我已经设法通过为每个单独的框手动编写代码来做到这一点，但我想学习如何使用循环来做到这一点。我似乎无法弄清楚。目前，我正在使用以下内容：

from matplotlib import gridspec

f, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row')

gs = gridspec.GridSpec(3, 3)
fig = plt.figure(figsize=(20,20))

fig.text(0.5, .95, 'Constant Slope for [O/Fe]/[Fe/H] for Various R and Z', ha='center', va='center', size = 50)
fig.text(0.5, 0.08, '[Fe/H]', ha='center', va='center', size = 60)
fig.text(0.09, 0.5, '[O/Fe]', ha='center', va='center', rotation='vertical', size = 60)

ax1 = plt.subplot(gs[0])
histogram1 = ax1.hist2d(fehsc, ofesc, bins=nbins, range=[[-1,.5],[0.2,0.4]])
counts = histogram1[0]
xpos = histogram1[1]
ypos = histogram1[2]
image = histogram1[3]
newcounts = counts #we're going to iterate over this

for i in range (nbins):
    xin = xpos[i]
    yin = ypos
    yline = m*xin + b
    reset = np.where(yin < yline) #anything less than yline we want to be 0
    #index = index[0:len(index)-1]  
    countout = counts[i]
    countout[reset] = 0
    newcounts[i] = countout
ax1.plot(xarr2, yarr2, color='w', linewidth='5', alpha = 0.3)
ax1.plot(xarr, yarr, color='r')
ax1.set_title('R in [5,7] kpc | Z in [1,2] kpc', size = 20)

ax2 = plt.subplot(gs[1])
histogram2 = ax2.hist2d(fehsc2, ofesc2, bins=nbins, range=[[-1,.5],[0.2,0.4]])
counts = histogram2[0]
xpos = histogram2[1]
ypos = histogram2[2]
image = histogram2[3]
newcounts = counts #we're going to iterate over this

for i in range (nbins):
    xin = xpos[i]
    yin = ypos
    yline = m*xin + b
    reset = np.where(yin < yline) #anything less than yline we want to be 0
    #index = index[0:len(index)-1]  
    countout = counts[i]
    countout[reset] = 0
    newcounts[i] = countout
ax2.plot(xarr2, yarr2, color='w', linewidth='5', alpha = 0.3)
ax2.plot(xarr, yarr, color='r')
ax2.set_title('R in [7,9] kpc | Z in [1,2] kpc', size = 20)

以此类推，直到 ax9。

我尝试做的事情如下：

for k in range(1,10):
    ax[k] = plt.subplot(gs[0])
    histogram1 = ax[k].hist2d(fehsc, ofesc, bins=nbins, range=[[-1,.5],[0.2,0.4]])
    counts = histogram1[0]
    xpos = histogram1[1]
    ypos = histogram1[2]
    image = histogram1[3]
    newcounts = counts #we're going to iterate over this

    for i in range (nbins):
        xin = xpos[i]
        yin = ypos
        yline = m*xin + b
        reset = np.where(yin < yline) #anything less than yline we want to be 0
        countout = counts[i]
        countout[reset] = 0
        newcounts[i] = countout
    ax[k].plot(xarr2, yarr2, color='w', linewidth='5', alpha = 0.3)
    ax[k].plot(xarr, yarr, color='r')
    ax[k].set_title('R in [5,7] kpc | Z in [1,2] kpc', size = 20)

因为我想在一个循环中使用 ax(k) 并一次运行所有九次迭代。但显然这不是这样做的方法，或者只是行不通。是否可以在调用 ax_ 时将其从 1 循环到 9？

【参考方案1】：

迭代许多子图的一种简单方法是通过

创建它们

fig, axes = plt.subplots(nrows=n, ncols=m)

axes 是n 行和m 列的数组。然后可以独立地遍历行和列，或者遍历轴数组的扁平版本。

本例中显示的是后者：

import matplotlib.pyplot as plt
import numpy as np

x = np.arange(11)
y = np.random.rand(len(x), 9)*10

fig, axes = plt.subplots(3,3, sharex=True, sharey=True)

for i, ax in enumerate(axes.flatten()):
    ax.bar(x, y[:,i], color=plt.cm.Paired(i/10.))

plt.show()

【参考方案2】：

import numpy as np
import matplotlib.pyplot as plt

# made some data for you. 
# xx and yy is list of numpy arrays. in your case it will [fehsc1, fehsc3, ...fehsc9] and [ofesc1, ofesc2...ofesc9]
xx =  [np.random.randn(100000) for i in range(9)]
yy =  [np.random.randn(100000) + 5 for i in range(9)]
nbins = 40
m = 1
b = 0.5

f, ax = plt.subplots(3, 3, sharex='col', sharey='row')
# because ax = [[ax0,ax1,ax2],[ax3,ax4,ax5],[ax7,ax8,ax9]], ie list of lists we have unravel it

axr = ax.ravel()
# now axr = [ax0,ax1,ax2,ax3,ax4,ax5,ax7,ax8,ax9]

fig = plt.figure(figsize=(20,20))

for i in range(len(xx)):
    histogram1 = axr[i].hist2d(xx[i], yy[i], bins=nbins)
    counts = histogram1[0]
    xpos = histogram1[1]
    ypos = histogram1[2]
    image = histogram1[3]
    newcounts = counts #we're going to iterate over this

    for k in range (nbins):
        xin = xpos[k]
        yin = ypos
        yline = m*xin + b
        reset = np.where(yin < yline) #anything less than yline we want to be 0
        #index = index[0:len(index)-1]  
        countout = counts[k]
        countout[reset] = 0
        newcounts[k] = countout
    #axr[i].plot(xarr2, yarr2, color='w', linewidth='5', alpha = 0.3)
plt.show()