Reduce parameter space for velocity tendency kernels (#3924)

* reduce parameter space

* change other kernels for symmetry

* only free surface forcing

* pass specific forcing in kernels

* bugfix

* bugfix

* reduce parameter space

* import on_architecture

* some changes

* put model_fields before the tracer loop

* remove tracers from here

* remove auxiliary fields

* revert to previous code

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl

@@ -1,5 +1,6 @@
 using Oceananigans.BuoyancyModels: g_Earth
 using Oceananigans.Grids: with_halo
+import Oceananigans.Grids: on_architecture
 
 struct SplitExplicitFreeSurface{H, U, M, FT, K , S, T} <: AbstractFreeSurface{H, FT}
     η :: H

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl

@@ -91,8 +91,8 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
                      model.tracers,
                      model.diffusivity_fields,
                      model.auxiliary_fields,
-                     c_forcing,
-                     model.clock)
+                     model.clock,
+                     c_forcing)
 
         launch!(arch, grid, kernel_parameters,
                 compute_hydrostatic_free_surface_Gc!,
@@ -137,20 +137,19 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para
                                 model.diffusivity_fields,
                                 model.pressure.pHY′,
                                 model.auxiliary_fields,
-                                model.forcing,
                                 model.clock)
 
     u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args...)
     v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args...)
 
     launch!(arch, grid, kernel_parameters,
             compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, 
-            active_cells_map, u_kernel_args;
+            active_cells_map, u_kernel_args, model.forcing.u;
             active_cells_map)
 
     launch!(arch, grid, kernel_parameters,
             compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, 
-            active_cells_map, v_kernel_args;
+            active_cells_map, v_kernel_args, model.forcing.v;
             active_cells_map)
 
     return nothing
@@ -180,27 +179,27 @@ end
 #####
 
 """ Calculate the right-hand-side of the u-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args)
+@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args, forcing)
     i, j, k = @index(Global, NTuple)
-    @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
+    @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args)
+@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args, forcing)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
-    @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
+    @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing)
 end
 
 """ Calculate the right-hand-side of the v-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args)
+@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args, forcing)
     i, j, k = @index(Global, NTuple)
-    @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
+    @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args)
+@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args, forcing)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
-    @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
+    @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing)
 end
 
 #####

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -38,8 +38,8 @@ implicitly during time-stepping.
                                                               diffusivities,
                                                               hydrostatic_pressure_anomaly,
                                                               auxiliary_fields,
-                                                              forcings,
-                                                              clock)
+                                                              clock, 
+                                                              forcing)
 
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
 
@@ -49,7 +49,7 @@ implicitly during time-stepping.
              - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
              - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
              - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields)
-             + forcings.u(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers)))
+             + forcing(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers)))
 end
 
 """
@@ -77,8 +77,8 @@ implicitly during time-stepping.
                                                               diffusivities,
                                                               hydrostatic_pressure_anomaly,
                                                               auxiliary_fields,
-                                                              forcings,
-                                                              clock)
+                                                              clock,
+                                                              forcing)
 
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
 
@@ -88,7 +88,7 @@ implicitly during time-stepping.
              - ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
              - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
              - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields)
-             + forcings.v(i, j, k, grid, clock, model_fields))
+             + forcing(i, j, k, grid, clock, model_fields))
 end
 
 """
@@ -116,8 +116,8 @@ where `c = C[tracer_index]`.
                                                           tracers,
                                                           diffusivities,
                                                           auxiliary_fields,
-                                                          forcing,
-                                                          clock) where tracer_index
+                                                          clock,
+                                                          forcing) where tracer_index
 
     @inbounds c = tracers[tracer_index]
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
@@ -135,4 +135,4 @@ where `c = C[tracer_index]`.
              - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
              + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
              + forcing(i, j, k, grid, clock, model_fields))
-end
+end

src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl

