(0.97.6) Move the computation of flux BC before the time-stepping (#4696)

* go ahead

* ad project

* fix tests

* start with compute

* more changes

* more changes

* more fixes

* Update cubed_sphere_exchange_bcs.jl

* Update compute_hydrostatic_flux_bcs.jl

* Update hydrostatic_free_surface_tendency_kernel_functions.jl

* Update explicit_free_surface.jl

---------

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

Author: simone-silvestri

Project.toml

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

src/BoundaryConditions/BoundaryConditions.jl

@@ -7,7 +7,7 @@ export
     FluxBoundaryCondition, ValueBoundaryCondition, GradientBoundaryCondition, DistributedCommunicationBoundaryCondition,
     validate_boundary_condition_topology, validate_boundary_condition_architecture,
     FieldBoundaryConditions,
-    apply_x_bcs!, apply_y_bcs!, apply_z_bcs!,
+    compute_x_bcs!, compute_y_bcs!, compute_z_bcs!,
     fill_halo_regions!
 
 using Adapt
@@ -35,7 +35,7 @@ include("fill_halo_regions_flux.jl")
 include("fill_halo_regions_nothing.jl")
 include("fill_halo_regions_zipper.jl")
 
-include("apply_flux_bcs.jl")
+include("compute_flux_bcs.jl")
 
 include("update_boundary_conditions.jl")
 include("polar_boundary_condition.jl")

src/BoundaryConditions/apply_flux_bcs.jl renamed to src/BoundaryConditions/compute_flux_bcs.jl

@@ -9,95 +9,95 @@ using Oceananigans.Grids: AbstractGrid
 #####
 
 # Unpack
-apply_x_bcs!(Gc, c, args...) = apply_x_bcs!(Gc, Gc.grid, c, c.boundary_conditions.west,   c.boundary_conditions.east, args...)
-apply_y_bcs!(Gc, c, args...) = apply_y_bcs!(Gc, Gc.grid, c, c.boundary_conditions.south,  c.boundary_conditions.north, args...)
-apply_z_bcs!(Gc, c, args...) = apply_z_bcs!(Gc, Gc.grid, c, c.boundary_conditions.bottom, c.boundary_conditions.top, args...)
+compute_x_bcs!(Gc, c, args...) = compute_x_bcs!(Gc, Gc.grid, c, c.boundary_conditions.west,   c.boundary_conditions.east, args...)
+compute_y_bcs!(Gc, c, args...) = compute_y_bcs!(Gc, Gc.grid, c, c.boundary_conditions.south,  c.boundary_conditions.north, args...)
+compute_z_bcs!(Gc, c, args...) = compute_z_bcs!(Gc, Gc.grid, c, c.boundary_conditions.bottom, c.boundary_conditions.top, args...)
 
 # Shortcuts for...
 #
 # Nothing tendencies.
-apply_x_bcs!(::Nothing, args...) = nothing
-apply_y_bcs!(::Nothing, args...) = nothing
-apply_z_bcs!(::Nothing, args...) = nothing
+compute_x_bcs!(::Nothing, args...) = nothing
+compute_y_bcs!(::Nothing, args...) = nothing
+compute_z_bcs!(::Nothing, args...) = nothing
 
 # Not-flux boundary conditions
 const NotFluxBC = Union{PBC, MCBC, DCBC, VBC, GBC, OBC, ZFBC, Nothing}
 
-apply_x_bcs!(Gc, ::AbstractGrid, c, ::NotFluxBC, ::NotFluxBC, ::AbstractArchitecture, args...) = nothing
-apply_y_bcs!(Gc, ::AbstractGrid, c, ::NotFluxBC, ::NotFluxBC, ::AbstractArchitecture, args...) = nothing
-apply_z_bcs!(Gc, ::AbstractGrid, c, ::NotFluxBC, ::NotFluxBC, ::AbstractArchitecture, args...) = nothing
+compute_x_bcs!(Gc, ::AbstractGrid, c, ::NotFluxBC, ::NotFluxBC, ::AbstractArchitecture, args...) = nothing
+compute_y_bcs!(Gc, ::AbstractGrid, c, ::NotFluxBC, ::NotFluxBC, ::AbstractArchitecture, args...) = nothing
+compute_z_bcs!(Gc, ::AbstractGrid, c, ::NotFluxBC, ::NotFluxBC, ::AbstractArchitecture, args...) = nothing
 
 # The real deal
 """
 Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
 the source term `Gc` at the left and right.
 """
-apply_x_bcs!(Gc, grid::AbstractGrid, c, west_bc, east_bc, arch::AbstractArchitecture, args...) =
-    launch!(arch, grid, :yz, _apply_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, Tuple(args))
+compute_x_bcs!(Gc, grid::AbstractGrid, c, west_bc, east_bc, arch::AbstractArchitecture, args...) =
+    launch!(arch, grid, :yz, _compute_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, Tuple(args))
 
 """
 Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
 the source term `Gc` at the left and right.
 """
-apply_y_bcs!(Gc, grid::AbstractGrid, c, south_bc, north_bc, arch::AbstractArchitecture, args...) =
-    launch!(arch, grid, :xz, _apply_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, Tuple(args))
+compute_y_bcs!(Gc, grid::AbstractGrid, c, south_bc, north_bc, arch::AbstractArchitecture, args...) =
+    launch!(arch, grid, :xz, _compute_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, Tuple(args))
 
 """
 Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
 the source term `Gc` at the top and bottom.
 """
-apply_z_bcs!(Gc, grid::AbstractGrid, c, bottom_bc, top_bc, arch::AbstractArchitecture, args...) =
-    launch!(arch, grid, :xy, _apply_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, Tuple(args))
+compute_z_bcs!(Gc, grid::AbstractGrid, c, bottom_bc, top_bc, arch::AbstractArchitecture, args...) =
+    launch!(arch, grid, :xy, _compute_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, Tuple(args))
 
 """
-    _apply_x_bcs!(Gc, grid, west_bc, east_bc, args...)
+    _compute_x_bcs!(Gc, grid, west_bc, east_bc, args...)
 
 Apply a west and/or east boundary condition to variable `c`.
 """
-@kernel function _apply_x_bcs!(Gc, loc, grid, west_bc, east_bc, args)
+@kernel function _compute_x_bcs!(Gc, loc, grid, west_bc, east_bc, args)
     j, k = @index(Global, NTuple)
-    apply_x_west_bc!(Gc, loc, west_bc, j, k, grid, args...)
-    apply_x_east_bc!(Gc, loc, east_bc, j, k, grid, args...)
+    compute_x_west_bc!(Gc, loc, west_bc, j, k, grid, args...)
+    compute_x_east_bc!(Gc, loc, east_bc, j, k, grid, args...)
 end
 
 """
-    _apply_y_bcs!(Gc, grid, south_bc, north_bc, args...)
+    _compute_y_bcs!(Gc, grid, south_bc, north_bc, args...)
 
 Apply a south and/or north boundary condition to variable `c`.
 """
-@kernel function _apply_y_bcs!(Gc, loc, grid, south_bc, north_bc, args)
+@kernel function _compute_y_bcs!(Gc, loc, grid, south_bc, north_bc, args)
     i, k = @index(Global, NTuple)
-    apply_y_south_bc!(Gc, loc, south_bc, i, k, grid, args...)
-    apply_y_north_bc!(Gc, loc, north_bc, i, k, grid, args...)
+    compute_y_south_bc!(Gc, loc, south_bc, i, k, grid, args...)
+    compute_y_north_bc!(Gc, loc, north_bc, i, k, grid, args...)
 end
 
 """
-    _apply_z_bcs!(Gc, grid, bottom_bc, top_bc, args...)
+    _compute_z_bcs!(Gc, grid, bottom_bc, top_bc, args...)
 
 Apply a top and/or bottom boundary condition to variable `c`.
 """
-@kernel function _apply_z_bcs!(Gc, loc, grid, bottom_bc, top_bc, args)
+@kernel function _compute_z_bcs!(Gc, loc, grid, bottom_bc, top_bc, args)
     i, j = @index(Global, NTuple)
-    apply_z_bottom_bc!(Gc, loc, bottom_bc, i, j, grid, args...)
-       apply_z_top_bc!(Gc, loc, top_bc,    i, j, grid, args...)
+    compute_z_bottom_bc!(Gc, loc, bottom_bc, i, j, grid, args...)
+       compute_z_top_bc!(Gc, loc, top_bc,    i, j, grid, args...)
 end
 
 # Shortcuts for zero flux or non-flux boundary conditions
-@inline   apply_x_east_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
-@inline   apply_x_west_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
-@inline  apply_y_north_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
-@inline  apply_y_south_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
-@inline    apply_z_top_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
-@inline apply_z_bottom_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
+@inline   compute_x_east_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
+@inline   compute_x_west_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
+@inline  compute_y_north_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
+@inline  compute_y_south_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
+@inline    compute_z_top_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
+@inline compute_z_bottom_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
 
 # shortcut for the zipper BC
-@inline apply_y_north_bc!(Gc, loc, ::ZBC, args...) = nothing
+@inline compute_y_north_bc!(Gc, loc, ::ZBC, args...) = nothing
 
 @inline flip(::Center) = Face()
 @inline flip(::Face) = Center()
 
 """
-    apply_x_west_bc!(Gc, loc, west_flux::BC{<:Flux}, j, k, grid, args...)
+    compute_x_west_bc!(Gc, loc, west_flux::BC{<:Flux}, j, k, grid, args...)
 
 Add the flux divergence associated with a west flux boundary condition on `c`.
 Note that because
@@ -111,26 +111,26 @@ If `west_bc.condition` is a function, the function must have the signature
 
 The same logic holds for south and bottom boundary conditions in `y`, and `z`, respectively.
 """
-@inline function apply_x_west_bc!(Gc, loc, west_flux::BC{<:Flux}, j, k, grid, args...)
+@inline function compute_x_west_bc!(Gc, loc, west_flux::BC{<:Flux}, j, k, grid, args...)
     LX, LY, LZ = loc
     @inbounds Gc[1, j, k] += getbc(west_flux, j, k, grid, args...) * Ax(1, j, k, grid, flip(LX), LY, LZ) / volume(1, j, k, grid, LX, LY, LZ)
     return nothing
 end
 
-@inline function apply_y_south_bc!(Gc, loc, south_flux::BC{<:Flux}, i, k, grid, args...)
+@inline function compute_y_south_bc!(Gc, loc, south_flux::BC{<:Flux}, i, k, grid, args...)
     LX, LY, LZ = loc
     @inbounds Gc[i, 1, k] += getbc(south_flux, i, k, grid, args...) * Ay(i, 1, k, grid, LX, flip(LY), LZ) / volume(i, 1, k, grid, LX, LY, LZ)
     return nothing
 end
 
-@inline function apply_z_bottom_bc!(Gc, loc, bottom_flux::BC{<:Flux}, i, j, grid, args...)
+@inline function compute_z_bottom_bc!(Gc, loc, bottom_flux::BC{<:Flux}, i, j, grid, args...)
     LX, LY, LZ = loc
     @inbounds Gc[i, j, 1] += getbc(bottom_flux, i, j, grid, args...) * Az(i, j, 1, grid, LX, LY, flip(LZ)) / volume(i, j, 1, grid, LX, LY, LZ)
     return nothing
 end
 
 """
-    apply_x_east_bc!(Gc, loc, east_flux::BC{<:Flux}, j, k, grid, args...)
+    compute_x_east_bc!(Gc, loc, east_flux::BC{<:Flux}, j, k, grid, args...)
 
 Add the part of flux divergence associated with a east boundary condition on `c`.
 Note that because
@@ -144,19 +144,19 @@ If `east_bc.condition` is a function, the function must have the signature
 
 The same logic holds for north and top boundary conditions in `y`, and `z`, respectively.
 """
-@inline function apply_x_east_bc!(Gc, loc, east_flux::BC{<:Flux}, j, k, grid, args...)
+@inline function compute_x_east_bc!(Gc, loc, east_flux::BC{<:Flux}, j, k, grid, args...)
     LX, LY, LZ = loc
     @inbounds Gc[grid.Nx, j, k] -= getbc(east_flux, j, k, grid, args...) * Ax(grid.Nx+1, j, k, grid, flip(LX), LY, LZ) / volume(grid.Nx, j, k, grid, LX, LY, LZ)
     return nothing
 end
 
-@inline function apply_y_north_bc!(Gc, loc, north_flux::BC{<:Flux}, i, k, grid, args...)
+@inline function compute_y_north_bc!(Gc, loc, north_flux::BC{<:Flux}, i, k, grid, args...)
     LX, LY, LZ = loc
     @inbounds Gc[i, grid.Ny, k] -= getbc(north_flux, i, k, grid, args...) * Ay(i, grid.Ny+1, k, grid, LX, flip(LY), LZ) / volume(i, grid.Ny, k, grid, LX, LY, LZ)
     return nothing
 end
 
-@inline function apply_z_top_bc!(Gc, loc, top_flux::BC{<:Flux}, i, j, grid, args...)
+@inline function compute_z_top_bc!(Gc, loc, top_flux::BC{<:Flux}, i, j, grid, args...)
     LX, LY, LZ = loc
     @inbounds Gc[i, j, grid.Nz] -= getbc(top_flux, i, j, grid, args...) * Az(i, j, grid.Nz+1, grid, LX, LY, flip(LZ)) / volume(i, j, grid.Nz, grid, LX, LY, LZ)
     return nothing

src/CubedSpheres/cubed_sphere_exchange_bcs.jl

@@ -1,27 +1,25 @@
 using Adapt
 
 using Oceananigans.Architectures: AbstractArchitecture
-using Oceananigans.Grids: AbstractGrid
 using Oceananigans.BoundaryConditions
 using Oceananigans.BoundaryConditions: AbstractBoundaryConditionClassification
+using Oceananigans.Grids: AbstractGrid
 
 import Base: show
-
-import Oceananigans.BoundaryConditions: bc_str
-
 import Oceananigans.BoundaryConditions:
-    apply_x_bcs!,
-    apply_y_bcs!,
-    apply_x_east_bc!,
-    apply_x_west_bc!,
-    apply_y_south_bc!,
-    apply_y_north_bc!
+    bc_str,
+    compute_x_bcs!,
+    compute_y_bcs!,
+    compute_x_east_bc!,
+    compute_x_west_bc!,
+    compute_y_south_bc!,
+    compute_y_north_bc!
 
 struct CubedSphereExchange <: AbstractBoundaryConditionClassification end
 
 const CubedSphereExchangeBC = BoundaryCondition{<:CubedSphereExchange}
 
-bc_str(::CubedSphereExchangeBC) ="CubedSphereExchange"
+bc_str(::CubedSphereExchangeBC) = "CubedSphereExchange"
 
 CubedSphereExchangeBoundaryCondition(val; kwargs...) = BoundaryCondition(CubedSphereExchange, val; kwargs...)
 
@@ -89,10 +87,10 @@ Adapt.adapt_structure(to, ::CubedSphereExchangeInformation) = nothing
 Adapt.adapt_structure(to, ::CubedSphereExchangeBC) = nothing
 
 # Don't "apply fluxes" across CubedSphere boundaries
-@inline apply_x_east_bc!(  Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
-@inline apply_x_west_bc!(  Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
-@inline apply_y_north_bc!( Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
-@inline apply_y_south_bc!( Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
+@inline compute_x_east_bc!(Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
+@inline compute_x_west_bc!(Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
+@inline compute_y_north_bc!(Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
+@inline compute_y_south_bc!(Gc, loc, ::CubedSphereExchangeBC, args...) = nothing
 
-apply_x_bcs!(Gc, ::AbstractGrid, c, ::CubedSphereExchangeBC, ::CubedSphereExchangeBC, ::AbstractArchitecture, args...) = nothing
-apply_y_bcs!(Gc, ::AbstractGrid, c, ::CubedSphereExchangeBC, ::CubedSphereExchangeBC, ::AbstractArchitecture, args...) = nothing
+compute_x_bcs!(Gc, ::AbstractGrid, c, ::CubedSphereExchangeBC, ::CubedSphereExchangeBC, ::AbstractArchitecture, args...) = nothing
+compute_y_bcs!(Gc, ::AbstractGrid, c, ::CubedSphereExchangeBC, ::CubedSphereExchangeBC, ::AbstractArchitecture, args...) = nothing

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -134,6 +134,7 @@ include("barotropic_pressure_correction.jl")
 include("hydrostatic_free_surface_tendency_kernel_functions.jl")
 include("compute_hydrostatic_free_surface_tendencies.jl")
 include("compute_hydrostatic_free_surface_buffers.jl")
+include("compute_hydrostatic_flux_bcs.jl")
 include("update_hydrostatic_free_surface_model_state.jl")
 include("hydrostatic_free_surface_ab2_step.jl")
 include("hydrostatic_free_surface_rk3_step.jl")

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_flux_bcs.jl

@@ -0,0 +1,37 @@
+import Oceananigans.TimeSteppers: compute_flux_bc_tendencies!
+
+#####
+##### Boundary contributions to hydrostatic free surface model
+#####
+
+function compute_flux_bcs!(Gcⁿ, c, arch, args)
+    compute_x_bcs!(Gcⁿ, c, arch, args...)
+    compute_y_bcs!(Gcⁿ, c, arch, args...)
+    compute_z_bcs!(Gcⁿ, c, arch, args...)
+    return nothing
+end
+
+""" Apply boundary conditions by adding flux divergences to the right-hand-side. """
+function compute_flux_bc_tendencies!(model::HydrostaticFreeSurfaceModel)
+
+    Gⁿ         = model.timestepper.Gⁿ
+    grid       = model.grid
+    arch       = architecture(grid)
+    velocities = model.velocities
+    tracers    = model.tracers
+    
+    args = (model.clock, fields(model), model.closure, model.buoyancy)
+
+    
+    # Velocity fields
+    for i in (:u, :v)
+        @apply_regionally compute_flux_bcs!(Gⁿ[i], velocities[i], arch, args)
+    end
+
+    # Tracer fields
+    for i in propertynames(tracers)
+        @apply_regionally compute_flux_bcs!(Gⁿ[i], tracers[i], arch, args)
+    end
+
+    return nothing
+end

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl

@@ -33,23 +33,11 @@ function compute_tendencies!(model::HydrostaticFreeSurfaceModel, callbacks)
     compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_parameters; active_cells_map)
     complete_communication_and_compute_buffer!(model, grid, arch)
 
-    # Calculate contributions to momentum and tracer tendencies from user-prescribed fluxes across the
-    # boundaries of the domain
-    compute_hydrostatic_boundary_tendency_contributions!(model.timestepper.Gⁿ,
-                                                         model.architecture,
-                                                         model.velocities,
-                                                         model.tracers,
-                                                         model.clock,
-                                                         fields(model),
-                                                         model.closure,
-                                                         model.buoyancy)
-
     for callback in callbacks
         callback.callsite isa TendencyCallsite && callback(model)
     end
 
     update_tendencies!(model.biogeochemistry, model)
-    multiply_by_grid_scaling!(model.timestepper.Gⁿ, model.tracers, model.grid)
 
     return nothing
 end
@@ -106,17 +94,6 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
     return nothing
 end
 
-#####
-##### Boundary condributions to hydrostatic free surface model
-#####
-
-function apply_flux_bcs!(Gcⁿ, c, arch, args)
-    apply_x_bcs!(Gcⁿ, c, arch, args...)
-    apply_y_bcs!(Gcⁿ, c, arch, args...)
-    apply_z_bcs!(Gcⁿ, c, arch, args...)
-    return nothing
-end
-
 """ Calculate momentum tendencies if momentum is not prescribed."""
 function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_parameters; active_cells_map=nothing)
 
@@ -157,25 +134,6 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para
     return nothing
 end
 
-
-""" Apply boundary conditions by adding flux divergences to the right-hand-side. """
-function compute_hydrostatic_boundary_tendency_contributions!(Gⁿ, arch, velocities, tracers, args...)
-
-    args = Tuple(args)
-
-    # Velocity fields
-    for i in (:u, :v)
-        apply_flux_bcs!(Gⁿ[i], velocities[i], arch, args)
-    end
-
-    # Tracer fields
-    for i in propertynames(tracers)
-        apply_flux_bcs!(Gⁿ[i], tracers[i], arch, args)
-    end
-
-    return nothing
-end
-
 #####
 ##### Tendency calculators for u, v
 #####