Merge pull request #94 from jbisits/jib-updateinsetting

Simplify initial condition setting
jbisits · Dec 5, 2024 · a9d88c2 · a9d88c2
2 parents 321879d + c02f9f0
commit a9d88c2
Show file tree

Hide file tree

Showing 18 changed files with 424 additions and 617 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "StaircaseShenanigans"
 uuid = "c2bb06a8-94f3-4279-b990-30bf3ab8ba6f"
 authors = ["Josef Bisits <jbisits@gmail.com>"]
-version = "0.5.0"
+version = "0.6.0"
 
 [deps]
 GibbsSeaWater = "9a22fb26-0b63-4589-b28e-8f9d0b5c3d05"

diff --git a/examples/single_interface.jl b/examples/single_interface.jl
@@ -3,9 +3,10 @@ using StaircaseShenanigans
 architecture = CPU() # or GPU()
 diffusivities = (ν = 1e-5, κ = (S = 1e-7, T = 1e-5))
 domain_extent = (Lx = 0.1, Ly = 0.1, Lz = -1.0)
+domain_topology = (x = Periodic, y = Periodic, z = Bounded)
 resolution = (Nx = 5, Ny = 5, Nz = 50)
 eos = CustomLinearEquationOfState(-0.5, 34.6)
-model_setup = (;architecture, diffusivities, domain_extent, resolution, eos)
+model_setup = (;architecture, diffusivities, domain_extent, domain_topology, resolution, eos)
 
 ## Initial conditions
 depth_of_interface = -0.5

diff --git a/examples/single_interface_periodic_tanh_background.jl b/examples/single_interface_periodic_tanh_background.jl
@@ -3,15 +3,16 @@ using StaircaseShenanigans
 architecture = CPU() # or GPU()
 diffusivities = (ν = 1e-5, κ = (S = 1e-7, T = 1e-5))
 domain_extent = (Lx = 0.1, Ly = 0.1, Lz = -1.0)
+domain_topology = (x = Periodic, y = Periodic, z = Periodic)
 resolution = (Nx = 5, Ny = 5, Nz = 50)
 eos = CustomLinearEquationOfState(-0.5, 34.6)
-model_setup = (;architecture, diffusivities, domain_extent, resolution, eos)
+model_setup = (;architecture, diffusivities, domain_extent, domain_topology, resolution, eos)
 
 ## Initial conditions
 depth_of_interface = -0.5
 salinity = [34.56, 34.70]
 temperature = [-1.5, 0.5]
-interface_ics = PeriodoicSingleInterfaceICs(eos, depth_of_interface, salinity, temperature)
+interface_ics = SingleInterfaceICs(eos, depth_of_interface, salinity, temperature)
 tracer_noise = TracerNoise(1e-6, 1e-6)
 
 ## setup model

diff --git a/src/StaircaseShenanigans.jl b/src/StaircaseShenanigans.jl
@@ -29,7 +29,6 @@ export StaircaseDNS, DNSModel, SDNS_simulation_setup
 
 export STStaircaseInitialConditions, StaircaseICs, SmoothSTStaircaseInitialConditions,
        STSingleInterfaceInitialConditions, SingleInterfaceICs,
-       PeriodicSTSingleInterfaceInitialConditions, PeriodoicSingleInterfaceICs,
        set_initial_conditions!
 
 export TanhInterfaceSmoothing, Tanh, NoSmoothing
@@ -49,14 +48,15 @@ export compute_R_ρ!
 export animate_tracers, animate_density, visualise_initial_conditions, visualise_initial_density,
        animate_tracers_anomaly, animate_density_anomaly, animate_profile_in_S_Θ_space
 
-include("interface_smoothing.jl")
 include("staircase_initial_conditions.jl")
 include("single_interfaces_initial_conditions.jl")
+include("interface_smoothing.jl")
+include("staircase_background.jl")
 include("staircase_noise.jl")
 include("staircase_diagnostics.jl")
 include("staircase_model.jl")
