Question

我目前在方法中使用以下代码来绘制numpy.array。可以根据绘图标签的选择来选择要绘制的轴。这段代码对我来说看起来非常多余，我想知道是否有办法让它更紧凑而没有（！）影响理解它的能力。

    if axes==1:    
        py.plot(par.x,np.real(self.psi[ni,:,0,0]),'ro-',
                      par.x,np.imag(self.psi[ni,:,0,0]),'bd-')
        py.xlabel('x')
        py.legend(('$\Re[\psi(i=%d,x,y=0,z=0)]$'%ni,
                            '$\Im[\psi(i=%d,x,y=0,z=0)]$'%ni),loc=0)

    elif axes==2:
        py.plot(par.y,np.real(self.psi[ni,0,:,0]),'ro-',
                      par.y,np.imag(self.psi[ni,0,:,0]),'bd-')
        py.xlabel('y')
        py.legend(('$\Re[\psi(i=%d,x=0,y,z=0)]$'%ni, 
                            '$\Im[\psi(i=%d,x=0,y,z=0)]$'%ni),loc=0)

    elif axes==3:
        py.plot(par.z,np.real(self.psi[ni,0,0,:]),'ro-',
                      par.z,np.imag(self.psi[ni,0,0,:]),'bd-')
        py.xlabel('z')
        py.legend(('$\Re[\psi(i=%d,x=0,y=0,z)]$'%ni, 
                            '$\Im[\psi(i=%d,x=0,y=0,z)]$'%ni),loc=0)

Answer 1

您可以将self.psi[ni,:,0,0]重写为：

idx = [ni, 0, 0, 0]
idx[axes] = slice(None)
self.psi[tuple(idx)]

甚至是：

idx = (ni,) + (0,)*(axes-1) + (slice(None),) + (0,)*(3-axes)
sel.psi[idx]

对于x轴的标记，我认为'xyz'[axes-1]会起作用，你可以重复使用这个想法来重写例如par.x为geattr(par, 'xyz'[axes-1])

Answer 2

您可以拥有要提供给plot，xlabel，范围和图例的值的字典，例如：

Axdict = {
    1:(par.x, 'x', 'ni,:,0,0', '(i=%d,x,y=0,z=0)'), 
    2:(par.y, 'y', 'ni,0,:,0', '(i=%d,x=0,y,z=0)'),
    3:(par.z, 'z', 'ni,0,0,:', '(i=%d,x=0,y=0,z)'),
}
p, lx, rs, l = Axdict.get(axes)   # Look up your subst
cdata = eval("self.psi[%s]" % rs) # Get the data you need
py.plot(p,np.real(cdata),'ro-',
        p,np.imag(cdata),'bd-')   # Plot it
py.xlabel(lx)                     # Use the label
lr = '$\Re[\psi()]$' % l          # Real Legend
li = '$\Im[\psi()]$' % l          # Imaginary Legend
py.legend((lr%ni, li%ni),loc=0)   # Add the legend

避免使用冗余代码绘制阵列的不同轴

2 个答案: