CliMA / Oceananigans.jl

🌊 Julia software for fast, friendly, flexible, ocean-flavored fluid dynamics on CPUs and GPUs
https://clima.github.io/OceananigansDocumentation/stable
MIT License
972 stars 193 forks source link

Broadcasting over regions for `MultiRegionField`s #3577

Open navidcy opened 5 months ago

navidcy commented 5 months ago

In for MultiRegionFields, we often forced to do, e.g.,

for region in 1:number_of_regions(grid)
    u[region] .= - ∂y(ψ[region])
    v[region] .= + ∂x(ψ[region])
end

We'd like to be able to simply do:

u .= - ∂y(ψ)
v .= + ∂x(ψ)

See, e.g., https://github.com/CliMA/Oceananigans.jl/blob/6730e6f6b2c8f1695e20b95ef467b5b14fdc4c5f/validation/multi_region/cubed_sphere_tracer_advection.jl#L59-L69

cc @siddharthabishnu, @simone-silvestri, @glwagner

simone-silvestri commented 5 months ago

I think wrapping Base.Broadcast.materialize! in https://github.com/CliMA/Oceananigans.jl/blob/6730e6f6b2c8f1695e20b95ef467b5b14fdc4c5f/src/Fields/broadcasting_abstract_fields.jl#L24 around @apply_regionally for multi region fields might do the job

glwagner commented 5 months ago

I think wrapping Base.Broadcast.materialize! in

https://github.com/CliMA/Oceananigans.jl/blob/6730e6f6b2c8f1695e20b95ef467b5b14fdc4c5f/src/Fields/broadcasting_abstract_fields.jl#L24

around @apply_regionally for multi region fields might do the job

Or we can extend materialize!(dest::MultiRegionField, ...)