Flat VerticallyStretchedRectilinearGrid plus new hydrostatic and flat…

… internal wave tests (#1865) * Expands internal wave test for hydrostatic model and flat topologies * Generalize Flat linear spacings to all rectilinear grids and fills in all linear metrics * Use advection=nothing in DistributedShallowWater test * Bump to 0.59.1 * Dispatch on different rectilinear grids specifically for flat vertical metrics * More verbose grid names + background_fields test * Tests FourierTridiagonalSolver with flat dimensions * Update fourier_tridiagonal_poisson_solver.jl * Update test_dynamics.jl * Update test_dynamics.jl * Refactor transform planning for VerticallyBoundedRectilinearGrid * Update src/Operators/spacings_and_areas_and_volumes.jl Co-authored-by: Navid C. Constantinou <[email protected]> * Instantiate unflattened topologies * a=b is better than if b, a = true, else a = false * Refactors transform planning for FourierTridiagonalSolver * Cosmetic improvements to test_poisson_solvers Co-authored-by: Navid C. Constantinou <[email protected]>
CliMA · Jul 18, 2021 · 4a0a87e · 4a0a87e
1 parent 3d9b87c
commit 4a0a87e
Show file tree

Hide file tree

Showing 10 changed files with 269 additions and 190 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.59.0"
+version = "0.59.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/Operators/spacings_and_areas_and_volumes.jl b/src/Operators/spacings_and_areas_and_volumes.jl
@@ -47,19 +47,57 @@ The operators in this file fall into three categories:
 @inline Δzᵃᵃᶜ(i, j, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.Δzᵃᵃᶜ[k]
 
 #####
-##### "Spacings" in Flat directions. Here we dispatch to `one`. This abuse of notation
-##### makes volumes correct as we want to multiply by 1, and avoids issues with
-##### derivatives such as those involved in the pressure correction step.
+##### "Spacings" in Flat directions for rectilinear grids.
+##### Here we dispatch all spacings to `one`. This abuse of notation
+##### makes volumes and areas correct. E.g., with `z` Flat we have `v = dx * dy * 1`.
+#####
+##### Note: Vertical metrics are specific to each rectilinear grid type; therefore
+##### we must dispatch on Flat vertical directions for each grid independently.
 #####
 
 using Oceananigans.Grids: Flat
 
-@inline Δx(   i, j, k, grid::RegularRectilinearGrid{FT, Flat})         where FT           = one(FT)
-@inline Δy(   i, j, k, grid::RegularRectilinearGrid{FT, TX, Flat})     where {FT, TX}     = one(FT)
-@inline ΔzC(  i, j, k, grid::RegularRectilinearGrid{FT, TX, TY, Flat}) where {FT, TX, TY} = one(FT)
-@inline ΔzF(  i, j, k, grid::RegularRectilinearGrid{FT, TX, TY, Flat}) where {FT, TX, TY} = one(FT)
-@inline Δzᵃᵃᶠ(i, j, k, grid::RegularRectilinearGrid{FT, TX, TY, Flat}) where {FT, TX, TY} = one(FT)
-@inline Δzᵃᵃᶜ(i, j, k, grid::RegularRectilinearGrid{FT, TX, TY, Flat}) where {FT, TX, TY} = one(FT)
+#####
+##### Horizontal metrics for AbstractRectilinearGrid
+#####
+
+const XFlatARG = AbstractRectilinearGrid{<:Any, <:Flat}
+const YFlatARG = AbstractRectilinearGrid{<:Any, <:Any, <:Flat}
+
+@inline Δx(i, j, k, grid::XFlatARG) = one(eltype(grid))
+@inline Δy(i, j, k, grid::YFlatARG) = one(eltype(grid))
+
+@inline Δxᶜᶜᵃ(i, j, k, grid::XFlatARG) = one(eltype(grid))
+@inline Δxᶜᶠᵃ(i, j, k, grid::XFlatARG) = one(eltype(grid))
+@inline Δxᶠᶠᵃ(i, j, k, grid::XFlatARG) = one(eltype(grid))
+@inline Δxᶠᶜᵃ(i, j, k, grid::XFlatARG) = one(eltype(grid))
+
+@inline Δyᶜᶜᵃ(i, j, k, grid::YFlatARG) = one(eltype(grid))
+@inline Δyᶠᶜᵃ(i, j, k, grid::YFlatARG) = one(eltype(grid))
+@inline Δyᶜᶠᵃ(i, j, k, grid::YFlatARG) = one(eltype(grid))
+@inline Δyᶠᶠᵃ(i, j, k, grid::YFlatARG) = one(eltype(grid))
+
+##### 
+##### Vertical metrics for RegularRectilinearGrid
+##### 
+
+const ZFlatRRG = RegularRectilinearGrid{<:Any, <:Any, <:Any, <:Flat}
+
+@inline ΔzC(  i, j, k, grid::ZFlatRRG) = one(eltype(grid))
+@inline ΔzF(  i, j, k, grid::ZFlatRRG) = one(eltype(grid))
+@inline Δzᵃᵃᶠ(i, j, k, grid::ZFlatRRG) = one(eltype(grid))
+@inline Δzᵃᵃᶜ(i, j, k, grid::ZFlatRRG) = one(eltype(grid))
+
+##### 
+##### Vertical metrics for VerticallyStretchedRectilinearGrid
+##### 
+
+const ZFlatVSRG = VerticallyStretchedRectilinearGrid{<:Any, <:Any, <:Any, <:Flat}
+
+@inline ΔzC(  i, j, k, grid::ZFlatVSRG) = one(eltype(grid))
+@inline ΔzF(  i, j, k, grid::ZFlatVSRG) = one(eltype(grid))
+@inline Δzᵃᵃᶠ(i, j, k, grid::ZFlatVSRG) = one(eltype(grid))
+@inline Δzᵃᵃᶜ(i, j, k, grid::ZFlatVSRG) = one(eltype(grid))
 
 #####
 ##### Areas for horizontally-regular algorithms

diff --git a/src/Solvers/discrete_transforms.jl b/src/Solvers/discrete_transforms.jl
@@ -27,7 +27,7 @@ normalization_factor(arch, topo, direction, N) = 1
 `FFTW.REDFT01` needs to be normalized by 1/2N.
 See: http://www.fftw.org/fftw3_doc/1d-Real_002deven-DFTs-_0028DCTs_0029.html#g_t1d-Real_002deven-DFTs-_0028DCTs_0029
 """
-normalization_factor(::CPU, ::Bounded, ::Backward, N) = 1/(2N)
+normalization_factor(::CPU, ::Bounded, ::Backward, N) = 1 / 2N
 
 #####
 ##### Twiddle factors

diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl
@@ -25,7 +25,7 @@ function FFTBasedPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT)
     transforms = plan_transforms(arch, grid, storage, planner_flag)
 
     # Need buffer for index permutations and transposes.
-    buffer_needed = arch isa GPU && Bounded in topo ? true : false
+    buffer_needed = arch isa GPU && Bounded in topo
     buffer = buffer_needed ? similar(storage) : nothing
 
     return FFTBasedPoissonSolver(arch, grid, eigenvalues, storage, buffer, transforms)

diff --git a/src/Solvers/fourier_tridiagonal_poisson_solver.jl b/src/Solvers/fourier_tridiagonal_poisson_solver.jl
@@ -1,4 +1,5 @@
 using Oceananigans.Operators: Δzᵃᵃᶜ, Δzᵃᵃᶠ
+using Oceananigans.Architectures: device_event
 
 struct FourierTridiagonalPoissonSolver{A, G, B, R, S, β, T}
                   architecture :: A
@@ -10,7 +11,7 @@ struct FourierTridiagonalPoissonSolver{A, G, B, R, S, β, T}
                     transforms :: T
 end
 
-@kernel function compute_diagonals!(D, grid, λx, λy)
+@kernel function compute_main_diagonals!(D, grid, λx, λy)
     i, j = @index(Global, NTuple)
     Nz = grid.Nz
 
@@ -28,11 +29,6 @@ function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT)
     TX, TY, TZ = topology(grid)
     TZ != Bounded && error("FourierTridiagonalPoissonSolver can only be used with a Bounded z topology.")
 
-    if grid isa VerticallyStretchedRectilinearGrid && any([T() isa Flat for T in (TX, TY)])
-        @warn "FourierTridiagonalPoissonSolver is probably wrong for topologies that contain " *
-              "Flat dimensions."
-    end
-
     Nx, Ny, Nz = size(grid)
 
     # Compute discrete Poisson eigenvalues
@@ -47,16 +43,13 @@ function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT)
     transforms = plan_transforms(arch, grid, sol_storage, planner_flag)
 
     # Lower and upper diagonals are the same
-    CUDA.allowscalar(true)
     lower_diagonal = CUDA.@allowscalar [1 / Δzᵃᵃᶠ(1, 1, k, grid) for k in 2:Nz]
     lower_diagonal = arch_array(arch, lower_diagonal)
     upper_diagonal = lower_diagonal
-    CUDA.allowscalar(false)
 
     # Compute diagonal coefficients for each grid point
     diagonal = arch_array(arch, zeros(Nx, Ny, Nz))
-    event = launch!(arch, grid, :xy, compute_diagonals!, diagonal, grid, λx, λy,
-                    dependencies=Event(device(arch)))
+    event = launch!(arch, grid, :xy, compute_main_diagonals!, diagonal, grid, λx, λy, dependencies=device_event(arch))
     wait(device(arch), event)
 
     # Set up batched tridiagonal solver
@@ -66,7 +59,7 @@ function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT)
                                          upper_diagonal = upper_diagonal)
 
     # Need buffer for index permutations and transposes.
-    buffer_needed = arch isa GPU && Bounded in (TX, TY) ? true : false
+    buffer_needed = arch isa GPU && Bounded in (TX, TY)
     buffer = buffer_needed ? similar(sol_storage) : nothing
 
     # Storage space for right hand side of Poisson equation

diff --git a/src/Solvers/plan_transforms.jl b/src/Solvers/plan_transforms.jl
@@ -46,9 +46,7 @@ plan_forward_transform(A::Union{Array, CuArray}, ::Flat, args...) = nothing
 
 batchable_GPU_topologies = ((Periodic, Periodic, Periodic),
                             (Periodic, Periodic, Bounded),
-                            (Bounded, Periodic, Periodic),
-                            (Periodic, Periodic, Flat),
-                            (Flat, Periodic, Periodic))
+                            (Bounded,  Periodic, Periodic))
 
 # In principle the order in which the transforms are applied does not matter of course,
 # but in practice we want to perform the `Bounded` forward transforms first because on
@@ -81,23 +79,24 @@ function plan_transforms(arch, grid::RegularRectilinearGrid, storage, planner_fl
     periodic_dims = findall(t -> t == Periodic, topo)
     bounded_dims = findall(t -> t == Bounded, topo)
 
-    if arch isa GPU && !(topo in batchable_GPU_topologies)
+    # Convert Flat to Bounded for inferring batchability and transform ordering
+    # Note that transforms are omitted in Flat directions.
+    unflattened_topo = Tuple(T() isa Flat ? Bounded : T for T in topo)
+
+    if arch isa GPU && !(unflattened_topo in batchable_GPU_topologies)
 
         rs_storage = reshape(storage, (Ny, Nx, Nz))
-        forward_plan_x = plan_forward_transform(storage   , topo[1](), [1], planner_flag)
+        forward_plan_x = plan_forward_transform(storage,    topo[1](), [1], planner_flag)
         forward_plan_y = plan_forward_transform(rs_storage, topo[2](), [1], planner_flag)
-        forward_plan_z = plan_forward_transform(storage   , topo[3](), [3], planner_flag)
+        forward_plan_z = plan_forward_transform(storage,    topo[3](), [3], planner_flag)
 
-        backward_plan_x = plan_backward_transform(storage   , topo[1](), [1], planner_flag)
+        backward_plan_x = plan_backward_transform(storage,    topo[1](), [1], planner_flag)
         backward_plan_y = plan_backward_transform(rs_storage, topo[2](), [1], planner_flag)
-        backward_plan_z = plan_backward_transform(storage   , topo[3](), [3], planner_flag)
+        backward_plan_z = plan_backward_transform(storage,    topo[3](), [3], planner_flag)
 
         forward_plans = (forward_plan_x, forward_plan_y, forward_plan_z)
         backward_plans = (backward_plan_x, backward_plan_y, backward_plan_z)
 
-        # Convert Flat to Bounded for ordering purposes (transforms are omitted in Flat directions anyways)
-        unflattened_topo = (T() isa Flat ? Bounded : T for T in topo)
-
         f_order = forward_orders(unflattened_topo...)
         b_order = backward_orders(unflattened_topo...)
 
@@ -117,9 +116,8 @@ function plan_transforms(arch, grid::RegularRectilinearGrid, storage, planner_fl
         # This is the case where batching transforms is possible. It's always possible on the CPU
         # since FFTW is awesome so it includes all topologies on the CPU.
         #
-        # On the GPU batching is possible when the topology is not one of non_batched_topologies
-        # (where an FFT is needed along dimension 2), so it includes (Periodic, Periodic, Periodic),
-        # (Periodic, Periodic, Bounded), and (Bounded, Periodic, Periodic).
+        # `batchable_GPU_topologies` occurs when there are two adjacent `Periodic` dimensions:
+        # (Periodic, Periodic, Periodic), (Periodic, Periodic, Bounded), and (Bounded, Periodic, Periodic).
 
         forward_periodic_plan = plan_forward_transform(storage, Periodic(), periodic_dims, planner_flag)
         forward_bounded_plan = plan_forward_transform(storage, Bounded(), bounded_dims, planner_flag)
@@ -138,12 +136,12 @@ function plan_transforms(arch, grid::RegularRectilinearGrid, storage, planner_fl
         )
     end
 
-    transforms = (forward = forward_transforms, backward = backward_transforms)
+    transforms = (forward=forward_transforms, backward=backward_transforms)
 
     return transforms
 end
 
-" Used by FourierTridiagonalPoissonSolver "
+""" Used by FourierTridiagonalPoissonSolver. """
 function plan_transforms(arch, grid::VerticallyStretchedRectilinearGrid, storage, planner_flag)
     Nx, Ny, Nz = size(grid)
     TX, TY, TZ = topo = topology(grid)
@@ -152,84 +150,66 @@ function plan_transforms(arch, grid::VerticallyStretchedRectilinearGrid, storage
     periodic_dims = findall(t -> t == Periodic, (TX, TY))
     bounded_dims = findall(t -> t == Bounded, (TX, TY))
 
-    if arch isa GPU && !(topo in batchable_GPU_topologies)
-        if (TX, TY) == (Periodic, Bounded)
-            forward_plan_x = plan_forward_transform(storage, Periodic(), [1], planner_flag)
-            forward_plan_y = plan_forward_transform(reshape(storage, (Ny, Nx, Nz)), Bounded(), [1], planner_flag)
-
-            backward_plan_x = plan_backward_transform(storage, Periodic(), [1], planner_flag)
-            backward_plan_y = plan_backward_transform(reshape(storage, (Ny, Nx, Nz)), Bounded(),  [1], planner_flag)
-
-            forward_transforms = (
-                DiscreteTransform(forward_plan_y, Forward(), arch, grid, [2]),
-                DiscreteTransform(forward_plan_x, Forward(), arch, grid, [1])
-            )
-
-            backward_transforms = (
-                DiscreteTransform(backward_plan_x, Backward(), arch, grid, [1]),
-                DiscreteTransform(backward_plan_y, Backward(), arch, grid, [2])
-            )
-
-        elseif (TX, TY) == (Bounded, Periodic)
-            forward_plan_x = plan_forward_transform(storage, Bounded(), [1], planner_flag)
-            forward_plan_y = plan_forward_transform(reshape(storage, (Ny, Nx, Nz)), Periodic(), [1], planner_flag)
-
-            backward_plan_x = plan_backward_transform(storage, Bounded(), [1], planner_flag)
-            backward_plan_y = plan_backward_transform(reshape(storage, (Ny, Nx, Nz)), Periodic(),  [1], planner_flag)
-
-            forward_transforms = (
-                DiscreteTransform(forward_plan_x, Forward(), arch, grid, [1]),
-                DiscreteTransform(forward_plan_y, Forward(), arch, grid, [2])
-            )
-
-            backward_transforms = (
-                DiscreteTransform(backward_plan_y, Backward(), arch, grid, [2]),
-                DiscreteTransform(backward_plan_x, Backward(), arch, grid, [1])
-            )
-
-        elseif (TX, TY) == (Bounded, Bounded)
-            forward_plan_x = plan_forward_transform(storage, Bounded(), [1], planner_flag)
-            forward_plan_y = plan_forward_transform(reshape(storage, (Ny, Nx, Nz)), Bounded(), [1], planner_flag)
-
-            backward_plan_x = plan_backward_transform(storage, Bounded(), [1], planner_flag)
-            backward_plan_y = plan_backward_transform(reshape(storage, (Ny, Nx, Nz)), Bounded(),  [1], planner_flag)
-
-            forward_transforms = (
-                DiscreteTransform(forward_plan_x, Forward(), arch, grid, [1]),
-                DiscreteTransform(forward_plan_y, Forward(), arch, grid, [2])
-            )
-
-            backward_transforms = (
-                DiscreteTransform(backward_plan_x, Backward(), arch, grid, [1]),
-                DiscreteTransform(backward_plan_y, Backward(), arch, grid, [2])
-            )
-        end
-    else
+    # Convert Flat to Bounded for ordering purposes (transforms are omitted in Flat directions anyways)
+
+    !(topo[3] === Bounded) && error("Cannot plan transforms on z-periodic VerticallyStretchedRectilinearGrids.")
+
+    if arch isa CPU
         # This is the case where batching transforms is possible. It's always possible on the CPU
         # since FFTW is awesome so it includes all topologies on the CPU.
         #
-        # On the GPU batching is possible when the topology is not one of non_batched_topologies
-        # (where an FFT is needed along dimension 2), so it includes (Periodic, Periodic, Periodic),
-        # (Periodic, Periodic, Bounded), and (Bounded, Periodic, Periodic).
-
+        # On the GPU and for vertically Bounded grids, batching is possible either in horizontally-periodic
+        # domains, or for domains that are `Bounded, Periodic, Bounded`.
+
         forward_periodic_plan = plan_forward_transform(storage, Periodic(), periodic_dims, planner_flag)
         forward_bounded_plan = plan_forward_transform(storage, Bounded(), bounded_dims, planner_flag)
 
-        forward_transforms = (
-            DiscreteTransform(forward_bounded_plan, Forward(), arch, grid, bounded_dims),
-            DiscreteTransform(forward_periodic_plan, Forward(), arch, grid, periodic_dims)
-        )
+        forward_transforms = (DiscreteTransform(forward_bounded_plan, Forward(), arch, grid, bounded_dims),
+                              DiscreteTransform(forward_periodic_plan, Forward(), arch, grid, periodic_dims))
 
         backward_periodic_plan = plan_backward_transform(storage, Periodic(), periodic_dims, planner_flag)
         backward_bounded_plan = plan_backward_transform(storage, Bounded(), bounded_dims, planner_flag)
 
-        backward_transforms = (
-            DiscreteTransform(backward_periodic_plan, Backward(), arch, grid, periodic_dims),
-            DiscreteTransform(backward_bounded_plan, Backward(), arch, grid, bounded_dims)
-        )
+        backward_transforms = (DiscreteTransform(backward_periodic_plan, Backward(), arch, grid, periodic_dims),
+                               DiscreteTransform(backward_bounded_plan, Backward(), arch, grid, bounded_dims))
+
+    elseif !(Bounded in (TX, TY))
+        # We're on the GPU and either (Periodic, Periodic), (Flat, Periodic), or (Periodic, Flat) in xy.
+        # So, we pretend like we need a 2D doubly-periodic transform (even if one dimension is Flat).
+        forward_periodic_plan = plan_forward_transform(storage, Periodic(), [1, 2], planner_flag)
+        backward_periodic_plan = plan_backward_transform(storage, Periodic(), [1, 2], planner_flag)
+
+        forward_transforms = tuple(DiscreteTransform(forward_periodic_plan, Forward(), arch, grid, [1, 2]))
+        backward_transforms = tuple(DiscreteTransform(backward_periodic_plan, Backward(), arch, grid, [1, 2]))
+
+    else # we are on the GPU and we cannot / should not batch!
+        rs_storage = reshape(storage, (Ny, Nx, Nz))
+
+        forward_plan_x = plan_forward_transform(storage,    topo[1](), [1], planner_flag)
+        forward_plan_y = plan_forward_transform(rs_storage, topo[2](), [1], planner_flag)
+
+        backward_plan_x = plan_backward_transform(storage,    topo[1](), [1], planner_flag)
+        backward_plan_y = plan_backward_transform(rs_storage, topo[2](), [1], planner_flag)
+
+        forward_plans = (forward_plan_x, forward_plan_y)
+        backward_plans = (backward_plan_x, backward_plan_y)
+
+        unflattened_topo = Tuple(T() isa Flat ? Bounded : T for T in topo)
+        f_order = forward_orders(unflattened_topo...)
+        b_order = backward_orders(unflattened_topo...)
+
+        # Extract third dimension
+        f_order = Tuple(f_order[i] for i in findall(d -> d != 3, f_order))
+        b_order = Tuple(b_order[i] for i in findall(d -> d != 3, b_order))
+
+        forward_transforms = (DiscreteTransform(forward_plans[f_order[1]], Forward(), arch, grid, [f_order[1]]),
+                              DiscreteTransform(forward_plans[f_order[2]], Forward(), arch, grid, [f_order[2]]))
+
+        backward_transforms = (DiscreteTransform(backward_plans[b_order[1]], Backward(), arch, grid, [b_order[1]]),
+                               DiscreteTransform(backward_plans[b_order[2]], Backward(), arch, grid, [b_order[2]]))
     end
 
-    transforms = (forward = forward_transforms, backward = backward_transforms)
+    transforms = (forward=forward_transforms, backward=backward_transforms)
 
     return transforms
 end
diff --git a/test/test_distributed_models.jl b/test/test_distributed_models.jl
@@ -496,9 +496,9 @@ end
 
     @testset "Time stepping ShallowWaterModel" begin
         topo = (Periodic, Periodic, Flat)
-        full_grid = RegularRectilinearGrid(topology=topo, size=(8, 8), extent=(1, 2), halo=(3,3))
+        full_grid = RegularRectilinearGrid(topology=topo, size=(8, 8), extent=(1, 2), halo=(3, 3))
         arch = MultiCPU(grid=full_grid, ranks=(1, 4, 1))
-        model = DistributedShallowWaterModel(architecture=arch, grid=full_grid, gravitational_acceleration=1)
+        model = DistributedShallowWaterModel(architecture=arch, advection=nothing, grid=full_grid, gravitational_acceleration=1)
 
         set!(model, h=1)
         time_step!(model, 1)