PartitionedSolvers refactor

.github/workflows/ci.yml

@@ -18,7 +18,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
           - '1'
         os:
           - ubuntu-latest
@@ -85,8 +84,7 @@ jobs:
       - run: |
           julia --project=HPCG -e '
             using Pkg
-            Pkg.develop(path=".")
-            Pkg.develop(path="./PartitionedSolvers")
+            Pkg.develop([Pkg.PackageSpec(path="."),Pkg.PackageSpec(path="./PartitionedSolvers")])
             Pkg.test("HPCG")'
   docs:
     name: Documentation

HPCG/Project.toml

@@ -23,7 +23,7 @@ DelimitedFiles = "1.9"
 JSON = "0.21"
 MPI = "0.20"
 PartitionedArrays = "0.5"
-PartitionedSolvers = "0.2"
+PartitionedSolvers = "0.3"
 Primes = "0.5"
 SparseMatricesCSR = "0.6"
 julia = "1.1"

HPCG/src/mg_preconditioner.jl

@@ -110,7 +110,7 @@ function generate_problem(ranks, npx, npy, npz, nx, ny, nz, solver)
     Axf = similar(r)
     Axf .= 0
     x .= 0
-    gs_state = setup(solver, x, A, r)
+    gs_state = solver(PartitionedSolvers.linear_problem(x, A, r))
     return A, r, x, Axf, gs_state
 end
 
@@ -139,11 +139,11 @@ function pc_setup(np, ranks, l, nx, ny, nz)
     r = Vector{PVector}(undef, l)
     x = Vector{PVector}(undef, l)
     Axf = Vector{PVector}(undef, l)
-    gs_states = Vector{PartitionedSolvers.Preconditioner}(undef, l)
+    gs_states = Vector{PartitionedSolvers.LinearSolver}(undef, l)
     npx, npy, npz = compute_optimal_shape_XYZ(np)
     nnz_vec = Vector{Int64}(undef, l)
     nrows_vec = Vector{Int64}(undef, l)
-    solver = PartitionedSolvers.additive_schwarz_correction_partition(gauss_seidel(; iters = 1))
+    solver = p -> PartitionedSolvers.gauss_seidel(p;iterations=1)
 
     # create top problem 
     A, r, x, Axf, gs_state = generate_problem(ranks, npx, npy, npz, nx, ny, nz, solver)
