Is there any solution to substitute loops in Numpy?

Question

I now have a simple loop

for i in range(4):
    Q[i, x[i] : x[i] + 2] = 1 / 2

where Q = np.zeros([4,4]) and x = [0,1,0,1].

The goal of this question is to obtain the result of this loop without using loop but using vectors.

One possible solution comes to my mind is to build slices by function np.s_. For example, define

S = np.s_[0:2, 1:3, 0:2, 1:3]

according to x. So the loop can be converted into

Q[np.arange(4),S] = 1 / 2

However, I have no idea how to automate S given the random values in x.

I need to get rid of the loop.

Asked By: Zuba Tupaki

||

Source

Answer 1

For reference, the Q you produce is:

In [190]: Q
Out[190]: 
array([[0.5, 0.5, 0. , 0. ],
       [0. , 0.5, 0.5, 0. ],
       [0.5, 0.5, 0. , 0. ],
       [0. , 0.5, 0.5, 0. ]])

The S doesn’t work:

In [191]: S = np.s_[0:2, 1:3, 0:2, 1:3]    
In [192]: S
Out[192]: (slice(0, 2, None), slice(1, 3, None), slice(0, 2, None), slice(1, 3, None))    
In [193]: Q[np.arange(4),S] = 1 / 2
IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

That is a list of slices, which does not qualify as an index

We can construct an advanced-indexing array from these slices:

In [194]: S = np.r_[0:2, 1:3, 0:2, 1:3]
In [195]: S
Out[195]: array([0, 1, 1, 2, 0, 1, 1, 2])
In [196]: Q[np.arange(4).repeat(2),S]
Out[196]: array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5])

Or making a (4,1) and (4,2) pair of indexing arrays that broadcast with each other:

In [199]: Q[np.arange(4)[:,None], S.reshape(4,2)]
Out[199]: 
array([[0.5, 0.5],
       [0.5, 0.5],
       [0.5, 0.5],
       [0.5, 0.5]])

Another way to construct that array of column indices is:

In [201]: np.arange(2)+np.array(x)[:,None]
Out[201]: 
array([[0, 1],
       [1, 2],
       [0, 1],
       [1, 2]])

In [202]: Q[np.arange(4)[:,None],np.arange(2)+np.array(x)[:,None]]
Out[202]: 
array([[0.5, 0.5],
       [0.5, 0.5],
       [0.5, 0.5],
       [0.5, 0.5]])

Of course this only works when slices for each row have the same length. If they differ in length, you have to use some sort of iteration.

Answered By: hpaulj

Is there any solution to substitute loops in Numpy?

Question:

Answers: