Merge pull request #45 from CliMA/ck/test_and_abstraction

Add profile abstraction, improve tests
CliMA · Jul 3, 2024 · 6aa54b3 · 6aa54b3
2 parents a86e0a5 + 0ab9f34
commit 6aa54b3
Show file tree

Hide file tree

Showing 17 changed files with 318 additions and 167 deletions.
diff --git a/src/AtmosphericProfilesLibrary.jl b/src/AtmosphericProfilesLibrary.jl
@@ -2,6 +2,32 @@ module AtmosphericProfilesLibrary
 
 import Dierckx
 
+abstract type AbstractProfile end
+struct TimeProfile{P} <: AbstractProfile
+  prof::P
+end
+@inline (prof::TimeProfile)(t) = prof.prof(t)
+
+struct TimeZProfile{P} <: AbstractProfile
+  prof::P
+end
+@inline (prof::TimeZProfile)(t, z) = prof.prof(t, z)
+
+struct ZProfile{P} <: AbstractProfile
+  prof::P
+end
+@inline (prof::ZProfile)(z) = prof.prof(z)
+
+struct ΠTimeZProfile{P} <: AbstractProfile
+  prof::P
+end
+@inline (prof::ΠTimeZProfile)(Π, t, z) = prof.prof(Π, t, z)
+
+struct ΠZProfile{P} <: AbstractProfile
+  prof::P
+end
+@inline (prof::ΠZProfile)(Π, z) = prof.prof(Π, z)
+
 # Large data-based profiles
 include("profiles/Soares.jl")
 include("profiles/Nieuwstadt.jl")

