move nh pressure to edmf precomputed quantities

Author: szy21

.buildkite/longruns_gpu/pipeline.yml

@@ -140,7 +140,7 @@ steps:
           CLIMACOMMS_DEVICE: "CUDA"
           JOB_NAME: "longrun_aquaplanet_allsky_diagedmf_0M"
 
-      - label: ":computer: aquaplanet equilmoist allsky radiation + prognostic edmf diffusion only + 0M microphysics"
+      - label: ":computer: aquaplanet equilmoist allsky radiation + prognostic edmf + 0M microphysics"
         command:
           - srun julia --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME
         artifact_paths: "$$JOB_NAME/output_active/*"

.buildkite/pipeline.yml

@@ -785,10 +785,6 @@ steps:
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
           --config_file $CONFIG_PATH/prognostic_edmfx_bomex_implicit_column.yml
           --job_id prognostic_edmfx_bomex_column_implicit
-
-          julia --color=yes --project=.buildkite reproducibility_tests/test_mse.jl
-          --job_id prognostic_edmfx_bomex_column_implicit
-          --out_dir prognostic_edmfx_bomex_column_implicit/output_active
         artifact_paths: "prognostic_edmfx_bomex_column_implicit/output_active/*"
         agents:
           slurm_mem: 20GB

config/longrun_configs/longrun_aquaplanet_allsky_progedmf_0M.yml

@@ -25,7 +25,6 @@ edmfx_entr_model: "Generalized"
 edmfx_detr_model: "Generalized"
 edmfx_nh_pressure: true
 edmfx_filter: true
-prognostic_tke: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 precip_model: "0M"

config/model_configs/prognostic_edmfx_aquaplanet.yml

@@ -17,16 +17,15 @@ edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 implicit_diffusion: true
 implicit_sgs_advection: true
-# This job crashes when setting the following to true. There is potentially a bug.
-# implicit_sgs_entr_detr: true
-# implicit_sgs_nh_pressure: true
-# implicit_sgs_mass_flux: true
+implicit_sgs_entr_detr: true
+implicit_sgs_nh_pressure: true
+implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 moist: equil
 cloud_model: "quadrature_sgs"
 precip_model: 0M
-dt: 10secs
+dt: 20secs
 t_end: 1hours
 dt_save_state_to_disk: 600secs
 toml: [toml/prognostic_edmfx.toml]

config/model_configs/prognostic_edmfx_aquaplanet_gpu.yml

@@ -17,16 +17,15 @@ edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
 implicit_diffusion: true
 implicit_sgs_advection: true
-# This job crashes when setting the following to true. There is potentially a bug.
-# implicit_sgs_entr_detr: true
-# implicit_sgs_nh_pressure: true
-# implicit_sgs_mass_flux: true
+implicit_sgs_entr_detr: true
+implicit_sgs_nh_pressure: true
+implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 moist: equil
 cloud_model: "quadrature_sgs"
 precip_model: 0M
-dt: 10secs
+dt: 20secs
 t_end: 1hours
 dt_save_state_to_disk: 600secs
 toml: [toml/prognostic_edmfx.toml]

config/model_configs/prognostic_edmfx_bomex_implicit_column.yml

@@ -31,7 +31,6 @@ perturb_initstate: false
 dt: "50secs"
 t_end: "6hours"
 dt_save_state_to_disk: "10mins"
-reproducibility_test: true
 toml: [toml/prognostic_edmfx_bomex.toml]
 netcdf_interpolation_num_points: [2, 2, 60]
 diagnostics:

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-224
+225
 
 # **README**
 #
@@ -21,6 +21,10 @@
 
 #=
 
+225
+- Move nonhydrostatic pressure drag calculation to precomputed quantities and 
+  remove one reproducibility job
+
 224
 - Machine precision differences due to https://github.com/CliMA/ClimaCore.jl/pull/2232
 

reproducibility_tests/reproducibility_test_job_ids.jl

