Allow users to mask output

src/Grids/inactive_node.jl

@@ -130,4 +130,3 @@ lies on the boundary between inactive and active cells in a `Bounded` direction.
 Return `true` when the location `(LX, LY, LZ)` lies on a boundary.
 """
 @inline boundary_node(i, j, k, grid, LX, LY, LZ) = peripheral_node(i, j, k, grid, LX, LY, LZ) & !inactive_node(i, j, k, grid, LX, LY, LZ)
-

src/ImmersedBoundaries/ImmersedBoundaries.jl

@@ -216,7 +216,7 @@ We then have
 
 * `inactive_node(i, 1, 1, grid, f, c, c) = false`
 
-As well as
+as well as
 
 * `inactive_node(i,   1, 1, grid, c, c, c) = false`
 * `inactive_node(i-1, 1, 1, grid, c, c, c) = true`
@@ -303,7 +303,7 @@ include("partial_cell_bottom.jl")
 include("conditional_fluxes.jl")
 include("immersed_boundary_condition.jl")
 include("conditional_differences.jl")
-include("mask_immersed_field.jl")
+include("mask_immersed.jl")
 include("immersed_reductions.jl")
 
 # Extension to distributed Immersed boundaries

src/ImmersedBoundaries/mask_immersed.jl

@@ -0,0 +1,58 @@
+using KernelAbstractions: @kernel, @index
+using Statistics
+using Oceananigans.Architectures: architecture
+using Oceananigans.Fields: location, ZReducedField, Field
+
+instantiate(T::Type) = T()
+instantiate(t) = t
+
+mask_immersed!(field, grid, loc, value) = nothing
+mask_immersed!(field::Field, value=zero(eltype(field.grid))) =
+    mask_immersed!(field, field.grid, location(field), value)
+
+"""
+    mask_immersed!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
+
+Mask `field` defined on `grid` with a `value` at `loc`ations at grid points where [`peripheral_node`](@ref)
+is to `true`.
+"""
+function mask_immersed!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
+    arch = architecture(field)
+    loc = instantiate.(loc)
+    launch!(arch, grid, :xyz, _mask_immersed!, field, loc, grid, value)
+    return nothing
+end
+
+
+@kernel function _mask_immersed!(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)
+end
+
+mask_immersed_horizontal!(field,     args...; kw...) = nothing
+mask_immersed_horizontal!(::Nothing, args...; kw...) = nothing
+mask_immersed_horizontal!(field, value=zero(eltype(field.grid)); k, mask = peripheral_node) =
+    mask_immersed_horizontal!(field, field.grid, location(field), value; k, mask)
+
+"""
+    mask_immersed_horizontal!(field::Field, grid::ImmersedBoundaryGrid, loc, value; k, mask)
+
+Mask `field` on `grid` with a `value` on the slices `[:, :, k]` where `mask` is `true`.
+"""
+function mask_immersed_horizontal!(field::Field, grid::ImmersedBoundaryGrid, loc, value; k, mask)
+    arch = architecture(field)
+    loc = instantiate.(loc)
+    return launch!(arch, grid, :xy, _mask_immersed_horizontal!, field, loc, grid, value, k, mask)
+end
+
+@kernel function _mask_immersed_horizontal!(field, loc, grid, value, k, mask)
+    i, j = @index(Global, NTuple)
+    @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field, mask)
+end
+
+#####
+##### Masking for GridFittedBoundary
+#####
+
+@inline scalar_mask(i, j, k, grid, ::AbstractGridFittedBoundary, LX, LY, LZ, value, field, mask=peripheral_node) =
+    @inbounds ifelse(mask(i, j, k, grid, LX, LY, LZ), value, field[i, j, k])

src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl

@@ -150,7 +150,7 @@ end
 
 function local_compute_integrated_volume_flux!(∫ᶻQ, velocities, arch)
 
-    foreach(mask_immersed_field!, velocities)
+    foreach(mask_immersed!, velocities)
 
     # Compute barotropic volume flux. Blocking.
     compute_vertically_integrated_volume_flux!(∫ᶻQ, velocities)

src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl

@@ -299,13 +299,14 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac
         set!(free_surface.η, free_surface.state.η̅)
     end
 
+
     fields_to_fill = (free_surface.state.U̅, free_surface.state.V̅)
     fill_halo_regions!(fields_to_fill; async = true)
 
     # Preparing velocities for the barotropic correction
     @apply_regionally begin 
-        mask_immersed_field!(model.velocities.u)
-        mask_immersed_field!(model.velocities.v)
+        mask_immersed!(model.velocities.u)
+        mask_immersed!(model.velocities.v)
     end
 
     return nothing
@@ -381,6 +382,11 @@ function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ)
     Guⁿ = model.timestepper.Gⁿ.u
     Gvⁿ = model.timestepper.Gⁿ.v
 
+    mask_immersed!(Gu⁻)
+    mask_immersed!(Gv⁻)
+    mask_immersed!(Guⁿ)
+    mask_immersed!(Gvⁿ)
+
     auxiliary = free_surface.auxiliary
 
     @apply_regionally setup_split_explicit_tendency!(auxiliary, free_surface_grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -4,7 +4,7 @@ using Oceananigans.BoundaryConditions
 using Oceananigans: UpdateStateCallsite
 using Oceananigans.Biogeochemistry: update_biogeochemical_state!
 using Oceananigans.TurbulenceClosures: compute_diffusivities!
-using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node
+using Oceananigans.ImmersedBoundaries: mask_immersed!, mask_immersed_horizontal!, inactive_node
 using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!, p_kernel_parameters
 using Oceananigans.Fields: replace_horizontal_vector_halos!
 
@@ -59,10 +59,10 @@ function mask_immersed_model_fields!(model, grid)
 
     foreach(fields_to_mask) do field
         if field !== η
-            mask_immersed_field!(field)
+            mask_immersed!(field)
         end
     end
-    mask_immersed_field_xy!(η, k=size(grid, 3)+1, mask = inactive_node)
+    mask_immersed_horizontal!(η, k=size(grid, 3)+1, mask = inactive_node)
 
     return nothing
 end

src/Models/NonhydrostaticModels/pressure_correction.jl

@@ -8,7 +8,7 @@ Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `vel
 function calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
 
     # Mask immersed velocities
-    foreach(mask_immersed_field!, model.velocities)
+    foreach(mask_immersed!, model.velocities)
 
     fill_halo_regions!(model.velocities, model.clock, fields(model))
 

src/Models/NonhydrostaticModels/set_nonhydrostatic_model.jl

@@ -42,8 +42,8 @@ function set!(model::NonhydrostaticModel; enforce_incompressibility=true, kwargs
     end
 
     # Apply a mask
-    foreach(mask_immersed_field!, model.tracers)
-    foreach(mask_immersed_field!, model.velocities)
+    foreach(mask_immersed!, model.tracers)
+    foreach(mask_immersed!, model.velocities)
     update_state!(model)
 
     if enforce_incompressibility

src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl

@@ -4,7 +4,7 @@ using Oceananigans.BoundaryConditions
 using Oceananigans.Biogeochemistry: update_biogeochemical_state!
 using Oceananigans.TurbulenceClosures: compute_diffusivities!
 using Oceananigans.Fields: compute!
-using Oceananigans.ImmersedBoundaries: mask_immersed_field!
+using Oceananigans.ImmersedBoundaries: mask_immersed!
 using Oceananigans.Models: update_model_field_time_series!
 
 import Oceananigans.TimeSteppers: update_state!
@@ -20,7 +20,7 @@ function update_state!(model::NonhydrostaticModel, callbacks=[]; compute_tendenc
 
     # Mask immersed tracers
     foreach(model.tracers) do tracer
-        @apply_regionally mask_immersed_field!(tracer)
+        @apply_regionally mask_immersed!(tracer)
     end
 
     # Update all FieldTimeSeries used in the model

src/Models/ShallowWaterModels/update_shallow_water_state.jl

@@ -1,5 +1,5 @@
 using Oceananigans.BoundaryConditions: fill_halo_regions!
-using Oceananigans.ImmersedBoundaries: mask_immersed_field!
+using Oceananigans.ImmersedBoundaries: mask_immersed!
 using Oceananigans.Models: update_model_field_time_series!
 
 import Oceananigans.TimeSteppers: update_state!
@@ -13,7 +13,7 @@ If `callbacks` are provided (in an array), they are called in the end.
 function update_state!(model::ShallowWaterModel, callbacks=[]; compute_tendencies = true)
 
     # Mask immersed fields
-    foreach(mask_immersed_field!, model.solution)
+    foreach(mask_immersed!, model.solution)
 
     # Update possible FieldTimeSeries used in the model
     update_model_field_time_series!(model, model.clock)

src/OutputWriters/fetch_output.jl

@@ -1,16 +1,28 @@
 using CUDA
 
 using Oceananigans.Fields: AbstractField, compute_at!, ZeroField
+using Oceananigans.ImmersedBoundaries: mask_immersed!
 using Oceananigans.Models.LagrangianParticleTracking: LagrangianParticles
 
 # Needed to support `fetch_output` with `model::Nothing`.
 time(model) = model.clock.time
 time(::Nothing) = nothing
 
-fetch_output(output, model) = output(model)
+# Default fetch_output with mask_immersed = nothing
+fetch_output(output, model) = fetch_output(output, model, nothing)
 
-function fetch_output(field::AbstractField, model)
+function fetch_output(output, model, mask_immersed)
+    fetched = output(model)
+    if fetched isa Field
+        !isnothing(mask_immersed) && mask_immersed!(fetched, mask_immersed)
+        return parent(fetched)
+    end
+    return fetched
+end
+
+function fetch_output(field::AbstractField, model, mask_immersed)
     compute_at!(field, time(model))
+    !isnothing(mask_immersed) && mask_immersed!(field, mask_immersed)
     return parent(field)
 end
 
@@ -38,7 +50,7 @@ convert_output(outputs::NamedTuple, writer) =
     NamedTuple(name => convert_output(outputs[name], writer) for name in keys(outputs))
 
 function fetch_and_convert_output(output, model, writer)
-    fetched = fetch_output(output, model)
+    fetched = fetch_output(output, model, writer.mask_immersed)
     return convert_output(fetched, writer)
 end
 

src/OutputWriters/jld2_output_writer.jl

@@ -9,7 +9,7 @@ default_included_properties(::NonhydrostaticModel) = [:grid, :coriolis, :buoyanc
 default_included_properties(::ShallowWaterModel) = [:grid, :coriolis, :closure]
 default_included_properties(::HydrostaticFreeSurfaceModel) = [:grid, :coriolis, :buoyancy, :closure]
 
-mutable struct JLD2OutputWriter{O, T, D, IF, IN, KW} <: AbstractOutputWriter
+mutable struct JLD2OutputWriter{O, T, D, IF, IN, KW, FT} <: AbstractOutputWriter
     filepath :: String
     outputs :: O
     schedule :: T
@@ -21,6 +21,7 @@ mutable struct JLD2OutputWriter{O, T, D, IF, IN, KW} <: AbstractOutputWriter
     overwrite_existing :: Bool
     verbose :: Bool
     jld2_kw :: KW
+    mask_immersed :: FT
 end
 
 noinit(args...) = nothing
@@ -32,6 +33,7 @@ ext(::Type{JLD2OutputWriter}) = ".jld2"
                           indices = (:, :, :),
                        with_halos = false,
                        array_type = Array{Float64},
+                    mask_immersed = nothing,
                      max_filesize = Inf,
                overwrite_existing = false,
                              init = noinit,
@@ -78,6 +80,9 @@ Keyword arguments
 - `array_type`: The array type to which output arrays are converted to prior to saving.
                 Default: `Array{Float64}`.
 
+- `mask_immersed`: The value with which immersed boundary regions are filled with before saving.
+                   You can either choose a numerical value or `NaN`. Default: `nothing`.
+
 ## File management
 
 - `max_filesize`: The writer will stop writing to the output file once the file size exceeds `max_filesize`,
@@ -163,6 +168,7 @@ function JLD2OutputWriter(model, outputs; filename, schedule,
                                indices = (:, :, :),
                             with_halos = false,
                             array_type = Array{Float64},
+                         mask_immersed = nothing,
                           max_filesize = Inf,
                     overwrite_existing = false,
                                   init = noinit,
@@ -175,7 +181,7 @@ function JLD2OutputWriter(model, outputs; filename, schedule,
     filename = auto_extension(filename, ".jld2")
     filepath = joinpath(dir, filename)
     overwrite_existing && isfile(filepath) && rm(filepath, force=true)
-    
+
     outputs = NamedTuple(Symbol(name) => construct_output(outputs[name], model.grid, indices, with_halos)
                          for name in keys(outputs))
 
@@ -184,8 +190,10 @@ function JLD2OutputWriter(model, outputs; filename, schedule,
 
     initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model)
 
-    return JLD2OutputWriter(filepath, outputs, schedule, array_type, init,
-                            including, part, max_filesize, overwrite_existing, verbose, jld2_kw)
+    mask_immersed = mask_immersed isa Number ? eltype(model.grid)(mask_immersed) : mask_immersed
+
+    return JLD2OutputWriter(filepath, outputs, schedule, array_type, init, including, part,
+                            max_filesize, overwrite_existing, verbose, jld2_kw, mask_immersed)
 end
 
 function initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model)