Add new Higher Level API

Author: avik-pal

Project.toml

@@ -1,6 +1,9 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
-authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
+authors = [
+    "Pankaj Mishra <pankajmishra1511@gmail.com>",
+    "Chris Rackauckas <contact@chrisrackauckas.com>",
+]
 version = "2.5.0"
 
 [deps]
@@ -21,12 +24,15 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [weakdeps]
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
+SparseDiffToolsSymbolicsExt = "Symbolics"
 SparseDiffToolsZygoteExt = "Zygote"
 
 [compat]
@@ -43,7 +49,9 @@ SciMLOperators = "0.2.11, 0.3"
 Setfield = "1"
 StaticArrayInterface = "1.3"
 StaticArrays = "1"
+Symbolics = "5.5"
 Tricks = "0.1.6"
+UnPack = "1"
 VertexSafeGraphs = "0.2"
 Zygote = "0.6"
 julia = "1.6"
@@ -52,6 +60,7 @@ julia = "1.6"
 ArrayInterfaceBandedMatrices = "2e50d22c-5be1-4042-81b1-c572ed69783d"
 ArrayInterfaceBlockBandedMatrices = "5331f1e9-51c7-46b0-a9b0-df4434785e0a"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -63,4 +72,18 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "ArrayInterfaceBandedMatrices", "ArrayInterfaceBlockBandedMatrices", "BandedMatrices", "BlockBandedMatrices", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
+test = [
+    "Test",
+    "ArrayInterfaceBandedMatrices",
+    "ArrayInterfaceBlockBandedMatrices",
+    "BandedMatrices",
+    "BlockBandedMatrices",
+    "IterativeSolvers",
+    "Pkg",
+    "Random",
+    "SafeTestsets",
+    "Symbolics",
+    "Zygote",
+    "StaticArrays",
+    "BenchmarkTools",
+]

README.md