+include("staircase_simulation.jl")
 include("set_initial_conditions.jl")
-include("staircase_background.jl")
 include("staircase_restoring.jl")
 include("alt_lineareos.jl")
 include("makie_plotting_functions.jl")

diff --git a/src/set_initial_conditions.jl b/src/set_initial_conditions.jl
@@ -33,9 +33,19 @@ function set_initial_conditions!(model, ics::SmoothSTStaircaseInitialConditions)
 
     return nothing
 end
-set_initial_conditions!(model, ics::Union{SingleInterfaceICs, PeriodoicSingleInterfaceICs}) =
-    set_initial_conditions!(model, ics, ics.interface_smoothing)
-function set_initial_conditions!(model, ics::SingleInterfaceICs, interface_smoothing::Type{<:NoSmoothing})
+"""
+    function set_initial_conditions!(model, ics::SingleInterfaceICs)
+Set the initial conditions for a single interface. When there is no smoothing the initial
+conditions are set as a step change at `ics.depth_of_interface`. Any `interface_smoothing`
+will be added if specified.
+When there is a `background_state` in the `SingleInterfaceICs`, nothing will be set in the
+model tracer fields (i.e. nothing set in `model.tracers.S` and `model.tracers.T`) but noise
+can be added.
+"""
+set_initial_conditions!(model, ics::SingleInterfaceICs) =
+    set_initial_conditions!(model, ics, ics.interface_smoothing, ics.background_state)
+function set_initial_conditions!(model, ics,
+                                 interface_smoothing::Type{<:NoSmoothing}, background_state::Type{<:NoBackground})
 
     depth_of_interface = ics.depth_of_interface
     z = znodes(model.grid, Center())
@@ -53,8 +63,8 @@ function set_initial_conditions!(model, ics::SingleInterfaceICs, interface_smoot
 
     return nothing
 end
-set_initial_conditions!(model, ics::PeriodoicSingleInterfaceICs, interface_smoothing::Type{<:NoSmoothing}) = nothing
-function set_initial_conditions!(model, ics, interface_smoothing::Type{<:Tanh})
+set_initial_conditions!(model, ics, interface_smoothing::Type{<:NoSmoothing}, background_state) = nothing
+function set_initial_conditions!(model, ics, interface_smoothing::Type{<:Tanh}, background_state)
 
     depth_of_interface = ics.depth_of_interface
     Lz = model.grid.Lz

diff --git a/src/single_interfaces_initial_conditions.jl b/src/single_interfaces_initial_conditions.jl
@@ -1,15 +1,11 @@
-abstract type AbstractSingleInterfaceInitialConditions <: AbstractInitialConditions end
-
-Base.iterate(sics::AbstractSingleInterfaceInitialConditions, state = 1) =
+Base.iterate(sics::AbstractInitialConditions, state = 1) =
     state > length(fieldnames(sics)) ? nothing : (getfield(sics, state), state + 1)
 """
     struct STSingleInterfaceInitialConditions
 Initial conditions for a single interface (i.e. two layers with uniform `S` and `T` seperated
-by a step change). The property `maintain_interface` is a `Boolean` which if set to `true` will
-set [reentrant_boundary_conditions](@ref) so that the interface will be maintained (by not
-letting the system run down).
+by a step change).
 """
-struct STSingleInterfaceInitialConditions{T, A, IS} <: AbstractSingleInterfaceInitialConditions
+struct STSingleInterfaceInitialConditions{T, A, BF, IS} <: AbstractInitialConditions
     "The depth of the interface"
      depth_of_interface :: T
     "Salinity values in each layer"
@@ -18,64 +14,36 @@ struct STSingleInterfaceInitialConditions{T, A, IS} <: AbstractSingleInterfaceIn
      temperature_values :: A
      "Initial smoothing function over interface"
     interface_smoothing :: IS
+    "`BackgroundField` about which anomaly is advected. Should be an `AbstractBackgroundFunction`."
+       background_state :: BF
      "Initial R_ρ at the interface"
                     R_ρ :: T
 end
 function STSingleInterfaceInitialConditions(model, depth_of_interface, salinity, temperature;
-                                            interface_smoothing = NoSmoothing)
+                                            interface_smoothing = NoSmoothing,
+                                            background_state = NoBackground)
 
     eos = model.buoyancy.model.equation_of_state
 
     R_ρ = compute_R_ρ(salinity, temperature, depth_of_interface, eos)
 
     return STSingleInterfaceInitialConditions(depth_of_interface, salinity, temperature,
-                                              interface_smoothing, R_ρ)
+                                              interface_smoothing, background_state, R_ρ)
 
 end
 function STSingleInterfaceInitialConditions(eos::BoussinesqEquationOfState, depth_of_interface,
-                                            salinity, temperature; interface_smoothing = NoSmoothing)
+                                            salinity, temperature;
+                                            interface_smoothing = NoSmoothing,
+                                            background_state = NoBackground)
 
     R_ρ = compute_R_ρ(salinity, temperature, depth_of_interface, eos)
 
     return STSingleInterfaceInitialConditions(depth_of_interface, salinity, temperature,
-                                              interface_smoothing, R_ρ)
+                                                interface_smoothing, background_state, R_ρ)
 
 end
 const SingleInterfaceICs = STSingleInterfaceInitialConditions # alias
