Merge pull request #3779 from CliMA/as/test-rad-update

Testing updated RRTMGPInterface following RRTMGP changes to deep-atmosphere metrics

Author: akshaysridhar

config/longrun_configs/amip_target_edonly.yml

@@ -2,7 +2,8 @@ h_elem: 16
 z_max: 60000.0
 z_elem: 63
 dz_bottom: 30.0
-deep_atmosphere: false
+deep_atmosphere: true
+mesh_warp_type: "SLEVE"
 rayleigh_sponge: true
 viscous_sponge: true
 dt_save_state_to_disk: "30days"

perf/jet_test_nfailures.jl

@@ -38,7 +38,7 @@ using Test
     # inference. By increasing this counter, we acknowledge that
     # we have introduced an inference failure. We hope to drive
     # this number down to 0.
-    n_allowed_failures = 248
+    n_allowed_failures = 249
     @show n
     if n < n_allowed_failures
         @info "Please update the n-failures to $n"

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-235
+236
 
 # **README**
 #
@@ -20,6 +20,10 @@
 
 
 #=
+236
+- Radiation fluxes in runs with the `deep_atmosphere` configuration now take into account the 
+  column expansion with height. Affects only cases with `DeepSphericalGlobalGeometry`.
+
 235
 - Related to #3775, the computation and update of non-orographic gravity wave (NOGW) 
   are now separated into callback and tendencies update, affecting the NOGW-related

src/parameterized_tendencies/radiation/RRTMGPInterface.jl

@@ -19,11 +19,13 @@ using Random
 abstract type AbstractRRTMGPMode end
 struct GrayRadiation <: AbstractRRTMGPMode
     add_isothermal_boundary_layer::Bool
+    deep_atmosphere::Bool
 end
 struct ClearSkyRadiation <: AbstractRRTMGPMode
     idealized_h2o::Bool
     add_isothermal_boundary_layer::Bool
     aerosol_radiation::Bool
+    deep_atmosphere::Bool
 end
 struct AllSkyRadiation{ACR <: AbstractCloudInRadiation} <: AbstractRRTMGPMode
     idealized_h2o::Bool
@@ -39,6 +41,7 @@ struct AllSkyRadiation{ACR <: AbstractCloudInRadiation} <: AbstractRRTMGPMode
     Disable this option when running production runs.
     """
     reset_rng_seed::Bool
+    deep_atmosphere::Bool
 end
 struct AllSkyRadiationWithClearSkyDiagnostics{
     ACR <: AbstractCloudInRadiation,
@@ -56,6 +59,7 @@ struct AllSkyRadiationWithClearSkyDiagnostics{
     Disable this option when running production runs.
     """
     reset_rng_seed::Bool
