Add StepChangeLinearGradient

examples/twolayer_example.jl

@@ -14,9 +14,10 @@ S₀ᵘ = (stable = 34.551, cabbeling = 34.58, unstable = 34.59)
 cabbeling = CabbelingUpperLayerInitialConditions(S₀ᵘ.cabbeling, T₀ᵘ)
 initial_conditions = TwoLayerInitialConditions(cabbeling)
 transition_depth = find_depth(model, INTERFACE_LOCATION)
-profile_function = StepChange(transition_depth)
+# profile_function = StepChange(transition_depth)
+profile_function = StepChangeLinearGradient(transition_depth, 1e-2, 0.0, model)
 noise_depth = find_depth(model, INTERFACE_LOCATION)
-initial_noise = SalinityNoise(noise_depth, 1e-2)
+initial_noise = SalinityNoise(noise_depth, 0.0)
 
 tldns = TwoLayerDNS(model, profile_function, initial_conditions; initial_noise)
 

src/TwoLayerDirectNumericalShenanigans.jl

@@ -45,7 +45,7 @@ export set_two_layer_initial_conditions!, find_depth
 
 export AbstractContinuousProfileFunction, HyperbolicTangent, Erf, MidPoint
 
-export AbststractStepChangeProfileFuncion, StepChange
+export AbststractStepChangeProfileFuncion, StepChange, StepChangeLinearGradient
 
 export AbstractTracerPerturbation, SalinityGaussianProfile, SalinityGaussianBlob,
        TemperatureGaussianProfile, TemperatureGaussianBlob

src/set_initialconditions.jl

