Merge pull request #3180 from CliMA/zs/aerosol_type

add different aerosol types to radiation

NEWS.md

@@ -3,6 +3,8 @@ ClimaAtmos.jl Release Notes
 
 Main
 -------
+- Allow different aerosol types for radiation.
+  PR [#3180](https://github.com/CliMA/ClimaAtmos.jl/pull/3180)
 
 v0.27.0
 -------

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,11 +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)
+            ᶜaero_type =
+                Fields.array2field(rrtmgp_model.center_aerosol_type, axes(Y.c))
+            @. ᶜaero_type =
+                set_aerosol_type(p.tracers.prescribed_aerosols_field)
             ᶜ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 = maximum(p.tracers.prescribed_aerosols_field) * ᶜΔz
         end
     end
 
@@ -246,6 +252,21 @@ function set_insolation_variables!(Y, p, t, ::TimeVaryingInsolation)
     end
 end
 
+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
+set_aerosol_type(NT) = set_aerosol_type(; NT...)
+
 NVTX.@annotate function save_state_to_disk_func(integrator, output_dir)
     (; t, u, p) = integrator
     Y = u

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 && !(any(
+            x -> x in aerosol_names,
+            ["DST01", "SSLT01", "SO4", "CB1", "CB2", "OC1", "OC2"],
+        ))
+            error(
+                "Need at least one aerosol type 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
                 )