Merge pull request #61 from MennoVeerman/main

updates and bug fixes to bw implementation

Author: Chiil

data/mie_lut_broadband.nc

@@ -145,12 +145,12 @@ class Gas_optics_rrtmgp_rt : public Gas_optics_rt
                 const Array_gpu<Float,2>& play,
                 const Array_gpu<Float,2>& plev,
                 const Array_gpu<Float,2>& tlay,
+                const Array_gpu<Float,2>& tlev,
                 const Array_gpu<Float,1>& tsfc,
                 const Gas_concs_gpu& gas_desc,
                 std::unique_ptr<Optical_props_arry_rt>& optical_props,
                 Source_func_lw_rt& sources,
-                const Array_gpu<Float,2>& col_dry,
-                const Array_gpu<Float,2>& tlev);
+                const Array_gpu<Float,2>& col_dry);
 
         // shortwave variant
         void gas_optics(
@@ -219,7 +219,7 @@ class Gas_optics_rrtmgp_rt : public Gas_optics_rt
         Array<Float,4> krayl;
 
         int idx_h2o;
-        
+
         Array_gpu<Float,1> solar_source_g;
         Array_gpu<Float,2> totplnk_gpu;
         Array_gpu<Float,4> planck_frac_gpu;
@@ -294,7 +294,7 @@ class Gas_optics_rrtmgp_rt : public Gas_optics_rt
                 const Float md_index, const Float sb_index);
 
         void compute_gas_taus(
-                const int col_s, const int ncol_block, const int ncol, const int nlay, 
+                const int col_s, const int ncol_block, const int ncol, const int nlay,
                 const int ngpt, const int nband, const int igpt,
                 const Array_gpu<Float,2>& play,
                 const Array_gpu<Float,2>& plev,
@@ -313,7 +313,7 @@ class Gas_optics_rrtmgp_rt : public Gas_optics_rt
                 std::unique_ptr<Optical_props_arry_rt>& optical_props);
 
         void source(
-                const int ncol, const int nlay, const int nband, const int ngpt, const int igpt,
+                const int col_s, const int ncol_sub, const int ncol, const int nlay, const int nband, const int ngpt, const int igpt,
                 const Array_gpu<Float,2>& play, const Array_gpu<Float,2>& plev,
                 const Array_gpu<Float,2>& tlay, const Array_gpu<Float,1>& tsfc,
                 const Array_gpu<int,2>& jtemp, const Array_gpu<int,2>& jpress,

data/mie_lut_visualisation.nc

@@ -61,12 +61,12 @@ class Gas_optics_rt : public Optical_props_rt
                 const Array_gpu<Float,2>& play,
                 const Array_gpu<Float,2>& plev,
                 const Array_gpu<Float,2>& tlay,
+                const Array_gpu<Float,2>& tlev,
                 const Array_gpu<Float,1>& tsfc,
                 const Gas_concs_gpu& gas_desc,
                 std::unique_ptr<Optical_props_arry_rt>& optical_props,
                 Source_func_lw_rt& sources,
-                const Array_gpu<Float,2>& col_dry,
-                const Array_gpu<Float,2>& tlev) = 0;
+                const Array_gpu<Float,2>& col_dry) = 0;
 
         // Shortwave variant.
         virtual void gas_optics(
@@ -80,7 +80,7 @@ class Gas_optics_rt : public Optical_props_rt
                 const Array_gpu<Float,2>& col_dry) = 0;
 
        virtual Float get_tsi() const = 0;
-       
+
        virtual Float band_source(const int gpt_start, const int gpt_end) const = 0;
 };
 #endif

include_rt/Gas_optics_rrtmgp_rt.h

@@ -24,7 +24,7 @@
 
 #ifndef SOURCE_FUNCTIONS_RT_H
 #define SOURCE_FUNCTIONS_RT_H
-#include "Optical_props_rt.h" 
+#include "Optical_props_rt.h"
 
 template<typename, int> class Array_gpu;
 
@@ -48,21 +48,18 @@ class Source_func_lw_rt : public Optical_props_rt
         Array_gpu<Float,1>& get_sfc_source()     { return sfc_source;     }
         Array_gpu<Float,1>& get_sfc_source_jac() { return sfc_source_jac; }
         Array_gpu<Float,2>& get_lay_source()     { return lay_source;     }