@@ -1090,3 +1090,42 @@ function set_two_layer_initial_conditions!(model::Oceananigans.AbstractModel,
     perturb_velocity!(model, initial_noise)
 
 end
+####
+#### Step Change with linear gradient
+####
+function set_two_layer_initial_conditions!(model::Oceananigans.AbstractModel,
+                                            initial_conditions::TwoLayerInitialConditions,
+                                            profile_function::StepChangeLinearGradient,
+                                            tracer_perturbation::Nothing,
+                                            initial_noise::Nothing)
+
+    S₀ᵘ, S₀ˡ = initial_conditions.S₀ᵘ, initial_conditions.S₀ˡ
+    T₀ᵘ, T₀ˡ = initial_conditions.T₀ᵘ, initial_conditions.T₀ˡ
+
+    initial_S_profile(x, y, z) = Heaviside_with_linear_gradient(z, S₀ˡ, S₀ᵘ, profile_function)
+    initial_T_profile(x, y, z) = Heaviside_with_linear_gradient(z, T₀ˡ, T₀ᵘ, profile_function, tracer = :T)
+
+    set!(model, S = initial_S_profile, T = initial_T_profile)
+
+return nothing
+
+end
+#### Step Change with linear gradient + salinity noise
+function set_two_layer_initial_conditions!(model::Oceananigans.AbstractModel,
+                                            initial_conditions::TwoLayerInitialConditions,
+                                            profile_function::StepChangeLinearGradient,
+                                            tracer_perturbation::Nothing,
+                                            initial_noise::SalinityNoise)
+
+S₀ᵘ, S₀ˡ = initial_conditions.S₀ᵘ, initial_conditions.S₀ˡ
+T₀ᵘ, T₀ˡ = initial_conditions.T₀ᵘ, initial_conditions.T₀ˡ
+
+initial_S_profile(x, y, z) = Heaviside_with_linear_gradient(z, S₀ˡ, S₀ᵘ, profile_function) +
+                             perturb_tracer(z, initial_noise)
+initial_T_profile(x, y, z) = Heaviside_with_linear_gradient(z, T₀ˡ, T₀ᵘ, profile_function, tracer = :T)
+
+set!(model, S = initial_S_profile, T = initial_T_profile)
+
+return nothing
+
+end

src/stepchangeprofilefunction.jl

@@ -1,8 +1,15 @@
 abstract type AbstractStepChangeProfileFunction <: AbstractProfileFunction end
 "`show` for `AbstractStepChangeProfileFunction`"
 function Base.show(io::IO, stpf::AbstractStepChangeProfileFunction)
-    println(io, "$(typeof(stpf))")
-    print(io, "┗━━ interface_location: z = $(stpf.interface_location)m")
+    if stpf isa StepChange
+        println(io, "$(typeof(stpf))")
+        print(io, "┗━━ interface_location: z = $(stpf.interface_location)m")
+    elseif stpf isa StepChangeLinearGradient
+        println(io, "$(typeof(stpf))")
+        println(io, "┣━ interface_location: z = $(stpf.interface_location)m ")
+        println(io, "┣━━━━━━━━━━━━━━━ dSdz: $(stpf.dSdz) ")
+          print(io, "┗━━━━━━━━━━━━━━━ dTdZ: $(stpf.dTdz)")
+    end
 end
 "`iterate` for `AbstractStepChangeProfileFunction`"
 Base.iterate(pf::AbstractStepChangeProfileFunction, state = 1) =
@@ -20,3 +27,61 @@ struct StepChange{T} <: AbstractStepChangeProfileFunction
 end
 Heaviside(z, Cˡ::Number, Cᵘ::Number, profile_function::StepChange) =
     z > profile_function.interface_location ? Cᵘ : Cˡ
+
+"""
+    struct StepChangeLinearGradient
+Containter with the depth of the `interface_location` between the upper and lower layers
+and values to construct a linear gradient for salinity (`dSdz`) and temperature (`dTdz`).
+The resulting profile is a step change between the upper and lower layer salinity and
+temperature that has a linear (with depth) gradient in each layer.
+**Note:** The `Cˡ` and `Cᵘ` values passed to [Heaviside_with_linear_gradient](@ref) are
+set at the *interface*, hence an offset to the salinity and temperature values is computed
+to ensure this is the behaviour.
+"""
+struct StepChangeLinearGradient{T} <: AbstractStepChangeProfileFunction
+    "Location of the interface between the two layers"
+    interface_location :: T
+    "Linear gradient value for salinity"
+                  dSdz :: T
+    "Salinity offset to ensure salinity values at interface are S₀ᵘ and S₀ˡ"
+       salinity_offset :: Tuple{T, T}
+    "Linear gradient value for temperature"
+                  dTdz :: T
+    "Temperature offset to ensure salinity values at interface are T₀ᵘ and T₀ˡ"
+    temperature_offset :: Tuple{T, T}
+end
+function StepChangeLinearGradient(interface_location, dSdz, dTdz, model)
+
+    offset_depth = model.architecture isa CPU ?  begin
+                                                z = znodes(model.grid, Center(), Center(), Center())
+                                                il_idx = findfirst(z .> interface_location)
+                                                z[il_idx], z[il_idx - 1]
+                                            end :
+                                                allowscalar() do
+                                                z = znodes(model.grid, Center(), Center(), Center())
+                                                il_idx = findfirst(z .> interface_location)
+                                                z[il_idx], z[il_idx - 1]
+                                            end
+
+    S_offset = (dSdz * offset_depth[1], dSdz * offset_depth[2])
+    T_offset = (-dTdz * offset_depth[1], -dTdz * offset_depth[2])
+
+    return StepChangeLinearGradient(interface_location, dSdz, S_offset, dTdz, T_offset)
+
+end
+function Heaviside_with_linear_gradient(z, Cˡ::Number, Cᵘ::Number,
+                                        profile_function::StepChangeLinearGradient;
+                                        tracer = :S)
+    interface_location = profile_function.interface_location
+    dCdz = tracer == :S ? -profile_function.dSdz : profile_function.dTdz
+    upper_offset, lower_offset = tracer == :S ? profile_function.salinity_offset :
+                                                profile_function.temperature_offset
+    Cᵘ_offset = upper_offset == 0 ? Cᵘ : Cᵘ + upper_offset
+    Cˡ_offset = lower_offset == 0 ? Cˡ : Cˡ + lower_offset
+    if z > interface_location
+        Cᵘ_offset + dCdz * z
+    else
+        Cˡ_offset + dCdz * z
+    end
+
+end

test/initialconditions_test.jl

@@ -91,3 +91,26 @@ function tracer_stepchange(dns::TwoLayerDNS)
     return S_upper, S_lower, T_upper, T_lower
 
 end
+function tracer_stepchangelineargradient(dns::TwoLayerDNS, z)
+
+    il = dns.profile_function.interface_location
+    il_idx = findfirst(z .> il)
+    depth = z[il_idx ]
+    depth2 = z[il_idx  - 1]
+    dSdz, dTdz = -dns.profile_function.dSdz, dns.profile_function.dTdz
+
+    S₀ˡ, S₀ᵘ = dns.initial_conditions.S₀ˡ, dns.initial_conditions.S₀ᵘ
+    S₀ˡ_offset, S₀ᵘ_offset = S₀ˡ - dSdz * depth2, S₀ᵘ - dSdz * depth
+    S_upper = S₀ᵘ_offset .+ dSdz .* z[z .> il]
+    S_lower = S₀ˡ_offset .+ dSdz .* z[z .≤ il]
+    S = vcat(S_lower, S_upper)
+
+    T₀ˡ, T₀ᵘ = dns.initial_conditions.T₀ˡ, dns.initial_conditions.T₀ᵘ
+    T₀ˡ_offset, T₀ᵘ_offset = T₀ˡ - dTdz * depth2, T₀ᵘ - dTdz * depth
+    T_upper = T₀ᵘ_offset .+ dTdz .* z[z .> il]
+    T_lower = T₀ˡ_offset .+ dTdz .* z[z .≤ il]
+    T = vcat(T_lower, T_upper)
+
+    return S, T
+
+end

test/runtests.jl

@@ -125,6 +125,39 @@ end
     @test isequal((true, true, true, true), tracer_stepchange(dns))
 end
 
+@testset "Step change linear gradient" begin
+
+    # without gradient, sanity check
+    model = DNSModel(architecture, DOMAIN_EXTENT, resolution, diffusivities)
+    dSdz, dTdz = 0.0, 0.0
+    profile_function = StepChangeLinearGradient(z[depth_idx], dSdz, dTdz, model)
+    initial_conditions = TwoLayerInitialConditions(34.551, -1.5, 34.7, 0.5)
+    dns = TwoLayerDNS(model, profile_function, initial_conditions)
+    set_two_layer_initial_conditions!(dns)
+    x_idx, y_idx = rand(1:model.grid.Nx), rand(1:model.grid.Ny)
+    S = interior(dns.model.tracers.S, x_idx, y_idx, :)
+    T = interior(dns.model.tracers.T, x_idx, y_idx, :)
+    S_test, T_test = tracer_stepchangelineargradient(dns, z)
+
+    @test all((S .== S_test))
+    @test all((T .== T_test))
+
+    # with gradient
+    model = DNSModel(architecture, DOMAIN_EXTENT, resolution, diffusivities)
+    dSdz, dTdz = 0.01, 0.05
+    profile_function = StepChangeLinearGradient(z[depth_idx], dSdz, dTdz, model)
+    initial_conditions = TwoLayerInitialConditions(34.551, -1.5, 34.7, 0.5)
+    dns = TwoLayerDNS(model, profile_function, initial_conditions)
+    set_two_layer_initial_conditions!(dns)
+    x_idx, y_idx = rand(1:model.grid.Nx), rand(1:model.grid.Ny)
+    S = interior(dns.model.tracers.S, x_idx, y_idx, :)
+    T = interior(dns.model.tracers.T, x_idx, y_idx, :)
+    S_test, T_test = tracer_stepchangelineargradient(dns, z)
+
+    @test all((S .== S_test))
+    @test all((T .== T_test))
+end
+
 @testset "Find depth" begin
     test_depth = rand(z)
     @test isequal(test_depth, find_depth(model, test_depth))