Array broadcasting with numpy

Question

How do I write the following loop using Python's implicit looping?

def kl(myA, myB, a, b):
    lots of stuff that assumes all inputs are scalars

x, y = meshgrid(inclusive_arange(0.0, xsize, 0.10),\
                inclusive_arange(0.0, ysize, 0.10))

for j in range(x.shape[0]):
    for i in range(x.shape[1]):
        z[j, i] = kl(x[j, i], y[j, i])

I want to do something like

z = kl(x, y)

but that gives:

TypeError: only length-1 arrays can be converted to Python scalars

Answer 1

The capability you're asking about only exists in Numpy, and it's called array broadcasting, not implicit looping. A function that broadcasts a scalar operation over an array is called a universal function, or ufunc. Many basic Numpy functions are of this type.

You can use numpy.frompyfunc to convert your existing function kl into a ufunc.

kl_ufunc = numpy.frompyfunc(kl, 4, 1)
...
z = kl_ufunc(x + 1.0, y + 1.0, myA, myB)

Of course, if you want, you could call the ufunc kl instead of kl_ufunc, but then the original definition of kl would be lost. That might be fine for your purposes.

Answer 2

There is a video series here which you might find useful:

http://showmedo.com/videotutorials/video?name=10370070&fromSeriesID=1037

Note that it is part of a tutorial series that discusses a broad range of numpy topics.

Just FYI.