Allow users to mask output

src/ImmersedBoundaries/mask_immersed_field.jl

@@ -6,24 +6,24 @@ using Oceananigans.Fields: location, ZReducedField, Field
 instantiate(T::Type) = T()
 instantiate(t) = t
 
-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, 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::Field, grid::ImmersedBoundaryGrid, loc, value)
+    mask_immersed!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
 
 masks `field` defined on `grid` with a value `val` at locations where `peripheral_node` evaluates to `true`
 """
-function mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
+function mask_immersed!(field::Field, grid::ImmersedBoundaryGrid, loc, value)
     arch = architecture(field)
     loc = instantiate.(loc)
-    launch!(arch, grid, :xyz, _mask_immersed_field!, field, loc, grid, value)
+    launch!(arch, grid, :xyz, _mask_immersed!, field, loc, grid, value)
     return nothing
 end
 
 
-@kernel function _mask_immersed_field!(field, loc, grid, value)
+@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

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

@@ -301,8 +301,8 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac
         set!(free_surface.η, free_surface.state.η̅)
 
         # Wait for predictor velocity update step to complete and mask it if immersed boundary.
-        mask_immersed_field!(model.velocities.u)
-        mask_immersed_field!(model.velocities.v)
+        mask_immersed!(model.velocities.u)
+        mask_immersed!(model.velocities.v)
     end
 
     fill_halo_regions!(free_surface.η)
@@ -336,8 +336,8 @@ function setup_split_explicit!(auxiliary, state, η, grid, Gu, Gv, Guⁿ, Gvⁿ,
     initialize_free_surface_state!(state, η)
 
     # Wait for predictor velocity update step to complete and mask it if immersed boundary.
-    mask_immersed_field!(Gu)
-    mask_immersed_field!(Gv)
+    mask_immersed!(Gu)
+    mask_immersed!(Gv)
 
     # Compute barotropic mode of tendency fields
     compute_barotropic_mode!(auxiliary.Gᵁ, auxiliary.Gⱽ, grid, 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: calculate_diffusivities!
-using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node
+using Oceananigans.ImmersedBoundaries: mask_immersed!, mask_immersed_field_xy!, inactive_node
 using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!
 
 import Oceananigans.TimeSteppers: update_state!
@@ -52,7 +52,7 @@ 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)

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: calculate_diffusivities!
 using Oceananigans.Fields: compute!
-using Oceananigans.ImmersedBoundaries: mask_immersed_field!
+using Oceananigans.ImmersedBoundaries: mask_immersed!
 
 import Oceananigans.TimeSteppers: update_state!
 
@@ -18,7 +18,7 @@ they are called in the end.
 function update_state!(model::NonhydrostaticModel, callbacks=[])
 
     # Mask immersed tracers
-    foreach(mask_immersed_field!, model.tracers)
+    foreach(mask_immersed!, model.tracers)
 
     # Fill halos for velocities and tracers
     fill_halo_regions!(merge(model.velocities, model.tracers),  model.clock, fields(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!
 
 import Oceananigans.TimeSteppers: update_state!
 
@@ -12,7 +12,7 @@ If `callbacks` are provided (in an array), they are called in the end.
 function update_state!(model::ShallowWaterModel, callbacks=[])
 
     # Mask immersed fields
-    foreach(mask_immersed_field!, model.solution)
+    foreach(mask_immersed!, model.solution)
 
     calculate_diffusivities!(model.diffusivity_fields, model.closure, model)
 

src/OutputWriters/fetch_output.jl

@@ -1,16 +1,25 @@
 using CUDA
 
 using Oceananigans.Fields: AbstractField, compute_at!
+using Oceananigans.ImmersedBoundaries: mask_immersed!
 using Oceananigans.LagrangianParticleTracking: LagrangianParticles
 
 # Needed to support `fetch_output` with `model::Nothing`.
 time(model) = model.clock.time
 time(::Nothing) = nothing
 
-fetch_output(output, model) = output(model)
+function fetch_output(output, model, mask_value)
+    fetched = output(model)
+    if fetched isa Field
+        !isnothing(mask_value) && mask_immersed!(fetched, mask_value)
+        return parent(fetched)
+    end
+    return fetched
+end
 
-function fetch_output(field::AbstractField, model)
+function fetch_output(field::AbstractField, model, mask_value)
     compute_at!(field, time(model))
+    !isnothing(mask_value) && mask_immersed!(field, mask_value)
     return parent(field)
 end
 
@@ -38,6 +47,6 @@ 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_value)
     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_value :: 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_value = 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_value)
 end
 
 function initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model)

src/OutputWriters/netcdf_output_writer.jl

@@ -10,7 +10,7 @@ using Oceananigans.Utils: versioninfo_with_gpu, oceananigans_versioninfo, pretty
 using Oceananigans.TimeSteppers: float_or_date_time
 using Oceananigans.Fields: reduced_dimensions, reduced_location, location, validate_indices
 
-mutable struct NetCDFOutputWriter{D, O, T, A} <: AbstractOutputWriter
+mutable struct NetCDFOutputWriter{D, O, T, A, FT} <: AbstractOutputWriter
     filepath :: String
     dataset :: D
     outputs :: O
@@ -19,6 +19,7 @@ mutable struct NetCDFOutputWriter{D, O, T, A} <: AbstractOutputWriter
     array_type :: A
     previous :: Float64
     verbose :: Bool
+    mask_value :: FT
 end
 
 ext(::Type{NetCDFOutputWriter}) = ".nc"
@@ -160,6 +161,7 @@ end
                                    array_type = Array{Float64},
                                       indices = nothing,
                                    with_halos = false,
+                                mask_immersed = nothing,
                             global_attributes = Dict(),
                             output_attributes = Dict(),
                                    dimensions = Dict(),
@@ -200,6 +202,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`.
+
 - `dimensions`: A `Dict` of dimension tuples to apply to outputs (required for function outputs).
 
 ## File management
@@ -330,6 +335,7 @@ function NetCDFOutputWriter(model, outputs; filename, schedule,
                                    array_type = Array{Float64},
                                       indices = (:, :, :),
                                    with_halos = false,
+                                mask_immersed = nothing,
                             global_attributes = Dict(),
                             output_attributes = Dict(),
                                    dimensions = Dict(),
@@ -425,7 +431,10 @@ function NetCDFOutputWriter(model, outputs; filename, schedule,
 
     close(dataset)
 
-    return NetCDFOutputWriter(filepath, dataset, outputs, schedule, overwrite_existing, array_type, 0.0, verbose)
+    mask_value = mask_immersed isa Number ? eltype(model.grid)(mask_immersed) : mask_immersed
+
+    return NetCDFOutputWriter(filepath, dataset, outputs, schedule, overwrite_existing,
+                              array_type, 0.0, verbose, mask_value)
 end
 
 get_default_dimension_attributes(grid::AbstractRectilinearGrid) =

src/OutputWriters/windowed_time_average.jl

@@ -182,7 +182,7 @@ function WindowedTimeAverage(operand, model=nothing; schedule, fetch_operand=tru
         result = similar(operand)
         result .= operand
     end
-        
+
     return WindowedTimeAverage(result, operand, 0.0, 0, 0.0, schedule, fetch_operand)
 end
 

test/test_dynamics.jl

@@ -1,7 +1,7 @@
 include("dependencies_for_runtests.jl")
 
 using Oceananigans.TurbulenceClosures: viscosity, ThreeDimensionalFormulation, HorizontalFormulation, VerticalFormulation
-using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom, mask_immersed_field!
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom, mask_immersed!
 
 function relative_error(u_num, u, time)
     u_ans = Field(location(u_num), u_num.grid)

test/test_immersed_advection.jl

@@ -1,6 +1,6 @@
 include("dependencies_for_runtests.jl")
 
-using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary, mask_immersed_field!
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary, mask_immersed!
 using Oceananigans.Advection: 
         _symmetric_interpolate_xᶠᵃᵃ,
         _symmetric_interpolate_xᶜᵃᵃ,
@@ -103,7 +103,7 @@ for arch in archs
 
         c = CenterField(ibg)
         set!(c, 1.0)
-        mask_immersed_field!(c)
+        mask_immersed!(c)
         fill_halo_regions!(c)
 
         for adv in advection_schemes, buffer in [1, 2, 3, 4, 5]
@@ -141,8 +141,8 @@ for arch in archs
         set!(u, 1.0)
         set!(v, 1.0)
 
-        mask_immersed_field!(u)
-        mask_immersed_field!(v)
+        mask_immersed!(u)
+        mask_immersed!(v)
 
         fill_halo_regions!(u)
         fill_halo_regions!(v)