-Base.summary(ics::SingleInterfaceICs) = "Single S-T interface at z = $(ics.depth_of_interface) on bounded z domain"
-
-"""
-    struct PeriodicSTSingleInterfaceInitialConditions
-Sets a `BackgroundField` according to `background_State` and uses a triply periodic domain
-to evolve salinity and temperature anomalies about the background state.
-"""
-struct PeriodicSTSingleInterfaceInitialConditions{T, A, IS, BF} <: AbstractSingleInterfaceInitialConditions
-    "The depth of the interface"
-     depth_of_interface :: T
-    "Salinity values in each layer"
-        salinity_values :: A
-     "Temperature values in each layer"
-     temperature_values :: A
-    "Initial smoothing function over interface"
-    interface_smoothing :: IS
-     "Initial R_ρ at the interface"
-                    R_ρ :: T
-    "`BackgroundField` about which anomaly is advected. Should be an `AbstractBackgroundFunction`."
-       background_state :: BF
-end
-function PeriodicSTSingleInterfaceInitialConditions(eos::BoussinesqEquationOfState, depth_of_interface,
-                                                    salinity, temperature;
-                                                    interface_smoothing = NoSmoothing,
-                                                    background_state = NoBackground)
-
-    R_ρ = compute_R_ρ(salinity, temperature, depth_of_interface, eos)
-
-    return PeriodicSTSingleInterfaceInitialConditions(depth_of_interface, salinity, temperature,
-                                                      interface_smoothing, R_ρ, background_state)
-
-end
-const PeriodoicSingleInterfaceICs = PeriodicSTSingleInterfaceInitialConditions # alias
-Base.summary(ics::PeriodoicSingleInterfaceICs) = "Single S-T interface at z = $(ics.depth_of_interface) on triply periodic domain"
+Base.summary(ics::SingleInterfaceICs) = "Single S-T interface at z = $(ics.depth_of_interface)"
 
 function compute_R_ρ(salinity, temperature, depth_of_interfaces::Number, eos)
 
@@ -92,21 +60,12 @@ function compute_R_ρ(salinity, temperature, depth_of_interfaces::Number, eos)
     return (0.5 * (ρ_f - ρ_u) + 0.5 * (ρ_l - ρ_g)) / (0.5 * (ρ_f - ρ_l) + 0.5 * (ρ_u - ρ_g))
 end
 
