diff --git a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl index a60ce110..f5df31f6 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -14,7 +14,7 @@ using ClimaSeaIce: SeaIceModel using Oceananigans: HydrostaticFreeSurfaceModel, architecture using Oceananigans.Units: Time -using Oceananigans.Grids: inactive_node, node, topology +using Oceananigans.Grids: inactive_node, node, topology, architecture using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Fields: ConstantField, interpolate, FractionalIndices using Oceananigans.Utils: launch!, KernelParameters @@ -272,6 +272,16 @@ function default_ao_specific_humidity(ocean) return ImpureSaturationSpecificHumidity(phase, x_H₂O) end +function default_ao_fluxes(FT) + stability_functions = atmosphere_ocean_stability_functions(FT) + return SimilarityTheoryFluxes(FT; stability_functions) +end + +function default_ai_fluxes(FT) + stability_functions = atmosphere_sea_ice_stability_functions(FT) + return SimilarityTheoryFluxes(FT; stability_functions) +end + """ ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; kwargs...) @@ -303,8 +313,8 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; exchange_grid = ocean.model.grid, radiation = Radiation(), freshwater_density = default_freshwater_density, - atmosphere_ocean_fluxes = SimilarityTheoryFluxes(eltype(exchange_grid)), - atmosphere_sea_ice_fluxes = atmosphere_sea_ice_similarity_theory(eltype(exchange_grid)), + atmosphere_ocean_fluxes = default_ao_fluxes(eltype(exchange_grid)), + atmosphere_sea_ice_fluxes = default_ai_fluxes(eltype(exchange_grid)), sea_ice_ocean_heat_flux = ThreeEquationHeatFlux(sea_ice), atmosphere_ocean_interface_temperature = BulkTemperature(), atmosphere_ocean_velocity_difference = RelativeVelocity(), @@ -319,7 +329,11 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; sea_ice_heat_capacity = heat_capacity(sea_ice), gravitational_acceleration = default_gravitational_acceleration) - FT = eltype(exchange_grid) + FT = eltype(exchange_grid) + arch = architecture(exchange_grid) + + atmosphere_ocean_fluxes = on_architecture(arch, atmosphere_ocean_fluxes) + atmosphere_sea_ice_fluxes = on_architecture(arch, atmosphere_sea_ice_fluxes) ocean_reference_density = convert(FT, ocean_reference_density) ocean_heat_capacity = convert(FT, ocean_heat_capacity) diff --git a/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl index 6229db1a..da8aeb43 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl @@ -6,6 +6,8 @@ using Thermodynamics: Liquid using KernelAbstractions.Extras.LoopInfo: @unroll using Statistics: norm +import Oceananigans.Utils: TabulatedFunction +import Oceananigans.Architectures: on_architecture import Thermodynamics as AtmosphericThermodynamics import Thermodynamics.Parameters: Rv_over_Rd @@ -113,6 +115,15 @@ function SimilarityTheoryFluxes(FT::DataType = Oceananigans.defaults.FloatType; solver_stop_criteria) end +on_architecture(arch, ψ::SimilarityTheoryFluxes) = + SimilarityTheoryFluxes(ψ.von_karman_constant, + ψ.turbulent_prandtl_number, + ψ.gustiness_parameter, + on_architecture(arch, ψ.stability_functions), + ψ.roughness_lengths, + ψ.similarity_form, + ψ.solver_stop_criteria) + ##### ##### Similarity profile types ##### @@ -275,12 +286,25 @@ Base.summary(ss::SimilarityScales) = Base.show(io::IO, ss::SimilarityScales) = print(io, summary(ss)) +Adapt.adapt_structure(to, ss::SimilarityScales) = + SimilarityScales(adapt(to, ss.momentum), + adapt(to, ss.temperature), + adapt(to, ss.water_vapor)) + +on_architecture(arch, ss::SimilarityScales) = + SimilarityScales(on_architecture(arch, ss.momentum), + on_architecture(arch, ss.temperature), + on_architecture(arch, ss.water_vapor)) + @inline stability_profile(ψ, ζ) = ψ(ζ) # Convenience abstract type AbstractStabilityFunction end + @inline (ψ::AbstractStabilityFunction)(ζ) = stability_profile(ψ, ζ) +on_architecture(arch, ψ::AbstractStabilityFunction) = ψ + """ EdsonMomentumStabilityFunction{FT} @@ -451,9 +475,23 @@ end end # Edson et al. (2013) -function atmosphere_ocean_stability_functions(FT=Oceananigans.defaults.FloatType) +function atmosphere_ocean_stability_functions(FT=Oceananigans.defaults.FloatType; + tabulate_stability_functions = true, + tabulation_ζ_range = (-100, 100), + tabulation_points = 100000) ψu = EdsonMomentumStabilityFunction{FT}() ψc = EdsonScalarStabilityFunction{FT}() + + if tabulate_stability_functions + ψu = TabulatedFunction(ψu, CPU(), FT; + range = tabulation_ζ_range, + points = tabulation_points) + + ψc = TabulatedFunction(ψc, CPU(), FT; + range = tabulation_ζ_range, + points = tabulation_points) + end + return SimilarityScales(ψu, ψc, ψc) end @@ -550,6 +588,8 @@ end Base.summary(ss::SplitStabilityFunction) = "SplitStabilityFunction" Base.show(io::IO, ss::SplitStabilityFunction) = print(io, "SplitStabilityFunction") +@inline (ψ::SplitStabilityFunction)(ζ) = stability_profile(ψ, ζ) + @inline function stability_profile(ψ::SplitStabilityFunction, ζ) Ψ_stable = stability_profile(ψ.stable, ζ) Ψ_unstable = stability_profile(ψ.unstable, ζ) @@ -557,16 +597,30 @@ Base.show(io::IO, ss::SplitStabilityFunction) = print(io, "SplitStabilityFunctio return ifelse(stable, Ψ_stable, Ψ_unstable) end -function atmosphere_sea_ice_stability_functions(FT=Oceananigans.defaults.FloatType) +function atmosphere_sea_ice_stability_functions(FT=Oceananigans.defaults.FloatType; + tabulate_stability_functions = false, + tabulation_ζ_range = (-100, 100), + tabulation_points = 100000) + unstable_momentum = PaulsonMomentumStabilityFunction{FT}() stable_momentum = ShebaMomentumStabilityFunction{FT}() - momentum = SplitStabilityFunction(stable_momentum, unstable_momentum) + ψu = SplitStabilityFunction(stable_momentum, unstable_momentum) unstable_scalar = PaulsonScalarStabilityFunction{FT}() stable_scalar = ShebaScalarStabilityFunction{FT}() - scalar = SplitStabilityFunction(stable_scalar, unstable_scalar) + ψc = SplitStabilityFunction(stable_scalar, unstable_scalar) + + if tabulate_stability_functions + ψu = TabulatedFunction(ψu, CPU(), FT; + range = tabulation_ζ_range, + points = tabulation_points) - return SimilarityScales(momentum, scalar, scalar) + ψc = TabulatedFunction(ψc, CPU(), FT; + range = tabulation_ζ_range, + points = tabulation_points) + end + + return SimilarityScales(ψu, ψc, ψc) end function atmosphere_sea_ice_similarity_theory(FT=Oceananigans.defaults.FloatType) diff --git a/test/test_surface_fluxes.jl b/test/test_surface_fluxes.jl index 779824bd..56420c4c 100644 --- a/test/test_surface_fluxes.jl +++ b/test/test_surface_fluxes.jl @@ -8,7 +8,9 @@ using NumericalEarth.OceanSeaIceModels.InterfaceComputations: surface_specific_humidity, SkinTemperature, BulkTemperature, - DiffusiveFlux + DiffusiveFlux, + atmosphere_ocean_stability_functions, + atmosphere_sea_ice_stability_functions using Thermodynamics using CUDA @@ -21,6 +23,7 @@ using ClimaSeaIce.SeaIceDynamics using ClimaSeaIce.Rheologies import NumericalEarth.OceanSeaIceModels.InterfaceComputations: surface_specific_humidity +import Oceananigans.Utils: TabulatedFunction using Statistics: mean, std @@ -257,6 +260,82 @@ end end end +@testset "Tabulated stability functions" begin + T_metadata = ECCOMetadatum(:temperature; date=DateTime(1993, 1, 1)) + S_metadata = ECCOMetadatum(:salinity; date=DateTime(1993, 1, 1)) + h_metadata = ECCOMetadatum(:sea_ice_thickness; date=DateTime(1993, 1, 1)) + ℵ_metadata = ECCOMetadatum(:sea_ice_concentration; date=DateTime(1993, 1, 1)) + + grid = Field(T_metadata).grid + ocean = ocean_simulation(grid, closure=nothing, momentum_advection=nothing, tracer_advection=nothing) + + atmosphere = JRA55PrescribedAtmosphere(architecture(grid); backend = JRA55NetCDFBackend(2)) + + set!(ocean.model; T=T_metadata, S=S_metadata) + + sea_ice = sea_ice_simulation(grid, ocean) + + set!(sea_ice.model, h=h_metadata, ℵ=ℵ_metadata) + + coupled_model_tabulated = OceanSeaIceModel(ocean, sea_ice; atmosphere) + + FT = eltype(grid) + + ao_stability_functions = atmosphere_ocean_stability_functions(FT, tabulate_stability_functions=false) + ai_stability_functions = atmosphere_sea_ice_stability_functions(FT, tabulate_stability_functions=false) + ao_exact_fluxes = SimilarityTheoryFluxes(eltype(grid); stability_functions = ao_stability_functions) + ai_exact_fluxes = SimilarityTheoryFluxes(eltype(grid); stability_functions = ai_stability_functions) + + interfaces_exact = ComponentInterfaces(atmosphere, ocean, sea_ice; + atmosphere_ocean_fluxes = ao_exact_fluxes, + atmosphere_sea_ice_fluxes = ai_exact_fluxes) + + coupled_model_exact = OceanSeaIceModel(ocean, sea_ice; atmosphere, interfaces=interfaces_exact) + + update_state!(coupled_model_tabulated) + update_state!(coupled_model_exact) + + fluxes_tabulated = coupled_model_tabulated.interfaces.atmosphere_ocean_interface.fluxes + fluxes_exact = coupled_model_exact.interfaces.atmosphere_ocean_interface.fluxes + + for flux_name in [:sensible_heat, :latent_heat, :water_vapor, :x_momentum, :y_momentum] + flux_tab = getfield(fluxes_tabulated, flux_name) + flux_exact = getfield(fluxes_exact, flux_name) + + flux_tab_data = interior(flux_tab, :, :, 1) + flux_exact_data = interior(flux_exact, :, :, 1) + + absolute_error = abs.(flux_tab_data .- flux_exact_data) + valid_mask = isfinite.(flux_exact_data) .& isfinite.(flux_tab_data) + mean_flux_data = mean(abs.(flux_exact_data[valid_mask])) + + abs_err_valid = absolute_error[valid_mask] + max_abs_error = maximum(abs_err_valid) + + @test max_abs_error < 0.01 * mean_flux_data + end + + fluxes_tabulated = coupled_model_tabulated.interfaces.atmosphere_sea_ice_interface.fluxes + fluxes_exact = coupled_model_exact.interfaces.atmosphere_sea_ice_interface.fluxes + + for flux_name in [:sensible_heat, :latent_heat, :water_vapor, :x_momentum, :y_momentum] + flux_tab = getfield(fluxes_tabulated, flux_name) + flux_exact = getfield(fluxes_exact, flux_name) + + flux_tab_data = interior(flux_tab, :, :, 1) + flux_exact_data = interior(flux_exact, :, :, 1) + + absolute_error = abs.(flux_tab_data .- flux_exact_data) + valid_mask = isfinite.(flux_exact_data) .& isfinite.(flux_tab_data) + mean_flux_data = mean(abs.(flux_exact_data[valid_mask])) + + abs_err_valid = absolute_error[valid_mask] + max_abs_error = maximum(abs_err_valid) + + @test max_abs_error < 0.01 * mean_flux_data + end +end + @testset "Fluxes regression" begin for arch in test_architectures @info "Testing fluxes regression..."