Refactor NOGW code structure and naming

- Move `input_uv` computation out of the nk loop for improved clarity.
- Consolidate multiple dispatch definitions for `non_orographic_gravity_wave_tendency`.
- Rename `NonOrographyGravityWave` to `NonOrographicGravityWave` for consistency.
- Recompute non-orographic gravity waves within callbacks at intervals defined by `dt_nogw`.

Author: ray-chew

config/default_configs/default_config.yml

@@ -220,6 +220,9 @@ orographic_gravity_wave:
 non_orographic_gravity_wave:
   help: "Apply parameterization for convective gravity wave forcing on horizontal mean flow [`false` (default), `true`]"
   value: false
+dt_nogw:
+  help: "Time between calling non-orographic gravity wave update"
+  value: "400secs"
 nh_poly:
   help: "Horizontal polynomial degree. Note: The number of quadrature points in 1D within each horizontal element is then Nq = <--nh_poly> + 1"
   value: 3

config/model_configs/single_column_nonorographic_gravity_wave.yml

@@ -7,6 +7,7 @@ dt: "400secs"
 output_default_diagnostics: true
 disable_surface_flux_tendency: true
 non_orographic_gravity_wave: true
+dt_nogw: "400secs"
 # diagnostics:
 #   - short_name: [ta, rhoa]
 #     period: 400secs

src/callbacks/callbacks.jl

@@ -399,6 +399,16 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     return nothing
 end
 
+
+
+NVTX.@annotate function nogw_model_callback!(integrator)
+    Y = integrator.u
+    p = integrator.p
+
+    non_orographic_gravity_wave_compute_tendency!(Y,p)
+end
+
+
 #Uniform insolation, magnitudes from Wing et al. (2018)
 #Note that the TOA downward shortwave fluxes won't be the same as the values in the paper if add_isothermal_boundary_layer is true
 function set_insolation_variables!(Y, p, t, ::RCEMIPIIInsolation)

src/callbacks/get_callbacks.jl

@@ -365,5 +365,22 @@ function get_callbacks(config, sim_info, atmos, params, Y, p)
             (callbacks..., call_every_dt(rrtmgp_model_callback!, dt_rad))
     end
 
+    if atmos.non_orographic_gravity_wave isa NonOrographicGravityWave
+        dt_nogw =
+            dt isa ITime ? ITime(time_to_seconds(parsed_args["dt_nogw"])) :
+            FT(time_to_seconds(parsed_args["dt_nogw"]))
+        dt_nogw, _, _, _ = promote(dt_nogw, t_start, dt, sim_info.t_end)
+        # We use Millisecond to support fractional seconds, eg. 0.1
+        dt_nogw_ms = Dates.Millisecond(1_000 * float(dt_nogw))
+        if parsed_args["dt_save_state_to_disk"] != "Inf" &&
+           !CA.isdivisible(dt_save_state_to_disk_dates, dt_nogw_ms)
+            @warn "Non-orographic gravity wave period ($(dt_nogw_ms)) is not an even divisor of the checkpoint frequency ($dt_save_state_to_disk_dates)"
+            @warn "This simulation will not be reproducible when restarted"
+        end
+
+        callbacks =
+            (callbacks..., call_every_dt(nogw_model_callback!, dt_nogw))
+    end
+
     return callbacks
 end

src/diagnostics/Diagnostics.jl

@@ -42,7 +42,7 @@ import ..PrognosticEDMFX
 import ..DiagnosticEDMFX
 
 # gravitywave_models
-import ..NonOrographyGravityWave
+import ..NonOrographicGravityWave
 import ..OrographicGravityWave
 
 # surface_model

src/diagnostics/gravitywave_diagnostics.jl

@@ -16,7 +16,7 @@ function compute_utendnogw!(
     cache,
     time,
     ::T,
-) where {T <: NonOrographyGravityWave}
+) where {T <: NonOrographicGravityWave}
     if isnothing(out)
         # u_waveforcing is not a field:
         # return Fields.array2field(
@@ -53,7 +53,7 @@ function compute_vtendnogw!(
     cache,
     time,
     ::T,
-) where {T <: NonOrographyGravityWave}
+) where {T <: NonOrographicGravityWave}
     if isnothing(out)
         # v_waveforcing is not a Field:
         # return Fields.array2field(

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

@@ -14,9 +14,8 @@ non_orographic_gravity_wave_cache(Y, atmos::AtmosModel) =
 
 non_orographic_gravity_wave_cache(Y, ::Nothing) = (;)
 
-non_orographic_gravity_wave_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
-function non_orographic_gravity_wave_cache(Y, gw::NonOrographyGravityWave)
+function non_orographic_gravity_wave_cache(Y, gw::NonOrographicGravityWave)
     if iscolumn(axes(Y.c))
         FT = Spaces.undertype(axes(Y.c))
         (; source_height, Bw, Bn, Bt_0, dc, cmax, c0, nk, cw, cn) = gw
@@ -134,12 +133,9 @@ function non_orographic_gravity_wave_cache(Y, gw::NonOrographyGravityWave)
     end
 end
 
