Test 1M EDMF (diag and prog) on GPU

NEWS.md

@@ -4,6 +4,11 @@ ClimaAtmos.jl Release Notes
 Main
 -------
 
+- Make prognostic and diagnostic EDMF work with 1-moment microphysics on GPU
+  PR [#3070](https://github.com/CliMA/ClimaAtmos.jl/pull/3070)
+- Add precipitation heating terms for 1-moment microphysics
+  PR [#3050](https://github.com/CliMA/ClimaAtmos.jl/pull/3050)
+
 v0.24.2
 -------
 

config/model_configs/gpu_diagnostic_edmfx_aquaplanet.yml

@@ -13,7 +13,7 @@ edmfx_nh_pressure: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 moist: equil
-precip_model: 0M
+precip_model: 1M
 dt: 100secs
 t_end: 6hours
-toml: [toml/diagnostic_edmfx_0M.toml]
+toml: [toml/diagnostic_edmfx.toml]

config/model_configs/gpu_prognostic_edmfx_aquaplanet.yml

@@ -1,22 +1,22 @@
-surface_setup: DefaultMoninObukhov 
+surface_setup: DefaultMoninObukhov
 rad: clearsky
-turbconv: prognostic_edmfx 
+turbconv: prognostic_edmfx
 prognostic_tke: true
 implicit_diffusion: true
 implicit_sgs_advection: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
-edmfx_upwinding: first_order 
-edmfx_entr_model: "Generalized" 
-edmfx_detr_model: "Generalized" 
-edmfx_nh_pressure: true 
+edmfx_upwinding: first_order
+edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
+edmfx_nh_pressure: true
 edmfx_filter: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
-moist: equil 
-precip_model: 0M
-dt: 10secs 
-t_end: 1hours 
+moist: equil
+precip_model: 1M
+dt: 10secs
+t_end: 1hours
 dt_save_to_disk: 600secs
 toml: [toml/prognostic_edmfx.toml]
 output_default_diagnostics: false

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -545,18 +545,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                 p_prev_level,
                 ρ_prev_level,
                 ρaʲ_prev_level,
-                draft_area(ρaʲ_prev_level, ρʲ_prev_level),
-                get_physical_w(
-                    u³ʲ_prev_halflevel,
-                    local_geometry_prev_halflevel,
-                ),
-                TD.relative_humidity(thermo_params, tsʲ_prev_level),
-                ᶜphysical_buoyancy(params, ρ_prev_level, ρʲ_prev_level),
-                get_physical_w(
-                    u³_prev_halflevel,
-                    local_geometry_prev_halflevel,
-                ),
-                TD.relative_humidity(thermo_params, ts_prev_level),
+                tsʲ_prev_level,
+                ρʲ_prev_level,
+                u³ʲ_prev_halflevel,
+                local_geometry_prev_halflevel,
+                u³_prev_halflevel,
+                ts_prev_level,
                 FT(0),
                 entrʲ_prev_level,
                 vert_div_level,

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -119,7 +119,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
         # Based on boundary conditions for updrafts we overwrite
         # the first interior point for EDMFX ᶜmseʲ...
         ᶜaʲ_int_val = p.scratch.temp_data_level
-        # TODO: replace this with the actual surface area fraction when 
+        # TODO: replace this with the actual surface area fraction when
         # using prognostic surface area
         @. ᶜaʲ_int_val = FT(turbconv_params.surface_area)
         ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1))

src/parameterized_tendencies/microphysics/microphysics_wrappers.jl

@@ -162,11 +162,13 @@ function compute_precipitation_sources!(
     mp,
     thp,
 )
-    FT = eltype(Sqₜᵖ)
-    @. Sqₜᵖ = FT(0)
-    @. Sqᵣᵖ = FT(0)
-    @. Sqₛᵖ = FT(0)
-    @. Seₜᵖ = FT(0)
+    FT = eltype(thp)
+    # @. Sqₜᵖ = FT(0) should work after fixing
+    # https://github.com/CliMA/ClimaCore.jl/issues/1786
+    @. Sqₜᵖ = ρ * FT(0)
+    @. Sqᵣᵖ = ρ * FT(0)
+    @. Sqₛᵖ = ρ * FT(0)
+    @. Seₜᵖ = ρ * FT(0)
 
     #! format: off
     # rain autoconversion: q_liq -> q_rain

src/prognostic_equations/edmfx_entr_detr.jl

@@ -111,6 +111,49 @@ function entrainment(
     return entr
 end
 
+function detrainment(
+    params,
+    z_prev_level::FT,
+    z_sfc_halflevel,
+    p_prev_level,
+    ρ_prev_level,
+    ρaʲ_prev_level,
+    tsʲ_prev_level,
+    ρʲ_prev_level,
+    u³ʲ_prev_halflevel,
+    local_geometry_prev_halflevel,
+    u³_prev_halflevel,
+    ts_prev_level,
+    ᶜbuoy⁰,
+    entrʲ_prev_level,
+    vert_div_level,
+    ᶜmassflux_vert_div, # mass flux divergence is not implemented for diagnostic edmf
+    tke_prev_level,
+    edmfx_detr_model,
+) where {FT}
+    thermo_params = CAP.thermodynamics_params(params)
+    detrainment(
+        params,
+        z_prev_level,
+        z_sfc_halflevel,
+        p_prev_level,
+        ρ_prev_level,
+        ρaʲ_prev_level,
+        draft_area(ρaʲ_prev_level, ρʲ_prev_level),
+        get_physical_w(u³ʲ_prev_halflevel, local_geometry_prev_halflevel),
+        TD.relative_humidity(thermo_params, tsʲ_prev_level),
+        ᶜphysical_buoyancy(params, ρ_prev_level, ρʲ_prev_level),
+        get_physical_w(u³_prev_halflevel, local_geometry_prev_halflevel),
+        TD.relative_humidity(thermo_params, ts_prev_level),
+        FT(0),
+        entrʲ_prev_level,
+        vert_div_level,
+        FT(0), # mass flux divergence is not implemented for diagnostic edmf
+        tke_prev_level,
+        edmfx_detr_model,
+    )
+end
+
 """
    Return detrainment rate [1/s].