Restructure OneScalar and TwoStream structs

CliMA · Feb 1, 2024 · 195b1c4 · 195b1c4
1 parent cf7211e
commit 195b1c4
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 57 deletions.
diff --git a/src/optics/GasOptics.jl b/src/optics/GasOptics.jl
@@ -144,21 +144,7 @@ Compute optical thickness, single scattering albedo, and asymmetry parameter.
     τ_major = interp3d(jfη..., jftemp..., jfpress..., kmajor, col_mix...) * col_dry
     # computing τ_minor
     lkp_minor = tropo == 1 ? lkp.minor_lower : lkp.minor_upper
-    τ_minor = compute_τ_minor(
-        lkp_minor,
-        tropo,
-        vmr,
-        vmr_h2o,
-        col_dry,
-        p_lay,
-        t_lay,
-        jftemp...,
-        jfη...,
-        igpt,
-        ibnd,
-        glay,
-        gcol,
-    )
+    τ_minor = compute_τ_minor(lkp_minor, vmr, vmr_h2o, col_dry, p_lay, t_lay, jftemp..., jfη..., igpt, ibnd, glay, gcol)
     # τ_Rayleigh is zero for longwave
     τ = τ_major + τ_minor
     return (τ, zero(τ), zero(τ), pfrac) # initializing asymmetry parameter
@@ -188,24 +174,10 @@ end
     τ_major = interp3d(jfη..., jftemp..., jfpress..., kmajor, col_mix...) * col_dry
     # computing τ_minor
     lkp_minor = tropo == 1 ? lkp.minor_lower : lkp.minor_upper
-    τ_minor = compute_τ_minor(
-        lkp_minor,
-        tropo,
-        vmr,
-        vmr_h2o,
-        col_dry,
-        p_lay,
-        t_lay,
-        jftemp...,
-        jfη...,
-        igpt,
-        ibnd,
-        glay,
-        gcol,
-    )
+    τ_minor = compute_τ_minor(lkp_minor, vmr, vmr_h2o, col_dry, p_lay, t_lay, jftemp..., jfη..., igpt, ibnd, glay, gcol)
     # compute τ_Rayleigh
     @inbounds rayleigh_coeff = tropo == 1 ? view(lkp.rayl_lower, :, :, igpt) : view(lkp.rayl_upper, :, :, igpt)
-    τ_ray = compute_τ_rayleigh(rayleigh_coeff, tropo, col_dry, vmr_h2o, jftemp..., jfη..., igpt)
+    τ_ray = compute_τ_rayleigh(rayleigh_coeff, col_dry, vmr_h2o, jftemp..., jfη..., igpt)
     τ = τ_major + τ_minor + τ_ray
     ssa = τ > 0 ? τ_ray / τ : zero(τ) # single scattering albedo
     return (τ, ssa, zero(τ)) # initializing asymmetry parameter
