Add support for aerosol optics. (#510)

docs/src/optics/AtmosphericStates.md

@@ -8,6 +8,7 @@ CurrentModule = RRTMGP.AtmosphericStates
 AtmosphericState
 GrayAtmosphericState
 CloudState
+AerosolState
 GrayOpticalThicknessSchneider2004
 GrayOpticalThicknessOGorman2008
 ```

docs/src/optics/LookUpTables.md

@@ -13,5 +13,6 @@ LookUpLW
 LookUpSW
 LookUpCld
 PadeCld
+LookUpAerosolMerra
 AbstractLookUp
 ```

docs/src/optics/Optics.md

@@ -9,11 +9,13 @@ AbstractOpticalProps
 OneScalar
 TwoStream
 compute_col_gas!
+compute_relative_humidity!
 compute_optical_props!
 ```
 
 ```@docs
 compute_col_gas_kernel!
+compute_relative_humidity_kernel!
 compute_interp_frac_temp
 compute_interp_frac_press
 compute_interp_frac_η
@@ -31,6 +33,11 @@ compute_pade_cld_props
 pade_eval
 ```
 
+```@docs
+add_aerosol_optics_1scalar!
+add_aerosol_optics_2stream!
+```
+
 ```@docs
 compute_gray_optical_thickness_lw
 compute_gray_optical_thickness_sw

ext/RRTMGPCUDAExt.jl

@@ -15,17 +15,20 @@ import RRTMGP.Optics
 import RRTMGP.Optics: compute_optical_props!
 import RRTMGP.Optics: compute_col_gas!
 import RRTMGP.Optics: compute_col_gas_kernel!
+import RRTMGP.Optics: compute_relative_humidity!
+import RRTMGP.Optics: compute_relative_humidity_kernel!
 import RRTMGP.AtmosphericStates: setup_gray_as_pr_grid!
 import RRTMGP.AtmosphericStates: setup_gray_as_pr_grid_kernel!
 import RRTMGP.AtmosphericStates: GrayAtmosphericState
 import RRTMGP.AtmosphericStates: AtmosphericState
 import RRTMGP.AtmosphericStates: CloudState
+import RRTMGP.AtmosphericStates: AerosolState
 import RRTMGP.AtmosphericStates
 import RRTMGP.GrayUtils: update_profile_lw!
 import RRTMGP.GrayUtils: compute_gray_heating_rate!
 import RRTMGP.GrayUtils: compute_gray_heating_rate_kernel!
 import RRTMGP.GrayUtils: update_profile_lw_kernel!
-import RRTMGP.LookUpTables: LookUpLW, LookUpCld, PadeCld, LookUpSW
+import RRTMGP.LookUpTables: LookUpLW, LookUpCld, PadeCld, LookUpSW, LookUpAerosolMerra
 import RRTMGP.RTESolver: rte_lw_noscat_solve!
 import RRTMGP.RTESolver: rte_lw_noscat_one_angle!
 import RRTMGP.RTESolver: rte_lw_2stream_solve!

ext/cuda/optics.jl

@@ -28,3 +28,33 @@ function compute_col_gas_CUDA!(col_dry, args...)
     end
     return nothing
 end
+
+function compute_relative_humidity!(
+    device::ClimaComms.CUDADevice,
+    rh::AbstractArray{FT, 2},
+    p_lay::AbstractArray{FT, 2},
+    t_lay::AbstractArray{FT, 2},
+    param_set::RP.ARP,
+    vmr_h2o::AbstractArray{FT, 2},
+) where {FT}
+    nlay, ncol = size(p_lay)
+    # ratio of water to dry air molecular weights
+    mwd = RP.molmass_water(param_set) / RP.molmass_dryair(param_set)
+    t_ref = FT(273.16) # reference temperature (K)
+    q_lay_min = FT(1e-7) # minimum water mass mixing ratio
+
+    args = (p_lay, t_lay, vmr_h2o, mwd, t_ref, q_lay_min)
+    tx, bx = _configure_threadblock(nlay * ncol)
+    @cuda always_inline = true threads = (tx) blocks = (bx) compute_relative_humidity_CUDA!(rh, args...)
+    return nothing
+end
+
+function compute_relative_humidity_CUDA!(rh, args...)
+    glx = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    nlay, ncol = size(rh)
+    if glx ≤ nlay * ncol
+        glay = (glx % nlay == 0) ? nlay : (glx % nlay)
+        gcol = cld(glx, nlay)
+        compute_relative_humidity_kernel!(rh, args..., glay, gcol)
+    end
+end

ext/cuda/rte_longwave_1scalar.jl

@@ -53,11 +53,12 @@ function rte_lw_noscat_solve!(
     as::AtmosphericState,
     lookup_lw::LookUpLW,
     lookup_lw_cld::Union{LookUpCld, Nothing} = nothing,
+    lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
 )
     nlay, ncol = AtmosphericStates.get_dims(as)
     nlev = nlay + 1
     tx, bx = _configure_threadblock(ncol)
-    args = (flux, flux_lw, src_lw, bcs_lw, op, angle_disc, nlay, ncol, as, lookup_lw, lookup_lw_cld)
+    args = (flux, flux_lw, src_lw, bcs_lw, op, angle_disc, nlay, ncol, as, lookup_lw, lookup_lw_cld, lookup_lw_aero)
     @cuda always_inline = true threads = (tx) blocks = (bx) rte_lw_noscat_solve_CUDA!(args...)
     return nothing
 end
@@ -73,7 +74,8 @@ function rte_lw_noscat_solve_CUDA!(
     ncol,
     as::AtmosphericState,
     lookup_lw::LookUpLW,
-    lookup_lw_cld::Union{LookUpCld, Nothing} = nothing,
+    lookup_lw_cld::Union{LookUpCld, Nothing},
+    lookup_lw_aero::Union{LookUpAerosolMerra, Nothing},
 )
     gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
     nlev = nlay + 1
@@ -97,7 +99,7 @@ function rte_lw_noscat_solve_CUDA!(
                 )
             end
             igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
-            compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
+            compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld, lookup_lw_aero)
             rte_lw_noscat_one_angle!(src_lw, bcs_lw, op, Ds, w_μ, gcol, flux, igpt, ibnd, nlay, nlev)
             add_to_flux!(flux_lw, flux, gcol)
         end

