Merge pull request #2932 from CliMA/glwagner-patch-1

Clean up ImmersedBoundaries module

src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl

@@ -2,6 +2,7 @@ using Adapt
 using CUDA: CuArray
 using OffsetArrays: OffsetArray
 using Oceananigans.Utils: getnamewrapper
+using Oceananigans.Grids: total_size
 using Oceananigans.Fields: fill_halo_regions!
 using Oceananigans.Architectures: arch_array
 using Oceananigans.BoundaryConditions: FBC
@@ -86,15 +87,12 @@ function ImmersedBoundaryGrid(grid, ib::AbstractGridFittedBottom{<:OffsetArray})
 end
 
 function validate_ib_size(grid, ib)
-    Nx, Ny, _ = size(grid)
-    Hx, Hy, _ = halo_size(grid)
+    bottom_height_size = total_size((Center, Center, Nothing), grid)[1:2]
 
-    bottom_height_size = (Nx, Ny) .+ 2 .* (Hx, Hy)
-
-    # Check that the size of a bottom field are 
-    # consistent with the size of the field
-    any(size(ib.bottom_height) .!= bottom_height_size) && 
+    size(ib.bottom_height) != bottom_height_size &&
         throw(ArgumentError("The dimensions of the immersed boundary $(size(ib.bottom_height)) do not match the grid size $(bottom_height_size)"))
+
+    return nothing
 end
 
 @inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom{<:Any, <:InterfaceImmersedCondition})
@@ -118,12 +116,12 @@ Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.
 ##### Implicit vertical diffusion
 #####
 
-####
-#### For a center solver we have to check the interface "solidity" at faces k+1 in both the
-#### Upper diagonal and the Lower diagonal 
-#### (because of tridiagonal convention where lower_diagonal on row k is found at k-1)
-#### Same goes for the face solver, where we check at centers k in both Upper and lower diagonal
-####
+#####
+##### For a center solver we have to check the interface "solidity" at faces k+1 in both the
+##### Upper diagonal and the Lower diagonal 
+##### (because of tridiagonal convention where lower_diagonal on row k is found at k-1)
+##### Same goes for the face solver, where we check at centers k in both Upper and lower diagonal
+#####
 
 @inline immersed_ivd_peripheral_node(i, j, k, ibg, LX, LY, ::Center) = immersed_peripheral_node(i, j, k+1, ibg, LX, LY, Face())
 @inline immersed_ivd_peripheral_node(i, j, k, ibg, LX, LY, ::Face)   = immersed_peripheral_node(i, j, k,   ibg, LX, LY, Center())
@@ -199,15 +197,15 @@ immersed_cell(i, j, k, grid, ib) = _immersed_cell(i, j, k, grid, ib)
 
 # support for Flat grids
 using Oceananigans.Grids: AbstractGrid
-for ImmBoundary in [:GridFittedBottom, :GridFittedBoundary]
+for IB in [:GridFittedBottom, :GridFittedBoundary]
     @eval begin
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, <:Any, <:Any}, ib::$ImmBoundary) = _immersed_cell(1, j, k, grid, ib)
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, Flat, <:Any}, ib::$ImmBoundary) = _immersed_cell(i, 1, k, grid, ib)
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, <:Any, Flat}, ib::$ImmBoundary) = _immersed_cell(i, j, 1, grid, ib)
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, Flat, <:Any},  ib::$ImmBoundary) = _immersed_cell(1, 1, k, grid, ib)
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, <:Any, Flat},  ib::$ImmBoundary) = _immersed_cell(1, j, 1, grid, ib)
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, Flat, Flat},  ib::$ImmBoundary) = _immersed_cell(i, 1, 1, grid, ib)
-        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, Flat, Flat},   ib::$ImmBoundary) = _immersed_cell(1, 1, 1, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, <:Any, <:Any}, ib::$IB) = _immersed_cell(1, j, k, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, Flat, <:Any}, ib::$IB) = _immersed_cell(i, 1, k, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, <:Any, Flat}, ib::$IB) = _immersed_cell(i, j, 1, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, Flat, <:Any},  ib::$IB) = _immersed_cell(1, 1, k, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, <:Any, Flat},  ib::$IB) = _immersed_cell(1, j, 1, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, Flat, Flat},  ib::$IB) = _immersed_cell(i, 1, 1, grid, ib)
+        @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, Flat, Flat},   ib::$IB) = _immersed_cell(1, 1, 1, grid, ib)
     end
 end