python - Add multiple values to one numpy array index

Question

Simple Version: if I do this:

import numpy as np
a = np.zeros(2)
a[[1, 1]] += np.array([1, 1])

I get [0, 1] as an output. but I would like [0, 2]. Is that possible somehow, using implicit numpy looping instead of looping over it myself?

What-I-actually-need-to-do version:

I have a structured array that contains an index, a value, and some boolean value. I would like to sum those values at those indices, based on the boolean. Clearly that can be done with a simple loop, but it seems like it should be possible with clever numpy indexing (as above).

For example, I have an array with 5 elements that I want to populate from the array with values, indices, and conditions:

import numpy as np
size = 5
nvalues = 10
np.random.seed(1)
a = np.zeros(nvalues, dtype=[('val', float), ('ix', int), ('cond', bool)])
a = np.rec.array(a)
a.val = np.random.rand(nvalues)
a.cond = (np.random.rand(nvalues) > 0.3)
a.ix = np.random.randint(size, size=nvalues)

# obvious solution
obvssum = np.zeros(size)
for i in a:
    if i.cond:
        obvssum[i.ix] += i.val

# is something this possible?
doesntwork = np.zeros(size)
doesntwork[a[a.cond].ix] += a[a.cond].val

print(doesntwork)
print(obvssum)

Output:

[ 0.          0.          0.61927097  0.02592623  0.29965467]
[ 0.          0.          1.05459336  0.02592623  1.27063303]

I think what's happening here is if a[a.cond].ix were guaranteed to be unique, my method would work just fine, as noted in the simple example.

Answer 1

This is what the at method of NumPy ufuncs is for:

output = numpy.zeros(size)
numpy.add.at(output, a[a.cond].ix, a[a.cond].val)