Adds methods to mask_immersed_field! for BinaryOperations

src/ImmersedBoundaries/mask_immersed_field.jl

@@ -1,5 +1,6 @@
 using KernelAbstractions: @kernel, @index
 using Statistics
+using Oceananigans.AbstractOperations: BinaryOperation
 using Oceananigans.Fields: location, ZReducedField, Field
 
 instantiate(T::Type) = T()
@@ -9,6 +10,26 @@ mask_immersed_field!(field, grid, loc, value) = nothing
 mask_immersed_field!(field::Field, value=zero(eltype(field.grid))) =
     mask_immersed_field!(field, field.grid, location(field), value)
 
+mask_immersed_field!(::Number, args...) = nothing
+
+function mask_immersed_field!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(+)}, value=zero(eltype(bop)))
+    a_value = ifelse(bop.b isa Number, -bop.b, value)
+    mask_immersed_field!(bop.a, a_value)
+
+    b_value = ifelse(bop.a isa Number, -bop.a, value)
+    mask_immersed_field!(bop.b, b_value)
+    return nothing
+end
+
+function mask_immersed_field!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(-)}, value=zero(eltype(bop)))
+    a_value = ifelse(bop.b isa Number, bop.b, value)
+    mask_immersed_field!(bop.a, a_value)
+
+    b_value = ifelse(bop.a isa Number, bop.a, value)
+    mask_immersed_field!(bop.b, b_value)
+    return nothing
+end
+
 """
     mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
 
@@ -21,7 +42,6 @@ function mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, val
     return nothing
 end
 
-
 @kernel function _mask_immersed_field!(field, loc, grid, value)
     i, j, k = @index(Global, NTuple)
     @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field)
@@ -32,6 +52,26 @@ mask_immersed_field_xy!(::Nothing, args...; kw...) = nothing
 mask_immersed_field_xy!(field, value=zero(eltype(field.grid)); k, mask = peripheral_node) =
     mask_immersed_field_xy!(field, field.grid, location(field), value; k, mask)
 
+mask_immersed_field_xy!(::Number, args...) = nothing
+
+function mask_immersed_field_xy!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(+)}, value=zero(eltype(bop)))
+    a_value = ifelse(bop.b isa Number, -bop.b, value)
+    mask_immersed_field_xy!(bop.a, a_value)
+
+    b_value = ifelse(bop.a isa Number, -bop.a, value)
+    mask_immersed_field_xy!(bop.b, b_value)
+    return nothing
+end
+
+function mask_immersed_field_xy!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(-)}, value=zero(eltype(bop)))
+    a_value = ifelse(bop.b isa Number, bop.b, value)
+    mask_immersed_field_xy!(bop.a, a_value)
+
+    b_value = ifelse(bop.a isa Number, bop.a, value)
+    mask_immersed_field_xy!(bop.b, b_value)
+    return nothing
+end
+
 """
     mask_immersed_field_xy!(field::Field, grid::ImmersedBoundaryGrid, loc, value; k, mask=peripheral_node)
 

validation/immersed_boundaries/mask_binary_operations.jl

@@ -0,0 +1,43 @@
+using Oceananigans, CairoMakie
+using Oceananigans.ImmersedBoundaries: mask_immersed_field!
+
+underlying_grid = RectilinearGrid(architecture; size = (32, 32, 32), x = (-20, 20), y = (-20, 20), z = (-20, 0))
+
+bathy(x, y) = 20 * exp(-(x^2 + y^2) / (2*10^2)) - 20
+
+grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bathy))
+
+c1 = CenterField(grid)
+c2 = CenterField(grid)
+
+Σc = c1 + c2
+Πc = c1 * c2
+ratio_c = c1 / c2
+pow_c = c1 ^ c2
+
+# Σc = 5, Πc = 6, ratio_c = 2/3, pow_c = 8
+
+fig = Figure()
+
+axs = [Axis(fig[1, 1], title = "+"),
+       Axis(fig[1, 2], title = "*"),
+       Axis(fig[2, 1], title = "/"),
+       Axis(fig[2, 2], title = "^")]
+
+xc, yc, zc = nodes(grid, Center(), Center(), Center())
+
+cos = []
+
+for (n, f) in enumerate([Σc, Πc, ratio_c, pow_c])
+    c1 isa Number || set!(c1, 2)
+    c2 isa Number || set!(c2, 3)
+
+    mask_immersed_field!(f)
+
+    @info f[1, 1, 16], f[16, 16, 16]
+
+    heatmap!(axs[n], xc, yc, [f[i, j, 16] for i in 1:grid.Nx, j in 1:grid.Ny], colorrange = (0, 8))
+end
+
+fig
+