-        Array_gpu<Float,2>& get_lev_source_inc() { return lev_source_inc; }
-        Array_gpu<Float,2>& get_lev_source_dec() { return lev_source_dec; }
+        Array_gpu<Float,2>& get_lev_source()     { return lev_source; }
 
         const Array_gpu<Float,1>& get_sfc_source()     const { return sfc_source;     }
         const Array_gpu<Float,1>& get_sfc_source_jac() const { return sfc_source_jac; }
         const Array_gpu<Float,2>& get_lay_source()     const { return lay_source;     }
-        const Array_gpu<Float,2>& get_lev_source_inc() const { return lev_source_inc; }
-        const Array_gpu<Float,2>& get_lev_source_dec() const { return lev_source_dec; }
+        const Array_gpu<Float,2>& get_lev_source()     const { return lev_source; }
 
     private:
         Array_gpu<Float,1> sfc_source;
         Array_gpu<Float,1> sfc_source_jac;
         Array_gpu<Float,2> lay_source;
-        Array_gpu<Float,2> lev_source_inc;
-        Array_gpu<Float,2> lev_source_dec;
+        Array_gpu<Float,2> lev_source;
 };
 
 #endif

include_rt/Gas_optics_rt.h

@@ -87,104 +87,62 @@ namespace Raytracer_functions
     }
 
     __device__
