Merge pull request #3796 from CliMA/sk/update_inputs

Restructure radiation callback

Author: akshaysridhar

src/callbacks/callbacks.jl

@@ -124,278 +124,18 @@ NVTX.@annotate function cloud_fraction_model_callback!(integrator)
     set_cloud_fraction!(Y, p, p.atmos.moisture_model, p.atmos.cloud_model)
 end
 
-# TODO: Move this somewhere else
-update_o3!(_, _, _) = nothing
-function update_o3!(p, t, ::PrescribedOzone)
-    evaluate!(p.tracers.o3, p.tracers.prescribed_o3_timevaryinginput, t)
-    return nothing
-end
-
-update_co2!(_, _, _) = nothing
-function update_co2!(p, t, ::MaunaLoaCO2)
-    evaluate!(p.tracers.co2, p.tracers.prescribed_co2_timevaryinginput, t)
-    return nothing
-end
-
-
 NVTX.@annotate function rrtmgp_model_callback!(integrator)
     Y = integrator.u
     p = integrator.p
     t = integrator.t
 
-    (; ᶜts, ᶜp, cloud_diagnostics_tuple, sfc_conditions) = p.precomputed
     (; params) = p
     (; ᶠradiation_flux, rrtmgp_model) = p.radiation
     (; radiation_mode) = p.atmos
 
-    # If we have prescribed ozone or aerosols, we need to update them
-    update_o3!(p, t, p.atmos.ozone)
-    update_co2!(p, t, p.atmos.co2)
-    if :prescribed_aerosols_field in propertynames(p.tracers)
-        for (key, tv) in pairs(p.tracers.prescribed_aerosol_timevaryinginputs)
-            field = getproperty(p.tracers.prescribed_aerosols_field, key)
-            evaluate!(field, tv, t)
-        end
-    end
-    if :prescribed_clouds_field in propertynames(p.radiation)
-        for (key, tv) in pairs(p.radiation.prescribed_cloud_timevaryinginputs)
-            field = getproperty(p.radiation.prescribed_clouds_field, key)
-            evaluate!(field, tv, t)
-        end
-    end
-
-    FT = Spaces.undertype(axes(Y.c))
-    thermo_params = CAP.thermodynamics_params(params)
-    cmc = CAP.microphysics_cloud_params(params)
-    T_min = CAP.optics_lookup_temperature_min(params)
-    T_max = CAP.optics_lookup_temperature_max(params)
-
-    sfc_ts = sfc_conditions.ts
-    sfc_T = Fields.array2field(rrtmgp_model.surface_temperature, axes(sfc_ts))
-    @. sfc_T = TD.air_temperature(thermo_params, sfc_ts)
-
-    rrtmgp_model.center_pressure .= Fields.field2array(ᶜp)
-    ᶜT = Fields.array2field(rrtmgp_model.center_temperature, axes(Y.c))
-    # TODO: move this to RRTMGP
-    @. ᶜT =
-        min(max(TD.air_temperature(thermo_params, ᶜts), FT(T_min)), FT(T_max))
-
-    if !(radiation_mode isa RRTMGPI.GrayRadiation)
-        ᶜrh =
-            Fields.array2field(rrtmgp_model.center_relative_humidity, axes(Y.c))
-        ᶜvmr_h2o = Fields.array2field(
-            rrtmgp_model.center_volume_mixing_ratio_h2o,
-            axes(Y.c),
-        )
-        if radiation_mode.idealized_h2o
-            # slowly increase the relative humidity from 0 to 0.6 to account for
-            # the fact that we have a very unrealistic initial condition
-            max_relative_humidity = FT(0.6)
-            t_increasing_humidity = FT(60 * 60 * 24 * 30)
-            if float(t) < t_increasing_humidity
-                max_relative_humidity *= float(t) / t_increasing_humidity
-            end
-            @. ᶜrh = max_relative_humidity
-
-            # temporarily store ᶜq_tot in ᶜvmr_h2o
-            ᶜq_tot = ᶜvmr_h2o
-            @. ᶜq_tot =
-                max_relative_humidity * TD.q_vap_saturation(thermo_params, ᶜts)
-
-            # filter ᶜq_tot so that it is monotonically decreasing with z
-            for i in 2:Spaces.nlevels(axes(ᶜq_tot))
-                level = Fields.field_values(Spaces.level(ᶜq_tot, i))
-                prev_level = Fields.field_values(Spaces.level(ᶜq_tot, i - 1))
-                @. level = min(level, prev_level)
-            end
-
-            # assume that ᶜq_vap = ᶜq_tot when computing ᶜvmr_h2o
-            @. ᶜvmr_h2o = TD.vol_vapor_mixing_ratio(
-                thermo_params,
-                TD.PhasePartition(ᶜq_tot),
-            )
-        else
-            @. ᶜvmr_h2o = TD.vol_vapor_mixing_ratio(
-                thermo_params,
-                TD.PhasePartition(thermo_params, ᶜts),
-            )
-            @. ᶜrh = min(max(TD.relative_humidity(thermo_params, ᶜts), 0), 1)
-        end
-    end
+    RRTMGPI.update_atmospheric_state!(integrator)
 
     set_insolation_variables!(Y, p, t, p.atmos.insolation)
