Refactor build_cloud_mask!. (#411)

CliMA · Dec 20, 2023 · 0d93d70 · 0d93d70
1 parent 4a09cae
commit 0d93d70
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 77 deletions.
diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl
@@ -44,7 +44,6 @@ 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},
 } <: AbstractAtmosphericState{FT, FTA1D}
@@ -76,10 +75,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"

diff --git a/src/optics/CloudOptics.jl b/src/optics/CloudOptics.jl
@@ -216,65 +216,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)
-
-    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)
-
-    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
-    end
+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 # 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
+    if start > 0
+        finish = _get_finish(cld_frac) # last cloudy layer
         # 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
-        # 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])
+        _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
-                local_cld_mask[ilay] = false
+                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
+
+function _get_finish(cld_frac)
+    @inbounds for ilay in reverse(eachindex(cld_frac))
+        cld_frac[ilay] > 0 && return ilay
+    end
+    return 0
+end
+
+function _get_start(cld_frac)
+    @inbounds for ilay in eachindex(cld_frac)
+        cld_frac[ilay] > 0 && return ilay
+    end
+    return 0
+end
+
+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
+        nlay = length(cld_mask)
+        for ilay in (finish + 1):nlay
+            @inbounds cld_mask[ilay] = false
         end
     end
     return nothing

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
@@ -60,11 +60,8 @@ function rte_sw_2stream_solve!(
                 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,
+                        view(as.cld_mask_sw, :, gcol),
+                        view(as.cld_frac, :, gcol),
                         as.cld_mask_type,
                     )
                     compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
@@ -100,7 +97,7 @@ function rte_sw_2stream_solve_CUDA!(
         @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
             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

diff --git a/test/read_all_sky.jl b/test/read_all_sky.jl
@@ -91,8 +91,6 @@ function setup_allsky_as(
     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(),

diff --git a/test/read_rfmip_clear_sky.jl b/test/read_rfmip_clear_sky.jl
@@ -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,