Add flux partitioning src + unit tests

implemented in AMIP driver, conserving deps amip runs + custom sfc state new pipeline run clean up clean up revs rev2 rev3 clean format
CliMA · Aug 7, 2023 · 1bb8fe7 · 1bb8fe7
1 parent 159252f
commit 1bb8fe7
Show file tree

Hide file tree

Showing 20 changed files with 970 additions and 278 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -139,6 +139,13 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: "Slabplanet: partitioned turbulent fluxes"
+        key: "slabplanet_partitioned_fluxes"
+        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --run_name slabplanet_partitioned_fluxes --FLOAT_TYPE Float64 --coupled true --turb_flux_partition PartitionedStateFluxes --surface_setup PrescribedSurface  --moist equil --vert_diff true --rad gray --energy_check true --mode_name slabplanet --t_end 10days --dt_save_to_sol 9days --dt_cpl 200 --dt 200secs --mono_surface true --h_elem 4 --precip_model 0M --anim true --job_id slabplanet_partitioned_fluxes"
+        artifact_paths: "experiments/AMIP/modular/output/slabplanet/slabplanet_partitioned_fluxes_artifacts/*"
+        agents:
+          slurm_mem: 20GB
+
       # Test: non-monotonous remapping for land mask
       - label: "Slabplanet: non-monotonous surface remap"
         key: "slabplanet_non-monotonous"

