Better folder naming, R_rho computation in initial conditions, add save_custom_output

Project.toml

@@ -4,6 +4,9 @@ authors = ["Josef Bisits <jbisits@gmail.com>"]
 version = "0.0.1"
 
 [deps]
+GibbsSeaWater = "9a22fb26-0b63-4589-b28e-8f9d0b5c3d05"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"

examples/four_stair_staircase.jl

@@ -1,7 +1,7 @@
 using StaircaseShenanigans
 
 architecture = CPU() # or GPU()
-diffusivities = (ν = 1e-4, κ = (S = 1e-5, T = 1e-5))
+diffusivities = (ν = 1e-4, κ = (S = 1e-7, T = 1e-5))
 domain_extent = (Lx = 0.1, Ly = 0.1, Lz = -1.0)
 resolution = (Nx = 10, Ny = 10, Nz = 100)
 
@@ -11,11 +11,17 @@ model = DNSModel(architecture, domain_extent, resolution, diffusivities)
 ## Initial conditions
 number_of_steps = 4
 depth_of_steps = [-0.2, -0.4, -0.6, -0.8]
-salinity = [34.58, 34.6, 34.62, 34.64, 34.66]
-temperature = [-1.5, -1.0, -0.5, 0.0, 0.5]
+salinity = [34.57, 34.6, 34.63, 34.66, 34.69]
+temperature = [-1.5, -1.45, -1.4, -1.35, -1.3]
 
 step_ics = StepInitialConditions(number_of_steps, depth_of_steps, salinity, temperature)
 
 sdns = StaircaseDNS(model, step_ics)
 
 set_staircase_initial_conditions!(sdns)
+
+Δt = 1e-2
+stop_time = 60 * 60 # seconds
+save_schedule = 60  # seconds
+output_path = joinpath(@__DIR__, "output")
+simulation = SDNS_simulation_setup(sdns, Δt, stop_time, save_schedule, save_computed_output!; output_path)

src/StaircaseShenanigans.jl

@@ -1,9 +1,11 @@
 module StaircaseShenanigans
 
 using Oceananigans, Reexport, Printf
+using Oceananigans: seawater_density
 using SeawaterPolynomials.TEOS10
 using SeawaterPolynomials.SecondOrderSeawaterPolynomials
 using SeawaterPolynomials: BoussinesqEquationOfState
+using NCDatasets, JLD2
 
 @reexport using Oceananigans, SeawaterPolynomials.TEOS10,
                 SeawaterPolynomials.SecondOrderSeawaterPolynomials
@@ -16,6 +18,8 @@ export StaircaseDNS, DNSModel, SDNS_simulation_setup
 
 export StepInitialConditions, SmoothStepInitialConditions, set_staircase_initial_conditions!
 
+export save_computed_output!
+
 include("staircase_model.jl")
 include("staircase_initial_conditions.jl")
 include("set_staircase_initial_conditions.jl")

src/staircase_initial_conditions.jl

@@ -10,6 +10,22 @@ struct StepInitialConditions{T} <: AbstractStaircaseInitialConditions
        salinity_values :: T
     "Temperature values in each layer"
     temperature_values :: T
