autodiff ahhhhh

Author: szy21

src/cache/precomputed_quantities.jl

@@ -464,13 +464,8 @@ function thermo_state_gs(
 )
     e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
     T = TD.air_temperature(thermo_params, e_int)
-    FT = eltype(thermo_params)
-    return (;
-        T = T,
-        q_tot_safe = zero(FT),
-        q_liq_rai = zero(FT),
-        q_ice_sno = zero(FT),
-    )
+    # Use zero(T) to match type of T (supports autodiff with dual numbers)
+    return ThermoState(T, zero(T), zero(T), zero(T))
 end
 
 function thermo_state_gs(
@@ -486,12 +481,7 @@ function thermo_state_gs(
     q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
     # Use saturation_adjustment once for all quantities
     sa_result = TD.saturation_adjustment(thermo_params, TD.ρe(), ρ, e_int, q_tot)
-    return (;
-        T = sa_result.T,
-        q_tot_safe = max(0, q_tot),
-        q_liq_rai = sa_result.q_liq,
-        q_ice_sno = sa_result.q_ice,
-    )
+    return ThermoState(sa_result.T, max(0, q_tot), sa_result.q_liq, sa_result.q_ice)
 end
 
 function thermo_state_gs(
@@ -508,12 +498,7 @@ function thermo_state_gs(
     q_liq_rai = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ)
     q_ice_sno = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ)
     T = TD.air_temperature(thermo_params, e_int, q_tot, q_liq_rai, q_ice_sno)
-    return (;
-        T = T,
-        q_tot_safe = max(0, q_tot),
-        q_liq_rai = max(0, q_liq_rai),
-        q_ice_sno = max(0, q_ice_sno),
-    )
+    return ThermoState(T, max(0, q_tot), max(0, q_liq_rai), max(0, q_ice_sno))
 end
 
 function eddy_diffusivity_coefficient_H(D₀, H, z_sfc, z)
@@ -526,6 +511,11 @@ function eddy_diffusivity_coefficient(C_E, norm_v_a, z_a, p)
     return p > p_pbl ? K_E : K_E * exp(-((p_pbl - p) / p_strato)^2)
 end
 
+@inline get_thermo_state_value(thermo_state::ThermoState, ::Val{:T}) = thermo_state.T
+@inline get_thermo_state_value(thermo_state::ThermoState, ::Val{:q_tot_safe}) = thermo_state.q_tot_safe
+@inline get_thermo_state_value(thermo_state::ThermoState, ::Val{:q_liq_rai}) = thermo_state.q_liq_rai
+@inline get_thermo_state_value(thermo_state::ThermoState, ::Val{:q_ice_sno}) = thermo_state.q_ice_sno
+
 """
     set_implicit_precomputed_quantities!(Y, p, t)
 
@@ -576,13 +566,13 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
         # @. ᶜK += Y.c.ρtke / Y.c.ρ
         # TODO: We should think more about these increments before we use them.
     end
-    # Compute all thermodynamic state variables in one call (avoids multiple saturation_adjustment calls)
-    ᶜthermo_state = p.scratch.ᶜtemp_thermo_state
-    @. ᶜthermo_state = thermo_state_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
-    @. ᶜT = ᶜthermo_state.T
-    @. ᶜq_tot_safe = ᶜthermo_state.q_tot_safe
-    @. ᶜq_liq_rai = ᶜthermo_state.q_liq_rai
-    @. ᶜq_ice_sno = ᶜthermo_state.q_ice_sno
+    # Compute thermodynamic state variables
+    # Note: For EquilMoistModel, this calls saturation_adjustment 3 times per grid point.
+    ᶜthermo_state = @. lazy(thermo_state_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ))
+    @. ᶜT = get_thermo_state_value(ᶜthermo_state, Val(:T))
+    @. ᶜq_tot_safe = get_thermo_state_value(ᶜthermo_state, Val(:q_tot_safe))
+    @. ᶜq_liq_rai = get_thermo_state_value(ᶜthermo_state, Val(:q_liq_rai))
+    @. ᶜq_ice_sno = get_thermo_state_value(ᶜthermo_state, Val(:q_ice_sno))
     ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
     @. ᶜh_tot = TD.total_enthalpy(thermo_params, ᶜe_tot, ᶜT, ᶜq_tot_safe, ᶜq_liq_rai, ᶜq_ice_sno)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜT, Y.c.ρ, ᶜq_tot_safe, ᶜq_liq_rai, ᶜq_ice_sno)

src/cache/temporary_quantities.jl

@@ -29,11 +29,6 @@ function temporary_quantities(Y, atmos)
 
     FT = Spaces.undertype(center_space)
     uvw_vec = UVW(FT(0), FT(0), FT(0))
-    # Type for thermo state: (T, q_tot_safe, q_liq_rai, q_ice_sno)
-    ThermoStateType = NamedTuple{
-        (:T, :q_tot_safe, :q_liq_rai, :q_ice_sno),
-        NTuple{4, FT},
-    }
     return (;
         ᶠtemp_scalar = Fields.Field(FT, face_space), # ᶠp, ᶠρK_h
         ᶠtemp_scalar_2 = Fields.Field(FT, face_space), # ᶠρK_u
@@ -44,7 +39,6 @@ function temporary_quantities(Y, atmos)
         ᶜtemp_scalar_5 = Fields.Field(FT, center_space),
         ᶜtemp_scalar_6 = Fields.Field(FT, center_space),
         ᶜtemp_scalar_7 = Fields.Field(FT, center_space),
-        ᶜtemp_thermo_state = Fields.Field(ThermoStateType, center_space), # thermo_state_gs result
         ᶠtemp_field_level = Fields.level(Fields.Field(FT, face_space), half),
         temp_field_level = Fields.level(Fields.Field(FT, center_space), 1),
         temp_field_level_2 = Fields.level(Fields.Field(FT, center_space), 1),

src/solver/types.jl

@@ -379,6 +379,16 @@ function MixingLength(master, wall, tke, buoy, l_grid)
     return MixingLength(promote(master, wall, tke, buoy, l_grid)...)
 end
 
+struct ThermoState{FT}
+    T::FT
+    q_tot_safe::FT
+    q_liq_rai::FT
+    q_ice_sno::FT
+end
+
+function ThermoState(T, q_tot_safe, q_liq_rai, q_ice_sno)
+    return ThermoState(promote(T, q_tot_safe, q_liq_rai, q_ice_sno)...)
+end
 
 abstract type AbstractEDMF end