Fix a bug in the adapt of PartialCellBottom (#3682)

* fix a bug in the adapt * first fix * bugfix * these should be all the bugs, still need to add the tests * another bugfix * Clamp bottom_height to top and bottom of grid * Bugfix * Implement clamping for both GF and PC bottom * Convert internal tide example to use extension * apply_regionally *sigh* * Disambiguate * using CairoMakie --------- Co-authored-by: Gregory L. Wagner <[email protected]>
CliMA · Sep 27, 2024 · 675e0ba · 675e0ba
1 parent 8700865
commit 675e0ba
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 36 deletions.
diff --git a/examples/internal_tide.jl b/examples/internal_tide.jl
@@ -14,6 +14,7 @@
 
 using Oceananigans
 using Oceananigans.Units
+using Oceananigans.ImmersedBoundaries: PartialCellBottom
 
 # ## Grid
 
@@ -46,7 +47,7 @@ width = 20kilometers
 hill(x) = h₀ * exp(-x^2 / 2width^2)
 bottom(x) = - H + hill(x)
 
-grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom))
+grid = ImmersedBoundaryGrid(underlying_grid, PartialCellBottom(bottom))
 
 # Let's see how the domain with the bathymetry is.
 
@@ -132,7 +133,7 @@ bᵢ(x, z) = Nᵢ² * z
 
 set!(model, u=uᵢ, b=bᵢ)
 
-# Now let's built a `Simulation`.
+# Now let's build a `Simulation`.
 
 Δt = 5minutes
 stop_time = 4days
@@ -234,9 +235,9 @@ n = Observable(1)
 title = @lift @sprintf("t = %1.2f days = %1.2f T₂",
                        round(times[$n] / day, digits=2) , round(times[$n] / T₂, digits=2))
 
-u′ₙ = @lift u′_t[$n]
- wₙ = @lift  w_t[$n]
-N²ₙ = @lift N²_t[$n]
+u′n = @lift u′_t[$n]
+ wn = @lift  w_t[$n]
+N²n = @lift N²_t[$n]
 
 axis_kwargs = (xlabel = "x [km]",
                ylabel = "z [km]",
@@ -248,15 +249,15 @@ fig = Figure(size = (700, 900))
 fig[1, :] = Label(fig, title, fontsize=24, tellwidth=false)
 
 ax_u = Axis(fig[2, 1]; title = "u'-velocity", axis_kwargs...)
-hm_u = heatmap!(ax_u, xu, zu, u′ₙ; colorrange = (-umax, umax), colormap = :balance)
+hm_u = heatmap!(ax_u, xu, zu, u′n; nan_color=:gray, colorrange=(-umax, umax), colormap=:balance)
 Colorbar(fig[2, 2], hm_u, label = "m s⁻¹")
 
 ax_w = Axis(fig[3, 1]; title = "w-velocity", axis_kwargs...)
-hm_w = heatmap!(ax_w, xw, zw, wₙ; colorrange = (-wmax, wmax), colormap = :balance)
+hm_w = heatmap!(ax_w, xw, zw, wn; nan_color=:gray, colorrange=(-wmax, wmax), colormap=:balance)
 Colorbar(fig[3, 2], hm_w, label = "m s⁻¹")
 
 ax_N² = Axis(fig[4, 1]; title = "stratification N²", axis_kwargs...)
-hm_N² = heatmap!(ax_N², xN², zN², N²ₙ; colorrange = (0.9Nᵢ², 1.1Nᵢ²), colormap = :thermal)
+hm_N² = heatmap!(ax_N², xN², zN², N²n; nan_color=:gray, colorrange=(0.9Nᵢ², 1.1Nᵢ²), colormap=:magma)
 Colorbar(fig[4, 2], hm_N², label = "s⁻²")
 
 fig

diff --git a/src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl b/src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl
@@ -53,3 +53,20 @@ const AGFB = AbstractGridFittedBoundary
     @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
+
+
diff --git a/src/ImmersedBoundaries/grid_fitted_bottom.jl b/src/ImmersedBoundaries/grid_fitted_bottom.jl
@@ -54,15 +54,15 @@ Keyword Arguments
 GridFittedBottom(bottom_height) = GridFittedBottom(bottom_height, CenterImmersedCondition())
 
 function Base.summary(ib::GridFittedBottom)
-    hmax  = maximum(ib.bottom_height)
-    hmin  = minimum(ib.bottom_height)
-    hmean = mean(ib.bottom_height)
+    zmax  = maximum(ib.bottom_height)
+    zmin  = minimum(ib.bottom_height)
+    zmean = mean(ib.bottom_height)
 
     summary1 = "GridFittedBottom("
 
-    summary2 = string("mean(z)=", prettysummary(hmean),
-                      ", min(z)=", prettysummary(hmin),
-                      ", max(z)=", prettysummary(hmax))
+    summary2 = string("mean(z)=", prettysummary(zmean),
+                      ", min(z)=", prettysummary(zmin),
+                      ", max(z)=", prettysummary(zmax))
 
     summary3 = ")"
 
@@ -87,6 +87,7 @@ Computes `ib.bottom_height` and wraps it in a Field.
 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)
     fill_halo_regions!(bottom_field)
     new_ib = GridFittedBottom(bottom_field, ib.immersed_condition)
     TX, TY, TZ = topology(grid)

