(0.96.32) Fix zspacing() for PartialCellBottom immersed boundaries and add tests for ImmersedGrids (#4572)

* Add missing Δz methods for PartialCellBottom immersed boundaries - Add all Δz function combinations for VSPCBIBG - Each Δz method maps to corresponding Δr method for vertically-static grids - Fixes zspacing() not adjusting for PartialCellBottom immersed boundaries

* bugfix

* add better tests for immersed boundaries

* add tests for gridFittedBoundary

* add rspacings and a test for it

* standardize notation

* leanup

* fix tests on GPU

* add test for mutableVerticalDiscretization

* bump patch release

---------

Co-authored-by: Tomas Chor <tomchor@login19.chn.perlmutter.nersc.gov>

Author: tomchor

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.96.31"
+version = "0.96.32"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/AbstractOperations/grid_metrics.jl

@@ -2,9 +2,9 @@ using Adapt
 using Oceananigans.Operators
 using Oceananigans.Grids: AbstractGrid
 using Oceananigans.Fields: AbstractField, default_indices, location
-using Oceananigans.Operators: Δx, Δy, Δz, Ax, Δλ, Δφ, Ay, Az, volume
+using Oceananigans.Operators: Δx, Δy, Δz, Δr, Ax, Δλ, Δφ, Ay, Az, volume
 
-import Oceananigans.Grids: xspacings, yspacings, zspacings, λspacings, φspacings
+import Oceananigans.Grids: xspacings, yspacings, zspacings, rspacings, λspacings, φspacings
 
 const AbstractGridMetric = Union{typeof(Δx),
                                  typeof(Δy),
@@ -159,6 +159,32 @@ function zspacings(grid, ℓx, ℓy, ℓz)
     return Δz_op
 end
 
+"""
+    rspacings(grid, ℓx, ℓy, ℓz)
+
+Return a `KernelFunctionOperation` that computes the grid spacings for `grid`
+in the ``r`` direction at location `ℓx, ℓy, ℓz`.
+
+Examples
+========
+```jldoctest
+julia> using Oceananigans
+
+julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1));
+
+julia> rspacings(grid, Center(), Center(), Face())
+KernelFunctionOperation at (Center, Center, Face)
+├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo
+├── kernel_function: Δr (generic function with 28 methods)
+└── arguments: ("Center", "Center", "Face")
+```
+"""
+function rspacings(grid, ℓx, ℓy, ℓz)
+    LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz))
+    Δr_op = KernelFunctionOperation{LX, LY, LZ}(Δr, grid, ℓx, ℓy, ℓz)
+    return Δr_op
+end
+
 """
     λspacings(grid, ℓx, ℓy, ℓz)
 

src/ImmersedBoundaries/grid_fitted_bottom.jl

@@ -72,7 +72,7 @@ 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, "└── bottom_height: ", prettysummary(ib.bottom_height), '\n')
 end
 
 on_architecture(arch, ib::GridFittedBottom) = GridFittedBottom(on_architecture(arch, ib.bottom_height), ib.immersed_condition)

src/ImmersedBoundaries/partial_cell_bottom.jl

@@ -1,9 +1,10 @@
 using Oceananigans.Utils: prettysummary
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Grids: bottommost_active_node
+using Oceananigans.Grids: bottommost_active_node, AbstractStaticGrid
 using Printf
 
-import Oceananigans.Operators: Δrᶜᶜᶜ, Δrᶜᶜᶠ, Δrᶜᶠᶜ, Δrᶜᶠᶠ, Δrᶠᶜᶜ, Δrᶠᶜᶠ, Δrᶠᶠᶜ, Δrᶠᶠᶠ
+import Oceananigans.Operators: Δrᶜᶜᶜ, Δrᶜᶜᶠ, Δrᶜᶠᶜ, Δrᶜᶠᶠ, Δrᶠᶜᶜ, Δrᶠᶜᶠ, Δrᶠᶠᶜ, Δrᶠᶠᶠ,
+                               Δzᶜᶜᶜ, Δzᶜᶜᶠ, Δzᶜᶠᶜ, Δzᶜᶠᶠ, Δzᶠᶜᶜ, Δzᶠᶜᶠ, Δzᶠᶠᶜ, Δzᶠᶠᶠ
 
 #####
 ##### PartialCellBottom
@@ -137,20 +138,20 @@ Criterion is zb ≥ z - ϵ Δz
 
 """
 @inline function _immersed_cell(i, j, k, underlying_grid, ib::PartialCellBottom)
