From 4f6ffd71933fc3dd55df8a6093fd52d420b06b1d Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Wed, 16 Oct 2024 09:49:24 -0600
Subject: [PATCH] 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,