Add flux module

docs/src/fieldexchanger.md

@@ -0,0 +1,35 @@
+# FieldExchanger
+
+This module contains general functions for the exchange of fields between the atmospheric and surface component models.
+
+The `FieldExchanger` needs to populate the coupler with
+- atmospheric fields (mostly fluxes), via the `import_atmos_fields!` function
+- average surface properties of each coupler gridpoint, via the `import_combined_surface_fields!` function
+
+The component models are updated by broadcasting the coupler fields, via the `update_model_sims!` function. For an update, this function requires that `update_field!` is defined for the particular variable and component model. Currently, we support the:
+- `AtmosModelSimulation`: `albedo`, `surface_temperature`
+    - if calculating fluxes in the atmospheric model: `roughness_momentum`, `roughness_buoyancy`, `beta`
+- `SurfaceModelSimulation`: `air_density`, `turbulent_energy_flux`, `turbulent_moisture_flux`, `radiative_energy_flux`, `liquid_precipitation`, `snow_precipitation`
+
+If an `update_field!` function is not defined for a particular component model, it will be ignored.
+
+Each component model is also required to define its own `step!` and `reinit!` functions, otherwise this will be ignored.
+
+## FieldExchanger API
+
+```@docs
+    ClimaCoupler.FieldExchanger.import_atmos_fields!
+    ClimaCoupler.FieldExchanger.import_combined_surface_fields!
+    ClimaCoupler.FieldExchanger.update_model_sims!
+    ClimaCoupler.FieldExchanger.update_sim!
+    ClimaCoupler.FieldExchanger.reinit_model_sims!
+    ClimaCoupler.FieldExchanger.step_model_sims!
+```
+
+
+## FieldExchanger Internal Functions
+
+```@docs
+    ClimaCoupler.FieldExchanger.step!
+    ClimaCoupler.FieldExchanger.reinit!
+```

docs/src/fluxcalculator.md

@@ -0,0 +1,20 @@
+# FluxCalculator
+
+This modules contains abstract types and functions to calculate surface fluxes in the coupler, or to call flux calculating functions from the component models.
+
+Fluxes over a heterogeneous surface (e.g., from a gridpoint where atmospheric cell is overlying both land and ocean) can be handled in two different ways:
+1. **Combined fluxes** (called with `CombinedAtmosGrid()`)
+  - these are calculated by averaging the surface properties for each gridpoint (e.g., land and ocean temperatures, albedos and roughness lengths are averaged, based on their respective area fractions), so the flux is calculated only once per gridpoint of the grid where we calculate fluxes. This is computationally faster, but it makes the fluxes on surface boundaries more diffuse. Currently, we use this method for calculating radiative fluxes in the atmosphere, and turbulent fluxes in the coupler (on the atmospheric grid). The fluxes are calculated in the atmospheric (host) model's cache, which can be setup to avoid allocating coupler fields.
+2. **Partitioned fluxes** (called with `PartitionedComponentModelGrid()`)
+  - these are calculated separately for each surface type. It is then the fluxes (rather than the surface states) that are combined and passed to the atmospheric model as a boundary condition. This method ensures that each surface model receives fluxes that correspond to its state properties, resulting in a more accurate model evolution. (TODO: implementation in the next PR)
+
+## FluxCalculator API
+
+```@docs
+    ClimaCoupler.FluxCalculator.TurbulentFluxPartition
+    ClimaCoupler.FluxCalculator.PartitionedComponentModelGrid
+    ClimaCoupler.FluxCalculator.CombinedAtmosGrid
+    ClimaCoupler.FluxCalculator.compute_combined_turbulent_fluxes!
+    ClimaCoupler.FluxCalculator.compute_atmos_turbulent_fluxes!
+
+```

experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl

@@ -1,5 +1,6 @@
 # atmos_init: for ClimaAtmos pre-AMIP interface
 import ClimaAtmos
+import ClimaCoupler.FluxCalculator: compute_atmos_turbulent_fluxes!
 
 driver_file = joinpath(pkgdir(ClimaAtmos), "examples", "hybrid", "driver.jl")
 ENV["CI_PERF_SKIP_RUN"] = true
@@ -11,7 +12,7 @@ catch err
     end
 end
 # the clima atmos `integrator` is now defined
-struct AtmosSimulation{P, Y, D, I} <: AtmosModelSimulation
+struct ClimaAtmosSimulation{P, Y, D, I} <: AtmosModelSimulation
     params::P
     Y_init::Y
     domain::D