@@ -47,6 +47,67 @@ function g(x) # out-of-place
 end
 ```
 
+## High Level API
+
+We need to perform the following steps to utilize SparseDiffTools:
+
+1. Specify a Sparsity Detection Algorithm. There are 3 possible choices currently:
+    1. `NoSparsityDetection`: This will ignore any AD choice and compute the dense Jacobian
+    2. `JacPrototypeSparsityDetection`: If you already know the sparsity pattern, you can
+       specify it as `JacPrototypeSparsityDetection(; jac_prototype=<sparsity pattern>)`.
+    3. `SymbolicsSparsityDetection`: This will use `Symbolics.jl` to automatically detect
+       the sparsity pattern. (Note that `Symbolics.jl` must be explicitly loaded before
+       using this functionality.)
+2. Now choose an AD backend from `ADTypes.jl`:
+    1. If using a Non `*Sparse*` type, then we will not use sparsity detection.
+    2. All other sparse AD types will internally compute the proper sparsity pattern, and
+       try to exploit that.
+3. Now there are 2 options:
+    1. Precompute the cache using `sparse_jacobian_cache` and use the `sparse_jacobian` or
+       `sparse_jacobian!` functions to compute the Jacobian. This option is recommended if
+       you are repeatedly computing the Jacobian for the same function.
+    2. Directly use `sparse_jacobian` or `sparse_jacobian!` to compute the Jacobian. This
+       option should be used if you are only computing the Jacobian once.
+
+```julia
+using Symbolics
+
+sd = SymbolicsSparsityDetection()
+adtype = AutoSparseFiniteDiff()
+x = rand(30)
+y = similar(x)
+
+# Option 1
+## OOP Function
+cache = sparse_jacobian_cache(adtype, sd, g, x; fx=y) # Passing `fx` is needed if size(y) != size(x)
+J = sparse_jacobian(adtype, cache, g, x)
+### Or
+J_preallocated = similar(J)
+sparse_jacobian!(J_preallocated, adtype, cache, g, x)
+
+## IIP Function
+cache = sparse_jacobian_cache(adtype, sd, f, y, x)
+J = sparse_jacobian(adtype, cache, f, y, x)
+### Or
+J_preallocated = similar(J)
+sparse_jacobian!(J_preallocated, adtype, cache, f, y, x)
+
+# Option 2
+## OOP Function
+J = sparse_jacobian(adtype, sd, g, x)
+### Or
+J_preallocated = similar(J)
+sparse_jacobian!(J_preallocated, adtype, sd, g, x)
+
+## IIP Function
+J = sparse_jacobian(adtype, sd, f, y, x)
+### Or
+J_preallocated = similar(J)
+sparse_jacobian!(J_preallocated, adtype, sd, f, y, x)
+```
+
+## Lower Level API
+
 For this function, we know that the sparsity pattern of the Jacobian is a
 `Tridiagonal` matrix. However, if we didn't know the sparsity pattern for
 the Jacobian, we could use the `Symbolics.jacobian_sparsity` function to automatically

ext/SparseDiffToolsSymbolicsExt.jl

@@ -0,0 +1,21 @@
+module SparseDiffToolsSymbolicsExt
+
+using SparseDiffTools, Symbolics
+import SparseDiffTools: AbstractSparseADType
+
+function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f, x; fx=nothing,
+    kwargs...)
+    fx = fx === nothing ? similar(f(x)) : dx
+    f!(y, x) = (y .= f(x))
+    J = Symbolics.jacobian_sparsity(f!, fx, x)
+    _alg = JacPrototypeSparsityDetection(J, alg.alg)
+    return _alg(ad, f, x; fx, kwargs...)
+end
+
+function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f!, fx, x; kwargs...)
+    J = Symbolics.jacobian_sparsity(f!, fx, x)
+    _alg = JacPrototypeSparsityDetection(J, alg.alg)
+    return _alg(ad, f!, fx, x; kwargs...)
+end
+
+end

ext/SparseDiffToolsZygoteExt.jl

@@ -9,6 +9,36 @@ import Tricks: static_hasmethod
 
 import SparseDiffTools: numback_hesvec!,
     numback_hesvec, autoback_hesvec!, autoback_hesvec, auto_vecjac!, auto_vecjac
+import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient!
+import ADTypes: AutoZygote, AutoSparseZygote
+
+## Satisfying High-Level Interface for Sparse Jacobians
+function __gradient(::Union{AutoSparseZygote, AutoZygote}, f, x, cols)
+    _, ∂x, _ = Zygote.gradient(__f̂, f, x, cols)
+    return vec(∂x)
+end
+
+function __gradient!(::Union{AutoSparseZygote, AutoZygote}, f!, fx, x, cols)
+    return error("Zygote.jl cannot differentiate in-place (mutating) functions.")
+end
+
+# Zygote doesn't provide a way to accumulate directly into `J`. So we modify the code from
+# https://github.com/FluxML/Zygote.jl/blob/82c7a000bae7fb0999275e62cc53ddb61aed94c7/src/lib/grad.jl#L140-L157C4
+import Zygote: _jvec, _eyelike, _gradcopy!
+
+@views function __jacobian!(J::AbstractMatrix, ::Union{AutoSparseZygote, AutoZygote}, f, x)
+    y, back = Zygote.pullback(_jvec ∘ f, x)
+    δ = _eyelike(y)
+    for k in LinearIndices(y)
+        grad = only(back(δ[:, k]))
+        _gradcopy!(J[k, :], grad)
+    end
+    return J
+end
+
+function __jacobian!(J, ::Union{AutoSparseZygote, AutoZygote}, f!, fx, x)
+    return error("Zygote.jl cannot differentiate in-place (mutating) functions.")
+end
 
 ### Jac, Hes products
 
@@ -117,7 +147,7 @@ function (L::AutoDiffVJP{<:AutoZygote, IIP, true})(dv, v, p, t;
     copy!(dv, _dv)
 end
 
-function (L::AutoDiffVJP{<:AutoZygote, true, false})(_, _, _, _; VJP_input = nothing)
+function (L::AutoDiffVJP{<:AutoZygote, true, false})(args...; kwargs...)
     error("Zygote requires an out of place method with signature f(u).")
 end
 

src/SparseDiffTools.jl

@@ -2,12 +2,15 @@ module SparseDiffTools
 
 # QoL/Helper Packages
 using Adapt, Compat, Reexport
+import UnPack: @unpack
 # Graph Coloring
 using Graphs, VertexSafeGraphs
 import Graphs: SimpleGraph
 # Differentiation
 using FiniteDiff, ForwardDiff
 @reexport using ADTypes
+import ADTypes: AbstractADType, AutoSparseZygote, AbstractSparseForwardMode,
+    AbstractSparseReverseMode, AbstractSparseFiniteDifferences, AbstractReverseMode
 import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 # Array Packages
 using ArrayInterface, SparseArrays
@@ -41,6 +44,13 @@ include("differentiation/compute_hessian_ad.jl")
 include("differentiation/jaches_products.jl")
 include("differentiation/vecjac_products.jl")
 
+# High Level Interface
+include("highlevel/common.jl")
+include("highlevel/coloring.jl")
+include("highlevel/forward_mode.jl")
+include("highlevel/reverse_mode.jl")
+include("highlevel/finite_diff.jl")
+
 Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
 parameterless_type(x) = parameterless_type(typeof(x))
 parameterless_type(x::Type) = __parameterless_type(x)
@@ -67,11 +77,18 @@ export auto_jacvec, auto_jacvec!, num_jacvec, num_jacvec!
 export num_vecjac, num_vecjac!, auto_vecjac, auto_vecjac!
 # HesVec Products
 export numauto_hesvec,
-    numauto_hesvec!, autonum_hesvec, autonum_hesvec!, numback_hesvec, numback_hesvec!
+    numauto_hesvec!, autonum_hesvec, autonum_hesvec!, numback_hesvec, numback_hesvec!,
+    num_hesvec, num_hesvec!, autoback_hesvec, autoback_hesvec!
 # HesVecGrad Products
 export num_hesvecgrad, num_hesvecgrad!, auto_hesvecgrad, auto_hesvecgrad!
 # Operators
 export JacVec, HesVec, HesVecGrad, VecJac
 export update_coefficients, update_coefficients!, value!
 
+# High Level Interface: sparse_jacobian
+export AutoSparseZygote  # FIXME: Remove once https://github.com/SciML/ADTypes.jl/pull/16 is merged
+export NoSparsityDetection,
+    SymbolicsSparsityDetection, JacPrototypeSparsityDetection, AutoSparsityDetection
+export sparse_jacobian, sparse_jacobian_cache, sparse_jacobian!
+
 end # module

src/coloring/high_level.jl

@@ -7,18 +7,31 @@ struct GreedyStar2Color <: SparseDiffToolsColoringAlgorithm end
 struct AcyclicColoring <: SparseDiffToolsColoringAlgorithm end
 
 """
-    matrix_colors(A, alg::ColoringAlgorithm = GreedyD1Color())
+    matrix_colors(A, alg::ColoringAlgorithm = GreedyD1Color();
+        partition_by_rows::Bool = false)
 
 Return the colorvec vector for the matrix A using the chosen coloring
 algorithm. If a known analytical solution exists, that is used instead.
 The coloring defaults to a greedy distance-1 coloring.
 
 Note that if A isa SparseMatrixCSC, the sparsity pattern is defined by structural nonzeroes,
 ie includes explicitly stored zeros.
+
+If `ArrayInterface.fast_matrix_colors(A)` is true, then uses
+`ArrayInterface.matrix_colors(A)` to compute the matrix colors.
 """
 function ArrayInterface.matrix_colors(A::AbstractMatrix,
     alg::SparseDiffToolsColoringAlgorithm = GreedyD1Color();
     partition_by_rows::Bool = false)
+
+    # If fast algorithm for matrix coloring exists use that
+    if !partition_by_rows
+        ArrayInterface.fast_matrix_colors(A) && return ArrayInterface.matrix_colors(A)
+    else
+        A_ = A'
+        ArrayInterface.fast_matrix_colors(A_) && return ArrayInterface.matrix_colors(A_)
+    end
+
     _A = A isa SparseMatrixCSC ? A : sparse(A) # Avoid the copy
     A_graph = matrix2graph(_A, partition_by_rows)
     return color_graph(A_graph, alg)

src/differentiation/compute_jacobian_ad.jl

@@ -312,16 +312,10 @@ function forwarddiff_color_jacobian_immutable(f, x::AbstractArray{<:Number},
     return J
 end
 
-function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
-    f,
-    x::AbstractArray{<:Number};
-    dx = similar(x, size(J, 1)),
-    colorvec = 1:length(x),
-    sparsity = ArrayInterface.has_sparsestruct(J) ? J :
-               nothing)
-    forwarddiff_color_jacobian!(J, f, x,
-        ForwardColorJacCache(f, x, dx = dx, colorvec = colorvec,
-            sparsity = sparsity))
+function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f,
+    x::AbstractArray{<:Number}; dx = similar(x, size(J, 1)), colorvec = 1:length(x),
+    sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing)
+    forwarddiff_color_jacobian!(J, f, x, ForwardColorJacCache(f, x; dx, colorvec, sparsity))
 end
 
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},

src/highlevel/coloring.jl

@@ -0,0 +1,26 @@
+struct MatrixColoringResult{C, J, NR, NC}
+    colorvec::C
+    jacobian_sparsity::J
+    nz_rows::NR
+    nz_cols::NC
+end
+
+struct NoMatrixColoring end
+
+# Using Non-Sparse AD / NoSparsityDetection results in NoMatrixColoring
+(::NoSparsityDetection)(::AbstractADType, args...; kwargs...) = NoMatrixColoring()
+
+## If no specialization is available, we don't perform sparsity detection
+(::AbstractMaybeSparsityDetection)(::AbstractADType, args...; kws...) = NoMatrixColoring()
+
+# Prespecified Jacobian Structure
+function (alg::JacPrototypeSparsityDetection)(ad::AbstractSparseADType, args...; kwargs...)
+    J = alg.jac_prototype
+    reverse_mode = ad isa AbstractSparseReverseMode
+    colorvec = matrix_colors(J, alg.alg; partition_by_rows = reverse_mode)
+    (nz_rows, nz_cols) = collect.(ArrayInterface.findstructralnz(J))
+    return MatrixColoringResult(colorvec, J, nz_rows, nz_cols)
+end
+
+# TODO: Heuristics to decide whether to use Sparse Differentiation or not
+#       Simple Idea: Check min(max(colorvec_cols), max(colorvec_rows))