-    z⁺ = rnode(i, j, k+1, underlying_grid, c, c, f)
+    r⁺ = rnode(i, j, k + 1, underlying_grid, c, c, f)
     ϵ  = ib.minimum_fractional_cell_height
-    Δz = Δrᶜᶜᶜ(i, j, k, underlying_grid)
-    z★ = z⁺ - Δz * ϵ
-    zb = @inbounds ib.bottom_height[i, j, 1]
-    return z★ < zb
+    Δr = Δrᶜᶜᶜ(i, j, k, underlying_grid)
+    r★ = r⁺ - Δr * ϵ
+    rᵇ = @inbounds ib.bottom_height[i, j, 1]
+    return r★ < rᵇ
 end
 
 @inline function Δrᶜᶜᶜ(i, j, k, ibg::PCBIBG)
     underlying_grid = ibg.underlying_grid
     ib = ibg.immersed_boundary
 
     # Get node at face above and defining nodes on c,c,f
-    rᶜᶜᶠₖ₊₁ = rnode(i, j, k+1, underlying_grid, c, c, f)
+    r⁺ = rnode(i, j, k + 1, underlying_grid, c, c, f)
 
     # Get bottom r-coordinate and fractional Δr parameter
     rᵇ = @inbounds ib.bottom_height[i, j, 1]
@@ -159,13 +160,13 @@ end
     at_the_bottom = bottommost_active_node(i, j, k, ibg, c, c, c)
 
     full_Δr    = Δrᶜᶜᶜ(i, j, k, ibg.underlying_grid)
-    partial_Δr = rᶜᶜᶠₖ₊₁ - rᵇ
+    partial_Δr = r⁺ - rᵇ
 
     return ifelse(at_the_bottom, partial_Δr, full_Δr)
 end
 
 @inline function Δrᶜᶜᶠ(i, j, k, ibg::PCBIBG)
-    just_above_bottom = bottommost_active_node(i, j, k, c, c, f)
+    just_above_bottom = bottommost_active_node(i, j, k, ibg, c, c, f)
     rᶜ = rnode(i, j, k, ibg.underlying_grid, c, c, c)
     rᶠ = rnode(i, j, k, ibg.underlying_grid, c, c, f)
 
@@ -195,3 +196,14 @@ YFlatPCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:Partial
 @inline Δrᶜᶠᶠ(i, j, k, ibg::YFlatPCBIBG) = Δrᶜᶜᶠ(i, j, k, ibg)
 @inline Δrᶠᶠᶜ(i, j, k, ibg::XFlatPCBIBG) = Δrᶜᶠᶜ(i, j, k, ibg)
 @inline Δrᶠᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δrᶠᶜᶜ(i, j, k, ibg)
+
+# Vertically-static, partial cell bottom, immergsed boundary grid
+VSPCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:AbstractStaticGrid, <:PartialCellBottom}
+@inline Δzᶜᶜᶜ(i, j, k, ibg::VSPCBIBG) = Δrᶜᶜᶜ(i, j, k, ibg)
+@inline Δzᶠᶜᶜ(i, j, k, ibg::VSPCBIBG) = Δrᶠᶜᶜ(i, j, k, ibg)
+@inline Δzᶜᶠᶜ(i, j, k, ibg::VSPCBIBG) = Δrᶜᶠᶜ(i, j, k, ibg)
+@inline Δzᶜᶜᶠ(i, j, k, ibg::VSPCBIBG) = Δrᶜᶜᶠ(i, j, k, ibg)
+@inline Δzᶠᶠᶜ(i, j, k, ibg::VSPCBIBG) = Δrᶠᶠᶜ(i, j, k, ibg)
+@inline Δzᶜᶠᶠ(i, j, k, ibg::VSPCBIBG) = Δrᶜᶠᶠ(i, j, k, ibg)
+@inline Δzᶠᶜᶠ(i, j, k, ibg::VSPCBIBG) = Δrᶠᶜᶠ(i, j, k, ibg)
+@inline Δzᶠᶠᶠ(i, j, k, ibg::VSPCBIBG) = Δrᶠᶠᶠ(i, j, k, ibg)