Add thermocline with variable N²(z) (#34)

Manifest.toml

@@ -699,10 +699,10 @@ uuid = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 version = "4.4.1"
 
 [[NNlib]]
-deps = ["Libdl", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
-git-tree-sha1 = "a8180fd1445e31c0b1add98dae8da694ac2c23fd"
+deps = ["Compat", "Libdl", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
+git-tree-sha1 = "1ae42464fea5258fd2ff49f1c4a40fc41cba3860"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.7.6"
+version = "0.7.7"
 
 [[NaNMath]]
 git-tree-sha1 = "c84c576296d0e2fbb3fc134d3e09086b3ea617cd"

Project.toml

@@ -13,6 +13,7 @@ Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

examples/three_layer_constant_fluxes.jl

@@ -27,7 +27,7 @@ using Oceananigans.Fields
 using Oceananigans.Fields: PressureField
 using Oceananigans.OutputWriters
 
-using LESbrary.Utils: SimulationProgressMessenger
+using LESbrary.Utils: SimulationProgressMessenger, fit_cubic, poly
 using LESbrary.NearSurfaceTurbulenceModels: SurfaceEnhancedModelConstant
 using LESbrary.TurbulenceStatistics: first_through_second_order, turbulent_kinetic_energy_budget
 using LESbrary.TurbulenceStatistics: subfilter_momentum_fluxes
@@ -88,12 +88,12 @@ function parse_command_line_arguments()
 
         "--thermocline-width"
             help = "The width of the thermocline in units of m."
-            default = 24.0
+            default = 96.0
             arg_type = Float64
 
         "--surface-layer-buoyancy-gradient"
             help = "The buoyancy gradient in the surface layer in units s⁻²."
-            default = 2e-6
+            default = 1e-6
             arg_type = Float64
 
         "--thermocline-buoyancy-gradient"
@@ -103,7 +103,7 @@ function parse_command_line_arguments()
 
         "--deep-buoyancy-gradient"
             help = "The buoyancy gradient below the thermocline in units s⁻²."
-            default = 2e-6
+            default = 1e-6
             arg_type = Float64
 
         "--hours"
@@ -246,7 +246,12 @@ model = IncompressibleModel(architecture = GPU(),
 
 ## Noise with 8 m decay scale
 Ξ(z) = rand() * exp(z / 8)
-
+
+p1 = (z_transition, θ_transition)
+p2 = (z_deep, θ_deep)
+coeffs = fit_cubic(p1, p2, dθdz_surface_layer, dθdz_deep)
+θ_thermocline(z) = poly(z, coeffs)
+
 """
     initial_temperature(x, y, z)
 
@@ -257,11 +262,11 @@ function initial_temperature(x, y, z)
 
     noise = 1e-6 * Ξ(z) * dθdz_surface_layer * grid.Lz
 
-    if z > z_transition
+    if z_transition < z <= 0
         return θ_surface + dθdz_surface_layer * z + noise
 
-    elseif z > z_deep
-        return θ_transition + dθdz_thermocline * (z - z_transition) + noise
+    elseif z_deep < z <= z_transition
+        return θ_thermocline(z) + noise
 
     else
         return θ_deep + dθdz_deep * (z - z_deep) + noise

src/Utils/Utils.jl

@@ -7,7 +7,9 @@ export
     select_device!,
 
     # progress_messenger.jl
-    SimulationProgressMessenger
+    SimulationProgressMessenger,
+
+    fit_cubic
 
 using Printf    
 
@@ -73,5 +75,6 @@ function prefix_tuple_names(prefix, tup)
 end
 
 include("progress_messenger.jl")
+include("fit_cubic.jl")
 
 end # module

src/Utils/fit_cubic.jl

@@ -0,0 +1,25 @@
+using NLsolve
+
+poly(x, p) = sum(p[n] * x^(n-1) for n in 1:length(p))
+∂poly(x, p) = sum(n * p[n+1] * x^(n-1) for n in 1:length(p)-1)
+
+"""
+        fit_cubic(p1, p2, s1, s2)
+
+Return the coeficients c such that f(x) = c₁ + c₂x + c₃x² + c₄x³ is a cubic polynomial that passes through the points p1 = (x1, y1) and p2 = (x2, y2) with slope s1 at p1 and slope s2 at p2. 
+"""
+function fit_cubic(p1, p2, s1, s2)
+    x1, y1 = p1
+    x2, y2 = p2
+
+    function f!(F, c)
+        F[1] = poly(x1, c) - y1
+        F[2] = poly(x2, c) - y2
+        F[3] = ∂poly(x1, c) - s1
+        F[4] = ∂poly(x2, c) - s2
+    end
+
+    results = nlsolve(f!, zeros(4), show_trace=true)
+    return results.zero
+end
+