@@ -91,15 +91,15 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
 
     u_kernel_args = tuple(start_momentum_kernel_args...,
                           u_immersed_bc, end_momentum_kernel_args...,
-                          forcings, hydrostatic_pressure, clock)
+                          hydrostatic_pressure, clock, forcings.u)
 
     v_kernel_args = tuple(start_momentum_kernel_args...,
                           v_immersed_bc, end_momentum_kernel_args...,
-                          forcings, hydrostatic_pressure, clock)
+                          hydrostatic_pressure, clock, forcings.v)
 
     w_kernel_args = tuple(start_momentum_kernel_args...,
                           w_immersed_bc, end_momentum_kernel_args...,
-                          forcings, hydrostatic_pressure, clock)
+                          hydrostatic_pressure, clock, forcings.w)
 
     exclude_periphery = true
     launch!(arch, grid, kernel_parameters, compute_Gu!, 
@@ -128,7 +128,7 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
                      start_tracer_kernel_args..., 
                      c_immersed_bc,
                      end_tracer_kernel_args...,
-                     forcing, clock)
+                     clock, forcing)
 
         launch!(arch, grid, kernel_parameters, compute_Gc!, 
                 c_tendency, grid, active_cells_map, args;

src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl

@@ -56,17 +56,17 @@ pressure anomaly.
                                      tracers,
                                      auxiliary_fields,
                                      diffusivities,
-                                     forcings,
                                      hydrostatic_pressure,
-                                     clock)
+                                     clock,
+                                     forcing)
 
     model_fields = merge(velocities, tracers, auxiliary_fields)
 
     total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
                         v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
                         w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))
 
-    total_velocities = with_advective_forcing(forcings.u, total_velocities)
+    total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_𝐯u(i, j, k, grid, advection, total_velocities, velocities.u)
              - div_𝐯u(i, j, k, grid, advection, velocities, background_fields.velocities.u)
@@ -77,7 +77,7 @@ pressure anomaly.
              - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields)
              + x_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
              + ∂t_uˢ(i, j, k, grid, stokes_drift, clock.time)
-             + forcings.u(i, j, k, grid, clock, model_fields))
+             + forcing(i, j, k, grid, clock, model_fields))
 end
 
 """
@@ -119,17 +119,17 @@ pressure anomaly.
                                      tracers,
                                      auxiliary_fields,
                                      diffusivities,
-                                     forcings,
                                      hydrostatic_pressure,
-                                     clock)
+                                     clock,
+                                     forcing)
 
     model_fields = merge(velocities, tracers, auxiliary_fields)
 
     total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
                         v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
                         w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))
 
-    total_velocities = with_advective_forcing(forcings.v, total_velocities)
+    total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_𝐯v(i, j, k, grid, advection, total_velocities, velocities.v)
              - div_𝐯v(i, j, k, grid, advection, velocities, background_fields.velocities.v)
@@ -140,7 +140,7 @@ pressure anomaly.
              - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields)
              + y_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
              + ∂t_vˢ(i, j, k, grid, stokes_drift, clock.time)
-             + forcings.v(i, j, k, grid, clock, model_fields))
+             + forcing(i, j, k, grid, clock, model_fields))
 end
 
 # Only add buoyancy if the hydrostatic pressure isa Nothing
@@ -185,17 +185,17 @@ velocity components, tracer fields, and precalculated diffusivities where applic
                                      tracers,
                                      auxiliary_fields,
                                      diffusivities,
-                                     forcings,
                                      hydrostatic_pressure,
-                                     clock)
+                                     clock,
+                                     forcing)
 
     model_fields = merge(velocities, tracers, auxiliary_fields)
 
     total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
                         v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
                         w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))
 
-    total_velocities = with_advective_forcing(forcings.w, total_velocities)
+    total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_𝐯w(i, j, k, grid, advection, total_velocities, velocities.w)
              - div_𝐯w(i, j, k, grid, advection, velocities, background_fields.velocities.w)
@@ -205,7 +205,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic
              - immersed_∂ⱼ_τ₃ⱼ(i, j, k, grid, velocities, w_immersed_bc, closure, diffusivities, clock, model_fields)
              + z_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
              + ∂t_wˢ(i, j, k, grid, stokes_drift, clock.time)
-             + forcings.w(i, j, k, grid, clock, model_fields))
+             + forcing(i, j, k, grid, clock, model_fields))
 end
 
 """
@@ -247,8 +247,8 @@ velocity components, tracer fields, and precalculated diffusivities where applic
                                  tracers,
                                  auxiliary_fields,
                                  diffusivities,
-                                 forcing,
-                                 clock) where tracer_index
+                                 clock,
+                                 forcing) where tracer_index
 
     @inbounds c = tracers[tracer_index]
     @inbounds background_fields_c = background_fields.tracers[tracer_index]

src/Models/ShallowWaterModels/compute_shallow_water_tendencies.jl

