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.

Sign up for GitHub

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

Jump to bottom

Restructure data layout in AtmosphericStates #455

Merged
merged 1 commit into from
Feb 23, 2024
Merged

Restructure data layout in AtmosphericStates #455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 20 additions & 10 deletions src/optics/AtmosphericStates.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,23 +38,19 @@ Atmospheric conditions, used to compute optical properties.
# Fields
$(DocStringExtensions.FIELDS)
"""
struct AtmosphericState{FTA1D, FTA1DN, FTA2D, VMR, CLD} <: AbstractAtmosphericState
struct AtmosphericState{FTA1D, FTA1DN, FTA2D, D, VMR, CLD} <: AbstractAtmosphericState
"longitude, in degrees (`ncol`), optional"
lon::FTA1DN
"latitude, in degrees (`ncol`), optional"
lat::FTA1DN
"Layer pressures `[Pa, mb]`; `(nlay,ncol)`"
p_lay::FTA2D
"storage for col_dry [`molecules per cm^2 of dry air`], play `[Pa, mb]`, tlay `[K]`; `(3, nlay, ncol)`"
layerdata::D
"Level pressures `[Pa, mb]`; `(nlay+1,ncol)`"
p_lev::FTA2D
"Layer temperatures `[K]`; `(nlay,ncol)`"
t_lay::FTA2D
"Level temperatures `[K]`; `(nlay+1,ncol)`"
t_lev::FTA2D
"Surface temperatures `[K]`; `(ncol)`"
t_sfc::FTA1D
"Number of molecules per cm^2 of dry air `(nlay, ncol)`"
col_dry::FTA2D
"volume mixing ratio of all relevant gases"
vmr::VMR
"cloud state"
Expand All @@ -63,11 +59,25 @@ end
Adapt.@adapt_structure AtmosphericState

# Number of layers
@inline get_nlay(as::AtmosphericState) = size(as.p_lay, 1)
@inline get_nlay(as::AtmosphericState) = size(as.layerdata, 2)
# Number of columns
@inline get_ncol(as::AtmosphericState) = size(as.p_lay, 2)
@inline get_ncol(as::AtmosphericState) = size(as.layerdata, 3)
# Number of layers and columns
@inline get_dims(as::AtmosphericState) = size(as.p_lay)
@inline get_dims(as::AtmosphericState) = size(as.layerdata, 2), size(as.layerdata, 3)

@inline getview_layerdata(as::AtmosphericState, gcol) = @inbounds view(as.layerdata, :, :, gcol)

# view of column amounts of dry air [molecules per cm^2 of dry air]
@inline getview_col_dry(as::AtmosphericState) = @inbounds view(as.layerdata, 1, :, :)
@inline getview_col_dry(as::AtmosphericState, gcol) = @inbounds view(as.layerdata, 1, :, gcol)

# view of layer pressures [Pa, mb]
@inline getview_p_lay(as::AtmosphericState) = @inbounds view(as.layerdata, 2, :, :)
@inline getview_p_lay(as::AtmosphericState, gcol) = @inbounds view(as.layerdata, 2, :, gcol)

# view of layer temperatures [K]
@inline getview_t_lay(as::AtmosphericState) = @inbounds view(as.layerdata, 3, :, :)
@inline getview_t_lay(as::AtmosphericState, gcol) = @inbounds view(as.layerdata, 3, :, gcol)

"""
CloudState{CD, CF, CM, CMT}
Expand Down
54 changes: 20 additions & 34 deletions src/optics/Optics.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,17 +169,13 @@ Computes optical properties for the longwave problem.
t_sfc = as.t_sfc[gcol]
ibnd = lkp.band_data.major_gpt2bnd[igpt]
totplnk = view(lkp.planck.tot_planck, :, ibnd)
col_dry_col = view(as.col_dry, :, gcol)
p_lay_col = view(as.p_lay, :, gcol)
t_lay_col = view(as.t_lay, :, gcol)
as_layerdata = AtmosphericStates.getview_layerdata(as, gcol)
t_lev_col = view(as.t_lev, :, gcol)
τ = view(op.τ, :, gcol)

