Can Numpy's frompyfunc() return an image?

Question

I want to build a NxN grid in the complex plane, and assign each point z in this grid a RGB color according to some rule f(z).

For example, the grid z is

x,y = numpy.ogrid[-1:1:N*1j,-1:1:N*1j]
z = x+y*1j

and the function setcolor() is

def setcolor(z):
    return (r,g,b) triple of z according to some rule

My problem occurs when I was going to display the image of f(z) with numpy's frompyfunc():

img = numpy.frompyfunc(setcolor,1,1)(z).astype(np.float)
plt.imshow(img)
plt.show()

The resulting imag is all blue (by default colormap 'jet'), no error nor warning, so it's obvious that imshow() has turned all (r,g,b) triple into a single float maybe 1 (or 0). I guess this is because of the astype(np.float) option, but it seems this is the only reasonable option.

Of course one can use two for loops to display img:

for i in range(N):
    for j in range(N):
        img[i,j] = f(z[i,j])

but this is not very efficient(my own opinion). I want to use the fromptfunc() directly.

So how I can display this image with frompyfunc()?

What does `numpy.frompyfunc(setcolor,1,1)(z)` look like (for small `N`)? `frompyfunc`, if working, might give as much as a 2x speed up, but nothing like a genuine vectorization. — hpaulj, Mar 29 '16 at 23:17
[`numpy.frompyfunc` seems to have issues with functions that already return sequences.](http://ideone.com/dPhCR9) Looks like it might be a NumPy bug. You probably shouldn't be using `frompyfunc` anyway, though. — user2357112, Mar 29 '16 at 23:32
@JoranBeasley: I compared the frompyfunc() and the "for loops", the former is significantly faster than the latter. May be there is some other way to speed up the code? — zhaoL, Mar 29 '16 at 23:50

score 0 · Accepted Answer · edited May 23 '17 at 12:24

In [588]: x,y=np.ogrid[-1:1:5*1j,-1:1:5*1j]
In [589]: z=x*y*1j
In [590]: def setcolor(z):
    return (z.real,z.imag,np.abs(z))
   .....: 
In [591]: fn=np.frompyfunc(setcolor,1,1)
In [592]: fn(z)
Out[592]: 
array([[None, None, None, None, None],
       [None, None, None, None, None],
       [None, None, None, None, None],
       [None, None, None, None, None],
       [None, None, None, None, None]], dtype=object)

which turns to an array of 0s with .astype(float). Here setcolor is just a simple way of getting 3 numbers from a complex input. You can choose something more realistic.

But if we give it the correct out count, we get a tuple of arrays.

In [593]: fn=np.frompyfunc(setcolor,1,3)
In [594]: fn(z)
Out[594]: 
(array([[0.0, 0.0, -0.0, -0.0, -0.0],
        [0.0, 0.0, -0.0, -0.0, -0.0],
        [-0.0, -0.0, 0.0, 0.0, 0.0],
        [-0.0, -0.0, 0.0, 0.0, 0.0],
        [-0.0, -0.0, 0.0, 0.0, 0.0]], dtype=object),
 array([[1.0, 0.5, 0.0, -0.5, -1.0],
        [0.5, 0.25, 0.0, -0.25, -0.5],
        [0.0, 0.0, 0.0, 0.0, 0.0],
        [-0.5, -0.25, 0.0, 0.25, 0.5],
        [-1.0, -0.5, 0.0, 0.5, 1.0]], dtype=object),
 array([[1.0, 0.5, 0.0, 0.5, 1.0],
        [0.5, 0.25, 0.0, 0.25, 0.5],
        [0.0, 0.0, 0.0, 0.0, 0.0],
        [0.5, 0.25, 0.0, 0.25, 0.5],
        [1.0, 0.5, 0.0, 0.5, 1.0]], dtype=object))

It could be turned in to a (3,N,N) array with:

np.array([A.astype(float) for A in fn(z)])

You'd have to do a transpose to make a (N,N,3) array.

For this small sample it isn't faster

In [599]: timeit np.array([A.astype(float) for A in fn(z)]).transpose([1,2,0])
1000 loops, best of 3: 664 µs per loop

In [603]: %%timeit 
   .....: img=np.zeros((5,5,3),float)
   .....: for i in range(5):
    for j in range(5):
        img[i,j]=setcolor(z[i,j])
   .....: 
1000 loops, best of 3: 316 µs per loop

In some other SO questions I found that frompyfunc is faster than vectorize, and a modest improvement over direct iteration.

This is faster - but it doesn't do the job

In [606]: timeit np.frompyfunc(setcolor,1,1)(z).astype(float)
10000 loops, best of 3: 25.6 µs per loop

The setcolor that I chose can take the array z directly, resulting in a fast 3d array creation.

In [608]: timeit np.array([A.astype(float) for A in setcolor(z)]).transpose([1,2,0])
10000 loops, best of 3: 47.1 µs per loop

From https://stackoverflow.com/a/29703463/901925, an earlier answer involving frompystack, I find that vstack is a faster way to assemble a tuple into a 3d array:

In [616]: timeit np.vstack(fn(z)).astype(float).reshape(3,5,5).transpose([1,2,0])
10000 loops, best of 3: 178 µs per loop

So done right, frompyfunc can give about a 2x speedup over explicit iteration, but still not as good as function that handles arrays directly.

score 0 · Answer 2 · answered Mar 29 '16 at 23:53

I left out the (rgb) tuple which is separate from the main question, on how to make an efficient assignment with numpy. The magic is all in using meshgrid:

import pylab as p
import numpy as np 

f= lambda z :  np.cos(z)  # example

X = np.arange(-5, 5, 0.1)
Y = np.arange(-5, 5, 0.1)
X, Y = np.meshgrid(X, Y)
xn, yn = X.shape
W = np.zeros_like(X)

for x in range(xn):
    for y in range(yn):
        z = complex(X[x,y],Y[x,y])
        w = float(f(z))
        W[x,y] = w

p.imshow(W)

Can Numpy's frompyfunc() return an image?

2 Answers2

Linked