From 6c829665e4f09eb4136b5f0a8747d660ec9d9a6a Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Sun, 29 Sep 2024 21:40:38 -0400
Subject: [PATCH 01/12] Reorganize utils and enforce gauge condition

---
 .../conjugate_gradient_poisson_solver.jl      | 75 +++++++++-----
 .../solve_for_pressure.jl                     | 97 ++++++++-----------
 src/Solvers/fft_based_poisson_solver.jl       |  3 +-
 3 files changed, 98 insertions(+), 77 deletions(-)

diff --git a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl b/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
index 3223b1e1ca..e0c257ad33 100644
--- a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
+++ b/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
@@ -80,47 +80,72 @@ function ConjugateGradientPoissonSolver(grid;
     return ConjugateGradientPoissonSolver(grid, rhs, conjugate_gradient_solver)
 end
 
-@kernel function compute_source_term!(rhs, grid, Δt, U★)
-    i, j, k = @index(Global, NTuple)
-    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
-    inactive = !inactive_cell(i, j, k, grid)
-    @inbounds rhs[i, j, k] = δ / Δt * inactive
-end
-
 function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, U★)
-    # We may want a criteria like this:
-    # min_Δt = eps(typeof(Δt))
-    # Δt <= min_Δt && return pressure
-
     rhs = solver.right_hand_side
     grid = solver.grid
     arch = architecture(grid)
-    launch!(arch, grid, :xyz, compute_source_term!, rhs, grid, Δt, U★)
-
-    # Solve pressure Pressure equation for pressure, given rhs
-    # @info "Δt before pressure solve: $(Δt)"
-    solve!(pressure, solver.conjugate_gradient_solver, rhs)
-
-    return pressure
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    return solve!(pressure, solver.conjugate_gradient_solver, rhs)
 end
 
 #####
 ##### A preconditioner based on the FFT solver
 #####
 
-@kernel function fft_preconditioner_right_hand_side!(preconditioner_rhs, rhs)
+@kernel function fft_preconditioner_rhs!(preconditioner_rhs, rhs)
     i, j, k = @index(Global, NTuple)
     @inbounds preconditioner_rhs[i, j, k] = rhs[i, j, k]
 end
 
-function precondition!(p, solver::FFTBasedPoissonSolver, rhs, args...)
+@kernel function fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::XDirection, grid, rhs)
+    i, j, k = @index(Global, NTuple)
+    @inbounds preconditioner_rhs[i, j, k] = Δxᶜᶜᶜ(i, j, k, grid) * rhs[i, j, k]
+end
+
+@kernel function fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::YDirection, grid, rhs)
+    i, j, k = @index(Global, NTuple)
+    @inbounds preconditioner_rhs[i, j, k] = Δyᶜᶜᶜ(i, j, k, grid) * rhs[i, j, k]
+end
+
+@kernel function fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::ZDirection, grid, rhs)
+    i, j, k = @index(Global, NTuple)
+    @inbounds preconditioner_rhs[i, j, k] = Δzᶜᶜᶜ(i, j, k, grid) * rhs[i, j, k]
+end
+
+function compute_preconditioner_rhs!(solver::FFTBasedPoissonSolver, rhs)
+    grid = solver.grid
+    arch = architecture(grid)
+    launch!(arch, grid, :xyz, fft_preconditioner_rhs!, solver.storage, rhs)
+    return nothing
+end
+
+function compute_preconditioner_rhs!(solver::FourierTridiagonalPoissonSolver, rhs)
+    grid = solver.grid
+    arch = architecture(grid)
+    tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, fourier_tridiagonal_preconditioner_rhs!,
+            solver.storage, tridiagonal_dir, rhs)
+    return nothing
+end
+
+function precondition!(p, solver, rhs, args...)
+    compute_preconditioner_rhs!(solver, rhs)
+    p = solve!(p, solver)
+
+    P = mean(p)
     grid = solver.grid
     arch = architecture(grid)
-    launch!(arch, grid, :xyz, fft_preconditioner_right_hand_side!, solver.storage, rhs)
-    p = solve!(p, solver, solver.storage)
+    launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, P)
+
     return p
 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
+end
+
 #####
 ##### The "DiagonallyDominantPreconditioner" used by MITgcm
 #####
@@ -133,6 +158,12 @@ Base.summary(::DiagonallyDominantPreconditioner) = "DiagonallyDominantPreconditi
     arch = architecture(p)
     fill_halo_regions!(r)
     launch!(arch, grid, :xyz, _diagonally_dominant_precondition!, p, grid, r)
+
+    P = mean(p)
+    grid = solver.grid
+    arch = architecture(grid)
+    launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, P)
+
     return p
 end
 
diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
index 5064a2c5d6..b496811994 100644
--- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl
+++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
@@ -7,88 +7,77 @@ using Oceananigans.Grids: XDirection, YDirection, ZDirection
 ##### Calculate the right-hand-side of the non-hydrostatic pressure Poisson equation.
 #####
 
-@kernel function calculate_pressure_source_term_fft_based_solver!(rhs, grid, Δt, U★)
+@kernel function _compute_source_term!(rhs, grid, Δt, U★)
     i, j, k = @index(Global, NTuple)
-    @inbounds rhs[i, j, k] = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w) / Δt
+    active = !inactive_cell(i, j, k, grid)
+    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    @inbounds rhs[i, j, k] = active * δ / Δt
 end
 
-@kernel function calculate_pressure_source_term_fourier_tridiagonal_solver!(rhs, grid, Δt, U★, ::XDirection)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::XDirection, grid, Δt, U★)
     i, j, k = @index(Global, NTuple)
-    @inbounds rhs[i, j, k] = Δxᶜᶜᶜ(i, j, k, grid) * divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w) / Δt
+    active = !inactive_cell(i, j, k, grid)
+    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    @inbounds rhs[i, j, k] = active * Δxᶜᶜᶜ(i, j, k, grid) * δ / Δt
 end
 
-@kernel function calculate_pressure_source_term_fourier_tridiagonal_solver!(rhs, grid, Δt, U★, ::YDirection)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::YDirection, grid, Δt, U★)
     i, j, k = @index(Global, NTuple)
-    @inbounds rhs[i, j, k] = Δyᶜᶜᶜ(i, j, k, grid) * divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w) / Δt
+    active = !inactive_cell(i, j, k, grid)
+    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    @inbounds rhs[i, j, k] = active * Δyᶜᶜᶜ(i, j, k, grid) * δ / Δt
 end
 
-@kernel function calculate_pressure_source_term_fourier_tridiagonal_solver!(rhs, grid, Δt, U★, ::ZDirection)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::ZDirection, grid, Δt, U★)
     i, j, k = @index(Global, NTuple)
-    @inbounds rhs[i, j, k] = Δzᶜᶜᶜ(i, j, k, grid) * divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w) / Δt
+    active = !inactive_cell(i, j, k, grid)
+    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    @inbounds rhs[i, j, k] = active * Δzᶜᶜᶜ(i, j, k, grid) * δ / Δt
 end
 
-#####
-##### Solve for pressure
-#####
-
-function solve_for_pressure!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, U★)
+function compute_source_term!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, U★)
     rhs  = solver.storage.zfield
     arch = architecture(solver)
     grid = solver.local_grid
