I am trying to index an np.array with another array so that I can have zeros everywhere after a certain index but it gives me the error
TypeError: only integer scalar arrays can be converted to a scalar index
Basically what I would like my code to do is that if I have:
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
d = np.array([2, 1, 3])
that I could do something like
a[d:] = 0
to give the output
a = [[ 1 2 3]
[ 4 0 6]
[ 0 0 9]
[ 0 0 0]]
It can be done with array indexing but it doesn't feel natural.
import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
d = np.array([2, 1, 3])
col_ix = [ 0, 0, 1, 1, 1, 2 ] # column ix for each item to change
row_ix = [ 2, 3, 1, 2, 3, 3 ] # row index for each item to change
a[ row_ix, col_ix ] = 0
a
# array([[1, 2, 3],
# [4, 0, 6],
# [0, 0, 9],
# [0, 0, 0]])
With a for loop
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
for ix_col, ix_row in enumerate( d ): # iterate across the columns
a[ ix_row:, ix_col ] = 0
a
# array([[1, 2, 3],
# [4, 0, 6],
# [0, 0, 9],
# [0, 0, 0]])
A widely used approach for this kind of problem is to construct a boolean mask, comparing the index array with the appropriate arange:
In [619]: mask = np.arange(4)[:,None]>=d
In [620]: mask
Out[620]:
array([[False, False, False],
[False, True, False],
[ True, True, False],
[ True, True, True]])
In [621]: a[mask]
Out[621]: array([ 5, 7, 8, 10, 11, 12])
In [622]: a[mask] = 0
In [623]: a
Out[623]:
array([[1, 2, 3],
[4, 0, 6],
[0, 0, 9],
[0, 0, 0]])
That's not necessarily faster than a row (or in this case column) iteration. Since slicing is basic indexing, it may be faster, even if done several times.
In [624]: for i,v in enumerate(d):
...: print(a[v:,i])
...:
[0 0]
[0 0 0]
[0]
Generally if a result involves multiple arrays or lists with different lengths, there isn't a "neat" multidimensional solution. Either iterate over those lists, or step back and "think outside the box".
d:expression,dhas to be a scalar, a single value, not an array!