在使用numpy绘图时使用修改过的arctan2的错误

时间:2013-11-05 18:04:06

标签: python python-2.7 numpy matplotlib

使用numpy函数atan2后,我得到下一个错误:

Traceback (most recent call last):
  File "<module1>", line 31, in <module>
  File "<module1>", line 22, in f
  File "<module1>", line 9, in ATN
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

这是我正在使用的代码

import numpy as np
from numpy import exp,arange,log
from pylab import meshgrid,cm,imshow,contour,clabel,colorbar,axis,title,show,streamplot


#Calculating ATN(Y / X)
def ATN(y, x):
    atn = np.arctan2(y, x)
    if atn == 0.0:
        corrected_atn = atn
    elif atn > 0.0:
        corrected_atn = atn
    elif atn < 0:
        corrected_atn = 2 * np.pi + atn
    else:
        print "Error, please check the input numbers"
    return corrected_atn

# PSI = streamline
def f(x, y, U = 8.0, h = 33.0, cs = 137.509870831):
    psi_1 = U * y
    psi_2 = cs * ATN(y - h, x)
    psi_3 = cs * ATN(y + h, x)
    psi = psi_1 + psi_2 + psi_3
    return psi

x = arange(-40.0,40.0,0.1)
y = arange(-40.0,40.0,0.1)
X,Y = meshgrid(x, y) # grid of point
#Z = z_func(X, Y) # evaluation of the function on the grid
Z= f(X, Y)
im = imshow(Z,cmap=cm.RdBu) # drawing the function
dx, dy = 1e-6, 1e-6
U = (f(X+dx, Y) - f(X, Y))/dx
V = (f(X, Y+dy) - f(X, Y))/dy
streamplot(X, Y, U, V, linewidth=1, color=(0, 0, 1, 0.3))

cset = contour(X, Y, Z,arange(-40,40,5.0),linewidths=2,cmap=cm.Set2)
clabel(cset,inline=True,fmt='%1.1f',fontsize=9)
colorbar(im) # adding the colobar on the right

# latex fashion title
title('$phi= 8.0 y + cs * ATN(y - h, x) + cs * ATN(y + h, x) $')

show()

我用函数ATN(y,x)数字得到的值是正确的,但我不知道如何处理这个歧义问题

2 个答案:

答案 0 :(得分:3)

您正在将数组传递给ATN函数,因此您需要处理调用np.arctan2时返回的数组,请尝试以下操作:

def ATN(y, x):
    atn = np.arctan2(y, x)
    atn[atn < 0] += 2*np.pi
    return atn

答案 1 :(得分:1)

另一种解决方案是使用模运算符:

def ATN(y, x):
    return np.arctan2(y,x) % (2*np.pi)

或等效,

def ATN(y, x):
    return np.mod(np.arctan2(y,x), 2*np.pi)