diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl index dfbbf3549..21bea0c7b 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl @@ -353,9 +353,10 @@ end ## Jacobian -struct FastDifferentiationOneArgJacobianPrep{SIG,Y,E1,E1!} <: DI.JacobianPrep{SIG} +struct FastDifferentiationOneArgJacobianPrep{SIG,Y,P,E1,E1!} <: DI.SparseJacobianPrep{SIG} _sig::Val{SIG} y_prototype::Y + sparsity::P jac_exe::E1 jac_exe!::E1! end @@ -376,14 +377,18 @@ function DI.prepare_jacobian_nokwarg( x_vec_var = myvec(x_var) context_vec_vars = map(myvec, context_vars) y_vec_var = myvec(y_var) - jac_var = if backend isa AutoSparse - sparse_jacobian(y_vec_var, x_vec_var) + if backend isa AutoSparse + jac_var = sparse_jacobian(y_vec_var, x_vec_var) + sparsity = DI.get_pattern(jac_var) else - jacobian(y_vec_var, x_vec_var) + jac_var = jacobian(y_vec_var, x_vec_var) + sparsity = nothing end jac_exe = make_function(jac_var, x_vec_var, context_vec_vars...; in_place=false) jac_exe! = make_function(jac_var, x_vec_var, context_vec_vars...; in_place=true) - return FastDifferentiationOneArgJacobianPrep(_sig, y_prototype, jac_exe, jac_exe!) + return FastDifferentiationOneArgJacobianPrep( + _sig, y_prototype, sparsity, jac_exe, jac_exe! + ) end function DI.jacobian( @@ -626,9 +631,10 @@ end ## Hessian -struct FastDifferentiationHessianPrep{SIG,G,E2,E2!} <: DI.HessianPrep{SIG} +struct FastDifferentiationHessianPrep{SIG,G,P,E2,E2!} <: DI.SparseHessianPrep{SIG} _sig::Val{SIG} gradient_prep::G + sparsity::P hess_exe::E2 hess_exe!::E2! end @@ -648,10 +654,12 @@ function DI.prepare_hessian_nokwarg( x_vec_var = myvec(x_var) context_vec_vars = map(myvec, context_vars) - hess_var = if backend isa AutoSparse - sparse_hessian(y_var, x_vec_var) + if backend isa AutoSparse + hess_var = sparse_hessian(y_var, x_vec_var) + sparsity = DI.get_pattern(hess_var) else - hessian(y_var, x_vec_var) + hess_var = hessian(y_var, x_vec_var) + sparsity = nothing end hess_exe = make_function(hess_var, x_vec_var, context_vec_vars...; in_place=false) hess_exe! = make_function(hess_var, x_vec_var, context_vec_vars...; in_place=true) @@ -659,7 +667,9 @@ function DI.prepare_hessian_nokwarg( gradient_prep = DI.prepare_gradient_nokwarg( strict, f, dense_ad(backend), x, contexts... ) - return FastDifferentiationHessianPrep(_sig, gradient_prep, hess_exe, hess_exe!) + return FastDifferentiationHessianPrep( + _sig, gradient_prep, sparsity, hess_exe, hess_exe! + ) end function DI.hessian( diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/twoarg.jl index fb4b560f8..8a899a27a 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/twoarg.jl @@ -289,8 +289,9 @@ end ## Jacobian -struct FastDifferentiationTwoArgJacobianPrep{SIG,E1,E1!} <: DI.JacobianPrep{SIG} +struct FastDifferentiationTwoArgJacobianPrep{SIG,P,E1,E1!} <: DI.SparseJacobianPrep{SIG} _sig::Val{SIG} + sparsity::P jac_exe::E1 jac_exe!::E1! end @@ -312,14 +313,16 @@ function DI.prepare_jacobian_nokwarg( x_vec_var = myvec(x_var) context_vec_vars = map(myvec, context_vars) y_vec_var = myvec(y_var) - jac_var = if backend isa AutoSparse - sparse_jacobian(y_vec_var, x_vec_var) + if backend isa AutoSparse + jac_var = sparse_jacobian(y_vec_var, x_vec_var) + sparsity = DI.get_pattern(jac_var) else - jacobian(y_vec_var, x_vec_var) + jac_var = jacobian(y_vec_var, x_vec_var) + sparsity = nothing end jac_exe = make_function(jac_var, x_vec_var, context_vec_vars...; in_place=false) jac_exe! = make_function(jac_var, x_vec_var, context_vec_vars...; in_place=true) - return FastDifferentiationTwoArgJacobianPrep(_sig, jac_exe, jac_exe!) + return FastDifferentiationTwoArgJacobianPrep(_sig, sparsity, jac_exe, jac_exe!) end function DI.value_and_jacobian( diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseArraysExt/DifferentiationInterfaceSparseArraysExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseArraysExt/DifferentiationInterfaceSparseArraysExt.jl index e4ea92486..cb43a1e28 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseArraysExt/DifferentiationInterfaceSparseArraysExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseArraysExt/DifferentiationInterfaceSparseArraysExt.jl @@ -3,7 +3,13 @@ module DifferentiationInterfaceSparseArraysExt using ADTypes: ADTypes using DifferentiationInterface import DifferentiationInterface as DI -using SparseArrays: sparse +using SparseArrays: SparseMatrixCSC, sparse, nonzeros + +function DI.get_pattern(M::SparseMatrixCSC) + S = similar(M, Bool) + nonzeros(S) .= true + return S +end include("sparsity_detector.jl") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/DifferentiationInterfaceSparseMatrixColoringsExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/DifferentiationInterfaceSparseMatrixColoringsExt.jl index 2de265807..99536977f 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/DifferentiationInterfaceSparseMatrixColoringsExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/DifferentiationInterfaceSparseMatrixColoringsExt.jl @@ -15,14 +15,21 @@ using SparseMatrixColorings: decompress! import SparseMatrixColorings as SMC -abstract type SparseJacobianPrep{SIG} <: DI.JacobianPrep{SIG} end +## SMC overloads -SMC.sparsity_pattern(prep::SparseJacobianPrep) = sparsity_pattern(prep.coloring_result) -SMC.column_colors(prep::SparseJacobianPrep) = column_colors(prep.coloring_result) -SMC.column_groups(prep::SparseJacobianPrep) = column_groups(prep.coloring_result) -SMC.row_colors(prep::SparseJacobianPrep) = row_colors(prep.coloring_result) -SMC.row_groups(prep::SparseJacobianPrep) = row_groups(prep.coloring_result) -SMC.ncolors(prep::SparseJacobianPrep) = ncolors(prep.coloring_result) +abstract type SMCSparseJacobianPrep{SIG} <: DI.SparseJacobianPrep{SIG} end + +SMC.sparsity_pattern(prep::DI.SparseJacobianPrep) = prep.sparsity +SMC.column_colors(prep::DI.SparseJacobianPrep) = column_colors(prep.coloring_result) +SMC.column_groups(prep::DI.SparseJacobianPrep) = column_groups(prep.coloring_result) +SMC.row_colors(prep::DI.SparseJacobianPrep) = row_colors(prep.coloring_result) +SMC.row_groups(prep::DI.SparseJacobianPrep) = row_groups(prep.coloring_result) +SMC.ncolors(prep::DI.SparseJacobianPrep) = ncolors(prep.coloring_result) + +SMC.sparsity_pattern(prep::DI.SparseHessianPrep) = prep.sparsity +SMC.column_colors(prep::DI.SparseHessianPrep) = column_colors(prep.coloring_result) +SMC.column_groups(prep::DI.SparseHessianPrep) = column_groups(prep.coloring_result) +SMC.ncolors(prep::DI.SparseHessianPrep) = ncolors(prep.coloring_result) include("jacobian.jl") include("jacobian_mixed.jl") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl index 8efeb3659..f3925c89d 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl @@ -1,15 +1,17 @@ -struct SparseHessianPrep{ +struct SMCSparseHessianPrep{ SIG, BS<:DI.BatchSizeSettings, + P<:AbstractMatrix, C<:AbstractColoringResult{:symmetric,:column}, M<:AbstractMatrix{<:Number}, S<:AbstractVector{<:NTuple}, R<:AbstractVector{<:NTuple}, E2<:DI.HVPPrep, E1<:DI.GradientPrep, -} <: DI.HessianPrep{SIG} +} <: DI.SparseHessianPrep{SIG} _sig::Val{SIG} batch_size_settings::BS + sparsity::P coloring_result::C compressed_matrix::M batched_seeds::S @@ -18,11 +20,6 @@ struct SparseHessianPrep{ gradient_prep::E1 end -SMC.sparsity_pattern(prep::SparseHessianPrep) = sparsity_pattern(prep.coloring_result) -SMC.column_colors(prep::SparseHessianPrep) = column_colors(prep.coloring_result) -SMC.column_groups(prep::SparseHessianPrep) = column_groups(prep.coloring_result) -SMC.ncolors(prep::SparseHessianPrep) = ncolors(prep.coloring_result) - ## Hessian, one argument function DI.prepare_hessian_nokwarg( @@ -39,13 +36,14 @@ function DI.prepare_hessian_nokwarg( N = length(column_groups(coloring_result)) batch_size_settings = DI.pick_batchsize(DI.outer(dense_backend), N) return _prepare_sparse_hessian_aux( - strict, batch_size_settings, coloring_result, f, backend, x, contexts... + strict, batch_size_settings, sparsity, coloring_result, f, backend, x, contexts... ) end function _prepare_sparse_hessian_aux( strict::Val, batch_size_settings::DI.BatchSizeSettings{B}, + sparsity::AbstractMatrix, coloring_result::AbstractColoringResult{:symmetric,:column}, f::F, backend::AutoSparse, @@ -68,9 +66,10 @@ function _prepare_sparse_hessian_aux( gradient_prep = DI.prepare_gradient_nokwarg( strict, f, DI.inner(dense_backend), x, contexts... ) - return SparseHessianPrep( + return SMCSparseHessianPrep( _sig, batch_size_settings, + sparsity, coloring_result, compressed_matrix, batched_seeds, @@ -83,7 +82,7 @@ end function DI.hessian!( f::F, hess, - prep::SparseHessianPrep{SIG,<:DI.BatchSizeSettings{B}}, + prep::SMCSparseHessianPrep{SIG,<:DI.BatchSizeSettings{B}}, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -128,7 +127,7 @@ function DI.hessian!( end function DI.hessian( - f::F, prep::SparseHessianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C} + f::F, prep::SMCSparseHessianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C} ) where {F,C} DI.check_prep(f, prep, backend, x, contexts...) hess = similar(sparsity_pattern(prep), eltype(x)) @@ -139,7 +138,7 @@ function DI.value_gradient_and_hessian!( f::F, grad, hess, - prep::SparseHessianPrep, + prep::SMCSparseHessianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -153,7 +152,7 @@ function DI.value_gradient_and_hessian!( end function DI.value_gradient_and_hessian( - f::F, prep::SparseHessianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C} + f::F, prep::SMCSparseHessianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C} ) where {F,C} DI.check_prep(f, prep, backend, x, contexts...) y, grad = DI.value_and_gradient( diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl index c5edc8cb2..6c0954db7 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl @@ -1,16 +1,18 @@ ## Preparation -struct PushforwardSparseJacobianPrep{ +struct SMCPushforwardSparseJacobianPrep{ SIG, BS<:DI.BatchSizeSettings, + P<:AbstractMatrix, C<:AbstractColoringResult{:nonsymmetric,:column}, M<:AbstractMatrix{<:Number}, S<:AbstractVector{<:NTuple}, R<:AbstractVector{<:NTuple}, E<:DI.PushforwardPrep, -} <: SparseJacobianPrep{SIG} +} <: SMCSparseJacobianPrep{SIG} _sig::Val{SIG} batch_size_settings::BS + sparsity::P coloring_result::C compressed_matrix::M batched_seeds::S @@ -18,17 +20,19 @@ struct PushforwardSparseJacobianPrep{ pushforward_prep::E end -struct PullbackSparseJacobianPrep{ +struct SMCPullbackSparseJacobianPrep{ SIG, BS<:DI.BatchSizeSettings, + P<:AbstractMatrix, C<:AbstractColoringResult{:nonsymmetric,:row}, M<:AbstractMatrix{<:Number}, S<:AbstractVector{<:NTuple}, R<:AbstractVector{<:NTuple}, E<:DI.PullbackPrep, -} <: SparseJacobianPrep{SIG} +} <: SMCSparseJacobianPrep{SIG} _sig::Val{SIG} batch_size_settings::BS + sparsity::P coloring_result::C compressed_matrix::M batched_seeds::S @@ -84,13 +88,22 @@ function _prepare_sparse_jacobian_aux( end batch_size_settings = DI.pick_batchsize(dense_backend, N) return _prepare_sparse_jacobian_aux_aux( - strict, batch_size_settings, coloring_result, y, f_or_f!y, backend, x, contexts... + strict, + batch_size_settings, + sparsity, + coloring_result, + y, + f_or_f!y, + backend, + x, + contexts..., ) end function _prepare_sparse_jacobian_aux_aux( strict::Val, batch_size_settings::DI.BatchSizeSettings{B}, + sparsity::AbstractMatrix, coloring_result::AbstractColoringResult{:nonsymmetric,:column}, y, f_or_f!y::FY, @@ -111,9 +124,10 @@ function _prepare_sparse_jacobian_aux_aux( pushforward_prep = DI.prepare_pushforward_nokwarg( strict, f_or_f!y..., dense_backend, x, batched_seeds[1], contexts... ) - return PushforwardSparseJacobianPrep( + return SMCPushforwardSparseJacobianPrep( _sig, batch_size_settings, + sparsity, coloring_result, compressed_matrix, batched_seeds, @@ -125,6 +139,7 @@ end function _prepare_sparse_jacobian_aux_aux( strict::Val, batch_size_settings::DI.BatchSizeSettings{B}, + sparsity::AbstractMatrix, coloring_result::AbstractColoringResult{:nonsymmetric,:row}, y, f_or_f!y::FY, @@ -145,9 +160,10 @@ function _prepare_sparse_jacobian_aux_aux( pullback_prep = DI.prepare_pullback_nokwarg( strict, f_or_f!y..., dense_backend, x, batched_seeds[1], contexts... ) - return PullbackSparseJacobianPrep( + return SMCPullbackSparseJacobianPrep( _sig, batch_size_settings, + sparsity, coloring_result, compressed_matrix, batched_seeds, @@ -161,7 +177,7 @@ end function DI.jacobian!( f::F, jac, - prep::SparseJacobianPrep, + prep::SMCSparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -171,7 +187,11 @@ function DI.jacobian!( end function DI.jacobian( - f::F, prep::SparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C} + f::F, + prep::SMCSparseJacobianPrep, + backend::AutoSparse, + x, + contexts::Vararg{DI.Context,C}, ) where {F,C} DI.check_prep(f, prep, backend, x, contexts...) jac = similar(sparsity_pattern(prep), eltype(x)) @@ -179,7 +199,11 @@ function DI.jacobian( end function DI.value_and_jacobian( - f::F, prep::SparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C} + f::F, + prep::SMCSparseJacobianPrep, + backend::AutoSparse, + x, + contexts::Vararg{DI.Context,C}, ) where {F,C} DI.check_prep(f, prep, backend, x, contexts...) return f(x, map(DI.unwrap, contexts)...), DI.jacobian(f, prep, backend, x, contexts...) @@ -188,7 +212,7 @@ end function DI.value_and_jacobian!( f::F, jac, - prep::SparseJacobianPrep, + prep::SMCSparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -204,7 +228,7 @@ function DI.jacobian!( f!::F, y, jac, - prep::SparseJacobianPrep, + prep::SMCSparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -216,7 +240,7 @@ end function DI.jacobian( f!::F, y, - prep::SparseJacobianPrep, + prep::SMCSparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -229,7 +253,7 @@ end function DI.value_and_jacobian( f!::F, y, - prep::SparseJacobianPrep, + prep::SMCSparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -244,7 +268,7 @@ function DI.value_and_jacobian!( f!::F, y, jac, - prep::SparseJacobianPrep, + prep::SMCSparseJacobianPrep, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -260,7 +284,7 @@ end function _sparse_jacobian_aux!( f_or_f!y::FY, jac, - prep::PushforwardSparseJacobianPrep{SIG,<:DI.BatchSizeSettings{B}}, + prep::SMCPushforwardSparseJacobianPrep{SIG,<:DI.BatchSizeSettings{B}}, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -306,7 +330,7 @@ end function _sparse_jacobian_aux!( f_or_f!y::FY, jac, - prep::PullbackSparseJacobianPrep{SIG,<:DI.BatchSizeSettings{B}}, + prep::SMCPullbackSparseJacobianPrep{SIG,<:DI.BatchSizeSettings{B}}, backend::AutoSparse, x, contexts::Vararg{DI.Context,C}, @@ -354,6 +378,6 @@ end ## Operator overloading -function DI.overloaded_input_type(prep::PushforwardSparseJacobianPrep) +function DI.overloaded_input_type(prep::SMCPushforwardSparseJacobianPrep) return DI.overloaded_input_type(prep.pushforward_prep) end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl index 42b4970db..463d5e173 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl @@ -1,9 +1,10 @@ ## Preparation -struct MixedModeSparseJacobianPrep{ +struct SMCMixedModeSparseJacobianPrep{ SIG, BSf<:DI.BatchSizeSettings, BSr<:DI.BatchSizeSettings, + P<:AbstractMatrix, C<:AbstractColoringResult{:nonsymmetric,:bidirectional}, M<:AbstractMatrix{<:Number}, Sf<:Vector{<:NTuple}, @@ -12,10 +13,11 @@ struct MixedModeSparseJacobianPrep{ Rr<:Vector{<:NTuple}, Ef<:DI.PushforwardPrep, Er<:DI.PullbackPrep, -} <: SparseJacobianPrep{SIG} +} <: SMCSparseJacobianPrep{SIG} _sig::Val{SIG} batch_size_settings_forward::BSf batch_size_settings_reverse::BSr + sparsity::P coloring_result::C compressed_matrix_forward::M compressed_matrix_reverse::M @@ -78,6 +80,7 @@ function _prepare_mixed_sparse_jacobian_aux( strict, batch_size_settings_forward, batch_size_settings_reverse, + sparsity, coloring_result, y, f_or_f!y, @@ -91,6 +94,7 @@ function _prepare_mixed_sparse_jacobian_aux_aux( strict::Val, batch_size_settings_forward::DI.BatchSizeSettings{Bf}, batch_size_settings_reverse::DI.BatchSizeSettings{Br}, + sparsity::AbstractMatrix, coloring_result::AbstractColoringResult{:nonsymmetric,:bidirectional}, y, f_or_f!y::FY, @@ -144,10 +148,11 @@ function _prepare_mixed_sparse_jacobian_aux_aux( contexts...; ) - return MixedModeSparseJacobianPrep( + return SMCMixedModeSparseJacobianPrep( _sig, batch_size_settings_forward, batch_size_settings_reverse, + sparsity, coloring_result, compressed_matrix_forward, compressed_matrix_reverse, @@ -165,7 +170,7 @@ end function _sparse_jacobian_aux!( f_or_f!y::FY, jac, - prep::MixedModeSparseJacobianPrep{ + prep::SMCMixedModeSparseJacobianPrep{ SIG,<:DI.BatchSizeSettings{Bf},<:DI.BatchSizeSettings{Br} }, backend::AutoSparse, diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl index 30c73ce23..720a8ad47 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl @@ -239,8 +239,9 @@ end ## Jacobian -struct SymbolicsOneArgJacobianPrep{SIG,E1,E1!} <: DI.JacobianPrep{SIG} +struct SymbolicsOneArgJacobianPrep{SIG,P,E1,E1!} <: DI.SparseJacobianPrep{SIG} _sig::Val{SIG} + sparsity::P jac_exe::E1 jac_exe!::E1! end @@ -255,16 +256,18 @@ function DI.prepare_jacobian_nokwarg( _sig = DI.signature(f, backend, x, contexts...; strict) x_var = variablize(x, :x) context_vars = variablize(contexts) - jac_var = if backend isa AutoSparse - sparsejacobian(vec(f(x_var, context_vars...)), vec(x_var)) + if backend isa AutoSparse + jac_var = sparsejacobian(vec(f(x_var, context_vars...)), vec(x_var)) + sparsity = DI.get_pattern(jac_var) else - jacobian(f(x_var, context_vars...), x_var) + jac_var = jacobian(f(x_var, context_vars...), x_var) + sparsity = nothing end erase_cache_vars!(context_vars, contexts) res = build_function(jac_var, x_var, context_vars...; expression=Val(false), cse=true) (jac_exe, jac_exe!) = res - return SymbolicsOneArgJacobianPrep(_sig, jac_exe, jac_exe!) + return SymbolicsOneArgJacobianPrep(_sig, sparsity, jac_exe, jac_exe!) end function DI.jacobian( @@ -317,9 +320,10 @@ end ## Hessian -struct SymbolicsOneArgHessianPrep{SIG,G,E2,E2!} <: DI.HessianPrep{SIG} +struct SymbolicsOneArgHessianPrep{SIG,G,P,E2,E2!} <: DI.SparseHessianPrep{SIG} _sig::Val{SIG} gradient_prep::G + sparsity::P hess_exe::E2 hess_exe!::E2! end @@ -335,10 +339,12 @@ function DI.prepare_hessian_nokwarg( x_var = variablize(x, :x) context_vars = variablize(contexts) # Symbolic.hessian only accepts vectors - hess_var = if backend isa AutoSparse - sparsehessian(f(x_var, context_vars...), vec(x_var)) + if backend isa AutoSparse + hess_var = sparsehessian(f(x_var, context_vars...), vec(x_var)) + sparsity = DI.get_pattern(hess_var) else - hessian(f(x_var, context_vars...), vec(x_var)) + hess_var = hessian(f(x_var, context_vars...), vec(x_var)) + sparsity = nothing end erase_cache_vars!(context_vars, contexts) @@ -350,7 +356,7 @@ function DI.prepare_hessian_nokwarg( gradient_prep = DI.prepare_gradient_nokwarg( strict, f, dense_ad(backend), x, contexts... ) - return SymbolicsOneArgHessianPrep(_sig, gradient_prep, hess_exe, hess_exe!) + return SymbolicsOneArgHessianPrep(_sig, gradient_prep, sparsity, hess_exe, hess_exe!) end function DI.hessian( diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/twoarg.jl index 41c80e95b..8ec5dc7a6 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/twoarg.jl @@ -180,8 +180,9 @@ end ## Jacobian -struct SymbolicsTwoArgJacobianPrep{SIG,E1,E1!} <: DI.JacobianPrep{SIG} +struct SymbolicsTwoArgJacobianPrep{SIG,P,E1,E1!} <: DI.SparseJacobianPrep{SIG} _sig::Val{SIG} + sparsity::P jac_exe::E1 jac_exe!::E1! end @@ -199,16 +200,18 @@ function DI.prepare_jacobian_nokwarg( y_var = variablize(y, :y) context_vars = variablize(contexts) f!(y_var, x_var, context_vars...) - jac_var = if backend isa AutoSparse - sparsejacobian(vec(y_var), vec(x_var)) + if backend isa AutoSparse + jac_var = sparsejacobian(vec(y_var), vec(x_var)) + sparsity = DI.get_pattern(jac_var) else - jacobian(y_var, x_var) + jac_var = jacobian(y_var, x_var) + sparsity = nothing end erase_cache_vars!(context_vars, contexts) res = build_function(jac_var, x_var, context_vars...; expression=Val(false), cse=true) (jac_exe, jac_exe!) = res - return SymbolicsTwoArgJacobianPrep(_sig, jac_exe, jac_exe!) + return SymbolicsTwoArgJacobianPrep(_sig, sparsity, jac_exe, jac_exe!) end function DI.jacobian( diff --git a/DifferentiationInterface/src/utils/linalg.jl b/DifferentiationInterface/src/utils/linalg.jl index b7dd4a42a..908c32c89 100644 --- a/DifferentiationInterface/src/utils/linalg.jl +++ b/DifferentiationInterface/src/utils/linalg.jl @@ -22,3 +22,14 @@ recursive_similar(x::AbstractArray, ::Type{T}) where {T} = similar(x, T) function recursive_similar(x::Union{Tuple,NamedTuple}, ::Type{T}) where {T} return map(xi -> recursive_similar(xi, T), x) end + +""" + get_pattern(M::AbstractMatrix) + +Return the `Bool`-valued sparsity pattern for a given matrix. + +Only specialized on `SparseMatrixCSC` because it is used with symbolic backends, and at the moment their sparse Jacobian/Hessian utilities return a `SparseMatrixCSC`. + +The trivial dense fallback is designed to protect against a change of format in these packages. +""" +get_pattern(M::AbstractMatrix) = trues(size(M)) diff --git a/DifferentiationInterface/src/utils/prep.jl b/DifferentiationInterface/src/utils/prep.jl index cb3bb6ff6..169afda86 100644 --- a/DifferentiationInterface/src/utils/prep.jl +++ b/DifferentiationInterface/src/utils/prep.jl @@ -45,6 +45,9 @@ struct NoJacobianPrep{SIG} <: JacobianPrep{SIG} _sig::Val{SIG} end +# assume the existence of a `sparsity` field +abstract type SparseJacobianPrep{SIG} <: JacobianPrep{SIG} end + """ $(docstring_preptype("HVPPrep", "hvp")) """ @@ -63,6 +66,9 @@ struct NoHessianPrep{SIG} <: HessianPrep{SIG} _sig::Val{SIG} end +# assume the existence of a `sparsity` field +abstract type SparseHessianPrep{SIG} <: HessianPrep{SIG} end + """ $(docstring_preptype("SecondDerivativePrep", "second_derivative")) """ diff --git a/DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl b/DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl index db6d2215b..43bf80cb8 100644 --- a/DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl +++ b/DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl @@ -2,6 +2,8 @@ using Pkg Pkg.add("FastDifferentiation") using DifferentiationInterface, DifferentiationInterfaceTest +using SparseMatrixColorings +using LinearAlgebra using FastDifferentiation: FastDifferentiation using Test @@ -29,3 +31,14 @@ test_differentiation( sparsity=true, logging=LOGGING, ); + +@testset "SparseMatrixColorings access" begin + x = rand(10) + backend = AutoSparse(AutoFastDifferentiation()) + jac_prep = prepare_jacobian(copy, backend, x) + jac!_prep = prepare_jacobian(copyto!, similar(x), backend, x) + hess_prep = prepare_hessian(x -> sum(abs2, x), backend, x) + @test sparsity_pattern(jac_prep) == Diagonal(trues(10)) + @test sparsity_pattern(jac!_prep) == Diagonal(trues(10)) + @test sparsity_pattern(hess_prep) == Diagonal(trues(10)) +end diff --git a/DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl b/DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl index 1ad15e8f1..6d3008004 100644 --- a/DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl +++ b/DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl @@ -2,6 +2,8 @@ using Pkg Pkg.add("Symbolics") using DifferentiationInterface, DifferentiationInterfaceTest +using SparseMatrixColorings +using LinearAlgebra using Symbolics: Symbolics using Test @@ -31,3 +33,14 @@ test_differentiation( sparsity=true, logging=LOGGING, ); + +@testset "SparseMatrixColorings access" begin + x = rand(10) + backend = AutoSparse(AutoSymbolics()) + jac_prep = prepare_jacobian(copy, backend, x) + jac!_prep = prepare_jacobian(copyto!, similar(x), backend, x) + hess_prep = prepare_hessian(x -> sum(abs2, x), backend, x) + @test sparsity_pattern(jac_prep) == Diagonal(trues(10)) + @test sparsity_pattern(jac!_prep) == Diagonal(trues(10)) + @test sparsity_pattern(hess_prep) == Diagonal(trues(10)) +end diff --git a/DifferentiationInterface/test/Core/Internals/linalg.jl b/DifferentiationInterface/test/Core/Internals/linalg.jl index 03798da87..0bb63d69d 100644 --- a/DifferentiationInterface/test/Core/Internals/linalg.jl +++ b/DifferentiationInterface/test/Core/Internals/linalg.jl @@ -1,9 +1,18 @@ -using DifferentiationInterface: recursive_similar +using DifferentiationInterface: recursive_similar, get_pattern +using SparseArrays using Test -@test recursive_similar(ones(Int, 2), Float32) isa Vector{Float32} -@test recursive_similar((ones(Int, 2), ones(Bool, 3, 4)), Float32) isa - Tuple{Vector{Float32},Matrix{Float32}} -@test recursive_similar((a=ones(Int, 2), b=(ones(Bool, 3, 4),)), Float32) isa - @NamedTuple{a::Vector{Float32}, b::Tuple{Matrix{Float32}}} -@test_throws MethodError recursive_similar(1, Float32) +@testset "Recursive similar" begin + @test recursive_similar(ones(Int, 2), Float32) isa Vector{Float32} + @test recursive_similar((ones(Int, 2), ones(Bool, 3, 4)), Float32) isa + Tuple{Vector{Float32},Matrix{Float32}} + @test recursive_similar((a=ones(Int, 2), b=(ones(Bool, 3, 4),)), Float32) isa + @NamedTuple{a::Vector{Float32}, b::Tuple{Matrix{Float32}}} + @test_throws MethodError recursive_similar(1, Float32) +end + +@testset "Sparsity pattern" begin + D = Diagonal(rand(10)) + @test_broken get_pattern(D) == Diagonal(trues(10)) + @test get_pattern(sparse(D)) == Diagonal(trues(10)) +end