Calculate tendencies on boundaries

src/TimeSteppers/TimeSteppers.jl

@@ -58,6 +58,7 @@ boundary_condition_function_arguments(model) =
      datatuple(model.tracers), model.parameters)
 
 include("generic_time_stepping.jl")
+include("velocity_and_tracer_tendencies.jl")
 include("time_stepping_kernels.jl")
 include("adams_bashforth.jl")
 

src/TimeSteppers/generic_time_stepping.jl

@@ -44,16 +44,31 @@ contribution from non-hydrostatic pressure.
 """
 function calculate_tendencies!(tendencies, velocities, tracers, pressures, diffusivities, model)
 
-    calculate_interior_source_terms!(
-        tendencies, model.architecture, model.grid, model.coriolis, model.buoyancy,
-        model.surface_waves, model.closure, velocities, tracers, pressures.pHY′,
-        diffusivities, model.forcing, model.parameters, model.clock.time
-    )
-
-    calculate_boundary_source_terms!(
+    # Note:
+    #
+    # "tendencies" is a NamedTuple of OffsetArrays corresponding to the tendency data for use
+    # in GPU computations.
+    #
+    # "model.timestepper.Gⁿ" is a NamedTuple of Fields, whose data also corresponds to 
+    # tendency data.
+
+    # Arguments needed to calculate tendencies for momentum and tracers
+    tendency_calculation_args = (tendencies, model.architecture, model.grid, model.coriolis, model.buoyancy,
+                                 model.surface_waves, model.closure, velocities, tracers, pressures.pHY′,
+                                 diffusivities, model.forcing, model.parameters, model.clock.time)
+
+    # Calculate contributions to momentum and tracer tendencies from fluxes and volume terms in the
+    # interior of the domain
+    calculate_interior_tendency_contributions!(tendency_calculation_args...)
+
+    # Calculate contributions to momentum and tracer tendencies from user-prescribed fluxes across the 
+    # boundaries of the domain
+    calculate_boundary_tendency_contributions!(
         model.timestepper.Gⁿ, model.architecture, model.velocities,
-        model.tracers, boundary_condition_function_arguments(model)...
-    )
+        model.tracers, boundary_condition_function_arguments(model)...)
+
+    # Calculate momentum tendencies on boundaries in `Bounded` directions.
+    calculate_velocity_tendencies_on_boundaries!(tendency_calculation_args...)
 
     return nothing
 end

src/TimeSteppers/time_stepping_kernels.jl

@@ -3,7 +3,9 @@
 #####
 
 """ Store previous value of the source term and calculate current source term. """
-function calculate_interior_source_terms!(G, arch, grid, coriolis, buoyancy, surface_waves, closure, U, C, pHY′, K, F, parameters, time)
+function calculate_interior_tendency_contributions!(G, arch, grid, coriolis, buoyancy, surface_waves, closure, 
+                                                    U, C, pHY′, K, F, parameters, time)
+
     # Manually choose thread-block layout here as it's ~20% faster.
     # See: https://github.com/climate-machine/Oceananigans.jl/pull/308
     Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
@@ -40,59 +42,170 @@ function calculate_interior_source_terms!(G, arch, grid, coriolis, buoyancy, sur
     return nothing
 end
 
-""" Calculate the right-hand-side of the u-momentum equation. """
+"""
+    calculate_velocity_tendencies_on_boundaries!(tendency_calculation_args...)
+
+Calculate the velocity tendencies *on* east, north, and top boundaries, when the
+x-, y-, or z- directions have a `Bounded` topology.
+"""
+function calculate_velocity_tendencies_on_boundaries!(tendency_calculation_args...)
+
+    calculate_east_boundary_Gu!(tendency_calculation_args...)
+    calculate_north_boundary_Gv!(tendency_calculation_args...)
+    calculate_top_boundary_Gw!(tendency_calculation_args...)
+
+    return nothing
+end
+
+# Fallbacks for non-bounded topologies in x-, y-, and z-directions
+@inline calculate_east_boundary_Gu!(args...) = nothing
+@inline calculate_north_boundary_Gv!(args...) = nothing
+@inline calculate_top_boundary_Gw!(args...) = nothing
+
+"""
+    calculate_east_boundary_Gu!(G, arch, grid::AbstractGrid{FT, <:Bounded},
+                                coriolis, buoyancy, surface_waves, closure,
+                                U, C, pHY′, K, F, parameters, time) where FT
+
+Calculate `Gu` on east boundaries when the x-direction has `Bounded` topology.
+"""
+function calculate_east_boundary_Gu!(G, arch, grid::AbstractGrid{FT, <:Bounded},
+                                     coriolis, buoyancy, surface_waves, closure,
+                                     U, C, pHY′, K, F, parameters, time) where FT
+
+    @launch(device(arch), config=launch_config(grid, :yz), 
+            _calculate_east_boundary_Gu!(G.u, grid, coriolis, surface_waves, 
+                                         closure, U, C, K, F, pHY′, parameters, time))
+
+    return nothing
+end
+
+"""
+    calculate_north_boundary_Gu!(G, arch, grid::AbstractGrid{FT, <:Bounded},
+                                 coriolis, buoyancy, surface_waves, closure,
+                                 U, C, pHY′, K, F, parameters, time) where FT
+
+Calculate `Gv` on north boundaries when the y-direction has `Bounded` topology.
+"""
+function calculate_north_boundary_Gv!(G, arch, grid::AbstractGrid{FT, TX, <:Bounded},
+                                      coriolis, buoyancy, surface_waves, closure,
+                                      U, C, pHY′, K, F, parameters, time) where {FT, TX}
+
+    @launch(device(arch), config=launch_config(grid, :xz), 
+            _calculate_north_boundary_Gv!(G.v, grid, coriolis, surface_waves, 
+                                          closure, U, C, K, F, pHY′, parameters, time))
+
+    return nothing
+end
+
+"""
+    calculate_top_boundary_Gw!(G, arch, grid::AbstractGrid{FT, <:Bounded},
+                               coriolis, buoyancy, surface_waves, closure,
+                               U, C, pHY′, K, F, parameters, time) where FT
+
+Calculate `Gw` on top boundaries when the z-direction has `Bounded` topology.
+"""
+function calculate_top_boundary_Gw!(G, arch, grid::AbstractGrid{FT, TX, TY, <:Bounded},
+                                    coriolis, buoyancy, surface_waves, closure,
+                                    U, C, pHY′, K, F, parameters, time) where {FT, TX, TY}
+
+    @launch(device(arch), config=launch_config(grid, :xy), 
+            _calculate_top_boundary_Gw!(G.w, grid, coriolis, surface_waves, 
+                                        closure, U, C, K, F, parameters, time))
+
+    return nothing
+end
+
+#####
+##### Tendency calculators for u-velocity, aka x-velocity
+#####
+
+""" Calculate the right-hand-side of the u-velocity equation. """
 function calculate_Gu!(Gu, grid, coriolis, surface_waves, closure, U, C, K, F, pHY′, parameters, time)
     @loop_xyz i j k grid begin
-        @inbounds Gu[i, j, k] = ( - div_ũu(i, j, k, grid, U)
-                                  - x_f_cross_U(i, j, k, grid, coriolis, U)
-                                  - ∂xᶠᵃᵃ(i, j, k, grid, pHY′)
-                                  + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, U, K)
-                                  + x_curl_Uˢ_cross_U(i, j, k, grid, surface_waves, U, time)
-                                  + ∂t_uˢ(i, j, k, grid, surface_waves, time)
-                                  + F.u(i, j, k, grid, time, U, C, parameters))
+        @inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                                    closure, U, C, K, F, pHY′, parameters, time)
+    end
+    return nothing
+end
+
+""" Calculate the right-hand-side of the u-velocity equation on the east boundary. """
+function _calculate_east_boundary_Gu!(Gu, grid, coriolis, surface_waves,
+                                      closure, U, C, K, F, pHY′, parameters, time)
+    i = grid.Nx + 1
+    @loop_yz j k grid begin
+        @inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                                    closure, U, C, K, F, pHY′, parameters, time)
     end
     return nothing
 end
 
-""" Calculate the right-hand-side of the v-momentum equation. """
+#####
+##### Tendency calculators for v-velocity
+#####
+
+""" Calculate the right-hand-side of the v-velocity equation. """
 function calculate_Gv!(Gv, grid, coriolis, surface_waves, closure, U, C, K, F, pHY′, parameters, time)
     @loop_xyz i j k grid begin
-        @inbounds Gv[i, j, k] = ( - div_ũv(i, j, k, grid, U)
-                                  - y_f_cross_U(i, j, k, grid, coriolis, U)
-                                  - ∂yᵃᶠᵃ(i, j, k, grid, pHY′)
-                                  + ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure, U, K)
-                                  + y_curl_Uˢ_cross_U(i, j, k, grid, surface_waves, U, time)
-                                  + ∂t_vˢ(i, j, k, grid, surface_waves, time)
-                                  + F.v(i, j, k, grid, time, U, C, parameters))
+        @inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                                    closure, U, C, K, F, pHY′, parameters, time)
+    end
+    return nothing
+end
+
+""" Calculate the right-hand-side of the v-velocity equation on the north boundary. """
+function _calculate_north_boundary_Gv!(Gv, grid, coriolis, surface_waves,
+                                       closure, U, C, K, F, pHY′, parameters, time)
+    j = grid.Ny + 1
+    @loop_xz i k grid begin
+        @inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                                    closure, U, C, K, F, pHY′, parameters, time)
     end
     return nothing
 end
 
-""" Calculate the right-hand-side of the w-momentum equation. """
+#####
+##### Tendency calculators for w-velocity
+#####
+
+""" Calculate the right-hand-side of the w-velocity equation. """
 function calculate_Gw!(Gw, grid, coriolis, surface_waves, closure, U, C, K, F, parameters, time)
     @loop_xyz i j k grid begin
-        @inbounds Gw[i, j, k] = ( - div_ũw(i, j, k, grid, U)
-                                  - z_f_cross_U(i, j, k, grid, coriolis, U)
-                                  + ∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure, U, K)
-                                  + z_curl_Uˢ_cross_U(i, j, k, grid, surface_waves, U, time)
-                                  + ∂t_wˢ(i, j, k, grid, surface_waves, time)
-                                  + F.w(i, j, k, grid, time, U, C, parameters))
+        @inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                                    closure, U, C, K, F, parameters, time)
+    end
+    return nothing
+end
+
+""" Calculate the right-hand-side of the w-velocity equation. """
+function _calculate_top_boundary_Gw!(Gw, grid, coriolis, surface_waves, closure, U, C, K, F, parameters, time)
+    k = grid.Nz + 1
+    @loop_xy i j grid begin
+        @inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                                    closure, U, C, K, F, parameters, time)
     end
     return nothing
 end
 
+#####
+##### Tracer(s)
+#####
+
 """ Calculate the right-hand-side of the tracer advection-diffusion equation. """
 function calculate_Gc!(Gc, grid, c, tracer_index, closure, buoyancy, U, C, K, Fc, parameters, time)
     @loop_xyz i j k grid begin
-        @inbounds Gc[i, j, k] = (- div_uc(i, j, k, grid, U, c)
-                                 + ∇_κ_∇c(i, j, k, grid, closure, c, tracer_index, K, C, buoyancy)
-                                 + Fc(i, j, k, grid, time, U, C, parameters))
+        @inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, c, tracer_index,
+                                                closure, buoyancy, U, C, K, Fc, parameters, time)
     end
     return nothing
 end
 
+#####
+##### Boundary contributions to tendencies due to user-prescribed fluxes
+#####
+
 """ Apply boundary conditions by adding flux divergences to the right-hand-side. """
-function calculate_boundary_source_terms!(Gⁿ, arch, U, C, args...)
+function calculate_boundary_tendency_contributions!(Gⁿ, arch, U, C, args...)
 
     # Velocity fields
     for i in 1:3
@@ -109,6 +222,7 @@ function calculate_boundary_source_terms!(Gⁿ, arch, U, C, args...)
     return nothing
 end
 
+
 #####
 ##### Vertical integrals
 #####

src/TimeSteppers/velocity_and_tracer_tendencies.jl

@@ -0,0 +1,130 @@
+"""
+    u_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                        closure, U, C, K, F, pHY′, parameters, time)
+
+Return the tendency for the horizontal velocity in the x-direction, or the east-west 
+direction, ``u``, at grid point `i, j, k`.
+
+The tendency for ``u`` is called ``G_u`` and defined via
+
+    ``∂_t u = G_u - ∂_x ϕ_n``
+
+where ∂_x ϕ_n is the non-hydrostatic pressure gradient in the x-direction.
+
+`coriolis`, `surface_waves`, and `closure` are types encoding information about Coriolis
+forces, surface waves, and the prescribed turbulence closure.
+
+The arguments `U`, `C`, and `K` are `NamedTuple`s with the three velocity components,
+tracer fields, and precalculated diffusivities where applicable. `F` is a named tuple of 
+forcing functions, `pHY′` is the hydrostatic pressure anomaly.
+
+`parameters` is a `NamedTuple` of scalar parameters for user-defined forcing functions 
+and `time` is the physical time of the model.
+"""
+@inline function u_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                     closure, U, C, K, F, pHY′, parameters, time)
+
+    return ( - div_ũu(i, j, k, grid, U)
+             - x_f_cross_U(i, j, k, grid, coriolis, U)
+             - ∂xᶠᵃᵃ(i, j, k, grid, pHY′)
+             + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, U, K)
+             + x_curl_Uˢ_cross_U(i, j, k, grid, surface_waves, U, time)
+             + ∂t_uˢ(i, j, k, grid, surface_waves, time)
+             + F.u(i, j, k, grid, time, U, C, parameters))
+end
+
+"""
+    v_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                        closure, U, C, K, F, pHY′, parameters, time)
+
+Return the tendency for the horizontal velocity in the y-direction, or the north-south 
+direction, ``v``, at grid point `i, j, k`.
+
+The tendency for ``v`` is called ``G_v`` and defined via
+
+    ``∂_t v = G_v - ∂_y ϕ_n``
+
+where ∂_y ϕ_n is the non-hydrostatic pressure gradient in the y-direction.
+
+`coriolis`, `surface_waves`, and `closure` are types encoding information about Coriolis
+forces, surface waves, and the prescribed turbulence closure.
+
+The arguments `U`, `C`, and `K` are `NamedTuple`s with the three velocity components,
+tracer fields, and precalculated diffusivities where applicable. `F` is a named tuple of 
+forcing functions, `pHY′` is the hydrostatic pressure anomaly.
+
+`parameters` is a `NamedTuple` of scalar parameters for user-defined forcing functions 
+and `time` is the physical time of the model.
+"""
+@inline function v_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                     closure, U, C, K, F, pHY′, parameters, time)
+
+    return ( - div_ũv(i, j, k, grid, U)
+             - y_f_cross_U(i, j, k, grid, coriolis, U)
+             - ∂yᵃᶠᵃ(i, j, k, grid, pHY′)
+             + ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure, U, K)
+             + y_curl_Uˢ_cross_U(i, j, k, grid, surface_waves, U, time)
+             + ∂t_vˢ(i, j, k, grid, surface_waves, time)
+             + F.v(i, j, k, grid, time, U, C, parameters))
+end
+
+"""
+    w_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                        closure, U, C, K, F, parameters, time)
+
+Return the tendency for the vertical velocity ``w`` at grid point `i, j, k`.
+The tendency for ``w`` is called ``G_w`` and defined via
+
+    ``∂_t w = G_w - ∂_z ϕ_n``
+
+where ∂_z ϕ_n is the non-hydrostatic pressure gradient in the z-direction.
+
+`coriolis`, `surface_waves`, and `closure` are types encoding information about Coriolis
+forces, surface waves, and the prescribed turbulence closure.
+
+The arguments `U`, `C`, and `K` are `NamedTuple`s with the three velocity components,
+tracer fields, and precalculated diffusivities where applicable. `F` is a named tuple of 
+forcing functions, `pHY′` is the hydrostatic pressure anomaly.
+
+`parameters` is a `NamedTuple` of scalar parameters for user-defined forcing functions 
+and `time` is the physical time of the model.
+"""
+@inline function w_velocity_tendency(i, j, k, grid, coriolis, surface_waves, 
+                                     closure, U, C, K, F, parameters, time)
+
+    return ( - div_ũw(i, j, k, grid, U)
+             - z_f_cross_U(i, j, k, grid, coriolis, U)
+             + ∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure, U, K)
+             + z_curl_Uˢ_cross_U(i, j, k, grid, surface_waves, U, time)
+             + ∂t_wˢ(i, j, k, grid, surface_waves, time)
+             + F.w(i, j, k, grid, time, U, C, parameters))
+end
+
+"""
+    tracer_tendency(i, j, k, grid, c, tracer_index, closure, buoyancy, U, C, K, Fc,
+                    parameters, time)
+
+Return the tendency for a tracer field `c` with index `tracer_index` 
+at grid point `i, j, k`.
+
+The tendency for ``c`` is called ``G_c`` and defined via
+
+    ``∂_t c = G_c``
+
+`closure` and `buoyancy` are types encoding information about the prescribed
+turbulence closure and buoyancy model.
+
+The arguments `U`, `C`, and `K` are `NamedTuple`s with the three velocity components, 
+tracer fields, and  precalculated diffusivities where applicable. 
+`Fc` is the user-defined forcing function for tracer `c`.
+
+`parameters` is a `NamedTuple` of scalar parameters for user-defined forcing functions 
+and `time` is the physical time of the model.
+"""
+@inline function tracer_tendency(i, j, k, grid, c, tracer_index, 
+                                 closure, buoyancy, U, C, K, Fc, parameters, time)
+
+    return ( - div_uc(i, j, k, grid, U, c)
+             + ∇_κ_∇c(i, j, k, grid, closure, c, tracer_index, K, C, buoyancy)
+             + Fc(i, j, k, grid, time, U, C, parameters))
+end