Add ConformalMapping Type to OrthogonalSphericalShellGrid for Special…

…ized Dispatch (#3832) * Embed ConformalMapping in OSSGrid * Apply suggestions from code review * Apply suggestions from code review * move ConformalCubedSpherePanel definition * define ξ, η * ConformalCubedSphereMapping -> CubedSphereConformalMapping --------- Co-authored-by: Navid C. Constantinou <[email protected]>
CliMA · Oct 9, 2024 · 989108a · 989108a
1 parent fe056fb
commit 989108a
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git a/src/Grids/orthogonal_spherical_shell_grid.jl b/src/Grids/orthogonal_spherical_shell_grid.jl
@@ -9,6 +9,12 @@ using Adapt: adapt_structure
 using Oceananigans
 using Oceananigans.Grids: prettysummary, coordinate_summary, BoundedTopology, length
 
+struct CubedSphereConformalMapping{FT, Rotation}
+    ξ :: Tuple{FT, FT}
+    η :: Tuple{FT, FT}
+    rotation :: Rotation
+end
+
 struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, C, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch}
     architecture :: Arch
     Nx :: Int
@@ -70,6 +76,7 @@ end
 const OSSG = OrthogonalSphericalShellGrid
 const ZRegOSSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number}
 const ZRegOrthogonalSphericalShellGrid = ZRegOSSG
+const ConformalCubedSpherePanel = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:CubedSphereConformalMapping}
 
 # convenience constructor for OSSG without any conformal_mapping properties
 OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz, Lz,
@@ -603,7 +610,7 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU()
                      Δzᵃᵃᶜ, Δzᵃᵃᶠ,
                      Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ)
 
-    conformal_mapping = (; ξ, η, rotation)
+    conformal_mapping = CubedSphereConformalMapping(ξ, η, rotation)
 
     grid = OrthogonalSphericalShellGrid{TX, TY, TZ}(CPU(), Nξ, Nη, Nz, Hx, Hy, Hz, Lz,
                                                     coordinate_arrays...,
@@ -876,7 +883,8 @@ function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = C
     φᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2])
     φᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2])
 
-    conformal_mapping = (ξ = (-1, 1), η = (-1, 1))
+    ξ, η = (-1, 1), (-1, 1)
+    conformal_mapping = CubedSphereConformalMapping(ξ, η, rotation)
 
     return OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz, Lz,
                                                      λᶜᶜᵃ,  λᶠᶜᵃ,  λᶜᶠᵃ,  λᶠᶠᵃ,

diff --git a/src/Operators/vorticity_operators.jl b/src/Operators/vorticity_operators.jl
@@ -1,11 +1,13 @@
+using Oceananigans.Grids: ConformalCubedSpherePanel
+
 """ Vertical circulation associated with horizontal velocities u, v. """
 @inline Γᶠᶠᶜ(i, j, k, grid, u, v) = δxᶠᶠᶜ(i, j, k, grid, Δy_qᶜᶠᶜ, v) - δyᶠᶠᶜ(i, j, k, grid, Δx_qᶠᶜᶜ, u)
 
 """
-     ζ₃ᶠᶠᶜ(i, j, k, grid, u, v)
+    ζ₃ᶠᶠᶜ(i, j, k, grid, u, v)
 
- The vertical vorticity associated with horizontal velocities ``u`` and ``v``.
- """
+The vertical vorticity associated with horizontal velocities ``u`` and ``v``.
+"""
 @inline ζ₃ᶠᶠᶜ(i, j, k, grid, u, v) = Γᶠᶠᶜ(i, j, k, grid, u, v) / Azᶠᶠᶜ(i, j, k, grid)
 
 #####
@@ -14,11 +16,11 @@
 #####
 
 """
-     Γᶠᶠᶜ(i, j, k, grid, u, v)
+    Γᶠᶠᶜ(i, j, k, grid, u, v)
 
- The vertical circulation associated with horizontal velocities ``u`` and ``v``.
- """
-@inline function Γᶠᶠᶜ(i, j, k, grid::OrthogonalSphericalShellGrid, u, v)
+The vertical circulation associated with horizontal velocities ``u`` and ``v``.
+"""
+@inline function Γᶠᶠᶜ(i, j, k, grid::ConformalCubedSpherePanel, u, v)
     # South-west corner
     if i == 1 && j == 1
         return Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u)