-
-    if radiation_mode isa RRTMGPI.AllSkyRadiation ||
-       radiation_mode isa RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics
-        if !radiation_mode.idealized_clouds
-            ᶜΔz = Fields.Δz_field(Y.c)
-            ᶜlwp = Fields.array2field(
-                rrtmgp_model.center_cloud_liquid_water_path,
-                axes(Y.c),
-            )
-            ᶜiwp = Fields.array2field(
-                rrtmgp_model.center_cloud_ice_water_path,
-                axes(Y.c),
-            )
-            ᶜfrac = Fields.array2field(
-                rrtmgp_model.center_cloud_fraction,
-                axes(Y.c),
-            )
-            ᶜreliq = Fields.array2field(
-                rrtmgp_model.center_cloud_liquid_effective_radius,
-                axes(Y.c),
-            )
-            ᶜreice = Fields.array2field(
-                rrtmgp_model.center_cloud_ice_effective_radius,
-                axes(Y.c),
-            )
-            # RRTMGP needs lwp and iwp in g/m^2
-            kg_to_g_factor = 1000
-            m_to_um_factor = FT(1e6)
-            cloud_liquid_water_content =
-                radiation_mode.cloud isa PrescribedCloudInRadiation ?
-                p.radiation.prescribed_clouds_field.clwc :
-                cloud_diagnostics_tuple.q_liq
-            cloud_ice_water_content =
-                radiation_mode.cloud isa PrescribedCloudInRadiation ?
-                p.radiation.prescribed_clouds_field.ciwc :
-                cloud_diagnostics_tuple.q_ice
-            cloud_fraction =
-                radiation_mode.cloud isa PrescribedCloudInRadiation ?
-                p.radiation.prescribed_clouds_field.cc :
-                cloud_diagnostics_tuple.cf
-            @. ᶜlwp =
-                kg_to_g_factor * Y.c.ρ * cloud_liquid_water_content * ᶜΔz /
-                max(cloud_fraction, eps(FT))
-            @. ᶜiwp =
-                kg_to_g_factor * Y.c.ρ * cloud_ice_water_content * ᶜΔz /
-                max(cloud_fraction, eps(FT))
-            @. ᶜfrac = cloud_fraction
-            # RRTMGP needs effective radius in microns
-
-            seasalt_aero_conc = p.scratch.ᶜtemp_scalar
-            dust_aero_conc = p.scratch.ᶜtemp_scalar_2
-            SO4_aero_conc = p.scratch.ᶜtemp_scalar_3
-            @. seasalt_aero_conc = 0
-            @. dust_aero_conc = 0
-            @. SO4_aero_conc = 0
-            # Get aerosol mass concentrations if available
-            seasalt_names = [:SSLT01, :SSLT02, :SSLT03, :SSLT04, :SSLT05]
-            dust_names = [:DST01, :DST02, :DST03, :DST04, :DST05]
-            SO4_names = [:SO4]
-            if :prescribed_aerosols_field in propertynames(p.tracers)
-                aerosol_field = p.tracers.prescribed_aerosols_field
-                for aerosol_name in propertynames(aerosol_field)
-                    if aerosol_name in seasalt_names
-                        data = getproperty(aerosol_field, aerosol_name)
-                        @. seasalt_aero_conc += data
-                    elseif aerosol_name in dust_names
-                        data = getproperty(aerosol_field, aerosol_name)
-                        @. dust_aero_conc += data
-                    elseif aerosol_name in SO4_names
-                        data = getproperty(aerosol_field, aerosol_name)
-                        @. SO4_aero_conc += data
-                    end
-                end
-            end
-
-            lwp_col = p.scratch.temp_field_level
-            ᶜliquid_water_mass_concentration =
-                @. lazy(cloud_liquid_water_content * Y.c.ρ)
-            Operators.column_integral_definite!(
-                lwp_col,
-                ᶜliquid_water_mass_concentration,
-            )
-
-            @. ᶜreliq = ifelse(
-                cloud_liquid_water_content > FT(0),
-                CM.CloudDiagnostics.effective_radius_Liu_Hallet_97(
-                    cmc.liquid,
-                    Y.c.ρ,
-                    cloud_liquid_water_content / max(eps(FT), cloud_fraction),
-                    ml_N_cloud_liquid_droplets(
-                        (cmc,),
-                        dust_aero_conc,
-                        seasalt_aero_conc,
-                        SO4_aero_conc,
-                        lwp_col,
-                    ),
-                    FT(0),
-                    FT(0),
-                ) * m_to_um_factor,
-                FT(0),
-            )
-
-            @. ᶜreice = ifelse(
-                cloud_ice_water_content > FT(0),
-                CM.CloudDiagnostics.effective_radius_const(cmc.ice) *
-                m_to_um_factor,
-                FT(0),
-            )
-        end
-    end
-
-    if !(radiation_mode isa RRTMGPI.GrayRadiation)
-        if radiation_mode.aerosol_radiation
-            ᶜΔz = Fields.Δz_field(Y.c)
-
-            more_aerosols = (
-                (:center_dust1_column_mass_density, :DST01),
-                (:center_dust2_column_mass_density, :DST02),
-                (:center_dust3_column_mass_density, :DST03),
-                (:center_dust4_column_mass_density, :DST04),
-                (:center_dust5_column_mass_density, :DST05),
-                (:center_ss1_column_mass_density, :SSLT01),
-                (:center_ss2_column_mass_density, :SSLT02),
-                (:center_ss3_column_mass_density, :SSLT03),
-                (:center_ss4_column_mass_density, :SSLT04),
-                (:center_ss5_column_mass_density, :SSLT05),
-            )
-
-            aerosol_names_pair = [
-                more_aerosols...,
-                (:center_so4_column_mass_density, :SO4),
-                (:center_bcpi_column_mass_density, :CB2),
-                (:center_bcpo_column_mass_density, :CB1),
-                (:center_ocpi_column_mass_density, :OC2),
-                (:center_ocpo_column_mass_density, :OC1),
-            ]
-
-            for (rrtmgp_aerosol_name, prescribed_aerosol_name) in
-                aerosol_names_pair
-                ᶜaero_conc = Fields.array2field(
-                    getproperty(rrtmgp_model, rrtmgp_aerosol_name),
-                    axes(Y.c),
-                )
-                if prescribed_aerosol_name in
-                   propertynames(p.tracers.prescribed_aerosols_field)
-                    aerosol_field = getproperty(
-                        p.tracers.prescribed_aerosols_field,
-                        prescribed_aerosol_name,
-                    )
-                    @. ᶜaero_conc = aerosol_field * Y.c.ρ * ᶜΔz
-                else
-                    @. ᶜaero_conc = 0
-                end
-            end
-
-        end
-        if :o3 in propertynames(p.tracers)
-            ᶜvmr_o3 = Fields.array2field(
-                rrtmgp_model.center_volume_mixing_ratio_o3,
-                axes(Y.c),
-            )
-            @. ᶜvmr_o3 = p.tracers.o3
-        end
-        if :co2 in propertynames(p.tracers)
-            if pkgversion(ClimaUtilities) < v"0.1.21"
-                rrtmgp_model.volume_mixing_ratio_co2 .= p.tracers.co2
-            else
-                rrtmgp_model.volume_mixing_ratio_co2 .= p.tracers.co2[]
-            end
-        end
-    end
-
     set_surface_albedo!(Y, p, t, p.atmos.surface_albedo)
 
     RRTMGPI.update_fluxes!(rrtmgp_model, UInt32(floor(t / integrator.p.dt)))