diff --git a/Project.toml b/Project.toml
@@ -4,6 +4,7 @@ authors = ["CliMA Contributors <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
 ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
@@ -18,6 +19,7 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"
 TempestRemap_jll = "8573a8c5-1df0-515e-a024-abad257ee284"
@@ -26,6 +28,7 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [compat]
+CLIMAParameters = "0.7"
 ClimaAtmos = "0.15"
 ClimaComms = "0.5"
 ClimaCore = "0.10"
@@ -39,6 +42,7 @@ OrdinaryDiffEq = "5, 6"
 Plots = "1"
 PrettyTables = "1, 2"
 SciMLBase = "1"
+StaticArrays = "1"
 SurfaceFluxes = "0.6"
 TempestRemap_jll = "2"
 TerminalLoggers = "0.1"

diff --git a/docs/src/fluxcalculator.md b/docs/src/fluxcalculator.md
@@ -3,18 +3,18 @@
 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()`)
+1. **Combined fluxes** (called with `CombinedStateFluxes()`)
   - 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)
+2. **Partitioned fluxes** (called with `PartitionedStateFluxes()`)
+  - 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. However, it is more computationally expensive, and requires more communication between the component models.
 
 ## FluxCalculator API
 
 ```@docs
     ClimaCoupler.FluxCalculator.TurbulentFluxPartition
-    ClimaCoupler.FluxCalculator.PartitionedComponentModelGrid
-    ClimaCoupler.FluxCalculator.CombinedAtmosGrid
-    ClimaCoupler.FluxCalculator.compute_combined_turbulent_fluxes!
-    ClimaCoupler.FluxCalculator.compute_atmos_turbulent_fluxes!
+    ClimaCoupler.FluxCalculator.PartitionedStateFluxes
+    ClimaCoupler.FluxCalculator.CombinedStateFluxes
+    ClimaCoupler.FluxCalculator.combined_turbulent_fluxes!
+    ClimaCoupler.FluxCalculator.atmos_turbulent_fluxes!
 
 ```
diff --git a/experiments/AMIP/modular/Manifest.toml b/experiments/AMIP/modular/Manifest.toml
@@ -239,7 +239,7 @@ uuid = "d934ef94-cdd4-4710-83d6-720549644b70"
 version = "0.3.5"
 
 [[deps.ClimaCoupler]]
-deps = ["ClimaAtmos", "ClimaComms", "ClimaCore", "ClimaCoreTempestRemap", "ClimaLSM", "Dates", "DocStringExtensions", "Insolation", "JLD2", "NCDatasets", "OrdinaryDiffEq", "Plots", "PrettyTables", "SciMLBase", "Statistics", "SurfaceFluxes", "TempestRemap_jll", "TerminalLoggers", "Thermodynamics", "UnPack"]
+deps = ["CLIMAParameters", "ClimaAtmos", "ClimaComms", "ClimaCore", "ClimaCoreTempestRemap", "ClimaLSM", "Dates", "DocStringExtensions", "Insolation", "JLD2", "NCDatasets", "OrdinaryDiffEq", "Plots", "PrettyTables", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "TempestRemap_jll", "TerminalLoggers", "Thermodynamics", "UnPack"]
 path = "../../.."
 uuid = "4ade58fe-a8da-486c-bd89-46df092ec0c7"
 version = "0.1.0"

diff --git a/experiments/AMIP/modular/cli_options.jl b/experiments/AMIP/modular/cli_options.jl
@@ -27,6 +27,10 @@ function argparse_settings()
         help = "Boolean flag indicating whether (1st order) monotone and conservative remapping is applied."
         arg_type = Bool
         default = false
+        "--turb_flux_partition"
+        help = "Method to partition turbulent fluxes. [`PartitionedStateFluxes`, `CombinedStateFluxes`]"
+        arg_type = String
+        default = "CombinedStateFluxes"
         "--albedo_from_file"
         help = "Access land surface albedo information from data file"
         arg_type = Bool

diff --git a/experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl b/experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl
@@ -1,10 +1,17 @@
 # atmos_init: for ClimaAtmos pre-AMIP interface
 import ClimaAtmos
 using ClimaAtmos: RRTMGPI
-import ClimaCoupler.FluxCalculator: compute_atmos_turbulent_fluxes!, calculate_surface_air_density
+import ClimaAtmos: CT1, CT2, CT12, CT3, C3, C12, unit_basis_vector_data, ⊗
+import ClimaCoupler.FluxCalculator:
+    atmos_turbulent_fluxes!,
+    calculate_surface_air_density,
+    PartitionedStateFluxes,
+    extrapolate_ρ_to_sfc,
+    get_surface_params
 using ClimaCore: Fields.level, Geometry
 import ClimaCoupler.FieldExchanger: get_thermo_params
 import ClimaCoupler.Interfacer: get_field, update_field!, name, get_model_state_vector
+using StaticArrays
 
 # the clima atmos `integrator` is now defined
 struct ClimaAtmosSimulation{P, Y, D, I} <: AtmosModelSimulation
@@ -39,7 +46,7 @@ function atmos_init(::Type{FT}, parsed_args::Dict) where {FT}
     ClimaAtmosSimulation(integrator.p.params, Y, spaces, integrator)
 end
 
-# required by the Interfacer
+# extensions required by the Interfacer
 get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux}) =
     Fields.level(sim.integrator.p.ᶠradiation_flux, half)
 get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) =
@@ -53,33 +60,74 @@ get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) =
 
 get_field(sim::ClimaAtmosSimulation, ::Val{:thermo_state_int}) = Spaces.level(sim.integrator.p.ᶜts, 1)
 
+# extensions required by FluxCalculator (partitioned fluxes)
+get_field(sim::ClimaAtmosSimulation, ::Val{:height_int}) =
+    Spaces.level(Fields.coordinate_field(sim.integrator.u.c).z, 1)
+get_field(sim::ClimaAtmosSimulation, ::Val{:height_sfc}) =
+    Spaces.level(Fields.coordinate_field(sim.integrator.u.f).z, half)
+function get_field(sim::ClimaAtmosSimulation, ::Val{:uv_int})
+    uₕ_int = Geometry.UVVector.(Spaces.level(sim.integrator.u.c.uₕ, 1))
+    return @. StaticArrays.SVector(uₕ_int.components.data.:1, uₕ_int.components.data.:2)
+end
+get_field(sim::ClimaAtmosSimulation, ::Val{:air_density}) = TD.air_density.(thermo_params, sim.integrator.p.ᶜts)
+get_field(sim::ClimaAtmosSimulation, ::Val{:air_temperature}) = TD.air_temperature.(thermo_params, sim.integrator.p.ᶜts)
+get_field(sim::ClimaAtmosSimulation, ::Val{:cv_m}) = TD.cv_m.(thermo_params, sim.integrator.p.ᶜts)
+get_field(sim::ClimaAtmosSimulation, ::Val{:gas_constant_air}) =
+    TD.gas_constant_air.(thermo_params, sim.integrator.p.ᶜts)
+
+get_surface_params(sim::ClimaAtmosSimulation) = CAP.surface_fluxes_params(sim.integrator.p.params)
 
-function update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, field)
-    sim.integrator.p.radiation_model.surface_temperature .= RRTMGPI.field2array(field)
+# extensions required by the Interfacer
+function update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, csf)
+    sim.integrator.p.radiation_model.surface_temperature .= RRTMGPI.field2array(csf.T_S)
 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
 
-step!(sim::ClimaAtmosSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
+# get_surface_params required by FluxCalculator (partitioned fluxes)
+function update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fluxes}, fields)
+    (; F_turb_energy, F_turb_moisture, F_turb_ρτxz, F_turb_ρτyz) = fields
+
+    Y = sim.integrator.u
+    surface_local_geometry = Fields.level(Fields.local_geometry_field(Y.f), Fields.half)
+    surface_normal = @. C3(unit_basis_vector_data(C3, surface_local_geometry))
+
+    # get template objects for the contravariant components of the momentum fluxes (required by Atmos boundary conditions)
+    vec_ct12_ct1 = @. CT12(CT2(unit_basis_vector_data(CT1, surface_local_geometry)), surface_local_geometry)
+    vec_ct12_ct2 = @. CT12(CT2(unit_basis_vector_data(CT2, surface_local_geometry)), surface_local_geometry)
+
+    sim.integrator.p.sfc_conditions.ρ_flux_uₕ .= (
+        surface_normal .⊗
+        C12.(
+            swap_space!(ones(axes(vec_ct12_ct1)), F_turb_ρτxz) .* vec_ct12_ct1 .+
+            swap_space!(ones(axes(vec_ct12_ct2)), F_turb_ρτyz) .* vec_ct12_ct2,
+            surface_local_geometry,
+        )
+    )
 
-reinit!(sim::ClimaAtmosSimulation) = reinit!(sim.integrator)
+    parent(sim.integrator.p.sfc_conditions.ρ_flux_h_tot) .= parent(F_turb_energy) .* parent(surface_normal) # (shf + lhf)
+    parent(sim.integrator.p.sfc_conditions.ρ_flux_q_tot) .= parent(F_turb_moisture) .* parent(surface_normal) # (evap)
 
-"""
-    update_sim!(atmos_sim::ClimaAtmosSimulation, csf)
+    # TODO: see if Atmos can rever to a simpler solution
+end
 
-Updates the surface fields for temperature and albedo.
+# extensions required by FieldExchanger
+step!(sim::ClimaAtmosSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
+
+reinit!(sim::ClimaAtmosSimulation) = reinit!(sim.integrator)
 
-# Arguments
-- `atmos_sim`: [ClimaAtmosSimulation] containing an atmospheric model simulation object.
-- `csf`: [NamedTuple] containing coupler fields.
-"""
 function update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
     update_field!(atmos_sim, Val(:albedo), csf.albedo)
-    update_field!(atmos_sim, Val(:surface_temperature), csf.T_S)
+    update_field!(atmos_sim, Val(:surface_temperature), csf)
+
+    if turbulent_fluxes isa PartitionedStateFluxes
+        update_field!(atmos_sim, Val(:turbulent_fluxes), csf)
+    end
 end
 
 
@@ -112,35 +160,48 @@ function coupler_surface_setup(
 end
 
 """
-    compute_atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
+    get_new_cache(atmos_sim::ClimaAtmosSimulation, csf)
 
-Computes turbulent surface fluxes using ClimaAtmos's `update_surface_conditions!`. This
-requires that we define a new temporary parameter Tuple, `new_p`, and save the new surface state
-in it. We do not want `new_p` to live in the atmospheric model permanently, because that would also
-trigger flux calculation during Atmos `step!`. We only want to trigger this once per coupling
-timestep from ClimaCoupler.
+Returns a new `p` with the updated surface conditions.
 """
+function get_new_cache(atmos_sim::ClimaAtmosSimulation, csf)
 
-function compute_atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
+    p = atmos_sim.integrator.p
 
-    function modified_atmos_p(atmos_sim, csf_sfc)
+    csf_sfc = (; T = csf.T_S, z0m = csf.z0m_S, z0b = csf.z0b_S, beta = csf.beta, q_vap = csf.q_sfc)
+    modified_atmos_cache(atmos_sim, csf_sfc)
+end
 
-        p = atmos_sim.integrator.p
+"""
+    modified_atmos_cache(atmos_sim, csf_sfc)
 
-        coupler_sfc_setup = coupler_surface_setup(CoupledMoninObukhov(), p; csf_sfc = csf_sfc)
+Returns a new `p` with the updated surface conditions.
+"""
+function modified_atmos_cache(atmos_sim, csf_sfc)
 
-        p_names = propertynames(p)
-        p_values = map(x -> x == :sfc_setup ? coupler_sfc_setup : getproperty(p, x), p_names) # TODO: use merge here
+    p = atmos_sim.integrator.p
 
-        (; zip(p_names, p_values)...)
-    end
+    coupler_sfc_setup = coupler_surface_setup(CoupledMoninObukhov(), p; csf_sfc = csf_sfc)
 
-    p = atmos_sim.integrator.p
+    p_names = propertynames(p)
+    p_values = map(x -> x == :sfc_setup ? coupler_sfc_setup : getproperty(p, x), p_names) # TODO: use merge here
 
-    csf_sfc = (; T = csf.T_S, z0m = csf.z0m_S, z0b = csf.z0b_S, beta = csf.beta, q_vap = csf.q_sfc)
-    new_p = modified_atmos_p(atmos_sim, csf_sfc)
+    (; zip(p_names, p_values)...)
+end
+
+"""
+    atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
+
+Computes turbulent surface fluxes using ClimaAtmos's `update_surface_conditions!`. This
+requires that we define a new temporary parameter Tuple, `new_p`, and save the new surface state
+in it. We do not want `new_p` to live in the atmospheric model permanently, because that would also
+trigger flux calculation during Atmos `step!`. We only want to trigger this once per coupling
+timestep from ClimaCoupler.
+"""
+function atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
+    new_p = get_new_cache(atmos_sim, csf)
     ClimaAtmos.SurfaceConditions.update_surface_conditions!(atmos_sim.integrator.u, new_p, atmos_sim.integrator.t)
-    p.sfc_conditions .= new_p.sfc_conditions
+    atmos_sim.integrator.p.sfc_conditions .= new_p.sfc_conditions
 end
 
 """

diff --git a/experiments/AMIP/modular/components/land/bucket_init.jl b/experiments/AMIP/modular/components/land/bucket_init.jl
@@ -21,6 +21,7 @@ using ClimaLSM:
 
 import ClimaCoupler.Interfacer: LandModelSimulation, get_field, update_field!, name
 import ClimaCoupler.FieldExchanger: step!, reinit!
+import ClimaCoupler.FluxCalculator: update_turbulent_fluxes_point!, surface_thermo_state
 
 """
     BucketSimulation{M, Y, D, I}
@@ -85,8 +86,6 @@ function ClimaLSM.Bucket.surface_fluxes(
     )
 end
 
-
-
 """
     net_radiation(radiation::CoupledRadiativeFluxes{FT},
                     model::BucketModel{FT},

diff --git a/experiments/AMIP/modular/components/land/bucket_utils.jl b/experiments/AMIP/modular/components/land/bucket_utils.jl
@@ -80,7 +80,7 @@ function make_lsm_domain(
     )
 end
 
-# required by Interfacer
+# extensions required by Interfacer
 get_field(sim::BucketSimulation, ::Val{:surface_temperature}) =
     ClimaLSM.surface_temperature(sim.model, sim.integrator.u, sim.integrator.p, sim.integrator.t)
 get_field(sim::BucketSimulation, ::Val{:surface_humidity}) =
@@ -92,11 +92,7 @@ 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.
+get_field(sim::BucketSimulation, ::Val{:air_density}) = sim.integrator.p.bucket.ρ_sfc
 
 function update_field!(sim::BucketSimulation, ::Val{:turbulent_energy_flux}, field)
     parent(sim.integrator.p.bucket.turbulent_energy_flux) .= parent(field)
@@ -123,6 +119,32 @@ step!(sim::BucketSimulation, t) = step!(sim.integrator, t - sim.integrator.t, tr
 
 reinit!(sim::BucketSimulation) = reinit!(sim.integrator)
 
+# extensions required by FluxCalculator (partitioned fluxes)
+function update_turbulent_fluxes_point!(sim::BucketSimulation, fields::NamedTuple, colidx::Fields.ColumnIndex)
+    (; F_turb_energy, F_turb_moisture) = fields
+    sim.integrator.p.bucket.turbulent_energy_flux[colidx] .= F_turb_energy
+    sim.integrator.p.bucket.evaporation[colidx] .=
+        F_turb_moisture ./ LSMP.ρ_cloud_liq(sim.model.parameters.earth_param_set)
+    return nothing
+end
+
+# extension of FluxCalculator.surface_thermo_state, overriding the saturated-surface default
+function surface_thermo_state(
+    sim::BucketSimulation,
+    thermo_params::TD.Parameters.ThermodynamicsParameters,
+    thermo_state_int,
+    colidx::ClimaCore.Fields.ColumnIndex,
+)
+
+    T_sfc = get_field(sim, Val(:surface_temperature), colidx)
+    # Note that the surface air density, ρ_sfc, 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.
+    ρ_sfc = get_field(sim, Val(:air_density), colidx)
+    q_sfc = get_field(sim, Val(:surface_humidity), colidx) # already calculated in rhs! (cache)
+    @. TD.PhaseEquil_ρTq.(thermo_params, ρ_sfc, T_sfc, q_sfc)
+end
+
 """
     get_model_state_vector(sim::BucketSimulation)
 

diff --git a/experiments/AMIP/modular/components/ocean/slab_ocean_init.jl b/experiments/AMIP/modular/components/ocean/slab_ocean_init.jl
@@ -2,6 +2,7 @@ using ClimaCore
 
 import ClimaCoupler.Interfacer: OceanModelSimulation, get_field, update_field!, name
 import ClimaCoupler.FieldExchanger: step!, reinit!
+import ClimaCoupler.FluxCalculator: update_turbulent_fluxes_point!
 
 """
     SlabOceanSimulation{P, Y, D, I}
@@ -17,6 +18,7 @@ struct SlabOceanSimulation{P, Y, D, I} <: OceanModelSimulation
     domain::D
     integrator::I
 end
+name(::SlabOceanSimulation) = "SlabOceanSimulation"
 
 # ocean parameters
 struct OceanSlabParameters{FT <: AbstractFloat}
@@ -103,14 +105,15 @@ Ensures that the space of the SST struct matches that of the land_fraction, and
 """
 clean_sst(SST, _info) = (swap_space!(zeros(axes(_info.land_fraction)), SST) .+ float_type_bcf(_info)(273.15))
 
-# required by Interfacer
+# extensions required by Interfacer
 get_field(sim::SlabOceanSimulation, ::Val{:surface_temperature}) = sim.integrator.u.T_sfc
 get_field(sim::SlabOceanSimulation, ::Val{:surface_humidity}) = sim.integrator.p.q_sfc
 get_field(sim::SlabOceanSimulation, ::Val{:roughness_momentum}) = sim.integrator.p.params.z0m
 get_field(sim::SlabOceanSimulation, ::Val{:roughness_buoyancy}) = sim.integrator.p.params.z0b
 get_field(sim::SlabOceanSimulation, ::Val{:beta}) = convert(eltype(sim.integrator.u), 1.0)
 get_field(sim::SlabOceanSimulation, ::Val{:albedo}) = sim.integrator.p.params.α
 get_field(sim::SlabOceanSimulation, ::Val{:area_fraction}) = sim.integrator.p.area_fraction
+get_field(sim::SlabOceanSimulation, ::Val{:air_density}) = sim.integrator.p.ρ_sfc
 
 function update_field!(sim::SlabOceanSimulation, ::Val{:area_fraction}, field::Fields.Field)
     sim.integrator.p.area_fraction .= field
@@ -126,10 +129,17 @@ function update_field!(sim::SlabOceanSimulation, ::Val{:air_density}, field)
     parent(sim.integrator.p.ρ_sfc) .= parent(field)
 end
 
+# extensions required by FieldExchanger
 step!(sim::SlabOceanSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)
 
 reinit!(sim::SlabOceanSimulation) = reinit!(sim.integrator)
 
+# extensions required by FluxCalculator (partitioned fluxes)
+function update_turbulent_fluxes_point!(sim::SlabOceanSimulation, fields::NamedTuple, colidx::Fields.ColumnIndex)
+    (; F_turb_energy) = fields
+    @. sim.integrator.p.F_turb_energy[colidx] = F_turb_energy
+end
+
 """
     get_model_state_vector(sim::SlabOceanSimulation)