Can't ldiv! a QR factorization with a BlockVector

[simonbyrne]

This came up while trying to use an implicit solver with OrdinaryDiffEq.

A simple reproducible example:

julia> using BlockArrays, LinearAlgebra

julia> x = BlockArray(rand(3))
1-blocked 3-element BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}}:
 0.46021136136833674
 0.5138272715735523
 0.581072254943535

julia> QR = qr(rand(3,3))
LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}
Q factor:
3×3 LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}}:
 -0.682104  -0.0521373  -0.729394
 -0.477444  -0.723754    0.498224
 -0.553878   0.688086    0.468783
R factor:
3×3 Matrix{Float64}:
 -0.998175  -0.835741  -0.730764
  0.0        0.429548   0.36884
  0.0        0.0        0.0893908

julia> ldiv!(QR, x)
ERROR: Overload materialize!(::Lmul{ArrayLayouts.AdjQRCompactWYQLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}})
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:33
 [2] materialize!(M::ArrayLayouts.Lmul{ArrayLayouts.AdjQRCompactWYQLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, BlockArrays.BlockLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, Adjoint{Float64, LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}}}, BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}}})
   @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/factorizations.jl:127
 [3] lmul!(A::Adjoint{Float64, LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}}}, B::BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}})
   @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/lmul.jl:47
 [4] materialize!(L::ArrayLayouts.Ldiv{ArrayLayouts.QRCompactWYLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, BlockArrays.BlockLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}, BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}}})
   @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/factorizations.jl:25
 [5] ldiv!
   @ ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:89 [inlined]
 [6] ldiv!(A::LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}, x::BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}})
   @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:135
 [7] top-level scope
   @ REPL[8]:1

Unfortunately I don't understand how ArrayLayouts works: can you give some pointers on what needs to be done?

[simonbyrne]

Actually it fails with just generic ldiv!:

