Support generic arrays

[amontoison]

close #605

using Krylov
using BlockArrays
K = rand(10,10)
B = rand(2, 10)

b = BlockVector(rand(Float64,12), [10,2])
2-blocked 12-element BlockVector{Float64}:
 0.8437110010150571  
 0.4530016441000567  
 0.686194576648592   
 0.34254162007097566 
 0.09544196921306347 
 0.5010761791056787  
 0.11907189994954914 
 0.17470076682895075 
 0.9550777571644686  
 0.4680632206046831  
 ────────────────────
 0.012863148806336877
 0.8840547265290272  

A = BlockArray(rand(Float64,12,12), [10,2], [10,2])
2×2-blocked 12×12 BlockMatrix{Float64}:
 0.116554   0.0703249  0.875842  0.506043   0.373476  0.23967    0.601902   0.860072   0.964966  0.436033    │  0.713543  0.470241
 0.799259   0.37769    0.346566  0.259132   0.635035  0.0446249  0.964017   0.92735    0.814636  0.669492    │  0.298945  0.921794
 0.816827   0.134661   0.127786  0.971469   0.840193  0.49803    0.892349   0.360627   0.160712  0.00137863  │  0.913815  0.846428
 0.20443    0.948703   0.137055  0.609844   0.164636  0.724305   0.90073    0.875758   0.862051  0.755453    │  0.145395  0.078032
 0.321294   0.687613   0.237468  0.0524443  0.833145  0.693929   0.237879   0.184222   0.625539  0.709003    │  0.515103  0.907079
 0.135291   0.879867   0.879918  0.279797   0.561364  0.963455   0.0180151  0.177779   0.994638  0.848967    │  0.21419   0.377067
 0.420909   0.864644   0.246007  0.709511   0.239352  0.279221   0.789019   0.526006   0.956405  0.581049    │  0.1384    0.642351
 0.709593   0.22967    0.274482  0.0517252  0.824723  0.934985   0.229917   0.514245   0.219793  0.583001    │  0.39413   0.196227
 0.350237   0.688092   0.994518  0.896093   0.103099  0.530616   0.420392   0.262814   0.622369  0.268439    │  0.169126  0.198431
 0.956745   0.696715   0.838499  0.576819   0.719354  0.250461   0.416511   0.963868   0.883324  0.630397    │  0.954737  0.9058  
 ────────────────────────────────────────────────────────────────────────────────────────────────────────────┼────────────────────
 0.438206   0.566904   0.495495  0.183546   0.625943  0.256114   0.116151   0.725794   0.93137   0.564458    │  0.233888  0.410463
 0.0936359  0.838417   0.210066  0.848525   0.834449  0.314861   0.849024   0.0837541  0.379385  0.96699     │  0.788724  0.133835

A[Block(1,1)] = K;
A[Block(1,2)] = B';
A[Block(2,1)] = B;

x, stats = Krylov.gmres(A, b)
([-1.6804307117446322, 1.2278956836087573, 0.7522500694970368, -1.1813882696427558, -1.7308386536458193, 1.3297218781185194, 0.7128256864623073, 0.43277609013448326, 1.819601362734177, 1.6564852399189884, -1.897804683660123, -0.6132195202643419], SimpleStats
 niter: 12
 solved: true
 inconsistent: false
 residuals: []
 Aresiduals: []
 κ₂(A): []
 status: solution good enough given atol and rtol
)

x
2-blocked 12-element BlockVector{Float64}:
 -1.6804307117446322 
  1.2278956836087573 
  0.7522500694970368 
 -1.1813882696427558 
 -1.7308386536458193 
  1.3297218781185194 
  0.7128256864623073 
  0.43277609013448326
  1.819601362734177  
  1.6564852399189884 
 ────────────────────
 -1.897804683660123  
 -0.6132195202643419 

norm(b - A * x)
2.8592284672503768e-15

[amontoison]

Thanks @amontoison for the PR, that looks great!

