Merge pull request #150 from fverdugo/solver_interface

Solver interface
fverdugo · May 26, 2024 · f167f73 · f167f73 · fverdugo · May 26, 2024
2 parents bdf5984 + 09409f1
commit f167f73
Show file tree

Hide file tree

Showing 9 changed files with 222 additions and 208 deletions.
diff --git a/PartitionedSolvers/CHANGELOG.md b/PartitionedSolvers/CHANGELOG.md
@@ -6,7 +6,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [0.1.0] - Unreleased
+## [0.2.0] - Unreleased
+
+### Changed
+
+- Linear solver interface.
+
+## [0.1.0]
 
 ### Added
 

diff --git a/PartitionedSolvers/Project.toml b/PartitionedSolvers/Project.toml
@@ -1,7 +1,7 @@
 name = "PartitionedSolvers"
 uuid = "11b65f7f-80ac-401b-9ef2-3db765482d62"
 authors = ["Francesc Verdugo <[email protected]>"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"

diff --git a/PartitionedSolvers/src/PartitionedSolvers.jl b/PartitionedSolvers/src/PartitionedSolvers.jl
@@ -6,14 +6,11 @@ using LinearAlgebra
 
 export setup
 export solve!
-export setup!
+export update!
 export finalize!
 export AbstractLinearSolver
 export linear_solver
-export attach_nullspace
 export default_nullspace
-export preconditioner
-export matrix
 export nullspace
 include("interfaces.jl")
 

diff --git a/PartitionedSolvers/src/amg.jl b/PartitionedSolvers/src/amg.jl
@@ -287,9 +287,7 @@ function smoothed_aggregation(;
     tentative_prolongator = tentative_prolongator_for_laplace,
     repartition_threshold = 2000,
     )
-    function coarsen(operator)
-        A = matrix(operator)
-        B = nullspace(operator)
+    function coarsen(A,B)
         diagA = dense_diag(A)
         node_to_aggregate, aggregates = aggregate(A,diagA;epsilon)
         n_nullspace_vecs = length(B)
@@ -299,14 +297,11 @@ function smoothed_aggregation(;
         R = transpose(P)
         Ac,cache = rap(R,A,P;reuse=true)
         Ac,Bc,R,P,cache,repartition_threshold = enhance_coarse_partition(A,Ac,Bc,R,P,cache,repartition_threshold)
-        coarse_operator = attach_nullspace(Ac,Bc)
-        coarse_operator,R,P,cache
+        Ac,Bc,R,P,cache
     end
-    function coarsen!(operator,coarse_operator,R,P,cache)
-        A = matrix(operator)
-        Ac = matrix(coarse_operator)
+    function coarsen!(A,Ac,R,P,cache)
         rap!(Ac,R,A,P,cache)
-        coarse_operator,R,P,cache
+        Ac,R,P,cache
     end
     (coarsen, coarsen!)
 end
@@ -341,14 +336,18 @@ function amg(;
         fine_params=amg_fine_params(),
         coarse_params=amg_coarse_params(),)
     amg_params = (;fine_params,coarse_params)
-    setup(x,O,b) = amg_setup(x,O,b,amg_params)
-    setup! = amg_setup!
+    setup(x,O,b,options) = amg_setup(x,O,b,nullspace(options),amg_params)
+    update! = amg_update!
     solve! = amg_solve!
     finalize! = amg_finalize!
-    linear_solver(;setup,setup!,solve!,finalize!)
+    linear_solver(;setup,update!,solve!,finalize!)
 end
 
-function amg_setup(x,operator,b,amg_params)
+function amg_setup(x,A,b,::Nothing,amg_params)
+    B = default_nullspace(A)
+    amg_setup(x,A,b,B,amg_params)
+end
+function amg_setup(x,A,b,B,amg_params)
     fine_params = amg_params.fine_params
     coarse_params = amg_params.coarse_params
     (;coarse_solver,coarse_size) = coarse_params
@@ -358,11 +357,10 @@ function amg_setup(x,operator,b,amg_params)
             return nothing
         end
         (;pre_smoother,pos_smoother,coarsening,cycle) = fine_level
-        pre_setup = setup(pre_smoother)(x,operator,b)
-        pos_setup = setup(pos_smoother)(x,operator,b)
+        pre_setup = setup(pre_smoother,x,A,b)
+        pos_setup = setup(pos_smoother,x,A,b)
         coarsen, _ = coarsening
-        coarse_operator,R,P,coarse_operator_setup = coarsen(operator)
-        Ac = matrix(coarse_operator)
+        Ac,Bc,R,P,Ac_setup = coarsen(A,B)
         nc = size(Ac,1)
         if nc <= coarse_size
             done = true
@@ -373,20 +371,21 @@ function amg_setup(x,operator,b,amg_params)
         e = similar(x)
         ec = similar(e,axes(Ac,2))
         ec2 = similar(e,axes(P,2)) # TODO
-        level_setup = (;R,P,r,rc,rc2,e,ec,ec2,operator,coarse_operator,pre_setup,pos_setup,coarse_operator_setup)
+        level_setup = (;R,P,r,rc,rc2,e,ec,ec2,A,B,Ac,Bc,pre_setup,pos_setup,Ac_setup)
         x = ec
         b = rc
-        operator = coarse_operator
+        A = Ac
+        B = Bc
         level_setup
     end
     n_fine_levels = count(i->i!==nothing,fine_levels)
     nlevels = n_fine_levels+1
-    coarse_solver_setup = setup(coarse_solver)(x,operator,b)
+    coarse_solver_setup = setup(coarse_solver,x,A,b)
     coarse_level = (;coarse_solver_setup)
     (;nlevels,fine_levels,coarse_level,amg_params)
 end
 
-function amg_solve!(x,setup,b)
+function amg_solve!(x,setup,b,options)
     level=1
     amg_cycle!(x,setup,b,level)
     x
@@ -395,16 +394,14 @@ end
 function amg_cycle!(x,setup,b,level)
     amg_params = setup.amg_params
     if level == setup.nlevels
-        coarse_solver = amg_params.coarse_params.coarse_solver
         coarse_solver_setup = setup.coarse_level.coarse_solver_setup
-        return solve!(coarse_solver)(x,coarse_solver_setup,b)
+        return solve!(x,coarse_solver_setup,b)
     end
     level_params = amg_params.fine_params[level]
     level_setup = setup.fine_levels[level]
-    (;pre_smoother,pos_smoother,cycle) = level_params
-    (;R,P,r,rc,rc2,e,ec,ec2,operator,coarse_operator,pre_setup,pos_setup) = level_setup
-    solve!(pre_smoother)(x,pre_setup,b)
-    A = matrix(operator)
+    (;cycle) = level_params
+    (;R,P,r,rc,rc2,e,ec,ec2,A,Ac,pre_setup,pos_setup) = level_setup
+    solve!(x,pre_setup,b)
     mul!(r,A,x)
     r .= b .- r
     mul!(rc2,R,r)
@@ -414,28 +411,29 @@ function amg_cycle!(x,setup,b,level)
     ec2 .= ec
     mul!(e,P,ec2)
     x .+= e
-    solve!(pos_smoother)(x,pos_setup,b)
+    solve!(x,pos_setup,b)
     x
 end
 
-function amg_statistics(setup)
+function amg_statistics(P::Preconditioner)
     # 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
     nlevels = setup.nlevels
     level_rows = zeros(Int,nlevels)
     level_nnz = zeros(Int,nlevels)
     for level in 1:(nlevels-1)
         level_setup = setup.fine_levels[level]
-        (;operator,) = level_setup
-        level_rows[level] = size(matrix(operator),1)
-        level_nnz[level] = nnz(matrix(operator))
+        (;A,) = level_setup
+        level_rows[level] = size(A,1)
+        level_nnz[level] = nnz(A)
     end
     level_setup = setup.fine_levels[nlevels-1]
-    (;coarse_operator) = level_setup
-    level_rows[end] = size(matrix(coarse_operator),1)
-    level_nnz[end] = nnz(matrix(coarse_operator))
+    (;Ac) = level_setup
+    level_rows[end] = size(Ac,1)
+    level_nnz[end] = nnz(Ac)
     nnz_total = sum(level_nnz)
     rows_total = sum(level_rows)
     level_id = collect(1:nlevels)
@@ -461,23 +459,22 @@ end
     amg_cycle!(args...)
 end
 
-function amg_setup!(setup,operator)
+function amg_update!(setup,A,options)
     amg_params = setup.amg_params
     nlevels = setup.nlevels
     for level in 1:(nlevels-1)
         level_params = amg_params.fine_params[level]
         level_setup = setup.fine_levels[level]
-        (;coarsening,pre_smoother,pos_smoother) = level_params
+        (;coarsening) = level_params
         _, coarsen! = coarsening
-        (;R,P,operator,coarse_operator,coarse_operator_setup,pre_setup,pos_setup) = level_setup
-        setup!(pre_smoother)(pre_setup,operator)
-        setup!(pos_smoother)(pos_setup,operator)
-        coarsen!(operator,coarse_operator,R,P,coarse_operator_setup)
-        operator = coarse_operator
+        (;R,P,A,Ac,Ac_setup,pre_setup,pos_setup) = level_setup
+        update!(pre_setup,A)
+        update!(pos_setup,A)
+        coarsen!(A,Ac,R,P,Ac_setup)
+        A = Ac
     end
-    coarse_solver = amg_params.coarse_params.coarse_solver
     coarse_solver_setup = setup.coarse_level.coarse_solver_setup
-    setup!(coarse_solver)(coarse_solver_setup,operator)
+    update!(coarse_solver_setup,A)
     setup
 end
 
@@ -487,14 +484,12 @@ function amg_finalize!(setup)
     for level in 1:(nlevels-1)
         level_params = amg_params.fine_params[level]
         level_setup = setup.fine_levels[level]
-        (;pre_smoother,pos_smoother) = level_params
         (;pre_setup,pos_setup) = level_setup
-        finalize!(pre_smoother)(pre_setup)
-        finalize!(pos_smoother)(pos_setup)
+        finalize!(pre_setup)
+        finalize!(pos_setup)
     end
-    coarse_solver = amg_params.coarse_params.coarse_solver
     coarse_solver_setup = setup.coarse_level.coarse_solver_setup
-    finalize!(coarse_solver)(coarse_solver_setup)
+    finalize!(coarse_solver_setup)
     nothing
 end
 
diff --git a/PartitionedSolvers/src/interfaces.jl b/PartitionedSolvers/src/interfaces.jl
@@ -1,18 +1,4 @@
 
-function attach_nullspace(A,ns)
-    MatrixWithNullspace(A,ns)
-end
-
-struct MatrixWithNullspace{A,B}
-    matrix::A
-    nullspace::B
-end
-matrix(a::MatrixWithNullspace) = a.matrix
-nullspace(a::MatrixWithNullspace) = a.nullspace
-
-matrix(a) = a
-nullspace(a) = default_nullspace(a)
-
 function default_nullspace(A)
     T = eltype(A)
     [ones(T,size(A,2))]
@@ -29,46 +15,61 @@ abstract type AbstractLinearSolver end
 function linear_solver(;
         setup,
         solve!,
-        setup!,
+        update!,
         finalize! = ls_setup->nothing,
     )
-    LinearSolver(setup,solve!,setup!,finalize!)
+    LinearSolver(setup,solve!,update!,finalize!)
 end
 
 struct LinearSolver{A,B,C,D} <: AbstractLinearSolver
     setup::A
     solve!::B
-    setup!::C
+    update!::C
     finalize!::D
 end
 
-setup(solver::LinearSolver) = solver.setup
-setup!(solver::LinearSolver) = solver.setup!
-solve!(solver::LinearSolver) = solver.solve!
-finalize!(solver::LinearSolver) = solver.finalize!
-
 struct Preconditioner{A,B}
     solver::A
     solver_setup::B
 end
 
-function preconditioner(solver,x,A,b)
-    solver_setup = setup(solver)(x,A,b)
+function setup_options(;nullspace=nothing)
+    options = (;nullspace)
+end
+
+function nullspace(options)
+    options.nullspace
+end
+
+function setup(solver::LinearSolver,x,A,b;kwargs...)
+    options = setup_options(;kwargs...)
+    solver_setup = solver.setup(x,A,b,options)
     Preconditioner(solver,solver_setup)
 end
 
-function preconditioner!(P::Preconditioner,A)
-    setup!(P.solver)(P.solver_setup,A)
+function update!(P::Preconditioner,A;kwargs...)
+    options = setup_options(;kwargs...)
+    P.solver.update!(P.solver_setup,A,options)
     P
 end
 
+function solve_options(;zero_guess=false)
+    options = (;zero_guess)
+end
+
+function solve!(x,P::Preconditioner,b;kwargs...)
+    options = solve_options(;kwargs...)
+    P.solver.solve!(x,P.solver_setup,b,options)
+    x
+end
+
 function LinearAlgebra.ldiv!(x,P::Preconditioner,b)
     fill!(x,zero(eltype(x)))
-    solve!(P.solver)(x,P.solver_setup,b)
+    solve!(x,P,b;zero_guess=true)
     x
 end
 
 function finalize!(P::Preconditioner)
-    PartitionedSolvers.finalize!(P.solver)(P.solver_setup)
+    P.solver.finalize!(P.solver_setup)
 end