julia> ldiv!(BlockArray(rand(3,3)), BlockArray(rand(3)))
ERROR: Overload materialize!(::Lmul{ArrayLayouts.AdjQRCompactWYQLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}})
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] materialize!(M::ArrayLayouts.Lmul{ArrayLayouts.AdjQRCompactWYQLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, BlockArrays.BlockLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, Adjoint{Float64, LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}}}, BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/factorizations.jl:127
  [3] lmul!(A::Adjoint{Float64, LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}}}, B::BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/lmul.jl:47
  [4] materialize!(L::ArrayLayouts.Ldiv{ArrayLayouts.QRCompactWYLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, BlockArrays.BlockLayout{ArrayLayouts.DenseColumnMajor, ArrayLayouts.DenseColumnMajor}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}, BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/factorizations.jl:25
  [5] ldiv!
    @ ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:89 [inlined]
  [6] ldiv!(A::LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}, x::BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:135
  [7] __ldiv!(#unused#::BlockMatrix{Float64, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, F::LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}, B::BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:76
  [8] _ldiv!
    @ ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:77 [inlined]
  [9] materialize!
    @ ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:95 [inlined]
 [10] ldiv!
    @ ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:89 [inlined]
 [11] ldiv!(A::BlockMatrix{Float64, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, x::BlockVector{Float64, Vector{Vector{Float64}}, Tuple{Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/3lnt1/src/ldiv.jl:130
 [12] top-level scope
    @ REPL[91]:1

[dlfivefifty]

This is unfortunately a non-trivial question as we need to think about breaking a QRCompactWY into blocks. In fact, geqrt allows specifying block size during factorization which might help:

http://www.netlib.org/lapack/explore-html/dd/d9a/group__double_g_ecomputational_gaddcf152e87deec6123a1899f6f51101e.html#gaddcf152e87deec6123a1899f6f51101e

Using a QRPacked factorization is more general but will be slow as it will default to generic solve.

[dlfivefifty]

Can you use a PseudoBlockVector? This should work

[jishnub]

Could you elaborate on how to use a PseudoBlockArray in this context?

I am trying to directly replace mortar with PseudoBlockArray in my code, which doesn't seem to work:

julia> A = PseudoBlockArray(reshape([rand(1,1), rand(1,1), rand(1,1), rand(1,1)],2,2))
1×1-blocked 2×2 PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
 [0.331438]  [0.624061]
 [0.919634]  [0.255869]

julia> B = rand(2, 2);

julia> A \ B
ERROR: MethodError: no method matching zero(::Type{Matrix{Float64}})
Closest candidates are:
  zero(::Union{Type{P}, P}) where P<:Dates.Period at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/Dates/src/periods.jl:53
  zero(::Zeros{T, N, Axes} where Axes) where {T, N} at /home/jishnu/.julia/packages/FillArrays/NBUnJ/src/FillArrays.jl:543
  zero(::SA) where SA<:StaticArrays.StaticArray at /home/jishnu/.julia/packages/StaticArrays/vxjOO/src/linalg.jl:93
  ...
Stacktrace:
  [1] _qreltype(#unused#::Type{Matrix{Float64}})
    @ LinearAlgebra /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:333
  [2] qr(::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ LinearAlgebra /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:415
  [3] qr(::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ LinearAlgebra /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:414
  [4] __qr(layout::ArrayLayouts.DenseColumnMajor, lengths::Tuple{Int64, Int64}, A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}; kwds::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:281
  [5] __qr(layout::ArrayLayouts.DenseColumnMajor, lengths::Tuple{Int64, Int64}, A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:281
  [6] _qr(layout::ArrayLayouts.DenseColumnMajor, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}; kwds::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:282
  [7] _qr(layout::ArrayLayouts.DenseColumnMajor, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:282
  [8] qr(::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}; kwds::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:315
  [9] qr(::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:315
 [10] _factorize(layout::ArrayLayouts.DenseColumnMajor, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:291
 [11] factorize(A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/factorizations.jl:320
 [12] _ldiv!
    @ ~/.julia/packages/ArrayLayouts/E7AfI/src/ldiv.jl:81 [inlined]
 [13] copyto!
    @ ~/.julia/packages/ArrayLayouts/E7AfI/src/ldiv.jl:98 [inlined]
 [14] copy
    @ ~/.julia/packages/ArrayLayouts/E7AfI/src/ldiv.jl:21 [inlined]
 [15] materialize
    @ ~/.julia/packages/ArrayLayouts/E7AfI/src/ldiv.jl:22 [inlined]
 [16] ldiv
    @ ~/.julia/packages/ArrayLayouts/E7AfI/src/ldiv.jl:86 [inlined]
 [17] \(A::PseudoBlockMatrix{Matrix{Float64}, Matrix{Matrix{Float64}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, x::Matrix{Float64})
    @ ArrayLayouts ~/.julia/packages/ArrayLayouts/E7AfI/src/ldiv.jl:141
 [18] top-level scope
    @ REPL[73]:1

[dlfivefifty]

are you aware you made a matrix of matrices, not a block matrix?

[simonbyrne]

Can you use a PseudoBlockVector? This should work

Not really: for my particular use case I have my own AbstractBlockVector type: https://github.com/CliMA/ClimaCore.jl/blob/main/src/Fields/fieldvector.jl

[dlfivefifty]

It's not clear how things are stored in memory in your use case, is it like BlockArray where they are not contiguous?

[simonbyrne]

Yes, it's basically laid out as a BlockVector, but with a bit of extra stuff for our specific use case

[dlfivefifty]

OK as I said it's a hard problem to fix properly. The easiest would be to make a copy into a PseudoBlockVector though not sure whether that extra allocation is a deal-breaker for you.

[simonbyrne]

That's what we've been doing at the moment. If we did want to define block factorizations, which methods would we need to extend?

[dlfivefifty]

It’s breaking at the lmul!(Q’, b) call (the stuff about materialize! is only relevant if you wanted to support “traits”, that is, multiple array types that share the same structure.)