Generalize assumed_field_location, plus allow FourierTridiagonalPoissonSolver on regular grid (#4535)

* Generalize assumed_field_location

* updates to support anelastic models

* generalize constructor to accomodate different boundary conditions / formulations

* update

* fix for stretched cases

* update interface to prssures corrections

* fixes

* try to fix distributed solver

* fix

* fixes

* clean up

Author: glwagner

ext/OceananigansReactantExt/TimeSteppers.jl

@@ -12,7 +12,7 @@ using Oceananigans.Utils: @apply_regionally, apply_regionally!
 using Oceananigans.TimeSteppers:
     update_state!,
     tick!,
-    calculate_pressure_correction!,
+    compute_pressure_correction!,
     correct_velocities_and_cache_previous_tendencies!,
     step_lagrangian_particles!,
     QuasiAdamsBashforth2TimeStepper
@@ -103,7 +103,7 @@ function time_step!(model::ReactantModel{<:QuasiAdamsBashforth2TimeStepper{FT}},
         model.clock.last_stage_Δt = Δt
     end
 
-    calculate_pressure_correction!(model, Δt)
+    compute_pressure_correction!(model, Δt)
     correct_velocities_and_cache_previous_tendencies!(model, Δt)
 
     update_state!(model, callbacks; compute_tendencies=true)

src/DistributedComputations/distributed_fft_tridiagonal_solver.jl

@@ -5,10 +5,14 @@ using Oceananigans.Operators: Δxᶠᵃᵃ, Δyᵃᶠᵃ, Δzᵃᵃᶠ
 
 using Oceananigans.Solvers: BatchedTridiagonalSolver,
                             stretched_direction,
+                            tridiagonal_direction,
+                            dimension,
+                            HomogeneousNeumannFormulation,
                             ZTridiagonalSolver,
                             YTridiagonalSolver,
                             XTridiagonalSolver,
-                            compute_main_diagonal!
+                            compute_main_diagonal!,
+                            compute_lower_diagonal!
 
 struct DistributedFourierTridiagonalPoissonSolver{G, L, B, P, R, S, β}
     plan :: P
@@ -146,30 +150,34 @@ Restrictions
     - Same as for two-dimensional decompositions with `Rx` (or `Ry`) equal to one
 
 """
-function DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid, planner_flag=FFTW.PATIENT; tridiagonal_direction = nothing)
+function DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid, planner_flag=FFTW.PATIENT; tridiagonal_formulation=nothing)
 
     validate_poisson_solver_distributed_grid(global_grid)
     validate_poisson_solver_configuration(global_grid, local_grid)
 
-    if isnothing(tridiagonal_direction)
-        tridiagonal_dim = stretched_dimensions(local_grid)[1]
-        tridiagonal_direction = stretched_direction(local_grid)
+    # Try to guess what direction should be tridiagonal
+    if isnothing(tridiagonal_formulation)
+        tridiagonal_dir = global_grid isa XYZRegularRG ? ZDirection() : stretched_direction(global_grid)
+        tridiagonal_formulation = HomogeneousNeumannFormulation(tridiagonal_dir)
     else
-        tridiagonal_dim = tridiagonal_direction == XDirection() ? 1 :
-                          tridiagonal_direction == YDirection() ? 2 : 3
+        tridiagonal_dir = tridiagonal_direction(tridiagonal_formulation)
     end
 
+    tridiagonal_dim = dimension(tridiagonal_dir)
+
     topology(global_grid, tridiagonal_dim) != Bounded &&
         error("`DistributedFourierTridiagonalPoissonSolver` requires that the stretched direction (dimension $tridiagonal_dim) is `Bounded`.")
 
-    FT         = Complex{eltype(local_grid)}
+    T = Complex{eltype(local_grid)}
     child_arch = child_architecture(local_grid)
-    storage    = TransposableField(CenterField(local_grid), FT)
+    storage_field = CenterField(local_grid)
+    storage = TransposableField(storage_field, T)
 
-    topo = (TX, TY, TZ) = topology(global_grid)
-    λx = dropdims(poisson_eigenvalues(global_grid.Nx, global_grid.Lx, 1, TX()), dims=(2, 3))
-    λy = dropdims(poisson_eigenvalues(global_grid.Ny, global_grid.Ly, 2, TY()), dims=(1, 3))
-    λz = dropdims(poisson_eigenvalues(global_grid.Nz, global_grid.Lz, 3, TZ()), dims=(1, 2))
+    TX, TY, TZ = topology(global_grid)
+    tx, ty, tz = TX(), TY(), TZ()
+    λx = dropdims(poisson_eigenvalues(global_grid.Nx, global_grid.Lx, 1, tx), dims=(2, 3))
+    λy = dropdims(poisson_eigenvalues(global_grid.Ny, global_grid.Ly, 2, ty), dims=(1, 3))
+    λz = dropdims(poisson_eigenvalues(global_grid.Nz, global_grid.Lz, 3, tz), dims=(1, 2))
 
     if tridiagonal_dim == 1
         arch = architecture(storage.xfield.grid)
@@ -191,39 +199,31 @@ function DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid, pla
     λ1 = on_architecture(child_arch, λ1)
     λ2 = on_architecture(child_arch, λ2)
 
-    plan = plan_distributed_transforms(global_grid, storage, planner_flag)
-
     # Lower and upper diagonals are the same
-    lower_diagonal = @allowscalar [ 1 / Δξᶠ(q, grid, Val(tridiagonal_dim)) for q in 2:size(grid, tridiagonal_dim) ]
-    lower_diagonal = on_architecture(child_arch, lower_diagonal)
-    upper_diagonal = lower_diagonal
+    main_diagonal = zeros(grid, size(grid)...)
 
-    # Compute diagonal coefficients for each grid point
-    diagonal = zeros(eltype(grid), size(grid)...)
-    diagonal = on_architecture(arch, diagonal)
-    launch_config = if tridiagonal_dim == 1
-                        :yz
-                    elseif tridiagonal_dim == 2
-                        :xz
-                    elseif tridiagonal_dim == 3
-                        :xy
-                    end
+    Nd = size(grid, tridiagonal_dim) - 1
+    lower_diagonal = zeros(grid, Nd) 
+    upper_diagonal = lower_diagonal
 
-    launch!(arch, grid, launch_config, compute_main_diagonal!, diagonal, grid, λ1, λ2, tridiagonal_direction)
+    compute_main_diagonal!(main_diagonal, tridiagonal_formulation, grid, λ1, λ2)
+    Nd > 0 && compute_lower_diagonal!(lower_diagonal, tridiagonal_formulation, grid)
 
     # Set up batched tridiagonal solver
-    btsolver = BatchedTridiagonalSolver(grid; lower_diagonal, diagonal, upper_diagonal, tridiagonal_direction)
+    btsolver = BatchedTridiagonalSolver(grid; lower_diagonal, upper_diagonal,
+                                        diagonal = main_diagonal,
+                                        tridiagonal_direction = tridiagonal_dir)
 
-    # We need to permute indices to apply bounded transforms on the GPU (r2r of r2c with twiddling)
+    # We need to permute indices to apply bounded transforms on the GPU (r2r or r2c with twiddling)
     x_buffer_needed = child_arch isa GPU && TX == Bounded
     z_buffer_needed = child_arch isa GPU && TZ == Bounded
 
     # We cannot really batch anything, so on GPUs we always have to permute indices in the y direction
     y_buffer_needed = child_arch isa GPU
 
-    buffer_x = x_buffer_needed ? on_architecture(child_arch, zeros(FT, size(storage.xfield)...)) : nothing
-    buffer_y = y_buffer_needed ? on_architecture(child_arch, zeros(FT, size(storage.yfield)...)) : nothing
-    buffer_z = z_buffer_needed ? on_architecture(child_arch, zeros(FT, size(storage.zfield)...)) : nothing
+    buffer_x = x_buffer_needed ? on_architecture(child_arch, zeros(T, size(storage.xfield)...)) : nothing
+    buffer_y = y_buffer_needed ? on_architecture(child_arch, zeros(T, size(storage.yfield)...)) : nothing
+    buffer_z = z_buffer_needed ? on_architecture(child_arch, zeros(T, size(storage.zfield)...)) : nothing
 
     buffer = if tridiagonal_dim == 1
         (; y = buffer_y, z = buffer_z)
@@ -233,6 +233,8 @@ function DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid, pla
         (; x = buffer_x, y = buffer_y)
     end
 
+    plan = plan_distributed_transforms(global_grid, storage, planner_flag)
+
     if tridiagonal_dim == 1
         forward  = (y! = plan.forward.y!,  z! = plan.forward.z!)
         backward = (y! = plan.backward.y!, z! = plan.backward.z!)

src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl

@@ -1,23 +1,23 @@
 using .SplitExplicitFreeSurfaces: barotropic_split_explicit_corrector!
-import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities!
+import Oceananigans.TimeSteppers: compute_pressure_correction!, make_pressure_correction!
 
-calculate_pressure_correction!(::HydrostaticFreeSurfaceModel, Δt) = nothing
+compute_pressure_correction!(::HydrostaticFreeSurfaceModel, Δt) = nothing
 
 #####
 ##### Barotropic pressure correction for models with a free surface
 #####
 
-pressure_correct_velocities!(model::HydrostaticFreeSurfaceModel, Δt; kwargs...) =
-    pressure_correct_velocities!(model, model.free_surface, Δt; kwargs...)
+make_pressure_correction!(model::HydrostaticFreeSurfaceModel, Δt; kwargs...) =
+    make_pressure_correction!(model, model.free_surface, Δt; kwargs...)
 
 # Fallback
-pressure_correct_velocities!(model, free_surface, Δt; kwargs...) = nothing
+make_pressure_correction!(model, free_surface, Δt; kwargs...) = nothing
 
 #####
 ##### Barotropic pressure correction for models with an Implicit free surface
 #####
 
-function pressure_correct_velocities!(model, ::ImplicitFreeSurface, Δt)
+function make_pressure_correction!(model, ::ImplicitFreeSurface, Δt)
 
     launch!(model.architecture, model.grid, :xyz,
             _barotropic_pressure_correction!,
@@ -30,7 +30,7 @@ function pressure_correct_velocities!(model, ::ImplicitFreeSurface, Δt)
     return nothing
 end
 
-function pressure_correct_velocities!(model, ::SplitExplicitFreeSurface, Δt)
+function make_pressure_correction!(model, ::SplitExplicitFreeSurface, Δt)
     u, v, _ = model.velocities
     grid = model.grid
     barotropic_split_explicit_corrector!(u, v, model.free_surface, grid)

src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl

@@ -160,7 +160,6 @@ end
     @inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
 end
 
-
 #####
 ##### Tracer(s)
 #####

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -222,7 +222,6 @@ function NonhydrostaticModel(; grid,
     model_fields = merge(velocities, tracers, auxiliary_fields)
     prognostic_fields = merge(velocities, tracers)
 
-
     # Instantiate timestepper if not already instantiated
     implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid)
     timestepper = TimeStepper(timestepper, grid, prognostic_fields; implicit_solver=implicit_solver)

src/Models/NonhydrostaticModels/pressure_correction.jl

@@ -1,19 +1,16 @@
-import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities!
+import Oceananigans.TimeSteppers: compute_pressure_correction!, make_pressure_correction!
 
 """
-    calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
+    compute_pressure_correction!(model::NonhydrostaticModel, Δt)
 
 Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `velocities`, and step size `Δt`.
 """
-function calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
+function compute_pressure_correction!(model::NonhydrostaticModel, Δt)
 
     # Mask immersed velocities
     foreach(mask_immersed_field!, model.velocities)
-
     fill_halo_regions!(model.velocities, model.clock, fields(model))
-
     solve_for_pressure!(model.pressures.pNHS, model.pressure_solver, Δt, model.velocities)
-
     fill_halo_regions!(model.pressures.pNHS)
 
     return nothing
@@ -28,7 +25,7 @@ Update the predictor velocities u, v, and w with the non-hydrostatic pressure mu
 
     `u^{n+1} = u^n - δₓp_{NH} * Δt / Δx`
 """
-@kernel function _pressure_correct_velocities!(U, grid, pNHSΔt)
+@kernel function _make_pressure_correction!(U, grid, pNHSΔt)
     i, j, k = @index(Global, NTuple)
 
     @inbounds U.u[i, j, k] -= ∂xᶠᶜᶜ(i, j, k, grid, pNHSΔt)
@@ -37,10 +34,10 @@ Update the predictor velocities u, v, and w with the non-hydrostatic pressure mu
 end
 
 "Update the solution variables (velocities and tracers)."
-function pressure_correct_velocities!(model::NonhydrostaticModel, Δt)
+function make_pressure_correction!(model::NonhydrostaticModel, Δt)
 
     launch!(model.architecture, model.grid, :xyz,
-            _pressure_correct_velocities!,
+            _make_pressure_correction!,
             model.velocities,
             model.grid,
             model.pressures.pNHS)

src/Models/NonhydrostaticModels/set_nonhydrostatic_model.jl

@@ -1,5 +1,5 @@
 using Oceananigans.BoundaryConditions: fill_halo_regions!
-using Oceananigans.TimeSteppers: update_state!, calculate_pressure_correction!, pressure_correct_velocities!
+using Oceananigans.TimeSteppers: update_state!, compute_pressure_correction!, make_pressure_correction!
 
 import Oceananigans.Fields: set!
 
@@ -51,8 +51,8 @@ function set!(model::NonhydrostaticModel; enforce_incompressibility=true, kwargs
 
     if enforce_incompressibility
         FT = eltype(model.grid)
-        calculate_pressure_correction!(model, one(FT))
-        pressure_correct_velocities!(model, one(FT))
+        compute_pressure_correction!(model, one(FT))
+        make_pressure_correction!(model, one(FT))
         update_state!(model; compute_tendencies = false)
     end
 

src/Models/NonhydrostaticModels/solve_for_pressure.jl

@@ -9,84 +9,84 @@ using Oceananigans.Solvers: solve!
 ##### Calculate the right-hand-side of the non-hydrostatic pressure Poisson equation.
 #####
 
-@kernel function _compute_source_term!(rhs, grid, Δt, Ũ)
+@kernel function _compute_source_term!(rhs, grid, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
+    u, v, w = Ũ
+    δ = divᶜᶜᶜ(i, j, k, grid, u, v, w)
     @inbounds rhs[i, j, k] = active * δ
 end
 
-@kernel function _fourier_tridiagonal_source_term!(rhs, ::XDirection, grid, Δt, Ũ)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::XDirection, grid, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
+    u, v, w = Ũ
+    δ = divᶜᶜᶜ(i, j, k, grid, u, v, w)
     @inbounds rhs[i, j, k] = active * Δxᶜᶜᶜ(i, j, k, grid) * δ
 end
 
-@kernel function _fourier_tridiagonal_source_term!(rhs, ::YDirection, grid, Δt, Ũ)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::YDirection, grid, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
+    u, v, w = Ũ
+    δ = divᶜᶜᶜ(i, j, k, grid, u, v, w)
     @inbounds rhs[i, j, k] = active * Δyᶜᶜᶜ(i, j, k, grid) * δ
 end
 
-@kernel function _fourier_tridiagonal_source_term!(rhs, ::ZDirection, grid, Δt, Ũ)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::ZDirection, grid, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
+    u, v, w = Ũ
+    δ = divᶜᶜᶜ(i, j, k, grid, u, v, w)
     @inbounds rhs[i, j, k] = active * Δzᶜᶜᶜ(i, j, k, grid) * δ
 end
 
-function compute_source_term!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, Ũ)
+function compute_source_term!(solver::DistributedFFTBasedPoissonSolver, Ũ)
     rhs  = solver.storage.zfield
     arch = architecture(solver)
     grid = solver.local_grid
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Ũ)
     return nothing
 end
 
-function compute_source_term!(pressure, solver::DistributedFourierTridiagonalPoissonSolver, Δt, Ũ)
+function compute_source_term!(solver::DistributedFourierTridiagonalPoissonSolver, Ũ)
     rhs = solver.storage.zfield
     arch = architecture(solver)
     grid = solver.local_grid
     tdir = solver.batched_tridiagonal_solver.tridiagonal_direction
-    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Δt, Ũ)
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Ũ)
     return nothing
 end
 
