From e67f1869f22162f8fd7e784601ce28a5aa4f836f Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Tue, 23 Apr 2024 15:50:11 -0700
Subject: [PATCH] Preallocate thermal state

---
 docs/src/APIs/shared_utilities.md |  1 +
 src/shared_utilities/drivers.jl   | 24 ++++++++++++++++++++----
 src/shared_utilities/models.jl    | 15 +++++++++++++++
 src/standalone/Bucket/Bucket.jl   |  1 +
 4 files changed, 37 insertions(+), 4 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..b35e775f30 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,21 @@ function construct_atmos_ts(
     return ts_in
 end
 
+"""
+Fill the pre-allocated ts_in `Field` with a thermodynamic state.
+"""
+function set_atmos_ts!(
+    ts_in,
+    atmos::PrescribedAtmosphere{FT},
+    p,
+    thermo_params,
+) where {FT}
+    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,
@@ -208,7 +224,7 @@ function turbulent_fluxes(
     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)
+    set_atmos_ts!(p.scratch.thermal_state, atmos, p, thermo_params)
     u_air = p.drivers.u
     h_air = atmos.h
 
@@ -220,7 +236,7 @@ function turbulent_fluxes(
         h_sfc,
         r_sfc,
         d_sfc,
-        ts_air,
+        p.scratch.thermal_state,
         u_air,
         h_air,
         atmos.gustiness,
@@ -518,8 +534,8 @@ 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)
+    set_atmos_ts!(p.scratch.thermal_state, atmos, p, thermo_params)
+    return compute_ρ_sfc.(thermo_params, p.scratch.thermal_state, T_sfc)
 end
 
 
diff --git a/src/shared_utilities/models.jl b/src/shared_utilities/models.jl
index f9ccbcf9c6..197ea71fa8 100644
--- a/src/shared_utilities/models.jl
+++ b/src/shared_utilities/models.jl
@@ -355,6 +355,7 @@ function initialize_auxiliary(model::AbstractModel{FT}, state) where {FT}
     )
     if :domain ∈ propertynames(model)
         p = add_dss_buffer_to_aux(p, model.domain)
+        p = add_scratch_to_aux(p, model.domain)
     else
         error(
             "Your model does not contain a domain. If this is intended, you will need a new method of initialize_auxiliary.",
@@ -363,6 +364,20 @@ function initialize_auxiliary(model::AbstractModel{FT}, state) where {FT}
     return p
 end
 
+function add_scratch_to_aux(p, domain)
+    hasproperty(domain.space, :surface) || error("Surface domain required")
+    FT = Spaces.undertype(domain.space.surface)
+    scratch = (;
+        # Thermal state on the surface
+        thermal_state = Fields.Field(
+            Thermodynamics.PhaseEquil{FT},
+            domain.space.surface,
+        )
+    )
+    return merge(p, (; scratch))
+end
+
+
 function initialize_vars(keys, types, domain_names, state, model_name)
     FT = eltype(state)
     if length(keys) == 0
diff --git a/src/standalone/Bucket/Bucket.jl b/src/standalone/Bucket/Bucket.jl
index bb629baba8..221923c461 100644
--- a/src/standalone/Bucket/Bucket.jl
+++ b/src/standalone/Bucket/Bucket.jl
@@ -27,6 +27,7 @@ using ClimaLand:
     turbulent_fluxes,
     net_radiation,
     construct_atmos_ts,
+    set_atmos_ts!,
     compute_ρ_sfc,
     AbstractExpModel,
     heaviside,