-    inline Float mie_sample_angle(const Float* mie_cdf, const Float* mie_lut, const Float random_number, const Float r_eff, const int n_mie)
+    inline int find_index(const float* mie_cdf, const int size, const float random_number)
     {
-        // interpolation over effective radius. Currently, r_eff should range between 2.5 and 21.5 (similar to RRTMGP) OR
-        // be exactly 100 micrometer for optical effects such as rainbows
-        const int r_idx = (r_eff == Float(100.)) ? 20 : min(max(int(r_eff-2.5), 0), 18);
-        const Float r_rest = fmod(r_eff-Float(2.5),Float(1.));
+        int left = 0;
+        int right = size - 1;
 
-        int i = 0;
-        while (random_number < mie_cdf[i])
-        {
-            ++i;
-        }
+        while (left < right) {
+            int mid = left + (right - left) / 2;
 
-        // sampled scattering angle
-        Float ang;
-        if (r_idx < 20)
-        {
-            if (i==0)
-            {
-                const Float ang_lwr = mie_lut[r_idx*n_mie]*(1-r_rest);
-                const Float ang_upr = mie_lut[(r_idx+1)*n_mie]*r_rest;
-                ang = ang_lwr + ang_upr;
-            }
-            else
-            {
-                const int midx_lwr = r_idx*n_mie;
-                const int midx_upr = (r_idx+1)*n_mie;
-                const Float dr = abs(mie_cdf[i] - mie_cdf[i-1]);
-
-                const Float ang_lwr = (abs(random_number - mie_cdf[i])*mie_lut[(i-1)+midx_lwr] + abs(mie_cdf[i-1]-random_number)*mie_lut[i+midx_lwr]) / dr;
-                const Float ang_upr = (abs(random_number - mie_cdf[i])*mie_lut[(i-1)+midx_upr] + abs(mie_cdf[i-1]-random_number)*mie_lut[i+midx_upr]) / dr;
-                ang = ang_lwr * (1-r_rest) + ang_upr * r_rest;
+            if (random_number >= mie_cdf[mid]) {
+                right = mid;
+            } else {
+                left = mid + 1;
             }
         }
-        else
-        {
-            if (i==0)
-            {
-                ang = mie_lut[r_idx*n_mie];
-            }
-            else
-            {
-                const int midx = r_idx*n_mie;
-                const Float dr = abs(mie_cdf[i] - mie_cdf[i-1]);
 
-                ang = (abs(random_number - mie_cdf[i])*mie_lut[(i-1)+midx] + abs(mie_cdf[i-1]-random_number)*mie_lut[i+midx]) / dr;
-            }
-        }
+        return left - 1;
+    }
+
+    __device__
+    inline Float mie_sample_angle(const Float* mie_cdf, const Float* mie_lut, const Float random_number, const Float r_eff, const int n_mie)
+    {
+        // interpolation over effective radius. Currently, r_eff should range between 2.5 and 21.5 (similar to RRTMGP)
+        const int r_idx = min(max(int(r_eff-2.5), 0), 18);
+        const Float r_rest = fmod(r_eff-Float(2.5),Float(1.));
+
+        const int i = min(max(0, find_index(mie_cdf, n_mie, random_number)), n_mie - 2);
+
+        const int midx_lwr = r_idx*n_mie;
+        const int midx_upr = (r_idx+1)*n_mie;
+        const Float dr = abs(mie_cdf[i+1] - mie_cdf[i]);
+
+        const Float ang_lwr = (abs(random_number - mie_cdf[i+1])*mie_lut[(i)+midx_lwr] + abs(mie_cdf[i]-random_number)*mie_lut[i+midx_lwr+1]) / dr;
+        const Float ang_upr = (abs(random_number - mie_cdf[i+1])*mie_lut[(i)+midx_upr] + abs(mie_cdf[i]-random_number)*mie_lut[i+midx_upr+1]) / dr;
+        const Float ang = ang_lwr * (1-r_rest) + ang_upr * r_rest;
         return ang;
     }
 
     __device__
     inline Float mie_interpolate_phase_table(const Float* mie_phase, const Float* mie_lut, const Float scat_ang, const Float r_eff, const int n_mie)
     {
-        // interpolation over effective radius. Currently, r_eff should range between 2.5 and 21.5 (similar to RRTMGP) OR
-        // be exactly 100 micrometer for optical effects such as rainbows
-        const int r_idx = (r_eff == Float(100.)) ? 20 : min(max(int(r_eff-2.5), 0), 18);
+        // interpolation over effective radius. Currently, r_eff should range between 2.5 and 21.5 (similar to RRTMGP)
+        const int r_idx = min(max(int(r_eff-2.5), 0), 18);
         const Float r_rest = fmod(r_eff-Float(2.5),Float(1.));
 
         // interpolation between 1800 equally spaced scattering angles between 0 and PI (both inclusive).
-        const Float d_pi = Float(1.74629942e-03);
-        const int i = min(max(0, int(1800-(scat_ang/d_pi+1))), 1798);
+        constexpr Float d_pi = Float(1.74629942e-03);
+        const int i = min(max(0, int(scat_ang/d_pi)), 1798);
 
-        // probability (of scattering at angle scat_ang)
-        Float prob;
-        if (r_idx < 20)
-        {
-            if (i==0)
-            {
-                const Float prob_lwr = mie_lut[r_idx*n_mie]*(1-r_rest);
-                const Float prob_upr = mie_lut[(r_idx+1)*n_mie]*r_rest;
-                prob = prob_lwr + prob_upr;
-            }
-            else
-            {
-                const int midx_lwr = r_idx*n_mie;
-                const int midx_upr = (r_idx+1)*n_mie;
-                const Float dr = abs(mie_phase[i] - mie_phase[i-1]);
-
-                const Float prob_lwr = (abs(scat_ang - mie_phase[i])*mie_lut[(i-1)+midx_lwr] + abs(mie_phase[i-1]-scat_ang)*mie_lut[i+midx_lwr]) / dr;
-                const Float prob_upr = (abs(scat_ang - mie_phase[i])*mie_lut[(i-1)+midx_upr] + abs(mie_phase[i-1]-scat_ang)*mie_lut[i+midx_upr]) / dr;
-                prob = prob_lwr * (1-r_rest) + prob_upr * r_rest;
-            }
-        }
-        else
-        {
-            if (i==0)
-            {
-                prob = mie_lut[r_idx*n_mie];
-            }
-            else
-            {
-                const int midx = r_idx*n_mie;
-                const Float dr = abs(mie_phase[i] - mie_phase[i-1]);
+        const int midx_lwr = r_idx*n_mie;
+        const int midx_upr = (r_idx+1)*n_mie;
+        const Float dr = abs(mie_phase[i+1] - mie_phase[i]);
+
+        const Float prob_lwr = (abs(scat_ang - mie_phase[i+1])*mie_lut[(i)+midx_lwr] + abs(mie_phase[i]-scat_ang)*mie_lut[i+1+midx_lwr]) / dr;
+        const Float prob_upr = (abs(scat_ang - mie_phase[i+1])*mie_lut[(i)+midx_upr] + abs(mie_phase[i]-scat_ang)*mie_lut[i+1+midx_upr]) / dr;
+        const Float prob = prob_lwr * (1-r_rest) + prob_upr * r_rest;
 
-                prob = (abs(scat_ang - mie_phase[i])*mie_lut[(i-1)+midx] + abs(mie_phase[i-1]-scat_ang)*mie_lut[i+midx]) / dr;
-            }
-        }
         return prob;
     }
 

include_rt/Source_functions_rt.h

@@ -33,10 +33,10 @@ namespace Gas_optics_rrtmgp_kernels_cuda_rt
     void reorder123x321(const int ni, const int nj, const int nk,
             const Float* arr_in, Float* arr_out);
 
-    
+
     void reorder12x21(const int ni, const int nj, const Float* arr_in, Float* arr_out);
 
-    
+
     void zero_array(const int ni, const int nj, const int nk, const int nn, Float* arr);
     void zero_array(const int ni, const int nj, const int nk, Float* arr);
     void zero_array(const int ni, const int nj, Float* arr);
@@ -67,7 +67,7 @@ namespace Gas_optics_rrtmgp_kernels_cuda_rt
 
     void combine_abs_and_rayleigh(
             const int col_s, const int ncol_sub, const int ncol, const int nlay,
-            const Float* tau_local, const Float* tau_rayleigh,
+            const Float* tau_rayleigh,
             Float* tau, Float* ssa, Float* g);
 
 
@@ -84,7 +84,7 @@ namespace Gas_optics_rrtmgp_kernels_cuda_rt
 
 
     void compute_tau_absorption(
-            const int col_s, const int ncol_sub, const int ncol, const int nlay, const int nband, 
+            const int col_s, const int ncol_sub, const int ncol, const int nlay, const int nband,
             const int ngpt, const int igpt,
             const int ngas, const int nflav, const int neta, const int npres, const int ntemp,
             const int nminorlower, const int nminorklower,
@@ -117,8 +117,8 @@ namespace Gas_optics_rrtmgp_kernels_cuda_rt
             Float* tau);
 
 
-    void Planck_source(
-            const int ncol, const int nlay, const int nbnd, const int ngpt, const int igpt,
+    void compute_planck_source(
+            const int col_s, const int ncol_sub, const int ncol, const int nlay, const int nbnd, const int ngpt, const int igpt,
             const int nflav, const int neta, const int npres, const int ntemp,
             const int nPlanckTemp,
             const Float* tlay,
@@ -137,8 +137,7 @@ namespace Gas_optics_rrtmgp_kernels_cuda_rt
             const Float* totplnk,
             Float* sfc_src,
             Float* lay_src,
-            Float* lev_src_inc,
-            Float* lev_src_dec,
+            Float* lev_src,
             Float* sfc_src_jac);
 }
 #endif

include_rt/raytracer_functions.h

@@ -17,8 +17,8 @@ using namespace Raytracer_functions;
 
 
 #ifdef RTE_USE_SP
-constexpr int bw_kernel_block= 512;
-constexpr int bw_kernel_grid = 1024;
+constexpr int bw_kernel_block= 128;
+constexpr int bw_kernel_grid = 2048;
 #else
 constexpr int bw_kernel_block = 256;
 constexpr int bw_kernel_grid = 256;

include_rt_kernels/gas_optics_rrtmgp_kernels_cuda_rt.h

@@ -32,25 +32,25 @@
 namespace Rte_solver_kernels_cuda_rt
 {
     void apply_BC(
-            const int ncol, const int nlay, const int ngpt, const Bool top_at_1,
+            const int ncol, const int nlay, const Bool top_at_1,
             const Float* inc_flux_dir, const Float* mu0, Float* gpt_flux_dir);
 
-    void apply_BC(const int ncol, const int nlay, const int ngpt, const Bool top_at_1, Float* gpt_flux_dn);
+    void apply_BC(const int ncol, const int nlay, const Bool top_at_1, Float* gpt_flux_dn);
 
-    void apply_BC(const int ncol, const int nlay, const int ngpt, const Bool top_at_1, const Float* inc_flux_dif, Float* gpt_flux_dn);
-    
-    void apply_BC(const int ncol, const int nlay, const int ngpt, const Bool top_at_1, const Float inc_flux, Float* gpt_flux_dn);
+    void apply_BC(const int ncol, const int nlay, const Bool top_at_1, const Float* inc_flux_dif, Float* gpt_flux_dn);
+
+    void apply_BC(const int ncol, const int nlay, const Bool top_at_1, const Float inc_flux, Float* gpt_flux_dn);
 
     void sw_solver_2stream(
-            const int ncol, const int nlay, const int ngpt, const Bool top_at_1,
+            const int ncol, const int nlay, const Bool top_at_1,
             const Float* tau, const Float* ssa, const Float* g,
             const Float* mu0, const Float* sfc_alb_dir, const Float* sfc_alb_dif,
             Float* flux_up, Float* flux_dn, Float* flux_dir);
 
     void lw_solver_noscat_gaussquad(
-            const int ncol, const int nlay, const int ngpt, const Bool top_at_1, const int nmus,
+            const int ncol, const int nlay, const Bool top_at_1, const int nmus,
             const Float* ds, const Float* weights, const Float* tau, const Float* lay_source,
-            const Float* lev_source_inc, const Float* lev_source_dec, const Float* sfc_emis,
+            const Float* lev_source, const Float* sfc_emis,
             const Float* sfc_src, Float* flux_up, Float* flux_dn,
             const Float* sfc_src_jac, Float* flux_up_jac);
 }

include_rt_kernels/raytracer_kernels_bw.h

@@ -45,18 +45,26 @@ class Radiation_solver_longwave
         #ifdef __CUDACC__
         void solve_gpu(
                 const bool switch_fluxes,
+                const bool switch_raytracing,
                 const bool switch_cloud_optics,
+                const bool switch_aerosol_optics,
                 const bool switch_single_gpt,
                 const int single_gpt,
+                const Int ray_count,
+                const Vector<int> grid_cells,
+                const Vector<Float> grid_d,
+                const Vector<int> kn_grid,
                 const Gas_concs_gpu& gas_concs,
+                Aerosol_concs_gpu& aerosol_concs,
                 const Array_gpu<Float,2>& p_lay, const Array_gpu<Float,2>& p_lev,
                 const Array_gpu<Float,2>& t_lay, const Array_gpu<Float,2>& t_lev,
                 Array_gpu<Float,2>& col_dry,
                 const Array_gpu<Float,1>& t_sfc, const Array_gpu<Float,2>& emis_sfc,
                 const Array_gpu<Float,2>& lwp, const Array_gpu<Float,2>& iwp,
                 const Array_gpu<Float,2>& rel, const Array_gpu<Float,2>& dei,
-                Array_gpu<Float,2>& tau, Array_gpu<Float,2>& lay_source,
-                Array_gpu<Float,2>& lev_source_inc, Array_gpu<Float,2>& lev_source_dec, Array_gpu<Float,1>& sfc_source,
+                const Array_gpu<Float,2>& rh,
+                Array_gpu<Float,2>& tot_tau_out, Array_gpu<Float,2>& cld_tau_out, Array_gpu<Float,2>& lay_source,
+                Array_gpu<Float,2>& lev_source, Array_gpu<Float,1>& sfc_source,
                 Array_gpu<Float,2>& lw_flux_up, Array_gpu<Float,2>& lw_flux_dn, Array_gpu<Float,2>& lw_flux_net,
                 Array_gpu<Float,2>& lw_gpt_flux_up, Array_gpu<Float,2>& lw_gpt_flux_dn, Array_gpu<Float,2>& lw_gpt_flux_net);
 
@@ -74,13 +82,16 @@ class Radiation_solver_longwave
         #ifdef __CUDACC__
         std::unique_ptr<Gas_optics_rrtmgp_rt> kdist_gpu;
         std::unique_ptr<Cloud_optics_rt> cloud_optics_gpu;
+        std::unique_ptr<Aerosol_optics_rt> aerosol_optics_gpu;
         Rte_lw_rt rte_lw;
 
         std::unique_ptr<Optical_props_arry_rt> optical_props;
 
         std::unique_ptr<Source_func_lw_rt> sources;
 
         std::unique_ptr<Optical_props_1scl_rt> cloud_optical_props;
+
+        std::unique_ptr<Optical_props_1scl_rt> aerosol_optical_props;
         #endif
 };
 
@@ -124,7 +135,7 @@ class Radiation_solver_shortwave
                 const Array_gpu<Float,2>& lwp, const Array_gpu<Float,2>& iwp,
                 const Array_gpu<Float,2>& rel, const Array_gpu<Float,2>& dei,
                 const Array_gpu<Float,2>& rh,
-                const Aerosol_concs_gpu& aerosol_concs,
+                Aerosol_concs_gpu& aerosol_concs,
                 Array_gpu<Float,2>& tot_tau_out, Array_gpu<Float,2>& tot_ssa_out,
                 Array_gpu<Float,2>& cld_tau_out, Array_gpu<Float,2>& cld_ssa_out, Array_gpu<Float,2>& cld_asy_out,
                 Array_gpu<Float,2>& aer_tau_out, Array_gpu<Float,2>& aer_ssa_out, Array_gpu<Float,2>& aer_asy_out,