numpy/numpy — pull request #2821 by jayvius (2012-12-14 18:43:18+00:00)
Title: in place fancy indexing for ufuncs

Currently when doing something like a[idx] += 1 where idx is a fancy index, and the fancy index lists an element more than once, the increment is performed on each element's original value instead of accumulating the increments. This is because the first operand is buffered, and the ufunc operates on the buffered operand.

This pull request adds a new method called 'at' to ufunc object for in place fancy indexing with no buffering. This allows a ufunc operation to be performed on an element more than once, with the result of each operation being stored in place before the next operation is performed. This is equivalent to doing a[idx] += b with no buffering for operand 'a'. For example:

a = np.array([0,1,2,3])
# second element is incremented twice
np.add.at(a, [0,1,1], 1)
a
>>> array([1, 3, 2, 3)]

The 'at' method also supports slicing. For example:

a = np.array([0,1,2,3])
np.add.at(a, (slice(None),), 1)
a
>>> array([1, 2, 3, 4)]

A more complicated example:

a = np.arange(27).reshape(3,3,3)
b = np.array([100,200,300])
np.add.at(a, (slice(None), slice(None), [1,2,1]), b)
a
>>> array([[[0,401,202],
[3,404,205],
[6,407,208]],

[[9, 410,211],
[12,413,214],
[15,416,217]],

[[18,419,220],
[21,422,223],
[24,425,226]]])

The 'at' method can be used with all ufuncs, including unary ufuncs:

a = np.array([0, 1, 2, 3])
np.negative.at(a, [0,1])
a
>>> array([0, -1, 2, 3])