Refactor SplitExplicitFreeSurface

ext/OceananigansMakieExt.jl

@@ -38,7 +38,7 @@ end
 convert_arguments(pl::Type{<:AbstractPlot}, f::Field) =
     convert_arguments(pl, convert_field_argument(f)...)
 
-function convert_arguments(pl::Type{<:AbstractPlot}, fop::AbstractOperation)
+function convert_arguments(pl::Type{<:AbstractPlot}, op::AbstractOperation)
     f = Field(op)
     compute!(f)
     return convert_arguments(pl, f)

geostrophic_adjustment.jl

@@ -0,0 +1,95 @@
+# Run this script with
+#
+# $ mpiexec -n 4 julia --project distributed_nonhydrostatic_two_dimensional_turbulence.jl
+#
+# for example.
+#
+# You also probably should set
+#
+# $ export JULIA_NUM_THREADS=1
+
+using MPI
+using Oceananigans
+using Oceananigans.DistributedComputations
+using Oceananigans.DistributedComputations: Sizes
+using Oceananigans.Grids: topology, architecture
+using Oceananigans.Units: kilometers, meters
+using Printf
+using JLD2
+
+arch = CPU()
+
+# Distribute problem irregularly
+Nx = 80
+Lh = 100kilometers
+Lz = 400meters
+
+topo = (Bounded, Periodic, Bounded)
+grid = RectilinearGrid(arch,
+                       size = (Nx, 3, 2),
+                       x = (0, Lh),
+                       y = (0, Lh),
+                       z = (-Lz, 0),
+                       topology = topo)
+
+# bottom(x, y) = x > 80kilometers && x < 90kilometers ? 100 : -500meters
+# grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom); active_cells_map = true)
+
+coriolis = FPlane(f=1e-4)
+
+model = HydrostaticFreeSurfaceModel(; grid,
+                                      coriolis,
+                                      free_surface = SplitExplicitFreeSurface(grid; substeps=10)) 
+
+gaussian(x, L) = exp(-x^2 / 2L^2)
+
+U = 0.1 # geostrophic velocity
+L = Lh / 40 # gaussian width
+x₀ = Lh / 4 # gaussian center
+vᵍ(x, y, z) = -U * (x - x₀) / L * gaussian(x - x₀, L)
+
+g = model.free_surface.gravitational_acceleration
+η = model.free_surface.η
+η₀ = coriolis.f * U * L / g # geostrophic free surface amplitude
+
+ηᵍ(x) = η₀ * gaussian(x - x₀, L)
+ηⁱ(x, y, z) = 2 * ηᵍ(x)
+
+set!(model, v = vᵍ)
+set!(model, η = ηⁱ)
+
+gravity_wave_speed = sqrt(g * grid.Lz) # hydrostatic (shallow water) gravity wave speed
+Δt = 2 * model.grid.Δxᶜᵃᵃ / gravity_wave_speed
+simulation = Simulation(model; Δt, stop_iteration = 1000)
+
+ut = []
+vt = []
+ηt = []
+
+save_u(sim) = push!(ut, deepcopy(sim.model.velocities.u))  
+save_v(sim) = push!(vt, deepcopy(sim.model.velocities.v))
+save_η(sim) = push!(ηt, deepcopy(sim.model.free_surface.η))
+
+function progress_message(sim) 
+    @info @sprintf("[%.2f%%], iteration: %d, time: %.3f, max|w|: %.2e",
+    100 * sim.model.clock.time / sim.stop_time, sim.model.clock.iteration,
+    sim.model.clock.time, maximum(abs, sim.model.velocities.u))
+end
+
+simulation.callbacks[:save_η]   = Callback(save_η, IterationInterval(1))
+simulation.callbacks[:save_v]   = Callback(save_v, IterationInterval(1))
+simulation.callbacks[:save_u]   = Callback(save_u, IterationInterval(1))
+simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(10))
+
+run!(simulation)
+
+iter  = Observable(1)
+η_img = @lift(interior(ηt[$iter], :, 1, 1)) 
+fig   = Figure()
+ax    = Axis(fig[1, 1])
+lines!(ax, η_img, color=:black)
+
+GLMakie.record(fig, "test.mp4", 1:length(ηt), framerate=2) do i
+    @info "Animating iteration $i/$(length(ηt))..."
+    iter[] = i
+end

