support sparse for LayoutArray and friends

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
-version = "1.15.1"
+version = "1.16"
 authors = ["Sheehan Olver <solver@mac.com>"]
 
 [deps]

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)

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)
 

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'),