diff --git a/src/profiles/ARM_SGP.jl b/src/profiles/ARM_SGP.jl
@@ -8,26 +8,26 @@ function ARM_SGP_θ_liq_ice(::Type{FT}) where {FT}
     z_in = ARM_SGP_z(FT)
     θ_liq_ice_in = FT[299.0, 301.5, 302.5, 303.53, 303.7, 307.13, 314.0, 343.2] # K
     profile = Dierckx.Spline1D(z_in, θ_liq_ice_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
 """ [Brown2002](@cite) """
-ARM_SGP_u(::Type{FT}) where {FT} = z -> FT(10)
+ARM_SGP_u(::Type{FT}) where {FT} = ZProfile(z -> FT(10))
 
 """ [Brown2002](@cite) """
 function ARM_SGP_q_tot(::Type{FT}) where {FT}
     z_in = ARM_SGP_z(FT)
     r_in = FT[15.2,15.17,14.98,14.8,14.7,13.5,3.0,3.0] ./ 1000 # qt should be in kg/kg
     q_tot_in = r_in ./ (1 .+ r_in)
     profile = Dierckx.Spline1D(z_in, q_tot_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
 """ [Brown2002](@cite) """
 function ARM_SGP_tke(::Type{FT}) where {FT}
-    return z -> if z <= 2500.0
-        FT(1) - z / 3000
-    else
-        FT(0)
-    end
+    return ZProfile(z -> if z <= 2500.0
+            FT(1) - z / 3000
+        else
+            FT(0)
+        end)
 end
 
 """ TMP TKE profile for testing """
@@ -50,7 +50,7 @@ function ARM_SGP_tke_prescribed(::Type{FT}) where {FT}
                 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
     not_type_stable_spline = Dierckx.Spline1D(z_in, tke_in; k = 1)
-    return x -> FT(not_type_stable_spline(x))
+    return ZProfile(x -> FT(not_type_stable_spline(x)))
 end
 
 """ [Brown2002](@cite) """
@@ -68,13 +68,14 @@ function ARM_SGP_dTdt(::Type{FT}) where {FT}
     RT_in = FT[-0.125, 0.0, 0.0, 0.0, 0.0, -0.1] ./ 3600
     dTdt_A = Dierckx.Spline1D(t_in, AT_in; k = 1)
     dTdt_R = Dierckx.Spline1D(t_in, RT_in; k = 1)
-    return (t, z) -> if z <= 1000.0
-        dTdt_A(t)+dTdt_R(t)
+    prof = (t, z) -> if z <= 1000.0
+        FT(dTdt_A(t)+dTdt_R(t))
     elseif z > 1000.0 && z <= 2000.0
-        (dTdt_A(t)+dTdt_R(t)) * (1 - (z - 1000) / 1000)
+        FT((dTdt_A(t)+dTdt_R(t)) * (1 - (z - 1000) / 1000))
     else
         FT(0)
     end
+    return TimeZProfile(prof)
 end
 
 """ [Brown2002](@cite) """
@@ -83,27 +84,28 @@ function ARM_SGP_dqtdt(::Type{FT}) where {FT}
     # Radiative forcing for qt converted to [kg/kg/sec]
     Rqt_in = FT[0.08, 0.02, 0.04, -0.1, -0.16, -0.3] ./ 1000 ./ 3600
     dqtdt = Dierckx.Spline1D(t_in, Rqt_in; k = 1)
-    return (Π, t, z) -> if z <= 1000.0
-        dqtdt(t) * Π
+    prof = (Π, t, z) -> if z <= 1000.0
+        FT(dqtdt(t) * Π)
     elseif z > 1000.0 && z <= 2000.0
-        dqtdt(t) * Π * (1 - (z - 1000) / 1000)
+        FT(dqtdt(t) * Π * (1 - (z - 1000) / 1000))
     else
         FT(0)
     end
+    return ΠTimeZProfile(prof)
 end
 
 """ [Brown2002](@cite) """
 function ARM_SGP_shf(::Type{FT}) where {FT}
     t_Sur_in = FT[0.0, 4.0, 6.5, 7.5, 10.0, 12.5, 14.5] .* 3600 #LES time is in sec
     shf = FT[-30.0, 90.0, 140.0, 140.0, 100.0, -10, -10] # W/m^2
     profile = Dierckx.Spline1D(t_Sur_in, shf; k = 1)
-    return profile
+    return TimeProfile(profile)
 end
 
 """ [Brown2002](@cite) """
 function ARM_SGP_lhf(::Type{FT}) where {FT}
     t_Sur_in = FT[0.0, 4.0, 6.5, 7.5, 10.0, 12.5, 14.5] .* 3600 #LES time is in sec
     lhf = FT[5.0, 250.0, 450.0, 500.0, 420.0, 180.0, 0.0] # W/m^2
     profile = Dierckx.Spline1D(t_Sur_in, lhf; k = 1)
-    return profile
+    return TimeProfile(profile)
 end
diff --git a/src/profiles/Bomex.jl b/src/profiles/Bomex.jl
@@ -1,5 +1,5 @@
 """ [Siebesma2003](@cite) """
-Bomex_θ_liq_ice(::Type{FT}) where {FT} = z -> if z <= 520.0
+Bomex_θ_liq_ice(::Type{FT}) where {FT} = ZProfile(z -> if z <= 520.0
         FT(298.7)
     elseif z > 520.0 && z <= 1480.0
         FT(298.7) + (z - 520) * (FT(302.4) - FT(298.7)) / (1480 - 520)
@@ -9,29 +9,29 @@ Bomex_θ_liq_ice(::Type{FT}) where {FT} = z -> if z <= 520.0
         FT(308.2) + (z - 2000) * (FT(311.85) - FT(308.2)) / (3000 - 2000)
     else
         FT(0)
-    end
+    end)
 """ [Siebesma2003](@cite) """
-Bomex_q_tot(::Type{FT}) where {FT} = z -> if z <= 520
+Bomex_q_tot(::Type{FT}) where {FT} = ZProfile(z -> if z <= 520
         (FT(17.0) + z * (FT(16.3) - 17) / 520) / 1000
     elseif z > 520.0 && z <= 1480.0
         (FT(16.3) + (z - 520) * (FT(10.7) - FT(16.3)) / (1480 - 520)) / 1000
     elseif z > 1480.0 && z <= 2000.0
         (FT(10.7) + (z - 1480) * (FT(4.2) - FT(10.7)) / (2000 - 1480)) / 1000
-    elseif z > 2000.0
+    else
         (FT(4.2) + (z - 2000) * (3 - FT(4.2)) / (3000 - 2000)) / 1000
-    end
+    end)
 """ [Siebesma2003](@cite) """
-Bomex_u(::Type{FT}) where {FT} = z -> if z <= 700.0
+Bomex_u(::Type{FT}) where {FT} = ZProfile(z -> if z <= 700.0
         FT(-8.75)
-    elseif z > 700.0
+    else
         FT(-8.75) + (z - 700) * (FT(-4.61) - FT(-8.75)) / (3000 - 700)
-    end
+    end)
 """ [Siebesma2003](@cite) """
-Bomex_tke(::Type{FT}) where {FT} = z -> if (z <= 2500.0)
+Bomex_tke(::Type{FT}) where {FT} = ZProfile(z -> if (z <= 2500.0)
         FT(1) - z / 3000
     else
         FT(0)
-    end
+    end)
 """ TMP TKE profile for testing """
 function Bomex_tke_prescribed(::Type{FT}) where {FT}
     z_in = FT[25.0, 75.0, 125.0, 175.0, 225.0, 275.0, 325.0, 375.0, 425.0, 475.0, 525.0,
@@ -46,38 +46,38 @@ function Bomex_tke_prescribed(::Type{FT}) where {FT}
                 0.0162, 0.0068, 0.0024, 0.0007, 0.0001, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
     not_type_stable_spline = Dierckx.Spline1D(z_in, tke_in; k = 1)
-    return x -> FT(not_type_stable_spline(x))
+    return ZProfile(x -> FT(not_type_stable_spline(x)))
 end
 
 # Geostrophic velocity profiles
 """ :( """
-Bomex_geostrophic_u(::Type{FT}) where {FT} = z -> -10 + FT(1.8e-3) * z
+Bomex_geostrophic_u(::Type{FT}) where {FT} = ZProfile(z -> -10 + FT(1.8e-3) * z)
 """ :( """
-Bomex_geostrophic_v(::Type{FT}) where {FT} = z -> FT(0)
+Bomex_geostrophic_v(::Type{FT}) where {FT} = ZProfile(z -> FT(0))
 
 """ dTdt(Π, z) """
-Bomex_dTdt(::Type{FT}) where {FT} = (Π, z) -> if z <= 1500.0
+Bomex_dTdt(::Type{FT}) where {FT} = ΠZProfile((Π, z) -> if z <= 1500.0
         (-2 / (3600 * 24)) * Π
     else
         (-2 / (3600 * 24) + (z - 1500) * (0 - -2 / (3600 * 24)) / (3000 - 1500)) * Π
-    end
+    end)
 
 # Large-scale drying
 """ :( """
-Bomex_dqtdt(::Type{FT}) where {FT} = z -> if z <= 300.0
+Bomex_dqtdt(::Type{FT}) where {FT} = ZProfile(z -> if z <= 300.0
         FT(-1.2e-8)   #kg/(kg * s)
     elseif z > 300.0 && z <= 500.0
         FT(-1.2e-8) + (z - 300) * (0 - FT(-1.2e-8)) / (500 - 300) #kg/(kg * s)
     else
         FT(0)
-    end
+    end)
 
 # Large scale subsidence
 """ :( """
-Bomex_subsidence(::Type{FT}) where {FT} = z -> if z <= 1500.0
+Bomex_subsidence(::Type{FT}) where {FT} = ZProfile(z -> if z <= 1500.0
         FT(0) + z * (FT(-0.65) / 100 - 0) / (1500 - 0)
     elseif z > 1500.0 && z <= 2100.0
         FT(-0.65) / 100 + (z - 1500) * (0 - FT(-0.65) / 100) / (2100 - 1500)
     else
         FT(0)
-    end
+    end)
diff --git a/src/profiles/DryBubble.jl b/src/profiles/DryBubble.jl
@@ -62,7 +62,7 @@ function DryBubble_θ_liq_ice(::Type{FT}) where {FT}
         299.9837, 299.9837
     ]
     profile = Dierckx.Spline1D(z_in, θ_liq_ice_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
 
 """ :( """
@@ -98,7 +98,7 @@ function DryBubble_updrafts_θ_liq_ice(::Type{FT}) where {FT}
         300.2587, 300.2216, 300.1782, 300.1452, 300.1143, 300.0859,
         300.0603, 300.0408, 300.0211, 300.0067, 299.9963, 299.9884]
     profile = Dierckx.Spline1D(z_in, θ_liq_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
 
 """ :( """
@@ -115,7 +115,7 @@ function DryBubble_updrafts_area(::Type{FT}) where {FT}
         0.13 , 0.13 , 0.125, 0.12 , 0.115, 0.115, 0.11 , 0.105, 0.1  ,
         0.095, 0.085, 0.08 , 0.07 , 0.055, 0.04]
     profile = Dierckx.Spline1D(z_in, Area_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
 
 """ :( """
@@ -133,7 +133,7 @@ function DryBubble_updrafts_w(::Type{FT}) where {FT}
         0.1422, 0.1425, 0.1424, 0.1419, 0.1361, 0.135 , 0.1335, 0.1316,
         0.1294, 0.1302, 0.1271, 0.1264, 0.1269, 0.1256]
     profile = Dierckx.Spline1D(z_in, W_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
 
 """ :( """
@@ -154,5 +154,5 @@ function DryBubble_updrafts_T(::Type{FT}) where {FT}
         267.2649, 266.7432, 266.2159, 265.698 , 265.1821, 264.6685,
         264.1574, 263.6518, 263.1461, 262.6451, 262.1476, 261.6524]
     profile = Dierckx.Spline1D(z_in, T_in; k = 1)
-    return profile
+    return ZProfile(profile)
 end
diff --git a/src/profiles/Dycoms_RF01.jl b/src/profiles/Dycoms_RF01.jl
@@ -1,36 +1,36 @@
 """ [Stevens2005](@cite) """
-Dycoms_RF01_q_tot(::Type{FT}) where {FT} = z -> if z <= 840.0
+Dycoms_RF01_q_tot(::Type{FT}) where {FT} = ZProfile(z -> if z <= 840.0
         FT(9.0) / FT(1000.0)
     else
         FT(1.5) / FT(1000.0)
-    end
+    end)
 
 """ [Stevens2005](@cite) """
-Dycoms_RF01_θ_liq_ice(::Type{FT}) where {FT} = z -> if z <= 840.0
+Dycoms_RF01_θ_liq_ice(::Type{FT}) where {FT} = ZProfile(z -> if z <= 840.0
         FT(289)
     else
         FT(297.5) + (z - FT(840))^FT(1.0 / 3.0)
-    end
+    end)
 
 """ [Stevens2005](@cite) """
-Dycoms_RF01_u0(::Type{FT}) where {FT} = z -> FT(7)
+Dycoms_RF01_u0(::Type{FT}) where {FT} = ZProfile(z -> FT(7))
 
 """ [Stevens2005](@cite) """
-Dycoms_RF01_v0(::Type{FT}) where {FT} = z -> FT(-5.5)
+Dycoms_RF01_v0(::Type{FT}) where {FT} = ZProfile(z -> FT(-5.5))
 
 """ [Stevens2005](@cite) """
-Dycoms_RF01_tke(::Type{FT}) where {FT} = z -> if z <= 800.0
+Dycoms_RF01_tke(::Type{FT}) where {FT} = ZProfile(z -> if z <= 800.0
         FT(1) - z / FT(1000)
     else
         FT(0)
-    end
+    end)
 """ TMP TKE profile for testing """
 function Dycoms_RF01_tke_prescribed(::Type{FT}) where {FT}
     z_in = FT[25.0, 75.0, 125.0, 175.0, 225.0, 275.0, 325.0, 375.0, 425.0, 475.0, 525.0, 575.0, 625.0, 675.0, 725.0,
               775.0, 825.0, 875.0, 925.0, 975.0, 1025.0, 1075.0, 1125.0, 1175.0, 1225.0, 1275.0, 1325.0, 1375.0, 1425.0, 1475.0]
     tke_in = FT[0.2726, 0.5479, 0.6597, 0.7079, 0.7285, 0.7343, 0.7319, 0.7252, 0.7166, 0.7064, 0.6887, 0.6317,
                 0.6362, 0.6266, 0.5832, 0.4633, 0.0504, 0.0001, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
     not_type_stable_spline = Dierckx.Spline1D(z_in, tke_in; k = 1)
-    return x -> FT(not_type_stable_spline(x))
+    return ZProfile(x -> FT(not_type_stable_spline(x)))
 end
 
diff --git a/src/profiles/Dycoms_RF02.jl b/src/profiles/Dycoms_RF02.jl
@@ -1,33 +1,33 @@
 """ [Ackerman2009](@cite) """
 Dycoms_RF02_q_tot(::Type{FT}) where {FT} =
-    z -> if z <= 795.0
+    ZProfile(z -> if z <= 795.0
         FT(9.45) / FT(1000.0)
     else
         (FT(5) - FT(3) * (FT(1) - exp(-(z - FT(795)) / FT(500)))) / FT(1000)
-    end
+    end)
 
 """ [Ackerman2009](@cite) """
 Dycoms_RF02_θ_liq_ice(::Type{FT}) where {FT} =
-    z -> if z <= 795.0
+    ZProfile(z -> if z <= 795.0
         FT(288.3)
     else
         FT(295.0) + (z - FT(795))^FT(1.0 / 3.0)
-    end
+    end)
 
 """ [Ackerman2009](@cite) """
 function Dycoms_RF02_u(::Type{FT}) where {FT}
-    return z -> FT(3) + FT(4.3) * z / FT(1000) - FT(5)
+    return ZProfile(z -> FT(3) + FT(4.3) * z / FT(1000) - FT(5))
 end
 
 """ [Ackerman2009](@cite) """
 function Dycoms_RF02_v(::Type{FT}) where {FT}
-    return z -> FT(-9) + FT(5.6) * z / FT(1000) - FT(-5.5)
+    return ZProfile(z -> FT(-9) + FT(5.6) * z / FT(1000) - FT(-5.5))
 end
 
 """ [Ackerman2009](@cite) """
 Dycoms_RF02_tke(::Type{FT}) where {FT} =
-    z -> if z <= 795.0
+    ZProfile(z -> if z <= 795.0
         FT(1) - z / FT(1000)
     else
         FT(0)
-    end
+    end)