Try #983:

.buildkite/pipeline.yml

@@ -161,6 +161,10 @@ steps:
         command: "julia --color=yes --project=examples examples/hybrid/driver.jl --z_max 45e3 --z_elem 25 --vert_diff true --surface_scheme bulk --moist equil --forcing held_suarez --microphy 0M --rayleigh_sponge true --alpha_rayleigh_uh 0 --viscous_sponge true --zd_rayleigh 30e3 --zd_viscous 30e3 --job_id sphere_held_suarez_rhoe_equilmoist_hightop_sponge --dt 450secs --t_end 1days"
         artifact_paths: "sphere_held_suarez_rhoe_equilmoist_hightop_sponge/*"
 
+      - label: ":computer: held suarez (ρe) equilmoist hightop gw parameterization"
+        command: "julia --color=yes --project=examples examples/hybrid/driver.jl --z_max 45e3 --z_elem 50 --dz_top 3e3 --vert_diff true --surface_scheme bulk --moist equil --forcing held_suarez --microphy 0M --non_orographic_gravity_wave true --job_id sphere_held_suarez_rhoe_equilmoist_hightop_gw --dt 450secs --t_end 1days --regression_test true"
+        artifact_paths: "sphere_held_suarez_rhoe_equilmoist_hightop_gw/*"
+
       - label: ":computer: held suarez (ρe_tot) equilmoist monin_obukhov"
         command: "julia --color=yes --project=examples examples/hybrid/driver.jl --vert_diff true --surface_scheme monin_obukhov --moist equil --forcing held_suarez --microphy 0M --job_id sphere_held_suarez_rhoe_equilmoist_mo --dt 600secs"
         artifact_paths: "sphere_held_suarez_rhoe_equilmoist_mo/*"

examples/hybrid/cli_options.jl

@@ -31,7 +31,7 @@ function parse_commandline()
         arg_type = String
         default = "6hours"
         "--config" # TODO: add box
-        help = "Spatial configuration [`sphere` (default), `column`]"
+        help = "Spatial configuration [`sphere` (default), `column`, `box`]"
         arg_type = String
         default = "sphere"
         "--moist"

examples/hybrid/driver.jl

@@ -160,8 +160,8 @@ function additional_cache(Y, params, model_spec, dt; use_tempest_mode = false)
             diffuse_momentum,
             coupled,
         ) : NamedTuple(),
