change prescribed aerosols to field of namedtuples

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", "DST01", "DST02", "DST03", "DST04", "OC1", "OC2", "SO4", "SOA", "SSLT01", "SSLT02", "SSLT03", "SSLT04"]
+prescribed_aerosols: ["CB1", "CB2", "SO4"]
 toml: [toml/sphere_aquaplanet_rhoe_equilmoist_allsky_gw_res.toml]
 diagnostics:
   - short_name: [edt, evu]

src/cache/tracer_cache.jl

@@ -2,8 +2,9 @@ using ClimaUtilities.ClimaArtifacts
 import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput
 
 function tracer_cache(Y, atmos, prescribed_aerosol_names, start_date)
-    prescribed_aerosol_timevaryinginputs = (;)
-    prescribed_aerosol_fields = (;)
+    if isempty(prescribed_aerosol_names)
+        return (;)
+    end
 
     # Take the aerosol2005/aero_2005.nc file, read the keys with names matching
     # the ones passed in the prescribed_aerosol_names option, and create a
@@ -13,30 +14,30 @@ function tracer_cache(Y, atmos, prescribed_aerosol_names, start_date)
     # The keys in the aero_2005.nc file have to match the ones passed with the
     # configuration. The file also has to be defined on the globe and provide
     # time series of lon-lat-z data.
-    if !isempty(prescribed_aerosol_names)
-        prescribed_aerosol_names_as_symbols = Symbol.(prescribed_aerosol_names)
-        target_space = axes(Y.c)
-        timevaryinginputs = [
-            TimeVaryingInput(
-                joinpath(
-                    @clima_artifact("aerosol2005", ClimaComms.context(Y.c)),
-                    "aero_2005.nc",
-                ),
-                name,
-                target_space;
-                reference_date = start_date,
-                regridder_type = :InterpolationsRegridder,
-            ) for name in prescribed_aerosol_names
-        ]
-        empty_fields =
-            [zero(Y.c.ρ) for _ in prescribed_aerosol_names_as_symbols]
-
-        prescribed_aerosol_timevaryinginputs =
-            (; zip(prescribed_aerosol_names_as_symbols, timevaryinginputs)...)
+    prescribed_aerosol_names_as_symbols = Symbol.(prescribed_aerosol_names)
+    target_space = axes(Y.c)
+    timevaryinginputs = [
+        TimeVaryingInput(
+            joinpath(
+                @clima_artifact("aerosol2005", ClimaComms.context(Y.c)),
+                "aero_2005.nc",
+            ),
+            name,
+            target_space;
+            reference_date = start_date,
+            regridder_type = :InterpolationsRegridder,
+        ) for name in prescribed_aerosol_names
+    ]
 
-        # We add empty Fields here. Fields are updated in the radiation callback
-        prescribed_aerosol_fields =
-            (; zip(prescribed_aerosol_names_as_symbols, empty_fields)...)
-    end
-    return (; prescribed_aerosol_fields, prescribed_aerosol_timevaryinginputs)
+    # Field is updated in the radiation callback
+    prescribed_aerosols_field = similar(
+        Y.c,
+        NamedTuple{
+            prescribed_aerosol_names_as_symbols,
+            NTuple{length(prescribed_aerosol_names_as_symbols), eltype(Y.c.ρ)},
+        },
+    )
+    prescribed_aerosol_timevaryinginputs =
+        (; zip(prescribed_aerosol_names_as_symbols, timevaryinginputs)...)
+    return (; prescribed_aerosols_field, prescribed_aerosol_timevaryinginputs)
 end

src/callbacks/callbacks.jl

@@ -64,9 +64,11 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     (; radiation_mode) = p.atmos
 
     # If we have prescribed aerosols, we need to update them
-    for (key, tv) in pairs(p.tracers.prescribed_aerosol_timevaryinginputs)
-        field = getfield(p.tracers.prescribed_aerosol_fields, key)
-        evaluate!(field, tv, t)
+    if !isempty(p.tracers)
+        for (key, tv) in pairs(p.tracers.prescribed_aerosol_timevaryinginputs)
+            field = getproperty(p.tracers.prescribed_aerosols_field, key)
+            evaluate!(field, tv, t)
+        end
     end
 
     FT = Spaces.undertype(axes(Y.c))
@@ -159,26 +161,15 @@ 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)
+                set_aerosol_type(p.tracers.prescribed_aerosols_field)
             ᶜaero_conc = Fields.array2field(
                 rrtmgp_model.center_aerosol_column_mass_density,
                 axes(Y.c),
             )
-            @. ᶜaero_conc =
-                set_aerosol_concentration(
-                    DST01,
-                    SSLT01,
-                    SO4,
-                    CB1,
-                    CB2,
-                    OC1,
-                    OC2,
-                ) * ᶜΔz
+            @. ᶜaero_conc = maximum(p.tracers.prescribed_aerosols_field) * ᶜΔz
         end
     end
 
@@ -261,17 +252,20 @@ 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)
+function set_aerosol_type(;
+    DST01 = 0,
+    SSLT01 = 0,
+    SO4 = 0,
+    CB1 = 0,
+    CB2 = 0,
+    OC1 = 0,
+    OC2 = 0,
+    _...,
+)
     _, 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
+set_aerosol_type(NT) = set_aerosol_type(; NT...)
 
 NVTX.@annotate function save_state_to_disk_func(integrator, output_dir)
     (; t, u, p) = integrator

src/parameterized_tendencies/radiation/radiation.jl

@@ -43,12 +43,12 @@ function radiation_model_cache(
     (; idealized_h2o, idealized_clouds) = radiation_mode
     if !(radiation_mode isa RRTMGPI.GrayRadiation)
         (; aerosol_radiation) = radiation_mode
-        if aerosol_radiation && !(all(
+        if aerosol_radiation && !(any(
             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",
+                "Need at least one aerosol type when aerosol radiation is turned on",
             )
         end
     end