diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index e9281eac32..e679a69f55 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -304,6 +304,12 @@ end isrectilinear(ibg::IBG) = isrectilinear(ibg.underlying_grid) +# TODO: eliminate when ImmersedBoundaries precedes Solvers +# This is messy because IBG does _not_ have a Poisson solver, only the underlying grid does +import Oceananigans.Solvers: has_fft_poisson_solver, fft_poisson_solver +has_fft_poisson_solver(ibg::IBG) = has_fft_poisson_solver(ibg.underlying_grid) +fft_poisson_solver(ibg::IBG) = fft_poisson_solver(ibg.underlying_grid) + @inline fractional_x_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_x_index(x, locs, grid.underlying_grid) @inline fractional_y_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_y_index(x, locs, grid.underlying_grid) @inline fractional_z_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_z_index(x, locs, grid.underlying_grid) diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl index bbad459ea9..fab7e36221 100644 --- a/src/Solvers/Solvers.jl +++ b/src/Solvers/Solvers.jl @@ -28,7 +28,7 @@ using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG Return the `M`th root of unity raised to the `k`th power. """ -@inline ω(M, k) = exp(-2im*π*k/M) +@inline ω(M, k) = exp(-2im * π * k / M) reshaped_size(N, dim) = dim == 1 ? (N, 1, 1) : dim == 2 ? (1, N, 1) : @@ -51,8 +51,13 @@ include("heptadiagonal_iterative_solver.jl") const GridWithFFTSolver = Union{XYZRegularRG, XYRegularRG, XZRegularRG, YZRegularRG} const GridWithFourierTridiagonalSolver = Union{XYRegularRG, XZRegularRG, YZRegularRG} +# TODO: will be unnecessary when ImmersedBoundaries precedes Solvers +has_fft_poisson_solver(grid) = false +has_fft_poisson_solver(::GridWithFFTSolver) = true + fft_poisson_solver(grid::XYZRegularRG) = FFTBasedPoissonSolver(grid) fft_poisson_solver(grid::GridWithFourierTridiagonalSolver) = FourierTridiagonalPoissonSolver(grid.underlying_grid) end # module + diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index 13ae437932..516b3cc6b5 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -1,4 +1,4 @@ -using Oceananigans.Operators: divᶜᶜᶜ, ∇²ᶜᶜᶜ +using Oceananigans.Operators using Statistics: mean using KernelAbstractions: @kernel, @index @@ -31,36 +31,54 @@ end @kernel function laplacian!(∇²ϕ, grid, ϕ) i, j, k = @index(Global, NTuple) - @inbounds ∇²ϕ[i, j, k] = ∇²ᶜᶜᶜ(i, j, k, grid, ϕ) + active = !inactive_cell(i, j, k, grid) + @inbounds ∇²ϕ[i, j, k] = ∇²ᶜᶜᶜ(i, j, k, grid, ϕ) * active +end + +struct RegularizedLaplacian{D} + δ :: D end -function compute_laplacian!(∇²ϕ, ϕ) +function (L::RegularizedLaplacian)(Lϕ, ϕ) grid = ϕ.grid arch = architecture(grid) fill_halo_regions!(ϕ) - launch!(arch, grid, :xyz, laplacian!, ∇²ϕ, grid, ϕ) + launch!(arch, grid, :xyz, laplacian!, Lϕ, grid, ϕ) + + if !isnothing(L.δ) + # Add regularization + ϕ̄ = mean(ϕ) + ΔLϕ = L.δ * ϕ̄ + grid = ϕ.grid + arch = architecture(grid) + launch!(arch, grid, :xyz, subtract_and_mask!, Lϕ, grid, ΔLϕ) + end + return nothing end struct DefaultPreconditioner end function ConjugateGradientPoissonSolver(grid; + regularization = nothing, preconditioner = DefaultPreconditioner(), reltol = sqrt(eps(grid)), abstol = sqrt(eps(grid)), kw...) if preconditioner isa DefaultPreconditioner # try to make a useful default - if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFTSolver - preconditioner = fft_poisson_solver(grid.underlying_grid) + if has_fft_poisson_solver(grid) + preconditioner = fft_poisson_solver(grid) else - preconditioner = DiagonallyDominantPreconditioner() + preconditioner = AsymptoticPoissonPreconditioner() end end rhs = CenterField(grid) + operator = RegularizedLaplacian(regularization) + preconditioner = RegularizedPreconditioner(preconditioner, rhs, regularization) - conjugate_gradient_solver = ConjugateGradientSolver(compute_laplacian!; + conjugate_gradient_solver = ConjugateGradientSolver(operator; reltol, abstol, preconditioner, @@ -110,16 +128,38 @@ function compute_preconditioner_rhs!(solver::FourierTridiagonalPoissonSolver, rh return nothing end -const FFTBasedPreconditioner = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} +struct RegularizedPreconditioner{P, R, D} + preconditioner :: P + rhs :: R + regularization :: D +end -function precondition!(p, preconditioner::FFTBasedPreconditioner, r, args...) - compute_preconditioner_rhs!(preconditioner, r) - p = solve!(p, preconditioner) +const SolverWithFFT = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} +const FFTBasedPreconditioner = RegularizedPreconditioner{<:SolverWithFFT} +function precondition!(p, regularized::FFTBasedPreconditioner, r, args...) + solver = regularized.preconditioner + compute_preconditioner_rhs!(solver, r) + solve!(p, solver) + regularize_poisson_solution!(p, regularized) + return p +end + +function regularize_poisson_solution!(p, regularized) + δ = regularized.regularization + rhs = regularized.rhs mean_p = mean(p) + + if !isnothing(δ) + mean_rhs = mean(rhs) + Δp = mean_p + mean_rhs / δ + else + Δp = mean_p + end + grid = p.grid arch = architecture(grid) - launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, mean_p) + launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, Δp) return p end @@ -127,25 +167,23 @@ end @kernel function subtract_and_mask!(a, grid, b) i, j, k = @index(Global, NTuple) active = !inactive_cell(i, j, k, grid) - a[i, j, k] = (a[i, j, k] - b) * active + @inbounds a[i, j, k] = (a[i, j, k] - b) * active end ##### -##### The "DiagonallyDominantPreconditioner" (Marshall et al 1997) +##### The "AsymptoticPoissonPreconditioner" (Marshall et al 1997) ##### -struct DiagonallyDominantPreconditioner end -Base.summary(::DiagonallyDominantPreconditioner) = "DiagonallyDominantPreconditioner" +struct AsymptoticPoissonPreconditioner end +const RegularizedNDPP = RegularizedPreconditioner{<:AsymptoticPoissonPreconditioner} +Base.summary(::AsymptoticPoissonPreconditioner) = "AsymptoticPoissonPreconditioner" -@inline function precondition!(p, ::DiagonallyDominantPreconditioner, r, args...) +@inline function precondition!(p, preconditioner::RegularizedNDPP, r, args...) grid = r.grid arch = architecture(p) fill_halo_regions!(r) - launch!(arch, grid, :xyz, _diagonally_dominant_precondition!, p, grid, r) - - mean_p = mean(p) - launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, mean_p) - + launch!(arch, grid, :xyz, _asymptotic_poisson_precondition!, p, grid, r) + regularize_poisson_solution!(p, preconditioner) return p end @@ -161,17 +199,23 @@ end Ay⁻(i, j, k, grid) - Ay⁺(i, j, k, grid) - Az⁻(i, j, k, grid) - Az⁺(i, j, k, grid) -@inline heuristic_residual(i, j, k, grid, r) = - @inbounds 1 / Ac(i, j, k, grid) * (r[i, j, k] - 2 * Ax⁻(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i-1, j, k, grid)) * r[i-1, j, k] - - 2 * Ax⁺(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i+1, j, k, grid)) * r[i+1, j, k] - - 2 * Ay⁻(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j-1, k, grid)) * r[i, j-1, k] - - 2 * Ay⁺(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j+1, k, grid)) * r[i, j+1, k] - - 2 * Az⁻(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j, k-1, grid)) * r[i, j, k-1] - - 2 * Az⁺(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j, k+1, grid)) * r[i, j, k+1]) - -@kernel function _diagonally_dominant_precondition!(p, grid, r) +@inline function heuristic_poisson_solution(i, j, k, grid, r) + @inbounds begin + a⁰⁰⁰ = r[i, j, k] + a⁻⁰⁰ = 2 * Ax⁻(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i-1, j, k, grid)) * r[i-1, j, k] + a⁺⁰⁰ = 2 * Ax⁺(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i+1, j, k, grid)) * r[i+1, j, k] + a⁰⁻⁰ = 2 * Ay⁻(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j-1, k, grid)) * r[i, j-1, k] + a⁰⁺⁰ = 2 * Ay⁺(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j+1, k, grid)) * r[i, j+1, k] + a⁰⁰⁻ = 2 * Az⁻(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j, k-1, grid)) * r[i, j, k-1] + a⁰⁰⁺ = 2 * Az⁺(i, j, k, grid) / (Ac(i, j, k, grid) + Ac(i, j, k+1, grid)) * r[i, j, k+1] + end + + return (a⁰⁰⁰ - a⁻⁰⁰ - a⁺⁰⁰ - a⁰⁻⁰ - a⁰⁺⁰ - a⁰⁰⁻ - a⁰⁰⁺) / Ac(i, j, k, grid) +end + +@kernel function _asymptotic_poisson_precondition!(p, grid, r) i, j, k = @index(Global, NTuple) active = !inactive_cell(i, j, k, grid) - @inbounds p[i, j, k] = heuristic_residual(i, j, k, grid, r) * active + @inbounds p[i, j, k] = heuristic_poisson_solution(i, j, k, grid, r) * active end diff --git a/src/Solvers/conjugate_gradient_solver.jl b/src/Solvers/conjugate_gradient_solver.jl index 848bba44f6..56448e7875 100644 --- a/src/Solvers/conjugate_gradient_solver.jl +++ b/src/Solvers/conjugate_gradient_solver.jl @@ -94,18 +94,18 @@ function ConjugateGradientSolver(linear_operation; FT = eltype(grid) return ConjugateGradientSolver(arch, - grid, - linear_operation, - FT(reltol), - FT(abstol), - maxiter, - 0, - zero(FT), - linear_operator_product, - search_direction, - residual, - preconditioner, - precondition_product) + grid, + linear_operation, + FT(reltol), + FT(abstol), + maxiter, + 0, + zero(FT), + linear_operator_product, + search_direction, + residual, + preconditioner, + precondition_product) end """ @@ -269,3 +269,4 @@ function Base.show(io::IO, solver::ConjugateGradientSolver) "├── abstol: ", prettysummary(solver.abstol), "\n", "└── maxiter: ", solver.maxiter) end + diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl index 3f5fae67e3..f6252947b6 100644 --- a/src/Solvers/fft_based_poisson_solver.jl +++ b/src/Solvers/fft_based_poisson_solver.jl @@ -79,7 +79,7 @@ end Solve the "generalized" Poisson equation, ```math -(∇² + m) ϕ = b, +(∇² + m) ϕ + μ ϕ̄ = b, ``` where ``m`` is a number, using a eigenfunction expansion of the discrete Poisson operator @@ -92,7 +92,7 @@ elements (typically the same type as `solver.storage`). Equation ``(∇² + m) ϕ = b`` is sometimes referred to as the "screened Poisson" equation when ``m < 0``, or the Helmholtz equation when ``m > 0``. """ -function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0) +function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0, μ=0) arch = architecture(solver) topo = TX, TY, TZ = topology(solver.grid) Nx, Ny, Nz = size(solver.grid) @@ -112,7 +112,11 @@ function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0) # If m === 0, the "zeroth mode" at `i, j, k = 1, 1, 1` is undetermined; # we set this to zero by default. Another slant on this "problem" is that # λx[1, 1, 1] + λy[1, 1, 1] + λz[1, 1, 1] = 0, which yields ϕ[1, 1, 1] = Inf or NaN. - m === 0 && CUDA.@allowscalar ϕc[1, 1, 1] = 0 + if m === μ === 0 + CUDA.@allowscalar ϕc[1, 1, 1] = 0 + elseif μ > 0 + CUDA.@allowscalar ϕc[1, 1, 1] = - b[1, 1, 1] / (μ - m) + end # Apply backward transforms in order for transform! in solver.transforms.backward