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/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/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/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl index f7104922be..0b20dbffbf 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. @@ -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 index 5064a2c5d6..c3db828018 100644 --- 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, Ũ) 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, Ũ.v, Ũ.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, Ũ) 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, Ũ.v, Ũ.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, Ũ) 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, Ũ.v, Ũ.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, Ũ) 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, Ũ.v, Ũ.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, Ũ) 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, Ũ) + 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, Ũ) + 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, Ũ) 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, Ũ) + 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, Ũ) 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, Ũ) + rhs = solver.storage arch = architecture(solver) grid = solver.grid + launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ) + 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, Ũ) + compute_source_term!(pressure, solver, Δt, Ũ) solve!(pressure, solver) + return pressure +end - return nothing +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, Ũ) + 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 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..bbad459ea9 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 @@ -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 diff --git a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl similarity index 57% rename from src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl rename to src/Solvers/conjugate_gradient_poisson_solver.jl index 3223b1e1ca..13ae437932 100644 --- a/src/Models/NonhydrostaticModels/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -1,16 +1,9 @@ -using Oceananigans.Architectures: device, architecture -using Oceananigans.Solvers: PreconditionedConjugateGradientSolver, FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver, 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.Solvers: precondition! -import ..Models: iteration +import Oceananigans.Architectures: architecture struct ConjugateGradientPoissonSolver{G, R, S} grid :: G @@ -18,7 +11,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)) @@ -27,7 +21,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', @@ -49,19 +42,17 @@ function compute_laplacian!(∇²ϕ, ϕ) return nothing end +struct DefaultPreconditioner end + function ConjugateGradientPoissonSolver(grid; - preconditioner = nothing, + preconditioner = DefaultPreconditioner(), reltol = sqrt(eps(grid)), - abstol = 0, + abstol = sqrt(eps(grid)), kw...) - if isnothing(preconditioner) # make a useful default - if grid isa ImmersedBoundaryGrid && grid.underlying_grid isa GridWithFFT - if grid.underlying_grid isa XYZRegularRG - preconditioner = FFTBasedPoissonSolver(grid.underlying_grid) - else # it's stretched in one direction - preconditioner = FourierTridiagonalPoissonSolver(grid.underlying_grid) - end + 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) else preconditioner = DiagonallyDominantPreconditioner() end @@ -69,60 +60,78 @@ function ConjugateGradientPoissonSolver(grid; rhs = CenterField(grid) - conjugate_gradient_solver = - PreconditionedConjugateGradientSolver(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 -@kernel function compute_source_term!(rhs, grid, Δt, U★) +##### +##### A preconditioner based on the FFT solver +##### + +@kernel function fft_preconditioner_rhs!(preconditioner_rhs, rhs) 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 + @inbounds preconditioner_rhs[i, j, k] = rhs[i, j, k] 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 +@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 - rhs = solver.right_hand_side +@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, 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) + launch!(arch, grid, :xyz, fft_preconditioner_rhs!, solver.storage, rhs) + return nothing +end - return pressure +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 -##### -##### A preconditioner based on the FFT solver -##### +const FFTBasedPreconditioner = Union{FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver} -@kernel function fft_preconditioner_right_hand_side!(preconditioner_rhs, rhs) - i, j, k = @index(Global, NTuple) - @inbounds preconditioner_rhs[i, j, k] = rhs[i, j, k] -end +function precondition!(p, preconditioner::FFTBasedPreconditioner, r, args...) + compute_preconditioner_rhs!(preconditioner, r) + p = solve!(p, preconditioner) -function precondition!(p, solver::FFTBasedPoissonSolver, rhs, args...) - grid = solver.grid + mean_p = mean(p) + grid = p.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, mean_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 +##### The "DiagonallyDominantPreconditioner" (Marshall et al 1997) ##### struct DiagonallyDominantPreconditioner end @@ -133,26 +142,25 @@ Base.summary(::DiagonallyDominantPreconditioner) = "DiagonallyDominantPreconditi arch = architecture(p) fill_halo_regions!(r) launch!(arch, grid, :xyz, _diagonally_dominant_precondition!, p, grid, r) + + mean_p = mean(p) + launch!(arch, grid, :xyz, subtract_and_mask!, p, grid, mean_p) + return p 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] - @@ -163,5 +171,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 + diff --git a/src/Solvers/preconditioned_conjugate_gradient_solver.jl b/src/Solvers/conjugate_gradient_solver.jl similarity index 82% rename from src/Solvers/preconditioned_conjugate_gradient_solver.jl rename to src/Solvers/conjugate_gradient_solver.jl index b856f57792..848bba44f6 100644 --- a/src/Solvers/preconditioned_conjugate_gradient_solver.jl +++ b/src/Solvers/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,26 @@ mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P} preconditioner_product :: P end -architecture(solver::PreconditionedConjugateGradientSolver) = solver.architecture +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 -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))), + reltol = sqrt(eps(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 +73,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 +93,7 @@ function PreconditionedConjugateGradientSolver(linear_operation; FT = eltype(grid) - return PreconditionedConjugateGradientSolver(arch, + return ConjugateGradientSolver(arch, grid, linear_operation, FT(reltol), @@ -108,7 +109,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 +157,7 @@ Loop: ρⁱ⁻¹ = ρ ``` """ -function solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...) +function solve!(x, solver::ConjugateGradientSolver, b, args...) # Initialize solver.iteration = 0 @@ -169,8 +170,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 +185,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 +193,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 +231,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 +259,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/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl index 81f8c1633d..3f5fae67e3 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) @@ -102,7 +102,9 @@ function solve!(ϕ, solver::FFTBasedPoissonSolver, b, 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, 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(ϕ)) @@ -131,3 +135,4 @@ end @inbounds ϕ[i′, j′, k′] = real(ϕc[i, j, k]) end + 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 + 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)