Integrate cubed sphere functionality with halo filling and core infrastructure updates (#4538)

* Make necessary modifications

* Update multi_region_split_explicit_free_surface.jl

* Fix bug

* Update src/Utils/multi_region_transformation.jl

Co-authored-by: Navid C. Constantinou <[email protected]>

* Update src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl

Co-authored-by: Simone Silvestri <[email protected]>

* Fix bug

* Compact code

* Specialize difference operators for CCSPanelGrid

* Fix bug

* Extend tupled_fill_halo_regions!

* Fix bug

* Remove replace_horizontal_vector_halos!

* Add explanatory comments and fix bug

* Update src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl

Co-authored-by: Simone Silvestri <[email protected]>

* Update src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl

Co-authored-by: Simone Silvestri <[email protected]>

* Update src/MultiRegion/unified_implicit_free_surface_solver.jl

Co-authored-by: Simone Silvestri <[email protected]>

* Fix bug

* Remove auxiliary_actions!

* Fix bugs

* Update src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

* comment the todo

* FieldBoundaryBuffers --> CommunicationBuffers

* Remove unused parameter in fill_halo_size function

* Fix bug

* Remove hardcoded situation-specific logic

* Incorporate modifications from PR #4566

* Incorporate modifications from PR #4568

* Fix typo

* Enhance code

* Reinstate unused parameter for consistency

* Fix out-of-bounds memory access bug

* Fix typo

* Fix bug

* Remove commented-out line

* Fix typo

* Fix bug

* correct getregion for an immersed boundary grid

* fix show for a multiregion grid

* active cells map for CSG

* add a small test for IBG construction on CSG

* including Immersed boundary grids in the PCG solver

* Fix out-of-bounds memory allocation error

* Update test_multi_region_cubed_sphere.jl

* Use ConformalCubedSphereGridOfSomeKind

* Fix out-of-bounds error

* Update test_multi_region_cubed_sphere.jl

* Fix out-of-bounds error

* Update test_multi_region_cubed_sphere.jl

* Leverage reduced_dimensions for kernel launches

* add a function inside apply regionally

* Enhance code

* Compact code

* switch the apply regionally inside the linear operation

* revert

* Update src/MultiRegion/multi_region_split_explicit_free_surface.jl

* Update src/MultiRegion/multi_region_boundary_conditions.jl

* Update src/MultiRegion/multi_region_boundary_conditions.jl

* Fix GPU error

* Define ConformalCubedSpherePanelGridOfSomeKind

* Fix GPU kernel type instability bug

---------

Co-authored-by: Navid C. Constantinou <[email protected]>
Co-authored-by: Simone Silvestri <[email protected]>

Author: 3 people

src/BoundaryConditions/BoundaryConditions.jl

@@ -20,9 +20,6 @@ using Oceananigans.Grids
 
 import Adapt: adapt_structure
 
-# Switch around halos for cubed sphere by exchanging buffer informations
-replace_horizontal_vector_halos!(velocities, grid::AbstractGrid; signed=true) = nothing
-
 include("boundary_condition_classifications.jl")
 include("boundary_condition.jl")
 include("discrete_boundary_function.jl")

src/Fields/field_tuples.jl

@@ -53,7 +53,10 @@ Fill halo regions for all `fields`. The algorithm:
   4. In every direction, the halo regions in each of the remaining `Field` tuple
      are filled simultaneously.
 """
-function fill_halo_regions!(maybe_nested_tuple::Union{NamedTuple, Tuple}, args...; kwargs...)
+function fill_halo_regions!(maybe_nested_tuple::Union{NamedTuple, Tuple}, args...; 
+                            signed = true,  # This kwarg is active only for a `ConformalCubedSphereGrid`, here we discard it.
+                            kwargs...)
+
     flattened = flattened_unique_values(maybe_nested_tuple)
 
     # Look for grid within the flattened field tuple:

src/ImmersedBoundaries/immersed_boundary_grid.jl

@@ -78,7 +78,7 @@ end
 const IBG = ImmersedBoundaryGrid
 
 @inline Base.getproperty(ibg::IBG, property::Symbol) = get_ibg_property(ibg, Val(property))
-@inline get_ibg_property(ibg::IBG, ::Val{property}) where property = getfield(getfield(ibg, :underlying_grid), property)
+@inline get_ibg_property(ibg::IBG, ::Val{property}) where property = getproperty(getfield(ibg, :underlying_grid), property)
 @inline get_ibg_property(ibg::IBG, ::Val{:immersed_boundary})      = getfield(ibg, :immersed_boundary)
 @inline get_ibg_property(ibg::IBG, ::Val{:underlying_grid})        = getfield(ibg, :underlying_grid)
 @inline get_ibg_property(ibg::IBG, ::Val{:interior_active_cells})  = getfield(ibg, :interior_active_cells)

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl

@@ -20,6 +20,7 @@ function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velociti
                                                grid, u, v)
 
     fill_halo_regions!((barotropic_velocities.U, barotropic_velocities.V))
+    fill_halo_regions!(sefs.η)
 
     return nothing
 end

src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl

@@ -5,7 +5,7 @@ using Oceananigans.Architectures
 using Oceananigans.Grids: with_halo, isrectilinear, halo_size
 using Oceananigans.Architectures: device
 
-import Oceananigans.Solvers: solve!, precondition!
+import Oceananigans.Solvers: solve!, precondition!, perform_linear_operation!
 import Oceananigans.Architectures: architecture
 
 """
@@ -50,7 +50,10 @@ function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, settings, gravitationa
     vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ)
 
     @apply_regionally compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas)
