(0.95.15) Fix bugs in ReactantExt for OrthogonalSphericalShellGrid (#4151)

* Fix bugs in ReactantExt for OrthogonalSphericalShellGrid

* Fix type info

* Add test for OSSG deconcretization

* Fix Bounded missing bug

* Thrashing around a bit

* Clean up architectures a bit

* Generalize types for OSSG and LLG plus bugfixes

* Bump to 0.95.15

* Make OSSG more friendly, fix bugs

* Mystical journey to fix broken dispatch

Author: glwagner

ext/OceananigansReactantExt/Architectures.jl

@@ -8,7 +8,9 @@ import Oceananigans.Architectures: device, architecture, array_type, on_architec
 const ReactantKernelAbstractionsExt = Base.get_extension(
     Reactant, :ReactantKernelAbstractionsExt
 )
+
 const ReactantBackend = ReactantKernelAbstractionsExt.ReactantBackend
+
 device(::ReactantState) = ReactantBackend()
 
 architecture(::Reactant.AnyConcretePJRTArray) = ReactantState
@@ -21,18 +23,14 @@ to_reactant_sharding(s::Sharding.AbstractSharding) = s
 to_reactant_sharding(::T) where {T} = error("Unsupported sharding type $T")
 
 on_architecture(::ReactantState, a::Reactant.AnyTracedRArray) = a
-function on_architecture(r::ReactantState, a::Array)
-    return Reactant.to_rarray(a; sharding=to_reactant_sharding(r.sharding))
-end
-function on_architecture(r::ReactantState, a::Reactant.AnyConcretePJRTArray)
-    return Reactant.to_rarray(a; sharding=to_reactant_sharding(r.sharding))
-end
-function on_architecture(r::ReactantState, a::BitArray)
-    return Reactant.to_rarray(a; sharding=to_reactant_sharding(r.sharding))
-end
-function on_architecture(r::ReactantState, a::SubArray{<:Any,<:Any,<:Array})
-    return Reactant.to_rarray(a; sharding=to_reactant_sharding(r.sharding))
-end
+
+const ArraysToRArray = Union{Array,
+                             Reactant.AnyConcretePJRTArray,
+                             BitArray,
+                             SubArray{<:Any, <:Any, <:Array}}
+
+on_architecture(r::ReactantState, a::ArraysToRArray) =
+    Reactant.to_rarray(a; sharding=to_reactant_sharding(r.sharding))
 
 unified_array(::ReactantState, a) = a
 

ext/OceananigansReactantExt/Grids.jl

@@ -8,7 +8,8 @@ using Oceananigans.Fields: Field
 using Oceananigans.ImmersedBoundaries: GridFittedBottom
 
 import ..OceananigansReactantExt: deconcretize
-import Oceananigans.Grids: LatitudeLongitudeGrid, RectilinearGrid
+import Oceananigans.Grids: LatitudeLongitudeGrid, RectilinearGrid, OrthogonalSphericalShellGrid
+import Oceananigans.OrthogonalSphericalShellGrids: RotatedLatitudeLongitudeGrid
 import Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 
 const ReactantGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:ReactantState}
@@ -40,29 +41,57 @@ function RectilinearGrid(arch::ReactantState, FT::DataType; kw...)
     return RectilinearGrid{TX, TY, TZ}(arch, other_properties...)
 end
 
