CODEWITHSUNDEEP
pythonpython-2.7numpymatplotlib
jenko_cp
jenko_cp

Reputation: 103

Errors using modified arctan2 in plotting with numpy

After using the numpy function atan2 I get the next error:

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()

This is the code that I'm using

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()

The values that I get numerically with the function ATN(y,x) are correct but I don't figure out how to deal with this issue of ambiguity

Upvotes: 1

Views: 365

Answers (2)

askewchan
askewchan

Reputation: 46530

An alternative solution is to use the modulo operator:

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

or equivalently,

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

Upvotes: 1

Jaime
Jaime

Reputation: 67417

You are passing arrays to your ATN function, so you need to handle the array returned by your call to np.arctan2, try the following:

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

Upvotes: 3

Related Questions