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?
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
2.7.9. This feature seems like pretty old one, as used here in 2013 question. Could you try reinstalling? I think this product with repeat would be quite useful for you anyway!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
nditertutorial page. At the end it has an example of usingnditerin Cython code. That is fast. But in Python code it isn't faster than other iteration methods. It's better to avoid iteration entirely.