diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 3a4570c4fb..087a9a4e8b 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -93,20 +93,24 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( FT = eltype(Y) n = n_mass_flux_subdomains(turbconv_model) (; ᶜΦ) = p.core - (; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed + (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed (; q_tot) = p.precomputed.ᶜspecific (; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) = p.precomputed.sfc_conditions (; ᶜρaʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed - (; ᶠu³⁰, ᶜK⁰) = p.precomputed + (; ᶜK⁰) = p.precomputed + (; params) = p thermo_params = CAP.thermodynamics_params(params) turbconv_params = CAP.turbconv_params(params) + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1)) uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1)) - u³_int_halflevel = Fields.field_values(Fields.level(ᶠu³, half)) + u³_int_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³), half)) h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1)) K_int_level = Fields.field_values(Fields.level(ᶜK, 1)) q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1)) @@ -194,7 +198,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( ρaʲs_int_level = Fields.field_values(Fields.level(ᶜρaʲs, 1)) u³ʲs_int_halflevel = Fields.field_values(Fields.level(ᶠu³ʲs, half)) - u³⁰_int_halflevel = Fields.field_values(Fields.level(ᶠu³⁰, half)) + u³⁰_int_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), half)) K⁰_int_level = Fields.field_values(Fields.level(ᶜK⁰, 1)) set_diagnostic_edmfx_env_quantities_level!( ρ_int_level, @@ -305,7 +309,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( (; dt) = p dt = float(dt) (; ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core - (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed + (; ᶜp, ᶜts, ᶜh_tot, ᶜK) = p.precomputed (; q_tot) = p.precomputed.ᶜspecific (; ᶜρaʲs, @@ -321,7 +325,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( ᶠnh_pressure³_buoyʲs, ᶠnh_pressure³_dragʲs, ) = p.precomputed - (; ᶠu³⁰, ᶜK⁰, ᶜtke⁰) = p.precomputed + (; ᶜK⁰, ᶜtke⁰) = p.precomputed if precip_model isa Microphysics1Moment ᶜq_liqʲs = p.precomputed.ᶜq_liqʲs @@ -343,6 +347,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( ᶜ∇Φ₃ = p.scratch.ᶜtemp_C3 @. ᶜ∇Φ₃ = ᶜgradᵥ(ᶠΦ) + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + z_sfc_halflevel = Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half)) @@ -350,7 +357,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( for i in 2:Spaces.nlevels(axes(Y.c)) ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i)) uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i)) - u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half)) + u³_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³), i - half)) K_level = Fields.field_values(Fields.level(ᶜK, i)) h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i)) q_tot_level = Fields.field_values(Fields.level(q_tot, i)) @@ -372,9 +379,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( ∇Φ₃_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ₃, i - 1)) ∇Φ₃_data_prev_level = ∇Φ₃_prev_level.components.data.:1 ρ_prev_level = Fields.field_values(Fields.level(Y.c.ρ, i - 1)) - u³_prev_halflevel = Fields.field_values(Fields.level(ᶠu³, i - 1 - half)) + u³_prev_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³), i - 1 - half)) u³⁰_prev_halflevel = - Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half)) + Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), i - 1 - half)) u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1 K_prev_level = Fields.field_values(Fields.level(ᶜK, i - 1)) h_tot_prev_level = Fields.field_values(Fields.level(ᶜh_tot, i - 1)) @@ -853,7 +860,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( end ρaʲs_level = Fields.field_values(Fields.level(ᶜρaʲs, i)) u³ʲs_halflevel = Fields.field_values(Fields.level(ᶠu³ʲs, i - half)) - u³⁰_halflevel = Fields.field_values(Fields.level(ᶠu³⁰, i - half)) + u³⁰_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), i - half)) K⁰_level = Fields.field_values(Fields.level(ᶜK⁰, i)) set_diagnostic_edmfx_env_quantities_level!( ρ_level, @@ -883,13 +890,14 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_top_bc!( ) n = n_mass_flux_subdomains(p.atmos.turbconv_model) (; ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p.precomputed - (; ᶠu³⁰, ᶠu³ʲs, ᶜuʲs, ᶠnh_pressure³_buoyʲs, ᶠnh_pressure³_dragʲs) = + (; ᶠu³ʲs, ᶜuʲs, ᶠnh_pressure³_buoyʲs, ᶠnh_pressure³_dragʲs) = p.precomputed (; precip_model) = p.atmos # set values for the top level i_top = Spaces.nlevels(axes(Y.c)) - u³⁰_halflevel = Fields.field_values(Fields.level(ᶠu³⁰, i_top + half)) + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + u³⁰_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), i_top + half)) @. u³⁰_halflevel = CT3(0) for j in 1:n @@ -960,7 +968,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! ᶜdz = Fields.Δz_field(axes(Y.c)) (; params) = p (; dt) = p - (; ᶜp, ᶜu, ᶜts) = p.precomputed + (; ᶜp, ᶜts) = p.precomputed (; ustar, obukhov_length) = p.precomputed.sfc_conditions (; ᶜtke⁰) = p.precomputed (; @@ -974,13 +982,16 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! thermo_params = CAP.thermodynamics_params(params) ᶜlg = Fields.local_geometry_field(Y.c) + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ,Y.c.ρ, Y.f.u₃) + if p.atmos.turbconv_model isa DiagnosticEDMFX - (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs, ᶠu³⁰, ᶜu⁰) = p.precomputed + (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs) = p.precomputed elseif p.atmos.turbconv_model isa EDOnlyEDMFX - ᶠu³⁰ = p.precomputed.ᶠu³ ᶜu⁰ = ᶜu end - @. ᶜu⁰ = C123(Y.c.uₕ) + ᶜinterp(C123(ᶠu³⁰)) # Set here, but used in a different function + + ᶜu⁰ = ᶜu_lazy(Y.c.uₕ, ᶠu³⁰) @. ᶜlinear_buoygrad = buoyancy_gradients( BuoyGradMean(), @@ -1039,7 +1050,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! ρatke_flux_values = Fields.field_values(ρatke_flux) ρ_int_values = Fields.field_values(Fields.level(Y.c.ρ, 1)) - u_int_values = Fields.field_values(Fields.level(ᶜu, 1)) + u_int_values = Fields.field_values(Fields.level(Base.materialize(ᶜu), 1)) ustar_values = Fields.field_values(ustar) int_local_geometry_values = Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1)) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index e35dcb7bca..158218dbfe 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -48,12 +48,14 @@ function set_precipitation_velocities!( moisture_model::NonEquilMoistModel, precip_model::Microphysics1Moment, ) - (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed + (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts) = p.precomputed (; ᶜΦ) = p.core cmc = CAP.microphysics_cloud_params(p.params) cmp = CAP.microphysics_1m_params(p.params) thp = CAP.thermodynamics_params(p.params) + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) + # compute the precipitation terminal velocity [m/s] @. ᶜwᵣ = CM1.terminal_velocity( cmp.pr, @@ -104,7 +106,7 @@ function set_precipitation_velocities!( moisture_model::NonEquilMoistModel, precip_model::Microphysics2Moment, ) - (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed + (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts) = p.precomputed (; ᶜΦ) = p.core cm1c = CAP.microphysics_cloud_params(p.params) @@ -112,6 +114,7 @@ function set_precipitation_velocities!( cm2p = CAP.microphysics_2m_params(p.params) thp = CAP.thermodynamics_params(p.params) + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) # compute the precipitation terminal velocity [m/s] # TODO sedimentation of snow is based on the 1M scheme @. ᶜwnᵣ = getindex( @@ -280,7 +283,7 @@ function set_precipitation_cache!( end function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) (; dt) = p - (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed + (; ᶜts, ᶜwᵣ, ᶜwₛ) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed (; q_tot, q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index fe74dd3d04..85acc04c71 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -43,8 +43,6 @@ function implicit_precomputed_quantities(Y, atmos) n = n_mass_flux_subdomains(turbconv_model) gs_quantities = (; ᶜspecific = Base.materialize(ᶜspecific_gs_tracers(Y)), - ᶜu = similar(Y.c, C123{FT}), - ᶠu³ = similar(Y.f, CT3{FT}), ᶠu = similar(Y.f, CT123{FT}), ᶜK = similar(Y.c, FT), ᶜts = similar(Y.c, TST), @@ -72,8 +70,6 @@ function implicit_precomputed_quantities(Y, atmos) ᶜmse⁰ = similar(Y.c, FT), ᶜq_tot⁰ = similar(Y.c, FT), ᶠu₃⁰ = similar(Y.f, C3{FT}), - ᶜu⁰ = similar(Y.c, C123{FT}), - ᶠu³⁰ = similar(Y.f, CT3{FT}), ᶜK⁰ = similar(Y.c, FT), ᶜts⁰ = similar(Y.c, TST), ᶜρ⁰ = similar(Y.c, FT), @@ -221,8 +217,6 @@ function precomputed_quantities(Y, atmos) ᶜturb_entrʲs = similar(Y.c, NTuple{n, FT}), ᶠnh_pressure³_buoyʲs = similar(Y.f, NTuple{n, CT3{FT}}), ᶠnh_pressure³_dragʲs = similar(Y.f, NTuple{n, CT3{FT}}), - ᶠu³⁰ = similar(Y.f, CT3{FT}), - ᶜu⁰ = similar(Y.c, C123{FT}), ᶜK⁰ = similar(Y.c, FT), ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}), ᶜK_u = similar(Y.c, FT), @@ -267,10 +261,13 @@ function precomputed_quantities(Y, atmos) ) end -# Interpolates the third contravariant component of Y.c.uₕ to cell faces. -function compute_ᶠuₕ³(ᶜuₕ, ᶜρ) - ᶜJ = Fields.local_geometry_field(ᶜρ).J - return @. lazy(ᶠwinterp(ᶜρ * ᶜJ, 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³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³, ρ) + ᶜu = ᶜu_lazy(ᶜuₕ, ᶠu₃) + ᶠu³ = ᶠu³_lazy(ᶜuₕ, ρ, ᶠu₃) + ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃) + return nothing end """ @@ -323,14 +320,6 @@ function set_velocity_at_top!(Y, turbconv_model) return nothing end -# 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₃) - ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃) - return nothing -end function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model) ρaʲs(sgsʲs) = map(sgsʲ -> sgsʲ.ρa, sgsʲs) @@ -455,21 +444,26 @@ quantities are updated. NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) (; turbconv_model, moisture_model, precip_model) = p.atmos (; ᶜΦ) = p.core - (; ᶜspecific, ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed + (; ᶜspecific, ᶠu, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed ᶠuₕ³ = p.scratch.ᶠtemp_CT3 n = n_mass_flux_subdomains(turbconv_model) thermo_params = CAP.thermodynamics_params(p.params) thermo_args = (thermo_params, moisture_model, precip_model) ᶜspecific .= ᶜspecific_gs_tracers(Y) - @. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ) + #@. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ) + @. ᶠuₕ³ = $ᶠuₕ³_lazy(Y.c.uₕ, Y.c.ρ) # TODO: We might want to move this to dss! (and rename dss! to something # like enforce_constraints!). set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model) set_velocity_at_top!(Y, turbconv_model) - set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶠuₕ³) + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + + set_velocity_quantities!(ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶠuₕ³, Y.c.ρ) + ᶜJ = Fields.local_geometry_field(Y.c).J @. ᶠu = CT123(ᶠwinterp(Y.c.ρ * ᶜJ, CT12(ᶜu))) + CT123(ᶠu³) if n > 0 @@ -517,7 +511,7 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t) (; turbconv_model, moisture_model, precip_model, cloud_model) = p.atmos (; vert_diff, call_cloud_diagnostics_per_stage) = p.atmos (; ᶜΦ) = p.core - (; ᶜu, ᶜts, ᶜp) = p.precomputed + (; ᶜts, ᶜp) = p.precomputed ᶠuₕ³ = p.scratch.ᶠtemp_CT3 # updated in set_implicit_precomputed_quantities! thermo_params = CAP.thermodynamics_params(p.params) diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 3cd1344850..8765242592 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -22,7 +22,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( (; turbconv_model) = p.atmos (; ᶜΦ,) = p.core (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed - (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = + (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment @@ -76,8 +76,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( turbconv_model, ) end - set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) - set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) + + set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) # Sets \ᶠu₃⁰ + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, ᶠu₃⁰) + set_velocity_quantities!(ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³, Y.c.ρ) + # @. ᶜK⁰ += ᶜtke⁰ if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment @@ -132,7 +135,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!( ᶜq_snoʲ = Y.c.sgsʲs.:($j).q_sno end - set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³) + ᶠu³ʲ = ᶠu³_lazy(Y.c.uₕ, ᶜρʲ, ᶠu₃ʲ) + set_velocity_quantities!(ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³, ᶜρʲ) + @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2 if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment @@ -367,7 +372,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos FT = eltype(params) n = n_mass_flux_subdomains(turbconv_model) - (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed + (; ᶜtke⁰, ᶜp, ᶜρa⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed (; ᶜmixing_length_tuple, ᶜmixing_length, @@ -398,6 +403,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ᶜvert_div = p.scratch.ᶜtemp_scalar ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2 ᶜw_vert_div = p.scratch.ᶜtemp_scalar_3 + + ᶜu = Base.materialize(ᶜu_lazy(Y.c.uₕ, Y.f.u₃)) + for j in 1:n # entrainment/detrainment @. ᶜentrʲs.:($$j) = entrainment( @@ -496,6 +504,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos # TODO: Make strain_rate_norm calculation a function in eddy_diffusion_closures # TODO: Currently the shear production only includes vertical gradients ᶠu⁰ = p.scratch.ᶠtemp_C123 + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰) ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰) diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl index 510619ab42..51bf319c32 100644 --- a/src/diagnostics/Diagnostics.jl +++ b/src/diagnostics/Diagnostics.jl @@ -58,6 +58,7 @@ import ..PrognosticSurfaceTemperature import ..draft_area import ..compute_gm_mixing_length! import ..horizontal_integral_at_boundary +import ..ᶜu_lazy, ..ᶠu³_lazy # We need the abbreviations for symbols like curl, grad, and so on include(joinpath("..", "utils", "abbreviations.jl")) diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl index b881b2bc4f..64c4cc0678 100644 --- a/src/diagnostics/core_diagnostics.jl +++ b/src/diagnostics/core_diagnostics.jl @@ -74,9 +74,9 @@ add_diagnostic_variable!( comments = "Eastward (zonal) wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - return copy(u_component.(Geometry.UVector.(cache.precomputed.ᶜu))) + return copy(u_component.(Geometry.UVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃)))) else - out .= u_component.(Geometry.UVector.(cache.precomputed.ᶜu)) + out .= u_component.(Geometry.UVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃))) end end, ) @@ -92,9 +92,9 @@ add_diagnostic_variable!( comments = "Northward (meridional) wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - return copy(v_component.(Geometry.VVector.(cache.precomputed.ᶜu))) + return copy(v_component.(Geometry.VVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃)))) else - out .= v_component.(Geometry.VVector.(cache.precomputed.ᶜu)) + out .= v_component.(Geometry.VVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃))) end end, ) @@ -199,7 +199,8 @@ add_diagnostic_variable!( units = "s^-1", comments = "Vertical component of relative vorticity", compute! = (out, state, cache, time) -> begin - vort = @. w_component.(Geometry.WVector(curlₕ(cache.precomputed.ᶜu))) + ᶜu = Base.materialize(ᶜu_lazy(state.c.uₕ, state.f.u₃)) + vort = @. w_component.(Geometry.WVector(curlₕ(ᶜu))) # We need to ensure smoothness, so we call DSS Spaces.weighted_dss!(vort) if isnothing(out) @@ -552,13 +553,13 @@ add_diagnostic_variable!( if isnothing(out) return copy( u_component.( - Geometry.UVector.(Fields.level(cache.precomputed.ᶜu, 1)), + Geometry.UVector.(Fields.level(Base.materialize(ᶜu_lazy(state.c.uₕ, state.f.u₃)), 1)), ), ) else out .= u_component.( - Geometry.UVector.(Fields.level(cache.precomputed.ᶜu, 1)), + Geometry.UVector.(Fields.level(Base.materialize(ᶜu_lazy(state.c.uₕ, state.f.u₃)), 1)), ) end end, @@ -577,13 +578,13 @@ add_diagnostic_variable!( if isnothing(out) return copy( v_component.( - Geometry.VVector.(Fields.level(cache.precomputed.ᶜu, 1)), - ), + Geometry.UVector.(Fields.level(Base.materialize(ᶜu_lazy(state.c.uₕ, state.f.u₃)), 1)), + ) ) else out .= v_component.( - Geometry.VVector.(Fields.level(cache.precomputed.ᶜu, 1)), + Geometry.UVector.(Fields.level(Base.materialize(ᶜu_lazy(state.c.uₕ, state.f.u₃)), 1)), ) end end, @@ -1369,11 +1370,11 @@ add_diagnostic_variable!( comments = "Error of steady-state eastward (zonal) wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - u_component.(Geometry.UVWVector.(cache.precomputed.ᶜu)) .- - u_component.(cache.steady_state_velocity.ᶜu) + u_component.(Geometry.UVWVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃))) .- + u_component.(cache.steady_state_velocity.ᶜu) else out .= - u_component.(Geometry.UVWVector.(cache.precomputed.ᶜu)) .- + u_component.(Geometry.UVWVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃))) .- u_component.(cache.steady_state_velocity.ᶜu) end end, @@ -1387,11 +1388,11 @@ add_diagnostic_variable!( comments = "Error of steady-state northward (meridional) wind component", compute! = (out, state, cache, time) -> begin if isnothing(out) - v_component.(Geometry.UVWVector.(cache.precomputed.ᶜu)) .- - v_component.(cache.steady_state_velocity.ᶜu) + v_component.(Geometry.UVWVector.(ᶜu_lazy(state.c.uₕ,state.f.u₃))) .- + v_component.(cache.steady_state_velocity.ᶜu) else out .= - v_component.(Geometry.UVWVector.(cache.precomputed.ᶜu)) .- + v_component.(Geometry.UVWVector.(ᶜu_lazy(state.c.uₕ, state.f.u₃))) .- v_component.(cache.steady_state_velocity.ᶜu) end end, diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index ec6aab3b8c..01e81c3002 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -715,10 +715,12 @@ function compute_waen!( time, turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) + ᶠu³⁰ = ᶠu³_lazy(state.c.uₕ,state.c.ρ, state.f.u₃) + ᶜu⁰ = ᶜu_lazy(state.c.uₕ, ᶠu³⁰) # Set here (lazy), but used elsewhere if isnothing(out) - return copy(w_component.(Geometry.WVector.(cache.precomputed.ᶜu⁰))) + return copy(w_component.(Geometry.WVector.(ᶜu⁰))) else - out .= w_component.(Geometry.WVector.(cache.precomputed.ᶜu⁰)) + out .= w_component.(Geometry.WVector.(ᶜu⁰)) end end diff --git a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl index 1ac289a006..872fcccdfa 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -30,7 +30,7 @@ function set_smagorinsky_lilly_precomputed_quantities!(Y, p) (; atmos, precomputed, scratch, params) = p c_smag = CAP.c_smag(params) Pr_t = CAP.Prandtl_number_0(CAP.turbconv_params(params)) - (; ᶜu, ᶠu³, ᶜts, ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶠD_smag) = precomputed + (; ᶜts, ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶠD_smag) = precomputed FT = eltype(Y) grav = CAP.grav(params) thermo_params = CAP.thermodynamics_params(params) @@ -46,6 +46,8 @@ function set_smagorinsky_lilly_precomputed_quantities!(Y, p) axis_uvw = (Geometry.UVWAxis(),) # Compute UVW velocities + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) ᶜu_uvw = @. ᶜtemp_UVW = UVW(ᶜu) ᶠu_uvw = @. ᶠtemp_UVW = UVW(ᶠinterp(Y.c.uₕ)) + UVW(ᶠu³) diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index a932682049..3ea2bc6ff2 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -25,7 +25,7 @@ Specifically, this function calculates: Arguments: - `Yₜ`: The tendency state vector, modified in place. - `Y`: The current state vector. -- `p`: Cache containing parameters, precomputed fields (e.g., velocities `ᶜu`, +- `p`: Cache containing parameters, precomputed fields (e.g., velocities `ᶜu⁰`, `ᶜuʲs`; pressure `ᶜp`; kinetic energy `ᶜK`; total enthalpy `ᶜh_tot`), and core components (e.g., geopotential `ᶜΦ`). - `t`: Current simulation time (not directly used in calculations). @@ -36,12 +36,13 @@ Modifies `Yₜ.c.ρ`, `Yₜ.c.ρe_tot`, `Yₜ.c.uₕ`, and EDMFX-related fields NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) n = n_mass_flux_subdomains(p.atmos.turbconv_model) (; ᶜΦ) = p.core - (; ᶜu, ᶜK, ᶜp) = p.precomputed + (; ᶜK, ᶜp) = p.precomputed if p.atmos.turbconv_model isa PrognosticEDMFX (; ᶜuʲs) = p.precomputed end + ᶜu = Base.materialize(ᶜu_lazy(Y.c.uₕ, Y.f.u₃)) @. Yₜ.c.ρ -= wdivₕ(Y.c.ρ * ᶜu) if p.atmos.turbconv_model isa PrognosticEDMFX for j in 1:n @@ -64,7 +65,8 @@ NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) if p.atmos.turbconv_model isa EDOnlyEDMFX ᶜu_for_tke_advection = ᶜu elseif p.atmos.turbconv_model isa AbstractEDMF - ᶜu_for_tke_advection = p.precomputed.ᶜu⁰ + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ,Y.c.ρ, Y.f.u₃) + ᶜu_for_tke_advection = Base.materialize(ᶜu_lazy(Y.c.uₕ, ᶠu³⁰)) else error( "Unsupported turbconv_model type for TKE advection: $(typeof(p.atmos.turbconv_model))", @@ -103,12 +105,12 @@ in `Yₜ.c.sgsʲs` if applicable. """ NVTX.@annotate function horizontal_tracer_advection_tendency!(Yₜ, Y, p, t) n = n_mass_flux_subdomains(p.atmos.turbconv_model) - (; ᶜu) = p.precomputed if p.atmos.turbconv_model isa PrognosticEDMFX (; ᶜuʲs) = p.precomputed end + ᶜu = Base.materialize(ᶜu_lazy(Y.c.uₕ, Y.f.u₃)) for ρχ_name in filter(is_tracer_var, propertynames(Y.c)) @. Yₜ.c.:($$ρχ_name) -= wdivₕ(Y.c.:($$ρχ_name) * ᶜu) end @@ -158,7 +160,7 @@ Arguments: - `Yₜ`: The tendency state vector, modified in place. - `Y`: The current state vector. - `p`: Cache containing parameters, core fields (e.g., `ᶜf³`, `ᶠf¹²`, `ᶜΦ`), - precomputed fields (e.g., `ᶜu`, `ᶠu³`, `ᶜK`, EDMF velocities/TKE if applicable), + precomputed fields (e.g., `ᶜK`, EDMF velocities/TKE if applicable), atmospheric model settings (`p.atmos.numerics` for upwinding schemes), and scratch space. - `t`: Current simulation time (not directly used in calculations). @@ -174,7 +176,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J (; ᶜf³, ᶠf¹², ᶜΦ) = p.core - (; ᶜu, ᶠu³, ᶜK) = p.precomputed + (; ᶜK) = p.precomputed (; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing (; ᶜuʲs, ᶜKʲs, ᶠKᵥʲs) = n > 0 ? p.precomputed : all_nothing (; energy_upwinding, tracer_upwinding) = p.atmos.numerics @@ -183,8 +185,8 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) ᶠu³⁰ = advect_tke ? ( - turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ : - p.precomputed.ᶠu³⁰ + turbconv_model isa EDOnlyEDMFX ? ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) : + ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) ) : nothing ᶜρa⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρa⁰ : Y.c.ρ) : nothing ᶜρ⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρ⁰ : Y.c.ρ) : nothing @@ -193,6 +195,8 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) ᶜω³ = p.scratch.ᶜtemp_CT3 ᶠω¹² = p.scratch.ᶠtemp_CT12 ᶠω¹²ʲs = p.scratch.ᶠtemp_CT12ʲs + + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) if point_type <: Geometry.Abstract3DPoint @. ᶜω³ = curlₕ(Y.c.uₕ) @@ -213,6 +217,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) ᶜρ = Y.c.ρ # Full vertical advection of passive tracers (like liq, rai, etc) ... + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name if !(ρχ_name in (@name(ρe_tot), @name(ρq_tot))) ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index a87568ec75..c401c119f2 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -40,9 +40,9 @@ function edmfx_sgs_mass_flux_tendency!( n = n_mass_flux_subdomains(turbconv_model) (; edmfx_sgsflux_upwinding) = p.atmos.numerics - (; ᶠu³, ᶜh_tot) = p.precomputed + (; ᶜh_tot) = p.precomputed (; ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = p.precomputed - (; ᶜρa⁰, ᶜρ⁰, ᶠu³⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed + (; ᶜρa⁰, ᶜρ⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed if ( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment @@ -52,6 +52,9 @@ function edmfx_sgs_mass_flux_tendency!( (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} # Enthalpy fluxes. First sum up the draft fluxes # TODO: Isolate assembly of flux term pattern to a function and @@ -236,12 +239,14 @@ function edmfx_sgs_mass_flux_tendency!( a_max = CAP.max_area(turbconv_params) n = n_mass_flux_subdomains(turbconv_model) (; edmfx_sgsflux_upwinding) = p.atmos.numerics - (; ᶠu³, ᶜh_tot) = p.precomputed + (; ᶜh_tot) = p.precomputed (; ᶜρaʲs, ᶜρʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs) = p.precomputed (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J FT = eltype(Y) + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} # Enthalpy fluxes ᶠu³_diff = p.scratch.ᶠtemp_CT3 @@ -345,7 +350,7 @@ function edmfx_sgs_diffusive_flux_tendency!( (; dt, params) = p turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜρa⁰, ᶜu⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜtke⁰, ᶜmixing_length) = p.precomputed + (; ᶜρa⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜtke⁰, ᶜmixing_length) = p.precomputed if ( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment @@ -429,6 +434,8 @@ function edmfx_sgs_diffusive_flux_tendency!( # Momentum diffusion ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ᶜu⁰ = ᶜu_lazy(Y.c.uₕ, ᶠu³⁰) ᶠstrain_rate .= compute_strain_rate_face(ᶜu⁰) @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ) end @@ -450,7 +457,7 @@ function edmfx_sgs_diffusive_flux_tendency!( (; dt, params) = p turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜu, ᶜh_tot, ᶜtke⁰, ᶜmixing_length) = p.precomputed + (; ᶜh_tot, ᶜtke⁰, ᶜmixing_length) = p.precomputed (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() @@ -505,6 +512,7 @@ function edmfx_sgs_diffusive_flux_tendency!( # Momentum diffusion ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) ᶠstrain_rate .= compute_strain_rate_face(ᶜu) @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ) end diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index ae394835a9..1f0cb51880 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -44,7 +44,7 @@ function edmfx_tke_tendency!( ) n = n_mass_flux_subdomains(turbconv_model) (; ᶜturb_entrʲs, ᶜentrʲs, ᶜdetrʲs, ᶠu³ʲs) = p.precomputed - (; ᶠu³⁰, ᶠu³, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰) = p.precomputed + (; ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰) = p.precomputed (; ᶜK_u, ᶜK_h) = p.precomputed ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ nh_pressure3_buoyʲs = @@ -55,6 +55,10 @@ function edmfx_tke_tendency!( p.precomputed.ᶠnh_pressure₃_dragʲs : p.precomputed.ᶠnh_pressure³_dragʲs ᶜtke_press = p.scratch.ᶜtemp_scalar @. ᶜtke_press = 0 + + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + for j in 1:n ᶜρaʲ = turbconv_model isa PrognosticEDMFX ? Y.c.sgsʲs.:($j).ρa : diff --git a/src/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl index 9a5b4b15b8..981390735a 100644 --- a/src/prognostic_equations/gm_sgs_closures.jl +++ b/src/prognostic_equations/gm_sgs_closures.jl @@ -67,7 +67,7 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p) ᶜdz = Fields.Δz_field(axes(Y.c)) ᶜlg = Fields.local_geometry_field(Y.c) - (; ᶜts, ᶠu³, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed + (; ᶜts, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed @. ᶜlinear_buoygrad = buoyancy_gradients( BuoyGradMean(), @@ -83,6 +83,7 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p) # TODO: move strain rate calculation to separate function ᶠu = p.scratch.ᶠtemp_C123 + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) @. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³) ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW ᶜstrain_rate .= compute_strain_rate_center(ᶠu) diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index 662691b662..3839a22b5b 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -108,10 +108,13 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end + # Conflicting broadcast rules + ᶜu = Base.materialize(ᶜu_lazy(Y.c.uₕ, Y.f.u₃)) + # Grid scale hyperdiffusion @. ᶜ∇²u = - C123(wgradₕ(divₕ(p.precomputed.ᶜu))) - - C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜu)))) + C123(wgradₕ(divₕ(ᶜu))) - + C123(wcurlₕ(C123(curlₕ(ᶜu)))) @. ᶜ∇²specific_energy = wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 0e1b0a3513..24391ddf30 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -139,8 +139,9 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) ᶜJ = Fields.local_geometry_field(Y.c).J ᶠJ = Fields.local_geometry_field(Y.f).J (; ᶠgradᵥ_ᶜΦ) = p.core - (; ᶜh_tot, ᶠu³, ᶜp) = p.precomputed + (; ᶜh_tot, ᶜp) = p.precomputed + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠu³) # Central vertical advection of active tracers (e_tot and q_tot) diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index a9de2af6a4..484240c390 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -337,7 +337,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) (; matrix) = cache (; params) = p (; ᶜΦ, ᶠgradᵥ_ᶜΦ) = p.core - (; ᶠu³, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed + (; ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, @@ -378,6 +378,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ᶠz = Fields.coordinate_field(Y.f).z zmax = z_max(axes(Y.f)) + ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃) + ᶜkappa_m = p.scratch.ᶜtemp_scalar @. ᶜkappa_m = TD.gas_constant_air(thermo_params, ᶜts) / TD.cv_m(thermo_params, ᶜts) diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 569e65cfc3..425e92d4f2 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -76,12 +76,12 @@ function vertical_diffusion_boundary_layer_tendency!( ) FT = eltype(Y) α_vert_diff_tracer = CAP.α_vert_diff_tracer(p.params) - (; ᶜu, ᶜh_tot, ᶜspecific, ᶜK_u, ᶜK_h) = p.precomputed + (; ᶜh_tot, ᶜspecific, ᶜK_u, ᶜK_h) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ if !disable_momentum_vertical_diffusion(p.atmos.vert_diff) ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW - ᶠstrain_rate .= compute_strain_rate_face(ᶜu) + ᶠstrain_rate = compute_strain_rate_face(ᶜu_lazy(Y.c.uₕ, Y.f.u₃)) @. Yₜ.c.uₕ -= C12( ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_u) * ᶠstrain_rate) / Y.c.ρ, ) diff --git a/src/surface_conditions/SurfaceConditions.jl b/src/surface_conditions/SurfaceConditions.jl index 441de9286d..31b183f21e 100644 --- a/src/surface_conditions/SurfaceConditions.jl +++ b/src/surface_conditions/SurfaceConditions.jl @@ -12,7 +12,7 @@ import ..PrescribedSurfaceTemperature import ..gcm_driven_timeseries import ..CT1, ..CT2, ..C12, ..CT12, ..C3 -import ..unit_basis_vector_data, ..projected_vector_data +import ..unit_basis_vector_data, ..projected_vector_data, ..ᶜu_lazy import ..get_wstar import ClimaCore: DataLayouts, Geometry, Fields diff --git a/src/surface_conditions/surface_conditions.jl b/src/surface_conditions/surface_conditions.jl index 107e480f46..2e2bac1ac6 100644 --- a/src/surface_conditions/surface_conditions.jl +++ b/src/surface_conditions/surface_conditions.jl @@ -15,13 +15,14 @@ function update_surface_conditions!(Y, p, t) ) int_local_geometry_values = Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1)) - (; ᶜts, ᶜu, sfc_conditions) = p.precomputed + (; ᶜts, sfc_conditions) = p.precomputed (; params, sfc_setup, atmos) = p thermo_params = CAP.thermodynamics_params(params) surface_fluxes_params = CAP.surface_fluxes_params(params) surface_temp_params = CAP.surface_temp_params(params) int_ts_values = Fields.field_values(Fields.level(ᶜts, 1)) - int_u_values = Fields.field_values(Fields.level(ᶜu, 1)) + ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃) + int_u_values = Fields.field_values(Fields.level(Base.materialize(ᶜu), 1)) int_z_values = Fields.field_values(Fields.level(Fields.coordinate_field(Y.c).z, 1)) sfc_conditions_values = Fields.field_values(sfc_conditions) diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl index 8ea41a7564..d1ecb1fa17 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -96,9 +96,20 @@ function compute_strain_rate_center(u::Fields.Field) ) / 2, ) end +function compute_strain_rate_center(u::Broadcast.Broadcasted) + @assert eltype(u) <: C123 + axis_uvw = Geometry.UVWAxis() + return @. lazy( + ( + Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(u))) + + adjoint(Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(u)))) + ) / 2, + ) +end """ ϵ .= compute_strain_rate_face(u::Field) + ϵ .= compute_strain_rate_face(u::Broadcast.Broadcasted) Compute the strain_rate at cell faces from velocity at cell centers. """ @@ -118,6 +129,22 @@ function compute_strain_rate_face(u::Fields.Field) ) / 2, ) end +function compute_strain_rate_face(u::Broadcast.Broadcasted) + @assert eltype(u) <: C123 + ∇ᵥuvw_boundary = + Geometry.outer(Geometry.WVector(0), Geometry.UVWVector(0, 0, 0)) + ᶠgradᵥ = Operators.GradientC2F( + bottom = Operators.SetGradient(∇ᵥuvw_boundary), + top = Operators.SetGradient(∇ᵥuvw_boundary), + ) + axis_uvw = Geometry.UVWAxis() + return @. lazy( + ( + Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u))) + + adjoint(Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u)))) + ) / 2, + ) +end """ g³³_field(space) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 443c5eb359..0ff39273ed 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -1,5 +1,19 @@ import ClimaCore.MatrixFields: @name +# Interpolates the third contravariant component of Y.c.uₕ to cell faces. +# Group velocity and kinetic energy computations. +@inline function ᶜu_lazy(ᶜuₕ, ᶠu₃) + return @. lazy(C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))) +end +@inline function ᶠu³_lazy(ᶜuₕ, ᶜρ, ᶠu₃) + ᶠuₕ³ = ᶠuₕ³_lazy(ᶜuₕ, ᶜρ) + return @. lazy(ᶠuₕ³ + CT3(ᶠu₃)) +end +@inline function ᶠuₕ³_lazy(ᶜuₕ, ᶜρ) + ᶜJ = Fields.local_geometry_field(ᶜρ).J + return @. lazy(ᶠwinterp(ᶜρ * ᶜJ, CT3(ᶜuₕ))) +end + """ specific(ρχ, ρ) specific(ρaχ, ρa, ρχ, ρ, turbconv_model)