ft: Add Kinematic Driver (KiD) model configuration

config/default_configs/default_config.yml

@@ -422,3 +422,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,25 @@
+## Model
+prescribed_flow: true
+initial_condition: "ShipwayHill2012"
+surface_setup: "ShipwayHill2012"
+tracer_upwinding: first_order
+moist: "nonequil"
+cloud_model: "grid_scale"
+call_cloud_diagnostics_per_stage: true
+precip_model: "1M"
+config: "column"
+## Simulation
+z_max: 4e3
+z_elem: 100
+z_stretch: false
+dt: "1secs"
+t_end: "1hours"
+# t_end: "10secs"
+dt_save_state_to_disk: "1hours"
+toml: [toml/kinematic_driver.toml]
+check_nan_every: 1
+## Diagnostics
+output_default_diagnostics: false
+diagnostics:
+  - short_name: [ta, thetaa, ha, rhoa, wa, hur, hus, cl, clw, cli, husra, hussn]
+    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,70 @@
+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"))
+ref_job_id = config.parsed_args["reference_job_id"]
+reference_job_id = isnothing(ref_job_id) ? simulation.job_id : ref_job_id
+make_plots(Val(Symbol(reference_job_id)), 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.
+=#

src/cache/precipitation_precomputed_quantities.jl

@@ -905,14 +905,10 @@ function set_precipitation_surface_fluxes!(
     @. ᶜq_sno = specific(Y.c.ρq_sno, Y.c.ρ)
     @. ᶜq_liq = specific(Y.c.ρq_liq, Y.c.ρ)
     @. ᶜq_ice = specific(Y.c.ρq_ice, Y.c.ρ)
-    sfc_qᵣ =
-        Fields.Field(Fields.field_values(Fields.level(ᶜq_rai, 1)), sfc_space)
-    sfc_qₛ =
-        Fields.Field(Fields.field_values(Fields.level(ᶜq_sno, 1)), sfc_space)
-    sfc_qₗ =
-        Fields.Field(Fields.field_values(Fields.level(ᶜq_liq, 1)), sfc_space)
-    sfc_qᵢ =
-        Fields.Field(Fields.field_values(Fields.level(ᶜq_ice, 1)), sfc_space)
+    sfc_qᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜq_rai, 1)), sfc_space)
+    sfc_qₛ = Fields.Field(Fields.field_values(Fields.level(ᶜq_sno, 1)), sfc_space)
+    sfc_qₗ = Fields.Field(Fields.field_values(Fields.level(ᶜq_liq, 1)), sfc_space)
+    sfc_qᵢ = Fields.Field(Fields.field_values(Fields.level(ᶜq_ice, 1)), sfc_space)
     sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space)
     sfc_wₛ = Fields.Field(Fields.field_values(Fields.level(ᶜwₛ, 1)), sfc_space)
     sfc_wₗ = Fields.Field(Fields.field_values(Fields.level(ᶜwₗ, 1)), sfc_space)

src/initial_conditions/initial_conditions.jl

