I came across this snippet of code in Jake Vanderplas's Data Science Handbook. The concept of using Broadcasting along with Fancy Indexing here wasn't clear to me. Please explain.
In[5]: X = np.arange(12).reshape((3, 4))
X
Out[5]: array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
In[6]: row = np.array([0, 1, 2])
col = np.array([2, 1, 3])
In[7]: X[row[:, np.newaxis], col]
Out[7]: array([[ 2, 1, 3],
[ 6, 5, 7],
[10, 9, 11]])
It says: "Here, each row value is matched with each column vector, exactly as we saw in broadcasting of arithmetic operations. For example:"
In[8]: row[:, np.newaxis] * col
Out[8]: array([[0, 0, 0],
[2, 1, 3],
[4, 2, 6]])
If you use an integer array to index another array you basically loop over the given indices and pick the respective elements (may still be an array) along the axis you are indexing and stack them together.
arr55 = np.arange(25).reshape((5, 5))
# array([[ 0, 1, 2, 3, 4],
# [ 5, 6, 7, 8, 9],
# [10, 11, 12, 13, 14],
# [15, 16, 17, 18, 19],
# [20, 21, 22, 23, 24]])
arr53 = arr55[:, [3, 3, 4]]
# pick the elements at (arr[:, 3], arr[:, 3], arr[:, 4])
# array([[ 3, 3, 4],
# [ 8, 8, 9],
# [13, 13, 14],
# [18, 18, 19],
# [23, 23, 24]])
So if you index an (m, n)
array with an row (or col) index of length k
(or length l
) the resulting shape is:
A_nm[row, :] -> A_km
A_nm[:, col] -> A_nl
If however you use two arrays row
and col
to index an array
you loop over both indices simultaneously and stack the elements (may still be arrays) at the respective position together.
Here it row
and col
must have the same length.
A_nm[row, col] -> A_k
array([ 3, 13, 24])
arr3 = arr55[[0, 2, 4], [3, 3, 4]]
# pick the element at (arr[0, 3], arr[2, 3], arr[4, 4])
Now finally for your question: it is possible to use broadcasting while indexing arrays. Sometimes it is not wanted that only the elements
(arr[0, 3], arr[2, 3], arr[4, 4])
are picked, but rather the expanded version:
(arr[0, [3, 3, 4]], arr[2, [3, 3, 4]], arr[4, [3, 3, 4]])
# each row value is matched with each column vector
This matching/broadcasting is exactly as in other arithmetic operations. But the example here might be bad in the sense, that not the result of the shown multiplication is of importance for the indexing. The focus here is on the combinations and the resulting shape:
row * col
# performs a element wise multiplication resulting in 3
numbers
row[:, np.newaxis] * col
# performs a multiplication where each row value is *matched* with each column vector
The example wanted to emphasis this matching of row
and col
.
We can have a look and play around with the different possibilities:
n = 3
m = 4
X = np.arange(n*m).reshape((n, m))
row = np.array([0, 1, 2]) # k = 3
col = np.array([2, 1, 3]) # l = 3
X[row, :] # A_nm[row, :] -> A_km
# array([[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
X[:, col] # A_nm[:, col] -> A_nl
# array([[ 2, 1, 3],
# [ 6, 5, 7],
# [10, 9, 11]])
X[row, col] # A_nm[row, col] -> A_l == A_k
# array([ 2, 5, 11]
X[row, :][:, col] # A_nm[row, :][:, col] -> A_km[:, col] -> A_kl
# == X[:, col][row, :]
# == X[row[:, np.newaxis], col] # A_nm[row[:, np.newaxis], col] -> A_kl
# array([[ 2, 1, 3],
# [ 6, 5, 7],
# [10, 9, 11]])
X[row, col[:, np.newaxis]]
# == X[row[:, np.newaxis], col].T
# array([[ 2, 6, 10],
# [ 1, 5, 9],
# [ 3, 7, 11]])