Store in place results from function returning multiple arrays

What is the best practice to store in place the results of a function returning multiple arrays?

A simple non-MWE:

f(x,y,z) = 1 ./ x, y'.*y, 2*z      # the actual function is long and complex, with matrix operations

n = 10
x = y = z = rand(n)
fx, fy, fz = Vector{Float64}(undef,n), Matrix{Float64}(undef,n, n), Vector{Float64}(undef,n)

fx, fy, fz .= f(x,y,z)   # ERROR: MethodError: no method matching copyto!

Should we define some f!!!() function that mutates in place the first three arguments? Example:

f!!!(fx,fy,fz,x,y,z) = @. begin fx = 1/x; fy = y'*y; fz = 2*z; return fx, fy, fz; end

f!!!(fx,fy,fz,x,y,z)

Thanks.

NB: edited original example because it had several typos.

StructArrays can help:

julia> f(x) = (x+1, x+2)

julia> sa = StructArray((zeros(10), zeros(10)))

julia> @btime map!(f, $sa, 1:10)
  11.008 ns (0 allocations: 0 bytes)
10-element StructArray(::Vector{Float64}, ::Vector{Float64}) with eltype Tuple{Float64, Float64}:
 (2.0, 3.0)
 (3.0, 4.0)
 (4.0, 5.0)
 (5.0, 6.0)
 (6.0, 7.0)
 (7.0, 8.0)
 (8.0, 9.0)
 (9.0, 10.0)
 (10.0, 11.0)
 (11.0, 12.0)

My suggestion: don’t return an array. Let the function work on scalars, and return a Tuple of scalars. And then mutate your arrays in-place in a loop.

1 Like

In my example, I would need an heterogeneous structure (the function returns vectors and a matrix):

sa = StructArray((zeros(n), zeros(n,n), zeros(n))) # ERRORS

but it seems that with StructArrays all field arrays must have same shape?

Of course, all subarrays need to have the same shape, as sa[i] = (sa.field_1[i], sa.field_2[i]).

Maybe I misunderstood your question: thought that f is a scalar, and you want to broadcast it. If f returns already allocated arrays, there’s no way to make it inplace.

1 Like

If you intend to avoid allocation, this seems impossible because you have already allocated new arrays inside the function. At that point you might as well just destructuringly assign those returned arrays to the variables: fx, fy, fz = f(x,y,z). If you want to do in-place mutation, you need to pass the fx, fy, fz instances into the function, like a function f!(xout, yout, zout, x, y, z) ... nothing end, and use in-place versions of whatever you do in the function.

fx, fy, fz .= f(x,y,z) errors because what you’re really doing is creating a Tuple (fx, fy, fz) and trying to broadcast!-reassign the Tuple’s elements. Even if a Tuple were mutable, the variables fx, fy, and fz aren’t being reassigned. An Array is mutable, so we can see what happens:

julia> x,y,z = (0,0,0) # destructuring assignment to variables x,y,z
(0, 0, 0)

julia> x,y,z
(0, 0, 0)

julia> [x,y,z] .= (1,2,3) # elementwise assignment to array [0,0,0]
3-element Array{Int64,1}:
 1
 2
 3

julia> x,y,z
(0, 0, 0)
1 Like

@DNF, thanks for the advice. Unfortunately, the function is quite long and complex involving matrix operations and it seems improbable one can make it scalar.

Thanks all for the explanations, it is a bit more clear now. We need scalar functions to have the dot fusion to work with broadcasting, or we have to pass pre-allocated arrays and mutate them inside the function.

Now what is the recommended syntax for my example: one exclamation mark (!) or three (!!!) in the function name? And the order of the function arguments that are mutated, should they appear all in front, or interleaved as in Benny’s response?

I believe the convention is to use 1 ! character always, and to put the (possibly) mutated arguments first, sorry about that, I’ll amend my post.

1 Like