JuliaSmoothOptimizers / Krylov.jl

A Julia Basket of Hand-Picked Krylov Methods
Other
338 stars 51 forks source link

Compatibility with BlockArrays #605

Open lijas opened 2 years ago

lijas commented 2 years ago

I have a linear system in the form: A = [K C'; C, 0], (C is relatively small) which requires a large amount of iterations with minres/gmres in order to converge. There is no "out of the box" preconditioner that helps with lowering the amount of iterations.

I dont know if this is the best way, but a few tutorials suggests creating a custom preconditioner, P^-1 = [ilu(K) 0; 0 C'inv(diag(K))C] or something similar (perhaps there is a better preconditioner for this case), for which you preferable want to use BlockArrays and overload ldiv! in for your costum preconditier:


function LinearAlgebra.ldiv!(y::BlockVector, A::CostumBlockPrecon, b::BlockVector)
    y[Block(1)] = A.K_ilu\b[Block(1)] #precomputed ilu preconditioner
    y[Block(2)] = A.B * (diag(A.K)\(A.B*b[Block(2)]))
end

This is the error you get if you try to use BlockArrays+Krylov:

using Krylov
using BlockArrays

K = rand(10,10); 
B = rand(2, 10)
b = rand(10 + 2);

b = BlockVector(rand(Float64,12), [10,2])
A = BlockArray(rand(Float64,12,12), [10,2], [10,2])
A[Block(1,1)] = K
A[Block(1,2)] = B'
A[Block(2,1)] = B

Krylov.gmres(A, b)
ERROR: MethodError: no method matching BlockVector{Float64, Vector{Vector{Float64}}, Tuple{BlockedUnitRange{Vector{Int64}}}}(::UndefInitializer, ::Int64)
Closest candidates are:
  BlockArray{T, N, R, BS}(::UndefInitializer, ::Tuple{Vararg{Int64, N}}) where {T, N, R<:(AbstractArray{<:AbstractArray{T, N}, N}), BS<:Tuple{Vararg{AbstractUnitRange{Int64}, N}}} at ~/.julia/packages/BlockArrays/UKehW/src/blockarray.jl:162
  BlockArray{T, N, R, BS}(::UndefInitializer, ::BS) where {T, N, R<:(AbstractArray{<:AbstractArray{T, N}, N}), BS<:Tuple{Vararg{AbstractUnitRange{Int64}, N}}} at ~/.julia/packages/BlockArrays/UKehW/src/blockarray.jl:149
Stacktrace:
 [1] GmresSolver(n::Int64, m::Int64, memory::Int64, S::Type{BlockVector{Float64, Vector{Vector{Float64}}, Tuple{BlockedUnitRange{Vector{Int64}}}}})
   @ Krylov ~/.julia/packages/Krylov/PApE9/src/krylov_solvers.jl:1483

I dont know if this is an issue that should be solved in BlockArrays or in this package.

amontoison commented 2 years ago

HI @lijas! Krylov.jl allocates a workspace based on the type of b. It's the best way that we find to support complex numbers, multiple precision and GPUs. The issue here is that Krylov.jl wants to use BlockVector for all internal vectors used by GMRES and they can't be initialized like Vector.

n = 10
Vector(undef, n)
CuVector(undef, n)
BlockVector(undef, n) # not working because n must be a tuple

But in that case, Krylov.jl should just use basic vectors (same problem in #573) and we just need to define ktypeof for this kind of AbstractVector. It's more efficient because BLAS routines are used for dense vectors. I will add a a note about it in the documentation.

import Krylov.ktypeof
ktypeof(v::BlockVector{T}) where T = Vector{T}

using Krylov
using BlockArrays

K = rand(10,10); 
B = rand(2, 10)
b = rand(10 + 2);

b = BlockVector(rand(Float64,12), [10,2])
A = BlockArray(rand(Float64,12,12), [10,2], [10,2])
A[Block(1,1)] = K
A[Block(1,2)] = B'
A[Block(2,1)] = B

x, stats = Krylov.gmres(A, b)  # it works !
norm(b - A * x)

For the preconditioner, it's a good idea to use a block diagonal preconditioner like you suggest with an incomplete LU of K and an approximation of the invert of the Schur complement B * K^{-1} * B'. Like the input operator A, the preconditioners (M or N) just need to implement mul! (and not ldiv!). Therefore, the M argument in GMRES should your preconditioner P.

using LinearOperators, IncompleteLU
n, m = size(B)
F_ilu = ilu(sparse(K))  # IncompleteLU factorization of K
M1 = LinearOperator(Float64, m, m, false, false, (y, v) -> ldiv!(y, F_ilu, v))
S =  B * inv(Diagonal(K)) * B' # Approximation of the Schur complement
F_lu = lu(S)
M2 = LinearOperator(Float64, n, n, false, false, (y, v) -> ldiv!(y, F_lu, v))
M = BlockDiagonalOperator(M1, M2)

x, stats = Krylov.gmres(A, b, M=M)
norm(b - A * x)

We can also define something like this to still use BlockArray:

# M.x and M.y can be created with similar(b)

function mul!(y::Vector, M::CostumBlockPrecon, x::Vector)
    M.x .= x
    M.y[Block(1)] = M.K_ilu \ M.x[Block(1)]
    M.y[Block(2)] = (M.B * (inv(Diagonal(M.K)) * M.B') \ *M.x[Block(2)]
    y .= M.y
end

Input and output vectors x and y are internal vectors of the workspace of GMRES, their type is Vector as explained before.

amontoison commented 2 years ago

If you are able to compute a complete factorization of K, you can also use the preconditioner P^{-1} = [K^{-1} 0; 0 diag(B'*K^{-1}*B)^{-1}] and use TriMR / GPMR instead of MINRES / GMRES. They are specialized for 2x2 block matrices and directly take as input the matrices K and B. The main advantage of these methods is that they will ALWAYS converge faster than MINRES or GMRES.

lijas commented 2 years ago

Hi! For fun, I tried digging in to the GMRES solver a bit to try to make it compatible BlockArrays. One thing i did was to change the initialization of vectors from Vector(undef, n) to similar(b) or zeros(b). Any reason why this would be worse/less general?

I also needed to remove type annotation of DenseArray in some places.

Doing these two changes, I was able to get a small BlockArray system to pass the solver and converge (I have no idea if it is the correct answer though :smile: ).

So with a bit of work it would be quite easy (I think) to have compatibility with BlockArrays. However, a bigger question is if it is worth It? It seems like iterative solvers + block arrays is a pretty uncommon thing to do... I cant really find that much information about it when searching for it. I mainly wanted it in order to to design the preconditioner in a better way, but perhaps there are better ways of doing that?

Another big issue is that matrix-vector multiplication in BlockArrays is extremely slow for some reason. I have know idea why, but that would need to fixed in BlockArrays as a first step.

Regarding the TriMR solver. I dont quite understand the use of it..The main point of using iterative solvers is when factorization of the K matrix is infeasible, right? So if I can get the factorization of K I can basically just solve the Schur compliment of the larger system much faster with a direct solution method. In my caseK is much larger than B, so is TriMR used when they are of similar size?

amontoison commented 2 years ago

The issue with similar(b) is that you want dense vectors (Vector) even if your right-hand side b is sparse (SparseVector). The other issue with BlockVector is that similar(b, n) returns a Vector and not a BlockVector. In some methods n != length(b) and we need the argument n.

I will find a way to support BlockVector as a right-hand side but it's important to still use DenseVector in the workspace for performances. BlockArrays are user-friendly to design the preconditioner but they require that vectors in the workspace have the same (block) structure, which we want to avoid. It should not be too hard to replace the blocks by some view if x and y are just basic vectors.

For TriMR, it depends of the size / structure of K and also if you need to solve multiple systems with the same K block but different B blocks.

lijas commented 2 years ago

Ahh, it makes sense to not use similar(b) for sparse arrays. And BlockMatrix*DenseVector can probably be made efficient, so no need to have BlockVectors in the workspace.