Using Julia's dot notation and in place operation

Viewed 300

How do I use both Julia's dot notation to do elementwise operations AND ensure that the result is saved in an already existing array?

function myfun(x, y)
    return x + y
end

a = myfun(1, 2)  # Results in a == 3

a = myfun.([1 2], [3; 4])  # Results in a == [4 5; 5 6]

function myfun!(x, y, out)
    out .= x + y
end

a = zeros(2, 2)
myfun!.([1 2], [3; 4], a)  # Results in a DimensionMismatch error

Also, does @. a = myfun([1 2], [3; 4]) write the result to a in the same way as I am trying to achieve with myfun!()? That is, does that line write the result directly to a without saving storing the result anywhere else first?

2 Answers

Your code should be:

julia> function myfun!(x, y, out)
           out .= x .+ y
       end
myfun! (generic function with 1 method)

julia> myfun!([1 2], [3; 4], a)
2×2 Matrix{Float64}:
 4.0  5.0
 5.0  6.0

julia> a
2×2 Matrix{Float64}:
 4.0  5.0
 5.0  6.0

As for @. a = myfun([1 2], [3; 4]) - the answer is yes, it does not create temporary arrays and operates in-place.

This isn't commonly required, and there are usually better ways to achieve this, but it's possible to broadcast on an output argument by using Reference values that point inside the output array.

julia> a = zeros(2, 2)
2×2 Matrix{Float64}:
 0.0  0.0
 0.0  0.0

julia> function myfun!(out, x, y)
          out[] = x + y
       end
myfun! (generic function with 1 method)

julia> myfun!.((Ref(a, i) for i in LinearIndices(a)), [1 2], [3; 4])
2×2 Matrix{Int64}:
 4  5
 5  6

julia> a
2×2 Matrix{Float64}:
 4.0  5.0
 5.0  6.0

Edit: Changed out to be the first parameter as per the Style guide - thanks to @phipsgabler for the reminder.

Related