@@ -26,5 +27,72 @@ function atmos_init(::Type{FT}, Y, integrator; params = nothing) where {FT}
         @warn("Running with ρe_int in coupled mode is not tested yet.")
     end
 
-    AtmosSimulation(params, Y, spaces, integrator)
+    ClimaAtmosSimulation(params, Y, spaces, integrator)
+end
+
+# required by the Interfacer
+get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux}) = level(sim.integrator.p.ᶠradiation_flux, half)
+get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) =
+    sim.integrator.p.col_integrated_rain .+ sim.integrator.p.col_integrated_snow # all fallen snow melts for now
+get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) = sim.integrator.p.col_integrated_snow .* FT(0)
+
+get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) = sim.integrator.p.ρ_dif_flux_h_tot
+get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) = sim.integrator.p.ρ_dif_flux_q_tot
+
+function update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, field)
+    sim.integrator.p.radiation_model.surface_temperature .= RRTMGPI.field2array(field)
+    parent(sim.integrator.p.T_sfc) .= parent(field)
+
+end
+function update_field!(sim::ClimaAtmosSimulation, ::Val{:albedo}, field)
+    sim.integrator.p.radiation_model.diffuse_sw_surface_albedo .=
+        reshape(RRTMGPI.field2array(field), 1, length(parent(field)))
+    sim.integrator.p.radiation_model.direct_sw_surface_albedo .=
+        reshape(RRTMGPI.field2array(field), 1, length(parent(field)))
+end
+function update_field!(sim::ClimaAtmosSimulation, ::Val{:roughness_momentum}, field)
+    parent(sim.integrator.p.sfc_inputs.z0m) .= parent(field)
+end
+function update_field!(sim::ClimaAtmosSimulation, ::Val{:roughness_buoyancy}, field)
+    parent(sim.integrator.p.sfc_inputs.z0b) .= parent(field)
+end
+
+step!(sim::ClimaAtmosSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
+
+reinit!(sim::ClimaAtmosSimulation) = reinit!(sim.integrator)
+
+
+# flux calculation borrowed from atmos
+function compute_atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
+    # TODO: dependence on atmos will be removed
+
+    thermo_params = CAP.thermodynamics_params(atmos_sim.integrator.p.params)
+
+    # calculate turbulent fluxes on atmos grid and save in atmos cache
+    update_field!(atmos_sim, Val(:surface_temperature), csf.T_S)
+
+    if :z0b in propertynames(atmos_sim.integrator.p.surface_scheme)
+        update_field!(atmos_sim, Val(:roughness_momentum), csf.z0m_S)
+        update_field!(atmos_sim, Val(:roughness_buoyancy), csf.z0b_S)
+    end
+
+    Fields.bycolumn(axes(atmos_sim.integrator.p.ts_sfc)) do colidx
+        ClimaAtmos.set_surface_thermo_state!(
+            ClimaAtmos.Decoupled(),
+            atmos_sim.integrator.p.surface_scheme.sfc_thermo_state_type,
+            atmos_sim.integrator.p.ts_sfc[colidx],
+            atmos_sim.integrator.p.T_sfc[colidx],
+            Spaces.level(atmos_sim.integrator.p.ᶜts[colidx], 1),
+            thermo_params,
+            atmos_sim.integrator.t,
+        )
+
+        get_surface_fluxes!(
+            atmos_sim.integrator.u,
+            atmos_sim.integrator.p,
+            atmos_sim.integrator.t,
+            colidx,
+            atmos_sim.integrator.p.atmos.vert_diff,
+        )
+    end
 end

experiments/AMIP/modular/components/land/bucket_init.jl

@@ -19,12 +19,15 @@ import ClimaLSM.Bucket:
 using ClimaLSM:
     make_ode_function, initialize, obtain_surface_space, make_set_initial_aux_state, surface_evaporative_scaling
 
+import ClimaCoupler.Interfacer: LandModelSimulation, get_field, update_field!
+import ClimaCoupler.FieldExchanger: step!, reinit!
+
 """
     BucketSimulation{M, Y, D, I}
 
 The bucket model simulation object.
 """
-struct BucketSimulation{M, Y, D, I, A} <: SurfaceModelSimulation
+struct BucketSimulation{M, Y, D, I, A} <: LandModelSimulation
     model::M
     Y_init::Y
     domain::D

experiments/AMIP/modular/components/land/bucket_utils.jl

@@ -152,3 +152,32 @@ get_field(sim::BucketSimulation, ::Val{:beta}) =
 get_field(sim::BucketSimulation, ::Val{:albedo}) =
     ClimaLSM.surface_albedo(sim.model, sim.integrator.u, sim.integrator.p)
 get_field(sim::BucketSimulation, ::Val{:area_fraction}) = sim.area_fraction
+
+
+# The surface air density is computed using the atmospheric state at the first level and making ideal gas
+# and hydrostatic balance assumptions. The land model does not compute the surface air density so this is
+# a reasonable stand-in.
+function update_field!(sim::BucketSimulation, ::Val{:air_density}, field)
+    parent(sim.integrator.p.bucket.ρ_sfc) .= parent(field)
+end
+
+function update_field!(sim::BucketSimulation, ::Val{:turbulent_energy_flux}, field)
+    parent(sim.integrator.p.bucket.turbulent_energy_flux) .= parent(field)
+end
+function update_field!(sim::BucketSimulation, ::Val{:turbulent_moisture_flux}, field)
+    ρ_liq = (LSMP.ρ_cloud_liq(sim.model.parameters.earth_param_set))
+    parent(sim.integrator.p.bucket.evaporation) .= parent(field ./ ρ_liq)
+end
+function update_field!(sim::BucketSimulation, ::Val{:radiative_energy_flux}, field)
+    parent(sim.integrator.p.bucket.R_n) .= parent(field)
+end
+function update_field!(sim::BucketSimulation, ::Val{:liquid_precipitation}, field)
+    parent(sim.integrator.p.bucket.P_liq) .= parent(field)
+end
+function update_field!(sim::BucketSimulation, ::Val{:snow_precipitation}, field)
+    parent(sim.integrator.p.bucket.P_snow) .= parent(field)
+end
+
+step!(sim::BucketSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
+
+reinit!(sim::BucketSimulation) = reinit!(sim.integrator)

experiments/AMIP/modular/components/ocean/slab_ocean_init.jl

@@ -1,14 +1,17 @@
 using ClimaCore
 
+import ClimaCoupler.Interfacer: OceanModelSimulation, get_field, update_field!
+import ClimaCoupler.FieldExchanger: step!, reinit!
+
 """
     SlabOceanSimulation{P, Y, D, I}
 
 Equation:
 
-    (h * ρ * c) dTdt = -(F_aero + F_rad)
+    (h * ρ * c) dTdt = -(F_turb_energy + F_radiative)
 
 """
-struct SlabOceanSimulation{P, Y, D, I} <: SurfaceModelSimulation
+struct SlabOceanSimulation{P, Y, D, I} <: OceanModelSimulation
     params::P
     Y_init::Y
     domain::D
@@ -52,9 +55,9 @@ end
 
 # ode
 function slab_ocean_rhs!(dY, Y, cache, t)
-    p, F_aero, F_rad, area_fraction = cache
+    p, F_turb_energy, F_radiative, area_fraction = cache
     FT = eltype(Y.T_sfc)
-    rhs = @. -(F_aero + F_rad) / (p.h * p.ρ * p.c)
+    rhs = @. -(F_turb_energy + F_radiative) / (p.h * p.ρ * p.c)
     parent(dY.T_sfc) .= parent(rhs .* Regridder.binary_mask.(area_fraction, threshold = eps()))
 end
 
@@ -70,8 +73,8 @@ function ocean_init(::Type{FT}; tspan, dt, saveat, space, area_fraction, stepper
     Y, space = slab_ocean_space_init(FT, space, params)
     cache = (
         params = params,
-        F_aero = ClimaCore.Fields.zeros(space),
-        F_rad = ClimaCore.Fields.zeros(space),
+        F_turb_energy = ClimaCore.Fields.zeros(space),
+        F_radiative = ClimaCore.Fields.zeros(space),
         area_fraction = area_fraction,
     )
     problem = OrdinaryDiffEq.ODEProblem(slab_ocean_rhs!, Y, tspan, cache)
@@ -98,3 +101,14 @@ get_field(sim::SlabOceanSimulation, ::Val{:area_fraction}) = sim.integrator.p.ar
 function update_field!(sim::SlabOceanSimulation, ::Val{:area_fraction}, field::Fields.Field)
     sim.integrator.p.area_fraction .= field
 end
+
+function update_field!(sim::SlabOceanSimulation, ::Val{:turbulent_energy_flux}, field)
+    parent(sim.integrator.p.F_turb_energy) .= parent(field)
+end
+function update_field!(sim::SlabOceanSimulation, ::Val{:radiative_energy_flux}, field)
+    parent(sim.integrator.p.F_radiative) .= parent(field)
+end
+
+step!(sim::SlabOceanSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
+
+reinit!(sim::SlabOceanSimulation) = reinit!(sim.integrator)

experiments/AMIP/modular/components/ocean/slab_seaice_init.jl

@@ -1,20 +1,26 @@
+import ClimaCoupler.Interfacer: SeaIceModelSimulation, get_field, update_field!
+import ClimaCoupler.FieldExchanger: step!, reinit!
+
 """
     PrescribedIceSimulation{P, Y, D, I}
 
 Ice concentration is prescribed, and we solve the following energy equation:
 
-    (h * ρ * c) d T_sfc dt = -(F_aero + F_rad) + F_conductive
+    (h * ρ * c) d T_sfc dt = -(F_turb_energy + F_radiativead) + F_conductive
 
     with
     F_conductive = k_ice (T_base - T_sfc) / (h)
 
     The surface temperature (`T_sfc`) is the prognostic variable which is being
-    modified by turbulent aerodynamic (`F_aero`) and radiative (`F_rad`) fluxes,
+    modified by turbulent aerodynamic (`F_turb_energy`) and radiative (`F_turb_energy`) fluxes,
     as well as a conductive flux that depends on the temperature difference
     across the ice layer (with `T_base` being prescribed).
 
+In the current version, the sea ice has a prescribed thickness, and we assume that it is not
+sublimating. That contribution has been zeroed out in the atmos fluxes.
+
 """
-struct PrescribedIceSimulation{P, Y, D, I} <: SurfaceModelSimulation
+struct PrescribedIceSimulation{P, Y, D, I} <: SeaIceModelSimulation
     params::P
     Y_init::Y
     domain::D
@@ -55,13 +61,13 @@ function ice_rhs!(du, u, p, _)
     FT = eltype(dY)
 
     params = p.params
-    F_aero = p.F_aero
-    F_rad = p.F_rad
+    F_turb_energy = p.F_turb_energy
+    F_radiative = p.F_radiative
     area_fraction = p.area_fraction
     T_freeze = params.T_freeze
 
     F_conductive = @. params.k_ice / (params.h) * (params.T_base - Y.T_sfc) # fluxes are defined to be positive when upward
-    rhs = @. (-F_aero - F_rad + F_conductive) / (params.h * params.ρ * params.c)
+    rhs = @. (-F_turb_energy - F_radiative + F_conductive) / (params.h * params.ρ * params.c)
 
     # do not count tendencies that lead to temperatures above freezing, and mask out no-ice areas
     unphysical = @. Regridder.binary_mask.(T_freeze - (Y.T_sfc + FT(rhs) * p.dt), threshold = FT(0)) .*
@@ -80,8 +86,8 @@ function ice_init(::Type{FT}; tspan, saveat, dt, space, area_fraction, stepper =
 
     Y = slab_ice_space_init(FT, space, params)
     additional_cache = (;
-        F_aero = ClimaCore.Fields.zeros(space),
-        F_rad = ClimaCore.Fields.zeros(space),
+        F_turb_energy = ClimaCore.Fields.zeros(space),
+        F_radiative = ClimaCore.Fields.zeros(space),
         area_fraction = area_fraction,
         dt = dt,
     )
@@ -115,3 +121,13 @@ get_field(sim::PrescribedIceSimulation, ::Val{:area_fraction}) = sim.integrator.
 function update_field!(sim::PrescribedIceSimulation, ::Val{:area_fraction}, field::Fields.Field)
     sim.integrator.p.area_fraction .= field
 end
+
+function update_field!(sim::PrescribedIceSimulation, ::Val{:turbulent_energy_flux}, field)
+    parent(sim.integrator.p.F_turb_energy) .= parent(field)
+end
+function update_field!(sim::PrescribedIceSimulation, ::Val{:radiative_energy_flux}, field)
+    parent(sim.integrator.p.F_radiative) .= parent(field)
+end
+
+step!(sim::PrescribedIceSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
+reinit!(sim::PrescribedIceSimulation) = reinit!(sim.integrator)