src/Grids/Grids.jl

@@ -18,6 +18,7 @@ export xnodes, ynodes, znodes, λnodes, φnodes
 export spacings
 export xspacings, yspacings, zspacings, xspacing, yspacing, zspacing
 export minimum_xspacing, minimum_yspacing, minimum_zspacing
+export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ
 export offset_data, new_data
 export on_architecture
 

src/Grids/grid_utils.jl

@@ -314,6 +314,15 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) =
              name, prettysummary(minimum(parent(Δ))),
              name, prettysummary(maximum(parent(Δ))))
 
+#####
+##### Static column depth
+#####
+
+@inline static_column_depthᶜᶜᵃ(i, j, grid) = grid.Lz
+@inline static_column_depthᶜᶠᵃ(i, j, grid) = grid.Lz
+@inline static_column_depthᶠᶜᵃ(i, j, grid) = grid.Lz
+@inline static_column_depthᶠᶠᵃ(i, j, grid) = grid.Lz
+
 #####
 ##### Spherical geometry
 #####

src/ImmersedBoundaries/ImmersedBoundaries.jl

@@ -17,15 +17,18 @@ import Base: show, summary
 import Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_y, cpu_face_constructor_z,
                            x_domain, y_domain, z_domain
 
-import Oceananigans.Grids: architecture, on_architecture, with_halo, inflate_halo_size_one_dimension,
+import Oceananigans.Grids: architecture, with_halo, inflate_halo_size_one_dimension,
                            xnode, ynode, znode, λnode, φnode, node,
                            ξnode, ηnode, rnode,
                            ξname, ηname, rname, node_names,
                            xnodes, ynodes, znodes, λnodes, φnodes, nodes,
                            ξnodes, ηnodes, rnodes,
+                           static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ,
                            inactive_cell
 
 
+import Oceananigans.Architectures: on_architecture
+
 import Oceananigans.Fields: fractional_x_index, fractional_y_index, fractional_z_index
 
 """
@@ -92,10 +95,6 @@ with_halo(halo, ibg::ImmersedBoundaryGrid) =
 inflate_halo_size_one_dimension(req_H, old_H, _, ::IBG)            = max(req_H + 1, old_H)
 inflate_halo_size_one_dimension(req_H, old_H, ::Type{Flat}, ::IBG) = 0
 
-# Defining the bottom
-@inline z_bottom(i, j, grid) = znode(i, j, 1, grid, c, c, f)
-@inline z_bottom(i, j, ibg::IBG) = error("The function `bottom` has not been defined for $(summary(ibg))!")
-
 function Base.summary(grid::ImmersedBoundaryGrid)
     FT = eltype(grid)
     TX, TY, TZ = topology(grid)

src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl

@@ -18,21 +18,3 @@ const AGFB = AbstractGridFittedBoundary
     @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, Flat, Flat},  ib::AGFB) = _immersed_cell(i, 1, 1, grid, ib)
     @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, Flat, Flat},   ib::AGFB) = _immersed_cell(1, 1, 1, grid, ib)
 end
-
-function clamp_bottom_height!(bottom_field, grid)
-    launch!(architecture(grid), grid, :xy, _clamp_bottom_height!, bottom_field, grid)
-    return nothing
-end
-
-const c = Center()
-const f = Face()
-
-@kernel function _clamp_bottom_height!(z, grid)
-    i, j = @index(Global, NTuple)
-    Nz = size(grid, 3)
-    zmin = znode(i, j, 1,    grid, c, c, f)
-    zmax = znode(i, j, Nz+1, grid, c, c, f)
-    @inbounds z[i, j, 1] = clamp(z[i, j, 1], zmin, zmax)
-end
-
-

src/ImmersedBoundaries/grid_fitted_bottom.jl

@@ -15,17 +15,17 @@ abstract type AbstractGridFittedBottom{H} <: AbstractGridFittedBoundary end
 
 # To enable comparison with PartialCellBottom in the limiting case that
 # fractional cell height is 1.0.
-struct CenterImmersedCondition end
 struct InterfaceImmersedCondition end
-
-Base.summary(::CenterImmersedCondition) = "CenterImmersedCondition"
-Base.summary(::InterfaceImmersedCondition) = "InterfaceImmersedCondition"
+struct CenterImmersedCondition end
 
 struct GridFittedBottom{H, I} <: AbstractGridFittedBottom{H}
     bottom_height :: H
     immersed_condition :: I
 end
 
+Base.summary(::CenterImmersedCondition) = "CenterImmersedCondition"
+Base.summary(::InterfaceImmersedCondition) = "InterfaceImmersedCondition"
+
 const GFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:GridFittedBottom}
 
 """
