Support for user-defined forcing functions.

Manifest.toml

@@ -54,9 +54,9 @@ uuid = "8bf52ea8-c179-5cab-976a-9e18b702a9bc"
 
 [[CUDAapi]]
 deps = ["Libdl", "Logging", "Test"]
-git-tree-sha1 = "350cde12f25d297609369a9acb4c6214211675db"
+git-tree-sha1 = "e1f551ad1c03b3fa2a966794ead05772bff4b064"
 uuid = "3895d2a7-ec45-59b8-82bb-cfc6a382f9b3"
-version = "0.5.4"
+version = "0.6.0"
 
 [[CUDAdrv]]
 deps = ["CUDAapi", "Libdl", "Printf", "Test"]
@@ -121,7 +121,7 @@ uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
 version = "0.0.10"
 
 [[Distributed]]
-deps = ["Random", "Serialization", "Sockets"]
+deps = ["LinearAlgebra", "Random", "Serialization", "Sockets"]
 uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
 [[FFTW]]
@@ -150,9 +150,9 @@ version = "0.3.5"
 
 [[ForwardDiff]]
 deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "InteractiveUtils", "LinearAlgebra", "NaNMath", "Random", "SparseArrays", "SpecialFunctions", "StaticArrays", "Test"]
-git-tree-sha1 = "e393bd3b9102659fb24fe88caedec41f2bc2e7de"
+git-tree-sha1 = "4c4d727f1b7e0092134fabfab6396b8945c1ea5b"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "0.10.2"
+version = "0.10.3"
 
 [[GPUArrays]]
 deps = ["Adapt", "FFTW", "FillArrays", "LinearAlgebra", "Printf", "Random", "Serialization", "StaticArrays", "Test"]
@@ -185,7 +185,7 @@ uuid = "d9be37ee-ecc9-5288-90f1-b9ca67657a75"
 version = "0.7.1"
 
 [[InteractiveUtils]]
-deps = ["Markdown"]
+deps = ["LinearAlgebra", "Markdown"]
 uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [[JLD]]
@@ -362,7 +362,7 @@ uuid = "30578b45-9adc-5946-b283-645ec420af67"
 version = "0.4.0"
 
 [[UUIDs]]
-deps = ["Random", "SHA"]
+deps = ["Random"]
 uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [[Unicode]]

examples/deep_convection_3d.jl

@@ -5,12 +5,26 @@ using Oceananigans
 function deep_convection_3d()
     Nx, Ny, Nz = 100, 100, 50
     Lx, Ly, Lz = 2000, 2000, 1000
-    Nt, Δt = 500, 20
+    Nt, Δt = 1000, 20
 
     model = Model((Nx, Ny, Nz), (Lx, Ly, Lz))
-    impose_initial_conditions!(model)
 
