Merge pull request #3824 from CliMA/cc/ed_cleanup

Refactor and document ED pieces

Author: costachris

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-242
+243
 
 # **README**
 #
@@ -20,6 +20,13 @@
 
 
 #=
+243
+- Consolidate eddy diffusivity logic in eddy_diffusion_closures.jl and mass flux
+  logic in mass_flux_closures.jl, update buoyancy calculations to be consistent throughout
+  (use geopotential gradient instead of g parameter), add helper functions for computing
+  diffusivity/viscosity, adds l_grid to MixingLength output struct. Includes more docstrings
+  and comments.
+
 242
 - Use ustar^3 surface tke boundary condition
 

src/ClimaAtmos.jl

@@ -94,8 +94,8 @@ include(
 include(joinpath("prognostic_equations", "hyperdiffusion.jl"))
 include(joinpath("prognostic_equations", "gm_sgs_closures.jl"))
 include(joinpath("prognostic_equations", "edmf_coriolis.jl"))
-include(joinpath("prognostic_equations", "buoyancy_gradients.jl"))
-include(joinpath("prognostic_equations", "edmfx_closures.jl"))
+include(joinpath("prognostic_equations", "eddy_diffusion_closures.jl"))
+include(joinpath("prognostic_equations", "mass_flux_closures.jl"))
 include(joinpath("prognostic_equations", "edmfx_entr_detr.jl"))
 include(joinpath("prognostic_equations", "edmfx_tke.jl"))
 include(joinpath("prognostic_equations", "edmfx_sgs_flux.jl"))

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -304,7 +304,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     (; params) = p
     (; dt) = p
     dt = float(dt)
-    (; ᶜΦ) = p.core
+    (; ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core
     (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
     (; q_tot) = p.precomputed.ᶜspecific
     (;
@@ -434,6 +434,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
             q_totʲ_prev_level =
                 Fields.field_values(Fields.level(ᶜq_totʲ, i - 1))
             ρʲ_prev_level = Fields.field_values(Fields.level(ᶜρʲ, i - 1))
+            ᶜgradᵥ_ᶠΦ_prev_level =
+                Fields.field_values(Fields.level(ᶜgradᵥ_ᶠΦ, i - 1))
             tsʲ_prev_level = Fields.field_values(Fields.level(ᶜtsʲ, i - 1))
             entrʲ_prev_level = Fields.field_values(Fields.level(ᶜentrʲ, i - 1))
             detrʲ_prev_level = Fields.field_values(Fields.level(ᶜdetrʲ, i - 1))
@@ -482,7 +484,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                     local_geometry_prev_halflevel,
                 ),
                 TD.relative_humidity(thermo_params, tsʲ_prev_level),
-                ᶜphysical_buoyancy(thermo_params, ρ_prev_level, ρʲ_prev_level),
+                vertical_buoyancy_acceleration(
+                    ρ_prev_level,
+                    ρʲ_prev_level,
+                    ᶜgradᵥ_ᶠΦ_prev_level,
+                    local_geometry_prev_halflevel,
+                ),
                 get_physical_w(
                     u³_prev_halflevel,
                     local_geometry_prev_halflevel,
@@ -644,6 +651,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                 FT(0), # mass flux divergence is not implemented for diagnostic edmf
                 w_vert_div_level,
                 tke_prev_level,
+                ᶜgradᵥ_ᶠΦ_prev_level,
                 p.atmos.edmfx_model.detr_model,
             )
 
@@ -952,6 +960,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
         ᶜmixing_length,
     ) = p.precomputed
     (; ᶜK_h, ᶜK_u, ρatke_flux) = p.precomputed
+    turbconv_params = CAP.turbconv_params(params)
     thermo_params = CAP.thermodynamics_params(params)
     ᶜlg = Fields.local_geometry_field(Y.c)
 
@@ -1015,10 +1024,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     )
     @. ᶜmixing_length = ᶜmixing_length_tuple.master
 
-    turbconv_params = CAP.turbconv_params(params)
-    c_m = CAP.tke_ed_coeff(turbconv_params)
-    @. ᶜK_u = c_m * ᶜmixing_length * sqrt(max(ᶜtke⁰, 0))
-    @. ᶜK_h = ᶜK_u / ᶜprandtl_nvec
+    @. ᶜ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))

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -360,7 +360,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
 
     (; params) = p
     (; dt) = p