ext/cuda/rte_longwave_2stream.jl

@@ -49,11 +49,12 @@ function rte_lw_2stream_solve!(
     as::AtmosphericState,
     lookup_lw::LookUpLW,
     lookup_lw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+    lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
 )
     nlay, ncol = AtmosphericStates.get_dims(as)
     nlev = nlay + 1
     tx, bx = _configure_threadblock(ncol)
-    args = (flux, flux_lw, src_lw, bcs_lw, op, nlay, ncol, as, lookup_lw, lookup_lw_cld)
+    args = (flux, flux_lw, src_lw, bcs_lw, op, nlay, ncol, as, lookup_lw, lookup_lw_cld, lookup_lw_aero)
     @cuda always_inline = true threads = (tx) blocks = (bx) rte_lw_2stream_solve_CUDA!(args...)
     return nothing
 end
@@ -68,7 +69,8 @@ function rte_lw_2stream_solve_CUDA!(
     ncol,
     as::AtmosphericState,
     lookup_lw::LookUpLW,
-    lookup_lw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+    lookup_lw_cld::Union{LookUpCld, PadeCld, Nothing},
+    lookup_lw_aero::Union{LookUpAerosolMerra, Nothing},
 )
     gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
     nlev = nlay + 1
@@ -89,7 +91,7 @@ function rte_lw_2stream_solve_CUDA!(
                     cloud_state.mask_type,
                 )
             end
-            compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
+            compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld, lookup_lw_aero)
             rte_lw_2stream!(op, flux, src_lw, bcs_lw, gcol, igpt, ibnd, nlev, ncol)
             if igpt == 1
                 map!(x -> x, view(flux_up_lw, :, gcol), view(flux_up, :, gcol))

ext/cuda/rte_shortwave_2stream.jl

@@ -66,12 +66,13 @@ function rte_sw_2stream_solve!(
     as::AtmosphericState,
     lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+    lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
 )
     nlay, ncol = AtmosphericStates.get_dims(as)
     nlev = nlay + 1
     n_gpt = length(lookup_sw.solar_src_scaled)
     tx, bx = _configure_threadblock(ncol)
-    args = (flux, flux_sw, op, bcs_sw, src_sw, nlay, ncol, as, lookup_sw, lookup_sw_cld)
+    args = (flux, flux_sw, op, bcs_sw, src_sw, nlay, ncol, as, lookup_sw, lookup_sw_cld, lookup_sw_aero)
     @cuda always_inline = true threads = (tx) blocks = (bx) rte_sw_2stream_solve_CUDA!(args...)
     return nothing
 end
@@ -87,6 +88,7 @@ function rte_sw_2stream_solve_CUDA!(
     as::AtmosphericState,
     lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+    lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
 )
     gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
     nlev = nlay + 1
