Arbitrary tracers

benchmark/benchmark_static_ocean.jl

@@ -19,7 +19,7 @@ for arch in archs, float_type in float_types, N in Ns
     Nx, Ny, Nz = N
     Lx, Ly, Lz = 1, 1, 1
 
-    model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), arch=arch, float_type=float_type)
+    model = Model(architecture=arch, float_type=float_type, grid=RegularCartesianGrid(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz)))
     time_step!(model, Ni, 1)
 
     bname =  benchmark_name(N, "", arch, float_type)

benchmark/benchmark_tracers.jl

@@ -0,0 +1,75 @@
+using Printf, TimerOutputs
+using Oceananigans
+
+include("benchmark_utils.jl")
+
+const timer = TimerOutput()
+
+####
+#### Benchmark parameters
+####
+
+Ni = 2  # Number of iterations before benchmarking starts.
+Nt = 5  # Number of iterations to use for benchmarking time stepping.
+
+         archs = [CPU()]             # Architectures to benchmark on.
+@hascuda archs = [CPU(), GPU()]      # Benchmark GPU on systems with CUDA-enabled GPUs.
+
+FT = Float64
+Nxyz(::CPU) = (32, 32, 32)
+Nxyz(::GPU) = (128, 128, 128)
+
+####
+#### Utility functions for generating tracer lists
+####
+
+function active_tracers(n)
+    n == 0 && return []
+    n == 1 && return [:b]
+    n == 2 && return [:T, :S]
+    throw(ArgumentError("Can't have more than 2 active tracers!"))
+end
+
+passive_tracers(n) = [Symbol("C" * string(n)) for n in 1:n]
+
+tracer_list(na, np) = Tuple(vcat(active_tracers(na), passive_tracers(np)))
+
+function na2buoyancy(n)
+    n == 0 && return nothing
+    n == 1 && return BuoyancyTracer()
+    n == 2 && return SeawaterBuoyancy()
+    throw(ArgumentError("Can't have more than 2 active tracers!"))
+end
+
+####
+#### Run benchmarks.
+####
+
+test_cases = [(0, 0), (0, 1), (0, 2), (1, 0), (2, 0), (2, 3), (2, 5), (2, 10), (2, 50)]
+
+for arch in archs, test_case in test_cases
+    N = Nxyz(arch)
+    Nx, Ny, Nz = N
+    Lx, Ly, Lz = 1, 1, 1
+
+    na, np = test_case
+    tracers = tracer_list(na, np)
+
+    bname =  benchmark_name(N, "$na active + $(lpad(np, 2)) passive", arch, FT)
+    @printf("Running benchmark: %s...\n", bname)
+
+    model = Model(architecture = arch,
+                    float_type = FT,
+                          grid = RegularCartesianGrid(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz)),
+                      buoyancy = na2buoyancy(na),
+                       tracers = tracers)
+    time_step!(model, Ni, 1)
+
+    for i in 1:Nt
+        @timeit timer bname time_step!(model, 1, 1)
+    end
+end
+
+print_timer(timer, title="Tracer benchmarks", sortby=:name)
+println("")
+

benchmark/benchmark_utils.jl

@@ -1,3 +1,5 @@
+using Oceananigans: AbstractArchitecture
+
 arch_name(::CPU) = "CPU"
 arch_name(::GPU) = "GPU"
 

examples/internal_wave.jl

@@ -42,7 +42,7 @@ f = 0.2     # inertial frequency
 U = k * ω   / (ω^2 - f^2)
 V = k * f   / (ω^2 - f^2)
 W = m * ω   / (ω^2 - N^2)
-Θ = m * N^2 / (ω^2 - N^2)
+B = m * N^2 / (ω^2 - N^2)
 
 # Finally, we set-up a small-amplitude, Gaussian envelope for the wave packet
 
@@ -56,7 +56,7 @@ a(x, z) = A * exp( -( (x - x₀)^2 + (z - z₀)^2 ) / 2δ^2 )
 u₀(x, y, z) = a(x, z) * U * cos(k*x + m*z)
 v₀(x, y, z) = a(x, z) * V * sin(k*x + m*z)
 w₀(x, y, z) = a(x, z) * W * cos(k*x + m*z)
-T₀(x, y, z) = a(x, z) * Θ * sin(k*x + m*z) + N^2 * z 
+b₀(x, y, z) = a(x, z) * B * sin(k*x + m*z) + N^2 * z 
 
 # We are now ready to instantiate our model on a uniform grid. 
 # We give the model a constant rotation rate with background vorticity `f`,
@@ -67,13 +67,14 @@ model = Model(
         grid = RegularCartesianGrid(N=(Nx, 1, Nx), L=(Lx, Lx, Lx)),
      closure = ConstantIsotropicDiffusivity(ν=1e-6, κ=1e-6),
     coriolis = FPlane(f=f), 
+     tracers = :b,
     buoyancy = BuoyancyTracer()
 )
 
