Is there an efficient way to reuse the sparsity pattern of a sparse matrix in Julia?

Question

I am wondering if there is an efficient way to reuse the sparsity pattern of a sparse matrix in Julia.

Let's say I have the vectors I, J and V and I create a sparse matrix using

using SparseArrays
A = sparse(I, J, V)

which seems to be very efficient even if I and J are not sorted and may occur several times.

I now want to exploit the matrix A, or more precisely its sparsity pattern, and fill it with new values W. An example would be the following code:

N = 10
I = vec([ 1, 1, 2, 2]' .+ (0:N))
J = vec([ 1, 2, 1, 2]' .+ (0:N))
V = vec([ 1,-1,-1, 1]' .* ones(N+1))
A = sparse(I, J, V)
W = vec([ 1,-1, 1,-1]' .* ones(N+1))
A = sparse(I, J, W)

Here, the last line should be replaced by something more efficient because the sparsity pattern of A stays the same.

I hoped that it would be as easy as

A.nzval .= 0
@inbounds A[I + size(A, 2) .* (J .- 1)] += W

I also tried several other (and way more complex) approaches like presorting the values of I and J and then explicitly setting the values of A.nzval. However, I always came to the conclusion that the most efficient way is still to recreate A from scratch using A = sparse(I, J, W), which obviously does not make sense since we have more information about A at this point.

How could I take advantage of the matrix A created earlier?

Answer 1

This question is relevant to the case when we have many sparse matrices which need to be initialized with the same I, J coordinates but different values. The exact sequence in initializing the sparse matrix is complicated by the possibility of repeating element coordinates.

The following tries to extract using I, J and a specialized sparse matrix a 'schedule' to create a sparse matrix which can then be used more efficiently for different values.

If the explanation isn't clear, maybe the code will be:

using SparseArrays

function sparse_sched(I,J,n = maximum(I), m = maximum(J))
    D1 = collect(1:length(I))
    Z = sparse(I,J,1:length(I),n, m,(x,y)->(D1[y] = x; x))
    D2 = similar(D1)
    for i in 1:nnz(Z)
        D2[Z.nzval[i]] = i
    end
    for i in eachindex(D2)
        D1[i] == 0 || (D2[i] = D2[D1[i]])
    end
    return D2
end

function apply_sched!(A,S,V,op = +)
    A.nzval .= zero(eltype(A))
    for i in eachindex(V)
        v = A.nzval[S[i]]
        A.nzval[S[i]] = iszero(v) ? V[i] : op(v, V[i])
    end
    return A
end

These functions can then be used as in the following code:

julia> N = 10;

julia> I = vec([ 1, 1, 2, 2]' .+ (0:N));

julia> J = vec([ 1, 2, 1, 2]' .+ (0:N));

julia> V = vec([ 1,-1,-1, 1]' .* ones(N+1));

julia> W = vec([ 1,-1, 1,-1]' .* ones(N+1));

julia> A = sparse(I, J, V);

julia> sched = sparse_sched(I,J);

julia> apply_sched!(similar(A), sched, W) == sparse(I,J,W)
true

julia> apply_sched!(similar(A), sched, V) == A
true

The last two lines show the new construction method is the same as the standard sparse matrix construction.

Benchmarking locally, it seems to offer some 3x speed improvement. If this is heavily used, then going through this trouble might be worth it.