Merge pull request #4104 from CliMA/ne/vanleer

Fix vertical_transport with vanleer_limiter on GPU

Author: szy21

config/model_configs/diagnostic_edmfx_aquaplanet_gpu.yml

@@ -9,6 +9,8 @@ edmfx_detr_model: "Generalized"
 edmfx_nh_pressure: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
+edmfx_sgsflux_upwinding: "vanleer_limiter"
+edmfx_tracer_upwinding: "vanleer_limiter"
 moist: equil
 cloud_model: "quadrature_sgs"
 precip_model: 0M

post_processing/jacobian_summary.jl

@@ -8,7 +8,7 @@ sparse approximations against the dense matrix.
 function print_jacobian_summary(integrator)
     Y = integrator.u
     t = integrator.t
-    (; p, dt) = integrator
+    (; p) = integrator
     timestepper_alg = integrator.alg
     tableau_coefficients =
         timestepper_alg isa CA.CTS.RosenbrockAlgorithm ?
@@ -19,7 +19,7 @@ function print_jacobian_summary(integrator)
 
     FT = eltype(Y)
     γs = filter(!iszero, CA.LinearAlgebra.diag(tableau_coefficients))
-    dtγ = FT(dt) * FT(γs[end])
+    dtγ = p.dt * FT(γs[end])
     scalar_names = CA.scalar_field_names(Y)
     block_keys = Iterators.product(scalar_names, scalar_names)
 

src/cache/cache.jl

