CODEWITHSUNDEEP
pythonarraysnumpy
MiamiBeach
MiamiBeach

Reputation: 3485

Python: How to rewrite function to be used with arrays?

I have simple function which looks like that:

def summer_scalar(x):
    a = np.array([1,2,3])
    b = np.array([7,8,9])
    return sum(a - b * x)

It works, but problems appear if I want to use it with arrays, for example Numpy arrays. If I submit numpy array I get error:

summer_scalar(np.array([7,8,9,7]))
ValueError: operands could not be broadcast together with shapes (3,) (4,)

I guess the reason is that it treats all my arrays as peers, but they are not: elements from arguments should be processed one after one and form an output array, arguments from a and b should be processed as they were before. So it tries to sum arrays of different size and fails.

What can I do?

I can rewrite my method:

def summer(x):
    a = np.array([1,2,3])
    b = np.array([7,8,9])
    s = 0
    for ai,bi in zip(a,b):
        s = s + ai - bi* x 
        #print ai,bi,s
    return s

This works.

A question is: is there a better way to modify my method for vector operations?

Upvotes: 0

Views: 75

Answers (1)

Divakar
Divakar

Reputation: 221514

Yes, there are better ways.

Method #1

Direct translation of the stated problem to a vectorized mindset leads us to leverage NumPy broacasting, like so -

(a-b*x[:,None]).sum(1)

That's an outer product, so we can also use NumPy ufunc to do the same -

(a-np.multiply.outer(x,b)).sum(1)

Method #2

Looking closely, it seems we won't need that explicit dimension-extension+broadcasting and we could just sum the elements in a and b and do something like this -

a.sum(0) - x*b.sum(0)

We could employ np.einsum here to replace x*b.sum(0), like so -

a.sum(0) - np.einsum('i,j->i', x,b)

Again, a.sum(0) could be replaced by one more einsum -

np.einsum('i->',a) - np.einsum('i,j->i', x,b)

So, lots of options to vectorize such an operation!

Upvotes: 1

Related Questions