I'm trying to plot some complex functions using numpy. Example of some working code:
import numpy as np
from PIL import Image
size = 1000
w = np.linspace(-10, 10, size)
x, y = np.meshgrid(w, w)
r = x + 1j*y
def f(q):
return np.angle(q)
z = f(r)
normalized = ((255/(np.amax(z) - np.amin(z)))*(z+abs(np.amin(z)))).astype(int)
data = [i for j in normalized for i in j]
img = Image.new('L', (size, size))
img.putdata(data[::-1]) #pixels are done bottom to top
img.show()
However, suppose I want the function f to have a simple comparison in it, like this:
def f(q):
if np.abs(q) < 4:
return 1
else:
return 0
I get the error
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
For the np.abs(q) < 4 check.
I did some digging and realized it's because Python is doing the operation on the entire r array, and it can't compare an array to an integer. So, I tried looking for ways to do element-wise comparisons.
This page looked promising: it says I can do element-wise comparisons by using np.less(a, b), so I tried
def f(q):
if np.less(np.abs(q), 4):
return 1
else:
return 0
and got the same ValueError. It seems as though both arguments for np.less() need to be arrays of the same size.
What I want is to compare each element of my array to a single, non-array quantity. I suppose I could make a dummy array of the same size filled with identical 4's, but there has to be a more elegant way of doing this.