Pressure solver for vertically stretched grids (#306)

Pressure solver for vertically stretched grids

src/Solvers/batched_tridiagonal_solver.jl

@@ -34,9 +34,9 @@ where `dl` stores the lower diagonal coefficients `a(i, j, k)`, `d` stores the d
 
 `params` is an optional named tuple of parameters that is accessible to functions.
 """
-function BatchedTridiagonalSolver(arch=CPU(); dl, d, du, f, grid, params=nothing)
+function BatchedTridiagonalSolver(arch=CPU(), FT=Float64; dl, d, du, f, grid, params=nothing)
     ArrayType = array_type(arch)
-    t = zeros(grid.Nx, grid.Ny, grid.Nz) |> ArrayType
+    t = zeros(FT, grid.Nx, grid.Ny, grid.Nz) |> ArrayType
 
     return BatchedTridiagonalSolver(dl, d, du, f, t, grid, params)
 end

test/test_pressure_solvers.jl

@@ -10,10 +10,10 @@ function ∇²!(grid, f, ∇²f)
     end
 end
 
-function fftw_planner_works(FT, Nx, Ny, Nz, planner_flag)
+function pressure_solver_instantiates(FT, Nx, Ny, Nz, planner_flag)
     grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(100, 100, 100))
-    solver = PressureSolver(CPU(), grid, HorizontallyPeriodicBCs())
-    true  # Just making sure the PressureSolver does not error/crash.
+    solver = PressureSolver(CPU(), grid, HorizontallyPeriodicBCs(), planner_flag)
+    return true  # Just making sure the PressureSolver does not error/crash.
 end
 
 function poisson_ppn_planned_div_free_cpu(FT, Nx, Ny, Nz, planner_flag)
@@ -187,16 +187,16 @@ end
 @testset "Pressure solvers" begin
     @info "Testing pressure solvers..."
 
-    # @testset "FFTW plans" begin
-    #     @info "  Testing FFTW planning..."
-    #
-    #     for FT in float_types
-    #         @test fftw_planner_works(FT, 32, 32, 32, FFTW.ESTIMATE)
-    #         @test fftw_planner_works(FT, 1,  32, 32, FFTW.ESTIMATE)
-    #         @test fftw_planner_works(FT, 32,  1, 32, FFTW.ESTIMATE)
-    #         @test fftw_planner_works(FT,  1,  1, 32, FFTW.ESTIMATE)
-    #     end
-    # end
+    @testset "Pressure solver instantiation" begin
+        @info "  Testing pressure solver instantiation..."
+
+        for FT in float_types
+            @test pressure_solver_instantiates(FT, 32, 32, 32, FFTW.ESTIMATE)
+            @test pressure_solver_instantiates(FT, 1,  32, 32, FFTW.ESTIMATE)
+            @test pressure_solver_instantiates(FT, 32,  1, 32, FFTW.ESTIMATE)
+            @test pressure_solver_instantiates(FT,  1,  1, 32, FFTW.ESTIMATE)
+        end
+    end
 
     @testset "Divergence-free solution [CPU]" begin
         @info "  Testing divergence-free solution [CPU]..."

test/test_solvers.jl

@@ -1,5 +1,7 @@
 using LinearAlgebra
-using Oceananigans: array_type
+using Oceananigans.Architectures: array_type
+using Oceananigans.BoundaryConditions: NoPenetrationBC
+using Oceananigans.TimeSteppers: _compute_w_from_continuity!
 
 function can_solve_single_tridiagonal_system(arch, N)
     ArrayType = array_type(arch)
@@ -118,6 +120,153 @@ function can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz
     return Array(ϕ) ≈ ϕ_correct
 end
 
+function vertically_stretched_poisson_solver_correct_answer(arch, Nx, Ny, zF)
+    Lx, Ly, Lz = 1, 1, zF[end]
+    Δx, Δy = Lx/Nx, Ly/Ny
+
+    #####
+    ##### Vertically stretched operators
+    #####
+
+    @inline δx_caa(i, j, k, f) = @inbounds f[i+1, j, k] - f[i, j, k]
+    @inline δy_aca(i, j, k, f) = @inbounds f[i, j+1, k] - f[i, j, k]
+    @inline δz_aac(i, j, k, f) = @inbounds f[i, j, k+1] - f[i, j, k]
+
+    @inline ∂x_caa(i, j, k, Δx,  f) = δx_caa(i, j, k, f) / Δx
+    @inline ∂y_aca(i, j, k, Δy,  f) = δy_aca(i, j, k, f) / Δy
+    @inline ∂z_aac(i, j, k, ΔzF, f) = δz_aac(i, j, k, f) / ΔzF[k]
+
+    @inline ∂x²(i, j, k, Δx, f)       = (∂x_caa(i, j, k, Δx, f)  - ∂x_caa(i-1, j, k, Δx, f))  / Δx
+    @inline ∂y²(i, j, k, Δy, f)       = (∂y_aca(i, j, k, Δy, f)  - ∂y_aca(i, j-1, k, Δy, f))  / Δy
+    @inline ∂z²(i, j, k, ΔzF, ΔzC, f) = (∂z_aac(i, j, k, ΔzF, f) - ∂z_aac(i, j, k-1, ΔzF, f)) / ΔzC[k]
+
+    @inline div_ccc(i, j, k, Δx, Δy, ΔzF, u, v, w) = ∂x_caa(i, j, k, Δx, u) + ∂y_aca(i, j, k, Δy, v) + ∂z_aac(i, j, k, ΔzF, w)
+
+    @inline ∇²(i, j, k, Δx, Δy, ΔzF, ΔzC, f) = ∂x²(i, j, k, Δx, f) + ∂y²(i, j, k, Δy, f) + ∂z²(i, j, k, ΔzF, ΔzC, f)
+
+    #####
+    ##### Generate "fake" vertically stretched grid
+    #####
+
+    function grid(zF)
+        Nz = length(zF) - 1
+        ΔzF = [zF[k+1] - zF[k] for k in 1:Nz]
+        zC = [(zF[k] + zF[k+1]) / 2 for k in 1:Nz]
+        ΔzC = [zC[k+1] - zC[k] for k in 1:Nz-1]
+        return zF, zC, ΔzF, ΔzC
+    end
+
+    Nz = length(zF) - 1
+    zF, zC, ΔzF, ΔzC = grid(zF)
+
+    # Need some halo regions.
+    ΔzF = OffsetArray([ΔzF[1], ΔzF...], 0:Nz)
+    ΔzC = [ΔzC..., ΔzC[end]]
+
+    # Useful for broadcasting z operations
+    ΔzC = reshape(ΔzC, (1, 1, Nz))
+
+    # Temporary hack: Useful for reusing fill_halo_regions! and BatchedTridiagonalSolver
+    # which only need Nx, Ny, Nz.
+    fake_grid = RegularCartesianGrid(size=(Nx, Ny, Nz), length=(Lx, Ly, Lz))
+
+    #####
+    ##### Generate batched tridiagonal system coefficients and solver
+    #####
+
+    function λi(Nx, Δx)
+        is = reshape(1:Nx, Nx, 1, 1)
+        @. (2sin((is-1)*π/Nx) / Δx)^2
+    end
+
+    function λj(Ny, Δy)
+        js = reshape(1:Ny, 1, Ny, 1)
+        @. (2sin((js-1)*π/Ny) / Δy)^2
+    end
+
+    kx² = λi(Nx, Δx)
+    ky² = λj(Ny, Δy)
+
+    # Lower and upper diagonals are the same
+    ld = [1/ΔzF[k] for k in 1:Nz-1]
+    ud = copy(ld)
+
+    # Diagonal (different for each i,j)
+    @inline δ(k, ΔzF, ΔzC, kx², ky²) = - (1/ΔzF[k-1] + 1/ΔzF[k]) - ΔzC[k] * (kx² + ky²)
+
+    d = zeros(Nx, Ny, Nz)
+    for i in 1:Nx, j in 1:Ny
+        d[i, j, 1] = -1/ΔzF[1] - ΔzC[1] * (kx²[i] + ky²[j])
+        d[i, j, 2:Nz-1] .= [δ(k, ΔzF, ΔzC, kx²[i], ky²[j]) for k in 2:Nz-1]
+        d[i, j, Nz] = -1/ΔzF[Nz-1] - ΔzC[Nz] * (kx²[i] + ky²[j])
+    end
+
+    #####
+    ##### Random right hand side
+    #####
+
+    # Random right hand side
+    Ru = CellField(Float64, arch, fake_grid)
+    Rv = CellField(Float64, arch, fake_grid)
+    Rw = CellField(Float64, arch, fake_grid)
+
+    interior(Ru) .= rand(Nx, Ny, Nz)
+    interior(Rv) .= rand(Nx, Ny, Nz)
+    interior(Rw) .= zeros(Nx, Ny, Nz)
+    U = (u=Ru, v=Rv, w=Rw)
+
+    uv_bcs = HorizontallyPeriodicBCs()
+    w_bcs = HorizontallyPeriodicBCs(top=NoPenetrationBC(), bottom=NoPenetrationBC())
+
+    fill_halo_regions!(Ru.data, uv_bcs, arch, fake_grid)
+    fill_halo_regions!(Rv.data, uv_bcs, arch, fake_grid)
+
+    _compute_w_from_continuity!(U, fake_grid)
+
+    fill_halo_regions!(Rw.data,  w_bcs, arch, fake_grid)
+
+    R = zeros(Nx, Ny, Nz)
+    for i in 1:Nx, j in 1:Ny, k in 1:Nz
+        R[i, j, k] = div_ccc(i, j, k, Δx, Δy, ΔzF, Ru.data, Rv.data, Rw.data)
+    end
+
+    # @show sum(R)  # should be zero by construction.
+
+    F = zeros(Nx, Ny, Nz)
+    F = ΔzC .* R  # RHS needs to be multiplied by ΔzC
+
+    #####
+    ##### Solve system
+    #####
+
+    F̃ = fft(F, [1, 2])
+
+    btsolver = BatchedTridiagonalSolver(arch, dl=ld, d=d, du=ud, f=F̃, grid=fake_grid)
+
+    ϕ̃ = zeros(Complex{Float64}, Nx, Ny, Nz)
+    solve_batched_tridiagonal_system!(ϕ̃, arch, btsolver)
+
+    ϕ = CellField(Float64, arch, fake_grid)
+    interior(ϕ) .= real.(ifft(ϕ̃, [1, 2]))
+    ϕ.data .= ϕ.data .- mean(interior(ϕ))
+
+    #####
+    ##### Compute Laplacian of solution ϕ to test that it's correct
+    #####
+
+    # Gotta fill halo regions
+    fbcs = HorizontallyPeriodicBCs()
+    fill_halo_regions!(ϕ.data, fbcs, arch, fake_grid)
+
+    ∇²ϕ = CellField(Float64, arch, fake_grid)
+
+    for i in 1:Nx, j in 1:Ny, k in 1:Nz
+        ∇²ϕ.data[i, j, k] = ∇²(i, j, k, Δx, Δy, ΔzF, ΔzC, ϕ.data)
+    end
+
+    return interior(∇²ϕ) ≈ R
+end
+
 @testset "Solvers" begin
     @info "Testing Solvers..."
 
@@ -134,4 +283,14 @@ end
             end
         end
     end
+
+    for arch in [CPU()]
+        @testset "Vertically stretched Poisson solver [FACR, $arch]" begin
+            @info "  Testing vertically stretched Poisson solver [FACR, $arch]..."
+
+            Nx = Ny = 8
+            zF = [1, 2, 4, 7, 11, 16, 22, 29, 37]
+            @test vertically_stretched_poisson_solver_correct_answer(arch, Nx, Ny, zF)
+        end
+    end
 end