From 4f4b5fc48e7537cb4ff3e055bb5f91ff834e3dd4 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 1 Dec 2023 11:01:35 -0800
Subject: [PATCH] Avoid allocations due to surface_ct3_unit

---
 src/cache/cache.jl                  |  7 +++++++
 src/diagnostics/core_diagnostics.jl | 12 +++---------
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/cache/cache.jl b/src/cache/cache.jl
index 81d45c9938e..354e1b2813a 100644
--- a/src/cache/cache.jl
+++ b/src/cache/cache.jl
@@ -134,6 +134,9 @@ function build_cache(Y, atmos, params, surface_setup, dt, start_date)
 
     numerics = (; limiter)
 
+    sfc_local_geometry =
+        Fields.level(Fields.local_geometry_field(Y.f), Fields.half)
+
     core = (
         ᶜΦ,
         ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(ᶜΦ),
@@ -142,6 +145,10 @@ function build_cache(Y, atmos, params, surface_setup, dt, start_date)
         ᶜT = similar(Y.c, FT),
         ᶜf,
         ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}),
+        # Used by diagnostics such as hfres, evspblw
+        surface_ct3_unit = CT3.(
+            unit_basis_vector_data.(CT3, sfc_local_geometry)
+        ),
     )
 
     sfc_setup = surface_setup(params)
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 6d511d68681..007b74da34b 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -419,9 +419,7 @@ add_diagnostic_variable!(
 # Eastward and northward surface drag component (2d)
 ###
 function compute_tau!(out, state, cache, component)
-    sfc_local_geometry =
-        Fields.level(Fields.local_geometry_field(state.f), Fields.half)
-    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    (; surface_ct3_unit) = cache.core
     (; ρ_flux_uₕ) = cache.precomputed.sfc_conditions
 
     if isnothing(out)
@@ -469,9 +467,7 @@ add_diagnostic_variable!(
 ###
 function compute_hfes!(out, state, cache, time)
     (; ρ_flux_h_tot) = cache.precomputed.sfc_conditions
-    sfc_local_geometry =
-        Fields.level(Fields.local_geometry_field(state.f), Fields.half)
-    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    (; surface_ct3_unit) = cache.core
     if isnothing(out)
         return dot.(ρ_flux_h_tot, surface_ct3_unit)
     else
@@ -503,9 +499,7 @@ function compute_evspsbl!(
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
     (; ρ_flux_q_tot) = cache.precomputed.sfc_conditions
-    sfc_local_geometry =
-        Fields.level(Fields.local_geometry_field(state.f), Fields.half)
-    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    (; surface_ct3_unit) = cache.core
 
     if isnothing(out)
         return dot.(ρ_flux_q_tot, surface_ct3_unit)