Merge pull request #3875 from CliMA/ne/atmos_model

Refactor AtmosModel

Author: nefrathenrici

src/diagnostics/Diagnostics.jl

@@ -2,10 +2,9 @@ module Diagnostics
 
 import Dates: Month, Day, Hour, DateTime, Period
 
-import ClimaComms
-
 import LinearAlgebra: dot
 
+import ClimaComms
 import ClimaCore:
     Fields, Geometry, InputOutput, Meshes, Spaces, Operators, Domains, Grids
 import ClimaCore.Utilities: half
@@ -15,6 +14,9 @@ import Thermodynamics as TD
 import ..lazy
 
 import ..AtmosModel
+import ..AtmosWater
+import ..AtmosRadiation
+import ..AtmosTurbconv
 import ..AtmosCallback
 import ..EveryNSteps
 

src/diagnostics/default_diagnostics.jl

@@ -214,9 +214,38 @@ function default_diagnostics(
     return [average_func(moist_diagnostics...; output_writer, start_date)...]
 end
 
+function default_diagnostics(
+    atmos_water::AtmosWater,
+    duration,
+    start_date;
+    output_writer,
+)
+    diagnostics = []
+
+    # Add moisture and precipitation model diagnostics
+    for model in (atmos_water.moisture_model, atmos_water.precip_model)
+        !isnothing(model) && append!(
+            diagnostics,
+            default_diagnostics(model, duration, start_date; output_writer),
+        )
+    end
+
+    return diagnostics
+end
+
 #######################
 # Precipitation model #
 #######################
+function default_diagnostics(
+    ::Microphysics0Moment,
+    duration,
+    start_date;
+    output_writer,
+)
+    # 0-moment microphysics doesn't have additional diagnostics beyond basic moisture
+    return []
+end
+
 function default_diagnostics(
     ::Microphysics1Moment,
     duration,
@@ -401,3 +430,41 @@ function default_diagnostics(::EDOnlyEDMFX, duration, start_date; output_writer)
         average_func(edonly_edmfx_diagnostics...; output_writer, start_date)...,
     ]
 end
+
+function default_diagnostics(
+    atmos_radiation::AtmosRadiation,
+    duration,
+    start_date;
+    output_writer,
+)
+    # Add radiation mode diagnostics
+    if !isnothing(atmos_radiation.radiation_mode)
+        return default_diagnostics(
+            atmos_radiation.radiation_mode,
+            duration,
+            start_date;
+            output_writer,
+        )
+    else
+        return []
+    end
+end
+
+function default_diagnostics(
+    atmos_turbconv::AtmosTurbconv,
+    duration,
+    start_date;
+    output_writer,
+)
+    # Add turbulence convection model diagnostics
+    if !isnothing(atmos_turbconv.turbconv_model)
+        return default_diagnostics(
+            atmos_turbconv.turbconv_model,
+            duration,
+            start_date;
+            output_writer,
+        )
+    else
+        return []
+    end
+end

src/diagnostics/radiation_diagnostics.jl

@@ -45,7 +45,7 @@ function compute_rsd!(
             cache.radiation.rrtmgp_model.face_sw_flux_dn,
             axes(state.f),
         )
-        @assert out !== nothing "Output field 'out' must not be `nothing` in this branch"
+        @assert !isnothing(out) "Output field 'out' must not be `nothing` in this branch"
         radiation_mode.deep_atmosphere &&
             apply_geometric_scaling!(out, z_lev, planet_radius, FT)
     end

src/parameterized_tendencies/radiation/radiation.jl

@@ -28,8 +28,12 @@ radiation_model_cache(Y, atmos::AtmosModel, args...) =
 ##### No Radiation
 #####
 
-radiation_model_cache(Y, radiation_mode::Nothing; args...) = (;)
-radiation_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
+radiation_model_cache(
+    Y,
+    radiation_mode::Union{Nothing, HeldSuarezForcing};
+    args...,
+) = (;)
+radiation_tendency!(Yₜ, Y, p, t, ::Union{Nothing, HeldSuarezForcing}) = nothing
 
 #####
 ##### RRTMGP Radiation

src/prognostic_equations/remaining_tendency.jl

@@ -141,7 +141,8 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     ᶜuₕ = Y.c.uₕ
     ᶠu₃ = Y.f.u₃
     ᶜρ = Y.c.ρ
-    (; forcing_type, moisture_model, rayleigh_sponge, viscous_sponge) = p.atmos
+    (; radiation_mode, moisture_model, rayleigh_sponge, viscous_sponge) =
+        p.atmos
     (; ls_adv, scm_coriolis) = p.atmos
     (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
@@ -151,7 +152,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)
     vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)
     rst_uₕ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)
