Search code examples
pythonarraysnumpyreduction

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


Solution

  • 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])