Update cld mask function

CliMA · Dec 18, 2023 · 756455d · 756455d
1 parent 1b9f74d
commit 756455d
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 72 deletions.
diff --git a/src/optics/CloudOptics.jl b/src/optics/CloudOptics.jl
@@ -216,65 +216,49 @@ 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, random_arr, gcol, ::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, cld_frac, ::MaxRandomOverlap)
     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)
+    nlay = size(cld_frac, 1)
 
-    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
+    start, finish = 0, 0
+    @inbounds for ilay in eachindex(cld_frac) # for GPU compatibility (start = findfirst(cld_frac .> FT(0)))
+        if cld_frac[ilay] > 0
+            start = start == 0 ? ilay : start
+            finish = ilay
+            cld_mask[ilay] = true # set mask for cloudy layers
+        else
+            cld_mask[ilay] = false
         end
     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
-        end
-        @inbounds for ilay in (finish + 1):nlay
-            local_cld_mask[ilay] = false
-        end
+    if start > 0
         # 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])
+        # RRTMG uses random_arr[finish] > (FT(1) - cld_frac[finish]), we change > to >= to address edge cases
+        cld_frac_ilayplus1 = cld_frac[finish]
+        random_ilayplus1 = Random.rand()
+        @inbounds cld_mask[finish] = cld_mask_ilayplus1 = random_ilayplus1 >= (FT(1) - cld_frac_ilayplus1)
         @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
+            cld_frac_ilay = cld_frac[ilay]
+            if cld_frac_ilay > FT(0)
+                if cld_mask_ilayplus1
+                    # use same random number from the layer above if layer above is cloudy
+                    cld_mask_ilay = random_ilayplus1 >= (FT(1) - cld_frac_ilay) # 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
+                    # update random numbers if layer above is not cloudy
+                    cld_mask_ilay = Random.rand() * (FT(1) - cld_frac_ilayplus1) >= (FT(1) - cld_frac_ilay)
                 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
+                # RRTMG uses random_arr[ilay] > (FT(1) - cld_frac[ilay]), we change > to >= to address edge cases
+                cld_mask[ilay] = cld_mask_ilay # = random_arr[ilay] >= (FT(1) - cld_frac_ilay)
             end
+            cld_frac_ilayplus1 = cld_frac_ilay
+            cld_mask_ilayplus1 = cld_mask_ilay
         end
     end
     return nothing

diff --git a/src/optics/Optics.jl b/src/optics/Optics.jl
@@ -194,20 +194,31 @@ end
 Computes optical properties for the shortwave problem.
 """
 function compute_optical_props!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
+    op::AbstractOpticalProps,
+    as::AtmosphericState,
     gcol::Int,
     igpt::Int,
-    lkp::Union{AbstractLookUp, Nothing} = nothing,
-    lkp_cld::Union{AbstractLookUp, Nothing} = nothing,
-) where {FT <: AbstractFloat}
+    lkp::LookUpSW,
+    ::Nothing,
+)
     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)
-        end
+        compute_optical_props_kernel!(op, as, ilay, gcol, igpt, lkp)
+    end
+    return nothing
+end
+
+function compute_optical_props!(
+    op::AbstractOpticalProps,
+    as::AtmosphericState,
+    gcol::Int,
+    igpt::Int,
+    lkp::LookUpSW,
+    lkp_cld::Union{LookUpCld, PadeCld},
+)
+    nlay = as.nlay
+    @inbounds for ilay in 1:nlay
+        compute_optical_props_kernel!(op, as, ilay, gcol, igpt, lkp, lkp_cld)
     end
     return nothing
 end

diff --git a/src/optics/OpticsKernels.jl b/src/optics/OpticsKernels.jl
@@ -37,14 +37,14 @@ function compute_optical_props_kernel!(
 end
 
 function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
+    op::AbstractOpticalProps,
+    as::AtmosphericState,
     glay,
     gcol,
     igpt,
-    lkp::LookUpSW{FT},
+    lkp::LookUpSW,
     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
@@ -79,13 +79,13 @@ function compute_optical_props_kernel!(
 end
 
 function compute_optical_props_kernel!(
-    op::AbstractOpticalProps{FT},
-    as::AtmosphericState{FT},
+    op::AbstractOpticalProps,
+    as::AtmosphericState,
     glay,
     gcol,
     igpt::Int,
-    lkp::LookUpSW{FT},
-) where {FT <: AbstractFloat}
+    lkp::LookUpSW,
+)
 
     vmr = as.vmr
     col_dry = as.col_dry[glay, gcol]
@@ -98,7 +98,7 @@ function compute_optical_props_kernel!(
 
     if op isa TwoStream
         @inbounds op.ssa[glay, gcol] = ssa
-        @inbounds op.g[glay, gcol] = FT(0) # initializing asymmetry parameter
+        @inbounds op.g[glay, gcol] = zero(p_lay) # initializing asymmetry parameter
     end
     return nothing
 end
diff --git 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/shortwave2stream.jl b/src/rte/shortwave2stream.jl
@@ -99,8 +99,11 @@ function rte_sw_2stream_solve!(
                 op_col = TwoStreamc(op, gcol)
                 src_sw_col = SourceSW2Strc(src_sw, gcol)
                 bcs_sw_col = SwBCsc(bcs_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)
+                bld_cld_mask && Optics.build_cloud_mask!(
+                    view(as.cld_mask_sw, :, gcol),
+                    view(as.cld_frac, :, gcol),
+                    as.cld_mask_type,
+                )
                 #igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
                 igpt == 1 && set_flux_to_zero!(flux_sw_col)
                 compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
@@ -116,7 +119,6 @@ function rte_sw_2stream_solve!(
     return nothing
 end
 
-
 function rte_sw_2stream_solve!(
     device::ClimaComms.CUDADevice,
     flux::FluxSW{FT},
@@ -169,7 +171,7 @@ function rte_sw_2stream_solve_CUDA!(
         flux_dn_col = flux_col.flux_dn
         @inbounds for igpt in 1:n_gpt
             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
             compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
             # Cell properties: transmittance and reflectance for direct and diffuse radiation