Could you give us a bit more details about the PR? If I understand correctly you:

* replace `S(undef, n)` by `similar(S, n)`

* remove some `axpy!(s, x, y)` and `kaxpy!(n, Complex{T}(s), x, dx, y, dy)` (not sure why though?)

* Add a specific constructor if the right-hand side is not a `DenseVector`

Yes, I can give more details about the pull request.
With the PR #705, I allowed more generic storage type for the Krylov workspaces (all structures that are a subtype of KrylovSolver).
In the past, only subtypes of DenseVector were allowed (Vector, CuVector, ROCVector, etc...).

With this PR, I change how the vectors in the Krylov workspaces are initialized and allocated.

• I replaced S(undef, n) by similar(S, n). It's exactly the same thing because similar(S, n) calls S(undef, n) here but I prefer this new syntax.

• I replaced m and n in all similar(S, n) and similar(S, m) by ixm and ixn. ixn and ixm are the axes returned the new functions kaxe and kaxes. Fancy vectors such as BlockVector or PartitionedVector that are not a subtype of DenseVector require them to be initialized.
  To not break previous implementations based on Krylov.jl, kaxe and kaxes calls length(b) and size(A) if b is a dense vector. The users only need to implement the function axes if it's mandatory (exotic storage types that are not a subtype of DenseVector).

• Because 0 is an integer, similar(S, 0) can't be used for some storage types that are not a subtype of a DenseVector. I must allocate all optional vectors required for warm-start, regularization or preconditioners in the Krylov workspaces.
  Do you have a better solution?

• Because FOM, GMRES, FGMRES and GPMR dynamically increase the size of Krylov basis, I updated how the basis vectors are allocated in these Krylov methods to be consistent with the modifications in src/krylov_solvers.Jl.

• axpy! and axpby! have a generic fallback for any AbstractVector but it requires scalar indexing. It was not an issue before this PR because I interfaced / implemented specialized methods in GPUArrays.jl, CUDA.jl, etc... But some packages such as PartitionedArrays deactivated scalar indexing, which means that I need to use the broadcast for this kind of storage types. I found a neat solution with the multiple dispatch.

It probably misses some unit tests showing/testing the new data types handled by these changes.

I will add a buildkite build that will test that the new data types are handled by these changes.

[codecov]

Codecov Report

Base: 98.22% // Head: 98.20% // Decreases project coverage by -0.03% ⚠️

Coverage data is based on head (a346c40) compared to base (49bdbec).
Patch coverage: 99.57% of modified lines in pull request are covered.

❗ Current head a346c40 differs from pull request most recent head 68d9f43. Consider uploading reports for the commit 68d9f43 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #706      +/-   ##
==========================================
- Coverage   98.22%   98.20%   -0.03%     
==========================================
  Files          38       38              
  Lines        6659     6695      +36     
==========================================
+ Hits         6541     6575      +34     
- Misses        118      120       +2     

Impacted Files	Coverage Δ
src/krylov_utils.jl	90.81% <77.77%> (-0.86%)	⬇️
src/fgmres.jl	100.00% <100.00%> (ø)
src/fom.jl	100.00% <100.00%> (ø)
src/gmres.jl	100.00% <100.00%> (ø)
src/gpmr.jl	100.00% <100.00%> (ø)
src/krylov_solvers.jl	99.88% <100.00%> (+<0.01%)	⬆️
src/cr.jl	66.35% <0.00%> (-0.16%)	⬇️
src/craigmr.jl	95.94% <0.00%> (-0.03%)	⬇️
src/lslq.jl	96.68% <0.00%> (-0.02%)	⬇️
src/lsqr.jl	97.81% <0.00%> (-0.02%)	⬇️
... and 30 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[dpo]

Suggested change

	`axes(A)` could not be defined for some linear operators and we only want to call it for fancy arrays.
	`axes(A)` may not be defined for some linear operators and we only want to call it for fancy arrays.

