Add type restrictions for debugging

CliMA · Nov 22, 2024 · 9a04d07 · 9a04d07
1 parent 7e00bf3
commit 9a04d07
Showing 1 changed file with 26 additions and 20 deletions.
diff --git a/src/parameterized_tendencies/radiation/held_suarez.jl b/src/parameterized_tendencies/radiation/held_suarez.jl
@@ -2,6 +2,8 @@
 ##### Held-Suarez
 #####
 
+import Thermodynamics as TD
+import Thermodynamics.Parameters as TDP
 import ClimaCore.Spaces as Spaces
 import ClimaCore.Fields as Fields
 
@@ -32,8 +34,6 @@ function held_suarez_ΔT_y_T_equator(
     return ΔT_y, T_equator
 end
 
-height_factor(σ, σ_b) = max(0, (σ - σ_b) / (1 - σ_b))
-
 struct HeldSuarezForcingParams{FT}
     day::FT
     σ_b::FT
@@ -49,12 +49,12 @@ end
 Base.Broadcast.broadcastable(x::HeldSuarezForcingParams) = tuple(x)
 
 function compute_ΔρT(
-    thermo_params,
-    ts_surf,
-    ρ,
-    p,
-    lat,
-    z_surface,
+    thermo_params::TDP.ThermodynamicsParameters,
+    ts_surf::TD.ThermodynamicState,
+    ρ::FT,
+    p::FT,
+    lat::FT,
+    z_surface::FT,
     s::HeldSuarezForcingParams,
 )
     σ = compute_σ(thermo_params, z_surface, p, ts_surf, s)
@@ -77,10 +77,10 @@ function compute_ΔρT(
 end
 
 function compute_σ(
-    thermo_params,
-    z_surface,
-    p,
-    ts_surf,
+    thermo_params::TDP.ThermodynamicsParameters,
+    z_surface::FT,
+    p::FT,
+    ts_surf::TD.ThermodynamicState,
     s::HeldSuarezForcingParams,
 )
     p / (
@@ -91,6 +91,15 @@ function compute_σ(
     )
 end
 
+height_factor(σ::FT, σ_b::FT) where {FT} = max(0, (σ - σ_b) / (1 - σ_b))
+height_factor(
+        thermo_params::TDP.ThermodynamicsParameters,
+        z_surface::FT,
+        p::FT,
+        ts_surf::TD.ThermodynamicState,
+        s::HeldSuarezForcingParams,
+    ) = height_factor(compute_σ(thermo_params,z_surface,p,ts_surf,s), s.σ_b)
+
 function forcing_tendency!(Yₜ, Y, p, t, ::HeldSuarezForcing)
     (; params) = p
     (; ᶜp, sfc_conditions) = p.precomputed
@@ -131,14 +140,11 @@ function forcing_tendency!(Yₜ, Y, p, t, ::HeldSuarezForcing)
     @. Yₜ.c.uₕ -=
         (
             k_f * height_factor(
-                compute_σ(
-                    thermo_params,
-                    z_surface,
-                    ᶜp,
-                    sfc_conditions.ts,
-                    hs_params,
-                ),
-                σ_b,
+                thermo_params,
+                z_surface,
+                ᶜp,
+                sfc_conditions.ts,
+                hs_params,
             )
         ) * Y.c.uₕ
     @. Yₜ.c.ρe_tot -=