Clean up immersed boundary interface (#4140)

* test it out

* add testing

* correct active map computation

* better naming

* bugfix

* bugfix

* bugfix

* bugfix

* bugfix

* comment

* materialize immersed bonudary

* this should work

* THIS SHOULD WORK

Author: simone-silvestri

ext/OceananigansReactantExt/Grids.jl

@@ -3,16 +3,17 @@ module Grids
 using Reactant
 using Oceananigans
 using Oceananigans.Architectures: ReactantState, CPU
-using Oceananigans.Grids: AbstractGrid, StaticVerticalDiscretization, MutableVerticalDiscretization
+using Oceananigans.Grids: AbstractGrid, AbstractUnderlyingGrid, StaticVerticalDiscretization, MutableVerticalDiscretization
 using Oceananigans.Fields: Field
-using Oceananigans.ImmersedBoundaries: GridFittedBottom
+using Oceananigans.ImmersedBoundaries: GridFittedBottom, AbstractImmersedBoundary
 
 import ..OceananigansReactantExt: deconcretize
 import Oceananigans.Grids: LatitudeLongitudeGrid, RectilinearGrid, OrthogonalSphericalShellGrid
 import Oceananigans.OrthogonalSphericalShellGrids: RotatedLatitudeLongitudeGrid
 import Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 
 const ReactantGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:ReactantState}
+const ReactantUnderlyingGrid{FT, TX, TY, TZ, CZ} = AbstractUnderlyingGrid{FT, TX, TY, TZ, CZ, <:ReactantState}
 
 deconcretize(z::StaticVerticalDiscretization) =
     StaticVerticalDiscretization(
@@ -61,13 +62,13 @@ end
 # This low-level constructor supports the external package OrthogonalSphericalShellGrids.jl.
 function OrthogonalSphericalShellGrid{TX, TY, TZ}(arch::ReactantState,
                                                   Nx, Ny, Nz, Hx, Hy, Hz,
-                                                  Lz :: FT,
+                                                     Lz :: FT,
                                                    λᶜᶜᵃ :: CC,  λᶠᶜᵃ :: FC,  λᶜᶠᵃ :: CF,  λᶠᶠᵃ :: FF,
                                                    φᶜᶜᵃ :: CC,  φᶠᶜᵃ :: FC,  φᶜᶠᵃ :: CF,  φᶠᶠᵃ :: FF, z :: Z,
                                                   Δxᶜᶜᵃ :: CC, Δxᶠᶜᵃ :: FC, Δxᶜᶠᵃ :: CF, Δxᶠᶠᵃ :: FF,
                                                   Δyᶜᶜᵃ :: CC, Δyᶠᶜᵃ :: FC, Δyᶜᶠᵃ :: CF, Δyᶠᶠᵃ :: FF, 
                                                   Azᶜᶜᵃ :: CC, Azᶠᶜᵃ :: FC, Azᶜᶠᵃ :: CF, Azᶠᶠᵃ :: FF,
-                                                  radius :: FT,
+                                                 radius :: FT,
                                                   conformal_mapping :: Map) where {TX, TY, TZ, FT, Z, Map,
                                                                                    CC, FC, CF, FF, C}
 
@@ -113,10 +114,7 @@ function reactant_immersed_boundary_grid(grid, ib; active_cells_map)
                                             ibg.interior_active_cells, ibg.active_z_columns)
 end
 
-ImmersedBoundaryGrid(grid::ReactantGrid, ib::GridFittedBottom; active_cells_map::Bool=true) =
-    reactant_immersed_boundary_grid(grid, ib; active_cells_map)
-
-ImmersedBoundaryGrid(grid::ReactantGrid, ib; active_cells_map::Bool=true) =
+ImmersedBoundaryGrid(grid::ReactantUnderlyingGrid, ib::AbstractImmersedBoundary; active_cells_map::Bool=true) =
     reactant_immersed_boundary_grid(grid, ib; active_cells_map)
 
 end # module

src/DistributedComputations/distributed_immersed_boundaries.jl

@@ -106,35 +106,35 @@ const DistributedActiveCellsIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:An
 # For the same reason we need to construct `south` and `north` maps if we partition the domain in the y-direction.
 # Therefore, the `interior_active_cells` in this case is a `NamedTuple` containing 5 elements.
 # Note that boundary-adjacent maps corresponding to non-partitioned directions are set to `nothing`