I think it would be reasonable to require axes() to be implemented. It's part of the array API. You can recover the dimensions with length.(axes(A)), which works for dense arrays too.

[amontoison]

A can be any structure that models a matrix of size m x n. I'm not sure that it's always relevant to define axes.
I have an example in mind:

structure MyJacobian{T} where {T}
  m::Int
  n::Int
  F::Function

MyJacobian is only defined such that mul!(y::Vector{T}, A::MyJacobian{T}, x::Vector{T}) can be overloaded and mul! performs Jacobian-vector products with automatic differentiation tools.
The user will probably never implement axes(A::MyJacobian).

[dpo]

What I'm saying is that it's not for Krylov.jl to try to guess what the implementation of axes() should be for a user type. We could simply document that users must implement axes().

[amontoison]

We can but it's a breaking change. The next release of Krylov.jl must be v"10.0" if we use axes instead of kaxes.

[dpo]

why? you just introduced kaxes, didn't you?

[amontoison]

Yes, I introduced kaxes to avoid that but you suggest to just use axes directly, am I wrong?
axes or kaxes are called when we create a KrylovSolver, which means every time that we solve a linear problem with a method that is not in-place such as cg(A, b).
Because only S <: DenseVector was supported before, I'm sure to not break any code if kaxes dispatches to size.

[tmigot]

Because 0 is an integer, similar(S, 0) can't be used for some storage types that are not a subtype of a DenseVector. I must allocate all optional vectors required for warm-start, regularization or preconditioners in the Krylov workspaces.
Do you have a better solution?

Thanks for the details. Isn't that strange the empty vector is not defined even for sparse arrays ? In the worst-case could you do similar(S, 1) instead to avoid depending on the size?

[tmigot]

A few more comments, and then looks good to me, thanks!

[amontoison]

Because 0 is an integer, similar(S, 0) can't be used for some storage types that are not a subtype of a DenseVector. I must allocate all optional vectors required for warm-start, regularization or preconditioners in the Krylov workspaces. Do you have a better solution?

Thanks for the details. Isn't that strange the empty vector is not defined even for sparse arrays ? In the worst-case could you do similar(S, 1) instead to avoid depending on the size?

I don't understand your comment Tangi. 1 is also an integer like 0. The problem is not the value 0, I just want to create a vector that satisfy ::S and take the minimum amount of memory. Some vectors require the constructor with the axes to be initialized.

For the sparse arrays, I will add a check to insure that S is never a sparse arrays. It's for that reason that I added S <: DenseVector everywhere in the past.

[tmigot]

don't understand your comment Tangi. 1 is also an integer like 0. The problem is not the value 0, I just want to create a vector that satisfy ::S and take the minimum amount of memory. Some vectors require the constructor with the axes to be initialized.

Ok, I didn't understand the problem at first. Do you have an example that would not work?

[amontoison]

don't understand your comment Tangi. 1 is also an integer like 0. The problem is not the value 0, I just want to create a vector that satisfy ::S and take the minimum amount of memory. Some vectors require the constructor with the axes to be initialized.

Ok, I didn't understand the problem at first. Do you have an example that would not work?

Yes, I do.

using BlockArrays, Krylov

T = Float64
n = 10
x = BlockVector(rand(T,n), [n-2,2])
S = Krylov.ktypeof(x)
similar(S, 0)

[tmigot]

don't understand your comment Tangi. 1 is also an integer like 0. The problem is not the value 0, I just want to create a vector that satisfy ::S and take the minimum amount of memory. Some vectors require the constructor with the axes to be initialized.

Ok, I didn't understand the problem at first. Do you have an example that would not work?

Yes, I do.

using BlockArrays, Krylov

T = Float64
n = 10
x = BlockVector(rand(T,n), [n-2,2])
S = Krylov.ktypeof(x)
similar(S, 0)

Well... that won't really help but you could create an issue on this package and on the packages where the empty vector is not defined. To me, it looks like a very natural property for any vector type,