@@ -1633,3 +1633,47 @@ 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).
+"""
+Base.@kwdef struct ShipwayHill2012 <: InitialCondition end
+
+function (initial_condition::ShipwayHill2012)(params)
+    FT = eltype(params)
+
+    ## Initialize the profile
+    z_0, z_1, z_2 = FT(0), FT(740), FT(3260)  # m
+    rv_0, rv_1, rv_2 = FT(0.015), FT(0.0138), FT(0.0024)  # kg/kg
+    θ_0, θ_1, θ_2 = FT(297.9), FT(297.9), FT(312.66)  # K
+
+    # profile of water vapour mixing ratio
+    linear_profile(z₀, z₁, x₀, x₁, z) = x₀ + (x₁ - x₀) / (z₁ - z₀) * (z - z₀)
+    rv(z) = if z < z_1
+        linear_profile(z_0, z_1, rv_0, rv_1, z)
+    else
+        linear_profile(z_1, z_2, rv_1, rv_2, z)
+    end
+    q_tot(z) = rv(z) / (1 + rv(z))
+    # profile of potential temperature
+    θ(z) = z < z_1 ? θ_0 : linear_profile(z_1, z_2, θ_1, θ_2, z)
+    ## 
+    thermo_params = CAP.thermodynamics_params(params)
+    p_0 = FT(100700) # 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

@@ -279,7 +279,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         )
         vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, FT(dt), energy_q_tot_upwinding)
         vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, FT(dt), Val(:none))
-        @. Yₜ.c.ρe_tot += vtt - vtt_central
+        # @. Yₜ.c.ρe_tot += vtt - vtt_central
     end
 
     if !(p.atmos.moisture_model isa DryModel) && energy_q_tot_upwinding != Val(:none)
@@ -289,35 +289,35 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         @. Yₜ.c.ρq_tot += vtt - vtt_central
     end
 
-    if isnothing(ᶠf¹²)
-        # shallow atmosphere
-        @. Yₜ.c.uₕ -=
-            ᶜinterp(ᶠω¹² × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) / (Y.c.ρ * ᶜJ) +
-            (ᶜf³ + ᶜω³) × CT12(ᶜu)
-        @. Yₜ.f.u₃ -= ᶠω¹² × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
-        for j in 1:n
-            @. Yₜ.f.sgsʲs.:($$j).u₃ -=
-                ᶠω¹²ʲs.:($$j) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
-                ᶠgradᵥ(ᶜKʲs.:($$j) - ᶜinterp(ᶠKᵥʲs.:($$j)))
-        end
-    else
-        # deep atmosphere
-        @. Yₜ.c.uₕ -=
-            ᶜinterp((ᶠf¹² + ᶠω¹²) × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) /
-            (Y.c.ρ * ᶜJ) + (ᶜf³ + ᶜω³) × CT12(ᶜu)
-        @. Yₜ.f.u₃ -= (ᶠf¹² + ᶠω¹²) × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
-        for j in 1:n
-            @. Yₜ.f.sgsʲs.:($$j).u₃ -=
-                (ᶠf¹² + ᶠω¹²ʲs.:($$j)) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
-                ᶠgradᵥ(ᶜKʲs.:($$j) - ᶜinterp(ᶠKᵥʲs.:($$j)))
-        end
-    end
-
-    if use_prognostic_tke(turbconv_model) # advect_tke triggers allocations
-        @. ᶜa_scalar = ᶜtke⁰ * draft_area(ᶜρa⁰, ᶜρ⁰)
-        vtt = vertical_transport(ᶜρ⁰, ᶠu³⁰, ᶜa_scalar, dt, edmfx_mse_q_tot_upwinding)
-        @. Yₜ.c.sgs⁰.ρatke += vtt
-    end
+    # if isnothing(ᶠf¹²)
+    #     # shallow atmosphere
+    #     @. Yₜ.c.uₕ -=
+    #         ᶜinterp(ᶠω¹² × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) / (Y.c.ρ * ᶜJ) +
+    #         (ᶜf³ + ᶜω³) × CT12(ᶜu)
+    #     @. Yₜ.f.u₃ -= ᶠω¹² × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
+    #     for j in 1:n
+    #         @. Yₜ.f.sgsʲs.:($$j).u₃ -=
+    #             ᶠω¹²ʲs.:($$j) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
+    #             ᶠgradᵥ(ᶜKʲs.:($$j) - ᶜinterp(ᶠKᵥʲs.:($$j)))
+    #     end
+    # else
+    #     # deep atmosphere
+    #     @. Yₜ.c.uₕ -=
+    #         ᶜinterp((ᶠf¹² + ᶠω¹²) × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) /
+    #         (Y.c.ρ * ᶜJ) + (ᶜf³ + ᶜω³) × CT12(ᶜu)
+    #     @. Yₜ.f.u₃ -= (ᶠf¹² + ᶠω¹²) × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
+    #     for j in 1:n
+    #         @. Yₜ.f.sgsʲs.:($$j).u₃ -=
+    #             (ᶠf¹² + ᶠω¹²ʲs.:($$j)) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
+    #             ᶠgradᵥ(ᶜKʲs.:($$j) - ᶜinterp(ᶠKᵥʲs.:($$j)))
+    #     end
+    # end
+
+    # if use_prognostic_tke(turbconv_model) # advect_tke triggers allocations
+    #     @. ᶜa_scalar = ᶜtke⁰ * draft_area(ᶜρa⁰, ᶜρ⁰)
+    #     vtt = vertical_transport(ᶜρ⁰, ᶠu³⁰, ᶜa_scalar, dt, edmfx_mse_q_tot_upwinding)
+    #     @. Yₜ.c.sgs⁰.ρatke += vtt
+    # end
 end
 
 """

src/prognostic_equations/dss.jl

