identify 1D and 2D coupler fields

Author: juliasloan25

docs/src/interfacer.md

@@ -54,9 +54,10 @@ is used for checkpointing the simulation.
 - `add_coupler_fields!(coupler_field_names::Set, ::ComponentModelSimulation)`:
 A function that adds names of quantities the coupler must exchange
 to support this component model. These will be added for each model
-in addition to the existing defaults: `z0m_sfc`, `z0b_sfc`, `beta`,
-`F_turb_energy`, `F_turb_moisture`, `F_turb_ρτxz`, `F_turb_ρτyz`,
-`temp1`, and `temp2`.
+in addition to the existing defaults: `z0m_sfc` (scalar), `z0b_sfc` (scalar),
+`beta` (scalar), `F_turb_energy`, `F_turb_moisture`, `F_turb_ρτxz`, `F_turb_ρτyz`,
+`temp1`, and `temp2`. They are defined as 2D fields at the surface, unless
+specified otherwise. When extending this function, if a variable
 
 ### ComponentModelSimulation - optional functions
 - `update_sim!(::ComponentModelSimulation, csf, turbulent_fluxes)`: A
@@ -197,9 +198,9 @@ overwritten or used as-is. These currently include the following:
 
 | Coupler name      | Description | Units | Default value |
 |-------------------|-------------|-------|---------------|
-| `beta` | factor that scales evaporation based on its estimated level of saturation | | 1 |
-| `emissivity` | measure of how much energy a surface radiates | | 1 |
-| `height_disp` | displacement height relative to the surface | m | 0 |
+| `beta` | factor that scales evaporation based on its estimated level of saturation (scalar) | | 1 |
+| `emissivity` | measure of how much energy a surface radiates (scalar) | | 1 |
+| `height_disp` | displacement height relative to the surface (scalar) | m | 0 |
 
 - `update_turbulent_fluxes!(::ComponentModelSimulation, fields::NamedTuple)`:
 This function updates the turbulent fluxes of the component model simulation

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -346,6 +346,8 @@ end
 Interfacer.reinit!(sim::ClimaAtmosSimulation) = Interfacer.reinit!(sim.integrator)
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::ClimaAtmosSimulation)
+
 Extend Interfacer.add_coupler_fields! to add the fields required for ClimaAtmosSimulation.
 
 The fields added are:
@@ -356,8 +358,10 @@ The fields added are:
 - `:q_sfc` (for moisture)
 """
 function Interfacer.add_coupler_fields!(coupler_field_names, ::ClimaAtmosSimulation)
+    # All fields are 2D; update the list of 2D coupler field names
     atmos_coupler_fields = [:surface_direct_albedo, :surface_diffuse_albedo, :ϵ_sfc, :T_sfc, :q_sfc]
-    push!(coupler_field_names, atmos_coupler_fields...)
+    push!(coupler_field_names[2], atmos_coupler_fields...)
+    return nothing
 end
 
 function FieldExchanger.update_sim!(sim::ClimaAtmosSimulation, csf, turbulent_fluxes)

experiments/ClimaEarth/components/land/climaland_bucket.jl

@@ -312,15 +312,19 @@ Interfacer.step!(sim::BucketSimulation, t) = Interfacer.step!(sim.integrator, t
 Interfacer.reinit!(sim::BucketSimulation) = Interfacer.reinit!(sim.integrator)
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::BucketSimulation)
+
 Extend Interfacer.add_coupler_fields! to add the fields required for BucketSimulation.
 
 The fields added are:
 - `:ρ_sfc`
 - `:F_radiative` (for radiation input)
 """
 function Interfacer.add_coupler_fields!(coupler_field_names, ::BucketSimulation)
+    # All fields are 2D; update the list of 2D coupler field names
     bucket_coupler_fields = [:ρ_sfc, :F_radiative]
-    push!(coupler_field_names, bucket_coupler_fields...)
+    push!(coupler_field_names[2], bucket_coupler_fields...)
+    return nothing
 end
 
 # extensions required by FluxCalculator (partitioned fluxes)

experiments/ClimaEarth/components/land/climaland_integrated.jl

@@ -416,8 +416,8 @@ end
 function Interfacer.update_field!(sim::ClimaLandSimulation, ::Val{:air_humidity}, field)
     parent(sim.integrator.p.drivers.q) .= parent(field)
 end
