Support ComponentArrays

[jmulcahy]

It would be nice if Krylov.jl supported ComponentArrays.jl (https://github.com/jonniedie/ComponentArrays.jl). ComponentArrays let you create vectors that for the purposes of linear algebra are 1D but also have a structured overlay. For example:

using ComponentArrays, LinearMaps

function grad2D(data)
    ∇row = diff(data.arr, dims=1)
    ∇col = diff(data.arr, dims=2)
    ComponentArray(∇row=∇row, ∇col=∇col)
end

data = ComponentArray(arr=randn(512, 512))
out = grad2D(data)

A = LinearMap(grad2D, length(out), length(data))

A * data == out

It's more natural to write grad2D in terms of the 2D array rather than some flattened form. It's also convenient to access ∇row and ∇col via out.∇row and out.∇col rather than raw indexing into the flat array. If you plug something like this into Krylov.jl it breaks somewhere around here: https://github.com/JuliaSmoothOptimizers/Krylov.jl/blob/main/src/krylov_solvers.jl#L1342. It looks like you'd need to expand the type specialization for every solver to at least support different types for Aᴴu and Av. I'm not sure what other stuff would break downstream.

[jmulcahy]

I was looking through the issues, and I came across this: #646. Not that my issue should hold up that one, but I think they're at least somewhat linked because the underlying question is about object storage and typing within the package. I'm guessing implementing either will run into similar problems related to properly initializing caches with the right types.

I'll add that I think differentiability of the solvers would be an amazing feature. :)

[amontoison]

Hi @jmulcahy!
I think that you should keep Vector{T} for the type of the vectors in the workspace. It will be more efficient for BLAS1 operations. Is it not to possible to do ComponentArray * Vector?

[jmulcahy]

Here's a complete example. I swapped diff for cumsum because the Hermitian is a little more straightforward:

using ComponentArrays, LinearMaps, Krylov

function cumsum2D(data)
    Σrow = cumsum(data.arr, dims=1)
    Σcol = cumsum(data.arr, dims=2)
    ComponentArray(Σrow=Σrow, Σcol=Σcol)
end

function cumsum2Dᴴ(sums)
    arr = reverse(cumsum(reverse(sums.Σrow, dims=1), dims=1), dims=1)
    arr .+= reverse(cumsum(reverse(sums.Σcol, dims=2), dims=2), dims=2)
    ComponentArray(arr=arr)
end

data = ComponentArray(arr=randn(8, 8))
out = cumsum2D(data)

A = LinearMap(cumsum2D, cumsum2Dᴴ, length(out), length(data))

b = A * data

lsmr(A, b)

Stacktrace:

ERROR: MethodError: Cannot `convert` an object of type UndefInitializer to an object of type Vector{Float64}
Closest candidates are:
  convert(::Type{T}, ::LinearAlgebra.Factorization) where T<:AbstractArray at C:\Users\John\.julia\juliaup\julia-1.8.5+0.x64.w64.mingw32\share\julia\stdlib\v1.8\LinearAlgebra\src\factorization.jl:58
  convert(::Type{<:Array}, ::ComponentArray) at C:\Users\John\.julia\packages\ComponentArrays\GsPr2\src\similar_convert_copy.jl:58
  convert(::Type{T}, ::AbstractArray) where T<:Array at array.jl:617
  ...
Stacktrace:
 [1] ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}}(data::UndefInitializer, axes::Int64)
   @ ComponentArrays C:\Users\John\.julia\packages\ComponentArrays\GsPr2\src\componentarray.jl:35
 [2] LsmrSolver(m::Int64, n::Int64, S::Type{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}}}; window::Int64)
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\krylov_solvers.jl:1340
 [3] LsmrSolver(A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}}; window::Int64)
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\krylov_solvers.jl:1358
 [4] lsmr(A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}}; window::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\lsmr.jl:125
 [5] lsmr(A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}})
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\lsmr.jl:124
 [6] top-level scope
   @ REPL[9]:1

I believe BLAS should still work with ComponentArrays since they're linear memory buffers. I believe that's one of the design goals of the package. My current workaround is to put only "normal" Vectors into Krylov.jl, but it requires some extra bookkeeping to properly reshape the data inside the maps and for final output. It's especially cumbersome because as it is now, there's no way to cumsum2D and cumsum2Dᴴ to know the array dimensions without adding some extra wrappers. Maybe there is another nice way to handle these sorts of problems?

[jmulcahy]

Example with all the reshaping inserted (copies not avoided for brevity):

using LinearMaps, Krylov

function cumsum2D_reshaping(arr)
    reshaped = reshape(arr, 8, 8)
    sums = similar(arr, 8, 8, 2)
	
    sums[:, :, 1] = cumsum(reshaped, dims=1)
    sums[:, :, 2] = cumsum(reshaped, dims=2)

    sums[:]