-    nc_writer = NetCDFFieldWriter(".", "deep_convection_3d", 50)
+    # impose_initial_conditions!(model)
+
+    # We will impose a surface heat flux of -800 W/m² ≈ -9e-6 K/s  in a square
+    # on the surface of the domain.
+    @inline function surface_cooling_disk(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
+        if k == 1 && (20 < i < 80) && (20 < j < 80)
+            return -9e-6
+        else
+            return 0
+        end
+    end
+
+    # Only impose surface_cooling_disk on the temperature field.
+    model.forcing = Forcing(nothing, nothing, nothing, surface_cooling_disk, nothing)
+
+    nc_writer = NetCDFFieldWriter(".", "deep_convection_3d", 20)
     push!(model.output_writers, nc_writer)
 
     time_step!(model; Nt=Nt, Δt=Δt)
@@ -55,6 +69,7 @@ function impose_cooling_disk!(model::Model)
     # source terms at each time step. Also convert surface heat flux [W/m²]
     # into a temperature tendency forcing [K/s].
     @. model.forcings.FT.data[:, :, 1] = (Q / cᵥ) * (g.Az / (model.eos.ρ₀ * g.V))
+    @show model.forcings.FT.data[50, 50, 1]
     nothing
 end
 

src/Oceananigans.jl

@@ -31,6 +31,8 @@ export
     OperatorTemporaryFields,
     StepperTemporaryFields,
 
+    Forcing,
+
     LinearEquationOfState,
     ρ!,
     δρ!,
@@ -104,6 +106,7 @@ include("planetary_constants.jl")
 include("grids.jl")
 include("fields.jl")
 include("fieldsets.jl")
+include("forcing.jl")
 
 include("operators/operators.jl")
 

src/forcing.jl

@@ -0,0 +1,28 @@
+"Dummy function and forcing default."
+@inline zero_func(args...) = 0
+
+"""
+    Forcing(Fu, Fv, Fw, FF, FS)
+
+    Forcing(; Fu=zero_func, Fv=zero_func, Fw=zero_func, FT=zero_func, FS=zero_func)
+
+Construct a `Forcing` to specify functions that force `u`, `v`, `w`, `T`, and `S`. 
+Forcing functions default to `zero_func`, which does nothing.
+"""
+struct Forcing{Tu,Tv,Tw,TT,TS}
+    u::Tu
+    v::Tv
+    w::Tw
+    T::TT
+    S::TS
+    function Forcing(Fu, Fv, Fw, FT, FS)
+        Fu = Fu === nothing ? zero_func : Fu
+        Fv = Fv === nothing ? zero_func : Fv
+        Fw = Fw === nothing ? zero_func : Fw
+        FT = FT === nothing ? zero_func : FT
+        FS = FS === nothing ? zero_func : FS
+        new{typeof(Fu),typeof(Fv),typeof(Fw),typeof(FT),typeof(FS)}(Fu, Fv, Fw, FT, FS)
+    end
+end
+
+Forcing(; Fu=nothing, Fv=nothing, Fw=nothing, FT=nothing, FS=nothing) = Forcing(Fu, Fv, Fw, FT, FS)

src/model.jl

@@ -11,6 +11,7 @@ mutable struct Model
     G::SourceTerms
     Gp::SourceTerms
     forcings::ForcingFields
+    forcing::Forcing
     stepper_tmp::StepperTemporaryFields
     operator_tmp::OperatorTemporaryFields
     ssp  # ::SpectralSolverParameters or ::SpectralSolverParametersGPU
@@ -22,7 +23,7 @@ end
 function Model(N, L, arch=:cpu, float_type=Float64)
     metadata = _ModelMetadata(arch, float_type)
     configuration = _ModelConfiguration(4e-2, 4e-2, 4e-2, 4e-2)
-    boundary_conditions = BoundaryConditions(:periodic, :periodic, :rigid_lid, :free_slip)
+    boundary_conditions = _BoundaryConditions(:periodic, :periodic, :rigid_lid, :free_slip)
 
     constants = Earth()
     eos = LinearEquationOfState()
@@ -38,6 +39,8 @@ function Model(N, L, arch=:cpu, float_type=Float64)
     stepper_tmp = StepperTemporaryFields(metadata, grid)
     operator_tmp = OperatorTemporaryFields(metadata, grid)
 
+    forcing = Forcing(nothing, nothing, nothing, nothing, nothing)
+
     time, time_step, Δt = 0, 0, 0
     clock = Clock(time, time_step, Δt)
 
@@ -71,6 +74,6 @@ function Model(N, L, arch=:cpu, float_type=Float64)
     ρ!(eos, grid, tracers)
 
     Model(metadata, configuration, boundary_conditions, constants, eos, grid,
-          velocities, tracers, pressures, G, Gp, forcings,
+          velocities, tracers, pressures, G, Gp, forcings, forcing,
           stepper_tmp, operator_tmp, ssp, clock, output_writers, diagnostics)
 end

src/time_steppers.jl

@@ -63,7 +63,7 @@ function time_step_kernel_cpu!(model::Model, Nt, Δt)
         update_source_terms!(Val(:CPU), fCor, χ, eos.ρ₀, cfg.κh, cfg.κv, cfg.𝜈h, cfg.𝜈v, Nx, Ny, Nz, Δx, Δy, Δz,
                              U.u.data, U.v.data, U.w.data, tr.T.data, tr.S.data, pr.pHY′.data,
                              G.Gu.data, G.Gv.data, G.Gw.data, G.GT.data, G.GS.data,
-                             Gp.Gu.data, Gp.Gv.data, Gp.Gw.data, Gp.GT.data, Gp.GS.data, F.FT.data)
+                             Gp.Gu.data, Gp.Gv.data, Gp.Gw.data, Gp.GT.data, Gp.GS.data, model.forcing)
 
         calculate_source_term_divergence_cpu!(Val(:CPU), Nx, Ny, Nz, Δx, Δy, Δz, G.Gu.data, G.Gv.data, G.Gw.data, RHS.data)
 