+    "Initial R_ρ at each step interface"
+                   R_ρ :: T
+end
+function StepInitialConditions(number_of_steps, depth_of_steps, salinity, temperature)
+
+    ΔT = diff(temperature)
+    ΔS = diff(salinity)
+    # Constant values take from Vallis, not sure if this is the best way to calculate
+    # these but will start with this.
+    α = 1.67e-4
+    β = 7.80e-4
+
+    R_ρ = @. (β * ΔS) / (α * ΔT)
+
+    return StepInitialConditions(number_of_steps, depth_of_steps, salinity, temperature, R_ρ)
+
 end
 """
     function set_steps(z, C, depth_of_steps)
@@ -55,7 +71,8 @@ function Base.show(io::IO, sics::AbstractStaircaseInitialConditions)
         println(io, "┣━━━━ number_of_steps: $(sics.number_of_steps)")
         println(io, "┣━━━━━ depth_of_steps: $(sics.depth_of_steps)")
         println(io, "┣━━━━ salinity_values: $(sics.salinity_values)")
-        print(io,   "┗━ temperature_values: $(sics.temperature_values)")
+        println(io, "┣━ temperature_values: $(sics.temperature_values)")
+        print(io,   "┗━━━━━━━━━━━━━━━━ R_ρ: $(round.(sics.R_ρ; digits = 2))")
     elseif sics isa SmoothStepInitialConditions
         println(io, "StepInitialConditions")
         println(io, "┣━━━━ number_of_steps: $(sics.number_of_steps)")

src/staircase_model.jl

@@ -19,7 +19,7 @@ StaircaseDNS(model, initial_conditions; initial_noise = nothing) =
 Base.iterate(sdns::AbstractStaircaseModel, state = 1) =
     state > length(fieldnames(sdns)) ? nothing : (getfield(sdns, state), state + 1)
 
-# Same DNS setup as in TLDNS but repeated here for easier changes that may be specific to
+# Same DNS setup as in sdns but repeated here for easier changes that may be specific to
 # setting up staircase experiments.
 """
     function DNSModel(architecture, domain_extent::NamedTuple, resolution::NamedTuple,
@@ -101,7 +101,10 @@ function grid_stretching(Lz::Number, Nz::Number, refinement::Number, stretching:
     return z_faces
 
 end
-
+"""
+    function SDNS_simulation_setup
+Build a `simulation` from `sdns`.
+"""
 function SDNS_simulation_setup(sdns::StaircaseDNS, Δt::Number,
                                 stop_time::Number, save_schedule::Number,
                                 save_custom_output!::Function=no_custom_output!;
@@ -140,8 +143,29 @@ function SDNS_simulation_setup(sdns::StaircaseDNS, Δt::Number,
     # Checkpointer setup
     checkpointer_setup!(simulation, model, output_dir, checkpointer_time_interval)
 
+    save_R_ρ!(simulation, sdns)
+
     return simulation
 
+end
+"""
+    save_R_ρ!(simulation::Simulation, sdns::StaircaseDNS)
+To the output file.
+"""
+function save_R_ρ!(simulation::Simulation, sdns::StaircaseDNS)
+
+    if simulation.output_writers[:tracers] isa NetCDFOutputWriter
+        NCDataset(simulation.output_writers[:computed_output].filepath, "a") do ds
+            ds.attrib["Step interface R_ρ"] = sdns.initial_conditions.R_ρ
+        end
+    elseif simulation.output_writers[:tracers] isa JLD2OutputWriter
+        jldopen(simulation.output_writers[:computed_output].filepath, "a+") do f
+            f["Step interface R_ρ"] = sdns.initial_conditions.R_ρ
+        end
+    end
+
+    return nothing
+
 end
 """
     function output_directory(sdns::StaircaseDNS, stop_time::Number, output_path)
@@ -153,9 +177,11 @@ function output_directory(sdns::StaircaseDNS, stop_time::Number, output_path)
     ic_type = typeof(sdns.initial_conditions)
     ic_string = ic_type <: StepInitialConditions ? "step_change" : "smoothed_step"
 
+    eos_string = is_linear_eos(sdns.model.buoyancy.model.equation_of_state.seawater_polynomial)
+
     stop_time_min = stop_time / 60 ≥ 1 ? string(round(Int, stop_time / 60)) :
                                          string(round(stop_time / 60; digits = 2))
-    expt_dir = ic_string *"_"* stop_time_min * "min"
+    expt_dir = eos_string *"_"* ic_string *"_"* stop_time_min * "min"
 
     output_dir = mkpath(joinpath(output_path, expt_dir))
     # make a simulation directory if one is not present
@@ -166,6 +192,14 @@ function output_directory(sdns::StaircaseDNS, stop_time::Number, output_path)
     return output_dir
 
 end
+is_linear_eos(eos::TEOS10SeawaterPolynomial) = "nonlineareos"
+function is_linear_eos(eos::SecondOrderSeawaterPolynomial)
+
+    non_linear_coefficients = (eos.R₀₁₁, eos.R₀₂₀, eos.R₁₀₁, eos.R₁₁₀, eos.R₂₀₀)
+    eos_type = all(non_linear_coefficients .== 0) ? "lineareos" : "nonlineareos"
+
+    return eos_type
+end
 """
     function save_tracers!(simulation, model, save_schedule, save_file, output_dir, overwrite_saved_output)
 Save `model.tracers` during a `Simulation` using an `OutputWriter`.
@@ -223,7 +257,44 @@ function save_velocities!(simulation, model, save_schedule, save_file, output_di
 end
 save_velocities!(simulation, model, save_info::Tuple) =
     save_velocities!(simulation, model, save_info...)
-"Default function for `save_custom_output!` in `TLDNS_simulation_setup`."
+    """
+    function save_computed_output!(simulation, sdns, save_schedule, save_file, output_dir,
+                                   overwrite_saved_output, reference_gp_height)
+Save selection of computed output during a `Simulation` using an `OutputWriter`.
+"""
+function save_computed_output!(simulation, sdns, save_schedule, save_file, output_dir,
+                               overwrite_saved_output, reference_gp_height)
+
+    model = sdns.model
+    σ = seawater_density(model, geopotential_height = reference_gp_height)
+    computed_outputs = Dict("σ" => σ)
+
+    oa = Dict(
+        "σ" => Dict("longname" => "Seawater potential density calculated using TEOS-10 at $(reference_gp_height)dbar",
+                    "units" => "kgm⁻³")
+        )
+
+    simulation.output_writers[:computed_output] =
+        save_file == :netcdf ? NetCDFOutputWriter(model, computed_outputs;
+                                                filename = "computed_output",
+                                                dir = output_dir,
+                                                overwrite_existing = overwrite_saved_output,
+                                                schedule = TimeInterval(save_schedule),
+                                                output_attributes = oa
+                                                ) :
+                                JLD2OutputWriter(model, computed_outputs;
+                                                filename = "computed_output",
+                                                dir = output_dir,
+                                                schedule = TimeInterval(save_schedule),
+                                                overwrite_existing = overwrite_saved_output)
+
+
+    return nothing
+
+end
+save_computed_output!(simulation, sdns, save_info::Tuple; reference_gp_height = 0) =
+    save_computed_output!(simulation, sdns, save_info..., reference_gp_height)
+"Default function for `save_custom_output!` in `sdns_simulation_setup`."
 no_custom_output!(simulation, model, save_info...) = nothing
 """
     function checkpointer_setup!(simulation, model, output_path, checkpointer_time_interval)