-
-    launch!(arch, grid, :xyz, calculate_pressure_source_term_fft_based_solver!,
-            rhs, grid, Δt, U★)
-
-    # Solve pressure Poisson equation for pressure, given rhs
-    solve!(pressure, solver)
-
-    return pressure
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    return nothing        
 end
 
-function solve_for_pressure!(pressure, solver::FFTBasedPoissonSolver, Δt, U★)
-
-    # Calculate right hand side:
-    rhs = solver.storage
+function compute_source_term!(pressure, solver::DistributedFourierTridiagonalPoissonSolver, Δt, U★)
+    rhs = solver.storage.zfield
     arch = architecture(solver)
-    grid = solver.grid
-
-    launch!(arch, grid, :xyz, calculate_pressure_source_term_fft_based_solver!,
-            rhs, grid, Δt, U★)
-
-    # Solve pressure Poisson given for pressure, given rhs
-    solve!(pressure, solver, rhs)
-
+    grid = solver.local_grid
+    tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!,
+            rhs, grid, Δt, U★, tridiagonal_dir)
     return nothing
 end
 
-function solve_for_pressure!(pressure, solver::DistributedFourierTridiagonalPoissonSolver, Δt, U★)
-
-    # Calculate right hand side:
-    rhs = solver.storage.zfield
+function compute_source_term!(pressure, solver::FourierTridiagonalPoissonSolver, Δt, U★)
+    rhs = solver.source_term
     arch = architecture(solver)
-    grid = solver.local_grid
-
-    launch!(arch, grid, :xyz, calculate_pressure_source_term_fourier_tridiagonal_solver!,
-            rhs, grid, Δt, U★, solver.batched_tridiagonal_solver.tridiagonal_direction)
-
-    # Pressure Poisson rhs, scaled by the spacing in the stretched direction at ᶜᶜᶜ, is stored in solver.source_term:
-    solve!(pressure, solver)
-
+    grid = solver.grid
+    tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!,
+            rhs, grid, Δt, U★, tridiagonal_dir)
     return nothing
 end
 
-function solve_for_pressure!(pressure, solver::FourierTridiagonalPoissonSolver, Δt, U★)
-
-    # Calculate right hand side:
-    rhs = solver.source_term
+function compute_source_term!(pressure, solver::FFTBasedPoissonSolver, Δt, U★)
+    rhs = solver.storage
     arch = architecture(solver)
     grid = solver.grid
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    return nothing
+end
 
-    launch!(arch, grid, :xyz, calculate_pressure_source_term_fourier_tridiagonal_solver!,
-            rhs, grid, Δt, U★, solver.batched_tridiagonal_solver.tridiagonal_direction)
+#####
+##### Solve for pressure
+#####
 
-    # Pressure Poisson rhs, scaled by the spacing in the stretched direction at ᶜᶜᶜ, is stored in solver.source_term:
+function solve_for_pressure!(pressure, solver, Δt, U★)
+    compute_source_term!(pressure, solver, Δt, U★)
     solve!(pressure, solver)
-
-    return nothing
+    return pressure
 end
+
diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl
index 81f8c1633d..25d1ef30cd 100644
--- a/src/Solvers/fft_based_poisson_solver.jl
+++ b/src/Solvers/fft_based_poisson_solver.jl
@@ -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, m=0)
+function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0)
     arch = architecture(solver)
     topo = TX, TY, TZ = topology(solver.grid)
     Nx, Ny, Nz = size(solver.grid)
@@ -131,3 +131,4 @@ end
 
     @inbounds ϕ[i′, j′, k′] = real(ϕc[i, j, k])
 end
+

From 0474a683b48244350cc10a79805fe09d938d9918 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 00:03:35 -0400
Subject: [PATCH 02/12] Fix formatting issue

---
 src/BoundaryConditions/field_boundary_conditions.jl |  9 +++++----
 src/Solvers/fft_based_poisson_solver.jl             |  8 ++++++--
 src/Solvers/fourier_tridiagonal_poisson_solver.jl   | 11 +++++++----
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl
index cbf0ef4beb..a42beeee2c 100644
--- a/src/BoundaryConditions/field_boundary_conditions.jl
+++ b/src/BoundaryConditions/field_boundary_conditions.jl
@@ -156,11 +156,12 @@ end
 # Friendly warning?
 function regularize_immersed_boundary_condition(ibc, grid, loc, field_name, args...)
     if !(ibc isa DefaultBoundaryCondition)
-        msg = """
-              $field_name was assigned an immersed $ibc, but this is not supported on
-              $(summary(grid))
+        msg = """$field_name was assigned an immersed boundary condition
+              $ibc ,
+              but this is not supported on
+              $(summary(grid)) .
               The immersed boundary condition on $field_name will have no effect.
-          """
+              """
 
         @warn msg
     end
diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl
index 25d1ef30cd..3f5fae67e3 100644
--- a/src/Solvers/fft_based_poisson_solver.jl
+++ b/src/Solvers/fft_based_poisson_solver.jl
@@ -102,7 +102,9 @@ function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0)
     ϕc = solver.storage
 
     # Transform b *in-place* to eigenfunction space
-    [transform!(b, solver.buffer) for transform! in solver.transforms.forward]
+    for transform! in solver.transforms.forward
+        transform!(b, solver.buffer)
+    end
 
     # Solve the discrete screened Poisson equation (∇² + m) ϕ = b.
     @. ϕc = - b / (λx + λy + λz - m)
@@ -113,7 +115,9 @@ function solve!(ϕ, solver::FFTBasedPoissonSolver, b=solver.storage, m=0)
     m === 0 && CUDA.@allowscalar ϕc[1, 1, 1] = 0
 
     # Apply backward transforms in order
-    [transform!(ϕc, solver.buffer) for transform! in solver.transforms.backward]
+    for transform! in solver.transforms.backward
+        transform!(ϕc, solver.buffer)
+    end
 
     launch!(arch, solver.grid, :xyz, copy_real_component!, ϕ, ϕc, indices(ϕ))
     
diff --git a/src/Solvers/fourier_tridiagonal_poisson_solver.jl b/src/Solvers/fourier_tridiagonal_poisson_solver.jl
index 648cca96b7..665489e120 100644
--- a/src/Solvers/fourier_tridiagonal_poisson_solver.jl
+++ b/src/Solvers/fourier_tridiagonal_poisson_solver.jl
@@ -119,13 +119,17 @@ function solve!(x, solver::FourierTridiagonalPoissonSolver, b=nothing)
     ϕ = solver.storage
 
     # Apply forward transforms in order
-    [transform!(solver.source_term, solver.buffer) for transform! in solver.transforms.forward]
+    for transform! in solver.transforms.forward
+        transform!(solver.source_term, solver.buffer)
+    end
 
     # Solve tridiagonal system of linear equations at every column.
     solve!(ϕ, solver.batched_tridiagonal_solver, solver.source_term)
 
     # Apply backward transforms in order
-    [transform!(ϕ, solver.buffer) for transform! in solver.transforms.backward]
+    for transform! in solver.transforms.backward
+        transform!(ϕ, solver.buffer)
+    end
 
     # Set the volume mean of the solution to be zero.
     # Solutions to Poisson's equation are only unique up to a constant (the global mean
