add coupler fields based on sim types

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -258,6 +258,21 @@ end
 
 Interfacer.reinit!(sim::ClimaAtmosSimulation) = Interfacer.reinit!(sim.integrator)
 
+"""
+Extend Interfacer.add_coupler_fields! to add the fields required for ClimaAtmosSimulation.
+
+The fields added are:
+- `:surface_direct_albedo` (for radiation)
+- `:surface_diffuse_albedo` (for radiation)
+- `:ϵ_sfc` (for radiation)
+- `:T_sfc` (for radiation)
+- `:q_sfc` (for moisture)
+"""
+function Interfacer.add_coupler_fields!(coupler_field_names::Set, ::ClimaAtmosSimulation)
+    atmos_coupler_fields = [:surface_direct_albedo, :surface_diffuse_albedo, :ϵ_sfc, :T_sfc, :q_sfc]
+    push!(coupler_field_names, atmos_coupler_fields...)
+end
+
 function FieldExchanger.update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
 
     u = atmos_sim.integrator.u

experiments/ClimaEarth/components/land/climaland_bucket.jl

@@ -300,6 +300,20 @@ end
 Interfacer.step!(sim::BucketSimulation, t) = Interfacer.step!(sim.integrator, t - sim.integrator.t, true)
 Interfacer.reinit!(sim::BucketSimulation) = Interfacer.reinit!(sim.integrator)
 
+"""
+Extend Interfacer.add_coupler_fields! to add the fields required for BucketSimulation.
+
+The fields added are:
+- `:ρ_sfc`
+- `:F_radiative` (for radiation input)
+- `:P_liq` (for precipitation input)
+- `:P_snow` (for precipitation input)
+"""
+function Interfacer.add_coupler_fields!(coupler_field_names::Set, ::BucketSimulation)
+    bucket_coupler_fields = [:ρ_sfc, :F_radiative, :P_liq, :P_snow]
+    push!(coupler_field_names, bucket_coupler_fields...)
+end
+
 # extensions required by FluxCalculator (partitioned fluxes)
 function FluxCalculator.update_turbulent_fluxes!(sim::BucketSimulation, fields::NamedTuple)
     (; F_turb_energy, F_turb_moisture) = fields

experiments/ClimaEarth/components/ocean/eisenman_seaice.jl

@@ -159,6 +159,18 @@ end
 Interfacer.step!(sim::EisenmanIceSimulation, t) = Interfacer.step!(sim.integrator, t - sim.integrator.t, true)
 Interfacer.reinit!(sim::EisenmanIceSimulation) = Interfacer.reinit!(sim.integrator)
 
+"""
+Extend Interfacer.add_coupler_fields! to add the fields required for EisenmanIceSimulation.
+
+The fields added are:
+- `:ρ_sfc` (for humidity calculation)
+- `:F_radiative` (for radiation input)
+"""
+function Interfacer.add_coupler_fields!(coupler_field_names::Set, ::EisenmanIceSimulation)
+    eisenman_coupler_fields = [:ρ_sfc, :F_radiative]
+    push!(coupler_field_names, eisenman_coupler_fields...)
+end
+
 # extensions required by FluxCalculator (partitioned fluxes)
 function FluxCalculator.update_turbulent_fluxes!(sim::EisenmanIceSimulation, fields::NamedTuple)
     (; F_turb_energy) = fields

experiments/ClimaEarth/components/ocean/prescr_seaice.jl

