More efficient sparsity handling (#222)

Author: gdalle

DifferentiationInterface/Project.toml

@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -54,6 +55,7 @@ ForwardDiff = "0.10.36"
 LinearAlgebra = "1"
 PolyesterForwardDiff = "0.1.1"
 ReverseDiff = "1.15.1"
+SparseArrays = "1"
 Symbolics = "5.27.1"
 Tapir = "0.1.2"
 Test = "1"
@@ -84,4 +86,25 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "Diffractor", "Documenter", "Enzyme", "FastDifferentiation", "FiniteDiff", "FiniteDifferences", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "PolyesterForwardDiff", "ReverseDiff", "SparseArrays", "Symbolics", "Tapir", "Test", "Tracker", "Zygote"]
+test = [
+    "ADTypes",
+    "Aqua",
+    "Diffractor",
+    "Documenter",
+    "Enzyme",
+    "FastDifferentiation",
+    "FiniteDiff",
+    "FiniteDifferences",
+    "ForwardDiff",
+    "JET",
+    "JuliaFormatter",
+    "Pkg",
+    "PolyesterForwardDiff",
+    "ReverseDiff",
+    "SparseArrays",
+    "Symbolics",
+    "Tapir",
+    "Test",
+    "Tracker",
+    "Zygote",
+]

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -30,7 +30,8 @@ using ADTypes:
     AutoZygote
 using DocStringExtensions
 using FillArrays: OneElement
-using LinearAlgebra: Symmetric, dot
+using LinearAlgebra: Symmetric, Transpose, dot, parent, transpose
+using SparseArrays: SparseMatrixCSC, nzrange, rowvals, sparse
 
 abstract type Extras end
 
@@ -54,6 +55,7 @@ include("second_order/hvp.jl")
 include("second_order/hessian.jl")
 
 include("sparse/detector.jl")
+include("sparse/matrices.jl")
 include("sparse/coloring.jl")
 include("sparse/compressed_matrix.jl")
 include("sparse/fallbacks.jl")

DifferentiationInterface/src/sparse/coloring.jl

@@ -10,9 +10,8 @@ end
 
 abstract type AbstractMatrixGraph end
 
-Base.size(g::AbstractMatrixGraph, args...) = size(g.A, args...)
-rows(g::AbstractMatrixGraph) = axes(g.A, 1)
-columns(g::AbstractMatrixGraph) = axes(g.A, 2)
+rows(g::AbstractMatrixGraph) = axes(g.A_colmajor, 1)
+columns(g::AbstractMatrixGraph) = axes(g.A_colmajor, 2)
 
 ## Jacobian coloring
 
@@ -25,22 +24,33 @@ This graph is defined as `G = (R, C, E)` where `R = 1:m` is the set of row indic
 
 # Fields
 
