faster transformation of Q of qr to sparse + pretty print

JuliaSparse · Jul 25, 2022 · 29a6b07 · 29a6b07
1 parent 134f67a
commit 29a6b07
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 6 deletions.
diff --git a/src/solvers/spqr.jl b/src/solvers/spqr.jl
@@ -132,7 +132,7 @@ struct QRSparseQ{Tv<:CHOLMOD.VTypes,Ti<:Integer} <: LinearAlgebra.AbstractQ{Tv}
     τ::Vector{Tv}
     n::Int # Number of columns in original matrix
 end
-
+Base.print_array(io::IO, Q::QRSparseQ) = Base.print_array(io, sparse(Q))
 Base.size(Q::QRSparseQ) = (size(Q.factors, 1), size(Q.factors, 1))
 Base.axes(Q::QRSparseQ) = map(Base.OneTo, size(Q))
 
@@ -168,10 +168,10 @@ 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
+ -0.707107    ⋅          ⋅        -0.707107
+   ⋅        -0.707107  -0.707107    ⋅ 
+   ⋅        -0.707107   0.707107    ⋅ 
+ -0.707107    ⋅          ⋅         0.707107
 R factor:
 2×2 SparseMatrixCSC{Float64, Int64} with 2 stored entries:
  -1.41421    ⋅
@@ -253,6 +253,7 @@ function LinearAlgebra.rmul!(A::StridedMatrix, Q::QRSparseQ)
     return A
 end
 
+
 function LinearAlgebra.lmul!(adjQ::Adjoint{<:Any,<:QRSparseQ}, A::StridedVecOrMat)
     Q = adjQ.parent
     if size(A, 1) != size(Q, 1)

diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl
@@ -748,6 +748,80 @@ 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} = _from_lmul(Tv, Ti, Q; sizehint=false)
+function _from_eachcol(Tv, Ti, M; sizehint=false)
+    colptr = zeros(Ti, size(M, 2) + 1)
+    nzval = Vector{Tv}(undef, 0)
+    rowval = Vector{Ti}(undef, 0)
+    if sizehint
+        nz = 0
+        for col in eachcol(M)
+            for i in col
+                if iszero(i) == false # should be _isnotzero(v)
+                    nz += 1
+                end
+            end
+        end
+        sizehint!(nzval, nz)
+        sizehint!(rowval, nz)
+    end
+    colptr[1] = 1
+    ind = 1
+    for (j, col) in enumerate(eachcol(M))
+        for (i, v) in enumerate(col)
+            if iszero(v) == false
+                push!(nzval, v)
+                push!(rowval, i)
+                ind += 1
+            end
+        end
+        colptr[j + 1] = ind
+    end
+    return SparseMatrixCSC{Tv, Ti}(size(M)..., colptr, rowval, nzval)
+end
+
+function _from_lmul(Tv, Ti, M; sizehint=false)
+
+    colptr = zeros(Ti, size(M, 2) + 1)
+    nzval = Vector{Tv}(undef, 0)
+    rowval = Vector{Ti}(undef, 0)
+    col = zeros(eltype(M), size(M, 1))
+
+    if sizehint
+        nz = 0
+        for j in axes(M, 2)
+            fill!(col, false)
+            col[j] = one(Tv)
+            lmul!(M, col)
+            for i in col
+                if iszero(i) == false # should be _isnotzero(v)
+                    nz += 1
+                end
+            end
+        end
+        sizehint!(nzval, nz)
+        sizehint!(rowval, nz)
+    end
+
+    colptr[1] = 1
+    ind = 1
+    for j in axes(M, 2)
+        fill!(col, false)
+        col[j] = one(Tv)
+        lmul!(M, col)
+        for (i, v) in enumerate(col)
+            if iszero(v) == false
+                push!(nzval, v)
+                push!(rowval, i)
+                ind += 1
+            end
+        end
+        colptr[j + 1] = ind
+    end
+    return SparseMatrixCSC{Tv, Ti}(size(M)..., colptr, rowval, nzval)
+end
+
 # converting from SparseMatrixCSC to other matrix types
 function Matrix(S::AbstractSparseMatrixCSC{Tv}) where Tv
     _checkbuffers(S)

diff --git a/test/spqr.jl b/test/spqr.jl
@@ -6,7 +6,7 @@ using Test
 using SparseArrays.SPQR
 using SparseArrays.CHOLMOD
 using LinearAlgebra: I, istriu, norm, qr, rank, rmul!, lmul!, Adjoint, Transpose
-using SparseArrays: sparse, sprandn, spzeros, SparseMatrixCSC
+using SparseArrays: sparse, sprandn, spzeros, SparseMatrixCSC, _from_eachcol, _from_lmul
 
 if Base.USE_GPL_LIBS
 
@@ -137,6 +137,18 @@ 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).Q
+    sQ = sparse(Q)
+    @test sQ == sparse(Matrix(Q))
+    for f in [_from_eachcol, _from_lmul],
+        sh in [true, false]
+        @test sQ == f(typeof(sQ).parameters..., Q; sizehint=sh)
+    end
+end
+
 end
 
 end # Base.USE_GPL_LIBS