Avoid allocations due to surface_ct3_unit

CliMA · Dec 4, 2023 · 4f4b5fc · 4f4b5fc
1 parent c2aa108
commit 4f4b5fc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git 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
@@ -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)