Add public same_sparsity_structure(A, B)

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "7.27.0"
+version = "7.28.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/sparsearrays.md

@@ -11,6 +11,7 @@ These routines allow for improving sparse iteration and indexing.
 ArrayInterface.isstructured
 ArrayInterface.findstructralnz
 ArrayInterface.has_sparsestruct
+ArrayInterface.same_sparsity_structure
 ```
 
 ## Matrix Coloring

ext/ArrayInterfaceCUDAExt.jl

@@ -25,6 +25,19 @@ const CuSparseArray = Union{
 }
 ArrayInterface.can_setindex(::Type{<:CuSparseArray}) = false
 
+# GPU CSC/CSR do not subtype `SparseArrays.AbstractSparseMatrixCSC`, so they need their own
+# structure comparison. The stored index arrays are `CuVector`s; `==` reduces on-device.
+function ArrayInterface.same_sparsity_structure(
+        A::CUSPARSE.CuSparseMatrixCSC, B::CUSPARSE.CuSparseMatrixCSC
+    )
+    size(A) == size(B) && A.colPtr == B.colPtr && A.rowVal == B.rowVal
+end
+function ArrayInterface.same_sparsity_structure(
+        A::CUSPARSE.CuSparseMatrixCSR, B::CUSPARSE.CuSparseMatrixCSR
+    )
+    size(A) == size(B) && A.rowPtr == B.rowPtr && A.colVal == B.colVal
+end
+
 function ArrayInterface.promote_eltype(
         ::Type{<:CUDA.CuArray{T, N, M}}, ::Type{T2}
     ) where {T, N, M, T2}

ext/ArrayInterfaceSparseArraysExt.jl

@@ -1,8 +1,9 @@
 module ArrayInterfaceSparseArraysExt
 
-import ArrayInterface: buffer, has_sparsestruct, issingular, findstructralnz, bunchkaufman_instance, DEFAULT_CHOLESKY_PIVOT, cholesky_instance, ldlt_instance, lu_instance, qr_instance
+import ArrayInterface: buffer, has_sparsestruct, issingular, findstructralnz, same_sparsity_structure, bunchkaufman_instance, DEFAULT_CHOLESKY_PIVOT, cholesky_instance, ldlt_instance, lu_instance, qr_instance
 using ArrayInterface.LinearAlgebra
 using SparseArrays
+using SparseArrays: AbstractSparseMatrixCSC, getcolptr, rowvals
 
 buffer(x::SparseMatrixCSC) = getfield(x, :nzval)
 buffer(x::SparseVector) = getfield(x, :nzval)
@@ -14,6 +15,13 @@ function findstructralnz(x::SparseMatrixCSC)
     (rowind, colind)
 end
 
+# Compares the stored index arrays directly (no `findnz` allocation), so it is cheap
+# enough for per-iteration reuse checks. Covers any CPU CSC type implementing the
+# `getcolptr`/`rowvals` interface.
+function same_sparsity_structure(A::AbstractSparseMatrixCSC, B::AbstractSparseMatrixCSC)
+    Base.size(A) == Base.size(B) && getcolptr(A) == getcolptr(B) && rowvals(A) == rowvals(B)
+end
+
 function bunchkaufman_instance(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
     bunchkaufman(SparseMatrixCSC{Tv, Ti}(similar(A, 1, 1)), check = false)
 end

src/ArrayInterface.jl

@@ -341,6 +341,50 @@ function findstructralnz(x::Union{Tridiagonal, SymTridiagonal})
     (rowind, colind)
 end
 
+"""
+    same_sparsity_structure(A, B) -> Bool
+
+Return `true` when `A` and `B` store their nonzeros in exactly the same structural
+positions, i.e. they share an identical sparsity pattern (and shape). This is the condition
+under which values may be written into `B` reusing `A`'s stored layout (or a cached symbolic
+factorization of `A` reused for `B`) without rebuilding the structure.
+
+Semantics by category:
+
+  - Dense arrays store every position, so two of them share a structure exactly when they
+    have the same shape.
+  - Structured (`Diagonal`, `Bidiagonal`, `Tridiagonal`, `SymTridiagonal`) and sparse
+    (`SparseMatrixCSC`, GPU CSC/CSR) matrices compare by their pattern — shape (plus
+    `uplo`/format) or the stored index arrays (e.g. `colptr`/`rowval` for CSC).
+  - Two arrays whose base types differ (e.g. a structured matrix and a dense one, or two
+    different structured types) return `false`: a layout is not reused across categories,
+    so they are treated as "not the same structure".
+  - Two arrays of the *same* base type with no specialized method are **unknown**, not
+    `false`: rather than fabricate an answer, a `MethodError` is thrown so the type can
+    define its own method. (`has_sparsestruct` can be used to check applicability first.)
+
+Unlike `findstructralnz`, this performs no allocation on its fast paths, so it is suitable
+for use in hot loops (for example, deciding per Jacobian evaluation whether a sparse
+`jac_prototype`'s structure needs rebuilding).
+"""
+function same_sparsity_structure(A::AbstractArray, B::AbstractArray)
+    # Same base type but no specialized method: we genuinely do not know -- error rather
+    # than return a possibly-wrong `false`. Different base types cannot share a structural
+    # category, so those are `false`.
+    parameterless_type(A) === parameterless_type(B) &&
+        throw(MethodError(same_sparsity_structure, (A, B)))
+    return false
+end
+same_sparsity_structure(A::DenseArray, B::DenseArray) = Base.size(A) == Base.size(B)
+same_sparsity_structure(A::Diagonal, B::Diagonal) = Base.size(A) == Base.size(B)
+function same_sparsity_structure(A::Bidiagonal, B::Bidiagonal)
+    Base.size(A) == Base.size(B) && A.uplo == B.uplo
+end
+# Base-type methods (not `A::T, B::T`) so differing type parameters, e.g.
+# `Tridiagonal{Float64}` vs `Tridiagonal{Int}`, still compare as the same structure.
+same_sparsity_structure(A::Tridiagonal, B::Tridiagonal) = Base.size(A) == Base.size(B)
+same_sparsity_structure(A::SymTridiagonal, B::SymTridiagonal) = Base.size(A) == Base.size(B)
+
 abstract type ColoringAlgorithm end
 
 """
@@ -1005,8 +1049,8 @@ end
         :cholesky_instance, :ldlt_instance, :lu_instance, :qr_instance,
         :svd_instance,
         # sparse arrays
-        :isstructured, :findstructralnz, :has_sparsestruct, :fast_matrix_colors,
-        :matrix_colors,
+        :isstructured, :findstructralnz, :has_sparsestruct, :same_sparsity_structure,
+        :fast_matrix_colors, :matrix_colors,
         # wrapping
         :is_forwarding_wrapper, :buffer, :parent_type,
         # tuples

test/core.jl

@@ -216,6 +216,44 @@ end
     @test [STri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4,5,6,7,5,6,7]
 end
 
+@testset "same_sparsity_structure" begin
+    sss = ArrayInterface.same_sparsity_structure
+
+    # structured matrices: pattern is fixed by shape (plus uplo)
+    @test sss(Diagonal([1,2,3]), Diagonal([4,5,6]))
+    @test !sss(Diagonal([1,2,3]), Diagonal([4,5]))
+    @test sss(Bidiagonal([1,2,3],[7,8],:U), Bidiagonal([4,5,6],[9,1],:U))
+    @test !sss(Bidiagonal([1,2,3],[7,8],:U), Bidiagonal([4,5,6],[9,1],:L))
+    @test sss(Tridiagonal([1,2],[1,2,3],[4,5]), Tridiagonal([6,7],[8,9,1],[2,3]))
+    @test !sss(Tridiagonal([1,2],[1,2,3],[4,5]), Tridiagonal([1],[1,2],[4]))
+    @test sss(SymTridiagonal([1,2,3],[4,5]), SymTridiagonal([6,7,8],[9,1]))
+
+    # same base type, differing type parameters ⇒ still the same structure
+    @test sss(Tridiagonal([1,2],[1,2,3],[4,5]), Tridiagonal([1.0,2],[1.0,2,3],[4.0,5]))
+
+    # dense arrays: every position is structural, so same shape ⇒ same structure
+    @test sss(rand(3,3), rand(3,3))
+    @test !sss(rand(3,3), rand(3,2))
+    @test sss(rand(4), rand(4))
+
+    # different base types ⇒ false (not compared position-by-position)
+    @test !sss(Diagonal([1,2,3]), Bidiagonal([1,2,3],[7,8],:U))
+    @test !sss(Diagonal([1,2,3]), rand(3,3))
+    @test !sss(Tridiagonal([1,2],[1,2,3],[4,5]), SymTridiagonal([1,2,3],[4,5]))
+
+    # same base type but no specialized method ⇒ unknown, error rather than fabricate false
+    @test_throws MethodError sss(UpperTriangular(rand(3,3)), UpperTriangular(rand(3,3)))
+
+    # sparse CSC: compares the stored index arrays
+    A = sparse([1,2,3],[1,2,3],[1.0,2.0,3.0])
+    B = sparse([1,2,3],[1,2,3],[4.0,5.0,6.0])   # same pattern, different values
+    C = sparse([1,2,3],[1,2,1],[4.0,5.0,6.0])   # different pattern
+    @test sss(A, B)
+    @test !sss(A, C)
+    @test sss(A, similar(A))                    # `similar` preserves the pattern
+    @test !sss(A, sparse([1,2],[1,2],[1.0,2.0]))  # different size
+end
+
 @testset "ndims_index" begin
     @test @inferred(ArrayInterface.ndims_index(CartesianIndices(()))) == 1
     @test @inferred(ArrayInterface.ndims_index(trues(2, 2))) == 2