fully explicit kid

Author: haakon-e

config/default_configs/default_config.yml

@@ -426,5 +426,5 @@ use_itime:
   2. Time dependent forcing/tendencies use time rounded to the nearest unit of time for dt"
   value: false
 prescribed_flow:
-  help: "Prescribe a flow field [`nothing` (default), `true`]"
+  help: "Prescribe a flow field [`nothing` (default), `ShipwayHill2012`]"
   value: ~

config/model_configs/kinematic_driver.yml

@@ -1,7 +1,7 @@
 ## Model
-prescribed_flow: true
-initial_condition: "ShipwayHill2012"
-surface_setup: "ShipwayHill2012"
+prescribed_flow: ShipwayHill2012
+initial_condition: ShipwayHill2012
+surface_setup: ShipwayHill2012
 energy_q_tot_upwinding: first_order
 tracer_upwinding: first_order
 # moist: "nonequil"
@@ -13,8 +13,10 @@ call_cloud_diagnostics_per_stage: true
 config: "column"
 hyperdiff: "false"
 ## Simulation
-z_max: 2e3
-z_elem: 128
+# z_max: 2e3  # TODO: Update top boundary condition for ρu₃qₜ to allow lower z_max
+# z_elem: 128
+z_max: 8e3
+z_elem: 512
 z_stretch: false
 use_auto_jacobian: true # Needed!! (+only 1 Newton iteration, which is the default)
 # ode_algo: "ARS111" # try: Explicit Euler

src/cache/temporary_quantities.jl

@@ -117,5 +117,6 @@ function temporary_quantities(Y, atmos)
         ᶜdiffusion_u_matrix = similar(Y.c, TridiagonalMatrixRow{FT}),
         ᶜtridiagonal_matrix_scalar = similar(Y.c, TridiagonalMatrixRow{FT}),
         ᶠtridiagonal_matrix_c3 = similar(Y.f, TridiagonalMatrixRow{C3{FT}}),
+        (!isnothing(atmos.prescribed_flow) ? (; temp_Yₜ_imp = similar(Y)) : (;))...,
     )
 end

src/prognostic_equations/advection.jl

@@ -314,9 +314,6 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), energy_q_tot_upwinding)
         vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), Val(:none))
         vtt_bc = ᶜρq_tot_vertical_transport_bc(prescribed_flow, thermo_params, t, ᶠu³)
-        if prescribed_flow isa PrescribedFlow
-            vtt_central = NullBroadcasted()  # turn off implicit advection
-        end
         @. Yₜ.c.ρq_tot += vtt - vtt_central + vtt_bc
     end
 

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -131,10 +131,6 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     if !(moisture_model isa DryModel)
         ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
         vtt = vertical_transport(Y.c.ρ, ᶠu³, ᶜq_tot, dt, Val(:none))
-        (; prescribed_flow) = p.atmos
-        if prescribed_flow isa PrescribedFlow
-            vtt = NullBroadcasted()  # Turn off implicit energy q_tot advection
-        end
         @. Yₜ.c.ρq_tot += vtt
     end
 

src/solver/model_getters.jl

@@ -660,7 +660,6 @@ function check_case_consistency(parsed_args)
     turbconv = parsed_args["turbconv"]
     topography = parsed_args["topography"]
     prescribed_flow = parsed_args["prescribed_flow"]
-    auto_jacobian = parsed_args["use_auto_jacobian"]
 
     ISDAC_mandatory = (ic, subs, surf, rad, extf)
     if "ISDAC" in ISDAC_mandatory
@@ -677,14 +676,20 @@ function check_case_consistency(parsed_args)
             "Implicit vertical diffusion is only supported when using a " *
             "turbulence convection model or vertical diffusion model.",
         )
-    elseif prescribed_flow === true
+    elseif !isnothing(prescribed_flow)
         @assert(topography == "NoWarp",
             "Prescribed flow elides `set_velocity_at_surface!` and `set_velocity_at_top!` \
              which is needed for topography. Thus, prescribed flow must have flat surface."
         )
