From 685d85839ae33e06c198315470ae3a0b72bc8fb7 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Sat, 10 May 2025 14:37:58 -0400 Subject: [PATCH] Use LazyBroadcast in set_velocity_quantities --- src/cache/precomputed_quantities.jl | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index f713f3fa34..ba61354c99 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -310,11 +310,27 @@ function set_velocity_at_top!(Y, turbconv_model) return nothing end +""" + compute_ᶜu(ᶜuₕ, ᶠu₃) +Computes the cell-centered velocity, given + - `ᶜuₕ` the cell-centered horizontal velocity + - `ᶠu₃` the cell-face `Covariant3` velocity +""" +compute_ᶜu(ᶜuₕ, ᶠu₃) = @. lazy(C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))) + +""" + compute_ᶠu³(ᶠuₕ³, ᶠu₃) +Computes the cell-face contravariant velocity, given + - `ᶠuₕ³` the cell-face horizontal `Contravariant3` velocity + - `ᶠu₃` the cell-face `Covariant3` velocity +""" +compute_ᶠu³(ᶠuₕ³, ᶠu₃) = ᶠuₕ³ + CT3(ᶠu₃) + # This is used to set the grid-scale velocity quantities ᶜu, ᶠu³, ᶜK based on # ᶠu₃, and it is also used to set the SGS quantities based on ᶠu₃⁰ and ᶠu₃ʲ. function set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³) - @. ᶜu = C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃)) - @. ᶠu³ = ᶠuₕ³ + CT3(ᶠu₃) + ᶜu .= compute_ᶜu(ᶜuₕ, ᶠu₃) + ᶠu³ .= compute_ᶠu³.(ᶠuₕ³, ᶠu₃) ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃) return nothing end @@ -451,7 +467,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) thermo_args = (thermo_params, moisture_model, precip_model) @. ᶜspecific = specific_gs(Y.c) - @. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ) + ᶠuₕ³ .= compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ) # TODO: We might want to move this to dss! (and rename dss! to something # like enforce_constraints!).