@@ -77,8 +77,22 @@ dss!(Y_state, params, t_current)
 """
 
 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
+    # 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)
+    FT = eltype(p.params)
+    # (; uₕ, u₃, ρ, p) = flow
+    (; prescribed_u₃) = flow
+    lg = Fields.local_geometry_field(Y.f)
+    @. Y.f.u₃ = C3(Geometry.WVector(prescribed_u₃(FT, t)), lg)
+    # @. Y.c.uₕ = uₕ(t)
+    # @. Y.c.ρ = ρ₀  # move to IC
+    # @. Y.c.p = p₀  # move to IC
+    return nothing
+end

src/prognostic_equations/remaining_tendency.jl

@@ -64,10 +64,10 @@ Returns:
 NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     Yₜ_lim .= zero(eltype(Yₜ_lim))
     Yₜ .= zero(eltype(Yₜ))
-    horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t)
+    # horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t)
     fill_with_nans!(p)  # TODO: would be better to limit this to debug mode (e.g., if p.debug_mode...)
-    horizontal_dynamics_tendency!(Yₜ, Y, p, t)
-    hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
+    # horizontal_dynamics_tendency!(Yₜ, Y, p, t)
+    # hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     vertical_advection_of_water_tendency!(Yₜ, Y, p, t)
     additional_tendency!(Yₜ, Y, p, t)

src/prognostic_equations/surface_flux.jl

@@ -109,18 +109,12 @@ function surface_flux_tendency!(Yₜ, Y, p, t)
     thermo_params = CAP.thermodynamics_params(params)
 
     if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion)
-        btt =
-            boundary_tendency_momentum(Y.c.ρ, Y.c.uₕ, sfc_conditions.ρ_flux_uₕ)
+        btt = boundary_tendency_momentum(Y.c.ρ, Y.c.uₕ, sfc_conditions.ρ_flux_uₕ)
         @. Yₜ.c.uₕ -= btt
     end
 
-    ᶜh_tot = @. lazy(
-        TD.total_specific_enthalpy(
-            thermo_params,
-            ᶜts,
-            specific(Y.c.ρe_tot, Y.c.ρ),
-        ),
-    )
+    ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
+    ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜe_tot))
     btt = boundary_tendency_scalar(ᶜh_tot, sfc_conditions.ρ_flux_h_tot)
     @. Yₜ.c.ρe_tot -= btt
     ρ_flux_χ = p.scratch.sfc_temp_C3

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -105,9 +105,7 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion)
         ᶠstrain_rate = compute_strain_rate_face_vertical(ᶜu)
-        @. Yₜ.c.uₕ -= C12(
-            ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠstrain_rate) / Y.c.ρ,
-        ) # assumes ᶜK_u = ᶜK_h
+        @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠstrain_rate) / Y.c.ρ)  # assumes ᶜK_u = ᶜK_h
     end
 
     ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(

src/solver/type_getters.jl

@@ -72,8 +72,7 @@ function get_atmos(config::AtmosConfig, params)
         forcing_type isa HeldSuarezForcing ? forcing_type : radiation_mode
     # Note: when disable_momentum_vertical_diffusion is true, the surface flux tendency
     # for momentum is not applied.
-    disable_momentum_vertical_diffusion =
-        final_radiation_mode isa HeldSuarezForcing
+    disable_momentum_vertical_diffusion = final_radiation_mode isa HeldSuarezForcing
 
     advection_test = parsed_args["advection_test"]
     @assert advection_test in (false, true)
@@ -108,12 +107,21 @@ function get_atmos(config::AtmosConfig, params)
     )
 
     vertical_diffusion = get_vertical_diffusion_model(
-        disable_momentum_vertical_diffusion,
-        parsed_args,
-        params,
-        FT,
+        disable_momentum_vertical_diffusion, parsed_args, params, FT,
     )
 
+    if parsed_args["prescribed_flow"]
+        function prescribed_u₃(FT::Type{<:Real}, _t)
+            t = FT(_t)
+            w1 = FT(2)
+            t1 = FT(600)  # 10 minutes
+            return t < t1 ? w1 * sin(π * t / t1) : FT(0)
+        end
+        prescribed_flow = PrescribedFlow(; prescribed_u₃)
+    else
+        prescribed_flow = nothing
+    end
+
     atmos = AtmosModel(;
         # AtmosWater - Moisture, Precipitation & Clouds
         moisture_model,
@@ -130,6 +138,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,
@@ -433,6 +444,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"])",