-        @assert(auto_jacobian === true, 
-            "Prescribed flow modifies the implicit treatment of sound waves, \
-             thus the manual sparse jacobian is not supported."
+        @assert(
+            !parsed_args["implicit_noneq_cloud_formation"] &&
+            !parsed_args["implicit_diffusion"] &&
+            !parsed_args["implicit_sgs_advection"] &&
+            !parsed_args["implicit_sgs_entr_detr"] &&
+            !parsed_args["implicit_sgs_nh_pressure"] &&
+            !parsed_args["implicit_sgs_vertdiff"] &&
+            !parsed_args["implicit_sgs_mass_flux"],
+            "Prescribed flow does not use the implicit solver."
         )
     end
 end

src/solver/type_getters.jl

@@ -114,7 +114,7 @@ function get_atmos(config::AtmosConfig, params)
         FT,
     )
 
-    if parsed_args["prescribed_flow"] === true
+    if parsed_args["prescribed_flow"] == "ShipwayHill2012"
         prescribed_flow = ShipwayHill2012VelocityProfile{FT}()
     else
         prescribed_flow = nothing
@@ -724,21 +724,45 @@ function get_sim_info(config::AtmosConfig)
     return sim
 end
 
+"""
+    fully_explicit_tendency!
+
+Experimental timestepping mode where all implicit tendencies are treated explicitly.
+"""
+function fully_explicit_tendency!(Yₜ, Yₜ_lim, Y, p, t)
+    (; temp_Yₜ_imp) = p.scratch
+    implicit_tendency!(temp_Yₜ_imp, Y, p, t)
+    remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
+    Yₜ .+= temp_Yₜ_imp
+end
+function fully_explicit_cache!(Y, p, t)
+    set_precomputed_quantities!(Y, p, t)
+    set_implicit_precomputed_quantities!(Y, p, t)
+end
+
+
 function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback)
     (; atmos, dt) = p
     s = @timed_str begin
-        T_imp! = SciMLBase.ODEFunction(
-            implicit_tendency!;
-            jac_prototype = get_jacobian(ode_algo, Y, atmos, parsed_args),
-            Wfact = update_jacobian!,
-        )
+        if isnothing(parsed_args["prescribed_flow"])
+            # This is the default case
+            T_exp_T_lim! = remaining_tendency!
+            T_imp! = SciMLBase.ODEFunction(implicit_tendency!;
+                jac_prototype = get_jacobian(ode_algo, Y, atmos, parsed_args),
+                Wfact = update_jacobian!,
+            )
+            cache! = set_precomputed_quantities!
+            cache_imp! = set_implicit_precomputed_quantities!
+        else
+            # `prescribed_flow` is an experimental case where the flow is prescribed,
+            # so implicit tendencies are treated explicitly to avoid treatment of sound waves
+            T_exp_T_lim! = fully_explicit_tendency!
+            T_imp! = nothing
+            cache! = fully_explicit_cache!
+            cache_imp! = nothing
+        end
         tendency_function = CTS.ClimaODEFunction(;
-            T_exp_T_lim! = remaining_tendency!,
-            T_imp!,
-            lim! = limiters_func!,
-            dss!,
-            cache! = set_precomputed_quantities!,
-            cache_imp! = set_implicit_precomputed_quantities!,
+            T_exp_T_lim!, T_imp!, cache!, cache_imp!, lim! = limiters_func!, dss!,
         )
     end
     @info "Define ode function: $s"
@@ -748,8 +772,7 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback)
     saveat = [t_begin, t_end]
     args = (problem, ode_algo)
     allow_custom_kwargs = (; kwargshandle = CTS.DiffEqBase.KeywordArgSilent)
-    kwargs =
-        (; saveat, callback, dt, adjustfinal = true, allow_custom_kwargs...)
+    kwargs = (; saveat, callback, dt, adjustfinal = true, allow_custom_kwargs...)
     return (args, kwargs)
 end