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

Support for user-defined forcing functions. #85

Merged
merged 5 commits into from
Feb 27, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
14 changes: 7 additions & 7 deletions Manifest.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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]]
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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]]
Expand Down Expand Up @@ -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]]
Expand Down
211 changes: 109 additions & 102 deletions examples/deep_convection_3d.jl
Original file line number Diff line number Diff line change
@@ -1,102 +1,109 @@
using Statistics: mean
using Plots
using Oceananigans

function deep_convection_3d()
Nx, Ny, Nz = 100, 100, 50
Lx, Ly, Lz = 2000, 2000, 1000
Nt, Δt = 500, 20

model = Model((Nx, Ny, Nz), (Lx, Ly, Lz))
impose_initial_conditions!(model)

nc_writer = NetCDFFieldWriter(".", "deep_convection_3d", 50)
push!(model.output_writers, nc_writer)

time_step!(model; Nt=Nt, Δt=Δt)
# make_temperature_movies(model, field_writer)
end

function impose_cooling_disk!(model::Model)
g = model.grid
c = model.constants

# Parameters for generating initial surface heat flux.
Rc = 600 # Radius of cooling disk [m].
Ts = 20 # Surface temperature [°C].
Q₀ = -800 # Cooling disk heat flux [W/m²].
Q₁ = 10 # Noise added to cooling disk heat flux [W/m²].
Ns = 5 * (c.f * Rc/g.Lz) # Stratification or Brunt–Väisälä frequency [s⁻¹].

αᵥ = 2.07e-4 # Volumetric coefficient of thermal expansion for water [K⁻¹].
cᵥ = 4181.3 # Isobaric mass heat capacity [J / kg·K].

Tz = Ns^2 / (c.g * αᵥ) # Vertical temperature gradient [K/m].

# Center horizontal coordinates so that (x,y) = (0,0) corresponds to the center
# of the domain (and the cooling disk).
x₀ = g.xC .- mean(g.xC)
y₀ = g.yC .- mean(g.yC)

# Calculate vertical temperature profile and convert to Kelvin.
T_ref = 273.15 .+ Ts .+ Tz .* (g.zC .- mean(Tz * g.zC))

# Impose reference temperature profile.
model.tracers.T.data .= repeat(reshape(T_ref, 1, 1, g.Nz), g.Nx, g.Ny, 1)

# Set surface heat flux to zero outside of cooling disk of radius Rᶜ.
r₀² = @. x₀*x₀ + y₀'*y₀'

# Generate surface heat flux field with small random fluctuations.
Q = Q₀ .+ Q₁ * (0.5 .+ rand(g.Nx, g.Ny))
Q[findall(r₀² .> Rc^2)] .= 0 # Set Q=0 outside cooling disk.

# Convert surface heat flux into 3D forcing term for use when calculating
# 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))
nothing
end

function impose_initial_conditions!(model::Model)
g = model.grid

impose_cooling_disk!(model)

@. model.tracers.S.data = model.eos.S₀

pHY_profile = [-model.eos.ρ₀ * model.constants.g * h for h in g.zC]
model.pressures.pHY.data .= repeat(reshape(pHY_profile, 1, 1, g.Nz), g.Nx, g.Ny, 1)
nothing
end

# function make_temperature_movies(model::Model, fw::FieldWriter)
# n_frames = Int(model.clock.time_step / fw.output_frequency)
#
# xC, yC, zC = model.grid.xC, model.grid.yC, model.grid.zC
# Δt = 20
#
# print("Creating temperature movie... ($n_frames frames)\n")
#
# Plots.gr()
# default(dpi=300)
# movie = @animate for tidx in 0:n_frames
# print("\rframe = $tidx / $n_frames ")
# temperature = read_output(model, fw, "T", tidx*fw.output_frequency*Δt)
# Plots.heatmap(xC, zC, rotl90(temperature[:, 50, :]) .- 293.15, color=:balance,
# clims=(-0.1, 0),
# title="T @ t=$(tidx*fw.output_frequency*Δt)")
# end
#
# mp4(movie, "deep_convection_3d_$(round(Int, time())).mp4", fps = 30)
#
# print("Creating surface temperature movie... ($n_frames frames)\n")
# movie = @animate for tidx in 0:n_frames
# print("\rframe = $tidx / $n_frames ")
# temperature = read_output(model, fw, "T", tidx*fw.output_frequency*Δt)
# Plots.heatmap(xC, yC, temperature[:, :, 1] .- 293.15, color=:balance,
# clims=(-0.1, 0),
# title="T @ t=$(tidx*fw.output_frequency*Δt)")
# end
# mp4(movie, "deep_convection_3d_$(round(Int, time())).mp4", fps = 30)
# end
using Statistics: mean
using Plots
using Oceananigans

function impose_cooling_disk!(model::Model)
g = model.grid
c = model.constants