@@ -111,7 +113,7 @@ function rte_sw_2stream_solve_CUDA!(
                     )
                 end
                 # compute optical properties
-                compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
+                compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld, lookup_sw_aero)
                 solar_frac = lookup_sw.solar_src_scaled[igpt]
                 ibnd = lookup_sw.band_data.major_gpt2bnd[igpt]
                 # rte shortwave solver

src/optics/AerosolOptics.jl

@@ -0,0 +1,198 @@
+"""
+    add_aerosol_optics_1scalar!(
+        τ,
+        aero_type,
+        aero_size,
+        aero_mass,
+        rel_hum,
+        lkp_aero::LookUpAerosolMerra,
+        ibnd,
+    )
+
+This function computes the `OneScalar` aerosol optics properties and adds them
+to the exising `OneScalar` optics properties.
+"""
+@inline function add_aerosol_optics_1scalar!(
+    τ,
+    aero_type,
+    aero_size,
+    aero_mass,
+    rel_hum,
+    lkp_aero::LookUpAerosolMerra,
+    ibnd,
+)
+    FT = eltype(τ)
+    @inbounds begin
+        nlay = length(τ)
+        for glay in 1:nlay
+            aero_type_lay = aero_type[glay]
+            if aero_type_lay ≠ 0
+                τ_aero, τ_ssa_aero, τ_ssag_aero = compute_lookup_aerosol(
+                    lkp_aero,
+                    ibnd,
+                    aero_type_lay,
+                    aero_mass[glay],
+                    aero_size[glay],
+                    rel_hum[glay],
+                )
+                τ[glay] += (τ_aero - τ_ssa_aero)
+            end
+        end
+    end
+    return nothing
+end
+
+"""
+    add_aerosol_optics_2stream!(
+        τ,
+        ssa,
+        g,
+        aero_type,
+        aero_size,
+        aero_mass,
+        rel_hum,
+        lkp_aero::LookUpAerosolMerra,
+        ibnd;
+        delta_scaling = false,
+    )
+
+This function computes the `TwoStream` aerosol optics properties and adds them
+to the exising `TwoStream` optics properties.
+"""
+@inline function add_aerosol_optics_2stream!(
+    τ,
+    ssa,
+    g,
+    aero_type,
+    aero_size,
+    aero_mass,
+    rel_hum,
+    lkp_aero::LookUpAerosolMerra,
+    ibnd;
+    delta_scaling = false,
+)
+    @inbounds begin
+        nlay = length(τ)
+        FT = eltype(τ)
+        for glay in 1:nlay
+            aero_type_lay = aero_type[glay]
+            if aero_type_lay > 0
+                τ_aero, τ_ssa_aero, τ_ssag_aero = compute_lookup_aerosol(
+                    lkp_aero,
+                    ibnd,
+                    aero_type_lay,
+                    aero_mass[glay],
+                    aero_size[glay],
+                    rel_hum[glay],
+                )
+                g_aero = τ_ssag_aero / max(eps(FT), τ_ssa_aero)
+                ssa_aero = τ_ssa_aero / max(eps(FT), τ_aero)
+                if delta_scaling # delta scaling is applied for shortwave problem
+                    τ_aero, ssa_aero, g_aero = delta_scale(τ_aero, ssa_aero, g_aero)
+                end
+                τ[glay], ssa[glay], g[glay] = increment_2stream(τ[glay], ssa[glay], g[glay], τ_aero, ssa_aero, g_aero)
+
+            end
+        end
+    end
+    return nothing
+end
+
+function compute_lookup_aerosol(lkp_aero, ibnd::Int, aerotype::Int, aeromass::FT, aerosize::FT, rh::FT) where {FT}
+    cargs = (lkp_aero, ibnd, aeromass) # common args
+    return aerotype == 1 ? compute_lookup_dust_props(cargs..., aerosize) :
+           (
+        aerotype == 2 ? compute_lookup_sea_salt_props(cargs..., aerosize, rh) :
+        (
+            aerotype == 3 ? compute_lookup_sulfate_props(cargs..., rh) :
+            (
+                aerotype == 4 ? compute_lookup_black_carbon_props(cargs...) :
+                (
+                    aerotype == 5 ? compute_lookup_black_carbon_props(cargs..., rh) :
+                    (
+                        aerotype == 6 ? compute_lookup_organic_carbon_props(cargs...) :
+                        (aerotype == 7 ? compute_lookup_organic_carbon_props(cargs..., rh) : (FT(0), FT(0), FT(0)))
+                    )
+                )
+            )
+        )
+    )
+end
+
+
+function compute_lookup_dust_props(lkp_aero, ibnd::Int, aeromass::FT, aerosize::FT) where {FT}
+    (; size_bin_limits, dust) = lkp_aero
+    bin = locate_merra_size_bin(size_bin_limits, aerosize)
+    τ = aeromass * dust[1, bin, ibnd]
+    τ_ssa = τ * dust[2, bin, ibnd]
+    τ_ssag = τ_ssa * dust[3, bin, ibnd]
+    return τ, τ_ssa, τ_ssag
+end
+
+function compute_lookup_sea_salt_props(lkp_aero, ibnd::Int, aeromass::FT, aerosize::FT, rh::FT) where {FT}
+    return (FT(0), FT(0), FT(0))
+end
+
+function compute_lookup_sulfate_props(lkp_aero, ibnd::Int, aeromass::FT, rh::FT) where {FT}
+    (; rh_levels, sulfate) = lkp_aero
+    @inbounds begin
+        loc, factor = interp1d_loc_factor(rh, rh_levels)
+        τ = aeromass * (sulfate[1, loc, ibnd] * (FT(1) - factor) + sulfate[1, loc + 1, ibnd] * factor)
+        τ_ssa = τ * (sulfate[2, loc, ibnd] * (FT(1) - factor) + sulfate[2, loc + 1, ibnd] * factor)
+        τ_ssag = τ_ssa * (sulfate[3, loc, ibnd] * (FT(1) - factor) + sulfate[3, loc + 1, ibnd] * factor)
+    end
+    return τ, τ_ssa, τ_ssag
+end
+
+function compute_lookup_black_carbon_props(lkp_aero, ibnd::Int, aeromass::FT) where {FT}
+    (; black_carbon) = lkp_aero
+    τ = aeromass * black_carbon[1, ibnd]
+    τ_ssa = τ * black_carbon[2, ibnd]
+    τ_ssag = τ_ssa * black_carbon[3, ibnd]
+    return τ, τ_ssa, τ_ssag
+end
+
+function compute_lookup_black_carbon_props(lkp_aero, ibnd::Int, aeromass::FT, rh::FT) where {FT}
+    (; rh_levels, black_carbon_rh) = lkp_aero
+    @inbounds begin
+        loc, factor = interp1d_loc_factor(rh, rh_levels)
+        τ = aeromass * (black_carbon_rh[1, loc, ibnd] * (FT(1) - factor) + black_carbon_rh[1, loc + 1, ibnd] * factor)
+        τ_ssa = τ * (black_carbon_rh[2, loc, ibnd] * (FT(1) - factor) + black_carbon_rh[2, loc + 1, ibnd] * factor)
+        τ_ssag = τ_ssa * (black_carbon_rh[3, loc, ibnd] * (FT(1) - factor) + black_carbon_rh[3, loc + 1, ibnd] * factor)
+    end
+    return τ, τ_ssa, τ_ssag
+end
+
+function compute_lookup_organic_carbon_props(lkp_aero, ibnd::Int, aeromass::FT) where {FT}
+    (; organic_carbon) = lkp_aero
+    τ = aeromass * organic_carbon[1, ibnd]
+    τ_ssa = τ * organic_carbon[2, ibnd]
+    τ_ssag = τ_ssa * organic_carbon[3, ibnd]
+    return τ, τ_ssa, τ_ssag
+end
+
+function compute_lookup_organic_carbon_props(lkp_aero, ibnd::Int, aeromass::FT, rh::FT) where {FT}
+    (; rh_levels, organic_carbon_rh) = lkp_aero
+    @inbounds begin
+        loc, factor = interp1d_loc_factor(rh, rh_levels)
+        τ =
+            aeromass *
+            (organic_carbon_rh[1, loc, ibnd] * (FT(1) - factor) + organic_carbon_rh[1, loc + 1, ibnd] * factor)
+        τ_ssa = τ * (organic_carbon_rh[2, loc, ibnd] * (FT(1) - factor) + organic_carbon_rh[2, loc + 1, ibnd] * factor)
+        τ_ssag =
+            τ_ssa * (organic_carbon_rh[3, loc, ibnd] * (FT(1) - factor) + organic_carbon_rh[3, loc + 1, ibnd] * factor)
+    end
+    return τ, τ_ssa, τ_ssag
+end
+
+function locate_merra_size_bin(size_bin_limits, aerosize)
+    nbins = size(size_bin_limits, 2)
+    bin = 1
+    for ibin in 1:nbins
+        if size_bin_limits[1, ibin] ≤ aerosize ≤ size_bin_limits[2, ibin]
+            bin = ibin
+            break
+        end
+    end
+    return bin
+end