Merge pull request #1207 from CliMA/js/csf

Sbozzolo · web-flow · commit b30a881fd007 · 2025-02-26T17:17:04.000-08:00
add coupler fields based on simulation types
diff --git a/NEWS.md b/NEWS.md
@@ -6,6 +6,15 @@ ClimaCoupler.jl Release Notes
 
 ### ClimaCoupler features
 
+#### Add coupler fields based on simulation type PR[#1207](https://github.com/CliMA/ClimaCoupler.jl/pull/1207)
+Previously, the coupler fields were hardcoded to be the same for all
+simulations, independent of what components were included. Now, each
+component model specifies the coupler fields it requires for coupling,
+and these are used to construct the set of coupler fields.
+TOA radiation and net precipitation are added only if conservation is enabled.
+The coupler fields are also now stored as a ClimaCore Field of NamedTuples,
+rather than as a NamedTuple of ClimaCore Fields.
+
 #### Remove extra `get_field` functions PR[#1203](https://github.com/CliMA/ClimaCoupler.jl/pull/1203)
 Removes the `get_field` functions for `air_density` for all models, which
 were unused except for the `BucketSimulation` method, which is replaced by a
diff --git a/docs/src/interfacer.md b/docs/src/interfacer.md
@@ -51,6 +51,12 @@ of SciMLBase.jl.
 - `get_model_prog_state(::ComponentModelSimulation)`: A function that
 returns the state vector of the simulation at its current state. This
 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`.
 
 ### ComponentModelSimulation - optional functions
 - `update_sim!(::ComponentModelSimulation, csf, turbulent_fluxes)`: A
diff --git a/experiments/ClimaEarth/components/atmosphere/climaatmos.jl b/experiments/ClimaEarth/components/atmosphere/climaatmos.jl
@@ -185,11 +185,11 @@ Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:thermo_state_int}) =
     CC.Spaces.level(sim.integrator.p.precomputed.ᶜts, 1)
 Interfacer.get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) =
     ρq_tot(atmos_sim.integrator.p.atmos.moisture_model, atmos_sim.integrator)
-function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, csf)
+function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, field)
     # note that this field is also being updated internally by the surface thermo state in ClimaAtmos
     # if turbulent fluxes are calculated, to ensure consistency. In case the turbulent fluxes are not
     # calculated, we update the field here.
-    sim.integrator.p.radiation.rrtmgp_model.surface_temperature .= CC.Fields.field2array(csf.T_sfc)
+    sim.integrator.p.radiation.rrtmgp_model.surface_temperature .= CC.Fields.field2array(field)
 end
 # extensions required by FluxCalculator (partitioned fluxes)
 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:height_int}) =
@@ -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, ::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
@@ -269,7 +284,7 @@ function FieldExchanger.update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbul
         !(p.atmos.insolation isa CA.IdealizedInsolation) && CA.set_insolation_variables!(u, p, t, p.atmos.insolation)
         Interfacer.update_field!(atmos_sim, Val(:surface_direct_albedo), csf.surface_direct_albedo)
         Interfacer.update_field!(atmos_sim, Val(:surface_diffuse_albedo), csf.surface_diffuse_albedo)
-        Interfacer.update_field!(atmos_sim, Val(:surface_temperature), csf)
+        Interfacer.update_field!(atmos_sim, Val(:surface_temperature), csf.T_sfc)
     end
 
     if turbulent_fluxes isa FluxCalculator.PartitionedStateFluxes
diff --git a/experiments/ClimaEarth/components/land/climaland_bucket.jl b/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, ::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
diff --git a/experiments/ClimaEarth/components/ocean/eisenman_seaice.jl b/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, ::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
diff --git a/experiments/ClimaEarth/components/ocean/prescr_seaice.jl b/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, ::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
diff --git a/experiments/ClimaEarth/components/ocean/slab_ocean.jl b/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, ::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
diff --git a/experiments/ClimaEarth/setup_run.jl b/experiments/ClimaEarth/setup_run.jl
@@ -459,34 +459,20 @@ 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
+    energy_check && push!(coupler_field_names, :radiative_energy_flux_toa, :P_net)
+
+    # allocate space for the coupler fields
+    coupler_fields = Interfacer.init_coupler_fields(FT, coupler_field_names, boundary_space)
+
     ## dates
     dates = (; date = [date], date0 = [date0])
 
diff --git a/experiments/ClimaEarth/test/debug_plots_tests.jl b/experiments/ClimaEarth/test/debug_plots_tests.jl
@@ -35,7 +35,7 @@ plot_field_names(sim::Interfacer.SurfaceStub) = (:stub_field,)
 @testset "import_atmos_fields!" begin
 
     boundary_space = TestHelper.create_space(FT)
-    coupler_names = (
+    coupler_names = [
         :surface_direct_albedo,
         :surface_diffuse_albedo,
         :F_radiative,
@@ -52,15 +52,15 @@ plot_field_names(sim::Interfacer.SurfaceStub) = (:stub_field,)
         :z0b_sfc,
         :z0m_sfc,
         :radiative_energy_flux_toa,
-    )
+    ]
     atmos_names = (:atmos_field,)
     surface_names = (:surface_field,)
     stub_names = (:stub_field,)
 
     atmos_fields = NamedTuple{atmos_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(atmos_names)))
     surface_fields = NamedTuple{surface_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(surface_names)))
     stub_fields = NamedTuple{stub_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(stub_names)))
-    coupler_fields = NamedTuple{coupler_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_names)))
+    coupler_fields = Interfacer.init_coupler_fields(FT, coupler_names, boundary_space)
 
     model_sims = (;
         atmos_sim = ClimaAtmosSimulation(atmos_fields),
diff --git a/experiments/ClimaEarth/user_io/debug_plots.jl b/experiments/ClimaEarth/user_io/debug_plots.jl
@@ -95,32 +95,15 @@ function debug(cs::Interfacer.CoupledSimulation, dir = "debug", cs_fields_ref =
 end
 
 """
-    debug(cs_fields::NamedTuple, dir, cs_fields_ref = nothing)
+    debug(cs_fields::CC.Fields.Field, dir, cs_fields_ref = nothing)
 
 Plot useful coupler fields (in `field_names`) and save plots to a directory.
 
 If `cs_fields_ref` is provided (e.g., using a copy of cs.fields from the initialization),
 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,
-    )
+function debug(cs_fields::CC.Fields.Field, dir, cs_fields_ref = nothing)
+    field_names = propertynames(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
diff --git a/src/ConservationChecker.jl b/src/ConservationChecker.jl
@@ -6,6 +6,7 @@ This module contains functions that check global conservation of energy and wate
 module ConservationChecker
 
 import ..Interfacer, ..Utilities
+import ClimaCore as CC
 
 export AbstractConservationCheck, EnergyConservationCheck, WaterConservationCheck, check_conservation!
 
@@ -75,11 +76,8 @@ function check_conservation!(
     ccs = cc.sums
     (; model_sims) = coupler_sim
 
-    boundary_space = coupler_sim.boundary_space # thin shell approx (boundary_space[z=0] = boundary_space[z_top])
-
-    FT = eltype(coupler_sim.fields[1])
-
-    total = 0
+    FT = CC.Spaces.undertype(coupler_sim.boundary_space)
+    total = FT(0)
 
     # save surfaces
     for sim in model_sims
diff --git a/src/FluxCalculator.jl b/src/FluxCalculator.jl
@@ -104,7 +104,7 @@ function calculate_surface_air_density(atmos_sim::Interfacer.AtmosModelSimulatio
 end
 
 """
-    partitioned_turbulent_fluxes!(model_sims::NamedTuple, fields::NamedTuple, boundary_space::CC.Spaces.AbstractSpace, surface_scheme, thermo_params::TD.Parameters.ThermodynamicsParameters)
+    partitioned_turbulent_fluxes!(model_sims::NamedTuple, fields::CC.Fields.Field, 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
@@ -127,15 +127,15 @@ TODO:
 """
 function partitioned_turbulent_fluxes!(
     model_sims::NamedTuple,
-    fields::NamedTuple,
+    fields::CC.Fields.Field,
     boundary_space::CC.Spaces.AbstractSpace,
     surface_scheme,
     thermo_params::TD.Parameters.ThermodynamicsParameters,
 )
 
     atmos_sim = model_sims.atmos_sim
     csf = fields
-    FT = eltype(csf[1])
+    FT = CC.Spaces.undertype(boundary_space)
 
     # reset coupler fields
     csf.F_turb_ρτxz .*= FT(0)
diff --git a/src/Interfacer.jl b/src/Interfacer.jl
@@ -84,6 +84,42 @@ 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 a list of default coupler fields needed to run a simulation.
+"""
+default_coupler_fields() = [
+    # 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,
+]
+
+"""
+    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.
+"""
+function init_coupler_fields(FT, coupler_field_names, boundary_space)
+    # First remove any duplicate field names
+    unique!(coupler_field_names)
+
+    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
+end
+
 """
     ComponentModelSimulation
 
@@ -209,6 +245,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, 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, sim::ComponentModelSimulation) = nothing
+
 """
     name(::ComponentModelSimulation)
 
diff --git a/src/surface_stub.jl b/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, ::AbstractSurfaceStub)
+    surface_coupler_fields = [:ρ_sfc]
+    push!(coupler_field_names, surface_coupler_fields...)
+end
+
 """
     step!(::AbstractSurfaceStub, t)
 
diff --git a/test/flux_calculator_tests.jl b/test/flux_calculator_tests.jl
