Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor build_cloud_mask! #411

Merged
merged 1 commit into from
Dec 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
5 changes: 0 additions & 5 deletions src/optics/AtmosphericStates.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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"
Expand Down
106 changes: 56 additions & 50 deletions src/optics/CloudOptics.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
sriharshakandala marked this conversation as resolved.
Show resolved Hide resolved
# 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
Expand Down
9 changes: 3 additions & 6 deletions src/rte/longwave2stream.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
9 changes: 3 additions & 6 deletions src/rte/shortwave2stream.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
5 changes: 0 additions & 5 deletions test/read_all_sky.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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),
}(
Expand All @@ -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(),
Expand Down
5 changes: 0 additions & 5 deletions test/read_rfmip_clear_sky.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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),
}(
Expand All @@ -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,
Expand Down