Merge #852

852: Add non-or0graphic gravity wave parameterization. r=charleskawczynski a=jiahe23

Co-authored-by: Zhaoyi Shen <[email protected]>
Co-authored-by: LenkaNovak <[email protected]>

## Purpose and Content
This PR implements the [Alexander and Dunkerton (1999)](https://doi.org/10.1175/1520-0469(1999)056%3C4167:ASPOMF%3E2.0.CO;2) paper for non-orographic gravity wave parameterization.

Examples are added to test agains the cases (Fig. 4 and 8) in the paper.

Help needed from software experts:
- [x] We need a way to read in the `.nc` data file for the tests;
- [ ]  The code may need considerable optimization for a good performance.
`@charleskawczynski` `@simonbyrne`  We appreciate if you could take a look

## Linked Issues
- Closes #783 

## PR Checklist
- [x] This PR has a corresponding issue OR is linked to an SDI.
- [ ] I have followed CliMA's codebase [contribution](https://clima.github.io/ClimateMachine.jl/latest/Contributing/) and [style](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) guidelines OR N/A.
- [ ] I have followed CliMA's [documentation policy](https://github.com/CliMA/policies/wiki/Documentation-Policy).
- [ ] I have checked all issues and PRs and I certify that this PR does not duplicate an open PR.
- [ ] I linted my code on my local machine prior to submission OR N/A.
- [ ] Unit tests are included OR N/A.
- [ ] Code used in an integration test OR N/A.
- [ ] All tests ran successfully on my local machine OR N/A.
- [ ] All classes, modules, and function contain docstrings OR N/A.
- [ ] Documentation has been added/updated OR N/A.


Co-authored-by: Jia He <[email protected]>

.buildkite/pipeline.yml

@@ -123,6 +123,18 @@ steps:
         command: "julia --color=yes --project=examples examples/hybrid/driver.jl --rad allskywithclear --idealized_h2o true --idealized_clouds true --hyperdiff false --config column --z_max 70e3 --z_elem 70 --dz_bottom 100 --dz_top 10000 --t_end 654days --dt 3hours --dt_save_to_sol 30hours --job_id sphere_single_column_radiative_equilibrium_allsky_idealized_clouds"
         artifact_paths: "sphere_single_column_radiative_equilibrium_allsky_idealized_clouds/*"
 
+      - label: ":computer: single column non-orographic gravity wave parameterization"
+        command: "julia --color=yes --project=examples examples/hybrid/driver.jl --hyperdiff false --config column --non_orographic_gravity_wave true --t_end 1500secs --dt 500secs --job_id single_column_nonorographic_gravity_wave --regression_test true"
+        artifact_paths: "single_column_nonorographic_gravity_wave/*"
+
+      - label: ":computer: non-orographic gravity wave parameterization unit test 3d"
+        command: "julia --color=yes --project=examples examples/hybrid/gravitywave_parameterization/gw_test_3d.jl"
+        artifact_paths: "nonorographic_gravity_wave_test_3d/*"
+
+      - label: ":computer: non-orographic gravity wave parameterization unit test single column"
+        command: "julia --color=yes --project=examples examples/hybrid/gravitywave_parameterization/gw_test_single_column.jl"
+        artifact_paths: "nonorographic_gravity_wave_test_single_column/*"
+
   - group: "MPI Examples"
     steps:
 

examples/hybrid/cli_options.jl

@@ -226,6 +226,10 @@ function parse_commandline()
         help = "Save most of the tc aux state to HDF5 file [`false` (default), `true`]"
         arg_type = Bool
         default = false
+        "--non_orographic_gravity_wave"
+        help = "Apply parameterization for convective gravity wave forcing on horizontal mean flow"
+        arg_type = Bool
+        default = false
     end
     parsed_args = ArgParse.parse_args(ARGS, s)
     return (s, parsed_args)

examples/hybrid/driver.jl

@@ -157,6 +157,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(),
         (;
             tendency_knobs = (;
                 hs_forcing = forcing_type isa HeldSuarezForcing,
@@ -166,6 +168,7 @@ function additional_cache(Y, params, model_spec, dt; use_tempest_mode = false)
                 rayleigh_sponge,
                 viscous_sponge,
                 hyperdiff,
+                non_orographic_gravity_wave = model_spec.non_orographic_gravity_wave,
                 has_turbconv = !isnothing(turbconv_model),
             )
         ),
@@ -182,6 +185,7 @@ function additional_tendency!(Yₜ, Y, p, t)
     (; rad_flux, vert_diff, hs_forcing) = p.tendency_knobs
     (; microphy_0M, hyperdiff, has_turbconv) = p.tendency_knobs
     (; rayleigh_sponge, viscous_sponge) = p.tendency_knobs
+    (; non_orographic_gravity_wave) = p.tendency_knobs
     hyperdiff && hyperdiffusion_tendency!(Yₜ, Y, p, t)
     rayleigh_sponge && rayleigh_sponge_tendency!(Yₜ, Y, p, t)
     viscous_sponge && viscous_sponge_tendency!(Yₜ, Y, p, t)
@@ -190,6 +194,7 @@ function additional_tendency!(Yₜ, Y, p, t)
     microphy_0M && zero_moment_microphysics_tendency!(Yₜ, Y, p, t)
     rad_flux && rrtmgp_model_tendency!(Yₜ, Y, p, t)
     has_turbconv && TCU.sgs_flux_tendency!(Yₜ, Y, p, t)
+    non_orographic_gravity_wave && gravity_wave_tendency!(Yₜ, Y, p, t)
 end
 
 ################################################################################
@@ -204,6 +209,12 @@ if parsed_args["trunc_stack_traces"]
 end
 
 include(joinpath("sphere", "baroclinic_wave_utilities.jl"))
+include(
+    joinpath(
+        "gravitywave_parameterization",
+        "gravity_wave_parameterization.jl",
+    ),
+)
 
 
 import ClimaCore: enable_threading

examples/hybrid/gravitywave_parameterization/gravity_wave_parameterization.jl

@@ -0,0 +1,240 @@
+function gravity_wave_cache(
+    Y,
+    ::Type{FT};
+    source_height = FT(15000),
+    Bm = FT(1.2),
+    F_S0 = FT(4e-3),
+    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]
+
+    return (;
+        gw_source_height = source_height,
+        gw_F_S0 = 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_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
+    ᶜρ = Y.c.ρ
+    # parameters
+    thermo_params = CAP.thermodynamics_params(params)
+    grav = FT(CAP.grav(params))
+
+    # compute buoyancy frequency
+    @. ᶜT = TD.air_temperature(thermo_params, ᶜts)
+
+    parent(ᶜdTdz) .= parent(Geometry.WVector.(ᶜgradᵥ.(ᶠinterp.(ᶜT))))
+
+    ᶜbuoyancy_frequency =
+        @. (grav / ᶜT) * (ᶜdTdz + grav / TD.cp_m(thermo_params, ᶜts))
+    ᶜbuoyancy_frequency = @. ifelse(
+        ᶜbuoyancy_frequency < FT(2.5e-5),
+        FT(sqrt(2.5e-5)),
+        sqrt(abs(ᶜbuoyancy_frequency)),
+    ) # to avoid small numbers
+    # alternative
+    # ᶜbuoyancy_frequency = [i > 2.5e-5 ? sqrt(i) : sqrt(2.5e-5)  for i in ᶜbuoyancy_frequency]
+    # TODO: create an extra layer at model top so that the gravity wave forcing
+    # .     occurring between the topmost model level and the upper boundary
+    # .     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)
+        ),
+    )
+    ᶠ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
+    u_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:1
+    v_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:2
+
+    uforcing = ones(axes(u_phy))
+    vforcing = similar(v_phy)
+    # bycolume
+    Fields.bycolumn(axes(ᶜρ)) do colidx
+        parent(uforcing[colidx]) .= gravity_wave_forcing(
+            parent(u_phy[colidx]),
+            source_level,
+            gw_F_S0,
+            gw_Bm,
+            gw_c,
+            gw_cw,
+            gw_c0,
+            gw_nk,
+            gw_k,
+            gw_k2,
+            parent(ᶜbuoyancy_frequency[colidx]),
+            parent(ᶜρ[colidx]),
+            parent(ᶠz[colidx]),
+        )
+        parent(vforcing[colidx]) .= gravity_wave_forcing(
+            parent(v_phy[colidx]),
+            source_level,
+            gw_F_S0,
+            gw_Bm,
+            gw_c,
+            gw_cw,
+            gw_c0,
+            gw_nk,
+            gw_k,
+            gw_k2,
+            parent(ᶜbuoyancy_frequency[colidx]),
+            parent(ᶜρ[colidx]),
+            parent(ᶠz[colidx]),
+        )
+
+    end
+
+    @. Yₜ.c.uₕ +=
+        Geometry.Covariant12Vector.(Geometry.UVVector.(uforcing, vforcing))
+
+end
+
+function gravity_wave_forcing(
+    ᶜu,
+    source_level,
+    F_S0,
+    Bm,
+    c,
+    cw,
+    c0,
+    nk,
+    kwv,
+    k2,
+    ᶜbf,
+    ᶜρ,
+    ᶠz,
+)
+
+    nc = length(c)
+
+    # define wave momentum flux (B0) at source level for each phase
+    # speed n, and the sum over all phase speeds (Bsum), which is needed
+    # to calculate the intermittency.
+
+    c_hat0 = c .- ᶜu[source_level]
+    # In GFDL code, flag is always 1: c = c0 * flag + c_hat0 * (1 - flag)
+    Bexp = @. exp(-log(2.0) * ((c - c0) / cw)^2)
+    B0 = @. sign(c_hat0) * Bm * Bexp
+    # In GFDL code, it is: B0 = @. sign(c_hat0) * (Bw * Bexp + Bn * Bexp)
+    # where Bw = Bm is the wide band and Bn is the narrow band
+    Bsum = sum(abs.(B0))
+    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 
+
+    ᶜdz = ᶠz[2:end] - ᶠz[1:(end - 1)]
+    wave_forcing = zeros(nc)
+    gwf = zeros(length(ᶜu))
+    for ink in 1:nk # loop over all wave lengths
+
+        mask = ones(nc)  # mask to determine which waves propagate upward
+        for k in source_level:length(ᶜu)
+            fac = FT(0.5) * (ᶜρ[k] / ᶜρ[source_level]) * kwv[ink] / ᶜbf[k]
+
+            ᶜHb = ᶜdz[k] / log(ᶜρ[k - 1] / ᶜρ[k])  # density scale height
+            alp2 = 0.25 / (ᶜHb * ᶜHb)
+            ω_r = sqrt((ᶜbf[k] * ᶜbf[k] * k2[ink]) / (k2[ink] + alp2)) # ω_r (critical frequency that marks total internal reflection)
+
+            fm = FT(0)
+            for n in 1:nc
+                # check only those waves which are still propagating, i.e., mask = 1.0
+                if (mask[n]) == 1.0
+                    c_hat = c[n] - ᶜu[k]
+                    # f phase speed matches the wind speed, remove c(n) from the set of propagating waves.
+                    if c_hat == 0.0
+                        mask[n] = 0.0
+                    else
+                        # define the criterion which determines if wave is reflected at this level (test).
+                        test = abs(c_hat) * kwv[ink] - ω_r
+                        if test >= 0.0
+                            # wave has undergone total internal reflection. remove it from the propagating set.
+                            mask[n] = 0.0
+                        else
+                            # if wave is not reflected at this level, determine if it is
+                            # breaking at this level (Foc >= 0), or if wave speed relative to
+                            # windspeed has changed sign from its value at the source level
+                            # (c_hat0[n] * c_hat <= 0). if it is above the source level and is
+                            # breaking, then add its momentum flux to the accumulated sum at
+                            # this level.
+                            # set mask=0.0 to remove phase speed band c[n] from the set of active 
+                            # waves moving upwards to the next level.
+                            if c_hat0[n] * c_hat <= 0.0
+                                mask[n] = 0.0
+                                if k > source_level
+                                    fm = fm + B0[n]
+                                end
+                            else
+                                Foc = B0[n] / (c_hat)^3 - fac
+                                if Foc >= 0.0
+                                    mask[n] = 0.0
+                                    if k > source_level
+                                        fm = fm + B0[n]
+                                    end
+                                end
+                            end
+                        end # (test >= 0.0)
+                    end #(c_hat == 0.0)
+                end # mask = 0
+
+            end # phase speed loop
+
+            # TODO: GFDL option to dump remaining flux at the top of the model
+
+            # compute the gravity wave momentum flux forcing 
+            # obtained across the entire wave spectrum at this level.
+            if k > source_level
+                rbh = sqrt(ᶜρ[k] * ᶜρ[k - 1])
+                wave_forcing[k] = (ᶜρ[source_level] / rbh) * fm * eps / ᶜdz[k]
+                wave_forcing[k - 1] =
+                    0.5 * (wave_forcing[k - 1] + wave_forcing[k])
+            else
+                wave_forcing[k] = 0.0
+            end
+
+        end # k
+
+        for k in source_level:length(ᶜu)
+            gwf[k] = gwf[k] + wave_forcing[k]
+        end
+
+    end # ink 
+
+    return gwf
+end