@@ -36,6 +36,7 @@ Return a bottom immersed boundary.
 Keyword Arguments
 =================
 
+
 * `bottom_height`: an array or function that gives the height of the
                    bottom in absolute ``z`` coordinates.
 
@@ -72,50 +73,77 @@ Base.summary(ib::GridFittedBottom{<:Function}) = @sprintf("GridFittedBottom(%s)"
 function Base.show(io::IO, ib::GridFittedBottom)
     print(io, summary(ib), '\n')
     print(io, "├── bottom_height: ", prettysummary(ib.bottom_height), '\n')
-    print(io, "└── immersed_condition: ", summary(ib.immersed_condition))
 end
 
-@inline z_bottom(i, j, ibg::GFBIBG) = @inbounds ibg.immersed_boundary.bottom_height[i, j, 1]
+on_architecture(arch, ib::GridFittedBottom) = GridFittedBottom(on_architecture(arch, ib.bottom_height), ib.immersed_condition)
+
+function on_architecture(arch, ib::GridFittedBottom{<:Field})
+    architecture(ib.bottom_height) == arch && return ib
+    arch_grid = on_architecture(arch, ib.bottom_height.grid)
+    new_bottom_height = Field{Center, Center, Nothing}(arch_grid)
+    set!(new_bottom_height, ib.bottom_height)
+    fill_halo_regions!(new_bottom_height)
+    return GridFittedBottom(new_bottom_height, ib.immersed_condition)
+end
+
+Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.bottom_height), adapt(to, ib.immersed_condition))
 
 """
     ImmersedBoundaryGrid(grid, ib::GridFittedBottom)
 
 Return a grid with `GridFittedBottom` immersed boundary (`ib`).
 
-Computes `ib.bottom_height` and wraps it in a Field.
+Computes `ib.bottom_height` and wraps it in a Field. `ib.bottom_height` is the z-coordinate of top-most interface
+of the last ``immersed`` cell in the column.
 """
 function ImmersedBoundaryGrid(grid, ib::GridFittedBottom)
     bottom_field = Field{Center, Center, Nothing}(grid)
     set!(bottom_field, ib.bottom_height)
-    @apply_regionally clamp_bottom_height!(bottom_field, grid)
+    @apply_regionally compute_numerical_bottom_height!(bottom_field, grid, ib)
     fill_halo_regions!(bottom_field)
-    new_ib = GridFittedBottom(bottom_field, ib.immersed_condition)
+    new_ib = GridFittedBottom(bottom_field)
     TX, TY, TZ = topology(grid)
     return ImmersedBoundaryGrid{TX, TY, TZ}(grid, new_ib)
 end
 