-function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:DistributedGrid})
+function map_interior_active_cells(grid::DistributedGrid, ib)
 
-    arch = architecture(ibg)
+    arch = architecture(grid)
 
     # If we using a synchronized architecture, nothing
     # changes with serial execution.
     if arch isa SynchronizedDistributed
-        return interior_active_indices(ibg; parameters = :xyz)
+        return interior_active_indices(grid, ib; parameters = :xyz)
     end
 
     Rx, Ry, _  = arch.ranks
-    Tx, Ty, _  = topology(ibg)
-    Nx, Ny, Nz = size(ibg)
-    Hx, Hy, _  = halo_size(ibg)
+    Tx, Ty, _  = topology(grid)
+    Nx, Ny, Nz = size(grid)
+    Hx, Hy, _  = halo_size(grid)
 
     west_boundary  = (1:Hx,       1:Ny, 1:Nz)
     east_boundary  = (Nx-Hx+1:Nx, 1:Ny, 1:Nz)
     south_boundary = (1:Nx, 1:Hy,       1:Nz)
     north_boundary = (1:Nx, Ny-Hy+1:Ny, 1:Nz)
 
-    include_west  = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected)
-    include_east  = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected)
-    include_south = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected)
-    include_north = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected)
+    include_west  = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected)
+    include_east  = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected)
+    include_south = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected)
+    include_north = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected)
 
-    west_halo_dependent_cells  = interior_active_indices(ibg; parameters = KernelParameters(west_boundary...))
-    east_halo_dependent_cells  = interior_active_indices(ibg; parameters = KernelParameters(east_boundary...))
-    south_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(south_boundary...))
-    north_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(north_boundary...))
+    west_halo_dependent_cells  = interior_active_indices(grid, ib; parameters = KernelParameters(west_boundary...))
+    east_halo_dependent_cells  = interior_active_indices(grid, ib; parameters = KernelParameters(east_boundary...))
+    south_halo_dependent_cells = interior_active_indices(grid, ib; parameters = KernelParameters(south_boundary...))
+    north_halo_dependent_cells = interior_active_indices(grid, ib; parameters = KernelParameters(north_boundary...))
 
     west_halo_dependent_cells  = ifelse(include_west,  west_halo_dependent_cells,  nothing)
     east_halo_dependent_cells  = ifelse(include_east,  east_halo_dependent_cells,  nothing)
@@ -147,7 +147,7 @@ function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any
     ox = Rx == 1 || Tx == RightConnected ? 0 : Hx
     oy = Ry == 1 || Ty == RightConnected ? 0 : Hy
 
