Reputation: 1
I was always thinking to avoid loop operation in my python code. Numpy really helps, but in some slightly complicated cases, I felt stuck on how to utilize numpy array wisely.
Below is an simple example illustrating my inability, a will be a parameter, and b is an numpy array.
def f(a,b):
return np.sum( a * b)
so there is no problem if I wish to evaluate this function by a given single parameter and an array.
x = 2
y = np.arange(3)
print (f(x,y))
But sometimes I want to evaluate different parameter value of the function altogether with a fixed array value.
I would try:
x2 = np.array([1,4,5,2,8])
print (f(x2,y))
What I wish to get is surely an array with value:
[f(1,y),f(4,y),f(5,y),f(2,y),f(8,y)]
However, python will try to evaluate the dot product of x and y, since now they are both np arrays and It will report
ValueError: operands could not be broadcast together with shapes (5,) (3,)
How should I overcome this, in numpy array-wise fashion, producing the sequence
[f(1,y),f(4,y),f(5,y),f(2,y),f(8,y)]
without using loops?
(In this example, I could resolved problem by modify f by:
def f(a,b):
return a * np.sum(b)
But in most general cases, we cannot factor the parameter out.)
Upvotes: 0
Views: 95
Reputation: 931
np.newaxis
is a very handy tool for cases like this, and gives you a bit more control over broadcasting. In this case, you'll want to add a new axes to give numpy some hints about where to broadcast:
>>> x = np.array([1,4,5,2,8])
>>> y = np.arange(3)
>>> x[:,np.newaxis] * y
array([[ 0, 1, 2],
[ 0, 4, 8],
[ 0, 5, 10],
[ 0, 2, 4],
[ 0, 8, 16]])
If you'd like the sums along the second axis, you can sum like this:
>>> (x[:, np.newaxis] * y).sum(axis=1)
array([ 3, 12, 15, 6, 24])
Upvotes: 2