-function Interfacer.update_field!(sim::ClimaLandSimulation, ::Val{:c_co2}, field)
-    sim.integrator.p.drivers.c_co2 .= field
+function Interfacer.update_field!(sim::ClimaLandSimulation, ::Val{:c_co2}, scalar)
+    sim.integrator.p.drivers.c_co2 .= scalar
 end
 function Interfacer.update_field!(sim::ClimaLandSimulation, ::Val{:liquid_precipitation}, field)
     # Arbitrarily take parameters from the soil (they are the same for all land sub-components)
@@ -471,29 +471,34 @@ function FieldExchanger.import_atmos_fields!(csf, sim::ClimaLandSimulation, atmo
     FieldExchanger.dummmy_remap!(csf.q_air, Interfacer.get_field(atmos_sim, Val(:specific_humidity)))
     FieldExchanger.dummmy_remap!(csf.P_liq, Interfacer.get_field(atmos_sim, Val(:liquid_precipitation)))
     FieldExchanger.dummmy_remap!(csf.P_snow, Interfacer.get_field(atmos_sim, Val(:snow_precipitation)))
-    # CO2 is a scalar so it doesn't need remapping TODO store CO2 as a scalar in the coupler fields
-    csf.c_co2 .= Interfacer.get_field(atmos_sim, Val(:co2))
+    # CO2 is a scalar so it doesn't need remapping, but we do need to access the Ref
+    csf.c_co2[] = Interfacer.get_field(atmos_sim, Val(:co2))
 end
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::ClimaLandSimulation)
+
 Extend Interfacer.add_coupler_fields! to add the fields required for ClimaLandSimulation.
 
 The fields added are:
+- `:c_co2` (for photosynthesis, biogeochemistry) - scalar
 - `:SW_d` (for radiative transfer)
 - `:LW_d` (for radiative transfer)
 - `:cos_zenith_angle` (for radiative transfer)
 - `:diffuse_fraction` (for radiative transfer)