# Parameters for generating initial surface heat flux.
Rc = 600 # Radius of cooling disk [m].
Ts = 20 # Surface temperature [°C].
Q₀ = -800 # Cooling disk heat flux [W/m²].
Q₁ = 10 # Noise added to cooling disk heat flux [W/m²].
Ns = 5 * (c.f * Rc/g.Lz) # Stratification or Brunt–Väisälä frequency [s⁻¹].

αᵥ = 2.07e-4 # Volumetric coefficient of thermal expansion for water [K⁻¹].
cᵥ = 4181.3 # Isobaric mass heat capacity [J / kg·K].

Tz = Ns^2 / (c.g * αᵥ) # Vertical temperature gradient [K/m].

# Center horizontal coordinates so that (x,y) = (0,0) corresponds to the center
# of the domain (and the cooling disk).
x₀ = g.xC .- mean(g.xC)
y₀ = g.yC .- mean(g.yC)

# Calculate vertical temperature profile and convert to Kelvin.
T_ref = 273.15 .+ Ts .+ Tz .* (g.zC .- mean(Tz * g.zC))

# Impose reference temperature profile.
model.tracers.T.data .= repeat(reshape(T_ref, 1, 1, g.Nz), g.Nx, g.Ny, 1)

# Set surface heat flux to zero outside of cooling disk of radius Rᶜ.
r₀² = @. x₀*x₀ + y₀'*y₀'

# Generate surface heat flux field with small random fluctuations.
Q = Q₀ .+ Q₁ * (0.5 .+ rand(g.Nx, g.Ny))
Q[findall(r₀² .> Rc^2)] .= 0 # Set Q=0 outside cooling disk.

# Convert surface heat flux into 3D forcing term for use when calculating
# 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

function impose_initial_conditions!(model::Model)
g = model.grid

impose_cooling_disk!(model)

@. model.tracers.S.data = model.eos.S₀

pHY_profile = [-model.eos.ρ₀ * model.constants.g * h for h in g.zC]
model.pressures.pHY.data .= repeat(reshape(pHY_profile, 1, 1, g.Nz), g.Nx, g.Ny, 1)
nothing
end

# function make_temperature_movies(model::Model, fw::FieldWriter)
# n_frames = Int(model.clock.time_step / fw.output_frequency)
#
# xC, yC, zC = model.grid.xC, model.grid.yC, model.grid.zC
# Δt = 20
#
# print("Creating temperature movie... ($n_frames frames)\n")
#
# Plots.gr()
# default(dpi=300)
# movie = @animate for tidx in 0:n_frames
# print("\rframe = $tidx / $n_frames ")
# temperature = read_output(model, fw, "T", tidx*fw.output_frequency*Δt)
# Plots.heatmap(xC, zC, rotl90(temperature[:, 50, :]) .- 293.15, color=:balance,
# clims=(-0.1, 0),
# title="T @ t=$(tidx*fw.output_frequency*Δt)")
# end
#
# mp4(movie, "deep_convection_3d_$(round(Int, time())).mp4", fps = 30)
#
# print("Creating surface temperature movie... ($n_frames frames)\n")
# movie = @animate for tidx in 0:n_frames
# print("\rframe = $tidx / $n_frames ")
# temperature = read_output(model, fw, "T", tidx*fw.output_frequency*Δt)
# Plots.heatmap(xC, yC, temperature[:, :, 1] .- 293.15, color=:balance,
# clims=(-0.1, 0),
# title="T @ t=$(tidx*fw.output_frequency*Δt)")
# end
# mp4(movie, "deep_convection_3d_$(round(Int, time())).mp4", fps = 30)
# end

Nx, Ny, Nz = 100, 100, 50
Lx, Ly, Lz = 2000, 2000, 1000
Nt, Δt = 1000, 20

model = Model((Nx, Ny, Nz), (Lx, Ly, Lz))

# 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 surface_cooling_disk(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) = ifelse(k == 1 && 20 < i < 80 && 20 < j < 80, -9e-6, 0)

# 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)
# make_temperature_movies(model, field_writer)
3 changes: 3 additions & 0 deletions src/Oceananigans.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@ export
OperatorTemporaryFields,
StepperTemporaryFields,

Forcing,

LinearEquationOfState,
ρ!,
δρ!,
Expand Down Expand Up @@ -104,6 +106,7 @@ include("planetary_constants.jl")
include("grids.jl")
include("fields.jl")
include("fieldsets.jl")
include("forcing.jl")

include("operators/operators.jl")

Expand Down
28 changes: 28 additions & 0 deletions src/forcing.jl
Original file line number Diff line number Diff line change
@@ -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)
7 changes: 5 additions & 2 deletions src/model.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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)

Expand Down Expand Up @@ -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
14 changes: 7 additions & 7 deletions src/time_steppers.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand All @@ -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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens if you try to pass grid instead of Nx, Ny, Nz, Δx, Δy, Δz?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the GPU it complains that the arguments are not isbitstype :(

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it. So both Fields and Grids will have to be worked on.

@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]
Expand Down
Loading