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), `true`]"
+  value: ~

config/model_configs/kinematic_driver.yml

@@ -0,0 +1,41 @@
+## Model
+prescribed_flow: true
+initial_condition: "ShipwayHill2012"
+surface_setup: "ShipwayHill2012"
+energy_q_tot_upwinding: first_order_kid
+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
+z_elem: 128
+z_stretch: false
+use_auto_jacobian: true # Needed!! (+only 1 Newton iteration, which is the default)
+# 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:
+  - short_name: [
+      # (z,t) fields
+      ta, thetaa, ha, rhoa, wa, hur, hus, cl, clw, cli, ua, va, ke,
+      # (t,) fields
+      lwp, pr,
+      # 1M
+      # (z,t) fields
+      # husra, hussn,
+      # (t,) fields
+      # rwp,
+    ]
+    period: 10secs

docs/bibliography.bib

@@ -283,3 +283,16 @@ @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},
+}

examples/hybrid/KiD_driver.jl

@@ -0,0 +1,68 @@
+import ClimaComms
+import ClimaAtmos as CA
+import ClimaCore: Fields
+import YAML
+import ClimaComms
+
+import Random
+# Random.seed!(Random.MersenneTwister())
+Random.seed!(1234)
+
+# --> get config
+configs_path = joinpath(pkgdir(CA), "config/model_configs/")
+pth = joinpath(configs_path, "kinematic_driver.yml");
+job_id = "kinematic_driver";
+config_dict = YAML.load_file(pth)
+# <--
+
+
+config = CA.AtmosConfig(config_dict; job_id)
+simulation = CA.get_simulation(config);
+
+sol_res = CA.solve_atmos!(simulation);  # solve!
+
+(; integrator) = simulation;
+(; p) = integrator;
+(; atmos, params) = p;
+
+# --> Make ci plots
+# ]add ClimaAnalysis, ClimaCoreSpectra
+include(joinpath(pkgdir(CA), "post_processing", "ci_plots.jl"))
+make_plots(Val(:kinematic_driver), simulation.output_dir)
+# <--
+
+# --> ClimaAnalysis
+import ClimaAnalysis
+# using ClimaAnalysis.Visualize
+import ClimaAnalysis.Visualize as viz
+using ClimaAnalysis.Utils: kwargs
+using CairoMakie;
+CairoMakie.activate!();
+# using GLMakie; GLMakie.activate!()
+simdir = ClimaAnalysis.SimDir(simulation.output_dir);
+
+# entr = get(simdir; short_name = "entr")
+# entr.dims  # (time, x, y, z)
+
+# fig = Figure();
+# viz.plot!(fig, entr, time=0, x=0, y=0, more_kwargs = Dict(:axis => kwargs(dim_on_y = true)))
+# viz.plot!(fig, entr, x=0, y=0);
+# fig
+# <--
+
+
+
+#=
+%use upwinding for the rain -- yes!
+
+qr, qs, all specific humidities,
+
+take 30% acoustic Courant number c=300m/s, divided by vertical res
+- for ~200
+
+ill make some plots
+
+is smag applied to qr??? check this!
+
+most likely precip would cause blow-ups.
+=#

post_processing/ci_plots.jl

@@ -1689,3 +1689,35 @@ 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", "rhoa", "ta", "thetaa", "wa", "cli", "hussn", "ua", "va", "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

@@ -480,8 +480,10 @@ NVTX.@annotate function set_implicit_precomputed_quantities_part1!(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/initial_conditions/initial_conditions.jl

@@ -1633,3 +1633,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 = PrecipStateMassNum(; q_rai = FT(0), q_sno = FT(0)),
+        )
+    end
+    return local_state
+end

src/prognostic_equations/advection.jl

@@ -286,6 +286,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         ᶜ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))
+        if energy_q_tot_upwinding == Val(:first_order_kid)
+            vtt = vertical_transport_kid(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), FT(t))
+            vtt_central = NullBroadcasted()  # turn off implicit advection
+        end
         @. Yₜ.c.ρq_tot += vtt - vtt_central
     end
 

src/prognostic_equations/dss.jl

@@ -74,11 +74,41 @@ 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)
     if do_dss(axes(Y.c))
         Spaces.weighted_dss!(Y.c => p.ghost_buffer.c, Y.f => p.ghost_buffer.f)
     end
+    prescribe_flow!(Y, p, t, p.atmos.prescribed_flow)
     return nothing
 end
+
+prescribe_flow!(Y, p, t, ::Nothing) = nothing
+function prescribe_flow!(Y, p, t, flow::PrescribedFlow)
+    ᶠlg = Fields.local_geometry_field(Y.f)
+    z = Fields.coordinate_field(Y.f).z
+    @. Y.f.u₃ = C3(Geometry.WVector(flow(z, t)), ᶠlg)
+
+    # return nothing  # comment out to try fixing energy
+
+    ### 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)
+    grav = CAP.grav(p.params)
+    z = Fields.coordinate_field(Y.c).z
+
+    @. 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₃)
+    ᶜe_pot = @. lazy(grav * z)
+    @. Y.c.ρe_tot = Y.c.ρ * TD.total_energy(thermo_params, ᶜts, ᶜe_kin, ᶜe_pot)
+    return nothing
+end

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -82,6 +82,22 @@ function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:first_order})
     ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠupwind1(ᶠu³, ᶜχ))))
 end
+vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:first_order_kid}) =
+    vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, Val(:first_order))  # transport e.g. ρe_tot as usual
+function vertical_transport_kid(ᶜρ, ᶠu³, ᶜχ, dt, t)
+    FT = eltype(ᶜρ)
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
+    ρ_sfc = FT(1.1650101)  # TODO: Get from initial conditions or state
+    q_tot_sfc = FT(0.014778325)  # TODO: Get from initial conditions or state
+    u₃_sfc = ShipwayHill2012VelocityProfile{FT}()(FT(0), t)
+
+    bottom_bc = Geometry.WVector(ρ_sfc * q_tot_sfc * u₃_sfc)
+    ᶜadvdivᵥ_kid = Operators.DivergenceF2C(
+        bottom = Operators.SetValue(bottom_bc), top = Operators.Extrapolate(),
+    )
+    return @. lazy(-(ᶜadvdivᵥ_kid(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠupwind1(ᶠu³, ᶜχ))))
+end
 @static if pkgversion(ClimaCore) ≥ v"0.14.22"
     function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:vanleer_limiter})
         ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
@@ -131,6 +147,9 @@ 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))
+        if p.atmos.numerics.energy_q_tot_upwinding == Val(:first_order_kid)
+            vtt = NullBroadcasted()  # Turn off implicit energy q_tot advection
+        end
         @. Yₜ.c.ρq_tot += vtt
     end
 

src/solver/type_getters.jl

@@ -114,6 +114,12 @@ function get_atmos(config::AtmosConfig, params)
         FT,
     )
 
+    if parsed_args["prescribed_flow"]
+        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,
@@ -434,6 +443,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"])",