Hackathon refactor

config/model_configs/prognostic_edmfx_gcmdriven_column.yml

@@ -41,7 +41,7 @@ diagnostics:
     period: 10mins
   - short_name: [arup, waup, taup, thetaaup, haup, husup, hurup, clwup, cliup, waen, taen, thetaaen, haen, husen, huren, clwen, clien, tke]
     period: 10mins
-  - short_name: [entr, detr, lmix, bgrad, strain, edt, evu]
+  - short_name: [entr, detr, lmix, bgrad, strain,] # edt, evu]
     period: 10mins
   - short_name: [rlut, rlutcs, rsut, rsutcs, clwvi, lwp, clivi, dsevi, clvi, prw, hurvi, husv]
     period: 10mins

src/cache/cloud_fraction.jl

@@ -39,9 +39,23 @@ NVTX.@annotate function set_cloud_fraction!(
 )
     (; params) = p
     (; turbconv_model) = p.atmos
-    (; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
+    (; ᶜts, cloud_diagnostics_tuple) = p.precomputed
     thermo_params = CAP.thermodynamics_params(params)
+    FT = eltype(p.params)
+    n = n_mass_flux_subdomains(turbconv_model)
     if isnothing(turbconv_model)
+        (; ustar, obukhov_length) = p.precomputed.sfc_conditions
+        (;
+            ᶜlinear_buoygrad,
+            ᶜstrain_rate_norm,
+            ᶠu³⁰,
+            ᶠu³,
+            ᶜtke⁰,
+            ᶜentrʲs,
+            ᶜdetrʲs,
+            ᶠu³ʲs,
+            ᶜρa⁰,
+        ) = p.precomputed
         if p.atmos.call_cloud_diagnostics_per_stage isa
            CallCloudDiagnosticsPerStage
             (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
@@ -53,6 +67,39 @@ NVTX.@annotate function set_cloud_fraction!(
             @. ᶜgradᵥ_θ_liq_ice =
                 ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
         end
+        ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
+        @. ᶜprandtl_nvec =
+            turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
+
+        sfc_tke = Fields.level(ᶜtke⁰, 1)
+        z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)
+        ᶜz = Fields.coordinate_field(Y.c).z
+        ᶜdz = Fields.Δz_field(axes(Y.c))
+
+        ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
+        @. ᶜtke_exch = 0
+        for j in 1:n
+            ᶠu³ʲ = ᶠu³ʲs.:($j)
+            @. ᶜtke_exch +=
+                Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰ *
+                (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
+        end
+
+        ᶜmixing_length = @. lazy(master_mixing_length(
+            params,
+            ustar,
+            ᶜz,
+            z_sfc,
+            ᶜdz,
+            max(sfc_tke, eps(FT)),
+            ᶜlinear_buoygrad,
+            max(ᶜtke⁰, 0),
+            obukhov_length,
+            ᶜstrain_rate_norm,
+            ᶜprandtl_nvec,
+            ᶜtke_exch,
+            p.atmos.edmfx_model.scale_blending_method,
+        ))
         compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     end
     if moist_model isa EquilMoistModel
@@ -64,12 +111,12 @@ NVTX.@annotate function set_cloud_fraction!(
     else
         @. cloud_diagnostics_tuple = make_named_tuple(
             ifelse(
-                p.precomputed.ᶜspecific.q_liq + p.precomputed.ᶜspecific.q_ice > 0,
+                specific(Y.c.ρq_liq, Y.c.ρ) + specific(Y.c.ρq_ice, Y.c.ρ) > 0,
                 1,
                 0,
             ),
-            p.precomputed.ᶜspecific.q_liq,
-            p.precomputed.ᶜspecific.q_ice,
+            specific(Y.c.ρq_liq, Y.c.ρ),
+            specific(Y.c.ρq_ice, Y.c.ρ),
         )
     end
 end
@@ -81,13 +128,26 @@ NVTX.@annotate function set_cloud_fraction!(
 )
     SG_quad = qc.SG_quad
     (; params) = p
+    (; ustar, obukhov_length) = p.precomputed.sfc_conditions
 
     FT = eltype(params)
+    n = n_mass_flux_subdomains(turbconv_model)
     thermo_params = CAP.thermodynamics_params(params)
-    (; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
+    (; ᶜts, cloud_diagnostics_tuple) = p.precomputed
     (; turbconv_model) = p.atmos
 
     if isnothing(turbconv_model)
+        (;
+            ᶜlinear_buoygrad,
+            ᶜstrain_rate_norm,
+            ᶠu³⁰,
+            ᶠu³,
+            ᶜtke⁰,
+            ᶜentrʲs, 
+            ᶜdetrʲs, 
+            ᶠu³ʲs,
+            ᶜρa⁰,
+        ) = p.precomputed
         if p.atmos.call_cloud_diagnostics_per_stage isa
            CallCloudDiagnosticsPerStage
             (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
@@ -98,7 +158,47 @@ NVTX.@annotate function set_cloud_fraction!(
                 ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts)))
             @. ᶜgradᵥ_θ_liq_ice =
                 ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
+
+            @. ᶜprandtl_nvec = turbulent_prandtl_number(
+                params,
+                ᶜlinear_buoygrad,
+                ᶜstrain_rate_norm,
+            )
         end
+
+        ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
+        @. ᶜprandtl_nvec =
+            turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
+
+        sfc_tke = Fields.level(ᶜtke⁰, 1)
+        z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)
+        ᶜz = Fields.coordinate_field(Y.c).z
+        ᶜdz = Fields.Δz_field(axes(Y.c))
+
+        ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
+        @. ᶜtke_exch = 0
+        for j in 1:n
+            ᶠu³ʲ = ᶠu³ʲs.:($j)
+            @. ᶜtke_exch +=
+                Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰ *
+                (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
+        end
+
+        ᶜmixing_length = @. lazy(master_mixing_length(
+            params,
+            ustar,
+            ᶜz,
+            z_sfc,
+            ᶜdz,
+            max(sfc_tke, eps(FT)),
+            ᶜlinear_buoygrad,
+            max(ᶜtke⁰, 0),
+            obukhov_length,
+            ᶜstrain_rate_norm,
+            ᶜprandtl_nvec,
+            ᶜtke_exch,
+            p.atmos.edmfx_model.scale_blending_method,
+        ))
         compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     end
 

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,
@@ -950,7 +950,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     (; params) = p
     (; dt) = p
     (; ᶜp, ᶜu, ᶜts) = p.precomputed
-    (; q_tot) = p.precomputed.ᶜspecific
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
     (; ᶜtke⁰) = p.precomputed
     (;
@@ -1024,8 +1023,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     )
     @. ᶜmixing_length = ᶜmixing_length_tuple.master
 
-    @. ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
-    @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
+    ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length))
+    ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
 
     ρatke_flux_values = Fields.field_values(ρatke_flux)
     ρ_int_values = Fields.field_values(Fields.level(Y.c.ρ, 1))
@@ -1068,14 +1067,13 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
     microphys_0m_params = CAP.microphysics_0m_params(p.params)
     (; dt) = p
     (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
-    (; q_tot) = p.precomputed.ᶜspecific
 
     # Environment precipitation sources (to be applied to grid mean)
     @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
         thermo_params,
         microphys_0m_params,
         dt,
-        q_tot,
+        specific(Y.c.ρq_tot, Y.c.ρ),
         ᶜts,
     )
     return nothing

src/cache/precipitation_precomputed_quantities.jl

@@ -105,7 +105,6 @@ function set_precipitation_velocities!(
     precip_model::Microphysics2Moment,
 )
     (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
-    (; q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific
     (; ᶜΦ) = p.core
 
     cm1c = CAP.microphysics_cloud_params(p.params)
@@ -116,23 +115,23 @@ function set_precipitation_velocities!(
     # compute the precipitation terminal velocity [m/s]
     # TODO sedimentation of snow is based on the 1M scheme
     @. ᶜwnᵣ = getindex(
-        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, q_rai, Y.c.ρ, Y.c.ρn_rai),
+        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, specific(Y.c.ρq_rai, Y.c.ρ), Y.c.ρ, Y.c.ρn_rai),
         1,
     )
     @. ᶜwᵣ = getindex(
-        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, q_rai, Y.c.ρ, Y.c.ρn_rai),
+        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, specific(Y.c.ρq_rai, Y.c.ρ), Y.c.ρ, Y.c.ρn_rai),
         2,
     )
-    @. ᶜwₛ = CM1.terminal_velocity(cm1p.ps, cm1p.tv.snow, Y.c.ρ, q_sno)
+    @. ᶜwₛ = CM1.terminal_velocity(cm1p.ps, cm1p.tv.snow, Y.c.ρ, specific(Y.c.ρq_sno, Y.c.ρ))
     # compute sedimentation velocity for cloud condensate [m/s]
     # TODO sedimentation velocities of cloud condensates are based 
     # on the 1M scheme. Sedimentation velocity of cloud number concentration
     # is equal to that of the mass.
     @. ᶜwnₗ =
-        CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, q_liq)
-    @. ᶜwₗ = CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, q_liq)
+        CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, specific(Y.c.ρq_liq, Y.c.ρ))
+    @. ᶜwₗ = CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, specific(Y.c.ρq_liq, Y.c.ρ))
     @. ᶜwᵢ =
