Mike Mike - 14 days ago 9
Python Question

Ensure parallel reduction of numpy array operation mapping to repeated position

Is there a way for numpy to ensure that an array operation mapping to repeated positions undergo a reduction, i.e. they are both performed on the result of each other?

a = numpy.zeros([4], int) # [0 0 0 0]
b = numpy.arange(0, 8) # [0 1 2 3 4 5 6 7]

positions = [0, 0, 1, 1, 2, 2, 3, 3]

a[positions] += b

# desired result: [0 + 1, 2 + 3, 4 + 5, 6 + 7]
# actual result: random crossover between [0, 2, 4, 6] and [1, 3, 5, 7]


as you can see both element 1 and 2 of b map to position 1 and so on, I need to make sure that += adds both, whereas by default it looks like it can randomly add 1 or 2 to zero at the same time, then store the result twice, which is in turn the result of only one of the operations

Answer

When there are repeated indices, the behavior of in-place addition in a numpy array is undefined. To ensure the behavior that you want, use numpy.add.at. (All numpy "ufuncs" have the at method.)

For example, here are your arrays:

In [21]: a
Out[21]: array([0, 0, 0, 0])

In [22]: b
Out[22]: array([0, 1, 2, 3, 4, 5, 6, 7])

In [23]: positions
Out[23]: [0, 0, 1, 1, 2, 2, 3, 3]

Use numpy.add.at to accumulate the values:

In [24]: np.add.at(a, positions, b)

In [25]: a
Out[25]: array([ 1,  5,  9, 13])