GCM integration with GW

.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/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,

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