-    hs_args = (ᶜuₕ, ᶜp, params, sfc_conditions.ts, moisture_model, forcing_type)
+    # For HeldSuarezForcing, the radiation_mode is used as the forcing parameter
+    forcing = radiation_mode isa HeldSuarezForcing ? radiation_mode : nothing
+    hs_args = (ᶜuₕ, ᶜp, params, sfc_conditions.ts, moisture_model, forcing)
     hs_tendency_uₕ = held_suarez_forcing_tendency_uₕ(hs_args...)
     hs_tendency_ρe_tot = held_suarez_forcing_tendency_ρe_tot(ᶜρ, hs_args...)
     edmf_cor_tend_uₕ = scm_coriolis_tendency_uₕ(ᶜuₕ, scm_coriolis)

src/solver/model_getters.jl

@@ -246,6 +246,7 @@ function get_radiation_mode(parsed_args, ::Type{FT}) where {FT}
         "gray",
         "allsky",
         "allskywithclear",
+        "held_suarez",
         "DYCOMS",
         "TRMM_LBA",
         "ISDAC",
@@ -286,6 +287,8 @@ function get_radiation_mode(parsed_args, ::Type{FT}) where {FT}
             reset_rng_seed,
             deep_atmosphere,
         )
+    elseif radiation_name == "held_suarez"
+        HeldSuarezForcing()
     elseif radiation_name == "DYCOMS"
         RadiationDYCOMS{FT}()
     elseif radiation_name == "TRMM_LBA"
@@ -352,11 +355,11 @@ end
 function get_forcing_type(parsed_args)
     forcing = parsed_args["forcing"]
     @assert forcing in (nothing, "held_suarez")
-    return if forcing == nothing
-        nothing
-    elseif forcing == "held_suarez"
-        HeldSuarezForcing()
+    if forcing == "held_suarez"
+        @warn "The 'held_suarez' forcing option is deprecated. Use rad='held_suarez' instead to set HeldSuarezForcing as a radiation mode."
+        return HeldSuarezForcing()  # Still return the object for backward compatibility
     end
+    return nothing
 end
 
 struct CallCloudDiagnosticsPerStage end
@@ -372,7 +375,7 @@ end
 
 function get_subsidence_model(parsed_args, radiation_mode, FT)
     subsidence = parsed_args["subsidence"]
-    subsidence == nothing && return nothing
+    isnothing(subsidence) && return nothing
 
     prof = if subsidence == "Bomex"
         APL.Bomex_subsidence(FT)
@@ -382,6 +385,8 @@ function get_subsidence_model(parsed_args, radiation_mode, FT)
         APL.Rico_subsidence(FT)
     elseif subsidence == "DYCOMS"
         @assert radiation_mode isa RadiationDYCOMS
+        # For DYCOMS case, subsidence is linearly proportional to height
+        # with slope equal to the divergence rate specified in radiation mode
         z -> -z * radiation_mode.divergence
     elseif subsidence == "ISDAC"
         APL.ISDAC_subsidence(FT)

src/solver/type_getters.jl

@@ -48,7 +48,11 @@ function get_atmos(config::AtmosConfig, params)
     (!isnothing(co2) && !with_rrtmgp) &&
         @warn ("$(co2) does nothing if RRTMGP is not used")
 
-    disable_momentum_vertical_diffusion = forcing_type isa HeldSuarezForcing
+    # HeldSuarezForcing can be set via radiation_mode or legacy forcing option for now
+    final_radiation_mode =
+        forcing_type isa HeldSuarezForcing ? forcing_type : radiation_mode
+    disable_momentum_vertical_diffusion =
+        final_radiation_mode isa HeldSuarezForcing
 
     advection_test = parsed_args["advection_test"]
     @assert advection_test in (false, true)
