CODEWITHSUNDEEP
pythonarraysnumpyvectorization
anon01
anon01

Reputation: 11171

vectorized addition of parameterized arrays

I'd like to sum a series of vectors defined by a (complicated) function of parameters. Is there a way to use vectorization in place of the nested for loop to get the same value for ans in the code below? Memory usage is not to be a constraint. Note that for the actual problem the nested for loop is known to limit performance, despite this not being the case here.

import numpy as np

a1_vec = np.array([0.3, 1])
a2_vec = np.array([3.3, 10])
b1_vec = np.array([0.5, 0.7])
b2_vec = np.array([1.5, 1.3])

x = np.arange(0, 10000)
ans = 0
for a1, b1 in zip(a1_vec, b1_vec):
    for a2, b2 in zip(a2_vec, b2_vec):
        ans += x*np.exp(- a1 - b2) + x**(1 / 2)*np.cos(b1) + x**(1 / 3)*np.sin(a2)

Upvotes: 1

Views: 77

Answers (1)

Magellan88
Magellan88

Reputation: 2573

So as @anon pointed out in his comment, you can use numpy's array broadcasting: using X[:,None,None] or the more explicit X[:,np.newaxis,np.newaxis] (which just is an more explicit alias for None). With this np.newaxis, you can introduce a new, empty dimension into your array. Your X.shape (dimensionality of the array) will be no longer be (10000) but (10000,1,1).

Any operation on two arrays needs the same dimensions on both arrays (e.g. (2,5) and (2,5) will work). One of the cool features of numpy is the broadcasting, which happens when an array operation is done between two arrays where one dimension that is larger than 1 is combined with a dimension of 1 (e.g. (2,5) and (2,1) where the second dimension is 5 in the first and 1 in the second vector). In this case numpy will broadcast the array with the dimension of 1 by simply looping over it (and it will do so in a c-compiled fast way). In the example it will take the of the (2,5) array with the (2,1) array it will pretend that the second array is also a (2,5) array just with the same two numbers 5 times repeated.

so as an easy example look at

a = np.arange(3)
print(a*a)
#[0,2,4]

this is just normal element by element multiplication

now if you introduce empty dimensions using broadcasting rules: 

print(a[:,None]*a[None,:])
#[[0 0 0]
# [0 1 2]
# [0 2 4]]

It's a really cool and the very key to understanding the power of numpy, but admittedly it also took me some time to get familiar with it.

Upvotes: 3

Related Questions