ft: Add Kinematic Driver (KiD) model configuration

config/default_configs/default_config.yml

@@ -425,3 +425,6 @@ use_itime:
   1. Surface conditions that explicitly depend on time (e.g. LifeCycleTan2018, TRMM_LBA, etc.),
   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), `ShipwayHill2012`]"
+  value: ~

config/model_configs/kinematic_driver.yml

@@ -0,0 +1,42 @@
+## Model
+prescribed_flow: ShipwayHill2012
+initial_condition: ShipwayHill2012
+surface_setup: ShipwayHill2012
+energy_q_tot_upwinding: first_order
+tracer_upwinding: first_order
+# moist: "nonequil"
+# precip_model: "1M"
+moist: "equil"
+precip_model: "0M"
+cloud_model: "grid_scale"
+call_cloud_diagnostics_per_stage: true
+config: "column"
+hyperdiff: "false"
+## Simulation
+# 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
+# ode_algo: "ARS111" # try: Explicit Euler
+dt: "1secs"
+# t_end: "1hours"
+t_end: "20mins"
+dt_save_state_to_disk: "1hours"
+toml: [toml/kinematic_driver.toml]
+check_nan_every: 1
+## Diagnostics
+output_default_diagnostics: false
+diagnostics:
+  # core: (z,t) fields
+  - short_name: [ta, thetaa, ha, rhoa, wa, hur, hus, cl, clw, cli, ua, va, ke]
+    period: 10secs
+  # core: (t,) fields
+  - short_name: [lwp, pr]
+    period: 10secs
+  # 1M: (z,t) fields
+  - short_name: [husra, hussn]
+    period: 10secs
+  # 1M: (t,) fields
+  - short_name: [rwp]
+    period: 10secs

docs/bibliography.bib

@@ -283,3 +283,15 @@ @article{Wing2018
 url = {https://gmd.copernicus.org/articles/11/793/2018/},
 doi = {10.5194/gmd-11-793-2018}
 }
+
+@article{ShipwayHill2012,
+	title = {Diagnosis of systematic differences between multiple parametrizations of warm rain microphysics using a kinematic framework},
+	volume = {138},
+	url = {https://onlinelibrary.wiley.com/doi/abs/10.1002/qj.1913},
+	doi = {10.1002/qj.1913},
+	pages = {2196--2211},
+	number = {669},
+	journaltitle = {Quarterly Journal of the Royal Meteorological Society},
+	author = {Shipway, B. J. and Hill, A. A.},
+	date = {2012},
+}

post_processing/ci_plots.jl

@@ -1550,3 +1550,40 @@ function make_plots(
         output_name = "summary_3D",
     )
 end
+
+function make_plots(::Val{:kinematic_driver}, output_paths::Vector{<:AbstractString})
+    function rescale_time_to_min(var)
+        if haskey(var.dims, "time")
+            var.dims["time"] .= var.dims["time"] ./ 60
+            var.dim_attributes["time"]["units"] = "min"
+        end
+        return var
+    end
+    simdirs = SimDir.(output_paths)
+    short_names = [
+        "hus", "clw", "husra", "ta", #"thetaa", "rhoa",
+        "wa",
+        # "cli", "hussn",
+        # "ke",
+    ]
+    short_names = short_names ∩ collect(keys(simdirs[1].vars))
+    vars = map_comparison(simdirs, short_names) do simdir, short_name
+        var = slice(get(simdir; short_name), x = 0, y = 0)
+        if short_name in ["hus", "clw", "husra", "cli", "hussn"]
+            var.data .= var.data .* 1000
+            var.attributes["units"] = "g/kg"
+        end
+        return rescale_time_to_min(var)
+    end
+    file_contour = make_plots_generic(output_paths, vars;
+        output_name = "tmp_contour",
+    )
+
+    short_names_lines = ["lwp", "rwp", "pr"]
+    short_names_lines = short_names_lines ∩ collect(keys(simdirs[1].vars))
+    vars_lines = map_comparison(simdirs, short_names_lines) do simdir, short_name
+        var = slice(get(simdir; short_name), x = 0, y = 0)
+        return rescale_time_to_min(var)
+    end
+    make_plots_generic(output_paths, vars_lines; summary_files = [file_contour])
+end

src/cache/precomputed_quantities.jl

@@ -468,8 +468,10 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
 
     # TODO: We might want to move this to dss! (and rename dss! to something
     # like enforce_constraints!).
-    set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
-    set_velocity_at_top!(Y, turbconv_model)
+    if !(p.atmos.prescribed_flow isa PrescribedFlow)
+        set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
+        set_velocity_at_top!(Y, turbconv_model)
+    end
 
     set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶠuₕ³)
     ᶜJ = Fields.local_geometry_field(Y.c).J

src/cache/temporary_quantities.jl

@@ -116,5 +116,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/initial_conditions/initial_conditions.jl

@@ -1859,3 +1859,43 @@ function (initial_condition::ISDAC)(params)
         )
     end
 end
