From 4ec78d33efc26866f4ecac0476a23bf47fd6064c Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Fri, 28 Jun 2024 11:31:34 -0700
Subject: [PATCH] add lookup table for aerosols

---
 .../radiation/RRTMGPInterface.jl              | 35 ++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl
index af7d2da186..8f8691109f 100644
--- a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl
+++ b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl
@@ -525,6 +525,16 @@ function RRTMGPModel(
                 end
                 lookups = (; lookups..., lookup_lw_cld)
             end
+            if radiation_mode.aerosol_radiation
+                local lookup_lw_aero
+                data_loader(joinpath("rrtmgp-aerosols-merra-lw.nc")) do ds
+                    lookup_lw_aero =
+                        RRTMGP.LookUpTables.LookUpAerosolMerra(ds, FT, DA)
+                end
+            else
+                lookup_lw_aero = nothing
+            end
+            lookups = (; lookups..., lookup_lw_aero)
         end
 
         src_lw = RRTMGP.Sources.source_func_longwave(
@@ -586,6 +596,17 @@ function RRTMGPModel(
                 end
                 lookups = (; lookups..., lookup_sw_cld)
             end
+
+            if radiation_mode.aerosol_radiation
+                local lookup_sw_aero
+                data_loader(joinpath("rrtmgp-aerosols-merra-sw.nc")) do ds
+                    lookup_sw_aero =
+                        RRTMGP.LookUpTables.LookUpAerosolMerra(ds, FT, DA)
+                end
+            else
+                lookup_sw_aero = nothing
+            end
+            lookups = (; lookups..., lookup_sw_aero)
         end
 
         src_sw =
@@ -791,7 +812,7 @@ function RRTMGPModel(
         end
 
         if radiation_mode.aerosol_radiation
-            aero_type = DA{FT}(undef, nlay, ncol)
+            aero_type = DA{Int}(undef, nlay, ncol)
             name = "center_aerosol_type"
             set_and_save!(aero_type, name, t..., dict)
             aero_size = DA{FT}(undef, nlay, ncol)
@@ -1133,6 +1154,8 @@ NVTX.@annotate update_lw_fluxes!(::ClearSkyRadiation, model) =
         model.lw_solver,
         model.as,
         model.lookups.lookup_lw,
+        nothing,
+        model.lookups.lookup_lw_aero,
     )
 NVTX.@annotate update_lw_fluxes!(::AllSkyRadiation, model) =
     RRTMGP.RTESolver.solve_lw!(
@@ -1140,6 +1163,7 @@ NVTX.@annotate update_lw_fluxes!(::AllSkyRadiation, model) =
         model.as,
         model.lookups.lookup_lw,
         model.lookups.lookup_lw_cld,
+        model.lookups.lookup_lw_aero,
     )
 NVTX.@annotate function update_lw_fluxes!(
     ::AllSkyRadiationWithClearSkyDiagnostics,
@@ -1149,6 +1173,8 @@ NVTX.@annotate function update_lw_fluxes!(
         model.lw_solver,
         model.as,
         model.lookups.lookup_lw,
+        nothing,
+        model.lookups.lookup_lw_aero,
     )
     parent(model.face_clear_lw_flux_up) .= parent(model.face_lw_flux_up)
     parent(model.face_clear_lw_flux_dn) .= parent(model.face_lw_flux_dn)
@@ -1158,6 +1184,7 @@ NVTX.@annotate function update_lw_fluxes!(
         model.as,
         model.lookups.lookup_lw,
         model.lookups.lookup_lw_cld,
+        model.lookups.lookup_lw_aero,
     )
 end
 
@@ -1168,6 +1195,8 @@ NVTX.@annotate update_sw_fluxes!(::ClearSkyRadiation, model) =
         model.sw_solver,
         model.as,
         model.lookups.lookup_sw,
+        nothing,
+        model.lookups.lookup_sw_aero,
     )
 NVTX.@annotate update_sw_fluxes!(::AllSkyRadiation, model) =
     RRTMGP.RTESolver.solve_sw!(
@@ -1175,6 +1204,7 @@ NVTX.@annotate update_sw_fluxes!(::AllSkyRadiation, model) =
         model.as,
         model.lookups.lookup_sw,
         model.lookups.lookup_sw_cld,
+        model.lookups.lookup_sw_aero,
     )
 NVTX.@annotate function update_sw_fluxes!(
     ::AllSkyRadiationWithClearSkyDiagnostics,
@@ -1184,6 +1214,8 @@ NVTX.@annotate function update_sw_fluxes!(
         model.sw_solver,
         model.as,
         model.lookups.lookup_sw,
+        nothing,
+        model.lookups.lookup_sw_aero,
     )
     parent(model.face_clear_sw_flux_up) .= parent(model.face_sw_flux_up)
     parent(model.face_clear_sw_flux_dn) .= parent(model.face_sw_flux_dn)
@@ -1195,6 +1227,7 @@ NVTX.@annotate function update_sw_fluxes!(
         model.as,
         model.lookups.lookup_sw,
         model.lookups.lookup_sw_cld,
+        model.lookups.lookup_sw_aero,
     )
 end