-    (; ᶠgradᵥ_ᶜΦ) = p.core
+    (; ᶠgradᵥ_ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core
     thermo_params = CAP.thermodynamics_params(params)
     turbconv_params = CAP.turbconv_params(params)
 
@@ -410,7 +410,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),
             get_physical_w(ᶜuʲs.:($$j), ᶜlg),
             TD.relative_humidity(thermo_params, ᶜtsʲs.:($$j)),
-            ᶜphysical_buoyancy(thermo_params, Y.c.ρ, ᶜρʲs.:($$j)),
+            vertical_buoyancy_acceleration(Y.c.ρ, ᶜρʲs.:($$j), ᶜgradᵥ_ᶠΦ, ᶜlg),
             get_physical_w(ᶜu, ᶜlg),
             TD.relative_humidity(thermo_params, ᶜts⁰),
             FT(0),
@@ -447,7 +447,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),
             get_physical_w(ᶜuʲs.:($$j), ᶜlg),
             TD.relative_humidity(thermo_params, ᶜtsʲs.:($$j)),
-            ᶜphysical_buoyancy(thermo_params, Y.c.ρ, ᶜρʲs.:($$j)),
+            vertical_buoyancy_acceleration(Y.c.ρ, ᶜρʲs.:($$j), ᶜgradᵥ_ᶠΦ, ᶜlg),
             get_physical_w(ᶜu, ᶜlg),
             TD.relative_humidity(thermo_params, ᶜts⁰),
             FT(0),
@@ -471,7 +471,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
         # term is still calculated here as it is used explicitly in the TKE equation.
         @. ᶠnh_pressure₃_buoyʲs.:($$j) = ᶠupdraft_nh_pressure_buoyancy(
             params,
-            ᶠbuoyancy(ᶠinterp(Y.c.ρ), ᶠinterp(ᶜρʲs.:($$j)), ᶠgradᵥ_ᶜΦ),
+            buoyancy(ᶠinterp(Y.c.ρ), ᶠinterp(ᶜρʲs.:($$j)), ᶠgradᵥ_ᶜΦ),
         )
     end
 
@@ -493,6 +493,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
         ᶜlg,
     )
 
+    # TODO: Make strain_rate_norm calculation a function in eddy_diffusion_closures
     # TODO: Currently the shear production only includes vertical gradients
     ᶠu⁰ = p.scratch.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
@@ -532,23 +533,18 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
 
     @. ᶜmixing_length = ᶜmixing_length_tuple.master
 
-    turbconv_params = CAP.turbconv_params(params)
-    c_m = CAP.tke_ed_coeff(turbconv_params)
-    @. ᶜK_u = c_m * ᶜmixing_length * sqrt(max(ᶜtke⁰, 0))
-    @. ᶜK_h = ᶜK_u / ᶜprandtl_nvec
+    @. ᶜ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(ᶜρa⁰, 1))
-    u_int_values = Fields.field_values(Fields.level(ᶜu, 1))
+    ρa_sfc_values = Fields.field_values(Fields.level(ᶜρa⁰, 1)) # TODO: replace by surface value
     ustar_values = Fields.field_values(ustar)
-    int_local_geometry_values =
-        Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1))
     sfc_local_geometry_values = Fields.field_values(
         Fields.level(Fields.local_geometry_field(Y.f), half),
     )
     @. ρatke_flux_values = surface_flux_tke(
         turbconv_params,
-        ρ_int_values,
+        ρa_sfc_values,
         ustar_values,
         sfc_local_geometry_values,
     )