Remove k_u, k_v, and mixing length from precomputed quantities

Author: costachris

config/model_configs/diagnostic_edmfx_aquaplanet_dense_autodiff.yml

@@ -19,6 +19,6 @@ edmfx_sgs_diffusive_flux: true
 moist: equil
 cloud_model: "quadrature_sgs"
 precip_model: 0M
-dt: 120secs
+dt: 80secs
 t_end: 1hours
 toml: [toml/diagnostic_edmfx_0M.toml]

perf/benchmark.jl

@@ -60,7 +60,7 @@ are_boundschecks_forced = Base.JLOptions().check_bounds == 1
         end
         @test compare_mem(trials, "Wfact", 0)
         @test compare_mem(trials, "ldiv!", 0)
-        @test compare_mem(trials, "T_imp!", 96)
+        @test compare_mem(trials, "T_imp!", 19680)
         @test compare_mem(trials, "T_exp_T_lim!", 190420)
         @test compare_mem(trials, "lim!", 0)
         @test compare_mem(trials, "dss!", 0)

perf/flame.jl

@@ -45,9 +45,9 @@ allocs_limit["flame_diagnostics"] = 10_677_144
 allocs_limit["flame_aquaplanet_diagedmf"] = 11_644_128
 allocs_limit["flame_aquaplanet_progedmf"] = 774_712
 allocs_limit["flame_aquaplanet_progedmf_dense_autodiff"] = 774_712
-allocs_limit["flame_diffusion"] = 138_432
+allocs_limit["flame_diffusion"] = 153_248
 allocs_limit["flame_threaded"] = 2047_736
-allocs_limit["flame_callbacks"] = 391_864
+allocs_limit["flame_callbacks"] = 450_336
 allocs_limit["flame_gravity_wave"] = 581_381_976
 # Ideally, we would like to track all the allocations, but this becomes too
 # expensive there is too many of them. Here, we set the default sample rate to

perf/jet_test_nfailures.jl

@@ -38,7 +38,7 @@ using Test
     # inference. By increasing this counter, we acknowledge that
     # we have introduced an inference failure. We hope to drive
     # this number down to 0.
-    n_allowed_failures = 253
+    n_allowed_failures = 265
     @show n
     if n < n_allowed_failures
         @info "Please update the n-failures to $n"

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-248
+249
 
 # **README**
 #
@@ -20,6 +20,9 @@
 
 
 #=
+249
+- Remove viscosity, diffusivity, and mixing length from precomputed quantities
+
 248
 - update deps: climacore 0.14.35
 

src/cache/cloud_fraction.jl

