Numpy broadcasting of fancy index

Question

How does the np.newaxis work within the index of numpy array in program 1? Why it works like this?

Program 1:

import numpy as np
x_id = np.array([0, 3])[:, np.newaxis]
y_id = np.array([1, 3, 4, 7])
A = np.zeros((6,8))
A[x_id, y_id] += 1 
print(A)

Result 1:

[[ 0.  1.  0.  1.  1.  0.  0.  1.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  1.  0.  1.  1.  0.  0.  1.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.]]

score 1 · Accepted Answer · answered Jul 22 '17 at 00:49

The newaxis turns the x_id array into a column vector, same as np.array([[0],[3]]).

So you are indexing A at the cartesian product of [0,3] and [1,3,4,7]. Or to put it another way, you end up with 1s for rows 0 and 3, columns 1,3,4,and 7.

Also look at np.ix_([0,3], [1,3,4,7])

or

In [832]: np.stack(np.meshgrid([0,3],[1,3,4,7],indexing='ij'),axis=2)
Out[832]: 
array([[[0, 1],
        [0, 3],
        [0, 4],
        [0, 7]],

       [[3, 1],
        [3, 3],
        [3, 4],
        [3, 7]]])

Be a little careful with the +=1 setting; if indices are duplicated you might not get what you expect, or would get with a loop.

Numpy broadcasting of fancy index

1 Answers1