diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 3a4570c4fb..ef15bdc4eb 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -93,12 +93,12 @@ 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 + (; ᶠu³, ᶜ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 + (; ᶠu³⁰, ᶜK⁰, ᶜts) = p.precomputed (; params) = p thermo_params = CAP.thermodynamics_params(params) @@ -107,11 +107,11 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( ρ_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)) - h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1)) + h_tot_int_level = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot(Y, thermo_params, ᶜts)), 1)) K_int_level = Fields.field_values(Fields.level(ᶜK, 1)) q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1)) - p_int_level = Fields.field_values(Fields.level(ᶜp, 1)) + p_int_level = Fields.field_values(Fields.level(Base.materialize(ᶜp(thermo_params, ᶜts)), 1)) Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1)) local_geometry_int_level = @@ -305,7 +305,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 + (; ᶠu³, ᶜts, ᶜK) = p.precomputed (; q_tot) = p.precomputed.ᶜspecific (; ᶜρaʲs, @@ -352,9 +352,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i)) u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half)) K_level = Fields.field_values(Fields.level(ᶜK, i)) - h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i)) + h_tot_level = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot(Y, thermo_params, ᶜts)), i)) q_tot_level = Fields.field_values(Fields.level(q_tot, i)) - p_level = Fields.field_values(Fields.level(ᶜp, i)) + p_level = Fields.field_values(Fields.level(Base.materialize(ᶜp(thermo_params,ᶜts)), i)) Φ_level = Fields.field_values(Fields.level(ᶜΦ, i)) local_geometry_level = Fields.field_values( Fields.level(Fields.local_geometry_field(Y.c), i), @@ -377,10 +377,10 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( Fields.field_values(Fields.level(ᶠ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)) + h_tot_prev_level = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot(Y, thermo_params, ᶜts)), i - 1)) q_tot_prev_level = Fields.field_values(Fields.level(q_tot, i - 1)) ts_prev_level = Fields.field_values(Fields.level(ᶜts, i - 1)) - p_prev_level = Fields.field_values(Fields.level(ᶜp, i - 1)) + p_prev_level = Fields.field_values(Fields.level(Base.materialize(ᶜp(thermo_params, ᶜts)), i - 1)) z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1)) dz_prev_level = Fields.field_values(Fields.level(ᶜdz, i - 1)) @@ -960,7 +960,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 + (; ᶜu, ᶜts) = p.precomputed (; ustar, obukhov_length) = p.precomputed.sfc_conditions (; ᶜtke⁰) = p.precomputed (; diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index fe74dd3d04..c93b17eea7 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -19,7 +19,6 @@ The following grid-scale quantities are treated implicitly: - `ᶜK`: kinetic energy on cell centers - `ᶜts`: thermodynamic state on cell centers - `ᶜp`: air pressure on cell centers - - `ᶜh_tot`: total enthalpy on cell centers If the `turbconv_model` is `PrognosticEDMFX`, there also two SGS versions of every quantity except for `ᶜp` (which is shared across all subdomains): - `_⁰`: value for the environment @@ -48,8 +47,6 @@ function implicit_precomputed_quantities(Y, atmos) ᶠu = similar(Y.f, CT123{FT}), ᶜK = similar(Y.c, FT), ᶜts = similar(Y.c, TST), - ᶜp = similar(Y.c, FT), - ᶜh_tot = similar(Y.c, FT), ) sgs_quantities = turbconv_model isa AbstractEDMF ? (; ᶜtke⁰ = similar(Y.c, FT)) : (;) @@ -455,7 +452,7 @@ 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, ᶠu³, ᶠu, ᶜK, ᶜts) = p.precomputed ᶠuₕ³ = p.scratch.ᶠtemp_CT3 n = n_mass_flux_subdomains(turbconv_model) thermo_params = CAP.thermodynamics_params(p.params) @@ -488,12 +485,6 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) # TODO: We should think more about these increments before we use them. end @. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ) - @. ᶜp = TD.air_pressure(thermo_params, ᶜts) - @. ᶜh_tot = TD.total_specific_enthalpy( - thermo_params, - ᶜts, - specific(Y.c.ρe_tot, Y.c.ρ), - ) if turbconv_model isa PrognosticEDMFX set_prognostic_edmf_precomputed_quantities_draft!(Y, p, ᶠuₕ³, t) @@ -517,7 +508,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 + (; ᶜu, ᶜts) = p.precomputed ᶠuₕ³ = p.scratch.ᶠtemp_CT3 # updated in set_implicit_precomputed_quantities! thermo_params = CAP.thermodynamics_params(p.params) @@ -581,7 +572,8 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t) @. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.ρ, vert_diff) elseif vert_diff isa VerticalDiffusion (; ᶜK_h) = p.precomputed - @. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vert_diff) + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + @. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp_lazy, vert_diff) end # TODO diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 3cd1344850..2ccf3eada7 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -21,20 +21,23 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( thermo_params = CAP.thermodynamics_params(p.params) (; turbconv_model) = p.atmos (; ᶜΦ,) = p.core - (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed - (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = + (; ᶜK) = p.precomputed + (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜts) = p.precomputed if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed end + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + ᶜh_tot_lazy = ᶜh_tot(Y,thermo_params, ᶜts) + @. ᶜρa⁰ = ρa⁰(Y.c) @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) @. ᶜmse⁰ = divide_by_ρa( - Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), + Y.c.ρ * (ᶜh_tot_lazy - ᶜK) - ρamse⁺(Y.c.sgsʲs), ᶜρa⁰, - Y.c.ρ * (ᶜh_tot - ᶜK), + Y.c.ρ * (ᶜh_tot_lazy - ᶜK), Y.c.ρ, turbconv_model, ) @@ -83,12 +86,15 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( p.atmos.precip_model isa Microphysics1Moment @. ᶜts⁰ = TD.PhaseNonEquil_phq( thermo_params, - ᶜp, + ᶜp_lazy, ᶜmse⁰ - ᶜΦ, TD.PhasePartition(ᶜq_tot⁰, ᶜq_liq⁰ + ᶜq_rai⁰, ᶜq_ice⁰ + ᶜq_sno⁰), ) else - @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, + ᶜp_lazy, + ᶜmse⁰ - ᶜΦ, + ᶜq_tot⁰) end @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) return nothing @@ -112,7 +118,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!( thermo_params = CAP.thermodynamics_params(p.params) (; ᶜΦ,) = p.core - (; ᶜp, ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed + (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs, ᶜts) = p.precomputed + + ᶜp_lazy = ᶜp(thermo_params, ᶜts) for j in 1:n ᶜuʲ = ᶜuʲs.:($j) @@ -138,7 +146,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!( p.atmos.precip_model isa Microphysics1Moment @. ᶜtsʲ = TD.PhaseNonEquil_phq( thermo_params, - ᶜp, + ᶜp_lazy, ᶜmseʲ - ᶜΦ, TD.PhasePartition( ᶜq_totʲ, @@ -147,7 +155,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!( ), ) else - @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) + @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, + ᶜp_lazy, + ᶜmseʲ - ᶜΦ, + ᶜq_totʲ) end @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) end @@ -173,9 +184,12 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( turbconv_params = CAP.turbconv_params(p.params) (; ᶜΦ,) = p.core - (; ᶜspecific, ᶜp, ᶜh_tot, ᶜK, ᶜtsʲs, ᶜρʲs) = p.precomputed + (; ᶜspecific, ᶜK, ᶜtsʲs, ᶜρʲs, ᶜts) = p.precomputed (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + ᶜh_tot_lazy = ᶜh_tot(Y, thermo_params, ᶜts) + for j in 1:n ᶜtsʲ = ᶜtsʲs.:($j) ᶜmseʲ = Y.c.sgsʲs.:($j).mse @@ -197,7 +211,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( Fields.level(Fields.coordinate_field(Y.f).z, Fields.half), ) ᶜρ_int_val = Fields.field_values(Fields.level(Y.c.ρ, 1)) - ᶜp_int_val = Fields.field_values(Fields.level(ᶜp, 1)) + ᶜp_int_val = Fields.field_values(Fields.level(Base.materialize(ᶜp_lazy), 1)) (; ρ_flux_h_tot, ρ_flux_q_tot, ustar, obukhov_length) = p.precomputed.sfc_conditions @@ -217,7 +231,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( # TODO: replace this with the actual surface area fraction when # using prognostic surface area @. ᶜaʲ_int_val = FT(turbconv_params.surface_area) - ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1)) + ᶜh_tot_int_val = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot_lazy), 1)) ᶜK_int_val = Fields.field_values(Fields.level(ᶜK, 1)) ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1)) @. ᶜmseʲ_int_val = sgs_scalar_first_interior_bc( @@ -367,7 +381,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⁰, ᶜu, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed (; ᶜmixing_length_tuple, ᶜmixing_length, @@ -386,6 +400,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ᶜdetrʲs, ᶜturb_entrʲs, ᶠnh_pressure₃_buoyʲs, + ᶜts, ) = p.precomputed (; ustar, obukhov_length) = p.precomputed.sfc_conditions @@ -398,6 +413,10 @@ 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 + + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + ᶜh_tot_lazy = ᶜp(thermo_params, ᶜts) + for j in 1:n # entrainment/detrainment @. ᶜentrʲs.:($$j) = entrainment( @@ -405,7 +424,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos turbconv_params, ᶜz, z_sfc, - ᶜp, + ᶜp_lazy, Y.c.ρ, draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)), get_physical_w(ᶜuʲs.:($$j), ᶜlg), @@ -441,7 +460,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos turbconv_params, ᶜz, z_sfc, - ᶜp, + ᶜp_lazy, Y.c.ρ, Y.c.sgsʲs.:($$j).ρa, draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)), diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl index b881b2bc4f..358d63534c 100644 --- a/src/diagnostics/core_diagnostics.jl +++ b/src/diagnostics/core_diagnostics.jl @@ -166,9 +166,9 @@ add_diagnostic_variable!( compute! = (out, state, cache, time) -> begin thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) - return TD.specific_enthalpy.(thermo_params, cache.precomputed.ᶜts) + return TD.air_pressure.(thermo_params, cache.precomputed.ᶜts) else - out .= TD.specific_enthalpy.(thermo_params, cache.precomputed.ᶜts) + out .= TD.air_pressure.(thermo_params, cache.precomputed.ᶜts) end end, ) @@ -181,10 +181,11 @@ add_diagnostic_variable!( long_name = "Pressure at Model Full-Levels", units = "Pa", compute! = (out, state, cache, time) -> begin + thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) - return copy(cache.precomputed.ᶜp) + return copy(TD.air_pressure.(thermo_params, cache.precomputed.ᶜts)) else - out .= cache.precomputed.ᶜp + out .= TD.air_pressure.(thermo_params, cache.precomputed.ᶜts) end end, ) @@ -1434,7 +1435,7 @@ function compute_cape!(out, state, cache, time) lazy.( TD.PhaseEquil_pθq.( thermo_params, - cache.precomputed.ᶜp, + TD.air_pressure.(thermo_params, cache.precomputed.ᶜts), surface_θ_liq_ice, surface_q, ), diff --git a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl index efb22bf5a5..4c392bcf1f 100644 --- a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl +++ b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl @@ -141,6 +141,7 @@ function non_orographic_gravity_wave_compute_tendency!( #unpack ᶜT = p.scratch.ᶜtemp_scalar (; ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) (; params) = p (; ᶜdTdz, @@ -159,7 +160,6 @@ function non_orographic_gravity_wave_compute_tendency!( ᶜz = Fields.coordinate_field(Y.c).z FT = Spaces.undertype(axes(Y.c)) # parameters - thermo_params = CAP.thermodynamics_params(params) grav = CAP.grav(params) # compute buoyancy frequency @@ -205,12 +205,11 @@ function non_orographic_gravity_wave_compute_tendency!( fill!(damp_level, Spaces.nlevels(axes(ᶜz))) elseif issphere(axes(Y.c)) - (; ᶜp) = p.precomputed (; gw_source_pressure, gw_damp_pressure, source_p_ρ_z_u_v_level) = p.non_orographic_gravity_wave # source level: the index of the highest level whose pressure is higher than source pressure - input = Base.Broadcast.broadcasted(tuple, ᶜp, ᶜρ, ᶜz, ᶜu, ᶜv, ᶜlevel) + input = Base.Broadcast.broadcasted(tuple, Base.materialize(ᶜp(thermo_params, ᶜts)), ᶜρ, ᶜz, ᶜu, ᶜv, ᶜlevel) Operators.column_reduce!( source_p_ρ_z_u_v_level, input, @@ -231,7 +230,7 @@ function non_orographic_gravity_wave_compute_tendency!( # damp level: the index of the lowest level whose pressure is lower than the damp pressure - input = Base.Broadcast.broadcasted(tuple, ᶜlevel, ᶜp) + input = Base.Broadcast.broadcasted(tuple, ᶜlevel, Base.materialize(ᶜp(thermo_params, ᶜts))) Operators.column_reduce!( damp_level, input; diff --git a/src/parameterized_tendencies/gravity_wave_drag/orographic_gravity_wave.jl b/src/parameterized_tendencies/gravity_wave_drag/orographic_gravity_wave.jl index 0633fb1214..8fb10f6703 100644 --- a/src/parameterized_tendencies/gravity_wave_drag/orographic_gravity_wave.jl +++ b/src/parameterized_tendencies/gravity_wave_drag/orographic_gravity_wave.jl @@ -80,7 +80,7 @@ end function orographic_gravity_wave_tendency!(Yₜ, Y, p, t, ::OrographicGravityWave) ᶜT = p.scratch.ᶜtemp_scalar (; params) = p - (; ᶜts, ᶜp) = p.precomputed + (; ᶜts) = p.precomputed (; ᶜdTdz) = p.orographic_gravity_wave (; topo_k_pbl, @@ -98,19 +98,21 @@ function orographic_gravity_wave_tendency!(Yₜ, Y, p, t, ::OrographicGravityWav FT = Spaces.undertype(axes(Y.c)) # parameters - thermo_params = CAP.thermodynamics_params(params) + thermo_params = CAP.thermodynamics_params(p.params) grav = FT(CAP.grav(params)) cp_d = FT(CAP.cp_d(params)) # z ᶜz = Fields.coordinate_field(Y.c).z ᶠz = Fields.coordinate_field(Y.f).z + + ᶜp_eager = Base.materialize(ᶜp(thermo_params, ᶜts)) # get PBL info @. ᶜT = TD.air_temperature(thermo_params, ᶜts) Fields.bycolumn(axes(Y.c.ρ)) do colidx parent(topo_k_pbl[colidx]) .= - get_pbl(ᶜp[colidx], ᶜT[colidx], ᶜz[colidx], grav, cp_d) + get_pbl(ᶜp_eager[colidx], ᶜT[colidx], ᶜz[colidx], grav, cp_d) end # buoyancy frequency at cell centers @@ -171,7 +173,7 @@ function orographic_gravity_wave_tendency!(Yₜ, Y, p, t, ::OrographicGravityWav u_phy[colidx], v_phy[colidx], Y.c.ρ[colidx], - ᶜp[colidx], + ᶜp_eager[colidx], Int(parent(topo_k_pbl[colidx])[1]), ) end @@ -200,7 +202,7 @@ function orographic_gravity_wave_tendency!(Yₜ, Y, p, t, ::OrographicGravityWav vforcing[colidx], ᶠN[colidx], topo_ᶠVτ[colidx], - ᶜp[colidx], + ᶜp_eager[colidx], topo_τ_x[colidx], topo_τ_y[colidx], topo_τ_l[colidx], @@ -226,7 +228,7 @@ function calc_nonpropagating_forcing!( ᶜvforcing, ᶠN, ᶠVτ, - ᶜp, + ᶜp_eager, τ_x, τ_y, τ_l, @@ -237,7 +239,7 @@ function calc_nonpropagating_forcing!( grav, ) FT = eltype(grav) - ᶠp = ᶠinterp.(ᶜp) + ᶠp = ᶠinterp.(ᶜp_eager) # compute k_ref: the upper bound for nonpropagating drag to function phase = FT(0) zlast = parent(Fields.level(ᶠz, k_pbl + half))[1] @@ -258,7 +260,7 @@ function calc_nonpropagating_forcing!( wtsum = FT(0) for k in k_pbl:k_ref tmp = Fields.level(weights, k) - parent(tmp) .= parent(ᶜp)[k] - parent(ᶠp)[k_ref] + parent(tmp) .= parent(ᶜp_eager)[k] - parent(ᶠp)[k_ref] wtsum += (parent(ᶠp)[k - 1] - parent(ᶠp)[k]) / parent(weights)[k] end @@ -275,14 +277,14 @@ function calc_propagate_forcing!(ᶜuforcing, ᶜvforcing, τ_x, τ_y, τ_l, τ_ return nothing end -function get_pbl(ᶜp, ᶜT, ᶜz, grav, cp_d) +function get_pbl(ᶜp_eager, ᶜT, ᶜz, grav, cp_d) FT = eltype(cp_d) idx = - (parent(ᶜp) .>= (FT(0.5) * parent(ᶜp)[1])) .& ( + (parent(ᶜp_eager) .>= (FT(0.5) * parent(ᶜp_eager)[1])) .& ( (parent(ᶜT)[1] + FT(1.5) .- parent(ᶜT)) .> (grav / cp_d * (parent(ᶜz) .- parent(ᶜz)[1])) ) - # parent(ᶜp) .>= (FT(0.5) * parent(ᶜp)[1]) follows the criterion in GFDL codes + # parent(ᶜp_eager) .>= (FT(0.5) * parent(ᶜp_eager)[1]) follows the criterion in GFDL codes # that the lowest layer that is geq to half of pressure at first face level; while # in our code, when interpolate from center to face, the first face level inherits # values at the first center level @@ -389,7 +391,7 @@ function calc_saturation_profile!( u_phy, v_phy, ᶜρ, - ᶜp, + ᶜp_eager, k_pbl, ) (; Fr_crit, topo_ρscale, topo_L0, topo_a0, topo_γ, topo_β, topo_ϵ) = @@ -451,7 +453,7 @@ function calc_saturation_profile!( # If the wave propagates to the top, the residual momentum flux is redistributed throughout the column weighted by pressure if parent(τ_sat)[end] > FT(0) - ᶠp = ᶠinterp.(ᶜp) + ᶠp = ᶠinterp.(ᶜp_eager) τ_sat .-= parent(τ_sat)[end] .* (parent(ᶠp)[1] .- ᶠp) ./ (parent(ᶠp)[1] .- parent(ᶠp)[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..27c9439f6b 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -96,7 +96,11 @@ horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) - (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜspecific, ᶜh_tot) = p.precomputed + (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜspecific) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) ## Momentum tendencies ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(Y.c.ρ) @@ -123,9 +127,13 @@ import UnrolledUtilities as UU function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) FT = eltype(Y) - (; sfc_temp_C3, ᶠtemp_scalar) = p.scratch - (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, ᶜspecific, ᶜh_tot, sfc_conditions) = + (; ᶠtemp_scalar) = p.scratch + (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, ᶜspecific, sfc_conditions) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions # Define operators diff --git a/src/parameterized_tendencies/radiation/update_inputs.jl b/src/parameterized_tendencies/radiation/update_inputs.jl index fdda502c4b..f84c72aee4 100644 --- a/src/parameterized_tendencies/radiation/update_inputs.jl +++ b/src/parameterized_tendencies/radiation/update_inputs.jl @@ -48,16 +48,19 @@ function update_temperature_pressure!((; u, p, t)::I) where {I} T_min = CAP.optics_lookup_temperature_min(params) T_max = CAP.optics_lookup_temperature_max(params) - (; ᶜts, ᶜp, sfc_conditions) = p.precomputed + (; ᶜts, sfc_conditions) = p.precomputed model = p.radiation.rrtmgp_model + # TODO: field2array compat with lazy ops + ᶜp = @. TD.air_pressure(thermo_params, ᶜts) + # update surface temperature sfc_ts = sfc_conditions.ts sfc_T = Fields.array2field(model.surface_temperature, axes(sfc_ts)) @. sfc_T = TD.air_temperature(thermo_params, sfc_ts) # update layer pressure - model.center_pressure .= Fields.field2array(p.precomputed.ᶜp) + model.center_pressure .= Fields.field2array(ᶜp) # compute layer temperature ᶜT = Fields.array2field(model.center_temperature, axes(u.c)) # TODO: move this to RRTMGP diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index a932682049..617f9f8851 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -26,7 +26,7 @@ Arguments: - `Yₜ`: The tendency state vector, modified in place. - `Y`: The current state vector. - `p`: Cache containing parameters, precomputed fields (e.g., velocities `ᶜu`, - `ᶜu⁰`, `ᶜuʲs`; pressure `ᶜp`; kinetic energy `ᶜK`; total enthalpy `ᶜh_tot`), + `ᶜu⁰`, `ᶜuʲs`; kinetic energy `ᶜK`), and core components (e.g., geopotential `ᶜΦ`). - `t`: Current simulation time (not directly used in calculations). @@ -36,7 +36,8 @@ 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 + (; ᶜts, ᶜu, ᶜK) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) if p.atmos.turbconv_model isa PrognosticEDMFX (; ᶜuʲs) = p.precomputed @@ -49,8 +50,8 @@ NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) end end - (; ᶜh_tot) = p.precomputed - @. Yₜ.c.ρe_tot -= wdivₕ(Y.c.ρ * ᶜh_tot * ᶜu) + ᶜh_tot_lazy = ᶜh_tot(Y, thermo_params, ᶜts) + @. Yₜ.c.ρe_tot -= wdivₕ(Y.c.ρ * ᶜh_tot_lazy * ᶜu) if p.atmos.turbconv_model isa PrognosticEDMFX for j in 1:n @@ -74,7 +75,8 @@ NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) end - @. Yₜ.c.uₕ -= C12(gradₕ(ᶜp) / Y.c.ρ + gradₕ(ᶜK + ᶜΦ)) + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + @. Yₜ.c.uₕ -= C12(gradₕ(ᶜp_lazy) / Y.c.ρ + gradₕ(ᶜK + ᶜΦ)) # Without the C12(), the right-hand side would be a C1 or C2 in 2D space. return nothing end @@ -104,6 +106,7 @@ 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 + thermo_params = CAP.thermodynamics_params(p.params) if p.atmos.turbconv_model isa PrognosticEDMFX (; ᶜuʲs) = p.precomputed @@ -178,7 +181,8 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) (; 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 - (; ᶜspecific) = p.precomputed + (; ᶜspecific, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) ᶠu³⁰ = advect_tke ? @@ -222,10 +226,11 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) end # ... and upwinding correction of energy and total water. # (The central advection of energy and total water is done implicitly.) + + ᶜh_tot_lazy = ᶜh_tot(Y,thermo_params,ᶜts) if energy_upwinding != Val(:none) - (; ᶜh_tot) = p.precomputed - vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) - vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot_lazy, float(dt), energy_upwinding) + vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot_lazy, float(dt), Val(:none)) @. Yₜ.c.ρe_tot += vtt - vtt_central end diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index a87568ec75..a1b5618653 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -40,9 +40,10 @@ function edmfx_sgs_mass_flux_tendency!( n = n_mass_flux_subdomains(turbconv_model) (; edmfx_sgsflux_upwinding) = p.atmos.numerics - (; ᶠu³, ᶜh_tot) = p.precomputed + (; ᶠu³) = p.precomputed (; ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = p.precomputed - (; ᶜρa⁰, ᶜρ⁰, ᶠu³⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed + (; ᶜρa⁰, ᶜρ⁰, ᶠu³⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) if ( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.precip_model isa Microphysics1Moment @@ -52,6 +53,9 @@ function edmfx_sgs_mass_flux_tendency!( (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + ᶜh_tot_lazy = ᶜp(thermo_params, ᶜts) + 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 @@ -62,7 +66,7 @@ function edmfx_sgs_mass_flux_tendency!( for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @. ᶜa_scalar = - (Y.c.sgsʲs.:($$j).mse + ᶜKʲs.:($$j) - ᶜh_tot) * + (Y.c.sgsʲs.:($$j).mse + ᶜKʲs.:($$j) - ᶜh_tot_lazy) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), @@ -75,7 +79,7 @@ function edmfx_sgs_mass_flux_tendency!( end # Add the environment fluxes @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ - @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) + @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot_lazy) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -231,17 +235,20 @@ function edmfx_sgs_mass_flux_tendency!( t, turbconv_model::DiagnosticEDMFX, ) - + thermo_params = CAP.thermodynamics_params(p.params) turbconv_params = CAP.turbconv_params(p.params) 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 + (; ᶠu³, ᶜts) = 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) + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + ᶜh_tot_lazy = ᶜh_tot(Y, thermo_params, ᶜts) + if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} # Enthalpy fluxes ᶠu³_diff = p.scratch.ᶠtemp_CT3 @@ -253,7 +260,7 @@ function edmfx_sgs_mass_flux_tendency!( # draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) # TODO: remove this filter when mass flux is treated implicitly @. ᶜa_scalar = - (ᶜmseʲs.:($$j) + ᶜKʲs.:($$j) - ᶜh_tot) * min( + (ᶜmseʲs.:($$j) + ᶜKʲs.:($$j) - ᶜh_tot_lazy) * min( min(draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)), a_max), FT(0.02) / max( Geometry.WVector(ᶜinterp(ᶠu³_diff)).components.data.:1, @@ -448,12 +455,16 @@ function edmfx_sgs_diffusive_flux_tendency!( # prognostic fluxes FT = Spaces.undertype(axes(Y.c)) (; dt, params) = p + thermo_params = CAP.thermodynamics_params(params) turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜu, ᶜh_tot, ᶜtke⁰, ᶜmixing_length) = p.precomputed - (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed + (; ᶜu, ᶜtke⁰, ᶜmixing_length) = p.precomputed + (; ᶜK_u, ᶜK_h, ρatke_flux, ᶜts) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() + ᶜp_lazy = ᶜp(thermo_params, ᶜts) + ᶜh_tot_lazy = ᶜh_tot(Y,thermo_params, ᶜts) + if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar @. ᶠρaK_h = ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) @@ -465,7 +476,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜh_tot))) + @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜh_tot_lazy))) if use_prognostic_tke(turbconv_model) # Turbulent TKE transport (diffusion) diff --git a/src/prognostic_equations/forcing/external_forcing.jl b/src/prognostic_equations/forcing/external_forcing.jl index 487a38e263..7f56ca584f 100644 --- a/src/prognostic_equations/forcing/external_forcing.jl +++ b/src/prognostic_equations/forcing/external_forcing.jl @@ -337,21 +337,12 @@ function external_forcing_tendency!( # horizontal advection, vertical fluctuation, nudging, subsidence (need to add), (; params) = p thermo_params = CAP.thermodynamics_params(params) - (; ᶜts, ᶜh_tot) = p.precomputed - (; - ᶜdTdt_fluc, - ᶜdqtdt_fluc, - ᶜdTdt_hadv, - ᶜdqtdt_hadv, - ᶜdTdt_rad, - ᶜT_nudge, - ᶜqt_nudge, - ᶜu_nudge, - ᶜv_nudge, - ᶜls_subsidence, - ᶜinv_τ_wind, - ᶜinv_τ_scalar, - ) = p.external_forcing + (; ᶜts) = p.precomputed + ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts)) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) + (; ᶜdTdt_fluc, ᶜdqtdt_fluc, ᶜdTdt_hadv, ᶜdqtdt_hadv, ᶜdTdt_rad, ᶜT_nudge, ᶜqt_nudge, ᶜu_nudge, ᶜv_nudge, ᶜls_subsidence, ᶜinv_τ_wind, ᶜinv_τ_scalar) = p.external_forcing ᶜlg = Fields.local_geometry_field(Y.c) ᶜuₕ_nudge = p.scratch.ᶜtemp_C12 @@ -574,7 +565,10 @@ function external_forcing_tendency!(Yₜ, Y, p, t, ::ISDACForcing) FT = Spaces.undertype(axes(Y.c)) (; params) = p thermo_params = CAP.thermodynamics_params(params) - (; ᶜts, ᶜh_tot, ᶜp) = p.precomputed + (; ᶜts) = p.precomputed + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) ᶜinv_τ_scalar = APL.ISDAC_inv_τ_scalar(FT) # s⁻¹ ᶜinv_τ_wind = APL.ISDAC_inv_τ_wind(FT) # s⁻¹ diff --git a/src/prognostic_equations/forcing/subsidence.jl b/src/prognostic_equations/forcing/subsidence.jl index c722a50e03..6af0ed33f7 100644 --- a/src/prognostic_equations/forcing/subsidence.jl +++ b/src/prognostic_equations/forcing/subsidence.jl @@ -68,7 +68,10 @@ subsidence_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing # No subsidence function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) (; moisture_model) = p.atmos subsidence_profile = p.atmos.subsidence.prof - (; ᶜh_tot) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) ᶠz = Fields.coordinate_field(axes(Y.f)).z ᶠlg = Fields.local_geometry_field(Y.f) diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index 662691b662..c56def8694 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -97,24 +97,26 @@ end # dss_hyperdiffusion_tendency_pairs NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos + thermo_params = CAP.thermodynamics_params(p.params) isnothing(hyperdiff) && return nothing n = n_mass_flux_subdomains(turbconv_model) diffuse_tke = use_prognostic_tke(turbconv_model) - (; ᶜp) = p.precomputed + (; ᶜts) = p.precomputed (; ᶜh_ref) = p.core (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end - + # Grid scale hyperdiffusion @. ᶜ∇²u = C123(wgradₕ(divₕ(p.precomputed.ᶜu))) - C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜu)))) + ᶜpressure = ᶜp(thermo_params, ᶜts) @. ᶜ∇²specific_energy = - wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) + wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜpressure / Y.c.ρ - ᶜh_ref)) if diffuse_tke (; ᶜtke⁰) = p.precomputed @@ -139,6 +141,7 @@ end # variables in dss_hyperdiffusion_tendency_pairs NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos + isnothing(hyperdiff) && return nothing (; divergence_damping_factor) = hyperdiff @@ -148,7 +151,8 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) diffuse_tke = use_prognostic_tke(turbconv_model) ᶜJ = Fields.local_geometry_field(Y.c).J point_type = eltype(Fields.coordinate_field(Y.c)) - (; ᶜp) = p.precomputed + (; ᶜts) = p.precomputed + (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX (; ᶜρa⁰) = p.precomputed diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 0e1b0a3513..0b6419801a 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -139,7 +139,10 @@ 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 + (; ᶠu³) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜp = @. lazy(TD.air_pressure(thermo_params, p.precomputed.ᶜts)) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, p.precomputed.ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠu³) diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index a9de2af6a4..2e06597ef4 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -337,7 +337,12 @@ 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 + (; ᶠu³, ᶜK, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(params) + ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts)) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, @@ -366,7 +371,6 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) # This term appears a few times in the Jacobian, and is technically # minus ∂e_int_∂q_tot ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - FT(CAP.e_int_v0(params)) - thermo_params = CAP.thermodynamics_params(params) ᶜρ = Y.c.ρ ᶜuₕ = Y.c.uₕ @@ -419,7 +423,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) MatrixFields.unrolled_foreach(tracer_info) do ρχ_name MatrixFields.has_field(Y, ρχ_name) || return ᶜχ = if ρχ_name === @name(c.ρe_tot) - p.precomputed.ᶜh_tot + ᶜh_tot else @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) end diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index cc9b8e7c2e..e9209295ff 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -1,4 +1,3 @@ - """ hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) @@ -137,7 +136,7 @@ dependencies or operator splitting assumptions. """ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) - (; ᶜh_tot, ᶜspecific) = p.precomputed + (; ᶜspecific) = p.precomputed ᶜuₕ = Y.c.uₕ ᶠu₃ = Y.f.u₃ ᶜρ = Y.c.ρ @@ -145,7 +144,12 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) (; ls_adv, scm_coriolis) = p.atmos (; params) = p thermo_params = CAP.thermodynamics_params(params) - (; ᶜp, sfc_conditions, ᶜts) = p.precomputed + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) + (; sfc_conditions, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts)) vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge) diff --git a/src/prognostic_equations/surface_flux.jl b/src/prognostic_equations/surface_flux.jl index da2d58c0fb..adca4d3bd7 100644 --- a/src/prognostic_equations/surface_flux.jl +++ b/src/prognostic_equations/surface_flux.jl @@ -104,7 +104,11 @@ function surface_flux_tendency!(Yₜ, Y, p, t) p.atmos.disable_surface_flux_tendency && return FT = eltype(Y) - (; ᶜh_tot, ᶜspecific, sfc_conditions) = p.precomputed + (; ᶜspecific, sfc_conditions) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) if !disable_momentum_vertical_diffusion(p.atmos.vert_diff) btt = diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 569e65cfc3..3f23dd4e18 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -76,7 +76,11 @@ 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 + (; ᶜu, ᶜspecific, ᶜK_u, ᶜK_h) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, + p.precomputed.ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ))) ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ if !disable_momentum_vertical_diffusion(p.atmos.vert_diff) diff --git a/src/utils/discrete_hydrostatic_balance.jl b/src/utils/discrete_hydrostatic_balance.jl index d6e9b06e96..97c379dfa4 100644 --- a/src/utils/discrete_hydrostatic_balance.jl +++ b/src/utils/discrete_hydrostatic_balance.jl @@ -12,22 +12,23 @@ Modify the energy variable in state `Y` given Y and the cache `p` so that function set_discrete_hydrostatic_balanced_state!(Y, p) FT = Spaces.undertype(axes(Y.c)) ᶠgradᵥ_ᶜp = similar(Y.f.u₃) + thermo_params = CAP.thermodynamics_params(p.params) + ᶜp = @. Base.materialize(TD.air_pressure(thermo_params, p.precomputed.ᶜts)) set_discrete_hydrostatic_balanced_pressure!( - p.precomputed.ᶜp, + ᶜp, ᶠgradᵥ_ᶜp, Y.c.ρ, p.core.ᶠgradᵥ_ᶜΦ, FT(CAP.MSLP(p.params)), ) - thermo_params = CAP.thermodynamics_params(p.params) if p.atmos.moisture_model isa DryModel @. p.precomputed.ᶜts = - TD.PhaseDry_ρp(thermo_params, Y.c.ρ, p.precomputed.ᶜp) + TD.PhaseDry_ρp(thermo_params, Y.c.ρ, ᶜp) elseif p.atmos.moisture_model isa EquilMoistModel @. p.precomputed.ᶜts = TD.PhaseEquil_ρpq( thermo_params, Y.c.ρ, - p.precomputed.ᶜp, + ᶜp, Y.c.ρq_tot / Y.c.ρ, ) else diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 443c5eb359..abafed968d 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -1,5 +1,34 @@ import ClimaCore.MatrixFields: @name + +""" + ᶜp(thermo_params, ᶜts) + +Return lazy evaluation of air pressure. +Args: + - thermo_params: Thermodynamic parameters accessible from ClimaAtmos.Parameters + - ᶜts: Thermodynamic state +""" +@inline function ᶜp(thermo_params, ᶜts) + return @. lazy(TD.air_pressure(thermo_params, ᶜts)) +end + +""" + ᶜh_tot(Y, thermo_params, ᶜts) + +Return lazy evaluation of total specific enthalpy. +Args: + - Y: Prognostic state variables + - thermo_params: Thermodynamic parameters accessible from ClimaAtmos.Parameters (CAP) + - ᶜts: Thermodynamic state +""" +@inline function ᶜh_tot(Y, thermo_params, ᶜts) + return @. lazy(TD.total_specific_enthalpy(thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ)) + ) +end + """ specific(ρχ, ρ) specific(ρaχ, ρa, ρχ, ρ, turbconv_model)