diff --git a/src/ImmersedBoundaries/partial_cell_bottom.jl b/src/ImmersedBoundaries/partial_cell_bottom.jl
@@ -11,18 +11,18 @@ struct PartialCellBottom{H, E} <: AbstractGridFittedBottom{H}
     minimum_fractional_cell_height :: E
 end
 
-const PCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:PartialCellBottom}
+const PCBIBG{FT, TX, TY, TZ} = ImmersedBoundaryGrid{FT, TX, TY, TZ, <:Any, <:PartialCellBottom} where {FT, TX, TY, TZ}
 
 function Base.summary(ib::PartialCellBottom)
-    hmax = maximum(parent(ib.bottom_height))
-    hmin = minimum(parent(ib.bottom_height))
-    hmean = mean(parent(ib.bottom_height))
+    zmax = maximum(parent(ib.bottom_height))
+    zmin = minimum(parent(ib.bottom_height))
+    zmean = mean(parent(ib.bottom_height))
 
     summary1 = "PartialCellBottom("
 
-    summary2 = string("mean(z)=", prettysummary(hmean),
-                      ", min(z)=", prettysummary(hmin),
-                      ", max(z)=", prettysummary(hmax),
+    summary2 = string("mean(zb)=", prettysummary(zmean),
+                      ", min(zb)=", prettysummary(zmin),
+                      ", max(zb)=", prettysummary(zmax),
                       ", ϵ=", prettysummary(ib.minimum_fractional_cell_height))
 
     summary3 = ")"
@@ -63,6 +63,7 @@ end
 function ImmersedBoundaryGrid(grid, ib::PartialCellBottom)
     bottom_field = Field{Center, Center, Nothing}(grid)
     set!(bottom_field, ib.bottom_height)
+    @apply_regionally clamp_bottom_height!(bottom_field, grid)
     fill_halo_regions!(bottom_field)
     new_ib = PartialCellBottom(bottom_field, ib.minimum_fractional_cell_height)
     TX, TY, TZ = topology(grid)
@@ -77,32 +78,49 @@ function on_architecture(arch, ib::PartialCellBottom{<:Field})
     return PartialCellBottom(new_bottom_height, ib.minimum_fractional_cell_height)
 end
 
-Adapt.adapt_structure(to, ib::PartialCellBottom) = PartialCellBottom(adapt(to, ib.bottom_height.data),
+Adapt.adapt_structure(to, ib::PartialCellBottom) = PartialCellBottom(adapt(to, ib.bottom_height),
                                                                      ib.minimum_fractional_cell_height)     
 
-on_architecture(to, ib::PartialCellBottom) = PartialCellBottom(on_architecture(to, ib.bottom_height.data),
+on_architecture(to, ib::PartialCellBottom) = PartialCellBottom(on_architecture(to, ib.bottom_height),
                                                                on_architecture(to, ib.minimum_fractional_cell_height))     
 
 """
-
-        --x--
-          ∘   k+1
-    k+1 --x--    ↑      <- node z
-          ∘   k  | Δz
-      k --x--    ↓
+    immersed     underlying
+
+      --x--        --x--
+            
+            
+        ∘   ↑        ∘   k+1
+            |
+            |               
+  k+1 --x-- |  k+1 --x--    ↑      <- node z
+        ∘   ↓               |
+   zb ⋅⋅x⋅⋅                 |
+                            |
+                     ∘   k  | Δz
+                            |
+                            |
+                 k --x--    ↓
       
-Criterion is h ≥ z - ϵ Δz
+Criterion is zb ≥ z - ϵ Δz
 
 """
 @inline function _immersed_cell(i, j, k, underlying_grid, ib::PartialCellBottom)
-    # Face node above current cell
-    z = znode(i, j, k+1, underlying_grid, c, c, f)
-    h = @inbounds ib.bottom_height[i, j, 1]
-    return z ≤ h
+    # Face node below current cell
+    z  = znode(i, j, k, underlying_grid, c, c, f)
+    zb = @inbounds ib.bottom_height[i, j, 1]
+    ϵ  = ib.minimum_fractional_cell_height
+    # z + Δz is equal to the face above the current cell
+    Δz = Δzᶜᶜᶜ(i, j, k, underlying_grid)
+    return (z + Δz * (1 - ϵ)) ≤ zb
 end
 
-@inline bottom_cell(i, j, k, ibg::PCBIBG) = !immersed_cell(i, j, k,   ibg.underlying_grid, ibg.immersed_boundary) &
-                                             immersed_cell(i, j, k-1, ibg.underlying_grid, ibg.immersed_boundary)
+@inline function bottom_cell(i, j, k, ibg::PCBIBG)
+    grid = ibg.underlying_grid
+    ib = ibg.immersed_boundary
+    # This one's not immersed, but the next one down is
+    return !immersed_cell(i, j, k, grid, ib) & immersed_cell(i, j, k-1, grid, ib)
+end
 
 @inline function Δzᶜᶜᶜ(i, j, k, ibg::PCBIBG)
     underlying_grid = ibg.underlying_grid
@@ -145,4 +163,18 @@ end
 @inline Δzᶜᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶠ(i, j-1, k, ibg), Δzᶜᶜᶠ(i, j, k, ibg))      
 @inline Δzᶠᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶠᶜᶠ(i, j-1, k, ibg), Δzᶠᶜᶠ(i, j, k, ibg))
 
+# Make sure Δz works for horizontally-Flat topologies.
+# (There's no point in using z-Flat with PartialCellBottom).
+XFlatPCBIBG = ImmersedBoundaryGrid{<:Any, <:Flat, <:Any, <:Any, <:Any, <:PartialCellBottom}
+YFlatPCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:PartialCellBottom}
+
+@inline Δzᶠᶜᶜ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶜᶜ(i, j, k, ibg)
+@inline Δzᶠᶜᶠ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶜᶠ(i, j, k, ibg)
+@inline Δzᶜᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δzᶜᶜᶜ(i, j, k, ibg)
+
+@inline Δzᶜᶠᶠ(i, j, k, ibg::YFlatPCBIBG) = Δzᶜᶜᶠ(i, j, k, ibg)
+@inline Δzᶠᶠᶜ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶠᶜ(i, j, k, ibg)
+@inline Δzᶠᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δzᶠᶜᶜ(i, j, k, ibg)
+
 @inline z_bottom(i, j, ibg::PCBIBG) = @inbounds ibg.immersed_boundary.bottom_height[i, j, 1]
+
diff --git a/test/test_immersed_boundary_grid.jl b/test/test_immersed_boundary_grid.jl
@@ -0,0 +1,16 @@
+include("dependencies_for_runtests.jl")
+
+grid = RectilinearGrid(; size=(2, 2, 2), extent = (1, 1, 1))
+
+@testset "Testing Immersed Boundaries" begin
+
+    @info "Testing the immersed boundary construction..."
+
+    bottom(x, y) = -1 + 0.5 * exp(-x^2 - y^2)
+    ibg = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom))
+
+    # Unit test (bottom is at the right position)
+
+    @info "Testing stably stratified initial conditions..."
+
+end