Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Arbitrary tracers #452

Merged
merged 66 commits into from
Oct 17, 2019
Merged
Show file tree
Hide file tree
Changes from 55 commits
Commits
Show all changes
66 commits
Select commit Hold shift + click to select a range
f0d1f4d
Adds tracer names as kwarg to Model constructor; changes fieldnames o…
glwagner Oct 3, 2019
fbbcca4
Better tracernames function
glwagner Oct 3, 2019
b6a9ca1
Replaces explicit references to tracers with loops over tracer fields…
glwagner Oct 3, 2019
80ca7cb
Implements arbitrary tracer functionality for forcings and boundary c…
glwagner Oct 3, 2019
e7821e8
Move model initialization functionality to new file
glwagner Oct 7, 2019
2e53235
Update boundary condition constructor for new with_tracers syntax
glwagner Oct 7, 2019
fff61fc
Extend constant diffusivity closures and smagorinsky for arbitrary tr…
glwagner Oct 7, 2019
cc415d8
Removes default from TracerFields constructor
glwagner Oct 7, 2019
acf8716
New syntax for TurbulentDiffusivities constructor that includes numbe…
glwagner Oct 7, 2019
676cb09
Change model initializtion to model initialization utils
glwagner Oct 7, 2019
32ddb70
Fixes bug in model constructor
glwagner Oct 7, 2019
6bc9c22
Sketching arbitrary tracers time stepping
glwagner Oct 7, 2019
602ac62
Bugfixes
glwagner Oct 7, 2019
216ccb4
Hacks to get tests passing in interim
glwagner Oct 7, 2019
5598d9e
Nukes typed keyword constructor
glwagner Oct 7, 2019
e56aa0e
Adds back default tracernames for test passing purposes
glwagner Oct 7, 2019
c33a1db
Test edits for arbitrary traacers
glwagner Oct 7, 2019
f98d952
Uncomments tests
glwagner Oct 7, 2019
0b4c045
Extends turbulence closures to allow arbitrary tracers, plus massive …
glwagner Oct 8, 2019
3be696b
Moves geometric mean function to turbulence closure utils
glwagner Oct 8, 2019
0059207
Implements loop over tracers for tracer right-hand-side calculation
glwagner Oct 8, 2019
b1f636e
Tweaks turbulence closure tests
glwagner Oct 8, 2019
809ea4d
Fixes cell diffusion timescale to work for arbitrary tracers / tupled…
glwagner Oct 8, 2019
31399d6
Updates output writers test for new tendency field names
glwagner Oct 8, 2019
8ed8726
Merge branch 'master' into arbitrary-tracers
glwagner Oct 8, 2019
85f6967
Docstrings and formatting update in constant diffusivity closures
glwagner Oct 8, 2019
4f51566
Removes markdown for simple diffusion and updates for arbitrary tracers
glwagner Oct 8, 2019
a3b961c
Merge branch 'arbitrary-tracers' of https://github.com/climate-machin…
glwagner Oct 8, 2019
1879453
Adds a validate buoyancy function to ensure that user-specified trace…
glwagner Oct 8, 2019
da2f801
Uses buoyancy tracer with internal wave and two dimensional turbulenc…
glwagner Oct 8, 2019
1fb83fa
Uses buoyancy and passive tracer "c" in regression, updates buoyancy …
glwagner Oct 8, 2019
c1f150e
Switches BuoyancyTracer to use name "b" for buoyancy tracer
glwagner Oct 8, 2019
1caa51c
Adds options for specifying single tracers or no tracer
glwagner Oct 8, 2019
cb55bd9
Tweaks to get internal wave test working
glwagner Oct 8, 2019
2d2f8ab
Try passing the length of solution / tracers as argument to kernels
glwagner Oct 8, 2019
7afe294
Changes
glwagner Oct 8, 2019
108f386
Attempting to move loops over tracers outside kernels
glwagner Oct 8, 2019
e080ea6
Change calc_diffusivities to calculate_diffusivities
glwagner Oct 9, 2019
9d709d7
Shenanigans and bugfixes
glwagner Oct 9, 2019
001ce4d
Nukes model initialization utils file until future consideration
glwagner Oct 9, 2019
9abc5bc
Updates docstrings in time steppers.jl
glwagner Oct 9, 2019
77c910d
Adds TracerFields placeholder to Oceananigans.jl
glwagner Oct 9, 2019
3c3e7ba
Moves around model initialization functionality
glwagner Oct 9, 2019
fbd7376
Tweaks in calculate_diffusivities for smagorinsky
glwagner Oct 9, 2019
856407f
Bugfix in rozema diffusivity function
glwagner Oct 9, 2019
e89aa1e
Updates pearson vortex and internal wave test for arbitrary tracers
glwagner Oct 9, 2019
77bbf1f
Adds module_suffix option to run_example to avoid duplicate names for…
glwagner Oct 9, 2019
fdf3194
Updates compute w continuity test for new syntax
glwagner Oct 9, 2019
05f7a85
Changes calc diffusivities to calculate and uses new non-kernel launc…
glwagner Oct 9, 2019
6f8d0a0
Reactivate all tests
glwagner Oct 9, 2019
d6506bb
Merge branch 'master' into arbitrary-tracers-outer-loops
glwagner Oct 11, 2019
f18deec
Bugfix in ModelForcing constructor
glwagner Oct 11, 2019
1b093d0
Use b for buoyancy and plankton for plankton in ocean convection example
glwagner Oct 11, 2019
8ade4f1
Add benchmark for active and passive tracers
ali-ramadhan Oct 12, 2019
98503ac
Fix static ocean benchmark
ali-ramadhan Oct 12, 2019
5c056c4
Uses Val type to call tracer flux divergence function in calculate_in…
glwagner Oct 15, 2019
606d404
Fixes bug in smagorinsky and amd associated with changeover to use of…
glwagner Oct 15, 2019
b1175c3
Fixes AMD to use Val type for computation of nonlinear diffusivities
glwagner Oct 15, 2019
f799a3a
Update tracer benchmark to work on V100 GPU
ali-ramadhan Oct 15, 2019
0613078
Merge branch 'arbitrary-tracers-outer-loops' of https://github.com/cl…
ali-ramadhan Oct 15, 2019
f9f2565
Merge branch 'master' into arbitrary-tracers-outer-loops
ali-ramadhan Oct 16, 2019
5e30ce0
Update stratified_couette_flow.jl to work with arbitrary tracers
ali-ramadhan Oct 16, 2019
684b989
Merge branch 'master' into arbitrary-tracers-outer-loops
glwagner Oct 17, 2019
b3629f6
Updates tracer names for rayleigh benard regression test
glwagner Oct 17, 2019
d55247c
Updates ocean large eddy simulation regression for arbitrary tracers …
glwagner Oct 17, 2019
5d796b5
Fixes for ocean LES regression test for the GPU
glwagner Oct 17, 2019
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion benchmark/benchmark_static_ocean.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
75 changes: 75 additions & 0 deletions benchmark/benchmark_tracers.jl
Original file line number Diff line number Diff line change
@@ -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("")

