In-place broadcasting for a specific function

My code is as follows:

function get_foo(x, y, args...)
    # specific algorithm that modifies x and y...
    return x + y
end

function gen_foo!(A::AbstractMatrix, args...)
    # is equivalent to:
    for coord in CartesianIndices(A)
        A[coord] = get_foo(coord.I..., args...)
    end
    # but faster due to optimizations based on how get_foo works.
end

This works fine with the following usage:

using OffsetArrays
A = OffsetArray{Int}(undef, -5:5, -5:5)
gen_foo!(A)

I would like to use syntax like this:

A .= get_foo.(A, args...)

which would be equivalent to gen_foo!(A, args...).

  1. Is it a good idea ?
  2. If yes, how can i implement it ? According to the julia documentation, the only function that takes the function as an argument is Base.Broadcast.broadcasted(f, args...). However, that seems to be intended for overriding broadcasting for specific objects, which doesn’t seem to match my case

Yes, this is completely normal and idiomatic.

Have you tried it? It should work already, with no effort on your part.

Ok, this seems to works:

function Base.broadcasted(
    style::Base.Broadcast.BroadcastStyle,
    f::typeof(get_foo),
    A::AbstractMatrix,
    args...,
)
    gen_foo!(A, args...)
    return A
end

But now for something like this:

B = get_foo.(A)

A === B is true, i.e. A has been changed in-place. The example from the doc:

broadcasted(::DefaultArrayStyle{1}, ::typeof(-), r::OrdinalRange) = range(-first(r), step=-step(r), length=length(r))

works because it does not modify anything in-place.

Again from the doc and reading the source code of broadcast.jl, Base.broadcasted should return a Broadcasted object to have a default behavior. But nothing is said about how to create one…

You can overload things, to make A .= get_foo.(A, args...) not use Base’s broadcasting, but IMO this is probably a bad idea. Your reader will find it much easier to understand if you just call gen_foo! when this is what you want.

julia> Meta.@lower A .= f.(A, b)
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─ %1 = A
│   %2 = f
│   %3 = A
│   %4 = b
│   %5 = Base.broadcasted(%2, %3, %4)
│   %6 = Base.materialize!(%1, %5)
└──      return %6
))))

julia> A = rand(2, 3);

julia> bc = Broadcast.broadcasted(get_foo, A, 4) |> Broadcast.instantiate;

julia> typeof(bc)
Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{2}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, typeof(get_foo), Tuple{Matrix{Float64}, Int64}}

julia> bc.
args
axes
f
style
julia> bc.args
([0.23129131980108864 0.23611011710743723 0.37511906216766255; 0.8143211624304202 0.6725483870746364 0.18429496255480815], 4)

julia> function Base.materialize!(A::AbstractMatrix, bc::Broadcast.Broadcasted{<:Any, <:Any, typeof(get_foo)})
         bc.args[1] === A || error("not handled... could call invoke?")
         println("calling Base.materialize! overload")
         gen_foo!(bc.args...)
         A
       end;

julia> A .= get_foo.(A, 4)
calling Base.materialize! overload
2Ă—3 Matrix{Float64}:
 2.0  3.0  4.0
 3.0  4.0  5.0
1 Like

I’m confused. Why do you need to define some broadcast function. Does it not automatically work without you defining anything?

No. Sorry if i expressed myself badly. What i want is overriding the default behavior (a raw loop), and use gen_foo! instead of the default broadcasting because there is some optimizations involved due to how get_foo works (if we know the nearby points, we can know an area of ​​points that have the same value so we can just fill them instead of computing get_biome each time).

Maybe you are looking for something like the GetindexArrays package:

julia> using Pkg

julia> pkg"add https://github.com/JuliaArrays/GetindexArrays.jl/"
    Updating git-repo `https://github.com/JuliaArrays/GetindexArrays.jl`

julia> using GetindexArrays, OffsetArrays

julia> A = OffsetArray{Int}(undef, -5:5, -5:5);

julia> A .= GetindexArray(axes(A)) do i,j  # this is `gen_foo`
       i+j   # this is `get_foo`
       end
11Ă—11 OffsetArray(::Matrix{Int64}, -5:5, -5:5) with ... with indices -5:5Ă—-5:5:
 -10  -9  -8  -7  -6  -5  -4  -3  -2  -1   0
  -9  -8  -7  -6  -5  -4  -3  -2  -1   0   1
  -8  -7  -6  -5  -4  -3  -2  -1   0   1   2
  -7  -6  -5  -4  -3  -2  -1   0   1   2   3
  -6  -5  -4  -3  -2  -1   0   1   2   3   4
  -5  -4  -3  -2  -1   0   1   2   3   4   5
  -4  -3  -2  -1   0   1   2   3   4   5   6
  -3  -2  -1   0   1   2   3   4   5   6   7
  -2  -1   0   1   2   3   4   5   6   7   8
  -1   0   1   2   3   4   5   6   7   8   9
   0   1   2   3   4   5   6   7   8   9  10

The documentation for GetindexArrays is sort of non-existent, and it isn’t an official package, but it is mostly self-explanatory.

I don’t want to find a way to write gen_foo!; I already have a gen_foo! function. Since it provides a “global” perspective, (unlike get_foo, which operates on a single point), I can implement specific optimizations with some maths. This makes it faster than a basic loop over the indices of the arrays. Broadcasting, by default, performs such loops, and I want to override this behavior to use my more efficient gen_foo! function instead.

But this is a cool package thanks for sharing!