Move aerosols to its own type

Aerosols were handles by parsing directly the YAML arguments. This led to implicit assumptions, and required aerosols to be handled differently. For example, we had to pass the list `aerosol_names` to `build_cache` even when no aerosol was required. This commit moves aerosol to its own type so that we can follow a more composable approach. It is still not fully possible to swap in a custom, user-defined type because the portion of code in RRTGMPInterface makes assumption on the names of the aerosols. I add error checking for that
CliMA · Dec 23, 2024 · 4e760dd · 4e760dd
1 parent b2b3e7f
commit 4e760dd
Show file tree

Hide file tree

Showing 9 changed files with 125 additions and 69 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -4,6 +4,15 @@ ClimaAtmos.jl Release Notes
 Main
 -------
 
+### Features
+
+### Aerosols are now its own type PR [3492](https://github.com/CliMA/ClimaAtmos.jl/pull/3492)
+
+Aerosols, currently only relevant for radiation interaction when using a RRTGMP
+model, are now described by a new type, `AbstractAerosols`. The one aerosol
+model currently implemented is `PrescribedCMIP5Aerosols`, where aerosol
+concentrations are read from CMIP5 forcing data. The YAML interface has not
+changed.
 
 v0.28.1
 -------

diff --git a/docs/src/tracers.md b/docs/src/tracers.md
@@ -1,4 +1,18 @@
-# Aerosols
+# Tracers
+
+## Aerosols
+
+`ClimaAtmos` can read aerosol profiles from the CIMP5 forcing data. The data is
+described in
+[`ClimaArtifacts`](https://github.com/CliMA/ClimaArtifacts/tree/main/aerosol_concentrations)
+and it is only relevant for the radiation transfer, and only when RRTMGP is
+used.
+
+To enable this mode, use the `PrescribedCMIP5Aerosols` aerosols model. Upon
+construction, `PrescribedCMIP5Aerosols` takes the list of aerosol names to be
+used. The names have to match what is in the file.
+
+## Ozone
 
 `ClimaAtmos` implements two modes for ozone profiles. These are only relevant
 for the radiation transfer, and only when RRTMGP is used.

diff --git a/src/cache/cache.jl b/src/cache/cache.jl
@@ -82,7 +82,7 @@ end
 
 # The model also depends on f_plane_coriolis_frequency(params)
 # This is a constant Coriolis frequency that is only used if space is flat
-function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
+function build_cache(Y, atmos, params, surface_setup, sim_info)
     (; dt, start_date, output_dir) = sim_info
     FT = eltype(params)
 
@@ -147,7 +147,7 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
 
     radiation_args =
         atmos.radiation_mode isa RRTMGPI.AbstractRRTMGPMode ?
-        (start_date, params, atmos.ozone, aerosol_names, atmos.insolation) : ()
+        (start_date, params, atmos.ozone, atmos.aerosols, atmos.insolation) : ()
 
     hyperdiff = hyperdiffusion_cache(Y, atmos)
     precipitation = precipitation_cache(Y, atmos)
@@ -157,7 +157,7 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
     non_orographic_gravity_wave = non_orographic_gravity_wave_cache(Y, atmos)
     orographic_gravity_wave = orographic_gravity_wave_cache(Y, atmos)
     radiation = radiation_model_cache(Y, atmos, radiation_args...)
-    tracers = tracer_cache(Y, atmos, aerosol_names, start_date)
+    tracers = tracer_cache(Y, atmos, start_date)
 
     args = (
         dt,

diff --git a/src/cache/tracer_cache.jl b/src/cache/tracer_cache.jl
@@ -23,53 +23,50 @@ function ozone_cache(::PrescribedOzone, Y, start_date)
     return (; o3, prescribed_o3_timevaryinginput)
 end
 
-function tracer_cache(Y, atmos, prescribed_aerosol_names, start_date)
-    if !isempty(prescribed_aerosol_names)
-        target_space = axes(Y.c)
+aerosols_cache(_, _, _) = (;)
+function aerosols_cache(aerosols::PrescribedCMIP5Aerosols, Y, start_date)
+    aerosol_names = aerosols.aerosol_names
+    target_space = axes(Y.c)
 
-        # Take the aerosol concentration file, read the keys with names matching
-        # the ones passed in the prescribed_aerosol_names option, and create a
-        # NamedTuple that uses the same keys and has as values the TimeVaryingInputs
-        # for those variables.
-        #
-        # The keys in the aerosol_concentrations.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.
-        prescribed_aerosol_names_as_symbols = Symbol.(prescribed_aerosol_names)
-        target_space = axes(Y.c)
-        extrapolation_bc = (Intp.Periodic(), Intp.Flat(), Intp.Flat())
-        timevaryinginputs = [
-            TimeVaryingInput(
-                AA.aerosol_concentration_file_path(;
-                    context = ClimaComms.context(Y.c),
-                ),
-                name,
-                target_space;
-                reference_date = start_date,
-                regridder_type = :InterpolationsRegridder,
-                regridder_kwargs = (; extrapolation_bc),
-                method = LinearPeriodFillingInterpolation(Year(1)),
-            ) for name in prescribed_aerosol_names
-        ]
+    # Take the aerosol concentration file, read the keys with names matching
+    # the ones passed in the prescribed_aerosol_names option, and create a
+    # NamedTuple that uses the same keys and has as values the TimeVaryingInputs
+    # for those variables.
+    #
+    # The keys in the aerosol_concentrations.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.
+    target_space = axes(Y.c)
+    extrapolation_bc = (Intp.Periodic(), Intp.Flat(), Intp.Flat())
+    timevaryinginputs = [
+        TimeVaryingInput(
+            AA.aerosol_concentration_file_path(;
+                context = ClimaComms.context(Y.c),
+            ),
+            name,
+            target_space;
+            reference_date = start_date,
+            regridder_type = :InterpolationsRegridder,
+            regridder_kwargs = (; extrapolation_bc),
+            method = LinearPeriodFillingInterpolation(Year(1)),
+        ) for name in aerosol_names
+    ]
+
+    # Field is updated in the radiation callback
+    prescribed_aerosols_field = similar(
+        Y.c,
+        NamedTuple{
+            aerosol_names,
+            NTuple{length(aerosol_names), eltype(Y.c.ρ)},
+        },
+    )
+    prescribed_aerosol_timevaryinginputs =
+        (; zip(aerosol_names, timevaryinginputs)...)
+    return (; prescribed_aerosols_field, prescribed_aerosol_timevaryinginputs)
+end
 
-        # 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)...)
-        aerosol_cache =
-            (; prescribed_aerosols_field, prescribed_aerosol_timevaryinginputs)
-    else
-        aerosol_cache = (;)
-    end
+function tracer_cache(Y, atmos, start_date)
+    aero_cache = aerosols_cache(atmos.aerosols, Y, start_date)
     o3_cache = ozone_cache(atmos.ozone, Y, start_date)
-    return (; aerosol_cache..., o3_cache...)
+    return (; aero_cache..., o3_cache...)
 end
diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
@@ -187,6 +187,9 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
 
     if !(radiation_mode isa RRTMGPI.GrayRadiation)
         if radiation_mode.aerosol_radiation
+            # This is because we hard-code the names of the symbols
+            @assert integrator.p.atmos.aerosols isa PrescribedCMIP5Aerosols "The only aerosol model currently supported is `PrescribedCMIP5Aerosols`"
+
             ᶜΔz = Fields.Δz_field(Y.c)
 
             ᶜaero_conc = Fields.array2field(

diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl
@@ -92,7 +92,7 @@ function radiation_model_cache(
     start_date,
     params,
     ozone,
-    aerosol_names,
+    aerosols,
     insolation_mode;
     interpolation = RRTMGPI.BestFit(),
     bottom_extrapolation = RRTMGPI.SameAsInterpolation(),
@@ -102,11 +102,13 @@ function radiation_model_cache(
     device = context.device
     if !(radiation_mode isa RRTMGPI.GrayRadiation)
         (; aerosol_radiation) = radiation_mode
-        if aerosol_radiation && !(any(
-            x -> x in aerosol_names,
-            ["DST01", "SSLT01", "SO4", "CB1", "CB2", "OC1", "OC2"],
-        ))
-            error(
+        if aerosol_radiation
+            aerosols isa PrescribedCMIP5Aerosols ||
+                error("Only `PrescribedCMIP5Aerosols` is currently supported")
+            any(
+                x -> x in aerosols.aerosol_names,
+                ["DST01", "SSLT01", "SO4", "CB1", "CB2", "OC1", "OC2"],
+            ) || error(
                 "Need at least one aerosol type when aerosol radiation is turned on",
             )
         end

diff --git a/src/solver/model_getters.jl b/src/solver/model_getters.jl
@@ -498,9 +498,10 @@ function get_surface_thermo_state_type(parsed_args)
     return dict[parsed_args["surface_thermo_state_type"]]
 end
 
-function get_tracers(parsed_args)
-    aerosol_names = Tuple(parsed_args["prescribed_aerosols"])
-    return (; aerosol_names)
+function get_aerosols(parsed_args)
+    isnothing(parsed_args["prescribed_aerosols"]) && return nothing
+    aerosol_names = Symbol.(Tuple(parsed_args["prescribed_aerosols"]))
+    return PrescribedCMIP5Aerosols(aerosol_names)
 end
 
 function get_tendency_model(parsed_args)

diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
@@ -20,6 +20,7 @@ function get_atmos(config::AtmosConfig, params)
     precip_model = get_precipitation_model(parsed_args)
     cloud_model = get_cloud_model(parsed_args)
     ozone = get_ozone(parsed_args)
+    aerosols = get_aerosols(parsed_args)
     radiation_mode = get_radiation_mode(parsed_args, FT)
     forcing_type = get_forcing_type(parsed_args)
     call_cloud_diagnostics_per_stage =
@@ -30,6 +31,10 @@ function get_atmos(config::AtmosConfig, params)
         ozone = IdealizedOzone()
     end
 
+    if !isnothing(aerosols) && !(radiation_mode isa RRTMGPI.AbstractRRTMGPMode)
+        @warn "$aerosols is effective only with an RRTGMP model"
+    end
+
     diffuse_momentum = !(forcing_type isa HeldSuarezForcing)
 
     advection_test = parsed_args["advection_test"]
@@ -56,6 +61,7 @@ function get_atmos(config::AtmosConfig, params)
     atmos = AtmosModel(;
         moisture_model,
         ozone,
+        aerosols,
         radiation_mode,
         subsidence = get_subsidence_model(parsed_args, radiation_mode, FT),
         ls_adv = get_large_scale_advection_model(parsed_args, FT),
@@ -677,17 +683,8 @@ function get_simulation(config::AtmosConfig)
         )
     end
 
-    tracers = get_tracers(config.parsed_args)
-
     s = @timed_str begin
-        p = build_cache(
-            Y,
-            atmos,
-            params,
-            surface_setup,
-            sim_info,
-            tracers.aerosol_names,
-        )
+        p = build_cache(Y, atmos, params, surface_setup, sim_info)
     end
     @info "Allocating cache (p): $s"
 

diff --git a/src/solver/types.jl b/src/solver/types.jl
@@ -103,6 +103,8 @@ struct GCMDrivenInsolation <: AbstractInsolation end
     AbstractOzone
 
 Describe how ozone concentration should be set.
+
+Currently, this is only relevant for RRTGMP.
 """
 abstract type AbstractOzone end
 
@@ -127,6 +129,33 @@ Refer to ClimaArtifacts for more information on how to obtain the artifact.
 """
 struct PrescribedOzone <: AbstractOzone end
 
+"""
+    AbstractAerosols
+
+Describe how aerosol concentrations should be set.
+
+Currently, this is only relevant for RRTGMP.
+"""
+abstract type AbstractAerosols end
+
+"""
+    PrescribedCMIP5Aerosols
+
+Implement a time-varying ozone profile as read from disk.
+
+The CMIP5 forcing dataset is used. For production runs, you should acquire the
+high-resolution, multi-year `aerosol_concentrations` artifact. If this is not
+available, a low resolution, single-year version will be used.
+
+Refer to ClimaArtifacts for more information on how to obtain the artifact.
+
+`aerosol_names` is the list of aerosols to use. The names have to match what is
+in the artifact.
+"""
+struct PrescribedCMIP5Aerosols{T} <: AbstractAerosols
+    aerosol_names::T
+end
+
 """
     AbstractCloudInRadiation
 
@@ -456,6 +485,7 @@ Base.@kwdef struct AtmosModel{
     F,
     S,
     OZ,
+    AERO,
     RM,
     LA,
     EXTFORCING,
@@ -489,6 +519,9 @@ Base.@kwdef struct AtmosModel{
     """What to do with ozone for radiation (when using RRTGMP)"""
     ozone::OZ = nothing
 
+    """What to do with aerosols for radiation (when using RRTGMP)"""
+    aerosols::AERO = nothing
+
     radiation_mode::RM = nothing
     ls_adv::LA = nothing
     external_forcing::EXTFORCING = nothing