@@ -148,9 +152,7 @@ function set_source_term!(solver::FourierTridiagonalPoissonSolver, source_term)
     grid = solver.grid
     arch = architecture(solver)
     solver.source_term .= source_term
-
     launch!(arch, grid, :xyz, multiply_by_stretched_spacing!, solver.source_term, grid)
-
     return nothing
 end
 
@@ -169,3 +171,4 @@ end
     i, j, k = @index(Global, NTuple)
     @inbounds a[i, j, k] *= Δzᵃᵃᶜ(i, j, k, grid)
 end
+

From 2be8ed2e52f96514ad4036f75c6b163f8b3518f7 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 00:03:57 -0400
Subject: [PATCH 03/12] Some small improvements

---
 .../conjugate_gradient_poisson_solver.jl      | 45 +++++++++----------
 1 file changed, 22 insertions(+), 23 deletions(-)

diff --git a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl b/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
index e0c257ad33..d230d29d39 100644
--- a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
+++ b/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
@@ -49,13 +49,15 @@ function compute_laplacian!(∇²ϕ, ϕ)
     return nothing
 end
 
+struct DefaultPreconditioner end
+
 function ConjugateGradientPoissonSolver(grid;
-                                        preconditioner = nothing,
+                                        preconditioner = DefaultPreconditioner(),
                                         reltol = sqrt(eps(grid)),
                                         abstol = 0,
                                         kw...)
 
-    if isnothing(preconditioner) # make a useful default
+    if preconditioner isa DefaultPreconditioner # try to make a useful default
         if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFT
             if grid.underlying_grid isa XYZRegularRG
                 preconditioner = FFTBasedPoissonSolver(grid.underlying_grid)
@@ -128,14 +130,16 @@ function compute_preconditioner_rhs!(solver::FourierTridiagonalPoissonSolver, rh
     return nothing
 end
 
-function precondition!(p, solver, rhs, args...)
-    compute_preconditioner_rhs!(solver, rhs)
-    p = solve!(p, solver)
+const FFTBasedPreconditioner = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver}
 
-    P = mean(p)
-    grid = solver.grid
+function precondition!(p, preconditioner::FFTBasedPreconditioner, r, args...)
+    compute_preconditioner_rhs!(preconditioner, r)
+    p = solve!(p, preconditioner)
+
+    mean_p = mean(p)
+    grid = p.grid
     arch = architecture(grid)
-    launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, P)
+    launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, mean_p)
 
     return p
 end
@@ -147,7 +151,7 @@ end
 end
 
 #####
-##### The "DiagonallyDominantPreconditioner" used by MITgcm
+##### The "DiagonallyDominantPreconditioner" (Marshall et al 1997)
 #####
 
 struct DiagonallyDominantPreconditioner end
@@ -159,10 +163,8 @@ Base.summary(::DiagonallyDominantPreconditioner) = "DiagonallyDominantPreconditi
     fill_halo_regions!(r)
     launch!(arch, grid, :xyz, _diagonally_dominant_precondition!, p, grid, r)
 
-    P = mean(p)
-    grid = solver.grid
-    arch = architecture(grid)
-    launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, P)
+    mean_p = mean(p)
+    launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, mean_p)
 
     return p
 end
@@ -170,20 +172,15 @@ end
 # Kernels that calculate coefficients for the preconditioner
 @inline Ax⁻(i, j, k, grid) = Axᶠᶜᶜ(i,   j, k, grid) / Δxᶠᶜᶜ(i,   j, k, grid) / Vᶜᶜᶜ(i, j, k, grid)
 @inline Ax⁺(i, j, k, grid) = Axᶠᶜᶜ(i+1, j, k, grid) / Δxᶠᶜᶜ(i+1, j, k, grid) / Vᶜᶜᶜ(i, j, k, grid)
-
 @inline Ay⁻(i, j, k, grid) = Ayᶜᶠᶜ(i, j,   k, grid) / Δyᶜᶠᶜ(i, j,   k, grid) / Vᶜᶜᶜ(i, j, k, grid)
 @inline Ay⁺(i, j, k, grid) = Ayᶜᶠᶜ(i, j+1, k, grid) / Δyᶜᶠᶜ(i, j+1, k, grid) / Vᶜᶜᶜ(i, j, k, grid)
-
 @inline Az⁻(i, j, k, grid) = Azᶜᶜᶠ(i, j, k,   grid) / Δzᶜᶜᶠ(i, j, k,   grid) / Vᶜᶜᶜ(i, j, k, grid)
 @inline Az⁺(i, j, k, grid) = Azᶜᶜᶠ(i, j, k+1, grid) / Δzᶜᶜᶠ(i, j, k+1, grid) / Vᶜᶜᶜ(i, j, k, grid)
 
-@inline Ac(i, j, k, grid) = - Ax⁻(i, j, k, grid) -
-                              Ax⁺(i, j, k, grid) -
-                              Ay⁻(i, j, k, grid) -
-                              Ay⁺(i, j, k, grid) -
-                              Az⁻(i, j, k, grid) -
-                              Az⁺(i, j, k, grid)
-
+@inline Ac(i, j, k, grid) = - Ax⁻(i, j, k, grid) - Ax⁺(i, j, k, grid) -
+                              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] -
@@ -194,5 +191,7 @@ end
 
 @kernel function _diagonally_dominant_precondition!(p, grid, r)
     i, j, k = @index(Global, NTuple)
-    @inbounds p[i, j, k] = heuristic_residual(i, j, k, grid, r)
+    active = !inactive_cell(i, j, k, grid)
+    @inbounds p[i, j, k] = heuristic_residual(i, j, k, grid, r) * active
 end
+

From 43420e607702eb4f1005678ae5f87047391dfadf Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 09:07:27 -0400
Subject: [PATCH 04/12] PreconditionedCGSolver -> CGSolver

---
 .../implicit_free_surface.jl                  |  4 +-
 .../pcg_implicit_free_surface_solver.jl       |  2 +-
 .../conjugate_gradient_poisson_solver.jl      |  4 +-
 src/MultiRegion/multi_region_models.jl        |  4 +-
 src/Solvers/Solvers.jl                        |  2 +-
 ...reconditioned_conjugate_gradient_solver.jl | 38 +++++++++----------
 ...reconditioned_conjugate_gradient_solver.jl |  8 ++--
 7 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl
index cbef7bf54a..e54ad5c49c 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl
@@ -54,7 +54,7 @@ where
 𝐮_⋆ = 𝐮^n + \\int_{t_n}^{t_{n+1}} 𝐆ᵤ \\, 𝖽t .
 ```
 
-This equation can be solved, in general, using the [`PreconditionedConjugateGradientSolver`](@ref) but 
+This equation can be solved, in general, using the [`ConjugateGradientSolver`](@ref) but 
 other solvers can be invoked in special cases.
 
 If ``H`` is constant, we divide through out to obtain
@@ -69,7 +69,7 @@ surface can be obtained using the [`FFTBasedPoissonSolver`](@ref).
 `solver_method` can be either of:
 * `:FastFourierTransform` for [`FFTBasedPoissonSolver`](@ref)
 * `:HeptadiagonalIterativeSolver`  for [`HeptadiagonalIterativeSolver`](@ref)
-* `:PreconditionedConjugateGradient` for [`PreconditionedConjugateGradientSolver`](@ref)
+* `:PreconditionedConjugateGradient` for [`ConjugateGradientSolver`](@ref)
 
 By default, if the grid has regular spacing in the horizontal directions then the `:FastFourierTransform` is chosen,
 otherwise the `:HeptadiagonalIterativeSolver`.
diff --git a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl
index ddc7c4ab8d..a6a05a7032 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl
@@ -65,7 +65,7 @@ function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, settings, gravitationa
     # TODO: reuse solver.storage for rhs when preconditioner isa FFTImplicitFreeSurfaceSolver?
     right_hand_side = ZFaceField(grid, indices = (:, :, size(grid, 3) + 1))
 
-    solver = PreconditionedConjugateGradientSolver(implicit_free_surface_linear_operation!;
+    solver = ConjugateGradientSolver(implicit_free_surface_linear_operation!;
                                                    template_field = right_hand_side,
                                                    settings...)
 
diff --git a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl b/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
index d230d29d39..a350921def 100644
--- a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
+++ b/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
@@ -1,5 +1,5 @@
 using Oceananigans.Architectures: device, architecture
-using Oceananigans.Solvers: PreconditionedConjugateGradientSolver, FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver, solve!
+using Oceananigans.Solvers: ConjugateGradientSolver, FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver, solve!
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Grids: inactive_cell
 using Oceananigans.Operators: divᶜᶜᶜ, ∇²ᶜᶜᶜ 
@@ -72,7 +72,7 @@ function ConjugateGradientPoissonSolver(grid;
     rhs = CenterField(grid)
 
     conjugate_gradient_solver =
-        PreconditionedConjugateGradientSolver(compute_laplacian!;
+        ConjugateGradientSolver(compute_laplacian!;
                                               reltol,
                                               abstol,
                                               preconditioner,
diff --git a/src/MultiRegion/multi_region_models.jl b/src/MultiRegion/multi_region_models.jl
index 3dabe07602..9d835342f7 100644
--- a/src/MultiRegion/multi_region_models.jl
+++ b/src/MultiRegion/multi_region_models.jl
@@ -7,7 +7,7 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization
 using Oceananigans.Advection: AbstractAdvectionScheme
 using Oceananigans.Advection: OnlySelfUpwinding, CrossAndSelfUpwinding
 using Oceananigans.ImmersedBoundaries: GridFittedBottom, PartialCellBottom, GridFittedBoundary
-using Oceananigans.Solvers: PreconditionedConjugateGradientSolver
+using Oceananigans.Solvers: ConjugateGradientSolver
 
 import Oceananigans.Advection: WENO, cell_advection_timescale
 import Oceananigans.Models.HydrostaticFreeSurfaceModels: build_implicit_step_solver, validate_tracer_advection
@@ -34,7 +34,7 @@ Types = (:HydrostaticFreeSurfaceModel,
          :SplitExplicitState,
          :SplitExplicitFreeSurface,
          :PrescribedVelocityFields,
-         :PreconditionedConjugateGradientSolver,
+         :ConjugateGradientSolver,
          :CrossAndSelfUpwinding,
          :OnlySelfUpwinding,
          :GridFittedBoundary,
diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl
index c5a5bf13ce..02a612fc81 100644
--- a/src/Solvers/Solvers.jl
+++ b/src/Solvers/Solvers.jl
@@ -4,7 +4,7 @@ export
     BatchedTridiagonalSolver, solve!,
     FFTBasedPoissonSolver,
     FourierTridiagonalPoissonSolver,
-    PreconditionedConjugateGradientSolver,
+    ConjugateGradientSolver,
     HeptadiagonalIterativeSolver
 
 using Statistics
diff --git a/src/Solvers/preconditioned_conjugate_gradient_solver.jl b/src/Solvers/preconditioned_conjugate_gradient_solver.jl
index b856f57792..d78e04d750 100644
--- a/src/Solvers/preconditioned_conjugate_gradient_solver.jl
+++ b/src/Solvers/preconditioned_conjugate_gradient_solver.jl
@@ -6,7 +6,7 @@ using LinearAlgebra
 
 import Oceananigans.Architectures: architecture
 
-mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P} 
+mutable struct ConjugateGradientSolver{A, G, L, T, F, M, P} 
     architecture :: A
     grid :: G
     linear_operation! :: L
@@ -22,25 +22,25 @@ mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P}
     preconditioner_product :: P
 end
 
-architecture(solver::PreconditionedConjugateGradientSolver) = solver.architecture
+architecture(solver::ConjugateGradientSolver) = solver.architecture
 
 initialize_precondition_product(preconditioner, template_field) = similar(template_field)
 initialize_precondition_product(::Nothing, template_field) = nothing
 
-Base.summary(::PreconditionedConjugateGradientSolver) = "PreconditionedConjugateGradientSolver"
+Base.summary(::ConjugateGradientSolver) = "ConjugateGradientSolver"
 
 # "Nothing" preconditioner
 @inline precondition!(z, ::Nothing, r, args...) = r
 
 """
-    PreconditionedConjugateGradientSolver(linear_operation;
+    ConjugateGradientSolver(linear_operation;
                                           template_field,
                                           maxiter = size(template_field.grid),
                                           reltol = sqrt(eps(eltype(template_field.grid))),
                                           abstol = 0,
                                           preconditioner = nothing)
 
-Returns a `PreconditionedConjugateGradientSolver` that solves the linear equation
+Returns a `ConjugateGradientSolver` that solves the linear equation
 ``A x = b`` using a iterative conjugate gradient method with optional preconditioning.
 
 The solver is used by calling
@@ -72,7 +72,7 @@ Arguments
 
 See [`solve!`](@ref) for more information about the preconditioned conjugate-gradient algorithm.
 """
-function PreconditionedConjugateGradientSolver(linear_operation;
+function ConjugateGradientSolver(linear_operation;
                                                template_field::AbstractField,
                                                maxiter = prod(size(template_field)),
                                                reltol = sqrt(eps(eltype(template_field.grid))),
@@ -92,7 +92,7 @@ function PreconditionedConjugateGradientSolver(linear_operation;
 
     FT = eltype(grid)
 
-    return PreconditionedConjugateGradientSolver(arch,
+    return ConjugateGradientSolver(arch,
                                                  grid,
                                                  linear_operation,
                                                  FT(reltol),
@@ -108,7 +108,7 @@ function PreconditionedConjugateGradientSolver(linear_operation;
 end
 
 """
-    solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)
+    solve!(x, solver::ConjugateGradientSolver, b, args...)
 
 Solve `A * x = b` using an iterative conjugate-gradient method, where `A * x` is
 determined by `solver.linear_operation`
@@ -156,7 +156,7 @@ Loop:
      ρⁱ⁻¹ = ρ
 ```
 """
-function solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)
+function solve!(x, solver::ConjugateGradientSolver, b, args...)
 
     # Initialize
     solver.iteration = 0
@@ -169,8 +169,8 @@ function solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)
     residual_norm = norm(solver.residual)
     tolerance = max(solver.reltol * residual_norm, solver.abstol)
 
-    @debug "PreconditionedConjugateGradientSolver, |b|: $(norm(b))"
-    @debug "PreconditionedConjugateGradientSolver, |A * x|: $(norm(q))"
+    @debug "ConjugateGradientSolver, |b|: $(norm(b))"
+    @debug "ConjugateGradientSolver, |A * x|: $(norm(q))"
 
     while iterating(solver, tolerance)
         iterate!(x, solver, b, args...)
@@ -184,7 +184,7 @@ function iterate!(x, solver, b, args...)
     p = solver.search_direction
     q = solver.linear_operator_product
 
-    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |r|: $(norm(r))"
+    @debug "ConjugateGradientSolver $(solver.iteration), |r|: $(norm(r))"
 
     # Preconditioned:   z = P * r
     # Unpreconditioned: z = r
@@ -192,15 +192,15 @@ function iterate!(x, solver, b, args...)
 
     ρ = dot(z, r)
 
-    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), ρ: $ρ"
-    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |z|: $(norm(z))"
+    @debug "ConjugateGradientSolver $(solver.iteration), ρ: $ρ"
+    @debug "ConjugateGradientSolver $(solver.iteration), |z|: $(norm(z))"
 
     @apply_regionally perform_iteration!(q, p, ρ, z, solver, args...)
 
     α = ρ / dot(p, q)
 
-    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |q|: $(norm(q))"
-    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), α: $α"
+    @debug "ConjugateGradientSolver $(solver.iteration), |q|: $(norm(q))"
+    @debug "ConjugateGradientSolver $(solver.iteration), α: $α"
         
     @apply_regionally update_solution_and_residuals!(x, r, q, p, α)
 
@@ -230,7 +230,7 @@ function perform_iteration!(q, p, ρ, z, solver, args...)
         β = ρ / solver.ρⁱ⁻¹
         pp .= zp .+ β .* pp
 
-        @debug "PreconditionedConjugateGradientSolver $(solver.iteration), β: $β"
+        @debug "ConjugateGradientSolver $(solver.iteration), β: $β"
     end
 
     # q = A * p
@@ -258,8 +258,8 @@ function iterating(solver, tolerance)
     return true
 end
 
-function Base.show(io::IO, solver::PreconditionedConjugateGradientSolver)
-    print(io, "PreconditionedConjugateGradientSolver on ", summary(solver.architecture), "\n",
+function Base.show(io::IO, solver::ConjugateGradientSolver)
+    print(io, "ConjugateGradientSolver on ", summary(solver.architecture), "\n",
               "├── template_field: ", summary(solver.residual), "\n",
               "├── grid: ", summary(solver.grid), "\n",
               "├── linear_operation!: ", prettysummary(solver.linear_operation!), "\n",
diff --git a/test/test_preconditioned_conjugate_gradient_solver.jl b/test/test_preconditioned_conjugate_gradient_solver.jl
index 771c7f50c2..4f4eb0e4c6 100644
--- a/test/test_preconditioned_conjugate_gradient_solver.jl
+++ b/test/test_preconditioned_conjugate_gradient_solver.jl
@@ -10,7 +10,7 @@ end
 
 function run_identity_operator_test(grid)
     b = CenterField(grid)
-    solver = PreconditionedConjugateGradientSolver(identity_operator!, template_field = b, reltol=0, abstol=10*sqrt(eps(eltype(grid))))
+    solver = ConjugateGradientSolver(identity_operator!, template_field = b, reltol=0, abstol=10*sqrt(eps(eltype(grid))))
     initial_guess = solution = similar(b)
     set!(initial_guess, (x, y, z) -> rand())
 
@@ -33,7 +33,7 @@ function run_poisson_equation_test(grid)
     ∇²ϕ = r = CenterField(grid)
     compute_∇²!(∇²ϕ, ϕ_truth, arch, grid)
 
-    solver = PreconditionedConjugateGradientSolver(compute_∇²!, template_field=ϕ_truth, reltol=eps(eltype(grid)), maxiter=Int(1e10))
+    solver = ConjugateGradientSolver(compute_∇²!, template_field=ϕ_truth, reltol=eps(eltype(grid)), maxiter=Int(1e10))
 
     # Solve Poisson equation
     ϕ_solution = CenterField(grid)
@@ -61,9 +61,9 @@ function run_poisson_equation_test(grid)
     return nothing
 end
 
-@testset "PreconditionedConjugateGradientSolver" begin
+@testset "ConjugateGradientSolver" begin
     for arch in archs
-        @info "Testing PreconditionedConjugateGradientSolver [$(typeof(arch))]..."
+        @info "Testing ConjugateGradientSolver [$(typeof(arch))]..."
         grid = RectilinearGrid(arch, size=(4, 8, 4), extent=(1, 3, 1))
         run_identity_operator_test(grid)
         run_poisson_equation_test(grid)

From 6f971ad6f67028b29f8813c04ce4ab5820fc18db Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 09:08:32 -0400
Subject: [PATCH 05/12] Change to conjugate gradient solver

---
 src/Solvers/Solvers.jl                                          | 2 +-
 ...onjugate_gradient_solver.jl => conjugate_gradient_solver.jl} | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename src/Solvers/{preconditioned_conjugate_gradient_solver.jl => conjugate_gradient_solver.jl} (100%)

diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl
index 02a612fc81..9c336af92a 100644
--- a/src/Solvers/Solvers.jl
+++ b/src/Solvers/Solvers.jl
@@ -33,13 +33,13 @@ reshaped_size(N, dim) = dim == 1 ? (N, 1, 1) :
                         dim == 3 ? (1, 1, N) : nothing
 
 include("batched_tridiagonal_solver.jl")
+include("conjugate_gradient_solver.jl")
 include("poisson_eigenvalues.jl")
 include("index_permutations.jl")
 include("discrete_transforms.jl")
 include("plan_transforms.jl")
 include("fft_based_poisson_solver.jl")
 include("fourier_tridiagonal_poisson_solver.jl")
-include("preconditioned_conjugate_gradient_solver.jl")
 include("sparse_approximate_inverse.jl")
 include("matrix_solver_utils.jl")
 include("sparse_preconditioners.jl")
diff --git a/src/Solvers/preconditioned_conjugate_gradient_solver.jl b/src/Solvers/conjugate_gradient_solver.jl
similarity index 100%
rename from src/Solvers/preconditioned_conjugate_gradient_solver.jl
rename to src/Solvers/conjugate_gradient_solver.jl

From ceb2f9da682bb6f9c9d00e756638bd77a2eb9faa Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 09:55:37 -0400
Subject: [PATCH 06/12] Move ConjugateGradientPoissonSolver to Solvers

---
 .../NonhydrostaticModels.jl                   | 27 +++++++------------
 .../solve_for_pressure.jl                     | 12 ++++++++-
 src/Solvers/Solvers.jl                        |  8 ++++++
 .../conjugate_gradient_poisson_solver.jl      | 14 +++-------
 4 files changed, 32 insertions(+), 29 deletions(-)
 rename src/{Models/NonhydrostaticModels => Solvers}/conjugate_gradient_poisson_solver.jl (94%)

diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
index f7104922be..544ab96469 100644
--- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
+++ b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
@@ -15,6 +15,7 @@ using Oceananigans.DistributedComputations: reconstruct_global_grid, Distributed
 using Oceananigans.DistributedComputations: DistributedFFTBasedPoissonSolver, DistributedFourierTridiagonalPoissonSolver
 using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
+using Oceananigans.Solvers: GridWithFFTSolver, GridWithFourierTridiagonalSolver 
 using Oceananigans.Utils: SumOfArrays
 
 import Oceananigans: fields, prognostic_fields
@@ -26,35 +27,24 @@ function nonhydrostatic_pressure_solver(::Distributed, local_grid::XYZRegularRG)
     return DistributedFFTBasedPoissonSolver(global_grid, local_grid)
 end
 
-function nonhydrostatic_pressure_solver(::Distributed, local_grid::XYRegularRG)
-    global_grid = reconstruct_global_grid(local_grid)
-    return DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid)
-end
-
-function nonhydrostatic_pressure_solver(::Distributed, local_grid::YZRegularRG)
-    global_grid = reconstruct_global_grid(local_grid)
-    return DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid)
-end
-
-function nonhydrostatic_pressure_solver(::Distributed, local_grid::XZRegularRG)
+function nonhydrostatic_pressure_solver(::Distributed, local_grid::GridWithFourierTridiagonalSolver)
     global_grid = reconstruct_global_grid(local_grid)
     return DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid)
 end
 
 nonhydrostatic_pressure_solver(arch, grid::XYZRegularRG) = FFTBasedPoissonSolver(grid)
-nonhydrostatic_pressure_solver(arch, grid::XYRegularRG)  = FourierTridiagonalPoissonSolver(grid)
-nonhydrostatic_pressure_solver(arch, grid::XZRegularRG)  = FourierTridiagonalPoissonSolver(grid)
-nonhydrostatic_pressure_solver(arch, grid::YZRegularRG)  = FourierTridiagonalPoissonSolver(grid)
+nonhydrostatic_pressure_solver(arch, grid::GridWithFourierTridiagonalSolver) =
+    FourierTridiagonalPoissonSolver(grid)
 
-const GridWithFFT = Union{XYZRegularRG, XYRegularRG, XZRegularRG, YZRegularRG}
+const IBGWithFFTSolver = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:GridWithFFTSolver}
 
-function nonhydrostatic_pressure_solver(arch, ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:GridWithFFT})
-    msg = """The FFT-based pressure_solver for `NonhydrostaticModel`s on `ImmersedBoundaryGrid`
+function nonhydrostatic_pressure_solver(arch, ibg::IBGWithFFTSolver)
+    msg = """The FFT-based pressure_solver for NonhydrostaticModels on ImmersedBoundaryGrid
           is approximate and will probably produce velocity fields that are divergent
           adjacent to the immersed boundary. An experimental but improved pressure_solver
           is available which may be used by writing
 
-              using Oceananigans.Models.NonhydrostaticModels: ConjugateGradientPoissonSolver
+              using Oceananigans.Solvers: ConjugateGradientPoissonSolver
               pressure_solver = ConjugateGradientPoissonSolver(grid)
 
           Please report issues to https://github.com/CliMA/Oceananigans.jl/issues.
@@ -121,3 +111,4 @@ include("compute_nonhydrostatic_tendencies.jl")
 include("compute_nonhydrostatic_boundary_tendencies.jl")
 
 end # module
+
diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
index b496811994..01b3f25719 100644
--- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl
+++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
@@ -1,7 +1,9 @@
 using Oceananigans.Operators
-using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver, solve!
 using Oceananigans.DistributedComputations: DistributedFFTBasedPoissonSolver
 using Oceananigans.Grids: XDirection, YDirection, ZDirection
+using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver
+using Oceananigans.Solvers: ConjugateGradientPoissonSolver
+using Oceananigans.Solvers: solve!
 
 #####
 ##### Calculate the right-hand-side of the non-hydrostatic pressure Poisson equation.
@@ -81,3 +83,11 @@ function solve_for_pressure!(pressure, solver, Δt, U★)
     return pressure
 end
 
+function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, U★)
+    rhs = solver.right_hand_side
+    grid = solver.grid
+    arch = architecture(grid)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    return solve!(pressure, solver.conjugate_gradient_solver, rhs)
+end
+
diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl
index 9c336af92a..bc3d46be74 100644
--- a/src/Solvers/Solvers.jl
+++ b/src/Solvers/Solvers.jl
@@ -20,6 +20,7 @@ using Oceananigans.BoundaryConditions
 using Oceananigans.Fields
 
 using Oceananigans.Grids: unpack_grid
+using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 
 """
     ω(M, k)