-function Base.show(io::IO, sics::AbstractSingleInterfaceInitialConditions)
-    if sics isa PeriodicSTSingleInterfaceInitialConditions
-        println(io, "PeriodicSTSingleInterfaceInitialConditions")
-        println(io, "┣━━ depth_of_interface: $(sics.depth_of_interface)")
-        println(io, "┣━━━━━ salinity_values: $(sics.salinity_values)")
-        println(io, "┣━━ temperature_values: $(sics.temperature_values)")
-        println(io, "┣━ interface_smoothing: $(summary(sics.interface_smoothing))")
-        println(io, "┣━━━━ background_state: $(summary(sics.background_state))")
-        print(io,   "┗━━━━━━━━━━━━━━━━━ R_ρ: $(round.(sics.R_ρ; digits = 2))")
-    elseif sics isa STSingleInterfaceInitialConditions
-        println(io, "STSingleInterfaceInitialConditions")
-        println(io, "┣━━ depth_of_interface: $(sics.depth_of_interface)")
-        println(io, "┣━━━━━ salinity_values: $(sics.salinity_values)")
-        println(io, "┣━━ temperature_values: $(sics.temperature_values)")
-        println(io, "┣━ interface_smoothing: $(summary(sics.interface_smoothing))")
-        print(io,   "┗━━━━━━━━━━━━━━━━━ R_ρ: $(round.(sics.R_ρ; digits = 2))")
-    end
+function Base.show(io::IO, sics::STSingleInterfaceInitialConditions)
+    println(io, "STSingleInterfaceInitialConditions")
+    println(io, "┣━━ depth_of_interface: $(sics.depth_of_interface)")
+    println(io, "┣━━━━━ salinity_values: $(sics.salinity_values)")
+    println(io, "┣━━ temperature_values: $(sics.temperature_values)")
+    println(io, "┣━ interface_smoothing: $(summary(sics.interface_smoothing))")
+    println(io, "┣━━━━ background_state: $(summary(sics.background_state))")
+    print(io,   "┗━━━━━━━━━━━━━━━━━ R_ρ: $(round.(sics.R_ρ; digits = 2))")
 end
diff --git a/src/staircase_background.jl b/src/staircase_background.jl
@@ -45,7 +45,7 @@ Base.summary(bt::BackgroundTanh) = "$(bt.func)"
 Base.summary(bl::BackgroundLinear) = "$(bl.func)"
 Base.summary(bn::Type{<:NoBackgroundFunction}) = "no background field"
 
-S_and_T_background_fields(ics::PeriodicSTSingleInterfaceInitialConditions, Lz) =
+S_and_T_background_fields(ics::STSingleInterfaceInitialConditions, Lz) =
     S_and_T_background_fields(ics, Lz, ics.background_state)
 "Return blank `NamedTuple` so that no `BackgroundField`s are set."
 S_and_T_background_fields(ics, Lz, background_state::Type{<:NoBackground}) = NamedTuple()
