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.