From 653609c5015a2353683f217a06bfa25590f46560 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Mon, 20 May 2024 13:45:04 -0700 Subject: [PATCH] rm bycolumn --- docs/src/fluxcalculator.md | 4 +- docs/src/interfacer.md | 2 +- .../atmosphere/climaatmos_extra_diags.jl | 13 +- .../components/land/climaland_bucket.jl | 21 +- .../components/ocean/eisenman_seaice.jl | 25 +- .../components/ocean/prescr_seaice.jl | 8 +- .../ClimaEarth/components/ocean/slab_ocean.jl | 8 +- src/FluxCalculator.jl | 248 ++++++++++-------- src/Interfacer.jl | 13 - src/surface_stub.jl | 5 - .../eisenman_seaice_tests.jl | 37 +-- test/flux_calculator_tests.jl | 56 ++-- test/interfacer_tests.jl | 32 +-- test/runtests.jl | 2 +- 14 files changed, 194 insertions(+), 280 deletions(-) diff --git a/docs/src/fluxcalculator.md b/docs/src/fluxcalculator.md index 83f0d6e45..cec79caea 100644 --- a/docs/src/fluxcalculator.md +++ b/docs/src/fluxcalculator.md @@ -22,8 +22,8 @@ Fluxes over a heterogeneous surface (e.g., from a gridpoint where atmospheric ce ClimaCoupler.FluxCalculator.get_surface_params ClimaCoupler.FluxCalculator.partitioned_turbulent_fluxes! ClimaCoupler.FluxCalculator.differentiate_turbulent_fluxes! - ClimaCoupler.FluxCalculator.get_surface_fluxes_point! - ClimaCoupler.FluxCalculator.update_turbulent_fluxes_point! + ClimaCoupler.FluxCalculator.get_surface_fluxes! + ClimaCoupler.FluxCalculator.update_turbulent_fluxes! ClimaCoupler.FluxCalculator.extrapolate_ρ_to_sfc ClimaCoupler.FluxCalculator.surface_thermo_state ClimaCoupler.FluxCalculator.water_albedo_from_atmosphere! diff --git a/docs/src/interfacer.md b/docs/src/interfacer.md index 471319a56..112c8b9be 100644 --- a/docs/src/interfacer.md +++ b/docs/src/interfacer.md @@ -153,7 +153,7 @@ following properties: | `surface_diffuse albedo` | bulk diffuse surface albedo; needed if calculated externally of the surface model (e.g. ocean albedo from the atmospheric state) | | ### SurfaceModelSimulation - optional functions -- `update_turbulent_fluxes_point!(::ComponentModelSimulation, fields::NamedTuple, colidx)`: +- `update_turbulent_fluxes!(::ComponentModelSimulation, fields::NamedTuple)`: This function updates the turbulent fluxes of the component model simulation at this point in horizontal space. The values are updated using the energy and moisture turbulent fluxes stored in fields which are calculated by the diff --git a/experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl b/experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl index 66cdd146a..013c23afe 100644 --- a/experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl +++ b/experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl @@ -65,16 +65,9 @@ CAD.add_diagnostic_variable!( (; C_E) = cache.atmos.vert_diff interior_uₕ = CC.Fields.level(state.c.uₕ, 1) ᶠp = ᶠK_E = cache.scratch.ᶠtemp_scalar - CC.Fields.bycolumn(axes(ᶜp)) do colidx - @. ᶠp[colidx] = CAD.ᶠinterp(ᶜp[colidx]) - ᶜΔz_surface = CC.Fields.Δz_field(interior_uₕ) - @. ᶠK_E[colidx] = CA.eddy_diffusivity_coefficient( - C_E, - CA.norm(interior_uₕ[colidx]), - ᶜΔz_surface[colidx] / 2, - ᶠp[colidx], - ) - end + @. ᶠp = CAD.ᶠinterp(ᶜp) + ᶜΔz_surface = CC.Fields.Δz_field(interior_uₕ) + @. ᶠK_E = CA.eddy_diffusivity_coefficient(C_E, CA.norm(interior_uₕ), ᶜΔz_surface / 2, ᶠp) if isnothing(out) return CAD.ᶜinterp.(ᶠK_E) else diff --git a/experiments/ClimaEarth/components/land/climaland_bucket.jl b/experiments/ClimaEarth/components/land/climaland_bucket.jl index 973710adc..9e476b933 100644 --- a/experiments/ClimaEarth/components/land/climaland_bucket.jl +++ b/experiments/ClimaEarth/components/land/climaland_bucket.jl @@ -29,8 +29,8 @@ Interfacer.name(::BucketSimulation) = "BucketSimulation" """ get_new_cache(p, Y, energy_check) -Returns a new `p` with an updated field to store e_per_area if energy conservation - checks are turned on. +Returns a new `p` with an updated field to store e_per_area if energy conservation + checks are turned on. """ function get_new_cache(p, Y, energy_check) if energy_check @@ -216,16 +216,12 @@ Interfacer.step!(sim::BucketSimulation, t) = Interfacer.step!(sim.integrator, t Interfacer.reinit!(sim::BucketSimulation) = Interfacer.reinit!(sim.integrator) # extensions required by FluxCalculator (partitioned fluxes) -function FluxCalculator.update_turbulent_fluxes_point!( - sim::BucketSimulation, - fields::NamedTuple, - colidx::CC.Fields.ColumnIndex, -) +function FluxCalculator.update_turbulent_fluxes!(sim::BucketSimulation, fields::NamedTuple) (; F_turb_energy, F_turb_moisture) = fields turbulent_fluxes = sim.integrator.p.bucket.turbulent_fluxes - turbulent_fluxes.shf[colidx] .= F_turb_energy + turbulent_fluxes.shf .= F_turb_energy earth_params = sim.model.parameters.earth_param_set - turbulent_fluxes.vapor_flux[colidx] .= F_turb_moisture ./ LP.ρ_cloud_liq(earth_params) + turbulent_fluxes.vapor_flux .= F_turb_moisture ./ LP.ρ_cloud_liq(earth_params) return nothing end @@ -234,15 +230,14 @@ function FluxCalculator.surface_thermo_state( sim::BucketSimulation, thermo_params::TD.Parameters.ThermodynamicsParameters, thermo_state_int, - colidx::CC.Fields.ColumnIndex, ) - T_sfc = Interfacer.get_field(sim, Val(:surface_temperature), colidx) + T_sfc = Interfacer.get_field(sim, Val(:surface_temperature)) # 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 = Interfacer.get_field(sim, Val(:air_density), colidx) - q_sfc = Interfacer.get_field(sim, Val(:surface_humidity), colidx) # already calculated in rhs! (cache) + ρ_sfc = Interfacer.get_field(sim, Val(:air_density)) + q_sfc = Interfacer.get_field(sim, Val(:surface_humidity)) # already calculated in rhs! (cache) @. TD.PhaseEquil_ρTq.(thermo_params, ρ_sfc, T_sfc, q_sfc) end diff --git a/experiments/ClimaEarth/components/ocean/eisenman_seaice.jl b/experiments/ClimaEarth/components/ocean/eisenman_seaice.jl index 4df68156f..5b22a7a0f 100644 --- a/experiments/ClimaEarth/components/ocean/eisenman_seaice.jl +++ b/experiments/ClimaEarth/components/ocean/eisenman_seaice.jl @@ -140,8 +140,8 @@ end function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:area_fraction}, field::CC.Fields.Field) sim.integrator.p.area_fraction .= field end -function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:∂F_turb_energy∂T_sfc}, field, colidx) - sim.integrator.p.Ya.∂F_turb_energy∂T_sfc[colidx] .= field +function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:∂F_turb_energy∂T_sfc}, field) + sim.integrator.p.Ya.∂F_turb_energy∂T_sfc .= field end function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:radiative_energy_flux_sfc}, field) parent(sim.integrator.p.Ya.F_rad) .= parent(field) @@ -155,13 +155,9 @@ Interfacer.step!(sim::EisenmanIceSimulation, t) = Interfacer.step!(sim.integrato Interfacer.reinit!(sim::EisenmanIceSimulation) = Interfacer.reinit!(sim.integrator) # extensions required by FluxCalculator (partitioned fluxes) -function FluxCalculator.update_turbulent_fluxes_point!( - sim::EisenmanIceSimulation, - fields::NamedTuple, - colidx::CC.Fields.ColumnIndex, -) +function FluxCalculator.update_turbulent_fluxes!(sim::EisenmanIceSimulation, fields::NamedTuple) (; F_turb_energy) = fields - @. sim.integrator.p.Ya.F_turb[colidx] = F_turb_energy + @. sim.integrator.p.Ya.F_turb = F_turb_energy end """ @@ -194,23 +190,22 @@ function FluxCalculator.differentiate_turbulent_fluxes!( fluxes; δT_sfc = 0.1, ) - (; thermo_state_int, surface_params, surface_scheme, colidx) = input_args - thermo_state_sfc_dT = - FluxCalculator.surface_thermo_state(sim, thermo_params, thermo_state_int, colidx, δT_sfc = δT_sfc) + (; thermo_state_int, surface_params, surface_scheme) = input_args + thermo_state_sfc_dT = FluxCalculator.surface_thermo_state(sim, thermo_params, thermo_state_int, δT_sfc = δT_sfc) input_args = merge(input_args, (; thermo_state_sfc = thermo_state_sfc_dT)) # set inputs based on whether the surface_scheme is `MoninObukhovScheme` or `BulkScheme` inputs = surface_inputs(surface_scheme, input_args) # calculate the surface fluxes - _, _, F_shf_δT_sfc, F_lhf_δT_sfc, _ = get_surface_fluxes_point!(inputs, surface_params) + _, _, F_shf_δT_sfc, F_lhf_δT_sfc, _ = get_surface_fluxes!(inputs, surface_params) (; F_shf, F_lhf) = fluxes # calculate the derivative ∂F_turb_energy∂T_sfc = @. (F_shf_δT_sfc + F_lhf_δT_sfc - F_shf - F_lhf) / δT_sfc - Interfacer.update_field!(sim, Val(:∂F_turb_energy∂T_sfc), ∂F_turb_energy∂T_sfc, colidx) + Interfacer.update_field!(sim, Val(:∂F_turb_energy∂T_sfc), ∂F_turb_energy∂T_sfc) end ### @@ -354,9 +349,7 @@ function ∑tendencies(dY, Y, cache, _) @. dY.T_sfc = -Y.T_sfc / Δt @. dY.q_sfc = -Y.q_sfc / Δt - CC.Fields.bycolumn(axes(Y.T_sfc)) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], p, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, p, thermo_params, Δt) # Get dY/dt @. dY.T_ml += Y.T_ml / Δt diff --git a/experiments/ClimaEarth/components/ocean/prescr_seaice.jl b/experiments/ClimaEarth/components/ocean/prescr_seaice.jl index ad2d2a727..4b46bbc93 100644 --- a/experiments/ClimaEarth/components/ocean/prescr_seaice.jl +++ b/experiments/ClimaEarth/components/ocean/prescr_seaice.jl @@ -131,13 +131,9 @@ Interfacer.step!(sim::PrescribedIceSimulation, t) = Interfacer.step!(sim.integra Interfacer.reinit!(sim::PrescribedIceSimulation) = Interfacer.reinit!(sim.integrator) # extensions required by FluxCalculator (partitioned fluxes) -function FluxCalculator.update_turbulent_fluxes_point!( - sim::PrescribedIceSimulation, - fields::NamedTuple, - colidx::CC.Fields.ColumnIndex, -) +function FluxCalculator.update_turbulent_fluxes!(sim::PrescribedIceSimulation, fields::NamedTuple) (; F_turb_energy) = fields - @. sim.integrator.p.F_turb_energy[colidx] = F_turb_energy + @. sim.integrator.p.F_turb_energy = F_turb_energy end """ diff --git a/experiments/ClimaEarth/components/ocean/slab_ocean.jl b/experiments/ClimaEarth/components/ocean/slab_ocean.jl index f6a1ed017..fa1277a8f 100644 --- a/experiments/ClimaEarth/components/ocean/slab_ocean.jl +++ b/experiments/ClimaEarth/components/ocean/slab_ocean.jl @@ -151,13 +151,9 @@ Interfacer.step!(sim::SlabOceanSimulation, t) = Interfacer.step!(sim.integrator, Interfacer.reinit!(sim::SlabOceanSimulation) = Interfacer.reinit!(sim.integrator) # extensions required by FluxCalculator (partitioned fluxes) -function FluxCalculator.update_turbulent_fluxes_point!( - sim::SlabOceanSimulation, - fields::NamedTuple, - colidx::CC.Fields.ColumnIndex, -) +function FluxCalculator.update_turbulent_fluxes!(sim::SlabOceanSimulation, fields::NamedTuple) (; F_turb_energy) = fields - @. sim.integrator.p.F_turb_energy[colidx] = F_turb_energy + @. sim.integrator.p.F_turb_energy = F_turb_energy end """ diff --git a/src/FluxCalculator.jl b/src/FluxCalculator.jl index 2d7919d2f..012b88b7d 100644 --- a/src/FluxCalculator.jl +++ b/src/FluxCalculator.jl @@ -24,7 +24,7 @@ export PartitionedStateFluxes, BulkScheme, partitioned_turbulent_fluxes!, get_surface_params, - update_turbulent_fluxes_point!, + update_turbulent_fluxes!, water_albedo_from_atmosphere! """ @@ -107,7 +107,7 @@ end 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_point!`, which calls `SurfaceFluxes.jl`. The coupler saves +using adapter function `get_surface_fluxes!`, which calls `SurfaceFluxes.jl`. The coupler saves the area-weighted sums of the fluxes. Args: @@ -143,70 +143,65 @@ function partitioned_turbulent_fluxes!( csf.F_turb_energy .*= FT(0) csf.F_turb_moisture .*= FT(0) - # iterate over all columns (when regridding, this will need to change) - CC.Fields.bycolumn(boundary_space) do colidx - # atmos state of center level 1 - z_int = Interfacer.get_field(atmos_sim, Val(:height_int), colidx) - uₕ_int = Interfacer.get_field(atmos_sim, Val(:uv_int), colidx) - thermo_state_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int), colidx) - z_sfc = Interfacer.get_field(atmos_sim, Val(:height_sfc), colidx) - - for sim in model_sims - # iterate over all surface models with non-zero area fractions - if sim isa Interfacer.SurfaceModelSimulation - # get area fraction (min = 0, max = 1) - area_fraction = Interfacer.get_field(sim, Val(:area_fraction), colidx) - # get area mask [0, 1], where area_mask = 1 if area_fraction > 0 - area_mask = Regridder.binary_mask.(area_fraction) - - if !iszero(parent(area_mask)) - - thermo_state_sfc = surface_thermo_state(sim, thermo_params, thermo_state_int, colidx) - - # set inputs based on whether the surface_scheme is `MoninObukhovScheme` or `BulkScheme` - surface_params = get_surface_params(atmos_sim) - scheme_properties = get_scheme_properties(surface_scheme, sim, colidx) - input_args = (; - thermo_state_sfc, - thermo_state_int, - uₕ_int, - z_int, - z_sfc, - scheme_properties, - surface_params, - surface_scheme, - colidx, - ) - inputs = surface_inputs(surface_scheme, input_args) - - # calculate the surface fluxes - fluxes = get_surface_fluxes_point!(inputs, surface_params) - (; F_turb_ρτxz, F_turb_ρτyz, F_shf, F_lhf, F_turb_moisture) = fluxes - - # perform additional diagnostics if required - differentiate_turbulent_fluxes!(sim, (thermo_params, input_args, fluxes)) - - # update fluxes in the coupler - fields = (; - F_turb_ρτxz = F_turb_ρτxz, - F_turb_ρτyz = F_turb_ρτyz, - F_turb_energy = F_shf .+ F_lhf, - F_turb_moisture = F_turb_moisture, - ) - - # update the fluxes of this surface model - update_turbulent_fluxes_point!(sim, fields, colidx) - - # add the flux contributing from this surface to the coupler field - @. csf.F_turb_ρτxz[colidx] += F_turb_ρτxz * area_fraction * area_mask - @. csf.F_turb_ρτyz[colidx] += F_turb_ρτyz * area_fraction * area_mask - @. csf.F_turb_energy[colidx] += (F_shf .+ F_lhf) * area_fraction * area_mask - @. csf.F_turb_moisture[colidx] += F_turb_moisture * area_fraction * area_mask - - end - end + # atmos state of center level 1 + z_int = Interfacer.get_field(atmos_sim, Val(:height_int)) + uₕ_int = Interfacer.get_field(atmos_sim, Val(:uv_int)) + thermo_state_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int)) + z_sfc = Interfacer.get_field(atmos_sim, Val(:height_sfc)) + + for sim in model_sims + # iterate over all surface models + if sim isa Interfacer.SurfaceModelSimulation + # get area fraction (min = 0, max = 1) + area_fraction = Interfacer.get_field(sim, Val(:area_fraction)) + # get area mask [0, 1], where area_mask = 1 if area_fraction > 0 + area_mask = Regridder.binary_mask.(area_fraction) + + thermo_state_sfc = FluxCalculator.surface_thermo_state(sim, thermo_params, thermo_state_int) + + # set inputs based on whether the surface_scheme is `MoninObukhovScheme` or `BulkScheme` + surface_params = FluxCalculator.get_surface_params(atmos_sim) + scheme_properties = FluxCalculator.get_scheme_properties(surface_scheme, sim) + + input_args = (; + thermo_state_sfc, + thermo_state_int, + uₕ_int, + z_int, + z_sfc, + scheme_properties, + surface_params, + surface_scheme, + boundary_space, + ) + inputs = FluxCalculator.surface_inputs(surface_scheme, input_args) + + # calculate the surface fluxes + fluxes = FluxCalculator.get_surface_fluxes!(inputs, surface_params) + (; F_turb_ρτxz, F_turb_ρτyz, F_shf, F_lhf, F_turb_moisture) = fluxes + + # perform additional diagnostics if required + FluxCalculator.differentiate_turbulent_fluxes!(sim, (thermo_params, input_args, fluxes)) + + # update fluxes in the coupler + fields = (; + F_turb_ρτxz = F_turb_ρτxz, + F_turb_ρτyz = F_turb_ρτyz, + F_turb_energy = F_shf .+ F_lhf, + F_turb_moisture = F_turb_moisture, + ) + + # update the fluxes of this surface model + FluxCalculator.update_turbulent_fluxes!(sim, fields) + + # add the flux contributing from this surface to the coupler field + # note that the fluxes are area-weighted, so if a surface model is + # not present at this point, the fluxes are zero + @. csf.F_turb_ρτxz += F_turb_ρτxz * area_fraction * area_mask + @. csf.F_turb_ρτyz += F_turb_ρτyz * area_fraction * area_mask + @. csf.F_turb_energy += (F_shf .+ F_lhf) * area_fraction * area_mask + @. csf.F_turb_moisture += F_turb_moisture * area_fraction * area_mask end - end # TODO: add allowable bounds here, check explicitly that all fluxes are equal @@ -218,25 +213,21 @@ struct BulkScheme <: AbstractSurfaceFluxScheme end struct MoninObukhovScheme <: AbstractSurfaceFluxScheme end """ - get_scheme_properties(scheme::AbstractSurfaceFluxScheme, sim::Interfacer.SurfaceModelSimulation, colidx::CC.Fields.ColumnIndex) + get_scheme_properties(scheme::AbstractSurfaceFluxScheme, sim::Interfacer.SurfaceModelSimulation) Returns the scheme-specific properties for the surface model simulation `sim`. """ -function get_scheme_properties(::BulkScheme, sim::Interfacer.SurfaceModelSimulation, colidx::CC.Fields.ColumnIndex) - Ch = Interfacer.get_field(sim, Val(:heat_transfer_coefficient), colidx) - Cd = Interfacer.get_field(sim, Val(:drag_coefficient), colidx) - beta = Interfacer.get_field(sim, Val(:beta), colidx) +function get_scheme_properties(::BulkScheme, sim::Interfacer.SurfaceModelSimulation) + Ch = Interfacer.get_field(sim, Val(:heat_transfer_coefficient)) + Cd = Interfacer.get_field(sim, Val(:drag_coefficient)) + beta = Interfacer.get_field(sim, Val(:beta)) FT = eltype(Ch) return (; z0b = FT(0), z0m = FT(0), Ch = Ch, Cd = Cd, beta = beta, gustiness = FT(1)) end -function get_scheme_properties( - ::MoninObukhovScheme, - sim::Interfacer.SurfaceModelSimulation, - colidx::CC.Fields.ColumnIndex, -) - z0m = Interfacer.get_field(sim, Val(:roughness_momentum), colidx) - z0b = Interfacer.get_field(sim, Val(:roughness_buoyancy), colidx) - beta = Interfacer.get_field(sim, Val(:beta), colidx) +function get_scheme_properties(::MoninObukhovScheme, sim::Interfacer.SurfaceModelSimulation) + z0m = Interfacer.get_field(sim, Val(:roughness_momentum)) + z0b = Interfacer.get_field(sim, Val(:roughness_buoyancy)) + beta = Interfacer.get_field(sim, Val(:beta)) FT = eltype(z0m) return (; z0b = z0b, z0m = z0m, Ch = FT(0), Cd = FT(0), beta = beta, gustiness = FT(1)) end @@ -248,61 +239,87 @@ Returns the inputs for the surface model simulation `sim`. """ function surface_inputs(::BulkScheme, input_args::NamedTuple) - (; thermo_state_sfc, thermo_state_int, uₕ_int, z_int, z_sfc, scheme_properties) = input_args + (; thermo_state_sfc, thermo_state_int, uₕ_int, z_int, z_sfc, scheme_properties, boundary_space) = input_args FT = CC.Spaces.undertype(axes(z_sfc)) + (; Ch, Cd, beta, gustiness) = scheme_properties + + # Extract the underlying data layouts of each field + # Note: this is a bit "dangerous" because it circumvents ClimaCore, but + # it allows us to broadcast over fields on slightly different spaces + fv = CC.Fields.field_values + z_int_fv = fv(z_int) + uₕ_int_fv = fv(uₕ_int) + thermo_state_int_fv = fv(thermo_state_int) + z_sfc_fv = fv(z_sfc) + thermo_state_sfc_fv = fv(thermo_state_sfc) + beta_fv = fv(beta) - (; z0b, z0m, Ch, Cd, beta, gustiness) = scheme_properties # wrap state values - return @. SF.Coefficients( - SF.StateValues(z_int, uₕ_int, thermo_state_int), # state_in + result = @. SF.Coefficients( + SF.StateValues(z_int_fv, uₕ_int_fv, thermo_state_int_fv), # state_in SF.StateValues( # state_sfc - z_sfc, + z_sfc_fv, StaticArrays.SVector(FT(0), FT(0)), - thermo_state_sfc, + thermo_state_sfc_fv, ), Cd, # Cd Ch, # Ch gustiness, # gustiness - beta, # beta + beta_fv, # beta ) + + # Put the result data layout back onto the surface space + return CC.Fields.Field(result, boundary_space) end function surface_inputs(::MoninObukhovScheme, input_args::NamedTuple) - (; thermo_state_sfc, thermo_state_int, uₕ_int, z_int, z_sfc, scheme_properties) = input_args + (; thermo_state_sfc, thermo_state_int, uₕ_int, z_int, z_sfc, scheme_properties, boundary_space) = input_args FT = CC.Spaces.undertype(axes(z_sfc)) - (; z0b, z0m, Ch, Cd, beta, gustiness) = scheme_properties - - # wrap state values - return @. SF.ValuesOnly( - SF.StateValues(z_int, uₕ_int, thermo_state_int), # state_in + (; z0b, z0m, beta, gustiness) = scheme_properties + + # Extract the underlying data layouts of each field + # Note: this is a bit "dangerous" because it circumvents ClimaCore, but + # it allows us to broadcast over fields on slightly different spaces + fv = CC.Fields.field_values + z_int_fv = fv(z_int) + uₕ_int_fv = uₕ_int isa CC.Fields.Field ? fv(uₕ_int) : uₕ_int + thermo_state_int_fv = fv(thermo_state_int) + z_sfc_fv = fv(z_sfc) + thermo_state_sfc_fv = fv(thermo_state_sfc) + beta_fv = beta isa CC.Fields.Field ? fv(beta) : beta + + # Compute state values + result = @. SF.ValuesOnly( + SF.StateValues(z_int_fv, uₕ_int_fv, thermo_state_int_fv), # state_in SF.StateValues( # state_sfc - z_sfc, + z_sfc_fv, StaticArrays.SVector(FT(0), FT(0)), - thermo_state_sfc, + thermo_state_sfc_fv, ), - z0m, # z0m - z0b, # z0b + z0m, # z0m + z0b, # z0b gustiness, # gustiness - beta, # beta + beta_fv, # beta ) + # Put the result data layout back onto the surface space + return CC.Fields.Field(result, boundary_space) end """ - surface_thermo_state(sim::Interfacer.SurfaceModelSimulation, thermo_params::TD.Parameters.ThermodynamicsParameters, thermo_state_int, colidx::CC.Fields.ColumnIndex) + surface_thermo_state(sim::Interfacer.SurfaceModelSimulation, thermo_params::TD.Parameters.ThermodynamicsParameters, thermo_state_int) Returns the surface parameters for the surface model simulation `sim`. The default is assuming saturated surfaces, unless an extension is defined for the given `SurfaceModelSimulation`. """ function surface_thermo_state( sim::Interfacer.SurfaceModelSimulation, thermo_params::TD.Parameters.ThermodynamicsParameters, - thermo_state_int, - colidx::CC.Fields.ColumnIndex; + thermo_state_int; δT_sfc = 0, ) FT = eltype(parent(thermo_state_int)) @warn("Simulation " * Interfacer.name(sim) * " uses the default thermo (saturated) surface state", maxlog = 10) # get surface temperature (or perturbed surface temperature for differentiation) - T_sfc = Interfacer.get_field(sim, Val(:surface_temperature), colidx) .+ FT(δT_sfc) + T_sfc = Interfacer.get_field(sim, Val(:surface_temperature)) .+ FT(δT_sfc) ρ_sfc = extrapolate_ρ_to_sfc.(thermo_params, thermo_state_int, T_sfc) # ideally the # calculate elsewhere, here just getter... q_sfc = TD.q_vap_saturation_generic.(thermo_params, T_sfc, ρ_sfc, TD.Liquid()) # default: saturated liquid surface @. TD.PhaseEquil_ρTq.(thermo_params, ρ_sfc, T_sfc, q_sfc) @@ -323,14 +340,13 @@ function extrapolate_ρ_to_sfc(thermo_params, ts_in, T_sfc) end """ - get_surface_fluxes_point!(inputs, surface_params::SF.Parameters.SurfaceFluxesParameters) + get_surface_fluxes!(inputs, surface_params::SF.Parameters.SurfaceFluxesParameters) Uses SurfaceFluxes.jl to calculate turbulent surface fluxes. It should be atmos model agnostic, and columnwise. """ -function get_surface_fluxes_point!(inputs, surface_params::SF.Parameters.SurfaceFluxesParameters) - +function get_surface_fluxes!(inputs, surface_params::SF.Parameters.SurfaceFluxesParameters) # calculate all fluxes (saturated surface conditions) - outputs = @. SF.surface_conditions(surface_params, inputs) + outputs = SF.surface_conditions.(surface_params, inputs) # drag F_turb_ρτxz = outputs.ρτxz @@ -341,7 +357,15 @@ function get_surface_fluxes_point!(inputs, surface_params::SF.Parameters.Surface F_lhf = outputs.lhf # moisture - F_turb_moisture = @. SF.evaporation(surface_params, inputs, outputs.Ch) + F_turb_moisture = SF.evaporation.(surface_params, inputs, outputs.Ch) + + # At locations where this surface model is not evaluated, we get `NaN` for + # surface fluxes. In that case, we replace the values with 0. + @. F_turb_ρτxz = ifelse(isnan(F_turb_ρτxz), zero(F_turb_ρτxz), F_turb_ρτxz) + @. F_turb_ρτyz = ifelse(isnan(F_turb_ρτyz), zero(F_turb_ρτyz), F_turb_ρτyz) + @. F_shf = ifelse(isnan(F_shf), zero(F_shf), F_shf) + @. F_lhf = ifelse(isnan(F_lhf), zero(F_lhf), F_lhf) + @. F_turb_moisture = ifelse(isnan(F_turb_moisture), zero(F_turb_moisture), F_turb_moisture) return (; F_turb_ρτxz = F_turb_ρτxz, @@ -368,22 +392,18 @@ function get_surface_params(atmos_sim::Interfacer.AtmosModelSimulation) end """ - update_turbulent_fluxes_point!(sim::Interfacer.SurfaceModelSimulation, fields::NamedTuple, colidx::CC.Fields.ColumnIndex) + update_turbulent_fluxes!(sim::Interfacer.SurfaceModelSimulation, fields::NamedTuple) Updates the fluxes in the surface model simulation `sim` with the fluxes in `fields`. """ -function update_turbulent_fluxes_point!( - sim::Interfacer.SurfaceModelSimulation, - fields::NamedTuple, - colidx::CC.Fields.ColumnIndex, -) +function update_turbulent_fluxes!(sim::Interfacer.SurfaceModelSimulation, fields::NamedTuple) return error( - "update_turbulent_fluxes_point! is required to be dispatched on" * - Interfacer.name(sim) * - ", but no method defined", + "update_turbulent_fluxes! is required to be dispatched on" * Interfacer.name(sim) * ", but no method defined", ) end +update_turbulent_fluxes!(sim::Interfacer.SurfaceStub, fields::NamedTuple) = nothing + """ differentiate_turbulent_fluxes!(sim::Interfacer.SurfaceModelSimulation, args) diff --git a/src/Interfacer.jl b/src/Interfacer.jl index 42b15575b..c707a54b3 100644 --- a/src/Interfacer.jl +++ b/src/Interfacer.jl @@ -168,19 +168,6 @@ get_field(sim::ComponentModelSimulation, val::Val) = get_field_error(sim, val) get_field_error(sim, val::Val{X}) where {X} = error("undefined field `$X` for " * name(sim)) -""" - get_field(::ComponentModelSimulation, ::Val, colidx::CC.Fields.ColumnIndex) - -Extension of `get_field(::ComponentModelSimulation, ::Val)`, indexing into the specified colum index. -""" -function get_field(sim::ComponentModelSimulation, val::Val, colidx::CC.Fields.ColumnIndex) - if get_field(sim, val) isa AbstractFloat - get_field(sim, val) - else - get_field(sim, val)[colidx] - end -end - """ update_field!(::AtmosModelSimulation, ::Val, _...) diff --git a/src/surface_stub.jl b/src/surface_stub.jl index b68480908..118092528 100644 --- a/src/surface_stub.jl +++ b/src/surface_stub.jl @@ -87,8 +87,3 @@ reinit!(::SurfaceStub) = nothing The stub surface simulation is not updated by this function. Extends `SciMLBase.step!`. """ step!(::SurfaceStub, _) = nothing - - -## Extensions of FluxCalculator.jl functions - -update_turbulent_fluxes_point!(sim::SurfaceStub, fields::NamedTuple, colidx::CC.Fields.ColumnIndex) = nothing diff --git a/test/component_model_tests/eisenman_seaice_tests.jl b/test/component_model_tests/eisenman_seaice_tests.jl index c1cf9e6f3..c0213de8e 100644 --- a/test/component_model_tests/eisenman_seaice_tests.jl +++ b/test/component_model_tests/eisenman_seaice_tests.jl @@ -36,9 +36,8 @@ for FT in (Float32, Float64) Ya.F_rad .= 0 Ya.F_turb .= 0 - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) + @test all(parent(Ya.e_base) .≈ 0) @test all(parent(Y.T_ml) .≈ params_ice.T_base) @test all(parent(Y.T_sfc) .≈ params_ice.T_base) @@ -62,9 +61,7 @@ for FT in (Float32, Float64) ∂F_atm∂T_sfc = get_∂F_rad_energy∂T_sfc(Y.T_sfc, params_ice) .+ Ya.∂F_turb_energy∂T_sfc @. Ya.F_rad = (1 - params_ice.α) * params_ice.σ * Y.T_sfc .^ 4 # outgoing longwave - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) F_atm = @. Ya.F_rad + Ya.F_turb @@ -93,9 +90,7 @@ for FT in (Float32, Float64) Y.T_sfc .= params_ice.T_base Y.T_ml .= params_ice.T_base - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) F_atm = @. Ya.F_rad + Ya.F_turb ∂F_atm∂T_sfc = 0 @@ -128,9 +123,7 @@ for FT in (Float32, Float64) Y.T_sfc .= params_ice.T_base .- 1 T_sfc_0 = deepcopy(Y.T_sfc) - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) F_atm = @. Ya.F_rad + Ya.F_turb ∂F_atm∂T_sfc = 0 @@ -163,9 +156,7 @@ for FT in (Float32, Float64) Y.T_ml .= params_ice.T_base T_ml_0 = deepcopy(Y.T_ml) - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) F_atm = @. Ya.F_rad + Ya.F_turb @@ -193,9 +184,7 @@ for FT in (Float32, Float64) Ya.F_turb .= 0 Ya.F_rad .= @. (1 - params_ice.α) * params_ice.σ * Y.T_sfc^4 - 300 - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) F_atm = @. Ya.F_rad + Ya.F_turb @@ -223,9 +212,7 @@ for FT in (Float32, Float64) # net outgoing longwave Ya.F_rad .= @. (1 - params_ice.α) * params_ice.σ * Y.T_sfc^4 - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) F_atm = @. Ya.F_rad + Ya.F_turb ∂F_atm∂T_sfc = get_∂F_rad_energy∂T_sfc(T_sfc_0, params_ice) .+ Ya.∂F_turb_energy∂T_sfc @@ -264,9 +251,7 @@ for FT in (Float32, Float64) Ya.F_turb .= 0 Ya.F_rad .= 0 - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) @test all(parent(Ya.e_base) .≈ params_ice.C0_base * ΔT_ml * FT(Δt)) # non-zero contribution from basal flux T_ml_new = T_ml_0 .- params_ice.C0_base .* ΔT_ml * FT(Δt) / (params_ocean.h * params_ocean.ρ * params_ocean.c) @@ -298,9 +283,7 @@ for FT in (Float32, Float64) Ya.F_turb .= 0 Ya.F_rad .= 0 - CC.Fields.bycolumn(boundary_space) do colidx - solve_eisenman_model!(Y[colidx], Ya[colidx], params, thermo_params, Δt) - end + solve_eisenman_model!(Y, Ya, params, thermo_params, Δt) @test all(parent(Ya.e_base) .≈ 0) # no contribution from basal flux T_ml_new = @. T_ml_0 + Ya.ocean_qflux * FT(Δt) / (params_ocean.h * params_ocean.ρ * params_ocean.c) diff --git a/test/flux_calculator_tests.jl b/test/flux_calculator_tests.jl index 0a4249df4..8bcbeccd2 100644 --- a/test/flux_calculator_tests.jl +++ b/test/flux_calculator_tests.jl @@ -80,25 +80,16 @@ Interfacer.get_field(sim::TestOcean, ::Val{:drag_coefficient}) = sim.integrator. Interfacer.get_field(sim::TestOcean, ::Union{Val{:surface_direct_albedo}, Val{:surface_diffuse_albedo}}) = sim.integrator.p.α -function FluxCalculator.surface_thermo_state( - sim::TestOcean, - thermo_params::ThermodynamicsParameters, - thermo_state_int, - colidx::CC.Fields.ColumnIndex, -) - T_sfc = Interfacer.get_field(sim, Val(:surface_temperature), colidx) +function FluxCalculator.surface_thermo_state(sim::TestOcean, thermo_params::ThermodynamicsParameters, thermo_state_int) + T_sfc = Interfacer.get_field(sim, Val(:surface_temperature)) ρ_sfc = thermo_state_int.ρ # arbitrary - q_sfc = Interfacer.get_field(sim, Val(:air_humidity), colidx) # read from cache + q_sfc = Interfacer.get_field(sim, Val(:air_humidity)) # read from cache @. TD.PhaseEquil_ρTq.(thermo_params, ρ_sfc, T_sfc, q_sfc) end -function FluxCalculator.update_turbulent_fluxes_point!( - sim::TestOcean, - fields::NamedTuple, - colidx::CC.Fields.ColumnIndex, -) +function FluxCalculator.update_turbulent_fluxes!(sim::TestOcean, fields::NamedTuple) (; F_turb_energy) = fields - @. sim.integrator.p.F_aero[colidx] = F_turb_energy + @. sim.integrator.p.F_aero = F_turb_energy end # simple surface sim object and extensions @@ -115,13 +106,8 @@ Interfacer.get_field(sim::DummySurfaceSimulation3, ::Val{:heat_transfer_coeffici Interfacer.get_field(sim::DummySurfaceSimulation3, ::Val{:drag_coefficient}) = sim.integrator.p.Cd Interfacer.get_field(sim::DummySurfaceSimulation3, ::Val{:beta}) = sim.integrator.p.beta -function surface_thermo_state( - sim::DummySurfaceSimulation3, - thermo_params::ThermodynamicsParameters, - thermo_state_int, - colidx::CC.Fields.ColumnIndex, -) - T_sfc = Interfacer.get_field(sim, Val(:surface_temperature), colidx) +function surface_thermo_state(sim::DummySurfaceSimulation3, thermo_params::ThermodynamicsParameters, thermo_state_int) + T_sfc = Interfacer.get_field(sim, Val(:surface_temperature)) FT = eltype(T_sfc) ρ_sfc = @. T_sfc * FT(0) .+ FT(1.2) # arbitrary @@ -240,10 +226,7 @@ for FT in (Float32, Float64) thermo_state_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int)) surface_thermo_states = similar(thermo_state_int) - CC.Fields.bycolumn(boundary_space) do colidx - surface_thermo_states[colidx] .= - FluxCalculator.surface_thermo_state(ocean_sim, thermo_params, thermo_state_int[colidx], colidx) - end + surface_thermo_states .= FluxCalculator.surface_thermo_state(ocean_sim, thermo_params, thermo_state_int) # analytical solution is possible for the BulkScheme() case if scheme isa FluxCalculator.BulkScheme @@ -259,16 +242,15 @@ for FT in (Float32, Float64) ρ_sfc * windspeed #-ρ_sfc * Ch * windspeed(sc) * (cp_m * ΔT + ΔΦ) - colidx = CC.Fields.ColumnIndex{2}((1, 1), 73) # arbitrary index # check the coupler field update - @test isapprox(parent(shf_analytical[colidx]), parent(fields.F_turb_energy[colidx]), rtol = 1e-6) + @test isapprox(parent(shf_analytical), parent(fields.F_turb_energy), rtol = 1e-6) # test the surface field update - @test parent(fields.F_turb_energy[colidx]) == parent(ocean_sim.integrator.p.F_aero[colidx]) + @test parent(fields.F_turb_energy) == parent(ocean_sim.integrator.p.F_aero) # test the atmos field update FieldExchanger.update_sim!(atmos_sim, fields, nothing) - @test parent(fields.F_turb_energy[colidx]) == -parent(atmos_sim.integrator.p.energy_bc[colidx]) + @test parent(fields.F_turb_energy) == -parent(atmos_sim.integrator.p.energy_bc) end @test Array(parent(fields.F_turb_moisture))[1] ≈ FT(0) @@ -285,14 +267,13 @@ for FT in (Float32, Float64) ) FluxCalculator.get_surface_params(DummySimulation([], [])) end - @testset "update_turbulent_fluxes_point! for FT=$FT" begin + @testset "update_turbulent_fluxes! for FT=$FT" begin sim = DummySurfaceSimulation3([], [], [], []) - colidx = CC.Fields.ColumnIndex{2}((1, 1), 73) # arbitrary index @test_throws ErrorException( - "update_turbulent_fluxes_point! is required to be dispatched on" * + "update_turbulent_fluxes! is required to be dispatched on" * Interfacer.name(sim) * ", but no method defined", - ) FluxCalculator.update_turbulent_fluxes_point!(sim, (;), colidx) == ErrorException + ) FluxCalculator.update_turbulent_fluxes!(sim, (;)) == ErrorException end @testset "surface_thermo_state for FT=$FT" begin @@ -308,9 +289,7 @@ for FT in (Float32, Float64) TestAtmos((; FT = FT), [], [], (; T = _ones .* FT(300), ρ = _ones .* FT(1.2), q = _ones .* FT(0.01))) thermo_params = get_thermo_params(atmos_sim) thermo_state_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int)) - colidx = CC.Fields.ColumnIndex{2}((1, 1), 73) # arbitrary index - @test FluxCalculator.surface_thermo_state(surface_sim, thermo_params, thermo_state_int[colidx], colidx).ρ == - thermo_state_int[colidx].ρ + @test FluxCalculator.surface_thermo_state(surface_sim, thermo_params, thermo_state_int).ρ == thermo_state_int.ρ end @testset "water_albedo_from_atmosphere!" begin @@ -348,3 +327,8 @@ for FT in (Float32, Float64) ) FluxCalculator.water_albedo_from_atmosphere!(atmos_sim2, ones(boundary_space), ones(boundary_space)) end end + +@testset "SurfaceStub update_turbulent_fluxes!" begin + FT = Float32 + @test isnothing(FluxCalculator.update_turbulent_fluxes!(Interfacer.SurfaceStub(FT(0)), (;))) +end diff --git a/test/interfacer_tests.jl b/test/interfacer_tests.jl index 968bd9fe5..b5452e296 100644 --- a/test/interfacer_tests.jl +++ b/test/interfacer_tests.jl @@ -56,10 +56,9 @@ for FT in (Float32, Float64) space = TestHelper.create_space(FT) for sim in (DummySimulation(space), DummySimulation2(space), DummySimulation3(space)) # field - colidx = CC.Fields.ColumnIndex{2}((1, 1), 73) - @test Array(parent(Interfacer.get_field(sim, Val(:var), colidx)))[1] == FT(1) + @test Array(parent(Interfacer.get_field(sim, Val(:var))))[1] == FT(1) # float - @test Interfacer.get_field(sim, Val(:var_float), colidx) == FT(2) + @test Interfacer.get_field(sim, Val(:var_float)) == FT(2) end end @@ -252,30 +251,3 @@ end FT = Float32 @test isnothing(Interfacer.reinit!(Interfacer.SurfaceStub(FT(0)))) end - -@testset "SurfaceStub update_turbulent_fluxes_point!" begin - FT = Float32 - colidx = CC.Fields.ColumnIndex{2}((1, 1), 73) # arbitrary index - @test isnothing(Interfacer.update_turbulent_fluxes_point!(Interfacer.SurfaceStub(FT(0)), (;), colidx)) -end - -# # Test that update_field! gives correct warnings for unextended fields -# for value in ( -# :air_density, -# :air_temperature, -# :energy, -# :height_int, -# :height_sfc, -# :liquid_precipitation, -# :radiative_energy_flux_sfc, -# :radiative_energy_flux_toa, -# :snow_precipitation, -# :turbulent_energy_flux, -# :turbulent_moisture_flux, -# :thermo_state_int, -# :uv_int, -# :water, -# ) -# val = Val(value) -# @test_throws ErrorException("undefined field `$value` for " * name(sim)) get_field(sim, val) -# end diff --git a/test/runtests.jl b/test/runtests.jl index c022715a2..3ea6fa6e3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -38,7 +38,7 @@ end @safetestset "FieldExchanger tests" begin include("field_exchanger_tests.jl") end -gpu_broken || @safetestset "FluxCalculator tests" begin +@safetestset "FluxCalculator tests" begin include("flux_calculator_tests.jl") end gpu_broken || @safetestset "Diagnostics tests" begin