@@ -483,7 +223,7 @@ function set_insolation_variables!(Y, p, t, tvi::TimeVaryingInsolation)
                 current_datetime,
                 date0,
                 insolation_params,
-            )
+            ),
         )
     bottom_coords = Fields.coordinate_field(Spaces.level(Y.c, 1))
     cos_zenith =

src/parameterized_tendencies/radiation/RRTMGPInterface.jl

@@ -10,7 +10,6 @@ import Adapt
 import ClimaComms
 using NVTX
 using Random
-
 # TODO: Move this file to RRTMGP.jl, once the interface has been settled.
 # It will be faster to do interface development in the same repo as experiment
 # development, but, since this is just a user-friendly wrapper for RRTMGP.jl, we
@@ -1025,6 +1024,7 @@ NVTX.@annotate function update_fluxes!(model, seedval)
     model.radiation_mode.add_isothermal_boundary_layer &&
         update_boundary_layer!(model)
     clip_values!(model)
+    # TODO : update_concentrations as part of update_atmospheric_state! 
     update_concentrations!(model.radiation_mode, model)
     update_lw_fluxes!(model.radiation_mode, model)
     update_sw_fluxes!(model.radiation_mode, model)
@@ -1262,4 +1262,5 @@ function update_net_fluxes!(::AllSkyRadiationWithClearSkyDiagnostics, model)
     model.face_flux .= model.face_lw_flux .+ model.face_sw_flux
 end
 
+include("update_inputs.jl")
 end # end module