-# We initialize the velocity and buoyancy (temperature) fields
+# We initialize the velocity and buoyancy fields
 # with our internal wave initial condition.
 
-set!(model, u=u₀, v=v₀, w=w₀, T=T₀)
+set!(model, u=u₀, v=v₀, w=w₀, b=b₀)
 
 # ## Some plotting utilities
 #

examples/ocean_convection_with_plankton.jl

@@ -39,8 +39,8 @@ end_time = 1day
 # Create boundary conditions. Note that temperature is buoyancy in our problem.
 #
 
-Tbcs = HorizontallyPeriodicBCs(   top = BoundaryCondition(Flux, Qb), 
-                               bottom = BoundaryCondition(Gradient, N²))
+buoyancy_bcs = HorizontallyPeriodicBCs(   top = BoundaryCondition(Flux, Qb), 
+                                       bottom = BoundaryCondition(Gradient, N²))
 
 # ## Define a forcing function
 #
@@ -55,15 +55,16 @@ model = Model(
                    grid = RegularCartesianGrid(N = (Nz, 1, Nz), L = (Lz, Lz, Lz)),
                 closure = ConstantIsotropicDiffusivity(ν=1e-4, κ=1e-4),
                coriolis = FPlane(f=1e-4), 
+                tracers = (:b, :plankton),
                buoyancy = BuoyancyTracer(),
-                forcing = ModelForcing(S=growth_and_decay),
-    boundary_conditions = BoundaryConditions(T=Tbcs)
+                forcing = ModelForcing(plankton=growth_and_decay),
+    boundary_conditions = BoundaryConditions(b=buoyancy_bcs)
 )
 
 ## Set initial condition. Initial velocity and salinity fluctuations needed for AMD.
 Ξ(z) = randn() * z / Lz * (1 + z / Lz) # noise
-T₀(x, y, z) = N² * z + N² * Lz * 1e-6 * Ξ(z)
-set!(model, T=T₀)
+b₀(x, y, z) = N² * z + N² * Lz * 1e-6 * Ξ(z)
+set!(model, b=b₀)
 
 ## A wizard for managing the simulation time-step.
 wizard = TimeStepWizard(cfl=0.2, Δt=1.0, max_change=1.1, max_Δt=90.0)
@@ -87,7 +88,7 @@ while model.clock.time < end_time
     ylabel("\$ z \$ (m)")
 
     sca(axs[2]); cla()
-    pcolormesh(xC, zC, data(model.tracers.S)[:, 1, :])
+    pcolormesh(xC, zC, data(model.tracers.plankton)[:, 1, :])
     title("Phytoplankton concentration")
     xlabel("\$ x \$ (m)")
     axs[2].tick_params(left=false, labelleft=false)

examples/simple_diffusion.jl

@@ -52,7 +52,7 @@ set!(model, T=Tᵢ)
 # `time_step!`, with a time-step size that ensures numerical stability.
 
 ## Time-scale for diffusion across a grid cell
-cell_diffusion_time_scale = model.grid.Δz^2 / model.closure.κ
+cell_diffusion_time_scale = model.grid.Δz^2 / model.closure.κ.T
 
 ## The function `time_step!` executes `Nt` time steps with step size `Δt`
 ## using a second-order Adams-Bashforth method

examples/two_dimensional_turbulence.jl

@@ -1,7 +1,14 @@
 # # Two dimensional turbulence example
 #
 # In this example, we initialize a random velocity field and observe its viscous,
-# turbulent decay in a two-dimensional domain.
+# turbulent decay in a two-dimensional domain. This example demonstrates:
+#
+#   * How to run a model with no buoyancy equation or tracers;
+#   * How to create user-defined fields
+#   * How to use differentiation functions
+#
+# For this example, we need `PyPlot` for plotting and `Statistics` for setting up
+# a random initial condition with zero mean velocity.
 
 using Oceananigans, PyPlot, Statistics
 
@@ -16,6 +23,8 @@ using Oceananigans.TurbulenceClosures: ∂x_faa, ∂y_afa
 
 model = Model(
         grid = RegularCartesianGrid(N=(128, 128, 1), L=(2π, 2π, 2π)),
+    buoyancy = nothing, 
+     tracers = nothing,
      closure = ConstantIsotropicDiffusivity(ν=1e-3, κ=1e-3)
 )
 
@@ -46,7 +55,7 @@ close("all")
 fig, ax = subplots()
 
 for i = 1:10
-    time_step!(model, Nt = 100, Δt = 1e-1)
+    time_step!(model, Nt=100, Δt=1e-1)
 
     vorticity!(ω, model.velocities.u, model.velocities.v)
 

src/Oceananigans.jl

@@ -229,9 +229,10 @@ end
 architecture(::Array) = CPU()
 @hascuda architecture(::CuArray) = GPU()
 
-# Place-holder buoyancy functions for use in TurbulenceClosures module
+# Place-holder functions for use in TurbulenceClosures module
 function buoyancy_perturbation end
 function buoyancy_frequency_squared end
