diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl
index 35ce03fc3..b497cbc01 100644
--- a/src/optics/AtmosphericStates.jl
+++ b/src/optics/AtmosphericStates.jl
@@ -10,6 +10,7 @@ using ..Vmrs
 
 export AbstractAtmosphericState,
     AtmosphericState,
+    AtmosphericStatec,
     GrayAtmosphericState,
     setup_gray_as_pr_grid,
     setup_gray_as_alt_grid,
@@ -44,9 +45,8 @@ struct AtmosphericState{
     FTA2D <: AbstractArray{FT, 2},
     CLDP <: Union{AbstractArray{FT, 2}, Nothing},
     CLDM <: Union{AbstractArray{Bool, 2}, Nothing},
-    RND <: Union{AbstractArray{FT, 2}, Nothing},
     CMASK <: Union{AbstractCloudMask, Nothing},
-    VMR <: AbstractVmr{FT},
+    VMR <: AbstractVmr,
 } <: AbstractAtmosphericState{FT, FTA1D}
     "longitude, in degrees (`ncol`), optional"
     lon::FTA1DN
@@ -76,10 +76,6 @@ struct AtmosphericState{
     cld_path_ice::CLDP
     "cloud fraction"
     cld_frac::CLDP
-    "random number storage for longwave bands `(nlay, ncol)`"
-    random_lw::RND
-    "random number storage for shortwave bands `(nlay, ncol)`"
-    random_sw::RND
     "cloud mask (longwave), = true if clouds are present"
     cld_mask_lw::CLDM
     "cloud mask (shortwave), = true if clouds are present"
@@ -96,4 +92,54 @@ struct AtmosphericState{
     ngas::Int
 end
 Adapt.@adapt_structure AtmosphericState
+
+struct AtmosphericStatec{FTVA1DN, FTVA1D, FTA1D, CLDP, CLDM, CMASK, VMR}
+    lon::FTVA1DN
+    lat::FTVA1DN
+    p_lay::FTA1D
+    p_lev::FTA1D
+    t_lay::FTA1D
+    t_lev::FTA1D
+    t_sfc::FTVA1D
+    col_dry::FTA1D
+    vmr::VMR
+    cld_r_eff_liq::CLDP
+    cld_r_eff_ice::CLDP
+    cld_path_liq::CLDP
+    cld_path_ice::CLDP
+    cld_frac::CLDP
+    cld_mask_lw::CLDM
+    cld_mask_sw::CLDM
+    cld_mask_type::CMASK
+    ice_rgh::Int
+    nlay::Int
+    ncol::Int
+    ngas::Int
+end
+Adapt.@adapt_structure AtmosphericStatec
+
+AtmosphericStatec(as::AtmosphericState, gcol::Int) = AtmosphericStatec(
+    isnothing(as.lon) ? nothing : view(as.lon, gcol:gcol),
+    isnothing(as.lat) ? nothing : view(as.lat, gcol:gcol),
+    view(as.p_lay, :, gcol),
+    view(as.p_lev, :, gcol),
+    view(as.t_lay, :, gcol),
+    view(as.t_lev, :, gcol),
+    view(as.t_sfc, gcol:gcol),
+    view(as.col_dry, :, gcol),
+    Vmrs.get_col_view(as.vmr, gcol),
+    isnothing(as.cld_r_eff_liq) ? nothing : view(as.cld_r_eff_liq, :, gcol),
+    isnothing(as.cld_r_eff_ice) ? nothing : view(as.cld_r_eff_ice, :, gcol),
+    isnothing(as.cld_path_liq) ? nothing : view(as.cld_path_liq, :, gcol),
+    isnothing(as.cld_path_ice) ? nothing : view(as.cld_path_ice, :, gcol),
+    isnothing(as.cld_frac) ? nothing : view(as.cld_frac, :, gcol),
+    isnothing(as.cld_mask_lw) ? nothing : view(as.cld_mask_lw, :, gcol),
+    isnothing(as.cld_mask_sw) ? nothing : view(as.cld_mask_sw, :, gcol),
+    as.cld_mask_type,
+    as.ice_rgh,
+    as.nlay,
+    as.ncol,
+    as.ngas,
+)
+
 end
diff --git a/src/optics/BCs.jl b/src/optics/BCs.jl
index 934268a62..ff4300672 100644
--- a/src/optics/BCs.jl
+++ b/src/optics/BCs.jl
@@ -3,7 +3,7 @@ module BCs
 using DocStringExtensions
 using Adapt
 
-export LwBCs, SwBCs
+export LwBCs, SwBCs, SwBCsc
 
 """
     LwBCs{FT,FTA1D,FTA2DN}
@@ -29,20 +29,20 @@ Shortwave boundary conditions
 $(DocStringExtensions.FIELDS)
 """
 struct SwBCs{FT, FTA1D, FTA1DN, FTA2D}
-    "zenith angle `(ncol)`"
-    zenith::FTA1D
+    "cosine of zenith angle `(ncol)`"
+    cos_zenith::FTA1D
     "top of atmosphere flux `(ncol)`"
     toa_flux::FTA1D
     "surface albedo for specular (direct) radiation `(nbnd, ncol)`"
     sfc_alb_direct::FTA2D
-    "incident diffuse flux at top of domain `[W/m2]` `(ncol, ngpt)`"
+    "incident diffuse flux at top of domain `[W/m2]` `(ngpt, ncol)`"
     inc_flux_diffuse::FTA1DN
     "surface albedo for diffuse radiation `(nbnd, ncol)`"
     sfc_alb_diffuse::FTA2D
 end
-SwBCs(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse) =
-    SwBCs{eltype(zenith), typeof(zenith), typeof(inc_flux_diffuse), typeof(sfc_alb_direct)}(
-        zenith,
+SwBCs(cos_zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse) =
+    SwBCs{eltype(cos_zenith), typeof(cos_zenith), typeof(inc_flux_diffuse), typeof(sfc_alb_direct)}(
+        cos_zenith,
         toa_flux,
         sfc_alb_direct,
         inc_flux_diffuse,
@@ -50,4 +50,24 @@ SwBCs(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse) =
     )
 Adapt.@adapt_structure SwBCs
 
+struct SwBCsc{FTVA1D, FTA1D, FTA1DN}
+    cos_zenith::FTVA1D
+    toa_flux::FTVA1D
+    sfc_alb_direct::FTA1D
+    inc_flux_diffuse::FTA1DN
+    sfc_alb_diffuse::FTA1D
+end
+Adapt.@adapt_structure SwBCsc
+
+function SwBCsc(swbcs::SwBCs, gcol::Int)
+    inc_flux_diffuse = swbcs.inc_flux_diffuse isa Nothing ? nothing : view(swbcs.inc_flux_diffuse, :, gcol)
+    return SwBCsc(
+        view(swbcs.cos_zenith, gcol:gcol),
+        view(swbcs.toa_flux, gcol:gcol),
+        view(swbcs.sfc_alb_direct, :, gcol),
+        inc_flux_diffuse,
+        view(swbcs.sfc_alb_diffuse, :, gcol),
+    )
+end
+
 end
diff --git a/src/optics/CloudOptics.jl b/src/optics/CloudOptics.jl
index f21bf0a9e..763de2884 100644
--- a/src/optics/CloudOptics.jl
+++ b/src/optics/CloudOptics.jl
@@ -1,77 +1,19 @@
-
-"""
-    add_cloud_optics_2stream(
-        op::TwoStream,
-        as::AtmosphericState,
-        lkp::LookUpLW,
-        lkp_cld,
-        glay, gcol,
-        ibnd,
-        igpt,
-    )
-
-This function computes the longwave TwoStream clouds optics properties and adds them
-to the TwoStream longwave gas optics properties.
-"""
-function add_cloud_optics_2stream(op::TwoStream, as::AtmosphericState, lkp::LookUpLW, lkp_cld, glay, gcol, ibnd, igpt)
-    if as.cld_mask_lw isa AbstractArray
-        cld_mask = as.cld_mask_lw[glay, gcol]
-        if cld_mask
-            τ_cl, τ_cl_ssa, τ_cl_ssag = compute_cld_props(lkp_cld, as, cld_mask, glay, gcol, ibnd, igpt)
-            op.τ[glay, gcol], op.ssa[glay, gcol], op.g[glay, gcol] =
-                increment_2stream(op.τ[glay, gcol], op.ssa[glay, gcol], op.g[glay, gcol], τ_cl, τ_cl_ssa, τ_cl_ssag)
-        end
-    end
-    return nothing
-end
-
-"""
-    add_cloud_optics_2stream(
-        op::TwoStream,
-        as::AtmosphericState,
-        lkp::LookUpSW,
-        lkp_cld,
-        glay, gcol,
-        ibnd,
-        igpt,
-    )
-
-This function computes the shortwave TwoStream clouds optics properties and adds them
-to the TwoStream shortwave gase optics properties.
-"""
-function add_cloud_optics_2stream(op::TwoStream, as::AtmosphericState, lkp::LookUpSW, lkp_cld, glay, gcol, ibnd, igpt)
-    if as.cld_mask_sw isa AbstractArray
-        cld_mask = as.cld_mask_sw[glay, gcol]
-        if cld_mask
-            τ_cl, τ_cl_ssa, τ_cl_ssag = compute_cld_props(lkp_cld, as, cld_mask, glay, gcol, ibnd, igpt)
-            τ_cl, τ_cl_ssa, τ_cl_ssag = delta_scale(τ_cl, τ_cl_ssa, τ_cl_ssag)
-            op.τ[glay, gcol], op.ssa[glay, gcol], op.g[glay, gcol] =
-                increment_2stream(op.τ[glay, gcol], op.ssa[glay, gcol], op.g[glay, gcol], τ_cl, τ_cl_ssa, τ_cl_ssag)
-        end
-    end
-    return nothing
-end
-"""
-    add_cloud_optics_2stream(op::OneScalar, args...)
-
-Cloud optics is currently only supported for the TwoStream solver.
-"""
-function add_cloud_optics_2stream(op::OneScalar, args...)
-    return nothing
-end
 """
     compute_cld_props(lkp_cld, as, glay, gcol, ibnd, igpt)
 
 This function computed the TwoSteam cloud optics properties using either the 
 lookup table method or pade method.
 """
-function compute_cld_props(lkp_cld::LookUpCld, as, cld_mask, glay, gcol, ibnd, igpt)
-    re_liq = as.cld_r_eff_liq[glay, gcol]
-    re_ice = as.cld_r_eff_ice[glay, gcol]
-    ice_rgh = as.ice_rgh
-    cld_path_liq = as.cld_path_liq[glay, gcol]
-    cld_path_ice = as.cld_path_ice[glay, gcol]
-    FT = eltype(re_liq)
+function compute_cld_props(
+    lkp_cld::LookUpCld,
+    re_liq::FT,
+    re_ice::FT,
+    ice_rgh::Int,
+    cld_path_liq::FT,
+    cld_path_ice::FT,
+    ibnd::UInt8,
+) where {FT}
+    #FT = eltype(re_liq)
     τl, τl_ssa, τl_ssag = FT(0), FT(0), FT(0)
     τi, τi_ssa, τi_ssag = FT(0), FT(0), FT(0)
     (;
@@ -95,18 +37,22 @@ function compute_cld_props(lkp_cld::LookUpCld, as, cld_mask, glay, gcol, ibnd, i
         loc = Int(max(min(unsafe_trunc(Int, (re_liq - radliq_lwr) / Δr_liq) + 1, nsize_liq - 1), 1))
         fac = (re_liq - radliq_lwr - (loc - 1) * Δr_liq) / Δr_liq
         fc1 = FT(1) - fac
-        τl = (fc1 * lut_extliq[loc, ibnd] + fac * lut_extliq[loc + 1, ibnd]) * cld_path_liq
-        τl_ssa = (fc1 * lut_ssaliq[loc, ibnd] + fac * lut_ssaliq[loc + 1, ibnd]) * τl
-        τl_ssag = (fc1 * lut_asyliq[loc, ibnd] + fac * lut_asyliq[loc + 1, ibnd]) * τl_ssa
+        @inbounds begin
+            τl = (fc1 * lut_extliq[loc, ibnd] + fac * lut_extliq[loc + 1, ibnd]) * cld_path_liq
+            τl_ssa = (fc1 * lut_ssaliq[loc, ibnd] + fac * lut_ssaliq[loc + 1, ibnd]) * τl
+            τl_ssag = (fc1 * lut_asyliq[loc, ibnd] + fac * lut_asyliq[loc + 1, ibnd]) * τl_ssa
+        end
     end
     # cloud ice particles
     if cld_path_ice > eps(FT)
         loc = Int(max(min(unsafe_trunc(Int, (re_ice - radice_lwr) / Δr_ice) + 1, nsize_ice - 1), 1))
         fac = (re_ice - radice_lwr - (loc - 1) * Δr_ice) / Δr_ice
         fc1 = FT(1) - fac
-        τi = (fc1 * lut_extice[loc, ibnd, ice_rgh] + fac * lut_extice[loc + 1, ibnd, ice_rgh]) * cld_path_ice
-        τi_ssa = (fc1 * lut_ssaice[loc, ibnd, ice_rgh] + fac * lut_ssaice[loc + 1, ibnd, ice_rgh]) * τi
-        τi_ssag = (fc1 * lut_asyice[loc, ibnd, ice_rgh] + fac * lut_asyice[loc + 1, ibnd, ice_rgh]) * τi_ssa
+        @inbounds begin
+            τi = (fc1 * lut_extice[loc, ibnd, ice_rgh] + fac * lut_extice[loc + 1, ibnd, ice_rgh]) * cld_path_ice
+            τi_ssa = (fc1 * lut_ssaice[loc, ibnd, ice_rgh] + fac * lut_ssaice[loc + 1, ibnd, ice_rgh]) * τi
+            τi_ssag = (fc1 * lut_asyice[loc, ibnd, ice_rgh] + fac * lut_asyice[loc + 1, ibnd, ice_rgh]) * τi_ssa
+        end
     end
 
     τ = τl + τi
@@ -117,13 +63,16 @@ function compute_cld_props(lkp_cld::LookUpCld, as, cld_mask, glay, gcol, ibnd, i
     return (τ, τ_ssa, τ_ssag)
 end
 
-function compute_cld_props(lkp_cld::PadeCld, as, cld_mask, glay, gcol, ibnd, igpt)
-    re_liq = as.cld_r_eff_liq[glay, gcol]
-    re_ice = as.cld_r_eff_ice[glay, gcol]
-    ice_rgh = as.ice_rgh
-    cld_path_liq = as.cld_path_liq[glay, gcol]
-    cld_path_ice = as.cld_path_ice[glay, gcol]
-    FT = eltype(re_liq)
+function compute_cld_props(
+    lkp_cld::PadeCld,
+    re_liq::FT,
+    re_ice::FT,
+    ice_rgh::Int,
+    cld_path_liq::FT,
+    cld_path_ice::FT,
+    ibnd::UInt8,
+) where {FT}
+    #FT = eltype(re_liq)
     τl, τl_ssa, τl_ssag = FT(0), FT(0), FT(0)
     τi, τi_ssa, τi_ssag = FT(0), FT(0), FT(0)
 
@@ -147,26 +96,30 @@ function compute_cld_props(lkp_cld::PadeCld, as, cld_mask, glay, gcol, ibnd, igp
     # This works only if there are precisely three size regimes (four bounds) and it's
     # previously guaranteed that size_bounds(1) <= size <= size_bounds(4)
     if cld_path_liq > eps(FT)
-        irad = Int(min(floor((re_liq - pade_sizreg_extliq[2]) / pade_sizreg_extliq[3]) + 2, 3))
-        τl = pade_eval(ibnd, re_liq, irad, m_ext, n_ext, pade_extliq) * cld_path_liq
+        @inbounds begin
+            irad = Int(min(floor((re_liq - pade_sizreg_extliq[2]) / pade_sizreg_extliq[3]) + 2, 3))
+            τl = pade_eval(ibnd, re_liq, irad, m_ext, n_ext, pade_extliq) * cld_path_liq
 
-        irad = Int(min(floor((re_liq - pade_sizreg_ssaliq[2]) / pade_sizreg_ssaliq[3]) + 2, 3))
-        τl_ssa = (FT(1) - max(FT(0), pade_eval(ibnd, re_liq, irad, m_ssa_g, n_ssa_g, pade_ssaliq))) * τl
+            irad = Int(min(floor((re_liq - pade_sizreg_ssaliq[2]) / pade_sizreg_ssaliq[3]) + 2, 3))
+            τl_ssa = (FT(1) - max(FT(0), pade_eval(ibnd, re_liq, irad, m_ssa_g, n_ssa_g, pade_ssaliq))) * τl
 
-        irad = Int(min(floor((re_liq - pade_sizreg_asyliq[2]) / pade_sizreg_asyliq[3]) + 2, 3))
-        τl_ssag = pade_eval(ibnd, re_liq, irad, m_ssa_g, n_ssa_g, pade_asyliq) * τl_ssa
+            irad = Int(min(floor((re_liq - pade_sizreg_asyliq[2]) / pade_sizreg_asyliq[3]) + 2, 3))
+            τl_ssag = pade_eval(ibnd, re_liq, irad, m_ssa_g, n_ssa_g, pade_asyliq) * τl_ssa
+        end
     end
 
     if cld_path_ice > eps(FT)
-        irad = Int(min(floor((re_ice - pade_sizreg_extice[2]) / pade_sizreg_extice[3]) + 2, 3))
+        @inbounds begin
+            irad = Int(min(floor((re_ice - pade_sizreg_extice[2]) / pade_sizreg_extice[3]) + 2, 3))
 
-        τi = pade_eval(ibnd, re_ice, irad, m_ext, n_ext, pade_extice, ice_rgh) * cld_path_ice
+            τi = pade_eval(ibnd, re_ice, irad, m_ext, n_ext, pade_extice, ice_rgh) * cld_path_ice
 
-        irad = Int(min(floor((re_ice - pade_sizreg_ssaice[2]) / pade_sizreg_ssaice[3]) + 2, 3))
-        τi_ssa = (FT(1) - max(FT(0), pade_eval(ibnd, re_ice, irad, m_ssa_g, n_ssa_g, pade_ssaice, ice_rgh))) * τi
+            irad = Int(min(floor((re_ice - pade_sizreg_ssaice[2]) / pade_sizreg_ssaice[3]) + 2, 3))
+            τi_ssa = (FT(1) - max(FT(0), pade_eval(ibnd, re_ice, irad, m_ssa_g, n_ssa_g, pade_ssaice, ice_rgh))) * τi
 
-        irad = Int(min(floor((re_ice - pade_sizreg_asyice[2]) / pade_sizreg_asyice[3]) + 2, 3))
-        τi_ssag = pade_eval(ibnd, re_ice, irad, m_ssa_g, n_ssa_g, pade_asyice, ice_rgh) * τi_ssa
+            irad = Int(min(floor((re_ice - pade_sizreg_asyice[2]) / pade_sizreg_asyice[3]) + 2, 3))
+            τi_ssag = pade_eval(ibnd, re_ice, irad, m_ssa_g, n_ssa_g, pade_asyice, ice_rgh) * τi_ssa
+        end
     end
 
     τ = τl + τi
@@ -216,65 +169,71 @@ function pade_eval(ibnd, re, irad, m, n, pade_coeffs, irgh::Union{Int, Nothing}
 end
 
 """
-    build_cloud_mask!(cld_mask, cld_frac, random_arr, gcol, igpt, ::MaxRandomOverlap)
+    build_cloud_mask!(cld_mask, cld_frac, ::MaxRandomOverlap)
 
 Builds McICA-sampled cloud mask from cloud fraction data for maximum-random overlap
 
 Reference: https://github.com/AER-RC/RRTMG_SW/
 """
-function build_cloud_mask!(cld_mask, cld_frac, random_arr, gcol, igpt, ::MaxRandomOverlap)
+function build_cloud_mask!(
+    cld_mask::AbstractArray{Bool, 1},
+    cld_frac::AbstractArray{FT, 1},
+    ::MaxRandomOverlap,
+) where {FT}
+    nlay = size(cld_frac, 1)
+    start = _get_start(cld_frac) # first cloudy layer
+
+    if start > 0
+        finish = _get_finish(cld_frac) # last cloudy layer
+        # set cloud mask for non-cloudy layers
+        _mask_outer_non_cloudy_layers!(cld_mask, start, finish)
+        # RRTMG uses random_arr[finish] > (FT(1) - cld_frac[finish]), 
+        # we change > to >= to address edge cases
+        @inbounds cld_frac_ilayplus1 = cld_frac[finish]
+        random_ilayplus1 = Random.rand()
+        @inbounds cld_mask[finish] = cld_mask_ilayplus1 = random_ilayplus1 >= (FT(1) - cld_frac_ilayplus1)
+        for ilay in (finish - 1):-1:start
+            @inbounds cld_frac_ilay = cld_frac[ilay]
+            if cld_frac_ilay > FT(0)
+                # use same random number from the layer above if layer above is cloudy
+                # update random numbers if layer above is not cloudy
+                random_ilay = cld_mask_ilayplus1 ? random_ilayplus1 : Random.rand() * (FT(1) - cld_frac_ilayplus1)
+                # RRTMG uses random_arr[ilay] > (FT(1) - cld_frac[ilay]), we change > to >= to address edge cases
+                cld_mask_ilay = random_ilay >= (FT(1) - cld_frac_ilay)
+                random_ilayplus1 = random_ilay
+            else
+                cld_mask_ilay = false
+            end
+            @inbounds cld_mask[ilay] = cld_mask_ilay
+            cld_frac_ilayplus1 = cld_frac_ilay
+            cld_mask_ilayplus1 = cld_mask_ilay
+        end
+    end
+    return nothing
+end
 
-    FT = eltype(cld_frac)
-    local_rand = view(random_arr, :, gcol)
-    local_cld_frac = view(cld_frac, :, gcol)
-    local_cld_mask = view(cld_mask, :, gcol)
-    nlay = size(local_rand, 1)
+function _get_finish(cld_frac)
+    @inbounds for ilay in reverse(eachindex(cld_frac))
+        cld_frac[ilay] > 0 && return ilay
+    end
+    return 0
+end
 
-    Random.rand!(local_rand)
-    start = 0
-    @inbounds for ilay in 1:nlay # for GPU compatibility (start = findfirst(local_cld_frac .> FT(0)))
-        if local_cld_frac[ilay] > FT(0)
-            start = ilay
-            break
-        end
+function _get_start(cld_frac)
+    @inbounds for ilay in eachindex(cld_frac)
+        cld_frac[ilay] > 0 && return ilay
     end
+    return 0
+end
 
-    if start == 0 # no clouds in entire column
-        @inbounds for ilay in 1:nlay
-            local_cld_mask[ilay] = false
-        end
-    else
-        # finish = findlast(local_cld_frac .> FT(0)) # for GPU compatibility
-        finish = start
-        @inbounds for ilay in nlay:-1:(start + 1)
-            if local_cld_frac[ilay] > FT(0)
-                finish = ilay
-                break
-            end
-        end
-        # set cloud mask for non-cloudy layers
-        @inbounds for ilay in 1:(start - 1)
-            local_cld_mask[ilay] = false
+function _mask_outer_non_cloudy_layers!(cld_mask, start, finish)
+    if start > 0
+        for ilay in 1:(start - 1)
+            @inbounds cld_mask[ilay] = false
         end
-        @inbounds for ilay in (finish + 1):nlay
-            local_cld_mask[ilay] = false
-        end
-        # set cloud mask for cloudy layers
-        # last layer
-        # RRTMG uses local_rand[finish] > (FT(1) - local_cld_frac[finish]), we change > to >= to address edge cases
-        @inbounds local_cld_mask[finish] = local_rand[finish] >= (FT(1) - local_cld_frac[finish])
-        @inbounds for ilay in (finish - 1):-1:start
-            if local_cld_frac[ilay] > FT(0)
-                if local_cld_mask[ilay + 1]
-                    local_rand[ilay] = local_rand[ilay + 1] # use same random number from the layer above if layer above is cloudy
-                else
-                    local_rand[ilay] *= (FT(1) - local_cld_frac[ilay + 1]) # update random numbers if layer above is not cloudy
-                end
-                # RRTMG uses local_rand[ilay] > (FT(1) - local_cld_frac[ilay]), we change > to >= to address edge cases
-                local_cld_mask[ilay] = local_rand[ilay] >= (FT(1) - local_cld_frac[ilay])
-            else
-                local_cld_mask[ilay] = false
-            end
+        nlay = length(cld_mask)
+        for ilay in (finish + 1):nlay
+            @inbounds cld_mask[ilay] = false
         end
     end
     return nothing
diff --git a/src/optics/Fluxes.jl b/src/optics/Fluxes.jl
index 55241078b..c58758832 100644
--- a/src/optics/Fluxes.jl
+++ b/src/optics/Fluxes.jl
@@ -3,7 +3,7 @@ module Fluxes
 using Adapt
 using DocStringExtensions
 
-export AbstractFlux, FluxLW, FluxSW, set_flux_to_zero!, add_to_flux!
+export AbstractFlux, FluxLW, FluxSW, set_flux_to_zero!, add_to_flux!, get_col_view, FluxSWc
 
 abstract type AbstractFlux{FT <: AbstractFloat, FTA2D <: AbstractArray{FT, 2}} end
 
@@ -64,6 +64,21 @@ function FluxSW(ncol::Int, nlay::Int, ::Type{FT}, ::Type{DA}) where {FT <: Abstr
     return FluxSW{FT, typeof(flux_net)}(flux_up, flux_dn, flux_net, flux_dn_dir)
 end
 
+struct FluxSWc{FTA1D}
+    flux_up::FTA1D
+    flux_dn::FTA1D
+    flux_net::FTA1D
+    flux_dn_dir::FTA1D
+end
+FluxSWc(flux_up, flux_dn, flux_net, flux_dn_dir) = FluxSWc{typeof(flux_up)}(flux_up, flux_dn, flux_net, flux_dn_dir)
+Adapt.@adapt_structure FluxSWc
+FluxSWc(flux::FluxSW, gcol::Int) = FluxSWc(
+    view(flux.flux_up, :, gcol),
+    view(flux.flux_dn, :, gcol),
+    view(flux.flux_net, :, gcol),
+    view(flux.flux_dn_dir, :, gcol),
+)
+
 """
     set_flux_to_zero!(flux::FluxLW{FT}) where {FT<:AbstractFloat}
 
@@ -123,6 +138,15 @@ function set_flux_to_zero!(flux::FluxSW{FT}, gcol::Int) where {FT <: AbstractFlo
     end
     return nothing
 end
+function set_flux_to_zero!(flux::FluxSWc)
+    (; flux_up, flux_dn, flux_net, flux_dn_dir) = flux
+    zval = zero(flux_up[1])
+    map!(x -> zval, flux_up, flux_up)
+    map!(x -> zval, flux_dn, flux_dn)
+    map!(x -> zval, flux_net, flux_net)
+    map!(x -> zval, flux_dn_dir, flux_dn_dir)
+    return nothing
+end
 
 """
     add_to_flux!(flux1::FluxLW, flux2::FluxLW)
@@ -190,4 +214,11 @@ function add_to_flux!(flux1::FluxSW, flux2::FluxSW, gcol)
     return nothing
 end
 
+get_col_view(flux::FluxSW, gcol::Int) = (
+    view(flux.flux_up, :, gcol),
+    view(flux.flux_dn, :, gcol),
+    view(flux.flux_dn_dir, :, gcol),
+    view(flux.flux_net, :, gcol),
+)
+
 end
diff --git a/src/optics/GasOptics.jl b/src/optics/GasOptics.jl
index 289f7e657..397950953 100644
--- a/src/optics/GasOptics.jl
+++ b/src/optics/GasOptics.jl
@@ -37,42 +37,15 @@ function compute_col_gas_kernel!(
     # Hydrostatic equation
     col_gas[glay, gcol] = (Δp * avogadro / (m2_to_cm2 * m_air * g0)) # molecules/cm^2
 end
-"""
-    compute_interp_fractions(
-        lkp::AbstractLookUp,
-        vmr,
-        p_lay,
-        t_lay,
-        tropo,
-        ibnd,
-        glay, gcol,
-    ) where {FT<:AbstractFloat}
 
-compute interpolation fractions for binary species parameter, pressure and temperature.
 """
-@inline function compute_interp_fractions(lkp::AbstractLookUp, vmr, p_lay, t_lay, tropo, ibnd, glay, gcol)
-    jftemp = compute_interp_frac_temp(lkp, t_lay, glay, gcol)
-    jfpress = compute_interp_frac_press(lkp, p_lay, tropo, glay, gcol)
-    jfη, col_mix = compute_interp_frac_η(lkp, vmr, tropo, jftemp[1], ibnd, glay, gcol)
-    return (jftemp, jfpress, jfη, col_mix)
-end
-
-"""
-    compute_interp_frac_temp(
-        lkp::AbstractLookUp,
-        t_lay,
-        glay,
-        gcol,
-    ) where {FT<:AbstractFloat}
+    compute_interp_frac_temp(Δ_t_ref, n_t_ref, t_ref, t_lay)
 
 compute interpolation fraction for temperature.
 """
-@inline function compute_interp_frac_temp(lkp::AbstractLookUp, t_lay, glay, gcol)
-    (; Δ_t_ref, n_t_ref, t_ref) = lkp
-
+@inline function compute_interp_frac_temp(Δ_t_ref, n_t_ref, t_ref, t_lay)
     @inbounds jtemp = loc_lower(t_lay, Δ_t_ref, n_t_ref, t_ref)
     @inbounds ftemp = (t_lay - t_ref[jtemp]) / Δ_t_ref
-
     return (jtemp, ftemp)
 end
 
@@ -87,9 +60,7 @@ end
 
 Compute interpolation fraction for pressure.
 """
-@inline function compute_interp_frac_press(lkp::AbstractLookUp, p_lay, tropo, glay, gcol)
-    (; Δ_ln_p_ref, p_ref, n_p_ref) = lkp
-
+@inline function compute_interp_frac_press(Δ_ln_p_ref, p_ref, n_p_ref, p_lay, tropo)
     log_p_lay = log(p_lay)
     @inbounds jpress = Int(min(max(fld(log(p_ref[1]) - log_p_lay, Δ_ln_p_ref) + 1, 1), n_p_ref - 1) + 1)
 
@@ -112,12 +83,7 @@ end
 
 Compute interpolation fraction for binary species parameter.
 """
-@inline function compute_interp_frac_η(lkp::AbstractLookUp, vmr, tropo, jtemp, ibnd, glay, gcol)
-    (; n_η, key_species, vmr_ref) = lkp
-    ig = view(key_species, :, tropo, ibnd)
-
-    vmr1 = get_vmr(vmr, ig[1], glay, gcol)
-    vmr2 = get_vmr(vmr, ig[2], glay, gcol)
+@inline function compute_interp_frac_η(n_η, ig, vmr_ref, (vmr1, vmr2), tropo, jtemp)
     FT = eltype(vmr1)
     itemp = 1
     @inbounds η_half = vmr_ref[tropo, ig[1] + 1, jtemp + itemp - 1] / vmr_ref[tropo, ig[2] + 1, jtemp + itemp - 1]
@@ -143,8 +109,8 @@ Compute interpolation fraction for binary species parameter.
 end
 
 """
-    compute_τ_ssa_lw_src!(
-        lkp::AbstractLookUp,
+    compute_gas_optics(
+        lkp::Union{LookUpLW, LookUpSW},
         vmr,
         col_dry,
         igpt,
@@ -152,30 +118,38 @@ end
         p_lay::FT,
         t_lay,
         glay, gcol,
-        src_args...,
     ) where {FT<:AbstractFloat}
 
-Compute optical thickness, single scattering albedo, asymmetry parameter 
-and longwave sources whenever applicable.
+Compute optical thickness, single scattering albedo, and asymmetry parameter.
 """
-@inline function compute_τ_ssa_lw_src!(
-    lkp::AbstractLookUp,
+@inline function compute_gas_optics(
+    lkp::Union{LookUpLW, LookUpSW},
     vmr,
     col_dry,
     igpt,
     ibnd,
     p_lay::FT,
-    t_lay,
+    t_lay::FT,
     glay,
     gcol,
-    src_args...,
 ) where {FT <: AbstractFloat}
     # upper/lower troposphere
     tropo = p_lay > lkp.p_ref_tropo ? 1 : 2
     # volume mixing ratio of h2o
-    vmr_h2o = get_vmr(vmr, lkp.idx_h2o, glay, gcol)
+    vmr_h2o = get_vmr(vmr, lkp.idx_h2o, glay)
+
+    (; Δ_t_ref, n_t_ref, t_ref) = lkp
+    jftemp = compute_interp_frac_temp(Δ_t_ref, n_t_ref, t_ref, t_lay)
+
+    (; Δ_ln_p_ref, p_ref, n_p_ref) = lkp
+    jfpress = compute_interp_frac_press(Δ_ln_p_ref, p_ref, n_p_ref, p_lay, tropo)
+
+    (; n_η, vmr_ref) = lkp
+    ig = view(lkp.key_species, 1:2, tropo, ibnd)
+    @inbounds vmr1 = get_vmr(vmr, ig[1], glay)
+    @inbounds vmr2 = get_vmr(vmr, ig[2], glay)
 
-    jftemp, jfpress, jfη, col_mix = compute_interp_fractions(lkp, vmr, p_lay, t_lay, tropo, ibnd, glay, gcol)
+    jfη, col_mix = compute_interp_frac_η(n_η, ig, vmr_ref, (vmr1, vmr2), tropo, jftemp[1])
 
     # computing τ_major
     τ_major = interp3d(jfη..., jftemp..., jfpress..., lkp.kmajor, igpt, col_mix...) * col_dry
@@ -183,15 +157,17 @@ and longwave sources whenever applicable.
     τ_minor =
         compute_τ_minor(lkp, tropo, vmr, vmr_h2o, col_dry, p_lay, t_lay, jftemp..., jfη..., igpt, ibnd, glay, gcol)
     # compute τ_Rayleigh
-    τ_ray = compute_τ_rayleigh(lkp, tropo, col_dry, vmr_h2o, jftemp..., jfη..., igpt)
+    τ_ray = FT(0)
+    if lkp isa LookUpSW
+        rayleigh_coeff = tropo == 1 ? lkp.rayl_lower : lkp.rayl_upper
+        τ_ray = compute_τ_rayleigh(rayleigh_coeff, tropo, col_dry, vmr_h2o, jftemp..., jfη..., igpt)
+    end
     τ = τ_major + τ_minor + τ_ray
     ssa = FT(0)
     if τ > 2 * eps(FT) # single scattering albedo
         ssa = τ_ray / τ
     end
-    # computing Planck sources for longwave problem
-    compute_lw_planck_src!(lkp, jfη..., jfpress..., jftemp..., t_lay, igpt, ibnd, glay, gcol, src_args...)
-    return (τ, ssa)
+    return (τ, ssa, zero(τ)) # initializing asymmetry parameter
 end
 
 """
@@ -262,7 +238,7 @@ Compute optical thickness contributions from minor gases.
     @inbounds loc_in_bnd = igpt - (lkp.bnd_lims_gpt[1, ibnd] - 1)
 
     @inbounds for i in minor_bnd_st[ibnd]:(minor_bnd_st[ibnd + 1] - 1)
-        vmr_imnr = get_vmr(vmr, idx_gases_minor[i], glay, gcol)
+        vmr_imnr = get_vmr(vmr, idx_gases_minor[i], glay)
         if vmr_imnr > eps(FT) * 2
             scaling = vmr_imnr * col_dry
 
@@ -271,9 +247,9 @@ Compute optical thickness contributions from minor gases.
                 sgas = idx_scaling_gas[i]
                 if sgas > 0
                     if scale_by_complement[i] == 1
-                        scaling *= (FT(1) - get_vmr(vmr, sgas, glay, gcol) * dry_fact)
+                        scaling *= (FT(1) - get_vmr(vmr, sgas, glay) * dry_fact)
                     else
-                        scaling *= get_vmr(vmr, sgas, glay, gcol) * dry_fact
+                        scaling *= get_vmr(vmr, sgas, glay) * dry_fact
                     end
                 end
             end
@@ -302,8 +278,8 @@ end
 
 Compute Rayleigh scattering optical depths for shortwave problem
 """
-@inline function compute_τ_rayleigh(
-    lkp::LookUpSW,
+@inline compute_τ_rayleigh(
+    rayleigh_coeff::AbstractArray{FT, 3},
     tropo::Int,
     col_dry::FT,
     vmr_h2o::FT,
@@ -314,19 +290,8 @@ Compute Rayleigh scattering optical depths for shortwave problem
     fη1::FT,
     fη2::FT,
     igpt::Int,
-) where {FT <: AbstractFloat}
-    if tropo == 1
-        τ_ray = interp2d(fη1, fη2, ftemp, lkp.rayl_lower, igpt, jη1, jη2, jtemp) * (vmr_h2o + FT(1)) * col_dry
-    else
-        τ_ray = interp2d(fη1, fη2, ftemp, lkp.rayl_upper, igpt, jη1, jη2, jtemp) * (vmr_h2o + FT(1)) * col_dry
-    end
-
-    return τ_ray
-end
-
-@inline function compute_τ_rayleigh(lkp::LookUpLW{FT}, args...) where {FT <: AbstractFloat}
-    return FT(0)
-end
+) where {FT <: AbstractFloat} =
+    interp2d(fη1, fη2, ftemp, rayleigh_coeff, igpt, jη1, jη2, jtemp) * (vmr_h2o + FT(1)) * col_dry
 
 """
     compute_lw_planck_src!(
@@ -352,31 +317,31 @@ end
 Computes Planck sources for the longwave problem.
 
 """
-@inline function compute_lw_planck_src!(
-    lkp::LookUpLW,
-    jη1,
-    jη2,
-    fη1,
-    fη2,
-    jpresst,
-    fpress,
-    jtemp,
-    ftemp,
-    t_lay,
-    igpt,
-    ibnd,
-    glay,
-    gcol,
-    sf,
-    t_lev,
-    t_sfc,
-)
+@inline function compute_lw_planck_src!(lkp::LookUpLW, vmr, p_lay, t_lay, igpt, ibnd, glay, gcol, t_target)
+    # upper/lower troposphere
+    tropo = p_lay > lkp.p_ref_tropo ? 1 : 2
+
+    # compute intepolation fractions
+    (; Δ_t_ref, n_t_ref, t_ref) = lkp
+    jtemp, ftemp = compute_interp_frac_temp(Δ_t_ref, n_t_ref, t_ref, t_lay)
+
+    (; Δ_ln_p_ref, p_ref, n_p_ref) = lkp
+    jpresst, fpress = compute_interp_frac_press(Δ_ln_p_ref, p_ref, n_p_ref, p_lay, tropo)
+
+    (; n_η, vmr_ref) = lkp
+    ig = view(lkp.key_species, 1:2, tropo, ibnd)
+    @inbounds vmr1 = get_vmr(vmr, ig[1], glay)
+    @inbounds vmr2 = get_vmr(vmr, ig[2], glay)
+
+    (jη1, jη2, fη1, fη2), (col_mix1, col_mix2) = compute_interp_frac_η(n_η, ig, vmr_ref, (vmr1, vmr2), tropo, jtemp)
+
     (; planck_fraction, t_planck, n_t_plnk, totplnk) = lkp
-    (; lay_source, lev_source_inc, lev_source_dec, sfc_source) = sf
+
     # compute Planck fraction
     p_frac = interp3d(jη1, jη2, fη1, fη2, jtemp, ftemp, jpresst, fpress, planck_fraction, igpt)
 
     planck_args = (t_planck, totplnk, ibnd)
+    #=
     # computing lay_source
     @inbounds lay_source[glay, gcol] = interp1d(t_lay, planck_args...) * p_frac
     # computing lev_source_inc
@@ -387,8 +352,30 @@ Computes Planck sources for the longwave problem.
         @inbounds sfc_source[gcol] = interp1d(t_sfc, planck_args...) * p_frac
     end
     return nothing
+    =#
+    return interp1d(t_target, planck_args...) * p_frac # source
 end
 
-@inline function compute_lw_planck_src!(lkp::LookUpSW, args...)
-    return nothing
+@inline function compute_lw_planck_fraction(lkp::LookUpLW, vmr, p_lay, t_lay, igpt, ibnd, glay, gcol)
+    # upper/lower troposphere
+    tropo = p_lay > lkp.p_ref_tropo ? 1 : 2
+
+    # compute intepolation fractions
+    (; Δ_t_ref, n_t_ref, t_ref) = lkp
+    jtemp, ftemp = compute_interp_frac_temp(Δ_t_ref, n_t_ref, t_ref, t_lay)
+
+    (; Δ_ln_p_ref, p_ref, n_p_ref) = lkp
+    jpresst, fpress = compute_interp_frac_press(Δ_ln_p_ref, p_ref, n_p_ref, p_lay, tropo)
+
+    (; n_η, vmr_ref) = lkp
+    ig = view(lkp.key_species, 1:2, tropo, ibnd)
+    @inbounds vmr1 = get_vmr(vmr, ig[1], glay)
+    @inbounds vmr2 = get_vmr(vmr, ig[2], glay)
+
+    (jη1, jη2, fη1, fη2), (col_mix1, col_mix2) = compute_interp_frac_η(n_η, ig, vmr_ref, (vmr1, vmr2), tropo, jtemp)
+
+    (; planck_fraction, t_planck, n_t_plnk, totplnk) = lkp
+
+    # compute Planck fraction
+    return interp3d(jη1, jη2, fη1, fη2, jtemp, ftemp, jpresst, fpress, planck_fraction, igpt)
 end
diff --git a/src/optics/Optics.jl b/src/optics/Optics.jl
index 539301df2..7cafb931a 100644
--- a/src/optics/Optics.jl
+++ b/src/optics/Optics.jl
@@ -15,7 +15,8 @@ using ..AngularDiscretizations
 import ..Parameters as RP
 #---------------------------------------
 
-export AbstractOpticalProps, OneScalar, TwoStream, compute_col_gas!, compute_optical_props!
+export AbstractOpticalProps,
+    OneScalar, TwoStream, compute_col_gas!, compute_optical_props!, OneScalarc, TwoStreamc, compute_optical_props
 
 """
     AbstractOpticalProps{FT,FTA2D}
@@ -49,6 +50,14 @@ function OneScalar(::Type{FT}, ncol::Int, nlay::Int, ::Type{DA}) where {FT <: Ab
     return OneScalar{eltype(τ), typeof(τ), typeof(ad)}(τ, ad)
 end
 
+struct OneScalarc{FTA1D, AD}
+    τ::FTA1D
+    angle_disc::AD
+end
+Adapt.@adapt_structure OneScalarc
+
+OneScalarc(op::OneScalar, gcol::Int) = OneScalarc(view(op.τ, :, gcol), op.angle_disc)
+
 """
     TwoStream{FT,FTA2D} <: AbstractOpticalProps{FT,FTA2D}
 
@@ -75,6 +84,14 @@ function TwoStream(::Type{FT}, ncol::Int, nlay::Int, ::Type{DA}) where {FT <: Ab
     return TwoStream{eltype(τ), typeof(τ)}(τ, ssa, g)
 end
 
+struct TwoStreamc{FTA1D}
+    τ::FTA1D
+    ssa::FTA1D
+    g::FTA1D
+end
+Adapt.@adapt_structure TwoStreamc
+TwoStreamc(op::TwoStream, gcol::Int) = TwoStreamc(view(op.τ, :, gcol), view(op.ssa, :, gcol), view(op.g, :, gcol))
+
 """
     compute_col_gas!(
         context,
@@ -151,15 +168,93 @@ function compute_optical_props!(
     sf::AbstractSourceLW{FT},
     gcol::Int,
     igpt::Int,
-    lkp::Union{AbstractLookUp, Nothing} = nothing,
-    lkp_cld::Union{AbstractLookUp, Nothing} = nothing,
+    lkp::LookUpLW,
+    ::Nothing,
 ) where {FT <: AbstractFloat}
-    nlay = as.nlay
-    @inbounds for ilay in 1:nlay
-        if lkp_cld isa Nothing
-            compute_optical_props_kernel!(op, as, ilay, gcol, sf, igpt, lkp)
-        else
-            compute_optical_props_kernel!(op, as, ilay, gcol, sf, igpt, lkp, lkp_cld)
+    (; nlay, vmr) = as
+    @inbounds ibnd = lkp.major_gpt2bnd[igpt]
+    @inbounds t_sfc = as.t_sfc[gcol]
+    is2stream = op isa TwoStream
+    planck_args = (lkp.t_planck, lkp.totplnk, ibnd)
+    vmr_col = Vmrs.get_col_view(vmr, gcol)
+    @inbounds for glay in 1:nlay
+        col_dry = as.col_dry[glay, gcol]
+        p_lay = as.p_lay[glay, gcol]
+        t_lay = as.t_lay[glay, gcol]
+        # gas optics
+        τ, ssa, g = compute_gas_optics(lkp, vmr_col, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
+        op.τ[glay, gcol] = τ
+        if is2stream
+            op.ssa[glay, gcol] = ssa
+            op.g[glay, gcol] = g
+        end
+        # compute Planck sources
+        p_frac = compute_lw_planck_fraction(lkp, vmr_col, p_lay, t_lay, igpt, ibnd, glay, gcol)
+
+        sf.lay_source[glay, gcol] = interp1d(t_lay, planck_args...) * p_frac
+        sf.lev_source_inc[glay, gcol] = interp1d(as.t_lev[glay + 1, gcol], planck_args...) * p_frac
+        sf.lev_source_dec[glay, gcol] = interp1d(as.t_lev[glay, gcol], planck_args...) * p_frac
+
+        if glay == 1
+            sf.sfc_source[gcol] = interp1d(t_sfc, planck_args...) * p_frac
+        end
+    end
+    return nothing
+end
+
+function compute_optical_props!(
+    op::TwoStream,
+    as::AtmosphericState,
+    sf::AbstractSourceLW,
+    gcol::Int,
+    igpt::Int,
+    lkp::LookUpLW,
+    lkp_cld::Union{LookUpCld, PadeCld},
+)
+    (; nlay, vmr, ice_rgh) = as
+    @inbounds ibnd = lkp.major_gpt2bnd[igpt]
+    @inbounds t_sfc = as.t_sfc[gcol]
+    planck_args = (lkp.t_planck, lkp.totplnk, ibnd)
+    op_col = TwoStreamc(op, gcol)
+    col_dry_col = view(as.col_dry, :, gcol)
+    p_lay_col = view(as.p_lay, :, gcol)
+    t_lay_col = view(as.t_lay, :, gcol)
+    re_liq_col = view(as.cld_r_eff_liq, :, gcol)
+    re_ice_col = view(as.cld_r_eff_ice, :, gcol)
+    cld_path_liq_col = view(as.cld_path_liq, :, gcol)
+    cld_path_ice_col = view(as.cld_path_ice, :, gcol)
+    cld_mask_lw_col = view(as.cld_mask_lw, :, gcol)
+    (; τ, ssa, g) = op_col
+    vmr_col = Vmrs.get_col_view(as.vmr, gcol)
+    @inbounds for glay in 1:nlay
+        col_dry = col_dry_col[glay]
+        p_lay = p_lay_col[glay]
+        t_lay = t_lay_col[glay]
+        cld_mask = cld_mask_lw_col[glay]
+        # compute gas optics
+        τ[glay], ssa[glay], g[glay] = compute_gas_optics(lkp, vmr_col, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
+        # compute Planck sources
+        p_frac = compute_lw_planck_fraction(lkp, vmr_col, p_lay, t_lay, igpt, ibnd, glay, gcol)
+
+        sf.lay_source[glay, gcol] = interp1d(t_lay, planck_args...) * p_frac
+        sf.lev_source_inc[glay, gcol] = interp1d(as.t_lev[glay + 1, gcol], planck_args...) * p_frac
+        sf.lev_source_dec[glay, gcol] = interp1d(as.t_lev[glay, gcol], planck_args...) * p_frac
+
+        if glay == 1
+            sf.sfc_source[gcol] = interp1d(t_sfc, planck_args...) * p_frac
+        end
+    end
+    @inbounds for glay in 1:nlay
+        cld_mask = cld_mask_lw_col[glay]
+        if cld_mask
+            re_liq = re_liq_col[glay]
+            re_ice = re_ice_col[glay]
+            cld_path_liq = cld_path_liq_col[glay]
+            cld_path_ice = cld_path_ice_col[glay]
+            # add cloud optics
+            τ_cl, ssa_cl, g_cl = compute_cld_props(lkp_cld, re_liq, re_ice, ice_rgh, cld_path_liq, cld_path_ice, ibnd)
+            τ_cl, ssa_cl, g_cl = delta_scale(τ_cl, ssa_cl, g_cl)
+            τ[glay], ssa[glay], g[glay] = increment_2stream(τ[glay], ssa[glay], g[glay], τ_cl, ssa_cl, g_cl)
         end
     end
     return nothing
@@ -170,6 +265,7 @@ end
         op::AbstractOpticalProps{FT},
         as::AtmosphericState{FT},
         gcol::Int,
+        glay::Int,
         igpt::Int,
         lkp::LookUpSW{FT},
         lkp_cld::Union{LookUpCld,PadeCld,Nothing} = nothing,
@@ -177,23 +273,42 @@ end
 
 Computes optical properties for the shortwave problem.
 """
-function compute_optical_props!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
+@inline function compute_optical_props(as::AtmosphericStatec, gcol::Int, glay::Int, igpt::Int, lkp::LookUpSW, ::Nothing)
+    (; nlay, vmr) = as
+    @inbounds begin
+        ibnd = lkp.major_gpt2bnd[igpt]
+        col_dry = as.col_dry[glay]
+        p_lay = as.p_lay[glay]
+        t_lay = as.t_lay[glay]
+        τ, ssa, g = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
+    end
+    return τ, ssa, g
+end
+
+@inline function compute_optical_props(
+    as::AtmosphericStatec,
     gcol::Int,
+    glay::Int,
     igpt::Int,
-    lkp::Union{AbstractLookUp, Nothing} = nothing,
-    lkp_cld::Union{AbstractLookUp, Nothing} = nothing,
-) where {FT <: AbstractFloat}
-    nlay = as.nlay
-    @inbounds for ilay in 1:nlay
-        if lkp_cld isa Nothing
-            compute_optical_props_kernel!(op, as, ilay, gcol, igpt, lkp)
-        else
-            compute_optical_props_kernel!(op, as, ilay, gcol, igpt, lkp, lkp_cld)
+    lkp::LookUpSW,
+    lkp_cld::Union{LookUpCld, PadeCld},
+)
+    (; nlay, vmr, ice_rgh, p_lay, t_lay, col_dry, cld_mask_sw) = as
+    @inbounds begin
+        ibnd = lkp.major_gpt2bnd[igpt]
+        # compute gas optics
+        τ, ssa, g = compute_gas_optics(lkp, vmr, col_dry[glay], igpt, ibnd, p_lay[glay], t_lay[glay], glay, gcol)
+        if cld_mask_sw[glay]
+            re_liq = as.cld_r_eff_liq[glay]
+            re_ice = as.cld_r_eff_ice[glay]
+            # compute cloud optics
+            τ_cl, ssa_cl, g_cl =
+                compute_cld_props(lkp_cld, re_liq, re_ice, ice_rgh, as.cld_path_liq[glay], as.cld_path_ice[glay], ibnd)
+            τ_cl, ssa_cl, g_cl = delta_scale(τ_cl, ssa_cl, g_cl) # delta scale
+            τ, ssa, g = increment_2stream(τ, ssa, g, τ_cl, ssa_cl, g_cl) # add to gas optics
         end
     end
-    return nothing
+    return τ, ssa, g
 end
 #-----------------------------------------------------------------------------
 """
@@ -248,6 +363,9 @@ function compute_optical_props!(
     return nothing
 end
 
-include("OpticsKernels.jl")
+include("OpticsUtils.jl")
+include("GasOptics.jl")
+include("CloudOptics.jl")
+include("GrayOpticsKernels.jl")
 
 end
diff --git a/src/optics/OpticsKernels.jl b/src/optics/OpticsKernels.jl
deleted file mode 100644
index 7c1bd2bef..000000000
--- a/src/optics/OpticsKernels.jl
+++ /dev/null
@@ -1,104 +0,0 @@
-
-include("OpticsUtils.jl")
-include("GasOptics.jl")
-include("CloudOptics.jl")
-include("GrayOpticsKernels.jl")
-
-function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
-    glay,
-    gcol,
-    sf::AbstractSourceLW{FT},
-    igpt,
-    lkp::LookUpLW{FT},
-    lkp_cld::Union{LookUpCld, PadeCld},
-) where {FT <: AbstractFloat}
-    ibnd = lkp.major_gpt2bnd[igpt]
-
-    compute_optical_props_kernel!(op, as, glay, gcol, sf, igpt, lkp) # longwave gas optics
-    add_cloud_optics_2stream(op, as, lkp, lkp_cld, glay, gcol, ibnd, igpt) # add cloud optics
-    return nothing
-end
-
-function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
-    glay,
-    gcol,
-    sf::AbstractSourceLW{FT},
-    igpt,
-    lkp::LookUpLW{FT},
-) where {FT <: AbstractFloat}
-    t_sfc = as.t_sfc[gcol]
-    t_lev = as.t_lev
-    compute_optical_props_kernel!(op, as, glay, gcol, igpt, lkp, sf, t_lev, t_sfc)
-    return nothing
-end
-
-function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
-    glay,
-    gcol,
-    igpt,
-    lkp::LookUpSW{FT},
-    lkp_cld::Union{LookUpCld, PadeCld},
-) where {FT <: AbstractFloat}
-    ibnd = lkp.major_gpt2bnd[igpt]
-
-    compute_optical_props_kernel!(op, as, glay, gcol, igpt, lkp) # shortwave gas optics
-    add_cloud_optics_2stream(op, as, lkp, lkp_cld, glay, gcol, ibnd, igpt) # add cloud optics
-    return nothing
-end
-
-function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
-    glay,
-    gcol,
-    igpt::Int,
-    lkp::LookUpLW{FT},
-    src_args...,
-) where {FT <: AbstractFloat}
-
-    vmr = as.vmr
-    col_dry = as.col_dry[glay, gcol]
-    p_lay = as.p_lay[glay, gcol]
-    t_lay = as.t_lay[glay, gcol]
-    ibnd = lkp.major_gpt2bnd[igpt]
-    τ, ssa = compute_τ_ssa_lw_src!(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol, src_args...)
-
-    @inbounds op.τ[glay, gcol] = τ
-
-    if op isa TwoStream
-        @inbounds op.ssa[glay, gcol] = ssa
-        @inbounds op.g[glay, gcol] = FT(0) # initializing asymmetry parameter
-    end
-    return nothing
-end
-
-function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
-    glay,
-    gcol,
-    igpt::Int,
-    lkp::LookUpSW{FT},
-) where {FT <: AbstractFloat}
-
-    vmr = as.vmr
-    col_dry = as.col_dry[glay, gcol]
-    p_lay = as.p_lay[glay, gcol]
-    t_lay = as.t_lay[glay, gcol]
-    ibnd = lkp.major_gpt2bnd[igpt]
-    τ, ssa = compute_τ_ssa_lw_src!(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
-
-    @inbounds op.τ[glay, gcol] = τ
-
-    if op isa TwoStream
-        @inbounds op.ssa[glay, gcol] = ssa
-        @inbounds op.g[glay, gcol] = FT(0) # initializing asymmetry parameter
-    end
-    return nothing
-end
diff --git a/src/optics/OpticsUtils.jl b/src/optics/OpticsUtils.jl
index ad3ddffb3..87e402340 100644
--- a/src/optics/OpticsUtils.jl
+++ b/src/optics/OpticsUtils.jl
@@ -128,7 +128,8 @@ Increment TwoStream optical properties `τ1`, `ssa1` and `g1`
 with `τ2`, `ssa2` and `g2`. Here `τ` is the optical thickness,
 `ssa` is the single-scattering albedo, and `g` is the symmetry parameter.
 """
-function increment_2stream(τ1::FT, ssa1::FT, g1::FT, τ2::FT, ssa2::FT, g2::FT) where {FT}
+@inline function increment_2stream(τ1::FT, ssa1::FT, g1::FT, τ2::FT, ssa2::FT, g2::FT) where {FT}
+    #FT = eltype(τ1)
     τ = τ1 + τ2
     ssa = τ1 * ssa1 + τ2 * ssa2
     ssag = (τ1 * ssa1 * g1 + τ2 * ssa2 * g2) / max(eps(FT), ssa)
@@ -140,8 +141,8 @@ end
 
 delta-scale two stream optical properties.
 """
-function delta_scale(τ, ssa, g)
-    FT = typeof(τ)
+@inline function delta_scale(τ::FT, ssa::FT, g::FT) where {FT}
+    #FT = typeof(τ)
     f = g * g
     wf = ssa * f
     τ_s = (FT(1) - wf) * τ
diff --git a/src/optics/Sources.jl b/src/optics/Sources.jl
index d23ec5f15..e6108bc19 100644
--- a/src/optics/Sources.jl
+++ b/src/optics/Sources.jl
@@ -5,7 +5,8 @@ using DocStringExtensions
 
 import ..Parameters as RP
 
-export AbstractSourceLW, SourceLWNoScat, SourceLW2Str, SourceSW2Str, source_func_longwave, source_func_shortwave
+export AbstractSourceLW,
+    SourceLWNoScat, SourceLW2Str, SourceSW2Str, source_func_longwave, source_func_shortwave, SourceSW2Strc
 
 """
     AbstractSourceLW{FT,FTA1D,FTA2D}
@@ -235,6 +236,34 @@ struct SourceSW2Str{FT <: AbstractFloat, FTA1D <: AbstractArray{FT, 1}, FTA2D <:
 end
 Adapt.@adapt_structure SourceSW2Str
 
+struct SourceSW2Strc{FTVA1D, FTA1D}
+    sfc_source::FTVA1D
+    albedo::FTA1D
+    src_up::FTA1D
+    src_dn::FTA1D
+    Rdif::FTA1D
+    Tdif::FTA1D
+    Rdir::FTA1D
+    Tdir::FTA1D
+    Tnoscat::FTA1D
+    src::FTA1D
+end
+Adapt.@adapt_structure SourceSW2Strc
+
+SourceSW2Strc(source::SourceSW2Str, gcol::Int) = SourceSW2Strc(
+    view(source.sfc_source, gcol:gcol),
+    view(source.albedo, :, gcol),
+    view(source.src_up, :, gcol),
+    view(source.src_dn, :, gcol),
+    view(source.Rdif, :, gcol),
+    view(source.Tdif, :, gcol),
+    view(source.Rdir, :, gcol),
+    view(source.Tdir, :, gcol),
+    view(source.Tnoscat, :, gcol),
+    view(source.src, :, gcol),
+)
+
+
 function SourceSW2Str(::Type{FT}, ::Type{DA}, nlay::Int, ncol::Int) where {FT <: AbstractFloat, DA}
     FTA1D = DA{FT, 1}
     FTA2D = DA{FT, 2}
diff --git a/src/optics/Vmrs.jl b/src/optics/Vmrs.jl
index d194cb0a3..916393cbd 100644
--- a/src/optics/Vmrs.jl
+++ b/src/optics/Vmrs.jl
@@ -3,14 +3,16 @@ using CUDA
 using DocStringExtensions
 using Adapt
 
-export AbstractVmr, Vmr, VmrGM, init_vmr, get_vmr
+export AbstractVmr, Vmr, VmrGM, init_vmr, get_vmr, Vmrc, VmrGMc
 
 """
-    AbstractVmr{FT}
+    AbstractVmr
 """
-abstract type AbstractVmr{FT <: AbstractFloat} end
+abstract type AbstractVmr end
+
+abstract type AbstractVmrc <: AbstractVmr end
 """
-    VmrGM{FT,FTA1D,FTA2D} <: AbstractVmr{FT}
+    VmrGM{FT,FTA1D,FTA2D} <: AbstractVmr
 
 Volume mixing ratios for various gases in the atmosphere. This struct can be used
 when only H₂O and O₃ concentrations vary spatially and a global mean is used 
@@ -19,7 +21,7 @@ for all other gases.
 # Fields
 $(DocStringExtensions.FIELDS)
 """
-struct VmrGM{FT <: AbstractFloat, FTA1D <: AbstractArray{FT, 1}, FTA2D <: AbstractArray{FT, 2}} <: AbstractVmr{FT}
+struct VmrGM{FT <: AbstractFloat, FTA1D <: AbstractArray{FT, 1}, FTA2D <: AbstractArray{FT, 2}} <: AbstractVmr
     "volume mixing ratio of H₂O"
     vmr_h2o::FTA2D
     "volume mixing ratio of Ozone"
@@ -30,6 +32,16 @@ end
 VmrGM(vmr_h2o, vmr_o3, vmr) = VmrGM{eltype(vmr_h2o), typeof(vmr), typeof(vmr_h2o)}(vmr_h2o, vmr_o3, vmr)
 Adapt.@adapt_structure VmrGM
 
+struct VmrGMc{FTA1DV, FTA1D} <: AbstractVmrc
+    vmr_h2o::FTA1DV
+    vmr_o3::FTA1DV
+    vmr::FTA1D
+end
+Adapt.@adapt_structure VmrGMc
+
+VmrGMc(vmr::VmrGM, gcol) = VmrGMc(view(vmr.vmr_h2o, :, gcol), view(vmr.vmr_o3, :, gcol), vmr.vmr)
+get_col_view(vmr::VmrGM, gcol) = VmrGMc(vmr, gcol)
+
 function VmrGM(
     ::Type{FT},
     ::Type{DA},
@@ -41,7 +53,7 @@ function VmrGM(
 end
 
 """
-    Vmr{FT,FTA3D} <: AbstractVmr{FT}
+    Vmr{FT,FTA3D} <: AbstractVmr
 
 Volume mixing ratios for various gases in the atmosphere. This struct can be used
 when concentrations vary spatially for all gases.
@@ -49,12 +61,21 @@ when concentrations vary spatially for all gases.
 # Fields
 $(DocStringExtensions.FIELDS)
 """
-struct Vmr{FT <: AbstractFloat, FTA3D <: AbstractArray{FT, 3}} <: AbstractVmr{FT}
+struct Vmr{FT <: AbstractFloat, FTA3D <: AbstractArray{FT, 3}} <: AbstractVmr
     "volume mixing ratio of all gases as a function of location and column"
     vmr::FTA3D
 end
 Vmr(vmr) = Vmr{eltype(vmr), typeof(vmr)}(vmr)
 Adapt.@adapt_structure Vmr
+
+struct Vmrc{FTA2D} <: AbstractVmrc
+    vmr::FTA2D
+end
+Adapt.@adapt_structure Vmrc
+
+Vmrc(vmr::Vmr, gcol) = Vmrc(view(vmr.vmr, :, :, gcol))
+get_col_view(vmr::Vmr, gcol) = Vmrc(vmr, gcol)
+
 """
     get_vmr(
         vmr::VmrGM{FT},
@@ -77,6 +98,19 @@ Obtain volume mixing ratio of gas `ig` for layer `ilay` of column `icol`.
     end
 end
 
+@inline function get_vmr(vmr::VmrGMc, ig::Int, ilay::Int)
+    if ig == 0
+        FT = eltype(vmr.vmr_h2o)
+        return FT(1)
+    elseif ig == 1 # h2o / h2o_foreign / h2o_self-continua
+        return @inbounds vmr.vmr_h2o[ilay]
+    elseif ig == 3 # ozone
+        return @inbounds vmr.vmr_o3[ilay]
+    else # other gases
+        return @inbounds vmr.vmr[ig]
+    end
+end
+
 """
     get_vmr(
         vmr::Vmr{FT},
@@ -87,11 +121,20 @@ end
 
 Obtain volume mixing ratio of gas `ig` for layer `ilay` of column `icol`.
 """
-@inline function get_vmr(vmr::Vmr{FT}, ig::Int, ilay::Int, icol::Int) where {FT <: AbstractFloat}
+@inline function get_vmr(vmr::Vmr{FT}, ig::Int, ilay::Int, icol::Int) where {FT}
+    if ig == 0
+        return FT(1)
+    else
+        return @inbounds vmr.vmr[ig, ilay, icol]
+    end
+end
+
+@inline function get_vmr(vmr::Vmrc, ig::Int, ilay::Int)
     if ig == 0
+        FT = eltype(vmr.vmr)
         return FT(1)
     else
-        return @inbounds vmr.vmr[ilay, icol, ig]
+        return @inbounds vmr.vmr[ig, ilay]
     end
 end
 
@@ -121,7 +164,7 @@ function init_vmr(
         return VmrGM{FT, FTA1D, FTA2D}(FTA2D(zeros(nlay, ncol)), FTA2D(zeros(nlay, ncol)), FTA1D(zeros(ngas)))
     else
         FTA3D = DA{FT, 3}
-        return Vmr{FT, FTA3D}(FTA3D(zeros(nlay, ncol, ngas)))
+        return Vmr{FT, FTA3D}(FTA3D(zeros(ngas, nlay, ncol)))
     end
 end
 
diff --git a/src/rte/RTESolver.jl b/src/rte/RTESolver.jl
index 0f404a07a..e7fc640dc 100644
--- a/src/rte/RTESolver.jl
+++ b/src/rte/RTESolver.jl
@@ -54,6 +54,60 @@ function solve_lw!(
     return nothing
 end
 
+"""
+    solve_sw!(
+        slv::Solver,
+        max_threads::Int,
+        lookup_lw::Union{LookUpLW, Nothing} = nothing,
+        lookup_lw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+    )
+
+Solver for the shortwave radiation problem
+"""
+solve_sw!(
+    slv::Solver,
+    max_threads::Int,
+    lookup_sw::Union{LookUpSW, Nothing} = nothing,
+    lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+) = solve_sw!(slv.context.device, slv, max_threads, lookup_sw, lookup_sw_cld)
+
+# non-scattering solver, gray optics
+solve_sw!(
+    device::ClimaComms.AbstractDevice,
+    (; flux_sw, op, bcs_sw, as)::Solver{<:Any, <:GrayAtmosphericState, <:OneScalar},
+    max_threads::Int,
+    ::Nothing,
+    ::Nothing,
+) = rte_sw_noscat_solve!(device, flux_sw, op, bcs_sw, max_threads, as) # non-scattering solver, gray optics
+
+# 2-stream solver, gray optics
+solve_sw!(
+    device::ClimaComms.AbstractDevice,
+    (; flux_sw, op, bcs_sw, src_sw, as)::Solver{<:Any, <:GrayAtmosphericState, <:TwoStream},
+    max_threads::Int,
+    ::Nothing,
+    ::Nothing,
+) = rte_sw_2stream_solve!(device, flux_sw, op, bcs_sw, src_sw, max_threads, as)
+
+# non-scattering solver, RRTMGP optics
+solve_sw!(
+    device::ClimaComms.AbstractDevice,
+    (; fluxb_sw, flux_sw, op, bcs_sw, as)::Solver{<:Any, <:AtmosphericState, <:OneScalar},
+    max_threads::Int,
+    lookup_sw::LookUpSW,
+    lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+) = rte_sw_noscat_solve!(device, fluxb_sw, flux_sw, op, bcs_sw, max_threads, as, slv.lookup_sw, slv.lookup_sw_cld)
+
+# 2-stream solver, RRTMGP optics
+solve_sw!(
+    device::ClimaComms.AbstractDevice,
+    (; fluxb_sw, flux_sw, op, bcs_sw, src_sw, as)::Solver{<:Any, <:AtmosphericState, <:TwoStream},
+    max_threads::Int,
+    lookup_sw::LookUpSW,
+    lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+) = rte_sw_2stream_solve!(device, fluxb_sw, flux_sw, op, bcs_sw, src_sw, max_threads, as, lookup_sw, lookup_sw_cld)
+#--------------------------------------------------------------------------------------------------
+#=
 """
     solve_sw!(
         slv::Solver,
@@ -72,50 +126,24 @@ function solve_sw!(
 )
     (; as, op, bcs_sw, src_sw, flux_sw, fluxb_sw) = slv
     context = RTE.context(slv)
+    device = context.device
 
     no_args = lookup_sw isa Nothing && lookup_sw_cld isa Nothing
     if no_args
         DA = RTE.array_type(slv)
         FT = RTE.float_type(slv)
-        major_gpt2bnd = DA(UInt8[1])
-        solar_src_scaled = DA(FT[1])
         flux = flux_sw
     else
-        (; major_gpt2bnd, solar_src_scaled) = lookup_sw
         flux = fluxb_sw
     end
 
     # solving radiative transfer equation
     if slv.op isa OneScalar
-        rte_sw_noscat_solve!(
-            context,
-            flux,
-            flux_sw,
-            op,
-            bcs_sw,
-            solar_src_scaled,
-            max_threads,
-            as,
-            lookup_sw,
-            lookup_sw_cld,
-        ) # no-scattering solver
+        rte_sw_noscat_solve!(device, flux, flux_sw, op, bcs_sw, max_threads, as, lookup_sw, lookup_sw_cld) # no-scattering solver
     else
-        rte_sw_2stream_solve!(
-            context,
-            flux,
-            flux_sw,
-            op,
-            bcs_sw,
-            src_sw,
-            major_gpt2bnd,
-            solar_src_scaled,
-            max_threads,
-            as,
-            lookup_sw,
-            lookup_sw_cld,
-        ) # 2-stream solver
+        rte_sw_2stream_solve!(device, flux, flux_sw, op, bcs_sw, src_sw, max_threads, as, lookup_sw, lookup_sw_cld) # 2-stream solver
     end
     return nothing
 end
-
+=#
 end
diff --git a/src/rte/longwave2stream.jl b/src/rte/longwave2stream.jl
index c29def4ea..016cf959f 100644
--- a/src/rte/longwave2stream.jl
+++ b/src/rte/longwave2stream.jl
@@ -52,11 +52,8 @@ function rte_lw_solve!(
                 ClimaComms.@threaded device for gcol in 1:ncol
                     igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
                     bld_cld_mask && Optics.build_cloud_mask!(
-                        as.cld_mask_lw,
-                        as.cld_frac,
-                        as.random_lw,
-                        gcol,
-                        igpt,
+                        view(as.cld_mask_lw, :, gcol),
+                        view(as.cld_frac, :, gcol),
                         as.cld_mask_type,
                     )
                     compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
@@ -91,7 +88,7 @@ function rte_lw_2stream_solve_CUDA!(
         @inbounds for igpt in 1:n_gpt
             igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
             if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
-                Optics.build_cloud_mask!(as.cld_mask_lw, as.cld_frac, as.random_lw, gcol, igpt, as.cld_mask_type)
+                Optics.build_cloud_mask!(view(as.cld_mask_lw, :, gcol), view(as.cld_frac, :, gcol), as.cld_mask_type)
             end
             compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
             rte_lw_2stream_combine_sources!(src_lw, gcol, nlev, ncol)
diff --git a/src/rte/shortwave1scalar.jl b/src/rte/shortwave1scalar.jl
index 742ebc0bf..eb283a12d 100644
--- a/src/rte/shortwave1scalar.jl
+++ b/src/rte/shortwave1scalar.jl
@@ -1,57 +1,120 @@
-"""
-    rte_sw_noscat_solve!(
-        context,
-        flux::FluxSW{FT},
-        flux_sw::FluxSW{FT},
-        op::OneScalar{FT},
-        bcs_sw::SwBCs{FT},
-        solar_src_scaled::AbstractArray{FT,1},
-        max_threads,
-        as::AbstractAtmosphericState{FT},
-        lookup_sw::Union{LookUpSW, Nothing} = nothing,
-        lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
-    ) where {FT<:AbstractFloat}
+function rte_sw_noscat_solve!(
+    device::ClimaComms.AbstractCPUDevice,
+    flux_sw::FluxSW{FT},
+    op::OneScalar{FT},
+    bcs_sw::SwBCs{FT},
+    max_threads,
+    as::GrayAtmosphericState{FT},
+) where {FT <: AbstractFloat}
+    (; nlay, ncol) = as
+    nlev = nlay + 1
+    n_gpt, igpt = 1, 1
+    solar_frac = FT(1)
+    # setting references for flux_sw
+    @inbounds begin
+        ClimaComms.@threaded device for gcol in 1:ncol
+            set_flux_to_zero!(flux_sw, gcol)
+            compute_optical_props!(op, as, gcol, igpt, nothing, nothing)
+            rte_sw_noscat_solve_kernel!(flux_sw, op, bcs_sw, igpt, n_gpt, solar_frac, gcol, nlev)
+        end
+    end
+    return nothing
+end
 
-No-scattering solver for the shortwave problem.
-(Extinction-only i.e. solar direct beam)
-"""
 function rte_sw_noscat_solve!(
-    context,
+    device::ClimaComms.CUDADevice,
+    flux_sw::FluxSW{FT},
+    op::OneScalar{FT},
+    bcs_sw::SwBCs{FT},
+    max_threads,
+    as::GrayAtmosphericState{FT},
+) where {FT <: AbstractFloat}
+    (; nlay, ncol) = as
+    nlev = nlay + 1
+    # setting references for flux_sw
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    args = (flux_sw, op, bcs_sw, nlay, ncol, as)
+    @cuda threads = (tx) blocks = (bx) rte_sw_noscat_solve_CUDA!(args...)
+    return nothing
+end
+
+function rte_sw_noscat_solve_CUDA!(
+    flux_sw::FluxSW{FT},
+    op::OneScalar{FT},
+    bcs_sw::SwBCs{FT},
+    nlay,
+    ncol,
+    as::GrayAtmosphericState{FT},
+) where {FT <: AbstractFloat}
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    nlev = nlay + 1
+    n_gpt, igpt = 1, 1
+    solar_frac = FT(1)
+    # setting references for flux_sw
+    if gcol ≤ ncol
+        @inbounds begin
+            set_flux_to_zero!(flux_sw, gcol)
+            compute_optical_props!(op, as, gcol, igpt, nothing, nothing)
+            rte_sw_noscat_solve_kernel!(flux_sw, op, bcs_sw, igpt, n_gpt, solar_frac, gcol, nlev)
+        end
+    end
+    return nothing
+end
+
+function rte_sw_noscat_solve!(
+    device::ClimaComms.AbstractCPUDevice,
     flux::FluxSW{FT},
     flux_sw::FluxSW{FT},
     op::OneScalar{FT},
     bcs_sw::SwBCs{FT},
-    solar_src_scaled::AbstractArray{FT, 1},
     max_threads,
-    as::AbstractAtmosphericState{FT},
-    lookup_sw::Union{LookUpSW, Nothing} = nothing,
+    as::AtmosphericState{FT},
+    lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
 ) where {FT <: AbstractFloat}
     (; nlay, ncol) = as
+    τ = op.τ
     nlev = nlay + 1
-    n_gpt = length(solar_src_scaled)
-    device = ClimaComms.device(context)
+    n_gpt = length(lookup_sw.solar_src_scaled)
     # setting references for flux_sw
-    if device isa ClimaComms.CUDADevice
-        tx = min(ncol, max_threads)
-        bx = cld(ncol, tx)
-        args = (flux, flux_sw, op, bcs_sw, nlay, ncol, solar_src_scaled, as, lookup_sw, lookup_sw_cld)
-        @cuda threads = (tx) blocks = (bx) rte_sw_noscat_solve_CUDA!(args...)
-    else
-        @inbounds begin
-            for igpt in 1:n_gpt
-                ClimaComms.@threaded device for gcol in 1:ncol
-                    igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
-                    compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
-                    rte_sw_noscat_solve_kernel!(flux, op, bcs_sw, igpt, solar_src_scaled, gcol, nlev)
-                    n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
+    @inbounds begin
+        for igpt in 1:n_gpt
+            ClimaComms.@threaded device for gcol in 1:ncol
+                igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
+                for glay in 1:nlay
+                    τ[glay, gcol], _, _ = compute_optical_props(as, gcol, glay, igpt, lookup_sw, lookup_sw_cld)
                 end
+                solar_frac = as isa GrayAtmosphericState ? FT(1) : lookup_sw.solar_src_scaled[igpt]
+                rte_sw_noscat_solve_kernel!(flux, op, bcs_sw, igpt, n_gpt, solar_frac, gcol, nlev)
+                n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
             end
         end
     end
     return nothing
 end
 
+function rte_sw_noscat_solve!(
+    device::ClimaComms.CUDADevice,
+    flux::FluxSW{FT},
+    flux_sw::FluxSW{FT},
+    op::OneScalar{FT},
+    bcs_sw::SwBCs{FT},
+    max_threads,
+    as::AtmosphericState{FT},
+    lookup_sw::LookUpSW,
+    lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+) where {FT <: AbstractFloat}
+    (; nlay, ncol) = as
+    nlev = nlay + 1
+    # setting references for flux_sw
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    args = (flux, flux_sw, op, bcs_sw, nlay, ncol, as, lookup_sw, lookup_sw_cld)
+    @cuda threads = (tx) blocks = (bx) rte_sw_noscat_solve_CUDA!(args...)
+    return nothing
+end
+
 function rte_sw_noscat_solve_CUDA!(
     flux::FluxSW{FT},
     flux_sw::FluxSW{FT},
@@ -59,20 +122,23 @@ function rte_sw_noscat_solve_CUDA!(
     bcs_sw::SwBCs{FT},
     nlay,
     ncol,
-    solar_src_scaled::AbstractArray{FT, 1},
-    as::AbstractAtmosphericState{FT},
-    lookup_sw::Union{LookUpSW, Nothing} = nothing,
+    as::AtmosphericState{FT},
+    lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
 ) where {FT <: AbstractFloat}
     gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
     nlev = nlay + 1
-    n_gpt = length(solar_src_scaled)
+    n_gpt = length(lookup_sw.solar_src_scaled)
     # setting references for flux_sw
     if gcol ≤ ncol
+        τ = op.τ
+        set_flux_to_zero!(flux_sw, gcol)
         @inbounds for igpt in 1:n_gpt
-            igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
-            compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
-            rte_sw_noscat_solve_kernel!(flux, op, bcs_sw, igpt, solar_src_scaled, gcol, nlev)
+            for glay in 1:nlay
+                τ[glay, gcol] = compute_optical_props(as, gcol, glay, igpt, lookup_sw, lookup_sw_cld)
+            end
+            solar_frac = as isa GrayAtmosphericState ? FT(1) : lookup_sw.solar_src_scaled[igpt]
+            rte_sw_noscat_solve_kernel!(flux, op, bcs_sw, igpt, n_gpt, solar_frac, gcol, nlev)
             n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
         end
     end
@@ -97,19 +163,20 @@ function rte_sw_noscat_solve_kernel!(
     op::OneScalar{FT},
     bcs_sw::SwBCs{FT},
     igpt::Int,
-    solar_src_scaled::AbstractArray{FT, 1},
+    n_gpt::Int,
+    solar_frac::FT,
     gcol::Int,
     nlev::Int,
 ) where {FT <: AbstractFloat}
-    solar_frac = solar_src_scaled[igpt]
-    (; toa_flux, zenith) = bcs_sw
-    n_gpt = length(solar_src_scaled)
+    #solar_frac = solar_src_scaled[igpt]
+    (; toa_flux, cos_zenith) = bcs_sw
+    #n_gpt = length(solar_src_scaled)
     τ = op.τ
     (; flux_dn_dir, flux_net) = flux
     # downward propagation
-    @inbounds flux_dn_dir[nlev, gcol] = toa_flux[gcol] * solar_frac * cos(zenith[gcol])
+    @inbounds flux_dn_dir[nlev, gcol] = toa_flux[gcol] * solar_frac * cos_zenith[gcol]
     @inbounds for ilev in (nlev - 1):-1:1
-        flux_dn_dir[ilev, gcol] = flux_dn_dir[ilev + 1, gcol] * exp(-τ[ilev, gcol] / cos(zenith[gcol]))
+        flux_dn_dir[ilev, gcol] = flux_dn_dir[ilev + 1, gcol] * exp(-τ[ilev, gcol] / cos_zenith[gcol])
         flux_net[ilev, gcol] = -flux_dn_dir[ilev, gcol]
     end
 end
diff --git a/src/rte/shortwave2stream.jl b/src/rte/shortwave2stream.jl
index 101d5e3c1..cd06e2cdc 100644
--- a/src/rte/shortwave2stream.jl
+++ b/src/rte/shortwave2stream.jl
@@ -1,84 +1,151 @@
-"""
-    rte_sw_2stream_solve!(
-        context,
-        flux::FluxSW{FT},
-        flux_sw::FluxSW{FT},
-        op::TwoStream{FT},
-        bcs_sw::SwBCs{FT},
-        src_sw::SourceSW2Str{FT},
-        major_gpt2bnd::AbstractArray{UInt8,1},
-        solar_src_scaled::AbstractArray{FT,1},
-        max_threads,
-        as::AbstractAtmosphericState{FT},
-        lookup_sw::Union{LookUpSW, Nothing} = nothing,
-        lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
-    ) where {FT<:AbstractFloat}
+function rte_sw_2stream_solve!(
+    device::ClimaComms.AbstractCPUDevice,
+    flux_sw::FluxSW{FT},
+    op::TwoStream{FT},
+    bcs_sw::SwBCs{FT},
+    src_sw::SourceSW2Str{FT},
+    max_threads,
+    as::GrayAtmosphericState{FT},
+) where {FT <: AbstractFloat}
+    (; nlay, ncol) = as
+    nlev = nlay + 1
+    n_gpt, igpt, ibnd = 1, 1, UInt8(1)
+    solar_frac = FT(1)
+    @inbounds begin
+        ClimaComms.@threaded device for gcol in 1:ncol
+            # get column views
+            flux_sw_col = FluxSWc(flux_sw, gcol)
+            op_col = TwoStreamc(op, gcol)
+            src_sw_col = SourceSW2Strc(src_sw, gcol)
+            bcs_sw_col = SwBCsc(bcs_sw, gcol)
+            set_flux_to_zero!(flux_sw_col)
+            compute_optical_props!(op, as, gcol, igpt, nothing, nothing)
+            # call shortwave rte solver
+            rte_sw!(op_col, src_sw_col, bcs_sw_col, flux_sw_col, solar_frac, igpt, n_gpt, ibnd, nlev)
+        end
+    end
+    return nothing
+end
+
+function rte_sw_2stream_solve!(
+    device::ClimaComms.CUDADevice,
+    flux_sw::FluxSW{FT},
+    op::TwoStream{FT},
+    bcs_sw::SwBCs{FT},
+    src_sw::SourceSW2Str{FT},
+    max_threads,
+    as::GrayAtmosphericState{FT},
+) where {FT <: AbstractFloat}
+    (; nlay, ncol) = as
+    nlev = nlay + 1
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    args = (flux_sw, op, bcs_sw, src_sw, nlay, ncol, as)
+    @cuda threads = (tx) blocks = (bx) rte_sw_2stream_solve_CUDA!(args...)
+    return nothing
+end
+
+function rte_sw_2stream_solve_CUDA!(
+    flux_sw::FluxSW{FT},
+    op::TwoStream{FT},
+    bcs_sw::SwBCs{FT},
+    src_sw::SourceSW2Str{FT},
+    nlay,
+    ncol,
+    as::GrayAtmosphericState{FT},
+) where {FT <: AbstractFloat}
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    nlev = nlay + 1
+    n_gpt, igpt, ibnd = 1, 1, UInt8(1)
+    solar_frac = FT(1)
+    if gcol ≤ ncol
+        @inbounds begin
+            # get column views
+            flux_sw_col = FluxSWc(flux_sw, gcol)
+            op_col = TwoStreamc(op, gcol)
+            src_sw_col = SourceSW2Strc(src_sw, gcol)
+            bcs_sw_col = SwBCsc(bcs_sw, gcol)
+            set_flux_to_zero!(flux_sw_col)
+            compute_optical_props!(op, as, gcol, igpt, nothing, nothing)
+            # call shortwave rte solver
+            rte_sw!(op_col, src_sw_col, bcs_sw_col, flux_sw_col, solar_frac, igpt, n_gpt, ibnd, nlev)
+        end
+    end
+    return nothing
+end
 
-Two stream solver for the shortwave problem.
-"""
 function rte_sw_2stream_solve!(
-    context,
+    device::ClimaComms.AbstractCPUDevice,
     flux::FluxSW{FT},
     flux_sw::FluxSW{FT},
     op::TwoStream{FT},
     bcs_sw::SwBCs{FT},
     src_sw::SourceSW2Str{FT},
-    major_gpt2bnd::AbstractArray{UInt8, 1},
-    solar_src_scaled::AbstractArray{FT, 1},
     max_threads,
-    as::AbstractAtmosphericState{FT},
-    lookup_sw::Union{LookUpSW, Nothing} = nothing,
+    as::AtmosphericState{FT},
+    lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
 ) where {FT <: AbstractFloat}
     (; nlay, ncol) = as
     nlev = nlay + 1
-    n_gpt = length(major_gpt2bnd)
-    device = ClimaComms.device(context)
-    if device isa ClimaComms.CUDADevice
-        tx = min(ncol, max_threads)
-        bx = cld(ncol, tx)
-        args = (
-            flux,
-            flux_sw,
-            op,
-            bcs_sw,
-            src_sw,
-            nlay,
-            ncol,
-            major_gpt2bnd,
-            solar_src_scaled,
-            as,
-            lookup_sw,
-            lookup_sw_cld,
-        )
-        @cuda threads = (tx) blocks = (bx) rte_sw_2stream_solve_CUDA!(args...)
-    else
-        @inbounds begin
-            bld_cld_mask = as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
-            # setting references for flux_sw
-            for igpt in 1:n_gpt
-                ClimaComms.@threaded device for gcol in 1:ncol
-                    igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
-                    bld_cld_mask && Optics.build_cloud_mask!(
-                        as.cld_mask_sw,
-                        as.cld_frac,
-                        as.random_sw,
-                        gcol,
-                        igpt,
-                        as.cld_mask_type,
-                    )
-                    compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
-                    sw_two_stream!(op, src_sw, bcs_sw, gcol, nlay) # Cell properties: transmittance and reflectance for direct and diffuse radiation
-                    sw_source_2str!(src_sw, bcs_sw, gcol, flux, igpt, solar_src_scaled, major_gpt2bnd, nlay) # Direct-beam and source for diffuse radiation
-                    adding_sw!(src_sw, bcs_sw, gcol, flux, igpt, major_gpt2bnd, nlev) # Transport
-                    n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
+    n_gpt = length(lookup_sw.solar_src_scaled)
+    @inbounds begin
+        bld_cld_mask = as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
+        for igpt in 1:n_gpt
+            ClimaComms.@threaded device for gcol in 1:ncol
+                # get column views
+                flux_col = FluxSWc(flux, gcol)
+                flux_sw_col = FluxSWc(flux_sw, gcol)
+                op_col = TwoStreamc(op, gcol)
+                as_col = AtmosphericStatec(as, gcol)
+                src_sw_col = SourceSW2Strc(src_sw, gcol)
+                bcs_sw_col = SwBCsc(bcs_sw, gcol)
+                bld_cld_mask && Optics.build_cloud_mask!(
+                    view(as.cld_mask_sw, :, gcol),
+                    view(as.cld_frac, :, gcol),
+                    as.cld_mask_type,
+                )
+                (; τ, ssa, g) = op_col
+                #igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
+                igpt == 1 && set_flux_to_zero!(flux_sw_col)
+                for glay in 1:nlay
+                    τ[glay], ssa[glay], g[glay] =
+                        compute_optical_props(as_col, gcol, glay, igpt, lookup_sw, lookup_sw_cld)
                 end
+                # Cell properties: transmittance and reflectance for direct and diffuse radiation
+                solar_frac = lookup_sw.solar_src_scaled[igpt]
+                ibnd = lookup_sw.major_gpt2bnd[igpt]
+                # call rte shortwave solver
+                rte_sw!(op_col, src_sw_col, bcs_sw_col, flux_col, solar_frac, igpt, n_gpt, ibnd, nlev)
+                n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
             end
         end
     end
     return nothing
 end
 
+function rte_sw_2stream_solve!(
+    device::ClimaComms.CUDADevice,
+    flux::FluxSW{FT},
+    flux_sw::FluxSW{FT},
+    op::TwoStream{FT},
+    bcs_sw::SwBCs{FT},
+    src_sw::SourceSW2Str{FT},
+    max_threads,
+    as::AtmosphericState{FT},
+    lookup_sw::LookUpSW,
+    lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+) where {FT <: AbstractFloat}
+    (; nlay, ncol) = as
+    nlev = nlay + 1
+    n_gpt = length(lookup_sw.solar_src_scaled)
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    args = (flux, flux_sw, op, bcs_sw, src_sw, nlay, ncol, as, lookup_sw, lookup_sw_cld)
+    @cuda threads = (tx) blocks = (bx) rte_sw_2stream_solve_CUDA!(args...)
+    return nothing
+end
+
 function rte_sw_2stream_solve_CUDA!(
     flux::FluxSW{FT},
     flux_sw::FluxSW{FT},
@@ -87,164 +154,128 @@ function rte_sw_2stream_solve_CUDA!(
     src_sw::SourceSW2Str{FT},
     nlay,
     ncol,
-    major_gpt2bnd::AbstractArray{UInt8, 1},
-    solar_src_scaled::AbstractArray{FT, 1},
-    as::AbstractAtmosphericState{FT},
-    lookup_sw::Union{LookUpSW, Nothing} = nothing,
+    as::AtmosphericState{FT},
+    lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
 ) where {FT <: AbstractFloat}
     gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
     nlev = nlay + 1
-    n_gpt = length(major_gpt2bnd)
+    n_gpt = length(lookup_sw.major_gpt2bnd)
     if gcol ≤ ncol
+        # get column views
+        flux_col = FluxSWc(flux, gcol)
+        flux_sw_col = FluxSWc(flux_sw, gcol)
+        op_col = TwoStreamc(op, gcol)
+        as_col = AtmosphericStatec(as, gcol)
+        src_sw_col = SourceSW2Strc(src_sw, gcol)
+        bcs_sw_col = SwBCsc(bcs_sw, gcol)
+        set_flux_to_zero!(flux_sw_col)
+        flux_up_col_sw = flux_sw_col.flux_up
+        flux_dn_col_sw = flux_sw_col.flux_dn
+        flux_net_col_sw = flux_sw_col.flux_net
+        flux_up_col = flux_col.flux_up
+        flux_dn_col = flux_col.flux_dn
+        (; τ, ssa, g) = op_col
         @inbounds for igpt in 1:n_gpt
-            igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
             if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
-                Optics.build_cloud_mask!(as.cld_mask_sw, as.cld_frac, as.random_sw, gcol, igpt, as.cld_mask_type)
+                Optics.build_cloud_mask!(view(as.cld_mask_sw, :, gcol), view(as.cld_frac, :, gcol), as.cld_mask_type)
+            end
+            for glay in 1:nlay
+                τ[glay], ssa[glay], g[glay] = compute_optical_props(as_col, gcol, glay, igpt, lookup_sw, lookup_sw_cld)
+            end
+            # Cell properties: transmittance and reflectance for direct and diffuse radiation
+            solar_frac = lookup_sw.solar_src_scaled[igpt]
+            ibnd = lookup_sw.major_gpt2bnd[igpt]
+            # rte shortwave solver
+            rte_sw!(op_col, src_sw_col, bcs_sw_col, flux_col, solar_frac, igpt, n_gpt, ibnd, nlev)
+            if n_gpt > 1
+                for i in eachindex(flux_up_col_sw)
+                    @inbounds flux_up_col_sw[i] += flux_up_col[i]
+                    @inbounds flux_dn_col_sw[i] += flux_dn_col[i]
+                end
             end
-            compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
-            sw_two_stream!(op, src_sw, bcs_sw, gcol, nlay) # Cell properties: transmittance and reflectance for direct and diffuse radiation
-            sw_source_2str!(src_sw, bcs_sw, gcol, flux, igpt, solar_src_scaled, major_gpt2bnd, nlay) # Direct-beam and source for diffuse radiation
-            adding_sw!(src_sw, bcs_sw, gcol, flux, igpt, major_gpt2bnd, nlev) # Transport
-            n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
+        end
+        for i in eachindex(flux_net_col_sw)
+            @inbounds flux_net_col_sw[i] = flux_up_col_sw[i] - flux_dn_col_sw[i]
         end
     end
     return nothing
 end
 
-"""
-    sw_two_stream!(
-        op::TwoStream{FT},
-        src_sw::SourceSW2Str{FT},
-        bcs_sw::SwBCs{FT},
-        gcol,
-        nlay,
-    ) where {FT<:AbstractFloat}
-
-Computes cell properties (transmittance and reflectance) for direct and diffuse radiation
-
-Two-stream solutions to direct and diffuse reflectance and transmittance for a layer
-with optical depth tau, single scattering albedo w0, and asymmetery parameter g.
- 
-Equations are developed in Meador and Weaver, 1980,
-doi:10.1175/1520-0469(1980)037<0630:TSATRT>2.0.CO;2
-"""
-function sw_two_stream!(
-    op::TwoStream{FT},
-    src_sw::SourceSW2Str{FT},
-    bcs_sw::SwBCs{FT},
-    gcol::Int,
-    nlay::Int,
-) where {FT <: AbstractFloat}
-    zenith = bcs_sw.zenith
-    (; τ, ssa, g) = op
-    (; Rdif, Tdif, Rdir, Tdir, Tnoscat) = src_sw
+function sw_2stream_coeffs(τ::FT, ssa::FT, g::FT, μ₀::FT) where {FT}
     k_min = FT(1e4 * eps(FT)) # Suggestion from Chiel van Heerwaarden
-    μ₀ = FT(cos(zenith[gcol]))
-
-    @inbounds for glay in 1:nlay
-        # Zdunkowski Practical Improved Flux Method "PIFM"
-        #  (Zdunkowski et al., 1980;  Contributions to Atmospheric Physics 53, 147-66)
-        γ1 = (FT(8) - ssa[glay, gcol] * (FT(5) + FT(3) * g[glay, gcol])) * FT(0.25)
-        γ2 = FT(3) * (ssa[glay, gcol] * (FT(1) - g[glay, gcol])) * FT(0.25)
-        γ3 = (FT(2) - (FT(3) * cos(zenith[gcol])) * g[glay, gcol]) * FT(0.25)
-        γ4 = FT(1) - γ3
-        α1 = γ1 * γ4 + γ2 * γ3                          # Eq. 16
-        α2 = γ1 * γ3 + γ2 * γ4                          # Eq. 17
-        k = sqrt(max((γ1 - γ2) * (γ1 + γ2), k_min))
+    # Zdunkowski Practical Improved Flux Method "PIFM"
+    #  (Zdunkowski et al., 1980;  Contributions to Atmospheric Physics 53, 147-66)
+    γ1 = (FT(8) - ssa * (FT(5) + FT(3) * g)) * FT(0.25)
+    γ2 = FT(3) * (ssa * (FT(1) - g)) * FT(0.25)
+    γ3 = (FT(2) - (FT(3) * μ₀) * g) * FT(0.25)
+    γ4 = FT(1) - γ3
+    α1 = γ1 * γ4 + γ2 * γ3                          # Eq. 16
+    α2 = γ1 * γ3 + γ2 * γ4                          # Eq. 17
+    k = sqrt(max((γ1 - γ2) * (γ1 + γ2), k_min))
 
-        exp_minusktau = exp(-τ[glay, gcol] * k)
-        exp_minus2ktau = exp_minusktau * exp_minusktau
+    exp_minusktau = exp(-τ * k)
+    exp_minus2ktau = exp_minusktau * exp_minusktau
 
-        # Refactored to avoid rounding errors when k, gamma1 are of very different magnitudes
-        RT_term = FT(1) / (k * (FT(1) + exp_minus2ktau) + γ1 * (FT(1) - exp_minus2ktau))
+    # Refactored to avoid rounding errors when k, gamma1 are of very different magnitudes
+    RT_term = FT(1) / (k * (FT(1) + exp_minus2ktau) + γ1 * (FT(1) - exp_minus2ktau))
 
-        @inbounds Rdif[glay, gcol] = RT_term * γ2 * (FT(1) - exp_minus2ktau) # Eqn. 25
-        @inbounds Tdif[glay, gcol] = RT_term * FT(2) * k * exp_minusktau     # Eqn. 26
+    Rdif = RT_term * γ2 * (FT(1) - exp_minus2ktau) # Eqn. 25
+    Tdif = RT_term * FT(2) * k * exp_minusktau     # Eqn. 26
 
-        # Transmittance of direct, unscattered beam. Also used below
-        @inbounds T₀ = Tnoscat[glay, gcol] = exp(-τ[glay, gcol] / cos(zenith[gcol]))
+    # Transmittance of direct, unscattered beam. Also used below
+    T₀ = Tnoscat = exp(-τ / μ₀)
 
-        # Direct reflect and transmission
-        k_μ = k * μ₀
-        k_γ3 = k * γ3
-        k_γ4 = k * γ4
-
-        # Equation 14, multiplying top and bottom by exp(-k*tau)
-        #   and rearranging to avoid div by 0.
-        if abs(FT(1) - k_μ * k_μ) ≥ eps(FT)
-            RT_term = ssa[glay, gcol] * RT_term / (FT(1) - k_μ * k_μ)
-        else
-            RT_term = ssa[glay, gcol] * RT_term / eps(FT)
-        end
+    # Direct reflect and transmission
+    k_μ = k * μ₀
+    k_γ3 = k * γ3
+    k_γ4 = k * γ4
 
-        @inbounds Rdir_unconstrained =
-            RT_term * (
-                (FT(1) - k_μ) * (α2 + k_γ3) - (FT(1) + k_μ) * (α2 - k_γ3) * exp_minus2ktau -
-                FT(2) * (k_γ3 - α2 * k_μ) * exp_minusktau * T₀
-            )
-        #
-        # Equation 15, multiplying top and bottom by exp(-k*tau),
-        #   multiplying through by exp(-tau/mu0) to
-        #   prefer underflow to overflow
-        # Omitting direct transmittance
-        #
-        @inbounds Tdir_unconstrained =
-            -RT_term * (
-                (FT(1) + k_μ) * (α1 + k_γ4) * T₀ - (FT(1) - k_μ) * (α1 - k_γ4) * exp_minus2ktau * T₀ -
-                FT(2) * (k_γ4 + α1 * k_μ) * exp_minusktau
-            )
-        # Final check that energy is not spuriously created, by recognizing that
-        # the beam can either be reflected, penetrate unscattered to the base of a layer, 
-        # or penetrate through but be scattered on the way - the rest is absorbed
-        # Makes the equations safer in single precision. Credit: Robin Hogan, Peter Ukkonen
-        Rdir[glay, gcol] = max(FT(0), min(Rdir_unconstrained, (FT(1) - T₀)))
-        Tdir[glay, gcol] = max(FT(0), min(Tdir_unconstrained, (FT(1) - T₀ - Rdir[glay, gcol])))
+    # Equation 14, multiplying top and bottom by exp(-k*tau)
+    #   and rearranging to avoid div by 0.
+    if abs(FT(1) - k_μ * k_μ) ≥ eps(FT)
+        RT_term = ssa * RT_term / (FT(1) - k_μ * k_μ)
+    else
+        RT_term = ssa * RT_term / eps(FT)
     end
-end
 
-"""
-    sw_source_2str!(
-        src_sw::SourceSW2Str{FT},
-        bcs_sw::SwBCs{FT},
-        gcol,
-        flux::FluxSW{FT},
-        solar_frac::FT,
-        ibnd,
-        nlay,
-    ) where {FT<:AbstractFloat}
-
-Direct-beam and source for diffuse radiation
-"""
-function sw_source_2str!(
-    src_sw::SourceSW2Str{FT},
-    bcs_sw::SwBCs{FT},
-    gcol::Int,
-    flux::FluxSW{FT},
-    igpt::Int,
-    solar_src_scaled::AbstractArray{FT, 1},
-    major_gpt2bnd::AbstractArray{UInt8, 1},
-    nlay::Int,
-) where {FT <: AbstractFloat}
-    ibnd = major_gpt2bnd[igpt]
-    solar_frac = solar_src_scaled[igpt]
-    (; toa_flux, zenith, sfc_alb_direct) = bcs_sw
-    (; Rdir, Tdir, Tnoscat, src_up, src_dn, sfc_source) = src_sw
-    flux_dn_dir = flux.flux_dn_dir
+    Rdir_unconstrained =
+        RT_term * (
+            (FT(1) - k_μ) * (α2 + k_γ3) - (FT(1) + k_μ) * (α2 - k_γ3) * exp_minus2ktau -
+            FT(2) * (k_γ3 - α2 * k_μ) * exp_minusktau * T₀
+        )
+    #
+    # Equation 15, multiplying top and bottom by exp(-k*tau),
+    #   multiplying through by exp(-tau/mu0) to
+    #   prefer underflow to overflow
+    # Omitting direct transmittance
+    #
+    Tdir_unconstrained =
+        -RT_term * (
+            (FT(1) + k_μ) * (α1 + k_γ4) * T₀ - (FT(1) - k_μ) * (α1 - k_γ4) * exp_minus2ktau * T₀ -
+            FT(2) * (k_γ4 + α1 * k_μ) * exp_minusktau
+        )
+    # Final check that energy is not spuriously created, by recognizing that
+    # the beam can either be reflected, penetrate unscattered to the base of a layer, 
+    # or penetrate through but be scattered on the way - the rest is absorbed
+    # Makes the equations safer in single precision. Credit: Robin Hogan, Peter Ukkonen
+    Rdir = max(FT(0), min(Rdir_unconstrained, (FT(1) - T₀)))
+    Tdir = max(FT(0), min(Tdir_unconstrained, (FT(1) - T₀ - Rdir)))
+    return (Rdir, Tdir, Tnoscat, Rdif, Tdif)
+end
 
-    # layer index = level index
-    # previous level is up (+1)
-    @inbounds flux_dn_dir[nlay + 1, gcol] = toa_flux[gcol] * solar_frac * cos(zenith[gcol])
-    @inbounds for ilev in nlay:-1:1
-        src_up[ilev, gcol] = Rdir[ilev, gcol] * flux_dn_dir[ilev + 1, gcol]
-        src_dn[ilev, gcol] = Tdir[ilev, gcol] * flux_dn_dir[ilev + 1, gcol]
-        flux_dn_dir[ilev, gcol] = Tnoscat[ilev, gcol] * flux_dn_dir[ilev + 1, gcol]
+# Direct-beam and source for diffuse radiation
+function get_flux_dn_dir(τ, μ₀, flux_dn_dir_top, lev)
+    nlay = length(τ)
+    τ_sum = zero(eltype(τ))
+    for ilev in nlay:-1:lev
+        τ_sum += τ[ilev]
     end
-    @inbounds sfc_source[gcol] = flux_dn_dir[1, gcol] * sfc_alb_direct[ibnd, gcol]
+    return flux_dn_dir_top * exp(-τ_sum / μ₀)
 end
 
 """
-    adding_sw!(
+    rte_sw!(
         src_sw::SourceSW2Str{FT},
         bcs_sw::SwBCs{FT},
         gcol,
@@ -259,64 +290,83 @@ Transport for the two stream shortwave problem.
 Transport of diffuse radiation through a vertically layered atmosphere.
 Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08)
 """
-function adding_sw!(
-    src_sw::SourceSW2Str{FT},
-    bcs_sw::SwBCs{FT},
-    gcol::Int,
-    flux::FluxSW{FT},
+function rte_sw!(
+    (; τ, ssa, g)::TwoStreamc,
+    (; albedo, src)::SourceSW2Strc,
+    bcs_sw::SwBCsc,
+    (; flux_up, flux_dn, flux_dn_dir)::FluxSWc,
+    solar_frac::FT,
     igpt::Int,
-    major_gpt2bnd::AbstractArray{UInt8, 1},
+    n_gpt::Int,
+    ibnd::UInt8,
     nlev::Int,
-) where {FT <: AbstractFloat}
-    ibnd = major_gpt2bnd[igpt]
+) where {FT}
     nlay = nlev - 1
-    n_gpt = length(major_gpt2bnd)
-    # destructuring
-    (; flux_up, flux_dn, flux_dn_dir, flux_net) = flux
-    (; albedo, sfc_source, Rdif, Tdif, src_up, src_dn, src) = src_sw
-
-    @inbounds for ilev in 1:nlev
-        flux_dn[ilev, gcol] = FT(0)
-        flux_up[ilev, gcol] = FT(0)
+    toa_flux = bcs_sw.toa_flux[1]
+    sfc_alb_direct = bcs_sw.sfc_alb_direct[ibnd]
+    μ₀ = bcs_sw.cos_zenith[1]
+    τ_sum = FT(0)
+    for ilay in 1:nlay
+        @inbounds τ_sum += τ[ilay]
     end
+    # Direct-beam and source for diffuse radiation
+    flux_dn_dir_top = toa_flux * solar_frac * μ₀
+    flux_dn_dir_bot = flux_dn_dir_top * exp(-τ_sum / μ₀) # store value at surface
+    @inbounds flux_dn_dir[1] = flux_dn_dir_bot # store value at surface
+    sfc_source = flux_dn_dir_bot * sfc_alb_direct
 
+    @inbounds flux_dn[nlev] = FT(0) # set to incoming flux when provided?
     # Albedo of lowest level is the surface albedo...
-    @inbounds albedo[1, gcol] = bcs_sw.sfc_alb_diffuse[ibnd, gcol]
+    @inbounds surface_albedo = albedo[1] = bcs_sw.sfc_alb_diffuse[ibnd]
     # ... and source of diffuse radiation is surface emission
-    @inbounds src[1, gcol] = sfc_source[gcol]
+    @inbounds src[1] = sfc_source
     # From bottom to top of atmosphere --
     #   compute albedo and source of upward radiation
+    τ_cum = τ_sum
+    albedo_ilev, src_ilev = surface_albedo, sfc_source
     @inbounds for ilev in 1:nlay
-        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
-        albedo[ilev + 1, gcol] = Rdif[ilev, gcol] + Tdif[ilev, gcol] * Tdif[ilev, gcol] * albedo[ilev, gcol] * denom # Equation 9
+        τ_ilev = τ[ilev]
+        (Rdir, Tdir, _, Rdif, Tdif) = sw_2stream_coeffs(τ_ilev, ssa[ilev], g[ilev], μ₀)
+        denom = FT(1) / (FT(1) - Rdif * albedo_ilev)  # Eq 10
+        albedo_ilevplus1 = Rdif + Tdif * Tdif * albedo_ilev * denom # Equation 9
         # 
         # Equation 11 -- source is emitted upward radiation at top of layer plus
         # radiation emitted at bottom of layer,
         # transmitted through the layer and reflected from layers below (Tdiff*src*albedo)
-        src[ilev + 1, gcol] =
-            src_up[ilev, gcol] + Tdif[ilev, gcol] * denom * (src[ilev, gcol] + albedo[ilev, gcol] * src_dn[ilev, gcol])
+        τ_cum -= τ_ilev
+        flux_dn_dir_ilevplus1 = flux_dn_dir_top * exp(-τ_cum / μ₀)
+        src_up_ilev = Rdir * flux_dn_dir_ilevplus1 #flux_dn_dir[ilev + 1]
+        src_dn_ilev = Tdir * flux_dn_dir_ilevplus1 #flux_dn_dir[ilev + 1]
+        src_ilev = src_up_ilev + Tdif * denom * (src_ilev + albedo_ilev * src_dn_ilev)
+        src[ilev + 1] = src_ilev
+        albedo[ilev + 1] = albedo_ilev = albedo_ilevplus1
     end
-
     # Eq 12, at the top of the domain upwelling diffuse is due to ...
-    @inbounds flux_up[nlev, gcol] =
-        flux_dn[nlev, gcol] * albedo[nlev, gcol] + # ... reflection of incident diffuse and
-        src[nlev, gcol]                          # scattering by the direct beam below
+    @inbounds flux_up[nlev] =
+        flux_dn[nlev] * albedo[nlev] + # ... reflection of incident diffuse and
+        src[nlev]                          # scattering by the direct beam below
 
     # From the top of the atmosphere downward -- compute fluxes
-    @inbounds for ilev in nlay:-1:1
-        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
-        flux_dn[ilev, gcol] =
-            (Tdif[ilev, gcol] * flux_dn[ilev + 1, gcol] + # Equation 13
-             Rdif[ilev, gcol] * src[ilev, gcol] +
-             src_dn[ilev, gcol]) * denom
-        flux_up[ilev, gcol] =
-            flux_dn[ilev, gcol] * albedo[ilev, gcol] + # Equation 12
-            src[ilev, gcol]
-    end
+    @inbounds flux_dn_ilevplus1 = flux_dn[nlev]
+    @inbounds flux_dn[nlev] += flux_dn_dir_top
 
-    @inbounds for ilev in 1:nlev
-        flux_dn[ilev, gcol] += flux_dn_dir[ilev, gcol]
-        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
+    τ_cum = FT(0)
+
+    @inbounds for ilev in nlay:-1:1
+        τ_ilev, albedo_ilev, src_ilev = τ[ilev], albedo[ilev], src[ilev]
+        (_, Tdir, _, Rdif, Tdif) = sw_2stream_coeffs(τ_ilev, ssa[ilev], g[ilev], μ₀)
+        denom = FT(1) / (FT(1) - Rdif * albedo_ilev)  # Eq 10
+        src_dn_ilev = Tdir * flux_dn_dir_top * exp(-τ_cum / μ₀)
+        τ_cum += τ_ilev
+        flux_dn_ilev = (Tdif * flux_dn_ilevplus1 + # Equation 13
+                        Rdif * src_ilev +
+                        src_dn_ilev) * denom
+        flux_up[ilev] =
+            flux_dn_ilev * albedo_ilev + # Equation 12
+            src_ilev
+        flux_dn[ilev] = flux_dn_ilev + flux_dn_dir_top * exp(-τ_cum / μ₀)
+        flux_dn_ilevplus1 = flux_dn_ilev
     end
+    return nothing
 end
 # ---------------------------------------------------------------
diff --git a/test/all_sky_utils.jl b/test/all_sky_utils.jl
index 6c3a82b23..737418628 100644
--- a/test/all_sky_utils.jl
+++ b/test/all_sky_utils.jl
@@ -69,7 +69,7 @@ function all_sky(
     close(ds_sw_cld)
     # reading input file 
     ds_in = Dataset(input_file, "r")
-    as, sfc_emis, sfc_alb_direct, sfc_alb_diffuse, zenith, toa_flux, bot_at_1 = setup_allsky_as(
+    as, sfc_emis, sfc_alb_direct, sfc_alb_diffuse, cos_zenith, toa_flux, bot_at_1 = setup_allsky_as(
         context,
         ds_in,
         idx_gases,
@@ -109,7 +109,7 @@ function all_sky(
     # setting up boundary conditions
     inc_flux_diffuse = nothing
 
-    bcs_sw = SwBCs(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
+    bcs_sw = SwBCs(cos_zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
 
     fluxb_sw = FluxSW(ncol, nlay, FT, DA) # flux storage for bandwise calculations
     flux_sw = FluxSW(ncol, nlay, FT, DA)  # shortwave fluxes for band calculations    
diff --git a/test/clear_sky_utils.jl b/test/clear_sky_utils.jl
index 409109db4..1a902dd8d 100644
--- a/test/clear_sky_utils.jl
+++ b/test/clear_sky_utils.jl
@@ -66,7 +66,7 @@ function clear_sky(
     # reading rfmip data to atmospheric state
     ds_lw_in = Dataset(input_file, "r")
 
-    (as, sfc_emis, sfc_alb_direct, zenith, toa_flux, usecol) =
+    (as, sfc_emis, sfc_alb_direct, cos_zenith, toa_flux, usecol) =
         setup_rfmip_as(context, ds_lw_in, idx_gases, exp_no, lookup_lw, FT, VMR, max_threads, param_set)
     close(ds_lw_in)
 
@@ -86,7 +86,7 @@ function clear_sky(
     src_sw = source_func_shortwave(FT, ncol, nlay, opc, DA)        # allocating shortwave source function object
     inc_flux_diffuse = nothing
     sfc_alb_diffuse = FTA2D(deepcopy(sfc_alb_direct))
-    bcs_sw = SwBCs(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
+    bcs_sw = SwBCs(cos_zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
 
     fluxb_sw = FluxSW(ncol, nlay, FT, DA) # flux storage for bandwise calculations
     flux_sw = FluxSW(ncol, nlay, FT, DA)  # shortwave fluxes for band calculations
diff --git a/test/gray_atm_utils.jl b/test/gray_atm_utils.jl
index eaa3de5b2..53075901d 100644
--- a/test/gray_atm_utils.jl
+++ b/test/gray_atm_utils.jl
@@ -159,7 +159,7 @@ function gray_atmos_sw_test(
     max_threads = 256            # maximum number of threads for KA kernels
     deg2rad = FT(π) / FT(180)
 
-    zenith = DA{FT, 1}(undef, ncol)   # solar zenith angle in radians
+    cos_zenith = DA{FT, 1}(undef, ncol)   # cosine of solar zenith angle in radians
     toa_flux = DA{FT, 1}(undef, ncol) # top of atmosphere flux
     sfc_alb_direct = DA{FT, 2}(undef, nbnd, ncol) # surface albedo (direct)
     sfc_alb_diffuse = DA{FT, 2}(undef, nbnd, ncol) # surface albedo (diffuse)
@@ -174,7 +174,7 @@ function gray_atmos_sw_test(
         lat = DA{FT}(range(FT(-90), FT(90), length = ncol)) # latitude
     end
 
-    zenith .= deg2rad * 52.95 #acsc(FT(1.66))               # corresponding to ~52.95 deg zenith angle    
+    cos_zenith .= cos(deg2rad * 52.95) #acsc(FT(1.66))               # corresponding to ~52.95 deg zenith angle    
     toa_flux .= FT(1407.679)
     sfc_alb_direct .= FT(0.1)
     sfc_alb_diffuse .= FT(0.1)
@@ -182,7 +182,7 @@ function gray_atmos_sw_test(
     as = setup_gray_as_pr_grid(context, nlay, lat, p0, pe, otp, param_set, DA) # init gray atmos state
     op = OPC(FT, ncol, nlay, DA)
     src_sw = source_func_shortwave(FT, ncol, nlay, opc, DA)
-    bcs_sw = SwBCs(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
+    bcs_sw = SwBCs(cos_zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
     fluxb_sw = FluxSW(ncol, nlay, FT, DA)
     flux_sw = FluxSW(ncol, nlay, FT, DA)
 
@@ -191,16 +191,16 @@ function gray_atmos_sw_test(
     solve_sw!(slv, max_threads)
 
     τ = Array(slv.op.τ)
-    zenith = Array(zenith)
-    flux_dn_dir = Array(slv.flux_sw.flux_dn_dir)
+    cos_zenith = Array(cos_zenith)
+    flux_dn_dir_bot = Array(slv.flux_sw.flux_dn_dir)[1]
     toa_flux = Array(toa_flux)
 
     # testing with exact solution
-    ot_tot = sum(τ[:, 1]) / cos(zenith[1])
-    exact = (toa_flux[1] * cos(zenith[1])) * exp(-ot_tot)
+    ot_tot = sum(τ[:, 1]) / cos_zenith[1]
+    exact = (toa_flux[1] * cos_zenith[1]) * exp(-ot_tot)
 
     rel_toler = FT(0.001)
-    rel_error = abs(flux_dn_dir[1] - exact) / exact
+    rel_error = abs(flux_dn_dir_bot - exact) / exact
     println("*************************************************")
     println("Running shortwave test for gray atmosphere model - $(opc); ncol = $ncol; context = $context")
     println("relative error = $rel_error")
@@ -208,7 +208,8 @@ function gray_atmos_sw_test(
 
     if device isa ClimaComms.CPUSingleThreaded
         JET.@test_opt solve_sw!(slv, max_threads)
-        @test_broken (@allocated solve_sw!(slv, max_threads)) == 0
+        @test (@allocated solve_sw!(slv, max_threads)) == 0
+        #@test_broken (@allocated solve_sw!(slv, max_threads)) == 0
         @test (@allocated solve_sw!(slv, max_threads)) ≤ 256
     end
 end
diff --git a/test/read_all_sky.jl b/test/read_all_sky.jl
index 7e6e59069..3a85ef786 100644
--- a/test/read_all_sky.jl
+++ b/test/read_all_sky.jl
@@ -41,13 +41,13 @@ function setup_allsky_as(
     sfc_emis = DA{FT, 2}(undef, nbnd_lw, ncol)
     sfc_alb_direct = DA{FT, 2}(undef, nbnd_sw, ncol)
     sfc_alb_diffuse = DA{FT, 2}(undef, nbnd_sw, ncol)
-    zenith = DA{FT, 1}(undef, ncol)
+    cos_zenith = DA{FT, 1}(undef, ncol)
     irrad = DA{FT, 1}(undef, ncol)
     # these values are taken from the example
     sfc_emis .= FT(0.98)
     sfc_alb_direct .= FT(0.06)
     sfc_alb_diffuse .= FT(0.06)
-    zenith .= acos(FT(0.86))
+    cos_zenith .= FT(0.86)
     irrad .= FT(lkp_sw.solar_src_tot)
 
 
@@ -72,27 +72,25 @@ function setup_allsky_as(
     #col_dry from the dataset not used in the FORTRAN RRTMGP example
 
     # Reading volume mixing ratios 
-    vmrat = zeros(FT, nlay, ncol, ngas)
+    vmrat = zeros(FT, ngas, nlay, ncol)
 
-    vmrat[:, 1, idx_gases["h2o"]] .= Array{FT}(Array(ds_in["vmr_h2o"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["o3"]] .= Array{FT}(Array(ds_in["vmr_o3"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["co2"]] .= Array{FT}(Array(ds_in["vmr_co2"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["n2o"]] .= Array{FT}(Array(ds_in["vmr_n2o"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["co"]] .= Array{FT}(Array(ds_in["vmr_co"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["ch4"]] .= Array{FT}(Array(ds_in["vmr_ch4"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["o2"]] .= Array{FT}(Array(ds_in["vmr_o2"])[1, lay_ind])
-    vmrat[:, 1, idx_gases["n2"]] .= Array{FT}(Array(ds_in["vmr_n2"])[1, lay_ind])
+    vmrat[idx_gases["h2o"], :, 1] .= Array{FT}(Array(ds_in["vmr_h2o"])[1, lay_ind])
+    vmrat[idx_gases["o3"], :, 1] .= Array{FT}(Array(ds_in["vmr_o3"])[1, lay_ind])
+    vmrat[idx_gases["co2"], :, 1] .= Array{FT}(Array(ds_in["vmr_co2"])[1, lay_ind])
+    vmrat[idx_gases["n2o"], :, 1] .= Array{FT}(Array(ds_in["vmr_n2o"])[1, lay_ind])
+    vmrat[idx_gases["co"], :, 1] .= Array{FT}(Array(ds_in["vmr_co"])[1, lay_ind])
+    vmrat[idx_gases["ch4"], :, 1] .= Array{FT}(Array(ds_in["vmr_ch4"])[1, lay_ind])
+    vmrat[idx_gases["o2"], :, 1] .= Array{FT}(Array(ds_in["vmr_o2"])[1, lay_ind])
+    vmrat[idx_gases["n2"], :, 1] .= Array{FT}(Array(ds_in["vmr_n2"])[1, lay_ind])
 
     for icol in 2:ncol
-        vmrat[:, icol, :] .= vmrat[:, 1, :]
+        vmrat[:, :, icol] .= vmrat[:, :, 1]
     end
     vmr = Vmr(DA(vmrat))
     col_dry = DA{FT, 2}(undef, nlay, ncol)
-    vmr_h2o = view(vmr.vmr, :, :, idx_gases["h2o"])
+    vmr_h2o = view(vmr.vmr, idx_gases["h2o"], :, :)
 
     cld_frac = zeros(FT, nlay, ncol)
-    random_lw = DA{FT}(undef, nlay, ncol)
-    random_sw = DA{FT}(undef, nlay, ncol)
     cld_mask_lw = zeros(Bool, nlay, ncol)
     cld_mask_sw = zeros(Bool, nlay, ncol)
     cld_r_eff_liq = zeros(FT, nlay, ncol)
@@ -155,7 +153,6 @@ function setup_allsky_as(
             typeof(p_lev),
             typeof(cld_r_eff_liq),
             typeof(cld_mask_lw),
-            typeof(random_lw),
             MaxRandomOverlap,
             typeof(vmr),
         }(
@@ -173,8 +170,6 @@ function setup_allsky_as(
             cld_path_liq,
             cld_path_ice,
             cld_frac,
-            random_lw,
-            random_sw,
             cld_mask_lw,
             cld_mask_sw,
             MaxRandomOverlap(),
@@ -186,7 +181,7 @@ function setup_allsky_as(
         sfc_emis,
         sfc_alb_direct,
         sfc_alb_diffuse,
-        zenith,
+        cos_zenith,
         irrad,
         bot_at_1,
     )
diff --git a/test/read_rfmip_clear_sky.jl b/test/read_rfmip_clear_sky.jl
index 48b1eb4ea..9fe713bc7 100644
--- a/test/read_rfmip_clear_sky.jl
+++ b/test/read_rfmip_clear_sky.jl
@@ -43,7 +43,7 @@ function setup_rfmip_as(
         end
     end
 
-    zenith = DA{FT, 1}(zenith)
+    cos_zenith = DA{FT, 1}(cos.(zenith))
     irrad = DA{FT, 1}(irrad)
     #--------------------------------------------------------------
 
@@ -118,8 +118,6 @@ function setup_rfmip_as(
     cld_path_liq = nothing
     cld_path_ice = nothing
     cld_frac = nothing
-    random_lw = nothing
-    random_sw = nothing
     cld_mask_lw = nothing
     cld_mask_sw = nothing
     ice_rgh = 1
@@ -132,7 +130,6 @@ function setup_rfmip_as(
             typeof(p_lev),
             typeof(cld_r_eff_liq),
             typeof(cld_mask_lw),
-            typeof(random_lw),
             Nothing,
             typeof(vmr),
         }(
@@ -150,8 +147,6 @@ function setup_rfmip_as(
             cld_path_liq,
             cld_path_ice,
             cld_frac,
-            random_lw,
-            random_sw,
             cld_mask_lw,
             cld_mask_sw,
             nothing,
@@ -162,7 +157,7 @@ function setup_rfmip_as(
         ),
         sfc_emis,
         sfc_alb,
-        zenith,
+        cos_zenith,
         irrad,
         usecol,
     )
diff --git a/test/rfmip_clear_sky_sw.jl b/test/rfmip_clear_sky_sw.jl
index 04b6c91bf..e563400b3 100644
--- a/test/rfmip_clear_sky_sw.jl
+++ b/test/rfmip_clear_sky_sw.jl
@@ -47,7 +47,7 @@ function sw_rfmip(context, ::Type{OPC}, ::Type{SRC}, ::Type{VMR}, ::Type{FT}) wh
     # reading rfmip data to atmospheric state
     ds_sw_in = Dataset(sw_input_file, "r")
 
-    (as, _, sfc_alb_direct, zenith, toa_flux, usecol) =
+    (as, _, sfc_alb_direct, cos_zenith, toa_flux, usecol) =
         setup_rfmip_as(context, ds_sw_in, idx_gases, exp_no, lookup_sw, FT, VMR, max_threads, param_set)
     close(ds_sw_in)
 
@@ -60,7 +60,7 @@ function sw_rfmip(context, ::Type{OPC}, ::Type{SRC}, ::Type{VMR}, ::Type{FT}) wh
     # setting up boundary conditions
     inc_flux_diffuse = nothing
     sfc_alb_diffuse = FTA2D(deepcopy(sfc_alb_direct))
-    bcs_sw = SwBCs{FT, FTA1D, Nothing, FTA2D}(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
+    bcs_sw = SwBCs{FT, FTA1D, Nothing, FTA2D}(cos_zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
     fluxb_sw = FluxSW(ncol, nlay, FT, DA) # flux storage for bandwise calculations
     flux_sw = FluxSW(ncol, nlay, FT, DA)  # shortwave fluxes for band calculations