-        CMNe.terminal_velocity(cm1c.ice, cm1c.Ch2022.small_ice, Y.c.ρ, q_ice)
+        CMNe.terminal_velocity(cm1c.ice, cm1c.Ch2022.small_ice, Y.c.ρ, specific(Y.c.ρq_ice, Y.c.ρ))
 
     # compute their contributions to energy and total water advection
     @. ᶜwₜqₜ =

src/cache/precomputed_quantities.jl

@@ -172,9 +172,7 @@ function precomputed_quantities(Y, atmos)
     advective_sgs_quantities =
         atmos.turbconv_model isa PrognosticEDMFX ?
         (;
-            ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
-            ᶜK_u = similar(Y.c, FT),
-            ᶜK_h = similar(Y.c, FT),
+            # ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
             bdmr_l = similar(Y.c, BidiagonalMatrixRow{FT}),
             bdmr_r = similar(Y.c, BidiagonalMatrixRow{FT}),
@@ -194,7 +192,6 @@ function precomputed_quantities(Y, atmos)
         (;
             ᶜ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),
         ) : (;)
@@ -398,34 +395,34 @@ function thermo_state(
     return get_ts(ρ, p, θ, e_int, q_tot, q_pt)
 end
 
-function thermo_vars(moisture_model, precip_model, specific, K, Φ)
-    energy_var = (; e_int = specific.e_tot - K - Φ)
+function thermo_vars(moisture_model, precip_model, ᶜY, K, Φ)
+    energy_var = (; e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ)
     moisture_var = if moisture_model isa DryModel
         (;)
     elseif moisture_model isa EquilMoistModel
-        (; specific.q_tot)
+        (; q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ))
     elseif moisture_model isa NonEquilMoistModel
         q_pt_args = (
-            specific.q_tot,
-            specific.q_liq + specific.q_rai,
-            specific.q_ice + specific.q_sno,
+            q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ),
+            q_liq = specific(ᶜY.ρq_liq, ᶜY.ρ),
+            q_ice = specific(ᶜY.ρq_ice, ᶜY.ρ),
         )
         (; q_pt = TD.PhasePartition(q_pt_args...))
     end
     return (; energy_var..., moisture_var...)
 end
 
-ts_gs(thermo_params, moisture_model, precip_model, specific, K, Φ, ρ) =
+ts_gs(thermo_params, moisture_model, precip_model, ᶜY, K, Φ, ρ) =
     thermo_state(
         thermo_params;
-        thermo_vars(moisture_model, precip_model, specific, K, Φ)...,
+        thermo_vars(moisture_model, precip_model, ᶜY, K, Φ)...,
         ρ,
     )
 
-ts_sgs(thermo_params, moisture_model, precip_model, specific, K, Φ, p) =
+ts_sgs(thermo_params, moisture_model, precip_model, ᶜY, K, Φ, p) =
     thermo_state(
         thermo_params;
-        thermo_vars(moisture_model, precip_model, specific, K, Φ)...,
+        thermo_vars(moisture_model, precip_model, ᶜY, K, Φ)...,
         p,
     )
 
@@ -489,9 +486,9 @@ 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..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
+    @. ᶜ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.e_tot)
+    @. ᶜ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)