+    deep_atmosphere::Bool
 end
 
 """
@@ -260,7 +264,7 @@ function extrap!(
     @. p = p⁺ * (T / T⁺)^(cₚ / R)
 end
 
-struct RRTMGPModel{R, I, B, L, P, LWS, SWS, AS, V}
+struct RRTMGPModel{R, I, B, L, P, LWS, SWS, AS, V, M}
     radiation_mode::R
     interpolation::I
     bottom_extrapolation::B
@@ -270,6 +274,7 @@ struct RRTMGPModel{R, I, B, L, P, LWS, SWS, AS, V}
     sw_solver::SWS
     as::AS  # Atmospheric state
     views::V  # user-friendly views into the solver
+    metric_scaling::M
 end
 
 # Allow cache to be moved on the CPU. Used by ClimaCoupler to save checkpoints
@@ -539,6 +544,7 @@ array.
   `requires_z(interpolation) || requires_z(bottom_extrapolation)`:
     - `center_z`: z-coordinate in m at cell centers
     - `face_z`: z-coordinate in m at cell faces
+    - `planet_radius`: planet radius (used to compute metric scaling factor)
 """
 RRTMGPModel(
     params::RRTMGP.Parameters.ARP,
@@ -694,6 +700,11 @@ function _RRTMGPModel(
     end
 
     p_lev = DA{FT}(undef, nlay + 1, ncol)
+    if radiation_mode.deep_atmosphere && :planet_radius in keys(dict)
+        metric_scaling = DA{FT}(undef, nlay + 1, ncol)
+    else
+        metric_scaling = nothing
+    end
     t_lev = DA{FT}(undef, nlay + 1, ncol)
     t_sfc = DA{FT}(undef, ncol)
     set_and_save!(t_sfc, "surface_temperature", t..., dict)
@@ -863,7 +874,6 @@ function _RRTMGPModel(
         else
             aerosol_state = nothing
         end
-
         as = RRTMGP.AtmosphericStates.AtmosphericState(
             lon,
             lat,
@@ -876,7 +886,6 @@ function _RRTMGPModel(
             cloud_state,
             aerosol_state,
         )
-
     end
 
     sw_solver = RRTMGP.RTE.TwoStreamSWRTE(
@@ -901,6 +910,14 @@ function _RRTMGPModel(
         set_and_save!(z_lay, "center_z", t..., dict)
         z_lev = DA{FT}(undef, nlay + 1, ncol)
         set_and_save!(z_lev, "face_z", t..., dict)
+        if radiation_mode.deep_atmosphere && :planet_radius in keys(dict)
+            planet_radius = pop!(dict, :planet_radius)
+            # Area ratio appears in denominator of RRTMGP scaling functions, 
+            # we therefore pass the multiplicative inverse from ClimaAtmos to 
+            # use mult ops instead of div in RRTMGP GPU kernels.
+            metric_scaling .=
+                inv.(((z_lev .+ planet_radius) ./ planet_radius) .^ (FT(2)))
+        end
     end
 
     if length(dict) > 0
@@ -921,6 +938,7 @@ function _RRTMGPModel(
         sw_solver,
         as,
         NamedTuple(views),
+        metric_scaling,
     )
 end
 
@@ -1143,14 +1161,15 @@ get_vmr_h2o(vmr::RRTMGP.Vmrs.Vmr, idx_gases_sw) =
     view(vmr.vmr, idx_gases_sw["h2o"], :, :)
 
 NVTX.@annotate update_lw_fluxes!(::GrayRadiation, model) =
-    RRTMGP.RTESolver.solve_lw!(model.lw_solver, model.as)
+    RRTMGP.RTESolver.solve_lw!(model.lw_solver, model.as, model.metric_scaling)
 NVTX.@annotate update_lw_fluxes!(::ClearSkyRadiation, model) =
     RRTMGP.RTESolver.solve_lw!(
         model.lw_solver,
         model.as,
         model.lookups.lookup_lw,
         nothing,
         model.lookups.lookup_lw_aero,
+        model.metric_scaling,
     )
 NVTX.@annotate update_lw_fluxes!(::AllSkyRadiation, model) =
     RRTMGP.RTESolver.solve_lw!(
@@ -1159,6 +1178,7 @@ NVTX.@annotate update_lw_fluxes!(::AllSkyRadiation, model) =
         model.lookups.lookup_lw,
         model.lookups.lookup_lw_cld,
         model.lookups.lookup_lw_aero,
+        model.metric_scaling,
     )
 NVTX.@annotate function update_lw_fluxes!(
     ::AllSkyRadiationWithClearSkyDiagnostics,
@@ -1170,6 +1190,7 @@ NVTX.@annotate function update_lw_fluxes!(
         model.lookups.lookup_lw,
         nothing,
         model.lookups.lookup_lw_aero,
+        model.metric_scaling,
     )
     parent(model.face_clear_lw_flux_up) .= parent(model.face_lw_flux_up)
     parent(model.face_clear_lw_flux_dn) .= parent(model.face_lw_flux_dn)
@@ -1180,18 +1201,20 @@ NVTX.@annotate function update_lw_fluxes!(
         model.lookups.lookup_lw,
         model.lookups.lookup_lw_cld,
         model.lookups.lookup_lw_aero,
+        model.metric_scaling,
     )
 end
 
 NVTX.@annotate update_sw_fluxes!(::GrayRadiation, model) =
-    RRTMGP.RTESolver.solve_sw!(model.sw_solver, model.as)
+    RRTMGP.RTESolver.solve_sw!(model.sw_solver, model.as, model.metric_scaling)
 NVTX.@annotate update_sw_fluxes!(::ClearSkyRadiation, model) =
     RRTMGP.RTESolver.solve_sw!(
         model.sw_solver,
         model.as,
         model.lookups.lookup_sw,
         nothing,
         model.lookups.lookup_sw_aero,
+        model.metric_scaling,
     )
 NVTX.@annotate update_sw_fluxes!(::AllSkyRadiation, model) =
     RRTMGP.RTESolver.solve_sw!(
@@ -1200,6 +1223,7 @@ NVTX.@annotate update_sw_fluxes!(::AllSkyRadiation, model) =
         model.lookups.lookup_sw,
         model.lookups.lookup_sw_cld,
         model.lookups.lookup_sw_aero,
+        model.metric_scaling,
     )
 NVTX.@annotate function update_sw_fluxes!(
     ::AllSkyRadiationWithClearSkyDiagnostics,
@@ -1211,6 +1235,7 @@ NVTX.@annotate function update_sw_fluxes!(
         model.lookups.lookup_sw,
         nothing,
         model.lookups.lookup_sw_aero,
+        model.metric_scaling,
     )
     parent(model.face_clear_sw_flux_up) .= parent(model.face_sw_flux_up)
     parent(model.face_clear_sw_flux_dn) .= parent(model.face_sw_flux_dn)
@@ -1223,6 +1248,7 @@ NVTX.@annotate function update_sw_fluxes!(
         model.lookups.lookup_sw,
         model.lookups.lookup_sw_cld,
         model.lookups.lookup_sw_aero,
+        model.metric_scaling,
     )
 end
 

src/parameterized_tendencies/radiation/radiation.jl

@@ -106,6 +106,9 @@ function rrtmgp_model_kwargs(
     ᶠspace = Spaces.face_space(space)
     ᶜz = Fields.coordinate_field(ᶜspace).z
     ᶠz = Fields.coordinate_field(ᶠspace).z
+    if ᶜspace.grid.global_geometry isa Geometry.AbstractSphericalGlobalGeometry
+        planet_radius = ᶜspace.grid.global_geometry.radius
+    end
     bottom_coords = Fields.coordinate_field(Spaces.level(ᶜspace, 1))
     latitude = if eltype(bottom_coords) <: Geometry.LatLongZPoint
         Fields.field2array(bottom_coords.lat)
@@ -117,8 +120,18 @@ function rrtmgp_model_kwargs(
         optical_thickness_parameter = (@. 7.2 + (1.8 - 7.2) * sind(latitude)^2),
         latitude,
     )
-    zkwargs =
-        (; center_z = Fields.field2array(ᶜz), face_z = Fields.field2array(ᶠz))
+    if ᶜspace.grid.global_geometry isa Geometry.AbstractSphericalGlobalGeometry
+        zkwargs = (;
+            center_z = Fields.field2array(ᶜz),
+            face_z = Fields.field2array(ᶠz),
+            planet_radius = planet_radius,
+        )
+    else
+        zkwargs = (;
+            center_z = Fields.field2array(ᶜz),
+            face_z = Fields.field2array(ᶠz),
+        )
+    end
     return include_z ? (; kwargs..., zkwargs...) : kwargs
 end
 
@@ -137,6 +150,9 @@ function rrtmgp_model_kwargs(
     ᶜΔz = Fields.Δz_field(ᶜspace)
     ᶜz = Fields.coordinate_field(ᶜspace).z
     ᶠz = Fields.coordinate_field(ᶠspace).z
+    if ᶜspace.grid.global_geometry isa Geometry.AbstractSphericalGlobalGeometry
+        planet_radius = ᶜspace.grid.global_geometry.radius
+    end
     latitude = if eltype(bottom_coords) <: Geometry.LatLongZPoint
         Fields.field2array(bottom_coords.lat)
     else
@@ -256,11 +272,21 @@ function rrtmgp_model_kwargs(
         end
 
         if include_z
-            kwargs = (;
-                kwargs...,
-                center_z = Fields.field2array(ᶜz),
-                face_z = Fields.field2array(ᶠz),
-            )
+            if ᶜspace.grid.global_geometry isa
+               Geometry.AbstractSphericalGlobalGeometry
+                kwargs = (;
+                    kwargs...,
+                    center_z = Fields.field2array(ᶜz),
+                    face_z = Fields.field2array(ᶠz),
+                    planet_radius = planet_radius,
+                )
+            else
+                kwargs = (;
+                    kwargs...,
+                    center_z = Fields.field2array(ᶜz),
+                    face_z = Fields.field2array(ᶠz),
+                )
+            end
         end
     end
     return kwargs

src/solver/model_getters.jl

@@ -238,6 +238,7 @@ function get_radiation_mode(parsed_args, ::Type{FT}) where {FT}
     reset_rng_seed = parsed_args["radiation_reset_rng_seed"]
     @assert reset_rng_seed in (true, false)
     radiation_name = parsed_args["rad"]
+    deep_atmosphere = parsed_args["deep_atmosphere"]
     @assert radiation_name in (
         nothing,
         "nothing",
@@ -257,12 +258,13 @@ function get_radiation_mode(parsed_args, ::Type{FT}) where {FT}
         @warn "prescribe_clouds_in_radiation does not have any effect with $radiation_name radiation option"
     end
     return if radiation_name == "gray"
-        RRTMGPI.GrayRadiation(add_isothermal_boundary_layer)
+        RRTMGPI.GrayRadiation(add_isothermal_boundary_layer, deep_atmosphere)
     elseif radiation_name == "clearsky"
         RRTMGPI.ClearSkyRadiation(
             idealized_h2o,
             add_isothermal_boundary_layer,
             aerosol_radiation,
+            deep_atmosphere,
         )
     elseif radiation_name == "allsky"
         RRTMGPI.AllSkyRadiation(
@@ -272,6 +274,7 @@ function get_radiation_mode(parsed_args, ::Type{FT}) where {FT}
             add_isothermal_boundary_layer,
             aerosol_radiation,
             reset_rng_seed,
+            deep_atmosphere,
         )
     elseif radiation_name == "allskywithclear"
         RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics(
@@ -281,6 +284,7 @@ function get_radiation_mode(parsed_args, ::Type{FT}) where {FT}
             add_isothermal_boundary_layer,
             aerosol_radiation,
             reset_rng_seed,
+            deep_atmosphere,
         )
     elseif radiation_name == "DYCOMS"
         RadiationDYCOMS{FT}()