@@ -45,4 +46,11 @@ include("matrix_solver_utils.jl")
 include("sparse_preconditioners.jl")
 include("heptadiagonal_iterative_solver.jl")
 
+const GridWithFFTSolver = Union{XYZRegularRG, XYRegularRG, XZRegularRG, YZRegularRG}
+const GridWithFourierTridiagonalSolver = Union{XYRegularRG, XZRegularRG, YZRegularRG}
+
+fft_poisson_solver(grid::XYZRegularRG) = FFTBasedPoissonSolver(grid)
+fft_poisson_solver(grid::GridWithFourierTridiagonalSolver) =
+    FourierTridiagonalPoissonSolver(grid.underlying_grid)
+
 end # module
diff --git a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl
similarity index 94%
rename from src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
rename to src/Solvers/conjugate_gradient_poisson_solver.jl
index a350921def..993d3c15d6 100644
--- a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl
+++ b/src/Solvers/conjugate_gradient_poisson_solver.jl
@@ -1,5 +1,7 @@
 using Oceananigans.Architectures: device, architecture
-using Oceananigans.Solvers: ConjugateGradientSolver, FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver, solve!
+using Oceananigans.Solvers: ConjugateGradientPoissonSolver
+using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver
+using Oceananigans.Solvers: fft_poisson_solver, solve!
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Grids: inactive_cell
 using Oceananigans.Operators: divᶜᶜᶜ, ∇²ᶜᶜᶜ 
