2

我目前在一个方法中使用以下代码来绘制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)
4

2 回答 2

2

你可以重写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])

于 2013-09-18T13:24:17.457 回答
1

您可以将值的字典提供给绘图、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
于 2013-09-18T11:44:33.793 回答