Vectorizing operation in NumPy

Question

Is there a general way to vectorize these kind of operations in NumPy?

In [2]: N = 8

In [3]: ll = np.arange(8)

In [4]: arr = np.zeros(ll.shape + (2, 2))

In [5]: ll.shape
Out[5]: (8,)

In [6]: arr.shape
Out[6]: (8, 2, 2)

In [7]: for ii in range(N):
   ...:     arr[ii, :, :] = np.array(...)  # 2 x 2 array function of ll[ii]

if that function is a linear operation on ll then this would be trivial, but is there a way to do it in the general case? Just to put an example:

In [8]: for ii in range(N):
   ...:     arr[ii, :, :] = np.array([
   ...:         [np.cos(ll[ii]) - 1, 0],
   ...:         [np.sin(ll[ii]), np.cos(ll[ii]) ** 2]
   ...:     ])

score 5 · Answer 1 · answered Jul 11 '13 at 13:42

5

The right way of assembling your arr array would be something like:

arr[:, 0, 0] = np.cos(ll) - 1
arr[:, 0, 1] = 0
arr[:, 1, 0] = np.sin(ll)
arr[:, 1, 1] = np.cos(ll) ** 2

You definitely shouldn't call np.array on a list of arrays that are going to be stored in an already existing array: it's wasteful intermediate array creation, which is a bad practice, and I doubt it adds any clarity to the code. A memory/performance conscious developer would probably do something like:

np.cos(ll, out=arr[:, 0, 0])
arr[:, 1, 1] = arr[:, 0, 0]
arr[:, 0, 0] -= 1
arr[:, 0, 1] = 0
np.sin(ll, out=arr[:, 1, 0])
arr[:, 1, 1] *= arr[:, 1, 1]

But this would fall under the premature optimization category more often than not.

You should also really not use ll as a variable name...

answered Jul 11 '13 at 13:42

Jaime

65,696
17
124
159

I expected some way of vectorizing the whole thing in a one liner, but now that you bring up the performance thing I might do some benchmarks with the solutions you suggest. Thanks! – astrojuanlu Jul 12 '13 at 06:46
You may want to take a look at the comment thread [here](http://stackoverflow.com/questions/17580666/improving-numpy-speed-for-gauss-seidel-jacobi-solver/17581581?noredirect=1#comment25611917_17581581). I am not sure what's going on, but it seems that using the `out` argument of ufuncs does not consistently perform better. Very weird. Oh, and your blog, *muy chulo, enhorabuena* (very nice, congrats). – Jaime Jul 12 '13 at 07:47

BrenBarn · Answer 2 · 2013-07-11T17:51:17.260

You can do it like this:

def func(x):
    return np.array([
        [np.cos(x)-1,np.repeat(0, len(x))],
        [np.sin(x), np.cos(x)**2]
    ])

Then func(x) will return an array of shape (2, 2, 8). You can get it in the orientation you want with func(x).T.

This only works when x is one-dimensional. I think you could work something out for higher dimensions using np.broadcast_arrays, but not sure exactly how at the moment. The basic thing, though, is that if you want to return an array, you can't use vectorized numpy functions like cos in some cells, while putting literal scalars (like 0) in other cells. You need to fill the scalar cells with an array whose shape is derived from the input array.

Thank you for your answer, I saw that the correct way to transpose it would be `arr.transpose(2, 0, 1)`. I am going to wait for other answers in case there is a clearer method, otherwise I will accept this. — astrojuanlu, Jul 11 '13 at 08:35

Vectorizing operation in NumPy

2 Answers2