end

function cumsum2Dᴴ_reshaping(sums)
    reshaped = reshape(sums, 8, 8, 2)
    arr = reverse(cumsum(reverse(reshaped[:, :, 1], dims=1), dims=1), dims=1)
    arr .+= reverse(cumsum(reverse(reshaped[:, :, 2], dims=2), dims=2), dims=2)

    arr[:]
end

data = randn(8, 8)
out = cumsum2D_reshaping(data[:])

A = LinearMap(cumsum2D_reshaping, cumsum2Dᴴ_reshaping, length(out), length(data))

b = A * data[:]

lsmr(A, b)

[amontoison]

Thanks for the examples, I'm currently traveling and I forgot my laptop charger so I can't test them until Thursday...
Based of the error returned, the issue is during the creation of the KrylovSolver.
Krylov.jl wants to use a ComponentArray for the storage type of vectors in the workspace but the default array constructor is not working for this array type.
Is it possible to test this code?

using ComponentArrays, LinearMaps, Krylov

function cumsum2D(data)
    Σrow = cumsum(data.arr, dims=1)
    Σcol = cumsum(data.arr, dims=2)
    ComponentArray(Σrow=Σrow, Σcol=Σcol)
end

function cumsum2Dᴴ(sums)
    arr = reverse(cumsum(reverse(sums.Σrow, dims=1), dims=1), dims=1)
    arr .+= reverse(cumsum(reverse(sums.Σcol, dims=2), dims=2), dims=2)
    ComponentArray(arr=arr)
end

data = ComponentArray(arr=randn(8, 8))
out = cumsum2D(data)

A = LinearMap(cumsum2D, cumsum2Dᴴ, length(out), length(data))
b = A * data

m, n = size(A)
solver = LsmrSolver(m, n, Vector{Float64})
x, stats = lsmr!(solver,A, b)

[jmulcahy]

Thanks for taking a look. That code hits this error on the last line:

ERROR: ktypeof(b) is not a subtype of Vector{Float64}
Stacktrace:
 [1] error(s::String)
   @ Base .\error.jl:35
 [2] lsmr!(solver::LsmrSolver{Float64, Float64, Vector{Float64}}, A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}}; M::LinearAlgebra.UniformScaling{Bool}, N::LinearAlgebra.UniformScaling{Bool}, ldiv::Bool, sqd::Bool, λ::Float64, radius::Float64, etol::Float64, axtol::Float64, btol::Float64, conlim::Float64, atol::Float64, rtol::Float64, itmax::Int64, verbose::Int64, history::Bool, callback::Krylov.var"#296#298", iostream::Core.CoreSTDOUT)
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\lsmr.jl:163
 [3] lsmr!(solver::LsmrSolver{Float64, Float64, Vector{Float64}}, A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}})
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\lsmr.jl:139
 [4] top-level scope
   @ REPL[11]:1

Looking at the LsmrSolver constructor, I think some of the internal members would need different typing. Because of the extra structured information, I think some would need to be typeof(data) and others would be typeof(b) even though everything is backed by a Vector{Float64}. cumsum2D and cumsum2Dᴴ can't accept Vector{Float64}.

[amontoison]

I would like to find a solution without multiple storage types for the vectors in the KrylovSolvers.
I will work on it asap (Thursday).

Your issue is similar to #605.

[jmulcahy]

That is similar., although I think BlockArrays only deals with the structured data problem and not N-dimensional arrays. ComponentArrays can still be jagged while aliasing contiguous memory, which I think BlockArrays can't do.

Along those lines, I found this: https://github.com/JeffFessler/LinearMapsAA.jl which creates linear mappings that return multi-dimensional arrays. I haven't looked closely, but I imagine it doesn't work with any of the standard solvers.

I have a workaround for this, so no rush. I thought I'd file an issue since it seems like this ergonomics problem with structured data and multi-dimensional arrays might be common. In doing some thinking and prototyping, I thought ComponentArray approach that it appears SciML has adopted broadly looked elegant. Thanks again for taking a look.

[amontoison]

Thanks for taking a look. That code hits this error on the last line:

ERROR: ktypeof(b) is not a subtype of Vector{Float64}
Stacktrace:
 [1] error(s::String)
   @ Base .\error.jl:35
 [2] lsmr!(solver::LsmrSolver{Float64, Float64, Vector{Float64}}, A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}}; M::LinearAlgebra.UniformScaling{Bool}, N::LinearAlgebra.UniformScaling{Bool}, ldiv::Bool, sqd::Bool, λ::Float64, radius::Float64, etol::Float64, axtol::Float64, btol::Float64, conlim::Float64, atol::Float64, rtol::Float64, itmax::Int64, verbose::Int64, history::Bool, callback::Krylov.var"#296#298", iostream::Core.CoreSTDOUT)
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\lsmr.jl:163
 [3] lsmr!(solver::LsmrSolver{Float64, Float64, Vector{Float64}}, A::LinearMaps.FunctionMap{Float64, typeof(cumsum2D), typeof(cumsum2Dᴴ)}, b::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(Σrow = ViewAxis(1:64, ShapedAxis((8, 8), NamedTuple())), Σcol = ViewAxis(65:128, ShapedAxis((8, 8), NamedTuple())))}}})
   @ Krylov C:\Users\John\.julia\packages\Krylov\0FaHW\src\lsmr.jl:139
 [4] top-level scope
   @ REPL[11]:1

Looking at the LsmrSolver constructor, I think some of the internal members would need different typing. Because of the extra structured information, I think some would need to be typeof(data) and others would be typeof(b) even though everything is backed by a Vector{Float64}. cumsum2D and cumsum2Dᴴ can't accept Vector{Float64}.

@jmulcahy
I have some issues with ComponentArrays, similar(v, n) where v is a ComponentVector returns a Vector.
I wanted to replace S(undef, n) by similar(S, n) but we still have the same issue.

[jmulcahy]

Maybe there is some way to utilize dispatch to instead hit similar(S) instead of similar(S, n)? I think the underlying issue is that a ComponentArray type must have fixed length. For example, this works:

data = ComponentArray(x=randn(16, 16), y=randn(32, 32), z=1)
similar(data)

Under the hood, it's making a Vector{Float64} of length 16^2 + 32^2 + 1. That behavior is very different than most Julia array types since most don't encode their length in their type.

[amontoison]

Ok but for solvers specialized for rectangular matrices such as LSMR, you need to allocate vectors of length n and other ones of length m where m, n = size(A).
Do you use LSMR just to illustrate the issue or is it the method that you need?

[jmulcahy]

I need it because some of my problems are indefinite and rectangular. The diff and cumsum examples are relevant, but there are also other mappings that have more complicated behavior than them. The diff one, for example, appears in image processing problems where you measure a gradient and want to retrieve the original image. I'd say, though, that this would be a nice feature across solvers, although potentially not easy to implement.

I think one way to do this would involve adjusting the type signature of LsmrSolver and KrylovSolver:

mutable struct LsmrSolver{T,FC,M,N} <: KrylovSolver{T,FC,M,N}
  m       :: Int
  n       :: Int
  x       :: N
  Nv      :: N
  Aᴴu     :: N
  h       :: N
  hbar    :: N
  Mu      :: M
  Av      :: M
  u       :: M
  v       :: N
  err_vec :: Vector{T}
  stats   :: LsmrStats{T}
end

So M and N are specifying the output types from multiplying by A and Aᴴ, respectively. I'm guessing here slightly since I'm not that familiar with the internals of LSMR, especially u and v since they're initialized as empty. In the cumsum example it would be M == typeof(out) and N == typeof(data). The equivalent to the current design would be M == N.

[amontoison]

Hi @jmulcahy! Sorry for the delay...
You don't need to add an additional parameter for LsmrSolver or KrylovSolver because M and N are both ComponentVector{T}.
I have the following example for you:

using ComponentArrays

x = ComponentArray(a=5.0, b=[(a=20.0, b=0.0), (a=33.0, b=0.0)])
y = ComponentArray(a=[(a=3.0, b=1.0, c=7.0)], b=20.0)

n = length(x)
m = length(y)
window = 5
T = FC = Float64
S = ComponentVector{Float64}

LsmrSolver{T,FC,S}(
	m,
	n,
	similar(x),
	similar(x),
	similar(x),
	similar(x),
	similar(x),
	similar(y),
	similar(y),
	similar(y),
	similar(x),
	zeros(T, window),
	LsmrStats(0, false, false, T[], T[], zero(T), zero(T), zero(T), zero(T), zero(T), "unknown")
)

The best solution is to definite a new LsmrSolver constructor for your application.
Because of the structure of the ComponentArrays, it's impossible to determine the output type for multiplications with Aᴴ AND A if the system is rectangular because we can only use the type of b .
We only have one structure (M or N).

It means that we can only use the in-place variant:

using ComponentArrays, LinearMaps, Krylov

function cumsum2D(data)
    Σrow = cumsum(data.arr, dims=1)
    Σcol = cumsum(data.arr, dims=2)
    ComponentArray(Σrow=Σrow, Σcol=Σcol)
end

