From 7ff391a6e456faa81255063b45b3f05e155660bc Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Mon, 6 May 2024 09:54:07 -0700
Subject: [PATCH] Preallocate thermal state

---
 docs/src/APIs/shared_utilities.md |  1 +
 src/shared_utilities/drivers.jl   | 29 +++++++++++++++++++++--------
 src/shared_utilities/models.jl    |  9 +++++++++
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/docs/src/APIs/shared_utilities.md b/docs/src/APIs/shared_utilities.md
index eb748711b0..fe80503575 100644
--- a/docs/src/APIs/shared_utilities.md
+++ b/docs/src/APIs/shared_utilities.md
@@ -77,6 +77,7 @@ ClimaLand.turbulent_fluxes
 ClimaLand.turbulent_fluxes_at_a_point
 ClimaLand.radiative_fluxes_at_a_point
 ClimaLand.construct_atmos_ts
+ClimaLand.set_atmos_ts!
 ClimaLand.surface_air_density
 ClimaLand.liquid_precipitation
 ClimaLand.snow_precipitation
diff --git a/src/shared_utilities/drivers.jl b/src/shared_utilities/drivers.jl
index f10e38456f..e3a826e5b6 100644
--- a/src/shared_utilities/drivers.jl
+++ b/src/shared_utilities/drivers.jl
@@ -16,6 +16,7 @@ export AbstractAtmosphericDrivers,
     CoupledRadiativeFluxes,
     compute_ρ_sfc,
     construct_atmos_ts,
+    set_atmos_ts!,
     turbulent_fluxes,
     net_radiation,
     turbulent_fluxes_at_a_point,
@@ -175,6 +176,20 @@ function construct_atmos_ts(
     return ts_in
 end
 
+"""
+    set_atmos_ts!(ts_in, atmos::PrescribedAtmosphere{FT}, p)
+
+Fill the pre-allocated ts_in `Field` with a thermodynamic state computed from the
+atmosphere.
+"""
+function set_atmos_ts!(ts_in, atmos::PrescribedAtmosphere{FT}, p) where {FT}
+    thermo_params = LP.thermodynamic_parameters(p.parameters.earth_param_set)
+    P = p.drivers.P
+    T = p.drivers.T
+    q = p.drivers.q
+    ts_in .= Thermodynamics.PhaseEquil_pTq.(thermo_params, P, T, q)
+    return nothing
+end
 
 """
     turbulent_fluxes(atmos::PrescribedAtmosphere,
@@ -206,9 +221,6 @@ function turbulent_fluxes(
     h_sfc = surface_height(model, Y, p)
     r_sfc = surface_resistance(model, Y, p, t)
     d_sfc = displacement_height(model, Y, p)
-    thermo_params =
-        LP.thermodynamic_parameters(model.parameters.earth_param_set)
-    ts_air = construct_atmos_ts(atmos, p, thermo_params)
     u_air = p.drivers.u
     h_air = atmos.h
 
@@ -220,7 +232,7 @@ function turbulent_fluxes(
         h_sfc,
         r_sfc,
         d_sfc,
-        ts_air,
+        p.drivers.thermal_state,
         u_air,
         h_air,
         atmos.gustiness,
@@ -518,8 +530,7 @@ function surface_air_density(
 )
     thermo_params =
         LP.thermodynamic_parameters(model.parameters.earth_param_set)
-    ts_in = construct_atmos_ts(atmos, p, thermo_params)
-    return compute_ρ_sfc.(thermo_params, ts_in, T_sfc)
+    return compute_ρ_sfc.(thermo_params, p.drivers.thermal_state, T_sfc)
 end
 
 
@@ -662,8 +673,9 @@ horizontal wind speed `u`, specific humidity `q`, and CO2 concentration
 `c_co2`.
 """
 function initialize_drivers(a::PrescribedAtmosphere{FT}, coords) where {FT}
-    keys = (:P_liq, :P_snow, :T, :P, :u, :q, :c_co2)
-    types = ([FT for k in keys]...,)
+    keys = (:P_liq, :P_snow, :T, :P, :u, :q, :c_co2, :thermal_state)
+    # The thermal state is a different type
+    types = ([FT for k in keys[1:(end - 1)]]..., Thermodynamics.PhaseEquil{FT})
     domain_names = ([:surface for k in keys]...,)
     model_name = :drivers
     # intialize_vars packages the variables as a named tuple,
@@ -813,6 +825,7 @@ function make_update_drivers(a::PrescribedAtmosphere{FT}) where {FT}
         evaluate!(p.drivers.u, a.u, t)
         evaluate!(p.drivers.q, a.q, t)
         evaluate!(p.drivers.c_co2, a.c_co2, t)
+        set_atmos_ts!(p.drivers.thermal_state, a, p)
     end
     return update_drivers!
 end
diff --git a/src/shared_utilities/models.jl b/src/shared_utilities/models.jl
index f9ccbcf9c6..045da339a1 100644
--- a/src/shared_utilities/models.jl
+++ b/src/shared_utilities/models.jl
@@ -192,6 +192,14 @@ driver variables in the cache.
 add_drivers_to_cache(p, model::AbstractModel) = p
 
 
+"""
+    add_parameters_to_cache(p, model::AbstractModel)
+
+Add `parameters` to `p`.
+"""
+add_parameters_to_cache(p, model::AbstractModel) =
+    merge(p, (; model.parameters))
+
 """
     make_imp_tendency(model::AbstractImExModel)
 
@@ -418,5 +426,6 @@ function initialize(model::AbstractModel{FT}) where {FT}
     Y = initialize_prognostic(model, coords)
     p = initialize_auxiliary(model, coords)
     p = add_drivers_to_cache(p, model, coords)
+    p = add_parameters_to_cache(p, model)
     return Y, p, coords
 end