Move density calculation to while model is running + inferred vertical diffusivity calculation

examples/twolayer_example.jl

@@ -3,9 +3,10 @@ using TwoLayerDirectNumericalShenanigans
 
 architecture = CPU() # or GPU()
 diffusivities = (ν = 1e-4, κ = (S = 1e-5, T = 1e-5))
+resolution = (Nx = 10, Ny = 10, Nz = 100)
 
 ## Setup the model
-model = DNS(architecture, DOMAIN_EXTENT, HIGH_RESOLUTION, diffusivities;
+model = DNS(architecture, DOMAIN_EXTENT, resolution, diffusivities;
             reference_density = REFERENCE_DENSITY)
 
 ## set initial conditions
@@ -18,17 +19,17 @@ profile_function = StepChange(transition_depth)#HyperbolicTangent(INTERFACE_LOCA
 tracer_perturbation_depth = find_depth(model, INTERFACE_LOCATION / 2)
 tracer_perturbation = SalinityGaussianProfile(tracer_perturbation_depth, 0.0, 1.5)
 noise_depth = find_depth(model, INTERFACE_LOCATION)
-initial_noise = SalinityNoise(noise_depth, 0.01)
+initial_noise = SalinityNoise(noise_depth, 1.0)
 
-dns = TwoLayerDNS(model, profile_function, initial_conditions)
+dns = TwoLayerDNS(model, profile_function, initial_conditions; initial_noise)
 
 set_two_layer_initial_conditions!(dns)
 
 ## build the simulation
 Δt = 1e-4
 stop_time = 60
 save_schedule = 0.5 # seconds
-simulation = DNS_simulation_setup(dns, Δt, stop_time, save_schedule)
+simulation = DNS_simulation_setup(dns, Δt, stop_time, save_schedule, save_velocities = true)
 
 ## Run the simulation
 run!(simulation)

src/TwoLayerDirectNumericalShenanigans.jl

@@ -1,8 +1,8 @@
 module TwoLayerDirectNumericalShenanigans
 
 using Oceananigans, Printf, Reexport, JLD2, Rasters, NCDatasets, GibbsSeaWater
+using Oceananigans: AbstractModel, Operators.ℑzᵃᵃᶠ
 using SeawaterPolynomials: TEOS10EquationOfState
-using Oceananigans: AbstractModel
 using SpecialFunctions: erf
 using Oceanostics: KineticEnergyDissipationRate
 using OceanRasterConversions: get_σₚ
@@ -61,6 +61,7 @@ include("initialconditions.jl")
 include("continuousprofilefunctions.jl")
 include("stepchangeprofilefunction.jl")
 include("tracerperturbations.jl")
+include("kernelfunctions.jl")
 include("twolayerdns.jl")
 include("noiseperturbations.jl")
 include("set_initialconditions.jl")

src/kernelfunctions.jl

@@ -0,0 +1,48 @@
+"""
+    C(i, j, k, grid, C)
+Get tracer `C` values for use in other function. There may be another way to do this for
+`KernelFunctionOperation`s but I have not found it so will use this for now. **Note:** this
+return the value of the tracer with no interpolation so if the tracer `C` is at
+`(Center, Center, Center)` the value extracted will be at (Center, Center, Center)`.
+"""
+C(i, j, k, grid, C) = C[i, j, k]
+"""
+    σ(i, j, k, grid, model, reference_pressure)
+Compute potential density `σ` at `reference_pressure` from salinity and temperature tracers
+in `model`.
+"""
+σ(i, j, k, grid, model, reference_pressure) = gsw_rho(C(i, j, k, grid, model.tracers.S),
+                                                      C(i, j, k, grid, model.tracers.T),
+                                                      reference_pressure)
+"""
+    DensityField(model, reference_pressure)
+Return an `KernelFunctionOperation` at `(Center, Center, Center)` that computes the
+potential density from the salinity and temperature tracers in `model` at `reference_pressure`.
+"""
+DensityField(model, reference_pressure) =
+    KernelFunctionOperation{Center, Center, Center}(σ, model.grid, model,
+                                                    reference_pressure)
+"""
+    σ_anomalyᶜᶜᶠ(i, j, k, grid, model, reference_pressure)
+Compute potential density anomaly at `reference_pressure` from salinity and temperature tracers
+in `model` that have been interpolated from `(Center, Center, Center)` to
+`(Center, Center, Face)`.
+"""
+σ_anomalyᶜᶜᶠ(i, j, k, grid, model, reference_pressure, σ_reference) =
+    gsw_rho(ℑzᵃᵃᶠ(i, j, k, grid, model.tracers.S), ℑzᵃᵃᶠ(i, j, k, grid, model.tracers.T),
+            reference_pressure) - σ_reference
+"""
+    function InterpolatedDensityAnomaly(model, reference_pressure)
+Return a `KernelFunctionOperation` at `(Center, Center, Face)` to compute the potential
+density anomaly from the salinity and temperature tracers from `model` that have been
+interpolated from `(Center, Center, Center)` to `(Center, Center, Face)`
+"""
+function InterpolatedDensityAnomaly(model, reference_pressure)
+
+    σ_reference = model.buoyancy.model.equation_of_state.reference_density
+    return KernelFunctionOperation{Center, Center, Face}(σ_anomalyᶜᶜᶠ,
+                                                         model.grid, model,
+                                                         reference_pressure,
+                                                         σ_reference)
+
+end

src/output_utils.jl

@@ -127,8 +127,9 @@ function field_ts_timemean(field_ts::FieldTimeSeries)
 
 end
 """
-    function predicted_maximum_density()
-Compute the predicted maximum density of water that can form along the mixing line.
+    function predicted_maximum_density
+Compute the predicted maximum density of water that can form along the mixing line between
+the salinity and temperature of the upper and lower layers.
 """
 function predicted_maximum_density!(simulation::Simulation, dns::TwoLayerDNS; reference_pressure = 0)
 
@@ -137,8 +138,15 @@ function predicted_maximum_density!(simulation::Simulation, dns::TwoLayerDNS; re
     S_mix = range(Sᵘ, Sˡ, step = 0.000001)
     T_mix = @. Tᵘ + slope * (Sᵘ - S_mix)
     ρ_mix_max = maximum(gsw_rho.(S_mix, T_mix, reference_pressure))
-    NCDataset(simulation.output_writers[:outputs].filepath, "a") do ds
-        ds.attrib["Predicted maximum density"] = ρ_mix_max
+    file_type = find_file_type(simulation.output_writers[:outputs].filepath)
+    if isequal(file_type, ".nc")
+        NCDataset(simulation.output_writers[:outputs].filepath, "a") do ds
+            ds.attrib["Predicted maximum density"] = ρ_mix_max
+        end
+    elseif isequal(file_type, ".jld2")
+        jldopen(simulation.output_writers[:outputs].filepath, "a+") do f
+            f["Predicted maximum density"] = ρ_mix_max
+        end
     end
 
     return nothing
@@ -208,15 +216,15 @@ function kolmogorov_and_batchelor_scale!(file::AbstractString)
     file_type = find_file_type(file)
     if isequal(file_type, ".nc")
 
-        ϵ = Raster(file, lazy = true, name = :ϵ)
-        ds = NCDataset(file, "a")
-        ν_str = ds.attrib["ν"]
-        ν = parse(Float64, ν_str[1:findfirst('m', ν_str)-1])
-        Sc = ds.attrib["Sc"]
-        η_min = minimum_η(ϵ; ν)
-        ds.attrib["η (min)"] = η_min # minimum space and time Kolmogorov scale
-        ds.attrib["λ_B"] = η_min / sqrt(Sc) # minimum space and time Batchelor scale
-        close(ds)
+        NCDataset(file, "a") do ds
+            ϵ = Raster(file, lazy = true, name = :ϵ)
+            ν_str = ds.attrib["ν"]
+            ν = parse(Float64, ν_str[1:findfirst('m', ν_str)-1])
+            Sc = ds.attrib["Sc"]
+            η_min = minimum_η(ϵ; ν)
+            ds.attrib["η (min)"] = η_min # minimum space and time Kolmogorov scale
+            ds.attrib["λ_B"] = η_min / sqrt(Sc) # minimum space and time Batchelor scale
+        end
 
     elseif isequal(file_type, ".jld2")
 
@@ -234,6 +242,35 @@ function kolmogorov_and_batchelor_scale!(file::AbstractString)
 
     return nothing
 
+end
+"""
+    function batchelor_scale!
+Compute and append the minimum space-time Batchelor scale and append it to save output.
+"""
+function batchelor_scale!(file::AbstractString)
+
+    file_type = find_file_type(file)
+    if isequal(file_type, ".nc")
+
+        NCDataset(file, "a") do ds
+            Sc = ds.attrib["Sc"]
+            η_min = minimum(ds["η_space"][:])
+            ds.attrib["λ_B"] = η_min / sqrt(Sc) # minimum space and time Batchelor scale
+        end
+
+    elseif isequal(file_type, ".jld2")
+
+        Sc = load(file, "Non_dimensional_numbers")["Sc"]
+        η = FieldTimeSeries(simulation.output_writers[:outputs].filepath, "η_space")
+        min_η = minimum(η)
+        jldopen(file, "a+") do f
+            f["minimum_batchelor_scale"] = min_η / sqrt(Sc)
+        end
+
+    end
+
+    return nothing
+
 end
 """
     function non_dimensional_numbers(simulation::Simulation, dns::TwoLayerDNS)
@@ -263,16 +300,16 @@ function non_dimensional_numbers!(simulation::Simulation, dns::TwoLayerDNS)
 
     if simulation.output_writers[:outputs] isa NetCDFOutputWriter
 
-        ds = NCDataset(simulation.output_writers[:outputs].filepath, "a")
-        ds.attrib["EOS"] = summary(model.buoyancy.model.equation_of_state.seawater_polynomial)
-        ds.attrib["Reference density"] = "$(model.buoyancy.model.equation_of_state.reference_density)kgm⁻³"
-        ds.attrib["ν"]  = "$(model.closure.ν) m²s⁻¹"
-        ds.attrib["κₛ"] = "$(model.closure.κ.S) m²s⁻¹"
-        ds.attrib["κₜ"] = "$(model.closure.κ.T) m²s⁻¹"
-        for key ∈ keys(nd_nums)
-            ds.attrib[key] = nd_nums[key]
+        NCDataset(simulation.output_writers[:outputs].filepath, "a") do ds
+            ds.attrib["EOS"] = summary(model.buoyancy.model.equation_of_state.seawater_polynomial)
+            ds.attrib["Reference density"] = "$(model.buoyancy.model.equation_of_state.reference_density)kgm⁻³"
+            ds.attrib["ν"]  = "$(model.closure.ν) m²s⁻¹"
+            ds.attrib["κₛ"] = "$(model.closure.κ.S) m²s⁻¹"
+            ds.attrib["κₜ"] = "$(model.closure.κ.T) m²s⁻¹"
+            for key ∈ keys(nd_nums)
+                ds.attrib[key] = nd_nums[key]
+            end
         end
-        close(ds)
 
     else
 

src/twolayerdns.jl

@@ -151,32 +151,69 @@ the course of a simulation;
 - `cfl` maximum cfl value used to determine the adaptive timestep size;
 - `diffusive_cfl` maximum diffusive cfl value used to determine the adaptive timestep size;
 - `max_change` maximum change in the timestep size;
-- `max_Δt` the maximum timestep.
+- `max_Δt` the maximum timestep;
+- `density_reference_pressure` for the seawater density calculation;
+- `save_velocities` defaults to `false`, if `true` model velocities will be saved to output.
 """
 function DNS_simulation_setup(dns::TwoLayerDNS, Δt::Number,
                               stop_time::Number, save_schedule::Number,
                               output_writer::Symbol=:netcdf;
                               cfl = 0.75,
                               diffusive_cfl = 0.75,
                               max_change = 1.2,
-                              max_Δt = 1e-1)
+                              max_Δt = 1e-1,
+                              density_reference_pressure = 0,
+                              save_velocities = false)
 
-    model, initial_conditions = dns.model, dns.initial_conditions
+    model = dns.model
     simulation = Simulation(model; Δt, stop_time)
 
     # time step adjustments
     wizard = TimeStepWizard(; cfl, diffusive_cfl, max_change, max_Δt)
     simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
 
-    # save output
+    # model tracers
+    S, T = model.tracers.S, model.tracers.T
+    # custom saved output
+
+    # Density
+    σ = DensityField(model, density_reference_pressure)
+
+    # Inferred vertical diffusivity
+    σ_anomaly_interpolated = InterpolatedDensityAnomaly(model, density_reference_pressure)
+    w = model.velocities.w
+    κᵥ = Integral((-w * σ_anomaly_interpolated) / σ)
+
+    # Minimum in space Kolmogorov length scale
     ϵ = KineticEnergyDissipationRate(model)
-    outputs = (S = model.tracers.S, T = model.tracers.T, ϵ = ϵ, w = model.velocities.w)
+    η_space(model) = minimum(model.closure.ν ./ ϵ)
+
+    # Dimensions and attributes for custom saved output
+    dims = Dict("η_space" => (), "σ" => ("xC", "xC", "zC"), "κᵥ" => ())
+    oa = Dict(
+        "σ" => Dict("longname" => "Seawater potential density calculated using TEOS-10 at $(density_reference_pressure)dbar",
+                    "units" => "kgm⁻³"),
+        "η_space" => Dict("longname" => "Minimum (in space) Kolmogorov length"),
+        "κᵥ" => Dict("longname" => "Inferred vertical diffusivity",
+                     "units" => "m²s⁻¹"))
+
+    # outputs to be saved during the simulation
+    outputs = Dict("S" => S, "T" => T, "η_space" => η_space, "σ" => σ, "κᵥ" => κᵥ)
+    if save_velocities
+        u, v = model.velocities.u, model.velocities.v
+        velocities = Dict("u" => u, "v" => v, "w" => w)
+        merge!(outputs, velocities)
+    end
+
     filename = form_filename(dns, stop_time, output_writer)
     simulation.output_writers[:outputs] = output_writer == :netcdf ?
                                             NetCDFOutputWriter(model, outputs,
                                                             filename = filename,
                                                             schedule = TimeInterval(save_schedule),
-                                                            overwrite_existing = true) :
+                                                            overwrite_existing = true,
+                                                            dimensions = dims,
+                                                            output_attributes = oa
+                                                            ) :
                                             JLD2OutputWriter(model, outputs,
                                                             filename = filename,
                                                             schedule = TimeInterval(save_schedule),
@@ -198,8 +235,6 @@ Create a filename for saved output based on the `profile_function`,`initial_cond
 """
 function form_filename(dns::TwoLayerDNS, stop_time::Number, output_writer::Symbol)
 
-    pf_string = dns.profile_function isa HyperbolicTangent ? "tanh" :
-                            dns.profile_function isa Erf ? "erf" : "midpoint"
     pf_string = lowercase(string(typeof(dns.profile_function))[1:findfirst('{', string(typeof(dns.profile_function))) - 1])
     ic_type = typeof(dns.initial_conditions)
     ic_string = ic_type <: StableTwoLayerInitialConditions ? "stable" :

test/kernelfunction_test.jl

@@ -0,0 +1,46 @@
+using TwoLayerDirectNumericalShenanigans: DensityField, InterpolatedDensityAnomaly
+using Oceananigans: Operators.ℑzᵃᵃᶠ
+using GibbsSeaWater
+
+architecture = CPU()
+diffusivities = (ν = 1e-4, κ = (S = 1e-5, T = 1e-5))
+resolution = (Nx = 10, Ny = 10, Nz = 100)
+
+model = DNS(architecture, DOMAIN_EXTENT, resolution, diffusivities;
+            reference_density = REFERENCE_DENSITY)
+
+z = znodes(model.grid, Center())
+
+T₀ᵘ = -1.5
+S₀ᵘ = 34.568
+cabbeling = StableUpperLayerInitialConditions(S₀ᵘ, T₀ᵘ)
+initial_conditions = TwoLayerInitialConditions(cabbeling)
+transition_depth = find_depth(model, INTERFACE_LOCATION)
+profile_function = StepChange(transition_depth)
+
+dns = TwoLayerDNS(model, profile_function, initial_conditions)
+
+set_two_layer_initial_conditions!(dns)
+
+reference_pressure = 0
+σ_field = Field(DensityField(dns.model, 0))
+compute!(σ_field)
+
+function test_density_profile(σ_field, computed_density_profile)
+
+    vertical_slice = interior(σ_field, rand(1:10), rand(1:10), :)
+
+    return vertical_slice .== computed_density_profile
+
+end
+
+σ_anomaly_interp_field = Field(InterpolatedDensityAnomaly(dns.model, reference_pressure))
+compute!(σ_anomaly_interp_field)
+
+function test_density_anom_profile(σ_anomaly_interp_field, computed_density_profile)
+
+    vertical_slice = interior(σ_anomaly_interp_field, rand(1:10), rand(1:10), :)
+
+    return vertical_slice .== computed_density_profile
+
+end

test/runtests.jl

@@ -1,7 +1,7 @@
 using TwoLayerDirectNumericalShenanigans, Test
 using TwoLayerDirectNumericalShenanigans: perturb_tracer
 
-include("twolayercontainers.jl")
+include("twolayercontainers_test.jl")
 
 @testset "TwoLayerDNS upper layer containers" begin
     for (i, uc) ∈ enumerate(upper_layer_containers)
@@ -128,3 +128,37 @@ end
     test_depth_range = vcat(z[10], z[20])
     @test isequal(z[10:20], find_depth(model, test_depth_range))
 end
+
+include("kernelfunction_test.jl")
+
+@testset "Density field computation" begin
+    d_idx = findfirst(z .== transition_depth)
+    σᵘ = gsw_rho.(interior(model.tracers.S, rand(1:10), rand(1:10), 1:d_idx),
+                  interior(model.tracers.T, rand(1:10), rand(1:10), 1:d_idx),
+                  reference_pressure)
+    σˡ = gsw_rho.(interior(model.tracers.S, rand(1:10), rand(1:10), d_idx+1:100),
+                  interior(model.tracers.T, rand(1:10), rand(1:10), d_idx+1:100),
+                  reference_pressure)
+    σ_computed_profile = vcat(σᵘ, σˡ)
+    @test isequal(trues(length(z)), test_density_profile(σ_field, σ_computed_profile))
+end
+
+@testset "Interpolated density anomaly" begin
+    d_idx = findfirst(z .== transition_depth)
+    S_interpolation = Field(KernelFunctionOperation{Center, Center, Face}(ℑzᵃᵃᶠ, model.grid,
+                                                                          model.tracers.S))
+    compute!(S_interpolation)
+    T_interpolation = Field(KernelFunctionOperation{Center, Center, Face}(ℑzᵃᵃᶠ, model.grid,
+                                                                          model.tracers.T))
+    compute!(T_interpolation)
+    σ_reference = dns.model.buoyancy.model.equation_of_state.reference_density
+    σᵘ = gsw_rho.(interior(S_interpolation, rand(1:10), rand(1:10), 1:d_idx),
+                  interior(T_interpolation, rand(1:10), rand(1:10), 1:d_idx),
+                  reference_pressure) .- σ_reference
+    σˡ = gsw_rho.(interior(S_interpolation, rand(1:10), rand(1:10), d_idx+1:101),
+                  interior(T_interpolation, rand(1:10), rand(1:10), d_idx+1:101),
+                  reference_pressure) .- σ_reference
+    σ_anom_computed_profile = vcat(σᵘ, σˡ)
+    @test isequal(trues(length(σ_anom_computed_profile)),
+                  test_density_anom_profile(σ_anomaly_interp_field, σ_anom_computed_profile))
+end