2 changes: 2 additions & 0 deletions benchmark/benchmark_utils.jl
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
using Oceananigans: AbstractArchitecture

arch_name(::CPU) = "CPU"
arch_name(::GPU) = "GPU"

Expand Down
9 changes: 5 additions & 4 deletions examples/internal_wave.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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`,
Expand All @@ -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
#
Expand Down
15 changes: 8 additions & 7 deletions examples/ocean_convection_with_plankton.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
#
Expand All @@ -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)
Expand All @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion examples/simple_diffusion.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
13 changes: 11 additions & 2 deletions examples/two_dimensional_turbulence.jl
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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)
)

Expand Down Expand Up @@ -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)

Expand Down
3 changes: 2 additions & 1 deletion src/Oceananigans.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
48 changes: 23 additions & 25 deletions src/TurbulenceClosures/TurbulenceClosures.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,16 @@ module TurbulenceClosures
export
IsotropicDiffusivity,
ConstantIsotropicDiffusivity,

ConstantAnisotropicDiffusivity,

ConstantSmagorinsky,
SmagorinskyLilly,
BlasiusSmagorinsky,

AnisotropicMinimumDissipation,
RozemaAnisotropicMinimumDissipation,
VerstappenAnisotropicMinimumDissipation,

TurbulentDiffusivities,
calc_diffusivities!,
calculate_diffusivities!,

∇_κ_∇c,
∇_κ_∇T,
Expand All @@ -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

Expand All @@ -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")

Expand Down
Loading