@@ -214,7 +186,6 @@ end
 """
     compute_τ_minor(
         lkp::AbstractLookUp,
-        tropo::Int,
         vmr,
         vmr_h2o::FT,
         col_dry,
@@ -236,7 +207,6 @@ Compute optical thickness contributions from minor gases.
 """
 @inline function compute_τ_minor(
     lkp_minor::LookUpMinor,
-    tropo::Int,
     vmr,
     vmr_h2o::FT,
     col_dry,
@@ -291,7 +261,6 @@ end
 """
     compute_τ_rayleigh(
         lkp::LookUpSW,
-        tropo::Int,
         col_dry::FT,
         vmr_h2o::FT,
         jtemp::Int,
@@ -307,7 +276,6 @@ Compute Rayleigh scattering optical depths for shortwave problem
 """
 @inline compute_τ_rayleigh(
     rayleigh_coeff::AbstractArray{FT, 2},
-    tropo::Int,
     col_dry::FT,
     vmr_h2o::FT,
     jtemp::Int,

diff --git a/src/optics/Optics.jl b/src/optics/Optics.jl
@@ -32,19 +32,22 @@ calculations accounting for extinction and emission
 # Fields
 $(DocStringExtensions.FIELDS)
 """
-struct OneScalar{FTA2D, AD} <: AbstractOpticalProps
-    "Optical Depth"
-    τ::FTA2D
+struct OneScalar{D, V, AD <: AngularDiscretization} <: AbstractOpticalProps
+    "storage for optical thickness"
+    layerdata::D
+    "view into optical depth"
+    τ::V
     "Angular discretization"
     angle_disc::AD
 end
 Adapt.@adapt_structure OneScalar
 
 function OneScalar(::Type{FT}, ncol::Int, nlay::Int, ::Type{DA}) where {FT <: AbstractFloat, DA}
-    τ = DA{FT, 2}(undef, nlay, ncol)
+    layerdata = DA{FT, 3}(undef, 1, nlay, ncol)
+    τ = view(layerdata, 1, :, :)
     ad = AngularDiscretization(FT, DA, 1)
-
-    return OneScalar{typeof(τ), typeof(ad)}(τ, ad)
+    V = typeof(τ)
+    return OneScalar{typeof(layerdata), V, typeof(ad)}(layerdata, τ, ad)
 end
 
 """
@@ -56,21 +59,27 @@ calculations accounting for extinction and emission
 # Fields
 $(DocStringExtensions.FIELDS)
 """
-struct TwoStream{FTA2D} <: AbstractOpticalProps
-    "Optical depth"
-    τ::FTA2D
-    "Single-scattering albedo"
-    ssa::FTA2D
-    "Asymmetry parameter"
-    g::FTA2D
+struct TwoStream{D, V} <: AbstractOpticalProps
+    "storage for optical depth, single scattering albedo and asymmerty parameter"
+    layerdata::D
+    "view into optical depth"
+    τ::V
+    "view into single scattering albedo"
+    ssa::V
+    "view into asymmetry parameter"
+    g::V
 end
 Adapt.@adapt_structure TwoStream
 
 function TwoStream(::Type{FT}, ncol::Int, nlay::Int, ::Type{DA}) where {FT <: AbstractFloat, DA}
-    τ = DA{FT, 2}(zeros(nlay, ncol))
-    ssa = DA{FT, 2}(zeros(nlay, ncol))
-    g = DA{FT, 2}(zeros(nlay, ncol))
-    return TwoStream{typeof(τ)}(τ, ssa, g)
+    layerdata = DA{FT, 3}(zeros(3, nlay, ncol))
+    V = typeof(view(layerdata, 1, :, :))
+    return TwoStream{typeof(layerdata), V}(
+        layerdata,
+        view(layerdata, 1, :, :),
+        view(layerdata, 2, :, :),
+        view(layerdata, 3, :, :),
+    )
 end
 
 """

diff --git a/test/all_sky_utils.jl b/test/all_sky_utils.jl
@@ -140,7 +140,7 @@ function all_sky(
 
     flux_up_lw = Array(slv.flux_lw.flux_up)
     flux_dn_lw = Array(slv.flux_lw.flux_dn)
-    flux_net_lw = flux_up_lw .- flux_dn_lw
+    flux_net_lw = Array(slv.flux_lw.flux_net)
 
     max_err_flux_up_lw = maximum(abs.(flux_up_lw .- comp_flux_up_lw))
     max_err_flux_dn_lw = maximum(abs.(flux_dn_lw .- comp_flux_dn_lw))
@@ -167,7 +167,7 @@ function all_sky(
     flux_up_sw = Array(slv.flux_sw.flux_up)
     flux_dn_sw = Array(slv.flux_sw.flux_dn)
     flux_dn_dir_sw = Array(slv.flux_sw.flux_dn_dir)
-    flux_net_sw = flux_up_sw .- flux_dn_sw
+    flux_net_sw = Array(slv.flux_sw.flux_net)
 
     max_err_flux_up_sw = maximum(abs.(flux_up_sw .- comp_flux_up_sw))
     max_err_flux_dn_sw = maximum(abs.(flux_dn_sw .- comp_flux_dn_sw))
@@ -195,7 +195,9 @@ function all_sky(
 
     @test max_err_flux_up_lw ≤ toler[FT]
     @test max_err_flux_dn_lw ≤ toler[FT]
+    @test max_err_flux_net_lw ≤ toler[FT]
 
     @test max_err_flux_up_sw ≤ toler[FT]
     @test max_err_flux_dn_sw ≤ toler[FT]
+    @test max_err_flux_net_sw ≤ toler[FT]
 end
diff --git a/test/clear_sky_utils.jl b/test/clear_sky_utils.jl
@@ -122,8 +122,7 @@ function clear_sky(
 
     flux_up_lw = Array(slv.flux_lw.flux_up)
     flux_dn_lw = Array(slv.flux_lw.flux_dn)
-
-    flux_net_lw = flux_up_lw .- flux_dn_lw
+    flux_net_lw = Array(slv.flux_lw.flux_net)
 
     max_err_flux_up_lw = maximum(abs.(flux_up_lw .- comp_flux_up_lw))
     max_err_flux_dn_lw = maximum(abs.(flux_dn_lw .- comp_flux_dn_lw))
@@ -160,7 +159,7 @@ function clear_sky(
 
     flux_up_sw = Array(slv.flux_sw.flux_up)
     flux_dn_sw = Array(slv.flux_sw.flux_dn)
-    flux_net_sw = flux_up_sw .- flux_dn_sw
+    flux_net_sw = Array(slv.flux_sw.flux_net)
 
     for i in 1:ncol
         if usecol[i] == 0
@@ -197,6 +196,8 @@ function clear_sky(
 
     @test max_err_flux_up_lw ≤ toler_lw[FT]
     @test max_err_flux_dn_lw ≤ toler_lw[FT]
+    @test max_err_flux_net_lw ≤ toler_lw[FT]
     @test max_err_flux_up_sw ≤ toler_sw[FT]
     @test max_err_flux_dn_sw ≤ toler_sw[FT]
+    @test max_err_flux_net_sw ≤ toler_sw[FT]
 end