Merge pull request #989 from CliMA/glw/unified-forcing-functions

New forcing function interface, defaults, and functionality

benchmark/benchmark_forcing_functions.jl

@@ -24,7 +24,7 @@ Nt = 10  # Number of iterations to use for benchmarking time stepping.
 ##### Forcing function definitions
 #####
 
-@inline function Fu_params(i, j, k, grid, time, U, C, params)
+@inline function Fu_params_func(i, j, k, grid, clock, model_fields, params)
     if k == 1
         return @inbounds -2*params.K/grid.Δz^2 * (U.u[i, j, 1] - 0)
     elseif k == grid.Nz
@@ -34,8 +34,15 @@ Nt = 10  # Number of iterations to use for benchmarking time stepping.
     end
 end
 
+@inline FT_params_func(i, j, k, grid, time, model_fields, params) = @inbounds ifelse(k == 1, -params.λ * (model_fields.T[i, j, 1] - 0), 0)
+
+Fu_params = Forcing(FT_params_func, discrete_form=true, parameters=(K=0.1,))
+FT_params = Forcing(FT_params_func, discrete_form=true, parameters=(λ=1e-4,))
+
+const λ = 1e-4
 const K = 0.1
-@inline function Fu_consts(i, j, k, grid, time, U, C, params)
+
+@inline function Fu_consts(i, j, k, grid, clock, model_forcing)
     if k == 1
         return @inbounds -2*K/grid.Δz^2 * (U.u[i, j, 1] - 0)
     elseif k == grid.Nz
@@ -45,10 +52,10 @@ const K = 0.1
     end
 end
 
-@inline FT_params(i, j, k, grid, time, U, C, params) = @inbounds ifelse(k == 1, -params.λ * (C.T[i, j, 1] - 0), 0)
+@inline FT_consts_func(i, j, k, grid, time, model_fields) = @inbounds ifelse(k == 1, -λ * (model_fields.T[i, j, 1] - 0), 0)
 
-const λ = 1e-4
-@inline FT_consts(i, j, k, grid, time, U, C, params) = @inbounds ifelse(k == 1, -λ * (C.T[i, j, 1] - 0), 0)
+Fu_consts = Forcing(FT_consts_func, discrete_form=true)
+FT_consts = Forcing(FT_consts_func, discrete_form=true)
 
 #####
 ##### Run benchmarks
@@ -58,9 +65,10 @@ for arch in archs, FT in float_types, N in Ns
     Nx, Ny, Nz = N
     Lx, Ly, Lz = 1, 1, 1
 
-    forced_model_params = Model(architecture = arch, float_type = FT,
-		                grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz)),
-                                forcing=ModelForcing(Fu=Fu_params, FT=FT_params), parameters=(K=0.1, λ=1e-4))
+    forced_model_params = Model(architecture = arch,
+                                float_type = FT,
+		                        grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz)),
+                                forcing = (Fl=Fu_params, FT=FT_params))
 
     time_step!(forced_model_params, Ni, 1)  # First 1-2 iterations usually slower.
 
@@ -70,9 +78,10 @@ for arch in archs, FT in float_types, N in Ns
         @timeit timer bn time_step!(forced_model_params, 1, 1)
     end
 
-    forced_model_consts = Model(architecture = arch, float_type = FT,
-		                grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz)),
-                                forcing=ModelForcing(Fu=Fu_consts, FT=FT_consts))
+    forced_model_consts = Model(architecture = arch,
+                                float_type = FT,
+		                        grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz)),
+                                forcing=(Fu=Fu_consts, FT=FT_consts))
 
     time_step!(forced_model_consts, Ni, 1)  # First 1-2 iterations usually slower.
 

docs/src/library.md

@@ -109,17 +109,18 @@ Pages   = [
 ]
 ```
 
-## Forcing
+## Forcings
 
 ```@autodocs