@@ -5,6 +5,5 @@ reproducibility_test_job_ids() = [
     "aquaplanet_equil_allsky_gw_raw_zonalasym",
     "diagnostic_edmfx_aquaplanet",
     "single_column_hydrostatic_balance_ft64",
-    "prognostic_edmfx_bomex_column_implicit",
     "prognostic_edmfx_aquaplanet",
 ]

src/cache/precomputed_quantities.jl

@@ -55,6 +55,7 @@ function implicit_precomputed_quantities(Y, atmos)
         turbconv_model isa AbstractEDMF ? (; ᶜtke⁰ = similar(Y.c, FT)) : (;)
     moisture_sgs_quantities =
         (
+            turbconv_model isa PrognosticEDMFX &&
             moisture_model isa NonEquilMoistModel &&
             precip_model isa Microphysics1Moment
         ) ?
@@ -82,6 +83,7 @@ function implicit_precomputed_quantities(Y, atmos)
             ᶠKᵥʲs = similar(Y.f, NTuple{n, FT}),
             ᶜtsʲs = similar(Y.c, NTuple{n, TST}),
             ᶜρʲs = similar(Y.c, NTuple{n, FT}),
+            ᶠnh_pressure₃ʲs = similar(Y.f, NTuple{n, C3{FT}}),
             moisture_sgs_quantities...,
         ) : (;)
     return (; gs_quantities..., sgs_quantities..., prognostic_sgs_quantities...)
@@ -167,7 +169,6 @@ function precomputed_quantities(Y, atmos)
             ᶜentrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜdetrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜturb_entrʲs = similar(Y.c, NTuple{n, FT}),
-            ᶠnh_pressure₃ʲs = similar(Y.f, NTuple{n, C3{FT}}),
             ᶜgradᵥ_θ_virt⁰ = Fields.Field(C3{FT}, cspace),
             ᶜgradᵥ_q_tot⁰ = Fields.Field(C3{FT}, cspace),
             ᶜgradᵥ_θ_liq_ice⁰ = Fields.Field(C3{FT}, cspace),

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -330,7 +330,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
     FT = eltype(params)
     n = n_mass_flux_subdomains(turbconv_model)
 
-    (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
+    (; ᶠgradᵥ_ᶜΦ) = p.core
+    (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶠu₃⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
     (;
         ᶜmixing_length_tuple,
         ᶜmixing_length,
@@ -340,13 +341,23 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
         ᶜK_h,
         ρatke_flux,
     ) = p.precomputed
-    (; ᶜuʲs, ᶜtsʲs, ᶠu³ʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p.precomputed
+    (;
+        ᶜuʲs,
+        ᶜtsʲs,
+        ᶠu³ʲs,
+        ᶜρʲs,
+        ᶜentrʲs,
+        ᶜdetrʲs,
+        ᶜturb_entrʲs,
+        ᶠnh_pressure₃ʲs,
+    ) = p.precomputed
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
 
     ᶜz = Fields.coordinate_field(Y.c).z
     z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)
     ᶜdz = Fields.Δz_field(axes(Y.c))
     ᶜlg = Fields.local_geometry_field(Y.c)
+    ᶠlg = Fields.local_geometry_field(Y.f)
 
     ᶜvert_div = p.scratch.ᶜtemp_scalar
     ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2
@@ -417,6 +428,24 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),
             dt,
         )
+
+        # nonhydrostatic pressure drag
+        scale_height =
+            CAP.R_d(params) * CAP.T_surf_ref(params) / CAP.grav(params)
+        for j in 1:n
+            if p.atmos.edmfx_model.nh_pressure isa Val{true}
+                @. ᶠnh_pressure₃ʲs.:($$j) = ᶠupdraft_nh_pressure(
+                    params,
+                    ᶠlg,
+                    ᶠbuoyancy(ᶠinterp(Y.c.ρ), ᶠinterp(ᶜρʲs.:($$j)), ᶠgradᵥ_ᶜΦ),
+                    Y.f.sgsʲs.:($$j).u₃,
+                    ᶠu₃⁰,
+                    scale_height,
+                )
+            else
+                @. ᶠnh_pressure₃ʲs.:($$j) = C3(0)
+            end
+        end
     end
 
     (; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed