Speed up SparseMatrixCSC(::AbstractQ) constructor and multiplication (

#196)

Co-authored-by: Daniel Karrasch <[email protected]>

src/solvers/spqr.jl

@@ -2,7 +2,7 @@
 
 module SPQR
 
-import Base: \
+import Base: \, *
 using Base: require_one_based_indexing
 using LinearAlgebra
 using ..LibSuiteSparse: SuiteSparseQR_C
@@ -167,11 +167,7 @@ julia> A = sparse([1,2,3,4], [1,1,2,2], [1.0,1.0,1.0,1.0])
 julia> qr(A)
 SparseArrays.SPQR.QRSparse{Float64, Int64}
 Q factor:
-4×4 SparseArrays.SPQR.QRSparseQ{Float64, Int64}:
- -0.707107   0.0        0.0       -0.707107
-  0.0       -0.707107  -0.707107   0.0
-  0.0       -0.707107   0.707107   0.0
- -0.707107   0.0        0.0        0.707107
+4×4 SparseArrays.SPQR.QRSparseQ{Float64, Int64}
 R factor:
 2×2 SparseMatrixCSC{Float64, Int64} with 2 stored entries:
  -1.41421    ⋅
@@ -284,6 +280,9 @@ function LinearAlgebra.rmul!(A::StridedMatrix, adjQ::Adjoint{<:Any,<:QRSparseQ})
     return A
 end
 
+(*)(Q::QRSparseQ, B::SparseMatrixCSC) = sparse(Q) * B
+(*)(A::SparseMatrixCSC, Q::QRSparseQ) = A * sparse(Q)
+
 @inline function Base.getproperty(F::QRSparse, d::Symbol)
     if d === :Q
         return QRSparseQ(F.factors, F.τ, size(F, 2))
@@ -312,6 +311,9 @@ function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, F::QRSparse)
     println(io, "\nColumn permutation:")
     show(io, mime, F.pcol)
 end
+function Base.show(io::IO, ::MIME{Symbol("text/plain")}, Q::QRSparseQ)
+    summary(io, Q)
+end
 
 # With a real lhs and complex rhs with the same precision, we can reinterpret
 # the complex rhs as a real rhs with twice the number of columns

src/sparsematrix.jl

@@ -748,6 +748,41 @@ SparseMatrixCSC{Tv}(M::Transpose{<:Any,<:AbstractSparseMatrixCSC}) where {Tv} =
 SparseMatrixCSC{Tv,Ti}(M::Adjoint{<:Any,<:AbstractSparseMatrixCSC}) where {Tv,Ti} = SparseMatrixCSC{Tv,Ti}(copy(M))
 SparseMatrixCSC{Tv,Ti}(M::Transpose{<:Any,<:AbstractSparseMatrixCSC}) where {Tv,Ti} = SparseMatrixCSC{Tv,Ti}(copy(M))
 
+# we can only view AbstractQs as columns
+SparseMatrixCSC{Tv,Ti}(Q::LinearAlgebra.AbstractQ) where {Tv,Ti} = sparse_with_lmul(Tv, Ti, Q)
+
+"""
+    sparse_with_lmul(Tv, Ti, Q) -> SparseMatrixCSC
+
+Helper function that creates a `SparseMatrixCSC{Tv,Ti}` representation of `Q`, where `Q` is
+supposed to not have fast `getindex` or not admit an iteration protocol at all, but instead
+a fast `lmul!(Q, v)` for dense vectors `v`. The prime example for such `Q`s is the Q factor
+of a (sparse) QR decomposition.
+"""
+function sparse_with_lmul(Tv, Ti, Q)
+    colptr = zeros(Ti, size(Q, 2) + 1)
+    nzval = Tv[]
+    rowval = Ti[]
+    col = zeros(eltype(Q), size(Q, 1))
+
+    colptr[1] = 1
+    ind = 1
+    for j in axes(Q, 2)
+        fill!(col, false)
+        col[j] = one(Tv)
+        lmul!(Q, col)
+        for (i, v) in enumerate(col)
+            if iszero(v) == false # should be _isnotzero(v)
+                push!(nzval, v)
+                push!(rowval, i)
+                ind += 1
+            end
+        end
+        colptr[j + 1] = ind
+    end
+    return SparseMatrixCSC{Tv,Ti}(size(Q)..., colptr, rowval, nzval)
+end
+
 # converting from SparseMatrixCSC to other matrix types
 function Matrix(S::AbstractSparseMatrixCSC{Tv}) where Tv
     _checkbuffers(S)
@@ -1137,9 +1172,9 @@ end
 """
     ftranspose!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}, f::Function) where {Tv,Ti}
 
-Transpose `A` and store it in `X` while applying the function `f` to the non-zero elements. 
+Transpose `A` and store it in `X` while applying the function `f` to the non-zero elements.
 Does not remove the zeros created by `f`. `size(X)` must be equal to `size(transpose(A))`.
-No additonal memory is allocated other than resizing the rowval and nzval of `X`, if needed. 
+No additonal memory is allocated other than resizing the rowval and nzval of `X`, if needed.
 
 See `halfperm!`
 """
@@ -1158,9 +1193,9 @@ end
 """
     transpose!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
 
-Transpose the matrix `A` and stores it in the matrix `X`. 
+Transpose the matrix `A` and stores it in the matrix `X`.
 `size(X)` must be equal to `size(transpose(A))`.
-No additonal memory is allocated other than resizing the rowval and nzval of `X`, if needed. 
+No additonal memory is allocated other than resizing the rowval and nzval of `X`, if needed.
 
 See `halfperm!`
 """
@@ -1170,8 +1205,8 @@ transpose!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti})
     adjoint!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
 
 Transpose the matrix `A` and stores the adjoint of the elements in the matrix `X`.
-`size(X)` must be equal to `size(transpose(A))`. 
-No additonal memory is allocated other than resizing the rowval and nzval of `X`, if needed. 
+`size(X)` must be equal to `size(transpose(A))`.
+No additonal memory is allocated other than resizing the rowval and nzval of `X`, if needed.
 
 See `halfperm!`
 """
@@ -3863,13 +3898,13 @@ end
 
 ## Uniform matrix arithmetic
 
-(+)(A::AbstractSparseMatrixCSC{Tv, Ti}, J::UniformScaling{T}) where {T<:Number, Tv, Ti} = 
+(+)(A::AbstractSparseMatrixCSC{Tv, Ti}, J::UniformScaling{T}) where {T<:Number, Tv, Ti} =
     A + sparse(T, Ti, J, size(A)...)
-(+)(J::UniformScaling{T}, A::AbstractSparseMatrixCSC{Tv, Ti}) where {T<:Number, Tv, Ti} = 
+(+)(J::UniformScaling{T}, A::AbstractSparseMatrixCSC{Tv, Ti}) where {T<:Number, Tv, Ti} =
     sparse(T, Ti, J, size(A)...) + A
-(-)(A::AbstractSparseMatrixCSC{Tv, Ti}, J::UniformScaling{T}) where {T<:Number, Tv, Ti} = 
+(-)(A::AbstractSparseMatrixCSC{Tv, Ti}, J::UniformScaling{T}) where {T<:Number, Tv, Ti} =
     A - sparse(T, Ti, J, size(A)...)
-(-)(J::UniformScaling{T}, A::AbstractSparseMatrixCSC{Tv, Ti}) where {T<:Number, Tv, Ti} = 
+(-)(J::UniformScaling{T}, A::AbstractSparseMatrixCSC{Tv, Ti}) where {T<:Number, Tv, Ti} =
     sparse(T, Ti, J, size(A)...) - A
 
 

test/spqr.jl

@@ -90,12 +90,14 @@ end
 @testset "Issue 26367" begin
     A = sparse([0.0 1 0 0; 0 0 0 0])
     @test Matrix(qr(A).Q) == Matrix(qr(Matrix(A)).Q) == Matrix(I, 2, 2)
+    @test sparse(qr(A).Q) == sparse(qr(Matrix(A)).Q) == Matrix(I, 2, 2)
+    @test (sparse(I, 2, 2) * qr(A).Q)::SparseMatrixCSC == sparse(qr(A).Q) == sparse(I, 2, 2)
 end
 
 @testset "Issue 26368" begin
     A = sparse([0.0 1 0 0; 0 0 0 0])
     F = qr(A)
-    @test F.Q*F.R == A[F.prow,F.pcol]
+    @test (F.Q*F.R)::SparseMatrixCSC == A[F.prow,F.pcol]
 end
 
 @testset "select ordering overdetermined" begin
@@ -137,6 +139,19 @@ end
     @test all(iszero, (rank(spzeros(10, i)) for i in 1:10))
 end
 
+
+@testset "sparse" begin
+    A = I + sprandn(100, 100, 0.01)
+    q = qr(A; ordering=SPQR.ORDERING_FIXED)
+    Q = q.Q
+    sQ = sparse(Q)
+    @test sQ == sparse(Matrix(Q))
+    Dq = qr(Matrix(A))
+    perm = inv(Matrix(I, size(A)...)[q.prow, :])
+    f = sum(q.R; dims=2) ./ sum(Dq.R; dims=2)
+    @test perm * (transpose(f) .* sQ) ≈ sparse(Dq.Q)
+end
+
 end
 
 end # Base.USE_GPL_LIBS