In-place add to an array using duplicate indices

New to Julia
1

My code modifies an array in place using indices that contain duplicates. An MWE is as follows:

out = zeros(5)
mock_addr = [1, 1, 1, 1]
out[mock_addr] += [1, 2, 3, 4]
out

The result is out = [4, 0, 0, 0, 0], which is the same as doing the same operation in Python using NumPy. NumPy provides np.add.at(out, mock_addr, [1, 2, 3, 4]), which yields my desired result of out = [0, 10, 0, 0, 0] (ignore the index shift).

Is there a similar function in Julia for the same functionality? I know one can write a for-loop, but I would like to keep the code concise.

2 
for (i,v) in zip(mock_addr, [1,2,3,4])
    out[i] += v
end
3

Thank you. I probably didn’t describe my question clearly. Is there any builtin function like np.add.at that does not require a for loop?

I use this operation on long expressions for multiple times, each on vectors of different sizes. I would rather not to use the for loop because a) eachindex need to be specified correctly, and b) the code would become less readable.

4

@views out[mock_addr] .+= [1, 2, 3, 4] gives you the desired result, but I’m not entirely sure whether this relies on it not noticing that the view(out, mock_addr) on the right aliases the one on the left. I think it fails on GPU arrays, which broadcast in parallel.

One function for this is:

NNlib.scatter!(+, zeros(5), [1,2,3,4.0], [1,1,1,1])
3 Likes
5

numpy functions are just loop in C.

well just make your own function that runs faster than numpy (maybe, untested):

julia> function add_at(out, addrs, v) # this is for scalar v
           for i in addrs
               out[i] += v
           end
           return out
       end
add_at (generic function with 2 methods)

julia> function add_at(out, addrs, vals::AbstractArray)
           length(addrs) == length(vals) || error("unmatching length")
           for (i,v) in zip(addrs, vals)
               out[i] += v
           end
           return out
       end
2 Likes
6

This is super interesting to know. Thank you.

7

Thanks again! I will go with this approach.

8

does that one use simd scatter gather ?

9

A note for others having a similar need:

To use the line above, both @views and the broadcast operation . are necessary. It won’t yield my desired result if either is missing.

1 Like