@@ -58,7 +60,7 @@ function ConjugateGradientPoissonSolver(grid;
                                         kw...)
 
     if preconditioner isa DefaultPreconditioner # try to make a useful default
-        if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFT
+        if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFTSolver
             if grid.underlying_grid isa XYZRegularRG
                 preconditioner = FFTBasedPoissonSolver(grid.underlying_grid)
             else # it's stretched in one direction
@@ -82,14 +84,6 @@ function ConjugateGradientPoissonSolver(grid;
     return ConjugateGradientPoissonSolver(grid, rhs, conjugate_gradient_solver)
 end
 
-function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, U★)
-    rhs = solver.right_hand_side
-    grid = solver.grid
-    arch = architecture(grid)
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
-    return solve!(pressure, solver.conjugate_gradient_solver, rhs)
-end
-
 #####
 ##### A preconditioner based on the FFT solver
 #####

From 0df16b78390b02fcced5d156755d2fc4f69610f2 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 13:46:04 -0400
Subject: [PATCH 07/12] More improvments

---
 src/Models/Models.jl                          |  1 +
 .../conjugate_gradient_poisson_solver.jl      | 29 ++++++++-----------
 src/Solvers/conjugate_gradient_solver.jl      |  3 +-
 3 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/src/Models/Models.jl b/src/Models/Models.jl