@@ -69,7 +69,7 @@ tanh_background(z, ΔC, Cₗ, Lz, z_interface, D) = Cₗ - 0.5 * ΔC * (1  + tan
 @inline linear_background(x, y, z, t, p) = p.Cᵤ - p.ΔC * z / p.Lz
 linear_background(z, ΔC, Cᵤ, Lz) = Cᵤ - ΔC * z / Lz
 
-function get_parameters!(ics::PeriodicSTSingleInterfaceInitialConditions, tracer::Symbol, Lz)
+function get_parameters!(ics::STSingleInterfaceInitialConditions, tracer::Symbol, Lz)
 
     z_interface = ics.depth_of_interface
     C = Array(getproperty(ics, tracer))
@@ -96,52 +96,50 @@ end
 Where there is `BackgroundField` (currently this assumes that is for periodic simulations)
 save the background state for the tracers and density so this can be used later
 """
-save_background_state!(simulation, sdns) = save_background_state!(simulation, sdns.model, sdns.initial_conditions)
-save_background_state!(simulation, model, initial_conditions) = nothing
-function save_background_state!(simulation, model, initial_conditions::PeriodoicSingleInterfaceICs)
-
-    if !isnothing(initial_conditions.background_state)
-        S_background = Field(model.background_fields.tracers.S)
-        compute!(S_background)
-        S_background_array = Array(interior(S_background, :, :, :))
-        T_background = Field(model.background_fields.tracers.T)
-        compute!(T_background)
-        T_background_array = Array(interior(T_background, :, :, :))
-        σ_background = Field(seawater_density(model, temperature = T_background, salinity = S_background,
-                                            geopotential_height = 0))
-        compute!(σ_background)
-        σ_background_array = Array(interior(σ_background, :, :, :))
-
-        if simulation.output_writers[:tracers] isa NetCDFOutputWriter
-
-            NCDataset(simulation.output_writers[:tracers].filepath, "a") do ds
-                defVar(ds, "S_background", S_background_array, ("xC", "yC", "zC"),
-                    attrib = Dict("longname" => "Background field for salinity",
-                                    "units" => "gkg⁻¹"))
-                defVar(ds, "T_background", T_background_array, ("xC", "yC", "zC"),
-                    attrib =  Dict("longname" => "Background field for temperature",
-                                    "units" => "°C"))
-            end
-
-            NCDataset(simulation.output_writers[:computed_output].filepath, "a") do ds
-                defVar(ds, "σ_background", σ_background_array, ("xC", "yC", "zC"),
-                    attrib = Dict("longname" => "Background field for potential density (0dbar) computed from the `S` and `T` background fields",
-                                    "units" => "kgm⁻³"))
-            end
-
-        elseif simulation.output_writers[:tracers] isa JLD2OutputWriter
-
-            jldopen(simulation.output_writers[:tracers].filepath, "a+") do f
-                f["S_background"] = S_background_array
-                f["T_background"] = T_background_array
-            end
-
-            jldopen(simulation.output_writers[:computed_output].filepath, "a+") do f
-                f["σ_background"] = σ_background_array
-            end
+save_background_state!(simulation, sdns) = save_background_state!(simulation, sdns.model, sdns.initial_conditions.background_state)
+save_background_state!(simulation, model, background_state::Type{<:NoBackground}) = nothing
+function save_background_state!(simulation, model, background_state)
+
+    S_background = Field(model.background_fields.tracers.S)
+    compute!(S_background)
+    S_background_array = Array(interior(S_background, :, :, :))
+    T_background = Field(model.background_fields.tracers.T)
+    compute!(T_background)
+    T_background_array = Array(interior(T_background, :, :, :))
+    σ_background = Field(seawater_density(model, temperature = T_background, salinity = S_background,
+                                        geopotential_height = 0))
+    compute!(σ_background)
+    σ_background_array = Array(interior(σ_background, :, :, :))
+
+    if simulation.output_writers[:tracers] isa NetCDFOutputWriter
+
+        NCDataset(simulation.output_writers[:tracers].filepath, "a") do ds
+            defVar(ds, "S_background", S_background_array, ("xC", "yC", "zC"),
+                attrib = Dict("longname" => "Background field for salinity",
+                                "units" => "gkg⁻¹"))
+            defVar(ds, "T_background", T_background_array, ("xC", "yC", "zC"),
+                attrib =  Dict("longname" => "Background field for temperature",
+                                "units" => "°C"))
+        end
+
+        NCDataset(simulation.output_writers[:computed_output].filepath, "a") do ds
+            defVar(ds, "σ_background", σ_background_array, ("xC", "yC", "zC"),
+                attrib = Dict("longname" => "Background field for potential density (0dbar) computed from the `S` and `T` background fields",
+                                "units" => "kgm⁻³"))
+        end
+
+    elseif simulation.output_writers[:tracers] isa JLD2OutputWriter
 
+        jldopen(simulation.output_writers[:tracers].filepath, "a+") do f
+            f["S_background"] = S_background_array
+            f["T_background"] = T_background_array
+        end
+
+        jldopen(simulation.output_writers[:computed_output].filepath, "a+") do f
+            f["σ_background"] = σ_background_array
         end
 
     end
+
     return nothing
 end