diff --git a/Project.toml b/Project.toml index 8e60927..6f723e6 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ArrayLayouts" uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" -version = "1.15.1" +version = "1.16" authors = ["Sheehan Olver "] [deps] diff --git a/ext/ArrayLayoutsSparseArraysExt.jl b/ext/ArrayLayoutsSparseArraysExt.jl index 32b543e..88212e2 100644 --- a/ext/ArrayLayoutsSparseArraysExt.jl +++ b/ext/ArrayLayoutsSparseArraysExt.jl @@ -1,15 +1,63 @@ module ArrayLayoutsSparseArraysExt using ArrayLayouts -using ArrayLayouts: copyto!_layout, equals_layout +using ArrayLayouts: copyto!_layout, equals_layout, sparse_layout, LayoutVecOrMat, UpperOrLowerTriangular using SparseArrays # Specifying the full namespace is necessary because of https://github.com/JuliaLang/julia/issues/48533 # See https://github.com/JuliaStats/LogExpFunctions.jl/pull/63 import ArrayLayouts.LinearAlgebra -using ArrayLayouts.LinearAlgebra: Adjoint, Transpose, Symmetric, Hermitian, Diagonal, UpperTriangular, UnitUpperTriangular, LowerTriangular, UnitLowerTriangular +using ArrayLayouts.LinearAlgebra: Adjoint, Transpose, Symmetric, Hermitian, Diagonal, UpperTriangular, UnitUpperTriangular, LowerTriangular, UnitLowerTriangular, + AdjOrTrans, HermOrSym import Base: copyto!, == +# use colsupport to deduce the sparsity pattern +ArrayLayouts.sparse_layout(_, A) = sparse_colsupport(A) + +function sparse_colsupport(A::AbstractMatrix{T}) where T + Base.require_one_based_indexing(A) + m, n = size(A) + colptr = Vector{Int}(undef, n+1) + rowval = Int[] + nzval = T[] + colptr[1] = 1 + for j = 1:n + for k in colsupport(A, j) + v = A[k, j] + if !iszero(v) + push!(rowval, k) + push!(nzval, v) + end + end + colptr[j+1] = length(rowval) + 1 + end + SparseMatrixCSC(m, n, colptr, rowval, nzval) +end + +function sparse_colsupport(a::AbstractVector{T}) where T + Base.require_one_based_indexing(a) + nzind = Int[] + nzval = T[] + for k in colsupport(a, 1) + v = a[k] + if !iszero(v) + push!(nzind, k) + push!(nzval, v) + end + end + SparseVector(length(a), nzind, nzval) +end + +const LayoutSparseVectors = Union{LayoutVector, SubArray{<:Any,1,<:LayoutArray}} +const LayoutSparseMatrixParents = Union{LayoutMatrix, SubArray{<:Any,2,<:LayoutArray}, AdjOrTrans{<:Any,<:LayoutVecOrMat}} +const LayoutSparseMatrices = Union{LayoutSparseMatrixParents, + HermOrSym{<:Any,<:LayoutMatrix}, + UpperOrLowerTriangular{<:Any,<:LayoutSparseMatrixParents}, + Diagonal{<:Any,<:LayoutVector}} + +SparseArrays.sparse(a::LayoutSparseVectors) = sparse_layout(a) +SparseArrays.sparse(A::LayoutSparseMatrices) = sparse_layout(A) + # ambiguity from sparsematrix.jl copyto!(dest::LayoutMatrix, src::SparseArrays.AbstractSparseMatrixCSC) = copyto!_layout(dest, src) diff --git a/src/ArrayLayouts.jl b/src/ArrayLayouts.jl index 353ab9d..e6eddf0 100644 --- a/src/ArrayLayouts.jl +++ b/src/ArrayLayouts.jl @@ -329,6 +329,16 @@ equals_layout(A, B) = equals_layout(MemoryLayout(A), MemoryLayout(B), A, B) ==(A::SubArray{<:Any,<:Any,<:LayoutArray}, B::AbstractArray) = equals_layout(A, B) ==(A::AbstractArray, B::SubArray{<:Any,<:Any,<:LayoutArray}) = equals_layout(A, B) +""" + sparse_layout(layout, A) + +Layout-aware hook for `SparseArrays.sparse(A)`, called with `layout = MemoryLayout(A)`. +Overload this to specialise conversion to a sparse array on memory layouts. The default +(defined in the SparseArrays extension) uses `colsupport` to deduce the sparsity pattern. +""" +function sparse_layout end +sparse_layout(A) = sparse_layout(MemoryLayout(A), A) + _fill_lmul!(β, A::AbstractArray) = iszero(β) ? zero!(A) : lmul!(β, A) diff --git a/test/test_layoutarray.jl b/test/test_layoutarray.jl index 3449cb1..ffa2cf6 100644 --- a/test/test_layoutarray.jl +++ b/test/test_layoutarray.jl @@ -834,6 +834,56 @@ ArrayLayouts.equals_layout(::MyEqualsLayout, _, A, B) = false end end +struct MySparseLayout <: MemoryLayout end +struct MyBidiagonalMatrix <: LayoutMatrix{Float64} end +Base.size(::MyBidiagonalMatrix) = (4,4) +Base.getindex(::MyBidiagonalMatrix, k::Int, j::Int) = k == j ? Float64(k) : k+1 == j ? 1.0 : error("outside colsupport") +MemoryLayout(::Type{MyBidiagonalMatrix}) = MySparseLayout() +ArrayLayouts.colsupport(::MySparseLayout, A, j) = max(1,minimum(j)-1):min(size(A,1),maximum(j)) +struct MyZeroVector <: LayoutVector{Float64} end +Base.size(::MyZeroVector) = (3,) +Base.getindex(::MyZeroVector, k::Int) = error("outside colsupport") +MemoryLayout(::Type{MyZeroVector}) = ZerosLayout() +struct MySparseOverloadLayout <: MemoryLayout end +struct MySparseOverloadVector <: LayoutVector{Float64} end +Base.size(::MySparseOverloadVector) = (3,) +Base.getindex(::MySparseOverloadVector, k::Int) = Float64(k) +MemoryLayout(::Type{MySparseOverloadVector}) = MySparseOverloadLayout() +ArrayLayouts.sparse_layout(::MySparseOverloadLayout, a) = sparsevec([2], [-1.0], 3) + +@testset "sparse" begin + A = [1.0 0 2; 0 3 4; 5 0 0] + M = MyMatrix(A) + for B in (M, view(M, 1:2, 2:3), M', transpose(M), Symmetric(M), Hermitian(M, :L), + UpperTriangular(M), UnitLowerTriangular(M), LowerTriangular(M'), UpperTriangular(view(M, :, :))) + S = sparse(B) + @test S isa SparseMatrixCSC{Float64,Int} + @test S == B + @test nnz(S) == count(!iszero, B) + end + v = MyVector([1.0, 0, 2]) + for b in (v, view(v, 2:3), view(M, :, 2)) + s = sparse(b) + @test s isa SparseVector{Float64,Int} + @test s == b + @test nnz(s) == count(!iszero, b) + end + D = sparse(Diagonal(v)) + @test D isa SparseMatrixCSC{Float64,Int} + @test D == Diagonal([1.0, 0, 2]) + @test nnz(D) == 2 + @test sparse(v') == v' + + B = MyBidiagonalMatrix() + S = sparse(B) + @test S == Bidiagonal([1.0:4;], ones(3), :U) + @test nnz(S) == 7 + s = sparse(MyZeroVector()) + @test s == zeros(3) + @test nnz(s) == 0 + @test sparse(MySparseOverloadVector()) == [0, -1.0, 0] +end + @testset "* for infinite layouts" begin tup = InfSymTridiagonal(), InfTridiagonal(), InfBidiagonal('U'), InfBidiagonal('L'),