@@ -191,6 +191,18 @@ Interfacer.update_field!(sim::PrescribedIceSimulation, ::Val{:turbulent_moisture
 Interfacer.step!(sim::PrescribedIceSimulation, t) = Interfacer.step!(sim.integrator, t - sim.integrator.t, true)
 Interfacer.reinit!(sim::PrescribedIceSimulation) = Interfacer.reinit!(sim.integrator)
 
+"""
+Extend Interfacer.add_coupler_fields! to add the fields required for PrescribedIceSimulation.
+
+The fields added are:
+- `:ρ_sfc` (for humidity calculation)
+- `:F_radiative` (for radiation input)
+"""
+function Interfacer.add_coupler_fields!(coupler_field_names::Set, ::PrescribedIceSimulation)
+    ice_coupler_fields = [:ρ_sfc, :F_radiative]
+    push!(coupler_field_names, ice_coupler_fields...)
+end
+
 # extensions required by FluxCalculator (partitioned fluxes)
 function FluxCalculator.update_turbulent_fluxes!(sim::PrescribedIceSimulation, fields::NamedTuple)
     (; F_turb_energy) = fields

experiments/ClimaEarth/components/ocean/slab_ocean.jl

@@ -149,6 +149,18 @@ end
 Interfacer.step!(sim::SlabOceanSimulation, t) = Interfacer.step!(sim.integrator, t - sim.integrator.t, true)
 Interfacer.reinit!(sim::SlabOceanSimulation) = Interfacer.reinit!(sim.integrator)
 
+"""
+Extend Interfacer.add_coupler_fields! to add the fields required for SlabOceanSimulation.
+
+The fields added are:
+- `:ρ_sfc` (for humidity calculation)
+- `:F_radiative` (for radiation input)
+"""
+function Interfacer.add_coupler_fields!(coupler_field_names::Set, ::SlabOceanSimulation)
+    ocean_coupler_fields = [:ρ_sfc, :F_radiative]
+    push!(coupler_field_names, ocean_coupler_fields...)
+end
+
 # extensions required by FluxCalculator (partitioned fluxes)
 function FluxCalculator.update_turbulent_fluxes!(sim::SlabOceanSimulation, fields::NamedTuple)
     (; F_turb_energy) = fields

experiments/ClimaEarth/setup_run.jl

@@ -459,34 +459,22 @@ function setup_and_run(config_dict::AbstractDict)
     global `CoupledSimulation` struct, `cs`, below.
     =#
 
-    ## coupler exchange fields
-    coupler_field_names = (
-        :T_sfc,
-        :z0m_sfc,
-        :z0b_sfc,
-        :ρ_sfc,
-        :q_sfc,
-        :surface_direct_albedo,
-        :surface_diffuse_albedo,
-        :beta,
-        :F_turb_energy,
-        :F_turb_moisture,
-        :F_turb_ρτxz,
-        :F_turb_ρτyz,
-        :F_radiative,
-        :P_liq,
-        :P_snow,
-        :radiative_energy_flux_toa,
-        :P_net,
-        :temp1,
-        :temp2,
-    )
-    coupler_fields = NamedTuple{coupler_field_names}(ntuple(i -> zeros(boundary_space), length(coupler_field_names)))
-    Utilities.show_memory_usage()
-
     ## model simulations
     model_sims = (atmos_sim = atmos_sim, ice_sim = ice_sim, land_sim = land_sim, ocean_sim = ocean_sim)
 
+    ## coupler exchange fields
+    coupler_field_names = Interfacer.default_coupler_fields()
+    for sim in model_sims
+        Interfacer.add_coupler_fields!(coupler_field_names, sim)
+    end
+    # add coupler fields required to track conservation, if specified
+    if energy_check
+        push!(coupler_field_names, :radiative_energy_flux_toa, :P_net)
+    end
+
+    coupler_fields = (; zip(coupler_field_names, ntuple(i -> zeros(boundary_space), length(coupler_field_names)))...)
+    Utilities.show_memory_usage()
+
     ## dates
     dates = (; date = [date], date0 = [date0])
 

experiments/ClimaEarth/user_io/debug_plots.jl

@@ -103,24 +103,7 @@ If `cs_fields_ref` is provided (e.g., using a copy of cs.fields from the initial
 plot the anomalies of the fields with respect to `cs_fields_ref`.
 """
 function debug(cs_fields::NamedTuple, dir, cs_fields_ref = nothing)
-    field_names = (
-        :surface_direct_albedo,
-        :surface_diffuse_albedo,
-        :F_radiative,
-        :F_turb_energy,
-        :F_turb_moisture,
-        :F_turb_ρτxz,
-        :F_turb_ρτyz,
-        :P_liq,
-        :P_snow,
-        :T_sfc,
-        :ρ_sfc,
-        :q_sfc,
-        :beta,
-        :z0b_sfc,
-        :z0m_sfc,
-        :radiative_energy_flux_toa,
-    )
+    field_names = keys(cs_fields)
     fig = Makie.Figure(size = (1500, 800))
     min_square_len = ceil(Int, sqrt(length(field_names)))
     for i in 1:min_square_len, j in 1:min_square_len

src/Interfacer.jl

@@ -84,6 +84,32 @@ Return the floating point type backing `T`: `T` can either be an object or a typ
 """
 float_type(::CoupledSimulation{FT}) where {FT} = FT
 
+"""
+    default_coupler_fields()
+
+Return the set of default coupler fields needed to run a simulation.
+"""
+default_coupler_fields() = Set{Symbol}([
+    # fields needed for flux calculations and exchange
+    :z0m_sfc,
+    :z0b_sfc,
+    :beta,
+    :F_turb_energy,
+    :F_turb_moisture,
+    :F_turb_ρτxz,
+    :F_turb_ρτyz,
+    # fields used for temporary storage during calculations
+    :temp1,
+    :temp2,
+])
+
+# fields needed for most surface models (not ClimaLandSimulation)
+# :ρ_sfc,
+# :F_radiative,
+# # fields needed for bucket/land
+# :P_liq,
+# :P_snow,
+
 """
     ComponentModelSimulation
 
@@ -209,6 +235,17 @@ update_field!(
 update_field_warning(sim, val::Val{X}) where {X} =
     @warn("`update_field!` is not extended for the `$X` field of " * name(sim) * ": skipping update.", maxlog = 1)
 
+
+"""
+    add_coupler_fields!(coupler_fields::Set, sim::ComponentModelSimulation, fields)
+
+A function to add fields to the set of coupler fields. This should be extended
+by component models that require coupler fields beyond the defaults.
+
+If this function isn't extended, no additional fields will be added.
+"""
+add_coupler_fields!(coupler_fields::Set, sim::ComponentModelSimulation) = nothing
+
 """
     name(::ComponentModelSimulation)
 

src/surface_stub.jl

@@ -97,6 +97,17 @@ The stub surface simulation is not updated by this function. Extends `SciMLBase.
 """
 reinit!(::AbstractSurfaceStub) = nothing
 
+"""
+Extend Interfacer.add_coupler_fields! to add the fields required for AbstractSurfaceStub.
+
+The fields added are:
+- `:ρ_sfc` (for humidity calculation)
+"""
+function Interfacer.add_coupler_fields!(coupler_field_names::Set, ::AbstractSurfaceStub)
+    surface_coupler_fields = [:ρ_sfc]
+    push!(coupler_field_names, surface_coupler_fields...)
+end
+
 """
     step!(::AbstractSurfaceStub, t)