-- `A::AbstractMatrix`
+- `A_colmajor::AbstractMatrix`: output of [`col_major`](@ref) applied to `A`
+- `A_rowmajor::AbstractMatrix`: output of [`row_major`](@ref) applied to `A`
 
 # Reference
 
 > [What Color Is Your Jacobian? Graph Coloring for Computing Derivatives](https://epubs.siam.org/doi/abs/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-struct BipartiteGraph{M<:AbstractMatrix} <: AbstractMatrixGraph
-    A::M
+struct BipartiteGraph{M1<:AbstractMatrix,M2<:AbstractMatrix} <: AbstractMatrixGraph
+    A_colmajor::M1
+    A_rowmajor::M2
+
+    function BipartiteGraph(A::AbstractMatrix)
+        A_colmajor = col_major(A)
+        A_rowmajor = row_major(A)
+        return new{typeof(A_colmajor),typeof(A_rowmajor)}(A_colmajor, A_rowmajor)
+    end
 end
 
-function neighbors_of_row(g::BipartiteGraph, i::Integer)
-    return filter(j -> !iszero(g.A[i, j]), columns(g))
-end
+neighbors_of_column(g::BipartiteGraph, j::Integer) = nz_in_col(g.A_colmajor, j)
+neighbors_of_row(g::BipartiteGraph, i::Integer) = nz_in_row(g.A_rowmajor, i)
 
-function neighbors_of_column(g::BipartiteGraph, j::Integer)
-    return filter(i -> !iszero(g.A[i, j]), rows(g))
+function colored_neighbors_of_column(
+    g::BipartiteGraph, j::Integer, colors::AbstractVector{<:Integer}
+)
+    return filter(neighbors_of_column(g, j)) do i
+        !iszero(colors[i])
+    end
 end
 
 function colored_neighbors_of_row(
@@ -51,14 +61,6 @@ function colored_neighbors_of_row(
     end
 end
 
-function colored_neighbors_of_column(
-    g::BipartiteGraph, j::Integer, colors::AbstractVector{<:Integer}
-)
-    return filter(neighbors_of_column(g, j)) do i
-        !iszero(colors[i])
-    end
-end
-
 function distance2_column_coloring(g::BipartiteGraph)
     n = length(columns(g))
     colors = zeros(Int, n)
@@ -130,18 +132,23 @@ This graph is defined as `G = (C, E)` where `C = 1:n` is the set of columns and
 
 # Fields
 
-- `A::AbstractMatrix`
+- `A_colmajor::AbstractMatrix`: output of [`col_major`](@ref) applied to `A`
 
 # Reference
 
 > [What Color Is Your Jacobian? Graph Coloring for Computing Derivatives](https://epubs.siam.org/doi/abs/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 struct AdjacencyGraph{M<:AbstractMatrix} <: AbstractMatrixGraph
-    A::M
+    A_colmajor::M
+
+    function AdjacencyGraph(A::AbstractMatrix)
+        A_colmajor = col_major(A)
+        return new{typeof(A_colmajor)}(A_colmajor)
+    end
 end
 
 function neighbors(g::AdjacencyGraph, j::Integer)
-    return filter(i -> (i != j) && !iszero(g.A[i, j]), columns(g))
+    return filter(!isequal(j), nz_in_col(g.A_colmajor, j))
 end
 
 function colored_neighbors(g::AdjacencyGraph, j::Integer, colors::AbstractVector{<:Integer})

DifferentiationInterface/src/sparse/matrices.jl

@@ -0,0 +1,42 @@
+## Conversion between row- and col-major
+
+"""
+    col_major(A::AbstractMatrix)
+
+Construct a column-major representation of the matrix `A`.
+"""
+col_major(A::M) where {M<:AbstractMatrix} = A
+col_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = M(A)
+
+"""
+    row_major(A::AbstractMatrix)
+
+Construct a row-major representation of the matrix `A`.
+"""
+row_major(A::M) where {M<:AbstractMatrix} = transpose(M(transpose(A)))
+row_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = A
+
+## Generic nz
+
+function nz_in_col(A_colmajor::AbstractMatrix, j::Integer)
+    return filter(i -> !iszero(A_colmajor[i, j]), axes(A_colmajor, 1))
+end
+
+function nz_in_row(A_rowmajor::AbstractMatrix, i::Integer)
+    return filter(j -> !iszero(A_rowmajor[i, j]), axes(A_rowmajor, 2))
+end
+
+## Sparse nz
+
+function nz_in_col(A_colmajor::SparseMatrixCSC{T}, j::Integer) where {T}
+    rv = rowvals(A_colmajor)
+    ind = nzrange(A_colmajor, j)
+    return view(rv, ind)
+end
+
+function nz_in_row(A_rowmajor::Transpose{T,<:SparseMatrixCSC{T}}, i::Integer) where {T}
+    A_transpose_colmajor = parent(A_rowmajor)
+    rv = rowvals(A_transpose_colmajor)
+    ind = nzrange(A_transpose_colmajor, i)
+    return view(rv, ind)
+end

DifferentiationInterface/test/chunk.jl …tiationInterface/test/internals/chunk.jlDifferentiationInterface/test/chunk.jl renamed to DifferentiationInterface/test/internals/chunk.jl DifferentiationInterface/test/chunk.jl renamed to DifferentiationInterface/test/internals/chunk.jl

@@ -0,0 +1,24 @@
+@testset "$(typeof(A))" for A in (
+    rand(0:1, 10, 10),
+    transpose(rand(0:1, 10, 10)),
+    sprand(100, 100, 0.1),
+    transpose(sprand(100, 100, 0.1)),
+)
+    A_colmajor = DI.col_major(A)
+    A_rowmajor = DI.row_major(A)
+
+    @test A_colmajor == A
+    @test A_rowmajor == A
+end
+
+@testset "$(typeof(A))" for A in (rand(0:1, 10, 10), sprand(100, 100, 0.1))
+    A_colmajor = DI.col_major(A)
+    A_rowmajor = DI.row_major(A)
+
+    for i in axes(A, 1)
+        @test DI.nz_in_row(A_rowmajor, i) == DI.nz_in_row(Matrix(A_rowmajor), i)
+    end
+    for j in axes(A, 2)
+        @test DI.nz_in_col(A_colmajor, j) == DI.nz_in_col(Matrix(A_colmajor), j)
+    end
+end

DifferentiationInterface/test/coloring.jl …tionInterface/test/internals/coloring.jlDifferentiationInterface/test/coloring.jl renamed to DifferentiationInterface/test/internals/coloring.jl DifferentiationInterface/test/coloring.jl renamed to DifferentiationInterface/test/internals/coloring.jl

@@ -51,15 +51,19 @@ include("test_imports.jl")
 
     @testset verbose = true "Internals" begin
         @testset verbose = true "Exception handling" begin
-            include("test_exceptions.jl")
+            include("internals/exceptions.jl")
         end
 
         @testset "Chunks" begin
-            include("chunk.jl")
+            include("internals/chunk.jl")
+        end
+
+        @testset "Matrices" begin
+            include("internals/matrices.jl")
         end
 
         @testset verbose = true "Coloring" begin
-            include("coloring.jl")
+            include("internals/coloring.jl")
         end
     end
 end;