Improve design of nonorographic gravity wave parameterization

[jiahe23]

Files

• src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl
• All tests in test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/

Notes

• Remove use of parent

• Remove use of copy and vec in copy(vec(parent(u_phy[colidx])))

  • if non_orographic_gravity_wave_forcing is purely functional, we don't need to make a copy of the inputs, we can pass pointers to the inputs (i.e., u_phy)

• Reduce allocations. And can check/test allocations with @allocated (see docs about usage)

  • ᶜbuoyancy_frequency = @. (grav / ᶜT) * (ᶜdTdz + grav / TD.cp_m(thermo_params, ᶜts)) should become
  • @. ᶜbuoyancy_frequency = (grav / ᶜT) * (ᶜdTdz + grav / TD.cp_m(thermo_params, ᶜts))
  • Move source_level = similar(Fields.level(Y.c.ρ, 1)) to cache
  • Move u_phy to cache in u_phy = Geometry.UVVector.(Y.c.uₕ).components.data.:1
  • Move uforcing in uforcing = ones(axes(u_phy)) to cache
  • preallocate Bw_exp in Bw_exp = @. exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cw)^2)
  • Apply the same to other variables that use similar patterns
  • Remove use of append!

• Remove use of Fields.bycolumn (this may require developing custom CPU/GPU kernels in ClimaCore). Ideally we can just use column_mapreduce!

• Remove use of findlast and findfirst. Hopefully we can just use column_mapreduce!? If not, then replace with ᶜupdraft_top in https://github.com/CliMA/ClimaAtmos.jl/pull/2819/files.

• Any use of getindex (e.g., [1] in parent(source_level[colidx]) .= findlast(parent(ᶜp[colidx]) .> gw_source_pressure)[1])

• Refactor non_orographic_gravity_wave_forcing to be purely functional (no mutation of input arguments). See the comments and cross-referenced issues / pull requests

[charleskawczynski]

Some concrete tasks:

• Use ClimaCore Fields and remove use of parent / reshape
• Make gravity_wave_tendency! use bycolumn for entire function to leverage threading

    source_level = argmin(
        abs.(
            unique(parent(Fields.coordinate_field(Y.c).z) .- gw_source_height)
        ),
    )

may require ClimaCore features

[charleskawczynski]

For example, it would be helpful to have

• findfirst/findlast/argmin/argmax defined for column fields

[charleskawczynski]

Note to self, current sketch of changes to gravity_wave_forcing:

function gravity_wave_forcing(::Type{FT}, args) where {FT}
    (; ᶜρ_k, ᶜρ_lev′, ᶜbf_k, kwv_ink, ᶜρ_km1, ᶜΔz_k) = args
    (; k2_ink, source_level) = args
    (; c, c_hat0, ϵ, B0) = args
    fac = FT(0.5) * (ᶜρ_k / ᶜρ_lev′) * kwv_ink / ᶜbf_k

    ᶜHb = ᶜΔz_k / log(ᶜρ_km1 / ᶜρ_k)  # density scale height
    alp2 = FT(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
            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
                mask[n] = 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
                    # wave has undergone total internal reflection. remove it from the propagating set.
                    mask[n] = 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
                        mask[n] = 0
                        if k > source_level
                            fm = fm + B0[n]
                        end
                    else
                        Foc = B0[n] / (c_hat)^3 - fac
                        if Foc >= 0
                            mask[n] = 0
                            if k > source_level
                                fm = fm + B0[n]
                            end
                        end
                    end
                end # (test >= 0)
            end #(c_hat == 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.
    return if k > source_level
        rbh = sqrt(ᶜρ_k * ᶜρ_km1)
        (ᶜρ_lev′ / rbh) * fm * ϵ / ᶜΔz_k
        # wave_forcing[k] = (ᶜρ_lev′ / rbh) * fm * ϵ / ᶜΔz_k
        # TODO: this just looks like interpolation after?
        # wave_forcing[k - 1] = FT(0.5) * (wave_forcing[k - 1] + wave_forcing[k])
    else
        # wave_forcing[k] = 0
        FT(0)
    end
end

[charleskawczynski]

We should also fix the types of the fields. For example ᶜdTdz is currently a Float32/Float64 field, and it should be a Covariant3Vector or WVector, depending on the equations. cc @xxy220022, @szy21, @dennisYatunin