CODEWITHSUNDEEP
pythonnumpymultidimensional-array
Peter Franek
Peter Franek

Reputation: 697

Defining an ndarray in python via a formula

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?

Upvotes: 3

Views: 94

Answers (2)

fivetentaylor
fivetentaylor

Reputation: 1287

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

Upvotes: 1

Divakar
Divakar

Reputation: 221574

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

Upvotes: 2

Related Questions