Python: An elegant/efficient way to evaluate function over bi-dimensional indexes?

Question

I am very new to Python (in the past I used Mathematica, Maple, or Matlab scripts). I am very impressed how NumPy can evaluate functions over arrays but having problems trying to implement it in several dimensions. My question is very simple (please don't laugh): is there a more elegant and efficient way to evaluate some function f (which is defined over R^2) without using loops?

import numpy    
M=numpy.zeros((10,10))
for i in range(0,10):
    for j in range(0,10):
        M[i,j]=f(i,j)
return M

Also see `numpy.vectorize` and `numpy.ufunc`. I've rarely used them and it's been ages, but I recall that they really came in handy (just remember to use built-in and numpy-provided functions as much as possible before customizing with your own functions). — mdscruggs, May 08 '13 at 04:00
thanks: that seems to do the job. I did not see it when looking for related questions. Should I delete my original question? — Sergio Parreiras, May 08 '13 at 04:04

score 1 · Accepted Answer · edited Nov 16 '14 at 09:21

The goal when coding with numpy is to implement your computation on the whole array, as much as possible. So if your function is, for example, f(x,y) = x**2 +2*y and you want to apply it to all integer pairs x,y in [0,10]x[0,10], do:

x,y = np.mgrid[0:10, 0:10]
fxy = x**2 + 2*y

If you don't find a way to express your function in such a way, then:

Ask how to do it (and state explicitly the function definition)
use numpy.vectorize

Same example using vectorize:

def f(x,y): return x**2 + 2*y

x,y = np.mgrid[0:10, 0:10]
fxy = np.vectorize(f)(x.ravel(),y.ravel()).reshape(x.shape)

Note that in practice I only use vectorize similarly to python map when the content of the arrays are not numbers. A typical example is to compute the length of all list in an array of lists:

# construct a sample list of lists 
list_of_lists = np.array([range(i) for i in range(1000)])
print np.vectorize(len)(list_of_lists)
# [0,1 ... 998,999]

score 0 · Answer 2 · answered May 08 '13 at 03:44

Yes, many numpy functions operate on N-dimensional arrays. Take this example:

>>> M = numpy.zeros((3,3))
>>> M[0][0] = 1
>>> M[2][2] = 1
>>> M
array([[ 1.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  1.]])
>>> M > 0.5
array([[ True, False, False],
       [False, False, False],
       [False, False,  True]], dtype=bool)
>>> numpy.sum(M)
2.0

Note the difference between numpy.sum, which operates on N-dimensional arrays, and sum, which only goes 1 level deep:

>>> sum(M)
array([ 1.,  0.,  1.])

So if you build your function f() out of operations that work on n-dimensional arrays, then f() itself will work on n-dimensional arrays.

I guess my question was ill formulated, sorry. To use your example: I want to find a fast way to construct M rather than evaluate the elements of M. Matt's comment -- i.e. to use itertools.product seems a promising direction. Should I delete my original question? — Sergio Parreiras, May 08 '13 at 04:02

mdscruggs · Answer 3 · 2013-05-08T03:55:05.813

You can also use numpy multi-dimension slicing, like below. You just provide slices for each dimension:

arr = np.zeros((5,5)) # 5 rows, 5 columns

# update only first column
arr[:,0] = 1

# update only last row ... same as arr[-1] = 1
arr[-1,:] = 1

# update center
arr[1:-1, 1:-1] = 1

print arr

output:

array([[ 1.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  0.],
       [ 1.,  1.,  1.,  1.,  0.],
       [ 1.,  1.,  1.,  1.,  0.],
       [ 1.,  1.,  1.,  1.,  1.]])

score 0 · Answer 4 · answered May 08 '13 at 03:55

A pure python answer, not depending upon numpy tools, is to make the Cartesian Product of two sequences:

from itertools import product
for i, j in product(range(0, 10), range(0, 10)):
    M[i,j]=f(i,j)

Edit: Actually, I should have read the question properly. This still uses loops, just one less loop.

Python: An elegant/efficient way to evaluate function over bi-dimensional indexes?

4 Answers4