From 4f6ffd71933fc3dd55df8a6093fd52d420b06b1d Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 16 Oct 2024 09:49:24 -0600 Subject: [PATCH 01/10] Add regularization to Poisson operator for CG solve --- .../conjugate_gradient_poisson_solver.jl | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index 13ae437932..0ef70f8a21 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -34,17 +34,29 @@ end @inbounds ∇²ϕ[i, j, k] = ∇²ᶜᶜᶜ(i, j, k, grid, ϕ) end -function compute_laplacian!(∇²ϕ, ϕ) +struct RegularizedLaplacian{D} + δ :: D +end + +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(ϕ) + parent(Lϕ) .+= L.δ * ϕ̄ + end + return nothing end struct DefaultPreconditioner end function ConjugateGradientPoissonSolver(grid; + regularization = nothing, preconditioner = DefaultPreconditioner(), reltol = sqrt(eps(grid)), abstol = sqrt(eps(grid)), @@ -60,7 +72,8 @@ function ConjugateGradientPoissonSolver(grid; rhs = CenterField(grid) - conjugate_gradient_solver = ConjugateGradientSolver(compute_laplacian!; + operator = RegularizedLaplacian(regularization) + conjugate_gradient_solver = ConjugateGradientSolver(operator; reltol, abstol, preconditioner, From 19f0d689be283fc080159141082d367add47364c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 16 Oct 2024 16:28:57 -0600 Subject: [PATCH 02/10] Add custom gauge condition to FFT solver --- src/Solvers/fft_based_poisson_solver.jl | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl index 3f5fae67e3..72f30d601b 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,13 @@ 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 + if μ === 0 + CUDA.@allowscalar ϕc[1, 1, 1] = 0 + else + CUDA.@allowscalar ϕc[1, 1, 1] = - b[1, 1, 1] / μ + end + end # Apply backward transforms in order for transform! in solver.transforms.backward From b62951ce5068af488634b286043115ad1a764c69 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 16 Oct 2024 16:33:27 -0600 Subject: [PATCH 03/10] Better implementation --- src/Solvers/fft_based_poisson_solver.jl | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl index 72f30d601b..dfcca1cedb 100644 --- a/src/Solvers/fft_based_poisson_solver.jl +++ b/src/Solvers/fft_based_poisson_solver.jl @@ -112,12 +112,10 @@ function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0, μ=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. - if m === 0 - if μ === 0 - CUDA.@allowscalar ϕc[1, 1, 1] = 0 - else - CUDA.@allowscalar ϕc[1, 1, 1] = - b[1, 1, 1] / μ - end + 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 From d382923a65236f1140cdd1e6162c9fc8d4434681 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 16 Oct 2024 16:34:18 -0600 Subject: [PATCH 04/10] Fix sign error --- src/Solvers/fft_based_poisson_solver.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl index dfcca1cedb..f6252947b6 100644 --- a/src/Solvers/fft_based_poisson_solver.jl +++ b/src/Solvers/fft_based_poisson_solver.jl @@ -115,7 +115,7 @@ function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0, μ=0) if m === μ === 0 CUDA.@allowscalar ϕc[1, 1, 1] = 0 elseif μ > 0 - CUDA.@allowscalar ϕc[1, 1, 1] = - b[1, 1, 1] / (m + μ) + CUDA.@allowscalar ϕc[1, 1, 1] = - b[1, 1, 1] / (μ - m) end # Apply backward transforms in order From f132a5f909e02eed534e4512a3e8f28cf3da52d3 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 16 Oct 2024 17:28:17 -0600 Subject: [PATCH 05/10] Cosmetic fix --- src/Solvers/conjugate_gradient_solver.jl | 25 ++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) 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 + From 42c3993fbde479a4680aa8ef9246f0fe27186d60 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 16 Oct 2024 17:28:38 -0600 Subject: [PATCH 06/10] Add regularization consistently in preconditioner --- .../conjugate_gradient_poisson_solver.jl | 60 ++++++++++++++----- 1 file changed, 46 insertions(+), 14 deletions(-) diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index 0ef70f8a21..e0ff9212b7 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 @@ -66,13 +66,14 @@ function ConjugateGradientPoissonSolver(grid; if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFTSolver preconditioner = fft_poisson_solver(grid.underlying_grid) else - preconditioner = DiagonallyDominantPreconditioner() + preconditioner = NearDiagonalPoissonPreconditioner() end end rhs = CenterField(grid) - operator = RegularizedLaplacian(regularization) + preconditioner = RegularizedPreconditioner(preconditioner, rhs, regularization) + conjugate_gradient_solver = ConjugateGradientSolver(operator; reltol, abstol, @@ -123,16 +124,34 @@ 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 + +const SolverWithFFT = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} +const FFTBasedPreconditioner = RegularizedPreconditioner{<:SolverWithFFT} -function precondition!(p, preconditioner::FFTBasedPreconditioner, r, args...) - compute_preconditioner_rhs!(preconditioner, r) - p = solve!(p, preconditioner) +function precondition!(p, regularized::FFTBasedPreconditioner, r, args...) + solver = regularized.preconditioner + compute_preconditioner_rhs!(solver, r) + p = solve!(p, solver) + δ = 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 @@ -140,24 +159,37 @@ 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 "NearDiagonalPoissonPreconditioner" (Marshall et al 1997) ##### -struct DiagonallyDominantPreconditioner end -Base.summary(::DiagonallyDominantPreconditioner) = "DiagonallyDominantPreconditioner" +struct NearDiagonalPoissonPreconditioner end +const RegularizedNDPP = RegularizedPreconditioner{<:NearDiagonalPoissonPreconditioner} +Base.summary(::NearDiagonalPoissonPreconditioner) = "NearDiagonalPoissonPreconditioner" -@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) + δ = preconditioner.regularization + rhs = preconditioner.rhs mean_p = mean(p) - launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, 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, Δp) return p end From b667af0c35b3191fa738894aea47cac0e3e82a5e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 17 Oct 2024 10:05:13 -0600 Subject: [PATCH 07/10] Add temporary interface for building fft solver --- src/ImmersedBoundaries/ImmersedBoundaries.jl | 6 ++++++ src/Solvers/Solvers.jl | 7 ++++++- src/Solvers/conjugate_gradient_poisson_solver.jl | 4 ++-- 3 files changed, 14 insertions(+), 3 deletions(-) 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 e0ff9212b7..be063dd806 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -63,8 +63,8 @@ function ConjugateGradientPoissonSolver(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 = NearDiagonalPoissonPreconditioner() end From c653542559c42d7e49d02501bd2f7df63423df9c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 17 Oct 2024 12:24:35 -0600 Subject: [PATCH 08/10] Change name of preconditioner --- .../conjugate_gradient_poisson_solver.jl | 38 +++++++++++-------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index be063dd806..d2b0f53245 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -66,7 +66,7 @@ function ConjugateGradientPoissonSolver(grid; if has_fft_poisson_solver(grid) preconditioner = fft_poisson_solver(grid) else - preconditioner = NearDiagonalPoissonPreconditioner() + preconditioner = AsymptoticPoissonPreconditioner() end end @@ -163,18 +163,18 @@ end end ##### -##### The "NearDiagonalPoissonPreconditioner" (Marshall et al 1997) +##### The "AsymptoticPoissonPreconditioner" (Marshall et al 1997) ##### -struct NearDiagonalPoissonPreconditioner end -const RegularizedNDPP = RegularizedPreconditioner{<:NearDiagonalPoissonPreconditioner} -Base.summary(::NearDiagonalPoissonPreconditioner) = "NearDiagonalPoissonPreconditioner" +struct AsymptoticPoissonPreconditioner end +const RegularizedNDPP = RegularizedPreconditioner{<:AsymptoticPoissonPreconditioner} +Base.summary(::AsymptoticPoissonPreconditioner) = "AsymptoticPoissonPreconditioner" @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) + launch!(arch, grid, :xyz, _asymptotic_poisson_precondition!, p, grid, r) δ = preconditioner.regularization rhs = preconditioner.rhs @@ -206,17 +206,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 From 3f62872b1dee461945c7afa9865619a9413ad473 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 18 Oct 2024 07:45:59 -0600 Subject: [PATCH 09/10] Change sign of regularization, plus cleaner code design --- .../conjugate_gradient_poisson_solver.jl | 33 ++++++++----------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index d2b0f53245..516b3cc6b5 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -31,7 +31,8 @@ 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} @@ -47,7 +48,10 @@ function (L::RegularizedLaplacian)(Lϕ, ϕ) if !isnothing(L.δ) # Add regularization ϕ̄ = mean(ϕ) - parent(Lϕ) .+= L.δ * ϕ̄ + ΔLϕ = L.δ * ϕ̄ + grid = ϕ.grid + arch = architecture(grid) + launch!(arch, grid, :xyz, subtract_and_mask!, Lϕ, grid, ΔLϕ) end return nothing @@ -136,15 +140,19 @@ const FFTBasedPreconditioner = RegularizedPreconditioner{<:SolverWithFFT} function precondition!(p, regularized::FFTBasedPreconditioner, r, args...) solver = regularized.preconditioner compute_preconditioner_rhs!(solver, r) - p = solve!(p, solver) + 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 / δ + Δp = mean_p + mean_rhs / δ else Δp = mean_p end @@ -175,22 +183,7 @@ Base.summary(::AsymptoticPoissonPreconditioner) = "AsymptoticPoissonPrecondition arch = architecture(p) fill_halo_regions!(r) launch!(arch, grid, :xyz, _asymptotic_poisson_precondition!, p, grid, r) - - δ = preconditioner.regularization - rhs = preconditioner.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, Δp) - + regularize_poisson_solution!(p, preconditioner) return p end From 130a667d3276518ee75d3d01d70c8b1ee4d4e6f1 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 22 Oct 2024 07:54:15 -0600 Subject: [PATCH 10/10] Regularizer --- .../conjugate_gradient_poisson_solver.jl | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index 516b3cc6b5..9ef8e117e5 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -46,7 +46,7 @@ function (L::RegularizedLaplacian)(Lϕ, ϕ) launch!(arch, grid, :xyz, laplacian!, Lϕ, grid, ϕ) if !isnothing(L.δ) - # Add regularization + # Add regularizer ϕ̄ = mean(ϕ) ΔLϕ = L.δ * ϕ̄ grid = ϕ.grid @@ -60,7 +60,7 @@ end struct DefaultPreconditioner end function ConjugateGradientPoissonSolver(grid; - regularization = nothing, + regularizer = nothing, preconditioner = DefaultPreconditioner(), reltol = sqrt(eps(grid)), abstol = sqrt(eps(grid)), @@ -75,8 +75,8 @@ function ConjugateGradientPoissonSolver(grid; end rhs = CenterField(grid) - operator = RegularizedLaplacian(regularization) - preconditioner = RegularizedPreconditioner(preconditioner, rhs, regularization) + operator = RegularizedLaplacian(regularizer) + preconditioner = RegularizedPoissonPreconditioner(preconditioner, rhs, regularizer) conjugate_gradient_solver = ConjugateGradientSolver(operator; reltol, @@ -128,17 +128,17 @@ function compute_preconditioner_rhs!(solver::FourierTridiagonalPoissonSolver, rh return nothing end -struct RegularizedPreconditioner{P, R, D} - preconditioner :: P +struct RegularizedPoissonPreconditioner{P, R, D} + unregularized_preconditioner :: P rhs :: R - regularization :: D + regularizer :: D end const SolverWithFFT = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} -const FFTBasedPreconditioner = RegularizedPreconditioner{<:SolverWithFFT} +const FFTBasedPreconditioner = RegularizedPoissonPreconditioner{<:SolverWithFFT} function precondition!(p, regularized::FFTBasedPreconditioner, r, args...) - solver = regularized.preconditioner + solver = regularized.unregularized_preconditioner compute_preconditioner_rhs!(solver, r) solve!(p, solver) regularize_poisson_solution!(p, regularized) @@ -146,13 +146,16 @@ function precondition!(p, regularized::FFTBasedPreconditioner, r, args...) end function regularize_poisson_solution!(p, regularized) - δ = regularized.regularization + δ = 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 @@ -175,10 +178,10 @@ end ##### struct AsymptoticPoissonPreconditioner end -const RegularizedNDPP = RegularizedPreconditioner{<:AsymptoticPoissonPreconditioner} +const RegularizedAPP = RegularizedPoissonPreconditioner{<:AsymptoticPoissonPreconditioner} Base.summary(::AsymptoticPoissonPreconditioner) = "AsymptoticPoissonPreconditioner" -@inline function precondition!(p, preconditioner::RegularizedNDPP, r, args...) +@inline function precondition!(p, preconditioner::RegularizedAPP, r, args...) grid = r.grid arch = architecture(p) fill_halo_regions!(r)