Separate compressed evaluation and decompression for sparse (#212)

Author: gdalle

.gitignore

@@ -8,4 +8,6 @@
 **/Manifest.toml
 **/docs/Manifest.toml
 
-*.csv
+*.csv
+
+playground.jl

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -3,6 +3,11 @@ module DifferentiationInterfaceEnzymeExt
 using ADTypes: ADTypes, AutoEnzyme
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    DerivativeExtras,
+    GradientExtras,
+    JacobianExtras,
+    PullbackExtras,
+    PushforwardExtras,
     NoDerivativeExtras,
     NoGradientExtras,
     NoJacobianExtras,

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -37,7 +37,7 @@ end
 
 ## Gradient
 
-struct EnzymeForwardGradientExtras{C,O}
+struct EnzymeForwardGradientExtras{C,O} <: GradientExtras
     shadow::O
 end
 
@@ -76,7 +76,7 @@ end
 
 ## Jacobian
 
-struct EnzymeForwardOneArgJacobianExtras{C,O}
+struct EnzymeForwardOneArgJacobianExtras{C,O} <: JacobianExtras
     shadow::O
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl

@@ -3,13 +3,18 @@ module DifferentiationInterfaceFiniteDiffExt
 using ADTypes: AutoFiniteDiff
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    DerivativeExtras,
+    GradientExtras,
+    HessianExtras,
+    JacobianExtras,
+    PullbackExtras,
+    PushforwardExtras,
     NoDerivativeExtras,
     NoGradientExtras,
     NoHessianExtras,
     NoJacobianExtras,
     NoPullbackExtras,
-    NoPushforwardExtras,
-    GradientExtras
+    NoPushforwardExtras
 using FiniteDiff:
     DerivativeCache,
     GradientCache,

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl

@@ -19,7 +19,7 @@ end
 
 ## Derivative
 
-struct FiniteDiffOneArgDerivativeExtras{C}
+struct FiniteDiffOneArgDerivativeExtras{C} <: DerivativeExtras
     cache::C
 end
 
@@ -128,7 +128,7 @@ end
 
 ## Jacobian
 
-struct FiniteDiffOneArgJacobianExtras{C}
+struct FiniteDiffOneArgJacobianExtras{C} <: JacobianExtras
     cache::C
 end
 
@@ -166,7 +166,7 @@ end
 
 ## Hessian
 
-struct FiniteDiffHessianExtras{C}
+struct FiniteDiffHessianExtras{C} <: HessianExtras
     cache::C
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl

@@ -19,7 +19,7 @@ end
 
 ## Derivative
 
-struct FiniteDiffTwoArgDerivativeExtras{C}
+struct FiniteDiffTwoArgDerivativeExtras{C} <: DerivativeExtras
     cache::C
 end
 
@@ -61,7 +61,7 @@ end
 
 ## Jacobian
 
-struct FiniteDiffTwoArgJacobianExtras{C}
+struct FiniteDiffTwoArgJacobianExtras{C} <: JacobianExtras
     cache::C
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/detector.jl

@@ -1,17 +1,15 @@
-struct SymbolicsSparsityDetector <: ADTypes.AbstractSparsityDetector end
-
-function ADTypes.jacobian_sparsity(f, x, ::SymbolicsSparsityDetector)
+function ADTypes.jacobian_sparsity(f, x, ::DI.SymbolicsSparsityDetector)
     y = similar(f(x))
     f!(y, x) = copyto!(y, f(x))
     return jacobian_sparsity(f!, y, x)
 end
 
-function ADTypes.jacobian_sparsity(f!, y, x, ::SymbolicsSparsityDetector)
+function ADTypes.jacobian_sparsity(f!, y, x, ::DI.SymbolicsSparsityDetector)
     f!_vec(y_vec, x_vec) = f!(reshape(y_vec, size(y)), reshape(x_vec, size(x)))
     return jacobian_sparsity(f!_vec, vec(y), vec(x))
 end
 
-function ADTypes.hessian_sparsity(f, x, ::SymbolicsSparsityDetector)
+function ADTypes.hessian_sparsity(f, x, ::DI.SymbolicsSparsityDetector)
     f_vec(x_vec) = f(reshape(x_vec, size(x)))
     return hessian_sparsity(f_vec, vec(x))
 end

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -52,8 +52,12 @@ include("second_order/second_derivative.jl")
 include("second_order/hvp.jl")
 include("second_order/hessian.jl")
 
+include("sparse/detector.jl")
 include("sparse/coloring.jl")
-include("sparse/sparse.jl")
+include("sparse/compressed_matrix.jl")
+include("sparse/fallbacks.jl")
+include("sparse/jacobian.jl")
+include("sparse/hessian.jl")
 
 export SecondOrder
 

DifferentiationInterface/src/sparse/coloring.jl

@@ -26,6 +26,10 @@ This graph is defined as `G = (R, C, E)` where `R = 1:m` is the set of row indic
 # Fields
 
 - `A::AbstractMatrix`
+
+# 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
@@ -127,6 +131,10 @@ This graph is defined as `G = (C, E)` where `C = 1:n` is the set of columns and
 # Fields
 
 - `A::AbstractMatrix`
+
+# 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
@@ -196,6 +204,23 @@ end
 
 ## ADTypes overloads
 
+"""
+    GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm
+
+Matrix coloring algorithm for sparse Jacobians and Hessians.
+
+Compatible with the [ADTypes.jl coloring framework](https://sciml.github.io/ADTypes.jl/stable/#Coloring-algorithm).
+
+# See also
+
+- `ADTypes.column_coloring`
+- `ADTypes.row_coloring`
+- `ADTypes.symmetric_coloring`
+
+# 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 GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm end
 
 function ADTypes.column_coloring(A, ::GreedyColoringAlgorithm)

DifferentiationInterface/src/sparse/compressed_matrix.jl

@@ -0,0 +1,73 @@
+"""
+    CompressedMatrix{dir}
+
+Compressed representation `B` of a sparse matrix `A ∈ ℝ^{m×n}` obtained by summing some of its columns (if `dir == :col`) or rows (if `dir == :row`), grouped by color.
+
+# Fields
+
+- `sparsity::AbstractMatrix{Bool}`: boolean sparsity pattern of the matrix `A`
+- `colors::Vector{Int}`: vector such that
+  - if `dir == `:col`, then `colors[j] ∈ 1:c` is the color of column `j`
+  - if `dir == `:row`, then `colors[i] ∈ 1:c` is the color of row `i`
+- `groups::Vector{Vector{Int}}`: vector of length `c` such that
+  - if `dir == :col`, then `groups[k]` is the vector of column indices assigned to the same color `k ∈ 1:c`
+  - if `dir == :row`, then `groups[k]` is the vector of row indices assigned to the same color `k ∈ 1:c`
+- `aggregates::AbstractMatrix`: matrix `B` such that
+  - if `dir == :col`, then `size(B) = (m, c)` and `B[:, c] = sum(A[:, k] for k in groups[c])`
+  - if `dir == :row`, then `size(B) = (c, n)` and `B[c, :] = sum(A[k, :] for k in groups[c])`
+"""
+mutable struct CompressedMatrix{dir,S<:AbstractMatrix{Bool},M<:AbstractMatrix}
+    sparsity::S
+    colors::Vector{Int}
+    groups::Vector{Vector{Int}}
+    aggregates::M
+end
+
+"""
+    CompressedMatrix{dir}(sparsity, groups, aggregates)
+
+Constructor for [`CompressedMatrix`](@ref).
+"""
+function CompressedMatrix{dir}(sparsity, colors, groups, aggregates) where {dir}
+    @assert dir in (:col, :row)
+    return CompressedMatrix{dir,typeof(sparsity),typeof(aggregates)}(
+        sparsity, colors, groups, aggregates
+    )
+end
+
+function decompress!(A::AbstractMatrix, compressed::CompressedMatrix{:col})
+    (; sparsity, colors, aggregates) = compressed
+    A .= zero(eltype(A))
+    @views for j in axes(A, 2)
+        k = colors[j]
+        nz_rows_j = (!iszero).(sparsity[:, j])
+        copyto!(A[nz_rows_j, j], aggregates[nz_rows_j, k])
+    end
+    return A
+end
+
+function decompress!(A::AbstractMatrix, compressed::CompressedMatrix{:row})
+    (; sparsity, colors, aggregates) = compressed
+    A .= zero(eltype(A))
+    @views for i in axes(A, 1)
+        k = colors[i]
+        nz_cols_i = (!iszero).(sparsity[i, :])
+        copyto!(A[i, nz_cols_i], aggregates[k, nz_cols_i])
+    end
+    return A
+end
+
+function decompress_symmetric!(A::AbstractMatrix, compressed::CompressedMatrix{:col})
+    (; sparsity, colors, groups, aggregates) = compressed
+    @views for j in axes(A, 2)
+        k = colors[j]
+        group = groups[k]
+        for i in axes(A, 1)
+            if (!iszero(sparsity[i, j]) && count(!iszero, sparsity[i, group]) == 1)
+                A[i, j] = aggregates[i, k]
+                A[j, i] = aggregates[i, k]
+            end
+        end
+    end
+    return A
+end