-- `:c_co2` (for photosynthesis, biogeochemistry)
 - `:P_air` (for canopy conductance)
 - `:T_air` (for canopy conductance)
 - `:q_air` (for canopy conductance)
 - `P_liq` (for moisture fluxes)
 - `P_snow` (for moisture fluxes)
 """
 function Interfacer.add_coupler_fields!(coupler_field_names, ::ClimaLandSimulation)
-    land_coupler_fields =
-        [:SW_d, :LW_d, :cos_zenith_angle, :diffuse_fraction, :c_co2, :P_air, :T_air, :q_air, :P_liq, :P_snow]
-    push!(coupler_field_names, land_coupler_fields...)
+    land_coupler_fields_1d = [:c_co2]
+    land_coupler_fields_2d =
+        [:SW_d, :LW_d, :cos_zenith_angle, :diffuse_fraction, :P_air, :T_air, :q_air, :P_liq, :P_snow]
+    push!(coupler_field_names[1], land_coupler_fields_1d...)
+    push!(coupler_field_names[2], land_coupler_fields_2d...)
+    return nothing
 end
 
 function Checkpointer.get_model_prog_state(sim::ClimaLandSimulation)

experiments/ClimaEarth/components/ocean/eisenman_seaice.jl

@@ -158,15 +158,19 @@ Interfacer.step!(sim::EisenmanIceSimulation, t) = Interfacer.step!(sim.integrato
 Interfacer.reinit!(sim::EisenmanIceSimulation) = Interfacer.reinit!(sim.integrator)
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::EisenmanIceSimulation)
+
 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, ::EisenmanIceSimulation)
+    # All fields are 2D; update the list of 2D coupler field names
     eisenman_coupler_fields = [:ρ_sfc, :F_radiative]
-    push!(coupler_field_names, eisenman_coupler_fields...)
+    push!(coupler_field_names[2], eisenman_coupler_fields...)
+    return nothing
 end
 
 # extensions required by FluxCalculator (partitioned fluxes)

experiments/ClimaEarth/components/ocean/prescr_seaice.jl

@@ -190,15 +190,19 @@ Interfacer.step!(sim::PrescribedIceSimulation, t) = Interfacer.step!(sim.integra
 Interfacer.reinit!(sim::PrescribedIceSimulation) = Interfacer.reinit!(sim.integrator)
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::PrescribedIceSimulation)
+
 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, ::PrescribedIceSimulation)
+    # All fields are 2D; update the list of 2D coupler field names
     ice_coupler_fields = [:ρ_sfc, :F_radiative]
-    push!(coupler_field_names, ice_coupler_fields...)
+    push!(coupler_field_names[2], ice_coupler_fields...)
+    return nothing
 end
 
 # extensions required by FluxCalculator (partitioned fluxes)

experiments/ClimaEarth/components/ocean/slab_ocean.jl

@@ -148,15 +148,19 @@ Interfacer.step!(sim::SlabOceanSimulation, t) = Interfacer.step!(sim.integrator,
 Interfacer.reinit!(sim::SlabOceanSimulation) = Interfacer.reinit!(sim.integrator)
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::SlabOceanSimulation)
+
 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, ::SlabOceanSimulation)
+    # All fields are 2D; update the list of 2D coupler field names
     ocean_coupler_fields = [:ρ_sfc, :F_radiative]
-    push!(coupler_field_names, ocean_coupler_fields...)
+    push!(coupler_field_names[2], ocean_coupler_fields...)
+    return nothing
 end
 
 # extensions required by FluxCalculator (partitioned fluxes)

experiments/ClimaEarth/setup_run.jl

@@ -490,8 +490,8 @@ function setup_and_run(config_dict::AbstractDict)
     for sim in model_sims
         Interfacer.add_coupler_fields!(coupler_field_names, sim)
     end
-    # add coupler fields required to track conservation, if specified
-    energy_check && push!(coupler_field_names, :radiative_energy_flux_toa, :P_net)
+    # add 2D coupler fields required to track conservation, if specified
+    energy_check && push!(coupler_field_names[2], :radiative_energy_flux_toa, :P_net)
 
     # allocate space for the coupler fields
     coupler_fields = Interfacer.init_coupler_fields(FT, coupler_field_names, boundary_space)

experiments/ClimaEarth/test/component_model_tests/climaland_tests.jl

@@ -77,7 +77,7 @@ end
 
     # Initialize the coupler fields so we can perform exchange
     coupler_field_names = Interfacer.default_coupler_fields()
-    map(sim -> Interfacer.add_coupler_fields!(coupler_field_names, sim), values(model_sims))
+    map(sim -> Interfacer.add_coupler_fields!(coupler_field_names, sim), model_sims)
     coupler_fields = Interfacer.init_coupler_fields(FT, coupler_field_names, boundary_space)
     flux_type = FluxCalculator.PartitionedStateFluxes()
 

src/FluxCalculator.jl

@@ -105,7 +105,7 @@ function calculate_surface_air_density(atmos_sim::Interfacer.AtmosModelSimulatio
 end
 
 """
-    partitioned_turbulent_fluxes!(model_sims::NamedTuple, fields::CC.Fields.Field, boundary_space::CC.Spaces.AbstractSpace, surface_scheme, thermo_params::TD.Parameters.ThermodynamicsParameters)
+    partitioned_turbulent_fluxes!(model_sims::NamedTuple, fields::NamedTuple, boundary_space::CC.Spaces.AbstractSpace, surface_scheme, thermo_params::TD.Parameters.ThermodynamicsParameters)
 
 The current setup calculates the aerodynamic fluxes in the coupler (assuming no regridding is needed)
 using adapter function `get_surface_fluxes!`, which calls `SurfaceFluxes.jl`. The coupler saves
@@ -128,7 +128,7 @@ TODO:
 """
 function partitioned_turbulent_fluxes!(
     model_sims::NamedTuple,
-    csf::CC.Fields.Field,
+    csf::NamedTuple,
     boundary_space::CC.Spaces.AbstractSpace,
     surface_scheme,
     thermo_params::TD.Parameters.ThermodynamicsParameters,

src/Interfacer.jl

@@ -99,42 +99,51 @@ float_type(::CoupledSimulation{FT}) where {FT} = FT
 """
     default_coupler_fields()
 