@@ -22,12 +22,9 @@ end
    Compute the grid scale cloud fraction based on sub-grid scale properties
 """
 NVTX.@annotate function set_cloud_fraction!(Y, p, ::DryModel, _)
-    (; ᶜmixing_length) = p.precomputed
     (; turbconv_model) = p.atmos
     FT = eltype(p.params)
-    if isnothing(turbconv_model)
-        compute_gm_mixing_length!(ᶜmixing_length, Y, p)
-    end
+
     p.precomputed.cloud_diagnostics_tuple .=
         ((; cf = FT(0), q_liq = FT(0), q_ice = FT(0)),)
 end
@@ -39,8 +36,9 @@ 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)
+
     if isnothing(turbconv_model)
         if p.atmos.call_cloud_diagnostics_per_stage isa
            CallCloudDiagnosticsPerStage
@@ -53,7 +51,6 @@ NVTX.@annotate function set_cloud_fraction!(
             @. ᶜgradᵥ_θ_liq_ice =
                 ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
         end
-        compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     end
     if moist_model isa EquilMoistModel
         @. cloud_diagnostics_tuple = make_named_tuple(
@@ -85,7 +82,7 @@ NVTX.@annotate function set_cloud_fraction!(
 
     FT = eltype(params)
     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)
@@ -100,9 +97,10 @@ NVTX.@annotate function set_cloud_fraction!(
             @. ᶜgradᵥ_θ_liq_ice =
                 ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
         end
-        compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     end
 
+    ᶜmixing_length = compute_gm_mixing_length(Y, p)
+
     diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
     @. cloud_diagnostics_tuple = quad_loop(
         SG_quad,
@@ -139,20 +137,23 @@ NVTX.@annotate function set_cloud_fraction!(
 
     FT = eltype(params)
     thermo_params = CAP.thermodynamics_params(params)
-    (; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
+    (; ᶜts, cloud_diagnostics_tuple) = p.precomputed
     (; turbconv_model) = p.atmos
 
     # TODO - we should make this default when using diagnostic edmf
     # environment
     diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
 
+    ᶜmixing_length_field = p.scratch.ᶜtemp_scalar
+    ᶜmixing_length_field .= ᶜmixing_length(Y, p)
+
     @. cloud_diagnostics_tuple = quad_loop(
         SG_quad,
         ᶜts,
         Geometry.WVector(p.precomputed.ᶜgradᵥ_q_tot),
         Geometry.WVector(p.precomputed.ᶜgradᵥ_θ_liq_ice),
         diagnostic_covariance_coeff,
-        ᶜmixing_length,
+        ᶜmixing_length_field,
         thermo_params,
     )
 
@@ -192,21 +193,24 @@ NVTX.@annotate function set_cloud_fraction!(
 
     FT = eltype(params)
     thermo_params = CAP.thermodynamics_params(params)
-    (; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
+    (; ᶜts⁰, cloud_diagnostics_tuple) = p.precomputed
     (; ᶜρʲs, ᶜtsʲs, ᶜρa⁰, ᶜρ⁰) = p.precomputed
     (; turbconv_model) = p.atmos
 
     # TODO - we should make this default when using diagnostic edmf
     # environment
     diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
 
+    ᶜmixing_length_field = p.scratch.ᶜtemp_scalar
+    ᶜmixing_length_field .= ᶜmixing_length(Y, p)
+
     @. cloud_diagnostics_tuple = quad_loop(
         SG_quad,
         ᶜts⁰,
         Geometry.WVector(p.precomputed.ᶜgradᵥ_q_tot⁰),
         Geometry.WVector(p.precomputed.ᶜgradᵥ_θ_liq_ice⁰),
         diagnostic_covariance_coeff,
-        ᶜmixing_length,
+        ᶜmixing_length_field,
         thermo_params,
     )
 

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -966,13 +966,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     (; ᶜp, ᶜu, ᶜts) = p.precomputed
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
     (; ᶜtke⁰) = p.precomputed
-    (;
-        ᶜlinear_buoygrad,
-        ᶜstrain_rate_norm,
-        ᶜmixing_length_tuple,
-        ᶜmixing_length,
-    ) = p.precomputed
-    (; ᶜK_h, ᶜK_u, ρatke_flux) = p.precomputed
+    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
+    (; ρatke_flux) = p.precomputed
     turbconv_params = CAP.turbconv_params(params)
     thermo_params = CAP.thermodynamics_params(params)
     ᶜlg = Fields.local_geometry_field(Y.c)
@@ -1004,42 +999,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
-    ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
-    @. ᶜprandtl_nvec =
-        turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
-
-    ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
-    @. ᶜtke_exch = 0
-    # using ᶜu⁰ would be more correct, but this is more consistent with the
-    # TKE equation, where using ᶜu⁰ results in allocation
-    for j in 1:n
-        @. ᶜtke_exch +=
-            ᶜρaʲs.:($$j) * ᶜdetrʲs.:($$j) / Y.c.ρ *
-            (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
-    end
-
-    sfc_tke = Fields.level(ᶜtke⁰, 1)
-    z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, half)
-    @. ᶜmixing_length_tuple = mixing_length(
-        params,
-        ustar,
-        ᶜz,
-        z_sfc,
-        ᶜdz,
-        max(sfc_tke, 0),
-        ᶜlinear_buoygrad,
-        max(ᶜtke⁰, 0),
-        obukhov_length,
-        ᶜstrain_rate_norm,
-        ᶜprandtl_nvec,
-        ᶜtke_exch,
-        p.atmos.edmfx_model.scale_blending_method,
-    )
-    @. ᶜmixing_length = ᶜmixing_length_tuple.master
-
-    @. ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
-    @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
-
     ρatke_flux_values = Fields.field_values(ρatke_flux)
     ρ_int_values = Fields.field_values(Fields.level(Y.c.ρ, 1))
     u_int_values = Fields.field_values(Fields.level(ᶜu, 1))

src/cache/precomputed_quantities.jl

@@ -114,7 +114,6 @@ function precomputed_quantities(Y, atmos)
     gs_quantities = (;
         ᶜwₜqₜ = similar(Y.c, Geometry.WVector{FT}),
         ᶜwₕhₜ = similar(Y.c, Geometry.WVector{FT}),
-        ᶜmixing_length = similar(Y.c, FT),
         ᶜlinear_buoygrad = similar(Y.c, FT),
         ᶜstrain_rate_norm = similar(Y.c, FT),
         sfc_conditions = similar(Spaces.level(Y.f, half), SCT),
@@ -173,9 +172,6 @@ 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),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
             bdmr_l = similar(Y.c, BidiagonalMatrixRow{FT}),
             bdmr_r = similar(Y.c, BidiagonalMatrixRow{FT}),
@@ -193,11 +189,8 @@ function precomputed_quantities(Y, atmos)
     edonly_quantities =
         atmos.turbconv_model isa EDOnlyEDMFX ?
         (;
-            ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
             ᶜtke⁰ = similar(Y.c, FT),
-            ᶜK_u = similar(Y.c, FT),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
-            ᶜK_h = similar(Y.c, FT),
         ) : (;)
 
     sgs_quantities = (;
@@ -225,20 +218,9 @@ function precomputed_quantities(Y, atmos)
             ᶠu³⁰ = similar(Y.f, CT3{FT}),
             ᶜu⁰ = similar(Y.c, C123{FT}),
             ᶜK⁰ = similar(Y.c, FT),
-            ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
-            ᶜK_u = similar(Y.c, FT),
-            ᶜK_h = similar(Y.c, FT),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
             precipitation_sgs_quantities...,
         ) : (;)
-    vert_diff_quantities =
-        if atmos.vertical_diffusion isa
-           Union{VerticalDiffusion, DecayWithHeightDiffusion}
-            ᶜK_h = similar(Y.c, FT)
-            (; ᶜK_u = ᶜK_h, ᶜK_h) # ᶜK_u aliases ᶜK_h because they are always equal.
-        else
-            (;)
-        end
     smagorinsky_lilly_quantities =
         if atmos.smagorinsky_lilly isa SmagorinskyLilly
             uvw_vec = UVW(FT(0), FT(0), FT(0))
@@ -259,7 +241,6 @@ function precomputed_quantities(Y, atmos)
         advective_sgs_quantities...,
         edonly_quantities...,
         diagnostic_sgs_quantities...,
-        vert_diff_quantities...,
         sedimentation_quantities...,
         precipitation_quantities...,
         surface_precip_fluxes...,
@@ -538,12 +519,6 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
             ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
     end
 
-    # TODO: It is too slow to calculate mixing length at every timestep
-    # if isnothing(turbconv_model)
-    #     (; ᶜmixing_length) = p.precomputed
-    #     compute_gm_mixing_length!(ᶜmixing_length, Y, p)
-    # end
-
     if turbconv_model isa PrognosticEDMFX
         set_prognostic_edmf_precomputed_quantities_bottom_bc!(Y, p, t)
         set_prognostic_edmf_precomputed_quantities_explicit_closures!(Y, p, t)
@@ -585,19 +560,6 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
     )
     set_precipitation_surface_fluxes!(Y, p, p.atmos.microphysics_model)
 
-    if vertical_diffusion isa DecayWithHeightDiffusion
-        (; ᶜK_h) = p.precomputed
-        @. ᶜK_h =
-            $compute_eddy_diffusivity_coefficient(Y.c.ρ, vertical_diffusion)
-    elseif vertical_diffusion isa VerticalDiffusion
-        (; ᶜK_h) = p.precomputed
-        @. ᶜK_h = $compute_eddy_diffusivity_coefficient(
-            Y.c.uₕ,
-            ᶜp,
-            vertical_diffusion,
-        )
-    end
-
     # TODO
     if call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage
         set_cloud_fraction!(Y, p, moisture_model, cloud_model)

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -368,15 +368,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     n = n_mass_flux_subdomains(turbconv_model)
 
     (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
-    (;
-        ᶜmixing_length_tuple,
-        ᶜmixing_length,
-        ᶜlinear_buoygrad,
-        ᶜstrain_rate_norm,
-        ᶜK_u,
-        ᶜK_h,
-        ρatke_flux,
-    ) = p.precomputed
+    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ρatke_flux) = p.precomputed
     (;
         ᶜuʲs,
         ᶜtsʲs,
@@ -501,41 +493,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
-    ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
-    @. ᶜprandtl_nvec =
-        turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
-
-    ᶜ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
-
-    sfc_tke = Fields.level(ᶜtke⁰, 1)
-    @. ᶜmixing_length_tuple = mixing_length(
-        p.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,
-    )
-
-    @. ᶜmixing_length = ᶜmixing_length_tuple.master
-
-    @. ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
-    @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
-
     ρatke_flux_values = Fields.field_values(ρatke_flux)
     ρa_sfc_values = Fields.field_values(Fields.level(ᶜρa⁰, 1)) # TODO: replace by surface value
     ustar_values = Fields.field_values(ustar)

src/diagnostics/Diagnostics.jl

@@ -56,8 +56,15 @@ import ..SlabOceanSST
 
 # functions used to calculate diagnostics
 import ..draft_area
-import ..compute_gm_mixing_length!
+import ..compute_gm_mixing_length
 import ..horizontal_integral_at_boundary
+import ..ᶜmixing_length
+import ..ᶜeddy_diffusivity
+import ..ᶜeddy_viscosity
+import ..turbulent_prandtl_number
+import ..smagorinsky_lilly_length
+import ..compute_eddy_diffusivity_coefficient
+
 
 # We need the abbreviations for symbols like curl, grad, and so on
 include(joinpath("..", "utils", "abbreviations.jl"))