+function OrthogonalSphericalShellGrid(arch::ReactantState, FT::DataType; kw...)
+    cpu_grid = OrthogonalSphericalShellGrid(CPU(), FT; kw...)
+    other_names = propertynames(cpu_grid)[2:end] # exclude architecture
+    other_properties = Tuple(getproperty(cpu_grid, name) for name in other_names)
+    TX, TY, TZ = Oceananigans.Grids.topology(cpu_grid)
+    return OrthogonalSphericalShellGrid{TX, TY, TZ}(arch, other_properties...)
+end
+
+# This is a kind of OrthogonalSphericalShellGrid
+function RotatedLatitudeLongitudeGrid(arch::ReactantState, FT::DataType; kw...)
+    cpu_grid = RotatedLatitudeLongitudeGrid(CPU(), FT; kw...)
+    other_names = propertynames(cpu_grid)[2:end] # exclude architecture
+    other_properties = Tuple(getproperty(cpu_grid, name) for name in other_names)
+    TX, TY, TZ = Oceananigans.Grids.topology(cpu_grid)
+    return OrthogonalSphericalShellGrid{TX, TY, TZ}(arch, other_properties...)
+end
+
+# This low-level constructor supports the external package OrthogonalSphericalShellGrids.jl.
 function OrthogonalSphericalShellGrid{TX, TY, TZ}(arch::ReactantState,
                                                   Nx, Ny, Nz, Hx, Hy, Hz,
                                                   Lz :: FT,
-                                                  λᶜᶜᵃ :: A, λᶠᶜᵃ :: A, λᶜᶠᵃ :: A, λᶠᶠᵃ :: A,
-                                                  φᶜᶜᵃ :: A, φᶠᶜᵃ :: A, φᶜᶠᵃ :: A, φᶠᶠᵃ :: A,
-                                                  z :: Z,
-                                                  Δxᶜᶜᵃ :: A, Δxᶠᶜᵃ :: A, Δxᶜᶠᵃ :: A, Δxᶠᶠᵃ :: A,
-                                                  Δyᶜᶜᵃ :: A, Δyᶜᶠᵃ :: A, Δyᶠᶜᵃ :: A, Δyᶠᶠᵃ :: A, 
-                                                  Azᶜᶜᵃ :: A, Azᶠᶜᵃ :: A, Azᶜᶠᵃ :: A, Azᶠᶠᵃ :: A,
+                                                   λᶜᶜᵃ :: CC,  λᶠᶜᵃ :: FC,  λᶜᶠᵃ :: CF,  λᶠᶠᵃ :: FF,
+                                                   φᶜᶜᵃ :: CC,  φᶠᶜᵃ :: FC,  φᶜᶠᵃ :: CF,  φᶠᶠᵃ :: FF, z :: Z,
+                                                  Δxᶜᶜᵃ :: CC, Δxᶠᶜᵃ :: FC, Δxᶜᶠᵃ :: CF, Δxᶠᶠᵃ :: FF,
+                                                  Δyᶜᶜᵃ :: CC, Δyᶠᶜᵃ :: FC, Δyᶜᶠᵃ :: CF, Δyᶠᶠᵃ :: FF, 
+                                                  Azᶜᶜᵃ :: CC, Azᶠᶜᵃ :: FC, Azᶜᶠᵃ :: CF, Azᶠᶠᵃ :: FF,
                                                   radius :: FT,
-                                                  conformal_mapping :: C) where {TX, TY, TZ, FT, Z, A, C}
-
-    args = (λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ,
-            φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ,
-            z,
-            Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ,
-            Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ,
-            Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ)
-
-    dargs = Tuple(deconcretize(a) for a in args)
-
-    return OrthogonalSphericalShellGrid{FT, TX, TY, TZ, CZ, A, C, Arch}(arch, Nx, Ny, Nz, Hx, Hy, Hz, Lz,
-                                                                        dargs..., radius, conformal_mapping)
+                                                  conformal_mapping :: Map) where {TX, TY, TZ, FT, Z, Map,
+                                                                                   CC, FC, CF, FF, C}
+
+    args1 = (λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ,
+             φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ)
+
+    args2 = (Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ,
+             Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ,
+             Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ)
+
+    dargs1 = Tuple(deconcretize(a) for a in args1)
+    dz = deconcretize(z)
+    dargs2 = Tuple(deconcretize(a) for a in args2)
+
+    Arch = typeof(arch)
+    DCC = typeof(dargs1[1]) # deconcretized
+    DFC = typeof(dargs1[2]) # deconcretized
+    DCF = typeof(dargs1[3]) # deconcretized
+    DFF = typeof(dargs1[4]) # deconcretized
+    DZ = typeof(dz) # deconcretized
+
+    return OrthogonalSphericalShellGrid{FT, TX, TY, TZ, DZ, Map,
+                                        DCC, DFC, DCF, DFF, Arch}(arch, Nx, Ny, Nz, Hx, Hy, Hz, Lz,
+                                                                  dargs1..., dz, dargs2..., radius, conformal_mapping)
 end
 
 deconcretize(gfb::GridFittedBottom) = GridFittedBottom(deconcretize(gfb.bottom_height),
@@ -89,8 +118,6 @@ ImmersedBoundaryGrid(grid::ReactantGrid, ib::GridFittedBottom; active_cells_map:
 
 ImmersedBoundaryGrid(grid::ReactantGrid, ib; active_cells_map::Bool=true) =
     reactant_immersed_boundary_grid(grid, ib; active_cells_map)
-    
-
 
 end # module
 

src/Grids/conformal_cubed_sphere_panel.jl

@@ -99,7 +99,7 @@ function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = C
                                                      φᶜᶜᵃ,  φᶠᶜᵃ,  φᶜᶠᵃ,  φᶠᶠᵃ,
                                                      z,
                                                     Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ,
-                                                    Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ,
+                                                    Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ,
                                                     Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ,
                                                     radius,
                                                     conformal_mapping)
@@ -213,23 +213,17 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU()
                                       halo = (1, 1, 1),
                                       rotation = nothing)
 
