Search code examples
pythonarraysnumpyindexingarray-broadcasting

How to specify columns when using repeated indices with numpy [for use with np.add.at()]


I'm trying to apply an addition operator to an array where I want repeated indices to indicate repeated addition operations. From a Python Data Science Book (https://jakevdp.github.io/PythonDataScienceHandbook/02.07-fancy-indexing.html), it seems that this is possible using np.add.at(original matrix, indices, thing to add), but I can't figure out how to specify the indices to operate on columns, not rows.

e.g. Dummy Example

# Create Array
A = np.arange(12)
A = A.reshape(4,3)
print(A)

gives

[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

and

# Create columns to add to A (in reality, all values won't be the same)
B = np.ones_like(A[:, [0,0]])
print(adder)

gives

[[1 1]
 [1 1]
 [1 1]
 [1 1]]

I want to perform the operation A[:, [0, 0]] += B but using the system where repeated indices indicate repeated operations (so in this case, both columns of B gets added to column 0). The result should thus be:

[[ 2  1  2]
 [ 5  4  5]
 [ 7  7  8]
 [ 11 10 11]]

This can be done using np.add.at(A, I, B) I believe, but how do I specify the indices I to correspond to [:, [0,0]] as this gives a syntax error (it seems that the indices matrix can't contain the : character?).

Thanks


Solution

  • In [12]: A = np.arange(12).reshape(4,3)
    In [13]: np.add.at(A, (slice(None), [0,0]), 1)
    In [14]: A
    Out[14]: 
    array([[ 2,  1,  2],
           [ 5,  4,  5],
           [ 8,  7,  8],
           [11, 10, 11]])
    

    This could also be written with s_ as

    np.add.at(A, np.s_[:, [0,0]], 1)
    

    s_ is a class object that lets us using indexing notation to create the necessary tuple. In an indexing context Python interpreter converts the : into a slice object.

    In [19]: np.s_[:, [0,0]]
    Out[19]: (slice(None, None, None), [0, 0])