+
+"""
+    ShipwayHill2012
+
+The `InitialCondition` described in [ShipwayHill2012](@cite), but with a hydrostatically
+balanced pressure profile.
+
+B. J. Shipway and A. A. Hill. 
+Diagnosis of systematic differences between multiple parametrizations of warm rain microphysics using a kinematic framework. 
+Quarterly Journal of the Royal Meteorological Society 138, 2196-2211 (2012).
+"""
+struct ShipwayHill2012 <: InitialCondition end
+function (initial_condition::ShipwayHill2012)(params)
+    FT = eltype(params)
+
+    ## Initialize the profile
+    z_values = FT[0, 740, 3260]
+    rv_values = FT[0.015, 0.0138, 0.0024]  # water vapor mixing ratio
+    θ_values = FT[297.9, 297.9, 312.66]    # potential temperature
+    linear_profile(zs, vals) = Intp.extrapolate(
+        Intp.interpolate((zs,), vals, Intp.Gridded(Intp.Linear())), Intp.Linear(),
+    )
+    # profile of water vapour mixing ratio
+    rv(z) = linear_profile(z_values, rv_values)(z)
+    q_tot(z) = rv(z) / (1 + rv(z))
+    # profile of potential temperature
+    θ(z) = linear_profile(z_values, θ_values)(z)
+    ## Hydrostatically balanced pressure profile
+    thermo_params = CAP.thermodynamics_params(params)
+    p_0 = FT(100_700) # Pa
+    p = hydrostatic_pressure_profile(; thermo_params, p_0, θ, q_tot)
+    function local_state(local_geometry)
+        (; z) = local_geometry.coordinates
+        return LocalState(; params, geometry = local_geometry,
+            thermo_state = TD.PhaseEquil_pθq(thermo_params, p(z), θ(z), q_tot(z)),
+            precip_state = NoPrecipState{typeof(z)}(),
+        )
+    end
+    return local_state
+end

src/prognostic_equations/advection.jl

@@ -164,6 +164,30 @@ NVTX.@annotate function horizontal_tracer_advection_tendency!(Yₜ, Y, p, t)
     return nothing
 end
 
+"""
+    ᶜρq_tot_vertical_transport_bc(flow, thermo_params, t, ᶠu³)
+
+Computes the vertical transport of `ρq_tot` at the surface due to prescribed flow.
+
+If the flow is not prescribed, this has no effect.
+
+# Arguments
+- `flow`: The prescribed flow model, see [`PrescribedFlow`](@ref).
+    - If `flow` is `nothing`, this has no effect.
+- `thermo_params`: The thermodynamic parameters, needed to compute surface air density.
+- `t`: The current time.
+- `ᶠu³`: The vertical velocity field.
+
+# Returns
+- The vertical transport of `ρq_tot` at the surface due to prescribed flow.
+"""
+ᶜρq_tot_vertical_transport_bc(::Nothing, _, _, _) = NullBroadcasted()
+function ᶜρq_tot_vertical_transport_bc(flow::PrescribedFlow, thermo_params, t, ᶠu³)
+    ρu₃qₜ_sfc_bc = get_ρu₃qₜ_surface(flow, thermo_params, t)
+    ᶜadvdivᵥ = Operators.DivergenceF2C(; bottom = Operators.SetValue(ρu₃qₜ_sfc_bc))
+    return @. lazy(-(ᶜadvdivᵥ(zero(ᶠu³))))
+end
+
 """
     explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
 
@@ -193,7 +217,7 @@ Modifies `Yₜ.c` (various tracers, `ρe_tot`, `ρq_tot`, `uₕ`), `Yₜ.f.u₃`
 `Yₜ.f.sgsʲs` (updraft `u₃`), and `Yₜ.c.sgs⁰.ρatke` as applicable.
 """
 NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
-    (; turbconv_model) = p.atmos
+    (; turbconv_model, prescribed_flow) = p.atmos
     n = n_prognostic_mass_flux_subdomains(turbconv_model)
     advect_tke = use_prognostic_tke(turbconv_model)
     point_type = eltype(Fields.coordinate_field(Y.c))
@@ -290,6 +314,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         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
+        if prescribed_flow isa PrescribedFlow
+            vtt_bc = ᶜρq_tot_vertical_transport_bc(prescribed_flow, thermo_params, t, ᶠu³)
+            @. Yₜ.c.ρq_tot += vtt_bc
+        end
     end
 
     if isnothing(ᶠf¹²)