-    halo_independent_cells = interior_active_indices(ibg; parameters = KernelParameters((nx, ny, Nz), (ox, oy, 0)))
+    halo_independent_cells = interior_active_indices(grid, ib; parameters = KernelParameters((nx, ny, Nz), (ox, oy, 0)))
 
     return (; halo_independent_cells, 
               west_halo_dependent_cells, 

src/ImmersedBoundaries/active_cells_map.jl

@@ -52,56 +52,38 @@ A tuple of indices corresponding to the linear index.
 """
 @inline active_linear_index_to_tuple(idx, active_cells_map) = @inbounds Base.map(Int, active_cells_map[idx])
 
-function ImmersedBoundaryGrid(grid, ib; active_cells_map::Bool=true) 
-    ibg = ImmersedBoundaryGrid(grid, ib)
-    TX, TY, TZ = topology(ibg)
-    
-    # Create the cells map on the CPU, then switch it to the GPU
-    if active_cells_map 
-        interior_active_cells = map_interior_active_cells(ibg)
-        active_z_columns = map_active_z_columns(ibg)
-    else
-        interior_active_cells = nothing
-        active_z_columns = nothing
-    end
-
-    return ImmersedBoundaryGrid{TX, TY, TZ}(ibg.underlying_grid, 
-                                            ibg.immersed_boundary, 
-                                            interior_active_cells,
-                                            active_z_columns)
-end
-
 with_halo(halo, ibg::ActiveCellsIBG) =
     ImmersedBoundaryGrid(with_halo(halo, ibg.underlying_grid), ibg.immersed_boundary; active_cells_map = true)
 
-@inline active_cell(i, j, k, ibg) = !immersed_cell(i, j, k, ibg)
-@inline active_column(i, j, k, grid, column) = column[i, j, k] != 0
+@inline active_cell(i, j, k, grid, ib) = !immersed_cell(i, j, k, grid, ib)
 
-@kernel function _set_active_indices!(active_cells_field, grid)
+@kernel function _set_active_indices!(active_cells_field, grid, ib)
     i, j, k = @index(Global, NTuple)
-    @inbounds active_cells_field[i, j, k] = active_cell(i, j, k, grid)
+    @inbounds active_cells_field[i, j, k] = active_cell(i, j, k, grid, ib)
 end
 
-function compute_interior_active_cells(ibg; parameters = :xyz)
-    active_cells_field = Field{Center, Center, Center}(ibg, Bool)
+function compute_interior_active_cells(grid, ib; parameters = :xyz)
+    active_cells_field = Field{Center, Center, Center}(grid, Bool)
     fill!(active_cells_field, false)
-    launch!(architecture(ibg), ibg, parameters, _set_active_indices!, active_cells_field, ibg)
+    launch!(architecture(grid), grid, parameters, _set_active_indices!, active_cells_field, grid, ib)
     return active_cells_field
 end
 
-function compute_active_z_columns(ibg)
-    one_field = OneField(Int)
-    condition = NotImmersed(truefunc)
-    mask = 0
+@kernel function _set_active_columns!(active_z_columns, grid, ib)
+    i, j = @index(Global, NTuple)
+    active_column = false
+    for k in 1:size(grid, 3)
+        active_column = active_column | active_cell(i, j, k, grid, ib)
+    end
+    @inbounds active_z_columns[i, j, 1] = active_column
+end
 
-    # Compute all the active cells in a z-column using a ConditionalOperation
-    conditional_active_cells = ConditionalOperation{Center, Center, Center}(one_field, identity, ibg, condition, mask)
-    active_cells_in_column   = sum(conditional_active_cells, dims = 3)
+function compute_active_z_columns(grid, ib)
+    active_z_columns = Field{Center, Center, Nothing}(grid, Bool)
+    fill!(active_z_columns, false)
 
-    # Check whether the column ``i, j`` is immersed, which would correspond to `active_cells_in_column[i, j, 1] == 0`
-    is_immersed_column = KernelFunctionOperation{Center, Center, Nothing}(active_column, ibg, active_cells_in_column)
-    active_z_columns = Field{Center, Center, Nothing}(ibg, Bool)
-    set!(active_z_columns, is_immersed_column)
+    # Compute the active cells in the column
+    launch!(architecture(grid), grid, :xy, _set_active_columns!, active_z_columns, grid, ib)
 
     return active_z_columns
 end
@@ -113,22 +95,23 @@ const MAXUInt16 = 2^16 - 1
 const MAXUInt32 = 2^32 - 1
 
 """
-    interior_active_indices(ibg; parameters = :xyz)
+    interior_active_indices(grid, ib; parameters = :xyz)
 
 Compute the indices of the active interior cells in the given immersed boundary grid within the indices
 specified by the `parameters` keyword argument
 
 # Arguments
-- `ibg`: The immersed boundary grid.
+- `grid`: The underlying grid.
+- `ib`: The immersed boundary.
 - `parameters`: (optional) The parameters to be used for computing the active cells. Default is `:xyz`.
 
 # Returns
 An array of tuples representing the indices of the active interior cells.
 """
-function interior_active_indices(ibg; parameters = :xyz)
-    active_cells_field = compute_interior_active_cells(ibg; parameters)
+function interior_active_indices(grid, ib; parameters = :xyz)
+    active_cells_field = compute_interior_active_cells(grid, ib; parameters)
 
-    N = maximum(size(ibg))
+    N = maximum(size(grid))
     IntType = N > MAXUInt8 ? (N > MAXUInt16 ? (N > MAXUInt32 ? UInt64 : UInt32) : UInt16) : UInt8
 
     IndicesType = Tuple{IntType, IntType, IntType}
@@ -137,18 +120,18 @@ function interior_active_indices(ibg; parameters = :xyz)
     # For this reason, we split the computation in vertical levels and `findall` the active indices in 
     # subsequent xy planes, then stitch them back together
     active_indices = IndicesType[]
-    active_indices = findall_active_indices!(active_indices, active_cells_field, ibg, IndicesType)
-    active_indices = on_architecture(architecture(ibg), active_indices)
+    active_indices = findall_active_indices!(active_indices, active_cells_field, grid, IndicesType)
+    active_indices = on_architecture(architecture(grid), active_indices)
 
     return active_indices
 end
 
 # Cannot `findall` on very large grids, so we split the computation in levels.
 # This makes the computation a little heavier but avoids OOM errors (this computation
 # is performed only once on setup)
-function findall_active_indices!(active_indices, active_cells_field, ibg, IndicesType)
+function findall_active_indices!(active_indices, active_cells_field, grid, IndicesType)
 
-    for k in 1:size(ibg, 3)
+    for k in 1:size(grid, 3)
         interior_indices = findall(on_architecture(CPU(), interior(active_cells_field, :, :, k:k)))
         interior_indices = convert_interior_indices(interior_indices, k, IndicesType)
         active_indices = vcat(active_indices, interior_indices)
@@ -169,22 +152,22 @@ end
 # In case of a serial grid, the interior computations are performed over the whole three-dimensional
 # domain. Therefore, the `interior_active_cells` field contains the indices of all the active cells in 
 # the range 1:Nx, 1:Ny and 1:Nz (i.e., we construct the map with parameters :xyz)
-map_interior_active_cells(ibg) = interior_active_indices(ibg; parameters = :xyz)
+map_interior_active_cells(grid, ib) = interior_active_indices(grid, ib; parameters = :xyz)
 
 # If we eventually want to perform also barotropic step, `w` computation and `p` 
 # computation only on active `columns`
-function map_active_z_columns(ibg)
-    active_cells_field = compute_active_z_columns(ibg)
+function map_active_z_columns(grid, ib)
+    active_cells_field = compute_active_z_columns(grid, ib)
     interior_cells     = on_architecture(CPU(), interior(active_cells_field, :, :, 1))
 
     full_indices = findall(interior_cells)
 
-    Nx, Ny, _ = size(ibg)
+    Nx, Ny, _ = size(grid)
     # Reduce the size of the active_cells_map (originally a tuple of Int64)
     N = max(Nx, Ny)
     IntType = N > MAXUInt8 ? (N > MAXUInt16 ? (N > MAXUInt32 ? UInt64 : UInt32) : UInt16) : UInt8
     surface_map = getproperty.(full_indices, Ref(:I)) .|> Tuple{IntType, IntType}
-    surface_map = on_architecture(architecture(ibg), surface_map)
+    surface_map = on_architecture(architecture(grid), surface_map)
 
     return surface_map
 end

src/ImmersedBoundaries/grid_fitted_bottom.jl

@@ -89,21 +89,20 @@ end
 Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.bottom_height), adapt(to, ib.immersed_condition))
 
 """
-    ImmersedBoundaryGrid(grid, ib::GridFittedBottom)
+    materialize_immersed_boundary(grid, ib)
 
-Return a grid with `GridFittedBottom` immersed boundary (`ib`).
-
-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.
+Returns a new `ib` wrapped around a Field that holds the numerical `immersed_boundary`. 
+If `ib` is an `AbstractGridFittedBottom`, `ib.bottom_height` is the z-coordinate of 
+top-most interface of the last ``immersed`` cell in the column. If `ib` is a `GridFittedBoundary`,
+`ib.mask` is a field of booleans that indicates whether a cell is immersed or not.
 """
-function ImmersedBoundaryGrid(grid, ib::GridFittedBottom)
+function materialize_immersed_boundary(grid, ib::GridFittedBottom)
     bottom_field = Field{Center, Center, Nothing}(grid)
     set!(bottom_field, ib.bottom_height)
     @apply_regionally compute_numerical_bottom_height!(bottom_field, grid, ib)
     fill_halo_regions!(bottom_field)
     new_ib = GridFittedBottom(bottom_field)
-    TX, TY, TZ = topology(grid)
-    return ImmersedBoundaryGrid{TX, TY, TZ}(grid, new_ib)
+    return new_ib
 end
 
 compute_numerical_bottom_height!(bottom_field, grid, ib) = 

src/ImmersedBoundaries/grid_fitted_boundary.jl

@@ -24,18 +24,9 @@ function compute_mask(grid, ib)
     return mask_field
 end
 
-function ImmersedBoundaryGrid(grid, ib::GridFittedBoundary; precompute_mask=true)
-    TX, TY, TZ = topology(grid)
-
-    # TODO: validate ib
-
-    if precompute_mask
-        mask_field = compute_mask(grid, ib)
-        new_ib = GridFittedBoundary(mask_field)
-        return ImmersedBoundaryGrid{TX, TY, TZ}(grid, new_ib)
-    else
-        return ImmersedBoundaryGrid{TX, TY, TZ}(grid, ib)
-    end
+function materialize_immersed_boundary(grid, ib::GridFittedBoundary)
+    mask_field = compute_mask(grid, ib)
+    return GridFittedBoundary(mask_field)
 end
 
 on_architecture(arch, ib::GridFittedBoundary{<:Field}) = GridFittedBoundary(compute_mask(on_architecture(arch, ib.mask.grid), ib))

src/ImmersedBoundaries/immersed_boundary_grid.jl

@@ -14,22 +14,41 @@ struct ImmersedBoundaryGrid{FT, TX, TY, TZ, G, I, M, S, Arch} <: AbstractGrid{FT
 end
 
 # Internal interface
-function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I, mi::M,
-                                          ms::S) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I, M, S}
+function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I, mi::M, ms::S) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I, M, S}
     FT = eltype(grid)
     arch = architecture(grid)
     Arch = typeof(arch)
     return ImmersedBoundaryGrid{FT, TX, TY, TZ, G, I, M, S, Arch}(arch, grid, ib, mi, ms)
 end
 
-# Constructor with no active map
-function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I}
-    FT = eltype(grid)
-    arch = architecture(grid)
-    Arch = typeof(arch)
-    return ImmersedBoundaryGrid{FT, TX, TY, TZ, G, I, Nothing, Nothing, Arch}(arch, grid, ib, nothing, nothing)
+"""
+    ImmersedBoundaryGrid(grid, ib::AbstractImmersedBoundary; active_cells_map::Bool=true)
+
+Return a grid with an `AbstractImmersedBoundary` immersed boundary (`ib`). If `active_cells_map` is `true`,
+the grid will also populate an `interior_active_cells` and `active_z_columns` fields that are a list of active indices in the 
+interior and on a reduced x-y plane, respectively.
+"""
+function ImmersedBoundaryGrid(grid::AbstractUnderlyingGrid, ib::AbstractImmersedBoundary; active_cells_map::Bool=false) 
+    new_ib = materialize_immersed_boundary(grid, ib)
+    
+    # Create the cells map on the CPU, then switch it to the GPU
+    if active_cells_map 
+        interior_active_cells = map_interior_active_cells(grid, new_ib)
+        active_z_columns = map_active_z_columns(grid, new_ib)
+    else
+        interior_active_cells = nothing
+        active_z_columns = nothing
+    end
+    
+    TX, TY, TZ = topology(grid)
+    
+    return ImmersedBoundaryGrid{TX, TY, TZ}(grid, 
+                                            new_ib, 
+                                            interior_active_cells,
+                                            active_z_columns)
 end
 
+
 const IBG = ImmersedBoundaryGrid
 
 @inline Base.getproperty(ibg::IBG, property::Symbol) = get_ibg_property(ibg, Val(property))

src/ImmersedBoundaries/partial_cell_bottom.jl

@@ -62,14 +62,12 @@ function PartialCellBottom(bottom_height; minimum_fractional_cell_height=0.2)
     return PartialCellBottom(bottom_height, minimum_fractional_cell_height)
 end
 
-function ImmersedBoundaryGrid(grid, ib::PartialCellBottom)
+function materialize_immersed_boundary(grid, ib::PartialCellBottom)
     bottom_field = Field{Center, Center, Nothing}(grid)
     set!(bottom_field, ib.bottom_height)
     @apply_regionally compute_numerical_bottom_height!(bottom_field, grid, ib)
     fill_halo_regions!(bottom_field)
-    new_ib = PartialCellBottom(bottom_field, ib.minimum_fractional_cell_height)
-    TX, TY, TZ = topology(grid)
-    return ImmersedBoundaryGrid{TX, TY, TZ}(grid, new_ib)
+    return PartialCellBottom(bottom_field, ib.minimum_fractional_cell_height)
 end
 
 @kernel function _compute_numerical_bottom_height!(bottom_field, grid, ib::PartialCellBottom)

test/test_active_cells_map.jl

@@ -70,7 +70,7 @@ Nz = 10
         end
 
         bottom_height = on_architecture(arch, bottom_height)
-        immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
+        immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height); active_cells_map = false)
         immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height); active_cells_map = true)
 
         @testset "Active cells map construction" begin