Generalization of OrthogonalSphericalShellGrid to be able to construct any grid with coords-metrics that vary in both i, j

src/Grids/Grids.jl

@@ -10,6 +10,7 @@ export XFlatGrid, YFlatGrid, ZFlatGrid
 export XRegRectilinearGrid, YRegRectilinearGrid, ZRegRectilinearGrid, HRegRectilinearGrid, RegRectilinearGrid
 export LatitudeLongitudeGrid, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid
 export OrthogonalSphericalShellGrid, ConformalCubedSphereGrid, ZRegOrthogonalSphericalShellGrid
+export conformal_cubed_sphere_panel
 export node, nodes
 export xnode, ynode, znode, λnode, φnode
 export xnodes, ynodes, znodes, λnodes, φnodes

src/Grids/orthogonal_spherical_shell_grid.jl

@@ -9,18 +9,14 @@ using Adapt: adapt_structure
 using Oceananigans
 using Oceananigans.Grids: prettysummary, coordinate_summary
 
-struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch}
+struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, C, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch}
     architecture :: Arch
     Nx :: Int
     Ny :: Int
     Nz :: Int
     Hx :: Int
     Hy :: Int
     Hz :: Int
-    ξₗ :: FT    # left-most domain for cube's ξ coordinate
-    ξᵣ :: FT    # right-most domain for cube's ξ coordinate
-    ηₗ :: FT    # left-most domain for cube's η coordinate
-    ηᵣ :: FT    # right-most domain for cube's η coordinate
     λᶜᶜᵃ :: A
     λᶠᶜᵃ :: A
     λᶜᶠᵃ :: A
@@ -46,33 +42,44 @@ struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, Arch} <: AbstractH
     Azᶜᶠᵃ :: A
     Azᶠᶠᵃ :: A
     radius :: FT
+    conformal_mapping :: C
 
     OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture::Arch,
                                              Nx, Ny, Nz,
-                                             Hx, Hy, Hz, ξₗ, ξᵣ, ηₗ, ηᵣ,
+                                             Hx, Hy, Hz,
                                               λᶜᶜᵃ :: A,  λᶠᶜᵃ :: A,  λᶜᶠᵃ :: A,  λᶠᶠᵃ :: A,
                                               φᶜᶜᵃ :: A,  φᶠᶜᵃ :: A,  φᶜᶠᵃ :: A,  φᶠᶠᵃ :: A, zᵃᵃᶜ :: R, zᵃᵃᶠ :: R,
                                              Δxᶜᶜᵃ :: A, Δxᶠᶜᵃ :: A, Δxᶜᶠᵃ :: A, Δxᶠᶠᵃ :: A,
                                              Δyᶜᶜᵃ :: A, Δyᶜᶠᵃ :: A, Δyᶠᶜᵃ :: A, Δyᶠᶠᵃ :: A, Δzᵃᵃᶜ :: FR, Δzᵃᵃᶠ :: FR,
                                              Azᶜᶜᵃ :: A, Azᶠᶜᵃ :: A, Azᶜᶠᵃ :: A, Azᶠᶠᵃ :: A,
-                                             radius :: FT) where {TX, TY, TZ, FT, A, R, FR, Arch} =
-        new{FT, TX, TY, TZ, A, R, FR, Arch}(architecture,
+                                             radius :: FT,
+                                             conformal_mapping :: C) where {TX, TY, TZ, FT, A, R, FR, C, Arch} =
+        new{FT, TX, TY, TZ, A, R, FR, C, Arch}(architecture,
                                             Nx, Ny, Nz,
                                             Hx, Hy, Hz,
-                                            ξₗ, ξᵣ, ηₗ, ηᵣ,
                                             λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ,
                                             φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ,
                                             Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ,
                                             Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δzᵃᵃᶜ, Δzᵃᵃᶠ,
-                                            Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius)
+                                            Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius, conformal_mapping)
 end
 
 const OSSG = OrthogonalSphericalShellGrid
 const ZRegOSSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number}
 const ZRegOrthogonalSphericalShellGrid = ZRegOSSG
 