src/prognostic_equations/dss.jl

@@ -74,11 +74,37 @@ See also:
 dss!(Y_state, params, t_current)
 # The ClimaCore.Field objects within Y_state.c and Y_state.f are now updated
 # with DSS applied, ensuring continuity across distributed elements.
+```
 """
 
-NVTX.@annotate function dss!(Y, p, t)
+NVTX.@annotate function dss!(Y, p, t)  # TODO: Rename to e.g. `apply_constraints!`
+    prescribe_flow!(Y, p, t, p.atmos.prescribed_flow)
     if do_dss(axes(Y.c))
         Spaces.weighted_dss!(Y.c => p.ghost_buffer.c, Y.f => p.ghost_buffer.f)
     end
     return nothing
 end
+
+prescribe_flow!(_, _, _, ::Nothing) = nothing
+function prescribe_flow!(Y, p, t, flow::PrescribedFlow)
+    (; ᶜΦ) = p.core
+    ᶠlg = Fields.local_geometry_field(Y.f)
+    z = Fields.coordinate_field(Y.f).z
+    @. Y.f.u₃ = C3(Geometry.WVector(flow(z, t)), ᶠlg)
+
+    ### Fix energy to initial temperature
+    ᶜlg = Fields.local_geometry_field(Y.c)
+    local_state = InitialConditions.ShipwayHill2012()(p.params)
+    get_ρ_init_dry(ls) = ls.thermo_state.ρ * (1 - ls.thermo_state.q_tot)
+    get_T_init(ls) = TD.air_temperature(ls.thermo_params, ls.thermo_state)
+    ᶜρ_init_dry = @. lazy(get_ρ_init_dry(local_state(ᶜlg)))
+    ᶜT_init = @. lazy(get_T_init(local_state(ᶜlg)))
+
+    thermo_params = CAP.thermodynamics_params(p.params)
+
+    @. Y.c.ρ = ᶜρ_init_dry + Y.c.ρq_tot
+    ᶜts = @. lazy(TD.PhaseEquil_ρTq(thermo_params, Y.c.ρ, ᶜT_init, Y.c.ρq_tot / Y.c.ρ))
+    ᶜe_kin = compute_kinetic(Y.c.uₕ, Y.f.u₃)
+    @. Y.c.ρe_tot = Y.c.ρ * TD.total_energy(thermo_params, ᶜts, ᶜe_kin, ᶜΦ)
+    return nothing
+end

src/solver/model_getters.jl

@@ -655,6 +655,8 @@ function check_case_consistency(parsed_args)
     imp_vert_diff = parsed_args["implicit_diffusion"]
     vert_diff = parsed_args["vert_diff"]
     turbconv = parsed_args["turbconv"]
+    topography = parsed_args["topography"]
+    prescribed_flow = parsed_args["prescribed_flow"]
 
     ISDAC_mandatory = (ic, subs, surf, rad, extf)
     if "ISDAC" in ISDAC_mandatory
@@ -671,5 +673,20 @@ function check_case_consistency(parsed_args)
             "Implicit vertical diffusion is only supported when using a " *
             "turbulence convection model or vertical diffusion model.",
         )
+    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(
+            !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,6 +114,12 @@ function get_atmos(config::AtmosConfig, params)
         FT,
     )
 
+    if parsed_args["prescribed_flow"] == "ShipwayHill2012"
+        prescribed_flow = ShipwayHill2012VelocityProfile{FT}()
+    else
+        prescribed_flow = nothing
+    end
+
     atmos = AtmosModel(;
         # AtmosWater - Moisture, Precipitation & Clouds
         moisture_model,
@@ -131,6 +137,9 @@ function get_atmos(config::AtmosConfig, params)
         advection_test,
         scm_coriolis = get_scm_coriolis(parsed_args, FT),
 
+        # PrescribedFlow
+        prescribed_flow,
+
         # AtmosRadiation
         radiation_mode = final_radiation_mode,
         ozone,
@@ -348,6 +357,8 @@ function get_initial_condition(parsed_args, atmos)
             parsed_args["start_date"],
             parsed_args["era5_initial_condition_dir"],
         )
+    elseif parsed_args["initial_condition"] == "ShipwayHill2012"
+        return ICs.ShipwayHill2012()
     else
         error(
             "Unknown `initial_condition`: $(parsed_args["initial_condition"])",
@@ -659,21 +670,40 @@ 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 args_integrator(parsed_args, Y, p, tspan, ode_algo, callback, dt_integrator)
     (; atmos) = 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_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_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! = set_precomputed_quantities!, cache_imp!,
+            lim! = limiters_func!, dss!,
         )
     end
     @info "Define ode function: $s"
@@ -685,8 +715,7 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback, dt_integr
     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