Defining an ndarray in python via a formula

Question

I have a multidimensional array, initiated as C=np.zeros([20,20,20,20]). Then I'm trying to assign some values to C via some formula (C(x)=(exp(-|x|^2) in this case). The following code works but is extremely slow.

it=np.nditer(C, flags=['multi_index'], op_flags=['readwrite'])
while not it.finished:
    diff=np.linalg.norm(np.array(it.multi_index))
    it[0]=np.exp(-diff**2)
    it.iternext()

Can this by done in a faster and possibly more pythonish way?

Judging from your example you have found the `nditer` tutorial page. At the end it has an example of using `nditer` in Cython code. That is fast. But in Python code it isn't faster than other iteration methods. It's better to avoid iteration entirely. — hpaulj, Oct 27 '15 at 21:23

Divakar · Accepted Answer · 2015-10-27T19:56:45.713

2

Here's one way to do it.

Step #1 Get all combinations corresponding to all indices being calculated with np.array(it.multi_index) in the code. On this, one can leverage product from itertools.

Step #2 Perform the L2 norm calculations across all combinations in a vectorized manner.

Step #3 Finally do C(x)=(exp(-|x|^2) in an elementwise manner.

# Get combinations using itertools.product
combs = np.array(list(product(range(N), repeat=4)))

# Perform L2 norm and elementwise exponential calculations to get final o/p 
out = np.exp(-np.sqrt((combs**2).sum(1))**2).reshape(N,N,N,N)

Runtime tests and verify output -

In [42]: def vectorized_app(N):
    ...:     combs = np.array(list(product(range(N), repeat=4)))
    ...:     return np.exp(-np.sqrt((combs**2).sum(1))**2).reshape(N,N,N,N)
    ...: 
    ...: def original_app(N):
    ...:     C=np.zeros([N,N,N,N])
    ...:     it=np.nditer(C, flags=['multi_index'], op_flags=['readwrite'])
    ...:     while not it.finished:
    ...:         diff_n=np.linalg.norm(np.array(it.multi_index))
    ...:         it[0]=np.exp(-diff_n**2)
    ...:         it.iternext()
    ...:     return C
    ...: 

In [43]: N = 10

In [44]: %timeit original_app(N)
1 loops, best of 3: 288 ms per loop

In [45]: %timeit vectorized_app(N)
100 loops, best of 3: 8.63 ms per loop

In [46]: np.allclose(vectorized_app(N),original_app(N))
Out[46]: True

edited Oct 27 '15 at 19:56

answered Oct 27 '15 at 19:25

Divakar

218,885
19
262
358

Thanks a lot. I just get the "product() got an unexpected keyword argument 'repeat'" but maybe it's an old versino of python (?) – Peter Franek Oct 27 '15 at 19:33
@PeterFranek What's your Python version? – Divakar Oct 27 '15 at 19:43
2.7.6.. (under Ubuntu)Should I reinstall it? – Peter Franek Oct 27 '15 at 19:49
@PeterFranek Ok, I am on `2.7.9`. This feature seems like pretty old one, as used here in [`2013 question`](http://stackoverflow.com/q/14615595/3293881). Could you try reinstalling? I think this `product with repeat` would be quite useful for you anyway! – Divakar Oct 27 '15 at 19:51
Just a last question, is your code really at least twice faster? If yes, I will reinstall :) – Peter Franek Oct 27 '15 at 19:52
@PeterFranek It's **`33x`** faster ;) – Divakar Oct 27 '15 at 19:57

fivetentaylor · Answer 2 · 2015-10-27T20:18:05.443

So it looks like you just wan't to apply your operation to the indices of each element? How about this:

x = np.exp(-np.linalg.norm(np.indices([20,20,20,20]), axis=0)**2)

np.indices is a really slick function. Also related are mgrid and meshgrid for more complex operations. In this case, since you have 4 dimensions, it returns an array with shape (4,20,20,20,20).

And pure numpy is a little faster :)

In [13]: timeit posted_code()
1 loops, best of 3: 843 ms per loop

In [14]: timeit np.exp(-np.linalg.norm(np.indices([20,20,20,20]), axis=0)**2)
100 loops, best of 3: 3.76 ms per loop

And it is exactly the same result:

In [26]: np.all(C == x)
Out[26]: True

Defining an ndarray in python via a formula

2 Answers2