Broadcasting slices of fields doesn't work as expected about oceananigans.jl HOT 8 CLOSED

Actually, I have to rectify, both commands do the same thing because c1[1, :, 1] on the LHS calls view(c1, 1, :, 1)

from oceananigans.jl.

It has to do with view(field, i, j, k) that returns a WindowedField with 3 dimensions while getindex(field, i, j, k) returns field.data[i, j, k]
broadcast follows this pattern:

julia> Meta.@lower c1[1, :, 1] .= c2[2, :, 3]
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─ %1 = Base.dotview(c1, 1, :, 1)
│   %2 = Base.getindex(c2, 2, :, 3)
│   %3 = Base.broadcasted(Base.identity, %2)
│   %4 = Base.materialize!(%1, %3)
└──      return %4
))))

and

julia> axes(Base.dotview(c1, 1, :, 1))
(1:1, Base.OneTo(3), 1:1)

julia> axes(Base.getindex(c2, 2, :, 3))
(OffsetArrays.IdOffsetRange(values=-2:6, indices=-2:6),)

This will work:

c1[1, :, 1] .= view(c2, 1, :, 1)

In general c1[1, :, 1] .= c2[1, :, 1] is not really GPU-compatible. Anyways, if we want to support this behavior we need to change getindex(field, i, j, k) to return a view instead of the underlying data in case of a Colon index.
We might need to think this a bit though, because it would be a big change and might break other implementations

from oceananigans.jl.

This in src/Fields/field.jl should allow that type of broadcasting

@propagate_inbounds Base.getindex(f::Field, ::Colon, j, k) = view(f, :, j, k)
@propagate_inbounds Base.getindex(f::Field, i, ::Colon, k) = view(f, i, :, k)
@propagate_inbounds Base.getindex(f::Field, i, j, ::Colon) = view(f, i, j, :)
@propagate_inbounds Base.getindex(f::Field, ::Colon, ::Colon, k) = view(f, :, :, k)
@propagate_inbounds Base.getindex(f::Field, ::Colon, j, ::Colon) = view(f, :, j, :)
@propagate_inbounds Base.getindex(f::Field, i, ::Colon, ::Colon) = view(f, i, :, :)
@propagate_inbounds Base.getindex(f::Field, ::Colon, ::Colon, ::Colon) = view(f, :, :, :)

from oceananigans.jl.

I see!

is this

c1[1, :, 1] .= view(c2, 1, :, 1)

or this

view(c1, 1, :, 1) .= view(c2, 1, :, 1)

preferred?

The second one is preferred because it's GPU-compatible.
I am not sure if we should add that behavior. That would essentially change the underlying functionality of getindex to return different types when using different indices, so it would be a big change. I am not sure allowing sliced broadcasts is enough to justify this change.

from oceananigans.jl.

I see!
is this

c1[1, :, 1] .= view(c2, 1, :, 1)

or this

view(c1, 1, :, 1) .= view(c2, 1, :, 1)

preferred?

The second one is preferred because it's GPU-compatible. I am not sure if we should add that behavior. That would essentially change the underlying functionality of getindex to return different types when using different indices, so it would be a big change. I am not sure allowing sliced broadcasts is enough to justify this change.

I think those are actually identical (or they could be --- for arrays broadcasting will call view under the hood)

from oceananigans.jl.

I see!

is this

c1[1, :, 1] .= view(c2, 1, :, 1)

or this

view(c1, 1, :, 1) .= view(c2, 1, :, 1)

preferred?

from oceananigans.jl.

This in src/Fields/field.jl should allow that type of broadcasting

@propagate_inbounds Base.getindex(f::Field, ::Colon, j, k) = view(f, :, j, k)
@propagate_inbounds Base.getindex(f::Field, i, ::Colon, k) = view(f, i, :, k)
@propagate_inbounds Base.getindex(f::Field, i, j, ::Colon) = view(f, i, j, :)
@propagate_inbounds Base.getindex(f::Field, ::Colon, ::Colon, k) = view(f, :, :, k)
@propagate_inbounds Base.getindex(f::Field, ::Colon, j, ::Colon) = view(f, :, j, :)
@propagate_inbounds Base.getindex(f::Field, i, ::Colon, ::Colon) = view(f, i, :, :)
@propagate_inbounds Base.getindex(f::Field, ::Colon, ::Colon, ::Colon) = view(f, :, :, :)

Would this be a good idea to add?

from oceananigans.jl.

Ok cool. Gotcha. I’m closing this.

from oceananigans.jl.