Integrate cubed sphere functionality with halo filling and core infrastructure updates

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/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, sefs.η)
 
     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!, auxiliary_actions!
 import Oceananigans.Architectures: architecture
 
 """
@@ -50,7 +50,12 @@ 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)
+
+    u = vertically_integrated_lateral_areas.xᶠᶜᶜ
+    v = vertically_integrated_lateral_areas.yᶜᶠᶜ
+
+    grid isa ConformalCubedSphereGrid ? fill_halo_regions!((u, v); signed=false) :
+                                        fill_halo_regions!(vertically_integrated_lateral_areas)
 
     # Set some defaults
     settings = Dict{Symbol, Any}(settings)
@@ -134,16 +139,16 @@ 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!)
+
+auxiliary_actions!(::ImplicitFreeSurfaceOperation, L_ηⁿ⁺¹, ηⁿ⁺¹, args...) = fill_halo_regions!(ηⁿ⁺¹)
+
 # 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,11 +3,12 @@ using Oceananigans.BoundaryConditions
 
 using Oceananigans: UpdateStateCallsite
 using Oceananigans.Biogeochemistry: update_biogeochemical_state!
-using Oceananigans.BoundaryConditions: update_boundary_condition!, replace_horizontal_vector_halos!
+using Oceananigans.BoundaryConditions: update_boundary_condition!
 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!
 using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!, p_kernel_parameters
+using Oceananigans.MultiRegion: ConformalCubedSphereGrid
 using Oceananigans.Fields: tupled_fill_halo_regions!
 
 import Oceananigans.Models.NonhydrostaticModels: compute_auxiliaries!
@@ -39,9 +40,20 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp
     # Update the boundary conditions
     @apply_regionally update_boundary_condition!(fields(model), model)
 
-    tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model), async=true)
+    if grid isa ConformalCubedSphereGrid
+        if model.tracers != nothing
+            for tracer in model.tracers
+                fill_halo_regions!(tracer)
+            end
+        end
+        if model.free_surface != nothing
+            fill_halo_regions!(model.free_surface.η)
+        end
+        fill_halo_regions!((model.velocities.u, model.velocities.v))
+    else
+        tupled_fill_halo_regions!(prognostic_fields(model), model.clock, fields(model); async = true)
+    end
 
-    @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)