function cumsum2Dᴴ(sums)
    arr = reverse(cumsum(reverse(sums.Σrow, dims=1), dims=1), dims=1)
    arr .+= reverse(cumsum(reverse(sums.Σcol, dims=2), dims=2), dims=2)
    ComponentArray(arr=arr)
end

data = ComponentArray(arr=randn(8, 8))
out = cumsum2D(data)

A = LinearMap(cumsum2D, cumsum2Dᴴ, length(out), length(data))
b = A * data

m, n = size(A)
window = 5
T = FC = eltype(b)
S = ComponentVector{T}

solver = LsmrSolver{T,FC,S}(
	m,
	n,
	similar(data),
	similar(data),
	similar(data),
	similar(data),
	similar(data),
	similar(out),
	similar(out),
	similar(out),
	similar(data),
	zeros(T, window),
	LsmrStats(0, false, false, T[], T[], zero(T), zero(T), zero(T), zero(T), zero(T), "unknown")
)
solver = lsmr!(solver, A, b)

┌──────────────────┬────────────────────────┬──────────────────┐
│        LsmrSolver│              nrows: 128│         ncols: 64│
├──────────────────┼────────────────────────┼──────────────────┤
│Precision: Float64│       Architecture: GPU│Storage: 218 bytes│
├──────────────────┼────────────────────────┼──────────────────┤
│         Attribute│                    Type│              Size│
├──────────────────┼────────────────────────┼──────────────────┤
│                 m│                   Int64│           8 bytes│
│                 n│                   Int64│           8 bytes│
│                 x│ComponentVector{Float64}│           8 bytes│
│                Nv│ComponentVector{Float64}│           8 bytes│
│               Aᴴu│ComponentVector{Float64}│           8 bytes│
│                 h│ComponentVector{Float64}│           8 bytes│
│              hbar│ComponentVector{Float64}│           8 bytes│
│                Mu│ComponentVector{Float64}│           8 bytes│
│                Av│ComponentVector{Float64}│           8 bytes│
│                 u│ComponentVector{Float64}│           8 bytes│
│                 v│ComponentVector{Float64}│           8 bytes│
│           err_vec│         Vector{Float64}│          40 bytes│
│             stats│      LsmrStats{Float64}│          90 bytes│
└──────────────────┴────────────────────────┴──────────────────┘

LsmrStats
 niter: 30
 solved: true
 inconsistent: false
 residuals: []
 Aresiduals: []
 residual: 1.8345314994833746e-6
 Aresidual: 1.718353653426171e-6
 κ₂(A): 5.973733856419315
 ‖A‖F: 22.260207728748714
 xNorm: 8.39488702040155
 status: found approximate zero-residual solution

julia> norm(b - A * solver.x)
1.8345314996956182e-6

[amontoison]

import Base.sizeof
import Krylov.ksizeof

function sizeof(v :: ComponentVector)
  type = typeof(v)
  nfields = fieldcount(type)
  storage = 0
  for i = 1:nfields
    field_i = getfield(v, i)
    size_i = Krylov.ksizeof(field_i)
    storage += size_i
  end
  return storage
end

solver
┌──────────────────┬────────────────────────┬──────────────────┐
│        LsmrSolver│              nrows: 128│         ncols: 64│
├──────────────────┼────────────────────────┼──────────────────┤
│Precision: Float64│       Architecture: GPU│Storage: 6.143 KiB│
├──────────────────┼────────────────────────┼──────────────────┤
│         Attribute│                    Type│              Size│
├──────────────────┼────────────────────────┼──────────────────┤
│                 m│                   Int64│           8 bytes│
│                 n│                   Int64│           8 bytes│
│                 x│ComponentVector{Float64}│         512 bytes│
│                Nv│ComponentVector{Float64}│         512 bytes│
│               Aᴴu│ComponentVector{Float64}│         512 bytes│
│                 h│ComponentVector{Float64}│         512 bytes│
│              hbar│ComponentVector{Float64}│         512 bytes│
│                Mu│ComponentVector{Float64}│        1024 bytes│
│                Av│ComponentVector{Float64}│        1024 bytes│
│                 u│ComponentVector{Float64}│        1024 bytes│
│                 v│ComponentVector{Float64}│         512 bytes│
│           err_vec│         Vector{Float64}│          40 bytes│
│             stats│      LsmrStats{Float64}│          90 bytes│
└──────────────────┴────────────────────────┴──────────────────┘

[jmulcahy]

Ahh, very cool. I thought the information about where elements were in the ComponentVectors were part of the type, so would then require a specialization for each one. I'll try this out.

[jmulcahy]

Sorry about the delay, but I did try this and it worked perfectly. Thanks for your help on this.

[amontoison]

Great! After #706, the updated constructors of the KrylovSolvers should work with ComponentArrays.