0

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.

dain
  • 672
  • 1
  • 7
  • 22

1 Answers1

0

The key is to return an array value instead of trying to coerce an array into a single bool, which is what if (some_array): keeps trying to do. There being no unambiguous way to decide what single boolean np.array([True, False]) should convert to, it doesn't even try.

So don't even branch:

def f(q):
    return abs(q) < 4

gives an array like

>>> f(np.array([1,3,5]))
array([ True,  True, False], dtype=bool)

which as numbers will behave like

>>> f(np.array([1,3,5])).astype(int)
array([1, 1, 0])

and give

circle

DSM
  • 342,061
  • 65
  • 592
  • 494
  • What if I need it to branch though? The real problem I'm trying to do is plot a Mandelbrot fractal. That requires checking whether or not each point diverges under iteration of some function, and I need it to do different things in either case. – dain Mar 01 '17 at 00:56
  • @dain: then you either need to rewrite your formulae in a way which doesn't need an individual branch (e.g. something like `np.where(np.array([1,2,3]) > 2, 10, 11)`), or write the equation in scalar form and then use `np.vectorize`. You'll lose some perf. with the latter. Aside: issues like these are why it's generally better to keep your post resembling your actual goal. If you simplify too early, then you get an answer to the question you actually asked ("What I want is to compare each element of my array to a single, non-array quantity"), instead of the one you want but didn't ask. – DSM Mar 01 '17 at 01:01