add different aerosol types to radiation

CliMA · Jul 8, 2024 · 6cfebc2 · 6cfebc2
1 parent d6fb433
commit 6cfebc2
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 5 deletions.
diff --git a/config/model_configs/sphere_aquaplanet_rhoe_equilmoist_allsky_gw_res.yml b/config/model_configs/sphere_aquaplanet_rhoe_equilmoist_allsky_gw_res.yml
@@ -16,7 +16,7 @@ rayleigh_sponge: true
 non_orographic_gravity_wave: true
 orographic_gravity_wave: "gfdl_restart"
 surface_setup: "DefaultMoninObukhov"
-prescribed_aerosols: ["CB1", "CB2", "SO4"]
+prescribed_aerosols: ["CB1", "CB2", "DST01", "DST02", "DST03", "DST04", "OC1", "OC2", "SO4", "SOA", "SSLT01", "SSLT02", "SSLT03", "SSLT04"]
 toml: [toml/sphere_aquaplanet_rhoe_equilmoist_allsky_gw_res.toml]
 diagnostics:
   - short_name: [edt, evu]

diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
@@ -159,11 +159,26 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     if !(radiation_mode isa RRTMGPI.GrayRadiation)
         if radiation_mode.aerosol_radiation
             ᶜΔz = Fields.Δz_field(Y.c)
+            (; DST01, SSLT01, SO4, CB1, CB2, OC1, OC2) =
+                p.tracers.prescribed_aerosol_fields
+            ᶜaero_type =
+                Fields.array2field(rrtmgp_model.center_aerosol_type, axes(Y.c))
+            @. ᶜaero_type =
+                set_aerosol_type(DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)
             ᶜaero_conc = Fields.array2field(
                 rrtmgp_model.center_aerosol_column_mass_density,
                 axes(Y.c),
             )
-            @. ᶜaero_conc = p.tracers.prescribed_aerosol_fields.:SO4 * ᶜΔz
+            @. ᶜaero_conc =
+                set_aerosol_concentration(
+                    DST01,
+                    SSLT01,
+                    SO4,
+                    CB1,
+                    CB2,
+                    OC1,
+                    OC2,
+                ) * ᶜΔz
         end
     end
 
@@ -246,6 +261,18 @@ function set_insolation_variables!(Y, p, t, ::TimeVaryingInsolation)
     end
 end
 
+# TODO: The order of the input arguments needs to be consistent with
+# the indexing of the rrtmgp lookup tables. We can generalize these functions 
+# and the mapping from tracers to index (maybe use ClimaCore fields?)
+function set_aerosol_type(DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)
+    _, index = findmax((DST01, SSLT01, SO4, CB1, CB2, OC1, OC2))
+    return index
+end
+
+function set_aerosol_concentration(DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)
+    return max(DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)
+end
+
 NVTX.@annotate function save_state_to_disk_func(integrator, output_dir)
     (; t, u, p) = integrator
     Y = u

diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl
@@ -43,8 +43,13 @@ function radiation_model_cache(
     (; idealized_h2o, idealized_clouds) = radiation_mode
     if !(radiation_mode isa RRTMGPI.GrayRadiation)
         (; aerosol_radiation) = radiation_mode
-        if aerosol_radiation && !("SO4" in aerosol_names)
-            error("Aerosol radiation is turned on but SO4 is not provided")
+        if aerosol_radiation && !(all(
+            x -> x in aerosol_names,
+            ["DST01", "SSLT01", "SO4", "CB1", "CB2", "OC1", "OC2"],
+        ))
+            error(
+                "All aerosol tpes must be included when aerosol radiation is turned on",
+            )
         end
     end
     FT = Spaces.undertype(axes(Y.c))
@@ -191,7 +196,7 @@ function radiation_model_cache(
             if aerosol_radiation
                 kwargs = (;
                     kwargs...,
-                    center_aerosol_type = 3, # assuming only sulfate
+                    center_aerosol_type = 0, # initialized in callback
                     center_aerosol_radius = 0.2, # assuming fixed aerosol radius
                     center_aerosol_column_mass_density = NaN, # initialized in callback
                 )