t_lev_dec = t_lev_col[1]
for glay in 1:nlay
col_dry = col_dry_col[glay]
p_lay = p_lay_col[glay]
t_lay = t_lay_col[glay]
col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay]
# compute gas optics
τ[glay], _, _, planckfrac = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
# compute longwave source terms
Expand Down Expand Up @@ -214,9 +210,7 @@ end
t_sfc = as.t_sfc[gcol]
ibnd = lkp.band_data.major_gpt2bnd[igpt]
totplnk = view(lkp.planck.tot_planck, :, ibnd)
col_dry_col = view(as.col_dry, :, gcol)
p_lay_col = view(as.p_lay, :, gcol)
t_lay_col = view(as.t_lay, :, gcol)
as_layerdata = AtmosphericStates.getview_layerdata(as, gcol)
t_lev_col = view(as.t_lev, :, gcol)
τ = view(op.τ, :, gcol)
ssa = view(op.ssa, :, gcol)
Expand All @@ -229,9 +223,7 @@ end
@inbounds begin
t_lev_dec = t_lev_col[1]
for glay in 1:nlay
col_dry = col_dry_col[glay]
p_lay = p_lay_col[glay]
t_lay = t_lay_col[glay]
col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay]
# compute gas optics
τ[glay], ssa[glay], g[glay], planckfrac =
compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
Expand Down Expand Up @@ -300,17 +292,14 @@ Computes optical properties for the shortwave problem.
)
nlay = AtmosphericStates.get_nlay(as)
(; vmr) = as
@inbounds ibnd = lkp.band_data.major_gpt2bnd[igpt]
@inbounds t_sfc = as.t_sfc[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)
τ = view(op.τ, :, gcol)

@inbounds begin
ibnd = lkp.band_data.major_gpt2bnd[igpt]
t_sfc = as.t_sfc[gcol]
as_layerdata = AtmosphericStates.getview_layerdata(as, gcol)
τ = view(op.τ, :, gcol)
end
@inbounds for glay in 1:nlay
col_dry = col_dry_col[glay]
p_lay = p_lay_col[glay]
t_lay = t_lay_col[glay]
col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay]
# compute gas optics
τ[glay], _, _ = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
end
Expand All @@ -327,19 +316,16 @@ end
)
nlay = AtmosphericStates.get_nlay(as)
(; vmr) = as
@inbounds ibnd = lkp.band_data.major_gpt2bnd[igpt]
@inbounds t_sfc = as.t_sfc[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)
τ = view(op.τ, :, gcol)
ssa = view(op.ssa, :, gcol)
g = view(op.g, :, gcol)

@inbounds begin
ibnd = lkp.band_data.major_gpt2bnd[igpt]
t_sfc = as.t_sfc[gcol]
as_layerdata = AtmosphericStates.getview_layerdata(as, gcol)
τ = view(op.τ, :, gcol)
ssa = view(op.ssa, :, gcol)
g = view(op.g, :, gcol)
end
@inbounds for glay in 1:nlay
col_dry = col_dry_col[glay]
p_lay = p_lay_col[glay]
t_lay = t_lay_col[glay]
col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay]
# compute gas optics
τ[glay], ssa[glay], g[glay] = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol)
end
Expand Down
7 changes: 6 additions & 1 deletion test/read_all_sky.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,11 @@ function setup_allsky_as(

compute_col_gas!(context, p_lev, col_dry, param_set, vmr_h2o, lat) # the example skips lat based gravity calculation

layerdata = similar(p_lay, 3, nlay, ncol)
layerdata[1, :, :] .= col_dry
layerdata[2, :, :] .= p_lay
layerdata[3, :, :] .= t_lay

t_sfc = DA(t_sfc)

cld_frac = DA(cld_frac)
Expand All @@ -155,7 +160,7 @@ function setup_allsky_as(
ice_rgh,
)
return (
AtmosphericState(lon, lat, p_lay, p_lev, t_lay, t_lev, t_sfc, col_dry, vmr, cloud_state),
AtmosphericState(lon, lat, layerdata, p_lev, t_lev, t_sfc, vmr, cloud_state),
sfc_emis,
sfc_alb_direct,
sfc_alb_diffuse,
Expand Down
7 changes: 6 additions & 1 deletion test/read_rfmip_clear_sky.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,10 +112,15 @@ function setup_rfmip_as(
# FORTRAN RRTMGP test case.
compute_col_gas!(context, p_lev, col_dry, param_set, vmr_h2o, lat) # the example skips lat based gravity calculation

layerdata = similar(p_lay, 3, nlay, ncol)
layerdata[1, :, :] .= col_dry
layerdata[2, :, :] .= p_lay
layerdata[3, :, :] .= t_lay

vmr = VMR(vmr_h2o, vmr_o3, FTA1D(vmrat))
#------------------
return (
AtmosphericState(lon, lat, p_lay, p_lev, t_lay, t_lev, t_sfc, col_dry, vmr, nothing),
AtmosphericState(lon, lat, layerdata, p_lev, t_lev, t_sfc, vmr, nothing),
sfc_emis,
sfc_alb,
cos_zenith,
Expand Down
Loading