Skip to content

Commit

Permalink
Merge pull request #85 from ali-ramadhan/forcings
Browse files Browse the repository at this point in the history
Support for user-defined forcing functions.
  • Loading branch information
ali-ramadhan committed Feb 27, 2019
2 parents f7ba1ce + 668eb27 commit 9d43e79
Show file tree
Hide file tree
Showing 7 changed files with 175 additions and 119 deletions.
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)
@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

0 comments on commit 9d43e79

Please sign in to comment.