Add type restrictions for debugging

CliMA · Nov 22, 2024 · ac1b5ec · ac1b5ec
1 parent 7e00bf3
commit ac1b5ec
Showing 1 changed file with 29 additions and 22 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,14 +49,14 @@ 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,
-)
+) where {FT}
     σ = compute_σ(thermo_params, z_surface, p, ts_surf, s)
     k_a = 1 / (40 * s.day)
     k_s = 1 / (4 * s.day)
@@ -77,12 +77,12 @@ 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,
-)
+) where {FT}
     p / (
         s.MSLP * exp(
             -s.grav * z_surface / s.R_d /
@@ -91,6 +91,16 @@ 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,
+) where {FT} =
+    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 +141,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 -=