index fa46608b96..9399a5b503 100644
--- a/src/Models/Models.jl
+++ b/src/Models/Models.jl
@@ -21,6 +21,7 @@ using Oceananigans.Utils: Time
 import Oceananigans: initialize!
 import Oceananigans.Architectures: architecture
 import Oceananigans.TimeSteppers: reset!
+import Oceananigans.Solvers: iteration
 
 # A prototype interface for AbstractModel.
 #
diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl
index 993d3c15d6..93f7580acd 100644
--- a/src/Solvers/conjugate_gradient_poisson_solver.jl
+++ b/src/Solvers/conjugate_gradient_poisson_solver.jl
@@ -11,8 +11,8 @@ using Statistics: mean
 
 using KernelAbstractions: @kernel, @index
 
+import Oceananigans.Architecture: architecture
 import Oceananigans.Solvers: precondition!
-import ..Models: iteration
 
 struct ConjugateGradientPoissonSolver{G, R, S}
     grid :: G
@@ -20,7 +20,8 @@ struct ConjugateGradientPoissonSolver{G, R, S}
     conjugate_gradient_solver :: S
 end
 
-iteration(cgps::ConjugateGradientPoissonSolver) = cgps.conjugate_gradient_solver.iteration
+architecture(solver::ConjugateGradientPoissonSolver) = architecture(cgps.grid)
+iteration(cgps::ConjugateGradientPoissonSolver) = iteration(cgps.conjugate_gradient_solver)
 
 Base.summary(ips::ConjugateGradientPoissonSolver) =
     summary("ConjugateGradientPoissonSolver on ", summary(ips.grid))