-@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom{<:Any, <:InterfaceImmersedCondition})
-    z = znode(i, j, k+1, underlying_grid, c, c, f)
-    h = @inbounds ib.bottom_height[i, j, 1]
-    return z ≤ h
+compute_numerical_bottom_height!(bottom_field, grid, ib) = 
+    launch!(architecture(grid), grid, :xy, _compute_numerical_bottom_height!, bottom_field, grid, ib)
+
+@kernel function _compute_numerical_bottom_height!(bottom_field, grid, ib::GridFittedBottom)
+    i, j = @index(Global, NTuple)
+    zb = @inbounds bottom_field[i, j, 1]
+    @inbounds bottom_field[i, j, 1] = znode(i, j, 1, grid, c, c, f)
+    condition = ib.immersed_condition
+    for k in 1:grid.Nz
+        z⁺ = znode(i, j, k+1, grid, c, c, f)
+        z  = znode(i, j, k,   grid, c, c, c)
+        bottom_cell = ifelse(condition isa CenterImmersedCondition, z ≤ zb, z⁺ ≤ zb)
+        @inbounds bottom_field[i, j, 1] = ifelse(bottom_cell, z⁺, zb)
+    end
 end
 
-@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom{<:Any, <:CenterImmersedCondition})
-    z = znode(i, j, k, underlying_grid, c, c, c)
-    h = @inbounds ib.bottom_height[i, j, 1]
-    return z ≤ h
+@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom)
+    z  = znode(i, j, k, underlying_grid, c, c, c)
+    zb = @inbounds ib.bottom_height[i, j, 1]
+    return z ≤ zb
 end
 
-on_architecture(arch, ib::GridFittedBottom) = GridFittedBottom(ib.bottom_height, ib.immersed_condition)
+#####
+##### Static column depth
+#####
 
-function on_architecture(arch, ib::GridFittedBottom{<:Field})
-    architecture(ib.bottom_height) == arch && return ib
-    arch_grid = on_architecture(arch, ib.bottom_height.grid)
-    new_bottom_height = Field{Center, Center, Nothing}(arch_grid)
-    set!(new_bottom_height, ib.bottom_height)
-    fill_halo_regions!(new_bottom_height)
-    return GridFittedBottom(new_bottom_height, ib.immersed_condition)
-end
+const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom}
+
+@inline static_column_depthᶜᶜᵃ(i, j, ibg::AGFBIBG) = @inbounds znode(i, j, ibg.Nz+1, ibg, c, c, f) - ibg.immersed_boundary.bottom_height[i, j, 1] 
+@inline static_column_depthᶜᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i, j-1, ibg), static_column_depthᶜᶜᵃ(i, j, ibg))
+@inline static_column_depthᶠᶜᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i-1, j, ibg), static_column_depthᶜᶜᵃ(i, j, ibg))
+@inline static_column_depthᶠᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶠᶜᵃ(i, j-1, ibg), static_column_depthᶠᶜᵃ(i, j, ibg))
+
+# Make sure column_height works for horizontally-Flat topologies.
+XFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Flat, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom}
+YFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:AbstractGridFittedBottom}
 
-Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.bottom_height),
-                                                                             ib.immersed_condition)
+@inline static_column_depthᶠᶜᵃ(i, j, ibg::XFlatAGFIBG) = static_column_depthᶜᶜᵃ(i, j, ibg)
+@inline static_column_depthᶜᶠᵃ(i, j, ibg::YFlatAGFIBG) = static_column_depthᶜᶜᵃ(i, j, ibg)
+@inline static_column_depthᶠᶠᵃ(i, j, ibg::XFlatAGFIBG) = static_column_depthᶜᶠᵃ(i, j, ibg)
+@inline static_column_depthᶠᶠᵃ(i, j, ibg::YFlatAGFIBG) = static_column_depthᶠᶜᵃ(i, j, ibg)