Draditate
Draditate

Reputation: 1

Very basic manipulation of numpy in python functions

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

Answers (1)

philngo
philngo

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

Related Questions