+# convenience constructor for OSSG without any conformal_mapping properties
+OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz,
+                             λᶜᶜᵃ,  λᶠᶜᵃ,  λᶜᶠᵃ,  λᶠᶠᵃ, φᶜᶜᵃ,  φᶠᶜᵃ,  φᶜᶠᵃ,  φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ,
+                             Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δzᵃᵃᶜ, Δzᵃᵃᶠ,
+                             Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) =
+    OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz,
+                                 λᶜᶜᵃ,  λᶠᶜᵃ,  λᶜᶠᵃ,  λᶠᶠᵃ, φᶜᶜᵃ,  φᶠᶜᵃ,  φᶜᶠᵃ,  φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ,
+                                 Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δzᵃᵃᶜ, Δzᵃᵃᶠ,
+                                 Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius, nothing)
+
 """
-    OrthogonalSphericalShellGrid(architecture::AbstractArchitecture = CPU(),
+    conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU(),
                                  FT::DataType = Float64;
                                  size,
                                  z,
@@ -125,15 +132,17 @@ Examples
 ```jldoctest
 julia> using Oceananigans
 
-julia> grid = OrthogonalSphericalShellGrid(size=(36, 34, 25), z=(-1000, 0))
+julia> using Oceananigans.Grids
+
+julia> grid = conformal_cubed_sphere_panel(size=(36, 34, 25), z=(-1000, 0))
 36×34×25 OrthogonalSphericalShellGrid{Float64, Bounded, Bounded, Bounded} on CPU with 1×1×1 halo and with precomputed metrics
 ├── centered at: North Pole, (λ, φ) = (0.0, 90.0)
 ├── longitude: Bounded  extent 90.0 degrees variably spaced with min(Δλ)=0.616164, max(Δλ)=2.58892
 ├── latitude:  Bounded  extent 90.0 degrees variably spaced with min(Δφ)=0.664958, max(Δφ)=2.74119
 └── z:         Bounded  z ∈ [-1000.0, 0.0]  regularly spaced with Δz=40.0
 ```
 """
-function OrthogonalSphericalShellGrid(architecture::AbstractArchitecture = CPU(),
+function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU(),
                                       FT::DataType = Float64;
                                       size,
                                       z,
@@ -568,10 +577,13 @@ function OrthogonalSphericalShellGrid(architecture::AbstractArchitecture = CPU()
                      Δzᵃᵃᶜ, Δzᵃᵃᶠ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ)
     metric_arrays = map(a -> arch_array(architecture, a), metric_arrays)
 
-    return OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz, ξ..., η...,
+    conformal_mapping = (; ξ, η)
+
+    return OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz,
                                                     coordinate_arrays...,
                                                     metric_arrays...,
-                                                    radius)
+                                                    radius,
+                                                    conformal_mapping)
 end
 
 function lat_lon_to_cartesian(lat, lon, radius)
@@ -586,7 +598,7 @@ lat_lon_to_z(lat, lon, radius) = radius * sind(lat)
 
 # architecture = CPU() default, assuming that a DataType positional arg
 # is specifying the floating point type.
-OrthogonalSphericalShellGrid(FT::DataType; kwargs...) = OrthogonalSphericalShellGrid(CPU(), FT; kwargs...)
+conformal_cubed_sphere_panel(FT::DataType; kwargs...) = conformal_cubed_sphere_panel(CPU(), FT; kwargs...)
 
 function load_and_offset_cubed_sphere_data(file, FT, arch, field_name, loc, topo, N, H)
 
@@ -607,7 +619,7 @@ function load_and_offset_cubed_sphere_data(file, FT, arch, field_name, loc, topo
     return offset_data(underlying_data, loc[1:2], topo[1:2], N[1:2], H[1:2])
 end
 
-function OrthogonalSphericalShellGrid(filepath::AbstractString, architecture = CPU(), FT = Float64;
+function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = CPU(), FT = Float64;
                                       panel, Nz, z,
                                       topology = (Bounded, Bounded, Bounded),
                                         radius = R_Earth,
@@ -621,7 +633,7 @@ function OrthogonalSphericalShellGrid(filepath::AbstractString, architecture = C
     ## The vertical coordinates can come out of the regular rectilinear grid!
 
     ξ, η = (-1, 1), (-1, 1)
-    ξη_grid = RectilinearGrid(architecture, FT; size=(1, 1, Nz), x=ξ, y=η, z, topology, halo)
+    ξη_grid = RectilinearGrid(architecture, FT; size = (1, 1, Nz), x = ξ, y = η, z, topology, halo)
 
      zᵃᵃᶠ = ξη_grid.zᵃᵃᶠ
      zᵃᵃᶜ = ξη_grid.zᵃᵃᶜ
@@ -674,14 +686,18 @@ function OrthogonalSphericalShellGrid(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])
 
-    return OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz, ξ..., η...,
+    conformal_mapping = (; ξ, η)
+
+    return OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz,
                                                      λᶜᶜᵃ,  λᶠᶜᵃ,  λᶜᶠᵃ,  λᶠᶠᵃ,
                                                      φᶜᶜᵃ,  φᶠᶜᵃ,  φᶜᶠᵃ,  φᶠᶠᵃ,
                                                      zᵃᵃᶜ,  zᵃᵃᶠ,
                                                     Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ,
                                                     Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ,
                                                     Δzᵃᵃᶜ, Δzᵃᵃᶠ,
-                                                    Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius)
+                                                    Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ,
+                                                    radius,
+                                                    conformal_mapping)
 end
 
 function on_architecture(arch::AbstractArchitecture, grid::OrthogonalSphericalShellGrid)