-        model_spec.non_orographic_gravity_wave ? gravity_wave_cache(Y, FT) :
-        NamedTuple(),
+        model_spec.non_orographic_gravity_wave ?
+        gravity_wave_cache(model_spec.model_config, Y, FT) : NamedTuple(),
         (;
             tendency_knobs = (;
                 hs_forcing = forcing_type isa HeldSuarezForcing,

examples/hybrid/gravitywave_parameterization/gravity_wave_parameterization.jl

@@ -1,4 +1,5 @@
 function gravity_wave_cache(
+    ::SingleColumnModel,
     Y,
     ::Type{FT};
     source_height = FT(15000),
@@ -13,10 +14,57 @@ function gravity_wave_cache(
 
     nc = Int(floor(FT(2 * cmax / dc + 1)))
     c = [FT((n - 1) * dc - cmax) for n in 1:nc]
-
+    gw_F_S0 = similar(Fields.level(Y.c.ρ, 1))  # there must be a better way...
+    gw_F_S0 .= F_S0
     return (;
         gw_source_height = source_height,
-        gw_F_S0 = F_S0,
+        gw_source_ampl = gw_F_S0,
+        gw_Bm = Bm,
+        gw_c = c,
+        gw_cw = cw,
+        gw_c0 = c0,
+        gw_nk = length(kwv),
+        gw_k = kwv,
+        gw_k2 = kwv .^ 2,
+        ᶜbuoyancy_frequency = similar(Y.c.ρ),
+        ᶜdTdz = similar(Y.c.ρ),
+    )
+end
+
+function gravity_wave_cache(
+    ::SphericalModel,
+    Y,
+    ::Type{FT};
+    source_pressure = FT(3e4),
+    Bm = FT(0.4), # as in GFDL code
+    Bt_0 = FT(0.0003),
+    Bt_n = FT(0.0003),
+    Bt_s = FT(0.0003),
+    ϕ0_n = FT(30),
+    ϕ0_s = FT(-30),
+    dϕ_n = FT(5),
+    dϕ_s = FT(-5),
+    dc = FT(0.6),
+    cmax = FT(99.6),
+    c0 = FT(0),
+    kwv = FT(2π / 100e5),
+    cw = FT(40.0),
+) where {FT}
+
+    nc = Int(floor(FT(2 * cmax / dc + 1)))
+    c = [FT((n - 1) * dc - cmax) for n in 1:nc]
+
+    ᶜlocal_geometry = Fields.local_geometry_field(Fields.level(Y.c, 1))
+    lat = ᶜlocal_geometry.coordinates.lat
+    source_ampl = @. Bt_0 +
+       Bt_n * FT(0.5) * (FT(1) + tanh((lat - ϕ0_n) / dϕ_n)) +
+       Bt_s * FT(0.5) * (FT(1) + tanh((lat - ϕ0_s) / dϕ_s))
+
+    # compute spatio variant source F_S0
+
+    return (;
+        gw_source_pressure = source_pressure,
+        gw_source_ampl = source_ampl,
         gw_Bm = Bm,
         gw_c = c,
         gw_cw = cw,
@@ -31,18 +79,13 @@ end
 
 function gravity_wave_tendency!(Yₜ, Y, p, t)
     #unpack
-    (;
-        gw_source_height,
-        gw_Bm,
-        gw_F_S0,
-        gw_c,
-        gw_cw,
-        gw_c0,
-        gw_nk,
-        gw_k,
-        gw_k2,
-    ) = p
-    (; ᶜts, ᶜT, ᶜdTdz, ᶜbuoyancy_frequency, params) = p
+    (; ᶜts, ᶜT, ᶜdTdz, ᶜbuoyancy_frequency, params, model_config) = p
+    (; gw_Bm, gw_source_ampl, gw_c, gw_cw, gw_c0, gw_nk, gw_k, gw_k2) = p
+    if model_config isa SingleColumnModel
+        (; gw_source_height) = p
+    elseif model_config isa SphericalModel
+        (; gw_source_pressure) = p
+    end
     ᶜρ = Y.c.ρ
     # parameters
     thermo_params = CAP.thermodynamics_params(params)
@@ -67,29 +110,49 @@ function gravity_wave_tendency!(Yₜ, Y, p, t)
     # .     may be calculated
     #
 
-    # source level: get the index of the level that is closest to the source height (GFDL uses the fist level below instead)
-    source_level = argmin(
-        abs.(
-            unique(parent(Fields.coordinate_field(Y.c).z) .- gw_source_height)
-        ),
-    )
+    # source level: get the index of the level that is closest to the source height/pressure (GFDL uses the fist level below instead)
+    if model_config isa SingleColumnModel
+        source_level = similar(Fields.level(Y.c.ρ, 1))
+        Fields.bycolumn(axes(ᶜρ)) do colidx
+            parent(source_level[colidx]) .= argmin(
+                abs.(
+                    parent(Fields.coordinate_field(Y.c).z[colidx]) .-
+                    gw_source_height
+                ),
+            )[1]
+        end
+    elseif model_config isa SphericalModel
+        (; ᶜK, ᶜp) = p
+        thermo_params = CAP.thermodynamics_params(params)
+        ᶜuₕ = Y.c.uₕ
+        ᶠw = Y.f.w
+        @. ᶜK = norm_sqr(C123(ᶜuₕ) + C123(ᶜinterp(ᶠw))) / 2
+        CA.thermo_state!(Y, p, ᶜinterp)
+        @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
+
+        source_level = similar(Fields.level(Y.c.ρ, 1))
+        Fields.bycolumn(axes(ᶜρ)) do colidx
+            parent(source_level[colidx]) .=
+                argmin(abs.(parent(ᶜp[colidx]) .- gw_source_pressure))[1]
+        end
+    end
     ᶠz = Fields.coordinate_field(Y.f).z
 
-    # TODO: prepare input variables for gravity_wave_forcing()
-    #       1. grab column wise data
-    #       2. convert coviant uv to physical uv
-    # need to make everything an array befere
+    # prepare physical uv input variables for gravity_wave_forcing()
     u_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:1
     v_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:2
 
+    # a place holder to store physical forcing on uv
     uforcing = ones(axes(u_phy))
     vforcing = similar(v_phy)
-    # bycolume
+
+    # GW parameterization applied bycolume
     Fields.bycolumn(axes(ᶜρ)) do colidx
         parent(uforcing[colidx]) .= gravity_wave_forcing(
+            model_config,
             parent(u_phy[colidx]),
-            source_level,
-            gw_F_S0,
+            Int(parent(source_level[colidx])[1]),
+            parent(gw_source_ampl[colidx])[1],
             gw_Bm,
             gw_c,
             gw_cw,
@@ -102,9 +165,10 @@ function gravity_wave_tendency!(Yₜ, Y, p, t)
             parent(ᶠz[colidx]),
         )
         parent(vforcing[colidx]) .= gravity_wave_forcing(
+            model_config,
             parent(v_phy[colidx]),
-            source_level,
-            gw_F_S0,
+            Int(parent(source_level[colidx])[1]),
+            parent(gw_source_ampl[colidx])[1],
             gw_Bm,
             gw_c,
             gw_cw,
@@ -119,15 +183,17 @@ function gravity_wave_tendency!(Yₜ, Y, p, t)
 
     end
 
+    # physical uv forcing converted to Covariant12Vector and added up to uₕ tendencies
     @. Yₜ.c.uₕ +=
         Geometry.Covariant12Vector.(Geometry.UVVector.(uforcing, vforcing))
 
 end
 
 function gravity_wave_forcing(
+    model_config,
     ᶜu,
     source_level,
-    F_S0,
+    source_ampl,  # F_S0 for single columne and source_ampl (F/ρ) as GFDL for sphere
     Bm,
     c,
     cw,
@@ -156,8 +222,8 @@ function gravity_wave_forcing(
     if (Bsum == 0.0)
         error("zero flux input at source level")
     end
-    # define the intermittency factor eps. Assumes constant Δc.
-    eps = (F_S0 / ᶜρ[source_level] / nk) / Bsum # in GFDL code: eps = (ampl * 1.5 / nk) / Bsum, ampl is equivalent to F_S)/ρ_source 
+    eps =
+        calc_intermitency(model_config, ᶜρ[source_level], source_ampl, nk, Bsum)
 
     ᶜdz = ᶠz[2:end] - ᶠz[1:(end - 1)]
     wave_forcing = zeros(nc)
@@ -238,3 +304,25 @@ function gravity_wave_forcing(
 
     return gwf
 end
+
+# calculate the intermittency factor eps -> assuming constant Δc.
+
+function calc_intermitency(
+    ::SingleColumnModel,
+    ρ_source_level,
+    source_ampl,
+    nk,
+    Bsum,
+)
+    return (source_ampl / ρ_source_level / nk) / Bsum
+end
+
+function calc_intermitency(
+    ::SphericalModel,
+    ρ_source_level,
+    source_ampl,
+    nk,
+    Bsum,
+)
+    return source_ampl * 1.5 / nk / Bsum
+end

examples/hybrid/gravitywave_parameterization/gw_test_3d.jl

@@ -3,6 +3,7 @@ using Dates
 using Interpolations
 using Statistics
 using Plots
+import ClimaAtmos: SingleColumnModel, SphericalModel
 include("gravity_wave_parameterization.jl")
 const FT = Float64
 
@@ -11,6 +12,7 @@ const FT = Float64
 
 face_z = 0:1e3:0.47e5
 center_z = 0.5 .* (face_z[1:(end - 1)] .+ face_z[2:end])
+model_config = SingleColumnModel()
 
 # compute the source parameters
 function gravity_wave_cache(
@@ -125,6 +127,7 @@ Jan_ρ = mean(center_ρ_zonalave[:, :, month .== 1], dims = 3)[:, :, 1]
 Jan_uforcing = zeros(length(lat), length(center_z))
 for j in 1:length(lat)
     Jan_uforcing[j, :] = gravity_wave_forcing(
+        model_config,
         Jan_u[j, :],
         source_level,
         params.gw_F_S0,

examples/hybrid/gravitywave_parameterization/gw_test_single_column.jl

@@ -3,6 +3,7 @@ using Dates
 using Interpolations
 using Statistics
 using Plots
+import ClimaAtmos: SingleColumnModel, SphericalModel
 include("gravity_wave_parameterization.jl")
 
 const FT = Float64
@@ -12,6 +13,7 @@ const FT = Float64
 
 face_z = FT.(0:1e3:0.5e5)
 center_z = FT(0.5) .* (face_z[1:(end - 1)] .+ face_z[2:end])
+model_config = SingleColumnModel()
 
 # compute the source parameters
 function gravity_wave_cache(
@@ -128,6 +130,7 @@ Jan_u = mean(center_u_mean[:, month .== 1], dims = 2)[:, 1]
 Jan_bf = mean(center_bf_mean[:, month .== 1], dims = 2)[:, 1]
 Jan_ρ = mean(center_ρ_mean[:, month .== 1], dims = 2)[:, 1]
 Jan_uforcing = gravity_wave_forcing(
+    model_config,
     Jan_u,
     source_level,
     params.gw_F_S0,
@@ -152,6 +155,7 @@ April_u = mean(center_u_mean[:, month .== 4], dims = 2)[:, 1]
 April_bf = mean(center_bf_mean[:, month .== 4], dims = 2)[:, 1]
 April_ρ = mean(center_ρ_mean[:, month .== 4], dims = 2)[:, 1]
 April_uforcing = gravity_wave_forcing(
+    model_config,
     April_u,
     source_level,
     params.gw_F_S0,
@@ -176,6 +180,7 @@ July_u = mean(center_u_mean[:, month .== 7], dims = 2)[:, 1]
 July_bf = mean(center_bf_mean[:, month .== 7], dims = 2)[:, 1]
 July_ρ = mean(center_ρ_mean[:, month .== 7], dims = 2)[:, 1]
 July_uforcing = gravity_wave_forcing(
+    model_config,
     July_u,
     source_level,
     params.gw_F_S0,
@@ -200,6 +205,7 @@ Oct_u = mean(center_u_mean[:, month .== 10], dims = 2)[:, 1]
 Oct_bf = mean(center_bf_mean[:, month .== 10], dims = 2)[:, 1]
 Oct_ρ = mean(center_ρ_mean[:, month .== 10], dims = 2)[:, 1]
 Oct_uforcing = gravity_wave_forcing(
+    model_config,
     Oct_u,
     source_level,
     params.gw_F_S0,

examples/hybrid/staggered_nonhydrostatic_model.jl

@@ -137,6 +137,7 @@ function default_cache(Y, params, model_spec, spaces, numerics, simulation)
         ),
         spaces,
         moisture_model = model_spec.moisture_model,
+        model_config = model_spec.model_config,
         Yₜ = similar(Y), # only needed when using increment formulation
         limiters,
         ᶜρh_kwargs...,

examples/hybrid/types.jl

@@ -14,7 +14,10 @@ import ClimaAtmos:
     Microphysics0Moment,
     HeldSuarezForcing,
     BulkSurfaceScheme,
-    MoninObukhovSurface
+    MoninObukhovSurface,
+    SingleColumnModel,
+    SphericalModel,
+    BoxModel
 
 import ClimaCore: InputOutput
 
@@ -31,6 +34,7 @@ function get_model_spec(::Type{FT}, parsed_args, namelist) where {FT}
 
     model_spec = (;
         moisture_model,
+        model_config = CA.model_config(parsed_args),
         energy_form = CA.energy_form(parsed_args),
         perturb_initstate = parsed_args["perturb_initstate"],
         idealized_h2o,

regression_tests/mse_tables.jl

@@ -149,6 +149,14 @@ all_best_mse["single_column_nonorographic_gravity_wave"][(:c, :uₕ, :components
 all_best_mse["single_column_nonorographic_gravity_wave"][(:c, :uₕ, :components, :data, 2)] = 0.0
 all_best_mse["single_column_nonorographic_gravity_wave"][(:f, :w, :components, :data, 1)] = 0.0
 #
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"] = OrderedCollections.OrderedDict()
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"][(:c, :ρ)] = 0
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"][(:c, :ρe_tot)] = 0
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"][(:c, :uₕ, :components, :data, 1)] = 0
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"][(:c, :uₕ, :components, :data, 2)] = 0
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"][(:c, :ρq_tot)] = 0
+all_best_mse["sphere_held_suarez_rhoe_equilmoist_hightop_gw"][(:f, :w, :components, :data, 1)] = 0
+#
 #! format: on
 #################################
 #################################

regression_tests/ref_counter.jl

@@ -1 +1 @@
-28
+29

src/model_getters.jl

@@ -10,6 +10,17 @@ function moisture_model(parsed_args)
     end
 end
 
+function model_config(parsed_args)
+    config = parsed_args["config"]
+    return if config == "sphere"
+        SphericalModel()
+    elseif config == "column"
+        SingleColumnModel()
+    elseif config == "box"
+        BoxModel()
+    end
+end
+
 function energy_form(parsed_args)
     energy_name = parsed_args["energy_name"]
     @assert energy_name in ("rhoe", "rhoe_int", "rhotheta")