Enforce gauge condition in preconditioners for ConjugateGradientPoiss…

…onSolver (#3802) * Reorganize utils and enforce gauge condition * Fix formatting issue * Some small improvements * PreconditionedCGSolver -> CGSolver * Change to conjugate gradient solver * Move ConjugateGradientPoissonSolver to Solvers * More improvments * Import ConjugateGradientPoissonSolver * Bugfixes * Bugfix * Fix function signature plus notation update * Missed a spot
CliMA · Oct 1, 2024 · 45838a5 · 45838a5
1 parent 3ea2545
commit 45838a5
Show file tree

Hide file tree

Showing 13 changed files with 199 additions and 181 deletions.
diff --git 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/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
@@ -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/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/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.
@@ -115,9 +105,9 @@ 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")
 
 end # module
+
diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl
@@ -1,94 +1,91 @@
 using Oceananigans.Operators
-using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver, solve!
 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!
 
 #####
 ##### 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
+    tdir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Δt, Ũ)
     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
+    tdir = solver.batched_tridiagonal_solver.tridiagonal_direction
+    launch!(arch, grid, :xyz, _fourier_tridiagonal_source_term!, rhs, tdir, grid, Δt, Ũ)
     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 pressure
+end
 
-    return nothing
+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/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
@@ -4,7 +4,7 @@ export
     BatchedTridiagonalSolver, solve!,
     FFTBasedPoissonSolver,
     FourierTridiagonalPoissonSolver,
-    PreconditionedConjugateGradientSolver,
+    ConjugateGradientSolver,
     HeptadiagonalIterativeSolver
 
 using Statistics
@@ -14,12 +14,14 @@ 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
 
 """
     ω(M, k)
@@ -33,16 +35,24 @@ 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("conjugate_gradient_poisson_solver.jl")
 include("sparse_approximate_inverse.jl")
 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