-function compute_source_term!(pressure, solver::FourierTridiagonalPoissonSolver, Δt, Ũ)
+function compute_source_term!(solver::FourierTridiagonalPoissonSolver, Ũ)
     rhs = solver.source_term
     arch = architecture(solver)
     grid = solver.grid
     tdir = solver.batched_tridiagonal_solver.tridiagonal_direction
-    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Δt, Ũ)
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Ũ)
     return nothing
 end
 
-function compute_source_term!(pressure, solver::FFTBasedPoissonSolver, Δt, Ũ)
+function compute_source_term!(solver::FFTBasedPoissonSolver, Ũ)
     rhs = solver.storage
     arch = architecture(solver)
     grid = solver.grid
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Ũ)
     return nothing
 end
 
 #####
 ##### Solve for pressure
 #####
 
-function solve_for_pressure!(pressure, solver, Δt, Ũ)
-    ϵ = eps(eltype(pressure))
-    Δt⁺ = max(ϵ, Δt)
-    Δt★ = Δt⁺ * isfinite(Δt)
-    pressure .*= Δt★
-
-    compute_source_term!(pressure, solver, Δt, Ũ)
+# Note that Δt is unused here.
+function solve_for_pressure!(pressure, solver, Δt, args...)
+    compute_source_term!(solver, args...)
     solve!(pressure, solver)
     return pressure
 end
 
-function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, Ũ)
+function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, args...)
     ϵ = eps(eltype(pressure))
     Δt⁺ = max(ϵ, Δt)
     Δt★ = Δt⁺ * isfinite(Δt)
@@ -95,7 +95,7 @@ function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, 
     rhs = solver.right_hand_side
     grid = solver.grid
     arch = architecture(grid)
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, args...)
     return solve!(pressure, solver.conjugate_gradient_solver, rhs)
 end