@@ -210,7 +210,7 @@ function update_buoyancy!(::Val{Dev}, gΔz, Nx, Ny, Nz, ρ, δρ, T, pHY′, ρ
     @synchronize
 end
 
-function update_source_terms!(::Val{Dev}, fCor, χ, ρ₀, κh, κv, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, u, v, w, T, S, pHY′, Gu, Gv, Gw, GT, GS, Gpu, Gpv, Gpw, GpT, GpS, FT) where Dev
+function update_source_terms!(::Val{Dev}, fCor, χ, ρ₀, κh, κv, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, u, v, w, T, S, pHY′, Gu, Gv, Gw, GT, GS, Gpu, Gpv, Gpw, GpT, GpS, F) where Dev
     @setup Dev
 
     @loop for k in (1:Nz; blockIdx().z)
@@ -222,12 +222,12 @@ function update_source_terms!(::Val{Dev}, fCor, χ, ρ₀, κh, κv, 𝜈h, 𝜈
                 @inbounds GpT[i, j, k] = GT[i, j, k]
                 @inbounds GpS[i, j, k] = GS[i, j, k]
 
-                @inbounds Gu[i, j, k] = -u∇u(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + fCor*avg_xy(v, Nx, Ny, i, j, k) - δx_c2f(pHY′, Nx, i, j, k) / (Δx * ρ₀) + 𝜈∇²u(u, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
-                @inbounds Gv[i, j, k] = -u∇v(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) - fCor*avg_xy(u, Nx, Ny, i, j, k) - δy_c2f(pHY′, Ny, i, j, k) / (Δy * ρ₀) + 𝜈∇²v(v, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
-                @inbounds Gw[i, j, k] = -u∇w(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)                                                                           + 𝜈∇²w(w, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
+                @inbounds Gu[i, j, k] = -u∇u(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + fCor*avg_xy(v, Nx, Ny, i, j, k) - δx_c2f(pHY′, Nx, i, j, k) / (Δx * ρ₀) + 𝜈∇²u(u, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + F.u(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
+                @inbounds Gv[i, j, k] = -u∇v(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) - fCor*avg_xy(u, Nx, Ny, i, j, k) - δy_c2f(pHY′, Ny, i, j, k) / (Δy * ρ₀) + 𝜈∇²v(v, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + F.v(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
+                @inbounds Gw[i, j, k] = -u∇w(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)                                                                           + 𝜈∇²w(w, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + F.w(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
 
-                @inbounds GT[i, j, k] = -div_flux(u, v, w, T, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + κ∇²(T, κh, κv, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + FT[i, j, k]
-                @inbounds GS[i, j, k] = -div_flux(u, v, w, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + κ∇²(S, κh, κv, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
+                @inbounds GT[i, j, k] = -div_flux(u, v, w, T, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + κ∇²(T, κh, κv, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + F.T(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
+                @inbounds GS[i, j, k] = -div_flux(u, v, w, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + κ∇²(S, κh, κv, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) + F.S(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
 
                 @inbounds Gu[i, j, k] = (1.5 + χ)*Gu[i, j, k] - (0.5 + χ)*Gpu[i, j, k]
                 @inbounds Gv[i, j, k] = (1.5 + χ)*Gv[i, j, k] - (0.5 + χ)*Gpv[i, j, k]

test/runtests.jl

@@ -385,4 +385,19 @@ using Oceananigans.Operators
         Oceananigans.Operators.𝜈∇²w!(g, w, 𝜈_lap_w, 𝜈h, 𝜈v, otmp)
         for idx in test_indices; @test 𝜈∇²w(w.data, 𝜈h, 𝜈v, g.Nx, g.Ny, g.Nz, g.Δx, g.Δy, g.Δz, idx...) ≈ 𝜈_lap_w.data[idx...]; end
     end
-end
+
+    @testset "Forcing" begin
+        add_one(args...) = 1.0
+        function test_forcing(fld)
+            kwarg = Dict(Symbol(:F, fld)=>add_one)
+            forcing = Forcing(; kwarg...)
+            f = getfield(forcing, fld)
+            f() == 1.0
+        end
+
+        for fld in [:u, :v, :w, :T, :S]
+            @test test_forcing(fld)
+        end
+    end
+
+end # Oceananigans tests