diff --git a/config/model_configs/diagnostic_edmfx_bomex_box.yml b/config/model_configs/diagnostic_edmfx_bomex_box.yml index 0ebc11cc8c1..6f172eb58e4 100644 --- a/config/model_configs/diagnostic_edmfx_bomex_box.yml +++ b/config/model_configs/diagnostic_edmfx_bomex_box.yml @@ -26,4 +26,59 @@ z_max: 3e3 z_stretch: false dt: "10secs" t_end: "6hours" -dt_save_to_disk: "10mins" \ No newline at end of file +dt_save_to_disk: "10mins" +diagnostics: + - short_name: ts + period: 10mins + - short_name: ta + period: 10mins + - short_name: thetaa + period: 10mins + - short_name: ha + period: 10mins + - short_name: pfull + period: 10mins + - short_name: rhoa + period: 10mins + - short_name: ua + period: 10mins + - short_name: va + period: 10mins + - short_name: wa + period: 10mins + - short_name: hur + period: 10mins + - short_name: hus + period: 10mins + - short_name: clw + period: 10mins + - short_name: cli + period: 10mins + - short_name: hussfc + period: 10mins + - short_name: evspsbl + period: 10mins + - short_name: arup + period: 10mins + - short_name: waup + period: 10mins + - short_name: taup + period: 10mins + - short_name: thetaaup + period: 10mins + - short_name: haup + period: 10mins + - short_name: husup + period: 10mins + - short_name: hurup + period: 10mins + - short_name: clwup + period: 10mins + - short_name: cliup + period: 10mins + - short_name: waen + period: 10mins + - short_name: tke + period: 10mins + - short_name: lmix + period: 10mins diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl index 17f3d4ff5f7..a44535fe354 100644 --- a/src/diagnostics/Diagnostics.jl +++ b/src/diagnostics/Diagnostics.jl @@ -32,6 +32,9 @@ import ClimaAtmos.RRTMGPInterface as RRTMGPI import ..EDMFX import ..DiagnosticEDMFX +# functions used to calculate diagnostics +import ..draft_area + # We need the abbreviations for symbols like curl, grad, and so on include(joinpath("..", "utils", "abbreviations.jl")) diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl index 45d070e732e..c338acb352a 100644 --- a/src/diagnostics/core_diagnostics.jl +++ b/src/diagnostics/core_diagnostics.jl @@ -47,14 +47,13 @@ error_diagnostic_variable(variable, model::T) where {T} = error_diagnostic_variable("Cannot compute $variable with model = $T") ### -# Rho (3d) +# Density (3d) ### add_diagnostic_variable!( short_name = "rhoa", long_name = "Air Density", standard_name = "air_density", units = "kg m^-3", - comments = "Density of air", compute! = (out, state, cache, time) -> begin if isnothing(out) return copy(state.c.ρ) @@ -129,7 +128,6 @@ add_diagnostic_variable!( long_name = "Air Temperature", standard_name = "air_temperature", units = "K", - comments = "Temperature of air", compute! = (out, state, cache, time) -> begin thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) @@ -148,7 +146,6 @@ add_diagnostic_variable!( long_name = "Air Potential Temperature", standard_name = "air_potential_temperature", units = "K", - comments = "Potential temperature of air", compute! = (out, state, cache, time) -> begin thermo_params = CAP.thermodynamics_params(cache.params) if isnothing(out) @@ -159,6 +156,23 @@ add_diagnostic_variable!( end, ) +### +# Enthalpy (3d) +### +add_diagnostic_variable!( + short_name = "ha", + long_name = "Air Specific Enthalpy", + units = "m^2 s^-2", + compute! = (out, state, cache, time) -> begin + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.specific_enthalpy.(thermo_params, cache.ᶜts) + else + out .= TD.specific_enthalpy.(thermo_params, cache.ᶜts) + end + end, +) + ### # Air pressure (3d) ### @@ -166,7 +180,6 @@ add_diagnostic_variable!( short_name = "pfull", long_name = "Pressure at Model Full-Levels", units = "Pa", - comments = "Pressure of air", compute! = (out, state, cache, time) -> begin if isnothing(out) return copy(cache.ᶜp) @@ -261,6 +274,70 @@ add_diagnostic_variable!( compute! = compute_hus!, ) +### +# Liquid water specific humidity (3d) +### +compute_clw!(out, state, cache, time) = + compute_clw!(out, state, cache, time, cache.atmos.moisture_model) +compute_clw!(_, _, _, _, model::T) where {T} = + error_diagnostic_variable("clw", model) + +function compute_clw!( + out, + state, + cache, + time, + moisture_model::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.liquid_specific_humidity.(thermo_params, cache.ᶜts) + else + out .= TD.liquid_specific_humidity.(thermo_params, cache.ᶜts) + end +end + +add_diagnostic_variable!( + short_name = "clw", + long_name = "Mass Fraction of Cloud Liquid Water", + standard_name = "mass_fraction_of_cloud_liquid_water_in_air", + units = "kg kg^-1", + comments = "Includes both large-scale and convective cloud. This is calculated as the mass of cloud liquid water in the grid cell divided by the mass of air (including the water in all phases) in the grid cells. ", + compute! = compute_clw!, +) + +### +# Ice water specific humidity (3d) +### +compute_cli!(out, state, cache, time) = + compute_cli!(out, state, cache, time, cache.atmos.moisture_model) +compute_cli!(_, _, _, _, model::T) where {T} = + error_diagnostic_variable("cli", model) + +function compute_cli!( + out, + state, + cache, + time, + moisture_model::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.ice_specific_humidity.(thermo_params, cache.ᶜts) + else + out .= TD.ice_specific_humidity.(thermo_params, cache.ᶜts) + end +end + +add_diagnostic_variable!( + short_name = "cli", + long_name = "Mass Fraction of Cloud Ice", + standard_name = "mass_fraction_of_cloud_ice_in_air", + units = "kg kg^-1", + comments = "Includes both large-scale and convective cloud. This is calculated as the mass of cloud ice in the grid cell divided by the mass of air (including the water in all phases) in the grid cell.", + compute! = compute_cli!, +) + ### # Surface specific humidity (2d) ### diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl index 8f2f6f54f31..00cf338726d 100644 --- a/src/diagnostics/default_diagnostics.jl +++ b/src/diagnostics/default_diagnostics.jl @@ -192,7 +192,8 @@ hourly_average(short_names; output_writer = HDF5Writer()) = # Core # ######## function core_default_diagnostics() - core_diagnostics = ["ts", "ta", "thetaa", "pfull", "rhoa", "ua", "va", "wa"] + core_diagnostics = + ["ts", "ta", "thetaa", "ha", "pfull", "rhoa", "ua", "va", "wa"] return [ daily_averages(core_diagnostics...)..., @@ -217,7 +218,7 @@ end function default_diagnostics( ::T, ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} - moist_diagnostics = ["hur", "hus", "hussfc", "evspsbl"] + moist_diagnostics = ["hur", "hus", "clw", "cli", "hussfc", "evspsbl"] return [daily_averages(moist_diagnostics...)...] end @@ -246,3 +247,61 @@ function default_diagnostics(::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics) return [daily_averages(clear_diagnostics...)...] end + +################## +# Turbconv model # +################## +function default_diagnostics(::EDMFX) + edmfx_draft_diagnostics = [ + "arup", + "rhoaup", + "waup", + "taup", + "thetaaup", + "haup", + "husup", + "hurup", + "clwup", + "cliup", + ] + edmfx_env_diagnostics = [ + "aren", + "rhoaen", + "waen", + "taen", + "thetaaen", + "husen", + "huren", + "clwen", + "clien", + "tke", + "lmix", + ] + + return [ + daily_averages(edmfx_draft_diagnostics...)..., + daily_averages(edmfx_env_diagnostics...)..., + ] +end + + +function default_diagnostics(::DiagnosticEDMFX) + diagnostic_edmfx_draft_diagnostics = [ + "arup", + "rhoaup", + "waup", + "taup", + "thetaaup", + "haup", + "husup", + "hurup", + "clwup", + "cliup", + ] + diagnostic_edmfx_env_diagnostics = ["waen", "tke", "lmix"] + + return [ + daily_averages(diagnostic_edmfx_draft_diagnostics...)..., + daily_averages(diagnostic_edmfx_env_diagnostics...)..., + ] +end diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl index cd83ee7f7d9..e93a084b0ed 100644 --- a/src/diagnostics/diagnostic.jl +++ b/src/diagnostics/diagnostic.jl @@ -182,6 +182,7 @@ end # Do you want to define more diagnostics? Add them here include("core_diagnostics.jl") include("radiation_diagnostics.jl") +include("edmfx_diagnostics.jl") # Default diagnostics and higher level interfaces include("default_diagnostics.jl") diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl new file mode 100644 index 00000000000..264dd927d2d --- /dev/null +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -0,0 +1,805 @@ +# This file is included in Diagnostics.jl + +# EDMFX diagnostics + +### +# Updraft area fraction (3d) +### +compute_arup!(out, state, cache, time) = + compute_arup!(out, state, cache, time, cache.atmos.turbconv_model) +compute_arup!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("arup", turbconv_model) + +function compute_arup!(out, state, cache, time, turbconv_model::EDMFX) + 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.ᶜtsʲs.:1), + ) + else + out .= + draft_area.( + (state.c.sgsʲs.:1).ρa, + TD.air_density.(thermo_params, cache.ᶜtsʲs.:1), + ) + end +end + +function compute_arup!(out, state, cache, time, turbconv_model::DiagnosticEDMFX) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return draft_area.( + cache.ᶜρaʲs.:1, + TD.air_density.(thermo_params, cache.ᶜtsʲs.:1), + ) + else + out .= + draft_area.( + cache.ᶜρaʲs.:1, + TD.air_density.(thermo_params, cache.ᶜtsʲs.:1), + ) + end +end + +add_diagnostic_variable!( + short_name = "arup", + long_name = "Updraft Area Fraction", + units = "", + comments = "Area fraction of the first updraft", + compute! = compute_arup!, +) + +### +# Updraft density (3d) +### +compute_rhoaup!(out, state, cache, time) = + compute_rhoaup!(out, state, cache, time, cache.atmos.turbconv_model) +compute_rhoaup!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("rhoaup", turbconv_model) + +function compute_rhoaup!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.air_density.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.air_density.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "rhoaup", + long_name = "Updraft Air Density", + units = "kg m^-3", + comments = "Density of the first updraft", + compute! = compute_rhoaup!, +) + +### +# Updraft w velocity (3d) +### +compute_waup!(out, state, cache, time) = + compute_waup!(out, state, cache, time, cache.atmos.turbconv_model) +compute_waup!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("waup", turbconv_model) + +function compute_waup!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + if isnothing(out) + return copy(Geometry.WVector.(cache.ᶜuʲs.:1).components.data.:1) + else + out .= Geometry.WVector.(cache.ᶜuʲs.:1).components.data.:1 + end +end + +add_diagnostic_variable!( + short_name = "waup", + long_name = "Updraft Upward Air Velocity", + units = "m s^-1", + comments = "Vertical wind component of the first updraft", + compute! = compute_waup!, +) + +### +# Updraft temperature (3d) +### +compute_taup!(out, state, cache, time) = + compute_taup!(out, state, cache, time, cache.atmos.turbconv_model) +compute_taup!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("taup", turbconv_model) + +function compute_taup!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.air_temperature.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.air_temperature.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "taup", + long_name = "Updraft Air Temperature", + units = "K", + comments = "Temperature of the first updraft", + compute! = compute_taup!, +) + +### +# Updraft potential temperature (3d) +### +compute_thetaaup!(out, state, cache, time) = + compute_thetaaup!(out, state, cache, time, cache.atmos.turbconv_model) +compute_thetaaup!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("thetaaup", turbconv_model) + +function compute_thetaaup!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.dry_pottemp.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.dry_pottemp.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "thetaaup", + long_name = "Updraft Air Potential Temperature", + units = "K", + comments = "Potential Temperature of the first updraft", + compute! = compute_thetaaup!, +) + +### +# Updraft specific enthalpy (3d) +### +compute_haup!(out, state, cache, time) = + compute_haup!(out, state, cache, time, cache.atmos.turbconv_model) +compute_haup!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("haup", turbconv_model) + +function compute_haup!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.specific_enthalpy.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.specific_enthalpy.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "haup", + long_name = "Updraft Air Specific Enthalpy", + units = "m^2 s^-2", + comments = "Specific enthalpy of the first updraft", + compute! = compute_haup!, +) + +### +# 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 and with a moist model and with EDMFX", +) + +function compute_husup!( + out, + state, + cache, + time, + moisture_model::Union{EquilMoistModel, NonEquilMoistModel}, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.total_specific_humidity.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.total_specific_humidity.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "husup", + long_name = "Updraft Specific Humidity", + units = "kg kg^-1", + comments = "Specific humidity of the first updraft", + compute! = compute_husup!, +) + +### +# 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{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.relative_humidity.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.relative_humidity.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "hurup", + long_name = "Updraft Relative Humidity", + units = "", + comments = "Relative humidity of the first updraft", + compute! = compute_hurup!, +) + +### +# Updraft liquid water specific humidity (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::Union{EquilMoistModel, NonEquilMoistModel}, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.liquid_specific_humidity.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.liquid_specific_humidity.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "clwup", + 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.", + compute! = compute_clwup!, +) + +### +# 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::Union{EquilMoistModel, NonEquilMoistModel}, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.ice_specific_humidity.(thermo_params, cache.ᶜtsʲs.:1) + else + out .= TD.ice_specific_humidity.(thermo_params, cache.ᶜtsʲs.:1) + end +end + +add_diagnostic_variable!( + short_name = "cliup", + 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.", + compute! = compute_cliup!, +) + +### +# Environment area fraction (3d) +### +compute_aren!(out, state, cache, time) = + compute_aren!(out, state, cache, time, cache.atmos.turbconv_model) +compute_aren!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("aren", turbconv_model) + +function compute_aren!(out, state, cache, time, turbconv_model::EDMFX) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return draft_area.( + cache.ᶜρa⁰, + TD.air_density.(thermo_params, cache.ᶜts⁰), + ) + else + out .= + draft_area.(cache.ᶜρa⁰, TD.air_density.(thermo_params, cache.ᶜts⁰)) + end +end + +function compute_aren!(out, state, cache, time, turbconv_model::DiagnosticEDMFX) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return draft_area.( + cache.ᶜρaʲs.:1, + TD.air_density.(thermo_params, cache.ᶜtsʲs.:1), + ) + else + out .= + draft_area.( + cache.ᶜρaʲs.:1, + TD.air_density.(thermo_params, cache.ᶜtsʲs.:1), + ) + end +end + +add_diagnostic_variable!( + short_name = "aren", + long_name = "Environment Area Fraction", + units = "", + compute! = compute_aren!, +) + +### +# Environment density (3d) +### +compute_rhoaen!(out, state, cache, time) = + compute_rhoaen!(out, state, cache, time, cache.atmos.turbconv_model) +compute_rhoaen!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("rhoaen", turbconv_model) + +function compute_rhoaen!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.air_density.(thermo_params, cache.ᶜts⁰) + else + out .= TD.air_density.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "rhoaen", + long_name = "Environment Air Density", + units = "kg m^-3", + compute! = compute_rhoaen!, +) + +### +# Environment w velocity (3d) +### +compute_waen!(out, state, cache, time) = + compute_waen!(out, state, cache, time, cache.atmos.turbconv_model) +compute_waen!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("waen", turbconv_model) + +function compute_waen!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + if isnothing(out) + return copy(Geometry.WVector.(cache.ᶜu⁰).components.data.:1) + else + out .= Geometry.WVector.(cache.ᶜu⁰).components.data.:1 + end +end + +add_diagnostic_variable!( + short_name = "waen", + long_name = "Environment Upward Air Velocity", + units = "m s^-1", + comments = "Vertical wind component of the environment", + compute! = compute_waen!, +) + +### +# Environment temperature (3d) +### +compute_taen!(out, state, cache, time) = + compute_taen!(out, state, cache, time, cache.atmos.turbconv_model) +compute_taen!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("taen", turbconv_model) + +function compute_taen!(out, state, cache, time, turbconv_model::EDMFX) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.air_temperature.(thermo_params, cache.ᶜts⁰) + else + out .= TD.air_temperature.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "taen", + long_name = "Environment Air Temperature", + units = "K", + compute! = compute_taen!, +) + +### +# Environment potential temperature (3d) +### +compute_thetaaen!(out, state, cache, time) = + compute_thetaaen!(out, state, cache, time, cache.atmos.turbconv_model) +compute_thetaaen!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("thetaaen", turbconv_model) + +function compute_thetaaen!(out, state, cache, time, turbconv_model::EDMFX) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.dry_pottemp.(thermo_params, cache.ᶜts⁰) + else + out .= TD.dry_pottemp.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "thetaaen", + long_name = "Environment Air Potential Temperature", + units = "K", + compute! = compute_thetaaen!, +) + +### +# Environment specific enthalpy (3d) +### +compute_haen!(out, state, cache, time) = + compute_haen!(out, state, cache, time, cache.atmos.turbconv_model) +compute_haen!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("haen", turbconv_model) + +function compute_haen!(out, state, cache, time, turbconv_model::EDMFX) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.dry_pottemp.(thermo_params, cache.ᶜts⁰) + else + out .= TD.dry_pottemp.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "haen", + 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::EDMFX, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.total_specific_humidity.(thermo_params, cache.ᶜts⁰) + else + out .= TD.total_specific_humidity.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "husen", + 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::EDMFX, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.relative_humidity.(thermo_params, cache.ᶜts⁰) + else + out .= TD.relative_humidity.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "huren", + long_name = "Environment Relative Humidity", + units = "", + compute! = compute_huren!, +) + +### +# Environment liquid water specific humidity (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::Union{EquilMoistModel, NonEquilMoistModel}, + turbconv_model::EDMFX, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.liquid_specific_humidity.(thermo_params, cache.ᶜts⁰) + else + out .= TD.liquid_specific_humidity.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "clwen", + 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.", + compute! = compute_clwen!, +) + +### +# 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::Union{EquilMoistModel, NonEquilMoistModel}, + turbconv_model::EDMFX, +) + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.ice_specific_humidity.(thermo_params, cache.ᶜts⁰) + else + out .= TD.ice_specific_humidity.(thermo_params, cache.ᶜts⁰) + end +end + +add_diagnostic_variable!( + short_name = "clien", + 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.", + compute! = compute_clien!, +) + +### +# Environment mixing length (3d) +### +compute_lmix!(out, state, cache, time) = + compute_lmix!(out, state, cache, time, cache.atmos.turbconv_model) +compute_lmix!(_, _, _, _, turbconv_model::T) where {T} = + error_diagnostic_variable("lmix", turbconv_model) + +function compute_lmix!( + out, + state, + cache, + time, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + if isnothing(out) + return copy(cache.ᶜmixing_length) + else + out .= cache.ᶜmixing_length + end +end + +add_diagnostic_variable!( + short_name = "lmix", + long_name = "Environment Mixing Length", + units = "m", + compute! = compute_lmix!, +) + +### +# Environment turbulent kinetic energy (3d) +### +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, + turbconv_model::Union{EDMFX, DiagnosticEDMFX}, +) + if isnothing(out) + return copy(cache.ᶜtke⁰) + else + out .= cache.ᶜtke⁰ + end +end + +add_diagnostic_variable!( + short_name = "tke", + long_name = "Environment Turbulent Kinetic Energy", + units = "m^2 s^-2", + compute! = compute_tke!, +)