-Modules = [Oceananigans.Forcing]
+Modules = [Oceananigans.Forcings]
 Private = false
 Pages   = [
-    "Forcing/Forcing.jl",
-    "Forcing/simple_forcing.jl",
-    "Forcing/model_forcing.jl",
-    "Forcing/parameterized_forcing.jl",
-    "Forcing/relaxation.jl"
+    "Forcings/Forcings.jl",
+    "Forcings/continuous_forcing.jl",
+    "Forcings/discrete_forcing.jl",
+    "Forcings/forcing.jl",
+    "Forcings/model_forcing.jl",
+    "Forcings/relaxation.jl"
 ]
 ```
 

docs/src/model_setup/boundary_conditions.md

@@ -146,22 +146,22 @@ Boundary conditions may also depend on model fields. For example, a linear drag
 is implemented with
 
 ```jldoctest
-julia> @inline linear_drag(i, j, grid, clock, state) = @inbounds - 0.2 * state.velocities.u[i, j, 1];
+julia> @inline linear_drag(i, j, grid, clock, model_fields) = @inbounds - 0.2 * model_fields.u[i, j, 1];
 
 julia> u_bottom_bc = FluxBoundaryCondition(linear_drag, discrete_form=true)
-BoundaryCondition: type=Flux, condition=linear_drag(i, j, grid, clock, state) in Main at none:1
+BoundaryCondition: type=Flux, condition=linear_drag(i, j, grid, clock, model_fields) in Main at none:1
 ```
 
 !!! info "The 'discrete form' for boundary condition functions"
     The argument `discrete_form=true` indicates to [`BoundaryCondition`](@ref) that `linear_drag`
     uses the 'discrete form'. Boundary condition functions that use the 'discrete form'
     are called with the signature 
     ```julia
-    f(i, j, grid, clock, state)
+    f(i, j, grid, clock, model_fields)
     ```
     where `i, j` are grid indices that vary along the boundary, `grid` is `model.grid`,
-    `clock` is the `model.clock`, and `state` is a `NamedTuple`
-    containing `state.velocities`, `state.tracers`, and `state.diffusivities`.
+    `clock` is the `model.clock`, and `model_fields` is a `NamedTuple`
+    containing `u, v, w`, the fields in `model.tracers`, and the fields in `model.diffusivities`.
     The signature is similar for $x$ and $y$ boundary conditions expect that `i, j` is replaced
     with `j, k` and `i, k` respectively.
 
@@ -170,17 +170,17 @@ BoundaryCondition: type=Flux, condition=linear_drag(i, j, grid, clock, state) in
 ```jldoctest
 julia> Cd = 0.2;  # drag coefficient
 
-julia> @inline linear_drag(i, j, grid, clock, state, Cd) = @inbounds - Cd * state.velocities.u[i, j, 1];
+julia> @inline linear_drag(i, j, grid, clock, model_fields, Cd) = @inbounds - Cd * model_fields.u[i, j, 1];
 
 julia> u_bottom_bc = BoundaryCondition(Flux, linear_drag, discrete_form=true, parameters=Cd)
-BoundaryCondition: type=Flux, condition=linear_drag(i, j, grid, clock, state, Cd) in Main at none:1
+BoundaryCondition: type=Flux, condition=linear_drag(i, j, grid, clock, model_fields, Cd) in Main at none:1
 ```
 
 !!! info "Inlining and avoiding bounds-checking in boundary condition functions"
     Boundary condition functions should be decorated with `@inline` when running on CPUs for performance reasons.
     On the GPU, all functions are force-inlined by default.
     In addition, the annotation `@inbounds` should be used when accessing the elements of an array
-    in a boundary condition function (such as `state.velocities[i, j, 1]` in the above example).
+    in a boundary condition function (such as `model_fields.u[i, j, 1]` in the above example).
     Using `@inbounds` will avoid a relatively expensive check that the index `i, j, 1` is 'in bounds'.
 
 ##### 7. A random, spatially-varying, constant-in-time temperature flux specified by an array