From 23640259b4f96424784587496f00e19205ada8b7 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Sun, 15 Dec 2024 17:53:27 +0100 Subject: [PATCH] let's go --- src/Advection/vector_invariant_advection.jl | 10 +- .../hydrostatic_spherical_coriolis.jl | 6 +- src/Operators/topology_aware_operators.jl | 96 ++++++++++++++++--- 3 files changed, 93 insertions(+), 19 deletions(-) diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index 20c6edcf2e..99a003e282 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -1,6 +1,6 @@ using Oceananigans.Operators using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ -using Oceananigans.Operators: mask_inactive_points_ℑxyᶠᶜᵃ, mask_inactive_points_ℑxyᶜᶠᵃ +using Oceananigans.Operators: ℑxyMᶠᶜᵃ, ℑxyMᶜᶠᵃ # These are also used in Coriolis/hydrostatic_spherical_coriolis.jl struct EnergyConserving{FT} <: AbstractAdvectionScheme{1, FT} end @@ -372,7 +372,7 @@ end Sζ = scheme.vorticity_stencil - @inbounds v̂ = mask_inactive_points_ℑxyᶠᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, v) / Δxᶠᶜᶜ(i, j, k, grid) + @inbounds v̂ = ℑxyMᶠᶜᶜ(i, j, k, grid, Δx_qᶜᶠᶜ, v) / Δxᶠᶜᶜ(i, j, k, grid) ζᴿ = _biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, scheme.vorticity_scheme, bias(v̂), ζ₃ᶠᶠᶜ, Sζ, u, v) return - v̂ * ζᴿ @@ -382,7 +382,7 @@ end Sζ = scheme.vorticity_stencil - @inbounds û = mask_inactive_points_ℑxyᶜᶠᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, u) / Δyᶜᶠᶜ(i, j, k, grid) + @inbounds û = ℑxyMᶜᶠᶜ(i, j, k, grid, Δy_qᶠᶜᶜ, u) / Δyᶜᶠᶜ(i, j, k, grid) ζᴿ = _biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, scheme.vorticity_scheme, bias(û), ζ₃ᶠᶠᶜ, Sζ, u, v) return + û * ζᴿ @@ -394,7 +394,7 @@ end @inline function U_dot_∇u(i, j, k, grid, advection::AbstractAdvectionScheme, U) - v̂ = ℑxᶠᵃᵃ(i, j, k, grid, ℑyᵃᶜᵃ, Δx_qᶜᶠᶜ, U.v) / Δxᶠᶜᶜ(i, j, k, grid) + v̂ = ℑxyMᶠᶜᶜ(i, j, k, grid, Δx_qᶜᶠᶜ, U.v) / Δxᶠᶜᶜ(i, j, k, grid) û = @inbounds U.u[i, j, k] return div_𝐯u(i, j, k, grid, advection, U, U.u) - @@ -404,7 +404,7 @@ end @inline function U_dot_∇v(i, j, k, grid, advection::AbstractAdvectionScheme, U) - û = ℑyᵃᶠᵃ(i, j, k, grid, ℑxᶜᵃᵃ, Δy_qᶠᶜᶜ, U.u) / Δyᶜᶠᶜ(i, j, k, grid) + û = ℑxyMᶜᶠᶜ(i, j, k, grid, Δy_qᶠᶜᶜ, U.u) / Δyᶜᶠᶜ(i, j, k, grid) v̂ = @inbounds U.v[i, j, k] return div_𝐯v(i, j, k, grid, advection, U, U.v) + diff --git a/src/Coriolis/hydrostatic_spherical_coriolis.jl b/src/Coriolis/hydrostatic_spherical_coriolis.jl index 9dd3cab861..8bc17fa783 100644 --- a/src/Coriolis/hydrostatic_spherical_coriolis.jl +++ b/src/Coriolis/hydrostatic_spherical_coriolis.jl @@ -1,5 +1,5 @@ using Oceananigans.Grids: LatitudeLongitudeGrid, OrthogonalSphericalShellGrid, peripheral_node, φnode -using Oceananigans.Operators: Δx_qᶜᶠᶜ, Δy_qᶠᶜᶜ, Δxᶠᶜᶜ, Δyᶜᶠᶜ, hack_sind, mask_inactive_points_ℑxyᶠᶜᵃ, mask_inactive_points_ℑxyᶜᶠᵃ +using Oceananigans.Operators: Δx_qᶜᶠᶜ, Δy_qᶠᶜᶜ, Δxᶠᶜᶜ, Δyᶜᶠᶜ, hack_sind, ℑxyMᶠᶜᶜ, ℑxyMᶜᶠᶜ using Oceananigans.Advection: EnergyConserving, EnstrophyConserving using Oceananigans.BoundaryConditions using Oceananigans.Fields @@ -69,11 +69,11 @@ const CoriolisActiveCellEnstrophyConserving = HydrostaticSphericalCoriolis{<:Act @inline x_f_cross_U(i, j, k, grid, coriolis::CoriolisActiveCellEnstrophyConserving, U) = @inbounds - ℑyᵃᶜᵃ(i, j, k, grid, fᶠᶠᵃ, coriolis) * - mask_inactive_points_ℑxyᶠᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, U[2]) / Δxᶠᶜᶜ(i, j, k, grid) + ℑxyMᶠᶜᶜ(i, j, k, grid, Δx_qᶜᶠᶜ, U[2]) / Δxᶠᶜᶜ(i, j, k, grid) @inline y_f_cross_U(i, j, k, grid, coriolis::CoriolisActiveCellEnstrophyConserving, U) = @inbounds + ℑxᶜᵃᵃ(i, j, k, grid, fᶠᶠᵃ, coriolis) * - mask_inactive_points_ℑxyᶜᶠᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, U[1]) / Δyᶜᶠᶜ(i, j, k, grid) + ℑxyMᶜᶠᶜ(i, j, k, grid, Δy_qᶠᶜᶜ, U[1]) / Δyᶜᶠᶜ(i, j, k, grid) ##### ##### Enstrophy-conserving scheme diff --git a/src/Operators/topology_aware_operators.jl b/src/Operators/topology_aware_operators.jl index 9d3e76bad1..4414e17623 100644 --- a/src/Operators/topology_aware_operators.jl +++ b/src/Operators/topology_aware_operators.jl @@ -66,18 +66,92 @@ const AGYL = AbstractUnderlyingGrid{FT, <:Any, LeftConnected} where FT @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) -# Masking interpolation operators + +#### +#### Masking interpolation operators +#### + +# Interpolation operators that mask inactive points. +# They assume that inactive points are zero. @inline not_peripheral_node(args...) = !peripheral_node(args...) -@inline function mask_inactive_points_ℑxyᶠᶜᵃ(i, j, k, grid, f::Function, args...) - neighboring_active_nodes = ℑxyᶠᶜᵃ(i, j, k, grid, not_peripheral_node, Center(), Face(), Center()) - return ifelse(neighboring_active_nodes == 0, zero(grid), - ℑxyᶠᶜᵃ(i, j, k, grid, f, args...) / neighboring_active_nodes) +@inline flip(::Type{Center}) = Face +@inline flip(::Type{Face}) = Center + +for LX in (:Center, :Face), LY in (:Center, :Face), LZ in (:Center, :Face) + LXe = @eval $LX + LYe = @eval $LY + LZe = @eval $LZ + + LXf = flip(LXe) + LYf = flip(LYe) + LZf = flip(LZe) + + ℑxˡᵃᵃ = Symbol(:ℑx, location_code(LXe, nothing, nothing)) + ℑyᵃˡᵃ = Symbol(:ℑy, location_code(nothing, LYe, nothing)) + ℑzᵃᵃˡ = Symbol(:ℑz, location_code(nothing, nothing, LZe)) + + ℑxMˡˡˡ = Symbol(:ℑxM, location_code(LXe, LYe, LZe)) + ℑyMˡˡˡ = Symbol(:ℑyM, location_code(LXe, LYe, LZe)) + ℑzMˡˡˡ = Symbol(:ℑzM, location_code(LXe, LYe, LZe)) + + @eval begin + @inline function $ℑxMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑxˡᵃᵃ(i, j, k, grid, not_peripheral_node, $LXf(), $LYe(), $LZe()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑxˡᵃᵃ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + + @inline function $ℑyMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑyᵃˡᵃ(i, j, k, grid, not_peripheral_node, $LXe(), $LYf(), $LZe()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑyᵃˡᵃ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + + @inline function $ℑzMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑzᵃᵃˡ(i, j, k, grid, not_peripheral_node, $LXe(), $LYe(), $LZf()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑzᵃᵃˡ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + end + + ℑxyˡˡᵃ = Symbol(:ℑxy, location_code(LXe, LYe, nothing)) + ℑyzᵃˡˡ = Symbol(:ℑyz, location_code(nothing, LYe, LZe)) + ℑxzˡᵃˡ = Symbol(:ℑxz, location_code(LXe, nothing, LZe)) + + ℑxyMˡˡˡ = Symbol(:ℑxyM, location_code(LXe, LYe, LZe)) + ℑyzMˡˡˡ = Symbol(:ℑyzM, location_code(LXe, LYe, LZe)) + ℑxzMˡˡˡ = Symbol(:ℑxzM, location_code(LXe, LYe, LZe)) + + @eval begin + @inline function $ℑxyMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑxyˡˡᵃ(i, j, k, grid, not_peripheral_node, $LXf(), $LYf(), $LZe()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑxyˡˡᵃ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + + @inline function $ℑyzMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑyzᵃˡˡ(i, j, k, grid, not_peripheral_node, $LXe(), $LYf(), $LZf()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑyzᵃˡˡ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + + @inline function $ℑxzMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑxzˡᵃˡ(i, j, k, grid, not_peripheral_node, $LXf(), $LYe(), $LZf()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑxzˡᵃˡ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + end + + ℑxyzˡˡˡ = Symbol(:ℑxyz, location_code(LXe, LYe, LZe)) + ℑxyzMˡˡˡ = Symbol(:ℑxyzM, location_code(LXe, LYe, LZe)) + + @eval begin + @inline function $ℑxyzMˡˡˡ(i, j, k, grid, f::Function, args...) + neighboring_active_nodes = $ℑxyzˡˡˡ(i, j, k, grid, not_peripheral_node, $LXf(), $LYf(), $LZf()) + return ifelse(neighboring_active_nodes == 0, zero(grid), + $ℑxyzˡˡˡ(i, j, k, grid, f, args...) / neighboring_active_nodes) + end + end end - -@inline function mask_inactive_points_ℑxyᶜᶠᵃ(i, j, k, grid, f::Function, args...) - neighboring_active_nodes = @inbounds ℑxyᶜᶠᵃ(i, j, k, grid, not_peripheral_node, Face(), Center(), Center()) - return ifelse(neighboring_active_nodes == 0, zero(grid), - ℑxyᶜᶠᵃ(i, j, k, grid, f, args...) / neighboring_active_nodes) -end \ No newline at end of file