Fix flux_dn_dir not updated in all sky two stream

Prior to this commit, `flux_dn_dir` was not set in the two stream case for non-gray radiation, leading to garbage values (whatever was in the memory at init). Now, the variable is accumulated over for ever g-point.
CliMA · Oct 7, 2024 · f225d73 · f225d73
1 parent 39c4891
commit f225d73
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 2 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -4,6 +4,16 @@ RRTMGP.jl Release Notes
 main
 ------
 
+### Bug fixes
+
+#### Fix  `flux_dn_dir` for non-gray radiation
+
+Prior to this release, `flux_dn_dir` was not correctly set in the two stream
+case for non-gray radiation, leading to incorrect values (whatever was in the
+memory at initialization). Now, the variable is correctly accumulated over for
+every g-point. Note, however, that only the value at the surface (`[1, :]`) is
+updated. PR [#550](https://github.com/CliMA/RRTMGP.jl/pull/550).
+
 v0.19.0
 -----
 - Compute aero_mask internally and store the array.

diff --git a/ext/cuda/rte_shortwave_1scalar.jl b/ext/cuda/rte_shortwave_1scalar.jl
@@ -70,8 +70,10 @@ function rte_sw_noscat_solve_CUDA!(
         flux_up_sw = flux_sw.flux_up
         flux_dn_sw = flux_sw.flux_dn
         flux_net_sw = flux_sw.flux_net
+        flux_dn_dir_sw = flux_sw.flux_dn_dir
         flux_up = flux.flux_up
         flux_dn = flux.flux_dn
+        flux_dn_dir = flux.flux_dn_dir
         μ₀ = bcs_sw.cos_zenith[gcol]
         @inbounds begin
             for igpt in 1:n_gpt
@@ -81,10 +83,12 @@ function rte_sw_noscat_solve_CUDA!(
                 if igpt == 1
                     map!(x -> x, view(flux_up_sw, :, gcol), view(flux_up, :, gcol))
                     map!(x -> x, view(flux_dn_sw, :, gcol), view(flux_dn, :, gcol))
+                    map!(x -> x, view(flux_dn_dir_sw, :, gcol), view(flux_dn_dir, :, gcol))
                 else
                     for ilev in 1:nlev
                         flux_up_sw[ilev, gcol] += flux_up[ilev, gcol]
                         flux_dn_sw[ilev, gcol] += flux_dn[ilev, gcol]
+                        flux_dn_dir_sw[ilev, gcol] += flux_dn_dir[ilev, gcol]
                     end
                 end
             end

diff --git a/ext/cuda/rte_shortwave_2stream.jl b/ext/cuda/rte_shortwave_2stream.jl
@@ -96,9 +96,11 @@ function rte_sw_2stream_solve_CUDA!(
     if gcol ≤ ncol
         flux_up_sw = flux_sw.flux_up
         flux_dn_sw = flux_sw.flux_dn
+        flux_dn_dir_sw = flux_sw.flux_dn_dir
         flux_net_sw = flux_sw.flux_net
         flux_up = flux.flux_up
         flux_dn = flux.flux_dn
+        flux_dn_dir = flux.flux_dn_dir
         FT = eltype(flux_up)
         (; cloud_state, aerosol_state) = as
         μ₀ = bcs_sw.cos_zenith[gcol]
@@ -127,10 +129,12 @@ function rte_sw_2stream_solve_CUDA!(
                 if igpt == 1
                     map!(x -> x, view(flux_up_sw, :, gcol), view(flux_up, :, gcol))
                     map!(x -> x, view(flux_dn_sw, :, gcol), view(flux_dn, :, gcol))
+                    map!(x -> x, view(flux_dn_dir_sw, :, gcol), view(flux_dn_dir, :, gcol))
                 else
                     for ilev in 1:nlev
                         flux_up_sw[ilev, gcol] += flux_up[ilev, gcol]
                         flux_dn_sw[ilev, gcol] += flux_dn[ilev, gcol]
+                        flux_dn_dir_sw[ilev, gcol] += flux_dn_dir[ilev, gcol]
                     end
                 end
             end

diff --git a/src/rte/shortwave1scalar.jl b/src/rte/shortwave1scalar.jl
@@ -38,8 +38,9 @@ function rte_sw_noscat_solve!(
     n_gpt = length(lookup_sw.solar_src_scaled)
     flux_up_sw = flux_sw.flux_up
     flux_dn_sw = flux_sw.flux_dn
+    flux_dn_dir_sw = flux_sw.flux_dn_dir
     flux_net_sw = flux_sw.flux_net
-    (; flux_up, flux_dn) = flux
+    (; flux_up, flux_dn, flux_dn_dir) = flux
     cos_zenith = bcs_sw.cos_zenith
     @inbounds begin
         for igpt in 1:n_gpt
@@ -51,10 +52,12 @@ function rte_sw_noscat_solve!(
                     if igpt == 1
                         map!(x -> x, view(flux_up_sw, :, gcol), view(flux_up, :, gcol))
                         map!(x -> x, view(flux_dn_sw, :, gcol), view(flux_dn, :, gcol))
+                        map!(x -> x, view(flux_dn_dir_sw, :, gcol), view(flux_dn_dir, :, gcol))
                     else
                         for ilev in 1:nlev
                             flux_up_sw[ilev, gcol] += flux_up[ilev, gcol]
                             flux_dn_sw[ilev, gcol] += flux_dn[ilev, gcol]
+                            flux_dn_dir_sw[ilev, gcol] += flux_dn_dir[ilev, gcol]
                         end
                     end
                 else

diff --git a/src/rte/shortwave2stream.jl b/src/rte/shortwave2stream.jl
@@ -53,8 +53,9 @@ function rte_sw_2stream_solve!(
         bld_cld_mask = cloud_state isa CloudState
         flux_up_sw = flux_sw.flux_up
         flux_dn_sw = flux_sw.flux_dn
+        flux_dn_dir_sw = flux_sw.flux_dn_dir
         flux_net_sw = flux_sw.flux_net
-        (; flux_up, flux_dn) = flux
+        (; flux_up, flux_dn, flux_dn_dir) = flux
         cos_zenith = bcs_sw.cos_zenith
         FT = eltype(flux_up)
         if aerosol_state isa AerosolState
@@ -82,10 +83,12 @@ function rte_sw_2stream_solve!(
                     if igpt == 1
                         map!(x -> x, view(flux_up_sw, :, gcol), view(flux_up, :, gcol))
                         map!(x -> x, view(flux_dn_sw, :, gcol), view(flux_dn, :, gcol))
+                        map!(x -> x, view(flux_dn_dir_sw, :, gcol), view(flux_dn_dir, :, gcol))
                     else
                         for ilev in 1:nlev
                             @inbounds flux_up_sw[ilev, gcol] += flux_up[ilev, gcol]
                             @inbounds flux_dn_sw[ilev, gcol] += flux_dn[ilev, gcol]
+                            @inbounds flux_dn_dir_sw[ilev, gcol] += flux_dn_dir[ilev, gcol]
                         end
                     end
                 else