-Return a list of default coupler fields needed to run a simulation.
-"""
-default_coupler_fields() = [
-    # fields needed for flux calculations
-    :z0m_sfc,
-    :z0b_sfc,
-    :beta,
-    :emissivity,
-    # fields used for flux exchange
-    :F_turb_energy,
-    :F_turb_moisture,
-    :F_turb_ρτxz,
-    :F_turb_ρτyz,
-    # fields used to track water conservation, and for water fluxes
-    :P_liq,
-    :P_snow,
-    # fields used for temporary storage during calculations
-    :temp1,
-    :temp2,
-]
+Return 2 vectors of default coupler fields needed to run a simulation,
+corresponding to 1D fields (scalars) and 2D fields.
+"""
+function default_coupler_fields()
+    keys_1d = [
+        # scalars needed for flux calculations
+        :z0m_sfc,
+        :z0b_sfc,
+        :beta,
+        :emissivity,
+    ]
+    keys_2d = [
+        # fields used for flux exchange
+        :F_turb_energy,
+        :F_turb_moisture,
+        :F_turb_ρτxz,
+        :F_turb_ρτyz,
+        # fields used to track water conservation, and for water fluxes
+        :P_liq,
+        :P_snow,
+        # fields used for temporary storage during calculations
+        :temp1,
+        :temp2,
+    ]
+    return (keys_1d, keys_2d)
+end
 
 """
     init_coupler_fields(FT, coupler_field_names, boundary_space)
 
-Allocate a Field of NamedTuples on the provided boundary space to store
-the provided coupler fields.
+Allocate a NamedTuple of Floats for scalars, and Fields on the provided
+boundary space to store the provided coupler fields.
 """
 function init_coupler_fields(FT, coupler_field_names, boundary_space)
+    keys_1d, keys_2d = coupler_field_names
     # First remove any duplicate field names
-    unique!(coupler_field_names)
+    unique!(keys_1d)
+    unique!(keys_2d)
+
+    # Allocate space for the 1D and 2D variables
+    vals_1d = ((Ref(zero(FT)) for _ in 1:length(keys_1d))...,)
+    vals_2d = ((zeros(boundary_space) for _ in 1:length(keys_2d))...,)
 
-    key_types = (coupler_field_names...,)
-    val_types = Tuple{(FT for _ in 1:length(coupler_field_names))...}
-    nt_type = NamedTuple{key_types, val_types}
-    coupler_fields = zeros(nt_type, boundary_space)
-    return coupler_fields
+    # Dict(zip((keys_1d..., keys_2d...), (vals_1d..., vals_2d...)))
+    return NamedTuple{(keys_1d..., keys_2d...)}((vals_1d..., vals_2d...))
 end
 
 """

src/surface_stub.jl

@@ -97,14 +97,17 @@ The stub surface simulation is not updated by this function. Extends `SciMLBase.
 reinit!(::AbstractSurfaceStub) = nothing
 
 """
+    Interfacer.add_coupler_fields!(coupler_field_names, ::AbstractSurfaceStub)
+
 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, ::AbstractSurfaceStub)
+    # Add to the 2D coupler fields
     surface_coupler_fields = [:ρ_sfc]
-    push!(coupler_field_names, surface_coupler_fields...)
+    push!(coupler_field_names[2], surface_coupler_fields...)
 end
 
 """

test/flux_calculator_tests.jl

@@ -199,20 +199,23 @@ for FT in (Float32, Float64)
 
             model_sims = (; atmos_sim, ocean_sim, ocean_sim2)
 
-            coupler_cache_names = [
-                :T_sfc,
-                :surface_direct_albedo,
-                :surface_diffuse_albedo,
-                :F_R_sfc,
-                :F_R_toa,
-                :P_liq,
-                :P_snow,
-                :P_net,
-                :F_turb_energy,
-                :F_turb_ρτxz,
-                :F_turb_ρτyz,
-                :F_turb_moisture,
-            ]
+            coupler_cache_names = (
+                [],
+                [
+                    :T_sfc,
+                    :surface_direct_albedo,
+                    :surface_diffuse_albedo,
+                    :F_R_sfc,
+                    :F_R_toa,
+                    :P_liq,
+                    :P_snow,
+                    :P_net,
+                    :F_turb_energy,
+                    :F_turb_ρτxz,
+                    :F_turb_ρτyz,
+                    :F_turb_moisture,
+                ],
+            )
             fields = Interfacer.init_coupler_fields(FT, coupler_cache_names, boundary_space)
 
             # calculate turbulent fluxes