From 5df20bc5f82d80bc5d901abca6037f3f13e5b5c1 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Thu, 30 Nov 2023 16:33:07 -0800 Subject: [PATCH] Avoid allocations in diagnostics with U/V/WVectors --- src/diagnostics/core_diagnostics.jl | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl index a12e5a7aca6..a45d69181fc 100644 --- a/src/diagnostics/core_diagnostics.jl +++ b/src/diagnostics/core_diagnostics.jl @@ -46,6 +46,11 @@ error_diagnostic_variable( error_diagnostic_variable(variable, model::T) where {T} = error_diagnostic_variable("Cannot compute $variable with model = $T") +# Helper functions to extract components of vectors +u_component(u::Geometry.LocalVector) = u.u +v_component(u::Geometry.LocalVector) = u.v +w_component(u::Geometry.LocalVector) = u.w + ### # Density (3d) ### @@ -74,9 +79,9 @@ add_diagnostic_variable!( comments = "Eastward (zonal) wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - return copy(Geometry.UVector.(cache.precomputed.ᶜu).components.data.:1) + return copy(u_component.(Geometry.UVector.(cache.precomputed.ᶜu))) else - out .= Geometry.UVector.(cache.precomputed.ᶜu).components.data.:1 + out .= u_component.(Geometry.UVector.(cache.precomputed.ᶜu)) end end, ) @@ -92,9 +97,9 @@ add_diagnostic_variable!( comments = "Northward (meridional) wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - return copy(Geometry.VVector.(cache.precomputed.ᶜu).components.data.:1) + return copy(v_component.(Geometry.VVector.(cache.precomputed.ᶜu))) else - out .= Geometry.VVector.(cache.precomputed.ᶜu).components.data.:1 + out .= v_component.(Geometry.VVector.(cache.precomputed.ᶜu)) end end, ) @@ -113,9 +118,9 @@ add_diagnostic_variable!( comments = "Vertical wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - return copy(Geometry.WVector.(cache.precomputed.ᶜu).components.data.:1) + return copy(w_component.(Geometry.WVector.(cache.precomputed.ᶜu))) else - out .= Geometry.WVector.(cache.precomputed.ᶜu).components.data.:1 + out .= w_component.(Geometry.WVector.(cache.precomputed.ᶜu)) end end, ) @@ -199,7 +204,7 @@ add_diagnostic_variable!( units = "s^-1", comments = "Vertical component of relative vorticity", compute! = (out, state, cache, time) -> begin - vort = @. Geometry.WVector(curlₕ(state.c.uₕ)).components.data.:1 + vort = @. w_component.(Geometry.WVector.(cache.precomputed.ᶜu)) # We need to ensure smoothness, so we call DSS Spaces.weighted_dss!(vort) if isnothing(out)