@@ -75,22 +75,22 @@ function compute_interior_tendency_contributions!(tendencies,
                                                   formulation)
 
     transport_args = (grid, gravitational_acceleration, advection.momentum, velocities, coriolis, closure, 
-                      bathymetry, solution, tracers, diffusivities, forcings, clock, formulation)
+                      bathymetry, solution, tracers, diffusivities, clock, formulation)
 
     h_args = (grid, gravitational_acceleration, advection.mass, coriolis, closure, 
-              solution, tracers, diffusivities, forcings, clock, formulation)
+              solution, tracers, diffusivities, clock, formulation)
 
-    launch!(arch, grid, :xyz, compute_Guh!, tendencies[1], transport_args...; exclude_periphery=true)
-    launch!(arch, grid, :xyz, compute_Gvh!, tendencies[2], transport_args...; exclude_periphery=true)
-    launch!(arch, grid, :xyz,  compute_Gh!, tendencies[3], h_args...)
+    launch!(arch, grid, :xyz, compute_Guh!, tendencies[1], transport_args..., forcings[1]; exclude_periphery=true)
+    launch!(arch, grid, :xyz, compute_Gvh!, tendencies[2], transport_args..., forcings[2]; exclude_periphery=true)
+    launch!(arch, grid, :xyz,  compute_Gh!, tendencies[3], h_args..., forcings.h)
 
     for (tracer_index, tracer_name) in enumerate(propertynames(tracers))
         @inbounds Gc = tendencies[tracer_index+3]
         @inbounds forcing = forcings[tracer_index+3]
         @inbounds c_advection = advection[tracer_name]
 
         launch!(arch, grid, :xyz, compute_Gc!, Gc, grid, Val(tracer_index),
-                c_advection, closure, solution, tracers, diffusivities, forcing, clock, formulation)
+                c_advection, closure, solution, tracers, diffusivities, clock, formulation, forcing)
     end
 
     return nothing

src/Models/ShallowWaterModels/solution_and_tracer_tendencies.jl

@@ -32,9 +32,9 @@ Compute the tendency for the x-directional transport, uh
                                       solution,
                                       tracers,
                                       diffusivities,
-                                      forcings,
                                       clock,
-                                      formulation)
+                                      formulation,
+                                      forcing)
 
     g = gravitational_acceleration
 
@@ -45,7 +45,7 @@ Compute the tendency for the x-directional transport, uh
              - x_f_cross_U(i, j, k, grid, coriolis, solution)
              - bathymetry_contribution_x(i, j, k, grid, g, solution.h, bathymetry, formulation)
              - sw_∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, formulation)
-             + forcings[1](i, j, k, grid, clock, merge(solution, tracers)))
+             + forcing(i, j, k, grid, clock, merge(solution, tracers)))
 end
 
 """
@@ -61,9 +61,9 @@ Compute the tendency for the y-directional transport, vh.
                                       solution,
                                       tracers,
                                       diffusivities,
-                                      forcings,
                                       clock,
-                                      formulation)
+                                      formulation,
+                                      forcing)
 
      g = gravitational_acceleration
 
@@ -74,7 +74,7 @@ Compute the tendency for the y-directional transport, vh.
              - y_f_cross_U(i, j, k, grid, coriolis, solution)
              - bathymetry_contribution_y(i, j, k, grid, g, solution.h, bathymetry, formulation)
              - sw_∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, formulation)
-             + forcings[2](i, j, k, grid, clock, merge(solution, tracers)))
+             + forcing(i, j, k, grid, clock, merge(solution, tracers)))
 end
 
 """
@@ -88,12 +88,12 @@ Compute the tendency for the height, h.
                                      solution,
                                      tracers,
                                      diffusivities,
-                                     forcings,
                                      clock,
-                                     formulation)
+                                     formulation,
+                                     forcing)
 
     return ( - div_Uh(i, j, k, grid, advection, solution, formulation)
-             + forcings.h(i, j, k, grid, clock, merge(solution, tracers)))
+             + forcing(i, j, k, grid, clock, merge(solution, tracers)))
 end
 
 @inline function tracer_tendency(i, j, k, grid,
@@ -103,9 +103,9 @@ end
                                  solution,
                                  tracers,
                                  diffusivities,
-                                 forcing,
                                  clock,
-                                 formulation) where tracer_index
+                                 formulation,
+                                 forcing) where tracer_index
 
     @inbounds c = tracers[tracer_index]