+function TracerFields end
 
 include("utils.jl")
 

src/TurbulenceClosures/TurbulenceClosures.jl

@@ -3,19 +3,16 @@ module TurbulenceClosures
 export
   IsotropicDiffusivity,
   ConstantIsotropicDiffusivity,
-
   ConstantAnisotropicDiffusivity,
-
   ConstantSmagorinsky,
   SmagorinskyLilly,
   BlasiusSmagorinsky,
-
   AnisotropicMinimumDissipation,
   RozemaAnisotropicMinimumDissipation,
   VerstappenAnisotropicMinimumDissipation,
 
   TurbulentDiffusivities,
-  calc_diffusivities!,
+  calculate_diffusivities!,
 
   ∇_κ_∇c,
   ∇_κ_∇T,
@@ -31,7 +28,10 @@ using
   Oceananigans.Operators,
   GPUifyLoops
 
-using Oceananigans: AbstractArchitecture, AbstractGrid, buoyancy_perturbation, buoyancy_frequency_squared
+import Oceananigans: with_tracers
+
+using Oceananigans: AbstractArchitecture, AbstractGrid, buoyancy_perturbation, buoyancy_frequency_squared, 
+                    TracerFields, device, launch_config
 
 @hascuda using CUDAdrv, CUDAnative
 
@@ -50,63 +50,61 @@ const κ₀ = 1.46e-7
 ####
 
 """
-    TurbulenceClosure{T}
+    TurbulenceClosure{FT}
 
 Abstract supertype for turbulence closures with model parameters stored as properties of
-type `T`.
+type `FT`.
 """
-abstract type TurbulenceClosure{T} end
+abstract type TurbulenceClosure{FT} end
 
 """
-    IsotropicDiffusivity{T} <: TurbulenceClosure{T}
+    IsotropicDiffusivity{FT} <: TurbulenceClosure{FT}
 
 Abstract supertype for turbulence closures that are defined by an isotropic viscosity
-and isotropic diffusivities with model parameters stored as properties of type `T`.
+and isotropic diffusivities with model parameters stored as properties of type `FT`.
 """
-abstract type IsotropicDiffusivity{T} <: TurbulenceClosure{T} end
+abstract type IsotropicDiffusivity{FT} <: TurbulenceClosure{FT} end
 
 """
-    TensorDiffusivity{T} <: TurbulenceClosure{T}
+    TensorDiffusivity{FT} <: TurbulenceClosure{FT}
 
 Abstract supertype for turbulence closures that are defined by a tensor viscosity and
-tensor diffusivities with model parameters stored as properties of type `T`.
+tensor diffusivities with model parameters stored as properties of type `FT`.
 """
-abstract type TensorDiffusivity{T} <: TurbulenceClosure{T} end
+abstract type TensorDiffusivity{FT} <: TurbulenceClosure{FT} end
 
 """
-    AbstractSmagorinsky{T}
+    AbstractSmagorinsky{FT}
 
 Abstract supertype for large eddy simulation models based off the model described
-by Smagorinsky with model parameters stored as properties of type `T`.
+by Smagorinsky with model parameters stored as properties of type `FT`.
 """
-abstract type AbstractSmagorinsky{T} <: IsotropicDiffusivity{T} end
+abstract type AbstractSmagorinsky{FT} <: IsotropicDiffusivity{FT} end
 
 """
-    AbstractAnisotropicMinimumDissipation{T}
+    AbstractAnisotropicMinimumDissipation{FT}
 
 Abstract supertype for large eddy simulation models based on the anisotropic minimum
-dissipation principle with model parameters stored as properties of type `T`.
+dissipation principle with model parameters stored as properties of type `FT`.
 """
-abstract type AbstractAnisotropicMinimumDissipation{T} <: IsotropicDiffusivity{T} end
+abstract type AbstractAnisotropicMinimumDissipation{FT} <: IsotropicDiffusivity{FT} end
 
 ####
 #### Include module code
 ####
 
-@inline ∇_κ_∇T(args...) = ∇_κ_∇c(args...)
-@inline ∇_κ_∇S(args...) = ∇_κ_∇c(args...)
-
 # Fallback constructor for diffusivity types without precomputed diffusivities:
 TurbulentDiffusivities(arch::AbstractArchitecture, grid::AbstractGrid, args...) = nothing
 
 include("turbulence_closure_utils.jl")
 include("closure_operators.jl")
 include("velocity_tracer_gradients.jl")
 
-include("constant_diffusivity_closures.jl")
+include("constant_isotropic_diffusivity.jl")
+include("constant_anisotropic_diffusivity.jl")
 include("smagorinsky.jl")
-include("rozema_anisotropic_minimum_dissipation.jl")
 include("verstappen_anisotropic_minimum_dissipation.jl")
+include("rozema_anisotropic_minimum_dissipation.jl")
 
 include("turbulence_closure_diagnostics.jl")