@@ -1,5 +1,5 @@
 struct AtmosCache{
-    TT,
+    FT,
     AM,
     NUM,
     CAP,
@@ -20,7 +20,7 @@ struct AtmosCache{
     CONSCHECK,
 }
     """Timestep of the simulation (in seconds). This is also used by callbacks and tendencies"""
-    dt::TT
+    dt::FT
 
     """AtmosModel"""
     atmos::AM
@@ -94,6 +94,7 @@ function build_cache(
 )
     (; dt, start_date, output_dir) = sim_info
     FT = eltype(params)
+    dt = FT(dt)
 
     ᶜcoord = Fields.local_geometry_field(Y.c).coordinates
     ᶠcoord = Fields.local_geometry_field(Y.f).coordinates

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -449,7 +449,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     ᶜdz = Fields.Δz_field(axes(Y.c))
     (; params) = p
     (; dt) = p
-    dt = FT(dt)
     (; ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core
     (; ᶜp, ᶠu³, ᶜts, ᶜK) = p.precomputed
     (;

src/cache/precipitation_precomputed_quantities.jl

@@ -580,8 +580,6 @@ sources from the sub-domains.
 set_precipitation_cache!(Y, p, _, _) = nothing
 function set_precipitation_cache!(Y, p, ::Microphysics0Moment, _)
     (; params, dt) = p
-    FT = eltype(p.params)
-    dt = FT(dt)
     (; ᶜts) = p.precomputed
     (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed
     (; ᶜΦ) = p.core

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -283,7 +283,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
                 buoyancy_flux_val < 0 || ᶜaʲ_int_val >= $(FT(turbconv_params.surface_area)),
                 entr_int_val,
                 detr_int_val +
-                ($(FT(turbconv_params.surface_area)) / ᶜaʲ_int_val - 1) / FT(dt),
+                ($(FT(turbconv_params.surface_area)) / ᶜaʲ_int_val - 1) / dt,
             ),
             ᶜaʲ_int_val,
             dt,

src/callbacks/callbacks.jl

@@ -124,7 +124,7 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     Y = integrator.u
     p = integrator.p
     t = integrator.t
-
+    FT = eltype(Y)
     (; params) = p
     (; ᶠradiation_flux, rrtmgp_model) = p.radiation
     (; radiation_mode) = p.atmos
@@ -134,7 +134,7 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     set_insolation_variables!(Y, p, t, p.atmos.insolation)
     set_surface_albedo!(Y, p, t, p.atmos.surface_albedo)
 
-    RRTMGPI.update_fluxes!(rrtmgp_model, UInt32(floor(t / integrator.p.dt)))
+    RRTMGPI.update_fluxes!(rrtmgp_model, UInt32(floor(FT(t) / integrator.p.dt)))
     Fields.field2array(ᶠradiation_flux) .= rrtmgp_model.face_flux
     return nothing
 end

src/parameterized_tendencies/microphysics/microphysics_wrappers.jl

@@ -25,15 +25,14 @@ end
 # Helper function to compute the limit of the tendency in the traingle limiter.
 # The limit is defined as the available amont / n times model time step.
 function limit(q, dt, n::Int)
-    FT = eltype(q)
-    return q / FT(dt) / n
+    return q / dt / n
 end
 
 function moisture_fixer(q, qᵥ, dt)
     FT = eltype(q)
     return triangle_inequality_limiter(
-        -min(FT(0), q / FT(dt)),
-        limit(qᵥ, FT(dt), 5),
+        -min(FT(0), q / dt),
+        limit(qᵥ, dt, 5),
         FT(0),
     )
 end
@@ -161,10 +160,9 @@ Returns the qₜ source term due to precipitation formation
 defined as Δm_tot / (m_dry + m_tot) for the 0-moment scheme
 """
 function q_tot_0M_precipitation_sources(thp, cmp::CMP.Parameters0M, dt, qₜ, ts)
-    FT = eltype(qₜ)
     return -triangle_inequality_limiter(
         -CM0.remove_precipitation(cmp, PP(thp, ts)),
-        qₜ / FT(dt),
+        qₜ / dt,
     )
 end
 
@@ -570,7 +568,7 @@ function aerosol_activation_sources(
     return ifelse(
         S_max < S || isnan(n_act) || n_act < nₗ,
         FT(0),
-        (n_act - nₗ) / FT(dt),
+        (n_act - nₗ) / dt,
     )
 end
 

src/prognostic_equations/advection.jl

@@ -265,7 +265,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name
             if !(ρχ_name in (@name(ρe_tot), @name(ρq_tot)))
                 ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-                vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, FT(dt), tracer_upwinding)
+                vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, tracer_upwinding)
                 @. ᶜρχₜ += vtt
             end
         end
@@ -280,15 +280,15 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
                 specific(Y.c.ρe_tot, Y.c.ρ),
             ),
         )
-        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, FT(dt), energy_q_tot_upwinding)
-        vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, FT(dt), Val(:none))
+        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, dt, energy_q_tot_upwinding)
+        vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, dt, Val(:none))
         @. Yₜ.c.ρe_tot += vtt - vtt_central
     end
 
     if !(p.atmos.moisture_model isa DryModel) && energy_q_tot_upwinding != Val(:none)
         ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
-        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), energy_q_tot_upwinding)
-        vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), Val(:none))
+        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, dt, energy_q_tot_upwinding)
+        vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, dt, Val(:none))
         @. Yₜ.c.ρq_tot += vtt - vtt_central
     end
 

src/prognostic_equations/edmfx_entr_detr.jl

@@ -664,7 +664,7 @@ function edmfx_first_interior_entr_tendency!(
         sgsʲs_ρaₜ_int_val = Fields.field_values(Fields.level(Yₜ.c.sgsʲs.:($j).ρa, 1))
         @. sgsʲs_ρaₜ_int_val += ifelse(buoyancy_flux_val < 0,
             0,
-            max(0, (sgsʲs_ρ_int_val * $(eps(FT)) - sgsʲs_ρa_int_val) / FT(dt)),
+            max(0, (sgsʲs_ρ_int_val * $(eps(FT)) - sgsʲs_ρa_int_val) / dt),
         )
 
         # Apply entrainment tendencies in the first model cell for moist static energy (mse) 
@@ -718,16 +718,16 @@ end
 limit_entrainment(entr::FT, a, dt) where {FT} = max(
     min(
         entr,
-        FT(0.9) * (1 - a) / max(a, eps(FT)) / FT(dt),
-        FT(0.9) * 1 / FT(dt),
+        FT(0.9) * (1 - a) / max(a, eps(FT)) / dt,
+        FT(0.9) * 1 / dt,
     ),
     0,
 )
 limit_detrainment(detr::FT, a, dt) where {FT} =
-    max(min(detr, FT(0.9) * 1 / FT(dt)), 0)
+    max(min(detr, FT(0.9) * 1 / dt), 0)
 
 function limit_turb_entrainment(dyn_entr::FT, turb_entr, dt) where {FT}
-    return max(min((FT(0.9) * 1 / FT(dt)) - dyn_entr, turb_entr), 0)
+    return max(min((FT(0.9) * 1 / dt) - dyn_entr, turb_entr), 0)
 end
 
 # limit entrainment and detrainment rates for diagnostic EDMF