ft: Add Kinematic Driver (KiD) model configuration

Author: haakon-e

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,39 @@
+## Model
+prescribed_flow: true
+initial_condition: "ShipwayHill2012"
+surface_setup: "ShipwayHill2012"
+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: 4e3
+z_elem: 200
+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:
+  - short_name: [
+      # (z,t) fields
+      ta, thetaa, ha, rhoa, wa, hur, hus, cl, clw, cli,
+      # (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", "ta", "thetaa", "wa", "cli", "hussn"]
+    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/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/dss.jl

@@ -74,11 +74,42 @@ 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)
+    FT = eltype(p.params)
+    (; prescribed_u₃) = flow
+    ᶠlg = Fields.local_geometry_field(Y.f)
+    @. Y.f.u₃ = C3(Geometry.WVector(prescribed_u₃(FT, 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(ls) = TD.air_density(ls.thermo_params, ls.thermo_state)
+    get_T_init(ls) = TD.air_temperature(ls.thermo_params, ls.thermo_state)
+    ᶜρ_init = @. lazy(get_ρ_init(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 + 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/solver/type_getters.jl

@@ -114,6 +114,18 @@ function get_atmos(config::AtmosConfig, 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 +142,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 +448,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"])",

src/solver/types.jl

@@ -6,6 +6,11 @@ import Dates
 import ClimaUtilities.ClimaArtifacts: @clima_artifact
 import LazyArtifacts
 
+abstract type AbstractPrescribedFlow end
+@kwdef struct PrescribedFlow
+    prescribed_u₃::Function
+end
+
 abstract type AbstractMoistureModel end
 abstract type AbstractMoistModel <: AbstractMoistureModel end
 struct DryModel <: AbstractMoistureModel end
@@ -681,11 +686,12 @@ Base.broadcastable(x::AtmosGravityWave) = tuple(x)
 Base.broadcastable(x::AtmosSponge) = tuple(x)
 Base.broadcastable(x::AtmosSurface) = tuple(x)
 
-struct AtmosModel{W, SCM, R, TC, GW, VD, SP, SU, NU}
+struct AtmosModel{W, SCM, R, TC, PF, GW, VD, SP, SU, NU}
     water::W
     scm_setup::SCM
     radiation::R
     turbconv::TC
+    prescribed_flow::PF
     gravity_wave::GW
     vertical_diffusion::VD
     sponge::SP
@@ -701,6 +707,7 @@ const ATMOS_MODEL_GROUPS = (
     (AtmosWater, :water),
     (AtmosRadiation, :radiation),
     (AtmosTurbconv, :turbconv),
+    (PrescribedFlow, :prescribed_flow),
     (AtmosGravityWave, :gravity_wave),
     (AtmosSponge, :sponge),
     (AtmosSurface, :surface),
@@ -918,11 +925,14 @@ function AtmosModel(; kwargs...)
     disable_surface_flux_tendency =
         get(atmos_model_kwargs, :disable_surface_flux_tendency, false)
 
+    prescribed_flow = get(atmos_model_kwargs, :prescribed_flow, nothing)
+
     return AtmosModel{
         typeof(water),
         typeof(scm_setup),
         typeof(radiation),
         typeof(turbconv),
+        typeof(prescribed_flow),
         typeof(gravity_wave),
         typeof(vertical_diffusion),
         typeof(sponge),
@@ -933,6 +943,7 @@ function AtmosModel(; kwargs...)
         scm_setup,
         radiation,
         turbconv,
+        prescribed_flow,
         gravity_wave,
         vertical_diffusion,
         sponge,

src/surface_conditions/surface_setups.jl

@@ -242,6 +242,27 @@ function (::TRMM_LBA)(params)
     return surface_state
 end
 
+struct ShipwayHill2012 end
+function (::ShipwayHill2012)(params)
+    function surface_state(surface_coordinates, interior_z, t)
+        FT = eltype(surface_coordinates)
+        T = FT(297.9)  # surface temperature (K)
+        p = FT(100700)  # surface pressure (Pa)
+        rv₀ = FT(0.015)  # water vapour mixing ratio at surface (kg/kg)
+        q_vap = rv₀ / (1 + rv₀)  # specific humidity at surface, assuming unsaturated surface air (kg/kg)
+        # Exchange coefficients, from Rico, above >>>
+        # Adjust the coefficients from 20 m to the actual value of z.
+        z0 = FT(1.5e-4)
+        adjustment = (log(20 / z0) / log(interior_z / z0))^2
+        Cd = FT(0.001229) * adjustment
+        Ch = FT(0.001094) * adjustment
+        parameterization = ExchangeCoefficients(; Cd, Ch)
+        # <<<
+        return SurfaceState(; parameterization, T, p, q_vap)
+    end
+    return surface_state
+end
+
 struct SimplePlume end
 function (::SimplePlume)(params)
     FT = eltype(params)