@@ -722,10 +738,11 @@ function on_architecture(arch::AbstractArchitecture, grid::OrthogonalSphericalSh
     new_grid = OrthogonalSphericalShellGrid{TX, TY, TZ}(arch,
                                                         grid.Nx, grid.Ny, grid.Nz,
                                                         grid.Hx, grid.Hy, grid.Hz,
-                                                        grid.ξₗ, grid.ξᵣ, grid.ηₗ, grid.ηᵣ,
                                                         coordinate_data...,
                                                         grid_spacing_data...,
-                                                        horizontal_area_data..., grid.radius)
+                                                        horizontal_area_data...,
+                                                        grid.radius,
+                                                        grid.conformal_mapping)
 
     return new_grid
 end
@@ -736,8 +753,6 @@ function Adapt.adapt_structure(to, grid::OrthogonalSphericalShellGrid)
     return OrthogonalSphericalShellGrid{TX, TY, TZ}(nothing,
                                                     grid.Nx, grid.Ny, grid.Nz,
                                                     grid.Hx, grid.Hy, grid.Hz,
-                                                    grid.ξₗ, grid.ξᵣ,
-                                                    grid.ηₗ, grid.ηᵣ,
                                                     adapt(to, grid.λᶜᶜᵃ),
                                                     adapt(to, grid.λᶠᶜᵃ),
                                                     adapt(to, grid.λᶜᶠᵃ),
@@ -762,7 +777,8 @@ function Adapt.adapt_structure(to, grid::OrthogonalSphericalShellGrid)
                                                     adapt(to, grid.Azᶠᶜᵃ),
                                                     adapt(to, grid.Azᶜᶠᵃ),
                                                     adapt(to, grid.Azᶠᶠᵃ),
-                                                    grid.radius)
+                                                    grid.radius,
+                                                    grid.conformal_mapping)
 end
 
 function Base.summary(grid::OrthogonalSphericalShellGrid)