@@ -29,7 +30,6 @@ function Base.show(io::IO, ips::ConjugateGradientPoissonSolver)
     A = architecture(ips.grid)
     print(io, "ConjugateGradientPoissonSolver:", '\n',
               "├── grid: ", summary(ips.grid), '\n',
-              "│   └── immersed_boundary: ", prettysummary(ips.grid.immersed_boundary), '\n',
               "└── conjugate_gradient_solver: ", summary(ips.conjugate_gradient_solver), '\n',
               "    ├── maxiter: ", prettysummary(ips.conjugate_gradient_solver.maxiter), '\n',
               "    ├── reltol: ", prettysummary(ips.conjugate_gradient_solver.reltol), '\n',
@@ -56,16 +56,12 @@ struct DefaultPreconditioner end
 function ConjugateGradientPoissonSolver(grid;
                                         preconditioner = DefaultPreconditioner(),
                                         reltol = sqrt(eps(grid)),
-                                        abstol = 0,
+                                        abstol = sqrt(eps(grid)),
                                         kw...)
 
     if preconditioner isa DefaultPreconditioner # try to make a useful default
         if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFTSolver
-            if grid.underlying_grid isa XYZRegularRG
-                preconditioner = FFTBasedPoissonSolver(grid.underlying_grid)
-            else # it's stretched in one direction
-                preconditioner = FourierTridiagonalPoissonSolver(grid.underlying_grid)
-            end
+            preconditioner = fft_poisson_solver(grid.underlying_grid)
         else
             preconditioner = DiagonallyDominantPreconditioner()
         end
@@ -73,14 +69,13 @@ function ConjugateGradientPoissonSolver(grid;
 
     rhs = CenterField(grid)
 
-    conjugate_gradient_solver =
-        ConjugateGradientSolver(compute_laplacian!;
-                                              reltol,
-                                              abstol,
-                                              preconditioner,
-                                              template_field = rhs,
-                                              kw...)
-
+    conjugate_gradient_solver = ConjugateGradientSolver(compute_laplacian!;
+                                                        reltol,
+                                                        abstol,
+                                                        preconditioner,
+                                                        template_field = rhs,
+                                                        kw...)
+        
     return ConjugateGradientPoissonSolver(grid, rhs, conjugate_gradient_solver)
 end
 
diff --git a/src/Solvers/conjugate_gradient_solver.jl b/src/Solvers/conjugate_gradient_solver.jl
index d78e04d750..848bba44f6 100644
--- a/src/Solvers/conjugate_gradient_solver.jl
+++ b/src/Solvers/conjugate_gradient_solver.jl
@@ -23,6 +23,7 @@ mutable struct ConjugateGradientSolver{A, G, L, T, F, M, P}
 end
 
 architecture(solver::ConjugateGradientSolver) = solver.architecture
+iteration(cgs::ConjugateGradientSolver) = cgs.iteration
 
 initialize_precondition_product(preconditioner, template_field) = similar(template_field)
 initialize_precondition_product(::Nothing, template_field) = nothing
@@ -36,7 +37,7 @@ Base.summary(::ConjugateGradientSolver) = "ConjugateGradientSolver"
     ConjugateGradientSolver(linear_operation;
                                           template_field,
                                           maxiter = size(template_field.grid),
-                                          reltol = sqrt(eps(eltype(template_field.grid))),
+                                          reltol = sqrt(eps(template_field.grid)),
                                           abstol = 0,
                                           preconditioner = nothing)
 

From c788cd1a3ca56dd0c870e80a1658361b30524c53 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 15:34:22 -0400
Subject: [PATCH 08/12] Import ConjugateGradientPoissonSolver

---
 src/Solvers/Solvers.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl
index bc3d46be74..bed7504b9c 100644
--- a/src/Solvers/Solvers.jl
+++ b/src/Solvers/Solvers.jl
@@ -41,6 +41,7 @@ include("discrete_transforms.jl")
 include("plan_transforms.jl")
 include("fft_based_poisson_solver.jl")
 include("fourier_tridiagonal_poisson_solver.jl")
+include("conjugate_gradient_poisson_solver.jl")
 include("sparse_approximate_inverse.jl")
 include("matrix_solver_utils.jl")
 include("sparse_preconditioners.jl")

From 41b1bd7b0f4ef7317b91bd36498568c6883b81fa Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Mon, 30 Sep 2024 16:35:33 -0400
Subject: [PATCH 09/12] Bugfixes

---
 .../NonhydrostaticModels/NonhydrostaticModels.jl      |  1 -
 src/Solvers/Solvers.jl                                |  3 ++-
 src/Solvers/conjugate_gradient_poisson_solver.jl      | 11 +----------
 3 files changed, 3 insertions(+), 12 deletions(-)

diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
index 544ab96469..0b20dbffbf 100644
--- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
+++ b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
@@ -105,7 +105,6 @@ include("solve_for_pressure.jl")
 include("update_hydrostatic_pressure.jl")
 include("update_nonhydrostatic_model_state.jl")
 include("pressure_correction.jl")
-include("conjugate_gradient_poisson_solver.jl")
 include("nonhydrostatic_tendency_kernel_functions.jl")
 include("compute_nonhydrostatic_tendencies.jl")
 include("compute_nonhydrostatic_boundary_tendencies.jl")
diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl
index bed7504b9c..bbad459ea9 100644
--- a/src/Solvers/Solvers.jl
+++ b/src/Solvers/Solvers.jl
@@ -14,12 +14,13 @@ using SparseArrays
 using KernelAbstractions
 
 using Oceananigans.Architectures: device, CPU, GPU, array_type, on_architecture
+using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Utils
 using Oceananigans.Grids
 using Oceananigans.BoundaryConditions
 using Oceananigans.Fields
 
-using Oceananigans.Grids: unpack_grid
+using Oceananigans.Grids: unpack_grid, inactive_cell
 using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 
 """
diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl
index 93f7580acd..13ae437932 100644
--- a/src/Solvers/conjugate_gradient_poisson_solver.jl
+++ b/src/Solvers/conjugate_gradient_poisson_solver.jl
@@ -1,18 +1,9 @@
-using Oceananigans.Architectures: device, architecture
-using Oceananigans.Solvers: ConjugateGradientPoissonSolver
-using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver
-using Oceananigans.Solvers: fft_poisson_solver, solve!
-using Oceananigans.BoundaryConditions: fill_halo_regions!
-using Oceananigans.Grids: inactive_cell
 using Oceananigans.Operators: divᶜᶜᶜ, ∇²ᶜᶜᶜ 
-using Oceananigans.Utils: launch!, prettysummary
-using Oceananigans.ImmersedBoundaries: mask_immersed_field!
 using Statistics: mean
 
 using KernelAbstractions: @kernel, @index
 
-import Oceananigans.Architecture: architecture
-import Oceananigans.Solvers: precondition!
+import Oceananigans.Architectures: architecture
 
 struct ConjugateGradientPoissonSolver{G, R, S}
     grid :: G

From e4e1b02f462289999331e4a58fa0ebbc20a82791 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Tue, 1 Oct 2024 05:33:08 -0400
Subject: [PATCH 10/12] Bugfix

---
 src/Models/NonhydrostaticModels/solve_for_pressure.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
index 01b3f25719..ffc603f0d7 100644
--- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl
+++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
@@ -1,6 +1,6 @@
 using Oceananigans.Operators
 using Oceananigans.DistributedComputations: DistributedFFTBasedPoissonSolver
-using Oceananigans.Grids: XDirection, YDirection, ZDirection
+using Oceananigans.Grids: XDirection, YDirection, ZDirection, inactive_cell
 using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver
 using Oceananigans.Solvers: ConjugateGradientPoissonSolver
 using Oceananigans.Solvers: solve!

From 19d87a01fb063093eeadabd209ca54432d60a359 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Tue, 1 Oct 2024 07:11:00 -0400
Subject: [PATCH 11/12] Fix function signature plus notation update

---
 .../solve_for_pressure.jl                     | 40 +++++++++----------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
index ffc603f0d7..dedc241ded 100644
--- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl
+++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
@@ -9,67 +9,67 @@ using Oceananigans.Solvers: solve!
 ##### Calculate the right-hand-side of the non-hydrostatic pressure Poisson equation.
 #####
 
-@kernel function _compute_source_term!(rhs, grid, Δt, U★)
+@kernel function _compute_source_term!(rhs, grid, Δt, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
     @inbounds rhs[i, j, k] = active * δ / Δt
 end
 
-@kernel function _fourier_tridiagonal_source_term!(rhs, ::XDirection, grid, Δt, U★)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::XDirection, grid, Δt, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
     @inbounds rhs[i, j, k] = active * Δxᶜᶜᶜ(i, j, k, grid) * δ / Δt
 end
 
-@kernel function _fourier_tridiagonal_source_term!(rhs, ::YDirection, grid, Δt, U★)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::YDirection, grid, Δt, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
     @inbounds rhs[i, j, k] = active * Δyᶜᶜᶜ(i, j, k, grid) * δ / Δt
 end
 
-@kernel function _fourier_tridiagonal_source_term!(rhs, ::ZDirection, grid, Δt, U★)
+@kernel function _fourier_tridiagonal_source_term!(rhs, ::ZDirection, grid, Δt, Ũ)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)
-    δ = divᶜᶜᶜ(i, j, k, grid, U★.u, U★.v, U★.w)
+    δ = divᶜᶜᶜ(i, j, k, grid, Ũ.u, Ũ.v, Ũ.w)
     @inbounds rhs[i, j, k] = active * Δzᶜᶜᶜ(i, j, k, grid) * δ / Δt
 end
 
-function compute_source_term!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, U★)
+function compute_source_term!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, Ũ)
     rhs  = solver.storage.zfield
     arch = architecture(solver)
     grid = solver.local_grid
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
     return nothing        
 end
 
-function compute_source_term!(pressure, solver::DistributedFourierTridiagonalPoissonSolver, Δt, U★)
+function compute_source_term!(pressure, solver::DistributedFourierTridiagonalPoissonSolver, Δt, Ũ)
     rhs = solver.storage.zfield
     arch = architecture(solver)
     grid = solver.local_grid
     tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
     launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!,
-            rhs, grid, Δt, U★, tridiagonal_dir)
+            rhs, tridiagonal_dir, grid, Δt, Ũ)
     return nothing
 end
 
-function compute_source_term!(pressure, solver::FourierTridiagonalPoissonSolver, Δt, U★)
+function compute_source_term!(pressure, solver::FourierTridiagonalPoissonSolver, Δt, Ũ)
     rhs = solver.source_term
     arch = architecture(solver)
     grid = solver.grid
     tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
     launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!,
-            rhs, grid, Δt, U★, tridiagonal_dir)
+            rhs, grid, Δt, Ũ, tridiagonal_dir)
     return nothing
 end
 
-function compute_source_term!(pressure, solver::FFTBasedPoissonSolver, Δt, U★)
+function compute_source_term!(pressure, solver::FFTBasedPoissonSolver, Δt, Ũ)
     rhs = solver.storage
     arch = architecture(solver)
     grid = solver.grid
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
     return nothing
 end
 
@@ -77,17 +77,17 @@ end
 ##### Solve for pressure
 #####
 
-function solve_for_pressure!(pressure, solver, Δt, U★)
-    compute_source_term!(pressure, solver, Δt, U★)
+function solve_for_pressure!(pressure, solver, Δt, Ũ)
+    compute_source_term!(pressure, solver, Δt, Ũ)
     solve!(pressure, solver)
     return pressure
 end
 
-function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, U★)
+function solve_for_pressure!(pressure, solver::ConjugateGradientPoissonSolver, Δt, Ũ)
     rhs = solver.right_hand_side
     grid = solver.grid
     arch = architecture(grid)
-    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, U★)
+    launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
     return solve!(pressure, solver.conjugate_gradient_solver, rhs)
 end
 

From 3b43e037ea75b7aab67d48089fd32362c3b37f69 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Tue, 1 Oct 2024 07:12:06 -0400
Subject: [PATCH 12/12] Missed a spot

---
 src/Models/NonhydrostaticModels/solve_for_pressure.jl | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
index dedc241ded..c3db828018 100644
--- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl
+++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
@@ -49,9 +49,8 @@ function compute_source_term!(pressure, solver::DistributedFourierTridiagonalPoi
     rhs = solver.storage.zfield
     arch = architecture(solver)
     grid = solver.local_grid
-    tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
-    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!,
-            rhs, tridiagonal_dir, grid, Δt, Ũ)
+    tdir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Δt, Ũ)
     return nothing
 end
 
@@ -59,9 +58,8 @@ function compute_source_term!(pressure, solver::FourierTridiagonalPoissonSolver,
     rhs = solver.source_term
     arch = architecture(solver)
     grid = solver.grid
-    tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction
-    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!,
-            rhs, grid, Δt, Ũ, tridiagonal_dir)
+    tdir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Δt, Ũ)
     return nothing
 end