From 974ef72fc43d91579b77e613f4ce751c4cce4fb3 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 7 Sep 2023 13:29:18 -0400 Subject: [PATCH 1/2] Allow specifying colorvecs directly --- Project.toml | 2 +- src/SparseDiffTools.jl | 5 +- src/highlevel/coloring.jl | 12 +++- src/highlevel/common.jl | 122 ++++++++++++++++++++++++++++++++--- test/test_sparse_jacobian.jl | 17 +++-- 5 files changed, 136 insertions(+), 22 deletions(-) diff --git a/Project.toml b/Project.toml index 2fada9b9..f3eda524 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" authors = ["Pankaj Mishra ", "Chris Rackauckas "] -version = "2.5.2" +version = "2.6.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 7a71c0ce..68633d78 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -88,8 +88,9 @@ export update_coefficients, update_coefficients!, value! # High Level Interface: sparse_jacobian export AutoSparseEnzyme -export NoSparsityDetection, - SymbolicsSparsityDetection, JacPrototypeSparsityDetection, AutoSparsityDetection +export NoSparsityDetection, SymbolicsSparsityDetection, JacPrototypeSparsityDetection, + PrecomputedJacobianColorvec, AutoSparsityDetection export sparse_jacobian, sparse_jacobian_cache, sparse_jacobian! +export init_jacobian end # module diff --git a/src/highlevel/coloring.jl b/src/highlevel/coloring.jl index 61815ae8..3e382600 100644 --- a/src/highlevel/coloring.jl +++ b/src/highlevel/coloring.jl @@ -16,8 +16,16 @@ struct NoMatrixColoring end # 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) + colorvec = matrix_colors(J, alg.alg; + partition_by_rows = ad isa AbstractSparseReverseMode) + (nz_rows, nz_cols) = ArrayInterface.findstructralnz(J) + return MatrixColoringResult(colorvec, J, nz_rows, nz_cols) +end + +# Prespecified Colorvecs +function (alg::PrecomputedJacobianColorvec)(ad::AbstractSparseADType, args...; kwargs...) + colorvec = _get_colorvec(alg, ad) + J = alg.jac_prototype (nz_rows, nz_cols) = ArrayInterface.findstructralnz(J) return MatrixColoringResult(colorvec, J, nz_rows, nz_cols) end diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index acd811f8..ee5a2ac8 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -9,18 +9,110 @@ abstract type AbstractSparsityDetection <: AbstractMaybeSparsityDetection end struct NoSparsityDetection <: AbstractMaybeSparsityDetection end +""" + SymbolicsSparsityDetection(; alg = GreedyD1Color()) + +Use Symbolics to compute the sparsity pattern of the Jacobian. This requires `Symbolics.jl` +to be explicitly loaded. + +## Keyword Arguments + + - `alg`: The algorithm used for computing the matrix colors + +See Also: [JacPrototypeSparsityDetection](@ref), [PrecomputedJacobianColorvec](@ref) +""" Base.@kwdef struct SymbolicsSparsityDetection{A <: ArrayInterface.ColoringAlgorithm} <: AbstractSparsityDetection alg::A = GreedyD1Color() end +""" + JacPrototypeSparsityDetection(; jac_prototype, alg = GreedyD1Color()) + +Use a pre-specified `jac_prototype` to compute the matrix colors of the Jacobian. + +## Keyword Arguments + + - `jac_prototype`: The prototype Jacobian used for computing the matrix colors + - `alg`: The algorithm used for computing the matrix colors + +See Also: [SymbolicsSparsityDetection](@ref), [PrecomputedJacobianColorvec](@ref) +""" Base.@kwdef struct JacPrototypeSparsityDetection{ - J, A <: ArrayInterface.ColoringAlgorithm, -} <: AbstractSparsityDetection + J, A <: ArrayInterface.ColoringAlgorithm} <: AbstractSparsityDetection jac_prototype::J alg::A = GreedyD1Color() end +""" + PrecomputedJacobianColorvec(jac_prototype, row_colorvec, col_colorvec) + +Use a pre-specified `colorvec` which can be directly used for sparse differentiation. Based +on whether a reverse mode or forward mode or finite differences is used, the corresponding +`row_colorvec` or `col_colorvec` is used. Atmost one of them can be set to `nothing`. + +## Arguments + + - `jac_prototype`: The prototype Jacobian used for computing structural nonzeros + - `row_colorvec`: The row colorvec of the Jacobian + - `col_colorvec`: The column colorvec of the Jacobian + +See Also: [SymbolicsSparsityDetection](@ref), [JacPrototypeSparsityDetection](@ref) +""" +struct PrecomputedJacobianColorvec{J, RC, CC} <: AbstractSparsityDetection + jac_prototype::J + row_colorvec::RC + col_colorvec::CC +end + +""" + PrecomputedJacobianColorvec(; jac_prototype, partition_by_rows::Bool = false, + colorvec = missing, row_colorvec = missing, col_colorvec = missing) + +Use a pre-specified `colorvec` which can be directly used for sparse differentiation. Based +on whether a reverse mode or forward mode or finite differences is used, the corresponding +`row_colorvec` or `col_colorvec` is used. Atmost one of them can be set to `nothing`. + +## Keyword Arguments + + - `jac_prototype`: The prototype Jacobian used for computing structural nonzeros + - `partition_by_rows`: Whether to partition the Jacobian by rows or columns (row + partitioning is used for reverse mode AD) + - `colorvec`: The colorvec of the Jacobian. If `partition_by_rows` is `true` then this + is the row colorvec, otherwise it is the column colorvec + - `row_colorvec`: The row colorvec of the Jacobian + - `col_colorvec`: The column colorvec of the Jacobian + +See Also: [SymbolicsSparsityDetection](@ref), [JacPrototypeSparsityDetection](@ref) +""" +function PrecomputedJacobianColorvec(; jac_prototype, partition_by_rows::Bool = false, + colorvec = missing, row_colorvec = missing, col_colorvec = missing) + if colorvec === missing + @assert row_colorvec !== missing||col_colorvec !== missing "Either `colorvec` or `row_colorvec` and `col_colorvec` must be specified!" + row_colorvec = row_colorvec === missing ? nothing : row_colorvec + col_colorvec = col_colorvec === missing ? nothing : col_colorvec + return PrecomputedJacobianColorvec(jac_prototype, row_colorvec, col_colorvec) + else + @assert row_colorvec === missing&&col_colorvec === missing "Specifying `colorvec` is incompatible with specifying `row_colorvec` or `col_colorvec`!" + row_colorvec = partition_by_rows ? colorvec : nothing + col_colorvec = partition_by_rows ? nothing : colorvec + return PrecomputedJacobianColorvec(jac_prototype, row_colorvec, col_colorvec) + end +end + +function _get_colorvec(alg::PrecomputedJacobianColorvec, ad) + cvec = alg.col_colorvec + @assert cvec!==nothing "`col_colorvec` is nothing, but Forward Mode AD or Finite Differences is being used!" + return cvec +end + +function _get_colorvec(alg::PrecomputedJacobianColorvec, ::AbstractReverseMode) + cvec = alg.row_colorvec + @assert cvec!==nothing "`row_colorvec` is nothing, but Reverse Mode AD is being used!" + return cvec +end + +# No one should be using this currently Base.@kwdef struct AutoSparsityDetection{A <: ArrayInterface.ColoringAlgorithm} <: AbstractSparsityDetection alg::A = GreedyD1Color() @@ -41,7 +133,8 @@ Inplace update the matrix `J` with the Jacobian of `f` at `x` using the AD backe function sparse_jacobian! end """ - sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, f, x; fx=nothing) + sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, f, x; + fx=nothing) sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, f!, fx, x) Takes the underlying AD backend `ad`, sparsity detection algorithm `sd`, function `f`, @@ -67,7 +160,7 @@ with the same cache to compute the jacobian. function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, args...; kwargs...) cache = sparse_jacobian_cache(ad, sd, args...; kwargs...) - J = __init_𝒥(cache) + J = init_jacobian(cache) return sparse_jacobian!(J, ad, cache, args...) end @@ -80,7 +173,7 @@ Jacobian at every function call """ function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, args...) - J = __init_𝒥(cache) + J = init_jacobian(cache) return sparse_jacobian!(J, ad, cache, args...) end @@ -106,7 +199,16 @@ function __gradient end function __gradient! end function __jacobian! end -function __init_𝒥 end +""" + init_jacobian(cache::AbstractMaybeSparseJacobianCache) + +Initialize the Jacobian based on the cache. Uses sparse jacobians if possible. + +!!! note + This function doesn't alias the provided jacobian prototype. It always initializes a + fresh jacobian that can be mutated without any side effects. +""" +function init_jacobian end # Misc Functions __chunksize(::AutoSparseForwardDiff{C}) where {C} = C @@ -123,12 +225,12 @@ end return :(nothing) end -function __init_𝒥(c::AbstractMaybeSparseJacobianCache) +function init_jacobian(c::AbstractMaybeSparseJacobianCache) T = promote_type(eltype(c.fx), eltype(c.x)) - return __init_𝒥(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x) + return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x) end -__init_𝒥(::Nothing, ::Type{T}, fx, x) where {T} = similar(fx, T, length(fx), length(x)) -__init_𝒥(J, ::Type{T}, _, _) where {T} = similar(J, T, size(J, 1), size(J, 2)) +init_jacobian(::Nothing, ::Type{T}, fx, x) where {T} = similar(fx, T, length(fx), length(x)) +init_jacobian(J, ::Type{T}, _, _) where {T} = similar(J, T, size(J, 1), size(J, 2)) __maybe_copy_x(_, x) = x __maybe_copy_x(_, ::Nothing) = nothing diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index 921a4c56..3ef4b206 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -26,9 +26,12 @@ J_true = ForwardDiff.jacobian(fdiff, x); # SparseDiffTools High-Level API J_sparsity = Symbolics.jacobian_sparsity(fdiff, similar(x), x); +row_colorvec = SparseDiffTools.matrix_colors(J_sparsity; partition_by_rows = true) +col_colorvec = SparseDiffTools.matrix_colors(J_sparsity; partition_by_rows = false) -SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_sparsity), - SymbolicsSparsityDetection(), NoSparsityDetection()] +SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_sparsity), + SymbolicsSparsityDetection(), NoSparsityDetection(), + PrecomputedJacobianColorvec(; jac_prototype = J_sparsity, row_colorvec, col_colorvec)] @testset "High-Level API" begin @testset "Sparsity Detection: $(nameof(typeof(sd)))" for sd in SPARSITY_DETECTION_ALGS @@ -40,7 +43,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme()) @testset "Cache & Reuse" begin cache = sparse_jacobian_cache(difftype, sd, fdiff, x) - J = SparseDiffTools.__init_𝒥(cache) + J = init_jacobian(cache) sparse_jacobian!(J, difftype, cache, fdiff, x) @@ -74,7 +77,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars @info "$(nameof(typeof(difftype)))() `sparse_jacobian` (complete) time: $(t₁)s" cache = sparse_jacobian_cache(difftype, sd, fdiff, x) - J = SparseDiffTools.__init_𝒥(cache) + J = init_jacobian(cache) sparse_jacobian!(J, difftype, sd, fdiff, x) @@ -95,7 +98,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x) @testset "Cache & Reuse" begin - J = SparseDiffTools.__init_𝒥(cache) + J = init_jacobian(cache) sparse_jacobian!(J, difftype, cache, fdiff, y, x) @test J ≈ J_true @@ -126,7 +129,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars t₁ = @elapsed sparse_jacobian(difftype, sd, fdiff, y, x) @info "$(nameof(typeof(difftype)))() `sparse_jacobian` (complete) time: $(t₁)s" - J = SparseDiffTools.__init_𝒥(cache) + J = init_jacobian(cache) sparse_jacobian!(J, difftype, sd, fdiff, y, x) @@ -142,7 +145,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars AutoZygote()) y = similar(x) cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x) - J = SparseDiffTools.__init_𝒥(cache) + J = init_jacobian(cache) @testset "Cache & Reuse" begin @test_throws Exception sparse_jacobian!(J, difftype, cache, fdiff, y, x) From a6731aaf2ee7d38c134640e7bbbd9a39e2f32db0 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 7 Sep 2023 13:32:29 -0400 Subject: [PATCH 2/2] Prevent breakage of BVP code --- src/highlevel/common.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index ee5a2ac8..6f6745c4 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -210,6 +210,10 @@ Initialize the Jacobian based on the cache. Uses sparse jacobians if possible. """ function init_jacobian end +# Never thought this was a useful function externally, but I ended up using it in quite a +# few places. Keeping this till I remove uses of those. +const __init_𝒥 = init_jacobian + # Misc Functions __chunksize(::AutoSparseForwardDiff{C}) where {C} = C