-    if architecture == GPU() && !has_cuda()
-        throw(ArgumentError("Cannot create a GPU grid. No CUDA-enabled GPU was detected!"))
-    end
-
     radius = FT(radius)
-
     TX, TY, TZ = topology
     Nξ, Nη, Nz = size
     Hx, Hy, Hz = halo
 
-    ## Use a regular rectilinear grid for the face of the cube
-
+    # Use a regular rectilinear grid for the face of the cube
     ξη_grid_topology = (Bounded, Bounded, topology[3])
 
     # construct the grid on CPU and convert to architecture later...
     ξη_grid = RectilinearGrid(CPU(), FT;
-                              size=(Nξ, Nη, Nz),
+                              size = (Nξ, Nη, Nz),
                               topology = ξη_grid_topology,
                               x=ξ, y=η, z, halo)
 
@@ -511,7 +505,6 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU()
             Azᶜᶠᵃ[i, j] = 2 * spherical_area_quadrilateral(a, b, c, d) * radius^2
         end
 
-
         # Azᶠᶠᵃ
 
         for j in 2:Nη, i in 2:Nξ
@@ -640,7 +633,7 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU()
                          zc)
 
     metric_arrays = (Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ,
-                     Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ,
+                     Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ,
                      Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ)
 
     conformal_mapping = CubedSphereConformalMapping(ξ, η, rotation)
@@ -660,7 +653,7 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU()
                          grid.z)
 
     metric_arrays = (grid.Δxᶜᶜᵃ, grid.Δxᶠᶜᵃ, grid.Δxᶜᶠᵃ, grid.Δxᶠᶠᵃ,
-                     grid.Δyᶜᶜᵃ, grid.Δyᶜᶠᵃ, grid.Δyᶠᶜᵃ, grid.Δyᶠᶠᵃ,
+                     grid.Δyᶜᶜᵃ, grid.Δyᶠᶜᵃ, grid.Δyᶜᶠᵃ, grid.Δyᶠᶠᵃ,
                      grid.Azᶜᶜᵃ, grid.Azᶠᶜᵃ, grid.Azᶜᶠᵃ, grid.Azᶠᶠᵃ)
 
     coordinate_arrays = map(a -> on_architecture(architecture, a), coordinate_arrays)

src/Grids/grid_generation.jl

@@ -78,7 +78,7 @@ function generate_coordinate(FT, topo::AT, N, H, node_generator, coordinate_name
     end
 
     Δᶜ = OffsetArray(on_architecture(arch, Δᶜ), -H)
-    Δᶠ = OffsetArray(on_architecture(arch, Δᶠ), -H-1)
+    Δᶠ = OffsetArray(on_architecture(arch, Δᶠ), -H - 1)
 
     F = OffsetArray(F, -H)
     C = OffsetArray(C, -H)