From 9ffbee31bc5a2fa38dd93fa1594b94cddaebba8c Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 18:41:23 +0100 Subject: [PATCH] Topology-aware operators (#3268) * new operators * comment * no disambiguation needed * bugfix * changed names * just to make the tests pass * Rename topology-aware operators * Rename topology-aware immersed boundary conditional differences * Update operator names in split-explicit free surface kernels * Fix some typos and functional call signatures --------- Co-authored-by: Ali Ramadhan --- .../conditional_differences.jl | 25 ++- .../split_explicit_free_surface_kernels.jl | 160 +++++------------- src/Operators/Operators.jl | 5 + src/Operators/topology_aware_operators.jl | 67 ++++++++ test/test_hydrostatic_regression.jl | 6 +- 5 files changed, 144 insertions(+), 119 deletions(-) create mode 100644 src/Operators/topology_aware_operators.jl diff --git a/src/ImmersedBoundaries/conditional_differences.jl b/src/ImmersedBoundaries/conditional_differences.jl index 1e9ae156bf..43706507d3 100644 --- a/src/ImmersedBoundaries/conditional_differences.jl +++ b/src/ImmersedBoundaries/conditional_differences.jl @@ -1,4 +1,4 @@ -import Oceananigans.Operators: +import Oceananigans.Operators: δxᶠᶜᶜ, δxᶠᶜᶠ, δxᶠᶠᶜ, δxᶠᶠᶠ, δxᶜᶜᶜ, δxᶜᶜᶠ, δxᶜᶠᶜ, δxᶜᶠᶠ, δyᶜᶠᶜ, δyᶜᶠᶠ, δyᶠᶠᶜ, δyᶠᶠᶠ, @@ -6,6 +6,10 @@ import Oceananigans.Operators: δzᶜᶜᶠ, δzᶜᶠᶠ, δzᶠᶜᶠ, δzᶠᶠᶠ, δzᶜᶜᶜ, δzᶜᶠᶜ, δzᶠᶜᶜ, δzᶠᶠᶜ +import Oceananigans.Operators: + δxTᶜᵃᵃ, δyTᵃᶜᵃ, + ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ + # Conditional differences that are "immersed boundary aware". # Here we return `zero(ibg)` rather than `δx` (for example) when _one_ of the # nodes involved in the difference is `immersed_inactive_node`. @@ -59,15 +63,30 @@ for (d, ξ) in enumerate((:x, :y, :z)) # `other_locs` contains locations in the two "other" directions not being differenced other_locs = [] - for l in 1:3 + for l in 1:3 if l != d push!(other_locs, loc[l]) end end - + @eval begin @inline $δξ(i, j, k, ibg::IBG, args...) = $conditional_δξ($(other_locs[1]), $(other_locs[2]), i, j, k, ibg, $δξ, args...) @inline $δξ(i, j, k, ibg::IBG, f::Function, args...) = $conditional_δξ($(other_locs[1]), $(other_locs[2]), i, j, k, ibg, $δξ, f::Function, args...) end end end + +# Topology-aware immersed boundary operators +# * Velocities are `0` on `peripheral_node`s and +# * Tracers should ensure no-flux on `inactive_node`s + +@inline conditional_U_fcc(i, j, k, grid, ibg::IBG, f::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, f, c, c), zero(ibg), f(i, j, k, grid, args...)) +@inline conditional_V_cfc(i, j, k, grid, ibg::IBG, f::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, c, f, c), zero(ibg), f(i, j, k, grid, args...)) + +@inline conditional_∂xTᶠᶜᶠ(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i-1, j, k, ibg, c, c, f), zero(ibg), ∂xTᶠᶜᶠ(i, j, k, ibg.underlying_grid, args...)) +@inline conditional_∂yTᶜᶠᶠ(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i, j-1, k, ibg, c, c, f), zero(ibg), ∂yTᶜᶠᶠ(i, j, k, ibg.underlying_grid, args...)) + +@inline δxTᶜᵃᵃ(i, j, k, ibg::IBG, f::Function, args...) = δxTᶜᵃᵃ(i, j, k, ibg.underlying_grid, conditional_U_fcc, ibg, f, args...) +@inline δyTᵃᶜᵃ(i, j, k, ibg::IBG, f::Function, args...) = δyTᵃᶜᵃ(i, j, k, ibg.underlying_grid, conditional_V_cfc, ibg, f, args...) +@inline ∂xTᶠᶜᶠ(i, j, k, ibg::IBG, f::Function, args...) = conditional_∂xTᶠᶜᶠ(i, j, k, ibg, f, args...) +@inline ∂yTᶜᶠᶠ(i, j, k, ibg::IBG, f::Function, args...) = conditional_∂yTᶜᶠᶠ(i, j, k, ibg, f, args...) diff --git a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl index 91ffed889f..98413d5eb8 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl @@ -1,6 +1,6 @@ using Oceananigans.Grids: topology using Oceananigans.Utils -using Oceananigans.AbstractOperations: Δz +using Oceananigans.AbstractOperations: Δz using Oceananigans.BoundaryConditions using Oceananigans.Operators using Oceananigans.Architectures: convert_args @@ -28,79 +28,15 @@ const μ = 1 - δ - γ - ϵ # # ∂t(η) = -∇⋅U # ∂t(U) = - gH∇η + f -# +# # the free surface field η and its average η̄ are located on `Face`s at the surface (grid.Nz +1). All other intermediate variables # (U, V, Ū, V̄) are barotropic fields (`ReducedField`) for which a k index is not defined -# Operators specific to the advancement of the Free surface and the Barotropic velocity. In particular, the base operators follow -# these rules: -# -# `δxᶠᵃᵃ_η` : Hardcodes Noflux or Periodic boundary conditions for the free surface η in x direction -# `δyᵃᶠᵃ_η` : Hardcodes Noflux or Periodic boundary conditions for the free surface η in y direction -# -# `δxᶜᵃᵃ_U` : Hardcodes NoPenetration or Periodic boundary conditions for the zonal barotropic velocity U in x direction -# `δyᵃᶜᵃ_V` : Hardcodes NoPenetration or Periodic boundary conditions for the meridional barotropic velocity V in y direction -# -# The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 -@inline δxᶠᵃᵃ_η(i, j, k, grid, T, η★::Function, args...) = δxᶠᵃᵃ(i, j, k, grid, η★, args...) -@inline δyᵃᶠᵃ_η(i, j, k, grid, T, η★::Function, args...) = δyᵃᶠᵃ(i, j, k, grid, η★, args...) -@inline δxᶜᵃᵃ_U(i, j, k, grid, T, U★::Function, args...) = δxᶜᵃᵃ(i, j, k, grid, U★, args...) -@inline δyᵃᶜᵃ_V(i, j, k, grid, T, V★::Function, args...) = δyᵃᶜᵃ(i, j, k, grid, V★, args...) - -@inline δxᶠᵃᵃ_η(i, j, k, grid, ::Type{Periodic}, η★::Function, args...) = ifelse(i == 1, η★(1, j, k, grid, args...) - η★(grid.Nx, j, k, grid, args...), δxᶠᵃᵃ(i, j, k, grid, η★, args...)) -@inline δyᵃᶠᵃ_η(i, j, k, grid, ::Type{Periodic}, η★::Function, args...) = ifelse(j == 1, η★(i, 1, k, grid, args...) - η★(i, grid.Ny, k, grid, args...), δyᵃᶠᵃ(i, j, k, grid, η★, args...)) - -@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{Periodic}, U★::Function, args...) = ifelse(i == grid.Nx, U★(1, j, k, grid, args...) - U★(grid.Nx, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...)) -@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{Periodic}, V★::Function, args...) = ifelse(j == grid.Ny, V★(i, 1, k, grid, args...) - V★(i, grid.Ny, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...)) - -# Enforce NoFlux conditions for `η★` - -@inline δxᶠᵃᵃ_η(i, j, k, grid, ::Type{Bounded}, η★::Function, args...) = ifelse(i == 1, zero(grid), δxᶠᵃᵃ(i, j, k, grid, η★, args...)) -@inline δyᵃᶠᵃ_η(i, j, k, grid, ::Type{Bounded}, η★::Function, args...) = ifelse(j == 1, zero(grid), δyᵃᶠᵃ(i, j, k, grid, η★, args...)) -@inline δxᶠᵃᵃ_η(i, j, k, grid, ::Type{RightConnected}, η★::Function, args...) = ifelse(i == 1, zero(grid), δxᶠᵃᵃ(i, j, k, grid, η★, args...)) -@inline δyᵃᶠᵃ_η(i, j, k, grid, ::Type{RightConnected}, η★::Function, args...) = ifelse(j == 1, zero(grid), δyᵃᶠᵃ(i, j, k, grid, η★, args...)) - -# Enforce Impenetrability conditions for `U★` and `V★` - -@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{Bounded}, U★::Function, args...) = ifelse(i == grid.Nx, - U★(i, j, k, grid, args...), - ifelse(i == 1, U★(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...))) -@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{Bounded}, V★::Function, args...) = ifelse(j == grid.Ny, - V★(i, j, k, grid, args...), - ifelse(j == 1, V★(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...))) - -@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{LeftConnected}, U★::Function, args...) = ifelse(i == grid.Nx, - U★(i, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...)) -@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{LeftConnected}, V★::Function, args...) = ifelse(j == grid.Ny, - V★(i, j, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...)) - -@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{RightConnected}, U★::Function, args...) = ifelse(i == 1, U★(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...)) -@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{RightConnected}, V★::Function, args...) = ifelse(j == 1, V★(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...)) - -# Derivative Operators - -@inline ∂xᶠᶜᶠ_η(i, j, k, grid, T, η★::Function, args...) = δxᶠᵃᵃ_η(i, j, k, grid, T, η★, args...) / Δxᶠᶜᶠ(i, j, k, grid) -@inline ∂yᶜᶠᶠ_η(i, j, k, grid, T, η★::Function, args...) = δyᵃᶠᵃ_η(i, j, k, grid, T, η★, args...) / Δyᶜᶠᶠ(i, j, k, grid) - -@inline div_xᶜᶜᶠ_U(i, j, k, grid, TX, U★, args...) = 1 / Azᶜᶜᶠ(i, j, k, grid) * δxᶜᵃᵃ_U(i, j, k, grid, TX, Δy_qᶠᶜᶠ, U★, args...) -@inline div_yᶜᶜᶠ_V(i, j, k, grid, TY, V★, args...) = 1 / Azᶜᶜᶠ(i, j, k, grid) * δyᵃᶜᵃ_V(i, j, k, grid, TY, Δx_qᶜᶠᶠ, V★, args...) - -# Immersed Boundary Operators (Velocities are `0` on `peripheral_node`s and the free surface should ensure no-flux on `inactive_node`s) - -@inline conditional_U_fcc(i, j, k, grid, ibg::IBG, U★::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, f, c, c), zero(ibg), U★(i, j, k, grid, args...)) -@inline conditional_V_cfc(i, j, k, grid, ibg::IBG, V★::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, c, f, c), zero(ibg), V★(i, j, k, grid, args...)) - -@inline conditional_∂xᶠᶜᶠ_η(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i-1, j, k, ibg, c, c, f), zero(ibg), ∂xᶠᶜᶠ_η(i, j, k, ibg.underlying_grid, args...)) -@inline conditional_∂yᶜᶠᶠ_η(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i, j-1, k, ibg, c, c, f), zero(ibg), ∂yᶜᶠᶠ_η(i, j, k, ibg.underlying_grid, args...)) - -@inline δxᶜᵃᵃ_U(i, j, k, ibg::IBG, T, U★::Function, args...) = δxᶜᵃᵃ_U(i, j, k, ibg.underlying_grid, T, conditional_U_fcc, ibg, U★, args...) -@inline δyᵃᶜᵃ_V(i, j, k, ibg::IBG, T, V★::Function, args...) = δyᵃᶜᵃ_V(i, j, k, ibg.underlying_grid, T, conditional_V_cfc, ibg, V★, args...) -@inline ∂xᶠᶜᶠ_η(i, j, k, ibg::IBG, T, η★::Function, args...) = conditional_∂xᶠᶜᶠ_η(i, j, k, ibg, T, η★, args...) -@inline ∂yᶜᶠᶠ_η(i, j, k, ibg::IBG, T, η★::Function, args...) = conditional_∂yᶜᶠᶠ_η(i, j, k, ibg, T, η★, args...) +# Special ``partial'' divergence for free surface evolution +@inline div_Txᶜᶜᶠ(i, j, k, grid, U★::Function, args...) = 1 / Azᶜᶜᶠ(i, j, k, grid) * δxTᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶠ, U★, args...) +@inline div_Tyᶜᶜᶠ(i, j, k, grid, V★::Function, args...) = 1 / Azᶜᶜᶠ(i, j, k, grid) * δyTᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶠ, V★, args...) -# Disambiguation -for Topo in [:Periodic, :Bounded, :RightConnected, :LeftConnected] - @eval begin - @inline δxᶜᵃᵃ_U(i, j, k, ibg::IBG, T::Type{$Topo}, U★::Function, args...) = δxᶜᵃᵃ_U(i, j, k, ibg.underlying_grid, T, conditional_U_fcc, ibg, U★, args...) - @inline δyᵃᶜᵃ_V(i, j, k, ibg::IBG, T::Type{$Topo}, V★::Function, args...) = δyᵃᶜᵃ_V(i, j, k, ibg.underlying_grid, T, conditional_V_cfc, ibg, V★, args...) - end -end +# The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 # Time stepping extrapolation U★, and η★ @@ -143,49 +79,48 @@ end free_surface_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², timestepper) end + @inline function free_surface_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², timestepper) k_top = grid.Nz+1 TX, TY, _ = topology(grid) - @inbounds begin + @inbounds begin advance_previous_free_surface!(i, j, k_top, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) - η[i, j, k_top] -= Δτ * (div_xᶜᶜᶠ_U(i, j, k_top-1, grid, TX, U★, timestepper, U, Uᵐ⁻¹, Uᵐ⁻²) + - div_yᶜᶜᶠ_V(i, j, k_top-1, grid, TY, U★, timestepper, V, Vᵐ⁻¹, Vᵐ⁻²)) + η[i, j, k_top] -= Δτ * (div_Txᶜᶜᶠ(i, j, k_top-1, grid, U★, timestepper, U, Uᵐ⁻¹, Uᵐ⁻²) + + div_Tyᶜᶜᶠ(i, j, k_top-1, grid, U★, timestepper, V, Vᵐ⁻¹, Vᵐ⁻²)) end return nothing end -@kernel function _split_explicit_barotropic_velocity!(averaging_weight, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², +@kernel function _split_explicit_barotropic_velocity!(averaging_weight, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², - η̅, U̅, V̅, Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + η̅, U̅, V̅, Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, timestepper) i, j = @index(Global, NTuple) - velocity_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + velocity_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², η̅, U̅, V̅, averaging_weight, - Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, timestepper) end -@inline function velocity_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², +@inline function velocity_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², η̅, U̅, V̅, averaging_weight, - Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, timestepper) k_top = grid.Nz+1 - - TX, TY, _ = topology(grid) - @inbounds begin + @inbounds begin advance_previous_velocity!(i, j, k_top-1, timestepper, U, Uᵐ⁻¹, Uᵐ⁻²) advance_previous_velocity!(i, j, k_top-1, timestepper, V, Vᵐ⁻¹, Vᵐ⁻²) # ∂τ(U) = - ∇η + G - U[i, j, k_top-1] += Δτ * (- g * Hᶠᶜ[i, j, 1] * ∂xᶠᶜᶠ_η(i, j, k_top, grid, TX, η★, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + Gᵁ[i, j, k_top-1]) - V[i, j, k_top-1] += Δτ * (- g * Hᶜᶠ[i, j, 1] * ∂yᶜᶠᶠ_η(i, j, k_top, grid, TY, η★, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + Gⱽ[i, j, k_top-1]) - + U[i, j, k_top-1] += Δτ * (- g * Hᶠᶜ[i, j] * ∂xTᶠᶜᶠ(i, j, k_top, grid, η★, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + Gᵁ[i, j, 1]) + V[i, j, k_top-1] += Δτ * (- g * Hᶜᶠ[i, j] * ∂yTᶜᶠᶠ(i, j, k_top, grid, η★, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + Gⱽ[i, j, 1]) + # time-averaging η̅[i, j, k_top] += averaging_weight * η[i, j, k_top] U̅[i, j, k_top-1] += averaging_weight * U[i, j, k_top-1] @@ -196,7 +131,7 @@ end # Barotropic Model Kernels # u_Δz = u * Δz @kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v) - i, j = @index(Global, NTuple) + i, j = @index(Global, NTuple) k_top = grid.Nz+1 @inbounds U[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] @@ -224,7 +159,7 @@ end end end -@inline function compute_barotropic_mode!(U, V, grid, u, v) +@inline function compute_barotropic_mode!(U, V, grid, u, v) active_cells_map = retrieve_surface_active_cells_map(grid) launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map) @@ -279,7 +214,7 @@ function barotropic_split_explicit_corrector!(u, v, free_surface, grid) arch = architecture(grid) - # take out "bad" barotropic mode, + # take out "bad" barotropic mode, # !!!! reusing U and V for this storage since last timestep doesn't matter compute_barotropic_mode!(U, V, grid, u, v) # add in "good" barotropic mode @@ -310,23 +245,23 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac wait_free_surface_communication!(free_surface, architecture(free_surface_grid)) # Calculate the substepping parameterers - settings = free_surface.settings + settings = free_surface.settings Nsubsteps = calculate_substeps(settings.substepping, Δt) - + # barotropic time step as fraction of baroclinic step and averaging weights - fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps) + fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps) Nsubsteps = length(weights) # barotropic time step in seconds - Δτᴮ = fractional_Δt * Δt - + Δτᴮ = fractional_Δt * Δt + # reset free surface averages - @apply_regionally begin + @apply_regionally begin initialize_free_surface_state!(free_surface.state, free_surface.η, settings.timestepper) - + # Solve for the free surface at tⁿ⁺¹ iterate_split_explicit!(free_surface, free_surface_grid, Δτᴮ, weights, Val(Nsubsteps)) - + # Reset eta for the next timestep set!(free_surface.η, free_surface.state.η̅) end @@ -335,7 +270,7 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac fill_halo_regions!(fields_to_fill; async = true) # Preparing velocities for the barotropic correction - @apply_regionally begin + @apply_regionally begin mask_immersed_field!(model.velocities.u) mask_immersed_field!(model.velocities.v) end @@ -347,7 +282,7 @@ end const FNS = FixedSubstepNumber const FTS = FixedTimeStepSize -# since weights can be negative in the first few substeps (as in the default averaging kernel), +# since weights can be negative in the first few substeps (as in the default averaging kernel), # we set a minimum number of substeps to execute to avoid numerical issues const MINIMUM_SUBSTEPS = 5 @@ -370,7 +305,7 @@ function iterate_split_explicit!(free_surface, grid, Δτᴮ, weights, ::Val{Nsu settings = free_surface.settings g = free_surface.gravitational_acceleration - # unpack state quantities, parameters and forcing terms + # unpack state quantities, parameters and forcing terms U, V = state.U, state.V Uᵐ⁻¹, Uᵐ⁻² = state.Uᵐ⁻¹, state.Uᵐ⁻² Vᵐ⁻¹, Vᵐ⁻² = state.Vᵐ⁻¹, state.Vᵐ⁻² @@ -385,13 +320,13 @@ function iterate_split_explicit!(free_surface, grid, Δτᴮ, weights, ::Val{Nsu free_surface_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_free_surface!) barotropic_velocity_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_barotropic_velocity!) - η_args = (grid, Δτᴮ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², - U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², + η_args = (grid, Δτᴮ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², timestepper) - U_args = (grid, Δτᴮ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U_args = (grid, Δτᴮ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², - η̅, U̅, V̅, Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + η̅, U̅, V̅, Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, timestepper) GC.@preserve η_args U_args begin @@ -422,10 +357,10 @@ end @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ) @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ) - for k in 2:grid.Nz + for k in 2:grid.Nz @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ) @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ) - end + end end # Calculate RHS for the barotropic time step.q @@ -437,10 +372,10 @@ end @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ) @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ) - for k in 2:grid.Nz + for k in 2:grid.Nz @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ) @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ) - end + end end @inline ab2_step_Gu(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT = @@ -452,8 +387,8 @@ end # Setting up the RHS for the barotropic step (tendencies of the barotropic velocity components) # This function is called after `calculate_tendency` and before `ab2_step_velocities!` function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ) - - # we start the time integration of η from the average ηⁿ + + # we start the time integration of η from the average ηⁿ Gu⁻ = model.timestepper.G⁻.u Gv⁻ = model.timestepper.G⁻.v Guⁿ = model.timestepper.Gⁿ.u @@ -469,14 +404,13 @@ function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ) return nothing end -@inline function setup_split_explicit_tendency!(auxiliary, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) +@inline function setup_split_explicit_tendency!(auxiliary, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) active_cells_map = retrieve_surface_active_cells_map(grid) - launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, auxiliary.Gᵁ, auxiliary.Gⱽ, grid, + launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, auxiliary.Gᵁ, auxiliary.Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ; active_cells_map) return nothing end - -wait_free_surface_communication!(free_surface, arch) = nothing +wait_free_surface_communication!(free_surface, arch) = nothing diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index 570d43dd6b..0a80bc183e 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -60,6 +60,10 @@ export ℑxᶜᵃᵃ, ℑxᶠᵃᵃ, ℑyᵃᶜᵃ, ℑyᵃᶠᵃ, ℑzᵃᵃᶜ export ℑxyᶜᶜᵃ, ℑxyᶠᶜᵃ, ℑxyᶠᶠᵃ, ℑxyᶜᶠᵃ, ℑxzᶜᵃᶜ, ℑxzᶠᵃᶜ, ℑxzᶠᵃᶠ, ℑxzᶜᵃᶠ, ℑyzᵃᶜᶜ, ℑyzᵃᶠᶜ, ℑyzᵃᶠᶠ, ℑyzᵃᶜᶠ export ℑxyzᶜᶜᶠ, ℑxyzᶜᶠᶜ, ℑxyzᶠᶜᶜ, ℑxyzᶜᶠᶠ, ℑxyzᶠᶜᶠ, ℑxyzᶠᶠᶜ, ℑxyzᶜᶜᶜ, ℑxyzᶠᶠᶠ +# Topology-aware operators +export δxTᶠᵃᵃ, δyTᵃᶠᵃ, δxTᶜᵃᵃ, δyTᵃᶜᵃ +export ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ + # Reference frame conversion export intrinsic_vector, extrinsic_vector @@ -86,6 +90,7 @@ include("products_between_fields_and_grid_metrics.jl") include("derivative_operators.jl") include("divergence_operators.jl") +include("topology_aware_operators.jl") include("vorticity_operators.jl") include("laplacian_operators.jl") diff --git a/src/Operators/topology_aware_operators.jl b/src/Operators/topology_aware_operators.jl new file mode 100644 index 0000000000..cab68c3b6e --- /dev/null +++ b/src/Operators/topology_aware_operators.jl @@ -0,0 +1,67 @@ +using Oceananigans.Grids: AbstractUnderlyingGrid + +const AGXB = AbstractUnderlyingGrid{FT, Bounded} where FT +const AGXP = AbstractUnderlyingGrid{FT, Periodic} where FT +const AGXR = AbstractUnderlyingGrid{FT, RightConnected} where FT +const AGXL = AbstractUnderlyingGrid{FT, LeftConnected} where FT + +const AGYB = AbstractUnderlyingGrid{FT, <:Any, Bounded} where FT +const AGYP = AbstractUnderlyingGrid{FT, <:Any, Periodic} where FT +const AGYR = AbstractUnderlyingGrid{FT, <:Any, RightConnected} where FT +const AGYL = AbstractUnderlyingGrid{FT, <:Any, LeftConnected} where FT + +# Topology-aware Operators with the following convention: +# +# `δxTᶠᵃᵃ` : Hardcodes `Noflux` or `Periodic` boundary conditions for a (Center, Center, Center) function `f` in the x-direction. +# `δyTᵃᶠᵃ` : Hardcodes `Noflux` or `Periodic` boundary conditions for a (Center, Center, Center) function `f` in the y-direction +# +# `δxTᶜᵃᵃ` : Hardcodes `NoPenetration` or `Periodic` boundary conditions for a (Face, Center, Center) function `U` in x direction +# `δyTᵃᶜᵃ` : Hardcodes `NoPenetration` or `Periodic` boundary conditions for a (Center, Face, Center) function `V` in y direction +# +# Note: The naming convention is that `T` denotes a topology-aware operator. So `δxTᶠᵃᵃ` is the topology-aware version of `δxᶠᵃᵃ`. + +# Fallback + +@inline δxTᶠᵃᵃ(i, j, k, grid, f::Function, args...) = δxᶠᵃᵃ(i, j, k, grid, f, args...) +@inline δyTᵃᶠᵃ(i, j, k, grid, f::Function, args...) = δyᵃᶠᵃ(i, j, k, grid, f, args...) +@inline δxTᶜᵃᵃ(i, j, k, grid, f::Function, args...) = δxᶜᵃᵃ(i, j, k, grid, f, args...) +@inline δyTᵃᶜᵃ(i, j, k, grid, f::Function, args...) = δyᵃᶜᵃ(i, j, k, grid, f, args...) + +# Enforce Periodic conditions + +@inline δxTᶠᵃᵃ(i, j, k, grid::AGXP, f::Function, args...) = ifelse(i == 1, f(1, j, k, grid, args...) - f(grid.Nx, j, k, grid, args...), δxᶠᵃᵃ(i, j, k, grid, f, args...)) +@inline δyTᵃᶠᵃ(i, j, k, grid::AGYP, f::Function, args...) = ifelse(j == 1, f(i, 1, k, grid, args...) - f(i, grid.Ny, k, grid, args...), δyᵃᶠᵃ(i, j, k, grid, f, args...)) + +@inline δxTᶜᵃᵃ(i, j, k, grid::AGXP, f::Function, args...) = ifelse(i == grid.Nx, f(1, j, k, grid, args...) - f(grid.Nx, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...)) +@inline δyTᵃᶜᵃ(i, j, k, grid::AGYP, f::Function, args...) = ifelse(j == grid.Ny, f(i, 1, k, grid, args...) - f(i, grid.Ny, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...)) + +# Enforce NoFlux conditions + +@inline δxTᶠᵃᵃ(i, j, k, grid::AGXB{FT}, f::Function, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...)) +@inline δyTᵃᶠᵃ(i, j, k, grid::AGYB{FT}, f::Function, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...)) + +@inline δxTᶠᵃᵃ(i, j, k, grid::AGXR{FT}, f::Function, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...)) +@inline δyTᵃᶠᵃ(i, j, k, grid::AGYR{FT}, f::Function, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...)) + +# Enforce Impenetrability conditions + +@inline δxTᶜᵃᵃ(i, j, k, grid::AGXB, f::Function, args...) = + ifelse(i == grid.Nx, - f(i, j, k, grid, args...), + ifelse(i == 1, f(2, j, k, grid, args...), + δxᶜᵃᵃ(i, j, k, grid, f, args...))) + +@inline δyTᵃᶜᵃ(i, j, k, grid::AGYB, f::Function, args...) = + ifelse(j == grid.Ny, - f(i, j, k, grid, args...), + ifelse(j == 1, f(i, 2, k, grid, args...), + δyᵃᶜᵃ(i, j, k, grid, f, args...))) + +@inline δxTᶜᵃᵃ(i, j, k, grid::AGXL, f::Function, args...) = ifelse(i == grid.Nx, - f(i, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...)) +@inline δyTᵃᶜᵃ(i, j, k, grid::AGYL, f::Function, args...) = ifelse(j == grid.Ny, - f(i, j, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...)) + +@inline δxTᶜᵃᵃ(i, j, k, grid::AGXR, f::Function, args...) = ifelse(i == 1, f(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...)) +@inline δyTᵃᶜᵃ(i, j, k, grid::AGYR, f::Function, args...) = ifelse(j == 1, f(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...)) + +# Derivative operators + +@inline ∂xTᶠᶜᶠ(i, j, k, grid, f::Function, args...) = δxTᶠᵃᵃ(i, j, k, grid, f, args...) / Δxᶠᶜᶠ(i, j, k, grid) +@inline ∂yTᶜᶠᶠ(i, j, k, grid, f::Function, args...) = δyTᵃᶠᵃ(i, j, k, grid, f, args...) / Δyᶜᶠᶠ(i, j, k, grid) diff --git a/test/test_hydrostatic_regression.jl b/test/test_hydrostatic_regression.jl index f7937d2eaa..62de3c8ffe 100644 --- a/test/test_hydrostatic_regression.jl +++ b/test/test_hydrostatic_regression.jl @@ -3,7 +3,7 @@ include("data_dependencies.jl") using Oceananigans.Grids: topology, XRegularLLG, YRegularLLG, ZRegularLLG -function show_hydrostatic_test(grid, free_surface, precompute_metrics) +function show_hydrostatic_test(grid, free_surface, precompute_metrics) typeof(grid) <: XRegularLLG ? gx = :regular : gx = :stretched typeof(grid) <: YRegularLLG ? gy = :regular : gy = :stretched @@ -82,8 +82,8 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") for free_surface in [explicit_free_surface, implicit_free_surface, split_explicit_free_surface] # GPU + ImplicitFreeSurface + precompute metrics cannot be tested on sverdrup at the moment - # because "uses too much parameter space (maximum 0x1100 bytes)" error - if !(precompute_metrics && free_surface isa ImplicitFreeSurface && arch isa GPU) && + # because "uses too much parameter space (maximum 0x1100 bytes)" error + if !(precompute_metrics && free_surface isa ImplicitFreeSurface && arch isa GPU) && !(free_surface isa ImplicitFreeSurface && arch isa Distributed) # Also no implicit free surface on distributed testset_str, info_str = show_hydrostatic_test(grid, free_surface, precompute_metrics)