@@ -86,23 +90,38 @@ function get_atmos(config::AtmosConfig, params)
     )
 
     atmos = AtmosModel(;
+        # AtmosWater - Moisture, Precipitation & Clouds
         moisture_model,
-        ozone,
-        co2,
-        radiation_mode,
-        subsidence = get_subsidence_model(parsed_args, radiation_mode, FT),
-        ls_adv = get_large_scale_advection_model(parsed_args, FT),
-        external_forcing = get_external_forcing_model(parsed_args, FT),
-        scm_coriolis = get_scm_coriolis(parsed_args, FT),
-        advection_test,
-        edmfx_model,
         precip_model,
         cloud_model,
         noneq_cloud_formation_mode = implicit_noneq_cloud_formation ?
                                      Implicit() : Explicit(),
-        forcing_type,
         call_cloud_diagnostics_per_stage,
+
+        # SCMSetup - Single-Column Model components
+        subsidence = get_subsidence_model(parsed_args, radiation_mode, FT),
+        external_forcing = get_external_forcing_model(parsed_args, FT),
+        ls_adv = get_large_scale_advection_model(parsed_args, FT),
+        advection_test,
+        scm_coriolis = get_scm_coriolis(parsed_args, FT),
+
+        # AtmosRadiation
+        radiation_mode = final_radiation_mode,
+        ozone,
+        co2,
+        insolation = get_insolation_form(parsed_args),
+
+        # AtmosTurbconv - Turbulence & Convection
+        edmfx_model,
         turbconv_model = get_turbconv_model(FT, parsed_args, turbconv_params),
+        sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
+        sgs_entr_detr_mode = implicit_sgs_entr_detr ? Implicit() : Explicit(),
+        sgs_nh_pressure_mode = implicit_sgs_nh_pressure ? Implicit() :
+                               Explicit(),
+        sgs_mf_mode = implicit_sgs_mass_flux ? Implicit() : Explicit(),
+        smagorinsky_lilly = get_smagorinsky_lilly_model(parsed_args),
+
+        # AtmosGravityWave
         non_orographic_gravity_wave = get_non_orographic_gravity_wave_model(
             parsed_args,
             FT,
@@ -111,24 +130,22 @@ function get_atmos(config::AtmosConfig, params)
             parsed_args,
             FT,
         ),
-        hyperdiff = get_hyperdiffusion_model(parsed_args, FT),
-        vert_diff,
-        diff_mode = implicit_diffusion ? Implicit() : Explicit(),
-        sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
-        sgs_entr_detr_mode = implicit_sgs_entr_detr ? Implicit() : Explicit(),
-        sgs_nh_pressure_mode = implicit_sgs_nh_pressure ? Implicit() :
-                               Explicit(),
-        sgs_mf_mode = implicit_sgs_mass_flux ? Implicit() : Explicit(),
+
+        # AtmosSponge
         viscous_sponge = get_viscous_sponge_model(parsed_args, params, FT),
-        smagorinsky_lilly = get_smagorinsky_lilly_model(parsed_args),
         rayleigh_sponge = get_rayleigh_sponge_model(parsed_args, params, FT),
+
+        # AtmosSurface
         sfc_temperature = get_sfc_temperature_form(parsed_args),
-        insolation = get_insolation_form(parsed_args),
-        disable_surface_flux_tendency = parsed_args["disable_surface_flux_tendency"],
         surface_model = get_surface_model(parsed_args),
         surface_albedo = get_surface_albedo_model(parsed_args, params, FT),
-        numerics = get_numerics(parsed_args),
+
+        # Top-level options (not grouped)
+        vert_diff,
+        numerics = get_numerics(parsed_args, FT),
+        disable_surface_flux_tendency = parsed_args["disable_surface_flux_tendency"],
     )
+    # TODO: Should this go in the AtmosModel constructor?
     @assert !@any_reltype(atmos, (UnionAll, DataType))
 
     @info "AtmosModel: \n$(summary(atmos))"
@@ -146,7 +163,7 @@ function get_scale_blending_method(parsed_args)
     end
 end
 
-function get_numerics(parsed_args)
+function get_numerics(parsed_args, FT)
     test_dycore =
         parsed_args["test_dycore_consistency"] ? TestDycoreConsistency() :
         nothing
@@ -173,13 +190,19 @@ function get_numerics(parsed_args)
     limiter = parsed_args["apply_limiter"] ? CA.QuasiMonotoneLimiter() : nothing
 
     # wrap each upwinding mode in a Val for dispatch
+    diff_mode = parsed_args["implicit_diffusion"] ? Implicit() : Explicit()
+
+    hyperdiff = get_hyperdiffusion_model(parsed_args, FT)
+
     numerics = AtmosNumerics(;
         energy_upwinding,
         tracer_upwinding,
         edmfx_upwinding,
         edmfx_sgsflux_upwinding,
         limiter,
         test_dycore_consistency = test_dycore,
+        diff_mode,
+        hyperdiff,
     )
     @info "numerics $(summary(numerics))"