diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index 2b45f4af9a..c58c5ce52a 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -26,9 +26,7 @@ import Oceananigans.Grids: architecture, with_halo, inflate_halo_size_one_dimens static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ, inactive_cell - import Oceananigans.Architectures: on_architecture - import Oceananigans.Fields: fractional_x_index, fractional_y_index, fractional_z_index include("immersed_boundary_grid.jl") diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl index c89a56e40e..b8f5c82d5a 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) end # module + diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index eb9ec9db21..bf3d4600de 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -1,5 +1,4 @@ using Oceananigans.Operators -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid using Statistics: mean using KernelAbstractions: @kernel, @index @@ -32,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 regularizer + ϕ̄ = 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; + regularizer = 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(regularizer) + preconditioner = RegularizedPoissonPreconditioner(preconditioner, rhs, regularizer) - conjugate_gradient_solver = ConjugateGradientSolver(compute_laplacian!; + conjugate_gradient_solver = ConjugateGradientSolver(operator; reltol, abstol, preconditioner, @@ -111,16 +128,41 @@ function compute_preconditioner_rhs!(solver::FourierTridiagonalPoissonSolver, rh return nothing end -const FFTBasedPreconditioner = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} +struct RegularizedPoissonPreconditioner{P, R, D} + unregularized_preconditioner :: P + rhs :: R + regularizer :: D +end -function precondition!(p, preconditioner::FFTBasedPreconditioner, r, args...) - compute_preconditioner_rhs!(preconditioner, r) - solve!(p, preconditioner) +const SolverWithFFT = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} +const FFTBasedPreconditioner = RegularizedPoissonPreconditioner{<:SolverWithFFT} +function precondition!(p, regularized::FFTBasedPreconditioner, r, args...) + solver = regularized.unregularized_preconditioner + compute_preconditioner_rhs!(solver, r) + solve!(p, solver) + regularize_poisson_solution!(p, regularized) + return p +end + +function regularize_poisson_solution!(p, regularized) + δ = regularized.regularizer + rhs = regularized.rhs mean_p = mean(p) + + if !isnothing(δ) + mean_rhs = mean(rhs) + Δp = mean_p + mean_rhs / δ + + # TODO: figure out if we should avoid zeroing the mean_p + # Δ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 @@ -128,25 +170,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 RegularizedAPP = RegularizedPoissonPreconditioner{<:AsymptoticPoissonPreconditioner} +Base.summary(::AsymptoticPoissonPreconditioner) = "AsymptoticPoissonPreconditioner" -@inline function precondition!(p, ::DiagonallyDominantPreconditioner, r, args...) +@inline function precondition!(p, preconditioner::RegularizedAPP, 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 @@ -162,17 +202,22 @@ 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