Oceananigans.jl: Implement halo regions.

[ali-ramadhan]

Even for single-core and single-GPU simulations. This might actually speed things up as we won't have to use incmod1 and decmod1 everywhere #57.

[glwagner]

Check out how OceanTurb.jl does boundary conditions:

# Forward Euler timestepping
function iterate!(model::AbstractModel{TS}, Δt) where TS <: ForwardEulerTimestepper

  for j in eachindex(model.solution)
    c, ∂c∂t, rhs, bcs = unpack(model, j)

    # Interior step
    for i in interior(c)
      @inbounds rhs.data[i] = ∂c∂t(model, i)
    end

    # Boundary conditions
    rhs.data[end] = ∂c∂t(model, bcs.top)
    rhs.data[1]   = ∂c∂t(model, bcs.bottom)
  end

  # Update solution
  for j in eachindex(model.solution)
    c, ∂c∂t, rhs, bcs = unpack(model, j)
    @. c.data += Δt*rhs.data
  end

  return nothing
end

The application of boundary conditions can be made even cleaner by using a function `apply boundary conditions'.

Boundary conditions are then implemented by defining a RHS (in OceanTurb as ∂c∂t) for each type of boundary condition. For example, for boundary conditions in diffusion.jl are implemented via

∇κ∇c_top(κ, c, flux)    = (     -flux      - κ∂z(κ, c, length(c)) ) / dzf(c, length(c))
∇κ∇c_bottom(κ, c, flux) = (  κ∂z(κ, c, 2)  +        flux          ) /    dzf(c, 1)

# Top and bottom flux estimates for constant (Dirichlet) boundary conditions
bottom_flux(κ, c, c_bndry, dzf) = -2*bottom(κ)*( bottom(c) - c_bndry ) / bottom(dzf) # -κ*∂c/∂z at the bottom
top_flux(κ, c, c_bndry, dzf)    = -2*  top(κ) *(  c_bndry  -  top(c) ) /   top(dzf)  # -κ*∂c/∂z at the top

# Flux Boundary conditions
∂c∂t(model, bc::FluxBC{Top})    = ∇κ∇c_top(   model.parameters.κ, model.solution.c, bc.flux(model))
∂c∂t(model, bc::FluxBC{Bottom}) = ∇κ∇c_bottom(model.parameters.κ, model.solution.c, bc.flux(model))

# Constant Boundary conditions
function ∂c∂t(model, bc::ValueBC{Bottom}) 
  flux = bottom_flux(model.parameters.κ, model.solution.c, bc.value(model), model.grid.dzf)
  return ∇κ∇c_bottom(model.parameters.κ, model.solution.c, flux)
end

function ∂c∂t(model, bc::ValueBC{Top}) 
  flux = top_flux(model.parameters.κ, model.solution.c, bc.value(model), model.grid.dzf)
  return ∇κ∇c_top(model.parameters.κ, model.solution.c, flux)
end

[vchuravy]

Other packages to be inspired by:

• https://github.com/JuliaArrays/OffsetArrays.jl
• https://github.com/JuliaArrays/TiledIteration.jl

[ali-ramadhan]

Thanks for the suggestions!

@glwagner Lots of nice ideas and cool syntax! Really like the way the boundary conditions are implemented, I'll see if I can translate it to the halo regions. Interested in some of the design choices you made, hope you don't mind if I ask you about them in person.

To me resolving this issue feels like implementing ghost cells for boundary conditions. I mean this is exactly what's happening for the periodic boundary conditions but it sounds like these halo regions are strictly for implementing periodic boundary conditions (without incmod1 and decmod1) and distributed parallelization? But might be good to make sure this won't blow up in our faces when we get to implementing bathymetry and continental/land boundaries. @jm-c @christophernhill (See #92 for decision not to use ghost cells.)

@vchuravy OffsetArrays.jl would be perfect for halo regions! Do you think this stuff would work for the GPU? I assume if e.g. isbitstype(OffsetArray{Float64}) == true then it can be used in a GPU kernel right out of the box? It's actually false but I might be checking/using it incorrectly. We might need to create our own OffsetCuArray maybe?

# 1D grid with 10 cells (and 1 ghost cell on each end of the domain).
julia> y = OffsetArray{Float64}(undef, 0:11)
OffsetArray(::Array{Float64,1}, 0:11) with eltype Float64 with indices 0:11:
 0.0
 0.0
 5.0e-324
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

julia> typeof(y)
OffsetArray{Float64,1,Array{Float64,1}}

julia> isbitstype(typeof(y))
false

I think it's important to implement #67 ASAP before we start changing data structures and abstractions so we know exactly when we've introduced a performance bottleneck.

[vchuravy]

then it can be used in a GPU kernel right out of the box?

Yes :) but you are missing the fact that Array{Float64, 1} is not isbitstype. I think to make it work out of the box you need

Adapt.adapt_structure(to, x::OffsetArray) = OffsetArray(adapt(to, parent(x)), x.offsets)
Base.Broadcast.BroadcastStyle(::Type{<:OffsetArray{<:Any, <:Any, AA}}) where AA = Base.Broadcast.BroadcastStyle(AA)

https://github.com/JuliaLabs/ShallowWaterBench/blob/aa8b4add55172f0c1920d6f8cebc8eccc424f2b5/src/GPUMeshing/src/GPUMeshing.jl#L79-L80

I should probably finish https://github.com/JuliaArrays/OffsetArrays.jl/pull/57

[ali-ramadhan]

Ah right, when broadcasting you wouldn't want the operations to act on the _ghost cells_ or _halo regions_ (well that's probably fine as long as they're overwritten before being used). So just use an @views maybe. Either way, don't think we use any broadcasts over CuArrayss but should probably play around in a GPU sandbox a bit before implementing anything.

[vchuravy]

Yeah, you still need the first line.

[glwagner]

@ali-ramadhan why not broadcast over CuArrays? Is there a Performance hit?

[ali-ramadhan]

Working on implementing halo regions with OffsetArrays.jl right now. Got a minimal working example running with GPUifyLoops so it should be a pretty straightforward upgrade from here. The performance benchmarks will be interesting. Will open a PR soon.

Some notes for myself:

• The binary and NetCDF output writers will have to be modified to avoid writing out the halo regions.
• We'll just need a kernel that fills in the halo regions at the end of each time step.
• The "halo size" will be part of the Grid and we'll probably have to make sure that the model runs with a large enough halo, e.g. halo of size 1 is enough for second-order advection schemes, but you need a halo of size 2 for fourth-order advection schemes.
• This should also make https://github.com/climate-machine/Oceananigans.jl/issues/57 redundant.