-    fill_halo_regions!(vertically_integrated_lateral_areas)
+
+    Ax = vertically_integrated_lateral_areas.xᶠᶜᶜ
+    Ay = vertically_integrated_lateral_areas.yᶜᶠᶜ
+    fill_halo_regions!((Ax, Ay); signed=false)
 
     # Set some defaults
     settings = Dict{Symbol, Any}(settings)
@@ -66,8 +69,8 @@ function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, settings, gravitationa
     right_hand_side = ZFaceField(grid, indices = (:, :, size(grid, 3) + 1))
 
     solver = ConjugateGradientSolver(implicit_free_surface_linear_operation!;
-                                                   template_field = right_hand_side,
-                                                   settings...)
+                                     template_field = right_hand_side,
+                                     settings...)
 
     return PCGImplicitFreeSurfaceSolver(vertically_integrated_lateral_areas, solver, right_hand_side)
 end
@@ -134,16 +137,19 @@ function implicit_free_surface_linear_operation!(L_ηⁿ⁺¹, ηⁿ⁺¹, ∫
     grid = L_ηⁿ⁺¹.grid
     arch = architecture(L_ηⁿ⁺¹)
 
-    # Note: because of `fill_halo_regions!` below, we cannot use `PCGImplicitFreeSurface` on a
-    # multi-region grid; `fill_halo_regions!` cannot be used within `@apply_regionally`
-    fill_halo_regions!(ηⁿ⁺¹)
-
     launch!(arch, grid, :xy, _implicit_free_surface_linear_operation!,
             L_ηⁿ⁺¹, grid,  ηⁿ⁺¹, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
 
     return nothing
 end
 
+ImplicitFreeSurfaceOperation = typeof(implicit_free_surface_linear_operation!)
+
+@inline function perform_linear_operation!(linear_operation!::ImplicitFreeSurfaceOperation, q, p, args...)
+    fill_halo_regions!(p)
+    @apply_regionally linear_operation!(q, p, args...)
+end
+
 # Kernels that act on vertically integrated / surface quantities
 @inline ∫ᶻ_Ax_∂x_ηᶠᶜᶜ(i, j, k, grid, ∫ᶻ_Axᶠᶜᶜ, η) = @inbounds ∫ᶻ_Axᶠᶜᶜ[i, j, k] * ∂xᶠᶜᶠ(i, j, k, grid, η)
 @inline ∫ᶻ_Ay_∂y_ηᶜᶠᶜ(i, j, k, grid, ∫ᶻ_Ayᶜᶠᶜ, η) = @inbounds ∫ᶻ_Ayᶜᶠᶜ[i, j, k] * ∂yᶜᶠᶠ(i, j, k, grid, η)

src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

@@ -63,11 +63,8 @@ function hydrostatic_velocity_fields(velocities::PrescribedVelocityFields, grid,
     v = wrap_prescribed_field(Center, Face, Center, velocities.v, grid; clock, parameters)
     w = wrap_prescribed_field(Center, Center, Face, velocities.w, grid; clock, parameters)
 
-    fill_halo_regions!(u)
-    fill_halo_regions!(v)
+    fill_halo_regions!((u, v))
     fill_halo_regions!(w)
-    prescribed_velocities = (; u, v, w)
-    @apply_regionally replace_horizontal_vector_halos!(prescribed_velocities, grid)
 
     return PrescribedVelocityFields(u, v, w, parameters)
 end

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -3,7 +3,7 @@ using Oceananigans.BoundaryConditions
 
 using Oceananigans: UpdateStateCallsite
 using Oceananigans.Biogeochemistry: update_biogeochemical_state!
-using Oceananigans.BoundaryConditions: update_boundary_conditions!, replace_horizontal_vector_halos!
+using Oceananigans.BoundaryConditions: update_boundary_conditions!
 using Oceananigans.TurbulenceClosures: compute_diffusivities!
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node
 using Oceananigans.Models: update_model_field_time_series!
@@ -39,9 +39,8 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp
     # Update the boundary conditions
     @apply_regionally update_boundary_conditions!(fields(model), model)
 
-    tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model), async=true)
+    tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model); async=true)
 
-    @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid)
     @apply_regionally compute_auxiliaries!(model)
 
     fill_halo_regions!(model.diffusivity_fields; only_local_halos = true)