From 46e774f73002a63aee27145cca62c4f0a5198f02 Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Tue, 10 Jun 2025 15:31:09 -0700 Subject: [PATCH 1/7] Added generic draft_sum (as comment for now) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove from ᶜρa⁰ Fixed syntax errors Removed ᶜρ⁰ from precomputed Remove ᶜq_tot⁰ (and fix forgotten ᶜρa⁰ Introduced helper functions to compute sums over draft, environmental volumetric variables, and specific env variables Use new helper functions to simplify calculation of ᶜq_tot⁰ Remove redundant draft sum functions Added specific_env_mse helper Removed ᶜmse⁰ from precomputed Remove ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰ from precomputed Remove redundant draft sum helper functions; docstrings Added TODO Renamed specific_gs to all_specific_gs; specific_sgs to all_specific_sgs; added docstrings Added new helper specific_sgs to cleanly extract specific SGS quantities; use it to get env TKE Change specific_sgs to use ClimaCore>MatrixFIelds for type stability Removed ᶜspecific (GS precomputed specific quantities); caveat lector [lots of changes] Correcting errors in previous commit removing gs precomputed quantities Removed a few more instances of ᶜspecific from precomputed quantities; removed ᶜtke⁰ Missing ᶜtke⁰ removal Another ᶜtke⁰ fix Introduced helper function for specific_tke and used it where needed Corrections of rebasing mistakes; updates to variable_manipulations for clarity. Removing some more ᶜspecific Added ᶜtke⁰ computations Removal of more specifics Remove h_tot Syntax corrections Fixes in cloud fraction Syntax error fix in jacobian Remove specific in precomputed_quantities; syntax error corrections --- src/ClimaAtmos.jl | 2 + src/cache/cloud_fraction.jl | 27 +- .../diagnostic_edmf_precomputed_quantities.jl | 24 +- .../precipitation_precomputed_quantities.jl | 44 +- src/cache/precomputed_quantities.jl | 70 +-- .../prognostic_edmf_precomputed_quantities.jl | 123 ++--- src/diagnostics/edmfx_diagnostics.jl | 34 +- .../les_sgs_models/smagorinsky_lilly.jl | 6 +- .../microphysics/precipitation.jl | 4 +- .../radiation/radiation.jl | 13 +- src/prognostic_equations/advection.jl | 22 +- src/prognostic_equations/edmfx_entr_detr.jl | 9 +- src/prognostic_equations/edmfx_sgs_flux.jl | 141 ++++-- src/prognostic_equations/edmfx_tke.jl | 16 +- .../forcing/external_forcing.jl | 3 +- .../forcing/subsidence.jl | 44 +- src/prognostic_equations/hyperdiffusion.jl | 12 +- .../implicit/implicit_tendency.jl | 24 +- .../implicit/manual_sparse_jacobian.jl | 85 ++-- .../remaining_tendency.jl | 4 +- src/prognostic_equations/surface_flux.jl | 3 +- .../vertical_diffusion_boundary_layer.jl | 18 +- src/solver/types.jl | 4 +- src/utils/variable_manipulations.jl | 437 ++++++++++++------ 24 files changed, 704 insertions(+), 465 deletions(-) diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 7ed4d5a5fc..f365631e80 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -8,6 +8,8 @@ import LazyBroadcast import LazyBroadcast: lazy import Thermodynamics as TD import Thermodynamics +import ClimaCore.MatrixFields: @name + include("compat.jl") include(joinpath("parameters", "Parameters.jl")) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index 06055c78ff..443c33205d 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -62,14 +62,16 @@ NVTX.@annotate function set_cloud_fraction!( TD.PhasePartition(thermo_params, ᶜts).ice, ) else + q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) @. cloud_diagnostics_tuple = make_named_tuple( ifelse( - specific(Y.c.ρq_liq, Y.c.ρ) + specific(Y.c.ρq_ice, Y.c.ρ) > 0, + q_liq + q_ice > 0, 1, 0, ), - specific(Y.c.ρq_liq, Y.c.ρ), - specific(Y.c.ρq_ice, Y.c.ρ), + q_liq, + q_ice, ) end end @@ -89,6 +91,16 @@ NVTX.@annotate function set_cloud_fraction!( (; turbconv_model) = p.atmos if isnothing(turbconv_model) + (; + ᶜlinear_buoygrad, + ᶜstrain_rate_norm, + ᶠu³⁰, + ᶠu³, + ᶜentrʲs, + ᶜdetrʲs, + ᶠu³ʲs, + ) = p.precomputed + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) if p.atmos.call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed @@ -193,8 +205,9 @@ NVTX.@annotate function set_cloud_fraction!( FT = eltype(params) thermo_params = CAP.thermodynamics_params(params) (; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed - (; ᶜρʲs, ᶜtsʲs, ᶜρa⁰, ᶜρ⁰) = p.precomputed + (; ᶜρʲs, ᶜtsʲs) = p.precomputed (; turbconv_model) = p.atmos + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) # TODO - we should make this default when using diagnostic edmf # environment @@ -213,9 +226,9 @@ NVTX.@annotate function set_cloud_fraction!( # weight cloud diagnostics by environmental area @. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}( tuple( - draft_area(ᶜρa⁰, ᶜρ⁰), - draft_area(ᶜρa⁰, ᶜρ⁰), - draft_area(ᶜρa⁰, ᶜρ⁰), + draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), ), ) # updrafts diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index f0cb9255b9..987d3f9b9b 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -61,7 +61,7 @@ NVTX.@annotate function set_diagnostic_edmfx_env_quantities_level!( local_geometry_halflevel, turbconv_model, ) - @. u³⁰_halflevel = divide_by_ρa( + @. u³⁰_halflevel = specific( ρ_level * u³_halflevel - unrolled_dotproduct(ρaʲs_level, u³ʲs_halflevel), ρ_level, @@ -93,16 +93,20 @@ 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 - (; q_tot) = p.precomputed.ᶜspecific + (; ᶜp, ᶠu³, ᶜK) = p.precomputed (; 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 (; params) = p + thermo_params = CAP.thermodynamics_params(params) turbconv_params = CAP.turbconv_params(params) + ᶜts = p.precomputed.ᶜts #TODO replace + + q_tot = specific(Y.c.ρq_tot, Y.c.ρ) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1)) uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1)) @@ -305,8 +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 - (; q_tot) = p.precomputed.ᶜspecific + (; ᶜp, ᶠu³, ᶜts, ᶜK) = p.precomputed (; ᶜρaʲs, ᶠu³ʲs, @@ -324,10 +327,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( (; ᶠu³⁰, ᶜK⁰, ᶜtke⁰) = p.precomputed if microphysics_model isa Microphysics1Moment - ᶜq_liqʲs = p.precomputed.ᶜq_liqʲs - ᶜq_iceʲs = p.precomputed.ᶜq_iceʲs - q_rai = p.precomputed.ᶜqᵣ - q_sno = p.precomputed.ᶜqₛ + q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) end thermo_params = CAP.thermodynamics_params(params) @@ -347,6 +348,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half)) # integral + q_tot = specific(Y.c.ρq_tot, Y.c.ρ) 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)) @@ -965,7 +967,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! (; dt) = p (; ᶜp, ᶜu, ᶜts) = p.precomputed (; ustar, obukhov_length) = p.precomputed.sfc_conditions - (; ᶜtke⁰) = p.precomputed (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, @@ -1083,6 +1084,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed # Environment precipitation sources (to be applied to grid mean) + q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources( thermo_params, microphys_0m_params, @@ -1126,4 +1128,4 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita # thermo_params, #) return nothing -end +end \ No newline at end of file diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index c7bb9c083b..fc5e770e96 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -112,6 +112,11 @@ function set_precipitation_velocities!( cm2p = CAP.microphysics_2m_params(p.params) thp = CAP.thermodynamics_params(p.params) + q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + # compute the precipitation terminal velocity [m/s] # TODO sedimentation of snow is based on the 1M scheme @. ᶜwnᵣ = getindex( @@ -253,10 +258,11 @@ function set_precipitation_cache!( ::PrognosticEDMFX, ) (; ᶜΦ) = p.core - (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs, ᶜρa⁰) = p.precomputed + (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs) = p.precomputed (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed (; ᶜts⁰, ᶜtsʲs) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) n = n_mass_flux_subdomains(p.atmos.turbconv_model) @@ -283,7 +289,8 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed - (; q_tot, q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific + q_rai = specific(Y.c.ρq_rai, Y.c.ρ) + q_sno = specific(Y.c.ρq_sno, Y.c.ρ) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2 @@ -357,6 +364,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _) (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed + q_liq = specific(Y.c.ρq_liq, Y.c.ρ) + q_rai = specific(Y.c.ρq_rai, Y.c.ρ) + n_liq = specific(Y.c.ρn_liq, Y.c.ρ) + n_rai = specific(Y.c.ρn_rai, Y.c.ρ) + ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2 @@ -469,14 +481,26 @@ function set_precipitation_surface_fluxes!( sfc_ρ = @. lazy(int_ρ * int_J / sfc_J) # Constant extrapolation to surface, consistent with simple downwinding - sfc_wᵣ = sfc_lev(ᶜwᵣ) - sfc_wₛ = sfc_lev(ᶜwₛ) - sfc_wₗ = sfc_lev(ᶜwₗ) - sfc_wᵢ = sfc_lev(ᶜwᵢ) - sfc_qᵣ = lazy.(specific.(sfc_lev(Y.c.ρq_rai), sfc_ρ)) - sfc_qₛ = lazy.(specific.(sfc_lev(Y.c.ρq_sno), sfc_ρ)) - sfc_qₗ = lazy.(specific.(sfc_lev(Y.c.ρq_liq), sfc_ρ)) - sfc_qᵢ = lazy.(specific.(sfc_lev(Y.c.ρq_ice), sfc_ρ)) + sfc_qᵣ = Fields.Field( + Fields.field_values(Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1)), + sfc_space, + ) + sfc_qₛ = Fields.Field( + Fields.field_values(Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1)), + sfc_space, + ) + sfc_qₗ = Fields.Field( + Fields.field_values(Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1)), + sfc_space, + ) + sfc_qᵢ = Fields.Field( + Fields.field_values(Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1)), + sfc_space, + ) + sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space) + sfc_wₛ = Fields.Field(Fields.field_values(Fields.level(ᶜwₛ, 1)), sfc_space) + sfc_wₗ = Fields.Field(Fields.field_values(Fields.level(ᶜwₗ, 1)), sfc_space) + sfc_wᵢ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵢ, 1)), sfc_space) @. surface_rain_flux = sfc_ρ * (sfc_qᵣ * (-sfc_wᵣ) + sfc_qₗ * (-sfc_wₗ)) @. surface_snow_flux = sfc_ρ * (sfc_qₛ * (-sfc_wₛ) + sfc_qᵢ * (-sfc_wᵢ)) diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 04d6e3bca6..704714cf85 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -11,28 +11,20 @@ Allocates precomputed quantities that are treated implicitly (i.e., updated on each iteration of the implicit solver). This includes all quantities related to velocity and thermodynamics that are used in the implicit tendency. -The following grid-scale quantities are treated implicitly: - - `ᶜspecific`: specific quantities on cell centers (for every prognostic - quantity `ρχ`, there is a corresponding specific quantity `χ`) +The following grid-scale quantities are treated implicitly and are precomputed: - `ᶜu`: covariant velocity on cell centers - `ᶠu`: contravariant velocity on cell faces - `ᶜ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 - `_ʲs`: a tuple of values for the mass-flux subdomains In addition, there are several other SGS quantities for `PrognosticEDMFX`: - - `ᶜtke⁰`: turbulent kinetic energy of the environment on cell centers - - `ᶜρa⁰`: area-weighted air density of the environment on cell centers - - `ᶜmse⁰`: moist static energy of the environment on cell centers - - `ᶜq_tot⁰`: total specific humidity of the environment on cell centers - - `ᶜρ⁰`: air density of the environment on cell centers - `ᶜρʲs`: a tuple of the air densities of the mass-flux subdomains on cell centers -For every other `AbstractEDMF`, only `ᶜtke⁰` is added as a precomputed quantity. + TODO: Rename `ᶜK` to `ᶜκ`. """ @@ -48,35 +40,18 @@ 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), + ᶜp = similar(Y.c, FT) ) sgs_quantities = - turbconv_model isa AbstractEDMF ? (; ᶜtke⁰ = similar(Y.c, FT)) : (;) - moisture_sgs_quantities = - ( - turbconv_model isa PrognosticEDMFX && - moisture_model isa NonEquilMoistModel && - microphysics_model isa Microphysics1Moment - ) ? - (; - ᶜq_liq⁰ = similar(Y.c, FT), - ᶜq_ice⁰ = similar(Y.c, FT), - ᶜq_rai⁰ = similar(Y.c, FT), - ᶜq_sno⁰ = similar(Y.c, FT), - ) : (;) + turbconv_model isa AbstractEDMF ? (;) : (;) prognostic_sgs_quantities = turbconv_model isa PrognosticEDMFX ? (; - ᶜρa⁰ = similar(Y.c, FT), - ᶜ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), ᶜuʲs = similar(Y.c, NTuple{n, C123{FT}}), ᶠu³ʲs = similar(Y.f, NTuple{n, CT3{FT}}), ᶜKʲs = similar(Y.c, NTuple{n, FT}), @@ -84,7 +59,6 @@ function implicit_precomputed_quantities(Y, atmos) ᶜtsʲs = similar(Y.c, NTuple{n, TST}), ᶜρʲs = similar(Y.c, NTuple{n, FT}), ᶠnh_pressure₃_dragʲs = similar(Y.f, NTuple{n, C3{FT}}), - moisture_sgs_quantities..., ) : (;) return (; gs_quantities..., sgs_quantities..., prognostic_sgs_quantities...) end @@ -194,7 +168,6 @@ function precomputed_quantities(Y, atmos) atmos.turbconv_model isa EDOnlyEDMFX ? (; ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}), - ᶜtke⁰ = similar(Y.c, FT), ᶜK_u = similar(Y.c, FT), ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}), ᶜK_h = similar(Y.c, FT), @@ -399,18 +372,23 @@ function thermo_state( return get_ts(ρ, p, θ, e_int, q_tot, q_pt) end -function thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ) - energy_var = (; e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ) +function thermo_vars(moisture_model, microphysics_model, Y_c, K, Φ) + # Compute specific quantities on-the-fly + e_tot = @. lazy(specific(Y_c.ρe_tot, Y_c.ρ)) + energy_var = (; e_int = e_tot - K - Φ) + moisture_var = if moisture_model isa DryModel (;) elseif moisture_model isa EquilMoistModel - (; q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)) + q_tot = @. lazy(specific(Y_c.ρq_tot, Y_c.ρ)) + (; q_tot) elseif moisture_model isa NonEquilMoistModel - q_pt_args = (; - q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ), - q_liq = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ), - q_ice = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ), - ) + q_tot = @. lazy(specific(Y_c.ρq_tot, Y_c.ρ)) + q_liq = @. lazy(specific(Y_c.ρq_liq, Y_c.ρ)) + q_ice = @. lazy(specific(Y_c.ρq_ice, Y_c.ρ)) + q_rai = @. lazy(specific(Y_c.ρq_rai, Y_c.ρ)) + q_sno = @. lazy(specific(Y_c.ρq_sno, Y_c.ρ)) + q_pt_args = (q_tot, q_liq + q_rai, q_ice + q_sno) (; q_pt = TD.PhasePartition(q_pt_args...)) end return (; energy_var..., moisture_var...) @@ -458,7 +436,7 @@ quantities are updated. NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) (; turbconv_model, moisture_model, microphysics_model) = p.atmos (; ᶜΦ) = p.core - (; ᶜspecific, ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed + (; ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp) = p.precomputed ᶠuₕ³ = p.scratch.ᶠtemp_CT3 n = n_mass_flux_subdomains(turbconv_model) thermo_params = CAP.thermodynamics_params(p.params) @@ -492,19 +470,15 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) 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) set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t) set_prognostic_edmf_precomputed_quantities_implicit_closures!(Y, p, t) - elseif turbconv_model isa AbstractEDMF - (; ᶜtke⁰) = p.precomputed - @. ᶜtke⁰ = Y.c.sgs⁰.ρatke / Y.c.ρ + elseif turbconv_model isa DiagnosticEDMFX + set_diagnostic_edmf_precomputed_quantities!(Y, p, t) + elseif !(isnothing(turbconv_model)) + # Do nothing for other turbconv models for now end end diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 131a780ccd..847a6350fc 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -21,76 +21,35 @@ 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⁰) = - p.precomputed - if p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed - end + (; ᶜp,ᶜK) = p.precomputed + (; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed - @. ᶜρ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), - ᶜρa⁰, - Y.c.ρ * (ᶜh_tot - ᶜK), - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_tot⁰ = divide_by_ρa( - Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_tot, - Y.c.ρ, - turbconv_model, - ) - if p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - @. ᶜq_liq⁰ = divide_by_ρa( - Y.c.ρq_liq - ρaq_liq⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_liq, - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_ice⁰ = divide_by_ρa( - Y.c.ρq_ice - ρaq_ice⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_ice, - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_rai⁰ = divide_by_ρa( - Y.c.ρq_rai - ρaq_rai⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_rai, - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_sno⁰ = divide_by_ρa( - Y.c.ρq_sno - ρaq_sno⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_sno, - Y.c.ρ, - turbconv_model, - ) - end + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + @. ᶜtke⁰ = specific_tke(Y.c.sgs⁰, Y.c, turbconv_model) set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ + ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment + ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) + ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) + ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) + ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) @. ᶜts⁰ = TD.PhaseNonEquil_phq( thermo_params, ᶜp, - ᶜmse⁰ - ᶜΦ, + specific_env_mse(Y.c, p) - ᶜΦ, 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, + specific_env_mse(Y.c, p) - ᶜΦ, + ᶜq_tot⁰, + ) end - @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) return nothing end @@ -173,7 +132,7 @@ 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 + (; ᶜp, ᶜK, ᶜtsʲs, ᶜρʲs) = p.precomputed (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions for j in 1:n @@ -217,6 +176,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 = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1)) ᶜK_int_val = Fields.field_values(Fields.level(ᶜK, 1)) ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1)) @@ -233,7 +193,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( ) # ... and the first interior point for EDMFX ᶜq_totʲ. - ᶜq_tot_int_val = Fields.field_values(Fields.level(ᶜspecific.q_tot, 1)) + ᶜq_tot_int_val = + Fields.field_values(Fields.level(specific(Y.c.ρq_tot, Y.c.ρ), 1)) ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1)) @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc( ᶜz_int_val - z_sfc_val, @@ -249,23 +210,27 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment # TODO - any better way to define the cloud and precip tracer flux? - ᶜq_liq_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_liq, 1)) + ᶜq_liq_int_val = Fields.field_values( + Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1), + ) ᶜq_liqʲ_int_val = Fields.field_values(Fields.level(ᶜq_liqʲ, 1)) @. ᶜq_liqʲ_int_val = ᶜq_liq_int_val - ᶜq_ice_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_ice, 1)) + ᶜq_ice_int_val = Fields.field_values( + Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1), + ) ᶜq_iceʲ_int_val = Fields.field_values(Fields.level(ᶜq_iceʲ, 1)) @. ᶜq_iceʲ_int_val = ᶜq_ice_int_val - ᶜq_rai_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_rai, 1)) + ᶜq_rai_int_val = Fields.field_values( + Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1), + ) ᶜq_raiʲ_int_val = Fields.field_values(Fields.level(ᶜq_raiʲ, 1)) @. ᶜq_raiʲ_int_val = ᶜq_rai_int_val - ᶜq_sno_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_sno, 1)) + ᶜq_sno_int_val = Fields.field_values( + Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1), + ) ᶜq_snoʲ_int_val = Fields.field_values(Fields.level(ᶜq_snoʲ, 1)) @. ᶜq_snoʲ_int_val = ᶜq_sno_int_val end @@ -367,7 +332,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 + (; ᶜu, ᶜp, ᶠu³⁰, ᶜts⁰) = p.precomputed (; ᶜmixing_length_tuple, ᶜmixing_length, @@ -478,6 +443,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos (; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed # First order approximation: Use environmental mean fields. @. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat + ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) @. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat @. ᶜgradᵥ_θ_liq_ice⁰ = ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat @@ -574,7 +540,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation (; params, dt) = p thp = CAP.thermodynamics_params(params) cmp = CAP.microphysics_0m_params(params) - (; ᶜts⁰, ᶜq_tot⁰, ᶜtsʲs, ᶜSqₜᵖʲs, ᶜSqₜᵖ⁰) = p.precomputed + (; ᶜts⁰, ᶜtsʲs, ᶜSqₜᵖʲs, ᶜSqₜᵖ⁰) = p.precomputed # Sources from the updrafts n = n_mass_flux_subdomains(p.atmos.turbconv_model) @@ -588,6 +554,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ) end # sources from the environment + ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, p.atmos.turbconv_model)) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰) return nothing end @@ -603,10 +570,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation thp = CAP.thermodynamics_params(params) cmp = CAP.microphysics_1m_params(params) cmc = CAP.microphysics_cloud_params(params) + (; turbconv_model) = p.atmos (; ᶜSqₗᵖʲs, ᶜSqᵢᵖʲs, ᶜSqᵣᵖʲs, ᶜSqₛᵖʲs, ᶜρʲs, ᶜtsʲs) = p.precomputed - (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜρ⁰, ᶜts⁰) = p.precomputed - (; ᶜq_tot⁰, ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed + (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜts⁰) = p.precomputed # TODO - can I re-use them between js and env? ᶜSᵖ = p.scratch.ᶜtemp_scalar @@ -677,6 +644,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation end # Precipitation sources and sinks from the environment + ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) + ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) + ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) + ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) compute_precipitation_sources!( ᶜSᵖ, ᶜSᵖ_snow, @@ -684,10 +655,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, - ᶜρ⁰, - ᶜq_tot⁰, - ᶜq_liq⁰, - ᶜq_ice⁰, + TD.air_density.(thp, ᶜts⁰), ᶜq_rai⁰, ᶜq_sno⁰, ᶜts⁰, @@ -699,10 +667,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ᶜSᵖ, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, - ᶜρ⁰, - ᶜq_tot⁰, - ᶜq_liq⁰, - ᶜq_ice⁰, + TD.air_density.(thp, ᶜts⁰), ᶜq_rai⁰, ᶜq_sno⁰, ᶜts⁰, diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index b6a8d6454b..d82653847a 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -633,15 +633,16 @@ compute_aren!(_, _, _, _, turbconv_model::T) where {T} = function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) + ᶜρa⁰ = @.lazy(ρa⁰(state.c)) if isnothing(out) return draft_area.( - cache.precomputed.ᶜρa⁰, + ᶜρa⁰, TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), ) else out .= draft_area.( - cache.precomputed.ᶜρa⁰, + ᶜρa⁰, TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), ) end @@ -943,6 +944,7 @@ function compute_clwen!( TD.liquid_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) end end + function compute_clwen!( out, state, @@ -951,11 +953,10 @@ function compute_clwen!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) - return cache.precomputed.ᶜq_liq⁰ + return specific_env_value(:q_liq, state.c, turbconv_model) else - out .= cache.precomputed.ᶜq_liq⁰ + out .= specific_env_value(:q_liq, state.c, turbconv_model) end end @@ -1007,6 +1008,7 @@ function compute_clien!( out .= TD.ice_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) end end + function compute_clien!( out, state, @@ -1015,11 +1017,10 @@ function compute_clien!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) - return cache.precomputed.ᶜq_ice⁰ + return specific_env_value(:q_ice, state.c, turbconv_model) else - out .= cache.precomputed.ᶜq_ice⁰ + out .= specific_env_value(:q_ice, state.c, turbconv_model) end end @@ -1064,11 +1065,10 @@ function compute_husraen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) - return cache.precomputed.ᶜq_rai⁰ + return specific_env_value(:q_rai, state.c, turbconv_model) else - out .= cache.precomputed.ᶜq_rai⁰ + out .= specific_env_value(:q_rai, state.c, turbconv_model) end end @@ -1113,11 +1113,10 @@ function compute_hussnen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) - return cache.precomputed.ᶜq_sno⁰ + return specific_env_value(:q_sno, state.c, turbconv_model) else - out .= cache.precomputed.ᶜq_sno⁰ + out .= specific_env_value(:q_sno, state.c, turbconv_model) end end @@ -1140,17 +1139,18 @@ compute_tke!(out, state, cache, time) = compute_tke!(_, _, _, _, turbconv_model::T) where {T} = error_diagnostic_variable("tke", turbconv_model) -function compute_tke!( +function compute_tke!( out, state, cache, time, turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX}, ) + if isnothing(out) - return copy(cache.precomputed.ᶜtke⁰) + return specific_tke(state.c.sgs⁰, state.c, turbconv_model) else - out .= cache.precomputed.ᶜtke⁰ + out .= specific_tke(state.c.sgs⁰, state.c, turbconv_model) 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..2c899fa93f 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -96,7 +96,7 @@ 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) = p.precomputed ## Momentum tendencies ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(Y.c.ρ) @@ -104,6 +104,7 @@ function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLill @. Yₜ.f.u₃ -= C3(wdivₕ(ᶠρ * ᶠτ_smag) / ᶠρ) ## Total energy tendency + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρe_tot += wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜh_tot)) ## Tracer diffusion and associated mass changes @@ -124,7 +125,7 @@ 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) = + (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, sfc_conditions) = p.precomputed (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions @@ -153,6 +154,7 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) @. Yₜ.f.u₃ -= C3(ᶠdivᵥ(Y.c.ρ * ᶜτ_smag) / ᶠρ) ## Total energy tendency + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜh_tot))) ## Tracer diffusion and associated mass changes diff --git a/src/parameterized_tendencies/microphysics/precipitation.jl b/src/parameterized_tendencies/microphysics/precipitation.jl index 3ed2eae58c..955e918359 100644 --- a/src/parameterized_tendencies/microphysics/precipitation.jl +++ b/src/parameterized_tendencies/microphysics/precipitation.jl @@ -148,7 +148,9 @@ function precipitation_tendency!( # Source terms from EDMFX updrafts (; ᶜSqₗᵖʲs, ᶜSqᵢᵖʲs, ᶜSqᵣᵖʲs, ᶜSqₛᵖʲs) = p.precomputed # Source terms from EDMFX environment - (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜρa⁰) = p.precomputed + (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰) = p.precomputed + + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) # Update from environment precipitation and cloud formation sources/sinks @. Yₜ.c.ρq_liq += ᶜρa⁰ * ᶜSqₗᵖ⁰ diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl index 18cc7771ef..9e97e933b7 100644 --- a/src/parameterized_tendencies/radiation/radiation.jl +++ b/src/parameterized_tendencies/radiation/radiation.jl @@ -448,7 +448,8 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS) @assert !(p.atmos.moisture_model isa DryModel) (; params) = p - (; ᶜκρq, ∫_0_∞_κρq, ᶠ∫_0_z_κρq, isoline_z_ρ_ρq, ᶠradiation_flux) = + (; ᶜts) = p.precomputed + (; ᶜκρq, ∫_0_∞_κρq, ᶠ∫_0_z_κρq, isoline_z_ρ_q, ᶠradiation_flux) = p.radiation (; ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(params) @@ -475,10 +476,10 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS) q_tot_isoline = FT(0.008) Operators.column_reduce!( (nt1, nt2) -> - abs(specific(nt1.ρq_tot, nt1.ρ) - q_tot_isoline) < - abs(specific(nt2.ρq_tot, nt2.ρ) - q_tot_isoline) ? nt1 : nt2, - isoline_z_ρ_ρq, - Base.broadcasted(NT ∘ tuple, ᶜz, Y.c.ρ, Y.c.ρq_tot), + abs(nt1.q_tot - q_tot_isoline) < abs(nt2.q_tot - q_tot_isoline) ? + nt1 : nt2, + isoline_z_ρ_q, + Base.broadcasted(NT ∘ tuple, ᶜz, Y.c.ρ, specific(Y.c.ρq_tot, Y.c.ρ)), ) zi = isoline_z_ρ_ρq.z @@ -565,4 +566,4 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationISDAC) )) # = -∂F/∂z = ρ cₚ ∂T/∂t (longwave radiation) return nothing -end +end \ No newline at end of file diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index fb31a95910..a8ad9e5f72 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -49,7 +49,7 @@ NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) end end - (; ᶜh_tot) = p.precomputed + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρe_tot -= wdivₕ(Y.c.ρ * ᶜh_tot * ᶜu) if p.atmos.turbconv_model isa PrognosticEDMFX @@ -178,7 +178,6 @@ 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 ᶠu³⁰ = advect_tke ? @@ -186,9 +185,20 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ : p.precomputed.ᶠu³⁰ ) : nothing - ᶜρa⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρa⁰ : Y.c.ρ) : nothing - ᶜρ⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρ⁰ : Y.c.ρ) : nothing - ᶜtke⁰ = advect_tke ? p.precomputed.ᶜtke⁰ : nothing + ᶜρa⁰ = advect_tke ? (n > 0 ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ) : nothing + ᶜρ⁰ = if advect_tke + if n > 0 + thermo_params = CAP.thermodynamics_params(p.params) + @. TD.air_density(thermo_params, p.precomputed.ᶜts⁰) + else + Y.c.ρ + end + else + nothing + end + ᶜtke⁰ = advect_tke ? + (@.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))) : + nothing ᶜa_scalar = p.scratch.ᶜtemp_scalar ᶜω³ = p.scratch.ᶜtemp_CT3 ᶠω¹² = p.scratch.ᶠtemp_CT12 @@ -223,7 +233,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) # ... and upwinding correction of energy and total water. # (The central advection of energy and total water is done implicitly.) if energy_upwinding != Val(:none) - (; ᶜh_tot) = p.precomputed + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) @. Yₜ.c.ρe_tot += vtt - vtt_central diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index c3eeef3734..47e8a895b3 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -528,11 +528,15 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF n = n_mass_flux_subdomains(turbconv_model) (; ᶜturb_entrʲs, ᶜentrʲs, ᶜdetrʲs) = p.precomputed - (; ᶜq_tot⁰, ᶜmse⁰, ᶠu₃⁰) = p.precomputed + (; ᶠu₃⁰) = p.precomputed + ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed + ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) + ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) + ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) + ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) end for j in 1:n @@ -544,6 +548,7 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * (ᶜmse⁰ - Y.c.sgsʲs.:($$j).mse) + ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) @. Yₜ.c.sgsʲs.:($$j).q_tot += (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * (ᶜq_tot⁰ - Y.c.sgsʲs.:($$j).q_tot) diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 5fced40005..5e42161b2c 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -40,15 +40,12 @@ 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 - if ( - p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - ) - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed - end + (; ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜρ⁰ = @. TD.air_density(thermo_params, ᶜts⁰) + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J @@ -59,6 +56,7 @@ function edmfx_sgs_mass_flux_tendency!( # [best after removal of precomputed quantities] ᶠu³_diff = p.scratch.ᶠtemp_CT3 ᶜa_scalar = p.scratch.ᶜtemp_scalar + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @. ᶜa_scalar = @@ -75,6 +73,7 @@ function edmfx_sgs_mass_flux_tendency!( end # Add the environment fluxes @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ + ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, @@ -93,7 +92,7 @@ function edmfx_sgs_mass_flux_tendency!( (Y.c.sgsʲs.:($$j).q_tot - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( - ᶜρʲs.:($j), + ᶜρʲs.:($$j), ᶠu³_diff, ᶜa_scalar, dt, @@ -102,9 +101,11 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_tot += vtt end # Add the environment fluxes + ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = - (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * + draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -119,13 +120,19 @@ function edmfx_sgs_mass_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) + ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) + ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) + ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) + ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) # Liquid, ice, rain and snow specific humidity fluxes for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @. ᶜa_scalar = - (Y.c.sgsʲs.:($$j).q_liq - specific(Y.c.ρq_liq, Y.c.ρ)) * - draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + ( + Y.c.sgsʲs.:($$j).q_liq - + specific(Y.c.ρq_liq, Y.c.ρ) + ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -136,8 +143,10 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_liq += vtt @. ᶜa_scalar = - (Y.c.sgsʲs.:($$j).q_ice - specific(Y.c.ρq_ice, Y.c.ρ)) * - draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + ( + Y.c.sgsʲs.:($$j).q_ice - + specific(Y.c.ρq_ice, Y.c.ρ) + ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -148,8 +157,10 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_ice += vtt @. ᶜa_scalar = - (Y.c.sgsʲs.:($$j).q_rai - specific(Y.c.ρq_rai, Y.c.ρ)) * - draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + ( + Y.c.sgsʲs.:($$j).q_rai - + specific(Y.c.ρq_rai, Y.c.ρ) + ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -160,8 +171,10 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_rai += vtt @. ᶜa_scalar = - (Y.c.sgsʲs.:($$j).q_sno - specific(Y.c.ρq_sno, Y.c.ρ)) * - draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + ( + Y.c.sgsʲs.:($$j).q_sno - + specific(Y.c.ρq_sno, Y.c.ρ) + ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -174,7 +187,8 @@ function edmfx_sgs_mass_flux_tendency!( @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = - (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * + draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -185,7 +199,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_liq += vtt @. ᶜa_scalar = - (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * + draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -196,7 +211,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_ice += vtt @. ᶜa_scalar = - (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * + draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -207,7 +223,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_rai += vtt @. ᶜa_scalar = - (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * + draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -236,14 +253,15 @@ 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 + (; ᶠu³) = 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) if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} - # Enthalpy fluxes + # energy + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) ᶠu³_diff = p.scratch.ᶠtemp_CT3 ᶜa_scalar = p.scratch.ᶜtemp_scalar for j in 1:n @@ -276,7 +294,7 @@ function edmfx_sgs_mass_flux_tendency!( for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ # @. ᶜa_scalar = - # (ᶜq_totʲs.:($$j) - specific(Y.c.ρq_tot, Y.c.ρ)) * + # (ᶜq_totʲs.:($$j) - specific(Y.c.ρq_tot, Y.c.ρ) * # draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) # TODO: remove this filter when mass flux is treated implicitly @. ᶜa_scalar = @@ -300,19 +318,51 @@ function edmfx_sgs_mass_flux_tendency!( end # TODO: add environment flux? end + # TODO: the following adds the environment flux to the tendency + # Make active and test later + # @. ᶠu³_diff = p.precomputed.ᶠu³⁰ - ᶠu³ + # ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + # ᶜρ⁰ = p.scratch.ᶜtemp_scalar_2 + # @. ᶜρ⁰ = TD.air_density( + # CAP.thermodynamics_params(p.params), + # p.precomputed.ᶜts⁰, + # ) + # ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) + # @. ᶜa_scalar = + # (ᶜmse⁰ + p.precomputed.ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) + # vtt = vertical_transport( + # ᶜρ⁰, + # ᶠu³_diff, + # ᶜa_scalar, + # dt, + # edmfx_sgsflux_upwinding, + # ) + # @. Yₜ.c.ρe_tot += vtt + # if !(p.atmos.moisture_model isa DryModel) + # ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + # @. ᶜa_scalar = + # (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * + # draft_area(ᶜρa⁰, ᶜρ⁰) + # vtt = vertical_transport( + # ᶜρ⁰, + # ᶠu³_diff, + # ᶜa_scalar, + # dt, + # edmfx_sgsflux_upwinding, + # ) + # @. Yₜ.c.ρq_tot += vtt + # end end - # TODO: add vertical momentum fluxes - - return nothing end """ edmfx_sgs_diffusive_flux_tendency!(Yₜ, Y, p, t, turbconv_model) -Computes and applies tendencies to the grid-mean prognostic variables due to the -divergence of subgrid-scale (SGS) diffusive fluxes, representing turbulent mixing -by the EDMFX environment component. +Computes and applies the tendency to the grid-mean state `Y` due to SGS +diffusive fluxes from the EDMFX environment. This involves calculating the +divergence of turbulent fluxes, which are parameterized using eddy diffusivity +and viscosity closures. This function parameterizes these fluxes using an eddy-diffusivity/viscosity approach (K-theory) for the environment (sgs⁰). Tendencies are calculated for @@ -326,9 +376,9 @@ in place. Arguments: - `Yₜ`: The tendency state vector for grid-mean variables. -- `Y`: The current state vector. -- `p`: Cache containing parameters, precomputed fields, atmospheric - model settings, and scratch space. +- `Y`: The current state vector (used for grid-mean and SGS properties). +- `p`: Cache containing parameters, precomputed fields, atmospheric model settings, + and scratch space. - `t`: Current simulation time. - `turbconv_model`: The turbulence convection model instance. """ @@ -345,15 +395,11 @@ 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 - if ( - p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - ) - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed - end - (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed + (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm,) = p.precomputed + (; ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar @@ -366,6 +412,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) + ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜmse⁰ + ᶜK⁰))) if use_prognostic_tke(turbconv_model) # Turbulent TKE transport (diffusion) @@ -394,6 +441,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) + ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜq_tot⁰))) @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion @. Yₜ.c.ρ -= ᶜρχₜ_diffusion # Effect of moisture diffusion on (moist) air mass @@ -402,6 +450,10 @@ function edmfx_sgs_diffusive_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) + ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) + ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) + ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) + ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) # Liquid, ice, rain and snow specific humidity diffusion α_vert_diff_tracer = CAP.α_vert_diff_tracer(params) @@ -449,10 +501,12 @@ function edmfx_sgs_diffusive_flux_tendency!( FT = Spaces.undertype(axes(Y.c)) (; dt, params) = p turbconv_params = CAP.turbconv_params(params) + thermo_params = CAP.thermodynamics_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜu, ᶜh_tot, ᶜtke⁰, ᶜmixing_length) = p.precomputed + (; ᶜu, ᶜmixing_length) = p.precomputed (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() + ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar @@ -465,6 +519,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜh_tot))) if use_prognostic_tke(turbconv_model) diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index ae394835a9..18ddf17512 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -44,9 +44,17 @@ 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 - (; ᶜK_u, ᶜK_h) = p.precomputed - ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ + (; + ᶠu³⁰, + ᶠu³, + ᶜstrain_rate_norm, + ᶜlinear_buoygrad + ) = p.precomputed + turbconv_params = CAP.turbconv_params(p.params) + FT = eltype(p.params) + + + ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ nh_pressure3_buoyʲs = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs @@ -73,6 +81,8 @@ function edmfx_tke_tendency!( # buoyancy production @. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰ * ᶜK_h * ᶜlinear_buoygrad + ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + # entrainment and detraiment # using ᶜu⁰ and local geometry results in allocation for j in 1:n diff --git a/src/prognostic_equations/forcing/external_forcing.jl b/src/prognostic_equations/forcing/external_forcing.jl index 0d1f023fde..64710ad3c0 100644 --- a/src/prognostic_equations/forcing/external_forcing.jl +++ b/src/prognostic_equations/forcing/external_forcing.jl @@ -403,7 +403,8 @@ function external_forcing_tendency!( Yₜ.c.ρq_tot, Y.c.ρ, ᶠls_subsidence³, - ᶜq_tot, + specific(Y.c.ρq_tot, Y.c.ρ), + specific(Y.c.ρq_tot, Y.c.ρ), Val{:first_order}(), ) diff --git a/src/prognostic_equations/forcing/subsidence.jl b/src/prognostic_equations/forcing/subsidence.jl index c722a50e03..8b996d8b91 100644 --- a/src/prognostic_equations/forcing/subsidence.jl +++ b/src/prognostic_equations/forcing/subsidence.jl @@ -65,10 +65,34 @@ If `subsidence_model` is `Nothing`, no subsidence tendency is applied. """ subsidence_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing # No subsidence -function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) + +""" + subsidence_tendency!(Yₜ, Y, p, t, subsidence_model::Subsidence) + +Applies subsidence tendencies to total energy (`ρe_tot`), total specific humidity +(`ρq_tot`), and other moisture species (`ρq_liq`, `ρq_ice`) if a `NonEquilMoistModel` +is used. + +The subsidence velocity profile `w_sub(z)` is obtained from `subsidence_model.prof`. +This profile is used to construct a face-valued vertical velocity field `ᶠsubsidence³`. +The `subsidence!` helper function is then called (currently with a first-order +upwind scheme) to compute and apply the vertical advective tendency for each relevant +scalar quantity `χ`. + +Arguments: +- `Yₜ`: The tendency state vector, modified in place. +- `Y`: The current state vector, used for density (`ρ`). +- `p`: Cache containing parameters, and the subsidence model object. +- `t`: Current simulation time. +- `subsidence`: The subsidence model object, containing the prescribed vertical + velocity profile `Dᵥ`. +""" +function subsidence_tendency!(Yₜ, Y, p, t, subsidence::Subsidence) + (; Dᵥ) = subsidence + ᶜρ = Y.c.ρ (; moisture_model) = p.atmos subsidence_profile = p.atmos.subsidence.prof - (; ᶜh_tot) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) ᶠz = Fields.coordinate_field(axes(Y.f)).z ᶠlg = Fields.local_geometry_field(Y.f) @@ -76,17 +100,23 @@ function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) @. ᶠsubsidence³ = subsidence_profile(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg)) - # Large-scale subsidence + # LS Subsidence + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}()) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - subsidence!(Yₜ.c.ρq_tot, Y.c.ρ, ᶠsubsidence³, ᶜq_tot, Val{:first_order}()) + subsidence!( + Yₜ.c.ρq_tot, + Y.c.ρ, + ᶠsubsidence³, + specific(Y.c.ρq_tot, Y.c.ρ), + Val{:first_order}(), + ) if moisture_model isa NonEquilMoistModel ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) subsidence!( Yₜ.c.ρq_liq, Y.c.ρ, ᶠsubsidence³, - ᶜq_liq, + specific(Y.c.ρq_liq, Y.c.ρ), Val{:first_order}(), ) ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) @@ -94,7 +124,7 @@ function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) Yₜ.c.ρq_ice, Y.c.ρ, ᶠsubsidence³, - ᶜq_ice, + specific(Y.c.ρq_ice, Y.c.ρ), Val{:first_order}(), ) end diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index c93adb5023..7d95c58729 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -117,7 +117,8 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) if diffuse_tke - (; ᶜtke⁰) = p.precomputed + ᶜtke⁰ = + @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) (; ᶜ∇²tke⁰) = p.hyperdiff @. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰)) end @@ -148,14 +149,14 @@ 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 (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX - (; ᶜρa⁰) = p.precomputed + ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end if use_prognostic_tke(turbconv_model) - (; ᶜtke⁰) = p.precomputed + ᶜtke⁰ = + @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) (; ᶜ∇²tke⁰) = p.hyperdiff end @@ -251,6 +252,7 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos isnothing(hyperdiff) && return nothing + ᶜspecific = all_specific_gs(Y.c) (; ᶜ∇²specific_tracers) = p.hyperdiff # TODO: Fix RecursiveApply bug in gradₕ to fuse this operation. @@ -292,6 +294,7 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) ν₄_scalar_for_precip = CAP.α_hyperdiff_tracer(p.params) * ν₄_scalar n = n_mass_flux_subdomains(turbconv_model) + ᶜspecific = all_specific_gs(Y.c) (; ᶜ∇²specific_tracers) = p.hyperdiff # TODO: Since we are not applying the limiter to density (or area-weighted @@ -335,3 +338,4 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) end return nothing end + diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 1468746c67..5e5f4d1f31 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -140,7 +140,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 + (; ᶠu³, ᶜp) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠu³) @@ -148,8 +150,13 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜh_tot, dt, Val(:none)) @. Yₜ.c.ρe_tot += vtt if !(moisture_model isa DryModel) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜq_tot, dt, Val(:none)) + vtt = vertical_transport( + Y.c.ρ, + ᶠu³, + specific(Y.c.ρq_tot, Y.c.ρ), + dt, + Val(:none), + ) @. Yₜ.c.ρq_tot += vtt end @@ -160,12 +167,14 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) if moisture_model isa NonEquilMoistModel (; ᶜwₗ, ᶜwᵢ) = p.precomputed @. Yₜ.c.ρq_liq -= ᶜprecipdivᵥ( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * + ᶠright_bias( Geometry.WVector(-(ᶜwₗ)) * specific(Y.c.ρq_liq, Y.c.ρ), ), ) @. Yₜ.c.ρq_ice -= ᶜprecipdivᵥ( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * + ᶠright_bias( Geometry.WVector(-(ᶜwᵢ)) * specific(Y.c.ρq_ice, Y.c.ρ), ), ) @@ -173,7 +182,8 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) if microphysics_model isa Microphysics1Moment (; ᶜwᵣ, ᶜwₛ) = p.precomputed @. Yₜ.c.ρq_rai -= ᶜprecipdivᵥ( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * + ᶠright_bias( Geometry.WVector(-(ᶜwᵣ)) * specific(Y.c.ρq_rai, Y.c.ρ), ), ) @@ -202,6 +212,8 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) ) @. Yₜ.c.ρq_sno -= ᶜprecipdivᵥ( ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * + ᶠright_bias( Geometry.WVector(-(ᶜwₛ)) * specific(Y.c.ρq_sno, Y.c.ρ), ), ) diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index 872a7185ea..d8cd65aa03 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -338,7 +338,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 + (; ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, @@ -366,8 +366,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) Δcp_v = FT(CAP.cp_v(params)) - cp_d # 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) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) ᶜρ = Y.c.ρ ᶜuₕ = Y.c.uₕ @@ -416,14 +416,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) - tracer_info = (@name(c.ρe_tot), @name(c.ρq_tot)) - MatrixFields.unrolled_foreach(tracer_info) do ρχ_name - MatrixFields.has_field(Y, ρχ_name) || return - ᶜχ = if ρχ_name === @name(c.ρe_tot) - p.precomputed.ᶜh_tot - else - @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - end + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) if use_derivative(topography_flag) ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] @. ∂ᶜρχ_err_∂ᶜuₕ = @@ -446,14 +440,16 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ᶠinterp_matrix() ) @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) + e_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) + q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) if MatrixFields.has_field(Y, @name(c.ρq_tot)) ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] @. ∂ᶠu₃_err_∂ᶜρq_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot + ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + specific(Y.c.ρe_tot, Y.c.ρ) - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) + cp_d * T_0 + e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * q_tot ) )) end @@ -518,15 +514,15 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρq_tot_err_∂ᶜρq_tot = matrix[@name(c.ρq_tot), @name(c.ρq_tot)] @. ∂ᶜρq_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρq_tot_err_∂ᶜρq_tot)) - (I,) - #TODO: tetsing explicit vs implicit + #TODO: testing explicit vs implicit #@. ∂ᶜρq_tot_err_∂ᶜρq_tot = # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ # DiagonalMatrixRow( # -1 / ᶜρ * ifelse( - # specific(Y.c.ρq_tot, Y.c.ρ) == 0, + # q_tot == 0, # (Geometry.WVector(FT(0)),), - # p.precomputed.ᶜwₜqₜ / specific(Y.c.ρq_tot, Y.c.ρ), + # p.precomputed.ᶜwₜqₜ / _tot, # ), # ) - (I,) @@ -567,8 +563,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρe_tot_err_∂ᶜρ = dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( ( - -(1 + ᶜkappa_m) * specific(Y.c.ρe_tot, Y.c.ρ) - - ᶜkappa_m * ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) + -(1 + ᶜkappa_m) * e_tot - + ᶜkappa_m * ∂e_int_∂q_tot * q_tot ) / ᶜρ, ) @. ∂ᶜρe_tot_err_∂ᶜρe_tot += @@ -581,28 +577,26 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + specific(Y.c.ρe_tot, Y.c.ρ) - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) + cp_d * T_0 + e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * q_tot ) )) @. ∂ᶜρq_tot_err_∂ᶜρ = dtγ * ᶜdiffusion_h_matrix ⋅ - DiagonalMatrixRow(-(specific(Y.c.ρq_tot, Y.c.ρ)) / ᶜρ) + DiagonalMatrixRow(-(q_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) end - MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, _) - MatrixFields.has_field(Y, ρχ_name) || return - ᶜρχ = MatrixFields.get_field(Y, ρχ_name) + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) ∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)] ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] - ᶜtridiagonal_matrix_scalar = ifelse( - ρχ_name in (@name(c.ρq_rai), @name(c.ρq_sno), @name(c.ρn_rai)), - ᶜdiffusion_h_matrix_scaled, - ᶜdiffusion_h_matrix, - ) + ᶜtridiagonal_matrix_scalar = if ρχ_name in (@name(c.ρq_rai), @name(c.ρq_sno), @name(c.ρn_rai)) + ᶜdiffusion_h_matrix_scaled + else + ᶜdiffusion_h_matrix + end @. ∂ᶜρχ_err_∂ᶜρ = dtγ * ᶜtridiagonal_matrix_scalar ⋅ DiagonalMatrixRow(-(ᶜχ) / ᶜρ) @. ∂ᶜρχ_err_∂ᶜρχ += @@ -613,10 +607,10 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) (; dt) = p - (; ᶜtke⁰, ᶜmixing_length) = p.precomputed - ᶜρa⁰ = - p.atmos.turbconv_model isa PrognosticEDMFX ? - p.precomputed.ᶜρa⁰ : ᶜρ + turbconv_model = p.atmos.turbconv_model + ᶜmixing_length = p.precomputed.ᶜmixing_length + ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ ᶜρatke⁰ = Y.c.sgs⁰.ρatke @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) = @@ -946,10 +940,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ DiagonalMatrixRow( ( - -(1 + ᶜkappa_m) * specific(Y.c.ρe_tot, Y.c.ρ) - - ᶜkappa_m * - ∂e_int_∂q_tot * - specific(Y.c.ρq_tot, Y.c.ρ) + -(1 + ᶜkappa_m) * e_tot - + ᶜkappa_m * ∂e_int_∂q_tot * q_tot ) / ᶜρ, ) @@ -958,8 +950,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + specific(Y.c.ρe_tot, Y.c.ρ) - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) + cp_d * T_0 + e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * q_tot ) )) @@ -975,7 +967,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ## grid-mean ρq_tot @. ∂ᶜρq_tot_err_∂ᶜρ += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(-(specific(Y.c.ρq_tot, Y.c.ρ)) / ᶜρ) + DiagonalMatrixRow(-(q_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ @@ -1014,10 +1006,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶠu₃ += dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( ᶠinterp( - ( - Y.c.sgsʲs.:(1).q_tot - - specific(Y.c.ρq_tot, Y.c.ρ) - ) * + (Y.c.sgsʲs.:(1).q_tot - q_tot) * ᶜρʲs.:(1) * ᶜJ * draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), @@ -1029,10 +1018,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶠu₃ʲ = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠinterp( - ( - Y.c.sgsʲs.:(1).q_tot - - specific(Y.c.ρq_tot, Y.c.ρ) - ) * + (Y.c.sgsʲs.:(1).q_tot - q_tot) * ᶜρʲs.:(1) * ᶜJ * draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), @@ -1052,9 +1038,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) matrix[@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)] @. ∂ᶜρq_tot_err_∂ᶜρa = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - (ᶠu³ʲs.:(1) - ᶠu³) * ᶠinterp(( - Y.c.sgsʲs.:(1).q_tot - specific(Y.c.ρq_tot, Y.c.ρ) - )) / ᶠJ, + (ᶠu³ʲs.:(1) - ᶠu³) * + ᶠinterp((Y.c.sgsʲs.:(1).q_tot - q_tot)) / ᶠJ, ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) end elseif rs isa RayleighSponge diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index ae94a52563..f557ae9bbc 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,6 @@ dependencies or operator splitting assumptions. """ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) - (; ᶜh_tot, ᶜspecific) = p.precomputed ᶜuₕ = Y.c.uₕ ᶠu₃ = Y.f.u₃ ᶜρ = Y.c.ρ @@ -148,6 +146,7 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) thermo_params = CAP.thermodynamics_params(params) (; ᶜp, sfc_conditions, ᶜts) = p.precomputed + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge) vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge) @@ -168,7 +167,6 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) @. Yₜ.f.u₃.components.data.:1 += vst_u₃ @. Yₜ.c.ρe_tot += vst_ρe_tot - # TODO: can we write this out explicitly? foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge) diff --git a/src/prognostic_equations/surface_flux.jl b/src/prognostic_equations/surface_flux.jl index 10b6b1ef90..8dc015373f 100644 --- a/src/prognostic_equations/surface_flux.jl +++ b/src/prognostic_equations/surface_flux.jl @@ -104,7 +104,7 @@ 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 + (; sfc_conditions) = p.precomputed if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion) btt = @@ -112,6 +112,7 @@ function surface_flux_tendency!(Yₜ, Y, p, t) @. Yₜ.c.uₕ -= btt end + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) btt = boundary_tendency_scalar(ᶜh_tot, sfc_conditions.ρ_flux_h_tot) @. Yₜ.c.ρe_tot -= btt ρ_flux_χ = p.scratch.sfc_temp_C3 diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index bca3cd65c2..c5828944e5 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -46,14 +46,20 @@ This function is dispatched based on the type of the vertical diffusion model for scalars. Zero-flux boundary conditions are explicitly applied. - **Note on mass conservation for `q_tot` diffusion**: The current implementation also modifies the tendency of total moist air density `Yₜ.c.ρ` based on the - diffusion tendency of total specific humidity `ρq_tot`: + diffusion tendency of total specific humidity `ρq_tot`: `Yₜ.c.ρ -= ᶜρχₜ_diffusion_for_q_tot`. -Arguments for all methods: -- `Yₜ`: The tendency state vector, modified in place. +This function is acting as a wrapper around the specific implementations +for different turbulence and convection models. + +The primary role of this function is to dispatch to the correct turbulence model's +tendency function. It operates on the state `Y` and its tendency `Yₜ`, using +the model-specific cache `p`. + +Arguments: +- `Yₜ`: The tendency state vector. - `Y`: The current state vector. - `p`: Cache containing parameters (e.g., `p.params` for `CAP.α_vert_diff_tracer`), - precomputed fields (e.g., `ᶜK_u`, `ᶜK_h`, `ᶜh_tot`, `ᶜspecific` tracer values), atmospheric model configurations (like `p.atmos.vertical_diffusion`), and scratch space. - `t`: Current simulation time (not directly used in diffusion calculations). - `vert_diff_model` (for dispatched methods): The specific vertical diffusion model instance. @@ -82,7 +88,8 @@ 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 + thermo_params = CAP.thermodynamics_params(p.params) + (; ᶜu, ᶜK_u, ᶜK_h) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion) @@ -97,6 +104,7 @@ function vertical_diffusion_boundary_layer_tendency!( top = Operators.SetValue(C3(0)), bottom = Operators.SetValue(C3(0)), ) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜh_tot))) diff --git a/src/solver/types.jl b/src/solver/types.jl index 3e2fbb8388..ce5bc047cc 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -358,13 +358,13 @@ This allows running simulations with TKE-based vertical diffusion. struct EDOnlyEDMFX <: AbstractEDMF end struct PrognosticEDMFX{N, TKE, FT} <: AbstractEDMF - a_half::FT # WARNING: this should never be used outside of divide_by_ρa + a_half::FT # WARNING: this should never be used outside of `specific` end PrognosticEDMFX{N, TKE}(a_half::FT) where {N, TKE, FT} = PrognosticEDMFX{N, TKE, FT}(a_half) struct DiagnosticEDMFX{N, TKE, FT} <: AbstractEDMF - a_half::FT # WARNING: this should never be used outside of divide_by_ρa + a_half::FT # WARNING: this should never be used outside of d`specific` end DiagnosticEDMFX{N, TKE}(a_half::FT) where {N, TKE, FT} = DiagnosticEDMFX{N, TKE, FT}(a_half) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 443c5eb359..80c88d65f8 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -178,198 +178,297 @@ foreach_gs_tracer(f::F, Y_or_similar_values...) where {F} = end """ + specific(ρχ, ρ) + specific(ρaχ, ρa, ρχ, ρ, turbconv_model) + +Calculates the specific quantity `χ` (per unit mass) from a density-weighted +quantity. This function uses multiple dispatch to select the appropriate +calculation method based on the number of arguments. + +**Grid-Scale Method (2 arguments)** + + specific(ρχ, ρ) + +Performs a direct division of the density-weighted quantity `ρχ` by the density `ρ`. +This method is used for grid-mean quantities where the density `ρ` is well-defined +and non-zero. + +**SGS Regularized Method (5 arguments)** + + specific(ρaχ, ρa, ρχ, ρ, turbconv_model) + +Calculates the specific quantity `χ` for a subgrid-scale (SGS) component by +dividing the density-area-weighted quantity `ρaχ` by the density-area +product `ρa`. + +This method includes regularization to handle cases where the SGS area fraction +(and thus `ρa`) is zero or vanishingly small. It performs a linear interpolation +between the SGS specific quantity (`ρaχ / ρa`) and the grid-mean specific +quantity (`ρχ / ρ`). The interpolation weight is computed by `sgs_weight_function` +to ensure a smooth and numerically stable transition, preventing division by zero. +Using this regularized version instead of directly computing `ρaχ / ρa` breaks the +assumption of domain decomposition (sum of SGS domains equals GS) when the approximated +area fraction `a` is small. + +Arguments: +- `ρaχ`: The density-area-weighted SGS quantity (e.g., `sgs.ρa * sgs.h_tot`). +- `ρa`: The density-area product of the SGS component. +- `ρχ`: The fallback grid-mean density-weighted quantity (e.g., `ρe_tot`, `ρq_tot`). +- `ρ`: The fallback grid-mean density. +- `turbconv_model`: The turbulence convection model, containing parameters for regularization (e.g., `a_half`). +""" +function specific(ρχ, ρ) + return ρχ / ρ +end + +function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) + # TODO: Replace turbconv_model struct by parameters, and include a_half in + # parameters, not in config + weight = sgs_weight_function(ρa / ρ, turbconv_model.a_half) + # If ρa is exactly zero, the weight function will be zero, causing the first + # term to be NaN (0 * ... / 0). The ifelse handles this case explicitly. + return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ +end + + """ sgs_weight_function(a, a_half) -Computes the weight of the SGS variables in the linear interpolation used in -`divide_by_ρa`. This is a continuously differentiable and monotonically -increasing function of `a` that is equal to 0 when `a ≤ 0`, is equal to 1 when -`a ≥ 1`, is equal to `1 / 2` when `a = a_half`, grows very rapidly near -`a = a_half`, and grows very slowly at all other values of `a`. If `a_half` is -sufficiently small, this function is essentially equal to 1 for all `a` more -than a few times larger than `a_half` (up to floating-point precision). - -We will now provide a description of how this function was constructed. We need -the function to be equal to 0 when `a ≤ 0` and equal to 1 when `a ≥ 1`. Since -the function must also be continuously differentiable, its derivative at these -values of `a` has to be 0. To obtain a function with these properties, we use a -piecewise definition: - - For all `a < 0`, the function is equal to 0. - - For all `a > 1`, the function is equal to 1. - - For all `0 ≤ a ≤ 1`, the function is a sigmoid that connects the point - `(0, 0)` to the point `(1, 1)`, with a derivative of 0 at these points. -Most well-known sigmoid functions connect the "points" `(-Inf, 0)` and -`(Inf, 1)`, not `(0, 0)` and `(1, 1)`. To obtain the desired sigmoid curve, we -begin with two simple sigmoid functions that go from `(-Inf, 0)` to `(Inf, 1)` -at different rates. In this case, we use two `tanh` functions, scaled and -translated so that they lie between 0 and 1: - - `fast_sigmoid(a) = (1 + tanh(a)) / 2` and - - `slow_sigmoid(a) = (1 + tanh(a / 2)) / 2`. -Note that the second sigmoid is commonly called the "logistic" function. We then -take the inverse of the sigmoid that grows more slowly, and we make that the -input of the sigmoid that grows more quickly: - - `sigmoid(a) = fast_sigmoid(slow_sigmoid⁻¹(a)) = - (1 + tanh(2 * atanh(2 * a - 1))) / 2`. -The resulting function goes from `(0, 0)` to `(1, 1)`, and, since the outer -sigmoid grows more quickly, it has the same asymptotic behavior as the outer -sigmoid, which means that its derivative at the boundary points is 0. If we had -instead put the sigmoid that grows more slowly on the outside, the asymptotic -behavior would come from the inverted inner sigmoid, which means that the -derivative at the boundary points would be `Inf`. - -The sigmoid function we have constructed reaches `1 / 2` when `a = 1 / 2`. More -generally, we need the weight function to reach `1 / 2` when `a` is some small -value `a_half`. To achieve this, we replace the input to the sigmoid function -with a smooth, monotonically increasing function that goes through `(0, 0)`, -`(a_half, 1 / 2)`, and `(1, 1)`. The simplest option is the power function - - `power(a) = a^(-1 / log2(a_half))`. -However, this function does not work well because, when `a_half < 1 / 2`, its -derivative at `a = 0` is `Inf`, and its derivative at `a = 1` is some positive -number. Making this power function the input to the sigmoid function causes the -derivative of the sigmoid to become `Inf` at `a = 0` when `a_half < 1 / 4`, and -it causes the sigmoid to grow too slowly from `1 / 2` to 1, only reaching 1 when -`a` is significantly larger than `a_half`. In order to fix this, we transform -the power function by replacing `a` with `1 - a`, `a_half` with `1 - a_half`, -and `power(a)` with `1 - power(a)`, which gives us - - `power(a) = 1 - (1 - a)^(-1 / log2(1 - a_half))`. -This transformed function works better because, when `a_half < 1 / 2`, its -derivative at `a = 0` is some positive number, and its derivative at `a = 1` is -0. When we make this the input to the sigmoid function, the result has a -continuous derivative and is essentially equal to 1 for all `a` more than a few -times larger than `a_half`. So, for all `0 ≤ a ≤ 1`, we define the weight -function as - - `weight(a) = sigmoid(power(a)) = - (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2`. -""" -sgs_weight_function(a, a_half) = - if a <= 0 # autodiff generates NaNs when a is 0 +Computes a smooth, monotonic weight function `w(a)` that ranges from 0 to 1. + +This function is used as the interpolation weight in the regularized `specific` +function. It ensures a numerically stable and smooth transition between a subgrid-scale +(SGS) quantity and its grid-mean counterpart, especially when the SGS area fraction `a` +is small. + +**Key Properties:** +- `w(a) = 0` for `a ≤ 0`. +- `w(a) = 1` for `a ≥ 1`. +- `w(a_half) = 0.5`. +- The function is continuously differentiable, with derivatives equal to zero at + `a = 0` and `a = 1`, which ensures smooth blending. +- The functions grows very rapidly near `a = a_half`, and grows very slowly at all other + values of `a`. +- For small `a_half`, the weight rapidly approaches 1 for values of `a` that are + a few times larger than `a_half`. + +**Construction Method:** +The function is piecewise. For `a` between 0 and 1, it is a custom sigmoid curve +constructed in two main steps to satisfy the key properties: +1. **Bounded Sigmoid Creation**: A base sigmoid is created that maps the interval + `(0, 1)` to `(0, 1)` with zero derivatives at the endpoints. This is achieved + by composing a standard `tanh` function with the inverse of a slower-growing + `tanh` function. +2. **Midpoint Control**: To ensure the function passes through the control point + `(a_half, 0.5)`, the input `a` is first transformed by a specially designed + power function (`1 - (1 - a)^k`) before being passed to the bounded sigmoid. + This transformation maps `a_half` to `0.5` while preserving differentiability + at the boundaries. + +Arguments: +- `a`: The input SGS area fraction (often approximated as `ρa / ρ`). +- `a_half`: The value of `a` at which the weight function should be 0.5, controlling + the transition point of the sigmoid curve. + +Returns: +- The computed weight, a value between 0 and 1. +""" +function sgs_weight_function(a, a_half) + if a < 0 zero(a) elseif a > min(1, 42 * a_half) # autodiff generates NaNs when a is large one(a) else (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2 end +end """ - divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model) + draft_sum(f, sgsʲs) -Computes `ρaχ / ρa`, regularizing the result to avoid issues when `a` is small. -This is done by performing a linear interpolation from `ρaχ / ρa` to `ρχ / ρ`, -using `sgs_weight_function(ρa / ρ, turbconv_model.a_half)` as the weight of -`ρaχ / ρa` in the interpolation. Note that `ρa / ρ` is the "anelastic -approximation" of `a`; we cannot directly use `a` to compute the weight because -this function needs to be called before `a` has been computed. Also, note that -using this function instead of directly computing `ρaχ / ρa` breaks the -assumption of domain decomposition when the approximated `a` is small. -""" -function divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model) - weight = sgs_weight_function(ρa / ρ, turbconv_model.a_half) - # If ρa = 0, we know that ρa / ρ = 0, which means that weight = 0. However, - # 0 * ρaχ / 0 = NaN, regardless of what ρaχ is, so the linear interpolation - # will always return NaN when ρa = 0. To avoid this problem, we need to add - # a special case for ρa = 0. - return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ -end +Computes the sum of a function `f` applied to each draft subdomain +state `sgsʲ` in the iterator `sgsʲs`. +Arguments: +- `f`: A function to apply to each element of `sgsʲs`. +- `sgsʲs`: An iterator over the draft subdomain states. """ - ρa⁺(gs) +draft_sum(f, sgsʲs) = mapreduce_with_init(f, +, sgsʲs) -Computes the total mass-flux subdomain area-weighted density, assuming that the -mass-flux subdomain states are stored in `gs.sgsʲs`. """ -ρa⁺(gs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa, +, gs.sgsʲs) + env_value(grid_scale_value, f_draft, gs) -""" - ρah_tot⁺(sgsʲs) +Computes the value of a quantity `ρaχ` in the environment subdomain by subtracting +the sum of its values in all draft subdomains from the grid-scale value. -Computes the total mass-flux subdomain area-weighted ρh_tot, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρah_tot⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.h_tot, +, sgsʲs) +This is based on the domain decomposition principle for density-area weighted +quantities: `GridMean(ρχ) = Env(ρaχ) + Sum(Drafts(ρaχ))`. +Arguments: +- `grid_scale_value`: The `ρa`-weighted grid-scale value of the quantity. +- `f_draft`: A function that extracts the corresponding value from a draft subdomain state. +- `gs`: The grid-scale state, which contains the draft subdomain states `gs.sgsʲs`. """ - ρamse⁺(sgsʲs) +function env_value(grid_scale_value, f_draft, gs) + return grid_scale_value - draft_sum(f_draft, gs.sgsʲs) +end -Computes the total mass-flux subdomain area-weighted ρmse, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. """ -ρamse⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.mse, +, sgsʲs) + specific_env_value(χ_name::Symbol, gs, turbconv_model) -""" - ρaq_tot⁺(sgsʲs) +Calculates the specific value of a quantity `χ` in the environment (`χ⁰`). -Computes the total mass-flux subdomain area-weighted ρq_tot, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_tot⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_tot, +, sgsʲs) +This function uses the domain decomposition principle to first find the +density-area-weighted environment value (`ρa⁰χ⁰`) and the environment +density-area (`ρa⁰`). It then computes the specific value using the +regularized `specific` function, which provides a stable result even when the +environment area fraction is very small. -""" - ρaq_liq⁺(sgsʲs) +Arguments: +- `χ_name`: The `Symbol` for the specific quantity `χ` (e.g., `:h_tot`, `:q_tot`). +- `gs`: The grid-scale state, containing grid-mean and draft subdomain states. +- `turbconv_model`: The turbulence convection model, containing parameters for regularization. + +Returns: +- The specific value of the quantity `χ` in the environment. +""" +function specific_env_value(χ_name::Symbol, gs, turbconv_model) + # Grid-scale density-weighted variable name, e.g., :ρq_tot + ρχ_name = Symbol(:ρ, χ_name) + + # Numerator: ρa⁰χ⁰ = (gs.ρχ) - (Σ sgsʲ.ρa * sgsʲ.χ) + ρaχ⁰ = env_value( + getproperty(gs, ρχ_name), + sgsʲ -> getproperty(sgsʲ, :ρa) * getproperty(sgsʲ, χ_name), + gs, + ) + + # Denominator: ρa⁰ = gs.ρ - Σ sgsʲ.ρa + ρa⁰_val = env_value(gs.ρ, sgsʲ -> sgsʲ.ρa, gs) + + # Call the 5-argument specific function for regularized division + return specific( + ρaχ⁰, # ρaχ for environment + ρa⁰_val, # ρa for environment + getproperty(gs, ρχ_name), # Fallback ρχ is the grid-mean value + gs.ρ, # Fallback ρ is the grid-mean value + turbconv_model, + ) +end -Computes the liquid water mass-flux subdomain area-weighted ρq_liq, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. """ -ρaq_liq⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_liq, +, sgsʲs) + specific_env_mse(gs, p) -""" - ρaq_ice⁺(sgsʲs) +Computes the specific moist static energy (`mse`) in the environment (`mse⁰`). -Computes the ice water mass-flux subdomain area-weighted ρq_ice, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_ice⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_ice, +, sgsʲs) +This is a specialized helper function because `mse` is not a grid-scale prognostic +variable. It first computes the grid-scale moist static energy density (`ρmse`) +from other grid-scale quantities (`ρ`, total specific enthalpy `h_tot`, specific +kinetic energy `K`). It then uses the `env_value` helper to compute the environment's +portion of `ρmse` and `ρa` via domain decomposition, and finally calculates the specific +value using the regularized `specific` function. -""" - ρaq_rai⁺(sgsʲs) +Arguments: +- `gs`: The grid-scale state (`Y.c`), containing `ρ` and `sgsʲs`. +- `p`: The cache, containing the `turbconv_model` + +Returns: +- A `ClimaCore.Fields.Field` containing the specific moist static energy of the + environment (`mse⁰`). +""" +function specific_env_mse(gs, p) + # Get necessary precomputed values from the cache `p` + (; ᶜK, ᶜts) = p.precomputed # TODO: replace by on-the-fly computation + (; turbconv_model) = p.atmos + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(gs.ρe_tot, gs.ρ))) + + # 1. Define the grid-scale moist static energy density `ρ * mse`. + grid_scale_ρmse = gs.ρ .* (ᶜh_tot .- ᶜK) + + # 2. Compute the environment's density-area-weighted mse (`ρa⁰mse⁰`). + ρa⁰mse⁰ = env_value(grid_scale_ρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, gs) + + # 3. Compute the environment's density-area product (`ρa⁰`). + ρa⁰_val = ρa⁰(gs) + + # 4. Compute and return the final specific environment mse (`mse⁰`). + return specific( + ρa⁰mse⁰, + ρa⁰_val, + grid_scale_ρmse, + gs.ρ, + turbconv_model, + ) +end -Computes the rain mass-flux subdomain area-weighted ρq_rai, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. """ -ρaq_rai⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_rai, +, sgsʲs) + ρa⁰(gs) -""" - ρaq_sno⁺(sgsʲs) +Computes the environment area-weighted density (`ρa⁰`). -Computes the snow mass-flux subdomain area-weighted ρq_sno, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_sno⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_sno, +, sgsʲs) +This function uses the `env_value` helper, which applies the domain +decomposition principle (`GridMean = Environment + Sum(Drafts)`) to calculate +the environment area-weighted density by subtracting the sum of all draft +subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`). -""" - ρa⁰(gs) +Arguments: +- `gs`: The grid-scale state, which contains the grid-mean density `gs.ρ` and + the draft subdomain states `gs.sgsʲs`. -Computes the environment area-weighted density, assuming that the mass-flux -subdomain states are stored in `gs.sgsʲs`. +Returns: +- The area-weighted density (`ρa⁰`). """ -ρa⁰(gs) = gs.ρ - mapreduce_with_init(sgsʲ -> sgsʲ.ρa, +, gs.sgsʲs) +ρa⁰(gs) = env_value(gs.ρ, sgsʲ -> sgsʲ.ρa, gs) """ - u₃⁺(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) + specific_tke(sgs⁰, gs, turbconv_model) -Computes the average mass-flux subdomain vertical velocity `u₃⁺` by dividing the -total momentum `ρaw⁺` by the total area-weighted density `ρa⁺`, both of which -are computed from the tuples of subdomain densities and velocities `ρaʲs` and -`u₃ʲs`. The division is computed using `divide_by_ρa` to avoid issues when `a⁺` -is small. -""" -u₃⁺(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( - unrolled_dotproduct(ρaʲs, u₃ʲs), - reduce(+, ρaʲs), - ρ * u₃, - ρ, - turbconv_model, -) +Computes the specific turbulent kinetic energy (`tke`) in the environment (`tke⁰`). + +This is a specialized helper that encapsulates the call to the regularized +`specific` function for the TKE variable. It provides `0` as the grid-scale +fallback value (`ρχ_fallback`) in the limit of small environmental area +fraction. + +Arguments: +- `sgs⁰`: The environment SGS state (`Y.c.sgs⁰`), containing `ρatke`. +- `gs`: The grid-scale state (`Y.c`), containing the grid-mean density `ρ`. +- `turbconv_model`: The turbulence convection model, for regularization parameters. + +Returns: +- The specific TKE of the environment (`tke⁰`). +""" +function specific_tke(sgs⁰, gs, turbconv_model) + ρa⁰_val = ρa⁰(gs) + + return specific( + sgs⁰.ρatke, # ρaχ for environment TKE + ρa⁰_val, # ρa for environment, now computed internally + 0, # Fallback ρχ is zero for TKE + gs.ρ, # Fallback ρ + turbconv_model, + ) +end """ u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) -Computes the environment vertical velocity `u₃⁰` by dividing the environment -momentum `ρaw⁰` by the environment area-weighted density `ρa⁰`, both of which -are computed from the domain decomposition of the grid-scale quantities `ρw` and -`ρ` into the mass-flux subdomain quantities `ρawʲs` and `ρaʲs` and the -environment quantities. The division is computed using `divide_by_ρa` to avoid -issues when `a⁰` is small. +Computes the environment vertical velocity `u₃⁰`. + +This function calculates the environment's total vertical momentum (`ρa⁰u₃⁰`) and +its total area-weighted density (`ρa⁰`) using the domain decomposition principle +(GridMean = Env + Sum(Drafts)). It then computes the final specific velocity `u₃⁰` +using the regularized `specific` function to ensure numerical stability when the +environment area fraction `a⁰` is small. + +Arguments: +- `ρaʲs`: A tuple of area-weighted densities for each draft subdomain. +- `u₃ʲs`: A tuple of vertical velocities for each draft subdomain. +- `ρ`: The grid-mean air density. +- `u₃`: The grid-mean vertical velocity. +- `turbconv_model`: The turbulence convection model, containing regularization parameters. """ -u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( +u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = specific( ρ * u₃ - unrolled_dotproduct(ρaʲs, u₃ʲs), ρ - reduce(+, ρaʲs), ρ * u₃, @@ -377,14 +476,50 @@ u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( turbconv_model, ) -import ClimaCore.RecursiveApply: ⊞, ⊠, rzero, rpromote_type +""" + mapreduce_with_init(f, op, iter...) + +A wrapper for Julia's `mapreduce` function that automatically determines +the initial value (`init`) for the reduction. + +This is useful for iterators whose elements are custom structs or +`ClimaCore.Geometry.AxisTensor`s, where the zero element cannot be inferred +as a simple scalar. It uses `ClimaCore.RecursiveApply` tools (`rzero`, +`rpromote_type`) to create a type-stable, correctly-structured zero element +based on the output of the function `f` applied to the first elements of the +iterators. + +Arguments: +- `f`: The function to apply to each element. +- `op`: The reduction operator (e.g., `+`, `*`). +- `iter...`: One or more iterators. +""" function mapreduce_with_init(f, op, iter...) r₀ = rzero(rpromote_type(typeof(f(map(first, iter)...)))) mapreduce(f, op, iter...; init = r₀) end -# Inference fails for certain mapreduce calls inside cuda -# kernels, so let's define a recursive unrolled dot product: +""" + unrolled_dotproduct(a::Tuple, b::Tuple) + +Computes the dot product of two `Tuple`s (`a` and `b`) using a recursive, +manually unrolled implementation. + +This function is designed to be type-stable and efficient for CUDA kernels, +where standard `mapreduce` implementations can otherwise suffer from type-inference +failures. + +It uses `ClimaCore.RecursiveApply` operators (`⊞` for addition, `⊠` for +multiplication), which allows it to handle dot products of tuples containing +complex, nested types such as `ClimaCore.Geometry.AxisTensor`s. + +Arguments: +- `a`: The first `Tuple`. +- `b`: The second `Tuple`, which must have the same length as `a`. + +Returns: +- The result of the dot product, `Σᵢ a[i] * b[i]`. +""" promote_type_mul(n::Number, x::Geometry.AxisTensor) = typeof(x) promote_type_mul(x::Geometry.AxisTensor, n::Number) = typeof(x) @inline function unrolled_dotproduct(a::Tuple, b::Tuple) From c28133090e0b255bf24930ac955560ec2d68f5ef Mon Sep 17 00:00:00 2001 From: costachris Date: Wed, 25 Jun 2025 14:55:49 -0700 Subject: [PATCH 2/7] debug --- src/cache/cloud_fraction.jl | 20 +---- .../diagnostic_edmf_precomputed_quantities.jl | 12 ++- src/cache/precomputed_quantities.jl | 24 +++--- .../prognostic_edmf_precomputed_quantities.jl | 41 +++++---- src/diagnostics/Diagnostics.jl | 2 + src/diagnostics/edmfx_diagnostics.jl | 10 +-- .../les_sgs_models/smagorinsky_lilly.jl | 22 +++-- .../radiation/radiation.jl | 2 +- src/prognostic_equations/advection.jl | 26 ++++-- src/prognostic_equations/edmfx_entr_detr.jl | 3 +- src/prognostic_equations/edmfx_sgs_flux.jl | 83 ++++++++++--------- src/prognostic_equations/edmfx_tke.jl | 10 +-- .../forcing/subsidence.jl | 30 +++---- src/prognostic_equations/hyperdiffusion.jl | 9 +- .../implicit/implicit_tendency.jl | 24 +++--- .../implicit/manual_sparse_jacobian.jl | 71 +++++++++------- .../remaining_tendency.jl | 8 +- src/prognostic_equations/surface_flux.jl | 12 ++- .../vertical_diffusion_boundary_layer.jl | 8 +- src/utils/variable_manipulations.jl | 26 +++--- 20 files changed, 256 insertions(+), 187 deletions(-) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index 443c33205d..5804a90442 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -64,15 +64,8 @@ NVTX.@annotate function set_cloud_fraction!( else q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) - @. cloud_diagnostics_tuple = make_named_tuple( - ifelse( - q_liq + q_ice > 0, - 1, - 0, - ), - q_liq, - q_ice, - ) + @. cloud_diagnostics_tuple = + make_named_tuple(ifelse(q_liq + q_ice > 0, 1, 0), q_liq, q_ice) end end @@ -91,15 +84,6 @@ NVTX.@annotate function set_cloud_fraction!( (; turbconv_model) = p.atmos if isnothing(turbconv_model) - (; - ᶜlinear_buoygrad, - ᶜstrain_rate_norm, - ᶠu³⁰, - ᶠu³, - ᶜentrʲs, - ᶜdetrʲs, - ᶠu³ʲs, - ) = p.precomputed ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) if p.atmos.call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 987d3f9b9b..9f2786cf4c 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -100,13 +100,19 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( (; ᶠu³⁰, ᶜK⁰) = p.precomputed (; params) = p - + thermo_params = CAP.thermodynamics_params(params) turbconv_params = CAP.turbconv_params(params) ᶜts = p.precomputed.ᶜts #TODO replace q_tot = specific(Y.c.ρq_tot, Y.c.ρ) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1)) uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1)) @@ -1128,4 +1134,4 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita # thermo_params, #) return nothing -end \ No newline at end of file +end diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 704714cf85..b560e7d7f3 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -3,6 +3,7 @@ ##### import Thermodynamics as TD import ClimaCore: Spaces, Fields +using Base.Broadcast: materialize """ implicit_precomputed_quantities(Y, atmos) @@ -40,10 +41,9 @@ 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) + ᶜp = similar(Y.c, FT), ) - sgs_quantities = - turbconv_model isa AbstractEDMF ? (;) : (;) + sgs_quantities = turbconv_model isa AbstractEDMF ? (;) : (;) prognostic_sgs_quantities = turbconv_model isa PrognosticEDMFX ? (; @@ -373,21 +373,20 @@ function thermo_state( end function thermo_vars(moisture_model, microphysics_model, Y_c, K, Φ) - # Compute specific quantities on-the-fly - e_tot = @. lazy(specific(Y_c.ρe_tot, Y_c.ρ)) + e_tot = materialize(@. lazy(specific(Y_c.ρe_tot, Y_c.ρ))) energy_var = (; e_int = e_tot - K - Φ) - + moisture_var = if moisture_model isa DryModel (;) elseif moisture_model isa EquilMoistModel - q_tot = @. lazy(specific(Y_c.ρq_tot, Y_c.ρ)) + q_tot = materialize(@. lazy(specific(Y_c.ρq_tot, Y_c.ρ))) (; q_tot) elseif moisture_model isa NonEquilMoistModel - q_tot = @. lazy(specific(Y_c.ρq_tot, Y_c.ρ)) - q_liq = @. lazy(specific(Y_c.ρq_liq, Y_c.ρ)) - q_ice = @. lazy(specific(Y_c.ρq_ice, Y_c.ρ)) - q_rai = @. lazy(specific(Y_c.ρq_rai, Y_c.ρ)) - q_sno = @. lazy(specific(Y_c.ρq_sno, Y_c.ρ)) + q_tot = materialize(@. lazy(specific(Y_c.ρq_tot, Y_c.ρ))) + q_liq = materialize(@. lazy(specific(Y_c.ρq_liq, Y_c.ρ))) + q_ice = materialize(@. lazy(specific(Y_c.ρq_ice, Y_c.ρ))) + q_rai = materialize(@. lazy(specific(Y_c.ρq_rai, Y_c.ρ))) + q_sno = materialize(@. lazy(specific(Y_c.ρq_sno, Y_c.ρ))) q_pt_args = (q_tot, q_liq + q_rai, q_ice + q_sno) (; q_pt = TD.PhasePartition(q_pt_args...)) end @@ -442,7 +441,6 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) thermo_params = CAP.thermodynamics_params(p.params) thermo_args = (thermo_params, moisture_model, microphysics_model) - ᶜspecific .= ᶜspecific_gs_tracers(Y) @. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ) # TODO: We might want to move this to dss! (and rename dss! to something diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 847a6350fc..aa5a6b17f0 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -21,15 +21,19 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( thermo_params = CAP.thermodynamics_params(p.params) (; turbconv_model) = p.atmos (; ᶜΦ,) = p.core - (; ᶜp,ᶜK) = p.precomputed + (; ᶜp, ᶜK) = p.precomputed (; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) - @. ᶜtke⁰ = specific_tke(Y.c.sgs⁰, Y.c, turbconv_model) + ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + + ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 + ᶜmse⁰ .= specific_env_mse(Y.c, p) + if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) @@ -39,16 +43,12 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( @. ᶜts⁰ = TD.PhaseNonEquil_phq( thermo_params, ᶜp, - specific_env_mse(Y.c, p) - ᶜΦ, + ᶜmse⁰ - ᶜΦ, TD.PhasePartition(ᶜq_tot⁰, ᶜq_liq⁰ + ᶜq_rai⁰, ᶜq_ice⁰ + ᶜq_sno⁰), ) else - @. ᶜts⁰ = TD.PhaseEquil_phq( - thermo_params, - ᶜp, - specific_env_mse(Y.c, p) - ᶜΦ, - ᶜq_tot⁰, - ) + + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) end return nothing end @@ -132,7 +132,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( turbconv_params = CAP.turbconv_params(p.params) (; ᶜΦ,) = p.core - (; ᶜp, ᶜK, ᶜtsʲs, ᶜρʲs) = p.precomputed + (; ᶜp, ᶜK, ᶜtsʲs, ᶜρʲs, ᶜts) = p.precomputed (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions for j in 1:n @@ -176,8 +176,15 @@ 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 = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) - ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1)) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) + ᶜh_tot_int_val = + Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), 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( @@ -193,8 +200,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( ) # ... and the first interior point for EDMFX ᶜq_totʲ. + + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) ᶜq_tot_int_val = - Fields.field_values(Fields.level(specific(Y.c.ρq_tot, Y.c.ρ), 1)) + Fields.field_values(Fields.level(Base.materialize(ᶜq_tot), 1)) ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1)) @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc( ᶜz_int_val - z_sfc_val, @@ -359,6 +368,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ᶜdz = Fields.Δz_field(axes(Y.c)) ᶜlg = Fields.local_geometry_field(Y.c) ᶠlg = Fields.local_geometry_field(Y.f) + ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜρa⁰ = @. lazy(ρa⁰(Y.c)) ᶜvert_div = p.scratch.ᶜtemp_scalar ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2 @@ -480,7 +491,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰) end - sfc_tke = Fields.level(ᶜtke⁰, 1) + sfc_tke = Fields.level(Base.materialize(ᶜtke⁰), 1) @. ᶜmixing_length_tuple = mixing_length( p.params, ustar, @@ -503,7 +514,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec) ρatke_flux_values = Fields.field_values(ρatke_flux) - ρa_sfc_values = Fields.field_values(Fields.level(ᶜρa⁰, 1)) # TODO: replace by surface value + ρa_sfc_values = Fields.field_values(Fields.level(Base.materialize(ᶜρa⁰), 1)) # TODO: replace by surface value ustar_values = Fields.field_values(ustar) sfc_local_geometry_values = Fields.field_values( Fields.level(Fields.local_geometry_field(Y.f), half), diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl index 059e868179..8e1589efe3 100644 --- a/src/diagnostics/Diagnostics.jl +++ b/src/diagnostics/Diagnostics.jl @@ -58,6 +58,8 @@ import ..SlabOceanSST import ..draft_area import ..compute_gm_mixing_length! import ..horizontal_integral_at_boundary +import ..ρa⁰ +import ..specific_tke # We need the abbreviations for symbols like curl, grad, and so on include(joinpath("..", "utils", "abbreviations.jl")) diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index d82653847a..ca369b0182 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -633,7 +633,7 @@ compute_aren!(_, _, _, _, turbconv_model::T) where {T} = function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - ᶜρa⁰ = @.lazy(ρa⁰(state.c)) + ᶜρa⁰ = @. lazy(ρa⁰(state.c)) if isnothing(out) return draft_area.( ᶜρa⁰, @@ -1139,18 +1139,18 @@ compute_tke!(out, state, cache, time) = compute_tke!(_, _, _, _, turbconv_model::T) where {T} = error_diagnostic_variable("tke", turbconv_model) -function compute_tke!( +function compute_tke!( out, state, cache, time, turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX}, ) - + ᶜtke = @. lazy(specific_tke(state.c.sgs⁰, state.c, turbconv_model)) if isnothing(out) - return specific_tke(state.c.sgs⁰, state.c, turbconv_model) + return Base.materialize(ᶜtke) else - out .= specific_tke(state.c.sgs⁰, state.c, turbconv_model) + out .= ᶜtke 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 2c899fa93f..2b57bd2bce 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -96,7 +96,7 @@ 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) = p.precomputed + (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜts) = p.precomputed ## Momentum tendencies ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(Y.c.ρ) @@ -104,7 +104,13 @@ function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLill @. Yₜ.f.u₃ -= C3(wdivₕ(ᶠρ * ᶠτ_smag) / ᶠρ) ## Total energy tendency - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot += wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜh_tot)) ## Tracer diffusion and associated mass changes @@ -125,9 +131,9 @@ 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, sfc_conditions) = - p.precomputed + (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, sfc_conditions) = p.precomputed (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions + (; ᶜts) = p.precomputed # Define operators ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ @@ -154,7 +160,13 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) @. Yₜ.f.u₃ -= C3(ᶠdivᵥ(Y.c.ρ * ᶜτ_smag) / ᶠρ) ## Total energy tendency - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜh_tot))) ## Tracer diffusion and associated mass changes diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl index 9e97e933b7..992054c4c7 100644 --- a/src/parameterized_tendencies/radiation/radiation.jl +++ b/src/parameterized_tendencies/radiation/radiation.jl @@ -566,4 +566,4 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationISDAC) )) # = -∂F/∂z = ρ cₚ ∂T/∂t (longwave radiation) return nothing -end \ No newline at end of file +end diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index a8ad9e5f72..7b7575fd2b 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -36,7 +36,9 @@ 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 + (; ᶜu, ᶜK, ᶜp, ᶜts) = p.precomputed + (; params) = p + thermo_params = CAP.thermodynamics_params(params) if p.atmos.turbconv_model isa PrognosticEDMFX (; ᶜuʲs) = p.precomputed @@ -49,7 +51,13 @@ NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) end end - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot -= wdivₕ(Y.c.ρ * ᶜh_tot * ᶜu) if p.atmos.turbconv_model isa PrognosticEDMFX @@ -174,7 +182,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 + (; ᶜu, ᶠu³, ᶜK, ᶜts) = 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 @@ -196,8 +204,8 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) else nothing end - ᶜtke⁰ = advect_tke ? - (@.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))) : + ᶜtke⁰ = + advect_tke ? (@.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))) : nothing ᶜa_scalar = p.scratch.ᶜtemp_scalar ᶜω³ = p.scratch.ᶜtemp_CT3 @@ -233,7 +241,13 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) # ... and upwinding correction of energy and total water. # (The central advection of energy and total water is done implicitly.) if energy_upwinding != Val(:none) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) @. Yₜ.c.ρe_tot += vtt - vtt_central diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index 47e8a895b3..8cfcc14d4a 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -530,7 +530,8 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF (; ᶜturb_entrʲs, ᶜentrʲs, ᶜdetrʲs) = p.precomputed (; ᶠu₃⁰) = p.precomputed - ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) + ᶜmse⁰ = p.scratch.ᶜtemp_scalar + ᶜmse⁰ .= specific_env_mse(Y.c, p) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 5e42161b2c..c93d7ce469 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -42,7 +42,7 @@ function edmfx_sgs_mass_flux_tendency!( (; edmfx_sgsflux_upwinding) = p.atmos.numerics (; ᶠu³) = p.precomputed (; ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = p.precomputed - (; ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed + (; ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) ᶜρ⁰ = @. TD.air_density(thermo_params, ᶜts⁰) ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) @@ -56,7 +56,13 @@ function edmfx_sgs_mass_flux_tendency!( # [best after removal of precomputed quantities] ᶠu³_diff = p.scratch.ᶠtemp_CT3 ᶜa_scalar = p.scratch.ᶜtemp_scalar - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @. ᶜa_scalar = @@ -73,7 +79,9 @@ function edmfx_sgs_mass_flux_tendency!( end # Add the environment fluxes @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ - ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) + + ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 + ᶜmse⁰ .= specific_env_mse(Y.c, p) @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, @@ -92,7 +100,7 @@ function edmfx_sgs_mass_flux_tendency!( (Y.c.sgsʲs.:($$j).q_tot - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( - ᶜρʲs.:($$j), + ᶜρʲs.:($j), ᶠu³_diff, ᶜa_scalar, dt, @@ -104,8 +112,7 @@ function edmfx_sgs_mass_flux_tendency!( ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = - (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * - draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -129,10 +136,8 @@ function edmfx_sgs_mass_flux_tendency!( @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @. ᶜa_scalar = - ( - Y.c.sgsʲs.:($$j).q_liq - - specific(Y.c.ρq_liq, Y.c.ρ) - ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + (Y.c.sgsʲs.:($$j).q_liq - specific(Y.c.ρq_liq, Y.c.ρ)) * + draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -143,10 +148,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_liq += vtt @. ᶜa_scalar = - ( - Y.c.sgsʲs.:($$j).q_ice - - specific(Y.c.ρq_ice, Y.c.ρ) - ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + (Y.c.sgsʲs.:($$j).q_ice - specific(Y.c.ρq_ice, Y.c.ρ)) * + draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -157,10 +160,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_ice += vtt @. ᶜa_scalar = - ( - Y.c.sgsʲs.:($$j).q_rai - - specific(Y.c.ρq_rai, Y.c.ρ) - ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + (Y.c.sgsʲs.:($$j).q_rai - specific(Y.c.ρq_rai, Y.c.ρ)) * + draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -171,10 +172,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_rai += vtt @. ᶜa_scalar = - ( - Y.c.sgsʲs.:($$j).q_sno - - specific(Y.c.ρq_sno, Y.c.ρ) - ) * draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) + (Y.c.sgsʲs.:($$j).q_sno - specific(Y.c.ρq_sno, Y.c.ρ)) * + draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)) vtt = vertical_transport( ᶜρʲs.:($j), ᶠu³_diff, @@ -187,8 +186,7 @@ function edmfx_sgs_mass_flux_tendency!( @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = - (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * - draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -199,8 +197,7 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_liq += vtt @. ᶜa_scalar = - (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * - draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -211,8 +208,7 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_ice += vtt @. ᶜa_scalar = - (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * - draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -223,8 +219,7 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_rai += vtt @. ᶜa_scalar = - (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * - draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -254,14 +249,20 @@ function edmfx_sgs_mass_flux_tendency!( n = n_mass_flux_subdomains(turbconv_model) (; edmfx_sgsflux_upwinding) = p.atmos.numerics (; ᶠu³) = p.precomputed - (; ᶜρaʲs, ᶜρʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs) = p.precomputed + (; ᶜρaʲs, ᶜρʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs, ᶜts) = p.precomputed (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J FT = eltype(Y) if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} # energy - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ᶠu³_diff = p.scratch.ᶠtemp_CT3 ᶜa_scalar = p.scratch.ᶜtemp_scalar for j in 1:n @@ -395,8 +396,8 @@ function edmfx_sgs_diffusive_flux_tendency!( (; dt, params) = p turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm,) = p.precomputed - (; ρatke_flux) = p.precomputed + (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed + (; ᶜmixing_length, ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) @@ -412,7 +413,9 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) + # ᶜmse⁰ = @. lazy(specific_env_mse(Y.c, p)) + ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 + ᶜmse⁰ .= specific_env_mse(Y.c, p) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜmse⁰ + ᶜK⁰))) if use_prognostic_tke(turbconv_model) # Turbulent TKE transport (diffusion) @@ -503,7 +506,7 @@ function edmfx_sgs_diffusive_flux_tendency!( turbconv_params = CAP.turbconv_params(params) thermo_params = CAP.thermodynamics_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜu, ᶜmixing_length) = p.precomputed + (; ᶜu, ᶜmixing_length, ᶜts) = p.precomputed (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) @@ -519,7 +522,13 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜtsʲ, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜh_tot))) if use_prognostic_tke(turbconv_model) diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index 18ddf17512..a8d6d6918d 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -44,16 +44,12 @@ 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 - ) = p.precomputed + (; ᶠu³⁰, ᶠu³, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜK_u, ᶜK_h) = + p.precomputed turbconv_params = CAP.turbconv_params(p.params) FT = eltype(p.params) - + ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ nh_pressure3_buoyʲs = turbconv_model isa PrognosticEDMFX ? diff --git a/src/prognostic_equations/forcing/subsidence.jl b/src/prognostic_equations/forcing/subsidence.jl index 8b996d8b91..8140a55f01 100644 --- a/src/prognostic_equations/forcing/subsidence.jl +++ b/src/prognostic_equations/forcing/subsidence.jl @@ -84,15 +84,13 @@ Arguments: - `Y`: The current state vector, used for density (`ρ`). - `p`: Cache containing parameters, and the subsidence model object. - `t`: Current simulation time. -- `subsidence`: The subsidence model object, containing the prescribed vertical - velocity profile `Dᵥ`. +- `subsidence`: The subsidence model object. """ -function subsidence_tendency!(Yₜ, Y, p, t, subsidence::Subsidence) - (; Dᵥ) = subsidence - ᶜρ = Y.c.ρ +function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) (; moisture_model) = p.atmos subsidence_profile = p.atmos.subsidence.prof thermo_params = CAP.thermodynamics_params(p.params) + (; ᶜts) = p.precomputed ᶠz = Fields.coordinate_field(axes(Y.f)).z ᶠlg = Fields.local_geometry_field(Y.f) @@ -100,23 +98,25 @@ function subsidence_tendency!(Yₜ, Y, p, t, subsidence::Subsidence) @. ᶠsubsidence³ = subsidence_profile(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg)) + ᶜρ = Y.c.ρ # LS Subsidence - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) - subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}()) - subsidence!( - Yₜ.c.ρq_tot, - Y.c.ρ, - ᶠsubsidence³, - specific(Y.c.ρq_tot, Y.c.ρ), - Val{:first_order}(), + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), ) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}()) + subsidence!(Yₜ.c.ρq_tot, Y.c.ρ, ᶠsubsidence³, ᶜq_tot, Val{:first_order}()) if moisture_model isa NonEquilMoistModel ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) subsidence!( Yₜ.c.ρq_liq, Y.c.ρ, ᶠsubsidence³, - specific(Y.c.ρq_liq, Y.c.ρ), + ᶜq_liq, Val{:first_order}(), ) ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) @@ -124,7 +124,7 @@ function subsidence_tendency!(Yₜ, Y, p, t, subsidence::Subsidence) Yₜ.c.ρq_ice, Y.c.ρ, ᶠsubsidence³, - specific(Y.c.ρq_ice, Y.c.ρ), + ᶜq_ice, Val{:first_order}(), ) end diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index 7d95c58729..8392424998 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -117,8 +117,7 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) if diffuse_tke - ᶜtke⁰ = - @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) (; ᶜ∇²tke⁰) = p.hyperdiff @. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰)) end @@ -155,8 +154,7 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end if use_prognostic_tke(turbconv_model) - ᶜtke⁰ = - @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) (; ᶜ∇²tke⁰) = p.hyperdiff end @@ -252,7 +250,6 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos isnothing(hyperdiff) && return nothing - ᶜspecific = all_specific_gs(Y.c) (; ᶜ∇²specific_tracers) = p.hyperdiff # TODO: Fix RecursiveApply bug in gradₕ to fuse this operation. @@ -294,7 +291,6 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) ν₄_scalar_for_precip = CAP.α_hyperdiff_tracer(p.params) * ν₄_scalar n = n_mass_flux_subdomains(turbconv_model) - ᶜspecific = all_specific_gs(Y.c) (; ᶜ∇²specific_tracers) = p.hyperdiff # TODO: Since we are not applying the limiter to density (or area-weighted @@ -338,4 +334,3 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) end return nothing end - diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 5e5f4d1f31..4ecfafe59e 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -140,9 +140,15 @@ 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 - (; ᶠu³, ᶜp) = p.precomputed + (; ᶠu³, ᶜp, ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠu³) @@ -150,10 +156,11 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜh_tot, dt, Val(:none)) @. Yₜ.c.ρe_tot += vtt if !(moisture_model isa DryModel) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) vtt = vertical_transport( Y.c.ρ, ᶠu³, - specific(Y.c.ρq_tot, Y.c.ρ), + ᶜq_tot, dt, Val(:none), ) @@ -167,14 +174,12 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) if moisture_model isa NonEquilMoistModel (; ᶜwₗ, ᶜwᵢ) = p.precomputed @. Yₜ.c.ρq_liq -= ᶜprecipdivᵥ( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * - ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( Geometry.WVector(-(ᶜwₗ)) * specific(Y.c.ρq_liq, Y.c.ρ), ), ) @. Yₜ.c.ρq_ice -= ᶜprecipdivᵥ( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * - ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( Geometry.WVector(-(ᶜwᵢ)) * specific(Y.c.ρq_ice, Y.c.ρ), ), ) @@ -182,8 +187,7 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) if microphysics_model isa Microphysics1Moment (; ᶜwᵣ, ᶜwₛ) = p.precomputed @. Yₜ.c.ρq_rai -= ᶜprecipdivᵥ( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * - ᶠright_bias( + ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( Geometry.WVector(-(ᶜwᵣ)) * specific(Y.c.ρq_rai, Y.c.ρ), ), ) @@ -212,8 +216,6 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) ) @. Yₜ.c.ρq_sno -= ᶜprecipdivᵥ( ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias( - ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * - ᶠright_bias( Geometry.WVector(-(ᶜwₛ)) * specific(Y.c.ρq_sno, Y.c.ρ), ), ) diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index d8cd65aa03..6168f13bfd 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -367,7 +367,14 @@ 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 thermo_params = CAP.thermodynamics_params(params) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - FT(CAP.e_int_v0(params)) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ᶜρ = Y.c.ρ ᶜuₕ = Y.c.uₕ @@ -416,8 +423,22 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) - foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name - ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) + tracer_info = (@name(c.ρe_tot), @name(c.ρq_tot)) + + MatrixFields.unrolled_foreach(tracer_info) do ρχ_name + MatrixFields.has_field(Y, ρχ_name) || return + ᶜχ = if ρχ_name === @name(c.ρe_tot) + @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) + else + @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + end + if use_derivative(topography_flag) ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] @. ∂ᶜρχ_err_∂ᶜuₕ = @@ -447,10 +468,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶠu₃_err_∂ᶜρq_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + e_tot - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * q_tot - ) + ᶜ∂kappa_m∂q_tot * + (cp_d * T_0 + e_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * q_tot) )) end @@ -562,10 +581,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] @. ∂ᶜρe_tot_err_∂ᶜρ = dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( - ( - -(1 + ᶜkappa_m) * e_tot - - ᶜkappa_m * ∂e_int_∂q_tot * q_tot - ) / ᶜρ, + (-(1 + ᶜkappa_m) * e_tot - ᶜkappa_m * ∂e_int_∂q_tot * q_tot) / + ᶜρ, ) @. ∂ᶜρe_tot_err_∂ᶜρe_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) @@ -576,27 +593,25 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρe_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + e_tot - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * q_tot - ) + ᶜ∂kappa_m∂q_tot * + (cp_d * T_0 + e_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * q_tot) )) @. ∂ᶜρq_tot_err_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ - DiagonalMatrixRow(-(q_tot) / ᶜρ) + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(q_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) end - foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + foreach_gs_tracer(Y) do ᶜρχ, ρχ_name + MatrixFields.has_field(Y, ρχ_name) || return ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) ∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)] ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] - ᶜtridiagonal_matrix_scalar = if ρχ_name in (@name(c.ρq_rai), @name(c.ρq_sno), @name(c.ρn_rai)) - ᶜdiffusion_h_matrix_scaled - else - ᶜdiffusion_h_matrix - end + ᶜtridiagonal_matrix_scalar = ifelse( + ρχ_name in (@name(c.ρq_rai), @name(c.ρq_sno), @name(c.ρn_rai)), + ᶜdiffusion_h_matrix_scaled, + ᶜdiffusion_h_matrix, + ) @. ∂ᶜρχ_err_∂ᶜρ = dtγ * ᶜtridiagonal_matrix_scalar ⋅ DiagonalMatrixRow(-(ᶜχ) / ᶜρ) @. ∂ᶜρχ_err_∂ᶜρχ += @@ -610,7 +625,9 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) turbconv_model = p.atmos.turbconv_model ᶜmixing_length = p.precomputed.ᶜmixing_length ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) - ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ + ᶜρa⁰ = + p.atmos.turbconv_model isa PrognosticEDMFX ? + (@.lazy(ρa⁰(Y.c))) : Y.c.ρ ᶜρatke⁰ = Y.c.sgs⁰.ρatke @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) = @@ -949,10 +966,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + e_tot - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * q_tot - ) + ᶜ∂kappa_m∂q_tot * + (cp_d * T_0 + e_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * q_tot) )) @. ∂ᶜρe_tot_err_∂ᶜρe_tot += diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index f557ae9bbc..3fe38dcdd5 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -146,7 +146,13 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) thermo_params = CAP.thermodynamics_params(params) (; ᶜp, sfc_conditions, ᶜts) = p.precomputed - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge) vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge) diff --git a/src/prognostic_equations/surface_flux.jl b/src/prognostic_equations/surface_flux.jl index 8dc015373f..e76e083040 100644 --- a/src/prognostic_equations/surface_flux.jl +++ b/src/prognostic_equations/surface_flux.jl @@ -104,7 +104,9 @@ function surface_flux_tendency!(Yₜ, Y, p, t) p.atmos.disable_surface_flux_tendency && return FT = eltype(Y) - (; sfc_conditions) = p.precomputed + (; params) = p + (; sfc_conditions, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(params) if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion) btt = @@ -112,7 +114,13 @@ function surface_flux_tendency!(Yₜ, Y, p, t) @. Yₜ.c.uₕ -= btt end - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) btt = boundary_tendency_scalar(ᶜh_tot, sfc_conditions.ρ_flux_h_tot) @. Yₜ.c.ρe_tot -= btt ρ_flux_χ = p.scratch.sfc_temp_C3 diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index c5828944e5..7822df474f 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -104,7 +104,13 @@ function vertical_diffusion_boundary_layer_tendency!( top = Operators.SetValue(C3(0)), bottom = Operators.SetValue(C3(0)), ) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(Y.c.ρe_tot, Y.c.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜh_tot))) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 80c88d65f8..6792bd080b 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -1,4 +1,5 @@ import ClimaCore.MatrixFields: @name +import ClimaCore.RecursiveApply: ⊞, ⊠, rzero, rpromote_type """ specific(ρχ, ρ) @@ -218,7 +219,7 @@ Arguments: - `turbconv_model`: The turbulence convection model, containing parameters for regularization (e.g., `a_half`). """ function specific(ρχ, ρ) - return ρχ / ρ + return ρχ / ρ end function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) @@ -230,7 +231,7 @@ function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ end - """ +""" sgs_weight_function(a, a_half) Computes a smooth, monotonic weight function `w(a)` that ranges from 0 to 1. @@ -380,24 +381,27 @@ function specific_env_mse(gs, p) (; ᶜK, ᶜts) = p.precomputed # TODO: replace by on-the-fly computation (; turbconv_model) = p.atmos thermo_params = CAP.thermodynamics_params(p.params) - ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, specific(gs.ρe_tot, gs.ρ))) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(gs.ρe_tot, gs.ρ), + ), + ) # 1. Define the grid-scale moist static energy density `ρ * mse`. grid_scale_ρmse = gs.ρ .* (ᶜh_tot .- ᶜK) # 2. Compute the environment's density-area-weighted mse (`ρa⁰mse⁰`). - ρa⁰mse⁰ = env_value(grid_scale_ρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, gs) + ρa⁰mse⁰ = p.scratch.ᶜtemp_scalar + @. ρa⁰mse⁰ = env_value(grid_scale_ρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, gs) # 3. Compute the environment's density-area product (`ρa⁰`). - ρa⁰_val = ρa⁰(gs) + ρa⁰_val = @. lazy(ρa⁰(gs)) # 4. Compute and return the final specific environment mse (`mse⁰`). - return specific( - ρa⁰mse⁰, - ρa⁰_val, - grid_scale_ρmse, - gs.ρ, - turbconv_model, + return @. lazy( + specific(ρa⁰mse⁰, ρa⁰_val, grid_scale_ρmse, gs.ρ, turbconv_model), ) end From 2e5bb37c095bd8bf61c644ddb3f2402f46b9a205 Mon Sep 17 00:00:00 2001 From: costachris Date: Thu, 10 Jul 2025 14:12:03 -0700 Subject: [PATCH 3/7] restructure env, specific helper functions --- src/cache/cloud_fraction.jl | 9 +- .../diagnostic_edmf_precomputed_quantities.jl | 50 ++- .../precipitation_precomputed_quantities.jl | 33 +- src/cache/precomputed_quantities.jl | 7 +- .../prognostic_edmf_precomputed_quantities.jl | 60 ++-- src/diagnostics/Diagnostics.jl | 4 +- src/diagnostics/edmfx_diagnostics.jl | 28 +- .../les_sgs_models/smagorinsky_lilly.jl | 1 + .../microphysics/precipitation.jl | 10 +- .../radiation/radiation.jl | 10 +- src/prognostic_equations/advection.jl | 15 +- src/prognostic_equations/edmfx_entr_detr.jl | 40 ++- src/prognostic_equations/edmfx_sgs_flux.jl | 54 +-- src/prognostic_equations/edmfx_tke.jl | 8 +- .../forcing/external_forcing.jl | 14 +- src/prognostic_equations/hyperdiffusion.jl | 8 +- .../implicit/implicit_tendency.jl | 8 +- .../implicit/manual_sparse_jacobian.jl | 63 ++-- .../vertical_diffusion_boundary_layer.jl | 2 +- src/solver/types.jl | 2 +- src/utils/variable_manipulations.jl | 313 ++++++++++-------- 21 files changed, 420 insertions(+), 319 deletions(-) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index 5804a90442..ca80896def 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -84,7 +84,6 @@ NVTX.@annotate function set_cloud_fraction!( (; turbconv_model) = p.atmos if isnothing(turbconv_model) - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) if p.atmos.call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed @@ -191,7 +190,7 @@ NVTX.@annotate function set_cloud_fraction!( (; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed (; ᶜρʲs, ᶜtsʲs) = p.precomputed (; turbconv_model) = p.atmos - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) # TODO - we should make this default when using diagnostic edmf # environment @@ -210,9 +209,9 @@ NVTX.@annotate function set_cloud_fraction!( # weight cloud diagnostics by environmental area @. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}( tuple( - draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), - draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), - draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)), ), ) # updrafts diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 9f2786cf4c..09a0272c99 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -105,21 +105,24 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( turbconv_params = CAP.turbconv_params(params) ᶜts = p.precomputed.ᶜts #TODO replace - q_tot = specific(Y.c.ρq_tot, Y.c.ρ) + q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) ᶜh_tot = @. lazy( TD.total_specific_enthalpy( thermo_params, ᶜts, - specific(Y.c.ρe_tot, Y.c.ρ), + ᶜe_tot, ), ) ρ_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), 1)) K_int_level = Fields.field_values(Fields.level(ᶜK, 1)) - q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1)) + q_tot_int_level = + Fields.field_values(Fields.level(Base.materialize(q_tot), 1)) p_int_level = Fields.field_values(Fields.level(ᶜp, 1)) Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1)) @@ -330,7 +333,9 @@ 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 + (; ᶠu³⁰, ᶜK⁰) = p.precomputed + + if microphysics_model isa Microphysics1Moment q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) @@ -354,14 +359,25 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half)) # integral - q_tot = specific(Y.c.ρq_tot, Y.c.ρ) + q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + 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)) 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)) + h_tot_level = + Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), i)) + q_tot_level = + Fields.field_values(Fields.level(Base.materialize(q_tot), i)) p_level = Fields.field_values(Fields.level(ᶜp, i)) Φ_level = Fields.field_values(Fields.level(ᶜΦ, i)) local_geometry_level = Fields.field_values( @@ -385,8 +401,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)) - q_tot_prev_level = Fields.field_values(Fields.level(q_tot, i - 1)) + h_tot_prev_level = + Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), i - 1)) + q_tot_prev_level = + Fields.field_values(Fields.level(Base.materialize(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)) z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1)) @@ -486,7 +504,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( Fields.field_values(Fields.level(ᶜq_iceʲ, i - 1)) end - tke_prev_level = Fields.field_values(Fields.level(ᶜtke⁰, i - 1)) + tke_prev_level = Fields.field_values( + Fields.level(Base.materialize(ᶜtke⁰), i - 1), + ) @. entrʲ_prev_level = entrainment( thermo_params, @@ -1015,6 +1035,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! @. ᶜprandtl_nvec = turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke_exch = p.scratch.ᶜtemp_scalar_2 @. ᶜtke_exch = 0 # using ᶜu⁰ would be more correct, but this is more consistent with the @@ -1025,7 +1047,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰) end - sfc_tke = Fields.level(ᶜtke⁰, 1) + sfc_tke = Fields.level(Base.materialize(ᶜtke⁰), 1) z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, half) @. ᶜmixing_length_tuple = mixing_length( params, @@ -1090,12 +1112,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed # Environment precipitation sources (to be applied to grid mean) - q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources( thermo_params, microphys_0m_params, dt, - specific(Y.c.ρq_tot, Y.c.ρ), + ᶜq_tot, ᶜts, ) return nothing diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index fc5e770e96..3187f6dfb4 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -262,14 +262,14 @@ function set_precipitation_cache!( (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed (; ᶜts⁰, ᶜtsʲs) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) n = n_mass_flux_subdomains(p.atmos.turbconv_model) - @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰ + @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰_vals @. ᶜS_ρe_tot = ᶜSqₜᵖ⁰ * - ᶜρa⁰ * + ᶜρa⁰_vals * e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts⁰, ᶜΦ) for j in 1:n @. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * Y.c.sgsʲs.:($$j).ρa @@ -289,8 +289,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed - q_rai = specific(Y.c.ρq_rai, Y.c.ρ) - q_sno = specific(Y.c.ρq_sno, Y.c.ρ) + q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2 @@ -364,10 +367,10 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _) (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed - q_liq = specific(Y.c.ρq_liq, Y.c.ρ) - q_rai = specific(Y.c.ρq_rai, Y.c.ρ) - n_liq = specific(Y.c.ρn_liq, Y.c.ρ) - n_rai = specific(Y.c.ρn_rai, Y.c.ρ) + q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + n_liq = @. lazy(specific(Y.c.ρn_liq, Y.c.ρ)) + n_rai = @. lazy(specific(Y.c.ρn_rai, Y.c.ρ)) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2 @@ -481,20 +484,24 @@ function set_precipitation_surface_fluxes!( sfc_ρ = @. lazy(int_ρ * int_J / sfc_J) # Constant extrapolation to surface, consistent with simple downwinding + ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) sfc_qᵣ = Fields.Field( - Fields.field_values(Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1)), + Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)), sfc_space, ) sfc_qₛ = Fields.Field( - Fields.field_values(Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1)), + Fields.field_values(Fields.level(Base.materialize(ᶜq_sno), 1)), sfc_space, ) sfc_qₗ = Fields.Field( - Fields.field_values(Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1)), + Fields.field_values(Fields.level(Base.materialize(ᶜq_liq), 1)), sfc_space, ) sfc_qᵢ = Fields.Field( - Fields.field_values(Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1)), + Fields.field_values(Fields.level(Base.materialize(ᶜq_ice), 1)), sfc_space, ) sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space) diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index b560e7d7f3..47d93a05a1 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -319,8 +319,9 @@ function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model) return nothing end -function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model) - @. ᶜK += ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ +function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰_vals, ᶠu₃⁰, turbconv_model) + @. ᶜK += + ᶜρa⁰_vals * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ for j in 1:n_mass_flux_subdomains(turbconv_model) ᶜρaʲ = Y.c.sgsʲs.:($j).ρa ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃ @@ -473,8 +474,6 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) set_prognostic_edmf_precomputed_quantities_draft!(Y, p, ᶠuₕ³, t) set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t) set_prognostic_edmf_precomputed_quantities_implicit_closures!(Y, p, t) - elseif turbconv_model isa DiagnosticEDMFX - set_diagnostic_edmf_precomputed_quantities!(Y, p, t) elseif !(isnothing(turbconv_model)) # Do nothing for other turbconv models for now end diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index aa5a6b17f0..b555813d62 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -24,22 +24,22 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( (; ᶜp, ᶜK) = p.precomputed (; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) - ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ - ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 ᶜmse⁰ .= specific_env_mse(Y.c, p) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment - ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) - ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) - ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) - ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) @. ᶜts⁰ = TD.PhaseNonEquil_phq( thermo_params, ᶜp, @@ -219,26 +219,31 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment # TODO - any better way to define the cloud and precip tracer flux? + + ᶜq_liq = @.lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + ᶜq_ice = @.lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + ᶜq_rai = @.lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + ᶜq_sno = @.lazy(specific(Y.c.ρq_sno, Y.c.ρ)) ᶜq_liq_int_val = Fields.field_values( - Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1), + Fields.level(Base.materialize(ᶜq_liq), 1), ) ᶜq_liqʲ_int_val = Fields.field_values(Fields.level(ᶜq_liqʲ, 1)) @. ᶜq_liqʲ_int_val = ᶜq_liq_int_val ᶜq_ice_int_val = Fields.field_values( - Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1), + Fields.level(Base.materialize(ᶜq_ice), 1), ) ᶜq_iceʲ_int_val = Fields.field_values(Fields.level(ᶜq_iceʲ, 1)) @. ᶜq_iceʲ_int_val = ᶜq_ice_int_val ᶜq_rai_int_val = Fields.field_values( - Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1), + Fields.level(Base.materialize(ᶜq_rai), 1), ) ᶜq_raiʲ_int_val = Fields.field_values(Fields.level(ᶜq_raiʲ, 1)) @. ᶜq_raiʲ_int_val = ᶜq_rai_int_val ᶜq_sno_int_val = Fields.field_values( - Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1), + Fields.level(Base.materialize(ᶜq_sno), 1), ) ᶜq_snoʲ_int_val = Fields.field_values(Fields.level(ᶜq_snoʲ, 1)) @. ᶜq_snoʲ_int_val = ᶜq_sno_int_val @@ -368,8 +373,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ᶜdz = Fields.Δz_field(axes(Y.c)) ᶜlg = Fields.local_geometry_field(Y.c) ᶠlg = Fields.local_geometry_field(Y.f) - ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) - ᶜρa⁰ = @. lazy(ρa⁰(Y.c)) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) ᶜvert_div = p.scratch.ᶜtemp_scalar ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2 @@ -454,7 +459,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos (; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed # First order approximation: Use environmental mean fields. @. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat - ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) @. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat @. ᶜgradᵥ_θ_liq_ice⁰ = ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat @@ -487,7 +492,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos for j in 1:n ᶠu³ʲ = ᶠu³ʲs.:($j) @. ᶜtke_exch += - Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰ * + Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰_vals * (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰) end @@ -514,7 +519,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec) ρatke_flux_values = Fields.field_values(ρatke_flux) - ρa_sfc_values = Fields.field_values(Fields.level(Base.materialize(ᶜρa⁰), 1)) # TODO: replace by surface value + ρa_sfc_values = + Fields.field_values(Fields.level(Base.materialize(ᶜρa⁰_vals), 1)) # TODO: replace by surface value ustar_values = Fields.field_values(ustar) sfc_local_geometry_values = Fields.field_values( Fields.level(Fields.local_geometry_field(Y.f), half), @@ -565,7 +571,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ) end # sources from the environment - ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, p.atmos.turbconv_model)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰) return nothing end @@ -655,10 +661,12 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation end # Precipitation sources and sinks from the environment - ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) - ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) - ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) - ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) + ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰)) compute_precipitation_sources!( ᶜSᵖ, ᶜSᵖ_snow, @@ -666,7 +674,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, - TD.air_density.(thp, ᶜts⁰), + ᶜρ⁰, + ᶜq_tot⁰, + ᶜq_liq⁰, + ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰, ᶜts⁰, @@ -678,7 +689,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ᶜSᵖ, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, - TD.air_density.(thp, ᶜts⁰), + ᶜρ⁰, + ᶜq_tot⁰, + ᶜq_liq⁰, + ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰, ᶜts⁰, diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl index 8e1589efe3..e08c1de0e9 100644 --- a/src/diagnostics/Diagnostics.jl +++ b/src/diagnostics/Diagnostics.jl @@ -58,8 +58,8 @@ import ..SlabOceanSST import ..draft_area import ..compute_gm_mixing_length! import ..horizontal_integral_at_boundary -import ..ρa⁰ -import ..specific_tke +import ..ᶜρa⁰ +import ..ᶜspecific_tke # We need the abbreviations for symbols like curl, grad, and so on include(joinpath("..", "utils", "abbreviations.jl")) diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index ca369b0182..54a41a6078 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -633,16 +633,16 @@ compute_aren!(_, _, _, _, turbconv_model::T) where {T} = function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - ᶜρa⁰ = @. lazy(ρa⁰(state.c)) + ᶜρa⁰_vals = ᶜρa⁰(state.c, cache) if isnothing(out) return draft_area.( - ᶜρa⁰, + ᶜρa⁰_vals, TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), ) else out .= draft_area.( - ᶜρa⁰, + ᶜρa⁰_vals, TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), ) end @@ -953,10 +953,11 @@ function compute_clwen!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), state.c, cache) if isnothing(out) - return specific_env_value(:q_liq, state.c, turbconv_model) + return Base.materialize(ᶜq_liq⁰) else - out .= specific_env_value(:q_liq, state.c, turbconv_model) + out .= ᶜq_liq⁰ end end @@ -1017,10 +1018,11 @@ function compute_clien!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), state.c, cache) if isnothing(out) - return specific_env_value(:q_ice, state.c, turbconv_model) + return Base.materialize(ᶜq_ice⁰) else - out .= specific_env_value(:q_ice, state.c, turbconv_model) + out .= ᶜq_ice⁰ end end @@ -1065,10 +1067,11 @@ function compute_husraen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), state.c, cache) if isnothing(out) - return specific_env_value(:q_rai, state.c, turbconv_model) + return Base.materialize(ᶜq_rai⁰) else - out .= specific_env_value(:q_rai, state.c, turbconv_model) + out .= ᶜq_rai⁰ end end @@ -1113,10 +1116,11 @@ function compute_hussnen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), state.c, cache) if isnothing(out) - return specific_env_value(:q_sno, state.c, turbconv_model) + return Base.materialize(ᶜq_sno⁰) else - out .= specific_env_value(:q_sno, state.c, turbconv_model) + out .= ᶜq_sno⁰ end end @@ -1146,7 +1150,7 @@ function compute_tke!( time, turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX}, ) - ᶜtke = @. lazy(specific_tke(state.c.sgs⁰, state.c, turbconv_model)) + ᶜtke = ᶜspecific_tke(state.c.sgs⁰, state.c, cache) if isnothing(out) return Base.materialize(ᶜtke) else diff --git a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl index 2b57bd2bce..69e7406a2c 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -97,6 +97,7 @@ vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) ## Momentum tendencies ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(Y.c.ρ) diff --git a/src/parameterized_tendencies/microphysics/precipitation.jl b/src/parameterized_tendencies/microphysics/precipitation.jl index 955e918359..b642880398 100644 --- a/src/parameterized_tendencies/microphysics/precipitation.jl +++ b/src/parameterized_tendencies/microphysics/precipitation.jl @@ -150,13 +150,13 @@ function precipitation_tendency!( # Source terms from EDMFX environment (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰) = p.precomputed - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) # Update from environment precipitation and cloud formation sources/sinks - @. Yₜ.c.ρq_liq += ᶜρa⁰ * ᶜSqₗᵖ⁰ - @. Yₜ.c.ρq_ice += ᶜρa⁰ * ᶜSqᵢᵖ⁰ - @. Yₜ.c.ρq_rai += ᶜρa⁰ * ᶜSqᵣᵖ⁰ - @. Yₜ.c.ρq_sno += ᶜρa⁰ * ᶜSqₛᵖ⁰ + @. Yₜ.c.ρq_liq += ᶜρa⁰_vals * ᶜSqₗᵖ⁰ + @. Yₜ.c.ρq_ice += ᶜρa⁰_vals * ᶜSqᵢᵖ⁰ + @. Yₜ.c.ρq_rai += ᶜρa⁰_vals * ᶜSqᵣᵖ⁰ + @. Yₜ.c.ρq_sno += ᶜρa⁰_vals * ᶜSqₛᵖ⁰ # Update from the updraft precipitation sources n = n_mass_flux_subdomains(p.atmos.turbconv_model) diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl index 992054c4c7..d1ae29f51a 100644 --- a/src/parameterized_tendencies/radiation/radiation.jl +++ b/src/parameterized_tendencies/radiation/radiation.jl @@ -449,7 +449,7 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS) (; params) = p (; ᶜts) = p.precomputed - (; ᶜκρq, ∫_0_∞_κρq, ᶠ∫_0_z_κρq, isoline_z_ρ_q, ᶠradiation_flux) = + (; ᶜκρq, ∫_0_∞_κρq, ᶠ∫_0_z_κρq, isoline_z_ρ_ρq, ᶠradiation_flux) = p.radiation (; ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(params) @@ -476,10 +476,10 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS) q_tot_isoline = FT(0.008) Operators.column_reduce!( (nt1, nt2) -> - abs(nt1.q_tot - q_tot_isoline) < abs(nt2.q_tot - q_tot_isoline) ? - nt1 : nt2, - isoline_z_ρ_q, - Base.broadcasted(NT ∘ tuple, ᶜz, Y.c.ρ, specific(Y.c.ρq_tot, Y.c.ρ)), + abs(specific.(nt1.ρq_tot, nt1.ρ) - q_tot_isoline) < + abs(specific.(nt2.ρq_tot, nt2.ρ) - q_tot_isoline) ? nt1 : nt2, + isoline_z_ρ_ρq, + Base.broadcasted(NT ∘ tuple, ᶜz, Y.c.ρ, specific.(Y.c.ρq_tot, Y.c.ρ)), ) zi = isoline_z_ρ_ρq.z diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 7b7575fd2b..6aef68de3a 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -186,6 +186,7 @@ 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 + thermo_params = CAP.thermodynamics_params(p.params) ᶠu³⁰ = advect_tke ? @@ -193,20 +194,18 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ : p.precomputed.ᶠu³⁰ ) : nothing - ᶜρa⁰ = advect_tke ? (n > 0 ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ) : nothing + ᶜρa⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y.c, p)) : lazy(Y.c.ρ)) : nothing ᶜρ⁰ = if advect_tke if n > 0 - thermo_params = CAP.thermodynamics_params(p.params) - @. TD.air_density(thermo_params, p.precomputed.ᶜts⁰) + (; ᶜts⁰) = p.precomputed.ᶜts⁰ + @. lazy(TD.air_density(thermo_params, ᶜts⁰)) else - Y.c.ρ + lazy(Y.c.ρ) end else nothing end - ᶜtke⁰ = - advect_tke ? (@.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))) : - nothing + ᶜtke⁰ = advect_tke ? (ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)) : nothing ᶜa_scalar = p.scratch.ᶜtemp_scalar ᶜω³ = p.scratch.ᶜtemp_CT3 ᶠω¹² = p.scratch.ᶠtemp_CT12 @@ -285,7 +284,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) end if use_prognostic_tke(turbconv_model) # advect_tke triggers allocations - @. ᶜa_scalar = ᶜtke⁰ * draft_area(ᶜρa⁰, ᶜρ⁰) + @. ᶜa_scalar = ᶜtke⁰ * draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport(ᶜρ⁰, ᶠu³⁰, ᶜa_scalar, dt, edmfx_upwinding) @. Yₜ.c.sgs⁰.ρatke += vtt end diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index 8cfcc14d4a..f5440ab001 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -534,44 +534,42 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF ᶜmse⁰ .= specific_env_mse(Y.c, p) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment - ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) - ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) - ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) - ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) end for j in 1:n + ᶜentrʲ = ᶜentrʲs.:($j) + ᶜdetrʲ = ᶜdetrʲs.:($j) + ᶜturb_entrʲ = ᶜturb_entrʲs.:($j) + ᶜmseʲ = Y.c.sgsʲs.:($j).mse + ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot - @. Yₜ.c.sgsʲs.:($$j).ρa += - Y.c.sgsʲs.:($$j).ρa * (ᶜentrʲs.:($$j) - ᶜdetrʲs.:($$j)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) - @. Yₜ.c.sgsʲs.:($$j).mse += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜmse⁰ - Y.c.sgsʲs.:($$j).mse) + @. Yₜ.c.sgsʲs.:($$j).ρa += Y.c.sgsʲs.:($$j).ρa * (ᶜentrʲ - ᶜdetrʲ) + + @. Yₜ.c.sgsʲs.:($$j).mse += (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜmse⁰ - ᶜmseʲ) - ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) @. Yₜ.c.sgsʲs.:($$j).q_tot += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_tot⁰ - Y.c.sgsʲs.:($$j).q_tot) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_tot⁰ - ᶜq_totʲ) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment @. Yₜ.c.sgsʲs.:($$j).q_liq += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_liq⁰ - Y.c.sgsʲs.:($$j).q_liq) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_liq⁰ - Y.c.sgsʲs.:($$j).q_liq) @. Yₜ.c.sgsʲs.:($$j).q_ice += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_ice⁰ - Y.c.sgsʲs.:($$j).q_ice) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_ice⁰ - Y.c.sgsʲs.:($$j).q_ice) @. Yₜ.c.sgsʲs.:($$j).q_rai += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_rai⁰ - Y.c.sgsʲs.:($$j).q_rai) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_rai⁰ - Y.c.sgsʲs.:($$j).q_rai) @. Yₜ.c.sgsʲs.:($$j).q_sno += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_sno⁰ - Y.c.sgsʲs.:($$j).q_sno) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_sno⁰ - Y.c.sgsʲs.:($$j).q_sno) end @. Yₜ.f.sgsʲs.:($$j).u₃ += - (ᶠinterp(ᶜentrʲs.:($$j)) .+ ᶠinterp(ᶜturb_entrʲs.:($$j))) * + (ᶠinterp(ᶜentrʲ) .+ ᶠinterp(ᶜturb_entrʲ)) * (ᶠu₃⁰ - Y.f.sgsʲs.:($$j).u₃) end return nothing diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index c93d7ce469..fde18ea572 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -45,7 +45,7 @@ function edmfx_sgs_mass_flux_tendency!( (; ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) ᶜρ⁰ = @. TD.air_density(thermo_params, ᶜts⁰) - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J @@ -82,7 +82,7 @@ function edmfx_sgs_mass_flux_tendency!( ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 ᶜmse⁰ .= specific_env_mse(Y.c, p) - @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) + @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -109,10 +109,11 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_tot += vtt end # Add the environment fluxes - ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = - (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * + draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -127,10 +128,10 @@ function edmfx_sgs_mass_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) - ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) - ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) - ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) - ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) # Liquid, ice, rain and snow specific humidity fluxes for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @@ -186,7 +187,8 @@ function edmfx_sgs_mass_flux_tendency!( @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = - (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * + draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -197,7 +199,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_liq += vtt @. ᶜa_scalar = - (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * + draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -208,7 +211,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_ice += vtt @. ᶜa_scalar = - (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * + draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -219,7 +223,8 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_rai += vtt @. ᶜa_scalar = - (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) + (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * + draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, ᶠu³_diff, @@ -255,6 +260,7 @@ function edmfx_sgs_mass_flux_tendency!( FT = eltype(Y) if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} + thermo_params = CAP.thermodynamics_params(p.params) # energy ᶜh_tot = @. lazy( TD.total_specific_enthalpy( @@ -340,7 +346,7 @@ function edmfx_sgs_mass_flux_tendency!( # ) # @. Yₜ.c.ρe_tot += vtt # if !(p.atmos.moisture_model isa DryModel) - # ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + # ᶜq_tot⁰ = @specific_env_value(:q_tot, Y.c, turbconv_model)) # @. ᶜa_scalar = # (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * # draft_area(ᶜρa⁰, ᶜρ⁰) @@ -399,21 +405,21 @@ function edmfx_sgs_diffusive_flux_tendency!( (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed (; ᶜmixing_length, ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) - ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar - @. ᶠρaK_h = ᶠinterp(ᶜρa⁰) * ᶠinterp(ᶜK_h) + @. ᶠρaK_h = ᶠinterp(ᶜρa⁰_vals) * ᶠinterp(ᶜK_h) ᶠρaK_u = p.scratch.ᶠtemp_scalar - @. ᶠρaK_u = ᶠinterp(ᶜρa⁰) * ᶠinterp(ᶜK_u) + @. ᶠρaK_u = ᶠinterp(ᶜρa⁰_vals) * ᶠinterp(ᶜK_u) # Total enthalpy diffusion ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - # ᶜmse⁰ = @. lazy(specific_env_mse(Y.c, p)) + ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 ᶜmse⁰ .= specific_env_mse(Y.c, p) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜmse⁰ + ᶜK⁰))) @@ -444,7 +450,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - ᶜq_tot⁰ = @.lazy(specific_env_value(:q_tot, Y.c, turbconv_model)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜq_tot⁰))) @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion @. Yₜ.c.ρ -= ᶜρχₜ_diffusion # Effect of moisture diffusion on (moist) air mass @@ -453,10 +459,10 @@ function edmfx_sgs_diffusive_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) - ᶜq_liq⁰ = @.lazy(specific_env_value(:q_liq, Y.c, turbconv_model)) - ᶜq_ice⁰ = @.lazy(specific_env_value(:q_ice, Y.c, turbconv_model)) - ᶜq_rai⁰ = @.lazy(specific_env_value(:q_rai, Y.c, turbconv_model)) - ᶜq_sno⁰ = @.lazy(specific_env_value(:q_sno, Y.c, turbconv_model)) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) # Liquid, ice, rain and snow specific humidity diffusion α_vert_diff_tracer = CAP.α_vert_diff_tracer(params) @@ -509,7 +515,7 @@ function edmfx_sgs_diffusive_flux_tendency!( (; ᶜu, ᶜmixing_length, ᶜts) = p.precomputed (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() - ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index a8d6d6918d..09d7f5ca2a 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -50,7 +50,7 @@ function edmfx_tke_tendency!( FT = eltype(p.params) - ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ + ᶜρa⁰_vals = turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y.c, p) : Y.c.ρ nh_pressure3_buoyʲs = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs @@ -73,11 +73,11 @@ function edmfx_tke_tendency!( if use_prognostic_tke(turbconv_model) # shear production - @. Yₜ.c.sgs⁰.ρatke += 2 * ᶜρa⁰ * ᶜK_u * ᶜstrain_rate_norm + @. Yₜ.c.sgs⁰.ρatke += 2 * ᶜρa⁰_vals * ᶜK_u * ᶜstrain_rate_norm # buoyancy production - @. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰ * ᶜK_h * ᶜlinear_buoygrad + @. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰_vals * ᶜK_h * ᶜlinear_buoygrad - ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) # entrainment and detraiment # using ᶜu⁰ and local geometry results in allocation diff --git a/src/prognostic_equations/forcing/external_forcing.jl b/src/prognostic_equations/forcing/external_forcing.jl index 64710ad3c0..36385ec154 100644 --- a/src/prognostic_equations/forcing/external_forcing.jl +++ b/src/prognostic_equations/forcing/external_forcing.jl @@ -337,7 +337,7 @@ 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 + (; ᶜts) = p.precomputed (; ᶜdTdt_fluc, ᶜdqtdt_fluc, @@ -358,6 +358,13 @@ function external_forcing_tendency!( @. ᶜuₕ_nudge = C12(Geometry.UVVector(ᶜu_nudge, ᶜv_nudge), ᶜlg) @. Yₜ.c.uₕ -= (Y.c.uₕ - ᶜuₕ_nudge) * ᶜinv_τ_wind + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) # nudging tendency ᶜdTdt_nudging = p.scratch.ᶜtemp_scalar ᶜdqtdt_nudging = p.scratch.ᶜtemp_scalar_2 @@ -403,8 +410,7 @@ function external_forcing_tendency!( Yₜ.c.ρq_tot, Y.c.ρ, ᶠls_subsidence³, - specific(Y.c.ρq_tot, Y.c.ρ), - specific(Y.c.ρq_tot, Y.c.ρ), + ᶜq_tot, Val{:first_order}(), ) @@ -581,7 +587,7 @@ 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) = p.precomputed ᶜinv_τ_scalar = APL.ISDAC_inv_τ_scalar(FT) # s⁻¹ ᶜinv_τ_wind = APL.ISDAC_inv_τ_wind(FT) # s⁻¹ diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index 8392424998..331dc914e2 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -117,7 +117,7 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) if diffuse_tke - ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) (; ᶜ∇²tke⁰) = p.hyperdiff @. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰)) end @@ -150,11 +150,11 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) point_type = eltype(Fields.coordinate_field(Y.c)) (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX - ᶜρa⁰ = @.lazy(ρa⁰(Y.c)) + ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end if use_prognostic_tke(turbconv_model) - ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) (; ᶜ∇²tke⁰) = p.hyperdiff end @@ -175,7 +175,7 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) # Sub-grid scale hyperdiffusion continued if (turbconv_model isa PrognosticEDMFX) && diffuse_tke - @. Yₜ.c.sgs⁰.ρatke -= ν₄_vorticity * wdivₕ(ᶜρa⁰ * gradₕ(ᶜ∇²tke⁰)) + @. Yₜ.c.sgs⁰.ρatke -= ν₄_vorticity * wdivₕ(ᶜρa⁰_vals * gradₕ(ᶜ∇²tke⁰)) end if turbconv_model isa PrognosticEDMFX for j in 1:n diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 4ecfafe59e..60f4f90b68 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -157,13 +157,7 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) @. Yₜ.c.ρe_tot += vtt if !(moisture_model isa DryModel) ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - vtt = vertical_transport( - Y.c.ρ, - ᶠu³, - ᶜq_tot, - dt, - Val(:none), - ) + vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜq_tot, dt, Val(:none)) @. Yₜ.c.ρq_tot += vtt end diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index 6168f13bfd..129b6a4f21 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -461,15 +461,15 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ᶠinterp_matrix() ) @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) - e_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) - q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] @. ∂ᶠu₃_err_∂ᶜρq_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot + ᶜ∂kappa_m∂q_tot * - (cp_d * T_0 + e_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * q_tot) + (cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * ᶜq_tot) )) end @@ -539,7 +539,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ # DiagonalMatrixRow( # -1 / ᶜρ * ifelse( - # q_tot == 0, + # ᶜq_tot == 0, # (Geometry.WVector(FT(0)),), # p.precomputed.ᶜwₜqₜ / _tot, # ), @@ -581,23 +581,24 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] @. ∂ᶜρe_tot_err_∂ᶜρ = dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( - (-(1 + ᶜkappa_m) * e_tot - ᶜkappa_m * ∂e_int_∂q_tot * q_tot) / + (-(1 + ᶜkappa_m) * ᶜe_tot - ᶜkappa_m * ∂e_int_∂q_tot * ᶜq_tot) / ᶜρ, ) @. ∂ᶜρe_tot_err_∂ᶜρe_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] ∂ᶜρq_tot_err_∂ᶜρ = matrix[@name(c.ρq_tot), @name(c.ρ)] @. ∂ᶜρe_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + ᶜ∂kappa_m∂q_tot * - (cp_d * T_0 + e_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * q_tot) + (cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * ᶜq_tot) )) @. ∂ᶜρq_tot_err_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(q_tot) / ᶜρ) + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(ᶜq_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) end @@ -624,10 +625,10 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) (; dt) = p turbconv_model = p.atmos.turbconv_model ᶜmixing_length = p.precomputed.ᶜmixing_length - ᶜtke⁰ = @.lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model)) - ᶜρa⁰ = - p.atmos.turbconv_model isa PrognosticEDMFX ? - (@.lazy(ρa⁰(Y.c))) : Y.c.ρ + ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜρa⁰_vals = + p.atmos.turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y.c, p) : + Y.c.ρ ᶜρatke⁰ = Y.c.sgs⁰.ρatke @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) = @@ -648,13 +649,13 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) dtγ * ( ᶜdiffusion_u_matrix - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) - ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰) / ᶜρa⁰) + ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰) / ᶜρa⁰_vals) @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = dtγ * ( ( ᶜdiffusion_u_matrix - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) - ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰) - DiagonalMatrixRow( + ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰_vals) - DiagonalMatrixRow( tke_dissipation_rate_tendency(ᶜtke⁰, ᶜmixing_length), ) ) - (I,) @@ -697,7 +698,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) / TD.cv_m(thermo_params, ᶜtsʲs.:(1)) - # Note this is the derivative of R_m / cp_m with respect to q_tot + # Note this is the derivative of R_m / cp_m with respect to ᶜq_tot # but we call it ∂kappa_m∂q_totʲ ᶜ∂kappa_m∂q_totʲ = p.scratch.ᶜtemp_scalar_2 @. ᶜ∂kappa_m∂q_totʲ = @@ -953,23 +954,29 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) Δcv_v * TD.gas_constant_air(thermo_params, ᶜts) ) / abs2(TD.cv_m(thermo_params, ᶜts)) + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + @. ∂ᶜρe_tot_err_∂ᶜρq_tot += + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow(( + ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + + ᶜ∂kappa_m∂q_tot * + (cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * ᶜq_tot) + )) + + else + ᶜq_tot = lazy(0) + end + @. ∂ᶜρe_tot_err_∂ᶜρ += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ DiagonalMatrixRow( ( - -(1 + ᶜkappa_m) * e_tot - - ᶜkappa_m * ∂e_int_∂q_tot * q_tot + -(1 + ᶜkappa_m) * ᶜe_tot - + ᶜkappa_m * ∂e_int_∂q_tot * ᶜq_tot ) / ᶜρ, ) - @. ∂ᶜρe_tot_err_∂ᶜρq_tot += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(( - ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * - (cp_d * T_0 + e_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * q_tot) - )) - @. ∂ᶜρe_tot_err_∂ᶜρe_tot += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) @@ -982,7 +989,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ## grid-mean ρq_tot @. ∂ᶜρq_tot_err_∂ᶜρ += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(-(q_tot) / ᶜρ) + DiagonalMatrixRow(-(ᶜq_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ @@ -1021,7 +1028,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶠu₃ += dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( ᶠinterp( - (Y.c.sgsʲs.:(1).q_tot - q_tot) * + (Y.c.sgsʲs.:(1).q_tot - ᶜq_tot) * ᶜρʲs.:(1) * ᶜJ * draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), @@ -1033,7 +1040,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶠu₃ʲ = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠinterp( - (Y.c.sgsʲs.:(1).q_tot - q_tot) * + (Y.c.sgsʲs.:(1).q_tot - ᶜq_tot) * ᶜρʲs.:(1) * ᶜJ * draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), @@ -1054,7 +1061,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶜρa = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( (ᶠu³ʲs.:(1) - ᶠu³) * - ᶠinterp((Y.c.sgsʲs.:(1).q_tot - q_tot)) / ᶠJ, + ᶠinterp((Y.c.sgsʲs.:(1).q_tot - ᶜq_tot)) / ᶠJ, ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) end elseif rs isa RayleighSponge diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 7822df474f..49a34edba2 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -89,7 +89,7 @@ function vertical_diffusion_boundary_layer_tendency!( FT = eltype(Y) α_vert_diff_tracer = CAP.α_vert_diff_tracer(p.params) thermo_params = CAP.thermodynamics_params(p.params) - (; ᶜu, ᶜK_u, ᶜK_h) = p.precomputed + (; ᶜu, ᶜK_u, ᶜK_h, ᶜts) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion) diff --git a/src/solver/types.jl b/src/solver/types.jl index ce5bc047cc..6062eb5d58 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -364,7 +364,7 @@ PrognosticEDMFX{N, TKE}(a_half::FT) where {N, TKE, FT} = PrognosticEDMFX{N, TKE, FT}(a_half) struct DiagnosticEDMFX{N, TKE, FT} <: AbstractEDMF - a_half::FT # WARNING: this should never be used outside of d`specific` + a_half::FT # WARNING: this should never be used outside of `specific` end DiagnosticEDMFX{N, TKE}(a_half::FT) where {N, TKE, FT} = DiagnosticEDMFX{N, TKE, FT}(a_half) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 6792bd080b..26fbafe362 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -43,9 +43,9 @@ Arguments: - `ρ_fallback`: The grid-mean density used for the fallback value. - `turbconv_model`: The turbulence convection model, containing parameters for regularization (e.g., `a_half`). """ -@inline specific(ρχ, ρ) = ρχ / ρ +specific(ρχ, ρ) = ρχ / ρ -@inline function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) +function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) # TODO: Replace turbconv_model struct by parameters, and include a_half in # parameters, not in config weight = sgs_weight_function(ρa / ρ, turbconv_model.a_half) @@ -178,58 +178,6 @@ foreach_gs_tracer(f::F, Y_or_similar_values...) where {F} = f(ρχ_or_χ_fields..., ρχ_name) end -""" - specific(ρχ, ρ) - specific(ρaχ, ρa, ρχ, ρ, turbconv_model) - -Calculates the specific quantity `χ` (per unit mass) from a density-weighted -quantity. This function uses multiple dispatch to select the appropriate -calculation method based on the number of arguments. - -**Grid-Scale Method (2 arguments)** - - specific(ρχ, ρ) - -Performs a direct division of the density-weighted quantity `ρχ` by the density `ρ`. -This method is used for grid-mean quantities where the density `ρ` is well-defined -and non-zero. - -**SGS Regularized Method (5 arguments)** - - specific(ρaχ, ρa, ρχ, ρ, turbconv_model) - -Calculates the specific quantity `χ` for a subgrid-scale (SGS) component by -dividing the density-area-weighted quantity `ρaχ` by the density-area -product `ρa`. - -This method includes regularization to handle cases where the SGS area fraction -(and thus `ρa`) is zero or vanishingly small. It performs a linear interpolation -between the SGS specific quantity (`ρaχ / ρa`) and the grid-mean specific -quantity (`ρχ / ρ`). The interpolation weight is computed by `sgs_weight_function` -to ensure a smooth and numerically stable transition, preventing division by zero. -Using this regularized version instead of directly computing `ρaχ / ρa` breaks the -assumption of domain decomposition (sum of SGS domains equals GS) when the approximated -area fraction `a` is small. - -Arguments: -- `ρaχ`: The density-area-weighted SGS quantity (e.g., `sgs.ρa * sgs.h_tot`). -- `ρa`: The density-area product of the SGS component. -- `ρχ`: The fallback grid-mean density-weighted quantity (e.g., `ρe_tot`, `ρq_tot`). -- `ρ`: The fallback grid-mean density. -- `turbconv_model`: The turbulence convection model, containing parameters for regularization (e.g., `a_half`). -""" -function specific(ρχ, ρ) - return ρχ / ρ -end - -function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) - # TODO: Replace turbconv_model struct by parameters, and include a_half in - # parameters, not in config - weight = sgs_weight_function(ρa / ρ, turbconv_model.a_half) - # If ρa is exactly zero, the weight function will be zero, causing the first - # term to be NaN (0 * ... / 0). The ifelse handles this case explicitly. - return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ -end """ sgs_weight_function(a, a_half) @@ -296,7 +244,7 @@ Arguments: draft_sum(f, sgsʲs) = mapreduce_with_init(f, +, sgsʲs) """ - env_value(grid_scale_value, f_draft, gs) + ᶜenv_value(grid_scale_value, f_draft, gs, turbconv_model) Computes the value of a quantity `ρaχ` in the environment subdomain by subtracting the sum of its values in all draft subdomains from the grid-scale value. @@ -304,58 +252,184 @@ the sum of its values in all draft subdomains from the grid-scale value. This is based on the domain decomposition principle for density-area weighted quantities: `GridMean(ρχ) = Env(ρaχ) + Sum(Drafts(ρaχ))`. +The function handles both PrognosticEDMFX and DiagnosticEDMFX models: +- For PrognosticEDMFX: Uses gs.sgsʲs to access draft subdomain states +- For DiagnosticEDMFX: Uses p.precomputed.ᶜρaʲs for draft area-weighted densities + Arguments: - `grid_scale_value`: The `ρa`-weighted grid-scale value of the quantity. - `f_draft`: A function that extracts the corresponding value from a draft subdomain state. -- `gs`: The grid-scale state, which contains the draft subdomain states `gs.sgsʲs`. -""" -function env_value(grid_scale_value, f_draft, gs) - return grid_scale_value - draft_sum(f_draft, gs.sgsʲs) +- `gs`: The grid-scale state, which contains the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX). +- `turbconv_model`: The turbulence convection model, used to determine how to access draft data. +""" +function ᶜenv_value( + grid_scale_value, + f_draft, + gs, + turbconv_model::PrognosticEDMFX, +) + return @. lazy(grid_scale_value - draft_sum(f_draft, gs.sgsʲs)) +end + +function ᶜenv_value( + grid_scale_value, + f_draft, + gs, + turbconv_model::DiagnosticEDMFX, + p, +) + # For DiagnosticEDMFX, we need to access precomputed quantities from p + (; ᶜρaʲs) = p.precomputed + n = n_mass_flux_subdomains(turbconv_model) + + ᶜdraft_sum_diag = p.scratch.ᶜtemp_scalar_4 + @. ᶜdraft_sum_diag = 0 + + for j in 1:n + ᶜρaʲ = ᶜρaʲs.:($j) + @. ᶜdraft_sum_diag += f_draft(ᶜρaʲ) + end + return @. lazy(grid_scale_value - ᶜdraft_sum_diag) end """ - specific_env_value(χ_name::Symbol, gs, turbconv_model) + ᶜspecific_env_value(::Val{χ_name}, gs, p) Calculates the specific value of a quantity `χ` in the environment (`χ⁰`). This function uses the domain decomposition principle to first find the density-area-weighted environment value (`ρa⁰χ⁰`) and the environment -density-area (`ρa⁰`). It then computes the specific value using the +density-weighted environmental area (`ρa⁰`). It then computes the specific value using the regularized `specific` function, which provides a stable result even when the environment area fraction is very small. Arguments: -- `χ_name`: The `Symbol` for the specific quantity `χ` (e.g., `:h_tot`, `:q_tot`). +- `::Val{χ_name}`: A `Val` type containing the symbol for the specific quantity `χ` (e.g., `Val(:h_tot)`, `Val(:q_tot)`). - `gs`: The grid-scale state, containing grid-mean and draft subdomain states. -- `turbconv_model`: The turbulence convection model, containing parameters for regularization. +- `p`: The cache, containing precomputed quantities and turbconv_model. Returns: - The specific value of the quantity `χ` in the environment. """ -function specific_env_value(χ_name::Symbol, gs, turbconv_model) +function ᶜspecific_env_value(::Val{χ_name}, gs, p) where {χ_name} + turbconv_model = p.atmos.turbconv_model + # Grid-scale density-weighted variable name, e.g., :ρq_tot ρχ_name = Symbol(:ρ, χ_name) - # Numerator: ρa⁰χ⁰ = (gs.ρχ) - (Σ sgsʲ.ρa * sgsʲ.χ) - ρaχ⁰ = env_value( - getproperty(gs, ρχ_name), - sgsʲ -> getproperty(sgsʲ, :ρa) * getproperty(sgsʲ, χ_name), - gs, - ) + gs_ρχ_field = getproperty(gs, ρχ_name) + + if turbconv_model isa PrognosticEDMFX + # Numerator: ρa⁰χ⁰ = (gs.ρχ) - (Σ sgsʲ.ρa * sgsʲ.χ) + ρaχ⁰ = ᶜenv_value( + gs_ρχ_field, + sgsʲ -> getproperty(sgsʲ, :ρa) * getproperty(sgsʲ, χ_name), + gs, + turbconv_model, + ) + # Denominator: ρa⁰ = gs.ρ - Σ sgsʲ.ρa + ᶜρa⁰_vals = ᶜρa⁰(gs, p) + else # DiagnosticEDMFX + # For DiagnosticEDMFX, we need to compute ρaʲ * χʲ for each draft + # Get the specific quantity values for all drafts + ᶜχʲs = getproperty(p.precomputed, Symbol(:ᶜ, χ_name, :ʲs)) + n = n_mass_flux_subdomains(turbconv_model) + + # Create combined ρaʲ * χʲ values for each draft + ᶜρaχʲs_combined = p.scratch.ᶜtemp_scalar_3 + @. ᶜρaχʲs_combined = 0 + for j in 1:n + ᶜρaʲ = p.precomputed.ᶜρaʲs.:($j) + ᶜχʲ = ᶜχʲs.:($j) + @. ᶜρaχʲs_combined += ᶜρaʲ * ᶜχʲ + end - # Denominator: ρa⁰ = gs.ρ - Σ sgsʲ.ρa - ρa⁰_val = env_value(gs.ρ, sgsʲ -> sgsʲ.ρa, gs) + # Numerator: ρa⁰χ⁰ = (gs.ρχ) - (Σ ρaʲ * χʲ) + ρaχ⁰ = gs_ρχ_field - ᶜρaχʲs_combined + + # Denominator: ρa⁰ = gs.ρ - Σ ρaʲ + ᶜρa⁰_vals = ᶜρa⁰(gs, p) + end # Call the 5-argument specific function for regularized division - return specific( - ρaχ⁰, # ρaχ for environment - ρa⁰_val, # ρa for environment - getproperty(gs, ρχ_name), # Fallback ρχ is the grid-mean value - gs.ρ, # Fallback ρ is the grid-mean value - turbconv_model, + return @. lazy( + specific( + ρaχ⁰, # ρaχ for environment + ᶜρa⁰_vals, # ρa for environment + gs_ρχ_field, # Fallback ρχ is the grid-mean value + gs.ρ, # Fallback ρ is the grid-mean value + turbconv_model, + ), ) end +""" + ρa⁰(gs, p) + +Computes the environment area-weighted density (`ρa⁰`). + +This function uses the `ᶜenv_value` helper, which applies the domain +decomposition principle (`GridMean = Environment + Sum(Drafts)`) to calculate +the environment area-weighted density by subtracting the sum of all draft +subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`). + +Arguments: +- `gs`: The grid-scale state, which contains the grid-mean density `gs.ρ` and + the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX). +- `p`: The cache, containing precomputed quantities and turbconv_model. + +Returns: +- The area-weighted density (`ρa⁰`). +""" + +function ᶜρa⁰(gs, p) + turbconv_model = p.atmos.turbconv_model + + if turbconv_model isa PrognosticEDMFX + return ᶜenv_value(gs.ρ, sgsʲ -> sgsʲ.ρa, gs, turbconv_model) + elseif turbconv_model isa DiagnosticEDMFX + return ᶜenv_value(gs.ρ, ᶜρaʲ -> ᶜρaʲ, gs, turbconv_model, p) + else + return gs.ρ + end +end + +""" + ᶜspecific_tke(sgs⁰, gs, p) + +Computes the specific turbulent kinetic energy (`tke`) in the environment (`tke⁰`). + +This is a specialized helper that encapsulates the call to the regularized +`specific` function for the TKE variable. It provides `0` as the grid-scale +fallback value (`ρχ_fallback`) in the limit of small environmental area +fraction. + +Arguments: +- `sgs⁰`: The environment SGS state (`Y.c.sgs⁰`), containing `ρatke`. +- `gs`: The grid-scale state (`Y.c`), containing the grid-mean density `ρ`. +- `p`: The cache, containing precomputed quantities and turbconv_model. + +Returns: +- The specific TKE of the environment (`tke⁰`). +""" +function ᶜspecific_tke(sgs⁰, gs, p) + turbconv_model = p.atmos.turbconv_model + ᶜρa⁰_vals = ᶜρa⁰(gs, p) + + # no sgs weighting function needed for EDOnlyEDMFX + if turbconv_model isa EDOnlyEDMFX + return @. lazy(specific(sgs⁰.ρatke, ᶜρa⁰_vals)) + else + return @. lazy(specific( + sgs⁰.ρatke, # ρaχ for environment TKE + ᶜρa⁰_vals, # ρa for environment, now computed internally + 0, # Fallback ρχ is zero for TKE + gs.ρ, # Fallback ρ + turbconv_model, + )) + end +end + """ specific_env_mse(gs, p) @@ -364,22 +438,23 @@ Computes the specific moist static energy (`mse`) in the environment (`mse⁰`). This is a specialized helper function because `mse` is not a grid-scale prognostic variable. It first computes the grid-scale moist static energy density (`ρmse`) from other grid-scale quantities (`ρ`, total specific enthalpy `h_tot`, specific -kinetic energy `K`). It then uses the `env_value` helper to compute the environment's +kinetic energy `K`). It then uses the `ᶜenv_value` helper to compute the environment's portion of `ρmse` and `ρa` via domain decomposition, and finally calculates the specific value using the regularized `specific` function. Arguments: - `gs`: The grid-scale state (`Y.c`), containing `ρ` and `sgsʲs`. -- `p`: The cache, containing the `turbconv_model` +- `p`: The cache, containing the turbconv_model and precomputed quantities. Returns: - A `ClimaCore.Fields.Field` containing the specific moist static energy of the environment (`mse⁰`). """ function specific_env_mse(gs, p) + turbconv_model = p.atmos.turbconv_model + # Get necessary precomputed values from the cache `p` (; ᶜK, ᶜts) = p.precomputed # TODO: replace by on-the-fly computation - (; turbconv_model) = p.atmos thermo_params = CAP.thermodynamics_params(p.params) ᶜh_tot = @. lazy( TD.total_specific_enthalpy( @@ -390,67 +465,37 @@ function specific_env_mse(gs, p) ) # 1. Define the grid-scale moist static energy density `ρ * mse`. - grid_scale_ρmse = gs.ρ .* (ᶜh_tot .- ᶜK) + grid_scale_ρmse = @. lazy(gs.ρ * (ᶜh_tot - ᶜK)) # 2. Compute the environment's density-area-weighted mse (`ρa⁰mse⁰`). ρa⁰mse⁰ = p.scratch.ᶜtemp_scalar - @. ρa⁰mse⁰ = env_value(grid_scale_ρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, gs) + + if turbconv_model isa PrognosticEDMFX + ρa⁰mse⁰ .= ᶜenv_value( + grid_scale_ρmse, + sgsʲ -> sgsʲ.ρa * sgsʲ.mse, + gs, + turbconv_model, + ) + else # DiagnosticEDMFX + # For DiagnosticEDMFX, compute ρaʲ * mseʲ for each draft manually + n = n_mass_flux_subdomains(turbconv_model) + ᶜρaχʲs_combined_mse = p.scratch.ᶜtemp_scalar_2 + @. ᶜρaχʲs_combined_mse = 0 + for j in 1:n + ᶜρaʲ = p.precomputed.ᶜρaʲs.:($j) + ᶜmseʲ = p.precomputed.ᶜmseʲs.:($j) + @. ᶜρaχʲs_combined_mse += ᶜρaʲ * ᶜmseʲ + end + @. ρa⁰mse⁰ = grid_scale_ρmse - ᶜρaχʲs_combined_mse + end # 3. Compute the environment's density-area product (`ρa⁰`). - ρa⁰_val = @. lazy(ρa⁰(gs)) + ᶜρa⁰_vals = ᶜρa⁰(gs, p) # 4. Compute and return the final specific environment mse (`mse⁰`). return @. lazy( - specific(ρa⁰mse⁰, ρa⁰_val, grid_scale_ρmse, gs.ρ, turbconv_model), - ) -end - -""" - ρa⁰(gs) - -Computes the environment area-weighted density (`ρa⁰`). - -This function uses the `env_value` helper, which applies the domain -decomposition principle (`GridMean = Environment + Sum(Drafts)`) to calculate -the environment area-weighted density by subtracting the sum of all draft -subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`). - -Arguments: -- `gs`: The grid-scale state, which contains the grid-mean density `gs.ρ` and - the draft subdomain states `gs.sgsʲs`. - -Returns: -- The area-weighted density (`ρa⁰`). -""" -ρa⁰(gs) = env_value(gs.ρ, sgsʲ -> sgsʲ.ρa, gs) - -""" - specific_tke(sgs⁰, gs, turbconv_model) - -Computes the specific turbulent kinetic energy (`tke`) in the environment (`tke⁰`). - -This is a specialized helper that encapsulates the call to the regularized -`specific` function for the TKE variable. It provides `0` as the grid-scale -fallback value (`ρχ_fallback`) in the limit of small environmental area -fraction. - -Arguments: -- `sgs⁰`: The environment SGS state (`Y.c.sgs⁰`), containing `ρatke`. -- `gs`: The grid-scale state (`Y.c`), containing the grid-mean density `ρ`. -- `turbconv_model`: The turbulence convection model, for regularization parameters. - -Returns: -- The specific TKE of the environment (`tke⁰`). -""" -function specific_tke(sgs⁰, gs, turbconv_model) - ρa⁰_val = ρa⁰(gs) - - return specific( - sgs⁰.ρatke, # ρaχ for environment TKE - ρa⁰_val, # ρa for environment, now computed internally - 0, # Fallback ρχ is zero for TKE - gs.ρ, # Fallback ρ - turbconv_model, + specific(ρa⁰mse⁰, ᶜρa⁰_vals, grid_scale_ρmse, gs.ρ, turbconv_model), ) end From ddeeb7b865ac7a407446b1f3267389edc67905e5 Mon Sep 17 00:00:00 2001 From: costachris Date: Tue, 15 Jul 2025 12:32:36 -0700 Subject: [PATCH 4/7] lazy thermo vars --- src/cache/cloud_fraction.jl | 4 +- .../diagnostic_edmf_precomputed_quantities.jl | 20 ++--- .../precipitation_precomputed_quantities.jl | 34 ++++---- src/cache/precomputed_quantities.jl | 79 ++++++++++++++++--- .../prognostic_edmf_precomputed_quantities.jl | 30 +++---- src/diagnostics/Diagnostics.jl | 1 + .../les_sgs_models/smagorinsky_lilly.jl | 1 + src/prognostic_equations/advection.jl | 8 +- .../forcing/external_forcing.jl | 2 +- .../forcing/subsidence.jl | 6 +- .../implicit/implicit_tendency.jl | 2 +- .../implicit/manual_sparse_jacobian.jl | 24 +++--- src/utils/variable_manipulations.jl | 18 +++-- 13 files changed, 140 insertions(+), 89 deletions(-) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index ca80896def..446ce04bf9 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -62,8 +62,8 @@ NVTX.@annotate function set_cloud_fraction!( TD.PhasePartition(thermo_params, ᶜts).ice, ) else - q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) - q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) + q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) @. cloud_diagnostics_tuple = make_named_tuple(ifelse(q_liq + q_ice > 0, 1, 0), q_liq, q_ice) end diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 09a0272c99..aeb16fd19d 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -105,15 +105,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( turbconv_params = CAP.turbconv_params(params) ᶜts = p.precomputed.ᶜts #TODO replace - q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) - ᶜh_tot = @. lazy( - TD.total_specific_enthalpy( - thermo_params, - ᶜts, - ᶜe_tot, - ), - ) + q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) + ᶜe_tot = ᶜspecific(Y.c.ρe_tot, Y.c.ρ) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜe_tot)) ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1)) uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1)) @@ -338,8 +332,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( if microphysics_model isa Microphysics1Moment - q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) - q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) + q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) end thermo_params = CAP.thermodynamics_params(params) @@ -359,7 +353,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half)) # integral - q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) ᶜh_tot = @. lazy( TD.total_specific_enthalpy( thermo_params, @@ -1112,7 +1106,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed # Environment precipitation sources (to be applied to grid mean) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources( thermo_params, microphys_0m_params, diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index 3187f6dfb4..7b32b6a8da 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -112,10 +112,10 @@ function set_precipitation_velocities!( cm2p = CAP.microphysics_2m_params(p.params) thp = CAP.thermodynamics_params(p.params) - q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) - q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) - q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) - q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) + q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) + q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) + q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) # compute the precipitation terminal velocity [m/s] # TODO sedimentation of snow is based on the 1M scheme @@ -289,11 +289,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed - q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) - q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) - q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) - q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) + q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) + q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) + q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) + q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2 @@ -367,10 +367,10 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _) (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed - q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) - q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) - n_liq = @. lazy(specific(Y.c.ρn_liq, Y.c.ρ)) - n_rai = @. lazy(specific(Y.c.ρn_rai, Y.c.ρ)) + q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) + q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) + n_liq = ᶜspecific(Y.c.ρn_liq, Y.c.ρ) + n_rai = ᶜspecific(Y.c.ρn_rai, Y.c.ρ) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2 @@ -484,10 +484,10 @@ function set_precipitation_surface_fluxes!( sfc_ρ = @. lazy(int_ρ * int_J / sfc_J) # Constant extrapolation to surface, consistent with simple downwinding - ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) - ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) - ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) - ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + ᶜq_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) + ᶜq_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) + ᶜq_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) + ᶜq_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) sfc_qᵣ = Fields.Field( Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)), sfc_space, diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 47d93a05a1..fa5315c7f1 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -373,29 +373,84 @@ function thermo_state( return get_ts(ρ, p, θ, e_int, q_tot, q_pt) end +function ᶜthermo_state( + thermo_params; + ρ = nothing, + p = nothing, + θ = nothing, + e_int = nothing, + q_tot = nothing, + q_pt = nothing, +) + + get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, ::Nothing) = + TD.PhaseDry_ρθ(thermo_params, ρ, θ) + get_ts(ρ, ::Nothing, θ, ::Nothing, q_tot, ::Nothing) = + TD.PhaseEquil_ρθq(thermo_params, ρ, θ, q_tot) + get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, q_pt) = + TD.PhaseNonEquil_ρθq(thermo_params, ρ, θ, q_pt) + get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, ::Nothing) = + TD.PhaseDry_ρe(thermo_params, ρ, e_int) + get_ts(ρ, ::Nothing, ::Nothing, e_int, q_tot, ::Nothing) = + TD.PhaseEquil_ρeq( + thermo_params, + ρ, + e_int, + q_tot, + 3, + eltype(thermo_params)(0.003), + ) + get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, q_pt) = + TD.PhaseNonEquil(thermo_params, e_int, ρ, q_pt) + get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, ::Nothing) = + TD.PhaseDry_pθ(thermo_params, p, θ) + get_ts(::Nothing, p, θ, ::Nothing, q_tot, ::Nothing) = + TD.PhaseEquil_pθq(thermo_params, p, θ, q_tot) + get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, q_pt) = + TD.PhaseNonEquil_pθq(thermo_params, p, θ, q_pt) + get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, ::Nothing) = + TD.PhaseDry_pe(thermo_params, p, e_int) + get_ts(::Nothing, p, ::Nothing, e_int, q_tot, ::Nothing) = + TD.PhaseEquil_peq(thermo_params, p, e_int, q_tot) + get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, q_pt) = + TD.PhaseNonEquil_peq(thermo_params, p, e_int, q_pt) + + return @. lazy(get_ts(ρ, p, θ, e_int, q_tot, q_pt)) +end + function thermo_vars(moisture_model, microphysics_model, Y_c, K, Φ) - e_tot = materialize(@. lazy(specific(Y_c.ρe_tot, Y_c.ρ))) - energy_var = (; e_int = e_tot - K - Φ) + e_tot = ᶜspecific(Y_c.ρe_tot, Y_c.ρ) + energy_var = (; e_int = @. lazy(e_tot - K - Φ)) moisture_var = if moisture_model isa DryModel (;) elseif moisture_model isa EquilMoistModel - q_tot = materialize(@. lazy(specific(Y_c.ρq_tot, Y_c.ρ))) + q_tot = ᶜspecific(Y_c.ρq_tot, Y_c.ρ) (; q_tot) elseif moisture_model isa NonEquilMoistModel - q_tot = materialize(@. lazy(specific(Y_c.ρq_tot, Y_c.ρ))) - q_liq = materialize(@. lazy(specific(Y_c.ρq_liq, Y_c.ρ))) - q_ice = materialize(@. lazy(specific(Y_c.ρq_ice, Y_c.ρ))) - q_rai = materialize(@. lazy(specific(Y_c.ρq_rai, Y_c.ρ))) - q_sno = materialize(@. lazy(specific(Y_c.ρq_sno, Y_c.ρ))) - q_pt_args = (q_tot, q_liq + q_rai, q_ice + q_sno) - (; q_pt = TD.PhasePartition(q_pt_args...)) + q_tot = ᶜspecific(Y_c.ρq_tot, Y_c.ρ) + q_liq = ᶜspecific(Y_c.ρq_liq, Y_c.ρ) + q_ice = ᶜspecific(Y_c.ρq_ice, Y_c.ρ) + q_rai = ᶜspecific(Y_c.ρq_rai, Y_c.ρ) + q_sno = ᶜspecific(Y_c.ρq_sno, Y_c.ρ) + (; + q_pt = @. lazy( + TD.PhasePartition(q_tot, q_liq + q_rai, q_ice + q_sno), + ) + ) end return (; energy_var..., moisture_var...) end ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) = - thermo_state( + ᶜthermo_state( + thermo_params; + thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)..., + ρ, + ) + +ᶜts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) = + ᶜthermo_state( thermo_params; thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)..., ρ, @@ -467,7 +522,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ # TODO: We should think more about these increments before we use them. end - @. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ) + ᶜts .= ᶜts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ) @. ᶜp = TD.air_pressure(thermo_params, ᶜts) if turbconv_model isa PrognosticEDMFX diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index b555813d62..8f45afc606 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -201,7 +201,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( # ... and the first interior point for EDMFX ᶜq_totʲ. - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) ᶜq_tot_int_val = Fields.field_values(Fields.level(Base.materialize(ᶜq_tot), 1)) ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1)) @@ -220,31 +220,27 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( p.atmos.microphysics_model isa Microphysics1Moment # TODO - any better way to define the cloud and precip tracer flux? - ᶜq_liq = @.lazy(specific(Y.c.ρq_liq, Y.c.ρ)) - ᶜq_ice = @.lazy(specific(Y.c.ρq_ice, Y.c.ρ)) - ᶜq_rai = @.lazy(specific(Y.c.ρq_rai, Y.c.ρ)) - ᶜq_sno = @.lazy(specific(Y.c.ρq_sno, Y.c.ρ)) - ᶜq_liq_int_val = Fields.field_values( - Fields.level(Base.materialize(ᶜq_liq), 1), - ) + ᶜq_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) + ᶜq_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) + ᶜq_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) + ᶜq_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) + ᶜq_liq_int_val = + Fields.field_values(Fields.level(Base.materialize(ᶜq_liq), 1)) ᶜq_liqʲ_int_val = Fields.field_values(Fields.level(ᶜq_liqʲ, 1)) @. ᶜq_liqʲ_int_val = ᶜq_liq_int_val - ᶜq_ice_int_val = Fields.field_values( - Fields.level(Base.materialize(ᶜq_ice), 1), - ) + ᶜq_ice_int_val = + Fields.field_values(Fields.level(Base.materialize(ᶜq_ice), 1)) ᶜq_iceʲ_int_val = Fields.field_values(Fields.level(ᶜq_iceʲ, 1)) @. ᶜq_iceʲ_int_val = ᶜq_ice_int_val - ᶜq_rai_int_val = Fields.field_values( - Fields.level(Base.materialize(ᶜq_rai), 1), - ) + ᶜq_rai_int_val = + Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)) ᶜq_raiʲ_int_val = Fields.field_values(Fields.level(ᶜq_raiʲ, 1)) @. ᶜq_raiʲ_int_val = ᶜq_rai_int_val - ᶜq_sno_int_val = Fields.field_values( - Fields.level(Base.materialize(ᶜq_sno), 1), - ) + ᶜq_sno_int_val = + Fields.field_values(Fields.level(Base.materialize(ᶜq_sno), 1)) ᶜq_snoʲ_int_val = Fields.field_values(Fields.level(ᶜq_snoʲ, 1)) @. ᶜq_snoʲ_int_val = ᶜq_sno_int_val end diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl index e08c1de0e9..f4c8bb975f 100644 --- a/src/diagnostics/Diagnostics.jl +++ b/src/diagnostics/Diagnostics.jl @@ -60,6 +60,7 @@ import ..compute_gm_mixing_length! import ..horizontal_integral_at_boundary import ..ᶜρa⁰ import ..ᶜspecific_tke +import ..ᶜspecific_env_value # We need the abbreviations for symbols like curl, grad, and so on include(joinpath("..", "utils", "abbreviations.jl")) diff --git a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl index 69e7406a2c..962f4fb767 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -135,6 +135,7 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, sfc_conditions) = p.precomputed (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions (; ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) # Define operators ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 6aef68de3a..e9528bc9ce 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -194,13 +194,13 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ : p.precomputed.ᶠu³⁰ ) : nothing - ᶜρa⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y.c, p)) : lazy(Y.c.ρ)) : nothing + ᶜρa⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y.c, p)) : @. lazy(Y.c.ρ)) : nothing ᶜρ⁰ = if advect_tke if n > 0 - (; ᶜts⁰) = p.precomputed.ᶜts⁰ + (; ᶜts⁰) = p.precomputed @. lazy(TD.air_density(thermo_params, ᶜts⁰)) else - lazy(Y.c.ρ) + @. lazy(Y.c.ρ) end else nothing @@ -253,7 +253,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) end if !(p.atmos.moisture_model isa DryModel) && tracer_upwinding != Val(:none) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), tracer_upwinding) vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), Val(:none)) @. Yₜ.c.ρq_tot += vtt - vtt_central diff --git a/src/prognostic_equations/forcing/external_forcing.jl b/src/prognostic_equations/forcing/external_forcing.jl index 36385ec154..9a397eb22f 100644 --- a/src/prognostic_equations/forcing/external_forcing.jl +++ b/src/prognostic_equations/forcing/external_forcing.jl @@ -405,7 +405,7 @@ function external_forcing_tendency!( ᶜh_tot, Val{:first_order}(), ) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) subsidence!( Yₜ.c.ρq_tot, Y.c.ρ, diff --git a/src/prognostic_equations/forcing/subsidence.jl b/src/prognostic_equations/forcing/subsidence.jl index 8140a55f01..3cf9428618 100644 --- a/src/prognostic_equations/forcing/subsidence.jl +++ b/src/prognostic_equations/forcing/subsidence.jl @@ -107,11 +107,11 @@ function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) specific(Y.c.ρe_tot, Y.c.ρ), ), ) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}()) subsidence!(Yₜ.c.ρq_tot, Y.c.ρ, ᶠsubsidence³, ᶜq_tot, Val{:first_order}()) if moisture_model isa NonEquilMoistModel - ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + ᶜq_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) subsidence!( Yₜ.c.ρq_liq, Y.c.ρ, @@ -119,7 +119,7 @@ function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) ᶜq_liq, Val{:first_order}(), ) - ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + ᶜq_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) subsidence!( Yₜ.c.ρq_ice, Y.c.ρ, diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 60f4f90b68..254c7b38dc 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -156,7 +156,7 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜh_tot, dt, Val(:none)) @. Yₜ.c.ρe_tot += vtt if !(moisture_model isa DryModel) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜq_tot, dt, Val(:none)) @. Yₜ.c.ρq_tot += vtt end diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index 129b6a4f21..3c54890c8d 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -436,7 +436,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ), ) else - @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜspecific(Y.c.ρq_tot, Y.c.ρ) end if use_derivative(topography_flag) @@ -461,9 +461,9 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ᶠinterp_matrix() ) @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) - ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) + ᶜe_tot = ᶜspecific(Y.c.ρe_tot, Y.c.ρ) if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] @. ∂ᶠu₃_err_∂ᶜρq_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( @@ -588,7 +588,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] ∂ᶜρq_tot_err_∂ᶜρ = matrix[@name(c.ρq_tot), @name(c.ρ)] @. ∂ᶜρe_tot_err_∂ᶜρq_tot += @@ -955,14 +955,16 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ) / abs2(TD.cv_m(thermo_params, ᶜts)) if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) @. ∂ᶜρe_tot_err_∂ᶜρq_tot += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(( - ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * - (cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * ᶜq_tot) - )) + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow(( + ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + + ᶜ∂kappa_m∂q_tot * ( + cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * ᶜq_tot + ) + )) else ᶜq_tot = lazy(0) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 26fbafe362..82663ea166 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -45,6 +45,8 @@ Arguments: """ specific(ρχ, ρ) = ρχ / ρ +ᶜspecific(ρχ, ρ) = @. lazy(ρχ / ρ) + function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) # TODO: Replace turbconv_model struct by parameters, and include a_half in # parameters, not in config @@ -320,14 +322,14 @@ function ᶜspecific_env_value(::Val{χ_name}, gs, p) where {χ_name} gs_ρχ_field = getproperty(gs, ρχ_name) if turbconv_model isa PrognosticEDMFX - # Numerator: ρa⁰χ⁰ = (gs.ρχ) - (Σ sgsʲ.ρa * sgsʲ.χ) - ρaχ⁰ = ᶜenv_value( + # Numerator: ρa⁰χ⁰ = ρχ - (Σ sgsʲ.ρa * sgsʲ.χ) + ᶜρaχ⁰ = ᶜenv_value( gs_ρχ_field, sgsʲ -> getproperty(sgsʲ, :ρa) * getproperty(sgsʲ, χ_name), gs, turbconv_model, ) - # Denominator: ρa⁰ = gs.ρ - Σ sgsʲ.ρa + # Denominator: ρa⁰ = ρ - Σ sgsʲ.ρa ᶜρa⁰_vals = ᶜρa⁰(gs, p) else # DiagnosticEDMFX # For DiagnosticEDMFX, we need to compute ρaʲ * χʲ for each draft @@ -344,17 +346,17 @@ function ᶜspecific_env_value(::Val{χ_name}, gs, p) where {χ_name} @. ᶜρaχʲs_combined += ᶜρaʲ * ᶜχʲ end - # Numerator: ρa⁰χ⁰ = (gs.ρχ) - (Σ ρaʲ * χʲ) - ρaχ⁰ = gs_ρχ_field - ᶜρaχʲs_combined + # Numerator: ρa⁰χ⁰ = ρχ - (Σ ρaʲ * χʲ) + ᶜρaχ⁰ = gs_ρχ_field - ᶜρaχʲs_combined - # Denominator: ρa⁰ = gs.ρ - Σ ρaʲ + # Denominator: ρa⁰ = ρ - Σ ρaʲ ᶜρa⁰_vals = ᶜρa⁰(gs, p) end # Call the 5-argument specific function for regularized division return @. lazy( specific( - ρaχ⁰, # ρaχ for environment + ᶜρaχ⁰, # ρaχ for environment ᶜρa⁰_vals, # ρa for environment gs_ρχ_field, # Fallback ρχ is the grid-mean value gs.ρ, # Fallback ρ is the grid-mean value @@ -418,7 +420,7 @@ function ᶜspecific_tke(sgs⁰, gs, p) # no sgs weighting function needed for EDOnlyEDMFX if turbconv_model isa EDOnlyEDMFX - return @. lazy(specific(sgs⁰.ρatke, ᶜρa⁰_vals)) + return ᶜspecific(sgs⁰.ρatke, ᶜρa⁰_vals) else return @. lazy(specific( sgs⁰.ρatke, # ρaχ for environment TKE From d4ae75d62c01db0bbf4dde19b2dd488b418c006a Mon Sep 17 00:00:00 2001 From: costachris Date: Thu, 17 Jul 2025 12:09:11 -0700 Subject: [PATCH 5/7] Simplify function args --- src/cache/cloud_fraction.jl | 3 +- .../diagnostic_edmf_precomputed_quantities.jl | 4 +- .../precipitation_precomputed_quantities.jl | 2 +- src/cache/precomputed_quantities.jl | 38 ++-- .../prognostic_edmf_precomputed_quantities.jl | 34 ++-- src/diagnostics/edmfx_diagnostics.jl | 12 +- .../microphysics/precipitation.jl | 2 +- src/prognostic_equations/advection.jl | 4 +- src/prognostic_equations/edmfx_entr_detr.jl | 12 +- src/prognostic_equations/edmfx_sgs_flux.jl | 34 ++-- src/prognostic_equations/edmfx_tke.jl | 4 +- src/prognostic_equations/hyperdiffusion.jl | 6 +- .../implicit/manual_sparse_jacobian.jl | 5 +- src/utils/variable_manipulations.jl | 166 ++++++++---------- 14 files changed, 150 insertions(+), 176 deletions(-) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index 446ce04bf9..589ce0fdc6 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -190,7 +190,7 @@ NVTX.@annotate function set_cloud_fraction!( (; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed (; ᶜρʲs, ᶜtsʲs) = p.precomputed (; turbconv_model) = p.atmos - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) # TODO - we should make this default when using diagnostic edmf # environment @@ -304,6 +304,7 @@ function quad_loop( FT = eltype(x1_hat) @assert(x1_hat >= FT(0)) @assert(x2_hat >= FT(0)) + # note: ᶜthermo_state is used as a pointwise function here _ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat) hc = TD.has_condensate(thermo_params, _ts) diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index aeb16fd19d..1adad5c124 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -361,7 +361,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( specific(Y.c.ρe_tot, Y.c.ρ), ), ) - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) for i in 2:Spaces.nlevels(axes(Y.c)) ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i)) @@ -1029,7 +1029,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! @. ᶜprandtl_nvec = turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm) - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) ᶜtke_exch = p.scratch.ᶜtemp_scalar_2 @. ᶜtke_exch = 0 diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index 7b32b6a8da..b267ace613 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -262,7 +262,7 @@ function set_precipitation_cache!( (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed (; ᶜts⁰, ᶜtsʲs) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) n = n_mass_flux_subdomains(p.atmos.turbconv_model) diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index fa5315c7f1..5b198a3132 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -43,7 +43,7 @@ function implicit_precomputed_quantities(Y, atmos) ᶜts = similar(Y.c, TST), ᶜp = similar(Y.c, FT), ) - sgs_quantities = turbconv_model isa AbstractEDMF ? (;) : (;) + sgs_quantities = (;) prognostic_sgs_quantities = turbconv_model isa PrognosticEDMFX ? (; @@ -373,6 +373,7 @@ function thermo_state( return get_ts(ρ, p, θ, e_int, q_tot, q_pt) end +const FieldOrValue = Union{Fields.Field, Base.AbstractBroadcasted, Real} function ᶜthermo_state( thermo_params; ρ = nothing, @@ -383,15 +384,15 @@ function ᶜthermo_state( q_pt = nothing, ) - get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, ::Nothing) = + get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, ::Nothing, ::Nothing) where {T <: FieldOrValue} = TD.PhaseDry_ρθ(thermo_params, ρ, θ) - get_ts(ρ, ::Nothing, θ, ::Nothing, q_tot, ::Nothing) = + get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, q_tot::T, ::Nothing) where {T <: FieldOrValue} = TD.PhaseEquil_ρθq(thermo_params, ρ, θ, q_tot) - get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, q_pt) = + get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, ::Nothing, q_pt) where {T <: FieldOrValue} = TD.PhaseNonEquil_ρθq(thermo_params, ρ, θ, q_pt) - get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, ::Nothing) = + get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, ::Nothing, ::Nothing) where {T <: FieldOrValue} = TD.PhaseDry_ρe(thermo_params, ρ, e_int) - get_ts(ρ, ::Nothing, ::Nothing, e_int, q_tot, ::Nothing) = + get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, q_tot::T, ::Nothing) where {T <: FieldOrValue} = TD.PhaseEquil_ρeq( thermo_params, ρ, @@ -400,19 +401,19 @@ function ᶜthermo_state( 3, eltype(thermo_params)(0.003), ) - get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, q_pt) = + get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, ::Nothing, q_pt) where {T <: FieldOrValue} = TD.PhaseNonEquil(thermo_params, e_int, ρ, q_pt) - get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, ::Nothing) = + get_ts(::Nothing, p::T, θ::T, ::Nothing, ::Nothing, ::Nothing) where {T <: FieldOrValue} = TD.PhaseDry_pθ(thermo_params, p, θ) - get_ts(::Nothing, p, θ, ::Nothing, q_tot, ::Nothing) = + get_ts(::Nothing, p::T, θ::T, ::Nothing, q_tot::T, ::Nothing) where {T <: FieldOrValue} = TD.PhaseEquil_pθq(thermo_params, p, θ, q_tot) - get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, q_pt) = + get_ts(::Nothing, p::T, θ::T, ::Nothing, ::Nothing, q_pt) where {T <: FieldOrValue} = TD.PhaseNonEquil_pθq(thermo_params, p, θ, q_pt) - get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, ::Nothing) = + get_ts(::Nothing, p::T, ::Nothing, e_int::T, ::Nothing, ::Nothing) where {T <: FieldOrValue} = TD.PhaseDry_pe(thermo_params, p, e_int) - get_ts(::Nothing, p, ::Nothing, e_int, q_tot, ::Nothing) = + get_ts(::Nothing, p::T, ::Nothing, e_int::T, q_tot::T, ::Nothing) where {T <: FieldOrValue} = TD.PhaseEquil_peq(thermo_params, p, e_int, q_tot) - get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, q_pt) = + get_ts(::Nothing, p::T, ::Nothing, e_int::T, ::Nothing, q_pt) where {T <: FieldOrValue} = TD.PhaseNonEquil_peq(thermo_params, p, e_int, q_pt) return @. lazy(get_ts(ρ, p, θ, e_int, q_tot, q_pt)) @@ -442,13 +443,6 @@ function thermo_vars(moisture_model, microphysics_model, Y_c, K, Φ) return (; energy_var..., moisture_var...) end -ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) = - ᶜthermo_state( - thermo_params; - thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)..., - ρ, - ) - ᶜts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) = ᶜthermo_state( thermo_params; @@ -456,8 +450,8 @@ ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) = ρ, ) -ts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) = - thermo_state( +ᶜts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) = + ᶜthermo_state( thermo_params; thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)..., p, diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 8f45afc606..8c2bb07357 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -24,22 +24,22 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( (; ᶜp, ᶜK) = p.precomputed (; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 - ᶜmse⁰ .= specific_env_mse(Y.c, p) + ᶜmse⁰ .= ᶜspecific_env_mse(Y, p) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment - ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) - ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) - ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) - ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) @. ᶜts⁰ = TD.PhaseNonEquil_phq( thermo_params, ᶜp, @@ -369,8 +369,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ᶜdz = Fields.Δz_field(axes(Y.c)) ᶜlg = Fields.local_geometry_field(Y.c) ᶠlg = Fields.local_geometry_field(Y.f) - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) ᶜvert_div = p.scratch.ᶜtemp_scalar ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2 @@ -455,7 +455,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos (; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed # First order approximation: Use environmental mean fields. @. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat @. ᶜgradᵥ_θ_liq_ice⁰ = ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat @@ -567,7 +567,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ) end # sources from the environment - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰) return nothing end @@ -657,11 +657,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation end # Precipitation sources and sinks from the environment - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) - ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) - ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) - ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) - ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰)) compute_precipitation_sources!( ᶜSᵖ, diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index 54a41a6078..6d392b96ae 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -633,7 +633,7 @@ compute_aren!(_, _, _, _, turbconv_model::T) where {T} = function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - ᶜρa⁰_vals = ᶜρa⁰(state.c, cache) + ᶜρa⁰_vals = ᶜρa⁰(state, cache) if isnothing(out) return draft_area.( ᶜρa⁰_vals, @@ -953,7 +953,7 @@ function compute_clwen!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) - ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), state.c, cache) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), state, cache) if isnothing(out) return Base.materialize(ᶜq_liq⁰) else @@ -1018,7 +1018,7 @@ function compute_clien!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) - ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), state.c, cache) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), state, cache) if isnothing(out) return Base.materialize(ᶜq_ice⁰) else @@ -1067,7 +1067,7 @@ function compute_husraen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), state.c, cache) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), state, cache) if isnothing(out) return Base.materialize(ᶜq_rai⁰) else @@ -1116,7 +1116,7 @@ function compute_hussnen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), state.c, cache) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), state, cache) if isnothing(out) return Base.materialize(ᶜq_sno⁰) else @@ -1150,7 +1150,7 @@ function compute_tke!( time, turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX}, ) - ᶜtke = ᶜspecific_tke(state.c.sgs⁰, state.c, cache) + ᶜtke = ᶜspecific_tke(state, cache) if isnothing(out) return Base.materialize(ᶜtke) else diff --git a/src/parameterized_tendencies/microphysics/precipitation.jl b/src/parameterized_tendencies/microphysics/precipitation.jl index b642880398..f5179d3a06 100644 --- a/src/parameterized_tendencies/microphysics/precipitation.jl +++ b/src/parameterized_tendencies/microphysics/precipitation.jl @@ -150,7 +150,7 @@ function precipitation_tendency!( # Source terms from EDMFX environment (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰) = p.precomputed - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) # Update from environment precipitation and cloud formation sources/sinks @. Yₜ.c.ρq_liq += ᶜρa⁰_vals * ᶜSqₗᵖ⁰ diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index e9528bc9ce..53be62196b 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -194,7 +194,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ : p.precomputed.ᶠu³⁰ ) : nothing - ᶜρa⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y.c, p)) : @. lazy(Y.c.ρ)) : nothing + ᶜρa⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y, p)) : @. lazy(Y.c.ρ)) : nothing ᶜρ⁰ = if advect_tke if n > 0 (; ᶜts⁰) = p.precomputed @@ -205,7 +205,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) else nothing end - ᶜtke⁰ = advect_tke ? (ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)) : nothing + ᶜtke⁰ = advect_tke ? (ᶜspecific_tke(Y, p)) : nothing ᶜa_scalar = p.scratch.ᶜtemp_scalar ᶜω³ = p.scratch.ᶜtemp_CT3 ᶠω¹² = p.scratch.ᶠtemp_CT12 diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index f5440ab001..0f33405d7c 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -531,13 +531,13 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF (; ᶠu₃⁰) = p.precomputed ᶜmse⁰ = p.scratch.ᶜtemp_scalar - ᶜmse⁰ .= specific_env_mse(Y.c, p) + ᶜmse⁰ .= ᶜspecific_env_mse(Y, p) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment - ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) - ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) - ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) - ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) end for j in 1:n @@ -547,7 +547,7 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF ᶜmseʲ = Y.c.sgsʲs.:($j).mse ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. Yₜ.c.sgsʲs.:($$j).ρa += Y.c.sgsʲs.:($$j).ρa * (ᶜentrʲ - ᶜdetrʲ) diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index fde18ea572..074280f3c1 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -45,7 +45,7 @@ function edmfx_sgs_mass_flux_tendency!( (; ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) ᶜρ⁰ = @. TD.air_density(thermo_params, ᶜts⁰) - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J @@ -81,7 +81,7 @@ function edmfx_sgs_mass_flux_tendency!( @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 - ᶜmse⁰ .= specific_env_mse(Y.c, p) + ᶜmse⁰ .= ᶜspecific_env_mse(Y, p) @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰_vals, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, @@ -109,7 +109,7 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_tot += vtt end # Add the environment fluxes - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * @@ -128,10 +128,10 @@ function edmfx_sgs_mass_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) - ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) - ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) - ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) - ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) # Liquid, ice, rain and snow specific humidity fluxes for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @@ -334,7 +334,7 @@ function edmfx_sgs_mass_flux_tendency!( # CAP.thermodynamics_params(p.params), # p.precomputed.ᶜts⁰, # ) - # ᶜmse⁰ = @.lazy(specific_env_mse(Y.c, p)) + # ᶜmse⁰ = @.lazy(ᶜspecific_env_mse(Y, p)) # @. ᶜa_scalar = # (ᶜmse⁰ + p.precomputed.ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) # vtt = vertical_transport( @@ -405,8 +405,8 @@ function edmfx_sgs_diffusive_flux_tendency!( (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed (; ᶜmixing_length, ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar @@ -421,7 +421,7 @@ function edmfx_sgs_diffusive_flux_tendency!( ) ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2 - ᶜmse⁰ .= specific_env_mse(Y.c, p) + ᶜmse⁰ .= ᶜspecific_env_mse(Y, p) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜmse⁰ + ᶜK⁰))) if use_prognostic_tke(turbconv_model) # Turbulent TKE transport (diffusion) @@ -450,7 +450,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜq_tot⁰))) @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion @. Yₜ.c.ρ -= ᶜρχₜ_diffusion # Effect of moisture diffusion on (moist) air mass @@ -459,10 +459,10 @@ function edmfx_sgs_diffusive_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) - ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p) - ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p) - ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p) - ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) # Liquid, ice, rain and snow specific humidity diffusion α_vert_diff_tracer = CAP.α_vert_diff_tracer(params) @@ -515,7 +515,7 @@ function edmfx_sgs_diffusive_flux_tendency!( (; ᶜu, ᶜmixing_length, ᶜts) = p.precomputed (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} ᶠρaK_h = p.scratch.ᶠtemp_scalar diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index 09d7f5ca2a..b9a22e36f4 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -50,7 +50,7 @@ function edmfx_tke_tendency!( FT = eltype(p.params) - ᶜρa⁰_vals = turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y.c, p) : Y.c.ρ + ᶜρa⁰_vals = turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y, p) : Y.c.ρ nh_pressure3_buoyʲs = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs @@ -77,7 +77,7 @@ function edmfx_tke_tendency!( # buoyancy production @. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰_vals * ᶜK_h * ᶜlinear_buoygrad - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) # entrainment and detraiment # using ᶜu⁰ and local geometry results in allocation diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index 331dc914e2..772ee69a7f 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -117,7 +117,7 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) if diffuse_tke - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) (; ᶜ∇²tke⁰) = p.hyperdiff @. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰)) end @@ -150,11 +150,11 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) point_type = eltype(Fields.coordinate_field(Y.c)) (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX - ᶜρa⁰_vals = ᶜρa⁰(Y.c, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end if use_prognostic_tke(turbconv_model) - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) (; ᶜ∇²tke⁰) = p.hyperdiff end diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index 3c54890c8d..847d18feed 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -625,10 +625,9 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) (; dt) = p turbconv_model = p.atmos.turbconv_model ᶜmixing_length = p.precomputed.ᶜmixing_length - ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p) + ᶜtke⁰ = ᶜspecific_tke(Y, p) ᶜρa⁰_vals = - p.atmos.turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y.c, p) : - Y.c.ρ + p.atmos.turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y, p) : Y.c.ρ ᶜρatke⁰ = Y.c.sgs⁰.ρatke @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) = diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 82663ea166..30807bbb6d 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -249,7 +249,8 @@ draft_sum(f, sgsʲs) = mapreduce_with_init(f, +, sgsʲs) ᶜenv_value(grid_scale_value, f_draft, gs, turbconv_model) Computes the value of a quantity `ρaχ` in the environment subdomain by subtracting -the sum of its values in all draft subdomains from the grid-scale value. +the sum of its values in all draft subdomains from the grid-scale value. Only available +for environmental area (ᶜρa⁰) for DiagnosticEDMFX. This is based on the domain decomposition principle for density-area weighted quantities: `GridMean(ρχ) = Env(ρaχ) + Sum(Drafts(ρaχ))`. @@ -261,7 +262,7 @@ The function handles both PrognosticEDMFX and DiagnosticEDMFX models: Arguments: - `grid_scale_value`: The `ρa`-weighted grid-scale value of the quantity. - `f_draft`: A function that extracts the corresponding value from a draft subdomain state. -- `gs`: The grid-scale state, which contains the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX). +- `gs`: The grid-scale iteration object, which contains the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX) from the state `Y.c`, or `ᶜρaʲs` in the cache for DiagnosticEDMFX. - `turbconv_model`: The turbulence convection model, used to determine how to access draft data. """ function ᶜenv_value( @@ -278,24 +279,12 @@ function ᶜenv_value( f_draft, gs, turbconv_model::DiagnosticEDMFX, - p, ) - # For DiagnosticEDMFX, we need to access precomputed quantities from p - (; ᶜρaʲs) = p.precomputed - n = n_mass_flux_subdomains(turbconv_model) - - ᶜdraft_sum_diag = p.scratch.ᶜtemp_scalar_4 - @. ᶜdraft_sum_diag = 0 - - for j in 1:n - ᶜρaʲ = ᶜρaʲs.:($j) - @. ᶜdraft_sum_diag += f_draft(ᶜρaʲ) - end - return @. lazy(grid_scale_value - ᶜdraft_sum_diag) + return @. lazy(grid_scale_value - draft_sum(f_draft, gs)) end """ - ᶜspecific_env_value(::Val{χ_name}, gs, p) + ᶜspecific_env_value(::Val{χ_name}, Y, p) Calculates the specific value of a quantity `χ` in the environment (`χ⁰`). @@ -307,66 +296,61 @@ environment area fraction is very small. Arguments: - `::Val{χ_name}`: A `Val` type containing the symbol for the specific quantity `χ` (e.g., `Val(:h_tot)`, `Val(:q_tot)`). -- `gs`: The grid-scale state, containing grid-mean and draft subdomain states. +- `Y`: The state, containing grid-mean and draft subdomain states. - `p`: The cache, containing precomputed quantities and turbconv_model. Returns: - The specific value of the quantity `χ` in the environment. """ -function ᶜspecific_env_value(::Val{χ_name}, gs, p) where {χ_name} +function ᶜspecific_env_value(::Val{χ_name}, Y, p) where {χ_name} turbconv_model = p.atmos.turbconv_model # Grid-scale density-weighted variable name, e.g., :ρq_tot ρχ_name = Symbol(:ρ, χ_name) - gs_ρχ_field = getproperty(gs, ρχ_name) + ᶜρχ = getproperty(Y.c, ρχ_name) + # environment density-area-weighted mse (`ρa⁰χ⁰`). + # Numerator: ρa⁰χ⁰ = ρχ - (Σ ρaʲ * χʲ) if turbconv_model isa PrognosticEDMFX # Numerator: ρa⁰χ⁰ = ρχ - (Σ sgsʲ.ρa * sgsʲ.χ) ᶜρaχ⁰ = ᶜenv_value( - gs_ρχ_field, + ᶜρχ, sgsʲ -> getproperty(sgsʲ, :ρa) * getproperty(sgsʲ, χ_name), - gs, + Y.c, turbconv_model, ) - # Denominator: ρa⁰ = ρ - Σ sgsʲ.ρa - ᶜρa⁰_vals = ᶜρa⁰(gs, p) - else # DiagnosticEDMFX - # For DiagnosticEDMFX, we need to compute ρaʲ * χʲ for each draft - # Get the specific quantity values for all drafts + + elseif turbconv_model isa DiagnosticEDMFX || turbconv_model isa EDOnlyEDMFX ᶜχʲs = getproperty(p.precomputed, Symbol(:ᶜ, χ_name, :ʲs)) n = n_mass_flux_subdomains(turbconv_model) - # Create combined ρaʲ * χʲ values for each draft - ᶜρaχʲs_combined = p.scratch.ᶜtemp_scalar_3 - @. ᶜρaχʲs_combined = 0 + # Σ ρaʲ * χʲ + ᶜρaχʲs_sum = p.scratch.ᶜtemp_scalar + @. ᶜρaχʲs_sum = 0 for j in 1:n ᶜρaʲ = p.precomputed.ᶜρaʲs.:($j) ᶜχʲ = ᶜχʲs.:($j) - @. ᶜρaχʲs_combined += ᶜρaʲ * ᶜχʲ + @. ᶜρaχʲs_sum += ᶜρaʲ * ᶜχʲ end - # Numerator: ρa⁰χ⁰ = ρχ - (Σ ρaʲ * χʲ) - ᶜρaχ⁰ = gs_ρχ_field - ᶜρaχʲs_combined - - # Denominator: ρa⁰ = ρ - Σ ρaʲ - ᶜρa⁰_vals = ᶜρa⁰(gs, p) + ᶜρaχ⁰ = @. lazy(ᶜρχ - ᶜρaχʲs_sum) end - # Call the 5-argument specific function for regularized division - return @. lazy( - specific( - ᶜρaχ⁰, # ρaχ for environment - ᶜρa⁰_vals, # ρa for environment - gs_ρχ_field, # Fallback ρχ is the grid-mean value - gs.ρ, # Fallback ρ is the grid-mean value - turbconv_model, - ), - ) + # Denominator: ρa⁰ = ρ - Σ ρaʲ + ᶜρa⁰_vals = ᶜρa⁰(Y, p) + + return @. lazy(specific( + ᶜρaχ⁰, # ρaχ for environment + ᶜρa⁰_vals, # ρa for environment + ᶜρχ, # Fallback ρχ is the grid-mean value + Y.c.ρ, # Fallback ρ is the grid-mean value + turbconv_model, + )) end """ - ρa⁰(gs, p) + ρa⁰(Y, p) Computes the environment area-weighted density (`ρa⁰`). @@ -376,28 +360,29 @@ the environment area-weighted density by subtracting the sum of all draft subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`). Arguments: -- `gs`: The grid-scale state, which contains the grid-mean density `gs.ρ` and - the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX). +- `Y`: The model state, which contains the grid-mean density `Y.c.ρ` and + the draft subdomain states `Y.c.sgsʲs` (for PrognosticEDMFX). - `p`: The cache, containing precomputed quantities and turbconv_model. Returns: - The area-weighted density (`ρa⁰`). """ -function ᶜρa⁰(gs, p) +function ᶜρa⁰(Y, p) turbconv_model = p.atmos.turbconv_model if turbconv_model isa PrognosticEDMFX - return ᶜenv_value(gs.ρ, sgsʲ -> sgsʲ.ρa, gs, turbconv_model) + return ᶜenv_value(Y.c.ρ, sgsʲ -> sgsʲ.ρa, Y.c, turbconv_model) elseif turbconv_model isa DiagnosticEDMFX - return ᶜenv_value(gs.ρ, ᶜρaʲ -> ᶜρaʲ, gs, turbconv_model, p) + (; ᶜρaʲs) = p.precomputed + return ᶜenv_value(Y.c.ρ, ᶜρaʲ -> ᶜρaʲ, ᶜρaʲs, turbconv_model) else - return gs.ρ + return Y.c.ρ end end """ - ᶜspecific_tke(sgs⁰, gs, p) + ᶜspecific_tke(Y, p) Computes the specific turbulent kinetic energy (`tke`) in the environment (`tke⁰`). @@ -407,33 +392,36 @@ fallback value (`ρχ_fallback`) in the limit of small environmental area fraction. Arguments: -- `sgs⁰`: The environment SGS state (`Y.c.sgs⁰`), containing `ρatke`. -- `gs`: The grid-scale state (`Y.c`), containing the grid-mean density `ρ`. +- `Y`: The state, containing the grid-mean density `ρ` and the environment SGS state `Y.c.sgs⁰`. - `p`: The cache, containing precomputed quantities and turbconv_model. Returns: - The specific TKE of the environment (`tke⁰`). """ -function ᶜspecific_tke(sgs⁰, gs, p) +function ᶜspecific_tke(Y, p) turbconv_model = p.atmos.turbconv_model - ᶜρa⁰_vals = ᶜρa⁰(gs, p) + ᶜρa⁰_vals = ᶜρa⁰(Y, p) + + sgs⁰ = Y.c.sgs⁰ # no sgs weighting function needed for EDOnlyEDMFX if turbconv_model isa EDOnlyEDMFX return ᶜspecific(sgs⁰.ρatke, ᶜρa⁰_vals) else - return @. lazy(specific( - sgs⁰.ρatke, # ρaχ for environment TKE - ᶜρa⁰_vals, # ρa for environment, now computed internally - 0, # Fallback ρχ is zero for TKE - gs.ρ, # Fallback ρ - turbconv_model, - )) + return @. lazy( + specific( + sgs⁰.ρatke, # ρaχ for environment TKE + ᶜρa⁰_vals, # ρa for environment, now computed internally + 0, # Fallback ρχ is zero for TKE + Y.c.ρ, # Fallback ρ + turbconv_model, + ), + ) end end """ - specific_env_mse(gs, p) + ᶜspecific_env_mse(Y, p) Computes the specific moist static energy (`mse`) in the environment (`mse⁰`). @@ -445,60 +433,52 @@ portion of `ρmse` and `ρa` via domain decomposition, and finally calculates th value using the regularized `specific` function. Arguments: -- `gs`: The grid-scale state (`Y.c`), containing `ρ` and `sgsʲs`. +- `Y`: The state containing `Y.c.ρ` and `Y.c.sgsʲs` (for PrognosticEDMFX). - `p`: The cache, containing the turbconv_model and precomputed quantities. Returns: - A `ClimaCore.Fields.Field` containing the specific moist static energy of the environment (`mse⁰`). """ -function specific_env_mse(gs, p) +function ᶜspecific_env_mse(Y, p) turbconv_model = p.atmos.turbconv_model - - # Get necessary precomputed values from the cache `p` - (; ᶜK, ᶜts) = p.precomputed # TODO: replace by on-the-fly computation + (; ᶜK, ᶜts) = p.precomputed thermo_params = CAP.thermodynamics_params(p.params) ᶜh_tot = @. lazy( TD.total_specific_enthalpy( thermo_params, ᶜts, - specific(gs.ρe_tot, gs.ρ), + specific(Y.c.ρe_tot, Y.c.ρ), ), ) - # 1. Define the grid-scale moist static energy density `ρ * mse`. - grid_scale_ρmse = @. lazy(gs.ρ * (ᶜh_tot - ᶜK)) + # grid-scale moist static energy density `ρ * mse`. + ᶜρmse = @. lazy(Y.c.ρ * (ᶜh_tot - ᶜK)) - # 2. Compute the environment's density-area-weighted mse (`ρa⁰mse⁰`). - ρa⁰mse⁰ = p.scratch.ᶜtemp_scalar + # environment density-area-weighted mse (`ρa⁰mse⁰`). + # Numerator: ρa⁰mse⁰ = ρmse - (Σ ρaʲ * mseʲ) if turbconv_model isa PrognosticEDMFX - ρa⁰mse⁰ .= ᶜenv_value( - grid_scale_ρmse, - sgsʲ -> sgsʲ.ρa * sgsʲ.mse, - gs, - turbconv_model, - ) - else # DiagnosticEDMFX - # For DiagnosticEDMFX, compute ρaʲ * mseʲ for each draft manually + ρa⁰mse⁰ = + ᶜenv_value(ᶜρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, Y.c, turbconv_model) + elseif turbconv_model isa DiagnosticEDMFX || turbconv_model isa EDOnlyEDMFX + n = n_mass_flux_subdomains(turbconv_model) - ᶜρaχʲs_combined_mse = p.scratch.ᶜtemp_scalar_2 - @. ᶜρaχʲs_combined_mse = 0 + ᶜρamseʲ_sum = p.scratch.ᶜtemp_scalar + @. ᶜρamseʲ_sum = 0 + # Numerator: ρa⁰mse⁰ = ρmse - (Σ ρaʲ * mseʲ) for j in 1:n ᶜρaʲ = p.precomputed.ᶜρaʲs.:($j) ᶜmseʲ = p.precomputed.ᶜmseʲs.:($j) - @. ᶜρaχʲs_combined_mse += ᶜρaʲ * ᶜmseʲ + @. ᶜρamseʲ_sum += ᶜρaʲ * ᶜmseʲ end - @. ρa⁰mse⁰ = grid_scale_ρmse - ᶜρaχʲs_combined_mse + ρa⁰mse⁰ = @. lazy(ᶜρmse - ᶜρamseʲ_sum) end - # 3. Compute the environment's density-area product (`ρa⁰`). - ᶜρa⁰_vals = ᶜρa⁰(gs, p) + # Denominator: ρa⁰ = ρ - Σ ρaʲ + ᶜρa⁰_vals = ᶜρa⁰(Y, p) - # 4. Compute and return the final specific environment mse (`mse⁰`). - return @. lazy( - specific(ρa⁰mse⁰, ᶜρa⁰_vals, grid_scale_ρmse, gs.ρ, turbconv_model), - ) + return @. lazy(specific(ρa⁰mse⁰, ᶜρa⁰_vals, ᶜρmse, Y.c.ρ, turbconv_model)) end """ From f763b36b23f3a0627037bec6a1431b797b6f4ed6 Mon Sep 17 00:00:00 2001 From: costachris Date: Thu, 17 Jul 2025 18:11:15 -0700 Subject: [PATCH 6/7] lazy levels --- .buildkite/pipeline.yml | 3 +++ .../diagnostic_edmf_precomputed_quantities.jl | 16 ++++++++-------- .../precipitation_precomputed_quantities.jl | 8 ++++---- .../prognostic_edmf_precomputed_quantities.jl | 16 ++++++++-------- 4 files changed, 23 insertions(+), 20 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 274e3767b0..eb280ac691 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -29,6 +29,9 @@ steps: - echo "--- Instantiate .buildkite" - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true); using CUDA; CUDA.precompile_runtime(); Pkg.status()'" + - echo "--- dev package" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"dy/lazy_field_levels\"))'" + agents: slurm_cpus_per_task: 8 slurm_gpus: 1 diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 1adad5c124..175be3529f 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -113,10 +113,10 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( 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(Base.materialize(ᶜh_tot), 1)) + 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(Base.materialize(q_tot), 1)) + Fields.field_values(Fields.level(q_tot, 1)) p_int_level = Fields.field_values(Fields.level(ᶜp, 1)) Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1)) @@ -369,9 +369,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( 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(Base.materialize(ᶜh_tot), i)) + Fields.field_values(Fields.level(ᶜh_tot, i)) q_tot_level = - Fields.field_values(Fields.level(Base.materialize(q_tot), i)) + Fields.field_values(Fields.level(q_tot, i)) p_level = Fields.field_values(Fields.level(ᶜp, i)) Φ_level = Fields.field_values(Fields.level(ᶜΦ, i)) local_geometry_level = Fields.field_values( @@ -396,9 +396,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( 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(Base.materialize(ᶜh_tot), i - 1)) + Fields.field_values(Fields.level(ᶜh_tot, i - 1)) q_tot_prev_level = - Fields.field_values(Fields.level(Base.materialize(q_tot), i - 1)) + 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)) z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1)) @@ -499,7 +499,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( end tke_prev_level = Fields.field_values( - Fields.level(Base.materialize(ᶜtke⁰), i - 1), + Fields.level(ᶜtke⁰, i - 1), ) @. entrʲ_prev_level = entrainment( @@ -1041,7 +1041,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰) end - sfc_tke = Fields.level(Base.materialize(ᶜtke⁰), 1) + sfc_tke = Fields.level(ᶜtke⁰, 1) z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, half) @. ᶜmixing_length_tuple = mixing_length( params, diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index b267ace613..ea64218d88 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -489,19 +489,19 @@ function set_precipitation_surface_fluxes!( ᶜq_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) ᶜq_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) sfc_qᵣ = Fields.Field( - Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)), + Fields.field_values(Fields.level(ᶜq_rai, 1)), sfc_space, ) sfc_qₛ = Fields.Field( - Fields.field_values(Fields.level(Base.materialize(ᶜq_sno), 1)), + Fields.field_values(Fields.level(ᶜq_sno, 1)), sfc_space, ) sfc_qₗ = Fields.Field( - Fields.field_values(Fields.level(Base.materialize(ᶜq_liq), 1)), + Fields.field_values(Fields.level(ᶜq_liq, 1)), sfc_space, ) sfc_qᵢ = Fields.Field( - Fields.field_values(Fields.level(Base.materialize(ᶜq_ice), 1)), + Fields.field_values(Fields.level(ᶜq_ice, 1)), sfc_space, ) sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space) diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 8c2bb07357..d19636a368 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -184,7 +184,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( ), ) ᶜh_tot_int_val = - Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), 1)) + Fields.field_values(Fields.level(ᶜh_tot, 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( @@ -203,7 +203,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ) ᶜq_tot_int_val = - Fields.field_values(Fields.level(Base.materialize(ᶜq_tot), 1)) + Fields.field_values(Fields.level(ᶜq_tot, 1)) ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1)) @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc( ᶜz_int_val - z_sfc_val, @@ -225,22 +225,22 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( ᶜq_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ) ᶜq_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) ᶜq_liq_int_val = - Fields.field_values(Fields.level(Base.materialize(ᶜq_liq), 1)) + Fields.field_values(Fields.level(ᶜq_liq, 1)) ᶜq_liqʲ_int_val = Fields.field_values(Fields.level(ᶜq_liqʲ, 1)) @. ᶜq_liqʲ_int_val = ᶜq_liq_int_val ᶜq_ice_int_val = - Fields.field_values(Fields.level(Base.materialize(ᶜq_ice), 1)) + Fields.field_values(Fields.level(ᶜq_ice, 1)) ᶜq_iceʲ_int_val = Fields.field_values(Fields.level(ᶜq_iceʲ, 1)) @. ᶜq_iceʲ_int_val = ᶜq_ice_int_val ᶜq_rai_int_val = - Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)) + Fields.field_values(Fields.level(ᶜq_rai, 1)) ᶜq_raiʲ_int_val = Fields.field_values(Fields.level(ᶜq_raiʲ, 1)) @. ᶜq_raiʲ_int_val = ᶜq_rai_int_val ᶜq_sno_int_val = - Fields.field_values(Fields.level(Base.materialize(ᶜq_sno), 1)) + Fields.field_values(Fields.level(ᶜq_sno, 1)) ᶜq_snoʲ_int_val = Fields.field_values(Fields.level(ᶜq_snoʲ, 1)) @. ᶜq_snoʲ_int_val = ᶜq_sno_int_val end @@ -492,7 +492,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰) end - sfc_tke = Fields.level(Base.materialize(ᶜtke⁰), 1) + sfc_tke = Fields.level(ᶜtke⁰, 1) @. ᶜmixing_length_tuple = mixing_length( p.params, ustar, @@ -516,7 +516,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ρatke_flux_values = Fields.field_values(ρatke_flux) ρa_sfc_values = - Fields.field_values(Fields.level(Base.materialize(ᶜρa⁰_vals), 1)) # TODO: replace by surface value + Fields.field_values(Fields.level(ᶜρa⁰_vals, 1)) # TODO: replace by surface value ustar_values = Fields.field_values(ustar) sfc_local_geometry_values = Fields.field_values( Fields.level(Fields.local_geometry_field(Y.f), half), From bc352c2e3aebc36787d86c2d78ce22eaa27f22f7 Mon Sep 17 00:00:00 2001 From: costachris Date: Fri, 18 Jul 2025 15:53:08 -0700 Subject: [PATCH 7/7] no Fields.Field in precip precomputed --- .../precipitation_precomputed_quantities.jl | 30 +++++++------------ 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index ea64218d88..a405c46da5 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -488,26 +488,16 @@ function set_precipitation_surface_fluxes!( ᶜq_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ) ᶜq_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ) ᶜq_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ) - sfc_qᵣ = Fields.Field( - Fields.field_values(Fields.level(ᶜq_rai, 1)), - sfc_space, - ) - sfc_qₛ = Fields.Field( - Fields.field_values(Fields.level(ᶜq_sno, 1)), - sfc_space, - ) - sfc_qₗ = Fields.Field( - Fields.field_values(Fields.level(ᶜq_liq, 1)), - sfc_space, - ) - sfc_qᵢ = Fields.Field( - Fields.field_values(Fields.level(ᶜq_ice, 1)), - sfc_space, - ) - sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space) - sfc_wₛ = Fields.Field(Fields.field_values(Fields.level(ᶜwₛ, 1)), sfc_space) - sfc_wₗ = Fields.Field(Fields.field_values(Fields.level(ᶜwₗ, 1)), sfc_space) - sfc_wᵢ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵢ, 1)), sfc_space) + sfc_qᵣ = Fields.field_values(Fields.level(ᶜq_rai, 1)) + sfc_qₛ = Fields.field_values(Fields.level(ᶜq_sno, 1)) + + sfc_qₗ = Fields.field_values(Fields.level(ᶜq_liq, 1)) + sfc_qᵢ = Fields.field_values(Fields.level(ᶜq_ice, 1)) + + sfc_wᵣ = Fields.field_values(Fields.level(ᶜwᵣ, 1)) + sfc_wₛ = Fields.field_values(Fields.level(ᶜwₛ, 1)) + sfc_wₗ = Fields.field_values(Fields.level(ᶜwₗ, 1)) + sfc_wᵢ = Fields.field_values(Fields.level(ᶜwᵢ, 1)) @. surface_rain_flux = sfc_ρ * (sfc_qᵣ * (-sfc_wᵣ) + sfc_qₗ * (-sfc_wₗ)) @. surface_snow_flux = sfc_ρ * (sfc_qₛ * (-sfc_wₛ) + sfc_qᵢ * (-sfc_wᵢ))