@@ -313,16 +313,16 @@ end
 """
 function pc_solve!(x, s::Mg_preconditioner, b, l; zero_guess = false)
     if l == 1
-        solve!(x, s.gs_states[l], b; zero_guess) # bottom solve
+        PartitionedSolvers.smooth!(x, s.gs_states[l], b; zero_guess) # bottom solve
     else
-        solve!(x, s.gs_states[l], b; zero_guess) # presmoother 
+        PartitionedSolvers.smooth!(x, s.gs_states[l], b; zero_guess) # presmoother 
         mul_no_lat!(s.Axf[l], s.A_vec[l], x)
         p_restrict!(s.r[l-1], b, s.Axf[l], s.f2c[l-1])
         s.x[l-1] .= 0.0
         pc_solve!(s.x[l-1], s, s.r[l-1], l - 1; zero_guess = true)
         p_prolongate!(x, s.x[l-1], s.f2c[l-1])
-        consistent!(x) |> wait
-        solve!(x, s.gs_states[l], b) # post smooth
+        #consistent!(x) |> wait #Already inside gauss_seidel
+        PartitionedSolvers.smooth!(x, s.gs_states[l], b) # post smooth
     end
     x
 end

PartitionedSolvers/Project.toml

@@ -1,18 +1,21 @@
 name = "PartitionedSolvers"
 uuid = "11b65f7f-80ac-401b-9ef2-3db765482d62"
 authors = ["Francesc Verdugo <f.verdugo.rojano@vu.nl>"]
-version = "0.2.2"
+version = "0.3.0"
 
 [deps]
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 
 [compat]
 IterativeSolvers = "0.9"
+NLsolve = "4"
 PartitionedArrays = "0.4.4, 0.5"
 SparseArrays = "1"
 julia = "1.6"

PartitionedSolvers/src/PartitionedSolvers.jl

@@ -5,39 +5,13 @@ using PartitionedArrays: val_parameter
 using SparseArrays
 using LinearAlgebra
 using IterativeSolvers
+using Printf
+import NLsolve
 using SparseMatricesCSR
 
-export setup
-export solve!
-export update!
-export finalize!
-export AbstractLinearSolver
-export linear_solver
-export default_nullspace
-export nullspace
-export uses_nullspace
-export uses_initial_guess
-export iterations!
 include("interfaces.jl")
-
-export lu_solver
-export jacobi_correction
-export richardson
-export jacobi
-export gauss_seidel
-export additive_schwarz_correction
-export additive_schwarz
+include("wrappers.jl")
 include("smoothers.jl")
-
-export amg
-export smoothed_aggregation
-export v_cycle
-export w_cycle
-export amg_level_params
-export amg_level_params_linear_elasticity
-export amg_fine_params
-export amg_coarse_params
-export amg_statistics
 include("amg.jl")
 
 end # module

PartitionedSolvers/src/amg.jl

@@ -1,4 +1,15 @@
 
+function default_nullspace(A)
+    T = eltype(A)
+    [ones(T,size(A,2))]
+end
+
+function default_nullspace(A::PSparseMatrix)
+    col_partition = partition(axes(A,2))
+    T = eltype(A)
+    [ pones(T,col_partition) ]
+end
+
 function aggregate(A,diagA=dense_diag(A);epsilon)
     # TODO It assumes CSC format for the moment
 
@@ -744,7 +755,7 @@ function dofs_to_node(dofs, block_size)
 end
 
 function amg_level_params_linear_elasticity(block_size;
-    pre_smoother = additive_schwarz(gauss_seidel(;iters=1);iters=1),
+    pre_smoother = p-> additive_schwarz(p;local_solver=q->gauss_seidel(q;iterations=1),iterations=1),
     coarsening = smoothed_aggregation_with_block_size(;approximate_omega = lambda_generic,
     tentative_prolongator = tentative_prolongator_with_block_size, block_size = block_size),
     cycle = v_cycle,
@@ -756,7 +767,7 @@ function amg_level_params_linear_elasticity(block_size;
 end
 
 function amg_level_params(;
-    pre_smoother = additive_schwarz(gauss_seidel(;iters=1);iters=1),
+    pre_smoother = p-> additive_schwarz(p;local_solver=q->gauss_seidel(q;iterations=1),iterations=1),
     coarsening = smoothed_aggregation(;),
     cycle = v_cycle,
     pos_smoother = pre_smoother,
@@ -774,23 +785,24 @@ end
 
 function amg_coarse_params(;
     #TODO more resonable defaults?
-    coarse_solver = lu_solver(),
+    coarse_solver = LinearAlgebra_lu,
     coarse_size = 10,
     )
     coarse_params = (;coarse_solver,coarse_size)
     coarse_params
 end
 
-function amg(;
+function amg(p;
         fine_params=amg_fine_params(),
         coarse_params=amg_coarse_params(),)
     amg_params = (;fine_params,coarse_params)
-    setup(x,O,b,options) = amg_setup(x,O,b,nullspace(options),amg_params)
-    update! = amg_update!
-    solve! = amg_solve!
-    finalize! = amg_finalize!
-    uses_nullspace = Val(true)
-    linear_solver(;setup,update!,solve!,finalize!,uses_nullspace)
+
+    x = solution(p)
+    b = rhs(p)
+    A = matrix(p)
+    B = nullspace(p)
+    setup = amg_setup(x,A,b,B,amg_params)
+    linear_solver(amg_update!,amg_step!,p,setup)
 end
 
 function amg_setup(x,A,b,::Nothing,amg_params)
@@ -807,8 +819,8 @@ function amg_setup(x,A,b,B,amg_params)
             return nothing
         end
         (;pre_smoother,pos_smoother,coarsening,cycle) = fine_level
-        pre_setup = setup(pre_smoother,x,A,b)
-        pos_setup = setup(pos_smoother,x,A,b)
+        pre_setup = pre_smoother(linear_problem(x,A,b))
+        pos_setup = pos_smoother(linear_problem(x,A,b))
         coarsen, _ = coarsening
         Ac,Bc,R,P,Ac_setup = coarsen(A,B)
         nc = size(Ac,1)
@@ -830,28 +842,29 @@ function amg_setup(x,A,b,B,amg_params)
     end
     n_fine_levels = count(i->i!==nothing,fine_levels)
     nlevels = n_fine_levels+1
-    coarse_solver_setup = setup(coarse_solver,x,A,b)
+    coarse_solver_setup = coarse_solver(linear_problem(x,A,b))
     coarse_level = (;coarse_solver_setup)
     (;nlevels,fine_levels,coarse_level,amg_params)
 end
 
-function amg_solve!(x,setup,b,options)
+function amg_step!(x,setup,b,phase=:start;kwargs...)
     level=1
     amg_cycle!(x,setup,b,level)
-    x
+    phase=:stop
+    x,setup,phase
 end
 
 function amg_cycle!(x,setup,b,level)
     amg_params = setup.amg_params
     if level == setup.nlevels
         coarse_solver_setup = setup.coarse_level.coarse_solver_setup
-        return solve!(x,coarse_solver_setup,b)
+        return smooth!(x,coarse_solver_setup,b)
     end
     level_params = amg_params.fine_params[level]
     level_setup = setup.fine_levels[level]
     (;cycle) = level_params
     (;R,P,r,rc,rc2,e,ec,ec2,A,Ac,pre_setup,pos_setup) = level_setup
-    solve!(x,pre_setup,b)
+    smooth!(x,pre_setup,b)
     mul!(r,A,x)
     r .= b .- r
     mul!(rc2,R,r)
@@ -861,16 +874,16 @@ function amg_cycle!(x,setup,b,level)
     ec2 .= ec
     mul!(e,P,ec2)
     x .+= e
-    solve!(x,pos_setup,b)
+    smooth!(x,pos_setup,b)
     x
 end
 
-function amg_statistics(P::Preconditioner)
+function amg_statistics(P)
     # Taken from: An Introduction to Algebraic Multigrid, R. D. Falgout, April 25, 2006
     # Grid complexity is the total number of grid points on all grids divided by the number
     # of grid points on the fine grid. Operator complexity is the total number of nonzeroes in the linear operators
     # on all grids divided by the number of nonzeroes in the fine grid operator
-    setup = P.solver_setup
+    setup = P.workspace
     nlevels = setup.nlevels
     level_rows = zeros(Int,nlevels)
     level_nnz = zeros(Int,nlevels)
@@ -909,7 +922,7 @@ end
     amg_cycle!(args...)
 end
 
-function amg_update!(setup,A,options)
+function amg_update!(setup,A)
     amg_params = setup.amg_params
     nlevels = setup.nlevels
     for level in 1:(nlevels-1)
@@ -918,13 +931,13 @@ function amg_update!(setup,A,options)
         (;coarsening) = level_params
         _, coarsen! = coarsening
         (;R,P,A,Ac,Ac_setup,pre_setup,pos_setup) = level_setup
-        update!(pre_setup,A)
-        update!(pos_setup,A)
+        update(pre_setup,matrix=A)
+        update(pos_setup,matrix=A)
         coarsen!(A,Ac,R,P,Ac_setup)
         A = Ac
     end
     coarse_solver_setup = setup.coarse_level.coarse_solver_setup
-    update!(coarse_solver_setup,A)
+    update(coarse_solver_setup,matrix=A)
     setup
 end