Reducing a Field on a ImmersedBoundaryGrid with additional masking

[tomchor]

I'm trying to calculate an average of a field on an ImmersedBoundaryGrid. The part that's a bit complicated is that I'm trying to average only the parts where another field satisfy a given condition. (So that's another mask on top of the immersed boundary mask.)

Here's a small example of what I'm trying to do:

using Oceananigans
grid_base = RectilinearGrid(GPU(), size=(4, 4, 4), extent=(1, 1, 1));
ibg = ImmersedBoundaryGrid(grid_base, GridFittedBottom((x, y) -> -1/2))
model = NonhydrostaticModel(grid=ibg, tracers=:b);
u, v, w = model.velocities; b=model.tracers.b
bz = Field(∂z(b));

Kb = Average(Field(-w*b/bz), mask=abs(bz)>1e-5)

But this throws me:

julia> Kb = Average(Field(-w*b/bz), mask=abs(bz)>1e-5)
ERROR: MethodError: no method matching isless(::Float64, ::Oceananigans.AbstractOperations.UnaryOperation{Center, Center, Face, typeof(abs), Field{Center, Center, Face, Oceananigans.AbstractOperations.Derivative{Center, Center, Face, typeof(∂zᶜᶜᶠ), Field{Center, Center, Center, Nothing, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, Matrix{Float64}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, CPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing}}, typeof(Oceananigans.Operators.identity3), typeof(∂z), ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, Matrix{Float64}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, CPU}, Float64}, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, Matrix{Float64}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, CPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Oceananigans.Fields.FieldStatus{Float64}, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing}}, typeof(Oceananigans.Operators.identity3), ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, Matrix{Float64}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, CPU}, Float64})

Closest candidates are:
  isless(::T, ::T) where T<:Union{Float16, Float32, Float64}
   @ Base float.jl:547
  isless(::AbstractFloat, ::AbstractFloat)
   @ Base operators.jl:177
  isless(::AbstractFloat, ::Static.StaticFloat64{Y}) where Y
   @ Static /glade/work/tomasc/.julia/packages/Static/s64Zx/src/float.jl:90
  ...

Which is kinda expected I guess since that operation with a field isn't defined. But I can't broadcast the operations either since I'm on the GPU and this would use scalar indices:

julia> Kb = Average(Field(-w*b/bz), mask=abs(bz).>1e-5)
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] assertscalar(op::String)
    @ GPUArraysCore /glade/work/tomasc/.julia/packages/GPUArraysCore/uOYfN/src/GPUArraysCore.jl:103
  [3] getindex(::CUDA.CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}, ::Int64, ::Int64, ::Int64)
    @ GPUArrays /glade/work/tomasc/.julia/packages/GPUArrays/5XhED/src/host/indexing.jl:9
...

And the examples above also don't take into account the immersed boundary masking, which I'm assuming I should also take into account manually, but not sure how.

What would be the best approach here?

[simone-silvestri]

For a conditional average you have to use the condition keyword argument.
The mask kwarg is the value you substitute where condition evaluates to false.
The immersed boundary condition is included by default.

For the moment we cannot condition on another field unfortunately, we can probably rework ConditionalOperation to make this work
For the moment a workaround is to generate a map to condition on

@inline large_stratification(i, j, k, grid, b) = abs(b[i, j, k]) > 1e-5
@compute condition = Field(KernelFunctionOperation{Center, Center, Face}(large_stratification, grid, bz))
Kb = Average(Field(-w*b/bz); condition)

This will average all cells with abs(bz) > 1e-5. If you want to average where abs(bz) < 1e-5 you need to do

@inline small_stratification(i, j, k, grid, b) = abs(b[i, j, k]) > 1e-5
@compute condition = Field(KernelFunctionOperation{Center, Center, Face}(small_stratification, grid, bz))
Kb = Average(Field(-w*b/bz); condition)

[tomchor]

Thanks @simone-silvestri. That makes sense.

Btw I think there's a typo in your second snippet. The first line should have a < sign I'm assuming?

[glwagner]

You can condition on a field because

Oceananigans.jl/src/AbstractOperations/conditional_operations.jl

Line 134 in f4d7a94

@inline get_condition(condition::AbstractArray, i, j, k, grid, args...) = @inbounds condition[i, j, k]

I think @simone-silvestri is referring to the fact that we don't have a streamlined interface for creating such "condition fields" (or condition operations). For example we don't support the operators | and & in abstract operations. It's easy to add though or do yourself.

[tomchor]

I'm getting an error using @simone-silvestri's suggestion. Maybe I'm missing something?

julia> using Oceananigans

julia> using Oceananigans.Fields: @compute

julia> grid_base = RectilinearGrid(GPU(), size=(4, 4, 4), extent=(1, 1, 1));

julia> ibg = ImmersedBoundaryGrid(grid_base, GridFittedBottom((x, y) -> -1/2));

julia> model = NonhydrostaticModel(grid=ibg, tracers=:b);

julia> u, v, w = model.velocities; b=model.tracers.b;

julia> bz = Field(∂z(b));

julia> @inline large_stratification(i, j, k, grid, bz) = abs(bz[i, j, k]) > 1e-10
large_stratification (generic function with 1 method)

julia> @compute condition = Field(KernelFunctionOperation{Center, Center, Face}(large_stratification, ibg, bz));

julia> Kb = Average(Field(-w*b/bz); condition);
ERROR: MethodError: condition_operand(::typeof(identity), ::Oceananigans.AbstractOperations.GridMetricOperation{Center, Center, Face, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, GPU}, Float64, typeof(Oceananigans.Operators.Vᶜᶜᶠ)}, ::Field{Center, Center, Face, KernelFunctionOperation{Center, Center, Face, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, GPU}, Float64, typeof(large_stratification), Tuple{Field{Center, Center, Face, Oceananigans.AbstractOperations.Derivative{Center, Center, Face, typeof(Oceananigans.Operators.∂zᶜᶜᶠ), Field{Center, Center, Center, Nothing, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing}}, typeof(Oceananigans.Operators.identity4), typeof(∂z), ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, GPU}, Float64}, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Oceananigans.Fields.FieldStatus{Float64}, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing}}}}, ImmersedBoundaryGrid{Float64, Periodic, Periodic, Bounded, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, GridFittedBottom{OffsetArrays.OffsetMatrix{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}, Oceananigans.ImmersedBoundaries.CenterImmersedCondition}, Nothing, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CUDA.CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Oceananigans.Fields.FieldStatus{Float64}, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing}}, ::Int64) is ambiguous.

Candidates:
  condition_operand(func::Function, operand::Oceananigans.Fields.AbstractField, condition::AbstractArray, mask)
    @ Oceananigans.AbstractOperations /glade/work/tomasc/.julia/packages/Oceananigans/ZGvk3/src/AbstractOperations/conditional_operations.jl:96
  condition_operand(func::Function, op::Oceananigans.Fields.AbstractField{<:Any, <:Any, <:Any, <:ImmersedBoundaryGrid}, cond, mask)
    @ Oceananigans.ImmersedBoundaries /glade/work/tomasc/.julia/packages/Oceananigans/ZGvk3/src/ImmersedBoundaries/immersed_reductions.jl:18

If I try

Kb = Average(Field(-w*b/bz); mask=condition);

it (apparently) works, but I'm guessing that overrides the immersed boundary masking? I'm a bit confused also about the difference between mask and condition, so I'm not sure what's appropriate here.