Tabulate stability functions for faster evaluation

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)

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,23 +588,39 @@ 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, ζ)
     stable = ζ > 0
     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)

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..."