From 6c1143d5caba5096f69626ff012fa7a144fee765 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Tue, 25 Nov 2025 15:56:09 -0800 Subject: [PATCH] rft: always return lazy edfmx diagnostics --- src/diagnostics/edmfx_diagnostics.jl | 1482 ++++++-------------------- 1 file changed, 315 insertions(+), 1167 deletions(-) diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index 474f0699d9..46656d45e2 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -7,45 +7,24 @@ ### compute_arup!(out, state, cache, time) = compute_arup!(out, state, cache, time, cache.atmos.turbconv_model) -compute_arup!(_, _, _, _, turbconv_model::T) where {T} = +compute_arup!(_, _, _, _, turbconv_model) = error_diagnostic_variable("arup", turbconv_model) -function compute_arup!(out, state, cache, time, turbconv_model::PrognosticEDMFX) +function compute_arup!(_, state, cache, _, turbconv_model::PrognosticEDMFX) + ρaʲ = (state.c.sgsʲs.:1).ρa thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return draft_area.( - (state.c.sgsʲs.:1).ρa, - TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1), - ) - else - out .= - draft_area.( - (state.c.sgsʲs.:1).ρa, - TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1), - ) - end + ρʲ = @. lazy(TD.air_density(thermo_params, cache.precomputed.ᶜtsʲs.:1)) + return @. lazy(draft_area(ρaʲ, ρʲ)) end - -function compute_arup!(out, state, cache, time, turbconv_model::DiagnosticEDMFX) +function compute_arup!(_, _, cache, _, turbconv_model::DiagnosticEDMFX) + ρaʲ = cache.precomputed.ᶜρaʲs.:1 thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return draft_area.( - cache.precomputed.ᶜρaʲs.:1, - TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1), - ) - else - out .= - draft_area.( - cache.precomputed.ᶜρaʲs.:1, - TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1), - ) - end + ρʲ = @. lazy(TD.air_density(thermo_params, cache.precomputed.ᶜtsʲs.:1)) + return @. lazy(draft_area(ρaʲ, ρʲ)) end -add_diagnostic_variable!( - short_name = "arup", +add_diagnostic_variable!(short_name = "arup", units = "", long_name = "Updraft Area Fraction", - units = "", comments = "Area fraction of the first updraft", compute! = compute_arup!, ) @@ -55,28 +34,16 @@ add_diagnostic_variable!( ### compute_rhoaup!(out, state, cache, time) = compute_rhoaup!(out, state, cache, time, cache.atmos.turbconv_model) -compute_rhoaup!(_, _, _, _, turbconv_model::T) where {T} = +compute_rhoaup!(_, _, _, _, turbconv_model) = error_diagnostic_variable("rhoaup", turbconv_model) -function compute_rhoaup!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) +function compute_rhoaup!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - else - out .= TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - end + return @. lazy(TD.air_density(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end -add_diagnostic_variable!( - short_name = "rhoaup", +add_diagnostic_variable!(short_name = "rhoaup", units = "kg m^-3", long_name = "Updraft Air Density", - units = "kg m^-3", comments = "Density of the first updraft", compute! = compute_rhoaup!, ) @@ -86,27 +53,14 @@ add_diagnostic_variable!( ### compute_waup!(out, state, cache, time) = compute_waup!(out, state, cache, time, cache.atmos.turbconv_model) -compute_waup!(_, _, _, _, turbconv_model::T) where {T} = +compute_waup!(_, _, _, _, turbconv_model) = error_diagnostic_variable("waup", turbconv_model) -function compute_waup!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - if isnothing(out) - return copy(w_component.(Geometry.WVector.(cache.precomputed.ᶜuʲs.:1))) - else - out .= w_component.(Geometry.WVector.(cache.precomputed.ᶜuʲs.:1)) - end -end +compute_waup!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + @. lazy(w_component(Geometry.WVector(cache.precomputed.ᶜuʲs.:1))) -add_diagnostic_variable!( - short_name = "waup", +add_diagnostic_variable!(short_name = "waup", units = "m s^-1", long_name = "Updraft Upward Air Velocity", - units = "m s^-1", comments = "Vertical wind component of the first updraft", compute! = compute_waup!, ) @@ -116,28 +70,16 @@ add_diagnostic_variable!( ### compute_taup!(out, state, cache, time) = compute_taup!(out, state, cache, time, cache.atmos.turbconv_model) -compute_taup!(_, _, _, _, turbconv_model::T) where {T} = +compute_taup!(_, _, _, _, turbconv_model) = error_diagnostic_variable("taup", turbconv_model) -function compute_taup!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) +function compute_taup!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.air_temperature.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - else - out .= TD.air_temperature.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - end + return @. lazy(TD.air_temperature(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end -add_diagnostic_variable!( - short_name = "taup", +add_diagnostic_variable!(short_name = "taup", units = "K", long_name = "Updraft Air Temperature", - units = "K", comments = "Temperature of the first updraft", compute! = compute_taup!, ) @@ -147,28 +89,16 @@ add_diagnostic_variable!( ### compute_thetaaup!(out, state, cache, time) = compute_thetaaup!(out, state, cache, time, cache.atmos.turbconv_model) -compute_thetaaup!(_, _, _, _, turbconv_model::T) where {T} = +compute_thetaaup!(_, _, _, _, turbconv_model) = error_diagnostic_variable("thetaaup", turbconv_model) -function compute_thetaaup!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) +function compute_thetaaup!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.dry_pottemp.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - else - out .= TD.dry_pottemp.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - end + return @. lazy(TD.dry_pottemp(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end -add_diagnostic_variable!( - short_name = "thetaaup", +add_diagnostic_variable!(short_name = "thetaaup", units = "K", long_name = "Updraft Air Potential Temperature", - units = "K", comments = "Potential Temperature of the first updraft", compute! = compute_thetaaup!, ) @@ -178,28 +108,16 @@ add_diagnostic_variable!( ### compute_haup!(out, state, cache, time) = compute_haup!(out, state, cache, time, cache.atmos.turbconv_model) -compute_haup!(_, _, _, _, turbconv_model::T) where {T} = +compute_haup!(_, _, _, _, turbconv_model) = error_diagnostic_variable("haup", turbconv_model) -function compute_haup!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) +function compute_haup!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.specific_enthalpy.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - else - out .= TD.specific_enthalpy.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - end + return @. lazy(TD.specific_enthalpy(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end -add_diagnostic_variable!( - short_name = "haup", +add_diagnostic_variable!(short_name = "haup", units = "m^2 s^-2", long_name = "Updraft Air Specific Enthalpy", - units = "m^2 s^-2", comments = "Specific enthalpy of the first updraft", compute! = compute_haup!, ) @@ -208,51 +126,22 @@ add_diagnostic_variable!( # Updraft total specific humidity (3d) ### compute_husup!(out, state, cache, time) = compute_husup!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_husup!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft specific humidity with a moist model and with EDMFX", -) - -function compute_husup!( - out, - state, - cache, - time, - moisture_model::Union{EquilMoistModel, NonEquilMoistModel}, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, + out, state, cache, time, cache.atmos.moisture_model, cache.atmos.turbconv_model, +) +compute_husup!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft specific humidity with a moist model and with EDMFX", + ) + +function compute_husup!(_, _, cache, _, + ::Union{EquilMoistModel, NonEquilMoistModel}, ::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.total_specific_humidity.( - thermo_params, - cache.precomputed.ᶜtsʲs.:1, - ) - else - out .= - TD.total_specific_humidity.( - thermo_params, - cache.precomputed.ᶜtsʲs.:1, - ) - end + return @. lazy(TD.total_specific_humidity(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end -add_diagnostic_variable!( - short_name = "husup", +add_diagnostic_variable!(short_name = "husup", units = "kg kg^-1", long_name = "Updraft Specific Humidity", - units = "kg kg^-1", comments = "Specific humidity of the first updraft", compute! = compute_husup!, ) @@ -261,44 +150,22 @@ add_diagnostic_variable!( # Updraft relative humidity (3d) ### compute_hurup!(out, state, cache, time) = compute_hurup!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_hurup!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft relative humidity and with a moist model and with EDMFX", -) - -function compute_hurup!( - out, - state, - cache, - time, - moisture_model::Union{EquilMoistModel, NonEquilMoistModel}, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, + out, state, cache, time, cache.atmos.moisture_model, cache.atmos.turbconv_model, +) +compute_hurup!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft relative humidity and with a moist model and with EDMFX", + ) + +function compute_hurup!(_, _, cache, _, + ::Union{EquilMoistModel, NonEquilMoistModel}, ::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.relative_humidity.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - else - out .= TD.relative_humidity.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - end + return @. lazy(TD.relative_humidity(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end -add_diagnostic_variable!( - short_name = "hurup", +add_diagnostic_variable!(short_name = "hurup", units = "", long_name = "Updraft Relative Humidity", - units = "", comments = "Relative humidity of the first updraft", compute! = compute_hurup!, ) @@ -307,79 +174,27 @@ add_diagnostic_variable!( # Updraft liquid water specific humidity and number mixing ratio (3d) ### compute_clwup!(out, state, cache, time) = compute_clwup!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_clwup!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft liquid water specific humidity and with a moist model and with EDMFX", -) - -function compute_clwup!( - out, - state, - cache, - time, - moisture_model::EquilMoistModel, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, + out, state, cache, time, cache.atmos.moisture_model, cache.atmos.turbconv_model, +) +compute_clwup!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft liquid water specific humidity and with a moist model and with EDMFX", + ) + +function compute_clwup!(_, _, cache, _, + ::EquilMoistModel, ::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.liquid_specific_humidity.( - thermo_params, - cache.precomputed.ᶜtsʲs.:1, - ) - else - out .= - TD.liquid_specific_humidity.( - thermo_params, - cache.precomputed.ᶜtsʲs.:1, - ) - end -end -function compute_clwup!( - out, - state, - cache, - time, - moisture_model::NonEquilMoistModel, - turbconv_model::PrognosticEDMFX, -) - if isnothing(out) - return (state.c.sgsʲs.:1).q_liq - else - out .= (state.c.sgsʲs.:1).q_liq - end -end -function compute_clwup!( - out, - state, - cache, - time, - moisture_model::NonEquilMoistModel, - turbconv_model::DiagnosticEDMFX, -) - if isnothing(out) - return (cache.precomputed.ᶜq_liqʲs.:1) - else - out .= (cache.precomputed.ᶜq_liqʲs.:1) - end + return @. lazy(TD.liquid_specific_humidity(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end +compute_clwup!(_, state, _, _, ::NonEquilMoistModel, ::PrognosticEDMFX) = + (state.c.sgsʲs.:1).q_liq +compute_clwup!(_, _, cache, _, ::NonEquilMoistModel, ::DiagnosticEDMFX) = + cache.precomputed.ᶜq_liqʲs.:1 add_diagnostic_variable!( - short_name = "clwup", + short_name = "clwup", units = "kg kg^-1", long_name = "Updraft Mass Fraction of Cloud Liquid Water", - units = "kg kg^-1", comments = """ This is calculated as the mass of cloud liquid water in the first updraft divided by the mass of air (including the water in all phases) in the first updraft. @@ -387,44 +202,19 @@ add_diagnostic_variable!( compute! = compute_clwup!, ) -compute_cdncup!(out, state, cache, time) = compute_cdncup!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_cdncup!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft rain number mixing ratio with a 2M precip model and with EDMFX", -) - -function compute_cdncup!( - out, - state, - cache, - time, - microphysics_model_model::Microphysics2Moment, - turbconv_model::PrognosticEDMFX, -) - if isnothing(out) - return (state.c.sgsʲs.:1).n_liq - else - out .= (state.c.sgsʲs.:1).n_liq - end -end +compute_cdncup!(out, state, cache, time) = compute_cdncup!(out, state, cache, time, + cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_cdncup!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft rain number mixing ratio with a 2M precip model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "cdncup", +compute_cdncup!(_, state, _, _, ::Microphysics2Moment, ::PrognosticEDMFX) = + (state.c.sgsʲs.:1).n_liq + +add_diagnostic_variable!(short_name = "cdncup", units = "kg^-1", long_name = "Updraft Number Mixing Ratio of Cloud Liquid Water", - units = "kg^-1", comments = """ This is calculated as the number of cloud water droplets in the updraft divided by the mass of air (including the water in all phases) in the updraft. @@ -435,77 +225,27 @@ add_diagnostic_variable!( ### # Updraft ice water specific humidity (3d) ### -compute_cliup!(out, state, cache, time) = compute_cliup!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_cliup!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft ice water specific humidity and with a moist model and with EDMFX", -) - -function compute_cliup!( - out, - state, - cache, - time, - moisture_model::EquilMoistModel, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, +compute_cliup!(out, state, cache, time) = compute_cliup!(out, state, cache, time, + cache.atmos.moisture_model, cache.atmos.turbconv_model, +) +compute_cliup!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft ice water specific humidity and with a moist model and with EDMFX", + ) + +function compute_cliup!(_, _, cache, _, + ::EquilMoistModel, ::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.ice_specific_humidity.( - thermo_params, - cache.precomputed.ᶜtsʲs.:1, - ) - else - out .= - TD.ice_specific_humidity.(thermo_params, cache.precomputed.ᶜtsʲs.:1) - end -end -function compute_cliup!( - out, - state, - cache, - time, - moisture_model::NonEquilMoistModel, - turbconv_model::PrognosticEDMFX, -) - if isnothing(out) - return (state.c.sgsʲs.:1).q_ice - else - out .= (state.c.sgsʲs.:1).q_ice - end -end -function compute_cliup!( - out, - state, - cache, - time, - moisture_model::NonEquilMoistModel, - turbconv_model::DiagnosticEDMFX, -) - if isnothing(out) - return (cache.precomputed.ᶜq_iceʲs.:1) - else - out .= (cache.precomputed.ᶜq_iceʲs.:1) - end + return @. lazy(TD.ice_specific_humidity(thermo_params, cache.precomputed.ᶜtsʲs.:1)) end +compute_cliup!(_, state, _, _, ::NonEquilMoistModel, ::PrognosticEDMFX) = + (state.c.sgsʲs.:1).q_ice +compute_cliup!(_, _, cache, _, ::NonEquilMoistModel, ::DiagnosticEDMFX) = + cache.precomputed.ᶜq_iceʲs.:1 -add_diagnostic_variable!( - short_name = "cliup", +add_diagnostic_variable!(short_name = "cliup", units = "kg kg^-1", long_name = "Updraft Mass Fraction of Cloud Ice", - units = "kg kg^-1", comments = """ This is calculated as the mass of cloud ice in the first updraft divided by the mass of air (including the water in all phases) in the first updraft. @@ -516,58 +256,22 @@ add_diagnostic_variable!( ### # Updraft rain water specific humidity and number mixing ratio (3d) ### -compute_husraup!(out, state, cache, time) = compute_husraup!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_husraup!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft rain water specific humidity with a 1M or 2M precip model and with EDMFX", -) - -function compute_husraup!( - out, - state, - cache, - time, - microphysics_model::Union{Microphysics1Moment, Microphysics2Moment}, - turbconv_model::PrognosticEDMFX, -) - if isnothing(out) - return (state.c.sgsʲs.:1).q_rai - else - out .= (state.c.sgsʲs.:1).q_rai - end -end -function compute_husraup!( - out, - state, - cache, - time, - moisture_model::Microphysics1Moment, - turbconv_model::DiagnosticEDMFX, -) - if isnothing(out) - return (cache.precomputed.ᶜq_raiʲs.:1) - else - out .= (cache.precomputed.ᶜq_raiʲs.:1) - end -end +compute_husraup!(out, state, cache, time) = compute_husraup!(out, state, cache, time, + cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_husraup!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft rain water specific humidity with a 1M or 2M precip model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "husraup", +compute_husraup!(_, state, _, _, + ::Union{Microphysics1Moment, Microphysics2Moment}, ::PrognosticEDMFX, +) = (state.c.sgsʲs.:1).q_rai +compute_husraup!(_, _, cache, _, ::Microphysics1Moment, ::DiagnosticEDMFX) = + cache.precomputed.ᶜq_raiʲs.:1 + +add_diagnostic_variable!(short_name = "husraup", units = "kg kg^-1", long_name = "Updraft Mass Fraction of Rain", - units = "kg kg^-1", comments = """ This is calculated as the mass of rain in the first updraft divided by the mass of air (including the water in all phases) in the first updraft. @@ -575,44 +279,19 @@ add_diagnostic_variable!( compute! = compute_husraup!, ) -compute_ncraup!(out, state, cache, time) = compute_ncraup!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_ncraup!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft rain number mixing ratio with a 2M precip model and with EDMFX", -) - -function compute_ncraup!( - out, - state, - cache, - time, - microphysics_model_model::Microphysics2Moment, - turbconv_model::PrognosticEDMFX, -) - if isnothing(out) - return (state.c.sgsʲs.:1).n_rai - else - out .= (state.c.sgsʲs.:1).n_rai - end -end +compute_ncraup!(out, state, cache, time) = compute_ncraup!(out, state, cache, time, + cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_ncraup!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft rain number mixing ratio with a 2M precip model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "ncraup", +compute_ncraup!(_, state, _, _, ::Microphysics2Moment, ::PrognosticEDMFX) = + (state.c.sgsʲs.:1).n_rai + +add_diagnostic_variable!(short_name = "ncraup", units = "kg^-1", long_name = "Updraft Number Mixing Ratio of Rain", - units = "kg^-1", comments = """ This is calculated as the number of raindrops in the updraft divided by the mass of air (including the water in all phases) in the updraft. @@ -623,58 +302,22 @@ add_diagnostic_variable!( ### # Updraft snow specific humidity (3d) ### -compute_hussnup!(out, state, cache, time) = compute_hussnup!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_hussnup!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft snow specific humidity with a 1M or 2M precip model and with EDMFX", -) - -function compute_hussnup!( - out, - state, - cache, - time, - microphysics_model::Union{Microphysics1Moment, Microphysics2Moment}, - turbconv_model::PrognosticEDMFX, -) - if isnothing(out) - return (state.c.sgsʲs.:1).q_sno - else - out .= (state.c.sgsʲs.:1).q_sno - end -end -function compute_hussnup!( - out, - state, - cache, - time, - moisture_model::Microphysics1Moment, - turbconv_model::DiagnosticEDMFX, -) - if isnothing(out) - return (cache.precomputed.ᶜq_snoʲs.:1) - else - out .= (cache.precomputed.ᶜq_snoʲs.:1) - end -end +compute_hussnup!(out, state, cache, time) = compute_hussnup!(out, state, cache, time, + cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_hussnup!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft snow specific humidity with a 1M or 2M precip model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "hussnup", +compute_hussnup!(_, state, _, _, + ::Union{Microphysics1Moment, Microphysics2Moment}, ::PrognosticEDMFX, +) = (state.c.sgsʲs.:1).q_sno +compute_hussnup!(_, _, cache, _, ::Microphysics1Moment, ::DiagnosticEDMFX) = + cache.precomputed.ᶜq_snoʲs.:1 + +add_diagnostic_variable!(short_name = "hussnup", units = "kg kg^-1", long_name = "Updraft Mass Fraction of Snow", - units = "kg kg^-1", comments = """ This is calculated as the mass of snow in the first updraft divided by the mass of air (including the water in all phases) in the first updraft. @@ -687,54 +330,28 @@ add_diagnostic_variable!( ### compute_entr!(out, state, cache, time) = compute_entr!(out, state, cache, time, cache.atmos.turbconv_model) -compute_entr!(_, _, _, _, turbconv_model::T) where {T} = +compute_entr!(_, _, _, _, turbconv_model) = error_diagnostic_variable("entr", turbconv_model) -function compute_entr!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - if isnothing(out) - return copy(cache.precomputed.ᶜentrʲs.:1) - else - out .= cache.precomputed.ᶜentrʲs.:1 - end -end +compute_entr!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + cache.precomputed.ᶜentrʲs.:1 -add_diagnostic_variable!( - short_name = "entr", +add_diagnostic_variable!(short_name = "entr", units = "s^-1", long_name = "Entrainment rate", - units = "s^-1", comments = "Entrainment rate of the first updraft", compute! = compute_entr!, ) compute_turbentr!(out, state, cache, time) = compute_turbentr!(out, state, cache, time, cache.atmos.turbconv_model) -compute_turbentr!(_, _, _, _, turbconv_model::T) where {T} = +compute_turbentr!(_, _, _, _, turbconv_model) = error_diagnostic_variable("turbentr", turbconv_model) -function compute_turbentr!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - if isnothing(out) - return copy(cache.precomputed.ᶜturb_entrʲs.:1) - else - out .= cache.precomputed.ᶜturb_entrʲs.:1 - end -end +compute_turbentr!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + cache.precomputed.ᶜturb_entrʲs.:1 -add_diagnostic_variable!( - short_name = "turbentr", +add_diagnostic_variable!(short_name = "turbentr", units = "s^-1", long_name = "Turbulent entrainment rate", - units = "s^-1", comments = "Turbulent entrainment rate of the first updraft", compute! = compute_turbentr!, ) @@ -744,27 +361,14 @@ add_diagnostic_variable!( ### compute_detr!(out, state, cache, time) = compute_detr!(out, state, cache, time, cache.atmos.turbconv_model) -compute_detr!(_, _, _, _, turbconv_model::T) where {T} = +compute_detr!(_, _, _, _, turbconv_model) = error_diagnostic_variable("detr", turbconv_model) -function compute_detr!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - if isnothing(out) - return copy(cache.precomputed.ᶜdetrʲs.:1) - else - out .= cache.precomputed.ᶜdetrʲs.:1 - end -end +compute_detr!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + cache.precomputed.ᶜdetrʲs.:1 -add_diagnostic_variable!( - short_name = "detr", +add_diagnostic_variable!(short_name = "detr", units = "s^-1", long_name = "Detrainment rate", - units = "s^-1", comments = "Detrainment rate of the first updraft", compute! = compute_detr!, ) @@ -774,48 +378,25 @@ add_diagnostic_variable!( ### compute_aren!(out, state, cache, time) = compute_aren!(out, state, cache, time, cache.atmos.turbconv_model) -compute_aren!(_, _, _, _, turbconv_model::T) where {T} = +compute_aren!(_, _, _, _, turbconv_model) = error_diagnostic_variable("aren", turbconv_model) function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model)) - if isnothing(out) - return draft_area.( - ᶜρa⁰, - TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), - ) - else - out .= - draft_area.( - ᶜρa⁰, - TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), - ) - end + ᶜρ⁰ = @. lazy(TD.air_density(thermo_params, cache.precomputed.ᶜts⁰)) + return @. lazy(draft_area(ᶜρa⁰, ᶜρ⁰)) end function compute_aren!(out, state, cache, time, turbconv_model::DiagnosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return 1 .- - draft_area.( - cache.precomputed.ᶜρaʲs.:1, - TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1), - ) - else - out .= - 1.0 - - draft_area.( - cache.precomputed.ᶜρaʲs.:1, - TD.air_density.(thermo_params, cache.precomputed.ᶜtsʲs.:1), - ) - end + ᶜρaʲ = cache.precomputed.ᶜρaʲs.:1 + ᶜρʲ = @. lazy(TD.air_density(thermo_params, cache.precomputed.ᶜts⁰)) + return @. lazy(1 - draft_area(ᶜρaʲ, ᶜρʲ)) end -add_diagnostic_variable!( - short_name = "aren", +add_diagnostic_variable!(short_name = "aren", units = "", long_name = "Environment Area Fraction", - units = "", compute! = compute_aren!, ) @@ -824,28 +405,16 @@ add_diagnostic_variable!( ### compute_rhoaen!(out, state, cache, time) = compute_rhoaen!(out, state, cache, time, cache.atmos.turbconv_model) -compute_rhoaen!(_, _, _, _, turbconv_model::T) where {T} = +compute_rhoaen!(_, _, _, _, turbconv_model) = error_diagnostic_variable("rhoaen", turbconv_model) -function compute_rhoaen!( - out, - state, - cache, - time, - turbconv_model::PrognosticEDMFX, -) +function compute_rhoaen!(_, _, cache, _, ::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰) - else - out .= TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.air_density(thermo_params, cache.precomputed.ᶜts⁰)) # ᶜρ⁰ end -add_diagnostic_variable!( - short_name = "rhoaen", +add_diagnostic_variable!(short_name = "rhoaen", units = "kg m^-3", long_name = "Environment Air Density", - units = "kg m^-3", compute! = compute_rhoaen!, ) @@ -854,27 +423,14 @@ add_diagnostic_variable!( ### compute_waen!(out, state, cache, time) = compute_waen!(out, state, cache, time, cache.atmos.turbconv_model) -compute_waen!(_, _, _, _, turbconv_model::T) where {T} = +compute_waen!(_, _, _, _, turbconv_model) = error_diagnostic_variable("waen", turbconv_model) -function compute_waen!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - if isnothing(out) - return copy(w_component.(Geometry.WVector.(cache.precomputed.ᶜu⁰))) - else - out .= w_component.(Geometry.WVector.(cache.precomputed.ᶜu⁰)) - end -end +compute_waen!(_, _, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + @. lazy(w_component(Geometry.WVector(cache.precomputed.ᶜu⁰))) -add_diagnostic_variable!( - short_name = "waen", +add_diagnostic_variable!(short_name = "waen", units = "m s^-1", long_name = "Environment Upward Air Velocity", - units = "m s^-1", comments = "Vertical wind component of the environment", compute! = compute_waen!, ) @@ -884,22 +440,16 @@ add_diagnostic_variable!( ### compute_taen!(out, state, cache, time) = compute_taen!(out, state, cache, time, cache.atmos.turbconv_model) -compute_taen!(_, _, _, _, turbconv_model::T) where {T} = +compute_taen!(_, _, _, _, turbconv_model) = error_diagnostic_variable("taen", turbconv_model) -function compute_taen!(out, state, cache, time, turbconv_model::PrognosticEDMFX) +function compute_taen!(_, _, cache, _, ::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.air_temperature.(thermo_params, cache.precomputed.ᶜts⁰) - else - out .= TD.air_temperature.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.air_temperature(thermo_params, cache.precomputed.ᶜts⁰)) end -add_diagnostic_variable!( - short_name = "taen", +add_diagnostic_variable!(short_name = "taen", units = "K", long_name = "Environment Air Temperature", - units = "K", compute! = compute_taen!, ) @@ -908,28 +458,16 @@ add_diagnostic_variable!( ### compute_thetaaen!(out, state, cache, time) = compute_thetaaen!(out, state, cache, time, cache.atmos.turbconv_model) -compute_thetaaen!(_, _, _, _, turbconv_model::T) where {T} = +compute_thetaaen!(_, _, _, _, turbconv_model) = error_diagnostic_variable("thetaaen", turbconv_model) -function compute_thetaaen!( - out, - state, - cache, - time, - turbconv_model::PrognosticEDMFX, -) +function compute_thetaaen!(_, _, cache, _, ::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.dry_pottemp.(thermo_params, cache.precomputed.ᶜts⁰) - else - out .= TD.dry_pottemp.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.dry_pottemp(thermo_params, cache.precomputed.ᶜts⁰)) end -add_diagnostic_variable!( - short_name = "thetaaen", +add_diagnostic_variable!(short_name = "thetaaen", units = "K", long_name = "Environment Air Potential Temperature", - units = "K", compute! = compute_thetaaen!, ) @@ -938,181 +476,87 @@ add_diagnostic_variable!( ### compute_haen!(out, state, cache, time) = compute_haen!(out, state, cache, time, cache.atmos.turbconv_model) -compute_haen!(_, _, _, _, turbconv_model::T) where {T} = +compute_haen!(_, _, _, _, turbconv_model) = error_diagnostic_variable("haen", turbconv_model) -function compute_haen!(out, state, cache, time, turbconv_model::PrognosticEDMFX) +function compute_haen!(_, _, cache, _, ::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.dry_pottemp.(thermo_params, cache.precomputed.ᶜts⁰) - else - out .= TD.dry_pottemp.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.dry_pottemp(thermo_params, cache.precomputed.ᶜts⁰)) end -add_diagnostic_variable!( - short_name = "haen", +add_diagnostic_variable!(short_name = "haen", units = "K", long_name = "Environment Air Specific Enthalpy", - units = "K", compute! = compute_haen!, ) ### # Environment total specific humidity (3d) ### -compute_husen!(out, state, cache, time) = compute_husen!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_husen!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft specific humidity and with a moist model and with EDMFX", -) - -function compute_husen!( - out, - state, - cache, - time, - moisture_model::Union{EquilMoistModel, NonEquilMoistModel}, - turbconv_model::PrognosticEDMFX, +compute_husen!(out, state, cache, time) = compute_husen!(out, state, cache, time, + cache.atmos.moisture_model, cache.atmos.turbconv_model, +) +compute_husen!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft specific humidity and with a moist model and with EDMFX", + ) + +function compute_husen!(_, _, cache, _, + ::Union{EquilMoistModel, NonEquilMoistModel}, ::PrognosticEDMFX, ) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.total_specific_humidity.( - thermo_params, - cache.precomputed.ᶜts⁰, - ) - else - out .= - TD.total_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.total_specific_humidity(thermo_params, cache.precomputed.ᶜts⁰)) end -add_diagnostic_variable!( - short_name = "husen", +add_diagnostic_variable!(short_name = "husen", units = "kg kg^-1", long_name = "Environment Specific Humidity", - units = "kg kg^-1", compute! = compute_husen!, ) ### # Environment relative humidity (3d) ### -compute_huren!(out, state, cache, time) = compute_huren!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_huren!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft relative humidity and with a moist model and with EDMFX", -) - -function compute_huren!( - out, - state, - cache, - time, - moisture_model::Union{EquilMoistModel, NonEquilMoistModel}, - turbconv_model::PrognosticEDMFX, +compute_huren!(out, state, cache, time) = compute_huren!(out, state, cache, time, + cache.atmos.moisture_model, cache.atmos.turbconv_model, +) +compute_huren!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft relative humidity and with a moist model and with EDMFX", + ) + +function compute_huren!(_, _, cache, _, + ::Union{EquilMoistModel, NonEquilMoistModel}, ::PrognosticEDMFX, ) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.relative_humidity.(thermo_params, cache.precomputed.ᶜts⁰) - else - out .= TD.relative_humidity.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.relative_humidity(thermo_params, cache.precomputed.ᶜts⁰)) end -add_diagnostic_variable!( - short_name = "huren", +add_diagnostic_variable!(short_name = "huren", units = "", long_name = "Environment Relative Humidity", - units = "", compute! = compute_huren!, ) ### # Environment liquid water specific humidity and number mixing ratio (3d) ### -compute_clwen!(out, state, cache, time) = compute_clwen!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_clwen!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft liquid water specific humidity and with a moist model and with EDMFX", -) - -function compute_clwen!( - out, - state, - cache, - time, - moisture_model::EquilMoistModel, - turbconv_model::PrognosticEDMFX, +compute_clwen!(out, state, cache, time) = compute_clwen!(out, state, cache, time, + cache.atmos.moisture_model, cache.atmos.turbconv_model, ) +compute_clwen!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft liquid water specific humidity and with a moist model and with EDMFX", + ) + +function compute_clwen!(_, _, cache, _, ::EquilMoistModel, ::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.liquid_specific_humidity.( - thermo_params, - cache.precomputed.ᶜts⁰, - ) - else - out .= - TD.liquid_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.liquid_specific_humidity(thermo_params, cache.precomputed.ᶜts⁰)) end -function compute_clwen!( - out, - state, - cache, - time, - moisture_model::NonEquilMoistModel, - turbconv_model::PrognosticEDMFX, -) - ᶜq_liq⁰ = ᶜspecific_env_value(@name(q_liq), state, cache) - if isnothing(out) - return Base.materialize(ᶜq_liq⁰) - else - out .= ᶜq_liq⁰ - end +function compute_clwen!(_, state, cache, _, ::NonEquilMoistModel, ::PrognosticEDMFX) + return ᶜspecific_env_value(@name(q_liq), state, cache) end -add_diagnostic_variable!( - short_name = "clwen", +add_diagnostic_variable!(short_name = "clwen", units = "kg kg^-1", long_name = "Envrionment Mass Fraction of Cloud Liquid Water", - units = "kg kg^-1", comments = """ This is calculated as the mass of cloud liquid water in the environment divided by the mass of air (including the water in all phases) in the environment. @@ -1120,45 +564,19 @@ add_diagnostic_variable!( compute! = compute_clwen!, ) -compute_cdncen!(out, state, cache, time) = compute_cdncen!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_cdncen!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft cloud liquid water number mixing ratio with a 2M model and with EDMFX", -) - -function compute_cdncen!( - out, - state, - cache, - time, - microphysics_model_model::Microphysics2Moment, - turbconv_model::PrognosticEDMFX, -) - ᶜn_liq⁰ = ᶜspecific_env_value(@name(n_liq), state, cache) - if isnothing(out) - return Base.materialize(ᶜn_liq⁰) - else - out .= ᶜn_liq⁰ - end -end +compute_cdncen!(out, state, cache, time) = compute_cdncen!(out, state, cache, time, + cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_cdncen!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft cloud liquid water number mixing ratio with a 2M model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "cdncen", +compute_cdncen!(_, state, cache, _, ::Microphysics2Moment, ::PrognosticEDMFX) = + ᶜspecific_env_value(@name(n_liq), state, cache) + +add_diagnostic_variable!(short_name = "cdncen", units = "kg^-1", long_name = "Environment Number Mixing Ratio of Cloud Liquid Water", - units = "kg^-1", comments = """ This is calculated as the number of cloud liquid droplets in the environment divided by the mass of air (including the water in all phases) in the environment. @@ -1170,60 +588,23 @@ add_diagnostic_variable!( # Environment ice water specific humidity (3d) ### compute_clien!(out, state, cache, time) = compute_clien!( - out, - state, - cache, - time, - cache.atmos.moisture_model, - cache.atmos.turbconv_model, -) -compute_clien!( - _, - _, - _, - _, - moisture_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft ice water specific humidity and with a moist model and with EDMFX", -) - -function compute_clien!( - out, - state, - cache, - time, - moisture_model::EquilMoistModel, - turbconv_model::PrognosticEDMFX, + out, state, cache, time, cache.atmos.moisture_model, cache.atmos.turbconv_model, ) +compute_clien!(_, _, _, _, moisture_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft ice water specific humidity and with a moist model and with EDMFX", + ) + +function compute_clien!(_, _, cache, _, ::EquilMoistModel, ::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) - if isnothing(out) - return TD.ice_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) - else - out .= TD.ice_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) - end + return @. lazy(TD.ice_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰)) end -function compute_clien!( - out, - state, - cache, - time, - moisture_model::NonEquilMoistModel, - turbconv_model::PrognosticEDMFX, -) - ᶜq_ice⁰ = ᶜspecific_env_value(@name(q_ice), state, cache) - if isnothing(out) - return Base.materialize(ᶜq_ice⁰) - else - out .= ᶜq_ice⁰ - end -end +compute_clien!(_, state, cache, _, ::NonEquilMoistModel, ::PrognosticEDMFX) = + ᶜspecific_env_value(@name(q_ice), state, cache) -add_diagnostic_variable!( - short_name = "clien", +add_diagnostic_variable!(short_name = "clien", units = "kg kg^-1", long_name = "Environment Mass Fraction of Cloud Ice", - units = "kg kg^-1", comments = """ This is calculated as the mass of cloud ice in the environment divided by the mass of air (including the water in all phases) in the environment. @@ -1235,44 +616,19 @@ add_diagnostic_variable!( # Environment rain water specific humidity and number mixing ratio (3d) ### compute_husraen!(out, state, cache, time) = compute_husraen!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_husraen!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft rain specific humidity with a 1M or 2M model and with EDMFX", -) - -function compute_husraen!( - out, - state, - cache, - time, - microphysics_model_model::Union{Microphysics1Moment, Microphysics2Moment}, - turbconv_model::PrognosticEDMFX, -) - ᶜq_rai⁰ = ᶜspecific_env_value(@name(q_rai), state, cache) - if isnothing(out) - return Base.materialize(ᶜq_rai⁰) - else - out .= ᶜq_rai⁰ - end -end + out, state, cache, time, cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_husraen!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft rain specific humidity with a 1M or 2M model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "husraen", +compute_husraen!(_, state, cache, _, + ::Union{Microphysics1Moment, Microphysics2Moment}, ::PrognosticEDMFX, +) = ᶜspecific_env_value(@name(q_rai), state, cache) + +add_diagnostic_variable!(short_name = "husraen", units = "kg kg^-1", long_name = "Environment Mass Fraction of Rain", - units = "kg kg^-1", comments = """ This is calculated as the mass of rain in the environment divided by the mass of air (including the water in all phases) in the environment. @@ -1281,44 +637,18 @@ add_diagnostic_variable!( ) compute_ncraen!(out, state, cache, time) = compute_ncraen!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_ncraen!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft rain number mixing ratio with a 2M model and with EDMFX", -) - -function compute_ncraen!( - out, - state, - cache, - time, - microphysics_model_model::Microphysics2Moment, - turbconv_model::PrognosticEDMFX, -) - ᶜn_rai⁰ = ᶜspecific_env_value(@name(n_rai), state, cache) - if isnothing(out) - return Base.materialize(ᶜn_rai⁰) - else - out .= ᶜn_rai⁰ - end -end + out, state, cache, time, cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_ncraen!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft rain number mixing ratio with a 2M model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "ncraen", +compute_ncraen!(_, state, cache, _, ::Microphysics2Moment, ::PrognosticEDMFX) = + ᶜspecific_env_value(@name(n_rai), state, cache) + +add_diagnostic_variable!(short_name = "ncraen", units = "kg^-1", long_name = "Environment Number Mixing Ratio of Rain", - units = "kg^-1", comments = """ This is calculated as the number of raindrops in the environment divided by the mass of air (including the water in all phases) in the environment. @@ -1330,44 +660,19 @@ add_diagnostic_variable!( # Environment snow water specific humidity (3d) ### compute_hussnen!(out, state, cache, time) = compute_hussnen!( - out, - state, - cache, - time, - cache.atmos.microphysics_model, - cache.atmos.turbconv_model, -) -compute_hussnen!( - _, - _, - _, - _, - microphysics_model::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute updraft snow specific humidity with a 1M or 2M model and with EDMFX", -) - -function compute_hussnen!( - out, - state, - cache, - time, - microphysics_model_model::Union{Microphysics1Moment, Microphysics2Moment}, - turbconv_model::PrognosticEDMFX, -) - ᶜq_sno⁰ = ᶜspecific_env_value(@name(q_sno), state, cache) - if isnothing(out) - return Base.materialize(ᶜq_sno⁰) - else - out .= ᶜq_sno⁰ - end -end + out, state, cache, time, cache.atmos.microphysics_model, cache.atmos.turbconv_model, +) +compute_hussnen!(_, _, _, _, microphysics_model, turbconv_model) = + error_diagnostic_variable( + "Can only compute updraft snow specific humidity with a 1M or 2M model and with EDMFX", + ) -add_diagnostic_variable!( - short_name = "hussnen", +compute_hussnen!(_, state, cache, _, + ::Union{Microphysics1Moment, Microphysics2Moment}, ::PrognosticEDMFX, +) = ᶜspecific_env_value(@name(q_sno), state, cache) + +add_diagnostic_variable!(short_name = "hussnen", units = "kg kg^-1", long_name = "Environment Mass Fraction of Snow", - units = "kg kg^-1", comments = """ This is calculated as the mass of snow in the environment divided by the mass of air (including the water in all phases) in the environment. @@ -1380,37 +685,20 @@ add_diagnostic_variable!( ### compute_tke!(out, state, cache, time) = compute_tke!(out, state, cache, time, cache.atmos.turbconv_model) -compute_tke!(_, _, _, _, turbconv_model::T) where {T} = - error_diagnostic_variable("tke", turbconv_model) - -function compute_tke!( - out, - state, - cache, - time, +compute_tke!(_, _, _, _, turbconv_model) = error_diagnostic_variable("tke", turbconv_model) + +function compute_tke!(_, state, _, _, turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX}, ) - if turbconv_model isa PrognosticEDMFX - sgsʲs = state.c.sgsʲs - ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, sgsʲs, turbconv_model)) - else - ᶜρa⁰ = state.c.ρ - end - - ᶜtke = @. lazy( - specific(state.c.sgs⁰.ρatke, state.c.ρ), - ) - if isnothing(out) - return Base.materialize(ᶜtke) - else - out .= ᶜtke - end + (; ρ, sgsʲs, sgs⁰) = state.c + ᶜρa⁰_pedmfx = @. lazy(ρa⁰(ρ, sgsʲs, turbconv_model)) + ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? ᶜρa⁰_pedmfx : ρ + ᶜtke⁰ = @. lazy(sgs⁰.ρatke, ρ) + return ᶜtke⁰ end -add_diagnostic_variable!( - short_name = "tke", +add_diagnostic_variable!(short_name = "tke", units = "m^2 s^-2", long_name = "Environment Turbulent Kinetic Energy", - units = "m^2 s^-2", compute! = compute_tke!, ) @@ -1419,29 +707,13 @@ add_diagnostic_variable!( ### compute_lmixw!(out, state, cache, time) = compute_lmixw!(out, state, cache, time, cache.atmos.turbconv_model) -compute_lmixw!(_, _, _, _, turbconv_model::T) where {T} = +compute_lmixw!(_, _, _, _, turbconv_model) = error_diagnostic_variable("lmixw", turbconv_model) +compute_lmixw!(_, state, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + ᶜmixing_length(state, cache, Val(:wall)) -function compute_lmixw!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - ᶜwall_mixing_length = ᶜmixing_length(state, cache, Val(:wall)) - - if isnothing(out) - return Base.materialize(ᶜwall_mixing_length) - else - out .= ᶜwall_mixing_length - end -end - -add_diagnostic_variable!( - short_name = "lmixw", +add_diagnostic_variable!(short_name = "lmixw", units = "m", long_name = "Environment Wall Constrained Mixing Length", - units = "m", compute! = compute_lmixw!, ) @@ -1450,30 +722,14 @@ add_diagnostic_variable!( ### compute_lmixtke!(out, state, cache, time) = compute_lmixtke!(out, state, cache, time, cache.atmos.turbconv_model) -compute_lmixtke!(_, _, _, _, turbconv_model::T) where {T} = +compute_lmixtke!(_, _, _, _, turbconv_model) = error_diagnostic_variable("lmixtke", turbconv_model) -function compute_lmixtke!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - - ᶜtke_mixing_length = ᶜmixing_length(state, cache, Val(:tke)) - - if isnothing(out) - return Base.materialize(ᶜtke_mixing_length) - else - out .= ᶜtke_mixing_length - end -end +compute_lmixtke!(_, state, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + ᶜmixing_length(state, cache, Val(:tke)) -add_diagnostic_variable!( - short_name = "lmixtke", +add_diagnostic_variable!(short_name = "lmixtke", units = "m", long_name = "Environment TKE Balanced Mixing Length", - units = "m", compute! = compute_lmixtke!, ) @@ -1482,29 +738,14 @@ add_diagnostic_variable!( ### compute_lmixb!(out, state, cache, time) = compute_lmixb!(out, state, cache, time, cache.atmos.turbconv_model) -compute_lmixb!(_, _, _, _, turbconv_model::T) where {T} = +compute_lmixb!(_, _, _, _, turbconv_model) = error_diagnostic_variable("lmixb", turbconv_model) -function compute_lmixb!( - out, - state, - cache, - time, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, -) - ᶜbuoy_mixing_length = ᶜmixing_length(state, cache, Val(:buoy)) - - if isnothing(out) - return Base.materialize(ᶜbuoy_mixing_length) - else - out .= ᶜbuoy_mixing_length - end -end +compute_lmixb!(_, state, cache, _, ::Union{PrognosticEDMFX, DiagnosticEDMFX}) = + ᶜmixing_length(state, cache, Val(:buoy)) -add_diagnostic_variable!( - short_name = "lmixb", +add_diagnostic_variable!(short_name = "lmixb", units = "m", long_name = "Environment Static Stability Mixing Length", - units = "m", compute! = compute_lmixb!, ) @@ -1512,59 +753,20 @@ add_diagnostic_variable!( # Diffusivity of heat (3d) ### compute_edt!(out, state, cache, time) = compute_edt!( - out, - state, - cache, - time, - cache.atmos.vertical_diffusion, - cache.atmos.turbconv_model, -) -compute_edt!( - _, - _, - _, - _, - vertical_diffusion::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute heat diffusivity with vertical diffusion or EDMFX", -) - -function compute_edt!( - out, - state, - cache, - time, - vertical_diffusion::Union{VerticalDiffusion, DecayWithHeightDiffusion}, - turbconv_model::Nothing, -) - (; vertical_diffusion) = cache.atmos - (; ᶜp) = cache.precomputed - - if vertical_diffusion isa DecayWithHeightDiffusion - ᶜK_h = - ᶜcompute_eddy_diffusivity_coefficient(state.c.ρ, vertical_diffusion) - elseif vertical_diffusion isa VerticalDiffusion - ᶜK_h = ᶜcompute_eddy_diffusivity_coefficient( - state.c.uₕ, - ᶜp, - vertical_diffusion, - ) - end - if isnothing(out) - return copy(ᶜK_h) - else - out .= ᶜK_h - end -end + out, state, cache, time, cache.atmos.vertical_diffusion, cache.atmos.turbconv_model, +) +compute_edt!(_, _, _, _, vertical_diffusion, turbconv_model) = + error_diagnostic_variable( + "Can only compute heat diffusivity with vertical diffusion or EDMFX", + ) + +compute_edt!(_, state, cache, _, model::VerticalDiffusion, ::Nothing) = + ᶜcompute_eddy_diffusivity_coefficient(state.c.uₕ, cache.precomputed.ᶜp, model) +compute_edt!(_, state, _, _, model::DecayWithHeightDiffusion, ::Nothing) = + ᶜcompute_eddy_diffusivity_coefficient(state.c.ρ, model) -function compute_edt!( - out, - state, - cache, - time, - vertical_diffusion::Nothing, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, +function compute_edt!(_, state, cache, _, + vertdiff::Nothing, turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) turbconv_params = CAP.turbconv_params(cache.params) (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = cache.precomputed @@ -1578,22 +780,15 @@ function compute_edt!( ) ᶜmixing_length_field = ᶜmixing_length(state, cache) ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field)) - ᶜprandtl_nvec = @. lazy( - turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm), - ) + ᶜprandtl_nvec = + @. lazy(turbulent_prandtl_number(cache.params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)) ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)) - if isnothing(out) - return Base.materialize(ᶜK_h) - else - out .= ᶜK_h - end + return ᶜK_h end -add_diagnostic_variable!( - short_name = "edt", +add_diagnostic_variable!(short_name = "edt", units = "m^2 s^-1", long_name = "Eddy Diffusivity Coefficient for Temperature", standard_name = "atmosphere_heat_diffusivity", - units = "m^2 s^-1", comments = "Vertical diffusion coefficient for temperature due to parameterized eddies", compute! = compute_edt!, ) @@ -1602,60 +797,20 @@ add_diagnostic_variable!( # Diffusivity of momentum (3d) ### compute_evu!(out, state, cache, time) = compute_evu!( - out, - state, - cache, - time, - cache.atmos.vertical_diffusion, - cache.atmos.turbconv_model, -) -compute_evu!( - _, - _, - _, - _, - vertical_diffusion::T1, - turbconv_model::T2, -) where {T1, T2} = error_diagnostic_variable( - "Can only compute momentum diffusivity with vertical diffusion or EDMFX", -) - -function compute_evu!( - out, - state, - cache, - time, - vertical_diffusion::Union{VerticalDiffusion, DecayWithHeightDiffusion}, - turbconv_model::Nothing, -) - (; vertical_diffusion) = cache.atmos - (; ᶜp) = cache.precomputed - - # this setup assumes ᶜK_u = ᶜK_h - if vertical_diffusion isa DecayWithHeightDiffusion - ᶜK_u = - ᶜcompute_eddy_diffusivity_coefficient(state.c.ρ, vertical_diffusion) - elseif vertical_diffusion isa VerticalDiffusion - ᶜK_u = ᶜcompute_eddy_diffusivity_coefficient( - state.c.uₕ, - ᶜp, - vertical_diffusion, - ) - end - if isnothing(out) - return copy(ᶜK_u) - else - out .= ᶜK_u - end -end + out, state, cache, time, cache.atmos.vertical_diffusion, cache.atmos.turbconv_model, +) +compute_evu!(_, _, _, _, vertical_diffusion, turbconv_model) = + error_diagnostic_variable( + "Can only compute momentum diffusivity with vertical diffusion or EDMFX", + ) + +compute_evu!(_, state, cache, _, model::VerticalDiffusion, ::Nothing) = + ᶜcompute_eddy_diffusivity_coefficient(state.c.uₕ, cache.precomputed.ᶜp, model) +compute_evu!(_, state, _, _, model::DecayWithHeightDiffusion, ::Nothing) = + ᶜcompute_eddy_diffusivity_coefficient(state.c.ρ, model) -function compute_evu!( - out, - state, - cache, - time, - vertical_diffusion::Nothing, - turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, +function compute_evu!(_, state, cache, _, + ::Nothing, turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) turbconv_params = CAP.turbconv_params(cache.params) @@ -1667,19 +822,12 @@ function compute_evu!( ) ᶜmixing_length_field = ᶜmixing_length(state, cache) ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field)) - - if isnothing(out) - return Base.materialize(ᶜK_u) - else - out .= ᶜK_u - end + return ᶜK_u end -add_diagnostic_variable!( - short_name = "evu", +add_diagnostic_variable!(short_name = "evu", units = "m^2 s^-1", long_name = "Eddy Viscosity Coefficient for Momentum", standard_name = "atmosphere_momentum_diffusivity", - units = "m^2 s^-1", comments = "Vertical diffusion coefficient for momentum due to parameterized eddies", compute! = compute_evu!, )