From 15a701f18b1dd6ee5ae64ad82dbc0ee46312bbcc Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Mon, 8 Jul 2024 14:19:25 -0700
Subject: [PATCH] add different aerosol types to radiation

---
 ...uaplanet_rhoe_equilmoist_allsky_gw_res.yml |  2 +-
 src/callbacks/callbacks.jl                    | 32 ++++++++++++++++++-
 .../radiation/radiation.jl                    | 11 +++++--
 3 files changed, 40 insertions(+), 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
index a35a1f85f4d..dc57724cf90 100644
--- 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
index 8bd6a118442..19a5f22bf74 100644
--- 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)
+            ᶜaero_type =
+                Fields.array2field(rrtmgp_model.center_aerosol_type, axes(Y.c))
+            (; DST01, SSLT01, SO4, CB1, CB2, OC1, OC2) =
+                p.tracers.prescribed_aerosol_fields
+            @. ᶜ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,21 @@ function set_insolation_variables!(Y, p, t, ::TimeVaryingInsolation)
     end
 end
 
+# TODO: generalize this function and mapping from tracers to index
+# (maybe use ClimaCore fields?)
+function set_aerosol_type(DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)
+    maxval = max(DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)
+    i = findfirst(
+        c -> c == maxval,
+        (DST01, SSLT01, SO4, CB1, CB2, OC1, OC2),
+    )::Int
+    return i
+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
index bd1c3ef6235..dfe580f3ff0 100644
--- 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
                 )