How can I make an in-place swap function?

[e3c6]

Suppose I have two arrays A, B of different sizes. Then A .+ B broadcasts them to a common size (if possible) and sums them. Suppose A here is “larger” (in the sense that all(size(B) .== size(A) .|| size(B) .== 1) is true).

Given an index i of A (A[i] ), how can I compute j such that B[j] gives the element of B that would be summed to A[i] in the broadcasted expression?

My motivation for this question, is that I am trying to code an in-place version of the following function:

function swap(conditions::Union{AbstractArray{Bool}, Bool}, x, y)
    x_new = ifelse.(conditions, y, x)
    y_new = ifelse.(conditions, x, y)
    return x_new, y_new
end

The following:

function swap!(x::AbstractArray, y::AbstractArray, conditions::AbstractArray{Bool})
    if size(x) == size(y) == size(conditions)
        for i in eachindex(x, y, conditions)
            if conditions[i]
                x[i], y[i] = y[i], x[i]
            else
                x[i], y[i] = x[i], y[i]
            end
        end
        return x, y
    else
        throw(DimensionMismatch("x, y and conditions must have the same size"))
    end
end

works if x, y, conditions are of the same size. I would like to relax that, requiring that x,y be of the same size, but conditions need only be “broadcastable” to a matching size.

[jules]

Hm good question, the closest I could come up with quickly was this, but it still allocates an array of nothings which is not nice

function swap!(x, y, conditions)
   broadcast(CartesianIndices(x), conditions) do i, c
       if c
           x[i], y[i] = y[i], x[i]
       end
   end
   return
end

[e3c6]

Maybe a FillArray could help, though I’m not sure how.

[jules]

I found some broadcasting internals that I could coax into doing this. Maybe there’s a different more obvious way

julia> function swap!(x, y, conditions)
           b = Broadcast.Broadcasted(tuple, (CartesianIndices(x), conditions))
           for (i, c) in b
               if c
                   x[i], y[i] = y[i], x[i]
               end
           end
           return
       end
swap! (generic function with 2 methods)

julia> x = fill(1, 3, 5); y = fill(2, 3, 5); conditions = [true, false, true, false, true]';

julia> swap!(x, y, conditions)

julia> x
3×5 Matrix{Int64}:
 2  1  2  1  2
 2  1  2  1  2
 2  1  2  1  2

julia> y
3×5 Matrix{Int64}:
 1  2  1  2  1
 1  2  1  2  1
 1  2  1  2  1

julia> @time swap!(x, y, conditions)
  0.000012 seconds

[mikmoore]

Here’s what I came up with.The first uses a view into conditions that repeatedly samples the dimensions that need broadcasting. The second figures out which dimensions need to be broadcast, then adjusts the index into conditions accordingly.

function swapif1!(x,y,cond)
	axes(x) == axes(y) || error("x and y must have the same axes")
	all(d -> size(cond,d) == 1 || axes(cond,d) == axes(x,d), ndims(x)) || error("could not broadcast cond to size of x, y")
	axc = ntuple(d -> ifelse(size(cond,d) == 1, StepRangeLen(firstindex(cond,d),0,size(x,d)), :), Val(ndims(x)))
	vcond = view(cond,axc...)
	for i in eachindex(x,y,vcond)
		if vcond[i]
			x[i],y[i] = y[i],x[i]
		end
	end
	return x,y
end

function swapif2!(x,y,cond)
	axes(x) == axes(y) || error("x and y must have the same axes")
	all(d -> size(cond,d) == 1 || axes(cond,d) == axes(x,d), ndims(x)) || error("could not broadcast cond to size of x, y")
	szc1 = ntuple(d -> size(cond,d) == 1, Val(ndims(x))) # find axes where size(cond,ax) == 1
	indc1 = ntuple(d -> firstindex(cond,d), Val(ndims(x))) # valid index for c on axes of size 1
	for i in eachindex(IndexCartesian(), x, y)
		ic = CartesianIndex(ifelse.(szc1, indc1, Tuple(i))...)
		if cond[ic]
			x[i],y[i] = y[i],x[i]
		end
	end
	return x,y
end

The first one ends up allocating a lot. I didn’t look into why, but there’s probably a type instability somewhere. So I’ll recommend the second.

But I think I like the suggestion above that uses tools from Broadcast to handle this. That seems clean.

[e3c6]

Very elegant, thanks!

@jules

Broadcast.Broadcasted

This is not public API, though, I think?

[jules]

Not sure. Not everything considered public/stable needs to be exported I think.

[mikmoore]

It looks like Broadcasted appears in the docs here. While it is not documented formally and explicitly, it is documented as part of the broadcasting interface. It is intended to be used to extend the broadcast machinery, so (in my reading) it is part of the broadcasting API and thus part of the Julia API, even though it is not exported.

However, I would use

Broadcast.broadcasted(tuple, CartesianIndices(x), conditions)

rather than

Broadcast.Broadcasted(tuple, (CartesianIndices(x), conditions))

The former produces the latter, but appears to be the more standard interface. Notice

julia> @code_lowered tuple.(CartesianIndices(x), conditions)
CodeInfo(
1 ─ %1 = Base.broadcasted(Main.tuple, x1, x2)
│   %2 = Base.materialize(%1)
└──      return %2
)

So what is being done here is equivalent to the broadcasted version, except that a normal call with dots also calls materialize to collect the result to a variable. But the non-materialized version can still be iterated and that’s why it works in this for loop. Broadcast.broadcasted can also be useful if you want to assemble a broadcasted expression in multiple parts or over multiple lines without materializing the intermediate results.

[e3c6]

Good point. So I’m writing it like this now:

function swap!(conditions::Union{AbstractArray{Bool}, Bool}, x::AbstractArray, y::AbstractArray)
    if size(x) == size(y)
        b = broadcasted(tuple, CartesianIndices(x), CartesianIndices(y), conditions)
        for (ix, iy, c) in b
            if c
                x[ix], y[iy] = y[ix], x[iy]
            end
        end
        return x, y
    else
        throw(DimensionMismatch("x, y must have the same size"))
    end
end

I am using different indices ix, iy in case x, y have different indexing patterns (though not sure this is the best approach).

[stevengj]

@e3c6

would like to relax that, requiring that x,y be of the same size, but conditions need only be “broadcastable” to a matching size.

You could just use ordinary broacasting via StructArrays.jl:

using StructArrays
swap(condition::Bool, (x, y)) = condition ? (y, x) : (x, y)
function swap!(x::AbstractArray, y::AbstractArray, conditions)
    a = StructArray((x, y))
    a .= swap.(conditions, a)
    return x, y
end