@@ -886,7 +902,7 @@ function with_halo(new_halo, old_grid::OrthogonalSphericalShellGrid; rotation=no
 
     z = cpu_face_constructor_z(old_grid)
 
-    new_grid = OrthogonalSphericalShellGrid(architecture(old_grid), eltype(old_grid);
+    new_grid = conformal_cubed_sphere_panel(architecture(old_grid), eltype(old_grid);
                                             size, z, ξ, η,
                                             topology = topo,
                                             radius = old_grid.radius,

src/MultiRegion/multi_region_cubed_sphere_grid.jl

@@ -1,5 +1,10 @@
 using Oceananigans.Architectures: architecture
-using Oceananigans.Grids: R_Earth, halo_size, size_summary, total_length, topology
+using Oceananigans.Grids: conformal_cubed_sphere_panel,
+                          R_Earth,
+                          halo_size,
+                          size_summary,
+                          total_length,
+                          topology
 
 import Oceananigans.Grids: grid_name
 
@@ -20,8 +25,8 @@ const ConformalCubedSphereGrid{FT, TX, TY, TZ} = MultiRegionGrid{FT, TX, TY, TZ,
                              partition = CubedSpherePartition(; R = 1),
                              devices = nothing)
 
-Return a `ConformalCubedSphereGrid` that comprises of six [`OrthogonalSphericalShellGrid`](@ref);
-we refer to each of these grids as a "panel". Each panel corresponds to a face of the cube.
+Return a `ConformalCubedSphereGrid` that comprises of six [`conformal_cubed_sphere_panel`](@ref)
+grids; we refer to each of these grids as a "panel". Each panel corresponds to a face of the cube.
 
 The keyword arguments prescribe the properties of each of the panels. Only the topology in
 the vertical direction can be prescribed and that's done via the `z_topology` keyword
@@ -189,7 +194,7 @@ function ConformalCubedSphereGrid(arch::AbstractArchitecture=CPU(), FT=Float64;
     region_η = Iterate(region_η)
     region_rotation = Iterate(region_rotation)
 
-    region_grids = construct_regionally(OrthogonalSphericalShellGrid, arch, FT;
+    region_grids = construct_regionally(conformal_cubed_sphere_panel, arch, FT;
                                         size = region_size,
                                         z,
                                         halo = region_halo,
@@ -447,7 +452,7 @@ function ConformalCubedSphereGrid(filepath::AbstractString, arch::AbstractArchit
     region_Nz = MultiRegionObject(Tuple(repeat([Nz], length(partition))), devices)
     region_panels = Iterate(Array(1:length(partition)))
 
-    region_grids = construct_regionally(OrthogonalSphericalShellGrid, filepath, arch, FT;
+    region_grids = construct_regionally(conformal_cubed_sphere_panel, filepath, arch, FT;
                                         Nz = region_Nz,
                                         z,
                                         panel = region_panels,

test/test_biogeochemistry.jl

@@ -130,7 +130,7 @@ end
         arch in archs,
         grid in (RectilinearGrid(arch; size = (2, 2, 2), extent = (2, 2, 2)), 
                  LatitudeLongitudeGrid(arch; size = (5, 5, 5), longitude = (-180, 180), latitude = (-85, 85), z = (-2, 0)),
-                 OrthogonalSphericalShellGrid(size = (3, 3, 3), z = (-2, 0)))
+                 conformal_cubed_sphere_panel(arch; size = (3, 3, 3), z = (-2, 0)))
 
         if !((model == NonhydrostaticModel) && ((grid isa LatitudeLongitudeGrid) | (grid isa OrthogonalSphericalShellGrid)))
             @info "Testing $bgc in $model on $grid..."

test/test_grids.jl

@@ -701,7 +701,8 @@ end
 #####
 
 function test_orthogonal_shell_grid_array_sizes_and_spacings(FT)
-    grid = OrthogonalSphericalShellGrid(CPU(), FT, size=(10, 10, 1), z=(0, 1))
+
+    grid = conformal_cubed_sphere_panel(CPU(), FT, size=(10, 10, 1), z=(0, 1))
 
     Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
     Hx, Hy, Hz = grid.Hx, grid.Hy, grid.Hz
@@ -913,7 +914,7 @@ end
         end
 
         # Testing show function
-        grid = OrthogonalSphericalShellGrid(CPU(), size=(10, 10, 1), z=(0, 1))
+        grid = conformal_cubed_sphere_panel(CPU(), size=(10, 10, 1), z=(0, 1))
 
         @test try
             show(grid); println()
@@ -932,7 +933,7 @@ end
                 radius = 234.5e6
 
                 Nx, Ny = 10, 8
-                grid = OrthogonalSphericalShellGrid(arch, FT, size=(Nx, Ny, 1); z, radius)
+                grid = conformal_cubed_sphere_panel(arch, FT, size=(Nx, Ny, 1); z, radius)
 
                 # the sum of area metrics Azᶜᶜᵃ is 1/6-th of the area of the sphere
                 @test sum(grid.Azᶜᶜᵃ) ≈ 4π * grid.radius^2 / 6
@@ -942,16 +943,16 @@ end
                 #
                 # (for odd number of grid points, the central grid points fall on great circles)
                 Nx, Ny = 11, 9
-                grid = OrthogonalSphericalShellGrid(arch, FT, size=(Nx, Ny, 1); z, radius)
+                grid = conformal_cubed_sphere_panel(arch, FT, size=(Nx, Ny, 1); z, radius)
                 @test sum(grid.Δxᶜᶜᵃ[:, (Ny+1)÷2]) ≈ 2π * grid.radius / 4
                 @test sum(grid.Δyᶜᶜᵃ[(Nx+1)÷2, :]) ≈ 2π * grid.radius / 4
 
                 Nx, Ny = 10, 9
-                grid = OrthogonalSphericalShellGrid(arch, FT, size=(Nx, Ny, 1); z, radius)
+                grid = conformal_cubed_sphere_panel(arch, FT, size=(Nx, Ny, 1); z, radius)
                 @test sum(grid.Δxᶜᶜᵃ[:, (Ny+1)÷2]) ≈ 2π * grid.radius / 4
 
                 Nx, Ny = 11, 8
-                grid = OrthogonalSphericalShellGrid(arch, FT, size=(Nx, Ny, 1); z, radius)
+                grid = conformal_cubed_sphere_panel(arch, FT, size=(Nx, Ny, 1); z, radius)
                 @test sum(grid.Δyᶜᶜᵃ[(Nx+1)÷2, :]) ≈ 2π * grid.radius / 4
             end
         end

test/test_multi_region_cubed_sphere.jl

@@ -100,7 +100,7 @@ end
     cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2"
 
     for panel in 1:6
-        grid = OrthogonalSphericalShellGrid(cs32_filepath; panel, Nz, z)
+        grid = conformal_cubed_sphere_panel(cs32_filepath; panel, Nz, z)
         @test grid isa OrthogonalSphericalShellGrid
     end