-function non_orographic_gravity_wave_tendency!(
-    Yₜ,
+function non_orographic_gravity_wave_compute_tendency!(
     Y,
-    p,
-    t,
-    ::NonOrographyGravityWave,
+    p
 )
     #unpack
     ᶜT = p.scratch.ᶜtemp_scalar
@@ -276,11 +272,39 @@ function non_orographic_gravity_wave_tendency!(
         p,
     )
 
+    # # update cache after computation
+    # # do I need to do this? Since these attributes were
+    # # unpacked via reference?
+    # p.non_orographic_gravity_wave.uforcing .= uforcing
+    # p.non_orographic_gravity_wave.vforcing .= vforcing
+
+end
+
+non_orographic_gravity_wave_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
+
+function non_orographic_gravity_wave_tendency!(
+    Yₜ,
+    Y,
+    p,
+    t,
+    ::NonOrographicGravityWave,
+)
+
+    #unpack
+    (;
+        uforcing,
+        vforcing,
+    ) = p.non_orographic_gravity_wave
+
     @. Yₜ.c.uₕ +=
         Geometry.Covariant12Vector.(Geometry.UVVector.(uforcing, vforcing))
 
 end
 
+
+
+
+
 function non_orographic_gravity_wave_forcing(
     ᶜu,
     ᶜv,
@@ -368,49 +392,50 @@ function non_orographic_gravity_wave_forcing(
     #StaticBitVector stores 8 boolean values in a UInt8, allowing efficient storage for up to 256 gravity wave break data.
     level_end = Spaces.nlevels(axes(ᶜρ))
 
-    # loop over all wave lengths
-    for ink in 1:gw_nk
+    # Collect all required fields in a broadcasted object
+    input_u = Base.Broadcast.broadcasted(
+        tuple,
+        ᶜu_p1,
+        ᶜu_source,
+        ᶜbf_p1,
+        ᶜρ,
+        ᶜρ_p1,
+        ᶜρ_source,
+        ᶜz_p1,
+        ᶜz,
+        source_level,
+        gw_Bw,
+        gw_Bn,
+        gw_cw,
+        gw_cn,
+        gw_flag,
+        ᶜlevel,
+        gw_source_ampl,
+    )
 
-        # Collect all required fields in a broadcasted object
-        input_u = Base.Broadcast.broadcasted(
-            tuple,
-            ᶜu_p1,
-            ᶜu_source,
-            ᶜbf_p1,
-            ᶜρ,
-            ᶜρ_p1,
-            ᶜρ_source,
-            ᶜz_p1,
-            ᶜz,
-            source_level,
-            gw_Bw,
-            gw_Bn,
-            gw_cw,
-            gw_cn,
-            gw_flag,
-            ᶜlevel,
-            gw_source_ampl,
-        )
+    input_v = Base.Broadcast.broadcasted(
+        tuple,
+        ᶜv_p1,
+        ᶜv_source,
+        ᶜbf_p1,
+        ᶜρ,
+        ᶜρ_p1,
+        ᶜρ_source,
+        ᶜz_p1,
+        ᶜz,
+        source_level,
+        gw_Bw,
+        gw_Bn,
+        gw_cw,
+        gw_cn,
+        gw_flag,
+        ᶜlevel,
+        gw_source_ampl,
+    )
 
-        input_v = Base.Broadcast.broadcasted(
-            tuple,
-            ᶜv_p1,
-            ᶜv_source,
-            ᶜbf_p1,
-            ᶜρ,
-            ᶜρ_p1,
-            ᶜρ_source,
-            ᶜz_p1,
-            ᶜz,
-            source_level,
-            gw_Bw,
-            gw_Bn,
-            gw_cw,
-            gw_cn,
-            gw_flag,
-            ᶜlevel,
-            gw_source_ampl,
-        )
+
+    # loop over all wave lengths
+    for ink in 1:gw_nk
 
         # Accumulate zonal wave forcing in every column
         waveforcing_column_accumulate!(

src/solver/model_getters.jl

@@ -183,9 +183,9 @@ function get_non_orographic_gravity_wave_model(
     @assert nogw_name in (true, false)
     return if nogw_name == true
         if parsed_args["config"] == "column"
-            NonOrographyGravityWave{FT}(; Bw = 1.2, Bn = 0.0, Bt_0 = 4e-3)
+            NonOrographicGravityWave{FT}(; Bw = 1.2, Bn = 0.0, Bt_0 = 4e-3)
         elseif parsed_args["config"] == "sphere"
-            NonOrographyGravityWave{FT}(;
+            NonOrographicGravityWave{FT}(;
                 Bw = 0.4,
                 Bn = 0.0,
                 cw = 35.0,

src/solver/types.jl

@@ -241,7 +241,7 @@ Base.@kwdef struct RayleighSponge{FT} <: AbstractSponge
 end
 
 abstract type AbstractGravityWave end
-Base.@kwdef struct NonOrographyGravityWave{FT} <: AbstractGravityWave
+Base.@kwdef struct NonOrographicGravityWave{FT} <: AbstractGravityWave
     source_pressure::FT = 31500
     damp_pressure::FT = 85
     source_height::FT = 15000