From e94cf92d7134fd3d79e0b65d10a34b01de05090b Mon Sep 17 00:00:00 2001 From: Ali Ramadhan Date: Tue, 26 Feb 2019 16:42:54 -0500 Subject: [PATCH 1/5] Fixed call to BoundaryConditions constructor. Resolves #84 --- src/model.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/model.jl b/src/model.jl index 47b3315c69..a29391c30b 100644 --- a/src/model.jl +++ b/src/model.jl @@ -22,7 +22,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() From f493a7f169430fd360eee8cf2653a7ceb25a348b Mon Sep 17 00:00:00 2001 From: Ali Ramadhan Date: Tue, 26 Feb 2019 17:27:57 -0500 Subject: [PATCH 2/5] Hacky support for user-defined forcing functions. --- src/Oceananigans.jl | 3 +++ src/forcing.jl | 28 ++++++++++++++++++++++++++++ src/model.jl | 5 ++++- src/time_steppers.jl | 14 +++++++------- 4 files changed, 42 insertions(+), 8 deletions(-) create mode 100644 src/forcing.jl diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 6e934504cb..f72adcb5ca 100644 --- a/src/Oceananigans.jl +++ b/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") diff --git a/src/forcing.jl b/src/forcing.jl new file mode 100644 index 0000000000..1746531198 --- /dev/null +++ b/src/forcing.jl @@ -0,0 +1,28 @@ +struct Forcing{Tu,Tv,Tw,TT,TS} + u::Tu + v::Tv + w::Tw + T::TT + S::TS +end + +@inline zero_func(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) = 0 + +function Forcing(Tu, Tv, Tw, TT, TS) + if Tu == nothing + Tu = zero_func + end + if Tv == nothing + Tv = zero_func + end + if Tw == nothing + Tw = zero_func + end + if TT == nothing + TT = zero_func + end + if TS == nothing + TS = zero_func + end + Forcing{typeof(Tu),typeof(Tv),typeof(Tw),typeof(TT),typeof(Tu)}(Tu, Tv, Tw, TT, TS) +end diff --git a/src/model.jl b/src/model.jl index a29391c30b..064dd01aca 100644 --- a/src/model.jl +++ b/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 @@ -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 diff --git a/src/time_steppers.jl b/src/time_steppers.jl index 84d48ed189..6c54931f2a 100644 --- a/src/time_steppers.jl +++ b/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] From b38c2b56f65be4223298d043af13871e1de35755 Mon Sep 17 00:00:00 2001 From: Ali Ramadhan Date: Tue, 26 Feb 2019 17:58:53 -0500 Subject: [PATCH 3/5] Hacked in a user-defined forcing function into the 3D deep convection example. --- examples/deep_convection_3d.jl | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/examples/deep_convection_3d.jl b/examples/deep_convection_3d.jl index 9bd6606648..51977adaf9 100644 --- a/examples/deep_convection_3d.jl +++ b/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 From 93a89e784a3445975906df8c58fb15b63db3ee2e Mon Sep 17 00:00:00 2001 From: Gregory LeClaire Wagner Date: Tue, 26 Feb 2019 23:10:51 -0500 Subject: [PATCH 4/5] Adds forcing tests and tweaks forcing constructors --- Manifest.toml | 14 +++++++------- src/forcing.jl | 48 ++++++++++++++++++++++++------------------------ test/runtests.jl | 17 ++++++++++++++++- 3 files changed, 47 insertions(+), 32 deletions(-) diff --git a/Manifest.toml b/Manifest.toml index 89b3523708..dcbc48026a 100644 --- a/Manifest.toml +++ b/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]] diff --git a/src/forcing.jl b/src/forcing.jl index 1746531198..9520101592 100644 --- a/src/forcing.jl +++ b/src/forcing.jl @@ -1,28 +1,28 @@ -struct Forcing{Tu,Tv,Tw,TT,TS} - u::Tu - v::Tv - w::Tw - T::TT - S::TS -end +"Dummy function and forcing default." +@inline zero_func(args...) = 0 -@inline zero_func(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) = 0 +""" + Forcing(Fu, Fv, Fw, FF, FS) -function Forcing(Tu, Tv, Tw, TT, TS) - if Tu == nothing - Tu = zero_func - end - if Tv == nothing - Tv = zero_func - end - if Tw == nothing - Tw = zero_func - end - if TT == nothing - TT = zero_func - end - if TS == nothing - TS = zero_func + 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 - Forcing{typeof(Tu),typeof(Tv),typeof(Tw),typeof(TT),typeof(Tu)}(Tu, Tv, Tw, TT, TS) end + +Forcing(; Fu=nothing, Fv=nothing, Fw=nothing, FT=nothing, FS=nothing) = Forcing(Fu, Fv, Fw, FT, FS) diff --git a/test/runtests.jl b/test/runtests.jl index ddbd3b232d..be27fc6b79 100644 --- a/test/runtests.jl +++ b/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 From 668eb27bcc1af07cd187f79374cd13f2e522cc52 Mon Sep 17 00:00:00 2001 From: ali-ramadhan Date: Wed, 27 Feb 2019 15:25:53 -0500 Subject: [PATCH 5/5] Refactored deep convection 3D example to use ifelse forcing. --- examples/deep_convection_3d.jl | 226 ++++++++++++++++----------------- 1 file changed, 109 insertions(+), 117 deletions(-) diff --git a/examples/deep_convection_3d.jl b/examples/deep_convection_3d.jl index 51977adaf9..c13ecb8b50 100644 --- a/examples/deep_convection_3d.jl +++ b/examples/deep_convection_3d.jl @@ -1,117 +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 = 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 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) - # 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)) - @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 +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) \ No newline at end of file